diff --git a/.gitignore b/.gitignore index f26ec1f897..c7a50afe69 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ /tests/test_metal_session_batch /tests/test_q4k_dot /tests/test_sampling +/tests/test_session_snapshot *.o *.dSYM/ __pycache__/ diff --git a/AGENT.md b/AGENT.md index 7a1387c9ea..6626a397ae 100644 --- a/AGENT.md +++ b/AGENT.md @@ -1,8 +1,9 @@ # Agent Notes -`ds4.c` is a DeepSeek V4 Flash specific inference engine. It is not a generic -GGUF runner. The goal is a small, readable, high-performance C codebase with -Objective-C only where Metal requires it and Metal kernels under `metal/`. +DS4 is a model-specific inference engine. It is not a generic GGUF runner. The +goal is a small, readable, high-performance C codebase with model integrations +under `models/`, Objective-C only where Metal requires it, and shared backend +infrastructure kept outside model directories. ## Goals @@ -30,13 +31,18 @@ Objective-C only where Metal requires it and Metal kernels under `metal/`. ## Layout -- `ds4.c`: model loading, tokenizer, CPU reference code, Metal graph scheduling, - sessions, disk-cache payload serialization. +- `ds4.c`: engine loading, tokenizer, sessions, placement, and disk-cache + orchestration. +- `ds4_model_provider.h`: whole-model lifecycle boundary used by the engine. +- `models//`: one model's provider, CPU/graph orchestration, and custom + CUDA, Metal, and ROCm implementations. - `ds4_cli.c`: command line, linenoise REPL, interactive transcript handling. - `ds4_server.c`: OpenAI/Anthropic compatible HTTP API, worker queue, streaming, tool-call mapping, disk KV cache policy. -- `ds4_metal.m`: Objective-C Metal runtime and kernel wrappers. -- `metal/*.metal`: compute kernels. +- `ds4_cuda.cu`, `ds4_metal.m`, `ds4_rocm.cu`: backend translation-unit entry + points. +- `cuda/`, `metal/`, `rocm/`, `kernels/`: shared backend runtime and reusable + low-level primitives. - `tests/`: unit and live integration tests. - `misc/`: ignored notes, experiments, and old planning material. diff --git a/Makefile b/Makefile index 81d4881b21..7708ffd395 100644 --- a/Makefile +++ b/Makefile @@ -15,8 +15,42 @@ OBJCFLAGS ?= -O3 -ffast-math $(DEBUG_FLAGS) $(NATIVE_CPU_FLAG) -Wall -Wextra -fo 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) +METAL_SRCS := $(wildcard metal/*.metal models/*/metal/shaders/*.metal) +ROCM_SRCS := $(wildcard rocm/*.cuh models/*/rocm/*.cuh) +MODEL_PROVIDER_OBJS := \ + ds4_model_provider.o \ + models/deepseek/provider.o \ + models/glm/provider.o +CUDA_IMPL_FRAGMENTS := \ + cuda/runtime.inc \ + models/deepseek/cuda/dense_attention.inc \ + models/deepseek/cuda/control.inc \ + cuda/common_dispatch.inc \ + models/deepseek/cuda/moe.inc \ + models/deepseek/cuda/hc.inc \ + cuda/runtime_services.inc \ + models/glm/cuda/kernels.inc +METAL_IMPL_FRAGMENTS := \ + metal/runtime.inc \ + metal/embedding.inc \ + metal/model_io.inc \ + metal/expert_streaming.inc \ + models/deepseek/metal/host/indexer.inc \ + metal/dense_norm.inc \ + models/deepseek/metal/host/attention.inc \ + metal/elementwise.inc \ + metal/moe_dispatch.inc \ + models/glm/metal/host/kernels.inc \ + models/deepseek/metal/host/moe.inc \ + models/deepseek/metal/host/hc.inc \ + metal/compat.inc +DS4_IMPL_FRAGMENTS := \ + kernels/cpu_quant.inc \ + kernels/cpu_matmul.inc \ + models/deepseek/cpu.inc \ + models/deepseek/graph.inc \ + models/glm/cpu.inc \ + models/glm/graph.inc DS4_TEST_MODEL ?= ds4flash.gguf DS4_TEST_MTP ?= gguf/DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf DS4_DSPARK_MODEL ?= $(DS4_TEST_MODEL) @@ -24,8 +58,8 @@ DS4_DSPARK_SUPPORT ?= gguf/DeepSeek-V4-Flash-DSpark-support.gguf ifeq ($(UNAME_S),Darwin) METAL_LDLIBS := $(LDLIBS) -framework Foundation -framework Metal -CORE_OBJS = ds4.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_metal.o ds4_layer_pack.o -CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o +CORE_OBJS = ds4.o $(MODEL_PROVIDER_OBJS) ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_metal.o ds4_layer_pack.o +CPU_CORE_OBJS = ds4_cpu.o $(MODEL_PROVIDER_OBJS) 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 @@ -35,8 +69,8 @@ ifneq ($(strip $(CUDA_ARCH)),) NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH) 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_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o -CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o +CORE_OBJS = ds4.o $(MODEL_PROVIDER_OBJS) ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o +CPU_CORE_OBJS = ds4_cpu.o $(MODEL_PROVIDER_OBJS) 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 @@ -47,7 +81,7 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS) METAL_LDLIBS := $(LDLIBS) endif -.PHONY: all help clean test test-metal-session-batch 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 +.PHONY: all help clean test test-session-snapshot test-metal-session-batch test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm ifeq ($(UNAME_S),Darwin) all: ds4 ds4-server ds4-bench ds4-eval ds4-agent @@ -129,7 +163,7 @@ cuda: strix-halo: $(MAKE) -B ds4 ds4-server ds4-bench ds4-eval ds4-agent \ - CORE_OBJS="ds4.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" \ + CORE_OBJS="ds4.o $(MODEL_PROVIDER_OBJS) 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) -DDS4_ROCM_BUILD" \ DS4_LINK="$(HIPCC) $(ROCM_CFLAGS)" \ DS4_LINK_LIBS="$(ROCM_LDLIBS)" @@ -165,9 +199,31 @@ cuda-regression: tests/cuda_long_context_smoke ./tests/cuda_long_context_smoke endif -ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h +tests/test_session_snapshot.o: tests/test_session_snapshot.c ds4.h + $(CC) $(CFLAGS) -I. -c -o $@ $< + +tests/test_session_snapshot: tests/test_session_snapshot.o $(CORE_OBJS) +ifeq ($(UNAME_S),Darwin) + $(CC) $(CFLAGS) -o $@ $^ $(METAL_LDLIBS) +else + $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) +endif + +test-session-snapshot: tests/test_session_snapshot + DS4_TEST_MODEL="$(DS4_TEST_MODEL)" ./tests/test_session_snapshot + +ds4.o: ds4.c $(DS4_IMPL_FRAGMENTS) ds4.h ds4_model_provider.h ds4_model_provider_builtin.h models/deepseek/provider.h models/glm/provider.h ds4_ssd.h ds4_distributed.h ds4_gpu.h $(CC) $(CFLAGS) -c -o $@ ds4.c +ds4_model_provider.o: ds4_model_provider.c ds4_model_provider.h ds4_model_provider_builtin.h ds4.h + $(CC) $(CFLAGS) -c -o $@ ds4_model_provider.c + +models/deepseek/provider.o: models/deepseek/provider.c models/deepseek/provider.h ds4_model_provider.h ds4_model_provider_builtin.h ds4.h + $(CC) $(CFLAGS) -I. -c -o $@ models/deepseek/provider.c + +models/glm/provider.o: models/glm/provider.c models/glm/provider.h ds4_model_provider.h ds4_model_provider_builtin.h ds4.h + $(CC) $(CFLAGS) -I. -c -o $@ models/glm/provider.c + ds4_ssd.o: ds4_ssd.c ds4_ssd.h $(CC) $(CFLAGS) -c -o $@ ds4_ssd.c @@ -219,7 +275,7 @@ rax.o: rax.c rax.h rax_malloc.h linenoise.o: linenoise.c linenoise.h $(CC) $(CFLAGS) -c -o $@ linenoise.c -ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h +ds4_cpu.o: ds4.c $(DS4_IMPL_FRAGMENTS) ds4.h ds4_model_provider.h ds4_model_provider_builtin.h models/deepseek/provider.h models/glm/provider.h ds4_ssd.h ds4_distributed.h ds4_gpu.h $(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -c -o $@ ds4.c ds4_cli_cpu.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h @@ -240,10 +296,10 @@ ds4_eval_cpu.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_agent.c -ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_SRCS) +ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_IMPL_FRAGMENTS) $(METAL_SRCS) $(CC) $(OBJCFLAGS) -c -o $@ ds4_metal.m -ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_iq2_tables_cuda.inc +ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_gpu_mgpu.h ds4_iq2_tables_cuda.inc $(CUDA_IMPL_FRAGMENTS) $(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS) @@ -270,13 +326,13 @@ tests/test_gpu_args.o: tests/test_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h 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_gpu.h ds4_gpu_mgpu.h ds4_layer_pack.h +ds4_cpu_test_hooks.o: ds4.c $(DS4_IMPL_FRAGMENTS) ds4.h ds4_model_provider.h ds4_model_provider_builtin.h models/deepseek/provider.h models/glm/provider.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_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o +tests/test_engine_mgpu_placement: tests/test_engine_mgpu_placement.o ds4_cpu_test_hooks.o $(MODEL_PROVIDER_OBJS) ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) ifneq ($(UNAME_S),Darwin) @@ -298,7 +354,7 @@ tests/test_gpu_lookup_cache_strict.o: tests/test_gpu_lookup_cache_strict.c ds4_g tests/test_gpu_lookup_cache_strict: tests/test_gpu_lookup_cache_strict.o ds4_cuda.o $(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 +ds4_cuda_test_hooks.o: ds4.c $(DS4_IMPL_FRAGMENTS) ds4.h ds4_model_provider.h ds4_model_provider_builtin.h models/deepseek/provider.h models/glm/provider.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 @@ -310,7 +366,7 @@ tests/test_engine_mgpu_refusal: tests/test_engine_mgpu_refusal.o ds4_gpu_args.o 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 +tests/test_engine_mgpu_runtime: tests/test_engine_mgpu_runtime.o ds4_cuda_test_hooks.o $(MODEL_PROVIDER_OBJS) 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 $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) tests/test_engine_correctness.o: tests/test_engine_correctness.c ds4.h ds4_gpu_mgpu.h @@ -322,7 +378,7 @@ tests/test_engine_correctness: tests/test_engine_correctness.o ds4_gpu_args.o ds tests/test_sampling.o: tests/test_sampling.c ds4.h $(CC) $(CFLAGS) -DDS4_TEST_HOOKS -I. -c -o $@ $< -tests/test_sampling: tests/test_sampling.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 +tests/test_sampling: tests/test_sampling.o ds4_cuda_test_hooks.o $(MODEL_PROVIDER_OBJS) 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 $(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 @@ -337,7 +393,7 @@ test-cuda-session-batch: 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 +tests/test_cuda_mixed_batch: tests/test_cuda_mixed_batch.o ds4_cuda_test_hooks.o $(MODEL_PROVIDER_OBJS) 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 $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) test-cuda-mixed-batch: tests/test_cuda_mixed_batch @@ -402,4 +458,4 @@ q4k-dot-test: tests/test_q4k_dot.c ./tests/test_q4k_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 tests/test_metal_session_batch 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 + 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 tests/test_session_snapshot tests/test_metal_session_batch 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 models/deepseek/provider.o models/glm/provider.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o diff --git a/cuda/README.md b/cuda/README.md new file mode 100644 index 0000000000..222df113de --- /dev/null +++ b/cuda/README.md @@ -0,0 +1,14 @@ +# CUDA implementation units + +This directory owns CUDA infrastructure and low-level paths reused by model +integrations: + +- `runtime.inc`: CUDA initialization, tensor storage, model caching, and + multi-GPU plumbing. +- `common_dispatch.inc`: concrete launch wrappers and shared dense dispatch. +- `runtime_services.inc`: device probing and streamed-expert cache loading. + +Model-specific CUDA implementations live under `models//cuda/`. +`ds4_cuda.cu` includes the shared and model-owned fragments exactly once and +still compiles as one CUDA translation unit. This split adds no kernel base +class, wrapper layer, dispatch table, or hot-path runtime boundary. diff --git a/cuda/common_dispatch.inc b/cuda/common_dispatch.inc new file mode 100644 index 0000000000..a6fdbd2fcc --- /dev/null +++ b/cuda/common_dispatch.inc @@ -0,0 +1,4084 @@ +/* GLM opt-in: batched q8_0 matmuls with blocks > 32 may run as a + * streaming dequant-to-f16 GEMM (exact-q8 native kernels only cover + * blocks <= 32). Never enabled on DS4 paths, keeping them byte-stable. */ +static int g_q8_dequant_gemm_enabled = 0; +extern "C" void ds4_gpu_enable_q8_dequant_gemm(void) { + g_q8_dequant_gemm_enabled = 1; +} + +__global__ static void q8_0_dequant_f16_kernel( + __half *out, + const unsigned char *w, + uint64_t total_blocks, + uint32_t blocks_per_row, + uint32_t in_dim) { + /* Two threads per q8_0 block; each converts 16 values with half2 + * stores so a warp writes 512B contiguously per block pair. */ + const uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t b = tid >> 1; + if (b >= total_blocks) return; + const uint32_t half_idx = (uint32_t)tid & 1u; + const unsigned char *blk = w + b * 34u; + const float d = __half2float(*(const __half *)blk); + const int8_t *q = (const int8_t *)(blk + 2) + half_idx * 16u; + const uint64_t row = b / blocks_per_row; + const uint32_t col = (uint32_t)(b - row * blocks_per_row) * 32u + + half_idx * 16u; + __half2 *dst = (__half2 *)(out + row * in_dim + col); + #pragma unroll + for (int k = 0; k < 8; k++) { + dst[k] = __floats2half2_rn(d * (float)q[2 * k], + d * (float)q[2 * k + 1]); + } +} + +static int cuda_matmul_q8_0_tensor_labeled(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_tok, const char *label) { + if (!out || !x || !model_map) return 0; + uint64_t blocks = (in_dim + 31) / 32; + if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0; + uint64_t weight_bytes = out_dim * blocks * 34; + if (weight_bytes > model_size - weight_offset) return 0; + if (x->bytes < n_tok * in_dim * sizeof(float) || + out->bytes < n_tok * out_dim * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const int physical_device = + (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) + ? g_gpu[logical_tier].device_id : 0; + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "q8_0"); + if (!wptr) return 0; + if (g_cublas_ready && n_tok > 1) { + const float *w_f32 = cuda_q8_f32_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label); + if (w_f32) { + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w_f32, + (int)in_dim, + (const float *)x->ptr, + (int)in_dim, + &beta, + (float *)out->ptr, + (int)out_dim); + return cublas_ok(st, "q8 fp32 matmul"); + } + const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label); + if (w_f16) { + const uint64_t xh_count = n_tok * in_dim; + __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "q8 f16 gemm activations"); + if (!xh) return 0; + f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count); + if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0; + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w_f16, + CUDA_R_16F, + (int)in_dim, + xh, + CUDA_R_16F, + (int)in_dim, + &beta, + out->ptr, + CUDA_R_32F, + (int)out_dim, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + if (st == CUBLAS_STATUS_SUCCESS) return 1; + fprintf(stderr, "ds4: cuBLAS q8 f16 matmul failed: status %d\n", (int)st); + cuda_q8_f16_cache_disable_after_failure("cuBLAS f16 matmul failure", + in_dim * out_dim * sizeof(__half)); + /* The F16 expansion cache is only an optimization. If cuBLAS + * rejects the cached path under memory pressure, retry the same + * operation through the native Q8 kernels below. */ + } + } + if (g_q8_dequant_gemm_enabled && g_cublas_ready && + n_tok >= 128u && blocks > 32u && (in_dim & 31u) == 0u) { + /* Streaming dequant + f16 GEMM: the exact-q8 batched kernels only + * cover blocks <= 32 (DS4 TP shard widths); the per-token fallback + * re-reads the full weight per token (~30x the bytes at GLM dims). + * Scratch layout: [w_f16][x_f16] in one arena grab. */ + const uint64_t wh_bytes = in_dim * out_dim * sizeof(__half); + const uint64_t xh_off = (wh_bytes + 255u) & ~255ull; + const uint64_t oo_off = + (xh_off + n_tok * in_dim * sizeof(__half) + 255u) & ~255ull; + const uint64_t gemm_tmp = oo_off + n_tok * out_dim * sizeof(float); + /* Scratch must live on the EXECUTING device: logical_tier is the + * out tensor's tier (0 for GLM graph buffers), and a GEMM reading + * its staged weights across PCIe costs ~20ms instead of ~0.1ms. */ + int exec_tier = logical_tier; + { + int cur_dev = -1; + if (cudaGetDevice(&cur_dev) == cudaSuccess) { + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == cur_dev) { exec_tier = t; break; } + } + } + } + void *tmp16 = cuda_tmp_alloc_on(exec_tier, gemm_tmp, "q8 dequant gemm"); + if (tmp16) { + __half *wh = (__half *)tmp16; + __half *xh = (__half *)((char *)tmp16 + xh_off); + /* GEMM into device-local scratch, then one bulk D2D to the + * (possibly peer-mapped) out tensor: scattered peer stores + * from GEMM kernels run at <1GB/s over PCIe. */ + float *olocal = (float *)((char *)tmp16 + oo_off); + const uint64_t total_blocks = out_dim * blocks; + q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256>>>( + wh, reinterpret_cast(wptr), + total_blocks, (uint32_t)blocks, (uint32_t)in_dim); + const uint64_t xh_count = n_tok * in_dim; + f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>( + xh, (const float *)x->ptr, xh_count); + if (cuda_ok(cudaGetLastError(), "q8 dequant gemm staging")) { + const int gemm_trace = getenv("DS4_GLM_GEMM_TRACE") != NULL; + cudaEvent_t ev0, ev1, ev2; + if (gemm_trace) { + cudaEventCreate(&ev0); cudaEventCreate(&ev1); cudaEventCreate(&ev2); + cudaEventRecord(ev0); + } + const float alpha = 1.0f; + const float beta = 0.0f; + if (gemm_trace) cudaEventRecord(ev1); + cublasStatus_t st = cublasGemmEx( + cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, CUBLAS_OP_N, + (int)out_dim, (int)n_tok, (int)in_dim, + &alpha, + wh, CUDA_R_16F, (int)in_dim, + xh, CUDA_R_16F, (int)in_dim, + &beta, + olocal, CUDA_R_32F, (int)out_dim, + CUDA_R_32F, CUBLAS_GEMM_DEFAULT); + if (st == CUBLAS_STATUS_SUCCESS) { + if (!cuda_ok(cudaMemcpyAsync(out->ptr, olocal, + n_tok * out_dim * sizeof(float), + cudaMemcpyDeviceToDevice, 0), + "q8 dequant gemm out copy")) { + st = CUBLAS_STATUS_INTERNAL_ERROR; + } + } + if (gemm_trace) { + cudaEventRecord(ev2); + cudaEventSynchronize(ev2); + float stage_ms = 0, gemm_ms = 0; + cudaEventElapsedTime(&stage_ms, ev0, ev1); + cudaEventElapsedTime(&gemm_ms, ev1, ev2); + fprintf(stderr, + "ds4: gemm trace in=%llu out=%llu n=%llu stage(before)=%.2f gemm=%.2f ms\n", + (unsigned long long)in_dim, (unsigned long long)out_dim, + (unsigned long long)n_tok, stage_ms, gemm_ms); + cudaEventDestroy(ev0); cudaEventDestroy(ev1); cudaEventDestroy(ev2); + } + if (st == CUBLAS_STATUS_SUCCESS) return 1; + fprintf(stderr, + "ds4: q8 dequant gemm failed: status %d; using native path\n", + (int)st); + } + } + } + const uint64_t xq_bytes = n_tok * blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1); + quantize_q8_0_f32_kernel<<>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0; + if (n_tok == 1) { + matmul_q8_0_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 warp launch"); + } + const bool force_decode_warp = + n_tok == 2u && g_glm_mtp_verify_mode; + if (n_tok > 1u && !force_decode_warp) { + /* T matches the reduction width of whichever reference kernel would + * have run: warp tree (32) for blocks <= 32, exact-thread tree + * otherwise. */ + const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks); + const int mma_rc = cuda_q8_mma_try_launch( + (float *)out->ptr, reinterpret_cast(wptr), + xq, xscale, in_dim, out_dim, n_tok, blocks, blocks, out_dim, mma_T); + if (mma_rc) return mma_rc > 0; + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL && + blocks <= 32u && + n_tok >= 8u) { + dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); + matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok8 warp launch"); + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL && + blocks <= 32u && + n_tok >= 4u) { + dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); + matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok4 warp launch"); + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + (blocks <= 32u || force_decode_warp)) { + if (force_decode_warp && + getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) { + matmul_q8_0_preq_batch_warp8_tok2_kernel + <<<((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), + "matmul_q8_0 batch tok2 warp launch"); + } + dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1); + matmul_q8_0_preq_batch_warp8_kernel<<>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch warp launch"); + } + const unsigned exact_threads = cuda_q8_exact_threads(blocks); + if (getenv("DS4_CUDA_NO_Q8_BATCH_EXACT_TOK2") == NULL && + n_tok >= 2u) { + dim3 bgrid((unsigned)out_dim, ((unsigned)n_tok + 1u) / 2u, 1); + matmul_q8_0_preq_batch_tok2_exact_kernel<<>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 exact tok2 launch"); + } + dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); + matmul_q8_0_preq_kernel<<>>((float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, out_dim, n_tok, blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 launch"); +} + +extern "C" int ds4_gpu_matmul_q8_0_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_tok) { + return cuda_matmul_q8_0_tensor_labeled(out, model_map, model_size, weight_offset, + in_dim, out_dim, x, n_tok, "q8_0"); +} + +extern "C" int ds4_gpu_matmul_q8_0_top1_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *values, + 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 index_offset) { + if (!selected || !values || !x || !model_map || + in_dim == 0 || out_dim == 0 || out_dim > UINT32_MAX) { + return 0; + } + const uint64_t blocks = (in_dim + 31u) / 32u; + if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34u)) { + return 0; + } + const uint64_t weight_bytes = out_dim * blocks * 34u; + if (weight_bytes > model_size - weight_offset || + x->bytes < in_dim * sizeof(float) || + selected->bytes < sizeof(uint32_t) || + values->bytes < sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(selected); + const char *wptr = cuda_resolve_weight_ptr(model_map, + weight_offset, + weight_bytes, + logical_tier, + "q8_0_top1"); + if (!wptr) return 0; + + const uint64_t xq_bytes = blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t key_offset = + (scale_offset + blocks * sizeof(float) + 7u) & ~7ull; + const uint64_t tmp_bytes = key_offset + sizeof(unsigned long long); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 top1 prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + unsigned long long *best_key = + (unsigned long long *)((char *)tmp + key_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + + if (!cuda_ok(cudaMemsetAsync(best_key, 0, sizeof(*best_key)), + "matmul_q8_0_top1 clear")) { + return 0; + } + quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>( + xq, + xscale, + (const float *)x->ptr, + in_dim, + blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 quantize launch")) return 0; + matmul_q8_0_top1_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + best_key, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + blocks, + index_offset, + use_dp4a); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 launch")) return 0; + matmul_q8_0_top1_unpack_kernel<<<1, 1>>>( + (uint32_t *)selected->ptr, + (float *)values->ptr, + best_key); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 unpack launch"); +} + +extern "C" int ds4_gpu_matmul_q8_0_kslice_rows_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, + uint64_t in_start, + uint64_t in_count, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || + in_count == 0 || n_tok == 0 || n_tok > 65535u) return 0; + if ((in_start % 32u) != 0 || (in_count % 32u) != 0 || + in_start > in_dim || in_count > in_dim - in_start) return 0; + const uint64_t full_blocks = (in_dim + 31u) / 32u; + const uint64_t block_start = in_start / 32u; + const uint64_t slice_blocks = in_count / 32u; + if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0; + const uint64_t weight_bytes = out_dim * full_blocks * 34u; + if (in_count > UINT64_MAX / n_tok || out_dim > UINT64_MAX / n_tok) { + return 0; + } + if (weight_bytes > model_size - weight_offset || + x->bytes < n_tok * in_count * sizeof(float) || + out->bytes < n_tok * out_dim * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const unsigned char *wptr = reinterpret_cast( + cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, + logical_tier, "q8_0_kslice")); + if (!wptr) return 0; + + const uint64_t xq_bytes = n_tok * slice_blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = + scale_offset + n_tok * slice_blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + const dim3 qgrid((unsigned)slice_blocks, (unsigned)n_tok, 1u); + quantize_q8_0_f32_kernel<<>>( + xq, + xscale, + (const float *)x->ptr, + in_count, + slice_blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice quantize launch")) return 0; + const dim3 grid(((unsigned)out_dim + 7u) / 8u, + (unsigned)n_tok, 1u); + matmul_q8_0_kslice_preq_warp8_kernel<<>>( + (float *)out->ptr, + wptr, + xq, + xscale, + in_count, + out_dim, + full_blocks, + block_start, + slice_blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice launch"); +} + +extern "C" int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + uint64_t in_start, + uint64_t in_count, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!out_hc || !block_out || !x || !block_add || !residual_hc || !split || + !model_map || in_dim == 0 || out_dim == 0 || in_count == 0 || + n_embd == 0 || n_hc == 0 || out_dim != (uint64_t)n_embd) { + return 0; + } + if ((in_start % 32u) != 0 || (in_count % 32u) != 0 || + in_start > in_dim || in_count > in_dim - in_start) return 0; + const uint64_t full_blocks = (in_dim + 31u) / 32u; + const uint64_t block_start = in_start / 32u; + const uint64_t slice_blocks = in_count / 32u; + if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0; + const uint64_t weight_bytes = out_dim * full_blocks * 34u; + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t split_bytes = + (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float); + if (weight_bytes > model_size - weight_offset || + x->bytes < in_count * sizeof(float) || + block_out->bytes < out_dim * sizeof(float) || + block_add->bytes < out_dim * sizeof(float) || + residual_hc->bytes < hc_bytes || + split->bytes < split_bytes || + out_hc->bytes < hc_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out_hc); + const unsigned char *wptr = reinterpret_cast( + cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, + logical_tier, "q8_0_kslice_hc_expand_add")); + if (!wptr) return 0; + + const uint64_t xq_bytes = slice_blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + slice_blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice hc expand prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + quantize_q8_0_f32_kernel<<<(unsigned)slice_blocks, 32>>>( + xq, + xscale, + (const float *)x->ptr, + in_count, + slice_blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add quantize launch")) return 0; + matmul_q8_0_kslice_hc_expand_add_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)out_hc->ptr, + (float *)block_out->ptr, + (const float *)block_add->ptr, + (const float *)residual_hc->ptr, + (const float *)split->ptr, + wptr, + xq, + xscale, + in_count, + out_dim, + full_blocks, + block_start, + slice_blocks, + n_embd, + n_hc, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add launch"); +} + +extern "C" int ds4_gpu_matmul_q8_0_pair_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, + uint64_t n_tok) { + if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out0_dim == 0 || out1_dim == 0 || n_tok == 0) { + return 0; + } + const uint64_t blocks = (in_dim + 31) / 32; + if (weight0_offset > model_size || weight1_offset > model_size || + out0_dim > UINT64_MAX / (blocks * 34) || + out1_dim > UINT64_MAX / (blocks * 34)) { + return 0; + } + const uint64_t weight0_bytes = out0_dim * blocks * 34; + const uint64_t weight1_bytes = out1_dim * blocks * 34; + if (weight0_bytes > model_size - weight0_offset || + weight1_bytes > model_size - weight1_offset || + x->bytes < in_dim * sizeof(float) || + out0->bytes < out0_dim * sizeof(float) || + out1->bytes < out1_dim * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out0); + const char *w0 = cuda_resolve_weight_ptr(model_map, weight0_offset, weight0_bytes, logical_tier, "q8_0_pair0"); + const char *w1 = cuda_resolve_weight_ptr(model_map, weight1_offset, weight1_bytes, logical_tier, "q8_0_pair1"); + if (!w0 || !w1) return 0; + + const bool force_decode_warp = + n_tok == 2u && g_glm_mtp_verify_mode; + if (n_tok != 1 && !force_decode_warp && !g_q8_cache_suppressed && + getenv("DS4_CUDA_Q8_PAIR_BATCH") == NULL) { + return cuda_matmul_q8_0_tensor_labeled(out0, model_map, model_size, weight0_offset, + in_dim, out0_dim, x, n_tok, "q8_0_pair0") && + cuda_matmul_q8_0_tensor_labeled(out1, model_map, model_size, weight1_offset, + in_dim, out1_dim, x, n_tok, "q8_0_pair1"); + } + + const uint64_t xq_bytes = n_tok * blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 pair prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1); + quantize_q8_0_f32_kernel<<>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) return 0; + if (n_tok != 1) { + if (force_decode_warp && + getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) { + matmul_q8_0_preq_batch_warp8_tok2_kernel + <<<((unsigned)out0_dim + 7u) / 8u, 256>>>( + (float *)out0->ptr, + reinterpret_cast(w0), + xq, xscale, in_dim, out0_dim, blocks, use_dp4a); + if (!cuda_ok(cudaGetLastError(), + "matmul_q8_0 pair0 tok2 warp launch")) { + return 0; + } + matmul_q8_0_preq_batch_warp8_tok2_kernel + <<<((unsigned)out1_dim + 7u) / 8u, 256>>>( + (float *)out1->ptr, + reinterpret_cast(w1), + xq, xscale, in_dim, out1_dim, blocks, use_dp4a); + return cuda_ok(cudaGetLastError(), + "matmul_q8_0 pair1 tok2 warp launch"); + } + const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks); + int mma_rc = cuda_q8_mma_try_launch( + (float *)out0->ptr, reinterpret_cast(w0), + xq, xscale, in_dim, out0_dim, n_tok, blocks, blocks, out0_dim, mma_T); + if (mma_rc < 0) return 0; + if (mma_rc > 0) { + mma_rc = cuda_q8_mma_try_launch( + (float *)out1->ptr, reinterpret_cast(w1), + xq, xscale, in_dim, out1_dim, n_tok, blocks, blocks, out1_dim, mma_T); + if (mma_rc > 0) return 1; + return 0; + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL && + blocks <= 32u && + n_tok >= 8u) { + dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); + matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( + (float *)out0->ptr, + reinterpret_cast(w0), + xq, + xscale, + in_dim, + out0_dim, + n_tok, + blocks, + use_dp4a); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok8 launch")) return 0; + dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); + matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( + (float *)out1->ptr, + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out1_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok8 launch"); + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL && + blocks <= 32u && + n_tok >= 4u) { + dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); + matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( + (float *)out0->ptr, + reinterpret_cast(w0), + xq, + xscale, + in_dim, + out0_dim, + n_tok, + blocks, + use_dp4a); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok4 launch")) return 0; + dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); + matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( + (float *)out1->ptr, + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out1_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok4 launch"); + } + if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && + blocks <= 32u) { + dim3 grid0(((unsigned)out0_dim + 7u) / 8u, (unsigned)n_tok, 1); + matmul_q8_0_preq_batch_warp8_kernel<<>>( + (float *)out0->ptr, + reinterpret_cast(w0), + xq, + xscale, + in_dim, + out0_dim, + n_tok, + blocks, + use_dp4a); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch warp launch")) return 0; + dim3 grid1(((unsigned)out1_dim + 7u) / 8u, (unsigned)n_tok, 1); + matmul_q8_0_preq_batch_warp8_kernel<<>>( + (float *)out1->ptr, + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out1_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch warp launch"); + } + if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT") == NULL) { + const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; + const unsigned exact_threads = cuda_q8_exact_threads(blocks); + if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT_TOK2") == NULL && + n_tok >= 2u) { + dim3 grid((unsigned)max_out_dim, ((unsigned)n_tok + 1u) / 2u, 1); + matmul_q8_0_pair_preq_batch_tok2_exact_kernel<<>>( + (float *)out0->ptr, + (float *)out1->ptr, + reinterpret_cast(w0), + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out0_dim, + out1_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair exact tok2 launch"); + } + dim3 grid((unsigned)max_out_dim, (unsigned)n_tok, 1); + matmul_q8_0_pair_preq_batch_kernel<<>>( + (float *)out0->ptr, + (float *)out1->ptr, + reinterpret_cast(w0), + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out0_dim, + out1_dim, + n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair batch exact launch"); + } + const unsigned exact_threads = cuda_q8_exact_threads(blocks); + dim3 grid0((unsigned)out0_dim, (unsigned)n_tok, 1); + matmul_q8_0_preq_kernel<<>>((float *)out0->ptr, + reinterpret_cast(w0), + xq, + xscale, + in_dim, out0_dim, n_tok, blocks, + use_dp4a); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch launch")) return 0; + dim3 grid1((unsigned)out1_dim, (unsigned)n_tok, 1); + matmul_q8_0_preq_kernel<<>>((float *)out1->ptr, + reinterpret_cast(w1), + xq, + xscale, + in_dim, out1_dim, n_tok, blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch launch"); + } + const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; + matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>( + (float *)out0->ptr, + (float *)out1->ptr, + reinterpret_cast(w0), + reinterpret_cast(w1), + xq, + xscale, + in_dim, + out0_dim, + out1_dim, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch"); +} + +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) { + if (!out || !x || !model_map || in_dim == 0u || out_dim == 0u || + n_rows == 0u || + x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || + out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) { + return 0; + } + const uint64_t blocks = (in_dim + 31u) / 32u; + if (weight_offset > model_size || + out_dim > UINT64_MAX / (blocks * 34u)) { + return 0; + } + const uint64_t weight_bytes = out_dim * blocks * 34u; + if (weight_bytes > model_size - weight_offset) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + if (logical_tier < 0 || logical_tier >= g_n_gpus || + ds4_tensor_device_idx(x) != logical_tier) { + return 0; + } + const char *wptr = cuda_resolve_weight_ptr( + model_map, weight_offset, weight_bytes, logical_tier, + "q8_0 decode rows exact"); + if (!wptr) return 0; + + const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = + scale_offset + (uint64_t)n_rows * blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on( + logical_tier, tmp_bytes, "q8_0 decode rows exact prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + dim3 qgrid((unsigned)blocks, n_rows, 1u); + quantize_q8_0_f32_kernel<<>>( + xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), + "q8_0 decode rows exact quantize launch")) { + return 0; + } + dim3 grid(((unsigned)out_dim + 7u) / 8u, n_rows, 1u); + matmul_q8_0_preq_warp8_kernel<<>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, xscale, in_dim, out_dim, blocks, cuda_q8_use_dp4a()); + return cuda_ok(cudaGetLastError(), + "q8_0 decode rows exact warp launch"); +} + +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) { + if (!out0 || !out1 || !x || !model_map || in_dim == 0u || + out0_dim == 0u || out1_dim == 0u || n_rows == 0u || + x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || + out0->bytes < (uint64_t)n_rows * out0_dim * sizeof(float) || + out1->bytes < (uint64_t)n_rows * out1_dim * sizeof(float)) { + return 0; + } + const uint64_t blocks = (in_dim + 31u) / 32u; + if (weight0_offset > model_size || weight1_offset > model_size || + out0_dim > UINT64_MAX / (blocks * 34u) || + out1_dim > UINT64_MAX / (blocks * 34u)) { + return 0; + } + const uint64_t weight0_bytes = out0_dim * blocks * 34u; + const uint64_t weight1_bytes = out1_dim * blocks * 34u; + if (weight0_bytes > model_size - weight0_offset || + weight1_bytes > model_size - weight1_offset) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out0); + if (logical_tier < 0 || logical_tier >= g_n_gpus || + ds4_tensor_device_idx(out1) != logical_tier || + ds4_tensor_device_idx(x) != logical_tier) { + return 0; + } + const char *w0 = cuda_resolve_weight_ptr( + model_map, weight0_offset, weight0_bytes, logical_tier, + "q8_0 pair decode rows exact gate"); + const char *w1 = cuda_resolve_weight_ptr( + model_map, weight1_offset, weight1_bytes, logical_tier, + "q8_0 pair decode rows exact up"); + if (!w0 || !w1) return 0; + + const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = + scale_offset + (uint64_t)n_rows * blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on( + logical_tier, tmp_bytes, "q8_0 pair decode rows exact prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + dim3 qgrid((unsigned)blocks, n_rows, 1u); + quantize_q8_0_f32_kernel<<>>( + xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), + "q8_0 pair decode rows exact quantize launch")) { + return 0; + } + const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; + dim3 grid(((unsigned)max_out + 7u) / 8u, n_rows, 1u); + matmul_q8_0_pair_preq_warp8_kernel<<>>( + (float *)out0->ptr, + (float *)out1->ptr, + reinterpret_cast(w0), + reinterpret_cast(w1), + xq, xscale, in_dim, out0_dim, out1_dim, blocks, + cuda_q8_use_dp4a()); + return cuda_ok(cudaGetLastError(), + "q8_0 pair decode rows exact warp launch"); +} + +static int cuda_matmul_q8_0_hc_expand_tensor_labeled( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_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, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *block_add2, + const ds4_gpu_tensor *owned_home_slots, + const ds4_gpu_tensor *owned_peer_packed, + const ds4_gpu_tensor *owned_selected, + uint32_t owned_expert_split, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc, + const char *label) { + if (!out_hc || !block_out || !x || !residual_hc || !split || !model_map || + in_dim == 0 || out_dim == 0 || n_embd == 0 || n_hc == 0 || + out_dim != (uint64_t)n_embd) { + return 0; + } + const uint64_t blocks = (in_dim + 31) / 32; + if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0; + const uint64_t weight_bytes = out_dim * blocks * 34; + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t split_bytes = (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float); + if (weight_bytes > model_size - weight_offset || + x->bytes < in_dim * sizeof(float) || + block_out->bytes < out_dim * sizeof(float) || + residual_hc->bytes < hc_bytes || + split->bytes < split_bytes || + out_hc->bytes < hc_bytes || + (block_add && block_add->bytes < out_dim * sizeof(float)) || + (block_add2 && block_add2->bytes < out_dim * sizeof(float)) || + ((owned_home_slots || owned_peer_packed || owned_selected) && + (!owned_home_slots || !owned_peer_packed || !owned_selected || + owned_expert_split == 0u || + owned_home_slots->bytes < 6u * out_dim * sizeof(float) || + owned_peer_packed->bytes < 4u * out_dim * sizeof(float) || + owned_selected->bytes < 6u * sizeof(int32_t)))) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out_hc); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, label ? label : "q8_0_hc_expand"); + if (!wptr) return 0; + + const uint64_t xq_bytes = blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 hc expand prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) return 0; + matmul_q8_0_hc_expand_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)out_hc->ptr, + (float *)block_out->ptr, + block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr, + block_add2 ? (const float *)block_add2->ptr : (const float *)block_out->ptr, + owned_home_slots ? (const float *)owned_home_slots->ptr : NULL, + owned_peer_packed ? (const float *)owned_peer_packed->ptr : NULL, + owned_selected ? (const int32_t *)owned_selected->ptr : NULL, + (const float *)residual_hc->ptr, + (const float *)split->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + n_embd, + n_hc, + blocks, + block_add ? 1 : 0, + block_add2 ? 1 : 0, + owned_home_slots ? 1 : 0, + owned_expert_split, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand launch"); +} + +extern "C" int ds4_gpu_matmul_f16_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_tok) { + if (!out || !x || !model_map) return 0; + if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0; + uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t); + if (weight_bytes > model_size - weight_offset) return 0; + if (x->bytes < n_tok * in_dim * sizeof(float) || + out->bytes < n_tok * out_dim * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f16"); + if (!wptr) return 0; + const __half *w = (const __half *)wptr; + const int serial_f16 = getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL; + const int router_shape = in_dim == 4096u && out_dim == 256u && n_tok == 1u; + const int serial_router = + !serial_f16 && + router_shape && + getenv("DS4_CUDA_SERIAL_ROUTER") != NULL; + const int ordered_router = + !serial_f16 && + !serial_router && + n_tok == 1u && + getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL; + const int small_out_one_token = + !serial_f16 && + !serial_router && + !g_quality_mode && + n_tok == 1u && + out_dim <= 32u && + in_dim >= 8192u && + getenv("DS4_CUDA_F16_SMALL_OUT") != NULL && + getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL && + getenv("DS4_CUDA_NO_F16_SMALL_OUT") == NULL; + if (small_out_one_token) { + matmul_f16_small_out_hx_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>( + (float *)out->ptr, + w, + (const float *)x->ptr, + in_dim, + out_dim); + return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_hx_ordered_chunks launch"); + } + const int small_out_batch = + !serial_f16 && + !serial_router && + n_tok > 1u && + out_dim <= 32u && + in_dim >= 4096u && + (g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) && + getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL; + if (small_out_batch) { + matmul_f16_small_out_batch_kernel<<<(unsigned)n_tok, 256>>>( + (float *)out->ptr, + w, + (const float *)x->ptr, + in_dim, + out_dim, + n_tok); + return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_batch launch"); + } + const int cublas_one_token = + n_tok == 1u && + getenv("DS4_CUDA_NO_F16_CUBLAS_ONE") == NULL && + (!g_quality_mode || getenv("DS4_CUDA_F16_CUBLAS_ONE") != NULL); + const int cublas_batch = + n_tok > 1u && getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL; + if (!serial_f16 && g_cublas_ready && (cublas_batch || cublas_one_token)) { + const uint64_t xh_count = n_tok * in_dim; + __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "f16 gemm activations"); + if (!xh) return 0; + f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count); + if (!cuda_ok(cudaGetLastError(), "f16 activation convert launch")) return 0; + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w, + CUDA_R_16F, + (int)in_dim, + xh, + CUDA_R_16F, + (int)in_dim, + &beta, + out->ptr, + CUDA_R_32F, + (int)out_dim, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + return cublas_ok(st, "f16 matmul"); + } + dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); + if (serial_f16 || serial_router) { + matmul_f16_serial_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); + return cuda_ok(cudaGetLastError(), serial_router ? "matmul_f16_router_serial launch" : "matmul_f16_serial launch"); + } + if (ordered_router) { + matmul_f16_ordered_chunks_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); + return cuda_ok(cudaGetLastError(), "matmul_f16_ordered_chunks launch"); + } + matmul_f16_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); + return cuda_ok(cudaGetLastError(), "matmul_f16 launch"); +} + +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) { + const uint64_t in_dim = 4096u; + const uint64_t out_dim = 256u; + if (!out || !x || !model_map || n_rows == 0u || + weight_offset > model_size) { + return 0; + } + const uint64_t weight_bytes = in_dim * out_dim * sizeof(uint16_t); + if (weight_bytes > model_size - weight_offset || + x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || + out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) { + return 0; + } + if (n_rows == 1u) { + return ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, x, 1); + } + const int logical_tier = ds4_tensor_device_idx(out); + if (ds4_tensor_device_idx(x) != logical_tier || !g_cublas_ready) return 0; + const __half *w = (const __half *)cuda_resolve_weight_ptr( + model_map, weight_offset, weight_bytes, logical_tier, + "f16_router_rows_exact"); + if (!w) return 0; + + const uint64_t xh_count = (uint64_t)n_rows * in_dim; + __half *xh = (__half *)cuda_tmp_alloc_on( + logical_tier, xh_count * sizeof(__half), + "f16 exact router batch activations"); + if (!xh) return 0; + f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>( + xh, (const float *)x->ptr, xh_count); + if (!cuda_ok(cudaGetLastError(), + "f16 exact router activation convert launch")) { + return 0; + } + const float alpha = 1.0f; + const float beta = 0.0f; + /* Larger batchCount values let cuBLAS select a different reduction and + * change logits. Four-row calls match the one-row decode bit for bit on + * this projection, while still replacing most per-session launches. */ + uint32_t row = 0; + for (; row + 4u <= n_rows; row += 4u) { + cublasStatus_t st = cublasGemmStridedBatchedEx( + cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + 1, + (int)in_dim, + &alpha, + w, + CUDA_R_16F, + (int)in_dim, + 0, + xh + (uint64_t)row * in_dim, + CUDA_R_16F, + (int)in_dim, + (long long int)in_dim, + &beta, + (float *)out->ptr + (uint64_t)row * out_dim, + CUDA_R_32F, + (int)out_dim, + (long long int)out_dim, + 4, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "f16 exact router row batch")) return 0; + } + for (; row < n_rows; row++) { + ds4_gpu_tensor out_row = *out; + ds4_gpu_tensor x_row = *x; + out_row.ptr = (float *)out->ptr + (uint64_t)row * out_dim; + out_row.bytes = out_dim * sizeof(float); + x_row.ptr = (float *)x->ptr + (uint64_t)row * in_dim; + x_row.bytes = in_dim * sizeof(float); + if (!ds4_gpu_matmul_f16_tensor( + &out_row, model_map, model_size, weight_offset, + in_dim, out_dim, &x_row, 1)) { + return 0; + } + } + return 1; +} + +extern "C" int ds4_gpu_matmul_f16_pair_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 out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) { + return 0; + } + if (getenv("DS4_CUDA_NO_F16_PAIR_MATMUL") != NULL || + getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL || + getenv("DS4_CUDA_SERIAL_ROUTER") != NULL || + getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") != NULL) { + return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, + in_dim, out_dim, x, n_tok) && + ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, + in_dim, out_dim, x, n_tok); + } + if (weight0_offset > model_size || weight1_offset > model_size || + out_dim > UINT64_MAX / in_dim || + n_tok > UINT64_MAX / in_dim || + n_tok > UINT64_MAX / out_dim) { + return 0; + } + const uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (weight_bytes > model_size - weight0_offset || + weight_bytes > model_size - weight1_offset || + x->bytes < x_bytes || + out0->bytes < out_bytes || + out1->bytes < out_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out0); + if (ds4_tensor_device_idx(out1) != logical_tier) { + return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, + in_dim, out_dim, x, n_tok) && + ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, + in_dim, out_dim, x, n_tok); + } + const __half *w0 = (const __half *)cuda_resolve_weight_ptr(model_map, weight0_offset, weight_bytes, logical_tier, "f16_pair0"); + const __half *w1 = (const __half *)cuda_resolve_weight_ptr(model_map, weight1_offset, weight_bytes, logical_tier, "f16_pair1"); + if (!w0 || !w1) return 0; + if (n_tok > 1) { + const bool small_out_batch_requested = + out_dim <= 32u && + in_dim >= 4096u && + (g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) && + getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL; + if (!small_out_batch_requested && + g_cublas_ready && + getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL) { + const uint64_t xh_count = n_tok * in_dim; + __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, + xh_count * sizeof(__half), + "f16 pair gemm activations"); + if (!xh) return 0; + f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>( + xh, (const float *)x->ptr, xh_count); + if (!cuda_ok(cudaGetLastError(), "f16 pair activation convert launch")) return 0; + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w0, + CUDA_R_16F, + (int)in_dim, + xh, + CUDA_R_16F, + (int)in_dim, + &beta, + out0->ptr, + CUDA_R_32F, + (int)out_dim, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "f16 pair matmul0")) return 0; + st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w1, + CUDA_R_16F, + (int)in_dim, + xh, + CUDA_R_16F, + (int)in_dim, + &beta, + out1->ptr, + CUDA_R_32F, + (int)out_dim, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + return cublas_ok(st, "f16 pair matmul1"); + } + return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, + in_dim, out_dim, x, n_tok) && + ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, + in_dim, out_dim, x, n_tok); + } + matmul_f16_pair_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>( + (float *)out0->ptr, + (float *)out1->ptr, + w0, + w1, + (const float *)x->ptr, + in_dim, + out_dim, + out_dim); + return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch"); +} + +extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor( + ds4_gpu_tensor *out_kv, + ds4_gpu_tensor *out_score, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t weight_kv_offset, + uint64_t weight_score_offset, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t in_dim, + uint32_t width, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos) { + (void)out_kv; + (void)out_score; + (void)state_kv; + (void)state_score; + (void)model_map; + (void)model_size; + (void)weight_kv_offset; + (void)weight_score_offset; + (void)ape_offset; + (void)ape_type; + (void)in_dim; + (void)width; + (void)x; + (void)ratio; + (void)pos; + return 0; +} + +extern "C" int ds4_gpu_matmul_f32_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_tok) { + if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) return 0; + if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0; + uint64_t weight_elems = out_dim * in_dim; + if (weight_elems > UINT64_MAX / sizeof(float)) return 0; + uint64_t weight_bytes = weight_elems * sizeof(float); + if (weight_bytes > model_size - weight_offset) return 0; + if (x->bytes < n_tok * in_dim * sizeof(float) || + out->bytes < n_tok * out_dim * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f32"); + if (!wptr) return 0; + const float *w = (const float *)wptr; + if (g_cublas_ready && n_tok > 1) { + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)out_dim, + (int)n_tok, + (int)in_dim, + &alpha, + w, + (int)in_dim, + (const float *)x->ptr, + (int)in_dim, + &beta, + (float *)out->ptr, + (int)out_dim); + return cublas_ok(st, "f32 matmul"); + } + dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); + matmul_f32_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); + return cuda_ok(cudaGetLastError(), "matmul_f32 launch"); +} + +extern "C" int ds4_gpu_repeat_hc_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *row, uint32_t n_embd, uint32_t n_hc) { + if (!out || !row || n_embd == 0 || n_hc == 0 || + row->bytes < (uint64_t)n_embd * sizeof(float) || + out->bytes < (uint64_t)n_embd * n_hc * sizeof(float)) { + return 0; + } + uint64_t n = (uint64_t)n_embd * n_hc; + repeat_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)row->ptr, n_embd, n_hc); + return cuda_ok(cudaGetLastError(), "repeat_hc launch"); +} + + +/* Non-causal batch attention over a raw KV ring for the DSpark draft block. + * Every query row attends over all n_raw visible rows plus the per-head sink, + * with the same exact one-block max/denominator/value accumulation order as + * the reference decode attention (scores in shared, sequential value pass). */ +__global__ static void attention_noncausal_raw_batch_heads_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t tok = blockIdx.x; + const uint32_t h = blockIdx.y; + if (tok >= n_tokens || h >= n_head) return; + extern __shared__ float sh_scores[]; /* n_raw floats */ + const float *qh = q + ((uint64_t)tok * n_head + h) * head_dim; + const float scale = rsqrtf((float)head_dim); + for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { + const uint32_t row = (raw_start + r) % raw_cap; + const float *kv = raw_kv + (uint64_t)row * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; + sh_scores[r] = dot * scale; + } + __syncthreads(); + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + float local_max = sinks[h]; + for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { + local_max = fmaxf(local_max, sh_scores[r]); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { + sh_scores[r] = expf(sh_scores[r] - max_s); + den_local += sh_scores[r]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + ((uint64_t)tok * n_head + h) * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + acc += raw_kv[(uint64_t)row * head_dim + d] * sh_scores[r]; + } + oh[d] = acc / denom; + } +} + +extern "C" int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (!heads || !q || !raw_kv || !model_map || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || n_head == 0 || head_dim == 0 || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, + "dspark_attn_sinks"); + if (!sinks) return 0; + const size_t shmem = (size_t)n_raw * sizeof(float); + if (shmem > 32768) return 0; /* draft blocks are tiny; guard anyway */ + dim3 grid(n_tokens, n_head, 1); + attention_noncausal_raw_batch_heads_kernel<<>>( + (float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_tokens, n_raw, raw_cap, raw_start, n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), "attention noncausal raw batch heads launch")) return 0; + static int verify_left = -1; + if (verify_left < 0) { + verify_left = getenv("DS4_DSPARK_VERIFY_NONCAUSAL") != NULL ? 3 : 0; + } + if (verify_left > 0) { + verify_left--; + (void)cudaDeviceSynchronize(); + const uint64_t qn = (uint64_t)n_tokens * n_head * head_dim; + const uint64_t kn = (uint64_t)raw_cap * head_dim; + std::vector hq(qn), hkv(kn), hout(qn), hsink(n_head); + (void)cudaMemcpy(hq.data(), q->ptr, qn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hkv.data(), raw_kv->ptr, kn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hout.data(), heads->ptr, qn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hsink.data(), sinks, (uint64_t)n_head * 4, cudaMemcpyDeviceToHost); + double max_abs = 0.0, max_rel = 0.0; + const double scale = 1.0 / sqrt((double)head_dim); + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t h = 0; h < n_head; h++) { + std::vector sc(n_raw); + double mx = (double)hsink[h]; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + double dot = 0.0; + for (uint32_t d = 0; d < head_dim; d++) { + dot += (double)hq[((uint64_t)t * n_head + h) * head_dim + d] * + (double)hkv[(uint64_t)row * head_dim + d]; + } + sc[r] = dot * scale; + if (sc[r] > mx) mx = sc[r]; + } + double den = exp((double)hsink[h] - mx); + for (uint32_t r = 0; r < n_raw; r++) den += exp(sc[r] - mx); + for (uint32_t d = 0; d < head_dim; d++) { + double acc = 0.0; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + acc += exp(sc[r] - mx) * (double)hkv[(uint64_t)row * head_dim + d]; + } + const double ref = acc / den; + const double got = (double)hout[((uint64_t)t * n_head + h) * head_dim + d]; + const double ad = fabs(ref - got); + if (ad > max_abs) max_abs = ad; + if (fabs(ref) > 1e-3 && ad / fabs(ref) > max_rel) max_rel = ad / fabs(ref); + } + } + } + fprintf(stderr, + "ds4: DSpark noncausal verify n_tok=%u n_raw=%u start=%u cap=%u " + "max_abs=%.3e max_rel=%.3e\n", + n_tokens, n_raw, raw_start, raw_cap, max_abs, max_rel); + } + return 1; +} + +extern "C" int ds4_gpu_repeat_hc_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) { + uint64_t rows_elems = 0; + uint64_t out_elems = 0; + if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0 || + (uint64_t)n_tokens > UINT64_MAX / n_embd || + (rows_elems = (uint64_t)n_tokens * n_embd) > UINT64_MAX / n_hc || + (out_elems = rows_elems * n_hc) > UINT64_MAX / sizeof(float) || + rows_elems > UINT64_MAX / sizeof(float) || + rows->bytes < rows_elems * sizeof(float) || + out->bytes < out_elems * sizeof(float)) { + return 0; + } + const uint64_t blocks = (out_elems + 255u) / 256u; + if (blocks > UINT32_MAX) return 0; + repeat_hc_rows_kernel<<<(unsigned)blocks, 256>>>((float *)out->ptr, (const float *)rows->ptr, n_tokens, n_embd, n_hc); + return cuda_ok(cudaGetLastError(), "repeat_hc_rows launch"); +} + +extern "C" int ds4_gpu_rms_norm_plain_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, float eps) { + if (!out || !x || out->bytes < (uint64_t)n * sizeof(float) || + x->bytes < (uint64_t)n * sizeof(float)) return 0; + if (n == 4096u) { + rms_norm_plain_fast4096_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); + } else if ((n & 2047u) == 0u) { + rms_norm_plain_batch8_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); + } else { + rms_norm_plain_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); + } + return cuda_ok(cudaGetLastError(), "rms_norm_plain launch"); +} +extern "C" int ds4_gpu_rms_norm_plain_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, uint32_t rows, float eps) { + if (!out || !x || out->bytes < (uint64_t)n * rows * sizeof(float) || + x->bytes < (uint64_t)n * rows * sizeof(float)) return 0; + if (n == 4096u) { + rms_norm_plain_fast4096_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); + } else if ((n & 2047u) == 0u) { + rms_norm_plain_batch8_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); + } else { + rms_norm_plain_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); + } + return cuda_ok(cudaGetLastError(), "rms_norm_plain launch"); +} +extern "C" int ds4_gpu_rms_norm_weight_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, float eps) { + if (!out || !x || !model_map || weight_offset > model_size || + model_size - weight_offset < (uint64_t)n * sizeof(float) || + out->bytes < (uint64_t)n * sizeof(float) || + x->bytes < (uint64_t)n * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight"); + if (!wptr) return 0; + const float *w = (const float *)wptr; + rms_norm_weight_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, w, n, 1, eps); + return cuda_ok(cudaGetLastError(), "rms_norm_weight launch"); +} +extern "C" int ds4_gpu_rms_norm_weight_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, uint32_t rows, float eps) { + if (!out || !x || !model_map || weight_offset > model_size || + model_size - weight_offset < (uint64_t)n * sizeof(float) || + out->bytes < (uint64_t)n * rows * sizeof(float) || + x->bytes < (uint64_t)n * rows * sizeof(float)) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight"); + if (!wptr) return 0; + const float *w = (const float *)wptr; + rms_norm_weight_kernel<<>>((float *)out->ptr, (const float *)x->ptr, w, n, rows, eps); + return cuda_ok(cudaGetLastError(), "rms_norm_weight launch"); +} +extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_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, + float eps) { + if (!g_cuda_disable_qkv_rms_fused) { + if (!q_out || !q || !kv_out || !kv || !model_map || + q_weight_offset > model_size || + kv_weight_offset > model_size || + model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) || + model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) || + q_out->bytes < (uint64_t)q_n * rows * sizeof(float) || + q->bytes < (uint64_t)q_n * rows * sizeof(float) || + kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) || + kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(q_out); + const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map, + q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight"); + const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map, + kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight"); + if (!q_w || !kv_w) return 0; + dim3 grid(rows, 2u, 1u); + dsv4_qkv_rms_norm_rows_kernel<<>>( + (float *)q_out->ptr, + (const float *)q->ptr, + q_w, + q_n, + (float *)kv_out->ptr, + (const float *)kv->ptr, + kv_w, + kv_n, + rows, + eps); + return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm rows launch"); + } + return ds4_gpu_rms_norm_weight_rows_tensor(q_out, q, model_map, model_size, + q_weight_offset, q_n, rows, eps) && + ds4_gpu_rms_norm_weight_rows_tensor(kv_out, kv, model_map, model_size, + kv_weight_offset, kv_n, rows, eps); +} + +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) { + if (g_cuda_disable_qkv_rms_fused) return 0; + if (!q_out || !q || !kv_out || !kv || !model_map || + q_weight_offset > model_size || + kv_weight_offset > model_size || + kv_n_head == 0 || kv_head_dim == 0 || + n_rot > kv_head_dim || (n_rot & 1u) || + kv_n != kv_n_head * kv_head_dim || + model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) || + model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) || + q_out->bytes < (uint64_t)q_n * rows * sizeof(float) || + q->bytes < (uint64_t)q_n * rows * sizeof(float) || + kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) || + kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(q_out); + const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map, + q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight"); + const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map, + kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight"); + if (!q_w || !kv_w) return 0; + dim3 grid(rows, 2u, 1u); + dsv4_qkv_rms_norm_rows_kv_rope_kernel<<>>( + (float *)q_out->ptr, + (const float *)q->ptr, + q_w, + q_n, + (float *)kv_out->ptr, + (const float *)kv->ptr, + kv_w, + kv_n, + rows, + kv_n_head, + kv_head_dim, + n_rot, + pos0, + n_ctx_orig, + inverse ? 1 : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + eps); + return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm kv rope launch"); +} + +extern "C" int ds4_gpu_head_rms_norm_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) { + if (!x || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; + head_rms_norm_kernel<<>>((float *)x->ptr, n_tok, n_head, head_dim, eps); + return cuda_ok(cudaGetLastError(), "head_rms_norm launch"); +} +extern "C" int ds4_gpu_head_rms_norm_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t 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) { + if (!x || n_rot > head_dim || (n_rot & 1u) || + x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; + head_rms_norm_rope_tail_kernel<<>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, eps); + return cuda_ok(cudaGetLastError(), "head_rms_norm_rope_tail launch"); +} +extern "C" int ds4_gpu_dsv4_fp8_kv_quantize_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) { + if (!x || n_rot > head_dim || x->bytes < (uint64_t)n_tok * head_dim * sizeof(float)) return 0; + fp8_kv_quantize_kernel<<>>((float *)x->ptr, n_tok, head_dim, n_rot); + return cuda_ok(cudaGetLastError(), "fp8_kv_quantize launch"); +} +extern "C" int ds4_gpu_dsv4_indexer_qat_tensor(ds4_gpu_tensor *x, uint32_t n_rows, uint32_t head_dim) { + if (!x || n_rows == 0 || head_dim != 128u || + x->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) { + return 0; + } + indexer_hadamard_fp4_kernel<<>>((float *)x->ptr, n_rows, head_dim); + return cuda_ok(cudaGetLastError(), "indexer_hadamard_fp4 launch"); +} +extern "C" int ds4_gpu_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t 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) { + if (!x || n_rot > head_dim || (n_rot & 1) || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; + uint32_t pairs = n_tok * n_head * (n_rot / 2); + rope_tail_kernel<<<(pairs + 255) / 256, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, 1, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "rope_tail launch"); +} +extern "C" int ds4_gpu_rope_tail_decode_rows_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_attention_decode_row *rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + 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) { + if (!x || !rows || n_rows == 0u || + n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u || + n_rot == 0u || n_rot > head_dim || (n_rot & 1u) != 0u || + x->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) { + return 0; + } + cuda_attention_decode_row_table table; + memset(&table, 0, sizeof(table)); + for (uint32_t i = 0; i < n_rows; i++) table.row[i].pos = rows[i].pos; + const uint32_t pairs = n_rows * n_head * (n_rot / 2u); + rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>( + (float *)x->ptr, table, n_rows, n_head, head_dim, n_rot, + n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "rope tail decode rows launch"); +} +extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim); +extern "C" int ds4_gpu_kv_fp8_store_raw_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t raw_row, + uint32_t head_dim, + uint32_t n_rot) { + if (!kv || !raw_cache || raw_cap == 0u || n_rot > head_dim || + kv->device_id != raw_cache->device_id || + kv->bytes < (uint64_t)head_dim * sizeof(float) || + raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) { + return 0; + } + cuda_attention_decode_row_table table; + memset(&table, 0, sizeof(table)); + table.row[0].raw_kv = (uint64_t)(uintptr_t)raw_cache->ptr; + table.row[0].raw_cap = raw_cap; + table.row[0].raw_start = raw_row % raw_cap; + fp8_kv_quantize_store_rows_kernel<<<1, 64>>>( + (float *)kv->ptr, table, 1u, head_dim, n_rot); + return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store launch"); +} + +extern "C" int ds4_gpu_kv_fp8_store_raw_decode_rows_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *const *raw_caches, + const uint32_t *raw_caps, + const uint32_t *raw_rows, + uint32_t n_rows, + uint32_t head_dim, + uint32_t n_rot) { + if (!kv || !raw_caches || !raw_caps || !raw_rows || n_rows == 0u || + n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_rot > head_dim || + kv->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) { + return 0; + } + cuda_attention_decode_row_table table; + memset(&table, 0, sizeof(table)); + for (uint32_t i = 0; i < n_rows; i++) { + const ds4_gpu_tensor *raw = raw_caches[i]; + if (!raw || raw_caps[i] == 0u || raw_rows[i] >= raw_caps[i] || + raw->device_id != kv->device_id || + raw->bytes < (uint64_t)raw_caps[i] * head_dim * sizeof(float)) { + return 0; + } + table.row[i].raw_kv = (uint64_t)(uintptr_t)raw->ptr; + table.row[i].raw_cap = raw_caps[i]; + table.row[i].raw_start = raw_rows[i]; + } + fp8_kv_quantize_store_rows_kernel<<>>( + (float *)kv->ptr, table, n_rows, head_dim, n_rot); + return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store rows launch"); +} +extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim) { + if (!raw_cache || !kv || raw_cap == 0 || + raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + kv->bytes < (uint64_t)head_dim * sizeof(float)) return 0; + store_raw_kv_batch_kernel<<<(head_dim + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, row, 1, head_dim); + return cuda_ok(cudaGetLastError(), "store_raw_kv launch"); +} +extern "C" int ds4_gpu_store_raw_kv_batch_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) { + if (!raw_cache || !kv || raw_cap == 0 || + raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float)) return 0; + uint64_t n = (uint64_t)n_tokens * head_dim; + store_raw_kv_batch_kernel<<<(n + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, pos0, n_tokens, head_dim); + return cuda_ok(cudaGetLastError(), "store_raw_kv_batch launch"); +} +extern "C" int ds4_gpu_compressor_store_batch_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens) { + if (!kv || !sc || !state_kv || !state_score || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + (ape_type != 0u && ape_type != 1u)) { + return 0; + } + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + kv->bytes < kv_bytes || sc->bytes < kv_bytes || + state_kv->bytes < state_bytes || state_score->bytes < state_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(state_kv); + const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); + if (!ape) return 0; + uint64_t n = (uint64_t)n_tokens * width; + compressor_store_kernel<<<(n + 255) / 256, 256>>>( + (const float *)kv->ptr, + (const float *)sc->ptr, + (float *)state_kv->ptr, + (float *)state_score->ptr, + ape, + 0, + ape_type, + head_dim, + ratio, + pos0, + n_tokens); + return cuda_ok(cudaGetLastError(), "compressor store launch"); +} + +extern "C" int ds4_gpu_compressor_update_tensor( + const ds4_gpu_tensor *kv_cur, + const ds4_gpu_tensor *sc_cur, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + ds4_gpu_tensor *comp_cache, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos, + uint32_t comp_row, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps, + bool state_already_stored) { + if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || + !model_map || head_dim == 0 || ratio == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { + return 0; + } + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset || + kv_cur->bytes < kv_bytes || sc_cur->bytes < kv_bytes || + state_kv->bytes < state_bytes || state_score->bytes < state_bytes || + (emit && comp_cache->bytes < comp_bytes)) { + return 0; + } + if (!state_already_stored) { + if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score, + model_map, model_size, ape_offset, ape_type, + head_dim, ratio, pos, 1)) { + return 0; + } + } + if (!emit) return 1; + ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( + comp_cache, + (uint64_t)comp_row * head_dim * sizeof(float), + (uint64_t)head_dim * sizeof(float)); + if (!comp_row_view) return 0; + compressor_update_pool_kernel<<<(head_dim + 255) / 256, 256>>>( + (float *)comp_row_view->ptr, + (const float *)state_kv->ptr, + (const float *)state_score->ptr, + head_dim, + ratio); + int ok = cuda_ok(cudaGetLastError(), "compressor update pool launch"); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view, + model_map, model_size, norm_offset, + head_dim, 1, rms_eps); + if (ok) ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, 1, head_dim, n_rot, + pos + 1u - ratio, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + ds4_gpu_tensor_free(comp_row_view); + if (ok && ratio == 4u) { + uint64_t half = 4ull * width; + compressor_shift_ratio4_kernel<<<(half + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, width); + ok = cuda_ok(cudaGetLastError(), "compressor ratio4 shift launch"); + } + return ok; +} +extern "C" int ds4_gpu_compressor_prefill_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { + return 0; + } + + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t n_comp = n_tokens / ratio; + const uint32_t cutoff = n_comp * ratio; + const uint32_t rem = n_tokens - cutoff; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset || + kv->bytes < kv_bytes || sc->bytes < kv_bytes || + state_kv->bytes < state_bytes || state_score->bytes < state_bytes || + (n_comp && comp_cache->bytes < comp_bytes)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(state_kv); + const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); + if (!ape) return 0; + + uint64_t state_n = (uint64_t)state_rows * width; + if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), + "compressor state kv zero")) return 0; + fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); + if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0; + + if (ratio == 4u) { + if (cutoff >= ratio) { + uint32_t prev_start = cutoff - ratio; + uint64_t n = (uint64_t)ratio * width; + compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, + (const float *)kv->ptr, (const float *)sc->ptr, + ape, 0, ape_type, width, ratio, pos0, + prev_start, 0, ratio); + if (!cuda_ok(cudaGetLastError(), "compressor prefill prev state launch")) return 0; + } + if (rem != 0) { + uint64_t n = (uint64_t)rem * width; + compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, + (const float *)kv->ptr, (const float *)sc->ptr, + ape, 0, ape_type, width, ratio, pos0, + cutoff, ratio, rem); + if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0; + } + } else if (rem != 0) { + uint64_t n = (uint64_t)rem * width; + compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, + (const float *)kv->ptr, (const float *)sc->ptr, + ape, 0, ape_type, width, ratio, pos0, + cutoff, 0, rem); + if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0; + } + if (n_comp != 0) { + dim3 grid((head_dim + 255) / 256, n_comp, 1); + compressor_prefill_pool_kernel<<>>( + (float *)comp_cache->ptr, + (const float *)kv->ptr, + (const float *)sc->ptr, + (const float *)state_kv->ptr, + (const float *)state_score->ptr, + ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 0); + if (!cuda_ok(cudaGetLastError(), "compressor prefill pool launch")) return 0; + if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache, + model_map, model_size, norm_offset, + head_dim, n_comp, rms_eps)) return 0; + if (n_rot != 0) { + const uint32_t pairs = n_comp * (n_rot / 2u); + rope_tail_kernel<<<(pairs + 255) / 256, 256>>>( + (float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot, + pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + if (!cuda_ok(cudaGetLastError(), "compressor prefill rope launch")) return 0; + } + if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0; + } + return 1; +} +extern "C" int ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { + return 0; + } + + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint32_t n_comp = n_tokens / ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset || + kv->bytes < kv_bytes || sc->bytes < kv_bytes || + state_kv->bytes < state_bytes || state_score->bytes < state_bytes || + comp_cache->bytes < comp_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(comp_cache); + const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); + if (!ape) return 0; + dim3 grid((head_dim + 255) / 256, n_comp, 1); + compressor_prefill_pool_kernel<<>>( + (float *)comp_cache->ptr, + (const float *)kv->ptr, + (const float *)sc->ptr, + (const float *)state_kv->ptr, + (const float *)state_score->ptr, + ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 1); + if (!cuda_ok(cudaGetLastError(), "compressor replay pool launch")) return 0; + if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache, + model_map, model_size, norm_offset, + head_dim, n_comp, rms_eps)) return 0; + if (n_rot != 0) { + const uint32_t pairs = n_comp * (n_rot / 2u); + rope_tail_kernel<<<(pairs + 255) / 256, 256>>>( + (float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot, + pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + if (!cuda_ok(cudaGetLastError(), "compressor replay rope launch")) return 0; + } + if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0; + + uint64_t state_n = (uint64_t)state_rows * width; + if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), + "compressor replay state kv zero")) return 0; + fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); + if (!cuda_ok(cudaGetLastError(), "compressor replay state score fill launch")) return 0; + uint32_t prev_start = n_tokens - ratio; + uint64_t n = (uint64_t)ratio * width; + compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, + (const float *)kv->ptr, (const float *)sc->ptr, + ape, 0, ape_type, width, ratio, pos0, + prev_start, 0, ratio); + return cuda_ok(cudaGetLastError(), "compressor replay state launch"); +} +extern "C" int ds4_gpu_compressor_prefill_state_ratio4_tensor( + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv_tail, + const ds4_gpu_tensor *sc_tail, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t pos0) { + if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map || + head_dim == 0 || (ape_type != 0u && ape_type != 1u)) { + return 0; + } + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape; + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + kv_tail->bytes < tail_bytes || sc_tail->bytes < tail_bytes || + state_kv->bytes < state_bytes || state_score->bytes < state_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(state_kv); + const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); + if (!ape) return 0; + uint64_t state_n = (uint64_t)state_rows * width; + if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), + "compressor state kv zero")) return 0; + fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); + if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0; + uint64_t n = (uint64_t)ratio * width; + compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( + (float *)state_kv->ptr, (float *)state_score->ptr, + (const float *)kv_tail->ptr, (const float *)sc_tail->ptr, + ape, 0, ape_type, width, ratio, pos0, + 0, 0, ratio); + return cuda_ok(cudaGetLastError(), "compressor state set launch"); +} + +/* perf-02 split-KV / flash-decode launch helper (opt-in, default OFF). + * Returns 1 if the split path handled the launch, 0 if the caller should fall + * through to the existing attention_decode_mixed_kernel path. + * + * Engages only for the single-token decode shape (n_tokens==1) and only when + * DS4_CUDA_SPLITKV_DECODE is set. S==1 is NOT handled here: the caller dispatches + * the old kernel as the bit-exact anchor when S would be 1. + */ +static int attention_decode_splitkv_launch( + int logical_tier, + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + /* n_tokens is fixed at 1 for the split path; compute the EXACT logical row + * count the kernel will use (raw_count + visible_comp) so S is sized to the + * real work. raw_count MUST apply the same window logic as the kernel / + * reference, otherwise a true-S==1 case (e.g. ratio=1, window=1, n_raw>=2, + * n_comp=0) could be over-estimated to S>1 and engage split-KV instead of + * the bit-exact old-kernel anchor. The count is head-independent. */ + const bool single_all = (ratio == 0u); + uint32_t qpos = pos0; /* t==0, n_tokens==1 */ + uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return 0; /* nothing to do; let old path handle it */ + const int manual_splitkv = cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0); + const uint32_t scoped_min_score = + (g_decode_fast_attention && !manual_splitkv) ? 512u : 0u; + uint32_t min_score = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_MIN_SCORE", + scoped_min_score, 0u, + DS4_CUDA_ATTENTION_SCORE_CAP, + NULL); + if (n_score < min_score) return 0; + /* S = clamp(ceil(n_score / CHUNK), 1, S_MAX); raise to S_FLOOR for short + * context to fill more SMs, but never exceed n_score (no empty chunks). + * Optional tuning knobs are guarded by min_needed so every block's chunk + * still fits the fixed shared score buffer. */ + const uint32_t split_cap = DS4_CUDA_SPLITKV_SCORE_CAP; + const uint32_t min_needed = (n_score + split_cap - 1u) / split_cap; + uint32_t chunk = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_CHUNK", + DS4_CUDA_SPLITKV_CHUNK, + 1u, split_cap, NULL); + uint32_t s_floor = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_FLOOR", + DS4_CUDA_SPLITKV_S_FLOOR, + 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); + uint32_t s_max = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_MAX", + DS4_CUDA_SPLITKV_S_MAX, + 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); + int exact_present = 0; + uint32_t S = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S", + 0u, 1u, DS4_CUDA_SPLITKV_S_MAX, + &exact_present); + if (!exact_present) { + S = (n_score + chunk - 1u) / chunk; + if (S < s_floor) S = s_floor < n_score ? s_floor : n_score; + if (S > s_max) S = s_max; + } + if (S < min_needed) S = min_needed; + if (S > n_score) S = n_score; + if (S <= 1u) return 0; /* S==1: caller uses the old kernel anchor */ + if (cuda_env_flag_enabled("DS4_CUDA_SPLITKV_GLOBAL_SOFTMAX", 0)) { + const uint64_t score_count = (uint64_t)n_head * n_score; + const uint64_t score_bytes = score_count * sizeof(float); + const uint64_t denom_offset = (score_bytes + 255u) & ~255ull; + const uint64_t denom_bytes = (uint64_t)n_head * sizeof(float); + const uint64_t partial_offset = (denom_offset + denom_bytes + 255u) & ~255ull; + const uint64_t partial_bytes = (uint64_t)n_head * S * head_dim * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, + partial_offset + partial_bytes, + "attention splitkv global softmax"); + if (!tmp) return 0; + float *scores = (float *)tmp; + float *denom = (float *)((char *)tmp + denom_offset); + float *partials = (float *)((char *)tmp + partial_offset); + dim3 score_grid(1, n_head, S); + attention_decode_score_split_scores_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim, S); + if (!cuda_ok(cudaGetLastError(), "attention splitkv global score launch")) return -1; + attention_decode_global_softmax_kernel<<>>( + scores, denom, sinks, n_score, n_head); + if (!cuda_ok(cudaGetLastError(), "attention splitkv global softmax launch")) return -1; + dim3 value_grid(1, n_head, S); + attention_decode_split_value_kernel<<>>( + partials, scores, raw_kv, comp_kv, raw_count, raw_first_idx, + raw_cap, raw_start, n_score, n_head, head_dim, S); + if (!cuda_ok(cudaGetLastError(), "attention splitkv global value launch")) return -1; + dim3 combine_grid(1, n_head, 1); + attention_decode_split_value_combine_kernel<<>>( + heads, partials, denom, n_head, head_dim, S); + if (!cuda_ok(cudaGetLastError(), "attention splitkv global combine launch")) return -1; + return 1; + } + /* Partials scratch: n_head * S * (head_dim + 2) floats (n_tokens==1). */ + uint64_t stride = (uint64_t)head_dim + 2u; + uint64_t count = (uint64_t)n_head * S * stride; + float *partials = (float *)cuda_tmp_alloc_on(logical_tier, count * sizeof(float), + "attention splitkv partials"); + if (!partials) return 0; + dim3 split_grid(1, n_head, S); + attention_decode_splitkv_kernel<<>>(partials, + q, + raw_kv, + comp_kv, + comp_mask, + use_comp_mask, + 1, pos0, n_raw, raw_cap, raw_start, + n_comp, window, ratio, n_head, head_dim, S); + if (!cuda_ok(cudaGetLastError(), "attention splitkv partial launch")) return -1; + dim3 combine_grid(1, n_head, 1); + attention_decode_splitkv_combine_kernel<<>>(heads, + sinks, + partials, + 1, n_head, head_dim, S); + if (!cuda_ok(cudaGetLastError(), "attention splitkv combine launch")) return -1; + return 1; +} + +extern "C" int ds4_gpu_attention_decode_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16 || + !heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || + (use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float))) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); + if (!sinks) return 0; + if (!cuda_attention_score_buffer_fits(n_comp)) { + if (!use_mask && head_dim == 512u && + !g_cuda_no_window_attention) { + const uint32_t synthetic_pos0 = n_raw - 1u; + dim3 online_grid(1, (n_head + 7u) / 8u, 1); + attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + 1, + synthetic_pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + 0, + 0, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode online launch"); + } + fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp); + return 0; + } + if (!use_mask && head_dim == 512u && + g_cuda_decode_heads8_online && + !g_cuda_no_window_attention) { + const uint32_t synthetic_pos0 = n_raw - 1u; + dim3 online_grid(1, (n_head + 7u) / 8u, 1); + attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + 1, + synthetic_pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + 0, + 0, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode heads8 online launch"); + } + const uint32_t score_lanes = + g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); + const uint32_t threads = + head_dim == 512u && score_lanes == 0u && + !g_cuda_no_decode_value512 ? 512u : 256u; + int score_split_rc = attention_decode_score_split_launch( + logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, + 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, + n_head, head_dim, threads, NULL); + if (score_split_rc == 1) { + return cuda_ok(cudaGetLastError(), "attention exact score split launch"); + } + if (score_split_rc < 0) return 0; + /* perf-02 split-KV opt-in (default OFF). n_tokens==1 here by construction. + * S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */ + if (cuda_splitkv_decode_requested()) { + int rc = attention_decode_splitkv_launch( + logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, + 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, n_head, head_dim); + if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv launch"); + if (rc < 0) return 0; + } + dim3 grid(1, n_head, 1); + attention_decode_mixed_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_mask ? (const float *)comp_mask->ptr : NULL, + use_mask, + 1, 0, n_raw, raw_cap, raw_start, n_comp, + 0, 0, n_head, head_dim, + score_lanes); + return cuda_ok(cudaGetLastError(), "attention decode launch"); +} + +extern "C" int ds4_gpu_attention_decode_heads_rope_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int *fused_inv_rope) { + if (fused_inv_rope) *fused_inv_rope = 0; + if (!g_cuda_exact_score_split_fuse_inv_rope || + n_rot == 0u || n_rot > head_dim || (n_rot & 1u) || + head_dim != 512u) { + return ds4_gpu_attention_decode_heads_tensor( + heads, model_map, model_size, sinks_offset, q, raw_kv, + n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, + comp_mask, use_mask, n_head, head_dim); + } + if (!use_mask && g_cuda_decode_heads8_online && !g_cuda_no_window_attention) { + return ds4_gpu_attention_decode_heads_tensor( + heads, model_map, model_size, sinks_offset, q, raw_kv, + n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, + comp_mask, use_mask, n_head, head_dim); + } + if (comp_kv_f16 || + !heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || + (use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float)) || + !cuda_attention_score_buffer_fits(n_comp)) { + return ds4_gpu_attention_decode_heads_tensor( + heads, model_map, model_size, sinks_offset, q, raw_kv, + n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, + comp_mask, use_mask, n_head, head_dim); + } + const uint32_t score_lanes = + g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); + const uint32_t threads = + score_lanes == 0u && !g_cuda_no_decode_value512 ? 512u : 256u; + if (threads < 512u) { + return ds4_gpu_attention_decode_heads_tensor( + heads, model_map, model_size, sinks_offset, q, raw_kv, + n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, + comp_mask, use_mask, n_head, head_dim); + } + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), + logical_tier, "attn_sinks"); + if (!sinks) return 0; + cuda_attention_inv_rope_params rope; + rope.n_rot = n_rot; + rope.pos0 = pos0; + rope.n_ctx_orig = n_ctx_orig; + rope.freq_base = freq_base; + rope.freq_scale = freq_scale; + rope.ext_factor = ext_factor; + rope.attn_factor = attn_factor; + rope.beta_fast = beta_fast; + rope.beta_slow = beta_slow; + int score_split_rc = attention_decode_score_split_launch( + logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, + 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, + n_head, head_dim, threads, &rope); + if (score_split_rc == 1) { + if (fused_inv_rope) *fused_inv_rope = 1; + return cuda_ok(cudaGetLastError(), + "attention exact score split fused inv rope launch"); + } + if (score_split_rc < 0) return 0; + return ds4_gpu_attention_decode_heads_tensor( + heads, model_map, model_size, sinks_offset, q, raw_kv, + n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, + comp_mask, use_mask, n_head, head_dim); +} + +extern "C" int ds4_gpu_attention_decode_rows_rope_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_attention_decode_row *rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!heads || !q || !rows || !model_map || n_rows < 2u || + n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u || + head_dim != 512u || n_rot == 0u || n_rot > head_dim || + (n_rot & 1u) != 0u || sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(heads); + if (logical_tier < 0 || logical_tier >= g_n_gpus || + ds4_tensor_device_idx(q) != logical_tier) { + return 0; + } + + /* This first grouped path mirrors the promoted default decode exactly. + * Alternative score kernels and graph/split-KV experiments retain the + * one-session dispatcher until they gain equivalent row-table variants. */ + if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0) || + !cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1) || + cuda_splitkv_decode_requested() || + g_cuda_decode_heads8_online || g_cuda_decode_score4 || + g_cuda_decode_score8 || g_cuda_no_decode_value512 || + g_cuda_exact_score_split_graph || g_cuda_exact_score_split_ldg || + g_cuda_exact_score_split_vec4 || + g_cuda_exact_score_split_vec4_plain || + g_cuda_exact_score_split_dim2 || + g_cuda_exact_score_split_fuse_inv_rope || + getenv("DS4_CUDA_NO_SCORE_TILE") != NULL || + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE") != NULL || + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK") != NULL || + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR") != NULL || + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX") != NULL || + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S") != NULL) { + return 0; + } + + cuda_attention_decode_row_table table; + memset(&table, 0, sizeof(table)); + uint32_t max_dense_score = 0u; + bool have_dense = false; + bool have_indexed = false; + for (uint32_t i = 0; i < n_rows; i++) { + const ds4_gpu_attention_decode_row r = rows[i]; + if (r.raw_kv == 0u || r.n_raw == 0u || r.raw_cap < r.n_raw || + r.raw_start >= r.raw_cap || (r.n_comp != 0u && r.comp_kv == 0u)) { + return 0; + } + if (r.indexed) { + if (r.comp_kv == 0u || r.topk == 0u || r.n_comp == 0u || + r.top_k == 0u || r.top_k > 512u || r.ratio == 0u) { + return 0; + } + have_indexed = true; + } else { + const uint32_t raw_count = r.n_raw > 256u ? 256u : r.n_raw; + const uint32_t n_score = raw_count + r.n_comp; + /* n_score==1 takes the legacy one-block kernel and is not a + * score-split shape. Decode after any nonempty prompt is >1. */ + if (n_score <= 1u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) { + return 0; + } + if (n_score > max_dense_score) max_dense_score = n_score; + have_dense = true; + } + table.row[i] = r; + } + + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), + logical_tier, "attn_sinks_rows"); + if (!sinks) return 0; + + if (have_dense) { + if ((uint64_t)n_rows > UINT64_MAX / n_head || + (uint64_t)n_rows * n_head > UINT64_MAX / max_dense_score) { + return 0; + } + const uint64_t score_count = + (uint64_t)n_rows * n_head * max_dense_score; + float *scores = (float *)cuda_tmp_alloc_on( + logical_tier, score_count * sizeof(float), + "attention exact decode rows"); + if (!scores) return 0; + + const size_t tile_shmem = + (size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) * + DS4_SCORE_TILE_STRIDE * sizeof(float); + static int tile_shmem_ready[DS4_MAX_GPUS] = {0}; + int physical_device = 0; + if (cudaGetDevice(&physical_device) != cudaSuccess || + physical_device < 0 || physical_device >= DS4_MAX_GPUS) { + return 0; + } + if (!tile_shmem_ready[physical_device]) { + if (!cuda_ok(cudaFuncSetAttribute( + attention_decode_score_split_scores_tile512_rows_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + (int)tile_shmem), + "attention score rows shared-memory opt-in")) { + return 0; + } + tile_shmem_ready[physical_device] = 1; + } + dim3 score_grid( + (max_dense_score + DS4_SCORE_TILE_ROWS - 1u) / + DS4_SCORE_TILE_ROWS, + (n_head + DS4_SCORE_TILE_HEADS - 1u) / + DS4_SCORE_TILE_HEADS, + n_rows); + attention_decode_score_split_scores_tile512_rows_kernel + <<>>( + scores, (const float *)q->ptr, table, n_rows, + max_dense_score, n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), + "attention exact score rows launch")) { + return 0; + } + dim3 final_grid(n_rows, n_head, 1u); + attention_decode_score_split_finalize_rows_kernel + <<>>( + (float *)heads->ptr, sinks, scores, table, n_rows, + max_dense_score, n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), + "attention exact finalize rows launch")) { + return 0; + } + } + if (have_indexed) { + dim3 indexed_grid(n_rows, n_head, 1u); + attention_indexed_mixed_decode_rows_kernel<<>>( + (float *)heads->ptr, sinks, (const float *)q->ptr, table, + n_rows, n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), + "attention indexed decode rows launch")) { + return 0; + } + } + + const uint32_t pairs = n_rows * n_head * (n_rot / 2u); + rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>( + (float *)heads->ptr, table, n_rows, n_head, head_dim, n_rot, + n_ctx_orig, 1, freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), + "attention decode rows inverse rope launch"); +} + +extern "C" int ds4_gpu_attention_prefill_raw_heads_tensor(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, uint64_t sinks_offset, const ds4_gpu_tensor *q, const ds4_gpu_tensor *raw_kv, uint32_t n_tokens, uint32_t window, uint32_t n_head, uint32_t head_dim) { + if (!heads || !q || !raw_kv || !model_map || sinks_offset > model_size || + model_size - sinks_offset < (uint64_t)n_head * sizeof(float) || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || + window > 256) return 0; + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); + if (!sinks) return 0; + if (n_tokens > 1 && head_dim == 512 && + getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL && + (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { + dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); + attention_static_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + (const float *)raw_kv->ptr, + n_tokens, + 0, + window, + 1, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention raw window launch"); + } + if (g_cublas_ready && n_tokens > 1 && head_dim == 512 && + getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) { + const uint32_t n_keys = n_tokens; + const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys; + const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim; + const uint64_t score_bytes = score_count * sizeof(float); + const uint64_t out_offset = (score_bytes + 255u) & ~255ull; + const uint64_t tmp_bytes = out_offset + out_count * sizeof(float); + float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention raw cublas"); + if (!tmp) return 0; + float *scores = tmp; + float *out_tmp = (float *)((char *)tmp + out_offset); + const float alpha = rsqrtf((float)head_dim); + const float beta = 0.0f; + cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_keys, + (int)n_tokens, + (int)head_dim, + &alpha, + (const float *)raw_kv->ptr, + (int)head_dim, + 0, + (const float *)q->ptr, + (int)(n_head * head_dim), + (long long)head_dim, + &beta, + scores, + (int)n_keys, + (long long)n_keys * n_tokens, + (int)n_head); + if (!cublas_ok(st, "attention raw score gemm")) return 0; + dim3 sgrid(n_tokens, n_head, 1); + attention_prefill_raw_softmax_kernel<<>>(scores, sinks, n_tokens, window, n_keys); + if (!cuda_ok(cudaGetLastError(), "attention raw softmax launch")) return 0; + const float one = 1.0f; + st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_N, + CUBLAS_OP_N, + (int)head_dim, + (int)n_tokens, + (int)n_keys, + &one, + (const float *)raw_kv->ptr, + (int)head_dim, + 0, + scores, + (int)n_keys, + (long long)n_keys * n_tokens, + &beta, + out_tmp, + (int)head_dim, + (long long)head_dim * n_tokens, + (int)n_head); + if (!cublas_ok(st, "attention raw value gemm")) return 0; + uint64_t n = (uint64_t)n_tokens * n_head * head_dim; + attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr, + out_tmp, + n_tokens, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention raw unpack launch"); + } + dim3 grid(n_tokens, n_head, 1); + attention_prefill_raw_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_tokens, window, n_head, head_dim); + return cuda_ok(cudaGetLastError(), "attention_prefill_raw launch"); +} +static int attention_decode_batch_launch( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16 || + !heads || !q || !raw_kv || !model_map || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + (n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || + (use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) { + return 0; + } + if (n_comp != 0 && ratio == 0) return 0; + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); + if (!sinks) return 0; + if (!cuda_attention_score_buffer_fits(n_comp)) { + if (!use_comp_mask && head_dim == 512u && + !g_cuda_no_window_attention) { + dim3 online_grid(n_tokens, (n_head + 7u) / 8u, 1); + attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode online launch"); + } + fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp); + return 0; + } + if (!use_comp_mask && n_tokens > 1 && head_dim == 512 && + !g_cuda_no_window_attention && + (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { + dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); + attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode window launch"); + } + if (!use_comp_mask && n_tokens == 1u && head_dim == 512 && + g_cuda_decode_heads8_online && + !g_cuda_no_window_attention) { + dim3 grid(1, (n_head + 7u) / 8u, 1); + attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode heads8 online batch launch"); + } + const uint32_t score_lanes = + g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); + const uint32_t threads = + n_tokens == 1u && head_dim == 512u && score_lanes == 0u && + !g_cuda_no_decode_value512 ? 512u : 256u; + if (n_tokens == 1u) { + int score_split_rc = attention_decode_score_split_launch( + logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim, threads, NULL); + if (score_split_rc == 1) { + return cuda_ok(cudaGetLastError(), "attention exact score split batch launch"); + } + if (score_split_rc < 0) return 0; + } + /* perf-02 split-KV opt-in (default OFF). Single-token decode only; multi- + * token batch shapes already fill the grid and fall through unchanged. + * S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */ + if (n_tokens == 1u && cuda_splitkv_decode_requested()) { + int rc = attention_decode_splitkv_launch( + logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, n_head, head_dim); + if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv batch launch"); + if (rc < 0) return 0; + } + dim3 grid(n_tokens, n_head, 1); + attention_decode_mixed_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_comp_mask ? (const float *)comp_mask->ptr : NULL, + use_comp_mask, n_tokens, pos0, n_raw, raw_cap, + raw_start, n_comp, window, ratio, n_head, head_dim, + score_lanes); + return cuda_ok(cudaGetLastError(), "attention decode batch launch"); +} + +extern "C" int ds4_gpu_attention_decode_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset, + q, raw_kv, NULL, 0, NULL, 0, n_tokens, pos0, + n_raw, raw_cap, raw_start, 0, window, 1, + n_head, head_dim); +} + +extern "C" int ds4_gpu_attention_decode_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16) return 0; + return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset, + q, raw_kv, comp_kv, comp_kv_f16, comp_mask, use_comp_mask, + n_tokens, pos0, n_raw, raw_cap, raw_start, + n_comp, window, ratio, n_head, head_dim); +} + +extern "C" int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *topk, + uint32_t n_tokens, + uint32_t pos0, + 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 n_head, + uint32_t head_dim) { + if (comp_kv_f16 || + !heads || !q || !raw_kv || !comp_kv || !topk || !model_map || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + n_comp == 0 || top_k == 0 || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || + comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float) || + topk->bytes < (uint64_t)n_tokens * top_k * sizeof(int32_t)) { + return 0; + } + if (top_k > 512u) return 0; + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); + if (!sinks) return 0; + const int32_t *topk_ptr = (const int32_t *)topk->ptr; + if (n_tokens > 1u && top_k == 512u && + getenv("DS4_CUDA_NO_INDEXED_TOPK_SORT") == NULL) { + const uint64_t sort_bytes = (uint64_t)n_tokens * top_k * sizeof(int32_t); + int32_t *sorted = (int32_t *)cuda_tmp_alloc_on(logical_tier, sort_bytes, "indexed attention topk sort"); + if (!sorted) return 0; + indexed_topk_sort_512_asc_kernel<<>>(sorted, topk_ptr, n_tokens); + if (!cuda_ok(cudaGetLastError(), "indexed attention topk sort launch")) return 0; + topk_ptr = sorted; + } + if (n_tokens > 1 && head_dim == 512 && top_k <= 512u && + getenv("DS4_CUDA_NO_INDEXED_HEADS8") == NULL) { + if (getenv("DS4_CUDA_INDEXED_TWOPASS") == NULL) { + dim3 grid(n_tokens, (n_head + 15u) / 16u, 1); + attention_indexed_mixed_heads8_online_kernel<8, 16><<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + (const float *)comp_kv->ptr, + topk_ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + top_k, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention indexed online launch"); + } + dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); + attention_indexed_mixed_heads8_rb4_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + (const float *)comp_kv->ptr, + topk_ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + top_k, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention indexed heads8 launch"); + } + dim3 grid(n_tokens, n_head, 1); + attention_indexed_mixed_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + (const float *)comp_kv->ptr, + topk_ptr, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + top_k, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention indexed mixed launch"); +} + +static int attention_prefill_mixed_launch( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || ratio == 0 || + (n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || + (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || + (use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(heads); + const float *sinks = (const float *)cuda_resolve_weight_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); + if (!sinks) return 0; + if (!use_comp_mask && n_tokens > 1 && head_dim == 512 && + getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL && + (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { + dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); + attention_static_mixed_heads8_online_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention mixed window launch"); + } + if (g_cublas_ready && n_tokens > 1 && head_dim == 512 && + getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) { + const uint32_t n_keys = n_tokens + n_comp; + const uint64_t kv_count = (uint64_t)n_keys * head_dim; + const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys; + const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim; + const uint64_t kv_bytes = kv_count * sizeof(float); + const uint64_t score_offset = (kv_bytes + 255u) & ~255ull; + const uint64_t score_bytes = score_count * sizeof(float); + const uint64_t out_offset = score_offset + ((score_bytes + 255u) & ~255ull); + const uint64_t tmp_bytes = out_offset + out_count * sizeof(float); + float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention mixed cublas"); + if (!tmp) return 0; + float *kv = tmp; + float *scores = (float *)((char *)tmp + score_offset); + float *out_tmp = (float *)((char *)tmp + out_offset); + attention_prefill_pack_mixed_kv_kernel<<<(kv_count + 255) / 256, 256>>>( + kv, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + n_tokens, + n_comp, + head_dim); + if (!cuda_ok(cudaGetLastError(), "attention mixed kv pack launch")) return 0; + const float alpha = rsqrtf((float)head_dim); + const float beta = 0.0f; + cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_keys, + (int)n_tokens, + (int)head_dim, + &alpha, + kv, + (int)head_dim, + 0, + (const float *)q->ptr, + (int)(n_head * head_dim), + (long long)head_dim, + &beta, + scores, + (int)n_keys, + (long long)n_keys * n_tokens, + (int)n_head); + if (!cublas_ok(st, "attention mixed score gemm")) return 0; + dim3 sgrid(n_tokens, n_head, 1); + attention_prefill_mixed_softmax_kernel<<>>( + scores, + sinks, + use_comp_mask ? (const float *)comp_mask->ptr : NULL, + use_comp_mask, + n_tokens, + n_comp, + window, + ratio, + n_keys); + if (!cuda_ok(cudaGetLastError(), "attention mixed softmax launch")) return 0; + const float one = 1.0f; + st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_N, + CUBLAS_OP_N, + (int)head_dim, + (int)n_tokens, + (int)n_keys, + &one, + kv, + (int)head_dim, + 0, + scores, + (int)n_keys, + (long long)n_keys * n_tokens, + &beta, + out_tmp, + (int)head_dim, + (long long)head_dim * n_tokens, + (int)n_head); + if (!cublas_ok(st, "attention mixed value gemm")) return 0; + uint64_t n = (uint64_t)n_tokens * n_head * head_dim; + attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr, + out_tmp, + n_tokens, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention mixed unpack launch"); + } + dim3 grid(n_tokens, n_head, 1); + attention_prefill_mixed_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, + use_comp_mask ? (const float *)comp_mask->ptr : NULL, + use_comp_mask, n_tokens, n_comp, window, ratio, + n_head, head_dim); + return cuda_ok(cudaGetLastError(), "attention prefill mixed launch"); +} + +extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16) return 0; + return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset, + q, raw_kv, comp_kv, NULL, 0, n_tokens, + n_comp, window, ratio, n_head, head_dim); +} + +extern "C" int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16) return 0; + return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset, + q, raw_kv, comp_kv, comp_mask, 1, n_tokens, + n_comp, window, ratio, n_head, head_dim); +} +extern "C" int ds4_gpu_attention_output_q8_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + 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, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + (void)group_tmp; + (void)low_tmp; + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0) { + return 0; + } + const uint64_t low_dim = (uint64_t)n_groups * rank; + const uint64_t blocks_a = (group_dim + 31) / 32; + const uint64_t blocks_b = (low_dim + 31) / 32; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * blocks_a * 34; + const uint64_t out_b_bytes = out_dim * blocks_b * 34; + if (out_a_offset > model_size || out_b_offset > model_size || + out_a_bytes > model_size - out_a_offset || + out_b_bytes > model_size - out_b_offset || + heads->bytes < (uint64_t)n_tokens * n_groups * group_dim * sizeof(float) || + low->bytes < (uint64_t)n_tokens * low_dim * sizeof(float) || + out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out); + const int physical_device = + (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) + ? g_gpu[logical_tier].device_id : 0; + const unsigned char *out_a = reinterpret_cast( + cuda_resolve_weight_ptr(model_map, out_a_offset, out_a_bytes, logical_tier, "attn_out_a")); + const unsigned char *out_b = reinterpret_cast( + cuda_resolve_weight_ptr(model_map, out_b_offset, out_b_bytes, logical_tier, "attn_out_b")); + if (!out_a || !out_b) return 0; + + const uint32_t profile = getenv("DS4_CUDA_ATTN_OUTPUT_PROFILE") != NULL; + cudaEvent_t prof_ev[3] = {NULL, NULL, NULL}; + if (profile) { + for (uint32_t i = 0; i < 3u; i++) { + if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) { + for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]); + memset(prof_ev, 0, sizeof(prof_ev)); + break; + } + } + if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0); + } + + const __half *out_a_f16 = NULL; + uint32_t out_a_cublas_min_tokens = 2u; + const char *out_a_min_env = getenv("DS4_CUDA_ATTENTION_OUTPUT_A_CUBLAS_MIN"); + if (out_a_min_env && out_a_min_env[0]) { + char *endp = NULL; + long v = strtol(out_a_min_env, &endp, 10); + if (endp != out_a_min_env && v > 1 && v < 4096) out_a_cublas_min_tokens = (uint32_t)v; + } + if (!g_quality_mode && + g_cublas_ready && + n_tokens >= out_a_cublas_min_tokens && + getenv("DS4_CUDA_NO_CUBLAS_ATTENTION_OUTPUT_A") == NULL) { + out_a_f16 = cuda_q8_f16_ptr(model_map, out_a_offset, out_a_bytes, group_dim, low_dim, physical_device, "attn_output_a"); + } + if (out_a_f16) { + const uint64_t heads_h_count = (uint64_t)n_groups * n_tokens * group_dim; + const uint64_t low_tmp_count = (uint64_t)n_groups * n_tokens * rank; + const uint64_t heads_h_bytes = heads_h_count * sizeof(__half); + const uint64_t low_tmp_offset = (heads_h_bytes + 255u) & ~255ull; + const uint64_t tmp_bytes = low_tmp_offset + low_tmp_count * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a cublas"); + if (!tmp) return 0; + __half *heads_h = (__half *)tmp; + float *low_packed = (float *)((char *)tmp + low_tmp_offset); + attention_pack_group_heads_f16_kernel<<<(heads_h_count + 255) / 256, 256>>>( + heads_h, + (const float *)heads->ptr, + n_tokens, + n_groups, + group_dim); + if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a pack launch")) return 0; + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasGemmStridedBatchedEx(cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)rank, + (int)n_tokens, + (int)group_dim, + &alpha, + out_a_f16, + CUDA_R_16F, + (int)group_dim, + (long long)rank * group_dim, + heads_h, + CUDA_R_16F, + (int)group_dim, + (long long)n_tokens * group_dim, + &beta, + low_packed, + CUDA_R_32F, + (int)rank, + (long long)rank * n_tokens, + (int)n_groups, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "attention output a gemm")) return 0; + attention_unpack_group_low_kernel<<<(low_tmp_count + 255) / 256, 256>>>( + (float *)low->ptr, + low_packed, + n_tokens, + n_groups, + rank); + if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a unpack launch")) return 0; + } else { + const uint64_t x_rows = (uint64_t)n_tokens * n_groups; + const uint64_t xq_bytes = x_rows * blocks_a * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a q8 prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1); + quantize_q8_0_f32_kernel<<>>(xq, + xscale, + (const float *)heads->ptr, + group_dim, + blocks_a); + if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a prequant launch")) return 0; + int grouped_mma_done = 0; + if (n_tokens >= 8u) { + /* One mma launch per group: T=32 matches the warp tree of the + * grouped reference kernels (multi-term slots for blocks > 32). */ + grouped_mma_done = 1; + for (uint32_t g = 0; g < n_groups && grouped_mma_done == 1; g++) { + const int rc = cuda_q8_mma_try_launch( + (float *)low->ptr + (uint64_t)g * rank, + reinterpret_cast(out_a) + + (uint64_t)g * rank * blocks_a * 34u, + xq + (uint64_t)g * blocks_a * 32u, + xscale + (uint64_t)g * blocks_a, + group_dim, rank, n_tokens, blocks_a, + (uint64_t)n_groups * blocks_a, low_dim, 32u); + if (rc < 0) return 0; + if (rc == 0) grouped_mma_done = 0; + } + } + if (grouped_mma_done) { + /* handled */ + } else if (getenv("DS4_CUDA_NO_ATTN_A_TOK2") == NULL && n_tokens >= 2u) { + dim3 grid_a(((unsigned)low_dim + 7u) / 8u, ((unsigned)n_tokens + 1u) / 2u, 1); + grouped_q8_0_a_preq_warp8_tok2_kernel<<>>((float *)low->ptr, + out_a, + xq, + xscale, + group_dim, + rank, + n_groups, + n_tokens, + blocks_a, + use_dp4a); + } else { + dim3 grid_a(((unsigned)low_dim + 7u) / 8u, (unsigned)n_tokens, 1); + grouped_q8_0_a_preq_warp8_kernel<<>>((float *)low->ptr, + out_a, + xq, + xscale, + group_dim, + rank, + n_groups, + n_tokens, + blocks_a, + use_dp4a); + } + if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a preq launch")) return 0; + } + + if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0); + (void)out_b; + int ok = cuda_matmul_q8_0_tensor_labeled(out, + model_map, + model_size, + out_b_offset, + low_dim, + out_dim, + low, + n_tokens, + "attn_output_b"); + if (prof_ev[2]) { + (void)cudaEventRecord(prof_ev[2], 0); + if (cudaEventSynchronize(prof_ev[2]) == cudaSuccess) { + float ms_a = 0.0f, ms_b = 0.0f, ms_total = 0.0f; + (void)cudaEventElapsedTime(&ms_a, prof_ev[0], prof_ev[1]); + (void)cudaEventElapsedTime(&ms_b, prof_ev[1], prof_ev[2]); + (void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[2]); + fprintf(stderr, + "ds4: CUDA attention output profile tokens=%u groups=%u group_dim=%llu rank=%llu low=%llu out=%llu A=%.3f B=%.3f total=%.3f ms\n", + n_tokens, + n_groups, + (unsigned long long)group_dim, + (unsigned long long)rank, + (unsigned long long)low_dim, + (unsigned long long)out_dim, + ms_a, + ms_b, + ms_total); + } + for (uint32_t i = 0; i < 3u; i++) (void)cudaEventDestroy(prof_ev[i]); + } + return ok; +} +extern "C" int ds4_gpu_attention_output_low_q8_rows_exact_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads, + uint32_t n_rows) { + if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || + n_groups_total == 0 || group_cnt == 0 || + group0 > n_groups_total || group_cnt > n_groups_total - group0 || + n_rows == 0 || (uint64_t)n_rows * group_cnt > 65535u) { + return 0; + } + const uint64_t low_dim = (uint64_t)group_cnt * rank; + const uint64_t blocks_a = (group_dim + 31) / 32; + const uint64_t row_a_bytes = blocks_a * 34u; + const uint64_t a_offset = + out_a_offset + (uint64_t)group0 * rank * row_a_bytes; + const uint64_t out_a_bytes = low_dim * row_a_bytes; + if (a_offset < out_a_offset || a_offset > model_size || + out_a_bytes > model_size - a_offset || + heads->bytes < (uint64_t)n_rows * n_groups_total * group_dim * + sizeof(float) || + low->bytes < (uint64_t)n_rows * low_dim * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(low); + const unsigned char *out_a = reinterpret_cast( + cuda_resolve_weight_ptr(model_map, a_offset, out_a_bytes, + logical_tier, "attn_out_a_rows")); + if (!out_a) return 0; + + const uint64_t x_rows = (uint64_t)n_rows * group_cnt; + const uint64_t xq_bytes = x_rows * blocks_a * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output low q8 prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = cuda_q8_use_dp4a(); + dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1); + quantize_q8_0_group_slice_rows_kernel<<>>( + xq, + xscale, + (const float *)heads->ptr, + group_dim, + blocks_a, + n_groups_total, + group0, + group_cnt); + if (!cuda_ok(cudaGetLastError(), + "attention_output_low_q8 rows prequant launch")) return 0; + dim3 grid_a(((unsigned)low_dim + 7u) / 8u, n_rows, 1u); + grouped_q8_0_a_preq_warp8_kernel<<>>((float *)low->ptr, + out_a, + xq, + xscale, + group_dim, + rank, + group_cnt, + n_rows, + blocks_a, + use_dp4a); + return cuda_ok(cudaGetLastError(), + "attention_output_low_q8 rows launch"); +} + +extern "C" int ds4_gpu_attention_output_low_q8_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads) { + return ds4_gpu_attention_output_low_q8_rows_exact_tensor( + low, model_map, model_size, out_a_offset, group_dim, rank, + n_groups, 0u, n_groups, heads, 1u); +} + +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) { + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups_total == 0 || + group_cnt == 0 || group0 > n_groups_total || + group_cnt > n_groups_total - group0 || out_dim == 0) { + return 0; + } + const uint64_t blocks_a = (group_dim + 31u) / 32u; + const uint64_t row_a_bytes = blocks_a * 34u; + const uint64_t low_dim_total = (uint64_t)n_groups_total * rank; + const uint64_t k_off = (uint64_t)group0 * rank; + const uint64_t k_cnt = (uint64_t)group_cnt * rank; + if ((k_off % 32u) != 0 || (k_cnt % 32u) != 0) return 0; + if (heads->bytes < (uint64_t)(group0 + group_cnt) * group_dim * sizeof(float) || + low->bytes < k_cnt * sizeof(float) || + out->bytes < out_dim * sizeof(float)) { + return 0; + } + + ds4_gpu_tensor heads_slice = *heads; + heads_slice.ptr = (char *)heads->ptr + (uint64_t)group0 * group_dim * sizeof(float); + heads_slice.bytes = (uint64_t)group_cnt * group_dim * sizeof(float); + heads_slice.owner = 0; + + const uint64_t a_off = out_a_offset + (uint64_t)group0 * rank * row_a_bytes; + return ds4_gpu_attention_output_low_q8_tensor(low, + model_map, + model_size, + a_off, + group_dim, + rank, + group_cnt, + &heads_slice) && + ds4_gpu_matmul_q8_0_kslice_rows_tensor(out, + model_map, + model_size, + out_b_offset, + low_dim_total, + out_dim, + k_off, + k_cnt, + low, + 1); +} +extern "C" int ds4_gpu_swiglu_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *gate, const ds4_gpu_tensor *up, uint32_t n, float clamp, float weight) { + if (!out || !gate || !up || + out->bytes < (uint64_t)n * sizeof(float) || + gate->bytes < (uint64_t)n * sizeof(float) || + up->bytes < (uint64_t)n * sizeof(float)) return 0; + swiglu_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)gate->ptr, (const float *)up->ptr, n, clamp, weight); + return cuda_ok(cudaGetLastError(), "swiglu launch"); +} +extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + if (getenv("DS4_CUDA_DISABLE_SHARED_GATE_UP_PAIR") == NULL) { + return ds4_gpu_matmul_q8_0_pair_tensor(gate, up, + model_map, model_size, + gate_offset, up_offset, + in_dim, out_dim, out_dim, + x, 1) && + ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f); + } + return ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size, + gate_offset, in_dim, out_dim, x, 1) && + ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size, + up_offset, in_dim, out_dim, x, 1) && + ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f); +} + +extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *prequant, + uint32_t expert_split, + bool home_rank) { + if (!mid || !x || !model_map || in_dim == 0u || out_dim == 0u || + x->bytes < in_dim * sizeof(float) || + mid->bytes < out_dim * sizeof(float) || + (selected && (selected->bytes < 6u * sizeof(int32_t) || + expert_split == 0u))) { + return 0; + } + const uint64_t blocks = (in_dim + 31u) / 32u; + if (gate_offset > model_size || up_offset > model_size || + out_dim > UINT64_MAX / (blocks * 34u)) { + return 0; + } + const uint64_t weight_bytes = out_dim * blocks * 34u; + if (weight_bytes > model_size - gate_offset || + weight_bytes > model_size - up_offset) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(x); + if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0; + if (selected && ds4_tensor_device_idx(selected) != logical_tier) return 0; + if (prequant && ds4_tensor_device_idx(prequant) != logical_tier) return 0; + const int mid_tier = ds4_tensor_device_idx(mid); + if (mid_tier != logical_tier && !g_gpu_peer_ok[logical_tier][mid_tier]) { + return 0; + } + const char *gate_w = cuda_resolve_weight_ptr( + model_map, gate_offset, weight_bytes, logical_tier, + "shared_mid_gate_exact"); + const char *up_w = cuda_resolve_weight_ptr( + model_map, up_offset, weight_bytes, logical_tier, + "shared_mid_up_exact"); + if (!gate_w || !up_w) return 0; + + const uint64_t xq_bytes = blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); + int8_t *xq; + float *xscale; + if (prequant) { + if (prequant->bytes < tmp_bytes) return 0; + xq = (int8_t *)prequant->ptr; + xscale = (float *)((char *)prequant->ptr + scale_offset); + } else { + void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, + "shared mid q8 exact prequant"); + if (!tmp) return 0; + xq = (int8_t *)tmp; + xscale = (float *)((char *)tmp + scale_offset); + quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>( + xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), + "shared mid q8 exact quantize launch")) { + return 0; + } + } + shared_mid_q8_0_preq_warp8_exact_kernel<<< + ((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)mid->ptr, + (const unsigned char *)gate_w, + (const unsigned char *)up_w, + xq, + xscale, + in_dim, + out_dim, + blocks, + clamp, + selected ? (const int32_t *)selected->ptr : NULL, + expert_split, + home_rank, + cuda_q8_use_dp4a()); + return cuda_ok(cudaGetLastError(), "shared mid q8 exact launch"); +} +extern "C" int ds4_gpu_add_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, uint32_t n) { + if (!out || !a || !b || + out->bytes < (uint64_t)n * sizeof(float) || + a->bytes < (uint64_t)n * sizeof(float) || + b->bytes < (uint64_t)n * sizeof(float)) return 0; + add_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)a->ptr, (const float *)b->ptr, n); + return cuda_ok(cudaGetLastError(), "add launch"); +} + +extern "C" 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) { + if (!out || !local || !remote || + out->bytes < (uint64_t)n * sizeof(float) || + local->bytes < (uint64_t)n * sizeof(float) || + remote->bytes < (uint64_t)n * sizeof(float)) return 0; + if (n == 0) return 1; + + const int od = ds4_tensor_device_idx(out); + const int ld = ds4_tensor_device_idx(local); + const int rd = ds4_tensor_device_idx(remote); + if (od != ld) return 0; + + const ds4_gpu_tensor *rhs = remote; + if (rd != od) { + if (!remote_tmp || + remote_tmp->bytes < (uint64_t)n * sizeof(float) || + ds4_tensor_device_idx(remote_tmp) != od) return 0; + if (!ds4_gpu_tensor_copy_xdev(remote_tmp, remote, + (uint64_t)n * sizeof(float))) return 0; + rhs = remote_tmp; + } + + int ok = 0; + WITH_DEVICE(g_gpu[od].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[od].stream; + add_kernel<<<(n + 255u) / 256u, 256, 0, s>>>( + (float *)out->ptr, + (const float *)local->ptr, + (const float *)rhs->ptr, + n); + ok = cuda_ok(cudaGetLastError(), "xdev add launch"); + cudaEvent_t e = (cudaEvent_t)g_gpu[od].boundary_event; + if (ok) ok = cuda_ok(cudaEventRecord(e, s), "xdev add event record"); + if (ok) ok = cuda_ok(cudaStreamWaitEvent(0, e, 0), "xdev add default wait"); + if (ok && g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s), "xdev add sync"); + } + } + return ok; +} +extern "C" int ds4_gpu_directional_steering_project_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_tensor *directions, + uint32_t layer, + uint32_t width, + uint32_t rows, + float scale) { + if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0; + const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float); + const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float); + if (x->bytes < x_bytes || directions->bytes < dir_bytes) return 0; + + uint32_t nth = 256u; + while (nth > width && nth > 1u) nth >>= 1; + directional_steering_project_kernel<<>>( + (float *)x->ptr, + (const float *)directions->ptr, + layer, + width, + rows, + scale); + return cuda_ok(cudaGetLastError(), "directional steering launch"); +} +extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t token, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits) { + if (!selected || !weights || !probs || !logits || !model_map || n_expert_groups > 1u || n_group_used > 0u) return 0; + if (n_expert != 256u || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0; + int32_t tok = (int32_t)token; + int ok = 1; + const float *bias = NULL; + const int32_t *hash = NULL; + const int logical_tier = ds4_tensor_device_idx(selected); + if (ok && has_bias && !hash_mode) { + if (bias_offset > model_size || model_size - bias_offset < 256u * sizeof(float)) ok = 0; + else bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, 256u * sizeof(float), logical_tier, "router_bias"); + if (!bias) ok = 0; + } + if (ok && hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t); + if (hash_offset > model_size || hash_bytes > model_size - hash_offset) ok = 0; + else hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash"); + if (!hash) ok = 0; + } + if (ok) { + if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL && + getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { + dim3 block(32, 4, 1); + router_select_warp_topk_kernel<<<1, block>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + has_bias && !hash_mode, hash_mode); + } else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { + router_select_parallel_kernel<<<1, 256>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + has_bias && !hash_mode, hash_mode); + } else { + router_select_kernel<<<1, 1>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, + bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, + has_bias && !hash_mode, hash_mode); + } + ok = cuda_ok(cudaGetLastError(), "router_select launch"); + } + return ok; +} +extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits, const ds4_gpu_tensor *tokens, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_tokens) { + if (n_expert != 256u || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0; + if (!selected || !weights || !probs || !logits || !tokens || !model_map || n_tokens == 0 || + n_expert_groups > 1u || n_group_used > 0u || + logits->bytes < (uint64_t)n_tokens * 256u * sizeof(float) || + probs->bytes < (uint64_t)n_tokens * 256u * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * 6u * sizeof(int32_t) || + weights->bytes < (uint64_t)n_tokens * 6u * sizeof(float)) { + return 0; + } + const float *bias = NULL; + const int32_t *hash = NULL; + const int logical_tier = ds4_tensor_device_idx(selected); + if (has_bias && !hash_mode) { + if (bias_offset > model_size || model_size - bias_offset < 256u * sizeof(float)) return 0; + bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, 256u * sizeof(float), logical_tier, "router_bias"); + if (!bias) return 0; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t); + if (hash_offset > model_size || hash_bytes > model_size - hash_offset) return 0; + hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash"); + if (!hash) return 0; + } + if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL && + getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { + dim3 block(32, 4, 1); + router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>((int32_t *)selected->ptr, + (float *)weights->ptr, + (float *)probs->ptr, + bias, + hash, + (const float *)logits->ptr, + (const int32_t *)tokens->ptr, + 0, + hash_rows, + n_tokens, + has_bias && !hash_mode, + hash_mode); + } else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { + router_select_parallel_kernel<<>>((int32_t *)selected->ptr, + (float *)weights->ptr, + (float *)probs->ptr, + bias, + hash, + (const float *)logits->ptr, + (const int32_t *)tokens->ptr, + 0, + hash_rows, + n_tokens, + has_bias && !hash_mode, + hash_mode); + } else { + router_select_kernel<<>>((int32_t *)selected->ptr, + (float *)weights->ptr, + (float *)probs->ptr, + bias, + hash, + (const float *)logits->ptr, + (const int32_t *)tokens->ptr, + 0, + hash_rows, + n_tokens, + has_bias && !hash_mode, + hash_mode); + } + return cuda_ok(cudaGetLastError(), "router_select launch"); +} diff --git a/cuda/runtime.inc b/cuda/runtime.inc new file mode 100644 index 0000000000..d2768e041f --- /dev/null +++ b/cuda/runtime.inc @@ -0,0 +1,3605 @@ +/* ========================================================================= + * Multi-GPU plumbing (device-aware CUDA). + * ========================================================================= */ + +static_assert(DS4_MAX_GPUS == 16, "DS4_MAX_GPUS stack tables sized for 16"); + +ds4_gpu_ctx g_gpu[DS4_MAX_GPUS]; +int g_n_gpus = 0; +int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS]; + +/* Per-pair pinned-host bounce buffers, indexed [src][dst]. Lazily grown + * to the largest copy seen for that pair. Each pair is its own allocation + * so concurrent fan-out copies from a single source GPU to multiple + * destinations cannot race for staging memory. */ +static void *g_xdev_bounce[DS4_MAX_GPUS][DS4_MAX_GPUS]; +static size_t g_xdev_bounce_bytes[DS4_MAX_GPUS][DS4_MAX_GPUS]; + +/* Internal helper: resolve a tensor's device index. -1 (untagged) is + * treated as device 0 for legacy callers. */ +static inline int ds4_tensor_device_idx(const ds4_gpu_tensor *t) { + if (!t) return 0; + int d = t->device_id; + if (d < 0) return 0; + return d; +} + +/* Debug/override flags are read once per CUDA init. The hot decode path calls + * the xdev helpers many times per token, so they must not re-enter getenv(). */ +static void cuda_xdev_env_refresh(void) { + g_xdev_sync_debug = getenv("DS4_CUDA_SYNC_XDEV") != NULL; + g_xdev_force_cuda_peer = getenv("DS4_FORCE_CUDA_PEER") != NULL; + g_xdev_force_host_bounce = getenv("DS4_FORCE_HOST_BOUNCE") != NULL; +} + +static void cuda_decode_dispatch_env_refresh(void) { + g_cuda_disable_qkv_rms_fused = getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL; + g_cuda_no_window_attention = getenv("DS4_CUDA_NO_WINDOW_ATTENTION") != NULL; + g_cuda_decode_heads8_online = getenv("DS4_CUDA_DECODE_HEADS8_ONLINE") != NULL; + g_cuda_decode_score4 = getenv("DS4_CUDA_DECODE_SCORE4") != NULL; + g_cuda_decode_score8 = getenv("DS4_CUDA_DECODE_SCORE8") != NULL; + g_cuda_no_decode_value512 = getenv("DS4_CUDA_NO_DECODE_VALUE512") != NULL; + g_cuda_no_top1 = getenv("DS4_CUDA_NO_TOP1") != NULL; + g_cuda_end_stream_sync = getenv("DS4_CUDA_END_STREAM_SYNC") != NULL; + g_cuda_no_setdevice_cache = getenv("DS4_CUDA_NO_SETDEVICE_CACHE") != NULL; + g_cuda_exact_score_split_graph = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_GRAPH") != NULL; + g_cuda_exact_score_split_ldg = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_LDG") != NULL; + g_cuda_exact_score_split_vec4 = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4") != NULL; + g_cuda_exact_score_split_vec4_plain = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4_PLAIN") != NULL; + g_cuda_exact_score_split_dim2 = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_DIM2") != NULL && + getenv("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DIM2") == NULL; + g_cuda_exact_score_split_fuse_inv_rope = + getenv("DS4_CUDA_EXACT_SCORE_SPLIT_FUSE_INV_ROPE") != NULL; + g_cuda_moe_decode_graph = getenv("DS4_CUDA_MOE_DECODE_GRAPH") != NULL; +} + +/* WITH_DEVICE(d) { ... } scope macro. + * + * Save the calling thread's current CUDA device, switch to device `d`, + * run the body exactly once, then restore the previous device. If the + * required CUDA calls fail, the body still runs (we don't have a clean + * way to early-exit a containing function from a macro), but the next + * CUDA call inside the body will surface the error naturally. + * + * Implementation: a for-loop with two synthetic variables. Iter 0 runs + * the body; on iter 1, the iteration step restores the previous device + * via cudaSetDevice and sets _wd_first = 0 so the loop exits. The + * single-statement-body restriction of for-loops is removed by the + * required `{ ... }` block in the call site. + */ +#define WITH_DEVICE(d) \ + for (int _wd_prev = -1, _wd_first = 1; \ + _wd_first; \ + _wd_first = 0, \ + (_wd_prev >= 0 ? (void)cudaSetDevice(_wd_prev) : (void)0)) \ + if (cudaGetDevice(&_wd_prev) != cudaSuccess) { /* leave */ } else \ + if (cudaSetDevice(d) != cudaSuccess) { /* leave */ } else + +/* ========================================================================= + * Per-device selective model cache (selective model cache). + * + * The public API in ds4_gpu.h declares ds4_tensor_range and the + * device_cache_tensors / lookup_cache entry points. ds4_cuda.cu does NOT + * include ds4_gpu.h historically (a pre-existing project convention), so + * we redeclare the struct here with the same layout the header uses. + * The implementation links by C linkage; struct compatibility is by + * field layout. */ +typedef struct { + uint64_t source_offset; + uint64_t bytes; + int target_device; +} ds4_tensor_range; + +struct cuda_device_cache { + void *base; /* device-side slab base */ + size_t bytes; + int present; +}; +static cuda_device_cache g_dev_cache[DS4_MAX_GPUS]; + +struct cache_range_entry { + uint64_t source_offset; + uint64_t bytes; + int device_id; + void *device_ptr; +}; +static std::vector g_cache_ranges; + +struct cuda_model_range { + const void *host_base; + uint64_t offset; + uint64_t bytes; + char *device_ptr; + void *registered_base; + char *registered_device_base; + uint64_t registered_bytes; + int host_registered; + int arena_allocated; +}; + +struct cuda_model_arena { + char *device_ptr; + uint64_t bytes; + uint64_t used; +}; + +struct cuda_q8_f16_range { + const void *host_base; + uint64_t offset; + uint64_t weight_bytes; + uint64_t in_dim; + uint64_t out_dim; + __half *device_ptr; + int device_id; /* physical CUDA device id; 0 in single-tier */ +}; + +struct cuda_q8_f32_range { + const void *host_base; + uint64_t offset; + uint64_t weight_bytes; + uint64_t in_dim; + uint64_t out_dim; + float *device_ptr; + int device_id; /* physical CUDA device id; 0 in single-tier */ +}; + +static std::vector g_model_ranges; +static std::vector g_model_arenas; +static std::unordered_map g_model_range_by_offset; +static std::vector g_q8_f16_ranges; +static std::unordered_map g_q8_f16_by_offset; +static std::vector g_q8_f32_ranges; +static std::unordered_map g_q8_f32_by_offset; +static uint64_t g_model_range_bytes; +static uint64_t g_q8_f16_bytes; +static uint64_t g_q8_f32_bytes; +static int g_q8_cache_suppressed; +static int g_q8_f16_disabled_after_oom; +static int g_q8_f16_budget_notice_printed; +static uint64_t g_model_load_progress_next; +static double g_model_load_progress_last; +static int g_model_load_progress_started; +static int g_model_load_progress_tty; +static void *g_cuda_tmp; +static uint64_t g_cuda_tmp_bytes; +static void *g_model_stage_raw[4]; +static void *g_model_stage[4]; +static cudaEvent_t g_model_stage_event[4]; +static uint64_t g_model_stage_bytes; +static void *g_stream_selected_stage_raw[4]; +static void *g_stream_selected_stage[4]; +static cudaEvent_t g_stream_selected_stage_event[4]; +static uint64_t g_stream_selected_stage_bytes; +static cudaStream_t g_stream_selected_upload_stream; + +static int cuda_ok(cudaError_t err, const char *what); +static const char *cuda_model_range_ptr_from_fd( + const void *model_map, + uint64_t offset, + uint64_t bytes, + const char *what); + +/* Forward declaration: defined later in this file. The resolver wrapper + * below uses it for multi-tier dispatch. */ +extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset, + uint64_t bytes, + int expected_device, + void **out_device_ptr); +__global__ static void dequant_q8_0_to_f16_kernel( + __half *out, + const unsigned char *w, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks); +__global__ static void dequant_q8_0_to_f32_kernel( + float *out, + const unsigned char *w, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks); + +static void *cuda_tmp_alloc(uint64_t bytes, const char *what) { + if (bytes == 0) return NULL; + if (g_cuda_tmp_bytes >= bytes) return g_cuda_tmp; + if (g_cuda_tmp) { + (void)cudaFree(g_cuda_tmp); + g_cuda_tmp = NULL; + g_cuda_tmp_bytes = 0; + } + void *ptr = NULL; + cudaError_t err = cudaMalloc(&ptr, (size_t)bytes); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA temp alloc failed for %s (%.2f MiB): %s\n", + what ? what : "scratch", (double)bytes / 1048576.0, cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + g_cuda_tmp = ptr; + g_cuda_tmp_bytes = bytes; + return g_cuda_tmp; +} + +/* Per-tier scratch accessor. + * + * Behavior: + * - At g_n_gpus <= 1 (single-tier), delegates to cuda_tmp_alloc which + * manages the legacy g_cuda_tmp slab. This guarantees byte-identical + * behavior to pre-task code for the gpu_cfg == NULL case. + * - For multi-tier, grows the per-device g_gpu[logical_tier].scratch + * slab on the corresponding physical device. Cleanup is already + * handled by ds4_gpu_cleanup (which walks g_gpu[i].scratch). + * + * The legacy g_cuda_tmp slab is untouched: init-time / preload callers + * still use cuda_tmp_alloc directly. No aliasing between g_cuda_tmp + * and g_gpu[0].scratch — they are independently owned and freed. + * + * Added for multi-GPU execution (multi-GPU execution), step A3 of the + * spec (sub-area 2). */ +static void *cuda_tmp_alloc_on(int logical_tier, uint64_t bytes, const char *what) { + if (bytes == 0) return NULL; + if (g_n_gpus <= 1) { + return cuda_tmp_alloc(bytes, what); + } + if (logical_tier < 0 || logical_tier >= g_n_gpus) { + fprintf(stderr, "ds4: cuda_tmp_alloc_on: bad tier %d (n_gpus=%d, what=%s)\n", + logical_tier, g_n_gpus, what ? what : "?"); + return NULL; + } + ds4_gpu_ctx *ctx = &g_gpu[logical_tier]; + if (ctx->scratch_bytes >= bytes) return ctx->scratch; + int prev = -1; + cudaError_t derr = cudaGetDevice(&prev); + if (derr != cudaSuccess) { + fprintf(stderr, + "ds4: cudaGetDevice failed before scratch alloc on tier %d (dev=%d, what=%s): %s\n", + logical_tier, ctx->device_id, what ? what : "scratch", + cudaGetErrorString(derr)); + (void)cudaGetLastError(); + return NULL; + } + derr = cudaSetDevice(ctx->device_id); + if (derr != cudaSuccess) { + fprintf(stderr, + "ds4: cudaSetDevice(%d) failed before scratch alloc on tier %d (what=%s): %s\n", + ctx->device_id, logical_tier, what ? what : "scratch", + cudaGetErrorString(derr)); + (void)cudaGetLastError(); + if (prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + if (ctx->scratch) { + (void)cudaFree(ctx->scratch); + ctx->scratch = NULL; + ctx->scratch_bytes = 0; + } + void *p = NULL; + cudaError_t err = cudaMalloc(&p, (size_t)bytes); + if (prev >= 0) (void)cudaSetDevice(prev); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA scratch alloc on tier %d (dev=%d) failed for %s (%.2f MiB): %s\n", + logical_tier, ctx->device_id, what ? what : "scratch", + (double)bytes / 1048576.0, cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + ctx->scratch = p; + ctx->scratch_bytes = (size_t)bytes; + return p; +} + +static int cuda_attention_score_buffer_fits(uint32_t n_comp) { + return n_comp <= DS4_CUDA_ATTENTION_SCORE_CAP - DS4_CUDA_ATTENTION_RAW_SCORE_CAP; +} + +static const char *cuda_model_ptr(const void *model_map, uint64_t offset) { + if (model_map == g_model_host_base && g_model_device_base) return g_model_device_base + offset; + return (const char *)model_map + offset; +} + +static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset, uint64_t bytes, const char *what) { + if (bytes == 0) return cuda_model_ptr(model_map, offset); + if (g_model_device_owned || g_model_registered) return cuda_model_ptr(model_map, offset); + if (g_model_hmm_direct && + getenv("DS4_CUDA_WEIGHT_CACHE") == NULL && + getenv("DS4_CUDA_WEIGHT_PRELOAD") == NULL) { + return cuda_model_ptr(model_map, offset); + } + const char *direct_env = getenv("DS4_CUDA_DIRECT_MODEL"); + if (direct_env && direct_env[0]) return cuda_model_ptr(model_map, offset); + + const uint64_t end = offset + bytes; + auto exact = g_model_range_by_offset.find(offset); + if (exact != g_model_range_by_offset.end()) { + const cuda_model_range &r = g_model_ranges[exact->second]; + if (r.host_base == model_map && end >= offset && bytes <= r.bytes) return r.device_ptr; + } + for (const cuda_model_range &r : g_model_ranges) { + if (r.host_base == model_map && offset >= r.offset && end >= offset && end <= r.offset + r.bytes) { + return r.device_ptr + (offset - r.offset); + } + if (r.host_base == model_map && r.host_registered && r.registered_base && r.registered_device_base) { + const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); + const uintptr_t h1 = h0 + bytes; + const uintptr_t r0 = (uintptr_t)r.registered_base; + const uintptr_t r1 = r0 + r.registered_bytes; + if (h1 >= h0 && h0 >= r0 && h1 <= r1) return r.registered_device_base + (h0 - r0); + } + } + + if (getenv("DS4_CUDA_NO_FD_CACHE") == NULL) { + const char *fd_ptr = cuda_model_range_ptr_from_fd(model_map, offset, bytes, what); + if (fd_ptr) return fd_ptr; + } + + cudaError_t err = cudaSuccess; + if (g_model_range_mapping_supported) { + const long page_sz_l = sysconf(_SC_PAGESIZE); + const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; + const uintptr_t host_addr = (uintptr_t)((const char *)model_map + offset); + const uintptr_t reg_addr = host_addr & ~(uintptr_t)(page_sz - 1u); + const uint64_t reg_delta = (uint64_t)(host_addr - reg_addr); + const uint64_t reg_bytes = (reg_delta + bytes + page_sz - 1u) & ~(page_sz - 1u); + void *reg_dev = NULL; + err = cudaHostRegister((void *)reg_addr, + (size_t)reg_bytes, + cudaHostRegisterMapped | cudaHostRegisterReadOnly); + if (err == cudaSuccess) { + err = cudaHostGetDevicePointer(®_dev, (void *)reg_addr, 0); + if (err == cudaSuccess && reg_dev) { + char *dev_ptr = (char *)reg_dev + reg_delta; + g_model_ranges.push_back({model_map, offset, bytes, dev_ptr, (void *)reg_addr, (char *)reg_dev, reg_bytes, 1, 0}); + g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA mapped %s %.2f MiB\n", + what ? what : "weights", + (double)bytes / 1048576.0); + } + return dev_ptr; + } + fprintf(stderr, "ds4: CUDA model range map pointer failed for %s: %s\n", + what ? what : "weights", cudaGetErrorString(err)); + (void)cudaHostUnregister((void *)reg_addr); + (void)cudaGetLastError(); + } else { + if (err == cudaErrorNotSupported || err == cudaErrorInvalidValue) g_model_range_mapping_supported = 0; + (void)cudaGetLastError(); + } + } + + void *dev = NULL; + err = cudaMalloc(&dev, (size_t)bytes); + if (err != cudaSuccess) { + (void)cudaGetLastError(); + fprintf(stderr, "ds4: CUDA model range alloc failed for %s (%.2f MiB): %s\n", + what ? what : "weights", (double)bytes / 1048576.0, cudaGetErrorString(err)); + return NULL; + } + + const char *src = (const char *)model_map + offset; + const uint64_t chunk = 64ull * 1024ull * 1024ull; + for (uint64_t done = 0; done < bytes; done += chunk) { + uint64_t n = bytes - done < chunk ? bytes - done : chunk; + err = cudaMemcpy((char *)dev + done, src + done, (size_t)n, cudaMemcpyHostToDevice); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f/%.2f MiB: %s\n", + what ? what : "weights", + (double)done / 1048576.0, + (double)bytes / 1048576.0, + cudaGetErrorString(err)); + (void)cudaFree(dev); + (void)cudaGetLastError(); + return NULL; + } + } + g_model_ranges.push_back({model_map, offset, bytes, (char *)dev, NULL, NULL, 0, 0, 0}); + g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; + g_model_range_bytes += bytes; + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA cached %s %.2f MiB (total %.2f GiB)\n", + what ? what : "weights", + (double)bytes / 1048576.0, + (double)g_model_range_bytes / 1073741824.0); + } + return (const char *)dev; +} + +/* Per-tier cuBLAS handle. Used by kernel-dispatch wrappers; returns the + * cuBLAS handle for the logical tier. The wrapper is expected to have + * cudaSetDevice'd to that tier's physical device already (kernels and + * cuBLAS calls ride the default stream and are naturally serialized). + * + * Added for multi-GPU execution (multi-GPU execution), sub-area 1. */ +static inline cublasHandle_t cuda_cublas_for_tier(int logical_tier) { + if (g_n_gpus <= 1) { + return (cublasHandle_t)g_gpu[0].cublas; + } + /* The executing device is authoritative: GLM per-layer switching runs + * generic launchers whose out tensors live on device 0 while the layer + * executes elsewhere. On DS4 paths the current device always equals the + * requested tier, so this is behavior-preserving there. */ + int cur_dev = -1; + if (cudaGetDevice(&cur_dev) == cudaSuccess) { + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == cur_dev) { + return (cublasHandle_t)g_gpu[t].cublas; + } + } + } + if (logical_tier < 0 || logical_tier >= g_n_gpus) { + return (cublasHandle_t)g_gpu[0].cublas; + } + return (cublasHandle_t)g_gpu[logical_tier].cublas; +} + +/* Multi-tier-aware weight pointer resolver. + * + * Used by kernel-dispatch wrappers in the per-layer execution path to + * obtain a device pointer for a weight slice on the layer's logical + * tier. Behavior: + * + * - When g_n_gpus <= 1 (single-tier engine), delegates to the existing + * cuda_model_range_ptr path. This short-circuit guarantees byte- + * identical behavior to pre-multi-tier code for the gpu_cfg == NULL + * case. + * + * - When g_n_gpus >= 2 (multi-tier engine), translates the logical + * tier index to the corresponding physical CUDA device id via + * g_gpu[logical_tier].device_id and looks up the slice strictly + * in the per-device selective cache via ds4_gpu_lookup_cache_strict. + * On miss, logs a diagnostic and returns NULL (no host-pointer + * fallback — a miss here is a placement/install bug). + * + * The caller is responsible for cudaSetDevice'ing to the right physical + * device before launching the kernel that consumes the returned pointer. + * Wrappers in this file thread `int logical_tier` from the dispatch + * caller; single-tier callers pass 0, which hits the short-circuit. + * + * Added for multi-GPU execution (multi-GPU execution), sub-area 3 of the + * spec. */ +/* Optional second (support) model map for speculative decoding. The strict + * multi-tier cache is keyed by source offset only, so support tensors are + * installed and resolved with a large disjoint offset bias. */ +static const void *g_support_host_base = NULL; +static uint64_t g_support_host_size = 0; +static uint64_t g_support_offset_bias = 0; + +extern "C" uint64_t ds4_gpu_tier_free_vram(int logical_tier) { + if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0; + int prev = -1; + if (cudaGetDevice(&prev) != cudaSuccess) prev = -1; + if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) return 0; + size_t free_b = 0, total_b = 0; + uint64_t out = 0; + if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) out = (uint64_t)free_b; + if (prev >= 0) (void)cudaSetDevice(prev); + return out; +} + +extern "C" int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias) { + if (!map || size == 0 || bias == 0) return 0; + g_support_host_base = map; + g_support_host_size = size; + g_support_offset_bias = bias; + return 1; +} + +static const char *cuda_resolve_weight_ptr(const void *model_map, + uint64_t offset, + uint64_t bytes, + int logical_tier, + const char *label) { + if (g_n_gpus <= 1) { + return cuda_model_range_ptr(model_map, offset, bytes, label); + } + if (g_support_host_base && model_map == g_support_host_base) { + offset += g_support_offset_bias; + } + if (logical_tier < 0 || logical_tier >= g_n_gpus) { + fprintf(stderr, + "ds4: cuda_resolve_weight_ptr: bad tier %d (n_gpus=%d, label=%s)\n", + logical_tier, g_n_gpus, label ? label : "?"); + return NULL; + } + const int physical_device = g_gpu[logical_tier].device_id; + void *dev_ptr = NULL; + if (ds4_gpu_lookup_cache_strict(offset, bytes, physical_device, &dev_ptr) + && dev_ptr) { + return (const char *)dev_ptr; + } + /* GLM multi-tier: generic launchers resolve by the OUT tensor's tier, + * but the executing device (set per layer) is where the weights were + * cached. Retry with the current device before declaring a miss; + * DS4 paths never reach this (out tier == current device). */ + int cur_dev = -1; + if (cudaGetDevice(&cur_dev) == cudaSuccess && + cur_dev != physical_device && + ds4_gpu_lookup_cache_strict(offset, bytes, cur_dev, &dev_ptr) && + dev_ptr) { + return (const char *)dev_ptr; + } + fprintf(stderr, + "ds4: selective-cache miss for offset=%llu bytes=%llu on " + "logical_tier=%d (physical_device=%d, current_device=%d, " + "label=%s); this is a placement/cache-install bug\n", + (unsigned long long)offset, (unsigned long long)bytes, + logical_tier, physical_device, cur_dev, label ? label : "?"); + return NULL; +} + +static int cuda_model_range_is_cached(const void *model_map, uint64_t offset, uint64_t bytes) { + if (bytes == 0) return 1; + if (g_model_device_owned || g_model_registered) return 1; + + const uint64_t end = offset + bytes; + if (end < offset) return 0; + for (const cuda_model_range &r : g_model_ranges) { + if (r.host_base == model_map && + offset >= r.offset && + end <= r.offset + r.bytes) { + return 1; + } + if (r.host_base == model_map && + r.host_registered && + r.registered_base && + r.registered_device_base) { + const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); + const uintptr_t h1 = h0 + bytes; + const uintptr_t r0 = (uintptr_t)r.registered_base; + const uintptr_t r1 = r0 + r.registered_bytes; + if (h1 >= h0 && h0 >= r0 && h1 <= r1) return 1; + } + } + return 0; +} + +static void cuda_q8_f16_cache_release_all(void) { + for (const cuda_q8_f16_range &r : g_q8_f16_ranges) { + (void)cudaFree(r.device_ptr); + } + g_q8_f16_ranges.clear(); + g_q8_f16_by_offset.clear(); + g_q8_f16_bytes = 0; +} + +static uint64_t cuda_parse_mib_env(const char *name, int *present) { + const char *env = getenv(name); + if (present) *present = 0; + if (!env || !env[0]) return 0; + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end == env || *end != '\0') return 0; + if (present) *present = 1; + if (v > UINT64_MAX / 1048576ull) return UINT64_MAX; + return (uint64_t)v * 1048576ull; +} + +static uint32_t cuda_parse_u32_env_clamped(const char *name, uint32_t fallback, + uint32_t min_value, uint32_t max_value, + int *present) { + const char *env = getenv(name); + if (present) *present = 0; + if (!env || !env[0]) return fallback; + errno = 0; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (errno != 0 || end == env || *end != '\0') return fallback; + if (present) *present = 1; + if (v < min_value) return min_value; + if (v > max_value) return max_value; + return (uint32_t)v; +} + +static int cuda_env_flag_enabled(const char *name, int fallback) { + const char *env = getenv(name); + if (!env || !env[0]) return fallback; + return strcmp(env, "0") != 0; +} + +extern "C" int ds4_gpu_set_decode_fast_attention(int enabled) { + const int old = g_decode_fast_attention; + g_decode_fast_attention = enabled != 0; + return old; +} + +extern "C" int ds4_gpu_set_decode_score_vec4(int enabled) { + const int old = g_decode_score_vec4; + g_decode_score_vec4 = enabled != 0; + return old; +} + +static bool cuda_splitkv_decode_requested(void) { + if (cuda_env_flag_enabled("DS4_CUDA_NO_SPLITKV_DECODE", 0)) return false; + return g_decode_fast_attention || + cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0); +} + +static uint64_t cuda_q8_f16_cache_limit_bytes(void) { + int present = 0; + const uint64_t limit = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_MB", &present); + return present ? limit : UINT64_MAX; +} + +static uint64_t cuda_q8_f16_cache_reserve_bytes(uint64_t total_bytes) { + int present = 0; + const uint64_t reserve = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", &present); + if (present) return reserve; + + if (total_bytes >= 112ull * 1024ull * 1024ull * 1024ull) { + return 512ull * 1048576ull; + } + + /* High-VRAM cards (>= 40 GiB, e.g. 48 GiB RTX 6000 Ada): use a small + * reserve so the selective Q8->F16 cache can actually engage at tight + * budgets (e.g. --gpu-vram 47,47, where the 81 GB model leaves only ~1.3 + * GiB free and the old 4 GiB floor rejected every cache allocation, + * forcing the scalar DP4A prefill kernel). + * + * NOTE: this 768 MiB value is a *bounded cache-growth guard*, not a hard + * guarantee that live free VRAM stays >= 768 MiB. cuda_q8_f16_cache_has_budget + * only blocks a *cache* allocation when free - request < reserve at that + * moment; allocations made outside cache accounting (cuda_tmp_alloc_on + * activation/prequant buffers, cuBLAS internal workspaces) can still dip + * below it. 768 MiB is chosen to leave headroom above the ~0.5 GiB + * memory-safety floor for those out-of-cache allocations; actual minimum + * free VRAM is verified by measurement, and the disable-after-failure path + * degrades gracefully if cuBLAS/alloc ever fails under pressure. Set + * DS4_CUDA_Q8_F16_CACHE_RESERVE_MB=4096 to restore the prior behavior. */ + if (total_bytes >= 40ull * 1024ull * 1024ull * 1024ull) { + const uint64_t hi_min_reserve = 768ull * 1048576ull; + const uint64_t hi_pct_reserve = total_bytes / 100u; /* 1% */ + return hi_pct_reserve > hi_min_reserve ? hi_pct_reserve : hi_min_reserve; + } + + /* Smaller cards (< 40 GiB): keep the conservative reserve. The expanded + * Q8->F16 cache is only an acceleration path; on a small card a sub-GiB + * reserve would be a large fraction of total VRAM, so keep enough free for + * cuBLAS workspaces, transient graph buffers, and driver bookkeeping. */ + const uint64_t min_reserve = 4096ull * 1048576ull; + const uint64_t pct_reserve = total_bytes / 20u; /* 5% */ + return pct_reserve > min_reserve ? pct_reserve : min_reserve; +} + +static void cuda_q8_f16_cache_budget_notice( + const char *reason, + uint64_t request_bytes, + uint64_t free_bytes, + uint64_t total_bytes, + uint64_t reserve_bytes, + uint64_t limit_bytes) { + if (g_q8_f16_budget_notice_printed && getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") == NULL) return; + g_q8_f16_budget_notice_printed = 1; + if (limit_bytes != UINT64_MAX && free_bytes == 0 && total_bytes == 0 && reserve_bytes == 0) { + fprintf(stderr, + "ds4: CUDA q8 fp16 cache %s; using q8 kernels " + "(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB)\n", + reason, + (double)request_bytes / 1048576.0, + (double)g_q8_f16_bytes / 1073741824.0, + (double)limit_bytes / 1073741824.0); + } else if (limit_bytes == UINT64_MAX) { + fprintf(stderr, + "ds4: CUDA q8 fp16 cache %s; using q8 kernels " + "(request=%.2f MiB cached=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n", + reason, + (double)request_bytes / 1048576.0, + (double)g_q8_f16_bytes / 1073741824.0, + (double)free_bytes / 1073741824.0, + (double)reserve_bytes / 1073741824.0, + (double)total_bytes / 1073741824.0); + } else { + fprintf(stderr, + "ds4: CUDA q8 fp16 cache %s; using q8 kernels " + "(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n", + reason, + (double)request_bytes / 1048576.0, + (double)g_q8_f16_bytes / 1073741824.0, + (double)limit_bytes / 1073741824.0, + (double)free_bytes / 1073741824.0, + (double)reserve_bytes / 1073741824.0, + (double)total_bytes / 1073741824.0); + } +} + +static int cuda_q8_f16_cache_has_budget(uint64_t request_bytes, const char *label) { + (void)label; + const uint64_t limit = cuda_q8_f16_cache_limit_bytes(); + if (limit == 0) return 0; + if (g_q8_f16_bytes > limit || request_bytes > limit - g_q8_f16_bytes) { + cuda_q8_f16_cache_budget_notice("limit reached", request_bytes, 0, 0, 0, limit); + return 0; + } + + size_t free_b = 0; + size_t total_b = 0; + cudaError_t err = cudaMemGetInfo(&free_b, &total_b); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA q8 fp16 cache memory query failed: %s; using q8 kernels\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + + const uint64_t free_bytes = (uint64_t)free_b; + const uint64_t total_bytes = (uint64_t)total_b; + const uint64_t reserve_bytes = cuda_q8_f16_cache_reserve_bytes(total_bytes); + if (request_bytes > free_bytes || + free_bytes - request_bytes < reserve_bytes) { + cuda_q8_f16_cache_budget_notice("budget exhausted", request_bytes, + free_bytes, total_bytes, + reserve_bytes, limit); + return 0; + } + return 1; +} + +static void cuda_q8_f16_cache_disable_after_failure(const char *what, uint64_t request_bytes) { + if (!g_q8_f16_disabled_after_oom) { + fprintf(stderr, + "ds4: CUDA q8 fp16 cache disabled after %s " + "(request=%.2f MiB cached=%.2f GiB); using q8 kernels\n", + what ? what : "allocation failure", + (double)request_bytes / 1048576.0, + (double)g_q8_f16_bytes / 1073741824.0); + } + g_q8_f16_disabled_after_oom = 1; + if (!g_q8_f16_ranges.empty()) { + (void)cudaDeviceSynchronize(); + cuda_q8_f16_cache_release_all(); + } + (void)cudaGetLastError(); +} + +static int cuda_q8_f16_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { + if (g_quality_mode) return 0; + if (g_q8_cache_suppressed) return 0; + if (g_q8_f16_disabled_after_oom) return 0; + if (getenv("DS4_CUDA_NO_Q8_F16_CACHE") != NULL) return 0; + if (cuda_q8_f16_cache_limit_bytes() == 0) return 0; + if (getenv("DS4_CUDA_Q8_F16_ALL") != NULL) return 1; + if (!label) return 0; + if (strstr(label, "attn_output_a") != NULL || + strstr(label, "attn_output_b") != NULL || + strstr(label, "attention_output_a") != NULL || + strstr(label, "attention_output_b") != NULL) { + return getenv("DS4_CUDA_NO_ATTENTION_OUTPUT_F16_CACHE") == NULL; + } + if (strstr(label, "attn_q_b") != NULL) { + return getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL; + } + if (strstr(label, "ffn_gate_shexp") != NULL || + strstr(label, "ffn_up_shexp") != NULL || + strstr(label, "ffn_down_shexp") != NULL) { + return 1; + } + return (in_dim == 4096u && out_dim == 2048u) || + (in_dim == 2048u && out_dim == 4096u) || + (in_dim == 4096u && out_dim == 1024u) || + (in_dim == 4096u && out_dim == 512u) || + (getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL && + in_dim == 1024u && out_dim == 32768u); +} + +static int cuda_q8_label_is_attention_output(const char *label) { + return label && + (strstr(label, "attn_output_a") != NULL || + strstr(label, "attn_output_b") != NULL || + strstr(label, "attention_output_a") != NULL || + strstr(label, "attention_output_b") != NULL); +} + +static int cuda_q8_use_dp4a(void) { + return getenv("DS4_CUDA_NO_Q8_DP4A") == NULL; +} + +static unsigned cuda_q8_exact_threads(uint64_t blocks) { + if (blocks <= 64u) return 64u; + if (blocks <= 128u) return 128u; + return 256u; +} + +static int cuda_q8_f16_preload_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { + if (cuda_q8_label_is_attention_output(label) && + getenv("DS4_CUDA_ATTENTION_OUTPUT_PRELOAD") == NULL && + getenv("DS4_CUDA_Q8_F16_ALL") == NULL) { + return 0; + } + return cuda_q8_f16_cache_allowed(label, in_dim, out_dim); +} + +static int cuda_q8_f32_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { + if (g_q8_cache_suppressed) return 0; + if (getenv("DS4_CUDA_NO_Q8_F32_CACHE") != NULL) return 0; + if (getenv("DS4_CUDA_Q8_F32_ALL") != NULL) return 1; + if (label && strstr(label, "attn_q_b") != NULL) { + return getenv("DS4_CUDA_ATTN_Q_B_F32_CACHE") != NULL; + } + return getenv("DS4_CUDA_Q8_F32_LARGE") != NULL && + in_dim == 1024u && out_dim == 32768u; +} + +/* Look up a per-device dequantized fp16 slice of the Q8_0 weight at + * (model_map, offset, weight_bytes, in_dim, out_dim). expected_device is a + * PHYSICAL CUDA device id (0 in single-tier; g_gpu[logical_tier].device_id in + * multi-tier). On hit returns the cached pointer for that device. On miss + * cudaSetDevice's to expected_device, allocates + dequants there, stamps the + * new entry with device_id == expected_device, restores the previous device, + * and returns the new pointer. + * + * Single-tier (g_n_gpus <= 1) uses the offset-keyed map for a fast path — + * legacy entries were stamped device_id=0, so the map remains authoritative. + * Multi-tier linear-scans the ranges vector filtering on device_id (the same + * offset may now legitimately map to multiple entries, one per device). + */ +static const __half *cuda_q8_f16_ptr( + const void *model_map, + uint64_t offset, + uint64_t weight_bytes, + uint64_t in_dim, + uint64_t out_dim, + int expected_device, + const char *label) { + if (g_n_gpus <= 1) { + auto exact = g_q8_f16_by_offset.find(offset); + if (exact != g_q8_f16_by_offset.end()) { + const cuda_q8_f16_range &r = g_q8_f16_ranges[exact->second]; + if (r.host_base == model_map && r.weight_bytes == weight_bytes && + r.in_dim == in_dim && r.out_dim == out_dim) { + return r.device_ptr; + } + } + } else { + for (const cuda_q8_f16_range &r : g_q8_f16_ranges) { + if (r.host_base == model_map && + r.offset == offset && + r.weight_bytes == weight_bytes && + r.in_dim == in_dim && + r.out_dim == out_dim && + r.device_id == expected_device) { + return r.device_ptr; + } + } + } + if (!cuda_q8_f16_cache_allowed(label, in_dim, out_dim)) return NULL; + + /* Source Q8 bytes: + * - Single-tier (g_n_gpus <= 1): cuda_model_range_ptr — preserves the + * legacy behavior (FD cache, host-register, or cudaMalloc-and-copy). + * - Multi-tier (g_n_gpus > 1): the per-device selective cache must + * already contain the weight on expected_device. Use the strict + * lookup; on miss this is a placement bug and we hard-fail. + */ + const char *q8; + if (g_n_gpus <= 1) { + q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, "q8_0"); + } else { + void *strict_ptr = NULL; + if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) || + !strict_ptr) { + fprintf(stderr, + "ds4: q8 fp16 cache miss: source bytes not in selective cache for " + "offset=%llu bytes=%llu device=%d (label=%s); placement bug\n", + (unsigned long long)offset, (unsigned long long)weight_bytes, + expected_device, label ? label : "?"); + return NULL; + } + q8 = (const char *)strict_ptr; + } + if (!q8) return NULL; + + if (in_dim != 0 && out_dim > UINT64_MAX / in_dim / sizeof(__half)) return NULL; + const uint64_t out_bytes = in_dim * out_dim * sizeof(__half); + if (!cuda_q8_f16_cache_has_budget(out_bytes, label)) return NULL; + + int prev = -1; + if (g_n_gpus > 1) { + cudaError_t derr = cudaGetDevice(&prev); + if (derr != cudaSuccess) { + fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp16 alloc on device %d: %s\n", + expected_device, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + return NULL; + } + derr = cudaSetDevice(expected_device); + if (derr != cudaSuccess) { + fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp16 alloc: %s\n", + expected_device, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + if (prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + } + __half *dev = NULL; + cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA q8 fp16 cache alloc failed on device %d (%.2f MiB): %s\n", + expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err)); + cuda_q8_f16_cache_disable_after_failure("allocation failure", out_bytes); + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + const uint64_t blocks = (in_dim + 31) / 32; + const uint64_t n = in_dim * out_dim; + dequant_q8_0_to_f16_kernel<<<(n + 255) / 256, 256>>>(dev, + (const unsigned char *)q8, + in_dim, + out_dim, + blocks); + if (!cuda_ok(cudaGetLastError(), "q8 fp16 dequant launch")) { + (void)cudaFree(dev); + cuda_q8_f16_cache_disable_after_failure("dequant launch failure", out_bytes); + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + g_q8_f16_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device}); + if (g_n_gpus <= 1) { + g_q8_f16_by_offset[offset] = g_q8_f16_ranges.size() - 1u; + } + g_q8_f16_bytes += out_bytes; + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA cached q8 fp16 %.2f MiB on device %d (total %.2f GiB)\n", + (double)out_bytes / 1048576.0, expected_device, + (double)g_q8_f16_bytes / 1073741824.0); + } + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return dev; +} + +/* Per-device dequantized fp32 cache. Same conventions as cuda_q8_f16_ptr. */ +static float *cuda_q8_f32_ptr( + const void *model_map, + uint64_t offset, + uint64_t weight_bytes, + uint64_t in_dim, + uint64_t out_dim, + int expected_device, + const char *label) { + if (g_n_gpus <= 1) { + auto exact = g_q8_f32_by_offset.find(offset); + if (exact != g_q8_f32_by_offset.end()) { + const cuda_q8_f32_range &r = g_q8_f32_ranges[exact->second]; + if (r.host_base == model_map && r.weight_bytes == weight_bytes && + r.in_dim == in_dim && r.out_dim == out_dim) { + return r.device_ptr; + } + } + } else { + for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { + if (r.host_base == model_map && + r.offset == offset && + r.weight_bytes == weight_bytes && + r.in_dim == in_dim && + r.out_dim == out_dim && + r.device_id == expected_device) { + return r.device_ptr; + } + } + } + if (!cuda_q8_f32_cache_allowed(label, in_dim, out_dim)) return NULL; + + /* Source Q8 bytes: legacy path in single-tier; strict per-device lookup + * in multi-tier (same rationale as cuda_q8_f16_ptr). */ + const char *q8; + if (g_n_gpus <= 1) { + q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, label ? label : "q8_0"); + } else { + void *strict_ptr = NULL; + if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) || + !strict_ptr) { + fprintf(stderr, + "ds4: q8 fp32 cache miss: source bytes not in selective cache for " + "offset=%llu bytes=%llu device=%d (label=%s); placement bug\n", + (unsigned long long)offset, (unsigned long long)weight_bytes, + expected_device, label ? label : "?"); + return NULL; + } + q8 = (const char *)strict_ptr; + } + if (!q8) return NULL; + + const uint64_t out_bytes = in_dim * out_dim * sizeof(float); + int prev = -1; + if (g_n_gpus > 1) { + cudaError_t derr = cudaGetDevice(&prev); + if (derr != cudaSuccess) { + fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp32 alloc on device %d: %s\n", + expected_device, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + return NULL; + } + derr = cudaSetDevice(expected_device); + if (derr != cudaSuccess) { + fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp32 alloc: %s\n", + expected_device, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + if (prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + } + float *dev = NULL; + cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA q8 fp32 cache alloc failed on device %d (%.2f MiB): %s\n", + expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err)); + (void)cudaGetLastError(); + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + const uint64_t blocks = (in_dim + 31) / 32; + const uint64_t n = in_dim * out_dim; + dequant_q8_0_to_f32_kernel<<<(n + 255) / 256, 256>>>(dev, + (const unsigned char *)q8, + in_dim, + out_dim, + blocks); + if (!cuda_ok(cudaGetLastError(), "q8 fp32 dequant launch")) { + (void)cudaFree(dev); + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return NULL; + } + g_q8_f32_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device}); + if (g_n_gpus <= 1) { + g_q8_f32_by_offset[offset] = g_q8_f32_ranges.size() - 1u; + } + g_q8_f32_bytes += out_bytes; + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA cached q8 fp32 %.2f MiB on device %d (total %.2f GiB)\n", + (double)out_bytes / 1048576.0, expected_device, + (double)g_q8_f32_bytes / 1073741824.0); + } + if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); + return dev; +} + +static int cuda_ok(cudaError_t err, const char *what) { + if (err == cudaSuccess) return 1; + fprintf(stderr, "ds4: CUDA %s failed: %s\n", what, cudaGetErrorString(err)); + return 0; +} + +static double cuda_wall_sec(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec + (double)ts.tv_nsec * 1.0e-9; +} + +static int cuda_model_load_progress_enabled(void) { + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") != NULL) return 0; + return 1; +} + +static void cuda_model_load_progress_reset(void) { + g_model_load_progress_next = 0; + g_model_load_progress_last = 0.0; + g_model_load_progress_started = 0; + g_model_load_progress_tty = 0; +} + +static void cuda_model_load_progress_note(uint64_t cached_bytes) { + if (!cuda_model_load_progress_enabled()) return; + + const double now = cuda_wall_sec(); + if (!g_model_load_progress_started) { + g_model_load_progress_started = 1; + g_model_load_progress_tty = isatty(STDERR_FILENO) != 0; + g_model_load_progress_next = (g_model_load_progress_tty ? 2ull : 16ull) * + 1024ull * 1024ull * 1024ull; + g_model_load_progress_last = now; + if (g_model_load_progress_tty) { + fprintf(stderr, "ds4: CUDA loading model tensors into device cache: 0.00 GiB"); + } else { + fprintf(stderr, "ds4: CUDA loading model tensors into device cache\n"); + } + } + + if (cached_bytes < g_model_load_progress_next && + now - g_model_load_progress_last < (g_model_load_progress_tty ? 2.0 : 10.0)) { + return; + } + + if (g_model_load_progress_tty) { + fprintf(stderr, "\rds4: CUDA loading model tensors into device cache: %.2f GiB", + (double)cached_bytes / 1073741824.0); + } else { + fprintf(stderr, "ds4: CUDA loading model tensors %.2f GiB cached\n", + (double)cached_bytes / 1073741824.0); + } + fflush(stderr); + g_model_load_progress_last = now; + const uint64_t step = (g_model_load_progress_tty ? 2ull : 16ull) * + 1024ull * 1024ull * 1024ull; + while (g_model_load_progress_next <= cached_bytes) { + g_model_load_progress_next += step; + } +} + +static int cuda_model_prefetch_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) { + if (!model_map || map_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0; + if (getenv("DS4_CUDA_NO_MODEL_PREFETCH") != NULL || + getenv("DS4_CUDA_COPY_MODEL") != NULL || + getenv("DS4_CUDA_WEIGHT_CACHE") != NULL || + getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) { + return 0; + } + + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + + int pageable = 0; + cudaError_t err = cudaDeviceGetAttribute(&pageable, cudaDevAttrPageableMemoryAccess, device); + if (err != cudaSuccess || !pageable) { + (void)cudaGetLastError(); + return 0; + } +#if CUDART_VERSION >= 13000 + cudaMemLocation loc; + memset(&loc, 0, sizeof(loc)); + loc.type = cudaMemLocationTypeDevice; + loc.id = device; +#else + int loc = device; +#endif + + const long page_sz_l = sysconf(_SC_PAGESIZE); + const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; + const uintptr_t host_addr = (uintptr_t)((const char *)model_map + map_offset); + const uintptr_t pre_addr = host_addr & ~(uintptr_t)(page_sz - 1u); + const uint64_t pre_delta = (uint64_t)(host_addr - pre_addr); + const uint64_t pre_bytes = (pre_delta + map_size + page_sz - 1u) & ~(page_sz - 1u); + void *pre_ptr = (void *)pre_addr; + + const double t0 = cuda_wall_sec(); + err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetReadMostly, loc); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model read-mostly advise skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetPreferredLocation, loc); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model preferred-location advise skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + + if (!g_model_prefetch_stream) { + err = cudaStreamCreateWithFlags(&g_model_prefetch_stream, cudaStreamNonBlocking); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model prefetch stream creation skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + } + +#if CUDART_VERSION >= 13000 + err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, 0, g_model_prefetch_stream); +#else + err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, g_model_prefetch_stream); +#endif + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model prefetch skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + if (getenv("DS4_CUDA_MODEL_PREFETCH_SYNC") != NULL) { + err = cudaStreamSynchronize(g_model_prefetch_stream); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model prefetch sync failed: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + } + const double t1 = cuda_wall_sec(); + fprintf(stderr, + "ds4: CUDA ATS/HMM prefetch queued %.2f GiB of model tensors in %.3fs\n", + (double)map_size / 1073741824.0, + t1 - t0); + g_model_hmm_direct = 1; + return 1; +} + +static uint64_t cuda_model_copy_chunk_bytes(void) { + uint64_t mb = 64; + const char *env = getenv("DS4_CUDA_MODEL_COPY_CHUNK_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env && v > 0) mb = (uint64_t)v; + } + if (mb < 16) mb = 16; + if (mb > 4096) mb = 4096; + return mb * 1048576ull; +} + +static void cuda_model_discard_source_pages(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes) { +#if defined(POSIX_MADV_DONTNEED) + if (getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || !model_map || bytes == 0 || offset > model_size) return; + if (bytes > model_size - offset) bytes = model_size - offset; + const long page_sz_l = sysconf(_SC_PAGESIZE); + const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; + const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); + const uintptr_t h1 = h0 + bytes; + const uintptr_t p0 = h0 & ~(uintptr_t)(page_sz - 1u); + const uintptr_t p1 = (h1 + page_sz - 1u) & ~(uintptr_t)(page_sz - 1u); + if (p1 > p0) (void)posix_madvise((void *)p0, (size_t)(p1 - p0), POSIX_MADV_DONTNEED); +#else + (void)model_map; + (void)model_size; + (void)offset; + (void)bytes; +#endif +} + +static void cuda_model_drop_file_pages(uint64_t offset, uint64_t bytes) { +#if defined(POSIX_FADV_DONTNEED) + if (g_model_fd < 0 || getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || bytes == 0) return; + (void)posix_fadvise(g_model_fd, (off_t)offset, (off_t)bytes, POSIX_FADV_DONTNEED); +#else + (void)offset; + (void)bytes; +#endif +} + +static uint64_t cuda_round_down(uint64_t v, uint64_t align) { + if (align <= 1) return v; + return (v / align) * align; +} + +static uint64_t cuda_round_up(uint64_t v, uint64_t align) { + if (align <= 1) return v; + const uint64_t rem = v % align; + return rem == 0 ? v : v + (align - rem); +} + +static void *cuda_align_ptr(void *ptr, uint64_t align) { + if (align <= 1) return ptr; + uintptr_t p = (uintptr_t)ptr; + uintptr_t a = (uintptr_t)align; + return (void *)(((p + a - 1u) / a) * a); +} + +static int cuda_model_stage_pool_alloc(uint64_t bytes) { + if (g_model_stage_bytes >= bytes) return 1; + for (size_t i = 0; i < 4; i++) { + if (g_model_stage_event[i]) { + (void)cudaEventDestroy(g_model_stage_event[i]); + g_model_stage_event[i] = NULL; + } + if (g_model_stage_raw[i]) { + (void)cudaFreeHost(g_model_stage_raw[i]); + g_model_stage_raw[i] = NULL; + g_model_stage[i] = NULL; + } + } + g_model_stage_bytes = 0; + if (!g_model_upload_stream) { + cudaError_t err = cudaStreamCreateWithFlags(&g_model_upload_stream, cudaStreamNonBlocking); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model upload stream creation failed: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + } + for (size_t i = 0; i < 4; i++) { + cudaError_t err = cudaMallocHost(&g_model_stage_raw[i], (size_t)bytes); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + g_model_stage[i] = cuda_align_ptr(g_model_stage_raw[i], g_model_direct_align); + err = cudaEventCreateWithFlags(&g_model_stage_event[i], cudaEventDisableTiming); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model staging event creation failed: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + } + g_model_stage_bytes = bytes; + return 1; +} + +static int cuda_pread_full(int fd, void *buf, uint64_t bytes, uint64_t offset) { + uint64_t done = 0; + while (done < bytes) { + const size_t n_req = (bytes - done > (uint64_t)SSIZE_MAX) ? (size_t)SSIZE_MAX : (size_t)(bytes - done); + ssize_t n = pread(fd, (char *)buf + done, n_req, (off_t)(offset + done)); + if (n < 0) { + if (errno == EINTR) continue; + return 0; + } + if (n == 0) return 0; + done += (uint64_t)n; + } + return 1; +} + +static int cuda_model_stage_read(void *stage, uint64_t stage_bytes, + uint64_t offset, uint64_t bytes, + const char **payload) { + *payload = (const char *)stage; +#if defined(__linux__) && defined(O_DIRECT) + if (g_model_direct_fd >= 0 && g_model_direct_align > 1 && g_model_file_size != 0) { + const uint64_t aligned_off = cuda_round_down(offset, g_model_direct_align); + const uint64_t delta = offset - aligned_off; + uint64_t read_size = cuda_round_up(delta + bytes, g_model_direct_align); + if (aligned_off <= g_model_file_size && + read_size <= stage_bytes && + read_size <= g_model_file_size - aligned_off) { + const int saved_errno = errno; + errno = 0; + if (cuda_pread_full(g_model_direct_fd, stage, read_size, aligned_off)) { + *payload = (const char *)stage + delta; + errno = saved_errno; + return 1; + } + const int direct_errno = errno; + if (direct_errno == EINVAL || direct_errno == EFAULT || direct_errno == ENOTSUP || direct_errno == EOPNOTSUPP) { + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA direct model read disabled: %s\n", strerror(direct_errno)); + } + (void)close(g_model_direct_fd); + g_model_direct_fd = -1; + g_model_direct_align = 1; + } + errno = direct_errno; + } + } +#else + (void)stage_bytes; +#endif + return cuda_pread_full(g_model_fd, stage, bytes, offset); +} + +static void cuda_stream_selected_stage_release(void) { + for (size_t i = 0; i < 4; i++) { + if (g_stream_selected_stage_event[i]) { + (void)cudaEventDestroy(g_stream_selected_stage_event[i]); + g_stream_selected_stage_event[i] = NULL; + } + if (g_stream_selected_stage_raw[i]) { + (void)cudaFreeHost(g_stream_selected_stage_raw[i]); + g_stream_selected_stage_raw[i] = NULL; + g_stream_selected_stage[i] = NULL; + } + } + g_stream_selected_stage_bytes = 0; + if (g_stream_selected_upload_stream) { + (void)cudaStreamDestroy(g_stream_selected_upload_stream); + g_stream_selected_upload_stream = NULL; + } +} + +static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { + if (g_stream_selected_stage_bytes >= bytes) return 1; + cuda_stream_selected_stage_release(); + cudaError_t err = cudaStreamCreateWithFlags( + &g_stream_selected_upload_stream, cudaStreamNonBlocking); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected upload stream creation failed: %s\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + for (size_t i = 0; i < 4; i++) { + err = cudaMallocHost(&g_stream_selected_stage_raw[i], (size_t)bytes); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected staging allocation failed: %s\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + cuda_stream_selected_stage_release(); + return 0; + } + g_stream_selected_stage[i] = cuda_align_ptr( + g_stream_selected_stage_raw[i], g_model_direct_align); + err = cudaEventCreateWithFlags(&g_stream_selected_stage_event[i], + cudaEventDisableTiming); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected staging event creation failed: %s\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + cuda_stream_selected_stage_release(); + return 0; + } + } + g_stream_selected_stage_bytes = bytes; + return 1; +} + +static int cuda_model_copy_to_device_streamed( + char *dst, + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t bytes, + const char *what) { + if (!dst || !model_map || offset > model_size || + bytes > model_size - offset) { + return 0; + } + if (bytes == 0) return 1; + if (g_model_fd < 0 || + (g_model_fd_host_base != NULL && model_map != g_model_fd_host_base)) { + return cuda_ok(cudaMemcpy(dst, + (const char *)model_map + offset, + (size_t)bytes, + cudaMemcpyHostToDevice), + what ? what : "stream selected expert copy"); + } + + const uint64_t chunk = cuda_model_copy_chunk_bytes(); + const uint64_t stage_bytes = + chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); + if (!cuda_stream_selected_stage_pool_alloc(stage_bytes)) return 0; + + uint64_t copied = 0; + uint64_t chunk_idx = 0; + while (copied < bytes) { + const uint64_t n = bytes - copied < chunk ? bytes - copied : chunk; + const uint64_t bi = chunk_idx % 4u; + cudaError_t err; + if (chunk_idx >= 4u) { + err = cudaEventSynchronize(g_stream_selected_stage_event[bi]); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected staging wait failed for %s: %s\n", + what ? what : "expert", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + } + const char *payload = NULL; + if (!cuda_model_stage_read(g_stream_selected_stage[bi], + g_stream_selected_stage_bytes, + offset + copied, n, &payload)) { + fprintf(stderr, + "ds4: CUDA streaming selected read failed for %s at %.2f MiB: %s\n", + what ? what : "expert", (double)copied / 1048576.0, + strerror(errno)); + return 0; + } + err = cudaMemcpyAsync(dst + copied, payload, (size_t)n, + cudaMemcpyHostToDevice, + g_stream_selected_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected copy failed for %s at %.2f MiB: %s\n", + what ? what : "expert", (double)copied / 1048576.0, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + err = cudaEventRecord(g_stream_selected_stage_event[bi], + g_stream_selected_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected staging record failed for %s: %s\n", + what ? what : "expert", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + cuda_model_drop_file_pages(offset + copied, n); + cuda_model_discard_source_pages(model_map, model_size, + offset + copied, n); + copied += n; + chunk_idx++; + } + + const cudaError_t err = + cudaStreamSynchronize(g_stream_selected_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming selected upload sync failed for %s: %s\n", + what ? what : "expert", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + return 1; +} + +static uint64_t cuda_model_cache_limit_bytes(void) { + uint64_t gb = 0; + const char *env = getenv("DS4_CUDA_WEIGHT_CACHE_LIMIT_GB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env) gb = (uint64_t)v; + } + if (gb == 0) return UINT64_MAX; + return gb * 1073741824ull; +} + +static uint64_t cuda_model_arena_chunk_bytes(uint64_t need) { + uint64_t mb = 1792; + const char *env = getenv("DS4_CUDA_WEIGHT_ARENA_CHUNK_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env && v > 0) mb = (uint64_t)v; + } + if (mb < 256) mb = 256; + if (mb > 8192) mb = 8192; + uint64_t bytes = mb * 1048576ull; + if (bytes < need) { + const uint64_t align = 256ull * 1048576ull; + bytes = (need + align - 1u) & ~(align - 1u); + } + return bytes; +} + +static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) { + if (bytes == 0) return NULL; + if (g_model_cache_full) return NULL; + const uint64_t align = 256u; + const uint64_t aligned = (bytes + align - 1u) & ~(align - 1u); + + for (cuda_model_arena &a : g_model_arenas) { + const uint64_t used = (a.used + align - 1u) & ~(align - 1u); + if (used <= a.bytes && aligned <= a.bytes - used) { + char *ptr = a.device_ptr + used; + a.used = used + aligned; + return ptr; + } + } + + const uint64_t limit = cuda_model_cache_limit_bytes(); + if (g_model_range_bytes > limit || aligned > limit - g_model_range_bytes) return NULL; + + const uint64_t chunk = cuda_model_arena_chunk_bytes(aligned); + void *dev = NULL; + cudaError_t err = cudaMalloc(&dev, (size_t)chunk); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model arena alloc failed for %s (%.2f MiB chunk): %s\n", + what ? what : "weights", + (double)chunk / 1048576.0, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + g_model_cache_full = 1; + return NULL; + } + g_model_arenas.push_back({(char *)dev, chunk, aligned}); + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + uint64_t arena_bytes = 0; + for (const cuda_model_arena &a : g_model_arenas) arena_bytes += a.bytes; + fprintf(stderr, "ds4: CUDA model arena allocated %.2f MiB (arenas %.2f GiB)\n", + (double)chunk / 1048576.0, + (double)arena_bytes / 1073741824.0); + } + return (char *)dev; +} + +static const char *cuda_model_range_ptr_from_fd( + const void *model_map, + uint64_t offset, + uint64_t bytes, + const char *what) { + if (g_model_fd < 0 || bytes == 0) return NULL; + if (g_model_fd_host_base != NULL && model_map != g_model_fd_host_base) return NULL; + const uint64_t limit = cuda_model_cache_limit_bytes(); + if (g_model_range_bytes > limit || bytes > limit - g_model_range_bytes) { + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA direct %s %.2f MiB (cache budget %.2f GiB exhausted)\n", + what ? what : "weights", + (double)bytes / 1048576.0, + (double)limit / 1073741824.0); + } + return cuda_model_ptr(model_map, offset); + } + + char *dev = cuda_model_arena_alloc(bytes, what); + if (!dev) { + if (getenv("DS4_CUDA_STRICT_WEIGHT_CACHE") != NULL) return NULL; + return cuda_model_ptr(model_map, offset); + } + cudaError_t err = cudaSuccess; + + const uint64_t chunk = cuda_model_copy_chunk_bytes(); + const uint64_t stage_bytes = chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); + if (!cuda_model_stage_pool_alloc(stage_bytes)) return NULL; + + uint64_t copied = 0; + uint64_t chunk_idx = 0; + while (copied < bytes) { + const uint64_t n = (bytes - copied < chunk) ? (bytes - copied) : chunk; + const uint64_t bi = chunk_idx % 4u; + if (chunk_idx >= 4u) { + err = cudaEventSynchronize(g_model_stage_event[bi]); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model staging wait failed for %s: %s\n", + what ? what : "weights", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + } + const char *payload = NULL; + if (!cuda_model_stage_read(g_model_stage[bi], g_model_stage_bytes, + offset + copied, n, &payload)) { + fprintf(stderr, "ds4: CUDA model range read failed for %s at %.2f MiB: %s\n", + what ? what : "weights", + (double)copied / 1048576.0, + strerror(errno)); + return NULL; + } + err = cudaMemcpyAsync(dev + copied, payload, (size_t)n, + cudaMemcpyHostToDevice, g_model_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f MiB: %s\n", + what ? what : "weights", + (double)copied / 1048576.0, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + err = cudaEventRecord(g_model_stage_event[bi], g_model_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model staging record failed for %s: %s\n", + what ? what : "weights", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + cuda_model_drop_file_pages(offset + copied, n); + cuda_model_discard_source_pages(model_map, g_model_registered_size, offset + copied, n); + copied += n; + cuda_model_load_progress_note(g_model_range_bytes + copied); + chunk_idx++; + } + err = cudaStreamSynchronize(g_model_upload_stream); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model range upload sync failed for %s: %s\n", + what ? what : "weights", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return NULL; + } + + g_model_ranges.push_back({model_map, offset, bytes, dev, NULL, NULL, 0, 0, 1}); + g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; + g_model_range_bytes += bytes; + cuda_model_load_progress_note(g_model_range_bytes); + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA fd-cached %s %.2f MiB (total %.2f GiB)\n", + what ? what : "weights", + (double)bytes / 1048576.0, + (double)g_model_range_bytes / 1073741824.0); + } + return (const char *)dev; +} + +static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) { + if (!model_map || model_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0; + if (getenv("DS4_CUDA_NO_MODEL_COPY") != NULL || + getenv("DS4_CUDA_DIRECT_MODEL") != NULL || + getenv("DS4_CUDA_WEIGHT_CACHE") != NULL || + getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) { + return 0; + } + if (g_model_device_owned || g_model_registered) return 1; + + void *dev = NULL; + const double t0 = cuda_wall_sec(); + cudaError_t err = cudaMalloc(&dev, (size_t)model_size); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + + fprintf(stderr, "ds4: CUDA chunk-copying %.2f GiB model image\n", + (double)model_size / 1073741824.0); + + const uint64_t chunk = cuda_model_copy_chunk_bytes(); + void *stage = NULL; + err = cudaMallocHost(&stage, (size_t)chunk); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err)); + (void)cudaFree(dev); + (void)cudaGetLastError(); + return 0; + } + + if (map_offset > 0) { + uint64_t copied_header = 0; + while (copied_header < map_offset) { + const uint64_t n = (map_offset - copied_header < chunk) ? (map_offset - copied_header) : chunk; + memcpy(stage, (const char *)model_map + copied_header, (size_t)n); + err = cudaMemcpy((char *)dev + copied_header, stage, (size_t)n, cudaMemcpyHostToDevice); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model header copy failed: %s\n", cudaGetErrorString(err)); + (void)cudaFreeHost(stage); + (void)cudaFree(dev); + (void)cudaGetLastError(); + return 0; + } + copied_header += n; + } + } + + uint64_t copied = 0; + double last_report = t0; + while (copied < map_size) { + const uint64_t n = (map_size - copied < chunk) ? (map_size - copied) : chunk; + const uint64_t off = map_offset + copied; + memcpy(stage, (const char *)model_map + off, (size_t)n); + err = cudaMemcpy((char *)dev + off, stage, (size_t)n, cudaMemcpyHostToDevice); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model chunk copy failed at %.2f GiB: %s\n", + (double)copied / 1073741824.0, cudaGetErrorString(err)); + (void)cudaFreeHost(stage); + (void)cudaFree(dev); + (void)cudaGetLastError(); + return 0; + } + cuda_model_discard_source_pages(model_map, model_size, off, n); + copied += n; + const double now = cuda_wall_sec(); + if (getenv("DS4_CUDA_MODEL_COPY_VERBOSE") != NULL && now - last_report >= 2.0) { + fprintf(stderr, "ds4: CUDA model chunk copy %.2f/%.2f GiB\n", + (double)copied / 1073741824.0, + (double)map_size / 1073741824.0); + last_report = now; + } + } + + (void)cudaFreeHost(stage); + g_model_device_base = (const char *)dev; + g_model_device_owned = 1; + g_model_hmm_direct = 0; + const double t1 = cuda_wall_sec(); + fprintf(stderr, + "ds4: CUDA model chunk copy complete in %.3fs (%.2f GiB tensors)\n", + t1 - t0, + (double)map_size / 1073741824.0); + return 1; +} + +static void cuda_model_range_release_all(void) { + for (const cuda_model_range &r : g_model_ranges) { + if (r.host_registered && r.registered_base) { + (void)cudaHostUnregister(r.registered_base); + } else if (r.device_ptr && !r.arena_allocated) { + (void)cudaFree(r.device_ptr); + } + } + for (const cuda_model_arena &a : g_model_arenas) { + if (a.device_ptr) (void)cudaFree(a.device_ptr); + } + g_model_arenas.clear(); + g_model_ranges.clear(); + g_model_range_by_offset.clear(); + g_model_range_bytes = 0; + cuda_model_load_progress_reset(); +} + +static int cublas_ok(cublasStatus_t st, const char *what) { + if (st == CUBLAS_STATUS_SUCCESS) return 1; + fprintf(stderr, "ds4: cuBLAS %s failed: status %d\n", what, (int)st); + return 0; +} + +extern "C" int ds4_gpu_init_multi(const ds4_gpu_config *cfg) { + if (!cfg || cfg->n_gpus < 1 || cfg->n_gpus > DS4_MAX_GPUS) return 0; + cuda_xdev_env_refresh(); + cuda_decode_dispatch_env_refresh(); + g_current_logical_tier = -1; + + /* g_n_gpus is published incrementally so ds4_gpu_cleanup() can unwind + * partial state on failure. We publish `i + 1` BEFORE allocating any + * resources for context `i`, so even if (e.g.) stream creation + * succeeds but event creation fails, cleanup still walks device `i` + * and destroys the stream. ds4_gpu_cleanup is null-safe per field — + * partial state is OK. */ + for (int i = 0; i < cfg->n_gpus; i++) { + ds4_gpu_ctx *c = &g_gpu[i]; + c->device_id = cfg->device_indices[i]; + if (c->device_id < 0) return 0; + /* Publish the in-progress device id so cleanup can target it on + * any later failure. cudaSetDevice is also required before + * cleanup's cudaEventDestroy / cudaStreamDestroy / cublasDestroy + * calls hit the right context. */ + g_n_gpus = i + 1; + if (!cuda_ok(cudaSetDevice(c->device_id), "init set device")) return 0; + cudaDeviceProp prop; + if (cudaGetDeviceProperties(&prop, c->device_id) == cudaSuccess) { + fprintf(stderr, "ds4: CUDA backend initialized on %s (sm_%d%d) dev=%d\n", + prop.name, prop.major, prop.minor, c->device_id); + } + /* Per-device stream. */ + cudaStream_t s = NULL; + if (!cuda_ok(cudaStreamCreate(&s), "init stream")) return 0; + c->stream = (void *)s; + /* Per-device boundary event (reusable, no timing). */ + cudaEvent_t ev = NULL; + if (!cuda_ok(cudaEventCreateWithFlags(&ev, cudaEventDisableTiming), + "init event")) return 0; + c->boundary_event = (void *)ev; + /* Per-device cuBLAS handle. */ + cublasHandle_t h = NULL; + if (!cublas_ok(cublasCreate(&h), "init cublas")) return 0; + c->cublas = (void *)h; + const cublasMath_t math_mode = + (g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL) + ? CUBLAS_DEFAULT_MATH + : CUBLAS_TF32_TENSOR_OP_MATH; + (void)cublasSetMathMode(h, math_mode); + c->cublas_ready = 1; + c->budget_bytes = cfg->vram_bytes[i]; + c->used_bytes = 0; + c->scratch = NULL; + c->scratch_bytes = 0; + } + + /* NxN peer-access matrix. + * + * Driver semantics: cudaDeviceCanAccessPeer + cudaDeviceEnablePeerAccess + * can both succeed even on hardware/drivers where cudaMemcpyPeerAsync + * silently delivers wrong data (notably RTX 6000 Ada under recent + * NVIDIA drivers, per the v0 design doc). The corruption is + * non-deterministic and can affect either or both directions of a + * pair. To guard against this, we run a multi-size, multi-iteration + * validation at init (see the loop below): write distinct known + * patterns, peer-copy them to the destination, read back, and only + * set peer_ok[i][j] if every probe round-trips byte-perfect. A + * single small probe is not sufficient — it can pass while realistic + * activation-sized copies still corrupt. On any failure the entry + * stays at 0 and cross-device copies fall back to the pinned-host + * bounce path automatically. */ + for (int i = 0; i < g_n_gpus; i++) { + for (int j = 0; j < g_n_gpus; j++) { + if (i == j) { g_gpu_peer_ok[i][j] = 1; continue; } + int can = 0; + (void)cudaDeviceCanAccessPeer(&can, g_gpu[i].device_id, + g_gpu[j].device_id); + if (!can) { g_gpu_peer_ok[i][j] = 0; continue; } + (void)cudaSetDevice(g_gpu[i].device_id); + cudaError_t e = cudaDeviceEnablePeerAccess(g_gpu[j].device_id, 0); + int enabled = (e == cudaSuccess || + e == cudaErrorPeerAccessAlreadyEnabled); + (void)cudaGetLastError(); + if (!enabled) { g_gpu_peer_ok[i][j] = 0; continue; } + + /* Runtime validation: peer copies on RTX 6000 Ada under recent + * NVIDIA drivers silently corrupt at realistic sizes even though + * the API returns success. cudaDeviceCanAccessPeer and + * cudaDeviceEnablePeerAccess can both report success while + * cudaMemcpyPeer delivers wrong data non-deterministically. + * Probe with multiple sizes and iterations; ALL must round-trip + * byte-perfect or we disable peer for this pair and silently fall + * back to the pinned-host bounce path. */ + static const size_t kValidateSizes[] = { + 4u * 1024u, + 256u * 1024u, + 1u * 1024u * 1024u, + 16u * 1024u * 1024u, + }; + const int kValidateIters = 4; + const int kNValidateSizes = (int)(sizeof(kValidateSizes) / + sizeof(kValidateSizes[0])); + const size_t kMaxValidate = kValidateSizes[kNValidateSizes - 1]; + + unsigned char *vh_src = (unsigned char *)malloc(kMaxValidate); + unsigned char *vh_dst = (unsigned char *)malloc(kMaxValidate); + if (!vh_src || !vh_dst) { + free(vh_src); free(vh_dst); + g_gpu_peer_ok[i][j] = 0; continue; + } + + void *src_dev = NULL; void *dst_dev = NULL; + (void)cudaSetDevice(g_gpu[i].device_id); + if (cudaMalloc(&src_dev, kMaxValidate) != cudaSuccess) { + (void)cudaGetLastError(); + free(vh_src); free(vh_dst); + g_gpu_peer_ok[i][j] = 0; continue; + } + (void)cudaSetDevice(g_gpu[j].device_id); + if (cudaMalloc(&dst_dev, kMaxValidate) != cudaSuccess) { + (void)cudaGetLastError(); + (void)cudaSetDevice(g_gpu[i].device_id); + (void)cudaFree(src_dev); + free(vh_src); free(vh_dst); + g_gpu_peer_ok[i][j] = 0; continue; + } + + int peer_validated = 1; + size_t failed_bytes = 0; + int failed_iter = -1; + for (int s_idx = 0; + s_idx < kNValidateSizes && peer_validated; + s_idx++) { + size_t n = kValidateSizes[s_idx]; + for (int it = 0; it < kValidateIters && peer_validated; it++) { + for (size_t k = 0; k < n; k++) { + vh_src[k] = (unsigned char) + ((k * 31u + (size_t)it * 17u + + (size_t)s_idx * 53u + 11u) & 0xffu); + } + (void)cudaSetDevice(g_gpu[i].device_id); + if (cudaMemcpy(src_dev, vh_src, n, + cudaMemcpyHostToDevice) != cudaSuccess) { + peer_validated = 0; failed_bytes = n; failed_iter = it; + break; + } + cudaError_t pc = cudaMemcpyPeer( + dst_dev, g_gpu[j].device_id, + src_dev, g_gpu[i].device_id, n); + if (pc != cudaSuccess) { + peer_validated = 0; failed_bytes = n; failed_iter = it; + break; + } + (void)cudaSetDevice(g_gpu[j].device_id); + if (cudaMemcpy(vh_dst, dst_dev, n, + cudaMemcpyDeviceToHost) != cudaSuccess) { + peer_validated = 0; failed_bytes = n; failed_iter = it; + break; + } + if (memcmp(vh_src, vh_dst, n) != 0) { + peer_validated = 0; failed_bytes = n; failed_iter = it; + break; + } + } + } + + (void)cudaSetDevice(g_gpu[j].device_id); + (void)cudaFree(dst_dev); + (void)cudaSetDevice(g_gpu[i].device_id); + (void)cudaFree(src_dev); + free(vh_src); + free(vh_dst); + + g_gpu_peer_ok[i][j] = peer_validated; + if (peer_validated) { + fprintf(stderr, + "ds4: peer access %d->%d validated across %d sizes x %d" + " iterations (max %zu MiB)\n", + g_gpu[i].device_id, g_gpu[j].device_id, + kNValidateSizes, kValidateIters, + kMaxValidate / (1024u * 1024u)); + } else { + fprintf(stderr, + "ds4: peer access %d->%d FAILED validation at" + " size=%zu iter=%d; falling back to pinned-host bounce\n", + g_gpu[i].device_id, g_gpu[j].device_id, + failed_bytes, failed_iter); + } + } + } + + g_cublas_ready = 1; + return 1; +} + +extern "C" int ds4_gpu_init(void) { + ds4_gpu_config cfg; + memset(&cfg, 0, sizeof(cfg)); + cfg.device_indices[0] = 0; + cfg.n_gpus = 1; + return ds4_gpu_init_multi(&cfg); +} + +extern "C" void ds4_gpu_cleanup(void) { + (void)cudaDeviceSynchronize(); + g_current_logical_tier = -1; + + /* Multi-GPU teardown: events, streams, cublas handles, scratch + * slabs, per-pair bounce buffers. */ + for (int i = 0; i < g_n_gpus; i++) { + ds4_gpu_ctx *c = &g_gpu[i]; + (void)cudaSetDevice(c->device_id); + attention_decode_score_split_graph_destroy_one(i); + routed_moe_decode_graph_destroy_one(i); + if (c->boundary_event) { + (void)cudaEventDestroy((cudaEvent_t)c->boundary_event); + c->boundary_event = NULL; + } + if (c->stream) { + (void)cudaStreamDestroy((cudaStream_t)c->stream); + c->stream = NULL; + } + if (c->cublas) { + (void)cublasDestroy((cublasHandle_t)c->cublas); + c->cublas = NULL; + c->cublas_ready = 0; + } + if (c->scratch) { + (void)cudaFree(c->scratch); + c->scratch = NULL; + c->scratch_bytes = 0; + } + } + for (int i = 0; i < DS4_MAX_GPUS; i++) { + for (int j = 0; j < DS4_MAX_GPUS; j++) { + if (g_xdev_bounce[i][j]) { + (void)cudaFreeHost(g_xdev_bounce[i][j]); + g_xdev_bounce[i][j] = NULL; + g_xdev_bounce_bytes[i][j] = 0; + } + } + } + cuda_stream_selected_cache_release(); + cuda_stream_selected_stage_release(); + g_n_gpus = 0; + g_cublas_ready = 0; + + /* Per-device selective cache teardown (selective model cache). */ + for (int d = 0; d < DS4_MAX_GPUS; d++) { + if (!g_dev_cache[d].present) continue; + int prev = -1; + (void)cudaGetDevice(&prev); + (void)cudaSetDevice(d); + if (g_dev_cache[d].base) (void)cudaFree(g_dev_cache[d].base); + g_dev_cache[d].base = NULL; + g_dev_cache[d].bytes = 0; + g_dev_cache[d].present = 0; + if (prev >= 0) (void)cudaSetDevice(prev); + } + g_cache_ranges.clear(); + + /* Continue with legacy global teardown below. */ + + cuda_model_range_release_all(); + cuda_q8_f16_cache_release_all(); + g_q8_f16_disabled_after_oom = 0; + g_q8_f16_budget_notice_printed = 0; + for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { + (void)cudaFree(r.device_ptr); + } + g_q8_f32_ranges.clear(); + g_q8_f32_by_offset.clear(); + g_q8_f32_bytes = 0; + if (g_cuda_tmp) { + (void)cudaFree(g_cuda_tmp); + g_cuda_tmp = NULL; + g_cuda_tmp_bytes = 0; + } + for (size_t i = 0; i < 4; i++) { + if (g_model_stage_event[i]) { + (void)cudaEventDestroy(g_model_stage_event[i]); + g_model_stage_event[i] = NULL; + } + if (g_model_stage_raw[i]) { + (void)cudaFreeHost(g_model_stage_raw[i]); + g_model_stage_raw[i] = NULL; + g_model_stage[i] = NULL; + } + } + g_model_stage_bytes = 0; + if (g_model_upload_stream) { + (void)cudaStreamDestroy(g_model_upload_stream); + g_model_upload_stream = NULL; + } + if (g_model_device_owned && g_model_device_base) { + (void)cudaFree((void *)g_model_device_base); + } + if (g_model_registered && g_model_host_base) { + (void)cudaHostUnregister((void *)g_model_host_base); + } + g_model_host_base = NULL; + g_model_device_base = NULL; + g_model_registered_size = 0; + g_model_registered = 0; + g_model_device_owned = 0; + g_model_range_mapping_supported = 1; + g_model_hmm_direct = 0; + g_model_fd = -1; + if (g_model_direct_fd >= 0) { + (void)close(g_model_direct_fd); + g_model_direct_fd = -1; + } + g_model_direct_align = 1; + g_model_file_size = 0; + g_model_cache_full = 0; + if (g_model_prefetch_stream) { + (void)cudaStreamDestroy(g_model_prefetch_stream); + g_model_prefetch_stream = NULL; + } +} + +__global__ static void fill_f32_kernel(float *x, uint64_t n, float v); + +extern "C" int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, + uint64_t bytes) { + if (!t) return 1; + if (device_id < 0 || device_id >= g_n_gpus) return 2; + if (bytes == 0) bytes = 1; + int ok = 0; + WITH_DEVICE(g_gpu[device_id].device_id) { + ok = cuda_ok(cudaMalloc(&t->ptr, (size_t)bytes), "tensor alloc"); + } + if (!ok) return 3; + t->bytes = bytes; + t->owner = 1; + t->device_id = device_id; + g_gpu[device_id].used_bytes += bytes; + return 0; +} + +/* Async D2D copy queued on the destination device's default stream — + * ordering against the producer comes from the caller's fence; the CPU + * does not block (unlike ds4_gpu_tensor_copy's sync cudaMemcpy). */ +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 cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes, + cudaMemcpyDeviceToDevice, 0), + "tensor copy async"); +} + +extern "C" void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t) { + if (!t) return; + int d = ds4_tensor_device_idx(t); + if (t->owner && t->ptr) { + WITH_DEVICE(g_gpu[d].device_id) { + (void)cudaFree(t->ptr); + } + } + t->ptr = NULL; + t->bytes = 0; + t->owner = 0; +} + +extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { + ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); + if (!t) return NULL; + if (ds4_gpu_tensor_alloc_on(t, 0, bytes) != 0) { + free(t); + return NULL; + } + return t; +} + +extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { + if (bytes == 0) bytes = 1; + ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); + if (!t) return NULL; + int ok = 0; + /* Managed memory is not device-bound, but we record device 0 so that + * subsequent ds4_gpu_tensor_free pairs with WITH_DEVICE(0) safely. */ + WITH_DEVICE(g_gpu[0].device_id) { + ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes), + "managed tensor alloc"); + } + if (!ok) { free(t); return NULL; } + t->bytes = bytes; + t->owner = 1; + t->device_id = 0; + return t; +} + +/* Heap-allocated tensor on a specific logical tier. + * + * Mirrors the legacy ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *) + * with an explicit tier argument. Internally calls + * ds4_gpu_tensor_alloc_on on a freshly malloc'd struct. + * + * The legacy ds4_gpu_tensor_alloc(bytes) (above) calls + * ds4_gpu_tensor_alloc_on(t, 0, bytes); ds4_gpu_tensor_alloc_ptr_on(0, + * bytes) is byte-equivalent. Single-tier callers MAY remain on the + * legacy 1-arg helper; new multi-tier callers in ds4.c use _ptr_on. */ +extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes) { + if (tier < 0 || tier >= g_n_gpus) { + fprintf(stderr, + "ds4: ds4_gpu_tensor_alloc_ptr_on: bad tier %d (n_gpus=%d)\n", + tier, g_n_gpus); + return NULL; + } + ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); + if (!t) return NULL; + if (ds4_gpu_tensor_alloc_on(t, tier, bytes) != 0) { + free(t); + return NULL; + } + return t; +} + +/* Heap-allocated managed-memory tensor on a specific logical tier. + * Differs from ds4_gpu_tensor_alloc_managed only in stamping + * tier instead of 0. Used by the per-layer KV cache when tier !=0. + * + * Managed-memory paging behavior: cudaMallocManaged pages between + * devices on first-touch. In a single-tier pipeline the page lives on + * tier 0; in a multi-tier pipeline the layer's kernels run on the + * layer's tier so the page lives there after first-touch and stays + * unless another device touches it. Stamping tier matches the home + * device for free-time accounting. */ +extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes) { + if (tier < 0 || tier >= g_n_gpus) { + fprintf(stderr, + "ds4: ds4_gpu_tensor_alloc_managed_on: bad tier %d (n_gpus=%d)\n", + tier, g_n_gpus); + return NULL; + } + if (bytes == 0) bytes = 1; + ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); + if (!t) return NULL; + int ok = 0; + /* Run the cudaMallocManaged call under the home tier's device so the + * first-touch home matches the stamped device_id; the page itself can + * migrate freely under managed-memory semantics. */ + WITH_DEVICE(g_gpu[tier].device_id) { + ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes), + "managed tensor alloc (tier)"); + } + if (!ok) { free(t); return NULL; } + t->bytes = bytes; + t->owner = 1; + t->device_id = tier; + return t; +} + +extern "C" int ds4_gpu_tensor_device(const ds4_gpu_tensor *t) { + return t ? t->device_id : -1; +} + +static uint64_t cuda_managed_kv_reserve_bytes(uint64_t total_bytes) { + const uint64_t min_reserve = 8ull * 1073741824ull; + const uint64_t max_reserve = 40ull * 1073741824ull; + uint64_t reserve = total_bytes / 4u; + if (reserve < min_reserve) reserve = min_reserve; + if (reserve > max_reserve) reserve = max_reserve; + return reserve; +} + +extern "C" int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) { + if (kv_cache_bytes == 0) return 0; + + /* Very large KV caches are where device-only cudaMalloc() can make a + * unified-memory machine unresponsive. Managed memory restores the old + * demand-paged behavior for this one long-lived allocation class only. */ + const uint64_t huge_kv = 8ull * 1073741824ull; + if (kv_cache_bytes >= huge_kv) return 1; + + const uint64_t large_context = 8ull * 1073741824ull; + if (context_bytes < large_context) return 0; + + size_t free_b = 0; + size_t total_b = 0; + cudaError_t err = cudaMemGetInfo(&free_b, &total_b); + if (err != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + + const uint64_t free_bytes = (uint64_t)free_b; + const uint64_t total_bytes = (uint64_t)total_b; + const uint64_t reserve_bytes = cuda_managed_kv_reserve_bytes(total_bytes); + if (context_bytes > free_bytes) return 1; + return free_bytes - context_bytes < reserve_bytes; +} + +extern "C" ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) { + if (!base || offset > base->bytes || bytes > base->bytes - offset) return NULL; + ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); + if (!t) return NULL; + t->ptr = (char *)base->ptr + offset; + t->bytes = bytes; + t->owner = 0; + t->device_id = base->device_id; /* inherit owning device */ + return t; +} + +extern "C" void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) { + if (!tensor) return; + int d = ds4_tensor_device_idx(tensor); + if (tensor->owner && tensor->ptr) { + WITH_DEVICE(g_gpu[d].device_id) { + (void)cudaFree(tensor->ptr); + } + } + free(tensor); +} + +extern "C" uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) { + return tensor ? tensor->bytes : 0; +} + +extern "C" void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) { + if (!tensor) return NULL; + /* Full-device sync preserves legacy semantics. */ + (void)cudaDeviceSynchronize(); + return tensor->ptr; +} + +extern "C" int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) { + if (!tensor || count > tensor->bytes / sizeof(float)) return 0; + if (count == 0) return 1; + int d = ds4_tensor_device_idx(tensor); + int ok = 0; + WITH_DEVICE(g_gpu[d].device_id) { + fill_f32_kernel<<<(count + 255u) / 256u, 256>>>((float *)tensor->ptr, count, value); + ok = cuda_ok(cudaGetLastError(), "tensor fill f32 launch"); + } + return ok; +} + +extern "C" int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) { + if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0; + int d = ds4_tensor_device_idx(tensor); + int ok = 0; + WITH_DEVICE(g_gpu[d].device_id) { + ok = cuda_ok(cudaMemcpy((char *)tensor->ptr + offset, data, (size_t)bytes, + cudaMemcpyHostToDevice), + "tensor write"); + } + return ok; +} + +extern "C" int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) { + if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0; + int d = ds4_tensor_device_idx(tensor); + int ok = 0; + WITH_DEVICE(g_gpu[d].device_id) { + ok = cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset, (size_t)bytes, + cudaMemcpyDeviceToHost), + "tensor read"); + } + return ok; +} + +extern "C" int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t bytes) { + if (!dst || !src || dst_offset > dst->bytes || src_offset > src->bytes || + bytes > dst->bytes - dst_offset || bytes > src->bytes - src_offset) { + return 0; + } + if (bytes == 0) return 1; + /* Same-device fast path; for cross-device, callers should use + * ds4_gpu_tensor_copy_xdev. We still tolerate cross-device callers + * here by routing to D2D copy on the destination's device. */ + int d = ds4_tensor_device_idx(dst); + int ok = 0; + WITH_DEVICE(g_gpu[d].device_id) { + ok = cuda_ok(cudaMemcpy((char *)dst->ptr + dst_offset, + (const char *)src->ptr + src_offset, + (size_t)bytes, + cudaMemcpyDeviceToDevice), + "tensor copy"); + } + return ok; +} + +__global__ static void moe_handoff_pack_kernel( + unsigned char *packed, + const float *ffn_norm, + const int32_t *selected, + const float *weights, + uint32_t n_embd, + uint32_t n_expert) { + const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + float *packed_norm = (float *)packed; + int32_t *packed_selected = (int32_t *)(packed + (uint64_t)n_embd * sizeof(float)); + float *packed_weights = (float *)(packed + (uint64_t)n_embd * sizeof(float) + + (uint64_t)n_expert * sizeof(int32_t)); + if (i < n_embd) packed_norm[i] = ffn_norm[i]; + if (i < n_expert) { + packed_selected[i] = selected[i]; + packed_weights[i] = weights[i]; + } +} + +extern "C" int ds4_gpu_moe_handoff_pack_tensor( + ds4_gpu_tensor *packed, + const ds4_gpu_tensor *ffn_norm, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_expert) { + if (!packed || !ffn_norm || !selected || !weights || + n_embd == 0 || n_expert == 0) { + return 0; + } + const uint64_t bytes = (uint64_t)n_embd * sizeof(float) + + (uint64_t)n_expert * sizeof(int32_t) + + (uint64_t)n_expert * sizeof(float); + if (packed->bytes < bytes || + ffn_norm->bytes < (uint64_t)n_embd * sizeof(float) || + selected->bytes < (uint64_t)n_expert * sizeof(int32_t) || + weights->bytes < (uint64_t)n_expert * sizeof(float)) { + return 0; + } + const uint32_t n = n_embd > n_expert ? n_embd : n_expert; + moe_handoff_pack_kernel<<<(n + 255u) / 256u, 256>>>( + (unsigned char *)packed->ptr, + (const float *)ffn_norm->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + n_embd, + n_expert); + return cuda_ok(cudaGetLastError(), "moe handoff pack launch"); +} + +/* Cross-device copy primitive. Path selection (highest priority first): + * DS4_FORCE_HOST_BOUNCE=1 -> always pinned-host bounce + * DS4_FORCE_CUDA_PEER=1 -> always cudaMemcpyPeerAsync (manual-testing + * override; bypasses g_gpu_peer_ok) + * otherwise -> peer if validation passed at init, else + * pinned-host bounce. */ +static int ds4_gpu_tensor_copy_xdev_impl(ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint64_t bytes, + bool order_dst_before_write) { + if (!dst || !src) return 0; + if (bytes == 0) return 1; + if (bytes > dst->bytes || bytes > src->bytes) return 0; + int sd = ds4_tensor_device_idx(src); + int dd = ds4_tensor_device_idx(dst); + + /* Same-device fast path. */ + if (sd == dd) { + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, bytes, + cudaMemcpyDeviceToDevice, s), + "xdev same-device copy"); + if (ok && g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s), "xdev same-device sync"); + } + } + return ok; + } + + int peer = g_gpu_peer_ok[sd][dd]; + if (g_xdev_force_cuda_peer) peer = 1; + if (g_xdev_force_host_bounce) peer = 0; + + if (peer) { + int ok = 0; + if (order_dst_before_write) { + WITH_DEVICE(g_gpu[dd].device_id) { + cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; + cudaEvent_t e2 = (cudaEvent_t)g_gpu[dd].boundary_event; + ok = cuda_ok(cudaEventRecord(e2, s2), "peer dst-ready event record"); + } + if (!ok) return 0; + } else { + ok = 1; + } + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; + if (order_dst_before_write) { + ok = cuda_ok(cudaStreamWaitEvent(s, (cudaEvent_t)g_gpu[dd].boundary_event, 0), + "peer src wait dst-ready"); + } + if (ok) ok = cuda_ok(cudaMemcpyPeerAsync( + dst->ptr, g_gpu[dd].device_id, + src->ptr, g_gpu[sd].device_id, + bytes, s), + "peer copy"); + if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer event record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; + (void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0); + if (g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s2), "peer dst sync"); + } + } + return ok; + } + + /* Per-pair pinned-host bounce buffer. */ + if (g_xdev_bounce_bytes[sd][dd] < bytes) { + if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]); + if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes), + "bounce alloc")) return 0; + g_xdev_bounce_bytes[sd][dd] = bytes; + } + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; + ok = cuda_ok(cudaMemcpyAsync(g_xdev_bounce[sd][dd], src->ptr, bytes, + cudaMemcpyDeviceToHost, s), + "bounce d2h"); + if (ok) ok = cuda_ok(cudaEventRecord(e, s), "bounce event record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; + (void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0); + ok = cuda_ok(cudaMemcpyAsync(dst->ptr, g_xdev_bounce[sd][dd], bytes, + cudaMemcpyHostToDevice, s2), + "bounce h2d"); + if (ok) ok = cuda_ok(cudaStreamSynchronize(s2), "bounce dst sync"); + } + return ok; +} + +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_xdev_impl(dst, src, bytes, false); +} + +static int ds4_gpu_tensor_copy_xdev_default_impl(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 == 0u) return 1; + const int sd = ds4_tensor_device_idx(src); + const int dd = ds4_tensor_device_idx(dst); + if (sd == dd) { + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes, + cudaMemcpyDeviceToDevice, 0), + "default-stream same-device copy"); + } + return ok; + } + + int peer = g_gpu_peer_ok[sd][dd]; + if (g_xdev_force_cuda_peer) peer = 1; + if (g_xdev_force_host_bounce) peer = 0; + if (peer) { + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + ok = cuda_ok(cudaMemcpyPeerAsync( + dst->ptr, g_gpu[dd].device_id, + src->ptr, g_gpu[sd].device_id, + (size_t)bytes, 0), + "default-stream peer copy"); + if (ok) { + ok = cuda_ok(cudaEventRecord( + (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "default-stream peer event record"); + } + } + if (ok) { + WITH_DEVICE(g_gpu[dd].device_id) { + ok = cuda_ok(cudaStreamWaitEvent( + 0, + (cudaEvent_t)g_gpu[sd].boundary_event, + 0), + "default-stream peer destination wait"); + } + } + return ok; + } + + if (g_xdev_bounce_bytes[sd][dd] < bytes) { + if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]); + if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes), + "default-stream bounce alloc")) return 0; + g_xdev_bounce_bytes[sd][dd] = bytes; + } + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + ok = cuda_ok(cudaMemcpy(g_xdev_bounce[sd][dd], src->ptr, (size_t)bytes, + cudaMemcpyDeviceToHost), + "default-stream bounce d2h"); + } + if (ok) { + WITH_DEVICE(g_gpu[dd].device_id) { + ok = cuda_ok(cudaMemcpy(dst->ptr, g_xdev_bounce[sd][dd], + (size_t)bytes, cudaMemcpyHostToDevice), + "default-stream bounce h2d"); + } + } + return ok; +} + +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_default_impl(dst, src, bytes); +} + +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) { + ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2}; + const ds4_gpu_tensor *srcs[3] = {src0, src1, src2}; + const uint64_t sizes[3] = {bytes0, bytes1, bytes2}; + int sd = -1; + int dd = -1; + for (int i = 0; i < 3; i++) { + if (sizes[i] == 0u) continue; + if (!dsts[i] || !srcs[i] || sizes[i] > dsts[i]->bytes || + sizes[i] > srcs[i]->bytes) { + return 0; + } + const int this_sd = ds4_tensor_device_idx(srcs[i]); + const int this_dd = ds4_tensor_device_idx(dsts[i]); + if (sd < 0) { + sd = this_sd; + dd = this_dd; + } else if (sd != this_sd || dd != this_dd) { + return 0; + } + } + if (sd < 0) return 1; + if (sd == dd) { + int ok = 1; + WITH_DEVICE(g_gpu[sd].device_id) { + for (int i = 0; ok && i < 3; i++) { + if (sizes[i] == 0u) continue; + ok = cuda_ok(cudaMemcpyAsync( + dsts[i]->ptr, srcs[i]->ptr, (size_t)sizes[i], + cudaMemcpyDeviceToDevice, 0), + "grouped default same-device copy"); + } + } + return ok; + } + + int peer = g_gpu_peer_ok[dd][sd]; + if (g_xdev_force_cuda_peer) peer = 1; + if (g_xdev_force_host_bounce) peer = 0; + if (!peer) { + for (int i = 0; i < 3; i++) { + if (sizes[i] != 0u && + !ds4_gpu_tensor_copy_xdev_default_impl( + dsts[i], srcs[i], sizes[i])) { + return 0; + } + } + return 1; + } + + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + ok = cuda_ok(cudaEventRecord( + (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "grouped default source-ready record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + ok = cuda_ok(cudaStreamWaitEvent( + 0, (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "grouped default destination wait"); + for (int i = 0; ok && i < 3; i++) { + if (sizes[i] == 0u) continue; + ok = cuda_ok(cudaMemcpyPeerAsync( + dsts[i]->ptr, g_gpu[dd].device_id, + srcs[i]->ptr, g_gpu[sd].device_id, + (size_t)sizes[i], 0), + "grouped destination-stream peer copy"); + } + } + return ok; +} + +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) { + ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2}; + const ds4_gpu_tensor *srcs[3] = {src0, src1, src2}; + uint64_t bytes[3] = {bytes0, bytes1, bytes2}; + int first = -1; + for (int i = 0; i < 3; i++) { + if (bytes[i] == 0) continue; + if (!dsts[i] || !srcs[i] || + bytes[i] > dsts[i]->bytes || bytes[i] > srcs[i]->bytes) { + return 0; + } + if (first < 0) first = i; + } + if (first < 0) return 1; + + const int sd = ds4_tensor_device_idx(srcs[first]); + const int dd = ds4_tensor_device_idx(dsts[first]); + for (int i = first + 1; i < 3; i++) { + if (bytes[i] == 0) continue; + if (ds4_tensor_device_idx(srcs[i]) != sd || + ds4_tensor_device_idx(dsts[i]) != dd) { + int ok = 1; + for (int j = 0; ok && j < 3; j++) { + if (bytes[j] == 0) continue; + ok = ds4_gpu_tensor_copy_xdev(dsts[j], srcs[j], bytes[j]); + } + return ok; + } + } + + if (sd == dd) { + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + ok = 1; + for (int i = 0; ok && i < 3; i++) { + if (bytes[i] == 0) continue; + ok = cuda_ok(cudaMemcpyAsync(dsts[i]->ptr, srcs[i]->ptr, bytes[i], + cudaMemcpyDeviceToDevice, s), + "xdev3 same-device copy"); + } + if (ok && g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s), "xdev3 same-device sync"); + } + } + return ok; + } + + int peer = g_gpu_peer_ok[sd][dd]; + if (g_xdev_force_cuda_peer) peer = 1; + if (g_xdev_force_host_bounce) peer = 0; + if (!peer) { + int ok = 1; + for (int i = 0; ok && i < 3; i++) { + if (bytes[i] == 0) continue; + ok = ds4_gpu_tensor_copy_xdev(dsts[i], srcs[i], bytes[i]); + } + return ok; + } + + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; + ok = 1; + for (int i = 0; ok && i < 3; i++) { + if (bytes[i] == 0) continue; + ok = cuda_ok(cudaMemcpyPeerAsync( + dsts[i]->ptr, g_gpu[dd].device_id, + srcs[i]->ptr, g_gpu[sd].device_id, + bytes[i], s), + "peer copy3"); + } + if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer copy3 event record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; + ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "peer copy3 dst wait"); + if (ok && g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s2), "peer copy3 dst sync"); + } + } + return ok; +} + +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_impl(dst, src, bytes, true); +} + +extern "C" int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) { + if (!src) return 0; + if (dst_tier < 0 || dst_tier >= g_n_gpus) return 0; + int sd = ds4_tensor_device_idx(src); + int dd = dst_tier; + if (sd == dd) return 1; + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; + cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; + ok = cuda_ok(cudaEventRecord(e, s), "xdev wait source event record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; + ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "xdev wait destination wait"); + if (ok && g_xdev_sync_debug) { + ok = cuda_ok(cudaStreamSynchronize(s2), "xdev wait dst sync"); + } + } + return ok; +} + +extern "C" int ds4_gpu_tensor_wait_xdev_default( + const ds4_gpu_tensor *src, + int dst_tier) { + if (!src || dst_tier < 0 || dst_tier >= g_n_gpus) return 0; + const int sd = ds4_tensor_device_idx(src); + const int dd = dst_tier; + if (sd == dd) return 1; + int ok = 0; + WITH_DEVICE(g_gpu[sd].device_id) { + ok = cuda_ok(cudaEventRecord( + (cudaEvent_t)g_gpu[sd].boundary_event, 0), + "default xdev wait source event record"); + } + if (!ok) return 0; + WITH_DEVICE(g_gpu[dd].device_id) { + ok = cuda_ok(cudaStreamWaitEvent( + 0, + (cudaEvent_t)g_gpu[sd].boundary_event, + 0), + "default xdev wait destination wait"); + } + return ok; +} + +extern "C" int ds4_gpu_q8_cache_suppressed(void) { + return g_q8_cache_suppressed; +} + +extern "C" void ds4_gpu_set_q8_cache_suppressed(int suppressed) { + g_q8_cache_suppressed = suppressed ? 1 : 0; +} + +__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap); + +extern "C" int ds4_gpu_pack_slot_rows_f32_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t n_rows, + uint32_t width, + uint32_t n_slots, + uint32_t slot_cap) { + uint64_t slot_rows = 0; + uint64_t slot_elems = 0; + uint64_t out_rows = 0; + uint64_t out_elems = 0; + if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 || + slot_cap == 0 || n_rows > slot_cap || + (uint64_t)n_slots > UINT64_MAX / slot_cap || + (slot_rows = (uint64_t)n_slots * slot_cap) > UINT64_MAX / width || + (slot_elems = slot_rows * width) > UINT64_MAX / sizeof(float) || + (uint64_t)n_rows > UINT64_MAX / n_slots || + (out_rows = (uint64_t)n_rows * n_slots) > UINT64_MAX / width || + (out_elems = out_rows * width) > UINT64_MAX / sizeof(float) || + slots->bytes < slot_elems * sizeof(float) || + out->bytes < out_elems * sizeof(float)) { + return 0; + } + const uint64_t blocks = (out_elems + 255u) / 256u; + if (blocks > UINT32_MAX) return 0; + pack_slot_rows_f32_kernel<<<(unsigned)blocks, 256>>>( + (float *)out->ptr, + (const float *)slots->ptr, + n_rows, + width, + n_slots, + slot_cap); + return cuda_ok(cudaGetLastError(), "pack_slot_rows_f32 launch"); +} + +extern "C" int ds4_gpu_begin_commands(void) { return 1; } +extern "C" int ds4_gpu_flush_commands(void) { return cuda_ok(cudaDeviceSynchronize(), "flush"); } +extern "C" int ds4_gpu_end_commands(void) { + if (g_cuda_end_stream_sync) { + return cuda_ok(cudaStreamSynchronize(0), "end commands stream"); + } + return cuda_ok(cudaDeviceSynchronize(), "end commands"); +} +extern "C" int ds4_gpu_synchronize(void) { return cuda_ok(cudaDeviceSynchronize(), "synchronize"); } + +extern "C" int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) { + if (!model_map || model_size == 0) return 0; + if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1; + cuda_stream_selected_cache_release(); + cuda_model_range_release_all(); + cuda_q8_f16_cache_release_all(); + g_q8_f16_disabled_after_oom = 0; + g_q8_f16_budget_notice_printed = 0; + for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { + (void)cudaFree(r.device_ptr); + } + g_q8_f32_ranges.clear(); + g_q8_f32_by_offset.clear(); + g_q8_f32_bytes = 0; + if (g_model_device_owned && g_model_device_base) { + (void)cudaFree((void *)g_model_device_base); + g_model_device_owned = 0; + } + if (g_model_registered && g_model_host_base) { + (void)cudaHostUnregister((void *)g_model_host_base); + g_model_registered = 0; + } + g_model_host_base = model_map; + g_model_device_base = (const char *)model_map; + g_model_registered_size = model_size; + g_model_range_mapping_supported = 1; + g_model_hmm_direct = 0; + g_model_cache_full = 0; + if (g_model_fd >= 0 && g_model_fd_host_base == NULL) { + g_model_fd_host_base = model_map; + } + + const char *copy_env = getenv("DS4_CUDA_COPY_MODEL"); + if (copy_env && copy_env[0]) { + void *dev = NULL; + const double t0 = clock() / (double)CLOCKS_PER_SEC; + cudaError_t err = cudaMalloc(&dev, (size_t)model_size); + if (err == cudaSuccess) { + fprintf(stderr, "ds4: CUDA copying %.2f GiB model to device memory\n", + (double)model_size / 1073741824.0); + err = cudaMemcpy(dev, model_map, (size_t)model_size, cudaMemcpyHostToDevice); + if (err == cudaSuccess) { + g_model_device_base = (const char *)dev; + g_model_device_owned = 1; + const double t1 = clock() / (double)CLOCKS_PER_SEC; + fprintf(stderr, "ds4: CUDA model copy complete in %.3fs\n", t1 - t0); + return 1; + } + fprintf(stderr, "ds4: CUDA model copy failed: %s\n", cudaGetErrorString(err)); + (void)cudaFree(dev); + (void)cudaGetLastError(); + } else { + fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + } + + cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size, + cudaHostRegisterMapped | cudaHostRegisterReadOnly); + if (err == cudaSuccess) { + void *dev = NULL; + err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0); + if (err == cudaSuccess && dev) { + g_model_device_base = (const char *)dev; + g_model_registered = 1; + fprintf(stderr, "ds4: CUDA registered %.2f GiB model mapping for device access\n", + (double)model_size / 1073741824.0); + } else { + fprintf(stderr, "ds4: CUDA host registration pointer lookup failed: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + } else { + fprintf(stderr, "ds4: CUDA host registration skipped: %s\n", cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + return 1; +} + +extern "C" int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) { + (void)max_tensor_bytes; + if (!ds4_gpu_register_model_map_no_copy(model_map, model_size)) return 0; + if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL && + !cuda_model_copy_chunked(model_map, model_size, map_offset, map_size)) { + (void)cuda_model_prefetch_range(model_map, model_size, map_offset, map_size); + } + return 1; +} + +/* Register the mmap'd host model pointer for selective-cache lookups WITHOUT + * triggering any device-side copy. Used by multi-GPU placement scaffolding's + * multi-tier path so DS4_CUDA_COPY_MODEL cannot reintroduce a full-model + * copy that defeats the per-device selective cache. + * + * This is the no-copy subset of ds4_gpu_set_model_map: same bookkeeping + * for the host pointer plus cudaHostRegister, but skipping the + * DS4_CUDA_COPY_MODEL branch that allocates and copies the entire model. */ +extern "C" int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size) { + if (!model_map || model_size == 0) return 0; + if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1; + + cuda_stream_selected_cache_release(); + cuda_model_range_release_all(); + cuda_q8_f16_cache_release_all(); + g_q8_f16_disabled_after_oom = 0; + g_q8_f16_budget_notice_printed = 0; + for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { + (void)cudaFree(r.device_ptr); + } + g_q8_f32_ranges.clear(); + g_q8_f32_by_offset.clear(); + g_q8_f32_bytes = 0; + if (g_model_device_owned && g_model_device_base) { + (void)cudaFree((void *)g_model_device_base); + g_model_device_owned = 0; + } + if (g_model_registered && g_model_host_base) { + (void)cudaHostUnregister((void *)g_model_host_base); + g_model_registered = 0; + } + g_model_host_base = model_map; + g_model_device_base = (const char *)model_map; + g_model_registered_size = model_size; + g_model_range_mapping_supported = 1; + g_model_hmm_direct = 0; + g_model_cache_full = 0; + if (g_model_fd >= 0 && g_model_fd_host_base == NULL) { + g_model_fd_host_base = model_map; + } + + /* No DS4_CUDA_COPY_MODEL branch — that is the entire point. */ + + cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size, + cudaHostRegisterMapped | cudaHostRegisterReadOnly); + if (err == cudaSuccess) { + void *dev = NULL; + err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0); + if (err == cudaSuccess && dev) { + g_model_device_base = (const char *)dev; + g_model_registered = 1; + fprintf(stderr, + "ds4: CUDA (no-copy) registered %.2f GiB model mapping for multi-tier selective cache\n", + (double)model_size / 1073741824.0); + } else { + fprintf(stderr, + "ds4: CUDA (no-copy) host registration pointer lookup failed: %s\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + } else { + fprintf(stderr, + "ds4: CUDA (no-copy) host registration skipped: %s\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + return 1; +} + +/* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1). + * Maps to the physical CUDA device id stored in g_gpu[].device_id. + * Added for multi-GPU placement scaffolding (multi-GPU CLI); first executed by + * multi-GPU execution (follow-up). */ +extern "C" int ds4_gpu_set_current_device(int logical_tier) { + if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1; + if (!g_cuda_no_setdevice_cache && g_current_logical_tier == logical_tier) { + return 0; + } + if (cudaSetDevice(g_gpu[logical_tier].device_id) == cudaSuccess) { + g_current_logical_tier = logical_tier; + return 0; + } + g_current_logical_tier = -1; + return -1; +} + +/* Fenced device switch for sequential cross-device pipelines (GLM + * per-layer placement): work queued on the next device's default stream + * waits for everything queued so far on the previous device's default + * stream. Async — no host sync. Falls back to a plain switch when the + * device does not change. */ +extern "C" int ds4_gpu_set_current_device_fenced(int logical_tier) { + if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1; + static cudaEvent_t fence_ev[DS4_MAX_GPUS]; + /* Resolve the ACTUAL current device: WITH_DEVICE blocks and direct + * cudaSetDevice calls can leave g_current_logical_tier stale, and a + * false "already there" here strands work on the wrong device. */ + int cur_dev = -1; + (void)cudaGetDevice(&cur_dev); + int prev = -1; + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == cur_dev) { prev = t; break; } + } + if (getenv("DS4_GLM_FENCE_TRACE")) { + fprintf(stderr, "ds4: fenced switch %d -> %d\n", prev, logical_tier); + } + if (prev == logical_tier) { + g_current_logical_tier = logical_tier; + return 0; + } + if (prev >= 0 && prev < g_n_gpus && prev != logical_tier) { + if (cudaSetDevice(g_gpu[prev].device_id) != cudaSuccess) return -1; + if (!fence_ev[prev] && + cudaEventCreateWithFlags(&fence_ev[prev], + cudaEventDisableTiming) != cudaSuccess) { + fence_ev[prev] = NULL; + } + if (fence_ev[prev]) { + (void)cudaEventRecord(fence_ev[prev], 0); + } + if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) { + g_current_logical_tier = -1; + return -1; + } + g_current_logical_tier = logical_tier; + if (fence_ev[prev]) { + (void)cudaStreamWaitEvent(0, fence_ev[prev], 0); + } + return 0; + } + return ds4_gpu_set_current_device(logical_tier); +} + +/* ========================================================================= + * Per-device selective model cache (selective model cache). + * + * ds4_gpu_device_cache_tensors copies the listed source ranges from the + * host mmap onto device_id's selective slab and appends sorted lookup + * entries. The legacy chunked-copy machinery (cuda_model_range_*) is + * NOT disturbed — it continues to drive all existing callers. New + * lookups fall back to it when no selective entry covers the range. + * + * Caller-context preference for overlap: when the same source range is + * cached on multiple devices, ds4_gpu_lookup_cache returns the entry + * whose device matches cudaGetDevice(). + * ========================================================================= */ + +extern "C" int ds4_gpu_device_cache_tensors(int device_id, + const ds4_tensor_range *ranges, + int n_ranges) { + if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1; + if (n_ranges < 0 || (!ranges && n_ranges > 0)) return 2; + if (n_ranges == 0) return 0; + + if (!g_model_host_base || g_model_registered_size == 0) return 3; + + /* Validate ranges against the mmap'd model bounds; reject ranges + * that overflow or extend past the mapped region. Done in a + * separate pass before any allocation so a bad input doesn't + * partially grow the slab. */ + uint64_t want_bytes = 0; + for (int i = 0; i < n_ranges; i++) { + if (ranges[i].target_device != device_id) continue; + const uint64_t off = ranges[i].source_offset; + const uint64_t nb = ranges[i].bytes; + /* Overflow-safe upper bound: off + nb must not exceed model + * size, and the sum must not wrap. */ + if (nb == 0) continue; + if (off > g_model_registered_size) return 8; + if (nb > g_model_registered_size - off) return 9; + /* Accumulate into want_bytes with overflow check. */ + if (want_bytes > UINT64_MAX - nb) return 10; + want_bytes += nb; + } + if (want_bytes == 0) return 0; + + cuda_device_cache &c = g_dev_cache[device_id]; + + int prev_device = -1; + if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1; + if (cudaSetDevice(device_id) != cudaSuccess) return 4; + + /* Allocate or grow the slab via cudaMalloc + d2d copy. */ + void *new_base = NULL; + size_t new_bytes = c.bytes + want_bytes; + + /* Refuse cleanly before cudaMalloc if the device clearly cannot hold + * the slab. The multi-tier packer reserves per-tier runtime scratch + * before placing tensors, but it cannot predict the cudaMalloc + * allocator's overhead (alignment, fragmentation after CUDA context + * init, default driver-side reservations). On a borderline budget + * that overhead pushes a "fits-by-packer-math" layout past the actual + * free pool and the cudaMalloc below OOMs after the engine already + * committed to the layout — same silent-late-OOM failure mode the + * upfront refusal path was added to eliminate. Catch it here too. */ + { + size_t free_b = 0, total_b = 0; + if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { + /* free_b already excludes the existing slab (it's still + * allocated), so the additional cudaMalloc only needs + * new_bytes free — not new_bytes + c.bytes. The old slab is + * freed AFTER the d2d copy succeeds. 2 GiB safety covers what + * the engine will allocate AFTER the cache slab in the same + * session_create: per-tier graph scratch (the planner can't + * predict its cumulative cudaMalloc alignment overhead), + * cuBLAS workspace beyond the 64 MiB the packer already + * reserves, and driver-side allocator slack. Without this + * headroom a borderline budget that fits the slab itself can + * still OOM at the per-tier tensor allocations a few moments + * later — same silent-late-OOM failure mode, one layer up. */ + const size_t safety = (size_t)2ull * 1024ull * 1024ull * 1024ull; + const size_t need = new_bytes + safety; + if (need > free_b) { + fprintf(stderr, + "ds4: device cache slab needs %.2f GiB on device %d " + "but only %.2f GiB free (slab=%.2f GiB + %.2f GiB safety). " + "Lower --gpu-vram / --ctx-max, or use --gpu-vram auto on " + "a host with more free VRAM. Refusing upfront to avoid " + "late OOM at cudaMalloc.\n", + (double)need / 1073741824.0, + device_id, + (double)free_b / 1073741824.0, + (double)new_bytes / 1073741824.0, + (double)safety / 1073741824.0); + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 5; + } + } + /* If cudaMemGetInfo itself failed, fall through; cudaMalloc's own + * error path still catches the late case, just with a less helpful + * message. */ + } + + if (!cuda_ok(cudaMalloc(&new_base, new_bytes), "device cache alloc")) { + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 5; + } + if (c.present && c.bytes > 0) { + cudaError_t e = cudaMemcpy(new_base, c.base, c.bytes, + cudaMemcpyDeviceToDevice); + if (e != cudaSuccess) { + cuda_ok(e, "device cache grow d2d"); + (void)cudaFree(new_base); + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 6; + } + /* Re-base existing entries on this device. */ + char *old_base = (char *)c.base; + char *grown = (char *)new_base; + for (size_t k = 0; k < g_cache_ranges.size(); k++) { + if (g_cache_ranges[k].device_id == device_id) { + g_cache_ranges[k].device_ptr = + grown + ((char *)g_cache_ranges[k].device_ptr - old_base); + } + } + (void)cudaFree(c.base); + } + c.base = new_base; + c.bytes = new_bytes; + c.present = 1; + + /* Copy ranges and append entries. */ + const char *host_base = (const char *)g_model_host_base; + size_t write_off = c.bytes - want_bytes; + for (int i = 0; i < n_ranges; i++) { + if (ranges[i].target_device != device_id) continue; + char *dev_ptr = (char *)c.base + write_off; + cudaError_t e = cudaMemcpy(dev_ptr, + host_base + ranges[i].source_offset, + (size_t)ranges[i].bytes, + cudaMemcpyHostToDevice); + if (e != cudaSuccess) { + cuda_ok(e, "device cache range h2d"); + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 7; + } + cache_range_entry ent; + ent.source_offset = ranges[i].source_offset; + ent.bytes = ranges[i].bytes; + ent.device_id = device_id; + ent.device_ptr = dev_ptr; + g_cache_ranges.push_back(ent); + write_off += ranges[i].bytes; + } + + /* Keep sorted by source_offset for binary-search lookup. */ + std::sort(g_cache_ranges.begin(), g_cache_ranges.end(), + [](const cache_range_entry &a, const cache_range_entry &b) { + if (a.source_offset != b.source_offset) + return a.source_offset < b.source_offset; + return a.device_id < b.device_id; + }); + + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 0; +} + +/* Install support-model tensor ranges into device_id's strict cache, + * copying from the registered support map and keying entries at + * source_offset + bias. Standalone slab (does not touch the main cache + * slab growth path). */ +extern "C" int ds4_gpu_device_cache_support_tensors(int device_id, + int entry_device_id, + const ds4_tensor_range *ranges, + int n_ranges, + int from_main_map) { + if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1; + if (entry_device_id < 0 || entry_device_id >= DS4_MAX_GPUS) return 1; + if (n_ranges <= 0 || !ranges) return 2; + const char *src_base; + uint64_t src_size; + uint64_t key_bias; + if (from_main_map) { + /* Auxiliary main-model ranges (e.g. the embedding bucket for the + * DSpark executor tier): standalone slab, unbiased offsets. */ + src_base = (const char *)g_model_host_base; + src_size = g_model_registered_size; + key_bias = 0; + } else { + src_base = (const char *)g_support_host_base; + src_size = g_support_host_size; + key_bias = g_support_offset_bias; + if (key_bias == 0) return 3; + } + if (!src_base || src_size == 0) return 3; + uint64_t want = 0; + for (int i = 0; i < n_ranges; i++) { + const uint64_t off = ranges[i].source_offset; + const uint64_t nb = ranges[i].bytes; + if (nb == 0) continue; + if (off > src_size || nb > src_size - off) return 8; + if (want > UINT64_MAX - nb) return 9; + want += nb; + } + if (want == 0) return 0; + int prev_device = -1; + if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1; + if (cudaSetDevice(device_id) != cudaSuccess) return 4; + void *base = NULL; + if (!cuda_ok(cudaMalloc(&base, (size_t)want), "support cache alloc")) { + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 5; + } + const char *host_base = src_base; + size_t write_off = 0; + for (int i = 0; i < n_ranges; i++) { + if (ranges[i].bytes == 0) continue; + char *dev_ptr = (char *)base + write_off; + cudaError_t e = cudaMemcpy(dev_ptr, + host_base + ranges[i].source_offset, + (size_t)ranges[i].bytes, + cudaMemcpyHostToDevice); + if (e != cudaSuccess) { + cuda_ok(e, "support cache range h2d"); + (void)cudaFree(base); + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 7; + } + cache_range_entry ent; + ent.source_offset = ranges[i].source_offset + key_bias; + ent.bytes = ranges[i].bytes; + /* Entries can claim a different (executor) device than the one the + * slab physically lives on: strict lookups filter by entry device, + * and peer access lets the executor's kernels dereference the + * spilled pointer directly. */ + ent.device_id = entry_device_id; + ent.device_ptr = dev_ptr; + g_cache_ranges.push_back(ent); + write_off += ranges[i].bytes; + } + std::sort(g_cache_ranges.begin(), g_cache_ranges.end(), + [](const cache_range_entry &a, const cache_range_entry &b) { + if (a.source_offset != b.source_offset) + return a.source_offset < b.source_offset; + return a.device_id < b.device_id; + }); + if (getenv("DS4_DSPARK_VERIFY_CACHE") != NULL) { + /* Read back every installed range and compare with the host copy. */ + int bad = 0; + write_off = 0; + for (int i = 0; i < n_ranges; i++) { + if (ranges[i].bytes == 0) continue; + char *dev_ptr = (char *)base + write_off; + std::vector tmp((size_t)ranges[i].bytes); + if (cudaMemcpy(tmp.data(), dev_ptr, (size_t)ranges[i].bytes, + cudaMemcpyDeviceToHost) != cudaSuccess || + memcmp(tmp.data(), host_base + ranges[i].source_offset, + (size_t)ranges[i].bytes) != 0) { + fprintf(stderr, + "ds4: support cache VERIFY MISMATCH offset=%llu bytes=%llu dev=%d\n", + (unsigned long long)ranges[i].source_offset, + (unsigned long long)ranges[i].bytes, device_id); + bad++; + } + write_off += ranges[i].bytes; + } + fprintf(stderr, "ds4: support cache verify dev=%d ranges=%d bad=%d\n", + device_id, n_ranges, bad); + } + if (prev_device >= 0) (void)cudaSetDevice(prev_device); + return 0; +} + +extern "C" int ds4_gpu_lookup_cache(uint64_t source_offset, uint64_t bytes, + int *out_device_id, void **out_device_ptr) { + int active_device = -1; + (void)cudaGetDevice(&active_device); + + if (!g_cache_ranges.empty()) { + /* upper_bound: first entry with source_offset > query. + * Candidates are at strictly earlier positions; scan all of + * them rather than breaking on the first non-covering entry, + * because the table allows overlap across devices. */ + auto it = std::upper_bound( + g_cache_ranges.begin(), g_cache_ranges.end(), + source_offset, + [](uint64_t off, const cache_range_entry &e) { + return off < e.source_offset; + }); + const cache_range_entry *match_any = NULL; + const cache_range_entry *match_pref = NULL; + while (it != g_cache_ranges.begin()) { + --it; + /* Overflow-safe coverage check: + * 1. source_offset >= it->source_offset + * 2. bytes <= it->bytes - (source_offset - it->source_offset) + * The second form computes only the remaining capacity inside + * the entry, so neither side can overflow even with bytes == + * UINT64_MAX. */ + if (source_offset >= it->source_offset) { + uint64_t into = source_offset - it->source_offset; + if (into <= it->bytes && bytes <= it->bytes - into) { + if (it->device_id == active_device) { + match_pref = &*it; + break; + } + if (!match_any) match_any = &*it; + } + } + /* Do NOT break on non-covering: an earlier entry may still + * cover if its bytes extend far enough. */ + } + const cache_range_entry *m = match_pref ? match_pref : match_any; + if (m) { + if (out_device_id) *out_device_id = m->device_id; + if (out_device_ptr) { + *out_device_ptr = + (char *)m->device_ptr + (source_offset - m->source_offset); + } + return 1; + } + } + + /* Legacy chunk-aware fallback (device 0 only). */ + const char *p = cuda_model_range_ptr_from_fd(g_model_host_base, + source_offset, bytes, + "lookup_cache"); + if (p) { + if (out_device_id) *out_device_id = 0; + if (out_device_ptr) *out_device_ptr = (void *)p; + return 1; + } + return 0; +} + +extern "C" int ds4_gpu_lookup_cache_device(uint64_t source_offset, uint64_t bytes) { + int d = -1; + if (!ds4_gpu_lookup_cache(source_offset, bytes, &d, NULL)) return -1; + return d; +} + +/* Strict per-device selective-cache lookup. + * + * Returns 1 only if a covering entry exists whose device_id matches the + * caller-supplied expected_device. Otherwise returns 0 with *out_device_ptr + * untouched. Unlike ds4_gpu_lookup_cache, this variant performs NO host- + * pointer fallback (no FD-cache, no model_range_ptr_from_fd) and NO + * different-device fallback. It is the canonical lookup for multi-tier + * dispatch where consuming a different device's pointer would be a + * correctness bug. expected_device is a PHYSICAL CUDA device id (the + * value stored in g_gpu[logical_tier].device_id, not the logical tier + * index). 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. Added for + * multi-GPU execution (multi-GPU execution). */ +extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset, + uint64_t bytes, + int expected_device, + void **out_device_ptr) { + if (g_cache_ranges.empty()) return 0; + + auto it = std::upper_bound( + g_cache_ranges.begin(), g_cache_ranges.end(), + source_offset, + [](uint64_t off, const cache_range_entry &e) { + return off < e.source_offset; + }); + while (it != g_cache_ranges.begin()) { + --it; + if (source_offset < it->source_offset) { + /* Should not happen given upper_bound semantics, but defensive. */ + continue; + } + uint64_t into = source_offset - it->source_offset; + if (into > it->bytes) continue; + if (bytes > it->bytes - into) continue; + if (it->device_id != expected_device) continue; + if (out_device_ptr) { + *out_device_ptr = + (char *)it->device_ptr + (source_offset - it->source_offset); + } + return 1; + } + return 0; +} + +extern "C" int ds4_gpu_set_model_fd(int fd) { + g_model_fd = fd; + g_model_fd_host_base = g_model_host_base; + g_model_file_size = 0; + if (g_model_direct_fd >= 0) { + (void)close(g_model_direct_fd); + g_model_direct_fd = -1; + } + g_model_direct_align = 1; + if (fd >= 0) { + struct stat st; + if (fstat(fd, &st) == 0 && st.st_size > 0) { + g_model_file_size = (uint64_t)st.st_size; + if (st.st_blksize > 1) g_model_direct_align = (uint64_t)st.st_blksize; + } +#if defined(__linux__) && defined(O_DIRECT) + if (getenv("DS4_CUDA_NO_DIRECT_IO") == NULL) { + char proc_path[64]; + snprintf(proc_path, sizeof(proc_path), "/proc/self/fd/%d", fd); + int direct_fd = open(proc_path, O_RDONLY | O_DIRECT); + if (direct_fd >= 0) { + g_model_direct_fd = direct_fd; + if (g_model_direct_align < 512) g_model_direct_align = 512; + if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA model direct I/O enabled (align=%llu)\n", + (unsigned long long)g_model_direct_align); + } + } else if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { + fprintf(stderr, "ds4: CUDA model direct I/O unavailable: %s\n", strerror(errno)); + } + } +#endif + } + return 1; +} + +extern "C" int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label) { + if (!model_map || bytes == 0) return 1; + if (offset > model_size || bytes > model_size - offset) return 0; + if (!cuda_model_range_ptr(model_map, offset, bytes, label ? label : "model_tensor")) return 0; + return cuda_model_range_is_cached(model_map, offset, bytes); +} + +extern "C" int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label) { + if (!model_map || bytes == 0) return 1; + if (offset > model_size || bytes > model_size - offset) return 0; + static int optional_q8_preload_disabled = 0; + if (optional_q8_preload_disabled) return 1; + const char *cache_label = label ? label : "q8_0"; + /* Preload runs before any multi-tier dispatch. The cache entries it creates + * are device-0 by construction; multi-tier callers in kernel wrappers will + * miss the linear scan (device_id filter) and allocate fresh per-device + * copies the first time they're consulted. */ + if (getenv("DS4_CUDA_Q8_F32_PRELOAD") != NULL && + cuda_q8_f32_cache_allowed(cache_label, in_dim, out_dim)) { + if (cuda_q8_f32_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1; + optional_q8_preload_disabled = 1; + return 1; + } + if (!cuda_q8_f16_preload_allowed(cache_label, in_dim, out_dim)) return 1; + if (cuda_q8_f16_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1; + optional_q8_preload_disabled = 1; + return 1; +} + +extern "C" void ds4_gpu_print_memory_report(const char *label) { + size_t free_b = 0, total_b = 0; + (void)cudaMemGetInfo(&free_b, &total_b); + fprintf(stderr, "ds4: CUDA memory report %s: free %.2f MiB total %.2f MiB\n", + label ? label : "", (double)free_b / 1048576.0, (double)total_b / 1048576.0); +} + +extern "C" void ds4_gpu_set_quality(bool quality) { + g_quality_mode = quality ? 1 : 0; + const cublasMath_t math_mode = + (g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL) + ? CUBLAS_DEFAULT_MATH + : CUBLAS_TF32_TENSOR_OP_MATH; + /* Walk every initialized per-tier handle. Single-tier (g_n_gpus == 1) + * walks exactly one entry. On any device-switch failure, + * skip the tier and continue — the function is void and the math-mode + * setting is advisory, but log so misconfiguration is visible. */ + for (int i = 0; i < g_n_gpus; i++) { + if (!g_gpu[i].cublas_ready || !g_gpu[i].cublas) continue; + int prev = -1; + cudaError_t derr = cudaGetDevice(&prev); + if (derr != cudaSuccess) { + fprintf(stderr, + "ds4: ds4_gpu_set_quality: cudaGetDevice failed before tier %d " + "(dev=%d): %s; skipping\n", + i, g_gpu[i].device_id, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + continue; + } + derr = cudaSetDevice(g_gpu[i].device_id); + if (derr != cudaSuccess) { + fprintf(stderr, + "ds4: ds4_gpu_set_quality: cudaSetDevice(%d) failed for tier %d: " + "%s; skipping\n", + g_gpu[i].device_id, i, cudaGetErrorString(derr)); + (void)cudaGetLastError(); + if (prev >= 0) (void)cudaSetDevice(prev); + continue; + } + cublasStatus_t st = cublasSetMathMode((cublasHandle_t)g_gpu[i].cublas, math_mode); + if (st != CUBLAS_STATUS_SUCCESS) { + fprintf(stderr, + "ds4: ds4_gpu_set_quality: cublasSetMathMode failed on tier %d " + "(dev=%d): status %d\n", + i, g_gpu[i].device_id, (int)st); + } + if (prev >= 0) (void)cudaSetDevice(prev); + } +} diff --git a/cuda/runtime_services.inc b/cuda/runtime_services.inc new file mode 100644 index 0000000000..f70b35afaa --- /dev/null +++ b/cuda/runtime_services.inc @@ -0,0 +1,308 @@ +/* --gpu-vram auto probe. Defined here (in the .cu unit) so the + * C-side parser (ds4_gpu_args.c) does not need to include + * . Returns 0 on success, nonzero on error + * (errbuf populated). See ds4_gpu_args.h. + * + * Side-effect-light: changes cudaSetDevice during probing; callers + * that care about the active device should reset it themselves + * before continuing. (The mgpu init path resets it anyway.) */ +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; + cudaError_t rc = cudaGetDeviceCount(&visible); + if (rc != cudaSuccess || visible <= 0) { + if (errbuf && errbuflen) { + snprintf(errbuf, errbuflen, + "cudaGetDeviceCount failed: %s", + rc == cudaSuccess ? "no devices" : cudaGetErrorString(rc)); + } + return 1; + } + /* Build the device list: either the explicit filter or 0..visible-1. */ + int devs[DS4_MAX_GPUS]; + int n_dev = 0; + if (device_filter && filter_len > 0) { + if (filter_len > DS4_MAX_GPUS) { + if (errbuf && errbuflen) { + snprintf(errbuf, errbuflen, + "--gpu-devices filter has %d entries (max %d)", + filter_len, DS4_MAX_GPUS); + } + return 1; + } + for (int i = 0; i < filter_len; i++) { + int d = device_filter[i]; + if (d < 0 || d >= visible) { + if (errbuf && errbuflen) { + snprintf(errbuf, errbuflen, + "--gpu-devices: device %d not in 0..%d", + d, visible - 1); + } + return 1; + } + devs[n_dev++] = d; + } + } else { + int cap = visible < DS4_MAX_GPUS ? visible : DS4_MAX_GPUS; + for (int i = 0; i < cap; i++) devs[n_dev++] = i; + } + out->n_gpus = n_dev; + out->safety_margin_bytes = safety_margin_bytes; + for (int i = 0; i < n_dev; i++) { + int d = devs[i]; + rc = cudaSetDevice(d); + if (rc != cudaSuccess) { + if (errbuf && errbuflen) { + snprintf(errbuf, errbuflen, + "cudaSetDevice(%d) failed: %s", + d, cudaGetErrorString(rc)); + } + return 1; + } + size_t free_b = 0, total_b = 0; + rc = cudaMemGetInfo(&free_b, &total_b); + if (rc != cudaSuccess) { + if (errbuf && errbuflen) { + snprintf(errbuf, errbuflen, + "cudaMemGetInfo on device %d failed: %s", + d, cudaGetErrorString(rc)); + } + return 1; + } + /* Auto-mode reserve. Auto-probe is the only place we override + * the user's stated budget, so this is where the conservative- + * on-the-user's-behalf reserve belongs. The + * engine path (engine_classify_multi_tier) still subtracts the + * user-supplied safety_margin_bytes + the cuBLAS workspace from + * whatever budget we hand back; that math is unchanged and + * applies on top of the reserve we trim here. + * + * Reserve = max(2 GiB, 5 % of free). Why these numbers: + * - 2 GiB floor covers runtime scratch / Q8 dequant caches / + * MTP optional state on small GPUs (8-12 GB cards) where + * 5 % is < 1 GiB and not enough headroom. + * - 5 % of free scales the reserve up on larger cards where + * workspace + KV growth needs proportionally more room. + * Explicit --gpu-vram 47,37 budgets do not go through this + * probe and are unaffected. */ + const size_t reserve_floor = (size_t)2ull * 1024ull * 1024ull * 1024ull; + const size_t reserve_pct = free_b / 20u; + const size_t reserve = reserve_floor > reserve_pct ? reserve_floor : reserve_pct; + const size_t budget = free_b > reserve ? (free_b - reserve) : 0; + (void)safety_margin_bytes; + out->device_indices[i] = d; + out->vram_bytes[i] = budget; + } + return 0; +} + +typedef struct ds4_gpu_stream_expert_table { + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; +} ds4_gpu_stream_expert_table; + +static int cuda_stream_selected_ensure_bytes( + char **ptr, uint64_t *capacity, uint64_t bytes, const char *label) { + if (*ptr && *capacity >= bytes) return 1; + if (*ptr) { + (void)cudaFree(*ptr); + *ptr = NULL; + *capacity = 0; + } + if (bytes == 0 || bytes > (uint64_t)SIZE_MAX) return 0; + cudaError_t err = cudaMalloc((void **)ptr, (size_t)bytes); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: CUDA streaming %s allocation failed for %.2f MiB: %s\n", + label, (double)bytes / 1048576.0, cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + *capacity = bytes; + return 1; +} + +static int cuda_stream_selected_ensure_i32(uint64_t count) { + if (count == 0 || count > UINT64_MAX / sizeof(int32_t)) return 0; + const uint64_t bytes = count * sizeof(int32_t); + return cuda_stream_selected_ensure_bytes( + (char **)&g_stream_selected_cache.slot_selected_ptr, + &g_stream_selected_cache.slot_selected_capacity, + bytes, + "selected-id remap"); +} + +static int cuda_stream_selected_ranges_valid( + const ds4_gpu_stream_expert_table *table) { + if (!table || !table->model_map || table->model_size == 0 || + table->n_total_expert == 0 || table->gate_expert_bytes == 0 || + table->down_expert_bytes == 0) { + return 0; + } + if ((uint64_t)table->n_total_expert > + UINT64_MAX / table->gate_expert_bytes || + (uint64_t)table->n_total_expert > + UINT64_MAX / table->down_expert_bytes) { + return 0; + } + const uint64_t gate_bytes = + (uint64_t)table->n_total_expert * table->gate_expert_bytes; + const uint64_t down_bytes = + (uint64_t)table->n_total_expert * table->down_expert_bytes; + return table->gate_offset <= table->model_size && + gate_bytes <= table->model_size - table->gate_offset && + table->up_offset <= table->model_size && + gate_bytes <= table->model_size - table->up_offset && + table->down_offset <= table->model_size && + down_bytes <= table->model_size - table->down_offset; +} + +static int cuda_stream_selected_cache_begin_load( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t slot_count) { + cuda_stream_selected_cache_invalidate(); + if (!g_ssd_streaming_mode) return 1; + if (!cuda_stream_selected_ranges_valid(table) || !selected_ids || + slot_count == 0) { + return 0; + } + if (g_n_gpus != 1) { + fprintf(stderr, + "ds4: CUDA SSD streaming requires single-GPU placement\n"); + return 0; + } + + std::vector expert_to_slot; + std::vector compact_ids; + std::vector slot_ids; + try { + expert_to_slot.assign(table->n_total_expert, -1); + compact_ids.reserve(slot_count < table->n_total_expert ? + slot_count : table->n_total_expert); + slot_ids.resize(slot_count); + } catch (...) { + return 0; + } + for (uint32_t i = 0; i < slot_count; i++) { + const int32_t expert = selected_ids[i]; + if (expert < 0 || (uint32_t)expert >= table->n_total_expert) { + fprintf(stderr, + "ds4: CUDA streaming expert id %d is outside 0..%u at layer %u\n", + expert, table->n_total_expert, table->layer); + return 0; + } + int32_t compact = expert_to_slot[(uint32_t)expert]; + if (compact < 0) { + compact = (int32_t)compact_ids.size(); + expert_to_slot[(uint32_t)expert] = compact; + compact_ids.push_back(expert); + } + slot_ids[i] = compact; + } + if (compact_ids.empty() || compact_ids.size() > UINT32_MAX) return 0; + const uint64_t compact_count = compact_ids.size(); + if (compact_count > UINT64_MAX / table->gate_expert_bytes || + compact_count > UINT64_MAX / table->down_expert_bytes) { + return 0; + } + const uint64_t gate_bytes = compact_count * table->gate_expert_bytes; + const uint64_t down_bytes = compact_count * table->down_expert_bytes; + const int logical_tier = 0; + if (g_stream_selected_cache.logical_tier != logical_tier && + (g_stream_selected_cache.gate_ptr || + g_stream_selected_cache.up_ptr || + g_stream_selected_cache.down_ptr || + g_stream_selected_cache.slot_selected_ptr)) { + cuda_stream_selected_cache_release(); + } + if (ds4_gpu_set_current_device(logical_tier) != 0 || + !cuda_stream_selected_ensure_bytes( + &g_stream_selected_cache.gate_ptr, + &g_stream_selected_cache.gate_capacity, + gate_bytes, "gate experts") || + !cuda_stream_selected_ensure_bytes( + &g_stream_selected_cache.up_ptr, + &g_stream_selected_cache.up_capacity, + gate_bytes, "up experts") || + !cuda_stream_selected_ensure_bytes( + &g_stream_selected_cache.down_ptr, + &g_stream_selected_cache.down_capacity, + down_bytes, "down experts") || + !cuda_stream_selected_ensure_i32(slot_count)) { + cuda_stream_selected_cache_invalidate(); + return 0; + } + + for (uint32_t i = 0; i < compact_ids.size(); i++) { + const uint64_t expert = (uint32_t)compact_ids[i]; + const uint64_t gate_src = + table->gate_offset + expert * table->gate_expert_bytes; + const uint64_t up_src = + table->up_offset + expert * table->gate_expert_bytes; + const uint64_t down_src = + table->down_offset + expert * table->down_expert_bytes; + const uint64_t gate_dst = (uint64_t)i * table->gate_expert_bytes; + const uint64_t down_dst = (uint64_t)i * table->down_expert_bytes; + if (!cuda_model_copy_to_device_streamed( + g_stream_selected_cache.gate_ptr + gate_dst, + table->model_map, table->model_size, + gate_src, table->gate_expert_bytes, + "stream gate expert copy") || + !cuda_model_copy_to_device_streamed( + g_stream_selected_cache.up_ptr + gate_dst, + table->model_map, table->model_size, + up_src, table->gate_expert_bytes, + "stream up expert copy") || + !cuda_model_copy_to_device_streamed( + g_stream_selected_cache.down_ptr + down_dst, + table->model_map, table->model_size, + down_src, table->down_expert_bytes, + "stream down expert copy")) { + cuda_stream_selected_cache_invalidate(); + return 0; + } + } + if (!cuda_ok(cudaMemcpy(g_stream_selected_cache.slot_selected_ptr, + slot_ids.data(), + (size_t)slot_count * sizeof(int32_t), + cudaMemcpyHostToDevice), + "stream selected-id remap copy")) { + cuda_stream_selected_cache_invalidate(); + return 0; + } + + g_stream_selected_cache.logical_tier = logical_tier; + g_stream_selected_cache.model_map = table->model_map; + g_stream_selected_cache.layer = table->layer; + g_stream_selected_cache.n_total_expert = table->n_total_expert; + g_stream_selected_cache.slot_count = slot_count; + g_stream_selected_cache.compact_count = (uint32_t)compact_count; + g_stream_selected_cache.gate_offset = table->gate_offset; + g_stream_selected_cache.up_offset = table->up_offset; + g_stream_selected_cache.down_offset = table->down_offset; + g_stream_selected_cache.gate_expert_bytes = table->gate_expert_bytes; + g_stream_selected_cache.down_expert_bytes = table->down_expert_bytes; + g_stream_selected_cache.slot_selected_tensor.ptr = + g_stream_selected_cache.slot_selected_ptr; + g_stream_selected_cache.slot_selected_tensor.bytes = + (uint64_t)slot_count * sizeof(int32_t); + g_stream_selected_cache.slot_selected_tensor.owner = 0; + g_stream_selected_cache.slot_selected_tensor.device_id = logical_tier; + g_stream_selected_cache.valid = 1; + return 1; +} diff --git a/ds4.c b/ds4.c index 2496319ace..55f484d2c9 100644 --- a/ds4.c +++ b/ds4.c @@ -1,12 +1,12 @@ /* ========================================================================= - * ds4.c - DeepSeek V4 inference engine. + * ds4.c - DwarfStar inference engine core. * ========================================================================= * - * This file is deliberately vertical: it owns GGUF loading, the fixed - * DeepSeek V4 tensor layouts, CPU reference kernels, the whole-model Metal - * graph driver, and tokenizer wiring. Model shape selection is intentionally - * narrow: validation accepts the known Flash and Pro layouts and fails early - * for anything else. + * This translation unit owns shared GGUF loading, engine/session state, + * tokenizer wiring, and public API dispatch. Concrete model pipelines and CPU + * kernels are grouped into implementation fragments under models/ and + * kernels/. They are included here to preserve private static linkage and + * whole-program optimization without introducing an operator abstraction. * * Loading is mmap based. The loader parses only the GGUF header, metadata * table, and tensor directory. Tensor data stays in the kernel page cache @@ -41,6 +41,10 @@ #include #include "ds4.h" +#include "ds4_model_provider.h" +#include "ds4_model_provider_builtin.h" +#include "models/deepseek/provider.h" +#include "models/glm/provider.h" #include "ds4_distributed.h" #include "ds4_tp.h" @@ -3077,1169 +3081,551 @@ static void model_warm_weights(const ds4_model *m) { t1 - t0, (unsigned long long)checksum); } +#include "kernels/cpu_quant.inc" + /* ========================================================================= - * Scalar Conversion and Quantized Tensor Kernels. + * Fixed Weight Binding and Model Validation. * ========================================================================= * - * These functions are the CPU reference math used by the C backend and by - * Metal diagnostics. They implement only the tensor formats present in the - * DeepSeek V4 Flash GGUF: F16, F32, Q8_0, Q2_K, IQ2_XXS, and Q8_K activation - * blocks used for expert dot products. + * The GGUF tensor directory is converted into a DS4-specific pointer table. + * After this section, the rest of the program addresses tensors by semantic + * fields such as layer->attn_q_a or layer->ffn_gate_exps rather than by string + * lookup. Shape validation is intentionally strict. */ -static inline float f16_to_f32(uint16_t h) { -#if defined(__ARM_NEON) - const float16x4_t hv = vreinterpret_f16_u16(vdup_n_u16(h)); - return vgetq_lane_f32(vcvt_f32_f16(hv), 0); -#else - uint32_t sign = (uint32_t)(h & 0x8000) << 16; - uint32_t exp = (h >> 10) & 0x1f; - uint32_t mant = h & 0x03ff; - uint32_t bits; - - if (exp == 0) { - if (mant == 0) { - bits = sign; - } else { - exp = 1; - while ((mant & 0x0400) == 0) { - mant <<= 1; - exp--; - } - mant &= 0x03ff; - bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); - } - } else if (exp == 31) { - bits = sign | 0x7f800000u | (mant << 13); - } else { - bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); +static uint32_t required_u32(const ds4_model *m, const char *key) { + uint32_t v = 0; + if (!model_get_u32(m, key, &v)) { + fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); + exit(1); } - - float f; - memcpy(&f, &bits, sizeof(f)); - return f; -#endif + return v; } -static inline uint16_t f32_to_f16(float f) { -#if defined(__ARM_NEON) - const float32x4_t fv = vdupq_n_f32(f); - const float16x4_t hv = vcvt_f16_f32(fv); - return vget_lane_u16(vreinterpret_u16_f16(hv), 0); -#else - uint32_t bits; - memcpy(&bits, &f, sizeof(bits)); - - const uint32_t sign = (bits >> 16) & 0x8000u; - int32_t exp = (int32_t)((bits >> 23) & 0xffu) - 127 + 15; - uint32_t mant = bits & 0x7fffffu; - - if (exp <= 0) { - if (exp < -10) return (uint16_t)sign; - mant |= 0x800000u; - const uint32_t shift = (uint32_t)(14 - exp); - uint32_t half_mant = mant >> shift; - const uint32_t round_bit = (mant >> (shift - 1)) & 1u; - const uint32_t sticky = mant & ((1u << (shift - 1)) - 1u); - if (round_bit && (sticky || (half_mant & 1u))) half_mant++; - return (uint16_t)(sign | half_mant); +static uint64_t required_u64_compat(const ds4_model *m, const char *key) { + uint64_t v = 0; + if (!model_get_u64_compat(m, key, &v)) { + fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); + exit(1); } + return v; +} - if (exp >= 31) { - if (((bits >> 23) & 0xffu) == 0xffu && mant != 0) { - return (uint16_t)(sign | 0x7e00u); - } - return (uint16_t)(sign | 0x7c00u); +static float required_f32(const ds4_model *m, const char *key) { + float v = 0.0f; + if (!model_get_f32_compat(m, key, &v)) { + fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); + exit(1); } + return v; +} - uint32_t half = sign | ((uint32_t)exp << 10) | (mant >> 13); - const uint32_t round = mant & 0x1fffu; - if (round > 0x1000u || (round == 0x1000u && (half & 1u))) half++; - return (uint16_t)half; -#endif +static bool required_bool(const ds4_model *m, const char *key) { + bool v = false; + if (!model_get_bool(m, key, &v)) { + fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); + exit(1); + } + return v; } -static void f16_round_inplace_cpu(float *x, uint32_t n) { - for (uint32_t i = 0; i < n; i++) x[i] = f16_to_f32(f32_to_f16(x[i])); +static ds4_tensor *required_tensor(const ds4_model *m, const char *name) { + ds4_tensor *t = model_find_tensor(m, name); + if (!t) { + fprintf(stderr, "ds4: required tensor is missing: %s\n", name); + exit(1); + } + return t; } -static float dsv4_e4m3fn_value_cpu(int i) { - static const float exp_scale[16] = { - 0.0f, 0.015625f, 0.03125f, 0.0625f, - 0.125f, 0.25f, 0.5f, 1.0f, - 2.0f, 4.0f, 8.0f, 16.0f, - 32.0f, 64.0f, 128.0f, 256.0f, - }; +static ds4_tensor *tensor_by_namef(const ds4_model *m, const char *fmt, uint32_t layer) { + char name[128]; + int n = snprintf(name, sizeof(name), fmt, layer); + if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); + return model_find_tensor(m, name); +} - const int exp = (i >> 3) & 0x0f; - const int mant = i & 0x07; - return exp == 0 - ? (float)mant * 0.001953125f - : (1.0f + (float)mant * 0.125f) * exp_scale[exp]; +static ds4_tensor *required_tensorf(const ds4_model *m, const char *fmt, uint32_t layer) { + char name[128]; + int n = snprintf(name, sizeof(name), fmt, layer); + if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); + return required_tensor(m, name); } -static float dsv4_e4m3fn_dequant_cpu(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float ax = fminf(fabsf(x), 448.0f); +static ds4_tensor *tensor_by_mtp_stage_suffix( + const ds4_model *m, + uint32_t stage, + const char *suffix) { + char name[160]; + int n = snprintf(name, sizeof(name), "mtp.%u.%s", stage, suffix); + if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); + return model_find_tensor(m, name); +} - int lo = 0; - int hi = 126; - while (lo < hi) { - const int mid = (lo + hi + 1) >> 1; - if (dsv4_e4m3fn_value_cpu(mid) <= ax) { - lo = mid; - } else { - hi = mid - 1; - } +static void tensor_expect_layout( + const ds4_tensor *t, + uint32_t type, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!t) ds4_die("internal error: missing tensor while validating layout"); + if (t->type != type) { + fprintf(stderr, + "ds4: tensor %.*s has type %s, expected %s\n", + (int)t->name.len, + t->name.ptr, + tensor_type_name(t->type), + tensor_type_name(type)); + exit(1); } - - int best = lo; - if (best < 126) { - const float best_diff = fabsf(ax - dsv4_e4m3fn_value_cpu(best)); - const float next_diff = fabsf(ax - dsv4_e4m3fn_value_cpu(best + 1)); - if (next_diff < best_diff || (next_diff == best_diff && ((best + 1) & 1) == 0 && (best & 1) != 0)) { - best++; - } + if (t->ndim != ndim) { + fprintf(stderr, + "ds4: tensor %.*s has %u dimensions, expected %u\n", + (int)t->name.len, + t->name.ptr, + t->ndim, + ndim); + exit(1); } - return sign * dsv4_e4m3fn_value_cpu(best); -} - -/* DeepSeek V4 stores the non-RoPE part of compressed KV through an E4M3-style - * round trip. Keeping this in the CPU reference makes cache values comparable - * to the Metal graph's compressed-cache behavior. */ -static void dsv4_fp8_kv_quantize_row_inplace_cpu(float *x, uint32_t head_dim, uint32_t n_rot) { - const uint32_t n_nope = head_dim - n_rot; - for (uint32_t off = 0; off < n_nope; off += 64) { - float amax = 0.0f; - for (uint32_t i = 0; i < 64; i++) { - const float av = fabsf(x[off + i]); - if (av > amax) amax = av; - } - - if (amax < 1.0e-4f) amax = 1.0e-4f; - const float scale = ldexpf(1.0f, (int)ceilf(log2f(amax / 448.0f))); - for (uint32_t i = 0; i < 64; i++) { - float v = x[off + i] / scale; - if (v > 448.0f) v = 448.0f; - if (v < -448.0f) v = -448.0f; - x[off + i] = dsv4_e4m3fn_dequant_cpu(v) * scale; - } + const uint64_t want[3] = { d0, d1, d2 }; + for (uint32_t i = 0; i < ndim; i++) { + if (t->dim[i] == want[i]) continue; + fprintf(stderr, + "ds4: tensor %.*s has dim[%u]=%" PRIu64 ", expected %" PRIu64 "\n", + (int)t->name.len, + t->name.ptr, + i, + t->dim[i], + want[i]); + exit(1); } } -static float dsv4_e2m1fn_value_cpu(int i) { - static const float values[8] = { - 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, - }; - return values[i & 7]; +static bool tensor_type_is_glm_dense_quant(uint32_t type) { + return type == DS4_TENSOR_Q8_0 || + type == DS4_TENSOR_Q4_K || + type == DS4_TENSOR_Q4_0; } -static float dsv4_e2m1fn_dequant_cpu(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float ax = fminf(fabsf(x), 6.0f); - int best = 0; - float best_diff = fabsf(ax - dsv4_e2m1fn_value_cpu(0)); - for (int i = 1; i < 8; i++) { - const float diff = fabsf(ax - dsv4_e2m1fn_value_cpu(i)); - if (diff < best_diff || (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { - best = i; - best_diff = diff; - } - } - return sign * dsv4_e2m1fn_value_cpu(best); +static bool tensor_type_is_dense_quant(uint32_t type) { + return type == DS4_TENSOR_Q8_0 || + type == DS4_TENSOR_Q4_K || + type == DS4_TENSOR_Q4_0; } -static void dsv4_hadamard128_inplace_cpu(float *x) { - for (uint32_t stride = 1; stride < 128; stride <<= 1) { - for (uint32_t base = 0; base < 128; base += 2u * stride) { - for (uint32_t i = 0; i < stride; i++) { - const float a = x[base + i]; - const float b = x[base + stride + i]; - x[base + i] = a + b; - x[base + stride + i] = a - b; - } - } +static void tensor_expect_glm_dense_quant_layout( + const ds4_tensor *t, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!t) ds4_die("internal error: missing tensor while validating GLM dense layout"); + if (!tensor_type_is_glm_dense_quant(t->type)) { + fprintf(stderr, + "ds4: tensor %.*s has type %s, expected q8_0, q4_K, or q4_0\n", + (int)t->name.len, + t->name.ptr, + tensor_type_name(t->type)); + exit(1); } - const float scale = 0.08838834764831845f; - for (uint32_t i = 0; i < 128; i++) x[i] *= scale; + tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } -static void dsv4_fp4_act_quantize_row_inplace_cpu(float *x, uint32_t n) { - if ((n % 32u) != 0) ds4_die("DSV4 FP4 activation quantization requires 32-aligned rows"); - for (uint32_t off = 0; off < n; off += 32) { - float amax = 0.0f; - for (uint32_t i = 0; i < 32; i++) { - const float av = fabsf(x[off + i]); - if (av > amax) amax = av; - } - - if (amax < 7.052966104933725e-38f) amax = 7.052966104933725e-38f; - const float scale = ldexpf(1.0f, (int)ceilf(log2f(amax / 6.0f))); - for (uint32_t i = 0; i < 32; i++) { - float v = x[off + i] / scale; - if (v > 6.0f) v = 6.0f; - if (v < -6.0f) v = -6.0f; - x[off + i] = dsv4_e2m1fn_dequant_cpu(v) * scale; - } +static void tensor_expect_dense_quant_layout( + const ds4_tensor *t, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!t) ds4_die("internal error: missing tensor while validating dense quant layout"); + if (!tensor_type_is_dense_quant(t->type)) { + fprintf(stderr, + "ds4: tensor %.*s has type %s, expected q8_0, q4_K, or q4_0\n", + (int)t->name.len, + t->name.ptr, + tensor_type_name(t->type)); + exit(1); } + tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } -/* The official DeepSeek V4 graph rotates indexer activations with a 128-wide - * Hadamard transform and immediately runs the FP4 activation-simulation - * round trip. This applies to both indexer Q and the indexer compressor KV; - * without it, the top-k compressed-row selection is not the model's graph. */ -static void dsv4_indexer_qat_row_inplace_cpu(float *x, uint32_t head_dim) { - if (head_dim != 128) ds4_die("DSV4 indexer QAT expects 128-wide indexer rows"); - dsv4_hadamard128_inplace_cpu(x); - dsv4_fp4_act_quantize_row_inplace_cpu(x, head_dim); +static void tensor_expect_optional( + const ds4_tensor *t, + uint32_t type, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (t) tensor_expect_layout(t, type, ndim, d0, d1, d2); } -static void dsv4_indexer_qat_rows_inplace_cpu(float *x, uint32_t rows, uint32_t head_dim) { - for (uint32_t r = 0; r < rows; r++) { - dsv4_indexer_qat_row_inplace_cpu(x + (uint64_t)r * head_dim, head_dim); +static void tensor_expect_plain_layout( + const ds4_tensor *t, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!t) ds4_die("internal error: missing tensor while validating layout"); + if (t->type != DS4_TENSOR_F16 && t->type != DS4_TENSOR_F32) { + fprintf(stderr, + "ds4: tensor %.*s has type %s, expected F16 or F32\n", + (int)t->name.len, + t->name.ptr, + tensor_type_name(t->type)); + exit(1); } + tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } -/* Quantize a float activation into Q8_K blocks so GGUF Q2_K/IQ2_XXS expert - * kernels can reuse the same activation for many expert rows. */ -static void ds4_quantize_row_q8_K(const float *x, block_q8_K *y, int64_t k) { - if (k % QK_K != 0) ds4_die("Q8_K quantization length is not QK_K aligned"); - const int64_t nb = k / QK_K; - - for (int64_t b = 0; b < nb; b++) { - float max = 0.0f; - float amax = 0.0f; - for (int j = 0; j < QK_K; j++) { - const float ax = fabsf(x[j]); - if (ax > amax) { - amax = ax; - max = x[j]; - } - } - - if (amax == 0.0f) { - y[b].d = 0.0f; - memset(y[b].qs, 0, sizeof(y[b].qs)); - memset(y[b].bsums, 0, sizeof(y[b].bsums)); - x += QK_K; - continue; - } - - const float iscale = -127.0f / max; - for (int j = 0; j < QK_K; j++) { - int v = (int)lrintf(iscale * x[j]); - if (v > 127) v = 127; - if (v < -128) v = -128; - y[b].qs[j] = (int8_t)v; - } - for (int j = 0; j < QK_K / 16; j++) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += y[b].qs[j * 16 + i]; - y[b].bsums[j] = (int16_t)sum; - } - y[b].d = 1.0f / iscale; - x += QK_K; - } +static bool tensor_type_is_f16_or_q8_0(uint32_t type) { + return type == DS4_TENSOR_F16 || type == DS4_TENSOR_Q8_0; } -static void ds4_vec_dot_q2_K_q8_K(int n, float *s, const block_q2_K *x, const block_q8_K *y) { - const int nb = n / QK_K; - -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - const uint8x16_t m3 = vdupq_n_u8(0x03); - const uint8x16_t m4 = vdupq_n_u8(0x0f); - const int32x4_t zero = vdupq_n_s32(0); - float sum = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = y[i].d * f16_to_f32(x[i].d); - const float dmin = -y[i].d * f16_to_f32(x[i].dmin); - - const uint8_t *q2 = x[i].qs; - const int8_t *q8 = y[i].qs; - const uint8_t *sc = x[i].scales; - - const uint8x16_t mins_and_scales = vld1q_u8(sc); - const uint8x16_t scales = vandq_u8(mins_and_scales, m4); - uint8_t scale_lanes[16]; - vst1q_u8(scale_lanes, scales); - - const uint8x16_t mins = vshrq_n_u8(mins_and_scales, 4); - const int16x8x2_t q8sums = vld1q_s16_x2(y[i].bsums); - const int16x8x2_t mins16 = {{ - vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(mins))), - vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(mins))), - }}; - const int32x4_t s0 = vaddq_s32( - vmull_s16(vget_low_s16(mins16.val[0]), vget_low_s16(q8sums.val[0])), - vmull_s16(vget_high_s16(mins16.val[0]), vget_high_s16(q8sums.val[0]))); - const int32x4_t s1 = vaddq_s32( - vmull_s16(vget_low_s16(mins16.val[1]), vget_low_s16(q8sums.val[1])), - vmull_s16(vget_high_s16(mins16.val[1]), vget_high_s16(q8sums.val[1]))); - sum += dmin * (float)vaddvq_s32(vaddq_s32(s0, s1)); - - int isum = 0; - int is = 0; - for (int j = 0; j < QK_K / 128; j++) { - const uint8x16x2_t q2bits = vld1q_u8_x2(q2); - q2 += 32; - -#define DS4_Q2_DOT_NOSHIFT(scale_index) do { \ - const int8x16x2_t q8bytes = vld1q_s8_x2(q8); \ - q8 += 32; \ - const int8x16_t q2lo = vreinterpretq_s8_u8(vandq_u8(q2bits.val[0], m3));\ - const int8x16_t q2hi = vreinterpretq_s8_u8(vandq_u8(q2bits.val[1], m3));\ - isum += vaddvq_s32(vdotq_s32(zero, q2lo, q8bytes.val[0])) * \ - scale_lanes[is + (scale_index)]; \ - isum += vaddvq_s32(vdotq_s32(zero, q2hi, q8bytes.val[1])) * \ - scale_lanes[is + 1 + (scale_index)]; \ - } while (0) - -#define DS4_Q2_DOT_SHIFT(shift, scale_index) do { \ - const int8x16x2_t q8bytes = vld1q_s8_x2(q8); \ - q8 += 32; \ - const int8x16_t q2lo = vreinterpretq_s8_u8( \ - vandq_u8(vshrq_n_u8(q2bits.val[0], (shift)), m3)); \ - const int8x16_t q2hi = vreinterpretq_s8_u8( \ - vandq_u8(vshrq_n_u8(q2bits.val[1], (shift)), m3)); \ - isum += vaddvq_s32(vdotq_s32(zero, q2lo, q8bytes.val[0])) * \ - scale_lanes[is + (scale_index)]; \ - isum += vaddvq_s32(vdotq_s32(zero, q2hi, q8bytes.val[1])) * \ - scale_lanes[is + 1 + (scale_index)]; \ - } while (0) - - DS4_Q2_DOT_NOSHIFT(0); - DS4_Q2_DOT_SHIFT(2, 2); - DS4_Q2_DOT_SHIFT(4, 4); - DS4_Q2_DOT_SHIFT(6, 6); - is += 8; - -#undef DS4_Q2_DOT_NOSHIFT -#undef DS4_Q2_DOT_SHIFT - } - - sum += d * (float)isum; - } - - *s = sum; -#else - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const uint8_t *q2 = x[i].qs; - const int8_t *q8 = y[i].qs; - const uint8_t *sc = x[i].scales; - - int summs = 0; - for (int j = 0; j < 16; j++) { - summs += y[i].bsums[j] * (sc[j] >> 4); - } - - const float dall = y[i].d * f16_to_f32(x[i].d); - const float dmin = y[i].d * f16_to_f32(x[i].dmin); - - int isum = 0; - int is = 0; - for (int k = 0; k < QK_K / 128; k++) { - int shift = 0; - for (int j = 0; j < 4; j++) { - int d = sc[is++] & 0x0f; - int isuml = dot_q2_16(q2, q8, shift); - isum += d * isuml; - - d = sc[is++] & 0x0f; - isuml = dot_q2_16(q2 + 16, q8 + 16, shift); - isum += d * isuml; - - shift += 2; - q8 += 32; - } - q2 += 32; - } - sumf += dall * (float)isum - dmin * (float)summs; +static void tensor_expect_f16_or_q8_0_layout( + const ds4_tensor *t, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!t) ds4_die("internal error: missing tensor while validating layout"); + if (!tensor_type_is_f16_or_q8_0(t->type)) { + fprintf(stderr, + "ds4: tensor %.*s has type %s, expected f16 or q8_0\n", + (int)t->name.len, + t->name.ptr, + tensor_type_name(t->type)); + exit(1); } - *s = sumf; -#endif -} - -static inline float q2_k_value_f32(const block_q2_K *blocks, uint32_t k) { - const uint32_t block = k / QK_K; - const uint32_t idx = k - block * QK_K; - const block_q2_K *xb = blocks + block; - const uint32_t group = idx / 16u; - const uint32_t l = idx - group * 16u; - const uint32_t q_base = 32u * (group / 8u) + 16u * (group & 1u); - const uint32_t shift = ((group / 2u) & 3u) * 2u; - const uint32_t q = ((uint32_t)xb->qs[q_base + l] >> shift) & 0x03u; - const uint32_t sc = xb->scales[group]; - return f16_to_f32(xb->d) * (float)(sc & 0x0fu) * (float)q - - f16_to_f32(xb->dmin) * (float)(sc >> 4u); + tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } -static float ds4_vec_dot_q2_K_f32(int n, const block_q2_K *x, const float *y) { - float sum = 0.0f; - for (int k = 0; k < n; k++) { - sum += q2_k_value_f32(x, (uint32_t)k) * y[k]; - } - return sum; +static bool tensor_is_routed_expert_type(uint32_t type) { + return type == DS4_TENSOR_Q8_0 || + type == DS4_TENSOR_IQ2_XXS || + type == DS4_TENSOR_Q2_K || + type == DS4_TENSOR_Q4_K || + type == DS4_TENSOR_Q5_K || + type == DS4_TENSOR_Q6_K; } -static inline void q4_k_get_scale_min(int j, const uint8_t *q, uint8_t *sc, uint8_t *m) { - if (j < 4) { - *sc = q[j] & 63; - *m = q[j + 4] & 63; - } else { - *sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); - *m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); +static DS4_MAYBE_UNUSED uint64_t routed_expert_block_bytes(uint32_t type) { + switch (type) { + case DS4_TENSOR_Q8_0: return 34; + case DS4_TENSOR_IQ2_XXS: return sizeof(block_iq2_xxs); + case DS4_TENSOR_Q2_K: return sizeof(block_q2_K); + case DS4_TENSOR_Q4_K: return sizeof(block_q4_K); + case DS4_TENSOR_Q5_K: return sizeof(block_q5_K); + case DS4_TENSOR_Q6_K: return sizeof(block_q6_K); + default: ds4_die("unsupported routed expert tensor type"); } + return 0; } -static void ds4_vec_dot_q4_K_q8_K(int n, float *s, const block_q4_K *x, const block_q8_K *y) { - const int nb = n / QK_K; - -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - const int32x4_t zero = vdupq_n_s32(0); - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = y[i].d * f16_to_f32(x[i].d); - const float dm = -y[i].d * f16_to_f32(x[i].dmin); - - const uint8_t *qs = x[i].qs; - const uint8_t *sc = x[i].scales; - const int8_t *q8 = y[i].qs; - - int32_t summs = 0; - for (int j = 0; j < QK_K / 32; j++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(j, sc, &sc_val, &m_val); - int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1]; - summs += m_val * gsum; - } - - int isum = 0; - for (int j = 0; j < QK_K / 32; j++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(j, sc, &sc_val, &m_val); - - const int byte_off = (j >> 1) * 32; - const int shift = (j & 1) * 4; - - /* Load 32 q8 values for this group */ - const int8x16x2_t q8v = vld1q_s8_x2(q8 + j * 32); - - /* Unpack 32 q4 values from 32 bytes at qs[byte_off] with shift */ - uint8_t q4_u[32]; - if (shift == 0) { - for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] & 0xF; - } else { - for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] >> 4; - } - - const int8x16_t q4a = vreinterpretq_s8_u8(vld1q_u8(q4_u)); - const int8x16_t q4b = vreinterpretq_s8_u8(vld1q_u8(q4_u + 16)); - - isum += vaddvq_s32(vdotq_s32(zero, q4a, q8v.val[0])) * sc_val; - isum += vaddvq_s32(vdotq_s32(zero, q4b, q8v.val[1])) * sc_val; - } +static DS4_MAYBE_UNUSED uint64_t routed_expert_row_bytes(const ds4_tensor *t) { + const gguf_type_info *info = tensor_type(t->type); + if (!info || info->block_elems == 0) ds4_die("unsupported routed expert tensor type"); + if ((t->dim[0] % info->block_elems) != 0) ds4_die("routed expert row is not quant block aligned"); + return (t->dim[0] / info->block_elems) * routed_expert_block_bytes(t->type); +} - sumf += d * (float)isum + dm * (float)summs; +static bool streaming_layer_routed_expert_bytes( + const ds4_layer_weights *layer, + uint64_t *per_expert_bytes_out) { + if (per_expert_bytes_out) *per_expert_bytes_out = 0; + if (!layer || + !per_expert_bytes_out || + !layer->ffn_gate_exps || + !layer->ffn_up_exps || + !layer->ffn_down_exps) { + return false; } - *s = sumf; -#else - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = y[i].d * f16_to_f32(x[i].d); - const float dm = -y[i].d * f16_to_f32(x[i].dmin); - - const uint8_t *qs = x[i].qs; - const uint8_t *sc = x[i].scales; - const int8_t *q8 = y[i].qs; - - int summs = 0; - for (int j = 0; j < QK_K / 32; j++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(j, sc, &sc_val, &m_val); - int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1]; - summs += m_val * gsum; - } - - int isum = 0; - for (int j = 0; j < QK_K / 32; j++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(j, sc, &sc_val, &m_val); - - const int byte_off = (j >> 1) * 32; - const int shift = (j & 1) * 4; - - for (int l = 0; l < 32; l++) { - isum += ((qs[byte_off + l] >> shift) & 0xF) * (int)q8[j * 32 + l] * sc_val; - } - } - - sumf += d * (float)isum + dm * (float)summs; + const uint64_t gate_row_bytes = + routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t up_row_bytes = + routed_expert_row_bytes(layer->ffn_up_exps); + const uint64_t down_row_bytes = + routed_expert_row_bytes(layer->ffn_down_exps); + if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_up_exps->dim[1] > UINT64_MAX / up_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + return false; } - *s = sumf; -#endif -} - -static void ds4_vec_dot_q5_K_q8_K(int n, float *s, const block_q5_K *x, const block_q8_K *y) { - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = y[i].d * f16_to_f32(x[i].d); - const float dmin = y[i].d * f16_to_f32(x[i].dmin); - const uint8_t *ql = x[i].qs; - const uint8_t *qh = x[i].qh; - const int8_t *q8 = y[i].qs; - const uint8_t *scales = x[i].scales; - - int64_t isum = 0; - int64_t summs = 0; - int is = 0; - uint8_t u1 = 1; - uint8_t u2 = 2; - - for (int j = 0; j < QK_K; j += 64) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(is, scales, &sc_val, &m_val); - summs += (int64_t)m_val * ((int32_t)y[i].bsums[2 * is] + (int32_t)y[i].bsums[2 * is + 1]); - for (int l = 0; l < 32; l++) { - const int q = (int)(ql[l] & 0x0F) + ((qh[l] & u1) ? 16 : 0); - isum += (int64_t)sc_val * q * (int)q8[j + l]; - } - - q4_k_get_scale_min(is + 1, scales, &sc_val, &m_val); - summs += (int64_t)m_val * ((int32_t)y[i].bsums[2 * (is + 1)] + (int32_t)y[i].bsums[2 * (is + 1) + 1]); - for (int l = 0; l < 32; l++) { - const int q = (int)(ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0); - isum += (int64_t)sc_val * q * (int)q8[j + 32 + l]; - } - - ql += 32; - is += 2; - u1 = (uint8_t)(u1 << 2); - u2 = (uint8_t)(u2 << 2); - } - - sumf += d * (float)isum - dmin * (float)summs; + const uint64_t gate_expert_bytes = + layer->ffn_gate_exps->dim[1] * gate_row_bytes; + const uint64_t up_expert_bytes = + layer->ffn_up_exps->dim[1] * up_row_bytes; + const uint64_t down_expert_bytes = + layer->ffn_down_exps->dim[1] * down_row_bytes; + if (gate_expert_bytes > UINT64_MAX - up_expert_bytes || + gate_expert_bytes + up_expert_bytes > + UINT64_MAX - down_expert_bytes) { + return false; } - *s = sumf; + const uint64_t per_expert_bytes = + gate_expert_bytes + up_expert_bytes + down_expert_bytes; + if (per_expert_bytes == 0) return false; + *per_expert_bytes_out = per_expert_bytes; + return true; } -static void ds4_vec_dot_q6_K_q8_K(int n, float *s, const block_q6_K *x, const block_q8_K *y) { - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = y[i].d * f16_to_f32(x[i].d); - const uint8_t *ql = x[i].ql; - const uint8_t *qh = x[i].qh; - const int8_t *scales = x[i].scales; - const int8_t *q8 = y[i].qs; - int64_t isum = 0; - - for (int n128 = 0; n128 < QK_K; n128 += 128) { - for (int l = 0; l < 32; l++) { - const int is = l / 16; - const int q1 = ((int)(ql[l + 0] & 0x0F) | (((qh[l] >> 0) & 3) << 4)) - 32; - const int q2 = ((int)(ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; - const int q3 = ((int)(ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; - const int q4 = ((int)(ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; - - isum += (int64_t)scales[is + 0] * q1 * (int)q8[n128 + l + 0]; - isum += (int64_t)scales[is + 2] * q2 * (int)q8[n128 + l + 32]; - isum += (int64_t)scales[is + 4] * q3 * (int)q8[n128 + l + 64]; - isum += (int64_t)scales[is + 6] * q4 * (int)q8[n128 + l + 96]; - } - - ql += 64; - qh += 32; - scales += 8; - } - - sumf += d * (float)isum; +static DS4_MAYBE_UNUSED bool streaming_layer_gate_down_expert_bytes( + const ds4_layer_weights *layer, + uint64_t *gate_expert_bytes, + uint64_t *down_expert_bytes) { + if (gate_expert_bytes) *gate_expert_bytes = 0; + if (down_expert_bytes) *down_expert_bytes = 0; + if (!layer || + !gate_expert_bytes || + !down_expert_bytes || + !layer->ffn_gate_exps || + !layer->ffn_down_exps) { + return false; } - *s = sumf; -} - -static float ds4_vec_dot_q4_K_f32(int n, const block_q4_K *x, const float *y) { - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d); - const float dmin = f16_to_f32(x[i].dmin); - const uint8_t *qs = x[i].qs; - const uint8_t *scales = x[i].scales; - const float *yb = y + (uint64_t)i * QK_K; - - for (int j = 0; j < QK_K / 32; j++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(j, scales, &sc_val, &m_val); - - const int byte_off = (j >> 1) * 32; - const int shift = (j & 1) * 4; - const float scale = d * (float)sc_val; - const float minv = dmin * (float)m_val; - for (int l = 0; l < 32; l++) { - const int q = (qs[byte_off + l] >> shift) & 0x0F; - sumf += (scale * (float)q - minv) * yb[j * 32 + l]; - } - } + const uint64_t gate_row_bytes = + routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t down_row_bytes = + routed_expert_row_bytes(layer->ffn_down_exps); + if (gate_row_bytes == 0 || + down_row_bytes == 0 || + layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + return false; } - return sumf; + *gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; + *down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + return *gate_expert_bytes != 0 && *down_expert_bytes != 0; } -static float ds4_vec_dot_q5_K_f32(int n, const block_q5_K *x, const float *y) { - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d); - const float dmin = f16_to_f32(x[i].dmin); - const uint8_t *ql = x[i].qs; - const uint8_t *qh = x[i].qh; - const uint8_t *scales = x[i].scales; - const float *yb = y + (uint64_t)i * QK_K; - - for (int group = 0; group < QK_K / 32; group++) { - uint8_t sc_val, m_val; - q4_k_get_scale_min(group, scales, &sc_val, &m_val); +static bool ds4_streaming_routed_expert_bytes( + const ds4_weights *weights, + uint64_t *per_expert_bytes_out) { + if (per_expert_bytes_out) *per_expert_bytes_out = 0; + if (!weights || !per_expert_bytes_out) return false; - const int ql_base = (group >> 1) * 32; - const int shift = (group & 1) * 4; - const uint8_t hmask = (uint8_t)(1u << group); - const float scale = d * (float)sc_val; - const float minv = dmin * (float)m_val; - for (int l = 0; l < 32; l++) { - const int q = ((ql[ql_base + l] >> shift) & 0x0F) + - ((qh[l] & hmask) ? 16 : 0); - sumf += (scale * (float)q - minv) * yb[group * 32 + l]; + /* Mixed-precision models can put an outlier quant at the first routed + * layer owned by a distributed slice. Choosing that first layer as the + * slab class makes every ordinary layer bypass the cache. Use the most + * common local size class instead (ties retain the earliest class). */ + uint64_t best_bytes = 0; + uint32_t best_count = 0; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + uint64_t candidate = 0; + if (!streaming_layer_routed_expert_bytes(&weights->layer[il], + &candidate)) { + continue; + } + uint32_t count = 0; + for (uint32_t jl = 0; jl < DS4_N_LAYER; jl++) { + uint64_t bytes = 0; + if (streaming_layer_routed_expert_bytes(&weights->layer[jl], + &bytes) && + bytes == candidate) { + count++; } } + if (count > best_count) { + best_bytes = candidate; + best_count = count; + } } - - return sumf; + if (best_count == 0) return false; + *per_expert_bytes_out = best_bytes; + return true; } -static float ds4_vec_dot_q6_K_f32(int n, const block_q6_K *x, const float *y) { - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d); - const uint8_t *ql = x[i].ql; - const uint8_t *qh = x[i].qh; - const int8_t *scales = x[i].scales; - const float *yb = y + (uint64_t)i * QK_K; +enum { DS4_STREAMING_PREFILL_HEADROOM_LAYERS = 2 }; - for (int n128 = 0; n128 < QK_K; n128 += 128) { - for (int l = 0; l < 32; l++) { - const int is = l / 16; - const int q1 = ((int)(ql[l + 0] & 0x0F) | (((qh[l] >> 0) & 3) << 4)) - 32; - const int q2 = ((int)(ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; - const int q3 = ((int)(ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; - const int q4 = ((int)(ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; +static bool ds4_streaming_cacheable_expert_count( + const ds4_weights *weights, + uint64_t *experts_out, + uint32_t *layers_out) { + if (experts_out) *experts_out = 0; + if (layers_out) *layers_out = 0; + if (!weights || !experts_out || DS4_N_EXPERT == 0) return false; - sumf += d * (float)scales[is + 0] * (float)q1 * yb[n128 + l + 0]; - sumf += d * (float)scales[is + 2] * (float)q2 * yb[n128 + l + 32]; - sumf += d * (float)scales[is + 4] * (float)q3 * yb[n128 + l + 64]; - sumf += d * (float)scales[is + 6] * (float)q4 * yb[n128 + l + 96]; - } + uint64_t slab_bytes = 0; + if (!ds4_streaming_routed_expert_bytes(weights, &slab_bytes)) return false; - ql += 64; - qh += 32; - scales += 8; + uint32_t layers = 0; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + uint64_t per_expert_bytes = 0; + if (!streaming_layer_routed_expert_bytes(&weights->layer[il], + &per_expert_bytes)) { + continue; } + if (per_expert_bytes == slab_bytes) layers++; } - return sumf; -} - -static inline float ds4_vec_dot_q5_q6_K_f32(uint32_t type, int n, const uint8_t *x, const float *y) { - if (type == DS4_TENSOR_Q5_K) { - return ds4_vec_dot_q5_K_f32(n, (const block_q5_K *)x, y); - } else if (type == DS4_TENSOR_Q6_K) { - return ds4_vec_dot_q6_K_f32(n, (const block_q6_K *)x, y); - } else { - ds4_die("expected a Q5_K or Q6_K tensor"); + if (layers == 0 || + (uint64_t)layers > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { + return false; } - return 0.0f; + *experts_out = (uint64_t)layers * (uint64_t)DS4_N_EXPERT; + if (layers_out) *layers_out = layers; + return true; } -static float ds4_vec_dot_iq2_xxs_f32(int n, const block_iq2_xxs *x, const float *y) { - pthread_once(&iq2xxs_signed_grid_once, iq2xxs_signed_grid_init); +static bool ds4_streaming_prefill_headroom_bytes( + const ds4_weights *weights, + uint64_t *bytes_out) { + if (bytes_out) *bytes_out = 0; + if (!weights || !bytes_out) return false; - const int nb = n / QK_K; - float sumf = 0.0f; - uint32_t aux32[2]; - const uint8_t *aux8 = (const uint8_t *)aux32; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d); - const uint16_t *q2 = x[i].qs; - const float *yb = y + (uint64_t)i * QK_K; - - for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { - memcpy(aux32, q2, 2 * sizeof(uint32_t)); - q2 += 4; - - const float scale = 0.125f * d * (float)(2u * (aux32[1] >> 28) + 1u); - const uint32_t base = (uint32_t)ib32 * 32u; - for (int l = 0; l < 4; l++) { - const uint32_t sign_idx = (aux32[1] >> (7 * l)) & 127u; - const int8_t *grid = iq2xxs_signed_grid[aux8[l]][sign_idx]; - const float *yf = yb + base + (uint32_t)l * 8u; - for (int j = 0; j < 8; j++) { - sumf += scale * (float)grid[j] * yf[j]; - } - } - } + uint64_t per_expert_bytes = 0; + uint64_t cacheable_experts = 0; + uint32_t cacheable_layers = 0; + if (!ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes) || + !ds4_streaming_cacheable_expert_count(weights, + &cacheable_experts, + &cacheable_layers)) { + return false; + } + (void)cacheable_experts; + const uint32_t reserve_layers = + cacheable_layers < DS4_STREAMING_PREFILL_HEADROOM_LAYERS ? + cacheable_layers : DS4_STREAMING_PREFILL_HEADROOM_LAYERS; + if (per_expert_bytes > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { + return false; } + const uint64_t layer_bytes = + per_expert_bytes * (uint64_t)DS4_N_EXPERT; + if (reserve_layers != 0 && + layer_bytes > UINT64_MAX / (uint64_t)reserve_layers) return false; - return sumf; + *bytes_out = layer_bytes * (uint64_t)reserve_layers; + return true; } -static void ds4_vec_dot_q8_K_q8_K(int n, float *s, - const block_q8_K *x, - const block_q8_K *y) { - const int nb = n / QK_K; - float sum = 0.0f; -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - for (int i = 0; i < nb; i++) { - int32x4_t isum = vdupq_n_s32(0); - for (int j = 0; j < QK_K; j += 16) { - isum = vdotq_s32(isum, vld1q_s8(x[i].qs + j), - vld1q_s8(y[i].qs + j)); - } - sum += x[i].d * y[i].d * (float)vaddvq_s32(isum); - } -#else - for (int i = 0; i < nb; i++) { - int isum = 0; - for (int j = 0; j < QK_K; j++) { - isum += (int)x[i].qs[j] * (int)y[i].qs[j]; - } - sum += x[i].d * y[i].d * (float)isum; - } -#endif - *s = sum; +/* + * Mixed-precision ("boosted") GGUFs upcast a few layers' routed experts to a + * bigger quant (e.g. Q4_K among IQ2 layers). The streaming expert cache is a + * single-size-class slab allocator sized from the dominant local routed-layer + * size class, so other layers can never be served from it: they must read + * expert weights through the mapped-model views instead. A layer is "uniform" + * iff its per-expert bytes match the slab class. + */ +static DS4_MAYBE_UNUSED bool weights_streaming_layer_experts_uniform( + const ds4_weights *w, + uint32_t il) { + uint64_t base = 0; + uint64_t bytes = 0; + if (!w || il >= DS4_N_LAYER) return true; + const ds4_layer_weights *l = &w->layer[il]; + if (!streaming_layer_routed_expert_bytes(l, &bytes)) return true; + if (!ds4_streaming_routed_expert_bytes(w, &base)) return true; + return bytes == base; } -static void ds4_vec_dot_q8_K_pair_q8_K( - int n, float *s0, float *s1, - const block_q8_K *x0, const block_q8_K *x1, - const block_q8_K *y) { - const int nb = n / QK_K; - float sum0 = 0.0f; - float sum1 = 0.0f; -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - for (int i = 0; i < nb; i++) { - int32x4_t isum0 = vdupq_n_s32(0); - int32x4_t isum1 = vdupq_n_s32(0); - for (int j = 0; j < QK_K; j += 16) { - const int8x16_t yv = vld1q_s8(y[i].qs + j); - isum0 = vdotq_s32(isum0, vld1q_s8(x0[i].qs + j), yv); - isum1 = vdotq_s32(isum1, vld1q_s8(x1[i].qs + j), yv); - } - sum0 += x0[i].d * y[i].d * (float)vaddvq_s32(isum0); - sum1 += x1[i].d * y[i].d * (float)vaddvq_s32(isum1); - } -#else - for (int i = 0; i < nb; i++) { - int isum0 = 0; - int isum1 = 0; - for (int j = 0; j < QK_K; j++) { - const int yv = (int)y[i].qs[j]; - isum0 += (int)x0[i].qs[j] * yv; - isum1 += (int)x1[i].qs[j] * yv; - } - sum0 += x0[i].d * y[i].d * (float)isum0; - sum1 += x1[i].d * y[i].d * (float)isum1; +static uint32_t ds4_streaming_cache_experts_for_byte_budget( + const ds4_weights *weights, + uint64_t bytes, + uint64_t *per_expert_bytes_out) { + uint64_t per_expert_bytes = 0; + if (per_expert_bytes_out) *per_expert_bytes_out = 0; + if (!weights || + bytes == 0 || + !ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes)) { + return 0; } -#endif - *s0 = sum0; - *s1 = sum1; + if (per_expert_bytes_out) *per_expert_bytes_out = per_expert_bytes; + return ds4_ssd_cache_experts_for_byte_budget(bytes, per_expert_bytes); } -static DS4_MAYBE_UNUSED void ds4_vec_dot_iq2_xxs_q8_K(int n, float *s, const block_iq2_xxs *x, const block_q8_K *y) { - const int nb = n / QK_K; - -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d) * y[i].d; - const uint16_t *q2 = x[i].qs; - const int8_t *q8 = y[i].qs; - float sumf1 = 0.0f; - float sumf2 = 0.0f; - - for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { - int8x16x4_t q8b = vld1q_s8_x4(q8); - q8 += 64; - - uint32_t aux32[4]; - memcpy(aux32, q2, sizeof(aux32)); - q2 += 8; - const uint8_t *aux8 = (const uint8_t *)aux32; - - int8x16_t q2u0 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[0])), - vld1_s8((const int8_t *)(iq2xxs_grid + aux8[1]))); - int8x16_t q2u1 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[2])), - vld1_s8((const int8_t *)(iq2xxs_grid + aux8[3]))); - int8x16_t q2u2 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[8])), - vld1_s8((const int8_t *)(iq2xxs_grid + aux8[9]))); - int8x16_t q2u3 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[10])), - vld1_s8((const int8_t *)(iq2xxs_grid + aux8[11]))); - - const int8x16_t q2s0 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[1] >> 0) & 127]), - vld1_s8(iq2xxs_signs[(aux32[1] >> 7) & 127])); - const int8x16_t q2s1 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[1] >> 14) & 127]), - vld1_s8(iq2xxs_signs[(aux32[1] >> 21) & 127])); - const int8x16_t q2s2 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[3] >> 0) & 127]), - vld1_s8(iq2xxs_signs[(aux32[3] >> 7) & 127])); - const int8x16_t q2s3 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[3] >> 14) & 127]), - vld1_s8(iq2xxs_signs[(aux32[3] >> 21) & 127])); - - q2u0 = vmulq_s8(q2u0, q2s0); - q2u1 = vmulq_s8(q2u1, q2s1); - q2u2 = vmulq_s8(q2u2, q2s2); - q2u3 = vmulq_s8(q2u3, q2s3); - - const int32x4_t p1 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), q2u0, q8b.val[0]), q2u1, q8b.val[1]); - const int32x4_t p2 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), q2u2, q8b.val[2]), q2u3, q8b.val[3]); - - sumf1 += (float)vaddvq_s32(p1) * (0.5f + (float)(aux32[1] >> 28)); - sumf2 += (float)vaddvq_s32(p2) * (0.5f + (float)(aux32[3] >> 28)); - } - - sumf += d * (sumf1 + sumf2); - } - - *s = 0.25f * sumf; -#else - uint32_t aux32[2]; - const uint8_t *aux8 = (const uint8_t *)aux32; - float sumf = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d = f16_to_f32(x[i].d) * y[i].d; - const uint16_t *q2 = x[i].qs; - const int8_t *q8 = y[i].qs; - int32_t bsum = 0; - - for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { - memcpy(aux32, q2, 2 * sizeof(uint32_t)); - q2 += 4; - - const uint32_t ls = 2 * (aux32[1] >> 28) + 1; - int32_t sumi = 0; - for (int l = 0; l < 4; l += 2) { - const uint32_t sign_idx0 = (aux32[1] >> (7 * l)) & 127; - const uint32_t sign_idx1 = (aux32[1] >> (7 * (l + 1))) & 127; - sumi += dot_iq2_pair_16(iq2xxs_signed_grid[aux8[l]][sign_idx0], - iq2xxs_signed_grid[aux8[l + 1]][sign_idx1], - q8); - q8 += 16; - } - bsum += sumi * (int32_t)ls; - } - sumf += d * (float)bsum; - } - *s = 0.125f * sumf; -#endif -} - -static void ds4_vec_dot_iq2_xxs_pair_q8_K( - int n, - float *s0, - float *s1, - const block_iq2_xxs *x0, - const block_iq2_xxs *x1, - const block_q8_K *y) { -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - const int nb = n / QK_K; - float total0 = 0.0f; - float total1 = 0.0f; - - for (int i = 0; i < nb; i++) { - const float d0 = f16_to_f32(x0[i].d) * y[i].d; - const float d1 = f16_to_f32(x1[i].d) * y[i].d; - const uint16_t *q20 = x0[i].qs; - const uint16_t *q21 = x1[i].qs; - const int8_t *q8 = y[i].qs; - float sum01 = 0.0f; - float sum02 = 0.0f; - float sum11 = 0.0f; - float sum12 = 0.0f; - - for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { - const int8x16x4_t q8b = vld1q_s8_x4(q8); - q8 += 64; - - uint32_t aux0[4]; - uint32_t aux1[4]; - memcpy(aux0, q20, sizeof(aux0)); - memcpy(aux1, q21, sizeof(aux1)); - q20 += 8; - q21 += 8; - const uint8_t *a0 = (const uint8_t *)aux0; - const uint8_t *a1 = (const uint8_t *)aux1; - -#define DS4_IQ2_PAIR_DOT(aux, aux8, accum_a, accum_b) do { \ - int8x16_t u0 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[0])), \ - vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[1]))); \ - int8x16_t u1 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[2])), \ - vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[3]))); \ - int8x16_t u2 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[8])), \ - vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[9]))); \ - int8x16_t u3 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[10])), \ - vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[11]))); \ - const int8x16_t sgn0 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[1] >> 0) & 127]), \ - vld1_s8(iq2xxs_signs[((aux)[1] >> 7) & 127])); \ - const int8x16_t sgn1 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[1] >> 14) & 127]), \ - vld1_s8(iq2xxs_signs[((aux)[1] >> 21) & 127])); \ - const int8x16_t sgn2 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[3] >> 0) & 127]), \ - vld1_s8(iq2xxs_signs[((aux)[3] >> 7) & 127])); \ - const int8x16_t sgn3 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[3] >> 14) & 127]), \ - vld1_s8(iq2xxs_signs[((aux)[3] >> 21) & 127])); \ - u0 = vmulq_s8(u0, sgn0); \ - u1 = vmulq_s8(u1, sgn1); \ - u2 = vmulq_s8(u2, sgn2); \ - u3 = vmulq_s8(u3, sgn3); \ - const int32x4_t p1 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), u0, q8b.val[0]), u1, q8b.val[1]); \ - const int32x4_t p2 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), u2, q8b.val[2]), u3, q8b.val[3]); \ - (accum_a) += (float)vaddvq_s32(p1) * (0.5f + (float)((aux)[1] >> 28)); \ - (accum_b) += (float)vaddvq_s32(p2) * (0.5f + (float)((aux)[3] >> 28)); \ - } while (0) - - DS4_IQ2_PAIR_DOT(aux0, a0, sum01, sum02); - DS4_IQ2_PAIR_DOT(aux1, a1, sum11, sum12); - -#undef DS4_IQ2_PAIR_DOT - } - - total0 += d0 * (sum01 + sum02); - total1 += d1 * (sum11 + sum12); - } - - *s0 = 0.25f * total0; - *s1 = 0.25f * total1; -#else - ds4_vec_dot_iq2_xxs_q8_K(n, s0, x0, y); - ds4_vec_dot_iq2_xxs_q8_K(n, s1, x1, y); -#endif -} - -typedef struct { - ds4_tensor *hc_attn_fn; - ds4_tensor *hc_attn_scale; - ds4_tensor *hc_attn_base; - ds4_tensor *attn_norm; - ds4_tensor *attn_q_a; - ds4_tensor *attn_q_a_norm; - ds4_tensor *attn_q_b; - ds4_tensor *attn_kv; - ds4_tensor *attn_kv_a_mqa; - ds4_tensor *attn_kv_a_norm; - ds4_tensor *attn_k_b; - ds4_tensor *attn_v_b; - ds4_tensor *attn_sinks; - ds4_tensor *attn_output; - ds4_tensor *attn_output_a; - ds4_tensor *attn_output_b; - ds4_tensor *attn_compressor_ape; - ds4_tensor *attn_compressor_kv; - ds4_tensor *attn_compressor_gate; - ds4_tensor *attn_compressor_norm; - ds4_tensor *indexer_attn_q_b; - ds4_tensor *indexer_attn_k; - ds4_tensor *indexer_k_norm; - ds4_tensor *indexer_k_norm_b; - ds4_tensor *indexer_proj; - ds4_tensor *indexer_compressor_ape; - ds4_tensor *indexer_compressor_kv; - ds4_tensor *indexer_compressor_gate; - ds4_tensor *indexer_compressor_norm; - ds4_tensor *hc_ffn_fn; - ds4_tensor *hc_ffn_scale; - ds4_tensor *hc_ffn_base; - ds4_tensor *ffn_norm; - ds4_tensor *ffn_gate_tid2eid; - ds4_tensor *ffn_gate; - ds4_tensor *ffn_up; - ds4_tensor *ffn_down; - ds4_tensor *ffn_gate_inp; - ds4_tensor *ffn_exp_probs_b; - ds4_tensor *ffn_gate_exps; - ds4_tensor *ffn_up_exps; - ds4_tensor *ffn_down_exps; - ds4_tensor *ffn_gate_shexp; - ds4_tensor *ffn_up_shexp; - ds4_tensor *ffn_down_shexp; - ds4_tensor *nextn_eh_proj; - ds4_tensor *nextn_enorm; - ds4_tensor *nextn_hnorm; - ds4_tensor *nextn_shared_head_norm; -} ds4_layer_weights; - -typedef struct { - ds4_tensor *token_embd; - ds4_tensor *output_hc_base; - ds4_tensor *output_hc_fn; - ds4_tensor *output_hc_scale; - ds4_tensor *output_norm; - ds4_tensor *output; - ds4_layer_weights layer[DS4_MAX_LAYER]; -} ds4_weights; - -typedef struct { - ds4_tensor *e_proj; - ds4_tensor *h_proj; - ds4_tensor *enorm; - ds4_tensor *hnorm; - ds4_tensor *norm; - ds4_tensor *hc_head_base; - ds4_tensor *hc_head_fn; - ds4_tensor *hc_head_scale; - ds4_layer_weights block; -} ds4_mtp_weights; - -typedef struct { - ds4_tensor *main_proj; - ds4_tensor *main_norm; - ds4_tensor *norm; - ds4_tensor *hc_head_base; - ds4_tensor *hc_head_fn; - ds4_tensor *hc_head_scale; - ds4_tensor *markov_w1; - ds4_tensor *markov_w2; - ds4_tensor *confidence_proj; - ds4_layer_weights block; -} ds4_dspark_stage_weights; - -typedef struct { - uint32_t n_stages; - uint32_t block_size; - uint32_t markov_rank; - uint32_t noise_token_id; - uint32_t target_layer_count; - uint32_t target_layers[DS4_DSPARK_MAX_TARGET_LAYERS]; - uint32_t present_tensors; - uint32_t missing_tensors; - uint32_t invalid_tensors; - uint32_t metadata_errors; - bool has_block_size; - bool has_markov_rank; - bool has_noise_token_id; - bool has_target_layers; - ds4_dspark_stage_weights stage[DS4_DSPARK_MAX_STAGES]; -} ds4_dspark_weights; - -/* ========================================================================= - * Fixed Weight Binding and Model Validation. - * ========================================================================= - * - * The GGUF tensor directory is converted into a DS4-specific pointer table. - * After this section, the rest of the program addresses tensors by semantic - * fields such as layer->attn_q_a or layer->ffn_gate_exps rather than by string - * lookup. Shape validation is intentionally strict. - */ - -static uint32_t required_u32(const ds4_model *m, const char *key) { - uint32_t v = 0; - if (!model_get_u32(m, key, &v)) { - fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); - exit(1); - } - return v; -} - -static uint64_t required_u64_compat(const ds4_model *m, const char *key) { - uint64_t v = 0; - if (!model_get_u64_compat(m, key, &v)) { - fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); - exit(1); - } - return v; -} - -static float required_f32(const ds4_model *m, const char *key) { - float v = 0.0f; - if (!model_get_f32_compat(m, key, &v)) { - fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); - exit(1); - } - return v; +#ifndef DS4_NO_GPU +static ds4_gpu_stream_expert_table graph_stream_expert_table_make( + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + ds4_gpu_stream_expert_table table; + memset(&table, 0, sizeof(table)); + if (!model || !layer) return table; + table.model_map = model->map; + table.model_size = model->size; + table.layer = il; + table.n_total_expert = DS4_N_EXPERT; + table.gate_offset = layer->ffn_gate_exps ? layer->ffn_gate_exps->abs_offset : 0; + table.up_offset = layer->ffn_up_exps ? layer->ffn_up_exps->abs_offset : 0; + table.down_offset = layer->ffn_down_exps ? layer->ffn_down_exps->abs_offset : 0; + table.gate_expert_bytes = gate_expert_bytes; + table.down_expert_bytes = down_expert_bytes; + return table; } +#endif -static bool required_bool(const ds4_model *m, const char *key) { - bool v = false; - if (!model_get_bool(m, key, &v)) { - fprintf(stderr, "ds4: required metadata key is missing: %s\n", key); - exit(1); - } - return v; -} +static uint64_t ds4_streaming_manual_cache_safe_bytes( + ds4_backend backend, + int ctx_size, + uint32_t prefill_chunk, + bool ssd_streaming) { +#ifdef DS4_NO_GPU + (void)backend; + (void)ctx_size; + (void)prefill_chunk; + (void)ssd_streaming; + return 0; +#else + const uint64_t gib = 1024ull * 1024ull * 1024ull; + const uint64_t recommended = ds4_gpu_recommended_working_set_size(); + if (recommended == 0) return 0; -static ds4_tensor *required_tensor(const ds4_model *m, const char *name) { - ds4_tensor *t = model_find_tensor(m, name); - if (!t) { - fprintf(stderr, "ds4: required tensor is missing: %s\n", name); - exit(1); - } - return t; + /* + * Explicit NGB budgets name only the routed expert cache. Keep that cache + * below the graph backend's working-set recommendation after accounting for + * the graph context/KV buffers. This is intentionally not an mlock-derived + * cap: crossing too close to the recommended working set makes short + * token-major prefill spend most of its time in VM/driver synchronization. + */ + uint64_t target = recommended > UINT64_MAX / 7ull ? + UINT64_MAX : (recommended * 7ull) / 8ull; + const ds4_context_memory ctx_mem = + ds4_context_memory_estimate_with_prefill_mode(backend, + ctx_size, + prefill_chunk, + ssd_streaming); + uint64_t safe = 0; + if (target > ctx_mem.total_bytes) safe = target - ctx_mem.total_bytes; + safe = (safe / gib) * gib; + if (safe == 0) safe = gib; + return safe; +#endif } -static ds4_tensor *tensor_by_namef(const ds4_model *m, const char *fmt, uint32_t layer) { - char name[128]; - int n = snprintf(name, sizeof(name), fmt, layer); - if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); - return model_find_tensor(m, name); +static uint64_t ds4_add_sat_u64(uint64_t a, uint64_t b) { + return a > UINT64_MAX - b ? UINT64_MAX : a + b; } -static ds4_tensor *required_tensorf(const ds4_model *m, const char *fmt, uint32_t layer) { - char name[128]; - int n = snprintf(name, sizeof(name), fmt, layer); - if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); - return required_tensor(m, name); +static uint64_t ds4_mul_sat_u64(uint64_t a, uint64_t b) { + if (a != 0 && b > UINT64_MAX / a) return UINT64_MAX; + return a * b; } -static ds4_tensor *tensor_by_mtp_stage_suffix( - const ds4_model *m, - uint32_t stage, - const char *suffix) { - char name[160]; - int n = snprintf(name, sizeof(name), "mtp.%u.%s", stage, suffix); - if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("tensor name is too long"); - return model_find_tensor(m, name); +static double ds4_bytes_to_gib(uint64_t bytes) { + return (double)bytes / 1073741824.0; } -static void tensor_expect_layout( +static void tensor_expect_routed_expert( const ds4_tensor *t, - uint32_t type, uint32_t ndim, uint64_t d0, uint64_t d1, uint64_t d2) { - if (!t) ds4_die("internal error: missing tensor while validating layout"); - if (t->type != type) { + if (!t) ds4_die("internal error: missing routed expert tensor while validating layout"); + if (!tensor_is_routed_expert_type(t->type)) { fprintf(stderr, - "ds4: tensor %.*s has type %s, expected %s\n", + "ds4: tensor %.*s has type %u (%s), expected a routed expert quant type\n", (int)t->name.len, t->name.ptr, - tensor_type_name(t->type), - tensor_type_name(type)); + t->type, + tensor_type_name(t->type)); exit(1); } if (t->ndim != ndim) { @@ -4266,42511 +3652,4040 @@ static void tensor_expect_layout( } } -static bool tensor_type_is_glm_dense_quant(uint32_t type) { - return type == DS4_TENSOR_Q8_0 || - type == DS4_TENSOR_Q4_K || - type == DS4_TENSOR_Q4_0; -} - -static bool tensor_type_is_dense_quant(uint32_t type) { - return type == DS4_TENSOR_Q8_0 || - type == DS4_TENSOR_Q4_K || - type == DS4_TENSOR_Q4_0; +static bool weights_have_output_head(const ds4_weights *w) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + return w && w->output_norm && w->output; + } + return w && + w->output_hc_base && + w->output_hc_fn && + w->output_hc_scale && + w->output_norm && + w->output; } -static void tensor_expect_glm_dense_quant_layout( - const ds4_tensor *t, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!t) ds4_die("internal error: missing tensor while validating GLM dense layout"); - if (!tensor_type_is_glm_dense_quant(t->type)) { - fprintf(stderr, - "ds4: tensor %.*s has type %s, expected q8_0, q4_K, or q4_0\n", - (int)t->name.len, - t->name.ptr, - tensor_type_name(t->type)); - exit(1); +static bool weights_have_partial_output_head(const ds4_weights *w) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + return w && (w->output_norm || w->output); } - tensor_expect_layout(t, t->type, ndim, d0, d1, d2); + return w && + (w->output_hc_base || + w->output_hc_fn || + w->output_hc_scale || + w->output_norm || + w->output); } -static void tensor_expect_dense_quant_layout( - const ds4_tensor *t, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!t) ds4_die("internal error: missing tensor while validating dense quant layout"); - if (!tensor_type_is_dense_quant(t->type)) { - fprintf(stderr, - "ds4: tensor %.*s has type %s, expected q8_0, q4_K, or q4_0\n", - (int)t->name.len, - t->name.ptr, - tensor_type_name(t->type)); - exit(1); +static bool weights_glm_dsa_layer_has_required(const ds4_layer_weights *l, uint32_t il) { + if (!l) return false; + if (!l->attn_norm || + !l->attn_q_a || + !l->attn_q_a_norm || + !l->attn_q_b || + !l->attn_kv_a_mqa || + !l->attn_kv_a_norm || + !l->attn_k_b || + !l->attn_v_b || + !l->attn_output || + !l->indexer_attn_q_b || + !l->indexer_attn_k || + !l->indexer_k_norm || + !l->indexer_k_norm_b || + !l->indexer_proj || + !l->ffn_norm) + { + return false; } - tensor_expect_layout(t, t->type, ndim, d0, d1, d2); -} -static void tensor_expect_optional( - const ds4_tensor *t, - uint32_t type, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (t) tensor_expect_layout(t, type, ndim, d0, d1, d2); -} - -static void tensor_expect_plain_layout( - const ds4_tensor *t, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!t) ds4_die("internal error: missing tensor while validating layout"); - if (t->type != DS4_TENSOR_F16 && t->type != DS4_TENSOR_F32) { - fprintf(stderr, - "ds4: tensor %.*s has type %s, expected F16 or F32\n", - (int)t->name.len, - t->name.ptr, - tensor_type_name(t->type)); - exit(1); + if (il < DS4_N_LEADING_DENSE) { + if (!l->ffn_gate || !l->ffn_up || !l->ffn_down) return false; + } else { + if (!l->ffn_gate_inp || + !l->ffn_exp_probs_b || + !l->ffn_gate_exps || + !l->ffn_up_exps || + !l->ffn_down_exps || + !l->ffn_gate_shexp || + !l->ffn_up_shexp || + !l->ffn_down_shexp) + { + return false; + } } - tensor_expect_layout(t, t->type, ndim, d0, d1, d2); -} -static bool tensor_type_is_f16_or_q8_0(uint32_t type) { - return type == DS4_TENSOR_F16 || type == DS4_TENSOR_Q8_0; -} - -static void tensor_expect_f16_or_q8_0_layout( - const ds4_tensor *t, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!t) ds4_die("internal error: missing tensor while validating layout"); - if (!tensor_type_is_f16_or_q8_0(t->type)) { - fprintf(stderr, - "ds4: tensor %.*s has type %s, expected f16 or q8_0\n", - (int)t->name.len, - t->name.ptr, - tensor_type_name(t->type)); - exit(1); + if (DS4_N_NEXTN_PREDICT != 0 && + il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER && + (!l->nextn_eh_proj || + !l->nextn_enorm || + !l->nextn_hnorm || + !l->nextn_shared_head_norm)) + { + return false; } - tensor_expect_layout(t, t->type, ndim, d0, d1, d2); -} -static bool tensor_is_routed_expert_type(uint32_t type) { - return type == DS4_TENSOR_Q8_0 || - type == DS4_TENSOR_IQ2_XXS || - type == DS4_TENSOR_Q2_K || - type == DS4_TENSOR_Q4_K || - type == DS4_TENSOR_Q5_K || - type == DS4_TENSOR_Q6_K; + return true; } -static DS4_MAYBE_UNUSED uint64_t routed_expert_block_bytes(uint32_t type) { - switch (type) { - case DS4_TENSOR_Q8_0: return 34; - case DS4_TENSOR_IQ2_XXS: return sizeof(block_iq2_xxs); - case DS4_TENSOR_Q2_K: return sizeof(block_q2_K); - case DS4_TENSOR_Q4_K: return sizeof(block_q4_K); - case DS4_TENSOR_Q5_K: return sizeof(block_q5_K); - case DS4_TENSOR_Q6_K: return sizeof(block_q6_K); - default: ds4_die("unsupported routed expert tensor type"); +static bool weights_layer_has_required(const ds4_layer_weights *l, uint32_t il) { + if (!l) return false; + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + return weights_glm_dsa_layer_has_required(l, il); } - return 0; -} - -static DS4_MAYBE_UNUSED uint64_t routed_expert_row_bytes(const ds4_tensor *t) { - const gguf_type_info *info = tensor_type(t->type); - if (!info || info->block_elems == 0) ds4_die("unsupported routed expert tensor type"); - if ((t->dim[0] % info->block_elems) != 0) ds4_die("routed expert row is not quant block aligned"); - return (t->dim[0] / info->block_elems) * routed_expert_block_bytes(t->type); -} - -static bool streaming_layer_routed_expert_bytes( - const ds4_layer_weights *layer, - uint64_t *per_expert_bytes_out) { - if (per_expert_bytes_out) *per_expert_bytes_out = 0; - if (!layer || - !per_expert_bytes_out || - !layer->ffn_gate_exps || - !layer->ffn_up_exps || - !layer->ffn_down_exps) { + if (!l->hc_attn_fn || + !l->hc_attn_scale || + !l->hc_attn_base || + !l->attn_norm || + !l->attn_q_a || + !l->attn_q_a_norm || + !l->attn_q_b || + !l->attn_kv || + !l->attn_kv_a_norm || + !l->attn_sinks || + !l->attn_output_a || + !l->attn_output_b || + !l->hc_ffn_fn || + !l->hc_ffn_scale || + !l->hc_ffn_base || + !l->ffn_norm || + !l->ffn_gate_inp || + !l->ffn_gate_exps || + !l->ffn_up_exps || + !l->ffn_down_exps || + !l->ffn_gate_shexp || + !l->ffn_up_shexp || + !l->ffn_down_shexp) + { return false; } - const uint64_t gate_row_bytes = - routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t up_row_bytes = - routed_expert_row_bytes(layer->ffn_up_exps); - const uint64_t down_row_bytes = - routed_expert_row_bytes(layer->ffn_down_exps); - if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_up_exps->dim[1] > UINT64_MAX / up_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0 && + (!l->attn_compressor_ape || + !l->attn_compressor_kv || + !l->attn_compressor_gate || + !l->attn_compressor_norm)) + { return false; } - - const uint64_t gate_expert_bytes = - layer->ffn_gate_exps->dim[1] * gate_row_bytes; - const uint64_t up_expert_bytes = - layer->ffn_up_exps->dim[1] * up_row_bytes; - const uint64_t down_expert_bytes = - layer->ffn_down_exps->dim[1] * down_row_bytes; - if (gate_expert_bytes > UINT64_MAX - up_expert_bytes || - gate_expert_bytes + up_expert_bytes > - UINT64_MAX - down_expert_bytes) { + if (ratio == 4 && + (!l->indexer_attn_q_b || + !l->indexer_proj || + !l->indexer_compressor_ape || + !l->indexer_compressor_kv || + !l->indexer_compressor_gate || + !l->indexer_compressor_norm)) + { return false; } - - const uint64_t per_expert_bytes = - gate_expert_bytes + up_expert_bytes + down_expert_bytes; - if (per_expert_bytes == 0) return false; - *per_expert_bytes_out = per_expert_bytes; + if (il < DS4_N_HASH_LAYER && !l->ffn_gate_tid2eid) return false; return true; } -static DS4_MAYBE_UNUSED bool streaming_layer_gate_down_expert_bytes( - const ds4_layer_weights *layer, - uint64_t *gate_expert_bytes, - uint64_t *down_expert_bytes) { - if (gate_expert_bytes) *gate_expert_bytes = 0; - if (down_expert_bytes) *down_expert_bytes = 0; - if (!layer || - !gate_expert_bytes || - !down_expert_bytes || - !layer->ffn_gate_exps || - !layer->ffn_down_exps) { - return false; - } - - const uint64_t gate_row_bytes = - routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t down_row_bytes = - routed_expert_row_bytes(layer->ffn_down_exps); - if (gate_row_bytes == 0 || - down_row_bytes == 0 || - layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { - return false; +static bool weights_layers_bound(const ds4_weights *w, uint32_t layer_start, uint32_t layer_end) { + if (!w || layer_start >= DS4_N_LAYER) return false; + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; + for (uint32_t il = layer_start; il <= layer_end; il++) { + if (!weights_layer_has_required(&w->layer[il], il)) return false; } - - *gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; - *down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; - return *gate_expert_bytes != 0 && *down_expert_bytes != 0; + return true; } -static bool ds4_streaming_routed_expert_bytes( - const ds4_weights *weights, - uint64_t *per_expert_bytes_out) { - if (per_expert_bytes_out) *per_expert_bytes_out = 0; - if (!weights || !per_expert_bytes_out) return false; - - /* Mixed-precision models can put an outlier quant at the first routed - * layer owned by a distributed slice. Choosing that first layer as the - * slab class makes every ordinary layer bypass the cache. Use the most - * common local size class instead (ties retain the earliest class). */ - uint64_t best_bytes = 0; - uint32_t best_count = 0; +static const ds4_layer_weights *weights_first_bound_layer(const ds4_weights *w) { + if (!w) return NULL; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - uint64_t candidate = 0; - if (!streaming_layer_routed_expert_bytes(&weights->layer[il], - &candidate)) { - continue; - } - uint32_t count = 0; - for (uint32_t jl = 0; jl < DS4_N_LAYER; jl++) { - uint64_t bytes = 0; - if (streaming_layer_routed_expert_bytes(&weights->layer[jl], - &bytes) && - bytes == candidate) { - count++; - } - } - if (count > best_count) { - best_bytes = candidate; - best_count = count; - } + if (weights_layer_has_required(&w->layer[il], il)) return &w->layer[il]; } - if (best_count == 0) return false; - *per_expert_bytes_out = best_bytes; - return true; + return NULL; } -enum { DS4_STREAMING_PREFILL_HEADROOM_LAYERS = 2 }; +/* Verify every tensor type and dimension used by the specialized pipeline. + * For distributed sliced GGUFs, only the advertised local layer range is + * required; token embedding and output head are validated when present. */ +static void weights_validate_glm_dsa_layout( + const ds4_weights *w, + uint32_t layer_start, + uint32_t layer_end, + bool require_token_embd, + bool require_output) { + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; + const uint64_t q_nope = DS4_N_KEY_MLA - DS4_N_ROT; + const uint64_t index_q_dim = + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; -static bool ds4_streaming_cacheable_expert_count( - const ds4_weights *weights, - uint64_t *experts_out, - uint32_t *layers_out) { - if (experts_out) *experts_out = 0; - if (layers_out) *layers_out = 0; - if (!weights || !experts_out || DS4_N_EXPERT == 0) return false; - - uint64_t slab_bytes = 0; - if (!ds4_streaming_routed_expert_bytes(weights, &slab_bytes)) return false; - - uint32_t layers = 0; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - uint64_t per_expert_bytes = 0; - if (!streaming_layer_routed_expert_bytes(&weights->layer[il], - &per_expert_bytes)) { - continue; - } - if (per_expert_bytes == slab_bytes) layers++; + if (!w) ds4_die("internal error: missing weights while validating GLM layout"); + if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in GLM weight layout validation"); + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) { + ds4_die("invalid layer range in GLM weight layout validation"); } - if (layers == 0 || - (uint64_t)layers > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { - return false; + if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); + if (w->token_embd) { + tensor_expect_glm_dense_quant_layout(w->token_embd, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); } - *experts_out = (uint64_t)layers * (uint64_t)DS4_N_EXPERT; - if (layers_out) *layers_out = layers; - return true; -} - -static bool ds4_streaming_prefill_headroom_bytes( - const ds4_weights *weights, - uint64_t *bytes_out) { - if (bytes_out) *bytes_out = 0; - if (!weights || !bytes_out) return false; - uint64_t per_expert_bytes = 0; - uint64_t cacheable_experts = 0; - uint32_t cacheable_layers = 0; - if (!ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes) || - !ds4_streaming_cacheable_expert_count(weights, - &cacheable_experts, - &cacheable_layers)) { - return false; - } - (void)cacheable_experts; - const uint32_t reserve_layers = - cacheable_layers < DS4_STREAMING_PREFILL_HEADROOM_LAYERS ? - cacheable_layers : DS4_STREAMING_PREFILL_HEADROOM_LAYERS; - if (per_expert_bytes > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { - return false; + const bool have_output = weights_have_output_head(w); + if (require_output && !have_output) ds4_die("required output head tensors are missing"); + if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); + if (have_output) { + tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_glm_dense_quant_layout(w->output, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); } - const uint64_t layer_bytes = - per_expert_bytes * (uint64_t)DS4_N_EXPERT; - if (reserve_layers != 0 && - layer_bytes > UINT64_MAX / (uint64_t)reserve_layers) return false; - *bytes_out = layer_bytes * (uint64_t)reserve_layers; - return true; -} + for (uint32_t il = layer_start; il <= layer_end; il++) { + const ds4_layer_weights *l = &w->layer[il]; + if (!weights_glm_dsa_layer_has_required(l, il)) { + fprintf(stderr, "ds4: required GLM tensors for layer %u are missing\n", il); + exit(1); + } -/* - * Mixed-precision ("boosted") GGUFs upcast a few layers' routed experts to a - * bigger quant (e.g. Q4_K among IQ2 layers). The streaming expert cache is a - * single-size-class slab allocator sized from the dominant local routed-layer - * size class, so other layers can never be served from it: they must read - * expert weights through the mapped-model views instead. A layer is "uniform" - * iff its per-expert bytes match the slab class. - */ -static DS4_MAYBE_UNUSED bool weights_streaming_layer_experts_uniform( - const ds4_weights *w, - uint32_t il) { - uint64_t base = 0; - uint64_t bytes = 0; - if (!w || il >= DS4_N_LAYER) return true; - const ds4_layer_weights *l = &w->layer[il]; - if (!streaming_layer_routed_expert_bytes(l, &bytes)) return true; - if (!ds4_streaming_routed_expert_bytes(w, &base)) return true; - return bytes == base; -} + tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_glm_dense_quant_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); + tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); + tensor_expect_glm_dense_quant_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0); + tensor_expect_glm_dense_quant_layout(l->attn_kv_a_mqa, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); + tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_KV_LORA, 0, 0); + tensor_expect_glm_dense_quant_layout(l->attn_k_b, 3, q_nope, DS4_N_KV_LORA, DS4_N_HEAD); + tensor_expect_glm_dense_quant_layout(l->attn_v_b, 3, DS4_N_KV_LORA, DS4_N_VALUE_MLA, DS4_N_HEAD); + tensor_expect_glm_dense_quant_layout(l->attn_output, 2, DS4_N_HEAD * DS4_N_VALUE_MLA, DS4_N_EMBD, 0); + tensor_expect_glm_dense_quant_layout(l->indexer_attn_k, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 0); + tensor_expect_glm_dense_quant_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, index_q_dim, 0); + tensor_expect_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); + tensor_expect_layout(l->indexer_k_norm_b, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); + tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); + tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); -static uint32_t ds4_streaming_cache_experts_for_byte_budget( - const ds4_weights *weights, - uint64_t bytes, - uint64_t *per_expert_bytes_out) { - uint64_t per_expert_bytes = 0; - if (per_expert_bytes_out) *per_expert_bytes_out = 0; - if (!weights || - bytes == 0 || - !ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes)) { - return 0; + if (il < DS4_N_LEADING_DENSE) { + tensor_expect_glm_dense_quant_layout(l->ffn_gate, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); + tensor_expect_glm_dense_quant_layout(l->ffn_up, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); + tensor_expect_glm_dense_quant_layout(l->ffn_down, 2, DS4_N_FF_DENSE, DS4_N_EMBD, 0); + } else { + tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); + tensor_expect_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); + tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); + if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { + fprintf(stderr, "ds4: GLM routed gate/up experts use different quant types in layer %u\n", il); + exit(1); + } + tensor_expect_glm_dense_quant_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_glm_dense_quant_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_glm_dense_quant_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); + } + + if (DS4_N_NEXTN_PREDICT != 0 && + il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { + tensor_expect_glm_dense_quant_layout(l->nextn_eh_proj, 2, 2u * DS4_N_EMBD, DS4_N_EMBD, 0); + tensor_expect_layout(l->nextn_enorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->nextn_hnorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->nextn_shared_head_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + } } - if (per_expert_bytes_out) *per_expert_bytes_out = per_expert_bytes; - return ds4_ssd_cache_experts_for_byte_budget(bytes, per_expert_bytes); } -#ifndef DS4_NO_GPU -static ds4_gpu_stream_expert_table graph_stream_expert_table_make( - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - ds4_gpu_stream_expert_table table; - memset(&table, 0, sizeof(table)); - if (!model || !layer) return table; - table.model_map = model->map; - table.model_size = model->size; - table.layer = il; - table.n_total_expert = DS4_N_EXPERT; - table.gate_offset = layer->ffn_gate_exps ? layer->ffn_gate_exps->abs_offset : 0; - table.up_offset = layer->ffn_up_exps ? layer->ffn_up_exps->abs_offset : 0; - table.down_offset = layer->ffn_down_exps ? layer->ffn_down_exps->abs_offset : 0; - table.gate_expert_bytes = gate_expert_bytes; - table.down_expert_bytes = down_expert_bytes; - return table; -} -#endif +static void weights_validate_layout( + const ds4_weights *w, + uint32_t layer_start, + uint32_t layer_end, + bool require_token_embd, + bool require_output) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + weights_validate_glm_dsa_layout(w, + layer_start, + layer_end, + require_token_embd, + require_output); + return; + } -static uint64_t ds4_streaming_manual_cache_safe_bytes( - ds4_backend backend, - int ctx_size, - uint32_t prefill_chunk, - bool ssd_streaming) { -#ifdef DS4_NO_GPU - (void)backend; - (void)ctx_size; - (void)prefill_chunk; - (void)ssd_streaming; - return 0; -#else - const uint64_t gib = 1024ull * 1024ull * 1024ull; - const uint64_t recommended = ds4_gpu_recommended_working_set_size(); - if (recommended == 0) return 0; + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; + const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - /* - * Explicit NGB budgets name only the routed expert cache. Keep that cache - * below the graph backend's working-set recommendation after accounting for - * the graph context/KV buffers. This is intentionally not an mlock-derived - * cap: crossing too close to the recommended working set makes short - * token-major prefill spend most of its time in VM/driver synchronization. - */ - uint64_t target = recommended > UINT64_MAX / 7ull ? - UINT64_MAX : (recommended * 7ull) / 8ull; - const ds4_context_memory ctx_mem = - ds4_context_memory_estimate_with_prefill_mode(backend, - ctx_size, - prefill_chunk, - ssd_streaming); - uint64_t safe = 0; - if (target > ctx_mem.total_bytes) safe = target - ctx_mem.total_bytes; - safe = (safe / gib) * gib; - if (safe == 0) safe = gib; - return safe; -#endif -} + if (!w) ds4_die("internal error: missing weights while validating layout"); + if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in weight layout validation"); + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) { + ds4_die("invalid layer range in weight layout validation"); + } -static uint64_t ds4_add_sat_u64(uint64_t a, uint64_t b) { - return a > UINT64_MAX - b ? UINT64_MAX : a + b; -} + if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); + if (w->token_embd) { + tensor_expect_layout(w->token_embd, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); + } -static uint64_t ds4_mul_sat_u64(uint64_t a, uint64_t b) { - if (a != 0 && b > UINT64_MAX / a) return UINT64_MAX; - return a * b; -} + const bool have_output = weights_have_output_head(w); + if (require_output && !have_output) ds4_die("required output head tensors are missing"); + if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); + if (have_output) { + tensor_expect_layout(w->output_hc_base, DS4_TENSOR_F32, 1, DS4_N_HC, 0, 0); + tensor_expect_layout(w->output_hc_fn, DS4_TENSOR_F16, 2, hc_dim, DS4_N_HC, 0); + tensor_expect_layout(w->output_hc_scale, DS4_TENSOR_F32, 1, 1, 0, 0); + tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_dense_quant_layout(w->output, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); + } -static double ds4_bytes_to_gib(uint64_t bytes) { - return (double)bytes / 1073741824.0; -} + for (uint32_t il = layer_start; il <= layer_end; il++) { + const ds4_layer_weights *l = &w->layer[il]; + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (!weights_layer_has_required(l, il)) { + fprintf(stderr, "ds4: required tensors for layer %u are missing\n", il); + exit(1); + } -static void tensor_expect_routed_expert( - const ds4_tensor *t, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!t) ds4_die("internal error: missing routed expert tensor while validating layout"); - if (!tensor_is_routed_expert_type(t->type)) { + tensor_expect_layout(l->hc_attn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); + tensor_expect_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); + tensor_expect_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); + tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_dense_quant_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); + tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); + tensor_expect_dense_quant_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0); + tensor_expect_dense_quant_layout(l->attn_kv, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); + tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); + tensor_expect_layout(l->attn_sinks, DS4_TENSOR_F32, 1, DS4_N_HEAD, 0, 0); + tensor_expect_dense_quant_layout(l->attn_output_a, 2, DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), out_low_dim, 0); + tensor_expect_dense_quant_layout(l->attn_output_b, 2, out_low_dim, DS4_N_EMBD, 0); + + if (ratio != 0) { + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint64_t comp_width = (uint64_t)coff * DS4_N_HEAD_DIM; + tensor_expect_layout(l->attn_compressor_ape, DS4_TENSOR_F16, 2, comp_width, ratio, 0); + tensor_expect_layout(l->attn_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); + tensor_expect_layout(l->attn_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); + tensor_expect_layout(l->attn_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); + } + if (ratio == 4) { + const uint64_t index_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + const uint64_t index_width = 2u * DS4_N_INDEXER_HEAD_DIM; + tensor_expect_f16_or_q8_0_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, index_q_dim, 0); + tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); + tensor_expect_layout(l->indexer_compressor_ape, DS4_TENSOR_F16, 2, index_width, ratio, 0); + tensor_expect_layout(l->indexer_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); + tensor_expect_layout(l->indexer_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); + tensor_expect_layout(l->indexer_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); + } + + tensor_expect_layout(l->hc_ffn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); + tensor_expect_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); + tensor_expect_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); + tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); + tensor_expect_optional(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); + tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); + if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { + fprintf(stderr, "ds4: routed gate/up experts use different quant types in layer %u\n", il); + exit(1); + } + tensor_expect_dense_quant_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_dense_quant_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_dense_quant_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); + if (il < DS4_N_HASH_LAYER) { + tensor_expect_layout(l->ffn_gate_tid2eid, DS4_TENSOR_I32, 2, DS4_N_EXPERT_USED, DS4_N_VOCAB, 0); + } + } +} + +static void mtp_weights_validate_layout(const ds4_mtp_weights *w) { + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; + const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; + const ds4_layer_weights *l = &w->block; + + tensor_expect_layout(w->hc_head_base, DS4_TENSOR_F32, 1, DS4_N_HC, 0, 0); + tensor_expect_plain_layout(w->hc_head_fn, 2, hc_dim, DS4_N_HC, 0); + tensor_expect_layout(w->hc_head_scale, DS4_TENSOR_F32, 1, 1, 0, 0); + tensor_expect_layout(w->e_proj, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_EMBD, 0); + tensor_expect_layout(w->h_proj, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_EMBD, 0); + tensor_expect_layout(w->enorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(w->hnorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(w->norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + + tensor_expect_plain_layout(l->hc_attn_fn, 2, hc_dim, hc_mix_dim, 0); + tensor_expect_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); + tensor_expect_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); + tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->attn_q_a, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); + tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); + tensor_expect_layout(l->attn_q_b, DS4_TENSOR_Q8_0, 2, DS4_N_LORA_Q, q_dim, 0); + tensor_expect_layout(l->attn_kv, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); + tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); + tensor_expect_layout(l->attn_sinks, DS4_TENSOR_F32, 1, DS4_N_HEAD, 0, 0); + tensor_expect_layout(l->attn_output_a, DS4_TENSOR_Q8_0, 2, DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), out_low_dim, 0); + tensor_expect_layout(l->attn_output_b, DS4_TENSOR_Q8_0, 2, out_low_dim, DS4_N_EMBD, 0); + + tensor_expect_plain_layout(l->hc_ffn_fn, 2, hc_dim, hc_mix_dim, 0); + tensor_expect_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); + tensor_expect_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); + tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_plain_layout(l->ffn_gate_inp, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); + tensor_expect_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); + tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); + if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { + ds4_die("MTP routed gate/up experts use different quant types"); + } + tensor_expect_layout(l->ffn_gate_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_layout(l->ffn_up_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_layout(l->ffn_down_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); +} + +typedef enum { + DS4_DSPARK_LAYOUT_F32, + DS4_DSPARK_LAYOUT_PLAIN, + DS4_DSPARK_LAYOUT_DENSE, + DS4_DSPARK_LAYOUT_ROUTED, +} ds4_dspark_layout_kind; + +static const char *dspark_layout_kind_name(ds4_dspark_layout_kind kind) { + switch (kind) { + case DS4_DSPARK_LAYOUT_F32: return "F32"; + case DS4_DSPARK_LAYOUT_PLAIN: return "F16 or F32"; + case DS4_DSPARK_LAYOUT_DENSE: return "F16, F32, or Q8_0"; + case DS4_DSPARK_LAYOUT_ROUTED: return "routed expert quant"; + } + return "unknown"; +} + +static bool dspark_tensor_type_matches(uint32_t type, + ds4_dspark_layout_kind kind) { + switch (kind) { + case DS4_DSPARK_LAYOUT_F32: + return type == DS4_TENSOR_F32; + case DS4_DSPARK_LAYOUT_PLAIN: + return type == DS4_TENSOR_F16 || type == DS4_TENSOR_F32; + case DS4_DSPARK_LAYOUT_DENSE: + return type == DS4_TENSOR_F16 || + type == DS4_TENSOR_F32 || + type == DS4_TENSOR_Q8_0; + case DS4_DSPARK_LAYOUT_ROUTED: + return tensor_is_routed_expert_type(type); + } + return false; +} + +static void dspark_validate_tensor_layout( + ds4_dspark_weights *dw, + const ds4_tensor *t, + const char *role, + ds4_dspark_layout_kind kind, + uint32_t ndim, + uint64_t d0, + uint64_t d1, + uint64_t d2) { + if (!dw || !t) return; + + bool ok = true; + if (!dspark_tensor_type_matches(t->type, kind)) { fprintf(stderr, - "ds4: tensor %.*s has type %u (%s), expected a routed expert quant type\n", + "ds4: DSpark tensor %.*s (%s) has type %s, expected %s\n", (int)t->name.len, t->name.ptr, - t->type, - tensor_type_name(t->type)); - exit(1); + role, + tensor_type_name(t->type), + dspark_layout_kind_name(kind)); + ok = false; } if (t->ndim != ndim) { fprintf(stderr, - "ds4: tensor %.*s has %u dimensions, expected %u\n", + "ds4: DSpark tensor %.*s (%s) has %u dimensions, expected %u\n", (int)t->name.len, t->name.ptr, + role, t->ndim, ndim); - exit(1); + ok = false; } const uint64_t want[3] = { d0, d1, d2 }; - for (uint32_t i = 0; i < ndim; i++) { + const uint32_t n = t->ndim < ndim ? t->ndim : ndim; + for (uint32_t i = 0; i < n; i++) { if (t->dim[i] == want[i]) continue; fprintf(stderr, - "ds4: tensor %.*s has dim[%u]=%" PRIu64 ", expected %" PRIu64 "\n", + "ds4: DSpark tensor %.*s (%s) has dim[%u]=%" PRIu64 + ", expected %" PRIu64 "\n", (int)t->name.len, t->name.ptr, + role, i, t->dim[i], want[i]); - exit(1); + ok = false; } + if (!ok) dw->invalid_tensors++; } -static bool weights_have_output_head(const ds4_weights *w) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - return w && w->output_norm && w->output; - } - return w && - w->output_hc_base && - w->output_hc_fn && - w->output_hc_scale && - w->output_norm && - w->output; +static void dspark_weights_note_metadata_error( + ds4_dspark_weights *dw, + const char *msg) { + if (!dw) return; + fprintf(stderr, "ds4: DSpark metadata error: %s\n", msg); + dw->metadata_errors++; } -static bool weights_have_partial_output_head(const ds4_weights *w) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - return w && (w->output_norm || w->output); +static void dspark_weights_validate_metadata(ds4_dspark_weights *dw) { + if (!dw) return; + if (!dw->has_block_size || dw->block_size == 0) { + dspark_weights_note_metadata_error(dw, "missing or zero block size"); + } else if (dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE) { + dspark_weights_note_metadata_error(dw, "block size exceeds runtime limit"); } - return w && - (w->output_hc_base || - w->output_hc_fn || - w->output_hc_scale || - w->output_norm || - w->output); -} - -static bool weights_glm_dsa_layer_has_required(const ds4_layer_weights *l, uint32_t il) { - if (!l) return false; - if (!l->attn_norm || - !l->attn_q_a || - !l->attn_q_a_norm || - !l->attn_q_b || - !l->attn_kv_a_mqa || - !l->attn_kv_a_norm || - !l->attn_k_b || - !l->attn_v_b || - !l->attn_output || - !l->indexer_attn_q_b || - !l->indexer_attn_k || - !l->indexer_k_norm || - !l->indexer_k_norm_b || - !l->indexer_proj || - !l->ffn_norm) - { - return false; + if (!dw->has_markov_rank || dw->markov_rank == 0) { + dspark_weights_note_metadata_error(dw, "missing or zero Markov rank"); } - - if (il < DS4_N_LEADING_DENSE) { - if (!l->ffn_gate || !l->ffn_up || !l->ffn_down) return false; - } else { - if (!l->ffn_gate_inp || - !l->ffn_exp_probs_b || - !l->ffn_gate_exps || - !l->ffn_up_exps || - !l->ffn_down_exps || - !l->ffn_gate_shexp || - !l->ffn_up_shexp || - !l->ffn_down_shexp) - { - return false; - } + if (!dw->has_noise_token_id || dw->noise_token_id >= DS4_N_VOCAB) { + dspark_weights_note_metadata_error(dw, "missing or out-of-range noise token"); + } + if (!dw->has_target_layers || dw->target_layer_count == 0) { + dspark_weights_note_metadata_error(dw, "missing target layer list"); + return; } - if (DS4_N_NEXTN_PREDICT != 0 && - il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER && - (!l->nextn_eh_proj || - !l->nextn_enorm || - !l->nextn_hnorm || - !l->nextn_shared_head_norm)) - { - return false; + uint32_t prev = UINT32_MAX; + for (uint32_t i = 0; i < dw->target_layer_count; i++) { + const uint32_t layer = dw->target_layers[i]; + if (layer >= DS4_N_LAYER) { + dspark_weights_note_metadata_error(dw, "target layer is outside the target model"); + } + if (i != 0 && layer <= prev) { + dspark_weights_note_metadata_error(dw, "target layers are not strictly increasing"); + } + prev = layer; } - - return true; } -static bool weights_layer_has_required(const ds4_layer_weights *l, uint32_t il) { - if (!l) return false; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - return weights_glm_dsa_layer_has_required(l, il); - } - if (!l->hc_attn_fn || - !l->hc_attn_scale || - !l->hc_attn_base || - !l->attn_norm || - !l->attn_q_a || - !l->attn_q_a_norm || - !l->attn_q_b || - !l->attn_kv || - !l->attn_kv_a_norm || - !l->attn_sinks || - !l->attn_output_a || - !l->attn_output_b || - !l->hc_ffn_fn || - !l->hc_ffn_scale || - !l->hc_ffn_base || - !l->ffn_norm || - !l->ffn_gate_inp || - !l->ffn_gate_exps || - !l->ffn_up_exps || - !l->ffn_down_exps || - !l->ffn_gate_shexp || - !l->ffn_up_shexp || - !l->ffn_down_shexp) - { - return false; - } - - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0 && - (!l->attn_compressor_ape || - !l->attn_compressor_kv || - !l->attn_compressor_gate || - !l->attn_compressor_norm)) - { - return false; - } - if (ratio == 4 && - (!l->indexer_attn_q_b || - !l->indexer_proj || - !l->indexer_compressor_ape || - !l->indexer_compressor_kv || - !l->indexer_compressor_gate || - !l->indexer_compressor_norm)) - { - return false; - } - if (il < DS4_N_HASH_LAYER && !l->ffn_gate_tid2eid) return false; - return true; -} +static void dspark_weights_validate_block_layout( + ds4_dspark_weights *dw, + const ds4_layer_weights *l) { + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; + const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; -static bool weights_layers_bound(const ds4_weights *w, uint32_t layer_start, uint32_t layer_end) { - if (!w || layer_start >= DS4_N_LAYER) return false; - if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; - if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; - for (uint32_t il = layer_start; il <= layer_end; il++) { - if (!weights_layer_has_required(&w->layer[il], il)) return false; - } - return true; -} + dspark_validate_tensor_layout(dw, l->hc_attn_fn, "hc_attn_fn", + DS4_DSPARK_LAYOUT_PLAIN, 2, + hc_dim, hc_mix_dim, 0); + dspark_validate_tensor_layout(dw, l->hc_attn_scale, "hc_attn_scale", + DS4_DSPARK_LAYOUT_F32, 1, 3, 0, 0); + dspark_validate_tensor_layout(dw, l->hc_attn_base, "hc_attn_base", + DS4_DSPARK_LAYOUT_F32, 1, + hc_mix_dim, 0, 0); + dspark_validate_tensor_layout(dw, l->attn_norm, "attn_norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_EMBD, 0, 0); + dspark_validate_tensor_layout(dw, l->attn_q_a, "attn_q_a", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_EMBD, DS4_N_LORA_Q, 0); + dspark_validate_tensor_layout(dw, l->attn_q_a_norm, "attn_q_a_norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_LORA_Q, 0, 0); + dspark_validate_tensor_layout(dw, l->attn_q_b, "attn_q_b", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_LORA_Q, q_dim, 0); + dspark_validate_tensor_layout(dw, l->attn_kv, "attn_kv", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_EMBD, DS4_N_HEAD_DIM, 0); + dspark_validate_tensor_layout(dw, l->attn_kv_a_norm, "attn_kv_a_norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_HEAD_DIM, 0, 0); + dspark_validate_tensor_layout(dw, l->attn_sinks, "attn_sinks", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_HEAD, 0, 0); + dspark_validate_tensor_layout(dw, l->attn_output_a, "attn_output_a", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), + out_low_dim, 0); + dspark_validate_tensor_layout(dw, l->attn_output_b, "attn_output_b", + DS4_DSPARK_LAYOUT_DENSE, 2, + out_low_dim, DS4_N_EMBD, 0); -static const ds4_layer_weights *weights_first_bound_layer(const ds4_weights *w) { - if (!w) return NULL; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (weights_layer_has_required(&w->layer[il], il)) return &w->layer[il]; + dspark_validate_tensor_layout(dw, l->hc_ffn_fn, "hc_ffn_fn", + DS4_DSPARK_LAYOUT_PLAIN, 2, + hc_dim, hc_mix_dim, 0); + dspark_validate_tensor_layout(dw, l->hc_ffn_scale, "hc_ffn_scale", + DS4_DSPARK_LAYOUT_F32, 1, 3, 0, 0); + dspark_validate_tensor_layout(dw, l->hc_ffn_base, "hc_ffn_base", + DS4_DSPARK_LAYOUT_F32, 1, + hc_mix_dim, 0, 0); + dspark_validate_tensor_layout(dw, l->ffn_norm, "ffn_norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_EMBD, 0, 0); + dspark_validate_tensor_layout(dw, l->ffn_gate_inp, "ffn_gate_inp", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_EMBD, DS4_N_EXPERT, 0); + dspark_validate_tensor_layout(dw, l->ffn_exp_probs_b, "exp_probs_b", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_EXPERT, 0, 0); + dspark_validate_tensor_layout(dw, l->ffn_gate_exps, "ffn_gate_exps", + DS4_DSPARK_LAYOUT_ROUTED, 3, + DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + dspark_validate_tensor_layout(dw, l->ffn_up_exps, "ffn_up_exps", + DS4_DSPARK_LAYOUT_ROUTED, 3, + DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + dspark_validate_tensor_layout(dw, l->ffn_down_exps, "ffn_down_exps", + DS4_DSPARK_LAYOUT_ROUTED, 3, + DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); + if (l->ffn_gate_exps && + l->ffn_up_exps && + l->ffn_gate_exps->type != l->ffn_up_exps->type) { + fprintf(stderr, + "ds4: DSpark routed gate/up experts use different quant types\n"); + dw->invalid_tensors++; } - return NULL; + dspark_validate_tensor_layout(dw, l->ffn_gate_shexp, "ffn_gate_shexp", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_EMBD, DS4_N_FF_EXP, 0); + dspark_validate_tensor_layout(dw, l->ffn_up_shexp, "ffn_up_shexp", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_EMBD, DS4_N_FF_EXP, 0); + dspark_validate_tensor_layout(dw, l->ffn_down_shexp, "ffn_down_shexp", + DS4_DSPARK_LAYOUT_DENSE, 2, + DS4_N_FF_EXP, DS4_N_EMBD, 0); } -/* Verify every tensor type and dimension used by the specialized pipeline. - * For distributed sliced GGUFs, only the advertised local layer range is - * required; token embedding and output head are validated when present. */ -static void weights_validate_glm_dsa_layout( - const ds4_weights *w, - uint32_t layer_start, - uint32_t layer_end, - bool require_token_embd, - bool require_output) { - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; - const uint64_t q_nope = DS4_N_KEY_MLA - DS4_N_ROT; - const uint64_t index_q_dim = - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - - if (!w) ds4_die("internal error: missing weights while validating GLM layout"); - if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in GLM weight layout validation"); - if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; - if (layer_end >= DS4_N_LAYER || layer_end < layer_start) { - ds4_die("invalid layer range in GLM weight layout validation"); - } - - if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); - if (w->token_embd) { - tensor_expect_glm_dense_quant_layout(w->token_embd, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); - } +static void dspark_weights_validate_layout(ds4_dspark_weights *dw) { + if (!dw) return; + dspark_weights_validate_metadata(dw); - const bool have_output = weights_have_output_head(w); - if (require_output && !have_output) ds4_die("required output head tensors are missing"); - if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); - if (have_output) { - tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_glm_dense_quant_layout(w->output, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + ds4_dspark_stage_weights *sw = &dw->stage[stage]; + dspark_weights_validate_block_layout(dw, &sw->block); + if (stage == 0) { + dspark_validate_tensor_layout(dw, sw->main_proj, "main_proj", + DS4_DSPARK_LAYOUT_DENSE, 2, + (uint64_t)dw->target_layer_count * + DS4_N_EMBD, + DS4_N_EMBD, 0); + dspark_validate_tensor_layout(dw, sw->main_norm, "main_norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_EMBD, 0, 0); + } } - for (uint32_t il = layer_start; il <= layer_end; il++) { - const ds4_layer_weights *l = &w->layer[il]; - if (!weights_glm_dsa_layer_has_required(l, il)) { - fprintf(stderr, "ds4: required GLM tensors for layer %u are missing\n", il); - exit(1); - } + if (dw->n_stages == 0) return; + ds4_dspark_stage_weights *final = &dw->stage[dw->n_stages - 1u]; + dspark_validate_tensor_layout(dw, final->norm, "norm", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_EMBD, 0, 0); + dspark_validate_tensor_layout(dw, final->hc_head_base, "hc_head_base", + DS4_DSPARK_LAYOUT_F32, 1, + DS4_N_HC, 0, 0); + dspark_validate_tensor_layout(dw, final->hc_head_fn, "hc_head_fn", + DS4_DSPARK_LAYOUT_PLAIN, 2, + (uint64_t)DS4_N_EMBD * DS4_N_HC, + DS4_N_HC, 0); + dspark_validate_tensor_layout(dw, final->hc_head_scale, "hc_head_scale", + DS4_DSPARK_LAYOUT_F32, 1, 1, 0, 0); + dspark_validate_tensor_layout(dw, final->markov_w1, "markov_w1", + DS4_DSPARK_LAYOUT_DENSE, 2, + dw->markov_rank, DS4_N_VOCAB, 0); + dspark_validate_tensor_layout(dw, final->markov_w2, "markov_w2", + DS4_DSPARK_LAYOUT_DENSE, 2, + dw->markov_rank, DS4_N_VOCAB, 0); + dspark_validate_tensor_layout(dw, final->confidence_proj, + "confidence_proj", + DS4_DSPARK_LAYOUT_DENSE, 2, + (uint64_t)DS4_N_EMBD + dw->markov_rank, + 1, 0); +} - tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_glm_dense_quant_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); - tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); - tensor_expect_glm_dense_quant_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0); - tensor_expect_glm_dense_quant_layout(l->attn_kv_a_mqa, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); - tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_KV_LORA, 0, 0); - tensor_expect_glm_dense_quant_layout(l->attn_k_b, 3, q_nope, DS4_N_KV_LORA, DS4_N_HEAD); - tensor_expect_glm_dense_quant_layout(l->attn_v_b, 3, DS4_N_KV_LORA, DS4_N_VALUE_MLA, DS4_N_HEAD); - tensor_expect_glm_dense_quant_layout(l->attn_output, 2, DS4_N_HEAD * DS4_N_VALUE_MLA, DS4_N_EMBD, 0); - tensor_expect_glm_dense_quant_layout(l->indexer_attn_k, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 0); - tensor_expect_glm_dense_quant_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, index_q_dim, 0); - tensor_expect_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); - tensor_expect_layout(l->indexer_k_norm_b, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); - tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); - tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - - if (il < DS4_N_LEADING_DENSE) { - tensor_expect_glm_dense_quant_layout(l->ffn_gate, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); - tensor_expect_glm_dense_quant_layout(l->ffn_up, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); - tensor_expect_glm_dense_quant_layout(l->ffn_down, 2, DS4_N_FF_DENSE, DS4_N_EMBD, 0); - } else { - tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); - tensor_expect_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); - tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); - if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { - fprintf(stderr, "ds4: GLM routed gate/up experts use different quant types in layer %u\n", il); - exit(1); - } - tensor_expect_glm_dense_quant_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_glm_dense_quant_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_glm_dense_quant_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); - } - - if (DS4_N_NEXTN_PREDICT != 0 && - il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { - tensor_expect_glm_dense_quant_layout(l->nextn_eh_proj, 2, 2u * DS4_N_EMBD, DS4_N_EMBD, 0); - tensor_expect_layout(l->nextn_enorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(l->nextn_hnorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(l->nextn_shared_head_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - } - } +static bool ds4_shape_matches_metadata( + const ds4_shape *s, + uint32_t n_layer, + uint32_t n_embd, + uint32_t n_vocab, + uint32_t n_head, + uint32_t n_head_kv, + uint32_t n_head_dim, + uint32_t n_value_dim, + uint32_t n_rot, + uint32_t n_lora_q, + uint32_t n_lora_o, + uint32_t n_out_group, + uint32_t n_expert, + uint32_t n_expert_used, + uint32_t n_ff_exp, + uint32_t n_expert_shared, + uint32_t n_hash_layer, + uint32_t n_swa, + uint32_t n_indexer_head, + uint32_t n_indexer_head_dim, + uint32_t n_indexer_top_k, + uint32_t n_hc, + uint32_t n_hc_sinkhorn_iter) { + return s->n_layer == n_layer && + s->n_embd == n_embd && + s->n_vocab == n_vocab && + s->n_head == n_head && + s->n_head_kv == n_head_kv && + s->n_head_dim == n_head_dim && + s->n_value_dim == n_value_dim && + s->n_rot == n_rot && + s->n_lora_q == n_lora_q && + s->n_lora_o == n_lora_o && + s->n_out_group == n_out_group && + s->n_expert == n_expert && + s->n_expert_used == n_expert_used && + s->n_ff_exp == n_ff_exp && + s->n_expert_shared == n_expert_shared && + s->n_hash_layer == n_hash_layer && + s->n_swa == n_swa && + s->n_indexer_head == n_indexer_head && + s->n_indexer_head_dim == n_indexer_head_dim && + s->n_indexer_top_k == n_indexer_top_k && + s->n_hc == n_hc && + s->n_hc_sinkhorn_iter == n_hc_sinkhorn_iter; } -static void weights_validate_layout( - const ds4_weights *w, - uint32_t layer_start, - uint32_t layer_end, - bool require_token_embd, - bool require_output) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - weights_validate_glm_dsa_layout(w, - layer_start, - layer_end, - require_token_embd, - require_output); +static void ds4_select_shape_from_metadata( + uint32_t n_layer, + uint32_t n_embd, + uint32_t n_vocab, + uint32_t n_head, + uint32_t n_head_kv, + uint32_t n_head_dim, + uint32_t n_value_dim, + uint32_t n_rot, + uint32_t n_lora_q, + uint32_t n_lora_o, + uint32_t n_out_group, + uint32_t n_expert, + uint32_t n_expert_used, + uint32_t n_ff_exp, + uint32_t n_expert_shared, + uint32_t n_hash_layer, + uint32_t n_swa, + uint32_t n_indexer_head, + uint32_t n_indexer_head_dim, + uint32_t n_indexer_top_k, + uint32_t n_hc, + uint32_t n_hc_sinkhorn_iter) { + if (ds4_shape_matches_metadata(&DS4_SHAPE_FLASH, + n_layer, n_embd, n_vocab, n_head, n_head_kv, + n_head_dim, n_value_dim, n_rot, n_lora_q, + n_lora_o, n_out_group, n_expert, + n_expert_used, n_ff_exp, n_expert_shared, + n_hash_layer, n_swa, n_indexer_head, + n_indexer_head_dim, n_indexer_top_k, n_hc, + n_hc_sinkhorn_iter)) { + g_ds4_shape = DS4_SHAPE_FLASH; return; } - - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; - const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - - if (!w) ds4_die("internal error: missing weights while validating layout"); - if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in weight layout validation"); - if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; - if (layer_end >= DS4_N_LAYER || layer_end < layer_start) { - ds4_die("invalid layer range in weight layout validation"); + if (ds4_shape_matches_metadata(&DS4_SHAPE_PRO, + n_layer, n_embd, n_vocab, n_head, n_head_kv, + n_head_dim, n_value_dim, n_rot, n_lora_q, + n_lora_o, n_out_group, n_expert, + n_expert_used, n_ff_exp, n_expert_shared, + n_hash_layer, n_swa, n_indexer_head, + n_indexer_head_dim, n_indexer_top_k, n_hc, + n_hc_sinkhorn_iter)) { + g_ds4_shape = DS4_SHAPE_PRO; + return; } - if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); - if (w->token_embd) { - tensor_expect_layout(w->token_embd, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); - } + fprintf(stderr, + "ds4: unsupported DeepSeek4 shape: layers=%u embd=%u heads=%u " + "q_lora=%u out_groups=%u experts=%u ff_exp=%u indexer_top_k=%u\n", + n_layer, + n_embd, + n_head, + n_lora_q, + n_out_group, + n_expert, + n_ff_exp, + n_indexer_top_k); + exit(1); +} - const bool have_output = weights_have_output_head(w); - if (require_output && !have_output) ds4_die("required output head tensors are missing"); - if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); - if (have_output) { - tensor_expect_layout(w->output_hc_base, DS4_TENSOR_F32, 1, DS4_N_HC, 0, 0); - tensor_expect_layout(w->output_hc_fn, DS4_TENSOR_F16, 2, hc_dim, DS4_N_HC, 0); - tensor_expect_layout(w->output_hc_scale, DS4_TENSOR_F32, 1, 1, 0, 0); - tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_dense_quant_layout(w->output, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); +static void validate_compress_ratio_metadata(const ds4_model *m) { + const char *key = "deepseek4.attention.compress_ratios"; + ds4_array_ref arr; + if (!model_get_array(m, key, &arr) || + (arr.type != GGUF_VALUE_UINT32 && arr.type != GGUF_VALUE_INT32)) { + fprintf(stderr, "ds4: required int32/uint32 array metadata key is missing: %s\n", key); + exit(1); + } + if (arr.len < DS4_N_LAYER) { + ds4_die("deepseek4.attention.compress_ratios is shorter than the layer count"); } - for (uint32_t il = layer_start; il <= layer_end; il++) { - const ds4_layer_weights *l = &w->layer[il]; - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (!weights_layer_has_required(l, il)) { - fprintf(stderr, "ds4: required tensors for layer %u are missing\n", il); - exit(1); - } - - tensor_expect_layout(l->hc_attn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); - tensor_expect_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); - tensor_expect_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); - tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_dense_quant_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); - tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); - tensor_expect_dense_quant_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0); - tensor_expect_dense_quant_layout(l->attn_kv, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); - tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); - tensor_expect_layout(l->attn_sinks, DS4_TENSOR_F32, 1, DS4_N_HEAD, 0, 0); - tensor_expect_dense_quant_layout(l->attn_output_a, 2, DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), out_low_dim, 0); - tensor_expect_dense_quant_layout(l->attn_output_b, 2, out_low_dim, DS4_N_EMBD, 0); - - if (ratio != 0) { - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint64_t comp_width = (uint64_t)coff * DS4_N_HEAD_DIM; - tensor_expect_layout(l->attn_compressor_ape, DS4_TENSOR_F16, 2, comp_width, ratio, 0); - tensor_expect_layout(l->attn_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); - tensor_expect_layout(l->attn_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); - tensor_expect_layout(l->attn_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); - } - if (ratio == 4) { - const uint64_t index_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - const uint64_t index_width = 2u * DS4_N_INDEXER_HEAD_DIM; - tensor_expect_f16_or_q8_0_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, index_q_dim, 0); - tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); - tensor_expect_layout(l->indexer_compressor_ape, DS4_TENSOR_F16, 2, index_width, ratio, 0); - tensor_expect_layout(l->indexer_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); - tensor_expect_layout(l->indexer_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); - tensor_expect_layout(l->indexer_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); + memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); + ds4_cursor c = cursor_at(m, arr.data_pos); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + uint32_t got = 0; + if (arr.type == GGUF_VALUE_UINT32) { + if (!cursor_u32(&c, &got)) ds4_die(c.error); + } else { + int32_t v = 0; + if (!cursor_read(&c, &v, sizeof(v))) ds4_die(c.error); + if (v < 0) ds4_die("metadata array contains a negative value"); + got = (uint32_t)v; } - tensor_expect_layout(l->hc_ffn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); - tensor_expect_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); - tensor_expect_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); - tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); - tensor_expect_optional(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); - tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); - if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { - fprintf(stderr, "ds4: routed gate/up experts use different quant types in layer %u\n", il); + const uint32_t expected = ds4_expected_layer_compress_ratio(il); + if (got != expected) { + fprintf(stderr, + "ds4: unexpected DeepSeek4 compression ratio at layer %u for %s: got %u, expected %u\n", + il, DS4_MODEL_SHAPE_NAME, got, expected); exit(1); } - tensor_expect_dense_quant_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_dense_quant_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_dense_quant_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); - if (il < DS4_N_HASH_LAYER) { - tensor_expect_layout(l->ffn_gate_tid2eid, DS4_TENSOR_I32, 2, DS4_N_EXPERT_USED, DS4_N_VOCAB, 0); - } + g_ds4_compress_ratios[il] = got; } } -static void mtp_weights_validate_layout(const ds4_mtp_weights *w) { - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; - const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - const ds4_layer_weights *l = &w->block; - - tensor_expect_layout(w->hc_head_base, DS4_TENSOR_F32, 1, DS4_N_HC, 0, 0); - tensor_expect_plain_layout(w->hc_head_fn, 2, hc_dim, DS4_N_HC, 0); - tensor_expect_layout(w->hc_head_scale, DS4_TENSOR_F32, 1, 1, 0, 0); - tensor_expect_layout(w->e_proj, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_EMBD, 0); - tensor_expect_layout(w->h_proj, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_EMBD, 0); - tensor_expect_layout(w->enorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(w->hnorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(w->norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - - tensor_expect_plain_layout(l->hc_attn_fn, 2, hc_dim, hc_mix_dim, 0); - tensor_expect_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); - tensor_expect_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); - tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_layout(l->attn_q_a, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); - tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); - tensor_expect_layout(l->attn_q_b, DS4_TENSOR_Q8_0, 2, DS4_N_LORA_Q, q_dim, 0); - tensor_expect_layout(l->attn_kv, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); - tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); - tensor_expect_layout(l->attn_sinks, DS4_TENSOR_F32, 1, DS4_N_HEAD, 0, 0); - tensor_expect_layout(l->attn_output_a, DS4_TENSOR_Q8_0, 2, DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), out_low_dim, 0); - tensor_expect_layout(l->attn_output_b, DS4_TENSOR_Q8_0, 2, out_low_dim, DS4_N_EMBD, 0); +static void config_expect_f32(const char *name, float got, float expected); - tensor_expect_plain_layout(l->hc_ffn_fn, 2, hc_dim, hc_mix_dim, 0); - tensor_expect_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); - tensor_expect_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); - tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); - tensor_expect_plain_layout(l->ffn_gate_inp, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); - tensor_expect_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); - tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); - if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { - ds4_die("MTP routed gate/up experts use different quant types"); +static void validate_swiglu_clamp_metadata(const ds4_model *m) { + const char *key = "deepseek4.swiglu_clamp_exp"; + ds4_array_ref arr; + if (!model_get_array(m, key, &arr) || + (arr.type != GGUF_VALUE_FLOAT32 && arr.type != GGUF_VALUE_FLOAT64)) { + fprintf(stderr, "ds4: required float array metadata key is missing: %s\n", key); + exit(1); } - tensor_expect_layout(l->ffn_gate_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_layout(l->ffn_up_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); - tensor_expect_layout(l->ffn_down_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); -} - -typedef enum { - DS4_DSPARK_LAYOUT_F32, - DS4_DSPARK_LAYOUT_PLAIN, - DS4_DSPARK_LAYOUT_DENSE, - DS4_DSPARK_LAYOUT_ROUTED, -} ds4_dspark_layout_kind; - -static const char *dspark_layout_kind_name(ds4_dspark_layout_kind kind) { - switch (kind) { - case DS4_DSPARK_LAYOUT_F32: return "F32"; - case DS4_DSPARK_LAYOUT_PLAIN: return "F16 or F32"; - case DS4_DSPARK_LAYOUT_DENSE: return "F16, F32, or Q8_0"; - case DS4_DSPARK_LAYOUT_ROUTED: return "routed expert quant"; + if (arr.len < DS4_N_LAYER) { + ds4_die("deepseek4.swiglu_clamp_exp is shorter than the layer count"); } - return "unknown"; -} -static bool dspark_tensor_type_matches(uint32_t type, - ds4_dspark_layout_kind kind) { - switch (kind) { - case DS4_DSPARK_LAYOUT_F32: - return type == DS4_TENSOR_F32; - case DS4_DSPARK_LAYOUT_PLAIN: - return type == DS4_TENSOR_F16 || type == DS4_TENSOR_F32; - case DS4_DSPARK_LAYOUT_DENSE: - return type == DS4_TENSOR_F16 || - type == DS4_TENSOR_F32 || - type == DS4_TENSOR_Q8_0; - case DS4_DSPARK_LAYOUT_ROUTED: - return tensor_is_routed_expert_type(type); + ds4_cursor c = cursor_at(m, arr.data_pos); + for (uint32_t i = 0; i < DS4_N_LAYER; i++) { + float got = 0.0f; + if (arr.type == GGUF_VALUE_FLOAT32) { + if (!cursor_read(&c, &got, sizeof(got))) ds4_die(c.error); + } else { + double v = 0.0; + if (!cursor_read(&c, &v, sizeof(v))) ds4_die(c.error); + got = (float)v; + } + config_expect_f32("swiglu_clamp_exp", got, DS4_SWIGLU_CLAMP_EXP); } - return false; } -static void dspark_validate_tensor_layout( - ds4_dspark_weights *dw, - const ds4_tensor *t, - const char *role, - ds4_dspark_layout_kind kind, - uint32_t ndim, - uint64_t d0, - uint64_t d1, - uint64_t d2) { - if (!dw || !t) return; - - bool ok = true; - if (!dspark_tensor_type_matches(t->type, kind)) { - fprintf(stderr, - "ds4: DSpark tensor %.*s (%s) has type %s, expected %s\n", - (int)t->name.len, - t->name.ptr, - role, - tensor_type_name(t->type), - dspark_layout_kind_name(kind)); - ok = false; - } - if (t->ndim != ndim) { - fprintf(stderr, - "ds4: DSpark tensor %.*s (%s) has %u dimensions, expected %u\n", - (int)t->name.len, - t->name.ptr, - role, - t->ndim, - ndim); - ok = false; - } +static void config_expect_u32(const char *name, uint32_t got, uint32_t expected) { + if (got == expected) return; + fprintf(stderr, "ds4: expected %s=%u for %s, got %u\n", + name, expected, DS4_MODEL_SHAPE_NAME, got); + exit(1); +} - const uint64_t want[3] = { d0, d1, d2 }; - const uint32_t n = t->ndim < ndim ? t->ndim : ndim; - for (uint32_t i = 0; i < n; i++) { - if (t->dim[i] == want[i]) continue; - fprintf(stderr, - "ds4: DSpark tensor %.*s (%s) has dim[%u]=%" PRIu64 - ", expected %" PRIu64 "\n", - (int)t->name.len, - t->name.ptr, - role, - i, - t->dim[i], - want[i]); - ok = false; - } - if (!ok) dw->invalid_tensors++; +static void config_expect_u64(const char *name, uint64_t got, uint64_t expected) { + if (got == expected) return; + fprintf(stderr, "ds4: expected %s=%" PRIu64 " for %s, got %" PRIu64 "\n", + name, expected, DS4_MODEL_SHAPE_NAME, got); + exit(1); } -static void dspark_weights_note_metadata_error( - ds4_dspark_weights *dw, - const char *msg) { - if (!dw) return; - fprintf(stderr, "ds4: DSpark metadata error: %s\n", msg); - dw->metadata_errors++; +static void config_expect_f32(const char *name, float got, float expected) { + const float scale = fabsf(expected) > 1.0f ? fabsf(expected) : 1.0f; + if (fabsf(got - expected) <= scale * 1.0e-6f) return; + fprintf(stderr, "ds4: expected %s=%.9g for %s, got %.9g\n", + name, (double)expected, DS4_MODEL_SHAPE_NAME, (double)got); + exit(1); } -static void dspark_weights_validate_metadata(ds4_dspark_weights *dw) { - if (!dw) return; - if (!dw->has_block_size || dw->block_size == 0) { - dspark_weights_note_metadata_error(dw, "missing or zero block size"); - } else if (dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE) { - dspark_weights_note_metadata_error(dw, "block size exceeds runtime limit"); - } - if (!dw->has_markov_rank || dw->markov_rank == 0) { - dspark_weights_note_metadata_error(dw, "missing or zero Markov rank"); - } - if (!dw->has_noise_token_id || dw->noise_token_id >= DS4_N_VOCAB) { - dspark_weights_note_metadata_error(dw, "missing or out-of-range noise token"); - } - if (!dw->has_target_layers || dw->target_layer_count == 0) { - dspark_weights_note_metadata_error(dw, "missing target layer list"); - return; - } +static void config_expect_bool(const char *name, bool got, bool expected) { + if (got == expected) return; + fprintf(stderr, "ds4: expected %s=%s for %s, got %s\n", + name, expected ? "true" : "false", DS4_MODEL_SHAPE_NAME, got ? "true" : "false"); + exit(1); +} - uint32_t prev = UINT32_MAX; - for (uint32_t i = 0; i < dw->target_layer_count; i++) { - const uint32_t layer = dw->target_layers[i]; - if (layer >= DS4_N_LAYER) { - dspark_weights_note_metadata_error(dw, "target layer is outside the target model"); - } - if (i != 0 && layer <= prev) { - dspark_weights_note_metadata_error(dw, "target layers are not strictly increasing"); - } - prev = layer; - } +static void config_validate_fixed_shape(uint32_t n_layer) { + config_expect_u32("block_count", n_layer, DS4_N_LAYER); } -static void dspark_weights_validate_block_layout( - ds4_dspark_weights *dw, - const ds4_layer_weights *l) { - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; - const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; +/* Validate metadata values that affect semantics: attention shape, HC count, + * expert routing, RoPE scaling, compression ratios, and SwiGLU clamp. */ +static void config_validate_deepseek4_model(const ds4_model *m) { + const uint32_t n_layer = required_u32(m, "deepseek4.block_count"); + const uint32_t n_embd = required_u32(m, "deepseek4.embedding_length"); + const uint32_t n_vocab = required_u32(m, "deepseek4.vocab_size"); + const uint32_t n_head = required_u32(m, "deepseek4.attention.head_count"); + const uint32_t n_head_kv = required_u32(m, "deepseek4.attention.head_count_kv"); + const uint32_t n_head_dim = required_u32(m, "deepseek4.attention.key_length"); + const uint32_t n_value_dim = required_u32(m, "deepseek4.attention.value_length"); + const uint32_t n_rot = required_u32(m, "deepseek4.rope.dimension_count"); + const uint32_t n_lora_q = required_u32(m, "deepseek4.attention.q_lora_rank"); + const uint32_t n_lora_o = required_u32(m, "deepseek4.attention.output_lora_rank"); + const uint32_t n_out_group = required_u32(m, "deepseek4.attention.output_group_count"); + const uint32_t n_expert = required_u32(m, "deepseek4.expert_count"); + const uint32_t n_expert_used = required_u32(m, "deepseek4.expert_used_count"); + const uint32_t n_ff_exp = required_u32(m, "deepseek4.expert_feed_forward_length"); + const uint32_t n_expert_shared = required_u32(m, "deepseek4.expert_shared_count"); + const uint32_t n_hash_layer = required_u32(m, "deepseek4.hash_layer_count"); + uint32_t n_expert_groups = 0; + uint32_t n_group_used = 0; + model_get_u32(m, "deepseek4.expert_group_count", &n_expert_groups); + model_get_u32(m, "deepseek4.expert_group_used_count", &n_group_used); + const uint32_t n_swa = required_u32(m, "deepseek4.attention.sliding_window"); + const uint32_t n_indexer_head = required_u32(m, "deepseek4.attention.indexer.head_count"); + const uint32_t n_indexer_head_dim = required_u32(m, "deepseek4.attention.indexer.key_length"); + const uint32_t n_indexer_top_k = required_u32(m, "deepseek4.attention.indexer.top_k"); + const uint32_t n_hc = required_u32(m, "deepseek4.hyper_connection.count"); + const uint32_t n_hc_sinkhorn_iter = required_u32(m, "deepseek4.hyper_connection.sinkhorn_iterations"); - dspark_validate_tensor_layout(dw, l->hc_attn_fn, "hc_attn_fn", - DS4_DSPARK_LAYOUT_PLAIN, 2, - hc_dim, hc_mix_dim, 0); - dspark_validate_tensor_layout(dw, l->hc_attn_scale, "hc_attn_scale", - DS4_DSPARK_LAYOUT_F32, 1, 3, 0, 0); - dspark_validate_tensor_layout(dw, l->hc_attn_base, "hc_attn_base", - DS4_DSPARK_LAYOUT_F32, 1, - hc_mix_dim, 0, 0); - dspark_validate_tensor_layout(dw, l->attn_norm, "attn_norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_EMBD, 0, 0); - dspark_validate_tensor_layout(dw, l->attn_q_a, "attn_q_a", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_EMBD, DS4_N_LORA_Q, 0); - dspark_validate_tensor_layout(dw, l->attn_q_a_norm, "attn_q_a_norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_LORA_Q, 0, 0); - dspark_validate_tensor_layout(dw, l->attn_q_b, "attn_q_b", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_LORA_Q, q_dim, 0); - dspark_validate_tensor_layout(dw, l->attn_kv, "attn_kv", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_EMBD, DS4_N_HEAD_DIM, 0); - dspark_validate_tensor_layout(dw, l->attn_kv_a_norm, "attn_kv_a_norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_HEAD_DIM, 0, 0); - dspark_validate_tensor_layout(dw, l->attn_sinks, "attn_sinks", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_HEAD, 0, 0); - dspark_validate_tensor_layout(dw, l->attn_output_a, "attn_output_a", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), - out_low_dim, 0); - dspark_validate_tensor_layout(dw, l->attn_output_b, "attn_output_b", - DS4_DSPARK_LAYOUT_DENSE, 2, - out_low_dim, DS4_N_EMBD, 0); - - dspark_validate_tensor_layout(dw, l->hc_ffn_fn, "hc_ffn_fn", - DS4_DSPARK_LAYOUT_PLAIN, 2, - hc_dim, hc_mix_dim, 0); - dspark_validate_tensor_layout(dw, l->hc_ffn_scale, "hc_ffn_scale", - DS4_DSPARK_LAYOUT_F32, 1, 3, 0, 0); - dspark_validate_tensor_layout(dw, l->hc_ffn_base, "hc_ffn_base", - DS4_DSPARK_LAYOUT_F32, 1, - hc_mix_dim, 0, 0); - dspark_validate_tensor_layout(dw, l->ffn_norm, "ffn_norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_EMBD, 0, 0); - dspark_validate_tensor_layout(dw, l->ffn_gate_inp, "ffn_gate_inp", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_EMBD, DS4_N_EXPERT, 0); - dspark_validate_tensor_layout(dw, l->ffn_exp_probs_b, "exp_probs_b", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_EXPERT, 0, 0); - dspark_validate_tensor_layout(dw, l->ffn_gate_exps, "ffn_gate_exps", - DS4_DSPARK_LAYOUT_ROUTED, 3, - DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - dspark_validate_tensor_layout(dw, l->ffn_up_exps, "ffn_up_exps", - DS4_DSPARK_LAYOUT_ROUTED, 3, - DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); - dspark_validate_tensor_layout(dw, l->ffn_down_exps, "ffn_down_exps", - DS4_DSPARK_LAYOUT_ROUTED, 3, - DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); - if (l->ffn_gate_exps && - l->ffn_up_exps && - l->ffn_gate_exps->type != l->ffn_up_exps->type) { - fprintf(stderr, - "ds4: DSpark routed gate/up experts use different quant types\n"); - dw->invalid_tensors++; - } - dspark_validate_tensor_layout(dw, l->ffn_gate_shexp, "ffn_gate_shexp", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_EMBD, DS4_N_FF_EXP, 0); - dspark_validate_tensor_layout(dw, l->ffn_up_shexp, "ffn_up_shexp", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_EMBD, DS4_N_FF_EXP, 0); - dspark_validate_tensor_layout(dw, l->ffn_down_shexp, "ffn_down_shexp", - DS4_DSPARK_LAYOUT_DENSE, 2, - DS4_N_FF_EXP, DS4_N_EMBD, 0); -} - -static void dspark_weights_validate_layout(ds4_dspark_weights *dw) { - if (!dw) return; - dspark_weights_validate_metadata(dw); - - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - ds4_dspark_stage_weights *sw = &dw->stage[stage]; - dspark_weights_validate_block_layout(dw, &sw->block); - if (stage == 0) { - dspark_validate_tensor_layout(dw, sw->main_proj, "main_proj", - DS4_DSPARK_LAYOUT_DENSE, 2, - (uint64_t)dw->target_layer_count * - DS4_N_EMBD, - DS4_N_EMBD, 0); - dspark_validate_tensor_layout(dw, sw->main_norm, "main_norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_EMBD, 0, 0); - } - } - - if (dw->n_stages == 0) return; - ds4_dspark_stage_weights *final = &dw->stage[dw->n_stages - 1u]; - dspark_validate_tensor_layout(dw, final->norm, "norm", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_EMBD, 0, 0); - dspark_validate_tensor_layout(dw, final->hc_head_base, "hc_head_base", - DS4_DSPARK_LAYOUT_F32, 1, - DS4_N_HC, 0, 0); - dspark_validate_tensor_layout(dw, final->hc_head_fn, "hc_head_fn", - DS4_DSPARK_LAYOUT_PLAIN, 2, - (uint64_t)DS4_N_EMBD * DS4_N_HC, - DS4_N_HC, 0); - dspark_validate_tensor_layout(dw, final->hc_head_scale, "hc_head_scale", - DS4_DSPARK_LAYOUT_F32, 1, 1, 0, 0); - dspark_validate_tensor_layout(dw, final->markov_w1, "markov_w1", - DS4_DSPARK_LAYOUT_DENSE, 2, - dw->markov_rank, DS4_N_VOCAB, 0); - dspark_validate_tensor_layout(dw, final->markov_w2, "markov_w2", - DS4_DSPARK_LAYOUT_DENSE, 2, - dw->markov_rank, DS4_N_VOCAB, 0); - dspark_validate_tensor_layout(dw, final->confidence_proj, - "confidence_proj", - DS4_DSPARK_LAYOUT_DENSE, 2, - (uint64_t)DS4_N_EMBD + dw->markov_rank, - 1, 0); -} - -static bool ds4_shape_matches_metadata( - const ds4_shape *s, - uint32_t n_layer, - uint32_t n_embd, - uint32_t n_vocab, - uint32_t n_head, - uint32_t n_head_kv, - uint32_t n_head_dim, - uint32_t n_value_dim, - uint32_t n_rot, - uint32_t n_lora_q, - uint32_t n_lora_o, - uint32_t n_out_group, - uint32_t n_expert, - uint32_t n_expert_used, - uint32_t n_ff_exp, - uint32_t n_expert_shared, - uint32_t n_hash_layer, - uint32_t n_swa, - uint32_t n_indexer_head, - uint32_t n_indexer_head_dim, - uint32_t n_indexer_top_k, - uint32_t n_hc, - uint32_t n_hc_sinkhorn_iter) { - return s->n_layer == n_layer && - s->n_embd == n_embd && - s->n_vocab == n_vocab && - s->n_head == n_head && - s->n_head_kv == n_head_kv && - s->n_head_dim == n_head_dim && - s->n_value_dim == n_value_dim && - s->n_rot == n_rot && - s->n_lora_q == n_lora_q && - s->n_lora_o == n_lora_o && - s->n_out_group == n_out_group && - s->n_expert == n_expert && - s->n_expert_used == n_expert_used && - s->n_ff_exp == n_ff_exp && - s->n_expert_shared == n_expert_shared && - s->n_hash_layer == n_hash_layer && - s->n_swa == n_swa && - s->n_indexer_head == n_indexer_head && - s->n_indexer_head_dim == n_indexer_head_dim && - s->n_indexer_top_k == n_indexer_top_k && - s->n_hc == n_hc && - s->n_hc_sinkhorn_iter == n_hc_sinkhorn_iter; -} - -static void ds4_select_shape_from_metadata( - uint32_t n_layer, - uint32_t n_embd, - uint32_t n_vocab, - uint32_t n_head, - uint32_t n_head_kv, - uint32_t n_head_dim, - uint32_t n_value_dim, - uint32_t n_rot, - uint32_t n_lora_q, - uint32_t n_lora_o, - uint32_t n_out_group, - uint32_t n_expert, - uint32_t n_expert_used, - uint32_t n_ff_exp, - uint32_t n_expert_shared, - uint32_t n_hash_layer, - uint32_t n_swa, - uint32_t n_indexer_head, - uint32_t n_indexer_head_dim, - uint32_t n_indexer_top_k, - uint32_t n_hc, - uint32_t n_hc_sinkhorn_iter) { - if (ds4_shape_matches_metadata(&DS4_SHAPE_FLASH, - n_layer, n_embd, n_vocab, n_head, n_head_kv, - n_head_dim, n_value_dim, n_rot, n_lora_q, - n_lora_o, n_out_group, n_expert, - n_expert_used, n_ff_exp, n_expert_shared, - n_hash_layer, n_swa, n_indexer_head, - n_indexer_head_dim, n_indexer_top_k, n_hc, - n_hc_sinkhorn_iter)) { - g_ds4_shape = DS4_SHAPE_FLASH; - return; - } - if (ds4_shape_matches_metadata(&DS4_SHAPE_PRO, - n_layer, n_embd, n_vocab, n_head, n_head_kv, - n_head_dim, n_value_dim, n_rot, n_lora_q, - n_lora_o, n_out_group, n_expert, - n_expert_used, n_ff_exp, n_expert_shared, - n_hash_layer, n_swa, n_indexer_head, - n_indexer_head_dim, n_indexer_top_k, n_hc, - n_hc_sinkhorn_iter)) { - g_ds4_shape = DS4_SHAPE_PRO; - return; - } - - fprintf(stderr, - "ds4: unsupported DeepSeek4 shape: layers=%u embd=%u heads=%u " - "q_lora=%u out_groups=%u experts=%u ff_exp=%u indexer_top_k=%u\n", - n_layer, - n_embd, - n_head, - n_lora_q, - n_out_group, - n_expert, - n_ff_exp, - n_indexer_top_k); - exit(1); -} - -static void validate_compress_ratio_metadata(const ds4_model *m) { - const char *key = "deepseek4.attention.compress_ratios"; - ds4_array_ref arr; - if (!model_get_array(m, key, &arr) || - (arr.type != GGUF_VALUE_UINT32 && arr.type != GGUF_VALUE_INT32)) { - fprintf(stderr, "ds4: required int32/uint32 array metadata key is missing: %s\n", key); - exit(1); - } - if (arr.len < DS4_N_LAYER) { - ds4_die("deepseek4.attention.compress_ratios is shorter than the layer count"); - } - - memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); - ds4_cursor c = cursor_at(m, arr.data_pos); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - uint32_t got = 0; - if (arr.type == GGUF_VALUE_UINT32) { - if (!cursor_u32(&c, &got)) ds4_die(c.error); - } else { - int32_t v = 0; - if (!cursor_read(&c, &v, sizeof(v))) ds4_die(c.error); - if (v < 0) ds4_die("metadata array contains a negative value"); - got = (uint32_t)v; - } - - const uint32_t expected = ds4_expected_layer_compress_ratio(il); - if (got != expected) { - fprintf(stderr, - "ds4: unexpected DeepSeek4 compression ratio at layer %u for %s: got %u, expected %u\n", - il, DS4_MODEL_SHAPE_NAME, got, expected); - exit(1); - } - g_ds4_compress_ratios[il] = got; - } -} - -static void config_expect_f32(const char *name, float got, float expected); - -static void validate_swiglu_clamp_metadata(const ds4_model *m) { - const char *key = "deepseek4.swiglu_clamp_exp"; - ds4_array_ref arr; - if (!model_get_array(m, key, &arr) || - (arr.type != GGUF_VALUE_FLOAT32 && arr.type != GGUF_VALUE_FLOAT64)) { - fprintf(stderr, "ds4: required float array metadata key is missing: %s\n", key); - exit(1); - } - if (arr.len < DS4_N_LAYER) { - ds4_die("deepseek4.swiglu_clamp_exp is shorter than the layer count"); - } - - ds4_cursor c = cursor_at(m, arr.data_pos); - for (uint32_t i = 0; i < DS4_N_LAYER; i++) { - float got = 0.0f; - if (arr.type == GGUF_VALUE_FLOAT32) { - if (!cursor_read(&c, &got, sizeof(got))) ds4_die(c.error); - } else { - double v = 0.0; - if (!cursor_read(&c, &v, sizeof(v))) ds4_die(c.error); - got = (float)v; - } - config_expect_f32("swiglu_clamp_exp", got, DS4_SWIGLU_CLAMP_EXP); - } -} - -static void config_expect_u32(const char *name, uint32_t got, uint32_t expected) { - if (got == expected) return; - fprintf(stderr, "ds4: expected %s=%u for %s, got %u\n", - name, expected, DS4_MODEL_SHAPE_NAME, got); - exit(1); -} - -static void config_expect_u64(const char *name, uint64_t got, uint64_t expected) { - if (got == expected) return; - fprintf(stderr, "ds4: expected %s=%" PRIu64 " for %s, got %" PRIu64 "\n", - name, expected, DS4_MODEL_SHAPE_NAME, got); - exit(1); -} - -static void config_expect_f32(const char *name, float got, float expected) { - const float scale = fabsf(expected) > 1.0f ? fabsf(expected) : 1.0f; - if (fabsf(got - expected) <= scale * 1.0e-6f) return; - fprintf(stderr, "ds4: expected %s=%.9g for %s, got %.9g\n", - name, (double)expected, DS4_MODEL_SHAPE_NAME, (double)got); - exit(1); -} - -static void config_expect_bool(const char *name, bool got, bool expected) { - if (got == expected) return; - fprintf(stderr, "ds4: expected %s=%s for %s, got %s\n", - name, expected ? "true" : "false", DS4_MODEL_SHAPE_NAME, got ? "true" : "false"); - exit(1); -} - -static void config_validate_fixed_shape(uint32_t n_layer) { - config_expect_u32("block_count", n_layer, DS4_N_LAYER); -} - -/* Validate metadata values that affect semantics: attention shape, HC count, - * expert routing, RoPE scaling, compression ratios, and SwiGLU clamp. */ -static void config_validate_deepseek4_model(const ds4_model *m) { - const uint32_t n_layer = required_u32(m, "deepseek4.block_count"); - const uint32_t n_embd = required_u32(m, "deepseek4.embedding_length"); - const uint32_t n_vocab = required_u32(m, "deepseek4.vocab_size"); - const uint32_t n_head = required_u32(m, "deepseek4.attention.head_count"); - const uint32_t n_head_kv = required_u32(m, "deepseek4.attention.head_count_kv"); - const uint32_t n_head_dim = required_u32(m, "deepseek4.attention.key_length"); - const uint32_t n_value_dim = required_u32(m, "deepseek4.attention.value_length"); - const uint32_t n_rot = required_u32(m, "deepseek4.rope.dimension_count"); - const uint32_t n_lora_q = required_u32(m, "deepseek4.attention.q_lora_rank"); - const uint32_t n_lora_o = required_u32(m, "deepseek4.attention.output_lora_rank"); - const uint32_t n_out_group = required_u32(m, "deepseek4.attention.output_group_count"); - const uint32_t n_expert = required_u32(m, "deepseek4.expert_count"); - const uint32_t n_expert_used = required_u32(m, "deepseek4.expert_used_count"); - const uint32_t n_ff_exp = required_u32(m, "deepseek4.expert_feed_forward_length"); - const uint32_t n_expert_shared = required_u32(m, "deepseek4.expert_shared_count"); - const uint32_t n_hash_layer = required_u32(m, "deepseek4.hash_layer_count"); - uint32_t n_expert_groups = 0; - uint32_t n_group_used = 0; - model_get_u32(m, "deepseek4.expert_group_count", &n_expert_groups); - model_get_u32(m, "deepseek4.expert_group_used_count", &n_group_used); - const uint32_t n_swa = required_u32(m, "deepseek4.attention.sliding_window"); - const uint32_t n_indexer_head = required_u32(m, "deepseek4.attention.indexer.head_count"); - const uint32_t n_indexer_head_dim = required_u32(m, "deepseek4.attention.indexer.key_length"); - const uint32_t n_indexer_top_k = required_u32(m, "deepseek4.attention.indexer.top_k"); - const uint32_t n_hc = required_u32(m, "deepseek4.hyper_connection.count"); - const uint32_t n_hc_sinkhorn_iter = required_u32(m, "deepseek4.hyper_connection.sinkhorn_iterations"); - - ds4_select_shape_from_metadata(n_layer, - n_embd, - n_vocab, - n_head, - n_head_kv, - n_head_dim, - n_value_dim, - n_rot, - n_lora_q, - n_lora_o, - n_out_group, - n_expert, - n_expert_used, - n_ff_exp, - n_expert_shared, - n_hash_layer, - n_swa, - n_indexer_head, - n_indexer_head_dim, - n_indexer_top_k, - n_hc, - n_hc_sinkhorn_iter); - - config_expect_u32("embedding_length", n_embd, DS4_N_EMBD); - config_expect_u32("vocab_size", n_vocab, DS4_N_VOCAB); - config_expect_u32("attention.head_count", n_head, DS4_N_HEAD); - config_expect_u32("attention.key_length", n_head_dim, DS4_N_HEAD_DIM); - config_expect_u32("attention.head_count_kv", n_head_kv, DS4_N_HEAD_KV); - config_expect_u32("attention.value_length", n_value_dim, DS4_N_VALUE_DIM); - config_expect_u32("rope.dimension_count", n_rot, DS4_N_ROT); - config_expect_u32("attention.output_group_count", n_out_group, DS4_N_OUT_GROUP); - config_expect_u32("attention.q_lora_rank", n_lora_q, DS4_N_LORA_Q); - config_expect_u32("attention.output_lora_rank", n_lora_o, DS4_N_LORA_O); - config_expect_u32("expert_count", n_expert, DS4_N_EXPERT); - config_expect_u32("expert_used_count", n_expert_used, DS4_N_EXPERT_USED); - config_expect_u32("expert_feed_forward_length", n_ff_exp, DS4_N_FF_EXP); - config_expect_u32("expert_shared_count", n_expert_shared, DS4_N_EXPERT_SHARED); - config_expect_u32("hash_layer_count", n_hash_layer, DS4_N_HASH_LAYER); - config_expect_u32("expert_group_count", n_expert_groups, 0); - config_expect_u32("expert_group_used_count", n_group_used, 0); - - config_expect_u32("attention.sliding_window", n_swa, DS4_N_SWA); - config_expect_u32("attention.indexer.head_count", n_indexer_head, DS4_N_INDEXER_HEAD); - config_expect_u32("attention.indexer.key_length", n_indexer_head_dim, DS4_N_INDEXER_HEAD_DIM); - config_expect_u32("attention.indexer.top_k", n_indexer_top_k, DS4_N_INDEXER_TOP_K); - config_expect_u32("hyper_connection.count", n_hc, DS4_N_HC); - config_expect_u32("hyper_connection.sinkhorn_iterations", n_hc_sinkhorn_iter, DS4_N_HC_SINKHORN_ITER); - - config_validate_fixed_shape(n_layer); - validate_compress_ratio_metadata(m); - - validate_swiglu_clamp_metadata(m); - - uint64_t rope_orig_ctx = DS4_ROPE_ORIG_CTX; - model_get_u64_compat(m, "deepseek4.rope.scaling.original_context_length", &rope_orig_ctx); - if (rope_orig_ctx != DS4_ROPE_ORIG_CTX) { - fprintf(stderr, "ds4: expected rope.scaling.original_context_length=%" PRIu64 - " for %s, got %" PRIu64 "\n", - (uint64_t)DS4_ROPE_ORIG_CTX, DS4_MODEL_SHAPE_NAME, rope_orig_ctx); - exit(1); - } - const float rope_freq_base = required_f32(m, "deepseek4.rope.freq_base"); - config_expect_f32("rope.freq_base", rope_freq_base, DS4_ROPE_FREQ_BASE); - float rope_scale_factor = DS4_ROPE_SCALE_FACTOR; - model_get_f32_compat(m, "deepseek4.rope.scaling.factor", &rope_scale_factor); - config_expect_f32("rope.scaling.factor", rope_scale_factor, DS4_ROPE_SCALE_FACTOR); - float rope_yarn_beta_fast = DS4_ROPE_YARN_BETA_FAST; - model_get_f32_compat(m, "deepseek4.rope.scaling.yarn_beta_fast", &rope_yarn_beta_fast); - config_expect_f32("rope.scaling.yarn_beta_fast", rope_yarn_beta_fast, DS4_ROPE_YARN_BETA_FAST); - float rope_yarn_beta_slow = DS4_ROPE_YARN_BETA_SLOW; - model_get_f32_compat(m, "deepseek4.rope.scaling.yarn_beta_slow", &rope_yarn_beta_slow); - config_expect_f32("rope.scaling.yarn_beta_slow", rope_yarn_beta_slow, DS4_ROPE_YARN_BETA_SLOW); - const float compress_rope_freq_base = required_f32(m, "deepseek4.attention.compress_rope_freq_base"); - config_expect_f32("attention.compress_rope_freq_base", compress_rope_freq_base, DS4_COMPRESS_ROPE_FREQ_BASE); - const float expert_weight_scale = required_f32(m, "deepseek4.expert_weights_scale"); - config_expect_f32("expert_weights_scale", expert_weight_scale, DS4_EXPERT_WEIGHT_SCALE); - const float rms_eps = required_f32(m, "deepseek4.attention.layer_norm_rms_epsilon"); - config_expect_f32("attention.layer_norm_rms_epsilon", rms_eps, DS4_RMS_EPS); - const float hc_eps = required_f32(m, "deepseek4.hyper_connection.epsilon"); - config_expect_f32("hyper_connection.epsilon", hc_eps, DS4_HC_EPS); - const bool expert_weight_norm = required_bool(m, "deepseek4.expert_weights_norm"); - config_expect_bool("expert_weights_norm", expert_weight_norm, true); -} - -static void config_validate_glm_dsa_model(const ds4_model *m) { - g_ds4_shape = DS4_SHAPE_GLM52; - memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); - - const uint32_t n_layer = required_u32(m, "glm-dsa.block_count"); - const uint64_t n_ctx = required_u64_compat(m, "glm-dsa.context_length"); - const uint32_t n_embd = required_u32(m, "glm-dsa.embedding_length"); - const uint32_t n_vocab = required_u32(m, "glm-dsa.vocab_size"); - const uint32_t n_ff_dense = required_u32(m, "glm-dsa.feed_forward_length"); - const uint32_t n_head = required_u32(m, "glm-dsa.attention.head_count"); - const uint32_t n_head_kv = required_u32(m, "glm-dsa.attention.head_count_kv"); - const uint32_t n_head_dim = required_u32(m, "glm-dsa.attention.key_length"); - const uint32_t n_value_dim = required_u32(m, "glm-dsa.attention.value_length"); - const uint32_t n_rot = required_u32(m, "glm-dsa.rope.dimension_count"); - const uint32_t n_lora_q = required_u32(m, "glm-dsa.attention.q_lora_rank"); - const uint32_t n_kv_lora = required_u32(m, "glm-dsa.attention.kv_lora_rank"); - const uint32_t n_key_mla = required_u32(m, "glm-dsa.attention.key_length_mla"); - const uint32_t n_value_mla = required_u32(m, "glm-dsa.attention.value_length_mla"); - const uint32_t n_expert = required_u32(m, "glm-dsa.expert_count"); - const uint32_t n_expert_used = required_u32(m, "glm-dsa.expert_used_count"); - const uint32_t n_ff_exp = required_u32(m, "glm-dsa.expert_feed_forward_length"); - const uint32_t n_expert_shared = required_u32(m, "glm-dsa.expert_shared_count"); - const uint32_t n_expert_group = required_u32(m, "glm-dsa.expert_group_count"); - const uint32_t n_expert_group_used = required_u32(m, "glm-dsa.expert_group_used_count"); - const uint32_t expert_gating_func = required_u32(m, "glm-dsa.expert_gating_func"); - const uint32_t n_leading_dense = required_u32(m, "glm-dsa.leading_dense_block_count"); - const uint32_t n_nextn = required_u32(m, "glm-dsa.nextn_predict_layers"); - const uint32_t n_indexer_head = required_u32(m, "glm-dsa.attention.indexer.head_count"); - const uint32_t n_indexer_head_dim = required_u32(m, "glm-dsa.attention.indexer.key_length"); - const uint32_t n_indexer_top_k = required_u32(m, "glm-dsa.attention.indexer.top_k"); - - config_expect_u32("block_count", n_layer, DS4_N_LAYER); - config_expect_u64("context_length", n_ctx, DS4_ROPE_ORIG_CTX); - config_expect_u32("embedding_length", n_embd, DS4_N_EMBD); - config_expect_u32("vocab_size", n_vocab, DS4_N_VOCAB); - config_expect_u32("feed_forward_length", n_ff_dense, DS4_N_FF_DENSE); - config_expect_u32("attention.head_count", n_head, DS4_N_HEAD); - config_expect_u32("attention.head_count_kv", n_head_kv, DS4_N_HEAD_KV); - config_expect_u32("attention.key_length", n_head_dim, DS4_N_HEAD_DIM); - config_expect_u32("attention.value_length", n_value_dim, DS4_N_VALUE_DIM); - config_expect_u32("rope.dimension_count", n_rot, DS4_N_ROT); - config_expect_u32("attention.q_lora_rank", n_lora_q, DS4_N_LORA_Q); - config_expect_u32("attention.kv_lora_rank", n_kv_lora, DS4_N_KV_LORA); - config_expect_u32("attention.key_length_mla", n_key_mla, DS4_N_KEY_MLA); - config_expect_u32("attention.value_length_mla", n_value_mla, DS4_N_VALUE_MLA); - config_expect_u32("expert_count", n_expert, DS4_N_EXPERT); - config_expect_u32("expert_used_count", n_expert_used, DS4_N_EXPERT_USED); - config_expect_u32("expert_feed_forward_length", n_ff_exp, DS4_N_FF_EXP); - config_expect_u32("expert_shared_count", n_expert_shared, DS4_N_EXPERT_SHARED); - config_expect_u32("expert_group_count", n_expert_group, 1); - config_expect_u32("expert_group_used_count", n_expert_group_used, 1); - config_expect_u32("expert_gating_func", expert_gating_func, 2); - config_expect_u32("leading_dense_block_count", n_leading_dense, DS4_N_LEADING_DENSE); - config_expect_u32("nextn_predict_layers", n_nextn, DS4_N_NEXTN_PREDICT); - config_expect_u32("attention.indexer.head_count", n_indexer_head, DS4_N_INDEXER_HEAD); - config_expect_u32("attention.indexer.key_length", n_indexer_head_dim, DS4_N_INDEXER_HEAD_DIM); - config_expect_u32("attention.indexer.top_k", n_indexer_top_k, DS4_N_INDEXER_TOP_K); - - const float rope_freq_base = required_f32(m, "glm-dsa.rope.freq_base"); - config_expect_f32("rope.freq_base", rope_freq_base, DS4_ROPE_FREQ_BASE); - const float rms_eps = required_f32(m, "glm-dsa.attention.layer_norm_rms_epsilon"); - config_expect_f32("attention.layer_norm_rms_epsilon", rms_eps, DS4_RMS_EPS); - const float expert_weight_scale = required_f32(m, "glm-dsa.expert_weights_scale"); - config_expect_f32("expert_weights_scale", expert_weight_scale, DS4_EXPERT_WEIGHT_SCALE); - const bool expert_weight_norm = required_bool(m, "glm-dsa.expert_weights_norm"); - config_expect_bool("expert_weights_norm", expert_weight_norm, true); -} - -static void config_validate_model(const ds4_model *m) { - ds4_str arch = {0}; - if (model_get_string(m, "general.architecture", &arch) && - ds4_streq(arch, "glm-dsa")) { - config_validate_glm_dsa_model(m); - return; - } - config_validate_deepseek4_model(m); -} - -static void weights_bind_output( - ds4_weights *w, - const ds4_model *m, - bool required, - bool optional) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - if (required) { - w->output_norm = required_tensor(m, "output_norm.weight"); - w->output = required_tensor(m, "output.weight"); - } else if (optional) { - w->output_norm = model_find_tensor(m, "output_norm.weight"); - w->output = model_find_tensor(m, "output.weight"); - } - } else if (required) { - w->output_hc_base = required_tensor(m, "output_hc_base.weight"); - w->output_hc_fn = required_tensor(m, "output_hc_fn.weight"); - w->output_hc_scale = required_tensor(m, "output_hc_scale.weight"); - w->output_norm = required_tensor(m, "output_norm.weight"); - w->output = required_tensor(m, "output.weight"); - } else if (optional) { - w->output_hc_base = model_find_tensor(m, "output_hc_base.weight"); - w->output_hc_fn = model_find_tensor(m, "output_hc_fn.weight"); - w->output_hc_scale = model_find_tensor(m, "output_hc_scale.weight"); - w->output_norm = model_find_tensor(m, "output_norm.weight"); - w->output = model_find_tensor(m, "output.weight"); - } - - if (optional && - weights_have_partial_output_head(w) && - !weights_have_output_head(w)) { - ds4_die("partial output head in GGUF"); - } -} - -static void weights_bind_glm_dsa_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { - l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); - l->attn_q_a = required_tensorf(m, "blk.%u.attn_q_a.weight", il); - l->attn_q_a_norm = required_tensorf(m, "blk.%u.attn_q_a_norm.weight", il); - l->attn_q_b = required_tensorf(m, "blk.%u.attn_q_b.weight", il); - l->attn_kv_a_mqa = required_tensorf(m, "blk.%u.attn_kv_a_mqa.weight", il); - l->attn_kv_a_norm = required_tensorf(m, "blk.%u.attn_kv_a_norm.weight", il); - l->attn_k_b = required_tensorf(m, "blk.%u.attn_k_b.weight", il); - l->attn_v_b = required_tensorf(m, "blk.%u.attn_v_b.weight", il); - l->attn_output = required_tensorf(m, "blk.%u.attn_output.weight", il); - l->indexer_attn_q_b = required_tensorf(m, "blk.%u.indexer.attn_q_b.weight", il); - l->indexer_attn_k = required_tensorf(m, "blk.%u.indexer.attn_k.weight", il); - l->indexer_k_norm = required_tensorf(m, "blk.%u.indexer.k_norm.weight", il); - l->indexer_k_norm_b = required_tensorf(m, "blk.%u.indexer.k_norm.bias", il); - l->indexer_proj = required_tensorf(m, "blk.%u.indexer.proj.weight", il); - l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); - - if (il < DS4_N_LEADING_DENSE) { - l->ffn_gate = required_tensorf(m, "blk.%u.ffn_gate.weight", il); - l->ffn_up = required_tensorf(m, "blk.%u.ffn_up.weight", il); - l->ffn_down = required_tensorf(m, "blk.%u.ffn_down.weight", il); - } else { - l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); - l->ffn_exp_probs_b = required_tensorf(m, "blk.%u.exp_probs_b.bias", il); - l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); - l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); - l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); - l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); - l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); - l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); - } - - if (DS4_N_NEXTN_PREDICT != 0 && - il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { - l->nextn_eh_proj = required_tensorf(m, "blk.%u.nextn.eh_proj.weight", il); - l->nextn_enorm = required_tensorf(m, "blk.%u.nextn.enorm.weight", il); - l->nextn_hnorm = required_tensorf(m, "blk.%u.nextn.hnorm.weight", il); - l->nextn_shared_head_norm = - required_tensorf(m, "blk.%u.nextn.shared_head_norm.weight", il); - } -} - -static void weights_bind_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - weights_bind_glm_dsa_layer(l, m, il); - return; - } - - const uint32_t compress_ratio = ds4_layer_compress_ratio(il); - - l->hc_attn_fn = required_tensorf(m, "blk.%u.hc_attn_fn.weight", il); - l->hc_attn_scale = required_tensorf(m, "blk.%u.hc_attn_scale.weight", il); - l->hc_attn_base = required_tensorf(m, "blk.%u.hc_attn_base.weight", il); - l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); - l->attn_q_a = required_tensorf(m, "blk.%u.attn_q_a.weight", il); - l->attn_q_a_norm = required_tensorf(m, "blk.%u.attn_q_a_norm.weight", il); - l->attn_q_b = required_tensorf(m, "blk.%u.attn_q_b.weight", il); - l->attn_kv = required_tensorf(m, "blk.%u.attn_kv.weight", il); - l->attn_kv_a_norm = required_tensorf(m, "blk.%u.attn_kv_a_norm.weight", il); - l->attn_sinks = required_tensorf(m, "blk.%u.attn_sinks.weight", il); - l->attn_output_a = required_tensorf(m, "blk.%u.attn_output_a.weight", il); - l->attn_output_b = required_tensorf(m, "blk.%u.attn_output_b.weight", il); - if (compress_ratio != 0) { - l->attn_compressor_ape = required_tensorf(m, "blk.%u.attn_compressor_ape.weight", il); - l->attn_compressor_kv = required_tensorf(m, "blk.%u.attn_compressor_kv.weight", il); - l->attn_compressor_gate = required_tensorf(m, "blk.%u.attn_compressor_gate.weight", il); - l->attn_compressor_norm = required_tensorf(m, "blk.%u.attn_compressor_norm.weight", il); - } - if (compress_ratio == 4) { - l->indexer_attn_q_b = required_tensorf(m, "blk.%u.indexer.attn_q_b.weight", il); - l->indexer_proj = required_tensorf(m, "blk.%u.indexer.proj.weight", il); - l->indexer_compressor_ape = required_tensorf(m, "blk.%u.indexer_compressor_ape.weight", il); - l->indexer_compressor_kv = required_tensorf(m, "blk.%u.indexer_compressor_kv.weight", il); - l->indexer_compressor_gate = required_tensorf(m, "blk.%u.indexer_compressor_gate.weight", il); - l->indexer_compressor_norm = required_tensorf(m, "blk.%u.indexer_compressor_norm.weight", il); - } - l->hc_ffn_fn = required_tensorf(m, "blk.%u.hc_ffn_fn.weight", il); - l->hc_ffn_scale = required_tensorf(m, "blk.%u.hc_ffn_scale.weight", il); - l->hc_ffn_base = required_tensorf(m, "blk.%u.hc_ffn_base.weight", il); - l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); - l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); - l->ffn_exp_probs_b = tensor_by_namef(m, "blk.%u.exp_probs_b.bias", il); - l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); - l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); - l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); - l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); - l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); - l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); - - if (il < DS4_N_HASH_LAYER) { - l->ffn_gate_tid2eid = required_tensorf(m, "blk.%u.ffn_gate_tid2eid.weight", il); - } -} - -/* Bind tensor names once into the fixed DS4 layer layout. This is the point - * where stringly GGUF metadata becomes direct model-specific pointers. */ -static void weights_bind( - ds4_weights *w, - const ds4_model *m, - bool load_slice, - uint32_t load_layer_start, - uint32_t load_layer_end, - bool require_output, - bool optional_output) { - memset(w, 0, sizeof(*w)); - - uint32_t executable_layers = DS4_N_LAYER; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - DS4_N_LAYER > DS4_N_NEXTN_PREDICT) { - executable_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; - } - uint32_t start = 0; - uint32_t end = executable_layers - 1u; - bool require_token_embd = true; - if (load_slice) { - if (load_layer_start >= executable_layers) ds4_die("invalid model load layer slice"); - start = load_layer_start; - end = load_layer_end == UINT32_MAX ? executable_layers - 1u : load_layer_end; - if (end >= executable_layers || end < start) ds4_die("invalid model load layer slice"); - require_token_embd = start == 0; - } else { - require_output = true; - optional_output = false; - } - - if (require_token_embd) { - w->token_embd = required_tensor(m, "token_embd.weight"); - } else { - w->token_embd = model_find_tensor(m, "token_embd.weight"); - } - weights_bind_output(w, m, require_output, optional_output); - - for (uint32_t il = start; il <= end; il++) { - weights_bind_layer(&w->layer[il], m, il); - } - /* GLM nextn/MTP block(s): excluded from the executable pass but bound - * so the drafter can run them. Only when the full model is loaded. */ - if (!load_slice && - DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - start == 0 && end == executable_layers - 1u) { - for (uint32_t il = executable_layers; il < DS4_N_LAYER; il++) { - weights_bind_layer(&w->layer[il], m, il); - } - } - - weights_validate_layout(w, start, end, require_token_embd, require_output); -} - -typedef struct { - uint64_t off; - uint64_t end; - bool isolate; -} ds4_model_map_span; - -typedef struct { - ds4_model_map_span *v; - uint32_t len; - uint32_t cap; - uint64_t max_tensor_bytes; -} ds4_model_map_span_vec; - -static void model_map_span_include_tensor( - const ds4_tensor *t, - uint64_t *lo, - uint64_t *hi, - uint64_t *max_tensor_bytes) { - if (!t || t->bytes == 0) return; - const uint64_t end = t->abs_offset + t->bytes; - if (*lo == UINT64_MAX || t->abs_offset < *lo) *lo = t->abs_offset; - if (end > *hi) *hi = end; - if (t->bytes > *max_tensor_bytes) *max_tensor_bytes = t->bytes; -} - -static void model_map_span_vec_append(ds4_model_map_span_vec *spans, uint64_t lo, uint64_t hi, bool isolate) { - if (!spans || lo == UINT64_MAX || hi <= lo) return; - if (spans->len == spans->cap) { - uint32_t new_cap = spans->cap ? spans->cap * 2u : 16u; - spans->v = xrealloc(spans->v, (size_t)new_cap * sizeof(spans->v[0])); - spans->cap = new_cap; - } - spans->v[spans->len++] = (ds4_model_map_span){lo, hi, isolate}; -} - -static uint32_t model_map_q4_pro_group_views(void) { - uint32_t views = 1; - const char *env = getenv("DS4_METAL_Q4_PRO_MAP_GROUPS"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0' && v > 0 && v <= 384 && (384u % (uint32_t)v) == 0) { - views = (uint32_t)v; - } - } - return views; -} - -static void model_map_span_vec_include_one(ds4_model_map_span_vec *spans, const ds4_tensor *t) { - if (!t || t->bytes == 0) return; - const uint64_t q4_isolated_min_bytes = 2ull * 1024ull * 1024ull * 1024ull; - const uint32_t q4_pro_group_views = model_map_q4_pro_group_views(); - if (t->type == DS4_TENSOR_Q4_K && - t->ndim == 3 && - t->dim[2] == 384 && - t->bytes >= q4_isolated_min_bytes && - (t->bytes % q4_pro_group_views) == 0) - { - /* - * PRO Q4 routed expert tensors are too large to hide inside broad - * layer spans. Isolate them so the default selected-expert path does - * not stack large aliases on top of layer-sized model views. Optional - * group splits are enabled by DS4_METAL_Q4_PRO_MAP_GROUPS for Metal - * experiments that bind stable grouped views. - */ - const uint64_t group_bytes = t->bytes / q4_pro_group_views; - if (group_bytes > spans->max_tensor_bytes) spans->max_tensor_bytes = group_bytes; - for (uint32_t i = 0; i < q4_pro_group_views; i++) { - const uint64_t lo = t->abs_offset + (uint64_t)i * group_bytes; - model_map_span_vec_append(spans, lo, lo + group_bytes, true); - } - return; - } - - uint64_t lo = UINT64_MAX, hi = 0; - model_map_span_include_tensor(t, &lo, &hi, &spans->max_tensor_bytes); - const bool isolate = t->type == DS4_TENSOR_Q4_K && - t->bytes >= q4_isolated_min_bytes; - model_map_span_vec_append(spans, lo, hi, isolate); -} - -static void model_map_span_vec_include_layer(ds4_model_map_span_vec *spans, const ds4_layer_weights *l) { -#define DS4_INCLUDE_TENSOR(t_) model_map_span_vec_include_one(spans, (t_)) - DS4_INCLUDE_TENSOR(l->hc_attn_fn); - DS4_INCLUDE_TENSOR(l->hc_attn_scale); - DS4_INCLUDE_TENSOR(l->hc_attn_base); - DS4_INCLUDE_TENSOR(l->attn_norm); - DS4_INCLUDE_TENSOR(l->attn_q_a); - DS4_INCLUDE_TENSOR(l->attn_q_a_norm); - DS4_INCLUDE_TENSOR(l->attn_q_b); - DS4_INCLUDE_TENSOR(l->attn_kv); - DS4_INCLUDE_TENSOR(l->attn_kv_a_mqa); - DS4_INCLUDE_TENSOR(l->attn_kv_a_norm); - DS4_INCLUDE_TENSOR(l->attn_k_b); - DS4_INCLUDE_TENSOR(l->attn_v_b); - DS4_INCLUDE_TENSOR(l->attn_sinks); - DS4_INCLUDE_TENSOR(l->attn_output); - DS4_INCLUDE_TENSOR(l->attn_output_a); - DS4_INCLUDE_TENSOR(l->attn_output_b); - DS4_INCLUDE_TENSOR(l->attn_compressor_ape); - DS4_INCLUDE_TENSOR(l->attn_compressor_kv); - DS4_INCLUDE_TENSOR(l->attn_compressor_gate); - DS4_INCLUDE_TENSOR(l->attn_compressor_norm); - DS4_INCLUDE_TENSOR(l->indexer_attn_q_b); - DS4_INCLUDE_TENSOR(l->indexer_attn_k); - DS4_INCLUDE_TENSOR(l->indexer_k_norm); - DS4_INCLUDE_TENSOR(l->indexer_k_norm_b); - DS4_INCLUDE_TENSOR(l->indexer_proj); - DS4_INCLUDE_TENSOR(l->indexer_compressor_ape); - DS4_INCLUDE_TENSOR(l->indexer_compressor_kv); - DS4_INCLUDE_TENSOR(l->indexer_compressor_gate); - DS4_INCLUDE_TENSOR(l->indexer_compressor_norm); - DS4_INCLUDE_TENSOR(l->hc_ffn_fn); - DS4_INCLUDE_TENSOR(l->hc_ffn_scale); - DS4_INCLUDE_TENSOR(l->hc_ffn_base); - DS4_INCLUDE_TENSOR(l->ffn_norm); - DS4_INCLUDE_TENSOR(l->ffn_gate_tid2eid); - DS4_INCLUDE_TENSOR(l->ffn_gate); - DS4_INCLUDE_TENSOR(l->ffn_up); - DS4_INCLUDE_TENSOR(l->ffn_down); - DS4_INCLUDE_TENSOR(l->ffn_gate_inp); - DS4_INCLUDE_TENSOR(l->ffn_exp_probs_b); - DS4_INCLUDE_TENSOR(l->ffn_gate_exps); - DS4_INCLUDE_TENSOR(l->ffn_up_exps); - DS4_INCLUDE_TENSOR(l->ffn_down_exps); - DS4_INCLUDE_TENSOR(l->ffn_gate_shexp); - DS4_INCLUDE_TENSOR(l->ffn_up_shexp); - DS4_INCLUDE_TENSOR(l->ffn_down_shexp); - DS4_INCLUDE_TENSOR(l->nextn_eh_proj); - DS4_INCLUDE_TENSOR(l->nextn_enorm); - DS4_INCLUDE_TENSOR(l->nextn_hnorm); - DS4_INCLUDE_TENSOR(l->nextn_shared_head_norm); -#undef DS4_INCLUDE_TENSOR -} - -static void model_map_span_vec_include_layer_decode_static(ds4_model_map_span_vec *spans, const ds4_layer_weights *l) { -#define DS4_INCLUDE_TENSOR(t_) model_map_span_vec_include_one(spans, (t_)) - DS4_INCLUDE_TENSOR(l->hc_attn_fn); - DS4_INCLUDE_TENSOR(l->hc_attn_scale); - DS4_INCLUDE_TENSOR(l->hc_attn_base); - DS4_INCLUDE_TENSOR(l->attn_norm); - DS4_INCLUDE_TENSOR(l->attn_q_a); - DS4_INCLUDE_TENSOR(l->attn_q_a_norm); - DS4_INCLUDE_TENSOR(l->attn_q_b); - DS4_INCLUDE_TENSOR(l->attn_kv); - DS4_INCLUDE_TENSOR(l->attn_kv_a_mqa); - DS4_INCLUDE_TENSOR(l->attn_kv_a_norm); - DS4_INCLUDE_TENSOR(l->attn_k_b); - DS4_INCLUDE_TENSOR(l->attn_v_b); - DS4_INCLUDE_TENSOR(l->attn_sinks); - DS4_INCLUDE_TENSOR(l->attn_output); - DS4_INCLUDE_TENSOR(l->attn_output_a); - DS4_INCLUDE_TENSOR(l->attn_output_b); - DS4_INCLUDE_TENSOR(l->attn_compressor_ape); - DS4_INCLUDE_TENSOR(l->attn_compressor_kv); - DS4_INCLUDE_TENSOR(l->attn_compressor_gate); - DS4_INCLUDE_TENSOR(l->attn_compressor_norm); - DS4_INCLUDE_TENSOR(l->indexer_attn_q_b); - DS4_INCLUDE_TENSOR(l->indexer_attn_k); - DS4_INCLUDE_TENSOR(l->indexer_k_norm); - DS4_INCLUDE_TENSOR(l->indexer_k_norm_b); - DS4_INCLUDE_TENSOR(l->indexer_proj); - DS4_INCLUDE_TENSOR(l->indexer_compressor_ape); - DS4_INCLUDE_TENSOR(l->indexer_compressor_kv); - DS4_INCLUDE_TENSOR(l->indexer_compressor_gate); - DS4_INCLUDE_TENSOR(l->indexer_compressor_norm); - DS4_INCLUDE_TENSOR(l->hc_ffn_fn); - DS4_INCLUDE_TENSOR(l->hc_ffn_scale); - DS4_INCLUDE_TENSOR(l->hc_ffn_base); - DS4_INCLUDE_TENSOR(l->ffn_norm); - DS4_INCLUDE_TENSOR(l->ffn_gate_tid2eid); - DS4_INCLUDE_TENSOR(l->ffn_gate); - DS4_INCLUDE_TENSOR(l->ffn_up); - DS4_INCLUDE_TENSOR(l->ffn_down); - DS4_INCLUDE_TENSOR(l->ffn_gate_inp); - DS4_INCLUDE_TENSOR(l->ffn_exp_probs_b); - DS4_INCLUDE_TENSOR(l->ffn_gate_shexp); - DS4_INCLUDE_TENSOR(l->ffn_up_shexp); - DS4_INCLUDE_TENSOR(l->ffn_down_shexp); - DS4_INCLUDE_TENSOR(l->nextn_eh_proj); - DS4_INCLUDE_TENSOR(l->nextn_enorm); - DS4_INCLUDE_TENSOR(l->nextn_hnorm); - DS4_INCLUDE_TENSOR(l->nextn_shared_head_norm); -#undef DS4_INCLUDE_TENSOR -} - -static bool glm_stream_resident_decode_layer_supported( - const ds4_layer_weights *l, - uint32_t il) { - if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || - !l || - il < DS4_N_LEADING_DENSE || - !l->ffn_gate_exps || - !l->ffn_up_exps || - !l->ffn_down_exps) { - return false; - } - if (l->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - l->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS || - l->ffn_down_exps->type == DS4_TENSOR_Q2_K)) { - return true; - } - return l->ffn_gate_exps->type == l->ffn_up_exps->type && - l->ffn_gate_exps->type == l->ffn_down_exps->type && - (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || - l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); -} - -static uint32_t g_glm_streaming_full_resident_start; -static uint32_t g_glm_streaming_full_resident_layers; - -static bool glm_stream_resident_decode_layer_enabled( - const ds4_layer_weights *l, - uint32_t il) { - if (!glm_stream_resident_decode_layer_supported(l, il)) return false; - return g_glm_streaming_full_resident_layers != 0 && - il >= g_glm_streaming_full_resident_start && - il - g_glm_streaming_full_resident_start < - g_glm_streaming_full_resident_layers; -} - -static bool glm_stream_expert_cache_addr_layout_supported( - const ds4_weights *w, - const ds4_layer_weights *l, - uint32_t il) { - if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || - !w || - !l || - il >= DS4_N_LAYER || - il < DS4_N_LEADING_DENSE || - !l->ffn_gate_exps || - !l->ffn_up_exps || - !l->ffn_down_exps || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > 8 || - DS4_N_EXPERT < 128 || - glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", - "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { - return false; - } - if (!weights_streaming_layer_experts_uniform(w, il)) return false; - - if (l->ffn_gate_exps->type != l->ffn_up_exps->type) return false; - const bool q2_addr = - l->ffn_gate_exps->type == DS4_TENSOR_Q2_K && - l->ffn_down_exps->type == DS4_TENSOR_Q2_K; - const bool q4_addr = - l->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - l->ffn_down_exps->type == DS4_TENSOR_Q4_K; - return q2_addr || q4_addr; -} - -static DS4_MAYBE_UNUSED bool glm_stream_expert_cache_addr_supported( - const ds4_weights *w, - const ds4_layer_weights *l, - uint32_t il) { - if (!glm_stream_expert_cache_addr_layout_supported(w, l, il)) { - return false; - } - -#ifdef DS4_NO_GPU - return false; -#else - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!streaming_layer_gate_down_expert_bytes(l, - &gate_expert_bytes, - &down_expert_bytes)) { - return false; - } - return ds4_gpu_stream_expert_cache_budget_for_expert_size( - gate_expert_bytes, - down_expert_bytes) >= DS4_N_EXPERT_USED; -#endif -} - -static bool glm_stream_selected_expert_cache_supported( - const ds4_layer_weights *l, - uint32_t il) { - if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || - !l || - il < DS4_N_LEADING_DENSE || - !l->ffn_gate_exps || - !l->ffn_up_exps || - !l->ffn_down_exps || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > 8 || - DS4_N_EXPERT < 128 || - glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") || - glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") || - glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", - "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { - return false; - } - - if (l->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || - l->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS) { - return false; - } - - if (l->ffn_down_exps->type == DS4_TENSOR_Q2_K) { - return !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", - "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); - } - if (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS) { - return !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_STREAM_ADDR_TABLE", - "DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE"); - } - return false; -} - -static bool glm_stream_decode_experts_are_streamed( - const ds4_weights *w, - const ds4_layer_weights *l, - uint32_t il) { - if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) return false; - return glm_stream_expert_cache_addr_layout_supported(w, l, il) || - glm_stream_selected_expert_cache_supported(l, il); -} - -/* - * Decode-time spans for one layer. The static set excludes routed expert - * tensors only when the streaming expert-cache path can really serve them. - * Boosted layers, mixed GLM quant layouts such as Q4 gate/up plus Q5 down, or - * undersized expert caches fall back to direct model-range reads. Include - * those expert tensors so cache-hit prefill extension and decode are covered. - */ -static void model_map_span_vec_include_layer_decode( - ds4_model_map_span_vec *spans, - const ds4_weights *w, - uint32_t il) { - const ds4_layer_weights *l = &w->layer[il]; - model_map_span_vec_include_layer_decode_static(spans, l); - if (!weights_streaming_layer_experts_uniform(w, il) || - (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - !glm_stream_decode_experts_are_streamed(w, l, il)) || - glm_stream_resident_decode_layer_enabled(l, il)) { - model_map_span_vec_include_one(spans, l->ffn_gate_exps); - model_map_span_vec_include_one(spans, l->ffn_up_exps); - model_map_span_vec_include_one(spans, l->ffn_down_exps); - } -} - -static void model_map_span_vec_include_output(ds4_model_map_span_vec *spans, const ds4_weights *w) { - model_map_span_vec_include_one(spans, w->output_hc_base); - model_map_span_vec_include_one(spans, w->output_hc_fn); - model_map_span_vec_include_one(spans, w->output_hc_scale); - model_map_span_vec_include_one(spans, w->output_norm); - model_map_span_vec_include_one(spans, w->output); -} - -static int model_map_span_cmp(const void *a, const void *b) { - const ds4_model_map_span *sa = a; - const ds4_model_map_span *sb = b; - if (sa->off < sb->off) return -1; - if (sa->off > sb->off) return 1; - if (sa->end < sb->end) return -1; - if (sa->end > sb->end) return 1; - return 0; -} - -static bool model_map_span_vec_finish(ds4_model_map_span_vec *spans) { - if (!spans || spans->len == 0 || spans->max_tensor_bytes == 0) return false; - - qsort(spans->v, spans->len, sizeof(spans->v[0]), model_map_span_cmp); - uint32_t out = 0; - for (uint32_t i = 0; i < spans->len; i++) { - if (out == 0 || - spans->v[i].off > spans->v[out - 1u].end || - spans->v[i].isolate || - spans->v[out - 1u].isolate) { - spans->v[out++] = spans->v[i]; - } else if (spans->v[i].end > spans->v[out - 1u].end) { - spans->v[out - 1u].end = spans->v[i].end; - } - } - spans->len = out; - return spans->len != 0; -} - -static DS4_MAYBE_UNUSED bool weights_model_map_spans( - const ds4_weights *w, - uint32_t layer_start, - uint32_t layer_end, - bool include_output, - ds4_model_map_span_vec *spans) { - if (!w || !spans) return false; - if (layer_start >= DS4_N_LAYER) return false; - if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; - if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; - - memset(spans, 0, sizeof(*spans)); - if (layer_start == 0) model_map_span_vec_include_one(spans, w->token_embd); - for (uint32_t il = layer_start; il <= layer_end; il++) { - model_map_span_vec_include_layer(spans, &w->layer[il]); - } - if (include_output) model_map_span_vec_include_output(spans, w); - return model_map_span_vec_finish(spans); -} - -static const uint8_t *tensor_expert_bytes( - const ds4_model *m, - const ds4_tensor *w, - uint32_t expert, - uint64_t *in_dim, - uint64_t *out_dim, - uint64_t *row_bytes); - -/* TP sharding keeps full layers but restricts every routed-expert blob to - * one contiguous rank range (rank 0 owns the lower expert ids, matching - * ds4_tp_owns_expert in metal/moe.metal). */ -static DS4_MAYBE_UNUSED bool weights_model_map_sharded_spans( - const ds4_weights *w, - const ds4_model *m, - int rank, - ds4_model_map_span_vec *spans) { - if (!w || !m || !spans || (rank != 0 && rank != 1)) return false; - memset(spans, 0, sizeof(*spans)); - model_map_span_vec_include_one(spans, w->token_embd); - for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { - const ds4_layer_weights *l = &w->layer[il]; - /* Dense/attention/router tensors only — the plain decode include - * would map the full expert blobs in non-streaming mode. */ - model_map_span_vec_include_layer_decode_static(spans, l); - const ds4_tensor *exps[3] = { l->ffn_gate_exps, l->ffn_up_exps, - l->ffn_down_exps }; - for (int t = 0; t < 3; t++) { - const ds4_tensor *x = exps[t]; - if (!x || x->ndim != 3 || x->dim[2] < 2) continue; - uint64_t in_dim = 0, out_dim = 0, row_bytes = 0; - (void)tensor_expert_bytes(m, x, 0, &in_dim, &out_dim, &row_bytes); - const uint64_t expert_bytes = out_dim * row_bytes; - const uint64_t low_experts = x->dim[2] / 2; - const uint64_t first_expert = rank == 1 ? low_experts : 0; - const uint64_t owned_experts = rank == 1 ? - x->dim[2] - low_experts : low_experts; - const uint64_t owned_bytes = owned_experts * expert_bytes; - const uint64_t lo = x->abs_offset + first_expert * expert_bytes; - /* Kernels index experts from the blob base, so the owned range - * must sit in one contiguous view. Rank 1 takes any remainder. */ - model_map_span_vec_append(spans, lo, lo + owned_bytes, true); - if (owned_bytes > spans->max_tensor_bytes) { - spans->max_tensor_bytes = owned_bytes; - } - } - } - model_map_span_vec_include_output(spans, w); - return model_map_span_vec_finish(spans); -} - -static DS4_MAYBE_UNUSED bool weights_model_map_decode_layer_spans( - const ds4_weights *w, - uint32_t il, - ds4_model_map_span_vec *spans) { - if (!w || !spans || il >= DS4_N_LAYER) return false; - memset(spans, 0, sizeof(*spans)); - model_map_span_vec_include_layer_decode(spans, w, il); - return model_map_span_vec_finish(spans); -} - -static DS4_MAYBE_UNUSED bool weights_model_map_decode_static_spans( - const ds4_weights *w, - bool include_token, - bool include_output, - ds4_model_map_span_vec *spans) { - if (!w || !spans) return false; - memset(spans, 0, sizeof(*spans)); - if (include_token) model_map_span_vec_include_one(spans, w->token_embd); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - model_map_span_vec_include_layer_decode(spans, w, il); - } - if (include_output) model_map_span_vec_include_output(spans, w); - return model_map_span_vec_finish(spans); -} - -static DS4_MAYBE_UNUSED bool weights_model_map_decode_static_slice_spans( - const ds4_weights *w, - uint32_t layer_start, - uint32_t layer_end, - bool include_token, - bool include_output, - ds4_model_map_span_vec *spans) { - if (!w || !spans) return false; - if (layer_start >= DS4_N_LAYER) return false; - if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; - if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; - - memset(spans, 0, sizeof(*spans)); - if (include_token) model_map_span_vec_include_one(spans, w->token_embd); - for (uint32_t il = layer_start; il <= layer_end; il++) { - model_map_span_vec_include_layer_decode(spans, w, il); - } - if (include_output) model_map_span_vec_include_output(spans, w); - return model_map_span_vec_finish(spans); -} - -static DS4_MAYBE_UNUSED uint64_t model_map_span_vec_total_bytes( - const ds4_model_map_span_vec *spans) { - if (!spans) return 0; - uint64_t total = 0; - for (uint32_t i = 0; i < spans->len; i++) { - const uint64_t bytes = spans->v[i].end - spans->v[i].off; - if (total > UINT64_MAX - bytes) return UINT64_MAX; - total += bytes; - } - return total; -} - -static DS4_MAYBE_UNUSED bool weights_streaming_non_routed_bytes( - const ds4_weights *w, - uint64_t *bytes_out) { - if (bytes_out) *bytes_out = 0; - if (!w || !bytes_out) return false; - - ds4_model_map_span_vec spans; - const bool include_token = - weights_layer_has_required(&w->layer[0], 0); - if (!weights_model_map_decode_static_spans(w, - include_token, - weights_have_output_head(w), - &spans)) { - return false; - } - *bytes_out = model_map_span_vec_total_bytes(&spans); - free(spans.v); - return true; -} - -static DS4_MAYBE_UNUSED bool weights_model_map_token_spans( - const ds4_weights *w, - ds4_model_map_span_vec *spans) { - if (!w || !spans) return false; - memset(spans, 0, sizeof(*spans)); - model_map_span_vec_include_one(spans, w->token_embd); - return model_map_span_vec_finish(spans); -} - -static DS4_MAYBE_UNUSED bool weights_model_map_output_spans( - const ds4_weights *w, - ds4_model_map_span_vec *spans) { - if (!w || !spans) return false; - memset(spans, 0, sizeof(*spans)); - model_map_span_vec_include_output(spans, w); - return model_map_span_vec_finish(spans); -} - -static void mtp_weights_bind(ds4_mtp_weights *w, const ds4_model *m) { - memset(w, 0, sizeof(*w)); - - w->hc_head_base = required_tensor(m, "mtp.0.hc_head_base.weight"); - w->hc_head_fn = required_tensor(m, "mtp.0.hc_head_fn.weight"); - w->hc_head_scale = required_tensor(m, "mtp.0.hc_head_scale.weight"); - w->e_proj = required_tensor(m, "mtp.0.e_proj.weight"); - w->h_proj = required_tensor(m, "mtp.0.h_proj.weight"); - w->enorm = required_tensor(m, "mtp.0.enorm.weight"); - w->hnorm = required_tensor(m, "mtp.0.hnorm.weight"); - w->norm = required_tensor(m, "mtp.0.norm.weight"); - - ds4_layer_weights *l = &w->block; - l->hc_attn_fn = required_tensor(m, "mtp.0.hc_attn_fn.weight"); - l->hc_attn_scale = required_tensor(m, "mtp.0.hc_attn_scale.weight"); - l->hc_attn_base = required_tensor(m, "mtp.0.hc_attn_base.weight"); - l->attn_norm = required_tensor(m, "mtp.0.attn_norm.weight"); - l->attn_q_a = required_tensor(m, "mtp.0.attn_q_a.weight"); - l->attn_q_a_norm = required_tensor(m, "mtp.0.attn_q_a_norm.weight"); - l->attn_q_b = required_tensor(m, "mtp.0.attn_q_b.weight"); - l->attn_kv = required_tensor(m, "mtp.0.attn_kv.weight"); - l->attn_kv_a_norm = required_tensor(m, "mtp.0.attn_kv_a_norm.weight"); - l->attn_sinks = required_tensor(m, "mtp.0.attn_sinks.weight"); - l->attn_output_a = required_tensor(m, "mtp.0.attn_output_a.weight"); - l->attn_output_b = required_tensor(m, "mtp.0.attn_output_b.weight"); - l->hc_ffn_fn = required_tensor(m, "mtp.0.hc_ffn_fn.weight"); - l->hc_ffn_scale = required_tensor(m, "mtp.0.hc_ffn_scale.weight"); - l->hc_ffn_base = required_tensor(m, "mtp.0.hc_ffn_base.weight"); - l->ffn_norm = required_tensor(m, "mtp.0.ffn_norm.weight"); - l->ffn_gate_inp = required_tensor(m, "mtp.0.ffn_gate_inp.weight"); - l->ffn_exp_probs_b = required_tensor(m, "mtp.0.exp_probs_b.bias"); - l->ffn_gate_exps = required_tensor(m, "mtp.0.ffn_gate_exps.weight"); - l->ffn_up_exps = required_tensor(m, "mtp.0.ffn_up_exps.weight"); - l->ffn_down_exps = required_tensor(m, "mtp.0.ffn_down_exps.weight"); - l->ffn_gate_shexp = required_tensor(m, "mtp.0.ffn_gate_shexp.weight"); - l->ffn_up_shexp = required_tensor(m, "mtp.0.ffn_up_shexp.weight"); - l->ffn_down_shexp = required_tensor(m, "mtp.0.ffn_down_shexp.weight"); - - mtp_weights_validate_layout(w); -} - -static ds4_tensor *dspark_bind_tensor( - ds4_dspark_weights *dw, - const ds4_model *m, - uint32_t stage, - const char *suffix, - bool required) { - ds4_tensor *t = tensor_by_mtp_stage_suffix(m, stage, suffix); - if (t) { - dw->present_tensors++; - } else if (required) { - dw->missing_tensors++; - } - return t; -} - -static void dspark_bind_block( - ds4_dspark_weights *dw, - ds4_layer_weights *l, - const ds4_model *m, - uint32_t stage) { - l->hc_attn_fn = dspark_bind_tensor(dw, m, stage, "hc_attn_fn.weight", true); - l->hc_attn_scale = dspark_bind_tensor(dw, m, stage, "hc_attn_scale.weight", true); - l->hc_attn_base = dspark_bind_tensor(dw, m, stage, "hc_attn_base.weight", true); - l->attn_norm = dspark_bind_tensor(dw, m, stage, "attn_norm.weight", true); - l->attn_q_a = dspark_bind_tensor(dw, m, stage, "attn_q_a.weight", true); - l->attn_q_a_norm = dspark_bind_tensor(dw, m, stage, "attn_q_a_norm.weight", true); - l->attn_q_b = dspark_bind_tensor(dw, m, stage, "attn_q_b.weight", true); - l->attn_kv = dspark_bind_tensor(dw, m, stage, "attn_kv.weight", true); - l->attn_kv_a_norm = dspark_bind_tensor(dw, m, stage, "attn_kv_a_norm.weight", true); - l->attn_sinks = dspark_bind_tensor(dw, m, stage, "attn_sinks.weight", true); - l->attn_output_a = dspark_bind_tensor(dw, m, stage, "attn_output_a.weight", true); - l->attn_output_b = dspark_bind_tensor(dw, m, stage, "attn_output_b.weight", true); - l->hc_ffn_fn = dspark_bind_tensor(dw, m, stage, "hc_ffn_fn.weight", true); - l->hc_ffn_scale = dspark_bind_tensor(dw, m, stage, "hc_ffn_scale.weight", true); - l->hc_ffn_base = dspark_bind_tensor(dw, m, stage, "hc_ffn_base.weight", true); - l->ffn_norm = dspark_bind_tensor(dw, m, stage, "ffn_norm.weight", true); - l->ffn_gate_inp = dspark_bind_tensor(dw, m, stage, "ffn_gate_inp.weight", true); - l->ffn_exp_probs_b = dspark_bind_tensor(dw, m, stage, "exp_probs_b.bias", true); - l->ffn_gate_exps = dspark_bind_tensor(dw, m, stage, "ffn_gate_exps.weight", true); - l->ffn_up_exps = dspark_bind_tensor(dw, m, stage, "ffn_up_exps.weight", true); - l->ffn_down_exps = dspark_bind_tensor(dw, m, stage, "ffn_down_exps.weight", true); - l->ffn_gate_shexp = dspark_bind_tensor(dw, m, stage, "ffn_gate_shexp.weight", true); - l->ffn_up_shexp = dspark_bind_tensor(dw, m, stage, "ffn_up_shexp.weight", true); - l->ffn_down_shexp = dspark_bind_tensor(dw, m, stage, "ffn_down_shexp.weight", true); -} - -static void dspark_weights_bind_optional( - ds4_dspark_weights *dw, - const ds4_model *m, - const ds4_dspark_summary *summary) { - memset(dw, 0, sizeof(*dw)); - if (!m || !summary) return; - - dw->n_stages = summary->stages < DS4_DSPARK_MAX_STAGES ? - summary->stages : DS4_DSPARK_MAX_STAGES; - dw->block_size = summary->block_size; - dw->markov_rank = summary->markov_rank; - dw->noise_token_id = summary->noise_token_id; - dw->target_layer_count = summary->target_layer_count; - dw->has_block_size = summary->has_block_size; - dw->has_markov_rank = summary->has_markov_rank; - dw->has_noise_token_id = summary->has_noise_token_id; - dw->has_target_layers = summary->has_target_layers; - memcpy(dw->target_layers, - summary->target_layers, - (size_t)dw->target_layer_count * sizeof(dw->target_layers[0])); - if (summary->stages > DS4_DSPARK_MAX_STAGES) dw->missing_tensors++; - - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - ds4_dspark_stage_weights *sw = &dw->stage[stage]; - dspark_bind_block(dw, &sw->block, m, stage); - if (stage == 0) { - sw->main_proj = dspark_bind_tensor(dw, m, stage, "main_proj.weight", true); - sw->main_norm = dspark_bind_tensor(dw, m, stage, "main_norm.weight", true); - } - } - - if (dw->n_stages != 0) { - const uint32_t final_stage = dw->n_stages - 1u; - ds4_dspark_stage_weights *sw = &dw->stage[final_stage]; - sw->norm = dspark_bind_tensor(dw, m, final_stage, "norm.weight", true); - sw->hc_head_base = - dspark_bind_tensor(dw, m, final_stage, "hc_head_base.weight", true); - sw->hc_head_fn = - dspark_bind_tensor(dw, m, final_stage, "hc_head_fn.weight", true); - sw->hc_head_scale = - dspark_bind_tensor(dw, m, final_stage, "hc_head_scale.weight", true); - sw->markov_w1 = - dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w1.weight", true); - sw->markov_w2 = - dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w2.weight", true); - sw->confidence_proj = - dspark_bind_tensor(dw, m, final_stage, "confidence_head.proj.weight", true); - } - - dspark_weights_validate_layout(dw); -} - -static void weights_free(ds4_weights *w) { - memset(w, 0, sizeof(*w)); -} - -/* Load one token embedding row and expand it to float activations. */ -static void embed_token_f16(const ds4_model *m, const ds4_weights *w, int token, float *out) { - ds4_tensor *te = w->token_embd; - if (te->type != DS4_TENSOR_F16 || te->ndim != 2) { - ds4_die("expected a 2D F16 token embedding tensor"); - } - if (token < 0 || (uint64_t)token >= te->dim[1]) { - ds4_die("token id is outside the embedding table"); - } - - const uint16_t *base = tensor_data(m, te); - const uint64_t stride = te->dim[0]; - const uint16_t *row = base + (uint64_t)token * stride; - - for (uint64_t i = 0; i < stride; i++) { - out[i] = f16_to_f32(row[i]); - } -} - -static void embed_token_q8_0(const ds4_model *m, const ds4_weights *w, int token, float *out) { - ds4_tensor *te = w->token_embd; - if (te->type != DS4_TENSOR_Q8_0 || te->ndim != 2) { - ds4_die("expected a 2D Q8_0 token embedding tensor"); - } - if (token < 0 || (uint64_t)token >= te->dim[1]) { - ds4_die("token id is outside the embedding table"); - } - - const uint64_t n = te->dim[0]; - const uint64_t blocks = (n + 31) / 32; - const uint8_t *row = (const uint8_t *)tensor_data(m, te) + - (uint64_t)token * blocks * 34; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); - const float scale = f16_to_f32(scale_bits); - const int8_t *qs = (const int8_t *)(row + b * 34 + 2); - const uint64_t i0 = b * 32; - const uint64_t bn = n - i0 < 32 ? n - i0 : 32; - for (uint64_t i = 0; i < bn; i++) { - out[i0 + i] = scale * (float)qs[i]; - } - } -} - -static void embed_token_any(const ds4_model *m, const ds4_weights *w, int token, float *out) { - if (!w->token_embd) ds4_die("token embedding tensor is missing"); - switch (w->token_embd->type) { - case DS4_TENSOR_F16: - embed_token_f16(m, w, token, out); - break; - case DS4_TENSOR_Q8_0: - embed_token_q8_0(m, w, token, out); - break; - default: - ds4_die("unsupported token embedding tensor type"); - } -} - -/* RMSNorm without a learned scale, used by hyper-connection control vectors. */ -static void rms_norm_no_weight(float *out, const float *x, uint64_t n, float eps) { - double ss = 0.0; - for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; - - const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); - for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale; -} - -/* Standard DS4 RMSNorm with learned per-channel scale. */ -static void rms_norm_weight(float *out, const float *x, const float *weight, uint64_t n, float eps) { - double ss = 0.0; - for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; - - const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); - for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale * weight[i]; -} - -/* Normalize each attention head independently after Q projection. */ -static void head_rms_norm_inplace(float *x, uint32_t n_head, uint32_t head_dim, float eps) { - for (uint32_t h = 0; h < n_head; h++) { - float *head = x + (uint64_t)h * head_dim; - double ss = 0.0; - for (uint32_t i = 0; i < head_dim; i++) ss += (double)head[i] * head[i]; - - const float scale = 1.0f / sqrtf((float)(ss / (double)head_dim) + eps); - for (uint32_t i = 0; i < head_dim; i++) head[i] *= scale; - } -} - -typedef struct { - float *out; - const uint16_t *data; - const float *x; - uint64_t in_dim; -} matvec_f16_ctx; - -static inline float dot_f16_row(const uint16_t *row, const float *x, uint64_t n) { -#if defined(__ARM_NEON) - uint64_t i = 0; - float32x4_t acc0 = vdupq_n_f32(0.0f); - float32x4_t acc1 = vdupq_n_f32(0.0f); - for (; i + 8 <= n; i += 8) { - const float16x8_t hv = vreinterpretq_f16_u16(vld1q_u16(row + i)); - const float32x4_t h0 = vcvt_f32_f16(vget_low_f16(hv)); - const float32x4_t h1 = vcvt_f32_f16(vget_high_f16(hv)); - acc0 = vfmaq_f32(acc0, h0, vld1q_f32(x + i)); - acc1 = vfmaq_f32(acc1, h1, vld1q_f32(x + i + 4)); - } - - float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); - for (; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; - return acc; -#else - float acc = 0.0f; - for (uint64_t i = 0; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; - return acc; -#endif -} - -static void matvec_f16_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_f16_ctx *ctx = vctx; - - for (uint64_t o = row0; o < row1; o++) { - const uint16_t *row = ctx->data + o * ctx->in_dim; - ctx->out[o] = dot_f16_row(row, ctx->x, ctx->in_dim); - } -} - -/* Dense F16 matvec for small control projections such as HC and router heads. */ -static void matvec_f16(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { - if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t out_dim = w->dim[1]; - matvec_f16_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .x = x, - .in_dim = in_dim, - }; - - const uint64_t ops = in_dim * out_dim; - const uint64_t min_rows = ops >= 262144 ? 1 : 512; - ds4_parallel_for_min_rows(out_dim, matvec_f16_worker, &ctx, min_rows); -} - -static void matvec_f16_serial(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { - if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t out_dim = w->dim[1]; - const uint16_t *data = tensor_data(m, w); - for (uint64_t o = 0; o < out_dim; o++) { - out[o] = dot_f16_row(data + o * in_dim, x, in_dim); - } -} - -typedef struct { - float *out; - const uint8_t *data; - const int8_t *xq; - const float *xscale; - uint64_t in_dim; - uint64_t row0; - uint64_t blocks; -} matvec_q8_0_ctx; - -typedef struct { - float *out0; - float *out1; - const uint8_t *data0; - const uint8_t *data1; - const int8_t *xq; - const float *xscale; - uint64_t in_dim; - uint64_t blocks; -} matvec_q8_0_pair_ctx; - -typedef struct { - float *out; - const uint8_t *data; - const int8_t *xq; - const float *xscale; - uint64_t in_dim; - uint64_t blocks; - uint64_t rank; -} matvec_q8_0_grouped_ctx; - -typedef struct { - float *out; - const uint8_t *data; - const int8_t *xq; - const float *xscale; - uint64_t n_tok; - uint64_t n_groups; - uint64_t group_dim; - uint64_t blocks; - uint64_t rank; -} matmul_q8_0_grouped_batch_ctx; - -typedef struct { - float *out; - const uint8_t *data; - const int8_t *xq; - const float *xscale; - uint64_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t blocks; -} matmul_q8_0_batch_ctx; - -typedef struct { - float *out0; - float *out1; - const uint8_t *data0; - const uint8_t *data1; - const int8_t *xq; - const float *xscale; - uint64_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t blocks; -} matmul_q8_0_pair_batch_ctx; - -typedef struct { - const float *x; - int8_t *xq; - float *xscale; - uint64_t in_dim; - uint64_t blocks; -} quantize_q8_0_batch_ctx; - -static inline int32_t dot_i8_32(const int8_t *a, const int8_t *b, uint64_t n) { -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - if (n == 32) { - int32x4_t acc = vdupq_n_s32(0); - acc = vdotq_s32(acc, vld1q_s8(a), vld1q_s8(b)); - acc = vdotq_s32(acc, vld1q_s8(a + 16), vld1q_s8(b + 16)); - return vaddvq_s32(acc); - } -#endif - int32_t sum = 0; - for (uint64_t i = 0; i < n; i++) sum += (int32_t)a[i] * (int32_t)b[i]; - return sum; -} - -static inline float dot_q8_0_row( - const uint8_t *row, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t blocks) { -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - if ((in_dim & 31u) == 0) { - float32x4_t accv0 = vdupq_n_f32(0.0f); - float32x4_t accv1 = vdupq_n_f32(0.0f); - - uint64_t b = 0; - for (; b + 1 < blocks; b += 2) { - uint16_t scale_bits0; - uint16_t scale_bits1; - memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); - memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); - - const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); - const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); - const int8_t *xq0 = xq + b * 32; - const int8_t *xq1 = xq + (b + 1) * 32; - - int32x4_t dot0 = vdupq_n_s32(0); - dot0 = vdotq_s32(dot0, vld1q_s8(qs0), vld1q_s8(xq0)); - dot0 = vdotq_s32(dot0, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + 16)); - - int32x4_t dot1 = vdupq_n_s32(0); - dot1 = vdotq_s32(dot1, vld1q_s8(qs1), vld1q_s8(xq1)); - dot1 = vdotq_s32(dot1, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + 16)); - - accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot0), f16_to_f32(scale_bits0) * xscale[b]); - accv1 = vfmaq_n_f32(accv1, vcvtq_f32_s32(dot1), f16_to_f32(scale_bits1) * xscale[b + 1]); - } - - if (b < blocks) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); - const int8_t *qs = (const int8_t *)(row + b * 34 + 2); - const int8_t *xqb = xq + b * 32; - int32x4_t dot = vdupq_n_s32(0); - dot = vdotq_s32(dot, vld1q_s8(qs), vld1q_s8(xqb)); - dot = vdotq_s32(dot, vld1q_s8(qs + 16), vld1q_s8(xqb + 16)); - accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot), f16_to_f32(scale_bits) * xscale[b]); - } - - return vaddvq_f32(vaddq_f32(accv0, accv1)); - } -#endif - - float acc = 0.0f; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); - const int8_t *qs = (const int8_t *)(row + b * 34 + 2); - - const uint64_t i0 = b * 32; - const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; - acc += f16_to_f32(scale_bits) * xscale[b] * (float)dot_i8_32(qs, xq + i0, n); - } - return acc; -} - -static inline void dot_q8_0_row_2( - const uint8_t *row, - const int8_t *xq0, - const float *xscale0, - const int8_t *xq1, - const float *xscale1, - uint64_t in_dim, - uint64_t blocks, - float *out0, - float *out1) { -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - if ((in_dim & 31u) == 0) { - float32x4_t acc00 = vdupq_n_f32(0.0f); - float32x4_t acc01 = vdupq_n_f32(0.0f); - float32x4_t acc10 = vdupq_n_f32(0.0f); - float32x4_t acc11 = vdupq_n_f32(0.0f); - - uint64_t b = 0; - for (; b + 1 < blocks; b += 2) { - uint16_t scale_bits0; - uint16_t scale_bits1; - memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); - memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); - - const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); - const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); - - int32x4_t d00 = vdupq_n_s32(0); - d00 = vdotq_s32(d00, vld1q_s8(qs0), vld1q_s8(xq0 + b * 32)); - d00 = vdotq_s32(d00, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + b * 32 + 16)); - int32x4_t d01 = vdupq_n_s32(0); - d01 = vdotq_s32(d01, vld1q_s8(qs1), vld1q_s8(xq0 + (b + 1) * 32)); - d01 = vdotq_s32(d01, vld1q_s8(qs1 + 16), vld1q_s8(xq0 + (b + 1) * 32 + 16)); - - int32x4_t d10 = vdupq_n_s32(0); - d10 = vdotq_s32(d10, vld1q_s8(qs0), vld1q_s8(xq1 + b * 32)); - d10 = vdotq_s32(d10, vld1q_s8(qs0 + 16), vld1q_s8(xq1 + b * 32 + 16)); - int32x4_t d11 = vdupq_n_s32(0); - d11 = vdotq_s32(d11, vld1q_s8(qs1), vld1q_s8(xq1 + (b + 1) * 32)); - d11 = vdotq_s32(d11, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + (b + 1) * 32 + 16)); - - const float s0 = f16_to_f32(scale_bits0); - const float s1 = f16_to_f32(scale_bits1); - acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), s0 * xscale0[b]); - acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), s1 * xscale0[b + 1]); - acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), s0 * xscale1[b]); - acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), s1 * xscale1[b + 1]); - } - - if (b < blocks) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); - const int8_t *qs = (const int8_t *)(row + b * 34 + 2); - - int32x4_t d0 = vdupq_n_s32(0); - d0 = vdotq_s32(d0, vld1q_s8(qs), vld1q_s8(xq0 + b * 32)); - d0 = vdotq_s32(d0, vld1q_s8(qs + 16), vld1q_s8(xq0 + b * 32 + 16)); - int32x4_t d1 = vdupq_n_s32(0); - d1 = vdotq_s32(d1, vld1q_s8(qs), vld1q_s8(xq1 + b * 32)); - d1 = vdotq_s32(d1, vld1q_s8(qs + 16), vld1q_s8(xq1 + b * 32 + 16)); - - const float s0 = f16_to_f32(scale_bits); - acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), s0 * xscale0[b]); - acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), s0 * xscale1[b]); - } - - *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); - *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); - return; - } -#endif - - *out0 = dot_q8_0_row(row, xq0, xscale0, in_dim, blocks); - *out1 = dot_q8_0_row(row, xq1, xscale1, in_dim, blocks); -} - -static inline DS4_MAYBE_UNUSED void dot_q8_0_row_pair( - const uint8_t *row0, - const uint8_t *row1, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t blocks, - float *out0, - float *out1) { -#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) - if ((in_dim & 31u) == 0) { - float32x4_t acc00 = vdupq_n_f32(0.0f); - float32x4_t acc01 = vdupq_n_f32(0.0f); - float32x4_t acc10 = vdupq_n_f32(0.0f); - float32x4_t acc11 = vdupq_n_f32(0.0f); - - uint64_t b = 0; - for (; b + 1 < blocks; b += 2) { - uint16_t s00, s01, s10, s11; - memcpy(&s00, row0 + b * 34, sizeof(s00)); - memcpy(&s01, row0 + (b + 1) * 34, sizeof(s01)); - memcpy(&s10, row1 + b * 34, sizeof(s10)); - memcpy(&s11, row1 + (b + 1) * 34, sizeof(s11)); - - const int8_t *xq0 = xq + b * 32; - const int8_t *xq1 = xq + (b + 1) * 32; - const int8x16_t xv00 = vld1q_s8(xq0); - const int8x16_t xv01 = vld1q_s8(xq0 + 16); - const int8x16_t xv10 = vld1q_s8(xq1); - const int8x16_t xv11 = vld1q_s8(xq1 + 16); - - const int8_t *q00 = (const int8_t *)(row0 + b * 34 + 2); - const int8_t *q01 = (const int8_t *)(row0 + (b + 1) * 34 + 2); - const int8_t *q10 = (const int8_t *)(row1 + b * 34 + 2); - const int8_t *q11 = (const int8_t *)(row1 + (b + 1) * 34 + 2); - - int32x4_t d00 = vdupq_n_s32(0); - d00 = vdotq_s32(d00, vld1q_s8(q00), xv00); - d00 = vdotq_s32(d00, vld1q_s8(q00 + 16), xv01); - int32x4_t d01 = vdupq_n_s32(0); - d01 = vdotq_s32(d01, vld1q_s8(q01), xv10); - d01 = vdotq_s32(d01, vld1q_s8(q01 + 16), xv11); - int32x4_t d10 = vdupq_n_s32(0); - d10 = vdotq_s32(d10, vld1q_s8(q10), xv00); - d10 = vdotq_s32(d10, vld1q_s8(q10 + 16), xv01); - int32x4_t d11 = vdupq_n_s32(0); - d11 = vdotq_s32(d11, vld1q_s8(q11), xv10); - d11 = vdotq_s32(d11, vld1q_s8(q11 + 16), xv11); - - acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), f16_to_f32(s00) * xscale[b]); - acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), f16_to_f32(s01) * xscale[b + 1]); - acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), f16_to_f32(s10) * xscale[b]); - acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), f16_to_f32(s11) * xscale[b + 1]); - } - - if (b < blocks) { - uint16_t s0, s1; - memcpy(&s0, row0 + b * 34, sizeof(s0)); - memcpy(&s1, row1 + b * 34, sizeof(s1)); - const int8_t *xqb = xq + b * 32; - const int8x16_t xv0 = vld1q_s8(xqb); - const int8x16_t xv1 = vld1q_s8(xqb + 16); - const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); - const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); - int32x4_t d0 = vdupq_n_s32(0); - d0 = vdotq_s32(d0, vld1q_s8(q0), xv0); - d0 = vdotq_s32(d0, vld1q_s8(q0 + 16), xv1); - int32x4_t d1 = vdupq_n_s32(0); - d1 = vdotq_s32(d1, vld1q_s8(q1), xv0); - d1 = vdotq_s32(d1, vld1q_s8(q1 + 16), xv1); - acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), f16_to_f32(s0) * xscale[b]); - acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), f16_to_f32(s1) * xscale[b]); - } - - *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); - *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); - return; - } -#endif - - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t s0_bits; - uint16_t s1_bits; - memcpy(&s0_bits, row0 + b * 34, sizeof(s0_bits)); - memcpy(&s1_bits, row1 + b * 34, sizeof(s1_bits)); - const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); - const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); - const uint64_t i0 = b * 32; - const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; - acc0 += f16_to_f32(s0_bits) * xscale[b] * (float)dot_i8_32(q0, xq + i0, n); - acc1 += f16_to_f32(s1_bits) * xscale[b] * (float)dot_i8_32(q1, xq + i0, n); - } - *out0 = acc0; - *out1 = acc1; -} - -static void quantize_q8_0_activation(const float *x, int8_t *xq, float *scale, uint64_t n) { - const uint64_t blocks = (n + 31) / 32; - for (uint64_t b = 0; b < blocks; b++) { - const uint64_t i0 = b * 32; - const uint64_t bn = n - i0 < 32 ? n - i0 : 32; - float amax = 0.0f; - for (uint64_t i = 0; i < bn; i++) { - const float ax = fabsf(x[i0 + i]); - if (ax > amax) amax = ax; - } - const float d = amax / 127.0f; - const float id = d != 0.0f ? 1.0f / d : 0.0f; - scale[b] = d; - for (uint64_t i = 0; i < bn; i++) { - int v = (int)lrintf(x[i0 + i] * id); - if (v > 127) v = 127; - if (v < -128) v = -128; - xq[i0 + i] = (int8_t)v; - } - for (uint64_t i = bn; i < 32 && i0 + i < blocks * 32; i++) { - xq[i0 + i] = 0; - } - } -} - -static void quantize_q8_0_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - quantize_q8_0_batch_ctx *ctx = vctx; - for (uint64_t t = t0; t < t1; t++) { - quantize_q8_0_activation(ctx->x + t * ctx->in_dim, - ctx->xq + t * ctx->blocks * 32, - ctx->xscale + t * ctx->blocks, - ctx->in_dim); - } -} - -static void quantize_q8_0_activation_batch( - const float *x, - int8_t *xq, - float *xscale, - uint64_t n_tok, - uint64_t in_dim) { - quantize_q8_0_batch_ctx ctx = { - .x = x, - .xq = xq, - .xscale = xscale, - .in_dim = in_dim, - .blocks = (in_dim + 31) / 32, - }; - ds4_parallel_for(n_tok, quantize_q8_0_batch_worker, &ctx); -} - -static void matvec_q8_0_worker(void *vctx, uint64_t r0, uint64_t r1) { - matvec_q8_0_ctx *ctx = vctx; - - for (uint64_t r = r0; r < r1; r++) { - const uint64_t o = ctx->row0 + r; - const uint8_t *row = ctx->data + o * ctx->blocks * 34; - ctx->out[r] = dot_q8_0_row(row, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks); - } -} - -static void matvec_q8_0_pair_worker(void *vctx, uint64_t r0, uint64_t r1) { - matvec_q8_0_pair_ctx *ctx = vctx; - - for (uint64_t r = r0; r < r1; r++) { - const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; - const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; - dot_q8_0_row_pair(row0, row1, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks, - ctx->out0 + r, ctx->out1 + r); - } -} - -static void matvec_q8_0_grouped_worker(void *vctx, uint64_t r0, uint64_t r1) { - matvec_q8_0_grouped_ctx *ctx = vctx; - - for (uint64_t idx = r0; idx < r1; idx++) { - const uint64_t group = idx / ctx->rank; - const uint64_t row_in_group = idx - group * ctx->rank; - const uint64_t tensor_row = group * ctx->rank + row_in_group; - const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; - const int8_t *xq = ctx->xq + group * ctx->blocks * 32; - const float *xscale = ctx->xscale + group * ctx->blocks; - ctx->out[idx] = dot_q8_0_row(row, xq, xscale, ctx->in_dim, ctx->blocks); - } -} - -static void matmul_q8_0_grouped_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { - matmul_q8_0_grouped_batch_ctx *ctx = vctx; - - for (uint64_t idx = r0; idx < r1; idx++) { - const uint64_t group = idx / ctx->rank; - const uint64_t row_in_group = idx - group * ctx->rank; - const uint64_t tensor_row = group * ctx->rank + row_in_group; - const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; - - uint64_t t = 0; - for (; t + 1 < ctx->n_tok; t += 2) { - const uint64_t xbase0 = (t * ctx->n_groups + group) * ctx->blocks; - const uint64_t xbase1 = ((t + 1) * ctx->n_groups + group) * ctx->blocks; - dot_q8_0_row_2(row, - ctx->xq + xbase0 * 32, - ctx->xscale + xbase0, - ctx->xq + xbase1 * 32, - ctx->xscale + xbase1, - ctx->group_dim, - ctx->blocks, - ctx->out + t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group, - ctx->out + (t + 1) * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group); - } - for (; t < ctx->n_tok; t++) { - const uint64_t xbase = (t * ctx->n_groups + group) * ctx->blocks; - ctx->out[t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group] = - dot_q8_0_row(row, - ctx->xq + xbase * 32, - ctx->xscale + xbase, - ctx->group_dim, - ctx->blocks); - } - } -} - -static void matmul_q8_0_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { - matmul_q8_0_batch_ctx *ctx = vctx; - - for (uint64_t r = r0; r < r1; r++) { - const uint8_t *row = ctx->data + r * ctx->blocks * 34; - uint64_t t = 0; - for (; t + 1 < ctx->n_tok; t += 2) { - dot_q8_0_row_2(row, - ctx->xq + t * ctx->blocks * 32, - ctx->xscale + t * ctx->blocks, - ctx->xq + (t + 1) * ctx->blocks * 32, - ctx->xscale + (t + 1) * ctx->blocks, - ctx->in_dim, - ctx->blocks, - ctx->out + t * ctx->out_dim + r, - ctx->out + (t + 1) * ctx->out_dim + r); - } - for (; t < ctx->n_tok; t++) { - ctx->out[t * ctx->out_dim + r] = - dot_q8_0_row(row, - ctx->xq + t * ctx->blocks * 32, - ctx->xscale + t * ctx->blocks, - ctx->in_dim, - ctx->blocks); - } - } -} - -static void matmul_q8_0_pair_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { - matmul_q8_0_pair_batch_ctx *ctx = vctx; - - for (uint64_t r = r0; r < r1; r++) { - const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; - const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; - uint64_t t = 0; - for (; t + 1 < ctx->n_tok; t += 2) { - const int8_t *xq0 = ctx->xq + t * ctx->blocks * 32; - const float *xscale0 = ctx->xscale + t * ctx->blocks; - const int8_t *xq1 = ctx->xq + (t + 1) * ctx->blocks * 32; - const float *xscale1 = ctx->xscale + (t + 1) * ctx->blocks; - dot_q8_0_row_2(row0, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, - ctx->out0 + t * ctx->out_dim + r, - ctx->out0 + (t + 1) * ctx->out_dim + r); - dot_q8_0_row_2(row1, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, - ctx->out1 + t * ctx->out_dim + r, - ctx->out1 + (t + 1) * ctx->out_dim + r); - } - for (; t < ctx->n_tok; t++) { - const int8_t *xq = ctx->xq + t * ctx->blocks * 32; - const float *xscale = ctx->xscale + t * ctx->blocks; - dot_q8_0_row_pair(row0, row1, xq, xscale, ctx->in_dim, ctx->blocks, - ctx->out0 + t * ctx->out_dim + r, - ctx->out1 + t * ctx->out_dim + r); - } - } -} - -/* Multiply selected Q8_0 rows by an activation that has already been quantized - * once. This avoids repeated activation quantization for paired projections. */ -static void matvec_q8_0_rows_prequant( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const int8_t * xq, - const float * xscale, - uint64_t row0, - uint64_t n_rows) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t out_dim = w->dim[1]; - if (row0 > out_dim || n_rows > out_dim - row0) ds4_die("Q8_0 row range is outside tensor"); - const uint64_t ctx_blocks = (in_dim + 31) / 32; - - matvec_q8_0_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .xq = xq, - .xscale = xscale, - .in_dim = in_dim, - .row0 = row0, - .blocks = ctx_blocks, - }; - ds4_parallel_for(n_rows, matvec_q8_0_worker, &ctx); -} - -static DS4_MAYBE_UNUSED void matvec_q8_0_prequant( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const int8_t * xq, - const float * xscale) { - matvec_q8_0_rows_prequant(out, m, w, xq, xscale, 0, w->dim[1]); -} - -static void matvec_q8_0_3d_slice_prequant( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const int8_t * xq, - const float * xscale, - uint64_t slice) { - if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); - if (slice >= w->dim[2]) ds4_die("Q8_0 slice is outside tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t out_dim = w->dim[1]; - const uint64_t blocks = (in_dim + 31) / 32; - const uint64_t slice_bytes = out_dim * blocks * 34; - const uint8_t *data = (const uint8_t *)tensor_data(m, w) + slice * slice_bytes; - - matvec_q8_0_ctx ctx = { - .out = out, - .data = data, - .xq = xq, - .xscale = xscale, - .in_dim = in_dim, - .row0 = 0, - .blocks = blocks, - }; - ds4_parallel_for(out_dim, matvec_q8_0_worker, &ctx); -} - -static DS4_MAYBE_UNUSED void matvec_q8_0_3d_slice( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint64_t slice) { - if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)blocks * 32); - float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); - - quantize_q8_0_activation(x, xq, xscale, in_dim); - matvec_q8_0_3d_slice_prequant(out, m, w, xq, xscale, slice); - - free(xscale); - free(xq); -} - -/* Compute two Q8_0 projections from the same input, used by gate/up and - * compressor kv/score pairs. */ -static void matvec_q8_0_pair_prequant( - float * out0, - float * out1, - const ds4_model * m, - const ds4_tensor * w0, - const ds4_tensor * w1, - const int8_t * xq, - const float * xscale) { - if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { - ds4_die("expected two 2D Q8_0 tensors"); - } - if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { - ds4_die("paired Q8_0 tensors do not have the same shape"); - } - - const uint64_t in_dim = w0->dim[0]; - matvec_q8_0_pair_ctx ctx = { - .out0 = out0, - .out1 = out1, - .data0 = tensor_data(m, w0), - .data1 = tensor_data(m, w1), - .xq = xq, - .xscale = xscale, - .in_dim = in_dim, - .blocks = (in_dim + 31) / 32, - }; - ds4_parallel_for(w0->dim[1], matvec_q8_0_pair_worker, &ctx); -} - -static void matmul_q8_0_batch_prequant( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const int8_t * xq, - const float * xscale, - uint64_t n_tok) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - - matmul_q8_0_batch_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .xq = xq, - .xscale = xscale, - .n_tok = n_tok, - .in_dim = w->dim[0], - .out_dim = w->dim[1], - .blocks = (w->dim[0] + 31) / 32, - }; - ds4_parallel_for(ctx.out_dim, matmul_q8_0_batch_worker, &ctx); -} - -static void matmul_q8_0_pair_batch_prequant( - float * out0, - float * out1, - const ds4_model * m, - const ds4_tensor * w0, - const ds4_tensor * w1, - const int8_t * xq, - const float * xscale, - uint64_t n_tok) { - if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { - ds4_die("expected two 2D Q8_0 tensors"); - } - if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { - ds4_die("paired Q8_0 tensors do not have the same shape"); - } - - matmul_q8_0_pair_batch_ctx ctx = { - .out0 = out0, - .out1 = out1, - .data0 = tensor_data(m, w0), - .data1 = tensor_data(m, w1), - .xq = xq, - .xscale = xscale, - .n_tok = n_tok, - .in_dim = w0->dim[0], - .out_dim = w0->dim[1], - .blocks = (w0->dim[0] + 31) / 32, - }; - ds4_parallel_for(ctx.out_dim, matmul_q8_0_pair_batch_worker, &ctx); -} - -/* Batched Q8_0 matmul for prefill: quantize all token activations, then scan - * weight rows once per output channel. */ -static void matmul_q8_0_batch( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint64_t n_tok) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); - float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); - - quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); - matmul_q8_0_batch_prequant(out, m, w, xq, xscale, n_tok); - - free(xscale); - free(xq); -} - -static void matmul_q8_0_pair_batch( - float * out0, - float * out1, - const ds4_model * m, - const ds4_tensor * w0, - const ds4_tensor * w1, - const float * x, - uint64_t n_tok) { - if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { - ds4_die("expected two 2D Q8_0 tensors"); - } - if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { - ds4_die("paired Q8_0 tensors do not have the same shape"); - } - - const uint64_t in_dim = w0->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); - float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); - - quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); - matmul_q8_0_pair_batch_prequant(out0, out1, m, w0, w1, xq, xscale, n_tok); - - free(xscale); - free(xq); -} - -static void matvec_q8_0_rows( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint64_t row0, - uint64_t n_rows) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - - const uint64_t in_dim = w->dim[0]; - const uint64_t ctx_blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)ctx_blocks * 32); - float *xscale = xmalloc((size_t)ctx_blocks * sizeof(xscale[0])); - - quantize_q8_0_activation(x, xq, xscale, in_dim); - matvec_q8_0_rows_prequant(out, m, w, xq, xscale, row0, n_rows); - - free(xscale); - free(xq); -} - -/* Single-token Q8_0 matvec, used heavily in decode. */ -static void matvec_q8_0(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { - matvec_q8_0_rows(out, m, w, x, 0, w->dim[1]); -} - -static inline float dot_q8_0_row_f32_ref( - const uint8_t *row, - const float *x, - uint64_t in_dim, - uint64_t blocks) { - float acc = 0.0f; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); - const int8_t *qs = (const int8_t *)(row + b * 34 + 2); - const float d = f16_to_f32(scale_bits); - const uint64_t i0 = b * 32; - const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; - for (uint64_t i = 0; i < n; i++) { - acc += d * (float)qs[i] * x[i0 + i]; - } - } - return acc; -} - -typedef struct { - float *out; - const uint8_t *data; - const float *x; - uint64_t in_dim; - uint64_t blocks; -} matvec_q8_0_f32_ref_ctx; - -static void matvec_q8_0_f32_ref_worker(void *vctx, uint64_t r0, uint64_t r1) { - matvec_q8_0_f32_ref_ctx *ctx = vctx; - const uint64_t row_bytes = ctx->blocks * 34; - for (uint64_t r = r0; r < r1; r++) { - ctx->out[r] = dot_q8_0_row_f32_ref(ctx->data + r * row_bytes, - ctx->x, - ctx->in_dim, - ctx->blocks); - } -} - -static void matvec_q8_0_f32_ref( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const float *x) { - if (w->type != DS4_TENSOR_Q8_0 || w->ndim < 2 || w->dim[0] == 0) { - ds4_die("expected a Q8_0 tensor with matrix rows"); - } - matvec_q8_0_f32_ref_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .x = x, - .in_dim = w->dim[0], - .blocks = (w->dim[0] + 31) / 32, - }; - ds4_parallel_for(w->elements / w->dim[0], matvec_q8_0_f32_ref_worker, &ctx); -} - -static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x); - -/* Decode scratch owns this temporary activation quantization so generation - * can assert that the hot path performs no malloc. */ -static void cpu_decode_quantize_q8_0( - ds4_cpu_decode_scratch * scratch, - const float * x, - uint64_t in_dim) { - if (in_dim > scratch->q8_cap) ds4_die("CPU decode Q8_0 scratch buffer is too small"); - quantize_q8_0_activation(x, scratch->q8_xq, scratch->q8_xscale, in_dim); -} - -static void matvec_q8_0_decode_scratch( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - ds4_cpu_decode_scratch * scratch) { - cpu_decode_quantize_q8_0(scratch, x, w->dim[0]); - matvec_q8_0_prequant(out, m, w, scratch->q8_xq, scratch->q8_xscale); -} - -static void matvec_q8_0_pair_decode_scratch( - float * out0, - float * out1, - const ds4_model * m, - const ds4_tensor * w0, - const ds4_tensor * w1, - const float * x, - ds4_cpu_decode_scratch * scratch) { - cpu_decode_quantize_q8_0(scratch, x, w0->dim[0]); - matvec_q8_0_pair_prequant(out0, out1, m, w0, w1, scratch->q8_xq, scratch->q8_xscale); -} - -static void matvec_any_decode_scratch( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - ds4_cpu_decode_scratch * scratch) { - if (w->type == 8) { - matvec_q8_0_decode_scratch(out, m, w, x, scratch); - } else { - matvec_any(out, m, w, x); - } -} - -static void matvec_q8_0_grouped_rows( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint32_t n_groups, - uint64_t group_dim, - uint64_t rank) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { - ds4_die("grouped Q8_0 tensor has an unexpected layout"); - } - - const uint64_t blocks = (group_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)n_groups * blocks * 32); - float *xscale = xmalloc((size_t)n_groups * blocks * sizeof(xscale[0])); - - for (uint32_t g = 0; g < n_groups; g++) { - quantize_q8_0_activation(x + (uint64_t)g * group_dim, - xq + (uint64_t)g * blocks * 32, - xscale + (uint64_t)g * blocks, - group_dim); - } - - matvec_q8_0_grouped_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .xq = xq, - .xscale = xscale, - .in_dim = group_dim, - .blocks = blocks, - .rank = rank, - }; - ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); - - free(xscale); - free(xq); -} - -static void matvec_q8_0_grouped_rows_decode_scratch( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint32_t n_groups, - uint64_t group_dim, - uint64_t rank, - ds4_cpu_decode_scratch * scratch) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { - ds4_die("grouped Q8_0 tensor has an unexpected layout"); - } - if ((uint64_t)n_groups * group_dim > scratch->q8_cap) { - ds4_die("CPU decode grouped Q8_0 scratch buffer is too small"); - } - - const uint64_t blocks = (group_dim + 31) / 32; - for (uint32_t g = 0; g < n_groups; g++) { - quantize_q8_0_activation(x + (uint64_t)g * group_dim, - scratch->q8_xq + (uint64_t)g * blocks * 32, - scratch->q8_xscale + (uint64_t)g * blocks, - group_dim); - } - - matvec_q8_0_grouped_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .xq = scratch->q8_xq, - .xscale = scratch->q8_xscale, - .in_dim = group_dim, - .blocks = blocks, - .rank = rank, - }; - ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); -} - -static void matmul_q8_0_grouped_batch( - float * out, - const ds4_model * m, - const ds4_tensor * w, - const float * x, - uint64_t n_tok, - uint32_t n_groups, - uint64_t group_dim, - uint64_t rank) { - if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); - if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { - ds4_die("grouped Q8_0 tensor has an unexpected layout"); - } - - const uint64_t blocks = (group_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)n_tok * n_groups * blocks * 32); - float *xscale = xmalloc((size_t)n_tok * n_groups * blocks * sizeof(xscale[0])); - - for (uint64_t t = 0; t < n_tok; t++) { - for (uint32_t g = 0; g < n_groups; g++) { - const uint64_t xbase = (t * n_groups + g) * blocks; - quantize_q8_0_activation(x + t * n_groups * group_dim + (uint64_t)g * group_dim, - xq + xbase * 32, - xscale + xbase, - group_dim); - } - } - - matmul_q8_0_grouped_batch_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .xq = xq, - .xscale = xscale, - .n_tok = n_tok, - .n_groups = n_groups, - .group_dim = group_dim, - .blocks = blocks, - .rank = rank, - }; - ds4_parallel_for((uint64_t)n_groups * rank, matmul_q8_0_grouped_batch_worker, &ctx); - - free(xscale); - free(xq); -} - -typedef struct { - float *out; - const float *data; - const float *x; - uint64_t in_dim; -} matvec_f32_ctx; - -static void matvec_f32_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_f32_ctx *ctx = vctx; - - for (uint64_t o = row0; o < row1; o++) { - double acc = 0.0; - const float *row = ctx->data + o * ctx->in_dim; - for (uint64_t i = 0; i < ctx->in_dim; i++) { - acc += (double)row[i] * ctx->x[i]; - } - ctx->out[o] = (float)acc; - } -} - -static void matvec_f32(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { - if (w->type != 0 || w->ndim != 2) ds4_die("expected a 2D F32 tensor"); - - matvec_f32_ctx ctx = { - .out = out, - .data = tensor_data(m, w), - .x = x, - .in_dim = w->dim[0], - }; - ds4_parallel_for(w->dim[1], matvec_f32_worker, &ctx); -} - -/* Dispatch for dense F32/F16/Q8_0 tensors used by auxiliary projections. */ -static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { - switch (w->type) { - case 0: matvec_f32(out, m, w, x); break; - case 1: matvec_f16(out, m, w, x); break; - case 8: matvec_q8_0(out, m, w, x); break; - default: - ds4_die("unsupported tensor type for dense matvec"); - } -} - -static float tensor_1d_value(const ds4_model *m, const ds4_tensor *t, uint64_t i) { - if (i >= t->elements) ds4_die("tensor scalar index is out of bounds"); - if (t->type == 0) { - const float *p = tensor_data(m, t); - return p[i]; - } - if (t->type == 1) { - const uint16_t *p = tensor_data(m, t); - return f16_to_f32(p[i]); - } - ds4_die("unsupported tensor scalar type"); - return 0.0f; -} - -static float tensor_2d_value(const ds4_model *m, const ds4_tensor *t, uint64_t x, uint64_t y) { - if (t->ndim != 2 || x >= t->dim[0] || y >= t->dim[1]) { - ds4_die("tensor 2D index is out of bounds"); - } - return tensor_1d_value(m, t, y * t->dim[0] + x); -} - -/* Locate one expert's 2D matrix inside a 3D GGUF expert tensor. */ -static const uint8_t *tensor_expert_bytes( - const ds4_model *m, - const ds4_tensor *w, - uint32_t expert, - uint64_t *in_dim, - uint64_t *out_dim, - uint64_t *row_bytes) { - if (w->ndim != 3) ds4_die("expected a 3D expert tensor"); - if (expert >= w->dim[2]) ds4_die("expert id is outside expert tensor"); - - *in_dim = w->dim[0]; - *out_dim = w->dim[1]; - - const gguf_type_info *info = tensor_type(w->type); - if (!info || info->block_elems == 0) ds4_die("unsupported expert tensor type"); - const uint64_t blocks = (*in_dim + info->block_elems - 1) / info->block_elems; - *row_bytes = blocks * info->block_bytes; - - const uint64_t expert_bytes = *out_dim * *row_bytes; - return (const uint8_t *)tensor_data(m, w) + (uint64_t)expert * expert_bytes; -} - -typedef struct { - float *out0; - float *out1; - const uint8_t *base0; - const uint8_t *base1; - const block_q8_K *xq; - uint64_t in_dim; - uint64_t row_bytes0; - uint64_t row_bytes1; -} matvec_iq2_xxs_pair_ctx; - -static void matvec_iq2_xxs_pair_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_iq2_xxs_pair_ctx *ctx = vctx; - for (uint64_t row = row0; row < row1; row++) { - const block_iq2_xxs *br0 = (const block_iq2_xxs *)(ctx->base0 + row * ctx->row_bytes0); - const block_iq2_xxs *br1 = (const block_iq2_xxs *)(ctx->base1 + row * ctx->row_bytes1); - ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &ctx->out0[row], &ctx->out1[row], br0, br1, ctx->xq); - } -} - -/* Project one routed expert's gate and up matrices. Both are IQ2_XXS and - * share the same Q8_K activation. */ -static void matvec_iq2_xxs_expert_pair_prequant( - float *out0, - float *out1, - const ds4_model *m, - const ds4_tensor *w0, - const ds4_tensor *w1, - const block_q8_K *xq, - uint32_t expert) { - if (w0->type != 16 || w1->type != 16) ds4_die("expected IQ2_XXS expert tensors"); - - uint64_t in_dim0, out_dim0, row_bytes0; - uint64_t in_dim1, out_dim1, row_bytes1; - const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &row_bytes0); - const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &row_bytes1); - if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired IQ2_XXS expert tensors do not match"); - if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); - - matvec_iq2_xxs_pair_ctx ctx = { - .out0 = out0, - .out1 = out1, - .base0 = base0, - .base1 = base1, - .xq = xq, - .in_dim = in_dim0, - .row_bytes0 = row_bytes0, - .row_bytes1 = row_bytes1, - }; - ds4_parallel_for(out_dim0, matvec_iq2_xxs_pair_worker, &ctx); -} - -static float silu(float x); - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_iq2_xxs_mid_ctx; - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const int8_t *xq; - const float *xscale; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t blocks; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q8_0_mid_ctx; - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q8_k_mid_ctx; - -static void matvec_iq2_xxs_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_iq2_xxs_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); - const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); - ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -static void matvec_q8_0_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_0_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - const uint8_t *gate_row = ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]; - const uint8_t *up_row = ctx->up_base[slot] + row * ctx->up_row_bytes[slot]; - dot_q8_0_row_pair(gate_row, up_row, ctx->xq, ctx->xscale, - ctx->in_dim, ctx->blocks, &gate, &up); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -static void matvec_q8_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_k_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); - const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); - ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -/* Build all selected expert hidden vectors: IQ2_XXS gate/up, clamp, SwiGLU, - * and router weight. The down projection runs later on the quantized mids. */ -static void matvec_iq2_xxs_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type != 16 || up_w->type != 16) ds4_die("expected IQ2_XXS expert tensors"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_iq2_xxs_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .clamp = clamp, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired IQ2_XXS expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("IQ2_XXS expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_iq2_xxs_mid_worker, &ctx); -} - -static DS4_MAYBE_UNUSED void matvec_q8_0_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const int8_t *xq, - const float *xscale, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type != DS4_TENSOR_Q8_0 || up_w->type != DS4_TENSOR_Q8_0) { - ds4_die("expected Q8_0 expert tensors"); - } - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q8_0_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .xscale = xscale, - .clamp = clamp, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired Q8_0 expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("Q8_0 expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ctx.blocks = in_dim0 / 32u; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_0_mid_worker, &ctx); -} - -static DS4_MAYBE_UNUSED void matvec_q8_k_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type != DS4_TENSOR_Q8_K || up_w->type != DS4_TENSOR_Q8_K) { - ds4_die("expected Q8_K expert tensors"); - } - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q8_k_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .clamp = clamp, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired Q8_K expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("Q8_K expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_k_mid_worker, &ctx); -} - -typedef struct { - float *out; - const uint8_t *base; - const block_q8_K *xq; - uint64_t in_dim; - uint64_t row_bytes; -} matvec_q2_k_ctx; - -static void matvec_q2_k_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q2_k_ctx *ctx = vctx; - for (uint64_t row = row0; row < row1; row++) { - const block_q2_K *br = (const block_q2_K *)(ctx->base + row * ctx->row_bytes); - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &ctx->out[row], br, ctx->xq); - } -} - -/* Single expert Q2_K down projection, kept mostly for tracing and diagnostics. */ -static void matvec_q2_k_expert( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const float *x, - uint32_t expert) { - if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); - - uint64_t in_dim, out_dim, row_bytes; - const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); - if (in_dim % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); - - block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); - ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); - - matvec_q2_k_ctx ctx = { - .out = out, - .base = base, - .xq = xq, - .in_dim = in_dim, - .row_bytes = row_bytes, - }; - ds4_parallel_for(out_dim, matvec_q2_k_worker, &ctx); - - free(xq); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q2_k_accum_ctx; - -static void matvec_q2_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q2_k_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - float v = 0.0f; - const block_q2_K *br = (const block_q2_K *)(ctx->base[i] + row * ctx->row_bytes[i]); - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); - acc += v; - } - ctx->out[row] = acc; - } -} - -/* Accumulate all selected experts' Q2_K down projections directly into the - * 4096-wide MoE output. */ -static void matvec_q2_k_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("Q2_K expert tensors do not share a layout"); - } - } - if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); - - const uint64_t n_blocks = in_dim0 / QK_K; - matvec_q2_k_accum_ctx ctx = { - .out = out, - .in_dim = in_dim0, - .n_expert = n_expert, - }; - for (int i = 0; i < n_expert; i++) { - ctx.base[i] = base[i]; - ctx.row_bytes[i] = row_bytes[i]; - ctx.xq[i] = xq + (uint64_t)i * n_blocks; - } - - ds4_parallel_for(out_dim0, matvec_q2_k_accum_worker, &ctx); -} - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q2_k_mid_ctx; - -static void matvec_q2_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q2_k_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); - - const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -static void matvec_q2_k_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type != DS4_TENSOR_Q2_K || up_w->type != DS4_TENSOR_Q2_K) { - ds4_die("expected Q2_K expert tensors"); - } - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q2_k_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .clamp = clamp, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired Q2_K expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("Q2_K expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q2_k_mid_worker, &ctx); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const int8_t *xq[DS4_MAX_EXPERT_USED]; - const float *xscale[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t blocks; - int n_expert; -} matvec_q8_0_accum_ctx; - -static void matvec_q8_0_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_0_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - const uint8_t *br = ctx->base[i] + row * ctx->row_bytes[i]; - acc += dot_q8_0_row(br, ctx->xq[i], ctx->xscale[i], - ctx->in_dim, ctx->blocks); - } - ctx->out[row] = acc; - } -} - -static void matvec_q8_0_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const int8_t *xq, - const float *xscale, - const int *selected, - int n_expert) { - if (w->type != DS4_TENSOR_Q8_0) ds4_die("expected a Q8_0 expert tensor"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q8_0_accum_ctx ctx = { - .out = out, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], - &in_dim, &out_dim, &ctx.row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("Q8_0 expert tensors do not share a layout"); - } - } - if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); - - const uint64_t blocks0 = in_dim0 / 32u; - ctx.in_dim = in_dim0; - ctx.blocks = blocks0; - for (int i = 0; i < n_expert; i++) { - ctx.xq[i] = xq + (uint64_t)i * blocks0 * 32u; - ctx.xscale[i] = xscale + (uint64_t)i * blocks0; - } - - ds4_parallel_for(out_dim0, matvec_q8_0_accum_worker, &ctx); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q8_k_accum_ctx; - -static void matvec_q8_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_k_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - float v = 0.0f; - const block_q8_K *br = (const block_q8_K *)(ctx->base[i] + row * ctx->row_bytes[i]); - ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); - acc += v; - } - ctx->out[row] = acc; - } -} - -static void matvec_q8_k_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type != DS4_TENSOR_Q8_K) ds4_die("expected a Q8_K expert tensor"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q8_k_accum_ctx ctx = { - .out = out, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], - &in_dim, &out_dim, &ctx.row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("Q8_K expert tensors do not share a layout"); - } - } - if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); - - const uint64_t n_blocks = in_dim0 / QK_K; - ctx.in_dim = in_dim0; - for (int i = 0; i < n_expert; i++) { - ctx.xq[i] = xq + (uint64_t)i * n_blocks; - } - - ds4_parallel_for(out_dim0, matvec_q8_k_accum_worker, &ctx); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_iq2_xxs_accum_ctx; - -static void matvec_iq2_xxs_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_iq2_xxs_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - float v = 0.0f; - const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[i] + row * ctx->row_bytes[i]); - ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); - acc += v; - } - ctx->out[row] = acc; - } -} - -static void matvec_iq2_xxs_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type != DS4_TENSOR_IQ2_XXS) ds4_die("expected an IQ2_XXS expert tensor"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("IQ2_XXS expert tensors do not share a layout"); - } - } - if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); - - const uint64_t n_blocks = in_dim0 / QK_K; - matvec_iq2_xxs_accum_ctx ctx = { - .out = out, - .in_dim = in_dim0, - .n_expert = n_expert, - }; - for (int i = 0; i < n_expert; i++) { - ctx.base[i] = base[i]; - ctx.row_bytes[i] = row_bytes[i]; - ctx.xq[i] = xq + (uint64_t)i * n_blocks; - } - - ds4_parallel_for(out_dim0, matvec_iq2_xxs_accum_worker, &ctx); -} - -typedef struct { - uint32_t token; - uint32_t slot; -} ds4_expert_pair; - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT]; - const uint8_t *up_base[DS4_MAX_EXPERT]; - const block_q8_K *xq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - const float *pair_weight; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT]; - uint64_t up_row_bytes[DS4_MAX_EXPERT]; - uint64_t xq_blocks; -} matvec_q2_k_batch_mid_ctx; - -static void matvec_q2_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_q2_k_batch_mid_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - - const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); - const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; - float gate = 0.0f; - float up = 0.0f; - - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - - ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; - } - } -} - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT]; - const uint8_t *up_base[DS4_MAX_EXPERT]; - const block_q8_K *xq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - const float *pair_weight; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT]; - uint64_t up_row_bytes[DS4_MAX_EXPERT]; - uint64_t xq_blocks; -} matvec_iq2_xxs_batch_mid_ctx; - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT]; - const uint8_t *up_base[DS4_MAX_EXPERT]; - const block_q8_K *xq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - const float *pair_weight; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t xq_blocks; - uint64_t gate_row_bytes[DS4_MAX_EXPERT]; - uint64_t up_row_bytes[DS4_MAX_EXPERT]; -} matvec_q8_k_batch_mid_ctx; - -static void matvec_iq2_xxs_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_iq2_xxs_batch_mid_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - - const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); - const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; - float gate = 0.0f; - float up = 0.0f; - - ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - - ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; - } - } -} - -static void matvec_q8_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_q8_k_batch_mid_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - - const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); - const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; - float gate = 0.0f; - float up = 0.0f; - - ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - - ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; - } - } -} - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT]; - const uint8_t *up_base[DS4_MAX_EXPERT]; - const int8_t *xq; - const float *xscale; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - const float *pair_weight; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t blocks; - uint64_t gate_row_bytes[DS4_MAX_EXPERT]; - uint64_t up_row_bytes[DS4_MAX_EXPERT]; -} matvec_q8_0_batch_mid_ctx; - -static void matvec_q8_0_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_q8_0_batch_mid_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - - const uint8_t *gate_row = ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]; - const uint8_t *up_row = ctx->up_base[expert] + row * ctx->up_row_bytes[expert]; - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - float gate = 0.0f; - float up = 0.0f; - - dot_q8_0_row_pair(gate_row, up_row, - ctx->xq + (uint64_t)pair.token * ctx->blocks * 32u, - ctx->xscale + (uint64_t)pair.token * ctx->blocks, - ctx->in_dim, ctx->blocks, &gate, &up); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - - ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; - } - } -} - -typedef struct { - const float *mid; - block_q8_K *midq; - uint64_t down_in_dim; - uint64_t down_blocks; -} quantize_mid_pairs_ctx; - -static void quantize_mid_pairs_worker(void *vctx, uint64_t p0, uint64_t p1) { - quantize_mid_pairs_ctx *ctx = vctx; - for (uint64_t p = p0; p < p1; p++) { - ds4_quantize_row_q8_K(ctx->mid + p * ctx->down_in_dim, - ctx->midq + p * ctx->down_blocks, - (int64_t)ctx->down_in_dim); - } -} - -typedef struct { - float *down_pair; - const uint8_t *base[DS4_MAX_EXPERT]; - const block_q8_K *midq; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t midq_blocks; -} matvec_q2_k_batch_down_ctx; - -static DS4_MAYBE_UNUSED void matvec_q2_k_batch_down_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_q2_k_batch_down_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, - ctx->down_pair + (uint64_t)pair_id * ctx->out_dim + row, - br, xq); - } - } -} - -typedef struct { - float *moe; - const uint8_t *base[DS4_MAX_EXPERT]; - const block_q8_K *midq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint32_t n_active; - uint32_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t midq_blocks; -} matvec_q2_k_batch_accum_rows_ctx; - -static void matvec_q2_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q2_k_batch_accum_rows_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - for (uint32_t t = 0; t < ctx->n_tok; t++) { - ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; - } - - for (uint32_t ai = 0; ai < ctx->n_active; ai++) { - const uint32_t expert = ctx->active_expert[ai]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; - float v = 0.0f; - - ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, xq); - ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; - } - } - } -} - -/* ========================================================================= - * Q4_K routed expert matrix-vector products. - * ========================================================================= */ - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q4_k_mid_ctx; - -static void matvec_q4_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q4_k_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); - - const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -static void matvec_q4_k_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type != DS4_TENSOR_Q4_K || up_w->type != DS4_TENSOR_Q4_K) - ds4_die("expected Q4_K expert tensors"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q4_k_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .clamp = clamp, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired Q4_K expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("Q4_K expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q4_k_mid_worker, &ctx); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - int n_expert; -} matvec_q4_k_accum_ctx; - -static void matvec_q4_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q4_k_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - float v = 0.0f; - const block_q4_K *br = (const block_q4_K *)(ctx->base[i] + row * ctx->row_bytes[i]); - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); - acc += v; - } - ctx->out[row] = acc; - } -} - -static void matvec_q4_k_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type != DS4_TENSOR_Q4_K) ds4_die("expected a Q4_K expert tensor"); - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("Q4_K expert tensors do not share a layout"); - } - } - if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); - - const uint64_t n_blocks = in_dim0 / QK_K; - matvec_q4_k_accum_ctx ctx = { - .out = out, - .in_dim = in_dim0, - .n_expert = n_expert, - }; - for (int i = 0; i < n_expert; i++) { - ctx.base[i] = base[i]; - ctx.row_bytes[i] = row_bytes[i]; - ctx.xq[i] = xq + (uint64_t)i * n_blocks; - } - - ds4_parallel_for(out_dim0, matvec_q4_k_accum_worker, &ctx); -} - -static inline void ds4_vec_dot_q5_q6_K_q8_K( - uint32_t type, - int n, - float *s, - const uint8_t *x, - const block_q8_K *y) { - if (type == DS4_TENSOR_Q5_K) { - ds4_vec_dot_q5_K_q8_K(n, s, (const block_q5_K *)x, y); - } else if (type == DS4_TENSOR_Q6_K) { - ds4_vec_dot_q6_K_q8_K(n, s, (const block_q6_K *)x, y); - } else { - ds4_die("expected a Q5_K or Q6_K tensor"); - } -} - -typedef struct { - float *out; - const uint8_t *base; - const block_q8_K *xq; - uint64_t in_dim; - uint64_t row_bytes; - uint32_t type; -} matvec_q5_q6_k_ctx; - -static void matvec_q5_q6_k_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q5_q6_k_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &ctx->out[row], - ctx->base + row * ctx->row_bytes, ctx->xq); - } -} - -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq; - float expert_weight[DS4_MAX_EXPERT_USED]; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - uint32_t gate_type; - uint32_t up_type; - int n_expert; -} matvec_q5_q6_k_mid_ctx; - -static void matvec_q5_q6_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q5_q6_k_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const int slot = (int)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - float gate = 0.0f; - float up = 0.0f; - - ds4_vec_dot_q5_q6_K_q8_K(ctx->gate_type, (int)ctx->in_dim, &gate, - ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot], - ctx->xq); - ds4_vec_dot_q5_q6_K_q8_K(ctx->up_type, (int)ctx->in_dim, &up, - ctx->up_base[slot] + row * ctx->up_row_bytes[slot], - ctx->xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -static void matvec_q5_q6_k_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if ((gate_w->type != DS4_TENSOR_Q5_K && gate_w->type != DS4_TENSOR_Q6_K) || - (up_w->type != DS4_TENSOR_Q5_K && up_w->type != DS4_TENSOR_Q6_K)) { - ds4_die("expected Q5_K/Q6_K expert tensors"); - } - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - matvec_q5_q6_k_mid_ctx ctx = { - .mid = mid, - .xq = xq, - .clamp = clamp, - .gate_type = gate_w->type, - .up_type = up_w->type, - .n_expert = n_expert, - }; - - for (int i = 0; i < n_expert; i++) { - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], - &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); - ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], - &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); - if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { - ds4_die("paired Q5_K/Q6_K expert tensors do not match"); - } - if (i == 0) { - in_dim0 = gate_in_dim; - out_dim0 = gate_out_dim; - } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { - ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); - } - ctx.expert_weight[i] = expert_weight[i]; - } - if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); - - ctx.in_dim = in_dim0; - ctx.out_dim = out_dim0; - ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q5_q6_k_mid_worker, &ctx); -} - -static void matvec_q5_q6_k_expert( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const float *x, - uint32_t expert) { - if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { - ds4_die("expected a Q5_K or Q6_K expert tensor"); - } - - uint64_t in_dim, out_dim, row_bytes; - const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); - if (in_dim % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); - - block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); - ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); - - matvec_q5_q6_k_ctx ctx = { - .out = out, - .base = base, - .xq = xq, - .in_dim = in_dim, - .row_bytes = row_bytes, - .type = w->type, - }; - ds4_parallel_for(out_dim, matvec_q5_q6_k_worker, &ctx); - - free(xq); -} - -typedef struct { - float *out; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - const block_q8_K *xq[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - uint32_t type; - int n_expert; -} matvec_q5_q6_k_accum_ctx; - -static void matvec_q5_q6_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q5_q6_k_accum_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (int i = 0; i < ctx->n_expert; i++) { - float v = 0.0f; - ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &v, - ctx->base[i] + row * ctx->row_bytes[i], - ctx->xq[i]); - acc += v; - } - ctx->out[row] = acc; - } -} - -static void matvec_q5_q6_k_experts_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { - ds4_die("expected a Q5_K or Q6_K expert tensor"); - } - if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); - - uint64_t in_dim0 = 0; - uint64_t out_dim0 = 0; - const uint8_t *base[DS4_MAX_EXPERT_USED]; - uint64_t row_bytes[DS4_MAX_EXPERT_USED]; - - for (int i = 0; i < n_expert; i++) { - uint64_t in_dim, out_dim; - base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); - if (i == 0) { - in_dim0 = in_dim; - out_dim0 = out_dim; - } else if (in_dim != in_dim0 || out_dim != out_dim0) { - ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); - } - } - if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); - - const uint64_t n_blocks = in_dim0 / QK_K; - matvec_q5_q6_k_accum_ctx ctx = { - .out = out, - .in_dim = in_dim0, - .type = w->type, - .n_expert = n_expert, - }; - for (int i = 0; i < n_expert; i++) { - ctx.base[i] = base[i]; - ctx.row_bytes[i] = row_bytes[i]; - ctx.xq[i] = xq + (uint64_t)i * n_blocks; - } - - ds4_parallel_for(out_dim0, matvec_q5_q6_k_accum_worker, &ctx); -} - -/* Q4_K batch mid worker: same structure as IQ2_XXS batch but uses Q4_K dot. */ -typedef struct { - float *mid; - const uint8_t *gate_base[DS4_MAX_EXPERT]; - const uint8_t *up_base[DS4_MAX_EXPERT]; - const block_q8_K *xq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - const float *pair_weight; - float clamp; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT]; - uint64_t up_row_bytes[DS4_MAX_EXPERT]; - uint64_t xq_blocks; -} matvec_q4_k_batch_mid_ctx; - -static void matvec_q4_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { - matvec_q4_k_batch_mid_ctx *ctx = vctx; - - for (uint64_t task = task0; task < task1; task++) { - const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); - const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; - const uint32_t expert = ctx->active_expert[active_idx]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - - const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); - const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; - float gate = 0.0f; - float up = 0.0f; - - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, xq); - - if (ctx->clamp > 1.0e-6f) { - if (gate > ctx->clamp) gate = ctx->clamp; - if (up > ctx->clamp) up = ctx->clamp; - if (up < -ctx->clamp) up = -ctx->clamp; - } - - ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; - } - } -} - -/* Q4_K batch down accum worker: same structure as Q2_K batch but uses Q4_K dot. */ -typedef struct { - float *moe; - const uint8_t *base[DS4_MAX_EXPERT]; - const block_q8_K *midq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint32_t n_active; - uint32_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t midq_blocks; -} matvec_q4_k_batch_accum_rows_ctx; - -static void matvec_q4_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q4_k_batch_accum_rows_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - for (uint32_t t = 0; t < ctx->n_tok; t++) { - ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; - } - - for (uint32_t ai = 0; ai < ctx->n_active; ai++) { - const uint32_t expert = ctx->active_expert[ai]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const block_q4_K *br = (const block_q4_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; - float v = 0.0f; - - ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, xq); - ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; - } - } - } -} - -typedef struct { - float *moe; - const uint8_t *base[DS4_MAX_EXPERT]; - const block_q8_K *midq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint32_t n_active; - uint32_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t midq_blocks; -} matvec_iq2_xxs_batch_accum_rows_ctx; - -static void matvec_iq2_xxs_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_iq2_xxs_batch_accum_rows_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - for (uint32_t t = 0; t < ctx->n_tok; t++) { - ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; - } - - for (uint32_t ai = 0; ai < ctx->n_active; ai++) { - const uint32_t expert = ctx->active_expert[ai]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[expert] + row * ctx->row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; - float v = 0.0f; - - ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, xq); - ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; - } - } - } -} - -/* Dispatch: call the right gate/up mid builder based on tensor type. */ -static void matvec_experts_mid_prequant( - float *mid, - const ds4_model *m, - const ds4_tensor *gate_w, - const ds4_tensor *up_w, - const block_q8_K *xq, - const int *selected, - const float *expert_weight, - int n_expert, - float clamp) { - if (gate_w->type == DS4_TENSOR_IQ2_XXS) { - matvec_iq2_xxs_experts_mid_prequant(mid, m, gate_w, up_w, xq, - selected, expert_weight, n_expert, clamp); - } else if (gate_w->type == DS4_TENSOR_Q2_K) { - matvec_q2_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, - selected, expert_weight, n_expert, clamp); - } else if (gate_w->type == DS4_TENSOR_Q4_K) { - matvec_q4_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, - selected, expert_weight, n_expert, clamp); - } else if (gate_w->type == DS4_TENSOR_Q5_K || gate_w->type == DS4_TENSOR_Q6_K) { - matvec_q5_q6_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, - selected, expert_weight, n_expert, clamp); - } else { - ds4_die("unsupported gate/up expert tensor type"); - } -} - -/* Dispatch: call the right down-projection accumulator based on tensor type. */ -static void matvec_experts_down_accum_prequant( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const block_q8_K *xq, - const int *selected, - int n_expert) { - if (w->type == DS4_TENSOR_IQ2_XXS) { - matvec_iq2_xxs_experts_accum_prequant(out, m, w, xq, selected, n_expert); - } else if (w->type == DS4_TENSOR_Q2_K) { - matvec_q2_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); - } else if (w->type == DS4_TENSOR_Q4_K) { - matvec_q4_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); - } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { - matvec_q5_q6_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); - } else { - ds4_die("unsupported down expert tensor type"); - } -} - -/* Dispatch: single-expert gate/up pair for tracing. */ -static void matvec_expert_pair_prequant( - float *out0, - float *out1, - const ds4_model *m, - const ds4_tensor *w0, - const ds4_tensor *w1, - const block_q8_K *xq, - uint32_t expert) { - if (w0->type == DS4_TENSOR_IQ2_XXS) { - matvec_iq2_xxs_expert_pair_prequant(out0, out1, m, w0, w1, xq, expert); - } else if (w0->type == DS4_TENSOR_Q2_K) { - uint64_t in_dim0, out_dim0, rb0; - uint64_t in_dim1, out_dim1, rb1; - const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); - const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); - if (w1->type != DS4_TENSOR_Q2_K || - in_dim0 != in_dim1 || - out_dim0 != out_dim1) { - ds4_die("paired Q2_K expert tensors do not match"); - } - - for (uint64_t row = 0; row < out_dim0; row++) { - const block_q2_K *gr = (const block_q2_K *)(base0 + row * rb0); - ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out0[row], gr, xq); - const block_q2_K *ur = (const block_q2_K *)(base1 + row * rb1); - ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out1[row], ur, xq); - } - } else if (w0->type == DS4_TENSOR_Q4_K) { - uint64_t in_dim0, out_dim0, rb0; - uint64_t in_dim1, out_dim1, rb1; - const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); - const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); - if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q4_K expert tensors do not match"); - - for (uint64_t row = 0; row < out_dim0; row++) { - const block_q4_K *gr = (const block_q4_K *)(base0 + row * rb0); - ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out0[row], gr, xq); - const block_q4_K *ur = (const block_q4_K *)(base1 + row * rb1); - ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out1[row], ur, xq); - } - } else if (w0->type == DS4_TENSOR_Q5_K || w0->type == DS4_TENSOR_Q6_K) { - uint64_t in_dim0, out_dim0, rb0; - uint64_t in_dim1, out_dim1, rb1; - const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); - const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); - if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q5_K/Q6_K expert tensors do not match"); - - for (uint64_t row = 0; row < out_dim0; row++) { - ds4_vec_dot_q5_q6_K_q8_K(w0->type, (int)in_dim0, &out0[row], base0 + row * rb0, xq); - ds4_vec_dot_q5_q6_K_q8_K(w1->type, (int)in_dim0, &out1[row], base1 + row * rb1, xq); - } - } else { - ds4_die("unsupported gate/up expert tensor type"); - } -} - -/* Dispatch: single-expert down projection for tracing. */ -static void matvec_expert_down( - float *out, - const ds4_model *m, - const ds4_tensor *w, - const float *x, - uint32_t expert) { - if (w->type == DS4_TENSOR_IQ2_XXS) { - uint64_t in_dim, out_dim, row_bytes; - const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); - if (in_dim % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); - - block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); - ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); - - for (uint64_t row = 0; row < out_dim; row++) { - const block_iq2_xxs *br = (const block_iq2_xxs *)(base + row * row_bytes); - ds4_vec_dot_iq2_xxs_q8_K((int)in_dim, &out[row], br, xq); - } - free(xq); - } else if (w->type == DS4_TENSOR_Q2_K) { - matvec_q2_k_expert(out, m, w, x, expert); - } else if (w->type == DS4_TENSOR_Q4_K) { - uint64_t in_dim, out_dim, row_bytes; - const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); - if (in_dim % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); - - block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); - ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); - - for (uint64_t row = 0; row < out_dim; row++) { - const block_q4_K *br = (const block_q4_K *)(base + row * row_bytes); - ds4_vec_dot_q4_K_q8_K((int)in_dim, &out[row], br, xq); - } - free(xq); - } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { - matvec_q5_q6_k_expert(out, m, w, x, expert); - } else { - ds4_die("unsupported down expert tensor type"); - } -} - -typedef struct { - float *moe; - const uint8_t *base[DS4_MAX_EXPERT]; - const block_q8_K *midq; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint32_t n_active; - uint32_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t midq_blocks; -} matvec_q8_k_batch_accum_rows_ctx; - -static void matvec_q8_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_k_batch_accum_rows_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - for (uint32_t t = 0; t < ctx->n_tok; t++) { - ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; - } - - for (uint32_t ai = 0; ai < ctx->n_active; ai++) { - const uint32_t expert = ctx->active_expert[ai]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const block_q8_K *br = (const block_q8_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; - float v = 0.0f; - - ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, xq); - ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; - } - } - } -} - -typedef struct { - float *moe; - const uint8_t *base[DS4_MAX_EXPERT]; - const int8_t *midq; - const float *midscale; - const ds4_expert_pair *pairs; - const uint32_t *pair_ids; - const uint32_t *expert_offset; - const uint32_t *active_expert; - uint32_t n_active; - uint32_t n_tok; - uint64_t in_dim; - uint64_t out_dim; - uint64_t row_bytes[DS4_MAX_EXPERT]; - uint64_t blocks; -} matvec_q8_0_batch_accum_rows_ctx; - -static void matvec_q8_0_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { - matvec_q8_0_batch_accum_rows_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - for (uint32_t t = 0; t < ctx->n_tok; t++) { - ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; - } - - for (uint32_t ai = 0; ai < ctx->n_active; ai++) { - const uint32_t expert = ctx->active_expert[ai]; - const uint32_t begin = ctx->expert_offset[expert]; - const uint32_t end = ctx->expert_offset[expert + 1]; - const uint8_t *br = ctx->base[expert] + row * ctx->row_bytes[expert]; - - for (uint32_t i = begin; i < end; i++) { - const uint32_t pair_id = ctx->pair_ids[i]; - const ds4_expert_pair pair = ctx->pairs[pair_id]; - const int8_t *xq = ctx->midq + (uint64_t)pair_id * ctx->blocks * 32u; - const float *xscale = ctx->midscale + (uint64_t)pair_id * ctx->blocks; - const float v = dot_q8_0_row(br, xq, xscale, ctx->in_dim, ctx->blocks); - ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; - } - } - } -} - -typedef struct { - float *moe; - const float *down_pair; - uint32_t n_tok; - uint64_t out_dim; -} sum_down_pairs_ctx; - -static DS4_MAYBE_UNUSED void sum_down_pairs_worker(void *vctx, uint64_t row0, uint64_t row1) { - sum_down_pairs_ctx *ctx = vctx; - for (uint64_t idx = row0; idx < row1; idx++) { - const uint32_t token = (uint32_t)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)token * ctx->out_dim; - float acc = 0.0f; - for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { - const uint64_t pair_id = (uint64_t)token * DS4_N_EXPERT_USED + slot; - acc += ctx->down_pair[pair_id * ctx->out_dim + row]; - } - ctx->moe[idx] = acc; - } -} - -/* ========================================================================= - * Hyper-Connection Transforms. - * ========================================================================= - * - * DeepSeek V4 Flash keeps four hyper-connection streams per token. Before - * attention or FFN, a learned small projection chooses how to reduce the HC - * state into the 4096-wide sublayer input. After the sublayer, the post and - * combine weights expand the result back into the four-stream HC state. - */ - -/* Decode the HC control projection. The output contains pre weights, post - * gates, and a small doubly-normalized combine matrix. */ -static void hc_split_sinkhorn_one( - float * out, - const float * mix, - const float * scale, - const float * base, - int n_hc, - int iters, - float eps) { - const float pre_scale = scale[0]; - const float post_scale = scale[1]; - const float comb_scale = scale[2]; - - for (int i = 0; i < n_hc; i++) { - const float z = mix[i] * pre_scale + base[i]; - out[i] = 1.0f / (1.0f + expf(-z)) + eps; - } - - for (int i = 0; i < n_hc; i++) { - const int off = n_hc + i; - const float z = mix[off] * post_scale + base[off]; - out[off] = 2.0f / (1.0f + expf(-z)); - } - - float c[16 * 16]; - - for (int dst = 0; dst < n_hc; dst++) { - float row_max = DS4_NEG_INF; - for (int src = 0; src < n_hc; src++) { - const int idx = src + dst * n_hc; - const int off = 2 * n_hc + idx; - const float v = mix[off] * comb_scale + base[off]; - c[idx] = v; - if (v > row_max) row_max = v; - } - - float row_sum = 0.0f; - for (int src = 0; src < n_hc; src++) { - const int idx = src + dst * n_hc; - const float v = expf(c[idx] - row_max); - c[idx] = v; - row_sum += v; - } - - const float inv = 1.0f / row_sum; - for (int src = 0; src < n_hc; src++) { - const int idx = src + dst * n_hc; - c[idx] = c[idx] * inv + eps; - } - } - - for (int src = 0; src < n_hc; src++) { - float sum = 0.0f; - for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; - - const float inv = 1.0f / (sum + eps); - for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; - } - - for (int iter = 1; iter < iters; iter++) { - for (int dst = 0; dst < n_hc; dst++) { - float sum = 0.0f; - for (int src = 0; src < n_hc; src++) sum += c[src + dst * n_hc]; - - const float inv = 1.0f / (sum + eps); - for (int src = 0; src < n_hc; src++) c[src + dst * n_hc] *= inv; - } - - for (int src = 0; src < n_hc; src++) { - float sum = 0.0f; - for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; - - const float inv = 1.0f / (sum + eps); - for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; - } - } - - for (int i = 0; i < n_hc * n_hc; i++) out[2 * n_hc + i] = c[i]; -} - -/* Reduce the four HC streams into the plain embedding vector consumed by a - * normal attention or FFN sublayer. */ -static void hc_weighted_sum_one( - float * out, - const float * x, - const float * weights, - uint32_t n_embd, - uint32_t n_hc) { - for (uint32_t d = 0; d < n_embd; d++) { - float acc = 0.0f; - for (uint32_t h = 0; h < n_hc; h++) { - acc += x[(uint64_t)h * n_embd + d] * weights[h]; - } - out[d] = acc; - } -} - -/* HC pre step for one token. It normalizes the HC state, projects the control - * vector, runs the Sinkhorn split, and emits the sublayer input plus post data. */ -static void hc_pre_from_state_one_scratch( - const ds4_model * model, - const ds4_tensor * fn, - const ds4_tensor * scale_tensor, - const ds4_tensor * base_tensor, - const float * residual_hc, - float * out, - float * post, - float * comb, - float * flat, - bool serial_fn) { - const uint32_t n_hc = DS4_N_HC; - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; - - float mix[24]; - float split[24]; - - rms_norm_no_weight(flat, residual_hc, hc_dim, DS4_RMS_EPS); - if (serial_fn) { - matvec_f16_serial(mix, model, fn, flat); - } else { - matvec_f16(mix, model, fn, flat); - } - - const float *scale = tensor_data(model, scale_tensor); - const float *base = tensor_data(model, base_tensor); - hc_split_sinkhorn_one(split, mix, scale, base, (int)n_hc, DS4_N_HC_SINKHORN_ITER, 1.0e-6f); - hc_weighted_sum_one(out, residual_hc, split, DS4_N_EMBD, n_hc); - - memcpy(post, split + n_hc, n_hc * sizeof(post[0])); - memcpy(comb, split + 2 * n_hc, n_hc * n_hc * sizeof(comb[0])); -} - -static void hc_pre_from_state_one( - const ds4_model * model, - const ds4_tensor * fn, - const ds4_tensor * scale_tensor, - const ds4_tensor * base_tensor, - const float * residual_hc, - float * out, - float * post, - float * comb) { - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; - float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); - - hc_pre_from_state_one_scratch(model, - fn, scale_tensor, base_tensor, - residual_hc, out, post, comb, - flat, false); - free(flat); -} - -static void layer_attn_pre_one( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * token_embd, - float * out, - float * residual_hc, - float * post, - float * comb) { - const uint32_t n_hc = DS4_N_HC; - - for (uint32_t h = 0; h < n_hc; h++) { - memcpy(residual_hc + (uint64_t)h * DS4_N_EMBD, token_embd, (size_t)DS4_N_EMBD * sizeof(token_embd[0])); - } - - hc_pre_from_state_one(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - residual_hc, out, post, comb); -} - -/* The input embedding starts all HC streams with the same token vector. */ -static void hc_from_plain_embedding(float *out_hc, const float *x, uint32_t n_embd, uint32_t n_hc) { - for (uint32_t h = 0; h < n_hc; h++) { - memcpy(out_hc + (uint64_t)h * n_embd, x, (size_t)n_embd * sizeof(x[0])); - } -} - -/* HC post step for one sublayer output. It injects the new block output and - * mixes the previous HC streams through the learned combine matrix. */ -static void hc_post_one( - float * out_hc, - const float * block_out, - const float * residual_hc, - const float * post, - const float * comb, - uint32_t n_embd, - uint32_t n_hc) { - for (uint32_t dst = 0; dst < n_hc; dst++) { - for (uint32_t d = 0; d < n_embd; d++) { - float acc = block_out[d] * post[dst]; - - for (uint32_t src = 0; src < n_hc; src++) { - /* The HC combine matrix is addressed as [dst_hc, src_hc]. */ - acc += comb[dst + src * n_hc] * residual_hc[(uint64_t)src * n_embd + d]; - } - - out_hc[(uint64_t)dst * n_embd + d] = acc; - } - } -} - -typedef struct { - float *out_hc; - const float *block_out; - const float *residual_hc; - const float *post; - const float *comb; - uint64_t hc_dim; - uint32_t n_embd; - uint32_t n_hc; -} hc_post_batch_ctx; - -static void hc_post_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - hc_post_batch_ctx *ctx = vctx; - for (uint64_t t = t0; t < t1; t++) { - hc_post_one(ctx->out_hc + t * ctx->hc_dim, - ctx->block_out + t * ctx->n_embd, - ctx->residual_hc + t * ctx->hc_dim, - ctx->post + t * ctx->n_hc, - ctx->comb + t * ctx->n_hc * ctx->n_hc, - ctx->n_embd, - ctx->n_hc); - } -} - -static void hc_post_batch( - float * out_hc, - const float * block_out, - const float * residual_hc, - const float * post, - const float * comb, - uint32_t n_tok, - uint32_t n_embd, - uint32_t n_hc) { - hc_post_batch_ctx ctx = { - .out_hc = out_hc, - .block_out = block_out, - .residual_hc = residual_hc, - .post = post, - .comb = comb, - .hc_dim = (uint64_t)n_hc * n_embd, - .n_embd = n_embd, - .n_hc = n_hc, - }; - ds4_parallel_for_min_rows(n_tok, hc_post_batch_worker, &ctx, 1); -} - -typedef struct { - float *out_hc; - const float *moe; - const float *shared; - const float *residual_hc; - const float *post; - const float *comb; - uint64_t hc_dim; - uint32_t n_embd; - uint32_t n_hc; -} hc_post_sum_batch_ctx; - -static void hc_post_sum_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - hc_post_sum_batch_ctx *ctx = vctx; - for (uint64_t t = t0; t < t1; t++) { - const float *moe = ctx->moe + t * ctx->n_embd; - const float *shared = ctx->shared + t * ctx->n_embd; - const float *residual = ctx->residual_hc + t * ctx->hc_dim; - const float *post = ctx->post + t * ctx->n_hc; - const float *comb = ctx->comb + t * ctx->n_hc * ctx->n_hc; - float *out = ctx->out_hc + t * ctx->hc_dim; - - for (uint32_t dst = 0; dst < ctx->n_hc; dst++) { - for (uint32_t d = 0; d < ctx->n_embd; d++) { - float acc = (moe[d] + shared[d]) * post[dst]; - for (uint32_t src = 0; src < ctx->n_hc; src++) { - acc += comb[dst + src * ctx->n_hc] * - residual[(uint64_t)src * ctx->n_embd + d]; - } - out[(uint64_t)dst * ctx->n_embd + d] = acc; - } - } - } -} - -static void hc_post_sum_batch( - float * out_hc, - const float * moe, - const float * shared, - const float * residual_hc, - const float * post, - const float * comb, - uint32_t n_tok, - uint32_t n_embd, - uint32_t n_hc) { - hc_post_sum_batch_ctx ctx = { - .out_hc = out_hc, - .moe = moe, - .shared = shared, - .residual_hc = residual_hc, - .post = post, - .comb = comb, - .hc_dim = (uint64_t)n_hc * n_embd, - .n_embd = n_embd, - .n_hc = n_hc, - }; - ds4_parallel_for_min_rows(n_tok, hc_post_sum_batch_worker, &ctx, 1); -} - -typedef struct { - const ds4_model *model; - const ds4_tensor *fn; - const ds4_tensor *scale; - const ds4_tensor *base; - const ds4_tensor *norm_w; - const float *inp_hc; - float *residual_hc; - float *cur; - float *norm; - float *post; - float *comb; - uint64_t hc_dim; - uint32_t n_hc; -} hc_pre_norm_batch_ctx; - -static void hc_pre_norm_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - hc_pre_norm_batch_ctx *ctx = vctx; - const float *norm_w = tensor_data(ctx->model, ctx->norm_w); - float *flat = xmalloc((size_t)ctx->hc_dim * sizeof(flat[0])); - - for (uint64_t t = t0; t < t1; t++) { - const float *residual = ctx->inp_hc + t * ctx->hc_dim; - if (ctx->residual_hc) { - float *dst = ctx->residual_hc + t * ctx->hc_dim; - memcpy(dst, residual, (size_t)ctx->hc_dim * sizeof(dst[0])); - residual = dst; - } - - hc_pre_from_state_one_scratch(ctx->model, - ctx->fn, - ctx->scale, - ctx->base, - residual, - ctx->cur + t * DS4_N_EMBD, - ctx->post + t * ctx->n_hc, - ctx->comb + t * ctx->n_hc * ctx->n_hc, - flat, - true); - rms_norm_weight(ctx->norm + t * DS4_N_EMBD, - ctx->cur + t * DS4_N_EMBD, - norm_w, - DS4_N_EMBD, - DS4_RMS_EPS); - } - - free(flat); -} - -/* Batched HC pre plus RMSNorm. Prefill uses this to keep the layer-major - * token batch in contiguous arrays. */ -static void hc_pre_norm_batch( - const ds4_model * model, - const ds4_tensor * fn, - const ds4_tensor * scale, - const ds4_tensor * base, - const ds4_tensor * norm_w, - const float * inp_hc, - float * residual_hc, - float * cur, - float * norm, - float * post, - float * comb, - uint32_t n_tok) { - hc_pre_norm_batch_ctx ctx = { - .model = model, - .fn = fn, - .scale = scale, - .base = base, - .norm_w = norm_w, - .inp_hc = inp_hc, - .residual_hc = residual_hc, - .cur = cur, - .norm = norm, - .post = post, - .comb = comb, - .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, - .n_hc = DS4_N_HC, - }; - ds4_parallel_for_min_rows(n_tok, hc_pre_norm_batch_worker, &ctx, 1); -} - -static void layer_attn_norm_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - const float *attn_norm = tensor_data(model, layer->attn_norm); - rms_norm_weight(out, x, attn_norm, DS4_N_EMBD, DS4_RMS_EPS); -} - -/* ========================================================================= - * Attention Projections, RoPE, and Attention Output. - * ========================================================================= - * - * This block performs the attention half of a transformer layer: HC pre, - * attention RMSNorm, Q and KV projections, layer-specific RoPE, sink-aware - * attention over raw and compressed KV rows, and the grouped LoRA output - * projection back to embedding width. - */ - -/* Q projection is low-rank: Q8_0 into the model-specific LoRA-Q rank, - * RMSNorm, then Q8_0 back to all attention heads. */ -static void layer_q_projection_normed_one( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * norm, - float * q) { - const uint32_t q_rank = DS4_N_LORA_Q; - float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); - float *qr_norm = xmalloc((size_t)q_rank * sizeof(qr_norm[0])); - - const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); - - matvec_q8_0(qr, model, layer->attn_q_a, norm); - rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); - matvec_q8_0(q, model, layer->attn_q_b, qr_norm); - head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); - - free(qr_norm); - free(qr); -} - -static void layer_q_projection_with_lora_one( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * norm, - float * q, - float * qr_norm) { - const uint32_t q_rank = DS4_N_LORA_Q; - float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); - const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); - - matvec_q8_0(qr, model, layer->attn_q_a, norm); - rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); - matvec_q8_0(q, model, layer->attn_q_b, qr_norm); - head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); - - free(qr); -} - -/* KV projection has one KV head of width 512, followed by a learned RMSNorm. */ -static void layer_kv_projection_normed_one( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * normed, - float * kv) { - float *raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(raw[0])); - - const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); - - matvec_q8_0(raw, model, layer->attn_kv, normed); - rms_norm_weight(kv, raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); - - free(raw); -} - -static void layer_q_projection_with_lora_one_decode_scratch( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * norm, - float * q, - float * qr_norm, - ds4_cpu_decode_scratch * scratch) { - const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); - - matvec_q8_0_decode_scratch(scratch->qr, model, layer->attn_q_a, norm, scratch); - rms_norm_weight(qr_norm, scratch->qr, q_a_norm, DS4_N_LORA_Q, DS4_RMS_EPS); - matvec_q8_0_decode_scratch(q, model, layer->attn_q_b, qr_norm, scratch); - head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); -} - -static void layer_kv_projection_normed_one_decode_scratch( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * normed, - float * kv, - ds4_cpu_decode_scratch * scratch) { - const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); - - matvec_q8_0_decode_scratch(scratch->kv_raw, model, layer->attn_kv, normed, scratch); - rms_norm_weight(kv, scratch->kv_raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); -} - -static float rope_yarn_ramp(float low, float high, int i0) { - const float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low); - return 1.0f - fminf(1.0f, fmaxf(0.0f, y)); -} - -static float rope_yarn_corr_dim(int n_dims, uint64_t n_ctx_orig, float n_rot, float base) { - return (float)n_dims * logf((float)n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) / (2.0f * logf(base)); -} - -static void rope_yarn_corr_dims(int n_dims, uint64_t n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]) { - const float start = floorf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); - const float end = ceilf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); - dims[0] = fmaxf(0.0f, start); - dims[1] = fminf((float)(n_dims - 1), end); -} - -/* Apply DS4 RoPE only to the tail of each head. Compressed layers use the - * long-context frequency base and scale; inverse mode rotates attention output - * back before the grouped output projection. */ -static void rope_tail_ext_inplace( - float * x, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos, - uint64_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool inverse) { - const uint32_t n_nope = head_dim - n_rot; - const float theta_scale = powf(freq_base, -2.0f / (float)n_rot); - const float sin_sign = inverse ? -1.0f : 1.0f; - float corr_dims[2] = { 0.0f, 0.0f }; - if (ext_factor != 0.0f) { - rope_yarn_corr_dims((int)n_rot, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - } - - for (uint32_t h = 0; h < n_head; h++) { - float *tail = x + (uint64_t)h * head_dim + n_nope; - float theta_extrap = (float)pos; - - for (uint32_t i = 0; i < n_rot; i += 2) { - const float theta_interp = freq_scale * theta_extrap; - float theta = theta_interp; - float mscale = attn_factor; - - if (ext_factor != 0.0f) { - const float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], (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); - } - - const float c = cosf(theta) * mscale; - const float s = sin_sign * sinf(theta) * mscale; - const float x0 = tail[i + 0]; - const float x1 = tail[i + 1]; - - tail[i + 0] = x0 * c - x1 * s; - tail[i + 1] = x0 * s + x1 * c; - - theta_extrap *= theta_scale; - } - } -} - -/* Dense layers and compressed layers use different RoPE bases. */ -static float layer_rope_freq_base(uint32_t il) { - return ds4_layer_compress_ratio(il) != 0 && DS4_COMPRESS_ROPE_FREQ_BASE > 0.0f - ? DS4_COMPRESS_ROPE_FREQ_BASE - : DS4_ROPE_FREQ_BASE; -} - -static float layer_rope_freq_scale(uint32_t il) { - if (ds4_layer_compress_ratio(il) == 0 || DS4_ROPE_SCALE_FACTOR <= 0.0f) { - return 1.0f; - } - return 1.0f / DS4_ROPE_SCALE_FACTOR; -} - -static void rope_tail_layer_inplace( - float * x, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos, - uint32_t il, - bool inverse) { - const bool compressed = ds4_layer_compress_ratio(il) != 0; - const float freq_base = layer_rope_freq_base(il); - const float freq_scale = layer_rope_freq_scale(il); - const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; - float attn_factor = 1.0f; - if (ext_factor != 0.0f && freq_scale > 0.0f) { - /* - * This YaRN helper applies magnitude scaling internally. DeepSeek V4 - * reference RoPE uses interpolation without that magnitude change, so - * pass the inverse factor here and let the helper cancel itself out. - */ - attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - - rope_tail_ext_inplace(x, n_head, head_dim, n_rot, pos, - compressed ? DS4_ROPE_ORIG_CTX : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - inverse); -} - -typedef struct { - float *x; - uint64_t stride; - uint32_t n_head; - uint32_t head_dim; - uint32_t n_rot; - uint32_t pos0; - uint32_t il; - bool inverse; -} rope_tail_batch_ctx; - -static void rope_tail_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - rope_tail_batch_ctx *ctx = vctx; - for (uint64_t tt = t0; tt < t1; tt++) { - rope_tail_layer_inplace(ctx->x + tt * ctx->stride, - ctx->n_head, - ctx->head_dim, - ctx->n_rot, - ctx->pos0 + (uint32_t)tt, - ctx->il, - ctx->inverse); - } -} - -static void rope_tail_layer_batch_inplace( - float *x, - uint64_t stride, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos0, - uint32_t il, - bool inverse, - uint32_t n_tok) { - rope_tail_batch_ctx ctx = { - .x = x, - .stride = stride, - .n_head = n_head, - .head_dim = head_dim, - .n_rot = n_rot, - .pos0 = pos0, - .il = il, - .inverse = inverse, - }; - ds4_parallel_for_min_rows(n_tok, rope_tail_batch_worker, &ctx, 1); -} - -static inline float dot_f32(const float *a, const float *b, uint32_t n) { -#if defined(__ARM_NEON) - uint32_t i = 0; - float32x4_t acc0 = vdupq_n_f32(0.0f); - float32x4_t acc1 = vdupq_n_f32(0.0f); - for (; i + 8 <= n; i += 8) { - acc0 = vfmaq_f32(acc0, vld1q_f32(a + i), vld1q_f32(b + i)); - acc1 = vfmaq_f32(acc1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); - } - float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); - for (; i < n; i++) acc += a[i] * b[i]; - return acc; -#else - float acc = 0.0f; - for (uint32_t i = 0; i < n; i++) acc += a[i] * b[i]; - return acc; -#endif -} - -static inline void axpy_f32(float *y, const float *x, float a, uint32_t n) { -#if defined(__ARM_NEON) - uint32_t i = 0; - const float32x4_t av = vdupq_n_f32(a); - for (; i + 8 <= n; i += 8) { - vst1q_f32(y + i, vfmaq_f32(vld1q_f32(y + i), av, vld1q_f32(x + i))); - vst1q_f32(y + i + 4, vfmaq_f32(vld1q_f32(y + i + 4), av, vld1q_f32(x + i + 4))); - } - for (; i < n; i++) y[i] += a * x[i]; -#else - for (uint32_t i = 0; i < n; i++) y[i] += a * x[i]; -#endif -} - -static inline void scale_f32(float *x, float a, uint32_t n) { -#if defined(__ARM_NEON) - uint32_t i = 0; - const float32x4_t av = vdupq_n_f32(a); - for (; i + 8 <= n; i += 8) { - vst1q_f32(x + i, vmulq_f32(vld1q_f32(x + i), av)); - vst1q_f32(x + i + 4, vmulq_f32(vld1q_f32(x + i + 4), av)); - } - for (; i < n; i++) x[i] *= a; -#else - for (uint32_t i = 0; i < n; i++) x[i] *= a; -#endif -} - -static float sigmoid_stable(float x) { - if (x >= 0.0f) { - const float e = expf(-x); - return 1.0f / (1.0f + e); - } else { - const float e = expf(x); - return e / (1.0f + e); - } -} - -/* Sink-aware attention over a set of KV rows. The learned sink logit is part - * of the softmax denominator but contributes no value vector. */ -static void layer_attention_rows_one( - float * out_heads, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * q, - const float * kv_rows, - uint32_t n_kv) { - const float *sinks = tensor_data(model, layer->attn_sinks); - const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); - float score_stack[512]; - float *score = n_kv <= 512 ? score_stack : xmalloc((size_t)n_kv * sizeof(score[0])); - - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; - - float max_score = sinks[h]; - for (uint32_t r = 0; r < n_kv; r++) { - const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; - score[r] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[r] > max_score) max_score = score[r]; - } - - float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; - memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); - - float denom = expf(sinks[h] - max_score); - for (uint32_t r = 0; r < n_kv; r++) { - const float weight = expf(score[r] - max_score); - const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - - const float inv = 1.0f / denom; - scale_f32(oh, inv, DS4_N_HEAD_DIM); - } - - if (score != score_stack) free(score); -} - -static void layer_attention_one( - float * out_heads, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * q, - const float * kv) { - layer_attention_rows_one(out_heads, model, layer, q, kv, 1); -} - -/* Attention output projection is grouped: each group first maps its heads to - * a 1024-rank low vector, then all groups are projected back to 4096. */ -static void layer_grouped_out_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * heads) { - const uint32_t n_groups = 8; - const uint32_t group_heads = DS4_N_HEAD / n_groups; - const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; - const uint32_t rank = 1024; - - float *low = xcalloc((size_t)n_groups * rank, sizeof(low[0])); - - matvec_q8_0_grouped_rows(low, model, layer->attn_output_a, heads, n_groups, group_dim, rank); - - matvec_q8_0(out, model, layer->attn_output_b, low); - free(low); -} - -static void layer_grouped_out_one_decode_scratch( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * heads, - ds4_cpu_decode_scratch * scratch) { - const uint32_t n_groups = 8; - const uint32_t group_heads = DS4_N_HEAD / n_groups; - const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; - const uint32_t rank = 1024; - - memset(scratch->attn_low, 0, (size_t)n_groups * rank * sizeof(scratch->attn_low[0])); - matvec_q8_0_grouped_rows_decode_scratch(scratch->attn_low, model, layer->attn_output_a, - heads, n_groups, group_dim, rank, scratch); - matvec_q8_0_decode_scratch(out, model, layer->attn_output_b, scratch->attn_low, scratch); -} - -static void layer_grouped_out_batch( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * heads, - uint32_t n_tok) { - const uint32_t n_groups = 8; - const uint32_t group_heads = DS4_N_HEAD / n_groups; - const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; - const uint32_t rank = 1024; - - float *low = xcalloc((size_t)n_tok * n_groups * rank, sizeof(low[0])); - - matmul_q8_0_grouped_batch(low, model, layer->attn_output_a, heads, - n_tok, n_groups, group_dim, rank); - matmul_q8_0_batch(out, model, layer->attn_output_b, low, n_tok); - - free(low); -} - -/* ========================================================================= - * Mixture-of-Experts FFN. - * ========================================================================= - * - * This is the FFN half of each layer. It includes the shared expert, routed - * expert selection, IQ2_XXS gate/up projections, SwiGLU, Q2_K down projection, - * and the HC post step that returns the result to four-stream state. - */ - -static float silu(float x) { - return x * sigmoid_stable(x); -} - -static float softplus_stable(float x) { - if (x > 20.0f) return x; - if (x < -20.0f) return expf(x); - return log1pf(expf(x)); -} - -static void swiglu(float *out, const float *gate, const float *up, uint64_t n, float clamp) { - for (uint64_t i = 0; i < n; i++) { - float g = gate[i]; - float u = up[i]; - if (clamp > 1.0e-6f) { - if (g > clamp) g = clamp; - if (u > clamp) u = clamp; - if (u < -clamp) u = -clamp; - } - out[i] = silu(g) * u; - } -} - -/* The shared expert is a normal Q8_0 SwiGLU MLP that runs for every token. */ -static void layer_shared_ffn_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - float *gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])); - float *up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])); - float *mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])); - const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)blocks * 32); - float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); - - if (layer->ffn_up_shexp->type != 8 || - layer->ffn_gate_shexp->type != 8 || - layer->ffn_up_shexp->dim[0] != in_dim) { - ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); - } - - quantize_q8_0_activation(x, xq, xscale, in_dim); - matvec_q8_0_pair_prequant(gate, up, model, - layer->ffn_gate_shexp, - layer->ffn_up_shexp, - xq, xscale); - swiglu(mid, gate, up, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP); - matvec_q8_0(out, model, layer->ffn_down_shexp, mid); - - free(xscale); - free(xq); - free(mid); - free(up); - free(gate); -} - -static void layer_shared_ffn_one_decode_scratch( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - ds4_cpu_decode_scratch * scratch) { - const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; - if (layer->ffn_up_shexp->type != 8 || - layer->ffn_gate_shexp->type != 8 || - layer->ffn_up_shexp->dim[0] != in_dim) { - ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); - } - - matvec_q8_0_pair_decode_scratch(scratch->shared_gate, - scratch->shared_up, - model, - layer->ffn_gate_shexp, - layer->ffn_up_shexp, - x, - scratch); - swiglu(scratch->shared_mid, scratch->shared_gate, scratch->shared_up, DS4_N_FF_EXP, - DS4_SWIGLU_CLAMP_EXP); - matvec_q8_0_decode_scratch(out, model, layer->ffn_down_shexp, scratch->shared_mid, scratch); -} - -typedef struct { - float *mid; - const float *gate; - const float *up; - uint64_t n; - float clamp; -} swiglu_batch_ctx; - -static void swiglu_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { - swiglu_batch_ctx *ctx = vctx; - for (uint64_t t = t0; t < t1; t++) { - swiglu(ctx->mid + t * ctx->n, - ctx->gate + t * ctx->n, - ctx->up + t * ctx->n, - ctx->n, - ctx->clamp); - } -} - -static void layer_shared_ffn_batch( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t n_tok) { - const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; - const uint64_t hidden = layer->ffn_gate_shexp->dim[1]; - - if (layer->ffn_up_shexp->type != 8 || - layer->ffn_gate_shexp->type != 8 || - layer->ffn_down_shexp->type != 8 || - layer->ffn_up_shexp->dim[0] != in_dim || - layer->ffn_up_shexp->dim[1] != hidden || - layer->ffn_down_shexp->dim[0] != hidden) { - ds4_die("shared expert tensors do not share the expected Q8_0 layout"); - } - - float *gate = xmalloc((size_t)n_tok * hidden * sizeof(gate[0])); - float *up = xmalloc((size_t)n_tok * hidden * sizeof(up[0])); - float *mid = xmalloc((size_t)n_tok * hidden * sizeof(mid[0])); - - matmul_q8_0_pair_batch(gate, up, model, - layer->ffn_gate_shexp, - layer->ffn_up_shexp, - x, - n_tok); - - swiglu_batch_ctx swiglu_ctx = { - .mid = mid, - .gate = gate, - .up = up, - .n = hidden, - .clamp = DS4_SWIGLU_CLAMP_EXP, - }; - ds4_parallel_for(n_tok, swiglu_batch_worker, &swiglu_ctx); - - matmul_q8_0_batch(out, model, layer->ffn_down_shexp, mid, n_tok); - - free(mid); - free(up); - free(gate); -} - -/* Early DS4 layers use token-id hash routing instead of top-k routing. */ -static void layer_hash_selected_experts( - int selected[DS4_MAX_EXPERT_USED], - const ds4_model *model, - const ds4_layer_weights *layer, - int token) { - ds4_tensor *t = layer->ffn_gate_tid2eid; - if (!t) ds4_die("hash routing table is missing for this layer"); - if (t->type != 26 || t->ndim != 2 || t->dim[0] != DS4_N_EXPERT_USED) { - ds4_die("ffn_gate_tid2eid.weight has an unexpected layout"); - } - if (token < 0 || (uint64_t)token >= t->dim[1]) { - ds4_die("token id is outside the hash routing table"); - } - - const int32_t *table = tensor_data(model, t); - const int32_t *row = table + (uint64_t)token * DS4_N_EXPERT_USED; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) selected[i] = row[i]; -} - -/* Router scores use sqrt(softplus(logit)); normalization happens only after - * the six selected experts are known. */ -static void layer_router_probs_one( - float probs[DS4_MAX_EXPERT], - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - float logits[DS4_MAX_EXPERT]; - - matvec_any(logits, model, layer->ffn_gate_inp, x); - for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { - probs[i] = sqrtf(softplus_stable(logits[i])); - } -} - -static void layer_hash_router_weights_from_probs( - float weights_out[DS4_MAX_EXPERT_USED], - const float probs[DS4_MAX_EXPERT], - const int selected[DS4_MAX_EXPERT_USED]) { - float sum = 0.0f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) ds4_die("hash-selected expert is outside router range"); - weights_out[i] = probs[selected[i]]; - sum += weights_out[i]; - } - - if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - weights_out[i] = weights_out[i] / sum * DS4_EXPERT_WEIGHT_SCALE; - } -} - -static void layer_hash_router_weights_one( - float weights_out[DS4_MAX_EXPERT_USED], - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - const int selected[DS4_MAX_EXPERT_USED]) { - float probs[DS4_MAX_EXPERT]; - - layer_router_probs_one(probs, model, layer, x); - layer_hash_router_weights_from_probs(weights_out, probs, selected); -} - -static void topk_desc(const float *score, int n, int k, int *idx) { - for (int i = 0; i < k; i++) idx[i] = -1; - - for (int i = 0; i < n; i++) { - for (int j = 0; j < k; j++) { - if (idx[j] < 0 || score[i] > score[idx[j]]) { - for (int m = k - 1; m > j; m--) idx[m] = idx[m - 1]; - idx[j] = i; - break; - } - } - } -} - -/* Later layers choose the six experts by biased top-k, but weight them using - * the unbiased router probabilities. */ -static void layer_topk_selected_experts_from_probs( - int selected[DS4_MAX_EXPERT_USED], - float expert_weight[DS4_MAX_EXPERT_USED], - const ds4_model *model, - const ds4_layer_weights *layer, - const float probs[DS4_MAX_EXPERT]); - -static void layer_topk_selected_experts( - int selected[DS4_MAX_EXPERT_USED], - float expert_weight[DS4_MAX_EXPERT_USED], - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x) { - float probs[DS4_MAX_EXPERT] = {0}; - - layer_router_probs_one(probs, model, layer, x); - layer_topk_selected_experts_from_probs(selected, expert_weight, model, layer, probs); -} - -static void layer_topk_selected_experts_from_probs( - int selected[DS4_MAX_EXPERT_USED], - float expert_weight[DS4_MAX_EXPERT_USED], - const ds4_model *model, - const ds4_layer_weights *layer, - const float probs[DS4_MAX_EXPERT]) { - float selection[DS4_MAX_EXPERT]; - - memcpy(selection, probs, sizeof(selection)); - - if (layer->ffn_exp_probs_b) { - const float *bias = tensor_data(model, layer->ffn_exp_probs_b); - for (uint32_t i = 0; i < DS4_N_EXPERT; i++) selection[i] += bias[i]; - } - - topk_desc(selection, (int)DS4_N_EXPERT, (int)DS4_N_EXPERT_USED, selected); - - float sum = 0.0f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - expert_weight[i] = probs[selected[i]]; - sum += expert_weight[i]; - } - if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - expert_weight[i] = expert_weight[i] / sum * DS4_EXPERT_WEIGHT_SCALE; - } -} - -static void print_vec_stats(const char *name, const float *x, uint64_t n); - -/* Single-token routed MoE. It selects six experts, runs IQ2_XXS gate/up, - * applies SwiGLU and router weights, then accumulates Q2_K down projections. */ -static void layer_routed_moe_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il, - int token, - float clamp, - bool trace) { - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - float *gate = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])) : NULL; - float *up = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])) : NULL; - float *mid = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])) : NULL; - float *mid_all = trace ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid_all[0])); - float *down = trace ? xmalloc((size_t)DS4_N_EMBD * sizeof(down[0])) : NULL; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const bool routed_q8_k = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; - if (routed_q8_0) { - if (trace) ds4_die("Q8_0 routed trace mode is not supported"); - if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); - if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { - ds4_die("Q8_0 expert input has an unexpected layout"); - } - } else { - if (routed_q8_k && trace) ds4_die("Q8_K routed trace mode is not supported"); - if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); - if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { - ds4_die("routed expert down input has an unexpected layout"); - } - } - block_q8_K *xq = routed_q8_0 ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0])); - block_q8_K *midq = (trace || routed_q8_0) ? NULL : - xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0])); - - memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); - if (!routed_q8_0) { - ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); - } - - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(selected, model, layer, token); - layer_hash_router_weights_one(expert_weight, model, layer, x, selected); - } else { - layer_topk_selected_experts(selected, expert_weight, model, layer, x); - } - - if (routed_q8_0) { - const uint64_t x_blocks = expert_in_dim / 32u; - int8_t *xq8 = xmalloc((size_t)x_blocks * 32u); - float *xscale8 = xmalloc((size_t)x_blocks * sizeof(float)); - quantize_q8_0_activation(x, xq8, xscale8, expert_in_dim); - matvec_q8_0_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq8, xscale8, selected, - expert_weight, - DS4_N_EXPERT_USED, clamp); - - const uint64_t mid_blocks = down_in_dim / 32u; - int8_t *midq8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * 32u); - float *midscale8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * sizeof(float)); - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, - midq8 + (uint64_t)i * mid_blocks * 32u, - midscale8 + (uint64_t)i * mid_blocks, - down_in_dim); - } - matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, - midq8, midscale8, selected, - DS4_N_EXPERT_USED); - free(midscale8); - free(midq8); - free(xscale8); - free(xq8); - } else if (routed_q8_k) { - matvec_q8_k_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, selected, expert_weight, - DS4_N_EXPERT_USED, clamp); - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, - midq + (uint64_t)i * (down_in_dim / QK_K), - (int64_t)down_in_dim); - } - matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, - midq, selected, DS4_N_EXPERT_USED); - } else if (!trace) { - matvec_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, - selected, - expert_weight, - DS4_N_EXPERT_USED, - clamp); - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, - midq + (uint64_t)i * (down_in_dim / QK_K), - (int64_t)down_in_dim); - } - matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); - } else { - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - const uint32_t expert = (uint32_t)selected[i]; - - matvec_expert_pair_prequant(gate, up, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, - expert); - char name[64]; - snprintf(name, sizeof(name), "blk.%u expert %u gate", il, expert); - print_vec_stats(name, gate, DS4_N_FF_EXP); - snprintf(name, sizeof(name), "blk.%u expert %u up", il, expert); - print_vec_stats(name, up, DS4_N_FF_EXP); - - /* - * DeepSeek V4 clamps routed expert gate/up values before SwiGLU and - * applies the router weight before the down projection. - */ - const float limit = clamp; - for (uint32_t j = 0; j < DS4_N_FF_EXP; j++) { - if (limit > 1.0e-6f) { - if (gate[j] > limit) gate[j] = limit; - if (up[j] > limit) up[j] = limit; - if (up[j] < -limit) up[j] = -limit; - } - mid[j] = silu(gate[j]) * up[j] * expert_weight[i]; - } - - snprintf(name, sizeof(name), "blk.%u expert %u mid", il, expert); - print_vec_stats(name, mid, DS4_N_FF_EXP); - - matvec_expert_down(down, model, layer->ffn_down_exps, mid, expert); - snprintf(name, sizeof(name), "blk.%u expert %u down", il, expert); - print_vec_stats(name, down, DS4_N_EMBD); - for (uint32_t j = 0; j < DS4_N_EMBD; j++) out[j] += down[j]; - } - } - - free(midq); - free(xq); - free(down); - free(mid_all); - free(mid); - free(up); - free(gate); -} - -/* Decode version of routed MoE: same math as layer_routed_moe_one(), but all - * large temporaries come from the persistent scratch arena. */ -static void layer_routed_moe_one_prealloc( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il, - int token, - float clamp, - float * mid_all, - block_q8_K * xq, - block_q8_K * midq, - int8_t * q8_xq, - float * q8_xscale, - int8_t * q8_midq, - float * q8_midscale) { - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const bool routed_q8_k = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; - - if (routed_q8_0) { - if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); - if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { - ds4_die("Q8_0 expert input has an unexpected layout"); - } - } else { - if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); - if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { - ds4_die("routed expert down input has an unexpected layout"); - } - } - - memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); - - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(selected, model, layer, token); - layer_hash_router_weights_one(expert_weight, model, layer, x, selected); - } else { - layer_topk_selected_experts(selected, expert_weight, model, layer, x); - } - - if (routed_q8_0) { - if (!q8_xq || !q8_xscale || !q8_midq || !q8_midscale) { - ds4_die("missing Q8_0 routed decode scratch"); - } - quantize_q8_0_activation(x, q8_xq, q8_xscale, expert_in_dim); - matvec_q8_0_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - q8_xq, q8_xscale, selected, - expert_weight, - DS4_N_EXPERT_USED, clamp); - const uint64_t mid_blocks = down_in_dim / 32u; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, - q8_midq + (uint64_t)i * mid_blocks * 32u, - q8_midscale + (uint64_t)i * mid_blocks, - down_in_dim); - } - matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, - q8_midq, q8_midscale, selected, - DS4_N_EXPERT_USED); - (void)il; - return; - } - - if (!mid_all || !xq || !midq) ds4_die("missing routed decode scratch"); - ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); - - if (routed_q8_k) { - matvec_q8_k_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, selected, expert_weight, - DS4_N_EXPERT_USED, clamp); - } else { - matvec_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, selected, expert_weight, - DS4_N_EXPERT_USED, clamp); - } - - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, - midq + (uint64_t)i * (down_in_dim / QK_K), - (int64_t)down_in_dim); - } - if (routed_q8_k) { - matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, - midq, selected, DS4_N_EXPERT_USED); - } else { - matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, - midq, selected, DS4_N_EXPERT_USED); - } - - (void)il; -} - -/* Prefill MoE groups token/expert pairs by expert so each active expert's - * rows are scanned once for the whole token batch. */ -static void layer_routed_moe_batch( - float * moe, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * norm, - const int * token_ids, - uint32_t n_tok, - uint32_t il, - float clamp) { - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t expert_out_dim = layer->ffn_gate_exps->dim[1]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const uint64_t down_out_dim = layer->ffn_down_exps->dim[1]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const bool routed_q8_k = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; - if (routed_q8_0) { - if ((expert_in_dim % 32u) != 0 || (down_in_dim % 32u) != 0) { - ds4_die("Q8_0 routed expert input is not QK8_0 aligned"); - } - } else { - if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); - if (down_in_dim % QK_K != 0) ds4_die("routed expert down input is not QK_K aligned"); - } - if (expert_out_dim != down_in_dim || down_out_dim != DS4_N_EMBD) { - ds4_die("routed expert tensor layout is unexpected"); - } - - const uint32_t total_pairs = n_tok * DS4_N_EXPERT_USED; - uint32_t counts[DS4_MAX_EXPERT + 1] = {0}; - uint32_t cursor[DS4_MAX_EXPERT] = {0}; - uint32_t active_expert[DS4_MAX_EXPERT]; - uint32_t n_active = 0; - - int *selected = xmalloc((size_t)total_pairs * sizeof(selected[0])); - float *pair_weight = xmalloc((size_t)total_pairs * sizeof(pair_weight[0])); - ds4_expert_pair *pairs = xmalloc((size_t)total_pairs * sizeof(pairs[0])); - - for (uint32_t t = 0; t < n_tok; t++) { - int sel[DS4_MAX_EXPERT_USED]; - float weights[DS4_MAX_EXPERT_USED]; - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(sel, model, layer, token_ids[t]); - layer_hash_router_weights_one(weights, model, layer, norm + (uint64_t)t * expert_in_dim, sel); - } else { - layer_topk_selected_experts(sel, weights, model, layer, norm + (uint64_t)t * expert_in_dim); - } - - for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { - const uint32_t pair_id = t * DS4_N_EXPERT_USED + slot; - selected[pair_id] = sel[slot]; - pair_weight[pair_id] = weights[slot]; - pairs[pair_id] = (ds4_expert_pair){ .token = t, .slot = slot }; - if (sel[slot] < 0 || (uint32_t)sel[slot] >= DS4_N_EXPERT) ds4_die("selected expert is outside range"); - counts[(uint32_t)sel[slot] + 1]++; - } - } - - for (uint32_t e = 0; e < DS4_N_EXPERT; e++) { - counts[e + 1] += counts[e]; - cursor[e] = counts[e]; - if (counts[e + 1] != counts[e]) active_expert[n_active++] = e; - } - - uint32_t *pair_ids = xmalloc((size_t)total_pairs * sizeof(pair_ids[0])); - for (uint32_t p = 0; p < total_pairs; p++) { - const uint32_t e = (uint32_t)selected[p]; - pair_ids[cursor[e]++] = p; - } - - if (routed_q8_0) { - const uint64_t x_blocks = expert_in_dim / 32u; - int8_t *xq8 = xmalloc((size_t)n_tok * x_blocks * 32u); - float *xscale8 = xmalloc((size_t)n_tok * x_blocks * sizeof(float)); - for (uint32_t t = 0; t < n_tok; t++) { - quantize_q8_0_activation(norm + (uint64_t)t * expert_in_dim, - xq8 + (uint64_t)t * x_blocks * 32u, - xscale8 + (uint64_t)t * x_blocks, - expert_in_dim); - } - - float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); - matvec_q8_0_batch_mid_ctx mid_ctx = { - .mid = mid, - .xq = xq8, - .xscale = xscale8, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .pair_weight = pair_weight, - .clamp = clamp, - .in_dim = expert_in_dim, - .out_dim = expert_out_dim, - .blocks = x_blocks, - }; - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, - &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); - mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, - &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); - if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || - gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { - ds4_die("Q8_0 batch expert tensor layout mismatch"); - } - } - ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_0_batch_mid_worker, &mid_ctx); - - const uint64_t mid_blocks = down_in_dim / 32u; - int8_t *midq8 = xmalloc((size_t)total_pairs * mid_blocks * 32u); - float *midscale8 = xmalloc((size_t)total_pairs * mid_blocks * sizeof(float)); - for (uint32_t p = 0; p < total_pairs; p++) { - quantize_q8_0_activation(mid + (uint64_t)p * down_in_dim, - midq8 + (uint64_t)p * mid_blocks * 32u, - midscale8 + (uint64_t)p * mid_blocks, - down_in_dim); - } - free(mid); - - matvec_q8_0_batch_accum_rows_ctx down_ctx = { - .moe = moe, - .midq = midq8, - .midscale = midscale8, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .n_active = n_active, - .n_tok = n_tok, - .in_dim = down_in_dim, - .out_dim = down_out_dim, - .blocks = mid_blocks, - }; - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t in_dim, out_dim; - down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, - &in_dim, &out_dim, &down_ctx.row_bytes[e]); - if (in_dim != down_in_dim || out_dim != down_out_dim) { - ds4_die("Q8_0 batch down expert tensor layout mismatch"); - } - } - ds4_parallel_for(down_out_dim, matvec_q8_0_batch_accum_rows_worker, &down_ctx); - - free(midscale8); - free(midq8); - free(xscale8); - free(xq8); - free(pair_ids); - free(pairs); - free(pair_weight); - free(selected); - (void)il; - return; - } - - if (routed_q8_k) { - const uint64_t xq_blocks = expert_in_dim / QK_K; - block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); - for (uint32_t t = 0; t < n_tok; t++) { - ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, - xq + (uint64_t)t * xq_blocks, - (int64_t)expert_in_dim); - } - - float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); - - matvec_q8_k_batch_mid_ctx mid_ctx = { - .mid = mid, - .xq = xq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .pair_weight = pair_weight, - .clamp = clamp, - .in_dim = expert_in_dim, - .out_dim = expert_out_dim, - .xq_blocks = xq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, - &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); - mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, - &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); - if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || - gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { - ds4_die("Q8_K batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_k_batch_mid_worker, &mid_ctx); - - const uint64_t midq_blocks = down_in_dim / QK_K; - block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); - quantize_mid_pairs_ctx quant_ctx = { - .mid = mid, - .midq = midq, - .down_in_dim = down_in_dim, - .down_blocks = midq_blocks, - }; - ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); - free(mid); - - matvec_q8_k_batch_accum_rows_ctx down_ctx = { - .moe = moe, - .midq = midq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .n_active = n_active, - .n_tok = n_tok, - .in_dim = down_in_dim, - .out_dim = down_out_dim, - .midq_blocks = midq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t in_dim, out_dim; - down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, - &in_dim, &out_dim, &down_ctx.row_bytes[e]); - if (in_dim != down_in_dim || out_dim != down_out_dim) { - ds4_die("Q8_K batch down expert tensor layout mismatch"); - } - } - - ds4_parallel_for(down_out_dim, matvec_q8_k_batch_accum_rows_worker, &down_ctx); - - free(midq); - free(pair_ids); - free(xq); - free(pairs); - free(pair_weight); - free(selected); - - (void)il; - return; - } - - const uint64_t xq_blocks = expert_in_dim / QK_K; - block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); - for (uint32_t t = 0; t < n_tok; t++) { - ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, - xq + (uint64_t)t * xq_blocks, - (int64_t)expert_in_dim); - } - - float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); - - const uint32_t gate_type = layer->ffn_gate_exps->type; - - /* Build mid vectors: dispatch based on gate/up tensor type. */ - if (gate_type == DS4_TENSOR_IQ2_XXS) { - matvec_iq2_xxs_batch_mid_ctx mid_ctx = { - .mid = mid, - .xq = xq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .pair_weight = pair_weight, - .clamp = clamp, - .in_dim = expert_in_dim, - .out_dim = expert_out_dim, - .xq_blocks = xq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, - &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); - mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, - &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); - if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || - gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_iq2_xxs_batch_mid_worker, &mid_ctx); - } else if (gate_type == DS4_TENSOR_Q2_K) { - matvec_q2_k_batch_mid_ctx mid_ctx = { - .mid = mid, - .xq = xq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .pair_weight = pair_weight, - .clamp = clamp, - .in_dim = expert_in_dim, - .out_dim = expert_out_dim, - .xq_blocks = xq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, - &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); - mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, - &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); - if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || - gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q2_k_batch_mid_worker, &mid_ctx); - } else if (gate_type == DS4_TENSOR_Q4_K) { - matvec_q4_k_batch_mid_ctx mid_ctx = { - .mid = mid, - .xq = xq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .pair_weight = pair_weight, - .clamp = clamp, - .in_dim = expert_in_dim, - .out_dim = expert_out_dim, - .xq_blocks = xq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t gate_in_dim, gate_out_dim; - uint64_t up_in_dim, up_out_dim; - mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, - &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); - mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, - &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); - if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || - gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q4_k_batch_mid_worker, &mid_ctx); - } else { - ds4_die("unsupported gate/up expert tensor type for batch"); - } - - const uint64_t midq_blocks = down_in_dim / QK_K; - block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); - quantize_mid_pairs_ctx quant_ctx = { - .mid = mid, - .midq = midq, - .down_in_dim = down_in_dim, - .down_blocks = midq_blocks, - }; - ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); - free(mid); - - /* Down projection: dispatch based on down tensor type. */ - const uint32_t down_type = layer->ffn_down_exps->type; - - if (down_type == DS4_TENSOR_IQ2_XXS) { - matvec_iq2_xxs_batch_accum_rows_ctx down_ctx = { - .moe = moe, - .midq = midq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .n_active = n_active, - .n_tok = n_tok, - .in_dim = down_in_dim, - .out_dim = down_out_dim, - .midq_blocks = midq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t in_dim, out_dim; - down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, - &in_dim, &out_dim, &down_ctx.row_bytes[e]); - if (in_dim != down_in_dim || out_dim != down_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for(down_out_dim, matvec_iq2_xxs_batch_accum_rows_worker, &down_ctx); - } else if (down_type == DS4_TENSOR_Q2_K) { - matvec_q2_k_batch_accum_rows_ctx down_ctx = { - .moe = moe, - .midq = midq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .n_active = n_active, - .n_tok = n_tok, - .in_dim = down_in_dim, - .out_dim = down_out_dim, - .midq_blocks = midq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t in_dim, out_dim; - down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, - &in_dim, &out_dim, &down_ctx.row_bytes[e]); - if (in_dim != down_in_dim || out_dim != down_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for(down_out_dim, matvec_q2_k_batch_accum_rows_worker, &down_ctx); - } else if (down_type == DS4_TENSOR_Q4_K) { - matvec_q4_k_batch_accum_rows_ctx down_ctx = { - .moe = moe, - .midq = midq, - .pairs = pairs, - .pair_ids = pair_ids, - .expert_offset = counts, - .active_expert = active_expert, - .n_active = n_active, - .n_tok = n_tok, - .in_dim = down_in_dim, - .out_dim = down_out_dim, - .midq_blocks = midq_blocks, - }; - - for (uint32_t ai = 0; ai < n_active; ai++) { - const uint32_t e = active_expert[ai]; - uint64_t in_dim, out_dim; - down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, - &in_dim, &out_dim, &down_ctx.row_bytes[e]); - if (in_dim != down_in_dim || out_dim != down_out_dim) { - ds4_die("batch expert tensor layout mismatch"); - } - } - - ds4_parallel_for(down_out_dim, matvec_q4_k_batch_accum_rows_worker, &down_ctx); - } else { - ds4_die("unsupported down expert tensor type for batch"); - } - - free(midq); - free(pair_ids); - free(xq); - free(pairs); - free(pair_weight); - free(selected); - - (void)il; -} - -static void print_vec_stats(const char *name, const float *x, uint64_t n); - -/* Full FFN sublayer for one token: HC pre, RMSNorm, routed MoE, shared expert, - * sum, and HC post. */ -static void layer_ffn_one( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - uint32_t il, - int token, - const float * steering_dirs, - float steering_scale, - bool trace) { - const uint32_t n_hc = DS4_N_HC; - const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; - const double t_start = profile ? now_sec() : 0.0; - double t_hc = 0.0; - double t_norm = 0.0; - double t_routed = 0.0; - double t_shared = 0.0; - double t_post = 0.0; - float *ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_cur[0])); - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); - float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); - float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); - float post[4]; - float comb[16]; - - double t0 = profile ? now_sec() : 0.0; - hc_pre_from_state_one(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - inp_hc, ffn_cur, post, comb); - if (profile) t_hc = now_sec() - t0; - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u ffn_cur", il); - print_vec_stats(name, ffn_cur, DS4_N_EMBD); - } - - t0 = profile ? now_sec() : 0.0; - const float *ffn_norm = tensor_data(model, layer->ffn_norm); - rms_norm_weight(norm, ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); - if (profile) t_norm = now_sec() - t0; - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u ffn_norm", il); - print_vec_stats(name, norm, DS4_N_EMBD); - } - - t0 = profile ? now_sec() : 0.0; - layer_routed_moe_one(moe, model, layer, norm, il, token, DS4_SWIGLU_CLAMP_EXP, trace); - if (profile) t_routed = now_sec() - t0; - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u routed_moe", il); - print_vec_stats(name, moe, DS4_N_EMBD); - } - t0 = profile ? now_sec() : 0.0; - layer_shared_ffn_one(shared, model, layer, norm); - if (profile) t_shared = now_sec() - t0; - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u shared_ffn", il); - print_vec_stats(name, shared, DS4_N_EMBD); - } - - t0 = profile ? now_sec() : 0.0; - for (uint32_t i = 0; i < DS4_N_EMBD; i++) { - ffn_out[i] = moe[i] + shared[i]; - } - cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, 1, steering_scale); - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u ffn_out", il); - print_vec_stats(name, ffn_out, DS4_N_EMBD); - } - - hc_post_one(out_hc, ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); - if (profile) t_post = now_sec() - t0; - if (trace) { - char name[64]; - snprintf(name, sizeof(name), "blk.%u ffn_post_hc", il); - print_vec_stats(name, out_hc, (uint64_t)n_hc * DS4_N_EMBD); - } - - if (profile) { - fprintf(stderr, - "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", - il, - t_hc * 1000.0, - t_norm * 1000.0, - t_routed * 1000.0, - t_shared * 1000.0, - t_post * 1000.0, - (now_sec() - t_start) * 1000.0); - } - - free(ffn_out); - free(shared); - free(moe); - free(norm); - free(ffn_cur); -} - -/* Allocation-free decode FFN using the persistent CPU scratch buffers. */ -static void layer_ffn_one_decode_scratch( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - uint32_t il, - int token, - const float * steering_dirs, - float steering_scale, - ds4_cpu_decode_scratch * scratch) { - const uint32_t n_hc = DS4_N_HC; - const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; - const double t_start = profile ? now_sec() : 0.0; - double t_hc = 0.0; - double t_norm = 0.0; - double t_routed = 0.0; - double t_shared = 0.0; - double t_post = 0.0; - float post[4]; - float comb[16]; - - double t0 = profile ? now_sec() : 0.0; - hc_pre_from_state_one_scratch(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - inp_hc, scratch->ffn_cur, post, comb, - scratch->hc_flat, - false); - if (profile) t_hc = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - const float *ffn_norm = tensor_data(model, layer->ffn_norm); - rms_norm_weight(scratch->ffn_norm, scratch->ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); - if (profile) t_norm = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_routed_moe_one_prealloc(scratch->ffn_moe, - model, - layer, - scratch->ffn_norm, - il, - token, - DS4_SWIGLU_CLAMP_EXP, - scratch->routed_mid_all, - scratch->routed_xq, - scratch->routed_midq, - scratch->routed_q8_xq, - scratch->routed_q8_xscale, - scratch->routed_q8_midq, - scratch->routed_q8_midscale); - if (profile) t_routed = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_shared_ffn_one_decode_scratch(scratch->ffn_shared, model, layer, scratch->ffn_norm, scratch); - if (profile) t_shared = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - for (uint32_t i = 0; i < DS4_N_EMBD; i++) { - scratch->ffn_out[i] = scratch->ffn_moe[i] + scratch->ffn_shared[i]; - } - cpu_directional_steering_project_rows(scratch->ffn_out, steering_dirs, il, 1, steering_scale); - hc_post_one(out_hc, scratch->ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); - if (profile) t_post = now_sec() - t0; - - if (profile) { - fprintf(stderr, - "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", - il, - t_hc * 1000.0, - t_norm * 1000.0, - t_routed * 1000.0, - t_shared * 1000.0, - t_post * 1000.0, - (now_sec() - t_start) * 1000.0); - } -} - -static void layer_ffn_batch( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - const int * token_ids, - uint32_t n_tok, - uint32_t il, - const float * steering_dirs, - float steering_scale) { - if (n_tok == 0) return; - const uint32_t n_hc = DS4_N_HC; - const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; - float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); - float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); - float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); - float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); - float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); - float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); - const float *ffn_norm = tensor_data(model, layer->ffn_norm); - - for (uint32_t t = 0; t < n_tok; t++) { - hc_pre_from_state_one(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - inp_hc + (uint64_t)t * hc_dim, - ffn_cur + (uint64_t)t * DS4_N_EMBD, - post + (uint64_t)t * n_hc, - comb + (uint64_t)t * n_hc * n_hc); - rms_norm_weight(norm + (uint64_t)t * DS4_N_EMBD, - ffn_cur + (uint64_t)t * DS4_N_EMBD, - ffn_norm, - DS4_N_EMBD, - DS4_RMS_EPS); - } - - layer_routed_moe_batch(moe, model, layer, norm, token_ids, n_tok, il, DS4_SWIGLU_CLAMP_EXP); - layer_shared_ffn_batch(shared, model, layer, norm, n_tok); - - if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { - float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); - for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { - ffn_out[i] = moe[i] + shared[i]; - } - cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); - hc_post_batch(out_hc, - ffn_out, - inp_hc, - post, - comb, - n_tok, - DS4_N_EMBD, - n_hc); - free(ffn_out); - } else { - hc_post_sum_batch(out_hc, - moe, - shared, - inp_hc, - post, - comb, - n_tok, - DS4_N_EMBD, - n_hc); - } - - free(comb); - free(post); - free(shared); - free(moe); - free(norm); - free(ffn_cur); -} - -typedef struct { - float *moe; - const ds4_model *model; - const ds4_layer_weights *layer; - const float *norm; - const int *token_ids; - uint64_t expert_in_dim; - uint64_t down_in_dim; - uint32_t il; - bool routed_q8_0; -} routed_moe_tokens_ctx; - -static void routed_moe_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { - routed_moe_tokens_ctx *ctx = vctx; - const uint64_t q8_x_blocks = ctx->expert_in_dim / 32u; - const uint64_t q8_mid_blocks = ctx->down_in_dim / 32u; - float *routed_mid = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); - block_q8_K *routed_xq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)(ctx->expert_in_dim / QK_K) * sizeof(routed_xq[0])); - block_q8_K *routed_midq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (ctx->down_in_dim / QK_K) * sizeof(routed_midq[0])); - int8_t *routed_q8_xq = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * 32u) : NULL; - float *routed_q8_xscale = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; - int8_t *routed_q8_midq = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * 32u) : NULL; - float *routed_q8_midscale = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; - - for (uint64_t t = t0; t < t1; t++) { - layer_routed_moe_one_prealloc(ctx->moe + t * DS4_N_EMBD, - ctx->model, - ctx->layer, - ctx->norm + t * DS4_N_EMBD, - ctx->il, - ctx->token_ids[t], - DS4_SWIGLU_CLAMP_EXP, - routed_mid, - routed_xq, - routed_midq, - routed_q8_xq, - routed_q8_xscale, - routed_q8_midq, - routed_q8_midscale); - } - - free(routed_q8_midscale); - free(routed_q8_midq); - free(routed_q8_xscale); - free(routed_q8_xq); - free(routed_midq); - free(routed_xq); - free(routed_mid); -} - -static void layer_routed_moe_tokens_parallel( - float * moe, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * norm, - const int * token_ids, - uint32_t n_tok, - uint32_t il) { - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const bool routed_q8_k = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; - if (routed_q8_k) { - if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); - if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { - ds4_die("Q8_K expert input has an unexpected layout"); - } - } - routed_moe_tokens_ctx ctx = { - .moe = moe, - .model = model, - .layer = layer, - .norm = norm, - .token_ids = token_ids, - .expert_in_dim = expert_in_dim, - .down_in_dim = down_in_dim, - .il = il, - .routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0, - }; - ds4_parallel_for_min_rows(n_tok, routed_moe_tokens_worker, &ctx, 1); -} - -/* Default prefill FFN path. HC and shared expert are batched, while routed - * experts can run either token-parallel or expert-grouped depending on size. */ -static void layer_ffn_shared_batch( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - const int * token_ids, - uint32_t n_tok, - uint32_t il, - const float * steering_dirs, - float steering_scale) { - const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; - const double t_start = profile ? now_sec() : 0.0; - double t_hc_norm = 0.0; - double t_routed = 0.0; - double t_shared = 0.0; - double t_post = 0.0; - const uint32_t n_hc = DS4_N_HC; - float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); - float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); - float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); - float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); - float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); - float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const bool routed_q8_k = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; - if (routed_q8_k) { - if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); - if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { - ds4_die("Q8_K expert input has an unexpected layout"); - } - } - const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; - const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; - const bool routed_token_parallel = - getenv("DS4_ROUTED_TOKEN_PARALLEL") != NULL || - (getenv("DS4_NO_ROUTED_TOKEN_PARALLEL") == NULL && n_tok >= 64); - float *routed_mid = routed_token_parallel ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); - block_q8_K *routed_xq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(routed_xq[0])); - block_q8_K *routed_midq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(routed_midq[0])); - int8_t *routed_q8_xq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; - float *routed_q8_xscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; - int8_t *routed_q8_midq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; - float *routed_q8_midscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; - - double t0 = profile ? now_sec() : 0.0; - hc_pre_norm_batch(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - layer->ffn_norm, - inp_hc, - NULL, - ffn_cur, - norm, - post, - comb, - n_tok); - if (profile) t_hc_norm = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - if (routed_token_parallel) { - layer_routed_moe_tokens_parallel(moe, model, layer, norm, token_ids, n_tok, il); - } else { - for (uint32_t t = 0; t < n_tok; t++) { - layer_routed_moe_one_prealloc(moe + (uint64_t)t * DS4_N_EMBD, - model, - layer, - norm + (uint64_t)t * DS4_N_EMBD, - il, - token_ids[t], - DS4_SWIGLU_CLAMP_EXP, - routed_mid, - routed_xq, - routed_midq, - routed_q8_xq, - routed_q8_xscale, - routed_q8_midq, - routed_q8_midscale); - } - } - if (profile) t_routed = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_shared_ffn_batch(shared, model, layer, norm, n_tok); - if (profile) t_shared = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { - float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); - for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { - ffn_out[i] = moe[i] + shared[i]; - } - cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); - hc_post_batch(out_hc, - ffn_out, - inp_hc, - post, - comb, - n_tok, - DS4_N_EMBD, - n_hc); - free(ffn_out); - } else { - hc_post_sum_batch(out_hc, - moe, - shared, - inp_hc, - post, - comb, - n_tok, - DS4_N_EMBD, - n_hc); - } - if (profile) t_post = now_sec() - t0; - - if (profile) { - fprintf(stderr, - "ds4: prefill detail layer %u ffn hc_norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f\n", - il, t_hc_norm, t_routed, t_shared, t_post, now_sec() - t_start); - } - - free(comb); - free(post); - free(routed_q8_midscale); - free(routed_q8_midq); - free(routed_q8_xscale); - free(routed_q8_xq); - free(routed_midq); - free(routed_xq); - free(routed_mid); - free(shared); - free(moe); - free(norm); - free(ffn_cur); -} - -typedef struct { - float *out_hc; - const ds4_model *model; - const ds4_layer_weights *layer; - const float *inp_hc; - const int *token_ids; - const float *steering_dirs; - float steering_scale; - uint64_t hc_dim; - uint32_t il; -} layer_ffn_tokens_ctx; - -static void layer_ffn_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { - layer_ffn_tokens_ctx *ctx = vctx; - for (uint64_t t = t0; t < t1; t++) { - layer_ffn_one(ctx->out_hc + t * ctx->hc_dim, - ctx->model, - ctx->layer, - ctx->inp_hc + t * ctx->hc_dim, - ctx->il, - ctx->token_ids[t], - ctx->steering_dirs, - ctx->steering_scale, - false); - } -} - -static void layer_ffn_tokens_parallel( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - const int * token_ids, - uint32_t n_tok, - uint32_t il, - const float * steering_dirs, - float steering_scale) { - layer_ffn_tokens_ctx ctx = { - .out_hc = out_hc, - .model = model, - .layer = layer, - .inp_hc = inp_hc, - .token_ids = token_ids, - .steering_dirs = steering_dirs, - .steering_scale = steering_scale, - .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, - .il = il, - }; - ds4_parallel_for(n_tok, layer_ffn_tokens_worker, &ctx); -} - -static void output_logits_one( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - const float * inp_hc); - -/* ========================================================================= - * KV Cache, Compressors, and CPU Layer Execution. - * ========================================================================= - * - * The CPU path is the correctness reference. It maintains raw SWA KV rows, - * optional compressed KV rows, the indexer mask for ratio-4 layers, and a - * reusable decode scratch arena so token generation does not allocate in the - * hot loop. - */ - -typedef struct { - float *raw_kv; - uint32_t n_raw; - uint32_t cap_raw; - - uint32_t compress_ratio; - uint32_t comp_cap; - uint32_t n_comp; - float *attn_comp_kv; - float *attn_state_kv; - float *attn_state_score; - - uint32_t n_index_comp; - float *index_comp_kv; - float *index_state_kv; - float *index_state_score; -} ds4_layer_cache; - -typedef struct { - ds4_layer_cache layer[DS4_MAX_LAYER]; - uint32_t head_dim; -} ds4_kv_cache; - -static uint32_t ds4_default_raw_cap(uint32_t ctx_size) { - uint32_t raw_cap = DS4_N_SWA; - if (raw_cap > ctx_size) raw_cap = ctx_size; - if (raw_cap == 0) raw_cap = 1; - return raw_cap; -} - -#define DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK 2048u - -static uint32_t ds4_effective_prefill_chunk(bool cuda_tensor_parallel, - uint32_t requested_chunk) { - if (requested_chunk != 0) return requested_chunk; - return cuda_tensor_parallel ? DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK : 0; -} - -static uint32_t ds4_prefill_cap_for_prompt(int prompt_len, - uint32_t requested_chunk) { - if (prompt_len <= 0) return 1; - uint32_t cap = (uint32_t)prompt_len; - - if (requested_chunk != 0) { - cap = requested_chunk; - } else { - const char *env = getenv("DS4_METAL_PREFILL_CHUNK"); - if (env && env[0]) { - char *endp = NULL; - const long v = strtol(env, &endp, 10); - if (endp != env) { - if (v <= 0) return cap; - cap = (uint32_t)v; - } - } else if (prompt_len > 4096) { - cap = DS4_MODEL_VARIANT == DS4_VARIANT_PRO ? 8192u : 4096u; - } - } - - if (cap == 0) cap = 1; - if (cap > (uint32_t)prompt_len) cap = (uint32_t)prompt_len; - return cap; -} - -/* Allocate all CPU decode temporaries once. This keeps generation deterministic - * from the VM's point of view and makes accidental hot-loop malloc visible. */ -static void cpu_decode_scratch_init(ds4_cpu_decode_scratch *scratch, uint32_t ctx_size) { - memset(scratch, 0, sizeof(*scratch)); - if (ctx_size == 0) ctx_size = 1; - const uint32_t raw_cap = ds4_default_raw_cap(ctx_size); - const uint32_t comp_cap = ctx_size / 4 + 2; - const uint32_t attn_score_cap = raw_cap + comp_cap; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t q8_cap = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t q8_blocks = (q8_cap + 31u) / 32u; - if ((DS4_N_EMBD % 32u) != 0 || (DS4_N_FF_EXP % 32u) != 0) { - ds4_die("Q8_0 routed decode scratch dimensions are not QK8_0 aligned"); - } - const uint64_t routed_q8_x_blocks = DS4_N_EMBD / 32u; - const uint64_t routed_q8_mid_blocks = DS4_N_FF_EXP / 32u; - - /* - * The CPU decode path used to malloc/free dozens of medium-sized buffers - * for every layer of every generated token. On macOS this can drive the VM - * system through repeated map/unmap bookkeeping while the huge model mmap is - * also being streamed, and we have observed kernel panics in VM accounting. - * Keep decode scratch resident for the whole generation instead. - */ - scratch->ctx_size = ctx_size; - scratch->comp_cap = comp_cap; - scratch->attn_score_cap = attn_score_cap; - scratch->q8_cap = (uint32_t)q8_cap; - - scratch->plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->cur = xmalloc((size_t)hc_dim * sizeof(float)); - scratch->next = xmalloc((size_t)hc_dim * sizeof(float)); - - scratch->attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->attn_residual = xmalloc((size_t)hc_dim * sizeof(float)); - scratch->q = xmalloc((size_t)q_dim * sizeof(float)); - scratch->qr = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); - scratch->qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); - scratch->kv_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - scratch->kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - scratch->heads = xmalloc((size_t)q_dim * sizeof(float)); - scratch->attn_low = xmalloc((size_t)DS4_N_OUT_GROUP * DS4_N_LORA_O * sizeof(float)); - scratch->attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - scratch->attn_score = xmalloc((size_t)attn_score_cap * sizeof(float)); - - scratch->comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - scratch->index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - scratch->comp_kv_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); - scratch->comp_sc_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); - scratch->comp_pooled = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - - scratch->index_allowed = xmalloc((size_t)comp_cap * sizeof(bool)); - scratch->index_q = xmalloc((size_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - scratch->index_weights = xmalloc((size_t)DS4_N_INDEXER_HEAD * sizeof(float)); - scratch->index_scores = xmalloc((size_t)comp_cap * sizeof(float)); - - scratch->ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->ffn_moe = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->ffn_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->shared_gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); - scratch->shared_up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); - scratch->shared_mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); - scratch->routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); - scratch->routed_xq = xmalloc((size_t)(DS4_N_EMBD / QK_K) * sizeof(block_q8_K)); - scratch->routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (DS4_N_FF_EXP / QK_K) * sizeof(block_q8_K)); - scratch->routed_q8_xq = xmalloc((size_t)routed_q8_x_blocks * 32u); - scratch->routed_q8_xscale = xmalloc((size_t)routed_q8_x_blocks * sizeof(float)); - scratch->routed_q8_midq = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u); - scratch->routed_q8_midscale = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(float)); - - scratch->q8_xq = xmalloc((size_t)q8_blocks * 32u); - scratch->q8_xscale = xmalloc((size_t)q8_blocks * sizeof(float)); - - scratch->hc_flat = xmalloc((size_t)hc_dim * sizeof(float)); - scratch->output_flat = xmalloc((size_t)hc_dim * sizeof(float)); - scratch->output_pre = xmalloc((size_t)DS4_N_HC * sizeof(float)); - scratch->output_weights = xmalloc((size_t)DS4_N_HC * sizeof(float)); - scratch->output_embd = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - scratch->output_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); -} - -static void cpu_decode_scratch_free(ds4_cpu_decode_scratch *scratch) { - if (!scratch) return; - free(scratch->output_norm); - free(scratch->output_embd); - free(scratch->output_weights); - free(scratch->output_pre); - free(scratch->output_flat); - free(scratch->hc_flat); - free(scratch->q8_xscale); - free(scratch->q8_xq); - free(scratch->routed_q8_midscale); - free(scratch->routed_q8_midq); - free(scratch->routed_q8_xscale); - free(scratch->routed_q8_xq); - free(scratch->routed_midq); - free(scratch->routed_xq); - free(scratch->routed_mid_all); - free(scratch->shared_mid); - free(scratch->shared_up); - free(scratch->shared_gate); - free(scratch->ffn_out); - free(scratch->ffn_shared); - free(scratch->ffn_moe); - free(scratch->ffn_norm); - free(scratch->ffn_cur); - free(scratch->index_scores); - free(scratch->index_weights); - free(scratch->index_q); - free(scratch->index_allowed); - free(scratch->comp_pooled); - free(scratch->comp_sc_cur); - free(scratch->comp_kv_cur); - free(scratch->index_comp); - free(scratch->comp); - free(scratch->attn_score); - free(scratch->after_attn_hc); - free(scratch->attn_out); - free(scratch->attn_low); - free(scratch->heads); - free(scratch->kv); - free(scratch->kv_raw); - free(scratch->qr_norm); - free(scratch->qr); - free(scratch->q); - free(scratch->attn_residual); - free(scratch->attn_norm); - free(scratch->attn_cur); - free(scratch->next); - free(scratch->cur); - free(scratch->plain); - memset(scratch, 0, sizeof(*scratch)); -} - -/* Allocate per-layer KV state: a raw sliding window for all layers, plus - * compressed attention/indexer caches for layers whose ratio is nonzero. */ -static void kv_cache_init(ds4_kv_cache *cache, uint32_t ctx_size, uint32_t raw_cap) { - memset(cache, 0, sizeof(*cache)); - if (raw_cap == 0) raw_cap = ds4_default_raw_cap(ctx_size); - if (raw_cap > ctx_size) raw_cap = ctx_size; - if (raw_cap == 0) raw_cap = 1; - - cache->head_dim = DS4_N_HEAD_DIM; - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - cache->layer[il].cap_raw = raw_cap; - cache->layer[il].raw_kv = xmalloc_zeroed((size_t)raw_cap * DS4_N_HEAD_DIM, sizeof(float)); - cache->layer[il].compress_ratio = ratio; - - if (ratio != 0) { - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint32_t comp_cap = ctx_size / ratio + 2; - const uint32_t attn_width = coff * DS4_N_HEAD_DIM; - const uint32_t attn_rows = coff * ratio; - - cache->layer[il].comp_cap = comp_cap; - cache->layer[il].attn_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_HEAD_DIM, sizeof(float)); - cache->layer[il].attn_state_kv = xmalloc_zeroed((size_t)attn_width * attn_rows, sizeof(float)); - cache->layer[il].attn_state_score = xmalloc((size_t)attn_width * attn_rows * sizeof(float)); - for (uint64_t i = 0; i < (uint64_t)attn_width * attn_rows; i++) { - cache->layer[il].attn_state_score[i] = DS4_NEG_INF; - } - - if (ratio == 4) { - const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; - const uint32_t index_rows = coff * ratio; - cache->layer[il].index_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_INDEXER_HEAD_DIM, sizeof(float)); - cache->layer[il].index_state_kv = xmalloc_zeroed((size_t)index_width * index_rows, sizeof(float)); - cache->layer[il].index_state_score = xmalloc((size_t)index_width * index_rows * sizeof(float)); - for (uint64_t i = 0; i < (uint64_t)index_width * index_rows; i++) { - cache->layer[il].index_state_score[i] = DS4_NEG_INF; - } - } - } - } -} - -static void kv_cache_free(ds4_kv_cache *cache) { - if (!cache) return; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - free(cache->layer[il].raw_kv); - free(cache->layer[il].attn_comp_kv); - free(cache->layer[il].attn_state_kv); - free(cache->layer[il].attn_state_score); - free(cache->layer[il].index_comp_kv); - free(cache->layer[il].index_state_kv); - free(cache->layer[il].index_state_score); - } - memset(cache, 0, sizeof(*cache)); -} - -/* Append to the raw SWA cache. Once full, it slides by one row. */ -static void kv_cache_push_raw(ds4_layer_cache *cache, const float *kv) { - if (cache->n_raw < cache->cap_raw) { - float *dst = cache->raw_kv + (uint64_t)cache->n_raw * DS4_N_HEAD_DIM; - for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); - cache->n_raw++; - return; - } - - memmove(cache->raw_kv, - cache->raw_kv + DS4_N_HEAD_DIM, - (size_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM * sizeof(cache->raw_kv[0])); - float *dst = cache->raw_kv + (uint64_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM; - for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); -} - -static void kv_cache_push_comp(float *rows, uint32_t *n_rows, uint32_t cap_rows, uint32_t row_dim, const float *kv) { - if (*n_rows >= cap_rows) ds4_die("compressed KV cache capacity exceeded"); - float *dst = rows + (uint64_t)(*n_rows) * row_dim; - for (uint32_t i = 0; i < row_dim; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); - (*n_rows)++; -} - -/* After prefill, clear unused compressor state rows so decode starts from the - * same partial-window state the streaming path would have produced. */ -static void compressor_finish_prefill_state_cpu( - float * state_kv, - float * state_score, - uint32_t head_dim, - uint32_t compress_ratio, - uint32_t n_tokens) { - if (!state_kv || !state_score || head_dim == 0 || compress_ratio == 0) return; - - const uint32_t coff = compress_ratio == 4 ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t rem = n_tokens % compress_ratio; - const uint32_t clear_start = compress_ratio == 4 ? compress_ratio + rem : rem; - const uint32_t clear_end = compress_ratio == 4 ? 2u * compress_ratio : compress_ratio; - - for (uint32_t row = clear_start; row < clear_end; row++) { - float *kv = state_kv + (uint64_t)row * width; - float *score = state_score + (uint64_t)row * width; - memset(kv, 0, (size_t)width * sizeof(kv[0])); - for (uint32_t i = 0; i < width; i++) score[i] = DS4_NEG_INF; - } -} - -static void kv_cache_finish_prefill_states(ds4_kv_cache *cache, uint32_t n_tokens) { - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_layer_cache *layer = &cache->layer[il]; - const uint32_t ratio = layer->compress_ratio; - if (ratio == 0) continue; - - compressor_finish_prefill_state_cpu(layer->attn_state_kv, - layer->attn_state_score, - DS4_N_HEAD_DIM, - ratio, - n_tokens); - if (ratio == 4) { - compressor_finish_prefill_state_cpu(layer->index_state_kv, - layer->index_state_score, - DS4_N_INDEXER_HEAD_DIM, - ratio, - n_tokens); - } - } -} - -/* Pool the current compression window with a softmax over per-dimension scores. - * Ratio-4 layers keep two lanes: attention compression and indexer compression. */ -static void compressor_pool_decode_state( - float * out, - float * state_kv, - float * state_score, - uint32_t head_dim, - uint32_t compress_ratio) { - const uint32_t coff = compress_ratio == 4 ? 2u : 1u; - const uint32_t width = coff * head_dim; - - for (uint32_t j = 0; j < head_dim; j++) { - float max_score = DS4_NEG_INF; - - if (compress_ratio == 4) { - for (uint32_t r = 0; r < compress_ratio; r++) { - const float sp = state_score[(uint64_t)r * width + j]; - const float sc = state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j]; - if (sp > max_score) max_score = sp; - if (sc > max_score) max_score = sc; - } - } else { - for (uint32_t r = 0; r < compress_ratio; r++) { - const float s = state_score[(uint64_t)r * width + j]; - if (s > max_score) max_score = s; - } - } - - if (max_score <= DS4_NEG_INF * 0.5f) { - out[j] = 0.0f; - continue; - } - - float denom = 0.0f; - float sum = 0.0f; - if (compress_ratio == 4) { - for (uint32_t r = 0; r < compress_ratio; r++) { - const float wp = expf(state_score[(uint64_t)r * width + j] - max_score); - const float wc = expf(state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j] - max_score); - denom += wp + wc; - sum += wp * state_kv[(uint64_t)r * width + j]; - sum += wc * state_kv[(uint64_t)(compress_ratio + r) * width + head_dim + j]; - } - } else { - for (uint32_t r = 0; r < compress_ratio; r++) { - const float w = expf(state_score[(uint64_t)r * width + j] - max_score); - denom += w; - sum += w * state_kv[(uint64_t)r * width + j]; - } - } - - out[j] = denom > 0.0f ? sum / denom : 0.0f; - } -} - -/* Streaming compressor update for one token. It projects kv/score rows, - * updates the rolling state, and emits a compressed KV row on ratio boundaries. */ -static bool compressor_decode_one( - float * out_comp, - const ds4_model * model, - const ds4_tensor * wkv, - const ds4_tensor * wgate, - const ds4_tensor * ape, - const ds4_tensor * norm, - const float * x, - float * state_kv, - float * state_score, - uint32_t head_dim, - uint32_t compress_ratio, - uint32_t il, - uint32_t pos) { - const uint32_t coff = compress_ratio == 4 ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t pos_mod = pos % compress_ratio; - const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; - const bool should_compress = ((pos + 1) % compress_ratio) == 0; - - float *kv_cur = xmalloc((size_t)width * sizeof(kv_cur[0])); - float *sc_cur = xmalloc((size_t)width * sizeof(sc_cur[0])); - if (wkv->type == 8 && - wgate->type == 8 && - wkv->ndim == 2 && - wgate->ndim == 2 && - wkv->dim[0] == wgate->dim[0]) { - const uint64_t in_dim = wkv->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - int8_t *xq = xmalloc((size_t)blocks * 32); - float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); - - quantize_q8_0_activation(x, xq, xscale, in_dim); - matvec_q8_0_pair_prequant(kv_cur, sc_cur, model, wkv, wgate, xq, xscale); - - free(xscale); - free(xq); - } else { - matvec_any(kv_cur, model, wkv, x); - matvec_any(sc_cur, model, wgate, x); - } - - for (uint32_t j = 0; j < width; j++) { - sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); - } - - memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); - memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); - - free(sc_cur); - free(kv_cur); - - if (!should_compress) { - return false; - } - - float *pooled = xmalloc((size_t)head_dim * sizeof(pooled[0])); - compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); - - double ss = 0.0; - for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; - const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); - for (uint32_t i = 0; i < head_dim; i++) { - out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); - } - - const uint32_t comp_pos = pos + 1 - compress_ratio; - rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); - if (head_dim == DS4_N_HEAD_DIM) { - dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); - } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { - dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); - } - - if (compress_ratio == 4) { - for (uint32_t r = 0; r < compress_ratio; r++) { - memcpy(state_kv + (uint64_t)r * width, - state_kv + (uint64_t)(compress_ratio + r) * width, - (size_t)width * sizeof(state_kv[0])); - memcpy(state_score + (uint64_t)r * width, - state_score + (uint64_t)(compress_ratio + r) * width, - (size_t)width * sizeof(state_score[0])); - } - for (uint32_t r = 0; r < compress_ratio; r++) { - memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, - state_kv + (uint64_t)r * width, - (size_t)width * sizeof(state_kv[0])); - memcpy(state_score + (uint64_t)(compress_ratio + r) * width, - state_score + (uint64_t)r * width, - (size_t)width * sizeof(state_score[0])); - } - } - - free(pooled); - return true; -} - -static bool compressor_decode_one_decode_scratch( - float * out_comp, - const ds4_model * model, - const ds4_tensor * wkv, - const ds4_tensor * wgate, - const ds4_tensor * ape, - const ds4_tensor * norm, - const float * x, - float * state_kv, - float * state_score, - uint32_t head_dim, - uint32_t compress_ratio, - uint32_t il, - uint32_t pos, - ds4_cpu_decode_scratch * scratch) { - const uint32_t coff = compress_ratio == 4 ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t pos_mod = pos % compress_ratio; - const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; - const bool should_compress = ((pos + 1) % compress_ratio) == 0; - - if (width > 2u * DS4_N_HEAD_DIM) ds4_die("compressor scratch width is outside the fixed model layout"); - float *kv_cur = scratch->comp_kv_cur; - float *sc_cur = scratch->comp_sc_cur; - - if (wkv->type == 8 && - wgate->type == 8 && - wkv->ndim == 2 && - wgate->ndim == 2 && - wkv->dim[0] == wgate->dim[0]) { - matvec_q8_0_pair_decode_scratch(kv_cur, sc_cur, model, wkv, wgate, x, scratch); - } else { - matvec_any_decode_scratch(kv_cur, model, wkv, x, scratch); - matvec_any_decode_scratch(sc_cur, model, wgate, x, scratch); - } - - for (uint32_t j = 0; j < width; j++) { - sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); - } - - memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); - memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); - - if (!should_compress) { - return false; - } - - float *pooled = scratch->comp_pooled; - compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); - - double ss = 0.0; - for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; - const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); - for (uint32_t i = 0; i < head_dim; i++) { - out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); - } - - const uint32_t comp_pos = pos + 1 - compress_ratio; - rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); - if (head_dim == DS4_N_HEAD_DIM) { - dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); - } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { - dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); - } - - if (compress_ratio == 4) { - for (uint32_t r = 0; r < compress_ratio; r++) { - memcpy(state_kv + (uint64_t)r * width, - state_kv + (uint64_t)(compress_ratio + r) * width, - (size_t)width * sizeof(state_kv[0])); - memcpy(state_score + (uint64_t)r * width, - state_score + (uint64_t)(compress_ratio + r) * width, - (size_t)width * sizeof(state_score[0])); - } - for (uint32_t r = 0; r < compress_ratio; r++) { - memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, - state_kv + (uint64_t)r * width, - (size_t)width * sizeof(state_kv[0])); - memcpy(state_score + (uint64_t)(compress_ratio + r) * width, - state_score + (uint64_t)r * width, - (size_t)width * sizeof(state_score[0])); - } - } - - return true; -} - -/* Attention over raw SWA rows plus optional compressed rows. Ratio-4 layers - * pass an indexer mask to hide compressed rows not selected for this token. */ -static void layer_attention_mixed_one( - float * out_heads, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * q, - const float * raw_kv, - uint32_t n_raw, - const float * comp_kv, - uint32_t n_comp, - const bool * comp_allowed) { - const float *sinks = tensor_data(model, layer->attn_sinks); - const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); - const uint32_t n_total = n_raw + n_comp; - float score_stack[512]; - float *score = n_total <= 512 ? score_stack : xmalloc((size_t)n_total * sizeof(score[0])); - - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; - float max_score = sinks[h]; - uint32_t idx = 0; - - for (uint32_t r = 0; r < n_raw; r++, idx++) { - const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; - score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[idx] > max_score) max_score = score[idx]; - } - for (uint32_t r = 0; r < n_comp; r++, idx++) { - if (comp_allowed && !comp_allowed[r]) { - score[idx] = DS4_NEG_INF; - continue; - } - const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; - score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[idx] > max_score) max_score = score[idx]; - } - - float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; - memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); - - float denom = expf(sinks[h] - max_score); - idx = 0; - for (uint32_t r = 0; r < n_raw; r++, idx++) { - const float weight = expf(score[idx] - max_score); - const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - for (uint32_t r = 0; r < n_comp; r++, idx++) { - if (score[idx] <= DS4_NEG_INF * 0.5f) continue; - const float weight = expf(score[idx] - max_score); - const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - - const float inv = 1.0f / denom; - scale_f32(oh, inv, DS4_N_HEAD_DIM); - } - - if (score != score_stack) free(score); -} - -static void layer_attention_mixed_one_decode_scratch( - float * out_heads, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * q, - const float * raw_kv, - uint32_t n_raw, - const float * comp_kv, - uint32_t n_comp, - const bool * comp_allowed, - ds4_cpu_decode_scratch * scratch) { - const float *sinks = tensor_data(model, layer->attn_sinks); - const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); - const uint32_t n_total = n_raw + n_comp; - if (n_total > scratch->attn_score_cap) ds4_die("CPU decode attention score scratch buffer is too small"); - float *score = scratch->attn_score; - - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; - float max_score = sinks[h]; - uint32_t idx = 0; - - for (uint32_t r = 0; r < n_raw; r++, idx++) { - const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; - score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[idx] > max_score) max_score = score[idx]; - } - for (uint32_t r = 0; r < n_comp; r++, idx++) { - if (comp_allowed && !comp_allowed[r]) { - score[idx] = DS4_NEG_INF; - continue; - } - const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; - score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[idx] > max_score) max_score = score[idx]; - } - - float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; - memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); - - float denom = expf(sinks[h] - max_score); - idx = 0; - for (uint32_t r = 0; r < n_raw; r++, idx++) { - const float weight = expf(score[idx] - max_score); - const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - for (uint32_t r = 0; r < n_comp; r++, idx++) { - if (score[idx] <= DS4_NEG_INF * 0.5f) continue; - const float weight = expf(score[idx] - max_score); - const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - - const float inv = 1.0f / denom; - scale_f32(oh, inv, DS4_N_HEAD_DIM); - } -} - -typedef struct { - float * out_heads; - const ds4_model * model; - const ds4_layer_weights * layer; - const float * q; - const float * raw_kv; - const float * comp_kv; - const uint32_t * comp_counts; - const uint8_t * allowed_mask; - const uint8_t * allowed_bits; - uint64_t allowed_stride; - uint32_t n_tok; - uint32_t raw_cap; -} layer_attention_prefix_batch_ctx; - -static inline bool attention_prefix_comp_allowed( - const layer_attention_prefix_batch_ctx *ctx, - uint32_t t, - uint32_t c) { - if (!ctx->allowed_bits || !ctx->allowed_mask || !ctx->allowed_mask[t]) return true; - const uint8_t *bits = ctx->allowed_bits + (uint64_t)t * ctx->allowed_stride; - return (bits[c >> 3] & (uint8_t)(1u << (c & 7u))) != 0; -} - -static void layer_attention_prefix_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { - layer_attention_prefix_batch_ctx *ctx = vctx; - const float *sinks = tensor_data(ctx->model, ctx->layer->attn_sinks); - const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); - const uint32_t max_comp = ctx->comp_counts ? ctx->comp_counts[ctx->n_tok - 1] : 0; - const uint32_t max_total = ctx->raw_cap + max_comp; - float score_stack[2048]; - float *score = max_total <= 2048 ? score_stack : xmalloc((size_t)max_total * sizeof(score[0])); - - for (uint64_t idx = r0; idx < r1; idx++) { - const uint32_t t = (uint32_t)(idx / DS4_N_HEAD); - const uint32_t h = (uint32_t)(idx - (uint64_t)t * DS4_N_HEAD); - const uint32_t raw_count = t + 1 < ctx->raw_cap ? t + 1 : ctx->raw_cap; - const uint32_t raw_start = t + 1 - raw_count; - const uint32_t comp_count = ctx->comp_counts ? ctx->comp_counts[t] : 0; - const float *qh = ctx->q + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; - - float max_score = sinks[h]; - uint32_t sidx = 0; - for (uint32_t r = 0; r < raw_count; r++, sidx++) { - const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; - score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[sidx] > max_score) max_score = score[sidx]; - } - for (uint32_t c = 0; c < comp_count; c++, sidx++) { - if (!attention_prefix_comp_allowed(ctx, t, c)) { - score[sidx] = DS4_NEG_INF; - continue; - } - const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; - score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; - if (score[sidx] > max_score) max_score = score[sidx]; - } - - float *oh = ctx->out_heads + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; - memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); - - float denom = expf(sinks[h] - max_score); - sidx = 0; - for (uint32_t r = 0; r < raw_count; r++, sidx++) { - const float weight = expf(score[sidx] - max_score); - const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - for (uint32_t c = 0; c < comp_count; c++, sidx++) { - if (score[sidx] <= DS4_NEG_INF * 0.5f) continue; - const float weight = expf(score[sidx] - max_score); - const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; - denom += weight; - axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); - } - - scale_f32(oh, 1.0f / denom, DS4_N_HEAD_DIM); - } - - if (score != score_stack) free(score); -} - -/* Prefix prefill attention for a fresh prompt. It computes each token's view - * of the raw window and compressed rows without running the decode loop. */ -static void layer_attention_prefix_batch( - float * out_heads, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * q, - const float * raw_kv, - const float * comp_kv, - const uint32_t * comp_counts, - const uint8_t * allowed_mask, - const uint8_t * allowed_bits, - uint64_t allowed_stride, - uint32_t n_tok, - uint32_t raw_cap) { - layer_attention_prefix_batch_ctx ctx = { - .out_heads = out_heads, - .model = model, - .layer = layer, - .q = q, - .raw_kv = raw_kv, - .comp_kv = comp_kv, - .comp_counts = comp_counts, - .allowed_mask = allowed_mask, - .allowed_bits = allowed_bits, - .allowed_stride = allowed_stride, - .n_tok = n_tok, - .raw_cap = raw_cap, - }; - ds4_parallel_for_min_rows((uint64_t)n_tok * DS4_N_HEAD, - layer_attention_prefix_batch_worker, - &ctx, - 1); -} - -/* Ratio-4 layers use an auxiliary indexer to select which compressed rows are - * visible to attention. This is the CPU allocation-owning helper. */ -static bool *indexer_allowed_decode_one( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * cur, - const float * qr_norm, - const float * index_comp, - uint32_t n_comp, - uint32_t il, - uint32_t pos) { - if (n_comp == 0) return NULL; - - bool *allowed = xcalloc(n_comp, sizeof(allowed[0])); - const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; - if (top_k == n_comp) { - for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; - return allowed; - } - - const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; - const uint32_t n_head = DS4_N_INDEXER_HEAD; - float *q = xmalloc((size_t)head_dim * n_head * sizeof(q[0])); - float *weights = xmalloc((size_t)n_head * sizeof(weights[0])); - float *scores = xmalloc((size_t)n_comp * sizeof(scores[0])); - - matvec_any(q, model, layer->indexer_attn_q_b, qr_norm); - rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); - dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); - - matvec_any(weights, model, layer->indexer_proj, cur); - const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); - for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; - - for (uint32_t c = 0; c < n_comp; c++) { - const float *kv = index_comp + (uint64_t)c * head_dim; - float s = 0.0f; - for (uint32_t h = 0; h < n_head; h++) { - const float *qh = q + (uint64_t)h * head_dim; - float dot = dot_f32(kv, qh, head_dim); - if (dot < 0.0f) dot = 0.0f; - s += dot * weights[h]; - } - scores[c] = s; - } - - for (uint32_t k = 0; k < top_k; k++) { - uint32_t best = 0; - float best_score = DS4_NEG_INF; - for (uint32_t c = 0; c < n_comp; c++) { - if (!allowed[c] && scores[c] > best_score) { - best = c; - best_score = scores[c]; - } - } - allowed[best] = true; - } - - free(scores); - free(weights); - free(q); - return allowed; -} - -/* Scratch-backed indexer selection for decode. */ -static bool *indexer_allowed_decode_one_decode_scratch( - const ds4_model * model, - const ds4_layer_weights * layer, - const float * cur, - const float * qr_norm, - const float * index_comp, - uint32_t n_comp, - uint32_t il, - uint32_t pos, - ds4_cpu_decode_scratch * scratch) { - if (n_comp == 0) return NULL; - if (n_comp > scratch->comp_cap) ds4_die("CPU decode indexer scratch buffer is too small"); - - bool *allowed = scratch->index_allowed; - memset(allowed, 0, (size_t)n_comp * sizeof(allowed[0])); - const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; - if (top_k == n_comp) { - for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; - return allowed; - } - - const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; - const uint32_t n_head = DS4_N_INDEXER_HEAD; - float *q = scratch->index_q; - float *weights = scratch->index_weights; - float *scores = scratch->index_scores; - - matvec_any_decode_scratch(q, model, layer->indexer_attn_q_b, qr_norm, scratch); - rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); - dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); - - matvec_any_decode_scratch(weights, model, layer->indexer_proj, cur, scratch); - const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); - for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; - - for (uint32_t c = 0; c < n_comp; c++) { - const float *kv = index_comp + (uint64_t)c * head_dim; - float s = 0.0f; - for (uint32_t h = 0; h < n_head; h++) { - const float *qh = q + (uint64_t)h * head_dim; - float dot = dot_f32(kv, qh, head_dim); - if (dot < 0.0f) dot = 0.0f; - s += dot * weights[h]; - } - scores[c] = s; - } - - for (uint32_t k = 0; k < top_k; k++) { - uint32_t best = 0; - float best_score = DS4_NEG_INF; - for (uint32_t c = 0; c < n_comp; c++) { - if (!allowed[c] && scores[c] > best_score) { - best = c; - best_score = scores[c]; - } - } - allowed[best] = true; - } - - return allowed; -} - -/* Single-token attention sublayer with raw SWA cache and DS4 compression. */ -static void layer_attention_raw_swa_one( - float * after_attn_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - ds4_layer_cache * cache, - const float * inp_hc, - uint32_t il, - uint32_t pos, - const float * steering_dirs, - float steering_scale) { - const uint32_t n_hc = DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - - float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); - float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); - float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); - float *q = xmalloc((size_t)q_dim * sizeof(q[0])); - float *qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(qr_norm[0])); - float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); - float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); - float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); - bool *comp_allowed = NULL; - float post[4]; - float comb[16]; - - memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); - hc_pre_from_state_one(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - attn_residual, attn_cur, post, comb); - - layer_attn_norm_one(attn_norm, model, layer, attn_cur); - layer_q_projection_with_lora_one(model, layer, attn_norm, q, qr_norm); - layer_kv_projection_normed_one(model, layer, attn_norm, kv); - - rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); - - kv_cache_push_raw(cache, kv); - - const uint32_t ratio = cache->compress_ratio; - if (ratio != 0) { - float *comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp[0])); - if (compressor_decode_one(comp, model, - layer->attn_compressor_kv, - layer->attn_compressor_gate, - layer->attn_compressor_ape, - layer->attn_compressor_norm, - attn_norm, - cache->attn_state_kv, - cache->attn_state_score, - DS4_N_HEAD_DIM, - ratio, - il, - pos)) { - kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); - } - free(comp); - - if (ratio == 4) { - float *index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp[0])); - if (compressor_decode_one(index_comp, model, - layer->indexer_compressor_kv, - layer->indexer_compressor_gate, - layer->indexer_compressor_ape, - layer->indexer_compressor_norm, - attn_norm, - cache->index_state_kv, - cache->index_state_score, - DS4_N_INDEXER_HEAD_DIM, - ratio, - il, - pos)) { - kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); - } - free(index_comp); - - comp_allowed = indexer_allowed_decode_one(model, layer, - attn_norm, qr_norm, - cache->index_comp_kv, - cache->n_index_comp, - il, pos); - } - - layer_attention_mixed_one(heads, model, layer, q, - cache->raw_kv, cache->n_raw, - cache->attn_comp_kv, cache->n_comp, - comp_allowed); - } else { - layer_attention_rows_one(heads, model, layer, q, cache->raw_kv, cache->n_raw); - } - - rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); - layer_grouped_out_one(attn_out, model, layer, heads); - cpu_directional_steering_project_rows(attn_out, steering_dirs, il, 1, steering_scale); - hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); - - free(comp_allowed); - free(attn_out); - free(heads); - free(kv); - free(qr_norm); - free(q); - free(attn_residual); - free(attn_norm); - free(attn_cur); -} - -/* Batched prefill attention. It projects Q/KV for all tokens, streams them - * through the same raw/compressed cache updates, then runs prefix attention. */ -static void layer_attention_raw_swa_batch( - float * after_attn_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - ds4_layer_cache * cache, - const float * inp_hc, - uint32_t n_tok, - uint32_t il, - uint32_t pos0, - const float * steering_dirs, - float steering_scale) { - const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; - const double t_start = profile ? now_sec() : 0.0; - double t_hc_norm = 0.0; - double t_q = 0.0; - double t_kv = 0.0; - double t_token_loop = 0.0; - double t_tl_rope_cache = 0.0; - double t_tl_compress = 0.0; - double t_tl_indexer = 0.0; - double t_tl_attn_rows = 0.0; - double t_tl_inv_rope = 0.0; - double t_out = 0.0; - const uint32_t n_hc = DS4_N_HC; - const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - - float *attn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_cur[0])); - float *attn_norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_norm[0])); - float *attn_residual = xmalloc((size_t)n_tok * hc_dim * sizeof(attn_residual[0])); - const uint32_t q_rank = DS4_N_LORA_Q; - float *qr = xmalloc((size_t)n_tok * q_rank * sizeof(qr[0])); - float *qr_norm = xmalloc((size_t)n_tok * q_rank * sizeof(qr_norm[0])); - float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); - float *kv_raw = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv_raw[0])); - float *kv = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv[0])); - float *heads = NULL; - float *attn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_out[0])); - float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); - float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); - - const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); - const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); - - double t0 = profile ? now_sec() : 0.0; - hc_pre_norm_batch(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - layer->attn_norm, - inp_hc, - attn_residual, - attn_cur, - attn_norm, - post, - comb, - n_tok); - if (profile) t_hc_norm = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - matmul_q8_0_batch(qr, model, layer->attn_q_a, attn_norm, n_tok); - for (uint32_t t = 0; t < n_tok; t++) { - rms_norm_weight(qr_norm + (uint64_t)t * q_rank, - qr + (uint64_t)t * q_rank, - q_a_norm, - q_rank, - DS4_RMS_EPS); - } - matmul_q8_0_batch(q, model, layer->attn_q_b, qr_norm, n_tok); - for (uint32_t t = 0; t < n_tok; t++) { - head_rms_norm_inplace(q + (uint64_t)t * q_dim, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_RMS_EPS); - } - if (profile) t_q = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - matmul_q8_0_batch(kv_raw, model, layer->attn_kv, attn_norm, n_tok); - for (uint32_t t = 0; t < n_tok; t++) { - rms_norm_weight(kv + (uint64_t)t * DS4_N_HEAD_DIM, - kv_raw + (uint64_t)t * DS4_N_HEAD_DIM, - kv_norm, - DS4_N_HEAD_DIM, - DS4_RMS_EPS); - } - if (profile) t_kv = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - const uint32_t ratio = cache->compress_ratio; - const bool prefer_parallel_attn = getenv("DS4_PARALLEL_ATTN_ROWS") != NULL; - const bool prefix_batch_attn = - prefer_parallel_attn && - getenv("DS4_NO_PARALLEL_ATTN_ROWS") == NULL && - cache->n_raw == 0 && - pos0 == 0; - if (!prefix_batch_attn) { - heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); - } - uint32_t batch_rope_max = 4096; - const char *batch_rope_max_env = getenv("DS4_BATCHED_ROPE_MAX"); - if (batch_rope_max_env && batch_rope_max_env[0]) { - long v = strtol(batch_rope_max_env, NULL, 10); - if (v >= 0 && v <= 65536) batch_rope_max = (uint32_t)v; - } - const bool batch_prefix_rope = - prefix_batch_attn && - getenv("DS4_NO_BATCHED_ROPE") == NULL && - n_tok <= batch_rope_max; - uint32_t *comp_counts = prefix_batch_attn ? - xcalloc((size_t)n_tok, sizeof(comp_counts[0])) : NULL; - uint8_t *allowed_mask = prefix_batch_attn && ratio == 4 ? - xcalloc((size_t)n_tok, sizeof(allowed_mask[0])) : NULL; - uint8_t *allowed_bits = NULL; - const uint64_t allowed_stride = ratio == 4 ? ((uint64_t)cache->comp_cap + 7u) / 8u : 0; - float *comp_scratch = NULL; - float *index_comp_scratch = NULL; - - if (ratio != 0) { - comp_scratch = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp_scratch[0])); - - if (ratio == 4) { - index_comp_scratch = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp_scratch[0])); - } - } - - if (batch_prefix_rope) { - double tx = profile ? now_sec() : 0.0; - rope_tail_layer_batch_inplace(q, - q_dim, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0, - il, - false, - n_tok); - rope_tail_layer_batch_inplace(kv, - DS4_N_HEAD_DIM, - DS4_N_HEAD_KV, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0, - il, - false, - n_tok); - if (profile) t_tl_rope_cache += now_sec() - tx; - } - - for (uint32_t t = 0; t < n_tok; t++) { - const uint32_t pos = pos0 + t; - float *q_t = q + (uint64_t)t * q_dim; - float *kv_t = kv + (uint64_t)t * DS4_N_HEAD_DIM; - bool *comp_allowed = NULL; - - double tx = profile ? now_sec() : 0.0; - if (!batch_prefix_rope) { - rope_tail_layer_inplace(q_t, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - rope_tail_layer_inplace(kv_t, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - } - dsv4_fp8_kv_quantize_row_inplace_cpu(kv_t, DS4_N_HEAD_DIM, DS4_N_ROT); - - kv_cache_push_raw(cache, kv_t); - if (profile) t_tl_rope_cache += now_sec() - tx; - - if (ratio != 0) { - tx = profile ? now_sec() : 0.0; - float *comp = comp_scratch; - const bool have_comp = compressor_decode_one(comp, model, - layer->attn_compressor_kv, - layer->attn_compressor_gate, - layer->attn_compressor_ape, - layer->attn_compressor_norm, - attn_norm + (uint64_t)t * DS4_N_EMBD, - cache->attn_state_kv, - cache->attn_state_score, - DS4_N_HEAD_DIM, - ratio, - il, - pos); - if (have_comp) { - kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); - } - - if (ratio == 4) { - float *index_comp = index_comp_scratch; - const bool have_index_comp = compressor_decode_one(index_comp, model, - layer->indexer_compressor_kv, - layer->indexer_compressor_gate, - layer->indexer_compressor_ape, - layer->indexer_compressor_norm, - attn_norm + (uint64_t)t * DS4_N_EMBD, - cache->index_state_kv, - cache->index_state_score, - DS4_N_INDEXER_HEAD_DIM, - ratio, - il, - pos); - if (have_index_comp) { - kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); - } - if (profile) t_tl_compress += now_sec() - tx; - - tx = profile ? now_sec() : 0.0; - comp_allowed = indexer_allowed_decode_one(model, layer, - attn_norm + (uint64_t)t * DS4_N_EMBD, - qr_norm + (uint64_t)t * q_rank, - cache->index_comp_kv, - cache->n_index_comp, - il, pos); - if (profile) t_tl_indexer += now_sec() - tx; - } else { - if (profile) t_tl_compress += now_sec() - tx; - } - - if (comp_counts) comp_counts[t] = cache->n_comp; - if (prefix_batch_attn && comp_allowed) { - if (!allowed_bits) { - allowed_bits = xcalloc((size_t)n_tok * allowed_stride, sizeof(allowed_bits[0])); - } - allowed_mask[t] = 1; - uint8_t *bits = allowed_bits + (uint64_t)t * allowed_stride; - for (uint32_t c = 0; c < cache->n_comp; c++) { - if (comp_allowed[c]) bits[c >> 3] |= (uint8_t)(1u << (c & 7u)); - } - } - - if (!prefix_batch_attn) { - tx = profile ? now_sec() : 0.0; - layer_attention_mixed_one(heads + (uint64_t)t * q_dim, model, layer, q_t, - cache->raw_kv, cache->n_raw, - cache->attn_comp_kv, cache->n_comp, - comp_allowed); - if (profile) t_tl_attn_rows += now_sec() - tx; - } - } else { - if (!prefix_batch_attn) { - tx = profile ? now_sec() : 0.0; - layer_attention_rows_one(heads + (uint64_t)t * q_dim, model, layer, q_t, cache->raw_kv, cache->n_raw); - if (profile) t_tl_attn_rows += now_sec() - tx; - } - } - - if (!prefix_batch_attn) { - tx = profile ? now_sec() : 0.0; - rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - il, - true); - if (profile) t_tl_inv_rope += now_sec() - tx; - } - - free(comp_allowed); - } - - if (prefix_batch_attn) { - double tx = profile ? now_sec() : 0.0; - const float *comp_kv_for_prefix = cache->attn_comp_kv ? cache->attn_comp_kv : kv; - if (!heads) { - heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); - } - layer_attention_prefix_batch(heads, model, layer, - q, - kv, - comp_kv_for_prefix, - comp_counts, - allowed_mask, - allowed_bits, - allowed_stride, - n_tok, - cache->cap_raw); - if (profile) t_tl_attn_rows += now_sec() - tx; - tx = profile ? now_sec() : 0.0; - if (batch_prefix_rope) { - rope_tail_layer_batch_inplace(heads, - q_dim, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0, - il, - true, - n_tok); - } else { - for (uint32_t t = 0; t < n_tok; t++) { - rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0 + t, - il, - true); - } - } - if (profile) t_tl_inv_rope += now_sec() - tx; - } - if (profile) t_token_loop = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_grouped_out_batch(attn_out, model, layer, heads, n_tok); - cpu_directional_steering_project_rows(attn_out, steering_dirs, il, n_tok, steering_scale); - - hc_post_batch(after_attn_hc, - attn_out, - attn_residual, - post, - comb, - n_tok, - DS4_N_EMBD, - n_hc); - if (profile) t_out = now_sec() - t0; - - if (profile) { - fprintf(stderr, - "ds4: prefill detail layer %u attn hc_norm=%.3f q=%.3f kv=%.3f token_loop=%.3f out=%.3f total=%.3f\n", - il, t_hc_norm, t_q, t_kv, t_token_loop, t_out, now_sec() - t_start); - if (getenv("DS4_PREFILL_PROFILE_TOKEN") != NULL) { - fprintf(stderr, - "ds4: prefill token detail layer %u rope_cache=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f\n", - il, t_tl_rope_cache, t_tl_compress, t_tl_indexer, t_tl_attn_rows, t_tl_inv_rope); - } - } - - free(allowed_bits); - free(allowed_mask); - free(comp_counts); - free(index_comp_scratch); - free(comp_scratch); - free(comb); - free(post); - free(attn_out); - free(heads); - free(kv); - free(kv_raw); - free(q); - free(qr_norm); - free(qr); - free(attn_residual); - free(attn_norm); - free(attn_cur); -} - -/* Full transformer layer for one decode token: attention sublayer followed by - * FFN sublayer, both operating on the HC state. */ -static void layer_forward_raw_swa_one( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - ds4_layer_cache * cache, - const float * inp_hc, - uint32_t il, - uint32_t pos, - int token, - const float * steering_dirs, - float steering_attn_scale, - float steering_ffn_scale, - ds4_cpu_decode_scratch * scratch) { - const uint32_t n_hc = DS4_N_HC; - const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; - const double t_start = profile ? now_sec() : 0.0; - double t_hc = 0.0; - double t_q = 0.0; - double t_kv = 0.0; - double t_rope_cache = 0.0; - double t_compress = 0.0; - double t_indexer = 0.0; - double t_attn_rows = 0.0; - double t_inv_rope = 0.0; - double t_out = 0.0; - double t_post = 0.0; - double t_ffn = 0.0; - - bool *comp_allowed = NULL; - float post[4]; - float comb[16]; - - double t0 = profile ? now_sec() : 0.0; - memcpy(scratch->attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); - hc_pre_from_state_one_scratch(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - scratch->attn_residual, scratch->attn_cur, post, comb, - scratch->hc_flat, - false); - if (profile) t_hc = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_attn_norm_one(scratch->attn_norm, model, layer, scratch->attn_cur); - const uint32_t ratio = cache->compress_ratio; - layer_q_projection_with_lora_one_decode_scratch(model, layer, - scratch->attn_norm, - scratch->q, - scratch->qr_norm, - scratch); - if (profile) t_q = now_sec() - t0; - t0 = profile ? now_sec() : 0.0; - layer_kv_projection_normed_one_decode_scratch(model, layer, - scratch->attn_norm, - scratch->kv, - scratch); - if (profile) t_kv = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - rope_tail_layer_inplace(scratch->q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - rope_tail_layer_inplace(scratch->kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - dsv4_fp8_kv_quantize_row_inplace_cpu(scratch->kv, DS4_N_HEAD_DIM, DS4_N_ROT); - - kv_cache_push_raw(cache, scratch->kv); - if (profile) t_rope_cache = now_sec() - t0; - - if (ratio != 0) { - t0 = profile ? now_sec() : 0.0; - if (compressor_decode_one_decode_scratch(scratch->comp, model, - layer->attn_compressor_kv, - layer->attn_compressor_gate, - layer->attn_compressor_ape, - layer->attn_compressor_norm, - scratch->attn_norm, - cache->attn_state_kv, - cache->attn_state_score, - DS4_N_HEAD_DIM, - ratio, - il, - pos, - scratch)) { - kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, scratch->comp); - } - - if (ratio == 4) { - if (compressor_decode_one_decode_scratch(scratch->index_comp, model, - layer->indexer_compressor_kv, - layer->indexer_compressor_gate, - layer->indexer_compressor_ape, - layer->indexer_compressor_norm, - scratch->attn_norm, - cache->index_state_kv, - cache->index_state_score, - DS4_N_INDEXER_HEAD_DIM, - ratio, - il, - pos, - scratch)) { - kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, - DS4_N_INDEXER_HEAD_DIM, scratch->index_comp); - } - if (profile) t_compress = now_sec() - t0; - } else if (profile) { - t_compress = now_sec() - t0; - } - } - if (ratio == 4) { - t0 = profile ? now_sec() : 0.0; - comp_allowed = indexer_allowed_decode_one_decode_scratch(model, layer, - scratch->attn_norm, - scratch->qr_norm, - cache->index_comp_kv, - cache->n_index_comp, - il, pos, - scratch); - if (profile) t_indexer = now_sec() - t0; - } - - t0 = profile ? now_sec() : 0.0; - if (ratio != 0) { - layer_attention_mixed_one_decode_scratch(scratch->heads, model, layer, scratch->q, - cache->raw_kv, cache->n_raw, - cache->attn_comp_kv, cache->n_comp, - comp_allowed, - scratch); - } else { - layer_attention_rows_one(scratch->heads, model, layer, scratch->q, cache->raw_kv, cache->n_raw); - } - if (profile) t_attn_rows = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - rope_tail_layer_inplace(scratch->heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); - if (profile) t_inv_rope = now_sec() - t0; - t0 = profile ? now_sec() : 0.0; - layer_grouped_out_one_decode_scratch(scratch->attn_out, model, layer, scratch->heads, scratch); - cpu_directional_steering_project_rows(scratch->attn_out, steering_dirs, il, 1, steering_attn_scale); - if (profile) t_out = now_sec() - t0; - t0 = profile ? now_sec() : 0.0; - hc_post_one(scratch->after_attn_hc, scratch->attn_out, scratch->attn_residual, post, comb, DS4_N_EMBD, n_hc); - if (profile) t_post = now_sec() - t0; - - t0 = profile ? now_sec() : 0.0; - layer_ffn_one_decode_scratch(out_hc, model, layer, scratch->after_attn_hc, il, token, - steering_dirs, steering_ffn_scale, scratch); - if (profile) t_ffn = now_sec() - t0; - - if (profile) { - fprintf(stderr, - "ds4: decode detail layer %u attn hc=%.3f q=%.3f kv=%.3f rope=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f out=%.3f post=%.3f ffn=%.3f total=%.3f ms\n", - il, - t_hc * 1000.0, - t_q * 1000.0, - t_kv * 1000.0, - t_rope_cache * 1000.0, - t_compress * 1000.0, - t_indexer * 1000.0, - t_attn_rows * 1000.0, - t_inv_rope * 1000.0, - t_out * 1000.0, - t_post * 1000.0, - t_ffn * 1000.0, - (now_sec() - t_start) * 1000.0); - } - -} - -static void output_logits_one_decode_scratch( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - const float * inp_hc, - ds4_cpu_decode_scratch * scratch); - -/* CPU decode for one token through all 43 layers. The caller owns scratch and - * cache lifetimes so no per-token allocations are needed. */ -static void forward_token_raw_swa_cpu_decode_scratch( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - ds4_kv_cache * cache, - int token, - uint32_t pos, - const float * steering_dirs, - float steering_attn_scale, - float steering_ffn_scale, - ds4_cpu_decode_scratch * scratch) { - float *cur = scratch->cur; - float *next = scratch->next; - - embed_token_f16(model, weights, token, scratch->plain); - hc_from_plain_embedding(cur, scratch->plain, DS4_N_EMBD, DS4_N_HC); - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - layer_forward_raw_swa_one(next, model, &weights->layer[il], &cache->layer[il], - cur, il, pos, token, - steering_dirs, - steering_attn_scale, - steering_ffn_scale, - scratch); - float *tmp = cur; - cur = next; - next = tmp; - } - - if (logits) { - output_logits_one_decode_scratch(logits, model, weights, cur, scratch); - } -} - -#ifndef DS4_NO_GPU -static void forward_token_raw_swa_cpu( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - ds4_kv_cache * cache, - int token, - uint32_t pos) { - ds4_cpu_decode_scratch scratch; - uint32_t ctx_guess = pos + 1; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = cache->layer[il].compress_ratio; - if (ratio != 0 && cache->layer[il].comp_cap > 2) { - const uint32_t ctx_from_comp = (cache->layer[il].comp_cap - 2u) * ratio; - if (ctx_guess < ctx_from_comp) ctx_guess = ctx_from_comp; - } - } - cpu_decode_scratch_init(&scratch, ctx_guess); - forward_token_raw_swa_cpu_decode_scratch(logits, model, weights, cache, token, pos, - NULL, 0.0f, 0.0f, &scratch); - cpu_decode_scratch_free(&scratch); -} -#endif - -/* CPU prefill in layer-major order. All prompt tokens pass through layer 0, - * then layer 1, etc., which exposes batch matmul opportunities. */ -static void prefill_layer_major_cpu( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - ds4_kv_cache * cache, - const token_vec * prompt, - const float * steering_dirs, - float steering_attn_scale, - float steering_ffn_scale) { - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t n_tok = (uint64_t)prompt->len; - float *cur = xmalloc((size_t)n_tok * hc_dim * sizeof(cur[0])); - float *next = xmalloc((size_t)n_tok * hc_dim * sizeof(next[0])); - float *attn = xmalloc((size_t)n_tok * hc_dim * sizeof(attn[0])); - float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); - uint32_t ffn_batch = 128; - const bool batched_attn = getenv("DS4_NO_BATCHED_ATTN") == NULL; - const bool batched_ffn = getenv("DS4_BATCHED_FFN") != NULL; - const bool parallel_ffn = getenv("DS4_PARALLEL_FFN") != NULL; - const bool shared_batch_ffn = getenv("DS4_NO_SHARED_BATCH_FFN") == NULL; - const char *batch_env = getenv("DS4_PREFILL_BATCH"); - ds4_cpu_decode_scratch decode_scratch; - bool decode_scratch_ready = false; - if (batch_env && batch_env[0]) { - long v = strtol(batch_env, NULL, 10); - if (v > 0 && v < 4096) ffn_batch = (uint32_t)v; - } - - for (uint64_t t = 0; t < n_tok; t++) { - embed_token_f16(model, weights, prompt->v[t], plain); - hc_from_plain_embedding(cur + t * hc_dim, plain, DS4_N_EMBD, DS4_N_HC); - } - - free(plain); - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - fprintf(stderr, "ds4: prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); - fflush(stderr); - - if (batched_attn) { - layer_attention_raw_swa_batch(attn, - model, - &weights->layer[il], - &cache->layer[il], - cur, - (uint32_t)n_tok, - il, - 0, - steering_dirs, - steering_attn_scale); - - if (batched_ffn) { - for (uint64_t t = 0; t < n_tok; t += ffn_batch) { - uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); - layer_ffn_batch(next + t * hc_dim, - model, - &weights->layer[il], - attn + t * hc_dim, - prompt->v + t, - nb, - il, - steering_dirs, - steering_ffn_scale); - } - } else if (shared_batch_ffn) { - layer_ffn_shared_batch(next, - model, - &weights->layer[il], - attn, - prompt->v, - (uint32_t)n_tok, - il, - steering_dirs, - steering_ffn_scale); - } else if (parallel_ffn) { - layer_ffn_tokens_parallel(next, - model, - &weights->layer[il], - attn, - prompt->v, - (uint32_t)n_tok, - il, - steering_dirs, - steering_ffn_scale); - } else { - for (uint64_t t = 0; t < n_tok; t++) { - layer_ffn_one(next + t * hc_dim, - model, - &weights->layer[il], - attn + t * hc_dim, - il, - prompt->v[t], - steering_dirs, - steering_ffn_scale, - false); - } - } - } else if (batched_ffn) { - for (uint64_t t = 0; t < n_tok; t++) { - layer_attention_raw_swa_one(attn + t * hc_dim, - model, - &weights->layer[il], - &cache->layer[il], - cur + t * hc_dim, - il, - (uint32_t)t, - steering_dirs, - steering_attn_scale); - } - - for (uint64_t t = 0; t < n_tok; t += ffn_batch) { - uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); - layer_ffn_batch(next + t * hc_dim, - model, - &weights->layer[il], - attn + t * hc_dim, - prompt->v + t, - nb, - il, - steering_dirs, - steering_ffn_scale); - } - } else { - if (!decode_scratch_ready) { - cpu_decode_scratch_init(&decode_scratch, (uint32_t)n_tok); - decode_scratch_ready = true; - } - for (uint64_t t = 0; t < n_tok; t++) { - layer_forward_raw_swa_one(next + t * hc_dim, - model, - &weights->layer[il], - &cache->layer[il], - cur + t * hc_dim, - il, - (uint32_t)t, - prompt->v[t], - steering_dirs, - steering_attn_scale, - steering_ffn_scale, - &decode_scratch); - } - } - - float *tmp = cur; - cur = next; - next = tmp; - } - - kv_cache_finish_prefill_states(cache, (uint32_t)n_tok); - - if (logits) { - output_logits_one(logits, model, weights, cur + (n_tok - 1) * hc_dim); - } - - if (decode_scratch_ready) cpu_decode_scratch_free(&decode_scratch); - free(next); - free(cur); - free(attn); -} - -/* Diagnostic first-token layer without cache history: the token attends only - * to itself, useful for checking a minimal end-to-end slice. */ -static void layer_forward_self_one( - float * out_hc, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * inp_hc, - uint32_t il, - uint32_t pos, - int token) { - const uint32_t n_hc = DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - - float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); - float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); - float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); - float *q = xmalloc((size_t)q_dim * sizeof(q[0])); - float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); - float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); - float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); - float *after_attn_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(after_attn_hc[0])); - float post[4]; - float comb[16]; - - memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); - hc_pre_from_state_one(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - attn_residual, attn_cur, post, comb); - - layer_attn_norm_one(attn_norm, model, layer, attn_cur); - layer_q_projection_normed_one(model, layer, attn_norm, q); - layer_kv_projection_normed_one(model, layer, attn_norm, kv); - rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); - dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); - f16_round_inplace_cpu(kv, DS4_N_HEAD_DIM); - - layer_attention_one(heads, model, layer, q, kv); - rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); - layer_grouped_out_one(attn_out, model, layer, heads); - hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); - - layer_ffn_one(out_hc, model, layer, after_attn_hc, il, token, - NULL, 0.0f, false); - - free(after_attn_hc); - free(attn_out); - free(heads); - free(kv); - free(q); - free(attn_residual); - free(attn_norm); - free(attn_cur); -} - -static void forward_first_token_cpu( - float * out_hc, - const ds4_model * model, - const ds4_weights * weights, - int token) { - float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); - float *cur = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(cur[0])); - float *next = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(next[0])); - - embed_token_f16(model, weights, token, plain); - hc_from_plain_embedding(cur, plain, DS4_N_EMBD, DS4_N_HC); - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - layer_forward_self_one(next, model, &weights->layer[il], cur, il, 0, token); - float *tmp = cur; - cur = next; - next = tmp; - } - - memcpy(out_hc, cur, (size_t)DS4_N_HC * DS4_N_EMBD * sizeof(out_hc[0])); - - free(next); - free(cur); - free(plain); -} - -/* Collapse final HC streams into the ordinary embedding vector before the - * output norm and vocabulary projection. */ -static void output_hc_head_one( - float * out, - const ds4_model * model, - const ds4_weights * weights, - const float * inp_hc) { - const uint32_t n_hc = DS4_N_HC; - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; - float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); - float *pre = xmalloc((size_t)n_hc * sizeof(pre[0])); - float *w = xmalloc((size_t)n_hc * sizeof(w[0])); - - rms_norm_no_weight(flat, inp_hc, hc_dim, DS4_RMS_EPS); - matvec_f16(pre, model, weights->output_hc_fn, flat); - - const float *scale = tensor_data(model, weights->output_hc_scale); - const float *base = tensor_data(model, weights->output_hc_base); - for (uint32_t i = 0; i < n_hc; i++) { - w[i] = sigmoid_stable(pre[i] * scale[0] + base[i]) + DS4_HC_EPS; - } - - hc_weighted_sum_one(out, inp_hc, w, DS4_N_EMBD, n_hc); - - free(w); - free(pre); - free(flat); -} - -/* Final language-model head: HC collapse, RMSNorm, and Q8_0 vocab projection. */ -static void output_logits_one( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - const float * inp_hc) { - float *embd = xmalloc((size_t)DS4_N_EMBD * sizeof(embd[0])); - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - - output_hc_head_one(embd, model, weights, inp_hc); - rms_norm_weight(norm, embd, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); - - matvec_q8_0(logits, model, weights->output, norm); - - free(norm); - free(embd); -} - -static void layer_glm_first_token_attention_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); - float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); - float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); - const uint64_t kv_blocks = (DS4_N_KV_LORA + 31) / 32; - int8_t *kvq = xmalloc((size_t)kv_blocks * 32); - float *kvscale = xmalloc((size_t)kv_blocks * sizeof(kvscale[0])); - - if (layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || - layer->attn_v_b->dim[0] != DS4_N_KV_LORA || - layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || - layer->attn_v_b->dim[2] != DS4_N_HEAD || - layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || - layer->attn_output->dim[1] != DS4_N_EMBD) { - ds4_die("GLM attention tensors have an unexpected layout"); - } - - rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0(kv_raw, model, layer->attn_kv_a_mqa, norm); - rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), - DS4_N_KV_LORA, DS4_RMS_EPS); - quantize_q8_0_activation(kv_norm, kvq, kvscale, DS4_N_KV_LORA); - - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - matvec_q8_0_3d_slice_prequant(heads + (uint64_t)h * DS4_N_VALUE_MLA, - model, - layer->attn_v_b, - kvq, - kvscale, - h); - } - matvec_q8_0(out, model, layer->attn_output, heads); - - free(kvscale); - free(kvq); - free(heads); - free(kv_norm); - free(kv_raw); - free(norm); -} - -static void layer_glm_first_token_attention_one_f32_ref( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); - float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); - float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); - - if (layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || - layer->attn_v_b->type != DS4_TENSOR_Q8_0 || - layer->attn_output->type != DS4_TENSOR_Q8_0 || - layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || - layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || - layer->attn_v_b->dim[0] != DS4_N_KV_LORA || - layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || - layer->attn_v_b->dim[2] != DS4_N_HEAD || - layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || - layer->attn_output->dim[1] != DS4_N_EMBD) { - ds4_die("GLM F32 attention reference found unexpected tensor layout"); - } - - rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0_f32_ref(kv_raw, model, layer->attn_kv_a_mqa, norm); - rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), - DS4_N_KV_LORA, DS4_RMS_EPS); - matvec_q8_0_f32_ref(heads, model, layer->attn_v_b, kv_norm); - matvec_q8_0_f32_ref(out, model, layer->attn_output, heads); - - free(heads); - free(kv_norm); - free(kv_raw); - free(norm); -} - -static void glm_k_b_project_f32_ref( - float * out, - const ds4_model * model, - const ds4_tensor * w, - const float * kv_norm) { - const uint32_t q_nope = DS4_N_KEY_MLA - DS4_N_ROT; - if (w->type != DS4_TENSOR_Q8_0 || - w->ndim != 3 || - w->dim[0] != q_nope || - w->dim[1] != DS4_N_KV_LORA || - w->dim[2] != DS4_N_HEAD) { - ds4_die("GLM k_b reference found unexpected tensor layout"); - } - - const uint8_t *data = tensor_data(model, w); - const uint64_t blocks = (q_nope + 31u) / 32u; - const uint64_t row_bytes = blocks * 34u; - memset(out, 0, (size_t)DS4_N_HEAD * q_nope * sizeof(out[0])); - - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - float *dst = out + (uint64_t)h * q_nope; - for (uint32_t j = 0; j < DS4_N_KV_LORA; j++) { - const uint8_t *row = - data + ((uint64_t)h * DS4_N_KV_LORA + j) * row_bytes; - const float xj = kv_norm[j]; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t scale_bits; - memcpy(&scale_bits, row + b * 34u, sizeof(scale_bits)); - const int8_t *qs = (const int8_t *)(row + b * 34u + 2u); - const float d = f16_to_f32(scale_bits) * xj; - const uint32_t i0 = (uint32_t)b * 32u; - const uint32_t n = q_nope - i0 < 32u ? q_nope - i0 : 32u; - for (uint32_t i = 0; i < n; i++) { - dst[i0 + i] += d * (float)qs[i]; - } - } - } - } -} - -static void layer_glm_attention_prefill_f32_ref( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t n_tok, - uint32_t pos0, - uint32_t il) { - if (n_tok == 0) return; - const uint32_t qk_dim = DS4_N_KEY_MLA; - const uint32_t q_nope = qk_dim - DS4_N_ROT; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * qk_dim; - const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; - const uint64_t kv_raw_dim = layer->attn_kv_a_mqa ? layer->attn_kv_a_mqa->dim[1] : 0; - - if (!layer->attn_norm || - !layer->attn_q_a || - !layer->attn_q_a_norm || - !layer->attn_q_b || - !layer->attn_kv_a_mqa || - !layer->attn_kv_a_norm || - !layer->attn_k_b || - !layer->attn_v_b || - !layer->attn_output || - layer->attn_q_a->type != DS4_TENSOR_Q8_0 || - layer->attn_q_b->type != DS4_TENSOR_Q8_0 || - layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || - layer->attn_k_b->type != DS4_TENSOR_Q8_0 || - layer->attn_v_b->type != DS4_TENSOR_Q8_0 || - layer->attn_output->type != DS4_TENSOR_Q8_0 || - layer->attn_norm->type != DS4_TENSOR_F32 || - layer->attn_q_a_norm->type != DS4_TENSOR_F32 || - layer->attn_kv_a_norm->type != DS4_TENSOR_F32 || - layer->attn_q_a->dim[0] != DS4_N_EMBD || - layer->attn_q_a->dim[1] != DS4_N_LORA_Q || - layer->attn_q_a_norm->dim[0] != DS4_N_LORA_Q || - layer->attn_q_b->dim[0] != DS4_N_LORA_Q || - layer->attn_q_b->dim[1] != q_dim || - layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || - kv_raw_dim < (uint64_t)DS4_N_KV_LORA + DS4_N_ROT || - layer->attn_kv_a_norm->dim[0] != DS4_N_KV_LORA || - layer->attn_k_b->dim[0] != q_nope || - layer->attn_k_b->dim[1] != DS4_N_KV_LORA || - layer->attn_k_b->dim[2] != DS4_N_HEAD || - layer->attn_v_b->dim[0] != DS4_N_KV_LORA || - layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || - layer->attn_v_b->dim[2] != DS4_N_HEAD || - layer->attn_output->dim[0] != heads_dim || - layer->attn_output->dim[1] != DS4_N_EMBD) { - ds4_die("GLM prefill attention reference found unexpected tensor layout"); - } - - float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); - float *q_rank = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank[0])); - float *q_rank_norm = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank_norm[0])); - float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); - float *kv_raw = xmalloc((size_t)n_tok * kv_raw_dim * sizeof(kv_raw[0])); - float *kv_norm = xmalloc((size_t)n_tok * DS4_N_KV_LORA * sizeof(kv_norm[0])); - float *k_nope = xmalloc((size_t)n_tok * DS4_N_HEAD * q_nope * sizeof(k_nope[0])); - float *key_cache = xmalloc((size_t)n_tok * DS4_N_HEAD * qk_dim * sizeof(key_cache[0])); - float *value_cache = xmalloc((size_t)n_tok * heads_dim * sizeof(value_cache[0])); - float *heads = xmalloc((size_t)n_tok * heads_dim * sizeof(heads[0])); - float *k_rot = xmalloc((size_t)DS4_N_ROT * sizeof(k_rot[0])); - float *scores = xmalloc((size_t)n_tok * sizeof(scores[0])); - - for (uint32_t t = 0; t < n_tok; t++) { - const float *xt = x + (uint64_t)t * DS4_N_EMBD; - float *norm_t = norm + (uint64_t)t * DS4_N_EMBD; - float *qr_t = q_rank + (uint64_t)t * DS4_N_LORA_Q; - float *qrn_t = q_rank_norm + (uint64_t)t * DS4_N_LORA_Q; - float *q_t = q + (uint64_t)t * q_dim; - float *raw_t = kv_raw + (uint64_t)t * kv_raw_dim; - float *kvn_t = kv_norm + (uint64_t)t * DS4_N_KV_LORA; - float *kn_t = k_nope + (uint64_t)t * DS4_N_HEAD * q_nope; - float *kc_t = key_cache + (uint64_t)t * DS4_N_HEAD * qk_dim; - float *vc_t = value_cache + (uint64_t)t * heads_dim; - - rms_norm_weight(norm_t, xt, tensor_data(model, layer->attn_norm), - DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0_f32_ref(qr_t, model, layer->attn_q_a, norm_t); - rms_norm_weight(qrn_t, qr_t, tensor_data(model, layer->attn_q_a_norm), - DS4_N_LORA_Q, DS4_RMS_EPS); - matvec_q8_0_f32_ref(q_t, model, layer->attn_q_b, qrn_t); - rope_tail_layer_inplace(q_t, DS4_N_HEAD, qk_dim, DS4_N_ROT, - pos0 + t, il, false); - - matvec_q8_0_f32_ref(raw_t, model, layer->attn_kv_a_mqa, norm_t); - rms_norm_weight(kvn_t, raw_t, tensor_data(model, layer->attn_kv_a_norm), - DS4_N_KV_LORA, DS4_RMS_EPS); - glm_k_b_project_f32_ref(kn_t, model, layer->attn_k_b, kvn_t); - matvec_q8_0_f32_ref(vc_t, model, layer->attn_v_b, kvn_t); - - memcpy(k_rot, raw_t + DS4_N_KV_LORA, (size_t)DS4_N_ROT * sizeof(k_rot[0])); - rope_tail_layer_inplace(k_rot, 1, DS4_N_ROT, DS4_N_ROT, - pos0 + t, il, false); - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - float *kd = kc_t + (uint64_t)h * qk_dim; - memcpy(kd, kn_t + (uint64_t)h * q_nope, - (size_t)q_nope * sizeof(kd[0])); - memcpy(kd + q_nope, k_rot, (size_t)DS4_N_ROT * sizeof(kd[0])); - } - } - - const float scale = 1.0f / sqrtf((float)qk_dim); - for (uint32_t t = 0; t < n_tok; t++) { - const uint32_t visible = t + 1u; - for (uint32_t h = 0; h < DS4_N_HEAD; h++) { - const float *q_h = q + ((uint64_t)t * DS4_N_HEAD + h) * qk_dim; - float max_score = -FLT_MAX; - for (uint32_t s = 0; s < visible; s++) { - const float *k_h = - key_cache + ((uint64_t)s * DS4_N_HEAD + h) * qk_dim; - float dot = 0.0f; - for (uint32_t i = 0; i < qk_dim; i++) dot += q_h[i] * k_h[i]; - scores[s] = dot * scale; - if (scores[s] > max_score) max_score = scores[s]; - } - - float denom = 0.0f; - for (uint32_t s = 0; s < visible; s++) { - scores[s] = expf(scores[s] - max_score); - denom += scores[s]; - } - if (denom < 1.0e-20f) denom = 1.0e-20f; - - float *head_out = - heads + ((uint64_t)t * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; - for (uint32_t d = 0; d < DS4_N_VALUE_MLA; d++) { - float acc = 0.0f; - for (uint32_t s = 0; s < visible; s++) { - const float *v_h = - value_cache + ((uint64_t)s * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; - acc += scores[s] * v_h[d]; - } - head_out[d] = acc / denom; - } - } - - matvec_q8_0_f32_ref(out + (uint64_t)t * DS4_N_EMBD, - model, - layer->attn_output, - heads + (uint64_t)t * heads_dim); - } - - free(scores); - free(k_rot); - free(heads); - free(value_cache); - free(key_cache); - free(k_nope); - free(kv_norm); - free(kv_raw); - free(q); - free(q_rank_norm); - free(q_rank); - free(norm); -} - -static void layer_glm_dense_ffn_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - const uint64_t hidden = layer->ffn_gate->dim[1]; - const uint64_t in_dim = layer->ffn_gate->dim[0]; - const uint64_t blocks = (in_dim + 31) / 32; - float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); - float *up = xmalloc((size_t)hidden * sizeof(up[0])); - float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); - int8_t *xq = xmalloc((size_t)blocks * 32); - float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); - - if (layer->ffn_gate->type != DS4_TENSOR_Q8_0 || - layer->ffn_up->type != DS4_TENSOR_Q8_0 || - layer->ffn_down->type != DS4_TENSOR_Q8_0 || - layer->ffn_up->dim[0] != in_dim || - layer->ffn_up->dim[1] != hidden || - layer->ffn_down->dim[0] != hidden || - layer->ffn_down->dim[1] != DS4_N_EMBD) { - ds4_die("GLM dense FFN tensors have an unexpected layout"); - } - - quantize_q8_0_activation(x, xq, xscale, in_dim); - matvec_q8_0_pair_prequant(gate, up, model, layer->ffn_gate, layer->ffn_up, xq, xscale); - swiglu(mid, gate, up, hidden, 0.0f); - matvec_q8_0(out, model, layer->ffn_down, mid); - - free(xscale); - free(xq); - free(mid); - free(up); - free(gate); -} - -static void layer_glm_dense_ffn_one_f32_ref( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x) { - const uint64_t hidden = layer->ffn_gate->dim[1]; - const uint64_t in_dim = layer->ffn_gate->dim[0]; - float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); - float *up = xmalloc((size_t)hidden * sizeof(up[0])); - float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); - - if (layer->ffn_gate->type != DS4_TENSOR_Q8_0 || - layer->ffn_up->type != DS4_TENSOR_Q8_0 || - layer->ffn_down->type != DS4_TENSOR_Q8_0 || - layer->ffn_up->dim[0] != in_dim || - layer->ffn_up->dim[1] != hidden || - layer->ffn_down->dim[0] != hidden || - layer->ffn_down->dim[1] != DS4_N_EMBD) { - ds4_die("GLM F32 dense FFN reference found unexpected tensor layout"); - } - - matvec_q8_0_f32_ref(gate, model, layer->ffn_gate, x); - matvec_q8_0_f32_ref(up, model, layer->ffn_up, x); - swiglu(mid, gate, up, hidden, 0.0f); - matvec_q8_0_f32_ref(out, model, layer->ffn_down, mid); - - free(mid); - free(up); - free(gate); -} - -static void layer_glm_router_selected_experts( - int selected[DS4_MAX_EXPERT_USED], - float expert_weight[DS4_MAX_EXPERT_USED], - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x) { - float logits[DS4_MAX_EXPERT]; - float probs[DS4_MAX_EXPERT]; - float selection[DS4_MAX_EXPERT]; - const float *bias = tensor_data(model, layer->ffn_exp_probs_b); - - matvec_any(logits, model, layer->ffn_gate_inp, x); - for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { - probs[i] = sigmoid_stable(logits[i]); - selection[i] = probs[i] + bias[i]; - } - - topk_desc(selection, (int)DS4_N_EXPERT, (int)DS4_N_EXPERT_USED, selected); - - float sum = 0.0f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) { - ds4_die("GLM selected expert is outside router range"); - } - expert_weight[i] = probs[selected[i]]; - sum += expert_weight[i]; - } - if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - expert_weight[i] = expert_weight[i] / sum * DS4_EXPERT_WEIGHT_SCALE; - } -} - -typedef struct { - float *mid; - const float *x; - const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; - const uint8_t *up_base[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t out_dim; - uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; - uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; - uint32_t gate_type; - uint32_t up_type; - uint32_t n_expert; -} glm_routed_moe_f32_mid_ctx; - -static bool glm_graph_gate_pair_type_supported(uint32_t gate_type, uint32_t up_type) { - return gate_type == up_type && - (gate_type == DS4_TENSOR_IQ2_XXS || - gate_type == DS4_TENSOR_Q2_K || - gate_type == DS4_TENSOR_Q4_K || - gate_type == DS4_TENSOR_Q5_K); -} - -static bool glm_graph_down_type_supported(uint32_t down_type) { - return down_type == DS4_TENSOR_IQ2_XXS || - down_type == DS4_TENSOR_Q2_K || - down_type == DS4_TENSOR_Q4_K || - down_type == DS4_TENSOR_Q5_K || - down_type == DS4_TENSOR_Q6_K; -} - -static float glm_routed_moe_dot_f32(uint32_t type, int n, const uint8_t *row, const float *x) { - if (type == DS4_TENSOR_IQ2_XXS) { - return ds4_vec_dot_iq2_xxs_f32(n, (const block_iq2_xxs *)row, x); - } - if (type == DS4_TENSOR_Q2_K) { - return ds4_vec_dot_q2_K_f32(n, (const block_q2_K *)row, x); - } - if (type == DS4_TENSOR_Q4_K) { - return ds4_vec_dot_q4_K_f32(n, (const block_q4_K *)row, x); - } - if (type == DS4_TENSOR_Q5_K || type == DS4_TENSOR_Q6_K) { - return ds4_vec_dot_q5_q6_K_f32(type, n, row, x); - } - ds4_die("GLM F32 routed-MoE reference encountered unsupported expert tensor type"); - return 0.0f; -} - -static void glm_routed_moe_f32_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { - glm_routed_moe_f32_mid_ctx *ctx = vctx; - - for (uint64_t idx = row0; idx < row1; idx++) { - const uint32_t slot = (uint32_t)(idx / ctx->out_dim); - const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; - const uint8_t *gate_row = ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]; - const uint8_t *up_row = ctx->up_base[slot] + row * ctx->up_row_bytes[slot]; - const float gate = glm_routed_moe_dot_f32(ctx->gate_type, (int)ctx->in_dim, gate_row, ctx->x); - const float up = glm_routed_moe_dot_f32(ctx->up_type, (int)ctx->in_dim, up_row, ctx->x); - ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; - } -} - -typedef struct { - float *out; - const float *mid; - const uint8_t *down_base[DS4_MAX_EXPERT_USED]; - uint64_t in_dim; - uint64_t out_dim; - uint64_t down_row_bytes[DS4_MAX_EXPERT_USED]; - uint32_t down_type; - uint32_t n_expert; -} glm_routed_moe_f32_down_ctx; - -static void glm_routed_moe_f32_down_worker(void *vctx, uint64_t row0, uint64_t row1) { - glm_routed_moe_f32_down_ctx *ctx = vctx; - - for (uint64_t row = row0; row < row1; row++) { - float acc = 0.0f; - for (uint32_t slot = 0; slot < ctx->n_expert; slot++) { - const uint8_t *down_row = ctx->down_base[slot] + row * ctx->down_row_bytes[slot]; - acc += glm_routed_moe_dot_f32(ctx->down_type, - (int)ctx->in_dim, - down_row, - ctx->mid + (uint64_t)slot * ctx->in_dim); - } - ctx->out[row] = acc; - } -} - -static void layer_glm_routed_moe_one_f32_ref( - float *out, - float *mid_all, - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x, - const int *selected, - const float *expert_weight) { - if (!layer->ffn_gate_exps || !layer->ffn_up_exps || !layer->ffn_down_exps || - !tensor_is_routed_expert_type(layer->ffn_gate_exps->type) || - !tensor_is_routed_expert_type(layer->ffn_up_exps->type) || - !glm_graph_gate_pair_type_supported(layer->ffn_gate_exps->type, - layer->ffn_up_exps->type) || - !glm_graph_down_type_supported(layer->ffn_down_exps->type)) { - ds4_die("GLM F32 routed-MoE reference expects supported matching routed gate/up tensors and down tensors"); - } - - glm_routed_moe_f32_mid_ctx mid_ctx = { - .mid = mid_all, - .x = x, - .gate_type = layer->ffn_gate_exps->type, - .up_type = layer->ffn_up_exps->type, - .n_expert = DS4_N_EXPERT_USED, - }; - glm_routed_moe_f32_down_ctx down_ctx = { - .out = out, - .mid = mid_all, - .down_type = layer->ffn_down_exps->type, - .n_expert = DS4_N_EXPERT_USED, - }; - - uint64_t gate_in0 = 0, gate_out0 = 0; - uint64_t down_in0 = 0, down_out0 = 0; - for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { - uint64_t gate_in, gate_out, up_in, up_out, down_in, down_out; - if (selected[slot] < 0 || (uint32_t)selected[slot] >= DS4_N_EXPERT) { - ds4_die("GLM F32 routed-MoE reference selected expert is outside range"); - } - mid_ctx.gate_base[slot] = - tensor_expert_bytes(model, layer->ffn_gate_exps, (uint32_t)selected[slot], - &gate_in, &gate_out, &mid_ctx.gate_row_bytes[slot]); - mid_ctx.up_base[slot] = - tensor_expert_bytes(model, layer->ffn_up_exps, (uint32_t)selected[slot], - &up_in, &up_out, &mid_ctx.up_row_bytes[slot]); - down_ctx.down_base[slot] = - tensor_expert_bytes(model, layer->ffn_down_exps, (uint32_t)selected[slot], - &down_in, &down_out, &down_ctx.down_row_bytes[slot]); - if (gate_in != up_in || gate_out != up_out || - down_in != gate_out || down_out != DS4_N_EMBD) { - ds4_die("GLM F32 routed-MoE reference found mismatched expert layouts"); - } - if (slot == 0) { - gate_in0 = gate_in; - gate_out0 = gate_out; - down_in0 = down_in; - down_out0 = down_out; - } else if (gate_in != gate_in0 || gate_out != gate_out0 || - down_in != down_in0 || down_out != down_out0) { - ds4_die("GLM F32 routed-MoE reference expert layouts are not uniform"); - } - mid_ctx.expert_weight[slot] = expert_weight[slot]; - } - - if (gate_in0 != DS4_N_EMBD || gate_in0 % QK_K != 0 || - down_in0 != DS4_N_FF_EXP || down_in0 % QK_K != 0 || - gate_out0 != DS4_N_FF_EXP || down_out0 != DS4_N_EMBD) { - ds4_die("GLM F32 routed-MoE reference found unexpected GLM expert dimensions"); - } - - mid_ctx.in_dim = gate_in0; - mid_ctx.out_dim = gate_out0; - down_ctx.in_dim = down_in0; - down_ctx.out_dim = down_out0; - ds4_parallel_for((uint64_t)DS4_N_EXPERT_USED * gate_out0, - glm_routed_moe_f32_mid_worker, - &mid_ctx); - ds4_parallel_for(down_out0, glm_routed_moe_f32_down_worker, &down_ctx); -} - -static void layer_glm_shared_ffn_one_f32_ref( - float *out, - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x) { - const uint64_t in_dim = layer->ffn_gate_shexp ? layer->ffn_gate_shexp->dim[0] : 0; - const uint64_t hidden = layer->ffn_gate_shexp ? layer->ffn_gate_shexp->dim[1] : 0; - if (!layer->ffn_gate_shexp || - !layer->ffn_up_shexp || - !layer->ffn_down_shexp || - layer->ffn_gate_shexp->type != DS4_TENSOR_Q8_0 || - layer->ffn_up_shexp->type != DS4_TENSOR_Q8_0 || - layer->ffn_down_shexp->type != DS4_TENSOR_Q8_0 || - layer->ffn_up_shexp->dim[0] != in_dim || - layer->ffn_up_shexp->dim[1] != hidden || - layer->ffn_down_shexp->dim[0] != hidden || - layer->ffn_down_shexp->dim[1] != DS4_N_EMBD || - in_dim != DS4_N_EMBD || - hidden != DS4_N_FF_EXP) { - ds4_die("GLM F32 shared expert reference found unexpected tensor layout"); - } - - float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); - float *up = xmalloc((size_t)hidden * sizeof(up[0])); - float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); - - matvec_q8_0_f32_ref(gate, model, layer->ffn_gate_shexp, x); - matvec_q8_0_f32_ref(up, model, layer->ffn_up_shexp, x); - swiglu(mid, gate, up, hidden, 0.0f); - matvec_q8_0_f32_ref(out, model, layer->ffn_down_shexp, mid); - - free(mid); - free(up); - free(gate); -} - -static void layer_glm_ffn_one_f32_ref( - float *out, - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x, - uint32_t il) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - - rms_norm_weight(norm, x, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); - if (il < DS4_N_LEADING_DENSE) { - layer_glm_dense_ffn_one_f32_ref(out, model, layer, norm); - } else { - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - float *mid = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid[0])); - float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); - float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); - - layer_glm_router_selected_experts(selected, expert_weight, model, layer, norm); - layer_glm_routed_moe_one_f32_ref(moe, mid, model, layer, norm, - selected, expert_weight); - layer_glm_shared_ffn_one_f32_ref(shared, model, layer, norm); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = moe[i] + shared[i]; - - free(shared); - free(moe); - free(mid); - } - - free(norm); -} - -static void layer_glm_first_token_one_f32_ref( - float *out, - const ds4_model *model, - const ds4_layer_weights *layer, - const float *x, - uint32_t il) { - float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); - float *after_attn = xmalloc((size_t)DS4_N_EMBD * sizeof(after_attn[0])); - float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); - - layer_glm_first_token_attention_one_f32_ref(attn_out, model, layer, x); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) after_attn[i] = x[i] + attn_out[i]; - - layer_glm_ffn_one_f32_ref(ffn_out, model, layer, after_attn, il); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = after_attn[i] + ffn_out[i]; - - free(ffn_out); - free(after_attn); - free(attn_out); -} - -static void forward_glm_first_token_cpu_f32_ref( - float *out_hidden, - const ds4_model *model, - const ds4_weights *weights, - int token) { - float *cur = xmalloc((size_t)DS4_N_EMBD * sizeof(cur[0])); - float *next = xmalloc((size_t)DS4_N_EMBD * sizeof(next[0])); - const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; - - embed_token_any(model, weights, token, cur); - for (uint32_t il = 0; il < normal_layers; il++) { - layer_glm_first_token_one_f32_ref(next, model, &weights->layer[il], cur, il); - float *tmp = cur; - cur = next; - next = tmp; - } - - memcpy(out_hidden, cur, (size_t)DS4_N_EMBD * sizeof(out_hidden[0])); - - free(next); - free(cur); -} - -static void output_logits_glm_one_f32_ref( - float *logits, - const ds4_model *model, - const ds4_weights *weights, - const float *hidden) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - - rms_norm_weight(norm, hidden, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0_f32_ref(logits, model, weights->output, norm); - - free(norm); -} - -static void layer_glm_routed_moe_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il) { - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - float *mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid_all[0])); - block_q8_K *xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0])); - block_q8_K *midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0])); - - if (expert_in_dim != DS4_N_EMBD || expert_in_dim % QK_K != 0 || - down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { - ds4_die("GLM routed expert tensors have an unexpected layout"); - } - - memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); - ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); - layer_glm_router_selected_experts(selected, expert_weight, model, layer, x); - - matvec_experts_mid_prequant(mid_all, model, - layer->ffn_gate_exps, - layer->ffn_up_exps, - xq, - selected, - expert_weight, - DS4_N_EXPERT_USED, - 0.0f); - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, - midq + (uint64_t)i * (down_in_dim / QK_K), - (int64_t)down_in_dim); - } - matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, - selected, DS4_N_EXPERT_USED); - - free(midq); - free(xq); - free(mid_all); - (void)il; -} - -static void layer_glm_sparse_ffn_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il) { - float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); - float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); - - layer_glm_routed_moe_one(moe, model, layer, x, il); - layer_shared_ffn_one(shared, model, layer, x); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = moe[i] + shared[i]; - - free(shared); - free(moe); -} - -static void layer_glm_ffn_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - - rms_norm_weight(norm, x, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); - if (il < DS4_N_LEADING_DENSE) { - layer_glm_dense_ffn_one(out, model, layer, norm); - } else { - layer_glm_sparse_ffn_one(out, model, layer, norm, il); - } - - free(norm); -} - -static void layer_glm_first_token_one( - float * out, - const ds4_model * model, - const ds4_layer_weights * layer, - const float * x, - uint32_t il) { - float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); - float *after_attn = xmalloc((size_t)DS4_N_EMBD * sizeof(after_attn[0])); - float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); - - layer_glm_first_token_attention_one(attn_out, model, layer, x); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) after_attn[i] = x[i] + attn_out[i]; - - layer_glm_ffn_one(ffn_out, model, layer, after_attn, il); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = after_attn[i] + ffn_out[i]; - - free(ffn_out); - free(after_attn); - free(attn_out); -} - -static void forward_glm_first_token_cpu( - float * out_hidden, - const ds4_model * model, - const ds4_weights * weights, - int token) { - float *cur = xmalloc((size_t)DS4_N_EMBD * sizeof(cur[0])); - float *next = xmalloc((size_t)DS4_N_EMBD * sizeof(next[0])); - const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; - - embed_token_any(model, weights, token, cur); - for (uint32_t il = 0; il < normal_layers; il++) { - layer_glm_first_token_one(next, model, &weights->layer[il], cur, il); - float *tmp = cur; - cur = next; - next = tmp; - } - - memcpy(out_hidden, cur, (size_t)DS4_N_EMBD * sizeof(out_hidden[0])); - - free(next); - free(cur); -} - -static void output_logits_glm_one( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - const float * hidden) { - float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); - - rms_norm_weight(norm, hidden, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0(logits, model, weights->output, norm); - - free(norm); -} - -/* Allocation-free logits head for CPU decode. */ -static void output_logits_one_decode_scratch( - float * logits, - const ds4_model * model, - const ds4_weights * weights, - const float * inp_hc, - ds4_cpu_decode_scratch * scratch) { - const uint32_t n_hc = DS4_N_HC; - const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; - - rms_norm_no_weight(scratch->output_flat, inp_hc, hc_dim, DS4_RMS_EPS); - matvec_f16(scratch->output_pre, model, weights->output_hc_fn, scratch->output_flat); - - const float *scale = tensor_data(model, weights->output_hc_scale); - const float *base = tensor_data(model, weights->output_hc_base); - for (uint32_t i = 0; i < n_hc; i++) { - scratch->output_weights[i] = sigmoid_stable(scratch->output_pre[i] * scale[0] + base[i]) + DS4_HC_EPS; - } - - hc_weighted_sum_one(scratch->output_embd, inp_hc, scratch->output_weights, DS4_N_EMBD, n_hc); - rms_norm_weight(scratch->output_norm, scratch->output_embd, - tensor_data(model, weights->output_norm), - DS4_N_EMBD, DS4_RMS_EPS); - matvec_q8_0_decode_scratch(logits, model, weights->output, scratch->output_norm, scratch); -} - -#ifndef DS4_NO_GPU -static int sample_argmax(const float *logits, uint32_t n_vocab); - -/* ========================================================================= - * Metal Reference Comparison Helpers. - * ========================================================================= - * - * These small scalar helpers are used only by diagnostics that compare the C - * reference path with the Metal executor. - */ - -static float max_abs_diff(const float *a, const float *b, uint64_t n) { - float max_diff = 0.0f; - for (uint64_t i = 0; i < n; i++) { - const float diff = fabsf(a[i] - b[i]); - if (diff > max_diff) max_diff = diff; - } - return max_diff; -} - -static float rms_abs_diff(const float *a, const float *b, uint64_t n) { - double ss = 0.0; - for (uint64_t i = 0; i < n; i++) { - const double d = (double)a[i] - (double)b[i]; - ss += d * d; - } - return n ? (float)sqrt(ss / (double)n) : 0.0f; -} - -static uint64_t argmax_f32(const float *x, uint64_t n) { - uint64_t best = 0; - for (uint64_t i = 1; i < n; i++) { - if (x[i] > x[best]) best = i; - } - return best; -} - -#endif - -static void print_vec_stats(const char *name, const float *x, uint64_t n) { - float minv = DS4_POS_INF; - float maxv = DS4_NEG_INF; - double ss = 0.0; - - for (uint64_t i = 0; i < n; i++) { - const float v = x[i]; - if (v < minv) minv = v; - if (v > maxv) maxv = v; - ss += (double)v * v; - } - - printf("%s: min=%g max=%g rms=%g\n", - name, minv, maxv, sqrt(ss / (double)n)); -} - -#ifndef DS4_NO_GPU -/* - * Apple Metal stores the persistent attention-compressed KV cache in F16. The - * compressor still pools, normalizes, RoPEs, and FP8-rounds rows in F32 staging - * before writing the cache, while checkpoints and debug dumps expand back to - * F32 for the stable external format. This is a storage optimization rather - * than a semantic approximation: all Metal attention consumers already run the - * compressed K/V rows through F16 FlashAttention/indexed-attention paths. - */ -#if defined(__APPLE__) -#define DS4_GPU_ATTN_COMP_CACHE_F16 1 -#else -#define DS4_GPU_ATTN_COMP_CACHE_F16 0 -#endif - -#define DS4_GPU_GLM_COMPACT_CACHE_F16 DS4_GPU_ATTN_COMP_CACHE_F16 - -/* ========================================================================= - * Metal Release Graph State. - * ========================================================================= - * - * The release Metal executor owns one fixed set of tensors for single-token - * decode and another for batched prefill. The structure is DS4-specific: - * tensor names follow the model stages rather than generic graph nodes. - */ - -enum { DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS = 64 }; - -typedef struct { - /* Class P — per-tier replicated kernel scratch buffers. - * Each used tier has its own copy; active_tier names the slot the - * current dispatch step reads/writes. Single-tier paths leave - * active_tier == 0; multi-tier dispatch updates active_tier in B6. - * - * Decode hidden-state buffers. A generated token enters as an embedding - * in cur_hc and leaves as logits after all 43 layers update their - * raw/compressed/indexer caches. The hc_pre / hc_post / hc_comb views - * are derived from hc_split per tier (see metal_graph_alloc_raw_cap). */ - ds4_gpu_tensor *cur_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *flat_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *hc_mix_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *hc_split_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *hc_pre_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ - ds4_gpu_tensor *hc_post_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ - ds4_gpu_tensor *hc_comb_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ - ds4_gpu_tensor *attn_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *attn_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *qr_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *qr_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *q_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *kv_raw_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *kv_by_tier[DS4_MAX_GPUS]; - int active_tier; - /* cached engine placement[] (length DS4_N_LAYER + 2) for the - * dispatch loops. NULL in single-tier mode — active_tier stays 0 and - * dispatch wrappers no-op the tier-switch + cross-device copy. The - * pointer aliases e->placement; the engine outlives the graph so this - * is safe. */ - const int *placement; - - /* Persistent KV state. Raw KV is a sliding-window ring per layer. Ratio-4 - * layers also keep an indexer-compressed cache; ratio-128 layers keep only - * the attention-compressed cache. The small state tensors are compressor - * frontiers for the next compressed row, so they must be snapshotted with - * the row counters whenever a checkpoint is saved or partially rewound. */ - ds4_gpu_tensor *layer_raw_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_comp_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_comp_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_raw_cache_tp[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_comp_cache_tp[DS4_MAX_LAYER]; - - /* Speculative decoding scratch. MTP is allowed to mutate graph state only - * if the target verifier can either commit it or restore the saved - * frontiers. The prefix1 buffers are the cheap partial-accept state for the - * common N=2 case. */ - ds4_gpu_tensor *spec_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_index_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_index_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_logits; - uint32_t layer_n_comp[DS4_MAX_LAYER]; - uint32_t layer_n_index_comp[DS4_MAX_LAYER]; - uint32_t spec_prefix1_n_comp[DS4_MAX_LAYER]; - uint32_t spec_prefix1_n_index_comp[DS4_MAX_LAYER]; - bool spec_capture_prefix1; - uint32_t raw_cap; - /* Maximum compressed-row capacity across layers. Shared work buffers use - * this worst-case size because ratio-4 indexer layers can still reach it. */ - uint32_t comp_cap; - /* Persistent compressed caches are per layer, so size them from the actual - * layer compression ratio instead of pessimistically using the ratio-4 cap - * for every ratio-128 layer. */ - uint32_t layer_comp_cap[DS4_MAX_LAYER]; - uint32_t attn_comp_stage_cap; - - /* Class P (per-layer work tensors). Each used tier has its - * own replica. They are reused in place by every layer instead of - * allocating a generic graph arena. This is why the code is verbose but - * predictable: each pointer names an actual DS4 stage. */ - ds4_gpu_tensor *comp_kv_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *comp_sc_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *attn_comp_stage_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *indexer_q_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *indexer_weights_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *indexer_scores_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *comp_mask_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *comp_selected_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *heads_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *attn_low_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *attn_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *after_attn_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *ffn_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *ffn_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *shared_gate_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *shared_up_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *shared_mid_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *shared_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *router_logits_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *router_probs_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *router_selected_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *router_weights_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *routed_gate_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *routed_up_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *routed_mid_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *routed_down_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *routed_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *ffn_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *after_ffn_hc_by_tier[DS4_MAX_GPUS]; - /* Class H — output-head buffers and logits live on the - * head tier only. head_tier is captured at metal_graph_alloc_raw_cap - * time from placement[DS4_N_LAYER + 1] (or 0 in single-tier / - * diagnostic paths). Non-head slots remain NULL. Readers go through - * the metal_graph_logits / metal_graph_output_* accessors below. */ - ds4_gpu_tensor *output_pre_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *output_weights_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *output_embd_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *output_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *logits_by_tier[DS4_MAX_GPUS]; - int head_tier; - - /* DSpark target features. The proposer consumes mean-over-HC rows from - * selected target layers; keeping them on-GPU avoids adding readbacks to - * the target path. */ - ds4_gpu_tensor *dspark_hc_mean_weights; - ds4_gpu_tensor *dspark_hc_mean_rows; - ds4_gpu_tensor *dspark_target_hidden; - ds4_gpu_tensor *dspark_target_hidden_batch; - ds4_gpu_tensor *dspark_stage0_packed; - ds4_gpu_tensor *dspark_stage0_proj; - ds4_gpu_tensor *dspark_main_x; - ds4_gpu_tensor *dspark_draft_tokens; - ds4_gpu_tensor *dspark_draft_hc; - ds4_gpu_tensor *dspark_target_hc; - ds4_gpu_tensor *dspark_stage_input_hc; - ds4_gpu_tensor *dspark_stage_output_hc; - ds4_gpu_tensor *dspark_position_ids; - ds4_gpu_tensor *dspark_raw_cache[DS4_DSPARK_MAX_STAGES]; - uint32_t dspark_cache_cap; - uint32_t dspark_cache_start; - uint32_t dspark_cache_token_start; - uint32_t dspark_cache_len; - uint32_t dspark_target_layer_count; - uint32_t dspark_block_size; - uint32_t dspark_target_layers[DS4_DSPARK_MAX_TARGET_LAYERS]; - uint32_t dspark_capture_mask; - uint32_t dspark_capture_checkpoint_len; - uint32_t dspark_capture_batch_mask; - uint32_t dspark_capture_batch_start; - uint32_t dspark_capture_batch_tokens; - bool dspark_capture_valid; - bool dspark_capture_batch_valid; - int dspark_exec_tier; - bool dspark_capture_enabled; - bool verify_small_batch_tp; - uint32_t pipeline_capture_chunk_start; - uint32_t pipeline_capture_chunk_len; - bool ssd_streaming; /* glm-branch SSD streaming; always false here */ - - /* Optional MTP model state. It has its own raw cache because the drafter - * runs on speculative future tokens; target KV state is updated only after - * verification accepts draft tokens. */ - ds4_gpu_tensor *mtp_embed; - ds4_gpu_tensor *mtp_enorm; - ds4_gpu_tensor *mtp_eproj; - ds4_gpu_tensor *mtp_eproj_hc; - ds4_gpu_tensor *mtp_hnorm_hc; - ds4_gpu_tensor *mtp_hproj_hc; - ds4_gpu_tensor *mtp_input_hc; - ds4_gpu_tensor *mtp_state_hc; - ds4_gpu_tensor *mtp_next_hc; - ds4_gpu_tensor *mtp_raw_cache; - uint32_t mtp_n_raw; - uint32_t prefill_cap; - uint32_t raw_window; - uint32_t batch_token_offset; - - /* Batched prefill tensors. Prefill is layer-major: a chunk of prompt - * tokens moves through layer 0, then layer 1, and so on, updating the same - * persistent caches used by decode. Keeping this separate from decode - * avoids a slow loop of one-token graph steps for long prompts. */ - /* Class E — embedding-tier-only prompt-token integer buffer. - * Captured at metal_graph_alloc_raw_cap time from placement[0] (or 0 in - * single-tier / diagnostic paths). Non-embedding slots stay NULL. Readers - * go through metal_graph_prefill_tokens() below. */ - ds4_gpu_tensor *prefill_tokens_by_tier[DS4_MAX_GPUS]; - int emb_tier; - /* Class P batch (chunked-prefill) scratch — per-tier - * replicated. The cur/next pair is ping-ponged per layer step on the - * layer's active tier; tier transitions copy the active buffer across - * boundaries via ds4_gpu_tensor_copy_xdev (handled in B6). */ - ds4_gpu_tensor *batch_cur_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_next_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_flat_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_hc_mix_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_hc_split_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_attn_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_attn_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_qr_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_qr_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_q_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_kv_raw_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_kv_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_comp_kv_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_comp_sc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_indexer_q_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_indexer_weights_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_heads_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_attn_low_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_attn_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_group_tmp_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_low_tmp_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_after_attn_hc_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_ffn_cur_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_ffn_norm_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_shared_gate_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_shared_up_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_shared_mid_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_shared_out_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_router_logits_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_router_probs_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_router_selected_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_router_weights_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_routed_gate_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_routed_up_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_routed_mid_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_routed_down_by_tier[DS4_MAX_GPUS]; - ds4_gpu_tensor *batch_routed_out_by_tier[DS4_MAX_GPUS]; - bool batch_routed_mid_is_f16; - ds4_gpu_tensor *batch_ffn_out_by_tier[DS4_MAX_GPUS]; - bool owns_prefill_workspace; - bool materialize_ffn_out; - /* Class P (replicated per tier — this is - * consumed in per-layer attn/FFN kernels, NOT embedding-only). Read-only - * after init; replicate by writing the same host directions buffer to - * every used tier's slot during session setup. */ - ds4_gpu_tensor *directional_steering_dirs_by_tier[DS4_MAX_GPUS]; - float directional_steering_attn_scale; - float directional_steering_ffn_scale; - bool cuda_tp_decode; - bool cuda_tp_attn; - bool cuda_tp_attn_peer_read; - bool cuda_tp_attn_heads; - bool cuda_tp_attn_cache_dup; - bool cuda_tp_moe; - bool cuda_tp_ep; - bool cuda_tp_ep_pack_exact; - bool cuda_tp_moe_delay_reduce; - bool cuda_tp_moe_copy3_handoff; - bool cuda_tp_moe_pack_handoff; - bool cuda_tp_moe_peer_read; - bool cuda_tp_moe_peer_router; - bool cuda_tp_shared; - bool cuda_tp_shared_fold; - bool cuda_tp_q; - bool cuda_tp_output; - bool cuda_tp_prefill_ffn; - bool cuda_tp_prefill_attn_output; - bool cuda_q_norm_rope_fuse; - bool cuda_qkv_kv_rope_fuse; - bool cuda_qkv_pair; - bool cuda_tp_attn_out_hc_fuse; - bool shared_gate_up_swiglu_fuse; - bool decode_stage_profile; - bool decode_index_stage_profile; - bool output_stage_profile; - ds4_gpu_tensor *tp_peer_tmp_by_tier[DS4_MAX_GPUS]; - uint32_t power_percent; - double prefill_layer_avg_sec[DS4_MAX_LAYER]; - double decode_token_avg_sec; - bool quality; - bool mtp_enabled; - /* Metal-only prefill helpers retained alongside the CUDA tiered workspace. */ - ds4_gpu_tensor *batch_q_half; - ds4_gpu_tensor *prefill_seed_router_selected; - uint32_t prefill_seed_tokens; - uint64_t prefill_selected_profile_rows; - uint64_t prefill_selected_profile_unique; - uint64_t prefill_selected_profile_selected_bytes; - uint64_t prefill_selected_profile_full_bytes; - uint32_t prefill_selected_profile_layers; - uint32_t prefill_selected_profile_min_unique; - uint32_t prefill_selected_profile_max_unique; - uint32_t streaming_preload_experts; - bool ssd_streaming_cold; - bool streaming_static_decode_map_current; - float *cpu_router_norm; - - /* Metal network tensor parallelism. These views alias engine-owned - * transport slabs except tp_logits_half, whose view object is session-owned. */ - uint32_t tp_world; - uint32_t tp_rank; - ds4_gpu_tensor **tp_out; - ds4_gpu_tensor **tp_in; - ds4_gpu_tensor **tp_batch_out; - ds4_gpu_tensor **tp_batch_in; - uint32_t tp_batch_rows; - ds4_gpu_tensor *tp_zero; - ds4_gpu_tensor *tp_logits_half; -} ds4_gpu_graph; - -/* Tensors that are temporary for chunked prefill and grouped multi-session - * decode. The batched server serializes every operation that uses them, so one - * engine-owned set can be aliased by all resident session graphs. */ -#define DS4_GPU_PREFILL_WORKSPACE_FIELDS(X) \ - X(prefill_tokens) \ - X(batch_ffn_out) \ - X(batch_routed_out) \ - X(batch_routed_down) \ - X(batch_routed_mid) \ - X(batch_routed_up) \ - X(batch_routed_gate) \ - X(batch_router_weights) \ - X(batch_router_selected) \ - X(batch_router_probs) \ - X(batch_router_logits) \ - X(batch_shared_out) \ - X(batch_shared_mid) \ - X(batch_shared_up) \ - X(batch_shared_gate) \ - X(batch_ffn_norm) \ - X(batch_ffn_cur) \ - X(batch_after_attn_hc) \ - X(batch_low_tmp) \ - X(batch_group_tmp) \ - X(batch_attn_out) \ - X(batch_attn_low) \ - X(batch_heads) \ - X(batch_indexer_weights) \ - X(batch_indexer_q) \ - X(batch_comp_sc) \ - X(batch_comp_kv) \ - X(batch_kv) \ - X(batch_kv_raw) \ - X(batch_q) \ - X(batch_qr_norm) \ - X(batch_qr) \ - X(batch_attn_norm) \ - X(batch_attn_cur) \ - X(batch_hc_split) \ - X(batch_hc_mix) \ - X(batch_flat_hc) \ - X(batch_next_hc) \ - X(batch_cur_hc) - -/* Class H accessors. All reader sites for the output-head - * tensors and the final logits route through these inlines, which read the - * head_tier slot captured at allocation time. Single-tier paths set - * head_tier == 0 and the slot is byte-identical to the legacy - * metal_graph_logits(g) / g->output_* pointers. Multi-tier paths set head_tier - * to placement[DS4_N_LAYER + 1]; other tier slots remain NULL. */ -static inline ds4_gpu_tensor *metal_graph_logits(const ds4_gpu_graph *g) { - return g->logits_by_tier[g->head_tier]; -} -static inline ds4_gpu_tensor *metal_graph_output_pre(const ds4_gpu_graph *g) { - return g->output_pre_by_tier[g->head_tier]; -} -static inline ds4_gpu_tensor *metal_graph_output_weights(const ds4_gpu_graph *g) { - return g->output_weights_by_tier[g->head_tier]; -} -static inline ds4_gpu_tensor *metal_graph_output_embd(const ds4_gpu_graph *g) { - return g->output_embd_by_tier[g->head_tier]; -} -static inline ds4_gpu_tensor *metal_graph_output_norm(const ds4_gpu_graph *g) { - return g->output_norm_by_tier[g->head_tier]; -} - -/* Class E accessor. The prompt-token integer buffer is - * consumed by the embedding kernel on the embedding tier only. Single-tier - * paths set emb_tier == 0 (byte-equivalent to the legacy single-tier - * pointer). Multi-tier paths set emb_tier = placement[0]. */ -static inline ds4_gpu_tensor *metal_graph_prefill_tokens(const ds4_gpu_graph *g) { - return g->prefill_tokens_by_tier[g->emb_tier]; -} - -/* Class P accessors. Each Class P kernel-scratch buffer is - * replicated across every tier the placement uses; the active_tier field - * names the slot the current dispatch step reads/writes. Single-tier paths - * leave active_tier == 0 (byte-equivalent to the legacy single-tier - * pointer). Multi-tier dispatch (wired up in B6) updates active_tier with - * the current layer's home tier before each kernel-dispatch wrapper runs. */ -#define DS4_GPU_GRAPH_CLASS_P_ACCESSOR(name) \ -static inline ds4_gpu_tensor *metal_graph_##name(const ds4_gpu_graph *g) { \ - return g->name##_by_tier[g->active_tier]; \ -} - -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(cur_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(flat_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_mix) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_split) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_pre) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_post) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_comb) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(qr) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(qr_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(q) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(kv_raw) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(kv) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_kv_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_sc_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_comp_stage) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_q) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_weights) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_scores) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_mask) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_selected) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(heads) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_low) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(after_attn_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_gate) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_up) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_mid) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_logits) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_probs) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_selected) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_weights) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_gate) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_up) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_mid) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_down) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(after_ffn_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_cur_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_next_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_flat_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_hc_mix) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_hc_split) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_qr) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_qr_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_q) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_kv_raw) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_kv) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_comp_kv) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_comp_sc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_indexer_q) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_indexer_weights) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_heads) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_low) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_group_tmp) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_low_tmp) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_after_attn_hc) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_cur) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_norm) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_gate) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_up) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_mid) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_logits) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_probs) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_selected) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_weights) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_gate) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_up) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_mid) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_down) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_out) -DS4_GPU_GRAPH_CLASS_P_ACCESSOR(directional_steering_dirs) - -/* dispatch-loop helpers for multi-tier per-layer execution. - * - * Single-tier (g->placement == NULL): all helpers are no-ops; active_tier - * stays 0 from memset; behavior is byte-equivalent to legacy. - * - * Multi-tier: each helper switches g->active_tier to the requested tier - * BEFORE the next kernel-dispatch wrapper reads any Class P accessor. If - * the source-tier Class P cur_hc (or batch_cur_hc) differs from the new - * tier's, ds4_gpu_tensor_copy_xdev ferries the active hidden state across - * the boundary. copy_xdev returns 1 on success, 0 on failure. The - * destination tensor's device_id was stamped at alloc_on time and is - * immutable. - * - * For decode (one token at a time): metal_graph_set_active_tier_decode - * swaps to the requested tier and copies cur_hc across the boundary. - * - * For batch (chunked prefill): metal_graph_set_active_tier_batch swaps - * tier and copies batch_cur_hc across the boundary. The next/cur pair - * is maintained per tier — after a copy, batch_next_hc on the destination - * tier becomes the swap target for the next layer step on that tier. - * - * Helpers always invoke ds4_gpu_set_current_device(tier) so the next - * kernel-launch sees the correct CUDA device. */ - -/* ds4_gpu_set_current_device is declared in ds4_gpu_mgpu.h — single-tier - * (g_n_gpus <= 1) callers no-op. Returns 0 on success. */ - -#ifdef DS4_NO_GPU -static inline int ds4_gpu_set_current_device(int tier) { (void)tier; return 0; } -static inline int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst, - const ds4_gpu_tensor *src, - uint64_t bytes) { - (void)dst; (void)src; (void)bytes; return 1; -} -static inline 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) { - (void)dst0; (void)src0; (void)bytes0; - (void)dst1; (void)src1; (void)bytes1; - (void)dst2; (void)src2; (void)bytes2; - return 1; -} -static inline int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst, - const ds4_gpu_tensor *src, - uint64_t bytes) { - (void)dst; (void)src; (void)bytes; return 1; -} -static inline int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) { - (void)src; (void)dst_tier; return 1; -} -static inline int ds4_gpu_moe_handoff_pack_tensor( - ds4_gpu_tensor *packed, - const ds4_gpu_tensor *ffn_norm, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_embd, - uint32_t n_expert) { - (void)packed; (void)ffn_norm; (void)selected; (void)weights; - (void)n_embd; (void)n_expert; return 1; -} -static inline int ds4_gpu_q8_cache_suppressed(void) { return 0; } -static inline void ds4_gpu_set_q8_cache_suppressed(int suppressed) { (void)suppressed; } -static inline int ds4_gpu_set_decode_fast_attention(int enabled) { - (void)enabled; - return 0; -} -static inline int ds4_gpu_set_decode_score_vec4(int enabled) { - (void)enabled; - return 0; -} -static inline 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) { - (void)q_out; (void)q; (void)model_map; (void)model_size; - (void)q_weight_offset; (void)q_n; (void)kv_out; (void)kv; - (void)kv_weight_offset; (void)kv_n; (void)rows; (void)kv_n_head; - (void)kv_head_dim; (void)n_rot; (void)pos0; (void)n_ctx_orig; - (void)inverse; (void)freq_base; (void)freq_scale; (void)ext_factor; - (void)attn_factor; (void)beta_fast; (void)beta_slow; (void)eps; - return 0; -} -#endif - -/* Returns true on success. Single-tier: no-op success. Multi-tier: - * sets the CUDA device, then if tier differs from current active_tier, - * copies cur_hc to the destination tier and updates active_tier. */ -static bool metal_graph_set_active_tier_decode(ds4_gpu_graph *g, int tier) { - if (!g->placement) { - /* Single-tier: just keep active_tier at 0; no device switch needed. */ - (void)tier; - return true; - } - if (tier < 0 || tier >= DS4_MAX_GPUS) return false; - if (tier == g->active_tier) return true; - if (ds4_gpu_set_current_device(tier) != 0) return false; - /* Boundary hop: copy cur_hc from source-tier to destination-tier slot. */ - if (g->active_tier >= 0) { - ds4_gpu_tensor *src = g->cur_hc_by_tier[g->active_tier]; - ds4_gpu_tensor *dst = g->cur_hc_by_tier[tier]; - if (src && dst) { - const uint64_t hc_bytes = (uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float); - if (!ds4_gpu_tensor_copy_xdev(dst, src, hc_bytes)) return false; - } - } - g->active_tier = tier; - return true; -} - -/* Returns true on success. Same semantics as the decode helper but ferries - * batch_cur_hc (which contains chunk_tokens * hc_dim floats — variable per - * prefill call). The caller passes the chunk size in tokens; single-tier - * paths ignore the argument. */ -static bool metal_graph_set_active_tier_batch(ds4_gpu_graph *g, int tier, uint32_t chunk_tokens) { - if (!g->placement) { - (void)tier; - (void)chunk_tokens; - return true; - } - if (tier < 0 || tier >= DS4_MAX_GPUS) return false; - if (tier == g->active_tier) return true; - if (ds4_gpu_set_current_device(tier) != 0) return false; - if (g->active_tier >= 0) { - ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[g->active_tier]; - ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[tier]; - if (src && dst) { - const uint64_t hc_bytes = - (uint64_t)chunk_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float); - if (!ds4_gpu_tensor_copy_xdev(dst, src, hc_bytes)) return false; - } - } - g->active_tier = tier; - return true; -} - -static bool metal_graph_set_active_tier_no_copy(ds4_gpu_graph *g, int tier) { - if (!g->placement) { - (void)tier; - return true; - } - if (tier < 0 || tier >= DS4_MAX_GPUS) return false; - if (tier == g->active_tier) return true; - if (ds4_gpu_set_current_device(tier) != 0) return false; - g->active_tier = tier; - return true; -} - -/* Upstream: --power N GPU duty-cycle throttling helpers. The single-tier - * --power=100 path is a no-op; multi-tier inherits the same helpers via - * graph_power_note_prefill_layer / graph_power_note_decode_token which we - * call from the shared encode / decode loops. */ - -static bool graph_power_throttle_enabled(const ds4_gpu_graph *g) { - return g && g->power_percent > 0 && g->power_percent < 100; -} - -static double graph_power_update_avg(double avg, double sample) { - if (sample <= 0.0 || !isfinite(sample)) return avg; - if (avg <= 0.0 || !isfinite(avg)) return sample; - return avg * 0.875 + sample * 0.125; -} - -static void graph_power_sleep(double work_sec, uint32_t power_percent) { - if (power_percent == 0 || power_percent >= 100) return; - /* Target duty cycle: work / (work + sleep) = power / 100. - * At --power 50 this sleeps for one measured work interval; at 25 it - * sleeps for three. */ - const double sleep = work_sec * (100.0 - (double)power_percent) / - (double)power_percent; - sleep_sec(sleep); -} - -static void graph_power_note_prefill_layer(ds4_gpu_graph *g, - uint32_t il, - double elapsed_sec) { - if (!graph_power_throttle_enabled(g)) return; - if (il >= DS4_N_LAYER) return; - g->prefill_layer_avg_sec[il] = - graph_power_update_avg(g->prefill_layer_avg_sec[il], elapsed_sec); - graph_power_sleep(g->prefill_layer_avg_sec[il], g->power_percent); -} - -static void graph_power_note_decode_token(ds4_gpu_graph *g, double elapsed_sec) { - if (!graph_power_throttle_enabled(g)) return; - g->decode_token_avg_sec = - graph_power_update_avg(g->decode_token_avg_sec, elapsed_sec); - graph_power_sleep(g->decode_token_avg_sec, g->power_percent); -} - -static void metal_graph_copy_prefill_workspace_pointers( - ds4_gpu_graph *dst, - const ds4_gpu_graph *src) { -#define DS4_COPY_PREFILL_FIELD(name) \ - memcpy(dst->name##_by_tier, src->name##_by_tier, \ - sizeof(dst->name##_by_tier)); - DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_COPY_PREFILL_FIELD) -#undef DS4_COPY_PREFILL_FIELD - dst->batch_q_half = src->batch_q_half; - dst->prefill_seed_router_selected = src->prefill_seed_router_selected; -} - -static void metal_graph_transfer_prefill_workspace( - ds4_gpu_graph *dst, - ds4_gpu_graph *src) { - memset(dst, 0, sizeof(*dst)); - dst->prefill_cap = src->prefill_cap; - dst->emb_tier = src->emb_tier; - dst->owns_prefill_workspace = true; - metal_graph_copy_prefill_workspace_pointers(dst, src); - src->owns_prefill_workspace = false; -} - -static uint64_t metal_graph_prefill_workspace_bytes(const ds4_gpu_graph *g) { - uint64_t total = 0; - for (int t = 0; t < DS4_MAX_GPUS; t++) { -#define DS4_COUNT_PREFILL_FIELD(name) \ - total += ds4_gpu_tensor_bytes(g->name##_by_tier[t]); - DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_COUNT_PREFILL_FIELD) -#undef DS4_COUNT_PREFILL_FIELD - } - total += ds4_gpu_tensor_bytes(g->batch_q_half); - total += ds4_gpu_tensor_bytes(g->prefill_seed_router_selected); - return total; -} - -static void metal_graph_free_prefill_workspace(ds4_gpu_graph *g) { - if (!g || !g->owns_prefill_workspace) return; - for (int t = 0; t < DS4_MAX_GPUS; t++) { -#define DS4_FREE_PREFILL_FIELD(name) \ - ds4_gpu_tensor_free(g->name##_by_tier[t]); \ - g->name##_by_tier[t] = NULL; - DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_FREE_PREFILL_FIELD) -#undef DS4_FREE_PREFILL_FIELD - } - ds4_gpu_tensor_free(g->batch_q_half); - ds4_gpu_tensor_free(g->prefill_seed_router_selected); - g->batch_q_half = NULL; - g->prefill_seed_router_selected = NULL; - g->owns_prefill_workspace = false; -} - -/* Release every Metal tensor owned by the whole-model graph runtime. */ -static void metal_graph_free(ds4_gpu_graph *g) { - /* free every Class P slot across all DS4_MAX_GPUS tier - * slots. Unallocated slots are NULL and ds4_gpu_tensor_free(NULL) is a - * no-op. The hc_pre / hc_post / hc_comb views must be freed BEFORE - * their parent hc_split — view destruction releases its own struct - * but does not touch the parent's memory. */ - metal_graph_free_prefill_workspace(g); - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(g->directional_steering_dirs_by_tier[t]); - g->directional_steering_dirs_by_tier[t] = NULL; - } - /* Class H free across all tier slots. Non-head slots are - * NULL and ds4_gpu_tensor_free(NULL) is a no-op. */ - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(g->logits_by_tier[t]); - g->logits_by_tier[t] = NULL; - } - ds4_gpu_tensor_free(g->mtp_raw_cache); - ds4_gpu_tensor_free(g->mtp_next_hc); - ds4_gpu_tensor_free(g->mtp_state_hc); - ds4_gpu_tensor_free(g->mtp_input_hc); - ds4_gpu_tensor_free(g->mtp_hproj_hc); - ds4_gpu_tensor_free(g->mtp_hnorm_hc); - ds4_gpu_tensor_free(g->mtp_eproj_hc); - ds4_gpu_tensor_free(g->mtp_eproj); - ds4_gpu_tensor_free(g->mtp_enorm); - ds4_gpu_tensor_free(g->mtp_embed); - ds4_gpu_tensor_free(g->spec_logits); - /* Class H output-head free across all tier slots. */ - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(g->output_norm_by_tier[t]); - g->output_norm_by_tier[t] = NULL; - ds4_gpu_tensor_free(g->output_embd_by_tier[t]); - g->output_embd_by_tier[t] = NULL; - ds4_gpu_tensor_free(g->output_weights_by_tier[t]); - g->output_weights_by_tier[t] = NULL; - ds4_gpu_tensor_free(g->output_pre_by_tier[t]); - g->output_pre_by_tier[t] = NULL; - } - /* Class P decode scratch + routed-FFN free across all - * tier slots. ffn_out is also a Class P field freed here. */ - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(g->after_ffn_hc_by_tier[t]); - ds4_gpu_tensor_free(g->ffn_out_by_tier[t]); - ds4_gpu_tensor_free(g->routed_out_by_tier[t]); - ds4_gpu_tensor_free(g->routed_down_by_tier[t]); - ds4_gpu_tensor_free(g->routed_mid_by_tier[t]); - ds4_gpu_tensor_free(g->routed_up_by_tier[t]); - ds4_gpu_tensor_free(g->routed_gate_by_tier[t]); - ds4_gpu_tensor_free(g->tp_peer_tmp_by_tier[t]); - ds4_gpu_tensor_free(g->router_weights_by_tier[t]); - ds4_gpu_tensor_free(g->router_selected_by_tier[t]); - ds4_gpu_tensor_free(g->router_probs_by_tier[t]); - ds4_gpu_tensor_free(g->router_logits_by_tier[t]); - ds4_gpu_tensor_free(g->shared_out_by_tier[t]); - ds4_gpu_tensor_free(g->shared_mid_by_tier[t]); - ds4_gpu_tensor_free(g->shared_up_by_tier[t]); - ds4_gpu_tensor_free(g->shared_gate_by_tier[t]); - ds4_gpu_tensor_free(g->ffn_norm_by_tier[t]); - ds4_gpu_tensor_free(g->ffn_cur_by_tier[t]); - ds4_gpu_tensor_free(g->after_attn_hc_by_tier[t]); - ds4_gpu_tensor_free(g->attn_out_by_tier[t]); - ds4_gpu_tensor_free(g->attn_low_by_tier[t]); - ds4_gpu_tensor_free(g->heads_by_tier[t]); - ds4_gpu_tensor_free(g->comp_sc_cur_by_tier[t]); - ds4_gpu_tensor_free(g->comp_kv_cur_by_tier[t]); - ds4_gpu_tensor_free(g->attn_comp_stage_by_tier[t]); - ds4_gpu_tensor_free(g->comp_mask_by_tier[t]); - ds4_gpu_tensor_free(g->comp_selected_by_tier[t]); - ds4_gpu_tensor_free(g->indexer_scores_by_tier[t]); - ds4_gpu_tensor_free(g->indexer_weights_by_tier[t]); - ds4_gpu_tensor_free(g->indexer_q_by_tier[t]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_raw_cache[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_raw_cache_tp[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_attn_comp_cache[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_attn_comp_cache_tp[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_attn_state_kv[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_attn_state_score[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_index_comp_cache[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_index_state_kv[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_index_state_score[il]); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - ds4_gpu_tensor_free(g->spec_attn_state_kv[il]); - ds4_gpu_tensor_free(g->spec_attn_state_score[il]); - ds4_gpu_tensor_free(g->spec_index_state_kv[il]); - ds4_gpu_tensor_free(g->spec_index_state_score[il]); - ds4_gpu_tensor_free(g->spec_prefix1_attn_state_kv[il]); - ds4_gpu_tensor_free(g->spec_prefix1_attn_state_score[il]); - ds4_gpu_tensor_free(g->spec_prefix1_index_state_kv[il]); - ds4_gpu_tensor_free(g->spec_prefix1_index_state_score[il]); - } - /* Class P decode-step scratch + decode HC group free across - * all tier slots. hc_pre / hc_post / hc_comb are VIEWS of hc_split — free - * them before hc_split so the view struct release happens with the parent - * still pointer-valid (view free does not touch parent memory). */ - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(g->kv_by_tier[t]); - ds4_gpu_tensor_free(g->kv_raw_by_tier[t]); - ds4_gpu_tensor_free(g->q_by_tier[t]); - ds4_gpu_tensor_free(g->qr_norm_by_tier[t]); - ds4_gpu_tensor_free(g->qr_by_tier[t]); - ds4_gpu_tensor_free(g->attn_norm_by_tier[t]); - ds4_gpu_tensor_free(g->attn_cur_by_tier[t]); - ds4_gpu_tensor_free(g->hc_comb_by_tier[t]); - ds4_gpu_tensor_free(g->hc_post_by_tier[t]); - ds4_gpu_tensor_free(g->hc_pre_by_tier[t]); - ds4_gpu_tensor_free(g->hc_split_by_tier[t]); - ds4_gpu_tensor_free(g->hc_mix_by_tier[t]); - ds4_gpu_tensor_free(g->flat_hc_by_tier[t]); - ds4_gpu_tensor_free(g->cur_hc_by_tier[t]); - } - ds4_gpu_tensor_free(g->dspark_position_ids); - ds4_gpu_tensor_free(g->dspark_stage_output_hc); - ds4_gpu_tensor_free(g->dspark_stage_input_hc); - ds4_gpu_tensor_free(g->dspark_target_hc); - ds4_gpu_tensor_free(g->dspark_draft_hc); - ds4_gpu_tensor_free(g->dspark_draft_tokens); - for (uint32_t stage = 0; stage < DS4_DSPARK_MAX_STAGES; stage++) { - ds4_gpu_tensor_free(g->dspark_raw_cache[stage]); - } - ds4_gpu_tensor_free(g->dspark_main_x); - ds4_gpu_tensor_free(g->dspark_stage0_proj); - ds4_gpu_tensor_free(g->dspark_stage0_packed); - ds4_gpu_tensor_free(g->dspark_target_hidden_batch); - ds4_gpu_tensor_free(g->dspark_target_hidden); - ds4_gpu_tensor_free(g->dspark_hc_mean_rows); - ds4_gpu_tensor_free(g->dspark_hc_mean_weights); - ds4_gpu_tensor_free(g->tp_logits_half); - free(g->cpu_router_norm); - memset(g, 0, sizeof(*g)); -} - -static bool metal_tensor_fill_f32(ds4_gpu_tensor *t, float v, uint64_t n) { - return ds4_gpu_tensor_fill_f32(t, v, n) != 0; -} - -/* ========================================================================= - * Directional Steering. - * ========================================================================= - * - * A steering file contains one normalized 4096-wide direction per layer. When - * enabled, the Metal graph edits selected block outputs in-place: - * - * y = y - scale * v * dot(v, y) - * - * Positive scales remove the represented direction from the activation. - * Negative scales add it. This is deliberately explicit and opt-in; with zero - * scales, the release graph does not allocate the direction tensor and follows - * the normal inference path. - */ - -/* directional_steering_dirs is Class P — replicated per tier. - * The same host directions buffer is written to every tier slot the engine's - * placement uses, then the load buffer is freed. Read-only after init, so - * the per-tier replicas stay byte-identical and never re-sync. */ -static bool metal_graph_load_directional_steering( - ds4_gpu_graph *g, - const char *path, - float attn_scale, - float ffn_scale) { - if (attn_scale == 0.0f && ffn_scale == 0.0f) return true; - - if (!path || !path[0]) { - fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); - return false; - } - - const uint64_t n = (uint64_t)DS4_N_LAYER * DS4_N_EMBD; - float *dirs = xmalloc((size_t)n * sizeof(dirs[0])); - bool ok = read_f32_binary_file(path, dirs, n); - if (ok) { - /* Replicate the directions buffer onto every Class P tier slot that - * has any other Class P scratch allocated (used_tier marker is the - * presence of g->cur_hc_by_tier[t]). Single-tier: only slot 0. */ - bool any = false; - for (int t = 0; ok && t < DS4_MAX_GPUS; t++) { - if (!g->cur_hc_by_tier[t]) continue; - g->directional_steering_dirs_by_tier[t] = - ds4_gpu_tensor_alloc_ptr_on(t, n * sizeof(dirs[0])); - ok = g->directional_steering_dirs_by_tier[t] != NULL && - ds4_gpu_tensor_write(g->directional_steering_dirs_by_tier[t], - 0, dirs, n * sizeof(dirs[0])) != 0; - if (ok) any = true; - } - if (ok && !any) { - /* No used tiers — graph not allocated yet. This shouldn't happen - * given the call site ordering, but bail cleanly. */ - ok = false; - } - } - free(dirs); - - if (!ok) { - fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); - return false; - } - g->directional_steering_attn_scale = attn_scale; - g->directional_steering_ffn_scale = ffn_scale; - fprintf(stderr, "ds4: directional steering enabled: %s attn=%g ffn=%g\n", - path, (double)attn_scale, (double)ffn_scale); - return true; -} - -static bool metal_graph_directional_steering_attn_enabled(const ds4_gpu_graph *g) { - return g && metal_graph_directional_steering_dirs(g) && - g->directional_steering_attn_scale != 0.0f; -} - -static bool metal_graph_directional_steering_ffn_enabled(const ds4_gpu_graph *g) { - return g && metal_graph_directional_steering_dirs(g) && - g->directional_steering_ffn_scale != 0.0f; -} - -static bool metal_graph_apply_directional_steering( - ds4_gpu_graph *g, - ds4_gpu_tensor *x, - uint32_t il, - uint32_t rows, - float scale) { - if (!g || !metal_graph_directional_steering_dirs(g) || scale == 0.0f) return true; - return ds4_gpu_directional_steering_project_tensor(x, - metal_graph_directional_steering_dirs(g), - il, - DS4_N_EMBD, - rows, - scale) != 0; -} - -static bool metal_graph_apply_directional_steering_attn( - ds4_gpu_graph *g, - ds4_gpu_tensor *x, - uint32_t il, - uint32_t rows) { - return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_attn_scale : 0.0f); -} - -static bool metal_graph_apply_directional_steering_ffn( - ds4_gpu_graph *g, - ds4_gpu_tensor *x, - uint32_t il, - uint32_t rows) { - return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f); -} - -static bool metal_graph_configure_dspark_capture( - ds4_gpu_graph *g, - const ds4_dspark_weights *dw) { - if (!g || !dw || dw->target_layer_count == 0) return true; - if (dw->target_layer_count > DS4_DSPARK_MAX_TARGET_LAYERS || - DS4_N_HC == 0 || - DS4_N_HC > DS4_MAX_HC) { - return false; - } - - g->dspark_hc_mean_weights = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); - g->dspark_hc_mean_rows = - ds4_gpu_tensor_alloc((uint64_t)g->prefill_cap * - DS4_N_HC * sizeof(float)); - g->dspark_target_hidden = - ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * - DS4_N_EMBD * sizeof(float)); - g->dspark_target_hidden_batch = - ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * - g->prefill_cap * - DS4_N_EMBD * sizeof(float)); - if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { - g->dspark_stage0_packed = - ds4_gpu_tensor_alloc(((uint64_t)dw->block_size + 1u) * - dw->target_layer_count * - DS4_N_EMBD * sizeof(float)); - } - g->dspark_stage0_proj = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->dspark_main_x = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - if (!g->dspark_hc_mean_weights || !g->dspark_hc_mean_rows || - !g->dspark_target_hidden || !g->dspark_target_hidden_batch || - !g->dspark_stage0_proj || !g->dspark_main_x) { - return false; - } - if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - g->dspark_draft_tokens = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * sizeof(int32_t)); - g->dspark_draft_hc = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * hc_dim * sizeof(float)); - g->dspark_target_hc = - ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->dspark_stage_input_hc = - ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * - hc_dim * sizeof(float)); - g->dspark_stage_output_hc = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * - hc_dim * sizeof(float)); - g->dspark_position_ids = - ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * - sizeof(int32_t)); - if (!g->dspark_draft_tokens || !g->dspark_draft_hc || - !g->dspark_target_hc || !g->dspark_stage_input_hc || - !g->dspark_stage_output_hc || !g->dspark_position_ids) { - return false; - } - if (dw->n_stages != 0 && g->raw_cap != 0) { - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - g->dspark_raw_cache[stage] = - ds4_gpu_tensor_alloc((uint64_t)g->raw_cap * - DS4_N_HEAD_DIM * sizeof(float)); - if (!g->dspark_raw_cache[stage]) return false; - } - g->dspark_cache_cap = g->raw_cap; - g->dspark_cache_start = 0; - g->dspark_cache_token_start = 0; - g->dspark_cache_len = 0; - } - g->dspark_block_size = dw->block_size; - } - - float mean[DS4_MAX_HC] = {0}; - const float inv_hc = 1.0f / (float)DS4_N_HC; - for (uint32_t i = 0; i < DS4_N_HC; i++) mean[i] = inv_hc; - if (ds4_gpu_tensor_write(g->dspark_hc_mean_weights, - 0, - mean, - (uint64_t)DS4_N_HC * sizeof(mean[0])) == 0) { - return false; - } - const uint64_t mean_rows_count = (uint64_t)g->prefill_cap * DS4_N_HC; - if (mean_rows_count == 0 || mean_rows_count > (uint64_t)SIZE_MAX / sizeof(float)) { - return false; - } - float *mean_rows = xmalloc((size_t)mean_rows_count * sizeof(mean_rows[0])); - for (uint64_t i = 0; i < mean_rows_count; i++) mean_rows[i] = inv_hc; - const bool mean_rows_ok = - ds4_gpu_tensor_write(g->dspark_hc_mean_rows, - 0, - mean_rows, - mean_rows_count * sizeof(mean_rows[0])) != 0; - free(mean_rows); - if (!mean_rows_ok) return false; - - g->dspark_target_layer_count = dw->target_layer_count; - memcpy(g->dspark_target_layers, - dw->target_layers, - (size_t)dw->target_layer_count * sizeof(g->dspark_target_layers[0])); - g->dspark_capture_mask = 0; - g->dspark_capture_checkpoint_len = 0; - g->dspark_capture_batch_mask = 0; - g->dspark_capture_batch_start = 0; - g->dspark_capture_batch_tokens = 0; - g->dspark_capture_valid = false; - g->dspark_capture_batch_valid = false; - g->dspark_capture_enabled = true; - return true; -} - -static uint64_t metal_graph_kv_cache_bytes_for_context(uint32_t ctx_size, uint32_t raw_cap) { - uint64_t bytes = (uint64_t)DS4_N_LAYER * - raw_cap * - DS4_N_HEAD_DIM * - sizeof(float); - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio == 0) continue; - const uint64_t comp_cap = (uint64_t)(ctx_size / ratio + 2u); - bytes += comp_cap * DS4_N_HEAD_DIM * - (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); - if (ratio == 4) { - bytes += comp_cap * DS4_N_INDEXER_HEAD_DIM * sizeof(float); - } - } - return bytes; -} - -static uint64_t metal_graph_context_bytes_for_kv_policy( - uint32_t ctx_size, - uint32_t raw_cap, - uint32_t prefill_cap, - uint64_t *kv_cache_bytes_out) { - uint32_t min_ratio = UINT32_MAX; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; - } - if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; - uint64_t comp_cap = (uint64_t)(ctx_size / min_ratio + 2u); - if (comp_cap < 2u) comp_cap = 2u; - const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context(ctx_size, raw_cap); - if (kv_cache_bytes_out) *kv_cache_bytes_out = kv_cache_bytes; - uint64_t bytes = kv_cache_bytes + - 2ull * comp_cap * prefill_cap * sizeof(float); - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - uint64_t attn_stage_cap = (uint64_t)(prefill_cap / min_ratio + 2u); - if (attn_stage_cap < 2u) attn_stage_cap = 2u; - bytes += attn_stage_cap * DS4_N_HEAD_DIM * sizeof(float); - } - return bytes; -} - -static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor_on( - bool managed, - int tier, - uint64_t bytes) { - if (g_n_gpus <= 1) { - return managed ? ds4_gpu_tensor_alloc_managed(bytes) - : ds4_gpu_tensor_alloc(bytes); - } - return managed ? ds4_gpu_tensor_alloc_managed_on(tier, bytes) - : ds4_gpu_tensor_alloc_ptr_on(tier, bytes); -} - -static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor(bool managed, uint64_t bytes) { - return metal_graph_alloc_kv_cache_tensor_on(managed, 0, bytes); -} - -/* ========================================================================= - * Metal Diagnostic Dump Hooks. - * ========================================================================= - * - * The release path calls these after important stages, but they are no-ops - * unless DS4_METAL_GRAPH_DUMP_PREFIX or DS4_ROCM_GRAPH_DUMP_PREFIX is set. - * Dumping synchronizes and restarts the command batch, so it is intentionally - * isolated here. - */ - -typedef struct { - int init; - const char *prefix; - const char *name; - int layer_set; - uint32_t layer; - int pos_set; - uint32_t pos; -} metal_graph_debug_config; - -static const metal_graph_debug_config *metal_graph_debug_get_config(void) { - static metal_graph_debug_config cfg; - if (!cfg.init) { - cfg.init = 1; - cfg.prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", - "DS4_METAL_GRAPH_DUMP_PREFIX"); - if (cfg.prefix && !cfg.prefix[0]) cfg.prefix = NULL; - cfg.name = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_NAME", - "DS4_METAL_GRAPH_DUMP_NAME"); - if (cfg.name && !cfg.name[0]) cfg.name = NULL; - - const char *layer_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_LAYER", - "DS4_METAL_GRAPH_DUMP_LAYER"); - if (layer_env && layer_env[0] && strcmp(layer_env, "all") != 0) { - cfg.layer_set = 1; - cfg.layer = (uint32_t)strtoul(layer_env, NULL, 10); - } - - const char *pos_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_POS", - "DS4_METAL_GRAPH_DUMP_POS"); - if (pos_env && pos_env[0]) { - cfg.pos_set = 1; - cfg.pos = (uint32_t)strtoul(pos_env, NULL, 10); - } - } - return &cfg; -} - -static const char *metal_graph_debug_prefix_for(const char *name, uint32_t il, uint32_t pos) { - const metal_graph_debug_config *cfg = metal_graph_debug_get_config(); - if (!cfg->prefix) return NULL; - if (cfg->name && strstr(cfg->name, name) == NULL) return NULL; - if (cfg->layer_set && cfg->layer != il) return NULL; - if (cfg->pos_set && cfg->pos != pos) return NULL; - return cfg->prefix; -} - -static bool metal_graph_debug_wants(const char *name, uint32_t il, uint32_t pos) { - return metal_graph_debug_prefix_for(name, il, pos) != NULL; -} - -static void metal_graph_debug_dump_tensor( - const char *name, - ds4_gpu_tensor *t, - uint64_t n_f32, - uint32_t il, - uint32_t pos) { - const char *prefix = metal_graph_debug_prefix_for(name, il, pos); - if (glm_graph_env_present("DS4_ROCM_GRAPH_DUMP_TRACE", - "DS4_METAL_GRAPH_DUMP_TRACE")) - fprintf(stderr, "ds4: dump? name=%s il=%u pos=%u t=%p n=%llu wants=%d\n", - name, il, pos, (void *)t, (unsigned long long)n_f32, - metal_graph_debug_wants(name, il, pos)); - if (!t || n_f32 == 0 || !metal_graph_debug_wants(name, il, pos)) return; - - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); - return; - } - - float *buf = xmalloc((size_t)n_f32 * sizeof(buf[0])); - if (ds4_gpu_tensor_read(t, 0, buf, n_f32 * sizeof(buf[0])) != 0) { - char path[1024]; - snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); - if (write_f32_binary_file(path, buf, n_f32)) { - fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); - } - } - free(buf); - - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); - } -} - -static void metal_graph_debug_dump_f16_tensor( - const char *name, - ds4_gpu_tensor *t, - uint64_t n_f16, - uint32_t il, - uint32_t pos) { - const char *prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", - "DS4_METAL_GRAPH_DUMP_PREFIX"); - if (!t || n_f16 == 0 || !metal_graph_debug_wants(name, il, pos)) return; - - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); - return; - } - - uint16_t *hbuf = xmalloc((size_t)n_f16 * sizeof(hbuf[0])); - float *fbuf = xmalloc((size_t)n_f16 * sizeof(fbuf[0])); - if (ds4_gpu_tensor_read(t, 0, hbuf, n_f16 * sizeof(hbuf[0])) != 0) { - for (uint64_t i = 0; i < n_f16; i++) fbuf[i] = f16_to_f32(hbuf[i]); - char path[1024]; - snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); - if (write_f32_binary_file(path, fbuf, n_f16)) { - fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); - } - } - free(fbuf); - free(hbuf); - - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); - } -} - -static void metal_graph_debug_dump_i32_tensor( - const char *name, - ds4_gpu_tensor *t, - uint64_t n_i32, - uint32_t il, - uint32_t pos) { - if (!t || n_i32 == 0) return; - const char *prefix = metal_graph_debug_prefix_for(name, il, pos); - if (!prefix) return; - - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); - return; - } - - int32_t *buf = xmalloc((size_t)n_i32 * sizeof(buf[0])); - if (ds4_gpu_tensor_read(t, 0, buf, n_i32 * sizeof(buf[0])) != 0) { - char path[1024]; - snprintf(path, sizeof(path), "%s_%s-%u_pos%u.i32", prefix, name, il, pos); - FILE *fp = fopen(path, "wb"); - if (fp) { - if (fwrite(buf, sizeof(buf[0]), (size_t)n_i32, fp) == (size_t)n_i32) { - fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); - } - fclose(fp); - } - } - free(buf); - - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); - } -} - -static bool metal_graph_needs_ffn_out(const ds4_gpu_graph *g, uint32_t il, uint32_t pos) { - return metal_graph_directional_steering_ffn_enabled(g) || - g->materialize_ffn_out || - metal_graph_debug_wants("ffn_out", il, pos); -} - -/* tier-aware lazy allocator. The Class P ffn_out scratch is - * created on demand the first time a layer that materializes ffn_out runs - * on a tier; subsequent visits to the same tier reuse the existing slot. - * Single-tier paths: active_tier == 0 always, behavior unchanged. */ -static bool metal_graph_ensure_ffn_out(ds4_gpu_graph *g) { - const int t = g->active_tier; - if (!g->ffn_out_by_tier[t]) { - g->ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( - t, (uint64_t)DS4_N_EMBD * sizeof(float)); - } - return g->ffn_out_by_tier[t] != NULL; -} - -static bool metal_graph_ensure_batch_ffn_out_on(ds4_gpu_graph *g, int t) { - if (t < 0 || t >= DS4_MAX_GPUS) return false; - if (!g->batch_ffn_out_by_tier[t]) { - g->batch_ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( - t, (uint64_t)g->prefill_cap * DS4_N_EMBD * sizeof(float)); - } - return g->batch_ffn_out_by_tier[t] != NULL; -} - -static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) { - return metal_graph_ensure_batch_ffn_out_on(g, g->active_tier); -} - -static bool metal_graph_tp_env_flag(const char *name, bool dflt) { - const char *env = getenv(name); - if (!env || !env[0]) return dflt; - return strcmp(env, "0") != 0; -} - -static bool metal_graph_cuda_tp_attn_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN", true); -#endif -} - -static bool metal_graph_cuda_tp_attn_peer_read_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_PEER_READ", true); -#endif -} - -static bool metal_graph_cuda_tp_attn_heads_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_HEADS", false); -#endif -} - -static bool metal_graph_cuda_tp_attn_cache_dup_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_CACHE_DUP", false); -#endif -} - -static bool metal_graph_cuda_tp_moe_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE", true); -#endif -} - -static bool metal_graph_cuda_tp_ep_pack_exact_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_PACK_EXACT", true); -#endif -} - -static bool metal_graph_cuda_tp_ep_direct_return_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DIRECT_RETURN", true); -#endif -} - -static bool metal_graph_cuda_tp_ep_delay_reduce_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DELAY_REDUCE", true); -#endif -} - -static bool metal_graph_cuda_tp_ep_fused_hc_reduce_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_HC_REDUCE", true); -#endif -} - -static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_fused_shared_mid_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_SHARED_MID", true); -#endif -} - -static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_balanced_shared_mid_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag( - "DS4_CUDA_TP_EP_BALANCED_SHARED_MID", true); -#endif -} - -static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_dual_prequant_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag( - "DS4_CUDA_TP_EP_DUAL_PREQUANT", true); -#endif -} - -static bool metal_graph_cuda_tp_moe_delay_reduce_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_DELAY_REDUCE", true); -#endif -} - -static bool metal_graph_cuda_tp_moe_pack_handoff_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PACK", false); -#endif -} - -static bool metal_graph_cuda_tp_moe_copy3_handoff_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_COPY3", false); -#endif -} - -static bool metal_graph_cuda_tp_moe_peer_read_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_READ", false); -#endif -} - -static bool metal_graph_cuda_tp_moe_peer_router_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_ROUTER", false); -#endif -} - -static bool metal_graph_cuda_tp_shared_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED", false); -#endif -} - -static bool metal_graph_cuda_tp_shared_fold_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED_FOLD", true); -#endif -} - -static bool metal_graph_cuda_tp_q_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_Q", false); -#endif -} - -static bool metal_graph_cuda_tp_output_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_OUTPUT", true); -#endif -} - -static bool metal_graph_cuda_greedy_split_top1_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLIT_TOP1"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLIT_TOP1", false); -#endif -} - -static bool metal_graph_cuda_output_fused_top1_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_OUTPUT_FUSED_TOP1", false); -#endif -} - -static bool metal_graph_cuda_verify_decode2_split_top1_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_VERIFY_DECODE2_SPLIT_TOP1"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_VERIFY_DECODE2_SPLIT_TOP1", false); -#endif -} - -static bool metal_graph_cuda_greedy_splitkv_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV", false); -#endif -} - -static bool metal_graph_cuda_greedy_vec4_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_VEC4", false); -#endif -} - -static bool metal_graph_cuda_splitkv_spec_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC", false); -#endif -} - -static bool metal_graph_cuda_splitkv_spec_toponly_row0_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_TOPONLY_ROW0"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_TOPONLY_ROW0", false); -#endif -} - -static bool metal_graph_cuda_splitkv_spec_batch_verify_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_BATCH_VERIFY"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_BATCH_VERIFY", false); -#endif -} - -static float metal_graph_cuda_greedy_vec4_margin_threshold(void) { -#if defined(__APPLE__) - return 0.0f; -#else - const char *env = getenv("DS4_CUDA_GREEDY_VEC4_MARGIN"); - if (env && env[0]) { - char *end = NULL; - double v = strtod(env, &end); - while (end && isspace((unsigned char)*end)) end++; - if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { - return (float)v; - } - fprintf(stderr, - "ds4: invalid DS4_CUDA_GREEDY_VEC4_MARGIN=%s; using 0.25\n", - env); - } - return 0.25f; -#endif -} - -static bool metal_graph_cuda_greedy_vec4_fallback_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4_FALLBACK"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_cuda_greedy_vec4_margin_threshold() > 0.0f; -#endif -} - -static float metal_graph_cuda_greedy_splitkv_margin_threshold(void) { -#if defined(__APPLE__) - return 0.0f; -#else - const char *env = getenv("DS4_CUDA_GREEDY_SPLITKV_MARGIN"); - if (env && env[0]) { - char *end = NULL; - double v = strtod(env, &end); - while (end && isspace((unsigned char)*end)) end++; - if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { - return (float)v; - } - fprintf(stderr, - "ds4: invalid DS4_CUDA_GREEDY_SPLITKV_MARGIN=%s; using 0.25\n", - env); - } - return 0.25f; -#endif -} - -static bool metal_graph_cuda_greedy_splitkv_fallback_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_FALLBACK"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_cuda_greedy_splitkv_margin_threshold() > 0.0f; -#endif -} - -static bool metal_graph_cuda_greedy_splitkv_top2_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_TOP2"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TOP2", true); -#endif -} - -static bool metal_graph_cuda_greedy_splitkv_trust_replay_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TRUST_REPLAY", false); -#endif -} - -static bool metal_graph_cuda_greedy_splitkv_pair_replay_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_PAIR_REPLAY"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_PAIR_REPLAY", false); -#endif -} - -static DS4_MAYBE_UNUSED uint32_t metal_graph_cuda_greedy_max_segment(const char *name) { - const char *env = getenv(name); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - while (end && isspace((unsigned char)*end)) end++; - if (end != env && end && *end == '\0' && v <= INT32_MAX) { - return (uint32_t)v; - } - fprintf(stderr, - "ds4: invalid %s=%s; expected 0..%d, using disabled\n", - name, - env, - INT32_MAX); - } - return 0; -} - -static uint32_t metal_graph_cuda_greedy_splitkv_max_segment(void) { -#if defined(__APPLE__) - return 0; -#else - return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_SPLITKV_MAX_SEGMENT"); -#endif -} - -static uint32_t metal_graph_cuda_greedy_vec4_max_segment(void) { -#if defined(__APPLE__) - return 0; -#else - return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_VEC4_MAX_SEGMENT"); -#endif -} - -static uint32_t metal_graph_cuda_greedy_splitkv_min_score(void) { - const char *env = getenv("DS4_CUDA_SPLITKV_MIN_SCORE"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - while (end && isspace((unsigned char)*end)) end++; - if (end != env && end && *end == '\0' && v <= UINT32_MAX) { - return (uint32_t)v; - } - } - return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_DECODE", false) ? 0u : 512u; -} - -static bool metal_graph_cuda_q_norm_rope_fuse_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_Q_NORM_ROPE_FUSE", true); -#endif -} - -static bool metal_graph_cuda_qkv_kv_rope_fuse_requested(void) { -#if defined(__APPLE__) - return false; -#else - const char *no = getenv("DS4_CUDA_NO_QKV_KV_ROPE_FUSE"); - if (no && no[0] && strcmp(no, "0") != 0) return false; - if (getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL) return false; - return metal_graph_tp_env_flag("DS4_CUDA_QKV_KV_ROPE_FUSE", true); -#endif -} - -static bool metal_graph_cuda_tp_prefill_ffn_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_FFN", true); -#endif -} - -static bool metal_graph_cuda_tp_prefill_attn_output_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_ATTN_OUTPUT", true); -#endif -} - -static bool metal_graph_cuda_prefill_pipeline_requested(const ds4_gpu_graph *g) { -#if defined(__APPLE__) - (void)g; - return false; -#else - const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE"); - if (env && env[0]) return strcmp(env, "0") != 0; - return g && g->cuda_tp_decode; -#endif -} - -static bool metal_graph_cuda_prefill_pipeline_q8_cache_requested(void) { -#if defined(__APPLE__) - return false; -#else - return metal_graph_tp_env_flag("DS4_CUDA_PREFILL_PIPELINE_Q8_CACHE", false); -#endif -} - -static uint32_t metal_graph_cuda_prefill_pipeline_microbatch(void) { - const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE_MB"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && v > 0 && v <= UINT32_MAX) return (uint32_t)v; - } - return 512; -} - -static int metal_graph_cuda_tp_partner_tier(int tier) { - if (g_n_gpus < 2 || (g_n_gpus & 1) != 0) return -1; - const int half = g_n_gpus / 2; - if (tier < 0 || tier >= half) return -1; - return tier + half; -} - -static uint32_t metal_graph_cuda_tp_output_tiers( - const ds4_gpu_graph *g, - int tiers[DS4_MAX_GPUS]) { - if (!g) return 0; - return metal_graph_cuda_tp_output_tiers_for_head(g->head_tier, - g->cuda_tp_output, - g_n_gpus, - tiers); -} - -static uint64_t metal_graph_q8_0_row_bytes(uint64_t in_dim) { - return ((in_dim + 31u) / 32u) * 34u; -} - -/* ========================================================================= - * Metal Release Graph Allocation. - * ========================================================================= */ - -/* Allocate the Metal graph state for a chosen raw-cache capacity. The model - * weights are not copied here; tensors reference the mapped GGUF. - * - * tier-aware per-layer allocation. - * placement: when non-NULL, an array of DS4_N_LAYER + 2 logical tiers - * (embedding, per-layer..., head). The per-layer KV / state allocations - * in this function use placement[il + 1] as the home tier for each - * layer il. When NULL (single-tier callers, diagnostic paths), all - * per-layer allocations land on tier 0 — byte-equivalent to legacy. - * - * Single-tier (g_n_gpus <= 1) is byte-equivalent regardless of placement, - * because metal_graph_alloc_kv_cache_tensor_on short-circuits to the - * legacy 1-arg helpers when g_n_gpus <= 1. */ -static bool metal_graph_alloc_raw_cap( - ds4_gpu_graph *g, - const ds4_weights *weights, - const ds4_layer_weights *layer, - uint32_t raw_cap, - uint32_t ctx_size, - uint32_t prefill_cap, - bool enable_mtp, - const int *placement, - bool cuda_tensor_parallel, - const ds4_gpu_graph *shared_prefill_workspace) { - const int saved_dspark_exec_tier = g->dspark_exec_tier; - memset(g, 0, sizeof(*g)); - g->dspark_exec_tier = saved_dspark_exec_tier; - g->owns_prefill_workspace = shared_prefill_workspace == NULL; - g->cpu_router_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])); - g->active_tier = placement ? -1 : 0; - /* cache placement on the graph so the dispatch loops can - * walk it without threading the engine pointer through every - * kernel-dispatch wrapper. NULL in single-tier callers (placement - * was already NULL on entry). */ - g->placement = placement; - g->cuda_tp_decode = placement && cuda_tensor_parallel; - g->cuda_tp_attn = g->cuda_tp_decode && metal_graph_cuda_tp_attn_requested(); - g->cuda_tp_attn_peer_read = metal_graph_cuda_tp_attn_peer_read_requested(); - g->cuda_tp_attn_heads = g->cuda_tp_decode && metal_graph_cuda_tp_attn_heads_requested(); - g->cuda_tp_attn_cache_dup = g->cuda_tp_attn_heads && - metal_graph_cuda_tp_attn_cache_dup_requested(); - g->cuda_tp_moe = g->cuda_tp_decode && metal_graph_cuda_tp_moe_requested(); - g->cuda_tp_ep = g->cuda_tp_moe && cuda_tensor_parallel; - g->cuda_tp_ep_pack_exact = - g->cuda_tp_ep && metal_graph_cuda_tp_ep_pack_exact_requested(); - g->cuda_tp_moe_delay_reduce = metal_graph_cuda_tp_moe_delay_reduce_requested(); - g->cuda_tp_moe_copy3_handoff = metal_graph_cuda_tp_moe_copy3_handoff_requested(); - g->cuda_tp_moe_pack_handoff = metal_graph_cuda_tp_moe_pack_handoff_requested(); - g->cuda_tp_moe_peer_read = metal_graph_cuda_tp_moe_peer_read_requested(); - g->cuda_tp_moe_peer_router = metal_graph_cuda_tp_moe_peer_router_requested(); - g->cuda_tp_shared = g->cuda_tp_decode && metal_graph_cuda_tp_shared_requested(); - g->cuda_tp_shared_fold = metal_graph_cuda_tp_shared_fold_requested(); - g->cuda_tp_q = g->cuda_tp_decode && metal_graph_cuda_tp_q_requested(); - g->cuda_tp_output = g->cuda_tp_decode && metal_graph_cuda_tp_output_requested(); - g->cuda_tp_prefill_ffn = g->cuda_tp_decode && metal_graph_cuda_tp_prefill_ffn_requested(); - g->cuda_tp_prefill_attn_output = - g->cuda_tp_decode && metal_graph_cuda_tp_prefill_attn_output_requested(); - g->cuda_q_norm_rope_fuse = metal_graph_cuda_q_norm_rope_fuse_requested(); - g->cuda_qkv_kv_rope_fuse = metal_graph_cuda_qkv_kv_rope_fuse_requested(); - g->cuda_qkv_pair = getenv("DS4_CUDA_NO_QKV_PAIR") == NULL; - g->cuda_tp_attn_out_hc_fuse = - getenv("DS4_CUDA_TP_ATTN_OUT_HC_FUSE") != NULL && - getenv("DS4_CUDA_NO_TP_ATTN_OUT_HC_FUSE") == NULL; - g->shared_gate_up_swiglu_fuse = - getenv("DS4_METAL_DISABLE_SHARED_GATE_UP_SWIGLU_FUSION") == NULL; - g->decode_stage_profile = getenv("DS4_METAL_DECODE_STAGE_PROFILE") != NULL; - g->decode_index_stage_profile = getenv("DS4_METAL_INDEXER_STAGE_PROFILE") != NULL; - g->output_stage_profile = getenv("DS4_METAL_OUTPUT_STAGE_PROFILE") != NULL; - const bool enable_splitkv_spec = metal_graph_cuda_splitkv_spec_requested(); - const bool enable_splitkv_batch_verify = - enable_splitkv_spec && metal_graph_cuda_splitkv_spec_batch_verify_requested(); - const bool enable_spec_logits = enable_mtp || enable_splitkv_batch_verify; - const bool enable_prefix1_snapshot = enable_mtp || enable_splitkv_spec; - const bool enable_frontier_snapshot = - enable_mtp || - enable_splitkv_spec || - (metal_graph_cuda_greedy_splitkv_requested() && - metal_graph_cuda_greedy_splitkv_fallback_requested()) || - (metal_graph_cuda_greedy_vec4_requested() && - metal_graph_cuda_greedy_vec4_fallback_requested()); - if (g->cuda_tp_decode && metal_graph_cuda_tp_partner_tier(0) < 0) { - fprintf(stderr, - "ds4: CUDA tensor parallelism requires an even multi-GPU placement; " - "have %d GPU tiers\n", - g_n_gpus); - metal_graph_free(g); - return false; - } - if (g->cuda_tp_ep && - (g_ds4_shape.family != DS4_MODEL_FAMILY_DEEPSEEK4 || - (DS4_N_EXPERT & 1u) != 0u)) { - fprintf(stderr, - "ds4: CUDA tensor parallelism requires an even-expert DeepSeek model\n"); - metal_graph_free(g); - return false; - } - if (g->cuda_tp_ep) { - fprintf(stderr, - "ds4: CUDA routed MoE expert ownership enabled " - "(half-resident decode and prefill)\n"); - } - g->mtp_enabled = enable_mtp; - if (raw_cap == 0) raw_cap = 1; - if (ctx_size == 0) ctx_size = raw_cap; - if (prefill_cap == 0) prefill_cap = 1; - uint32_t raw_window = DS4_N_SWA; - if (raw_window > ctx_size) raw_window = ctx_size; - if (raw_window == 0) raw_window = 1; - if (raw_cap < raw_window) raw_cap = raw_window; - if (raw_cap > ctx_size) raw_cap = ctx_size; - if (raw_cap == 0) raw_cap = 1; - g->raw_cap = raw_cap; - g->raw_window = raw_window; - g->prefill_cap = prefill_cap; - uint32_t min_ratio = UINT32_MAX; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (!weights_layer_has_required(&weights->layer[il], il)) continue; - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; - } - if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; - g->comp_cap = ctx_size / min_ratio + 2u; - if (g->comp_cap < 2u) g->comp_cap = 2u; - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - g->attn_comp_stage_cap = prefill_cap / min_ratio + 2u; - if (g->attn_comp_stage_cap < 2u) g->attn_comp_stage_cap = 2u; - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (!weights_layer_has_required(&weights->layer[il], il)) { - g->layer_comp_cap[il] = 0; - continue; - } - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio == 0) { - g->layer_comp_cap[il] = 0; - } else { - g->layer_comp_cap[il] = ctx_size / ratio + 2u; - if (g->layer_comp_cap[il] < 2u) g->layer_comp_cap[il] = 2u; - } - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); - const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; - const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; - /* Distributed coordinators do not normally own the output head. The - * logits workspace still has a fixed model-vocabulary shape, while the - * actual head is encoded only on a node that bound its tensors. */ - const uint64_t vocab_dim = - weights->output ? weights->output->dim[1] : DS4_N_VOCAB; - const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM - ? DS4_N_HEAD_DIM - : DS4_N_INDEXER_HEAD_DIM); - const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - const uint64_t pc = prefill_cap; - uint64_t kv_cache_bytes = 0; - const uint64_t context_bytes = - metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes); - const bool managed_kv_cache = - ds4_gpu_should_use_managed_kv_cache(kv_cache_bytes, context_bytes) != 0; - if (managed_kv_cache) { - /* - * CUDA device allocations are fastest, but a million-token KV cache is - * large enough to starve DGX Spark's unified CPU/GPU memory once the - * model cache and driver allocations are present. For this one - * long-lived cache class, managed memory restores the old demand-paged - * behavior. It can be slower, but it keeps oversized contexts from - * turning memory pressure into a machine-wide lockup. - */ - fprintf(stderr, - "ds4: CUDA using managed KV cache for ctx=%u " - "(kv cache %.2f GiB, context buffers %.2f GiB); " - "this may degrade performance but is needed for very large contexts\n", - ctx_size, - (double)kv_cache_bytes / 1073741824.0, - (double)context_bytes / 1073741824.0); - } - - /* Class P decode HC scratch — replicated across every tier - * the placement uses (per-tier kernel-scratch). Single-tier path - * (placement == NULL) collapses to tier 0 only; _ptr_on(0, ...) short- - * circuits to legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte- - * equivalent. The hc_pre/hc_post/hc_comb buffers are VIEWS of hc_split - * and therefore allocated per tier alongside their parent. */ - bool used_tier[DS4_MAX_GPUS] = {0}; - used_tier[0] = true; /* single-tier baseline always uses tier 0 */ - if (placement) { - for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER + 2u; i++) { - const int p = placement[i]; - if (p >= 0 && p < DS4_MAX_GPUS) used_tier[p] = true; - } - } - if (g->cuda_tp_decode) { - const int half = g_n_gpus / 2; - for (int t = half; t < g_n_gpus; t++) { - if (used_tier[t]) { - fprintf(stderr, - "ds4: CUDA tensor parallelism expects layer homes in lower-half " - "tiers; placement already uses tier %d\n", - t); - metal_graph_free(g); - return false; - } - } - for (int t = 0; t < half; t++) { - if (used_tier[t]) used_tier[t + half] = true; - } - fprintf(stderr, - "ds4: CUDA decode TP enabled: pairing lower-half tiers with " - "upper-half tiers\n"); - } - for (int t = 0; t < DS4_MAX_GPUS; t++) { - if (!used_tier[t]) continue; - g->cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); - g->flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); - g->hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); - g->hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); - g->hc_pre_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], - 0, - (uint64_t)DS4_N_HC * sizeof(float)); - g->hc_post_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], - (uint64_t)DS4_N_HC * sizeof(float), - (uint64_t)DS4_N_HC * sizeof(float)); - g->hc_comb_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], - 2ull * DS4_N_HC * sizeof(float), - (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); - g->attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); - g->qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); - g->q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_dim * sizeof(float)); - g->kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - g->kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - } - bool state_init_ok = true; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - /* A distributed process owns only its bound layer slice. Persistent - * KV state must follow that ownership just like the model tensors; - * allocating every model layer here defeats split-model residency. */ - if (!weights_layer_has_required(&weights->layer[il], il)) continue; - /* per-layer Class L allocations land on the layer's - * home tier. placement is NULL on single-tier / diagnostic paths - * (all-tier-0); non-NULL on the engine path that opted into - * multi-tier. layer_tier == 0 in single-tier mode is the - * byte-equivalent path through metal_graph_alloc_kv_cache_tensor_on - * and ds4_gpu_tensor_alloc_ptr_on. */ - const int layer_tier = placement ? placement[il + 1] : 0; - g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor_on( - managed_kv_cache, - layer_tier, - (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); - const int layer_tp_partner = g->cuda_tp_attn_cache_dup - ? metal_graph_cuda_tp_partner_tier(layer_tier) : -1; - if (layer_tp_partner >= 0) { - g->layer_raw_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( - managed_kv_cache, - layer_tp_partner, - (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); - } - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0) { - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; - const uint64_t attn_rows = (uint64_t)coff * ratio; - g->layer_attn_comp_cache[il] = metal_graph_alloc_kv_cache_tensor_on( - managed_kv_cache, - layer_tier, - (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * - (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); - if (layer_tp_partner >= 0) { - g->layer_attn_comp_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( - managed_kv_cache, - layer_tp_partner, - (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * - (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); - } - g->layer_attn_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - g->layer_attn_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - if (enable_frontier_snapshot) { - g->spec_attn_state_kv[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - g->spec_attn_state_score[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - if (enable_prefix1_snapshot) { - g->spec_prefix1_attn_state_kv[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - g->spec_prefix1_attn_state_score[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); - } - } - if (g->layer_attn_state_kv[il]) { - state_init_ok = state_init_ok && - metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows); - } - if (g->layer_attn_state_score[il]) { - state_init_ok = state_init_ok && - metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows); - } - - if (ratio == 4) { - const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM; - const uint64_t index_rows = (uint64_t)coff * ratio; - g->layer_index_comp_cache[il] = metal_graph_alloc_kv_cache_tensor_on( - managed_kv_cache, - layer_tier, - (uint64_t)g->layer_comp_cap[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - g->layer_index_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - g->layer_index_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - if (enable_frontier_snapshot) { - g->spec_index_state_kv[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - g->spec_index_state_score[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - if (enable_prefix1_snapshot) { - g->spec_prefix1_index_state_kv[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - g->spec_prefix1_index_state_score[il] = - ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); - } - } - if (g->layer_index_state_kv[il]) { - state_init_ok = state_init_ok && - metal_tensor_fill_f32(g->layer_index_state_kv[il], 0.0f, index_width * index_rows); - } - if (g->layer_index_state_score[il]) { - state_init_ok = state_init_ok && - metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows); - } - } - } - } - /* Class P per-layer decode scratch + routed-expert state — - * replicated across every used tier. ffn_out is lazily allocated by - * metal_graph_ensure_ffn_out (per-tier on first touch). */ - for (int t = 0; t < DS4_MAX_GPUS; t++) { - if (!used_tier[t]) continue; - g->comp_kv_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, comp_width_max * sizeof(float)); - g->comp_sc_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, comp_width_max * sizeof(float)); - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - /* Upstream's F16-compressed attn staging buffer. Only allocated when - * the F16-cache mode is enabled (the non-F16 path stages in-place). */ - g->attn_comp_stage_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)g->attn_comp_stage_cap * DS4_N_HEAD_DIM * sizeof(float)); - } - g->indexer_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, indexer_q_dim * sizeof(float)); - g->indexer_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float)); - g->indexer_scores_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)g->comp_cap * pc * sizeof(float)); - g->comp_mask_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)g->comp_cap * pc * sizeof(float)); - g->comp_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)(DS4_N_INDEXER_TOP_K ? DS4_N_INDEXER_TOP_K : 1u) * pc * sizeof(uint32_t)); - g->heads_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_dim * sizeof(float)); - g->attn_low_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, low_dim * sizeof(float)); - g->attn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->after_attn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); - g->ffn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->ffn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->shared_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); - g->shared_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); - g->shared_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); - g->shared_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->router_logits_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT * sizeof(float)); - g->router_probs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT * sizeof(float)); - g->router_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT_USED * sizeof(int)); - g->router_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT_USED * sizeof(float)); - g->routed_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_down_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); - g->routed_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); - if (g->cuda_tp_decode) { - g->tp_peer_tmp_by_tier[t] = - ds4_gpu_tensor_alloc_ptr_on(t, DS4_CUDA_TP_PEER_TMP_BYTES); - } - g->after_ffn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); - } - /* Class H — head_tier captured from placement[DS4_N_LAYER + 1] - * (or 0 in single-tier / diagnostic paths). Output-head tensors and the - * final logits buffer allocate on head_tier only; other tier slots stay - * NULL. The _ptr_on(0, ...) path short-circuits to the legacy - * ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte-equivalent. */ - g->head_tier = placement ? placement[DS4_N_LAYER + 1] : 0; - int output_tp_tiers[DS4_MAX_GPUS] = {0}; - const uint32_t output_tp_ways = g->cuda_tp_output - ? metal_graph_cuda_tp_output_tiers(g, output_tp_tiers) : 0; - if (g->cuda_tp_output && output_tp_ways < 2u) { - fprintf(stderr, - "ds4: CUDA output TP requires output head tier %d to be in " - "the lower half of the CUDA placement\n", - g->head_tier); - metal_graph_free(g); - return false; - } - uint64_t output_logits_elems = vocab_dim; - if (enable_spec_logits && output_tp_ways >= 2u) { - const uint64_t max_shard_vocab = - (vocab_dim + output_tp_ways - 1u) / output_tp_ways; - const uint64_t spec_shard_elems = - (uint64_t)DS4_DSPARK_MAX_BLOCK_SIZE * max_shard_vocab; - if (spec_shard_elems > output_logits_elems) { - output_logits_elems = spec_shard_elems; - } - } - g->output_pre_by_tier[g->head_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_HC * sizeof(float)); - g->output_weights_by_tier[g->head_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_HC * sizeof(float)); - g->output_embd_by_tier[g->head_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->output_norm_by_tier[g->head_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_EMBD * sizeof(float)); - g->logits_by_tier[g->head_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->head_tier, - output_logits_elems * sizeof(float)); - for (uint32_t i = 1; i < output_tp_ways; i++) { - const int t = output_tp_tiers[i]; - if (t < 0 || t >= DS4_MAX_GPUS || t == g->head_tier) continue; - g->output_norm_by_tier[t] = - ds4_gpu_tensor_alloc_ptr_on(t, - (uint64_t)DS4_N_EMBD * sizeof(float)); - g->logits_by_tier[t] = - ds4_gpu_tensor_alloc_ptr_on(t, - output_logits_elems * sizeof(float)); - } - /* - * MTP is deliberately outside the normal graph footprint. A session that - * does not opt in with --mtp must allocate and execute exactly the same - * buffers as the plain decoder: no support-model mapping, no draft logits, - * and no MTP scratch hidden behind otherwise unused tensors. - */ - if (enable_mtp) { - g->mtp_embed = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_enorm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_eproj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_eproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_hnorm_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_hproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor( - managed_kv_cache, - (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); - g->mtp_n_raw = 0; - } - if (enable_spec_logits) { - const int spec_tier = - g->dspark_exec_tier > 0 && g->dspark_exec_tier < DS4_MAX_GPUS - ? g->dspark_exec_tier : 0; - g->spec_logits = spec_tier - ? ds4_gpu_tensor_alloc_ptr_on(spec_tier, (uint64_t)16 * DS4_N_VOCAB * sizeof(float)) - : ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); - } - - /* Class E — emb_tier captured from placement[0] (or 0 in - * single-tier / diagnostic paths). _ptr_on(0, ...) short-circuits to the - * legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte-equivalent. */ - g->emb_tier = placement ? placement[0] : 0; - /* Class P chunked-prefill batch scratch — replicated across - * every used tier. The cur/next pair (batch_cur_hc / batch_next_hc) is - * ping-ponged per layer step on each tier; tier transitions copy via - * ds4_gpu_tensor_copy_xdev (handled in B6). batch_ffn_out is lazily - * allocated by metal_graph_ensure_batch_ffn_out (per-tier on first touch) - * and included in the CUDA scratch estimate because TP prefill can use it - * as the combined routed+shared FFN buffer. */ - if (shared_prefill_workspace) { - if (shared_prefill_workspace->prefill_cap < prefill_cap || - shared_prefill_workspace->emb_tier != g->emb_tier) { - fprintf(stderr, - "ds4: shared prefill workspace is incompatible " - "(capacity %u/%u, embedding tier %d/%d)\n", - shared_prefill_workspace->prefill_cap, - prefill_cap, - shared_prefill_workspace->emb_tier, - g->emb_tier); - } else { - metal_graph_copy_prefill_workspace_pointers( - g, shared_prefill_workspace); - } - } else { - g->prefill_tokens_by_tier[g->emb_tier] = - ds4_gpu_tensor_alloc_ptr_on(g->emb_tier, pc * sizeof(int32_t)); - for (int t = 0; t < DS4_MAX_GPUS; t++) { - if (!used_tier[t]) continue; - g->batch_cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); - g->batch_next_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); - g->batch_flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); - g->batch_hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); - g->batch_hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); - g->batch_attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); - g->batch_qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); - g->batch_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); - g->batch_kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); - g->batch_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); - g->batch_comp_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); - g->batch_comp_sc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); - g->batch_indexer_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * indexer_q_dim * sizeof(float)); - g->batch_indexer_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_INDEXER_HEAD * sizeof(float)); - g->batch_heads_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); - g->batch_attn_low_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * low_dim * sizeof(float)); - g->batch_attn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_group_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * group_dim * sizeof(float)); - g->batch_low_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_LORA_O * sizeof(float)); - g->batch_after_attn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); - g->batch_ffn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_ffn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_shared_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); - g->batch_shared_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); - g->batch_shared_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); - g->batch_shared_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - g->batch_router_logits_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); - g->batch_router_probs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); - g->batch_router_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(int)); - g->batch_router_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(float)); - g->batch_routed_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_down_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); - g->batch_routed_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); - } - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - g->batch_q_half = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(uint16_t)); - } - g->prefill_seed_router_selected = ds4_gpu_tensor_alloc( - (uint64_t)DS4_N_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * sizeof(int32_t)); - } - - bool layer_cache_ok = true; - for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { - if (!weights_layer_has_required(&weights->layer[il], il)) continue; - layer_cache_ok = g->layer_raw_cache[il] != NULL; - if (layer_cache_ok && g->cuda_tp_attn_cache_dup) { - layer_cache_ok = g->layer_raw_cache_tp[il] != NULL; - } - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (layer_cache_ok && ratio != 0) { - layer_cache_ok = g->layer_attn_comp_cache[il] != NULL && - (!g->cuda_tp_attn_cache_dup || - g->layer_attn_comp_cache_tp[il] != NULL) && - g->layer_attn_state_kv[il] != NULL && - g->layer_attn_state_score[il] != NULL && - (!enable_frontier_snapshot || - (g->spec_attn_state_kv[il] != NULL && - g->spec_attn_state_score[il] != NULL)) && - (!enable_prefix1_snapshot || - (g->spec_prefix1_attn_state_kv[il] != NULL && - g->spec_prefix1_attn_state_score[il] != NULL)); - } - if (layer_cache_ok && ratio == 4) { - layer_cache_ok = g->layer_index_comp_cache[il] != NULL && - g->layer_index_state_kv[il] != NULL && - g->layer_index_state_score[il] != NULL && - (!enable_frontier_snapshot || - (g->spec_index_state_kv[il] != NULL && - g->spec_index_state_score[il] != NULL)) && - (!enable_prefix1_snapshot || - (g->spec_prefix1_index_state_kv[il] != NULL && - g->spec_prefix1_index_state_score[il] != NULL)); - } - } - - /* Class P validation — check every used tier's slot. */ - bool class_p_ok = true; - for (int t = 0; class_p_ok && t < DS4_MAX_GPUS; t++) { - if (!used_tier[t]) continue; - class_p_ok = - g->cur_hc_by_tier[t] && g->flat_hc_by_tier[t] && g->hc_mix_by_tier[t] && g->hc_split_by_tier[t] && - g->hc_pre_by_tier[t] && g->hc_post_by_tier[t] && g->hc_comb_by_tier[t] && - g->attn_cur_by_tier[t] && g->attn_norm_by_tier[t] && g->qr_by_tier[t] && g->qr_norm_by_tier[t] && - g->q_by_tier[t] && g->kv_raw_by_tier[t] && g->kv_by_tier[t] && - g->comp_kv_cur_by_tier[t] && g->comp_sc_cur_by_tier[t] && - (!DS4_GPU_ATTN_COMP_CACHE_F16 || g->attn_comp_stage_by_tier[t]) && - g->indexer_q_by_tier[t] && g->indexer_weights_by_tier[t] && g->indexer_scores_by_tier[t] && - g->comp_mask_by_tier[t] && g->comp_selected_by_tier[t] && - g->heads_by_tier[t] && g->attn_low_by_tier[t] && g->attn_out_by_tier[t] && - g->after_attn_hc_by_tier[t] && g->ffn_cur_by_tier[t] && g->ffn_norm_by_tier[t] && - g->shared_gate_by_tier[t] && g->shared_up_by_tier[t] && g->shared_mid_by_tier[t] && - g->shared_out_by_tier[t] && - g->router_logits_by_tier[t] && g->router_probs_by_tier[t] && - g->router_selected_by_tier[t] && g->router_weights_by_tier[t] && - g->routed_gate_by_tier[t] && g->routed_up_by_tier[t] && g->routed_mid_by_tier[t] && - g->routed_down_by_tier[t] && g->routed_out_by_tier[t] && - (!g->cuda_tp_decode || g->tp_peer_tmp_by_tier[t]) && - g->after_ffn_hc_by_tier[t] && - g->batch_cur_hc_by_tier[t] && g->batch_next_hc_by_tier[t] && g->batch_flat_hc_by_tier[t] && - g->batch_hc_mix_by_tier[t] && g->batch_hc_split_by_tier[t] && - g->batch_attn_cur_by_tier[t] && g->batch_attn_norm_by_tier[t] && - g->batch_qr_by_tier[t] && g->batch_qr_norm_by_tier[t] && g->batch_q_by_tier[t] && - g->batch_kv_raw_by_tier[t] && g->batch_kv_by_tier[t] && - g->batch_comp_kv_by_tier[t] && g->batch_comp_sc_by_tier[t] && - g->batch_indexer_q_by_tier[t] && g->batch_indexer_weights_by_tier[t] && - g->batch_heads_by_tier[t] && g->batch_attn_low_by_tier[t] && g->batch_attn_out_by_tier[t] && - g->batch_group_tmp_by_tier[t] && g->batch_low_tmp_by_tier[t] && g->batch_after_attn_hc_by_tier[t] && - g->batch_ffn_cur_by_tier[t] && g->batch_ffn_norm_by_tier[t] && - g->batch_shared_gate_by_tier[t] && g->batch_shared_up_by_tier[t] && - g->batch_shared_mid_by_tier[t] && g->batch_shared_out_by_tier[t] && - g->batch_router_logits_by_tier[t] && g->batch_router_probs_by_tier[t] && - g->batch_router_selected_by_tier[t] && g->batch_router_weights_by_tier[t] && - g->batch_routed_gate_by_tier[t] && g->batch_routed_up_by_tier[t] && - g->batch_routed_mid_by_tier[t] && g->batch_routed_down_by_tier[t] && - g->batch_routed_out_by_tier[t]; - } - bool output_tp_ok = true; - for (uint32_t i = 1; i < output_tp_ways; i++) { - const int t = output_tp_tiers[i]; - if (t < 0 || t >= DS4_MAX_GPUS || t == g->head_tier) continue; - output_tp_ok = output_tp_ok && - g->output_norm_by_tier[t] != NULL && - g->logits_by_tier[t] != NULL; - } - const bool ok = state_init_ok && layer_cache_ok && class_p_ok && - /* Class H — validate the head_tier slot - * (single-tier: head_tier == 0, byte-equivalent). */ - metal_graph_output_pre(g) && metal_graph_output_weights(g) && - metal_graph_output_embd(g) && metal_graph_output_norm(g) && - metal_graph_logits(g) && output_tp_ok && - (!enable_mtp || - (g->mtp_embed && g->mtp_enorm && g->mtp_eproj && - g->mtp_eproj_hc && g->mtp_hnorm_hc && g->mtp_hproj_hc && - g->mtp_input_hc && g->mtp_state_hc && g->mtp_next_hc && - g->mtp_raw_cache)) && - (!enable_spec_logits || g->spec_logits) && - /* Class E — validate the emb_tier slot. */ - metal_graph_prefill_tokens(g) && - g->cpu_router_norm && - (!DS4_GPU_ATTN_COMP_CACHE_F16 || g->batch_q_half) && - g->prefill_seed_router_selected; - if (!ok) metal_graph_free(g); - return ok; -} - -static bool metal_graph_alloc( - ds4_gpu_graph *g, - const ds4_weights *weights, - const ds4_layer_weights *layer) { - /* single-tier convenience wrapper; placement=NULL routes - * all per-layer allocations to tier 0. */ - return metal_graph_alloc_raw_cap(g, weights, layer, DS4_N_SWA, DS4_N_SWA, - 1, false, NULL, false, NULL); -} - -static bool metal_graph_install_model_spans( - const ds4_model *model, - const ds4_model_map_span_vec *spans, - const char *label) { - if (!model || !spans || spans->len == 0) return false; - - uint64_t *offsets = xmalloc((size_t)spans->len * sizeof(offsets[0])); - uint64_t *sizes = xmalloc((size_t)spans->len * sizeof(sizes[0])); - for (uint32_t i = 0; i < spans->len; i++) { - offsets[i] = spans->v[i].off; - sizes[i] = spans->v[i].end - spans->v[i].off; - } - - const bool ok = ds4_gpu_set_model_map_spans(model->map, - model->size, - offsets, - sizes, - spans->len, - spans->max_tensor_bytes) != 0; - if (!ok) { - fprintf(stderr, - "ds4: Metal SSD streaming failed to map %s model spans\n", - label ? label : "requested"); - } - free(offsets); - free(sizes); - return ok; -} - -static bool metal_graph_stream_readahead_enabled(void) { - return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_READAHEAD", - "DS4_METAL_ENABLE_STREAMING_READAHEAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_READAHEAD"); -} - -static bool metal_graph_stream_madvise_willneed_enabled(void) { - return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_MADVISE_WILLNEED", - "DS4_METAL_ENABLE_STREAMING_MADVISE_WILLNEED") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_MADVISE_WILLNEED", - "DS4_METAL_DISABLE_STREAMING_MADVISE_WILLNEED"); -} - -static bool metal_graph_stream_decode_static_map_enabled(void) { - if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_STATIC_DECODE_MAP", - "DS4_METAL_DISABLE_STREAMING_STATIC_DECODE_MAP")) { - return false; - } -#ifdef DS4_ROCM_BUILD - return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_STATIC_DECODE_MAP", - "DS4_METAL_ENABLE_STREAMING_STATIC_DECODE_MAP"); -#else - return true; -#endif -} - -static bool metal_graph_stream_decode_static_map_state_cache_enabled(void) { - return !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_STATIC_MAP_STATE_CACHE", - "DS4_METAL_DISABLE_STREAMING_STATIC_MAP_STATE_CACHE"); -} - -static bool metal_graph_stream_decode_layer_batch_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - !g_expert_profile.active && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_LAYER_BATCH", - "DS4_METAL_DISABLE_STREAMING_LAYER_BATCH") && - (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE", - "DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE", - "DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE")) && - !glm_graph_env_present("DS4_ROCM_DECODE_STAGE_PROFILE", - "DS4_METAL_DECODE_STAGE_PROFILE") && - !glm_graph_env_present("DS4_ROCM_GRAPH_DUMP_PREFIX", - "DS4_METAL_GRAPH_DUMP_PREFIX"); -} - -static void metal_graph_stream_readahead_range_impl( - const ds4_model *model, - uint64_t offset, - uint64_t size, - bool enabled) { - if (!enabled || - !model || - model->fd < 0 || - !model->map || - offset > model->size || - size == 0 || - size > model->size - offset) { - return; - } - -#if defined(F_RDADVISE) - uint64_t pos = offset; - uint64_t rem = size; - while (rem > 0) { - const uint64_t chunk64 = - rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem; - if (pos > (uint64_t)LLONG_MAX) break; - - struct radvisory ra; - ra.ra_offset = (off_t)pos; - ra.ra_count = (int)chunk64; - (void)fcntl(model->fd, F_RDADVISE, &ra); - - pos += chunk64; - rem -= chunk64; - } -#else - (void)model; - (void)offset; - (void)size; -#endif -} - -static bool metal_graph_stream_madvise_willneed_range_impl( - const ds4_model *model, - uint64_t offset, - uint64_t size, - bool enabled, - uint64_t *advised) { - if (!enabled || - !model || - !model->map || - offset > model->size || - size == 0 || - size > model->size - offset) { - return !enabled; - } - -#if defined(POSIX_MADV_WILLNEED) - const uint64_t page = (uint64_t)getpagesize(); - if (page == 0) return false; - const uint64_t page_offset = offset & ~(page - 1u); - const uint64_t leading = offset - page_offset; - if (size > UINT64_MAX - leading || - leading + size > UINT64_MAX - (page - 1u)) { - return false; - } - uint64_t advise_bytes = align_up(leading + size, page); - if (advise_bytes > model->size - page_offset) { - advise_bytes = model->size - page_offset; - } - if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) { - return false; - } - uint8_t *base = (uint8_t *)model->map; - const int rc = posix_madvise((void *)(base + page_offset), - (size_t)advise_bytes, - POSIX_MADV_WILLNEED); - if (rc != 0) return false; - if (advised) { - if (*advised > UINT64_MAX - advise_bytes) { - *advised = UINT64_MAX; - } else { - *advised += advise_bytes; - } - } - return true; -#else - (void)model; - (void)offset; - (void)size; - (void)advised; - return true; -#endif -} - -static void metal_graph_stream_readahead_range( - const ds4_model *model, - uint64_t offset, - uint64_t size) { - metal_graph_stream_readahead_range_impl(model, - offset, - size, - metal_graph_stream_readahead_enabled()); - metal_graph_stream_madvise_willneed_range_impl( - model, - offset, - size, - metal_graph_stream_madvise_willneed_enabled(), - NULL); -} - -static void metal_graph_stream_readahead_spans( - const ds4_model *model, - const ds4_model_map_span_vec *spans) { - if (!spans) return; - for (uint32_t i = 0; i < spans->len; i++) { - metal_graph_stream_readahead_range(model, - spans->v[i].off, - spans->v[i].end - spans->v[i].off); - } -} - -typedef struct { - uint64_t off; - uint64_t size; -} metal_graph_stream_pagein_range; - -typedef struct { - pthread_t thread; - const ds4_model *model; - metal_graph_stream_pagein_range *ranges; - pthread_t *threads; - struct metal_graph_stream_pagein_worker *workers; - uint32_t n_ranges; - uint32_t n_threads; - uint32_t layer; - uint32_t n_tokens; - uint32_t unique; - uint64_t bytes; - uint64_t touched; - double read_ms; - double thread_ms; - bool profile; - bool madvise_only; - bool pread_only; - bool readahead_only; - bool started; - bool ok; - uint8_t sink; -} metal_graph_stream_pagein_job; - -typedef struct metal_graph_stream_pagein_worker { - metal_graph_stream_pagein_job *job; - uint32_t first; - uint32_t stride; - uint64_t touched; - double thread_ms; - bool ok; - uint8_t sink; -} metal_graph_stream_pagein_worker; - -static bool metal_graph_stream_prefill_selected_pagein_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_PAGEIN", - "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_PAGEIN") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_PAGEIN", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_PAGEIN"); -} - -static bool metal_graph_stream_prefill_selected_madvise_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_MADVISE", - "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_MADVISE") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_MADVISE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_MADVISE"); -} - -static bool metal_graph_stream_prefill_layer_pagein_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN"); -} - -static bool metal_graph_stream_prefill_layer_readahead_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); -} - -static bool metal_graph_stream_prefill_layer_pread_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); -} - -static bool metal_graph_stream_prefill_layer_madvise_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE"); -} - -static uint32_t metal_graph_stream_prefill_batch_selected_addr_auto_max(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX", - "DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX"); - if (env && env[0]) { - char *end = NULL; - const long v = strtol(env, &end, 10); - if (end != env) { - if (v <= 0) return 0; - if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; - return (uint32_t)v; - } - } -#ifdef DS4_ROCM_BUILD - if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO || - DS4_MODEL_VARIANT == DS4_VARIANT_FLASH || - DS4_MODEL_VARIANT == DS4_VARIANT_GLM52) return UINT32_MAX; -#endif - if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO) return 800u; - if (DS4_MODEL_VARIANT == DS4_VARIANT_FLASH) return 760u; - return 0; -} - -static uint32_t metal_graph_stream_prefill_batch_selected_addr_auto_min(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN", - "DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN"); - if (env && env[0]) { - char *end = NULL; - const long v = strtol(env, &end, 10); - if (end != env) { - if (v <= 0) return 0; - if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; - return (uint32_t)v; - } - } -#ifdef DS4_ROCM_BUILD - if (DS4_MODEL_VARIANT == DS4_VARIANT_GLM52) return 2u; -#endif - if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO || - DS4_MODEL_VARIANT == DS4_VARIANT_FLASH) return 2u; - return 0; -} - -static bool metal_graph_stream_prefill_batch_selected_addr_enabled( - const ds4_gpu_graph *g, - const ds4_weights *weights, - uint32_t n_tokens) { - if (!g || - !g->ssd_streaming || - g->quality || - !weights || - n_tokens <= 1 || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR", - "DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_EXPERT_ADDR_TABLE", - "DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") || - glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") || - glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") || - DS4_N_LAYER == 0) { - return false; - } - const uint32_t routed_il = - DS4_N_LEADING_DENSE < DS4_N_LAYER ? DS4_N_LEADING_DENSE : 0u; - const ds4_layer_weights *layer = &weights->layer[routed_il]; - if (!layer->ffn_gate_exps || !layer->ffn_up_exps || !layer->ffn_down_exps) { - return false; - } -#ifdef DS4_ROCM_BUILD - const bool selected_iq2 = - glm_stream_selected_expert_cache_supported(layer, routed_il); - const bool selected_q2 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q2_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q2_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && - glm_stream_expert_cache_addr_layout_supported(weights, layer, routed_il); - if (!selected_iq2 && !selected_q2) return false; -#else - if (DS4_N_EXPERT_USED != 6 || - layer->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || - layer->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || - layer->ffn_down_exps->type != DS4_TENSOR_Q2_K) { - return false; - } -#endif - - const uint32_t cache_configured = - ds4_gpu_stream_expert_cache_configured_count(); -#ifdef DS4_ROCM_BUILD - if (cache_configured == 0) { - return false; - } -#else - if (cache_configured < DS4_N_EXPERT) { - return false; - } -#endif - - if (glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR", - "DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR")) { - return true; - } - - const uint32_t max_tokens = - metal_graph_stream_prefill_batch_selected_addr_auto_max(); - const uint32_t min_tokens = - metal_graph_stream_prefill_batch_selected_addr_auto_min(); - return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens; -} - -static bool metal_graph_cuda_stream_prefill_batch_selected_addr_enabled( - const ds4_gpu_graph *g, - const ds4_weights *weights, - uint32_t n_tokens) { -#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) - if (!g || - !g->ssd_streaming || - g->quality || - !weights || - n_tokens <= 1 || - getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || - getenv("DS4_CUDA_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL || - DS4_N_LAYER == 0 || - DS4_N_EXPERT < 128 || - DS4_N_EXPERT_USED != 6) { - return false; - } - - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const ds4_layer_weights *layer = &weights->layer[il]; - if (!layer->ffn_gate_exps || !layer->ffn_up_exps || - !layer->ffn_down_exps) { - continue; - } - const bool q4 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q4_K; - const bool iq2 = - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K; - if (q4 || iq2) return true; - } - return false; -#else - (void)g; - (void)weights; - (void)n_tokens; - return false; -#endif -} - -#ifdef DS4_ROCM_BUILD -enum { DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 1024 }; -enum { DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS = 8 }; - -typedef struct rocm_graph_stream_layer_expert_load { - pthread_t thread; - bool active; - bool ok; - const ds4_model *model; - const ds4_layer_weights *layer; - uint32_t il; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; -} rocm_graph_stream_layer_expert_load; - -static bool rocm_graph_stream_prefill_full_layer_enabled( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens) { - return g && - g->ssd_streaming && - !g->quality && - layer && - n_tokens >= DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS && - glm_stream_resident_decode_layer_supported(layer, il); -} - -static uint32_t rocm_graph_stream_prefill_full_layer_seed_tokens(void) { - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_count(); - const uint64_t entries_per_token = - (uint64_t)DS4_N_LAYER * (uint64_t)DS4_N_EXPERT_USED; - if (entries_per_token == 0) return 1; - uint32_t seed_tokens = budget == 0 ? 1 : (uint32_t)(budget / entries_per_token); - if (seed_tokens < 1) seed_tokens = 1; - if (seed_tokens > DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS) { - seed_tokens = DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS; - } - return seed_tokens; -} - -static bool rocm_graph_stream_layer_expert_bytes( - const ds4_layer_weights *layer, - uint64_t *gate_expert_bytes, - uint64_t *down_expert_bytes) { - return streaming_layer_gate_down_expert_bytes(layer, - gate_expert_bytes, - down_expert_bytes); -} - -static bool rocm_graph_stream_layer_expert_load_sync( - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - return model && - layer && - ds4_gpu_stream_expert_cache_load_layer(&table) != 0; -} - -static void *rocm_graph_stream_layer_expert_load_thread_main(void *arg) { - rocm_graph_stream_layer_expert_load *job = arg; - if (!job) return NULL; - job->ok = rocm_graph_stream_layer_expert_load_sync(job->model, - job->layer, - job->il, - job->gate_expert_bytes, - job->down_expert_bytes); - return NULL; -} - -static bool rocm_graph_stream_layer_expert_load_join( - rocm_graph_stream_layer_expert_load *job) { - if (!job || !job->active) return true; - const int rc = pthread_join(job->thread, NULL); - const bool ok = rc == 0 && job->ok; - if (rc != 0) { - fprintf(stderr, - "ds4: ROCm streaming full-layer expert load join failed: %s\n", - strerror(rc)); - } - memset(job, 0, sizeof(*job)); - return ok; -} - -static bool rocm_graph_stream_layer_expert_load_start( - rocm_graph_stream_layer_expert_load *job, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!job || job->active || !model || !layer) return false; - memset(job, 0, sizeof(*job)); - job->model = model; - job->layer = layer; - job->il = il; - job->gate_expert_bytes = gate_expert_bytes; - job->down_expert_bytes = down_expert_bytes; - const int rc = pthread_create(&job->thread, - NULL, - rocm_graph_stream_layer_expert_load_thread_main, - job); - if (rc != 0) { - fprintf(stderr, - "ds4: failed to start ROCm streaming full-layer expert load " - "thread for layer %u: %s\n", - il, - strerror(rc)); - memset(job, 0, sizeof(*job)); - return false; - } - job->active = true; - return true; -} - -static bool rocm_graph_stream_layer_expert_load_start_next( - rocm_graph_stream_layer_expert_load *job, - const ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il, - uint32_t n_tokens) { - if (!job || - !model || - !weights || - il >= DS4_N_LAYER || - !rocm_graph_stream_prefill_full_layer_enabled(g, - &weights->layer[il], - il, - n_tokens)) { - return true; - } - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!rocm_graph_stream_layer_expert_bytes(&weights->layer[il], - &gate_expert_bytes, - &down_expert_bytes)) { - return false; - } - return rocm_graph_stream_layer_expert_load_start(job, - model, - &weights->layer[il], - il, - gate_expert_bytes, - down_expert_bytes); -} - -static bool rocm_graph_stream_layer_expert_load_ready( - rocm_graph_stream_layer_expert_load *job, - const ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il, - uint32_t n_tokens) { - if (!model || !weights || il >= DS4_N_LAYER) return false; - if (!rocm_graph_stream_prefill_full_layer_enabled(g, - &weights->layer[il], - il, - n_tokens)) { - return true; - } - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!rocm_graph_stream_layer_expert_bytes(&weights->layer[il], - &gate_expert_bytes, - &down_expert_bytes)) { - return false; - } - if (job && job->active) { - if (job->il != il) { - fprintf(stderr, - "ds4: ROCm streaming full-layer expert load expected layer " - "%u but pending job is layer %u\n", - il, - job->il); - return false; - } - return rocm_graph_stream_layer_expert_load_join(job); - } - return rocm_graph_stream_layer_expert_load_sync(model, - &weights->layer[il], - il, - gate_expert_bytes, - down_expert_bytes); -} - -static bool rocm_graph_stream_seed_full_layer_selected( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens) { - if (!rocm_graph_stream_prefill_full_layer_enabled(g, layer, il, n_tokens)) { - return true; - } - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!rocm_graph_stream_layer_expert_bytes(layer, - &gate_expert_bytes, - &down_expert_bytes)) { - return false; - } - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - if (ds4_gpu_stream_expert_cache_seed_from_layer_selected( - &table, - metal_graph_batch_router_selected(g), - n_tokens, - rocm_graph_stream_prefill_full_layer_seed_tokens(), - DS4_N_EXPERT_USED) == 0) { - static bool warned = false; - if (!warned) { - fprintf(stderr, - "ds4: ROCm streaming full-layer prefill seed skipped; " - "decode may start with a colder expert cache\n"); - warned = true; - } - } - return true; -} -#endif - -static bool metal_graph_stream_prefill_selected_profile_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_PROFILE", - "DS4_METAL_STREAMING_PREFILL_SELECTED_PROFILE") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_PROFILE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_PROFILE"); -} - -static void metal_graph_stream_prefill_selected_profile_reset( - ds4_gpu_graph *g) { - if (!g) return; - g->prefill_selected_profile_rows = 0; - g->prefill_selected_profile_unique = 0; - g->prefill_selected_profile_selected_bytes = 0; - g->prefill_selected_profile_full_bytes = 0; - g->prefill_selected_profile_layers = 0; - g->prefill_selected_profile_min_unique = UINT32_MAX; - g->prefill_selected_profile_max_unique = 0; -} - -static uint64_t metal_graph_stream_prefill_selected_profile_add_bytes( - uint64_t a, - uint64_t b) { - return a > UINT64_MAX - b ? UINT64_MAX : a + b; -} - -static bool metal_graph_selected_profile_layer_impl( - ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - const char *label) { - if (!layer || !metal_graph_batch_router_selected(g) || n_tokens == 0 || - DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || - DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return false; - } - - const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; - if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; - int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); - const bool read_ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), - 0, - selected, - n_ids * sizeof(selected[0])) != 0; - if (!read_ok) { - free(selected); - return false; - } - - bool seen[DS4_MAX_EXPERT] = { false }; - uint32_t unique = 0; - for (uint64_t i = 0; i < n_ids; i++) { - const int32_t expert = selected[i]; - if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: Metal streaming prefill selected profile expert id %d is outside 0..%u at layer %u\n", - expert, - (uint32_t)DS4_N_EXPERT, - il); - free(selected); - return false; - } - if (!seen[expert]) { - seen[expert] = true; - unique++; - } - } - free(selected); - - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { - fprintf(stderr, "ds4: Metal streaming prefill selected profile byte size overflow at layer %u\n", il); - return false; - } - const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; - const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; - if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes || - gate_expert_bytes + gate_expert_bytes > UINT64_MAX - down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming prefill selected profile byte size overflow at layer %u\n", il); - return false; - } - const uint64_t per_expert_bytes = gate_expert_bytes + gate_expert_bytes + - down_expert_bytes; - const uint64_t selected_bytes = - unique > UINT64_MAX / per_expert_bytes ? - UINT64_MAX : (uint64_t)unique * per_expert_bytes; - const uint64_t full_bytes = - (uint64_t)DS4_N_EXPERT > UINT64_MAX / per_expert_bytes ? - UINT64_MAX : (uint64_t)DS4_N_EXPERT * per_expert_bytes; - const double ratio = full_bytes == 0 ? 0.0 : - (double)selected_bytes / (double)full_bytes; - - g->prefill_selected_profile_layers++; - g->prefill_selected_profile_rows = - metal_graph_stream_prefill_selected_profile_add_bytes( - g->prefill_selected_profile_rows, - n_ids); - g->prefill_selected_profile_unique = - metal_graph_stream_prefill_selected_profile_add_bytes( - g->prefill_selected_profile_unique, - unique); - g->prefill_selected_profile_selected_bytes = - metal_graph_stream_prefill_selected_profile_add_bytes( - g->prefill_selected_profile_selected_bytes, - selected_bytes); - g->prefill_selected_profile_full_bytes = - metal_graph_stream_prefill_selected_profile_add_bytes( - g->prefill_selected_profile_full_bytes, - full_bytes); - if (unique < g->prefill_selected_profile_min_unique) { - g->prefill_selected_profile_min_unique = unique; - } - if (unique > g->prefill_selected_profile_max_unique) { - g->prefill_selected_profile_max_unique = unique; - } - - fprintf(stderr, - "ds4: %s layer=%u " - "tokens=%u unique=%u/%u selected=%.2f GiB full=%.2f GiB ratio=%.3f\n", - label ? label : "selected expert profile", - il, - n_tokens, - unique, - (uint32_t)DS4_N_EXPERT, - (double)selected_bytes / (1024.0 * 1024.0 * 1024.0), - (double)full_bytes / (1024.0 * 1024.0 * 1024.0), - ratio); - return true; -} - -static bool metal_graph_stream_prefill_selected_profile_layer( - ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens) { - if (!metal_graph_stream_prefill_selected_profile_enabled(g)) return true; - return metal_graph_selected_profile_layer_impl( - g, - layer, - il, - n_tokens, - "Metal streaming prefill selected profile"); -} - -static void metal_graph_selected_profile_summary_impl( - const ds4_gpu_graph *g, - const char *label) { - if (!g || g->prefill_selected_profile_layers == 0) { - return; - } - const double layers = (double)g->prefill_selected_profile_layers; - const double avg_unique = (double)g->prefill_selected_profile_unique / layers; - const double ratio = g->prefill_selected_profile_full_bytes == 0 ? 0.0 : - (double)g->prefill_selected_profile_selected_bytes / - (double)g->prefill_selected_profile_full_bytes; - fprintf(stderr, - "ds4: %s summary " - "layers=%u avg_unique=%.1f min_unique=%u max_unique=%u " - "selected=%.2f GiB full=%.2f GiB ratio=%.3f rows=%" PRIu64 "\n", - label ? label : "selected expert profile", - g->prefill_selected_profile_layers, - avg_unique, - g->prefill_selected_profile_min_unique == UINT32_MAX ? - 0 : g->prefill_selected_profile_min_unique, - g->prefill_selected_profile_max_unique, - (double)g->prefill_selected_profile_selected_bytes / - (1024.0 * 1024.0 * 1024.0), - (double)g->prefill_selected_profile_full_bytes / - (1024.0 * 1024.0 * 1024.0), - ratio, - g->prefill_selected_profile_rows); -} - -static void metal_graph_stream_prefill_selected_profile_summary( - const ds4_gpu_graph *g) { - if (!metal_graph_stream_prefill_selected_profile_enabled(g)) return; - metal_graph_selected_profile_summary_impl( - g, - "Metal streaming prefill selected profile"); -} - -static bool metal_graph_stream_pagein_touch_range( - const ds4_model *model, - uint64_t offset, - uint64_t size, - uint64_t *touched, - uint8_t *sink) { - if (!model || - !model->map || - model->size == 0 || - offset > model->size || - size == 0 || - size > model->size - offset) { - return false; - } - - const uint64_t page = (uint64_t)getpagesize(); - const uint64_t page_offset = offset & ~(page - 1u); - const uint64_t leading = offset - page_offset; - if (size > UINT64_MAX - leading || - leading + size > UINT64_MAX - (page - 1u)) { - return false; - } - uint64_t touch_bytes = align_up(leading + size, page); - if (touch_bytes > model->size - page_offset) { - touch_bytes = model->size - page_offset; - } - if (touch_bytes == 0 || touch_bytes > (uint64_t)SIZE_MAX) { - return false; - } - - const uint8_t *base = (const uint8_t *)model->map; - const volatile uint8_t *p = - (const volatile uint8_t *)(base + page_offset); - -#if defined(POSIX_MADV_WILLNEED) - (void)posix_madvise((void *)(base + page_offset), - (size_t)touch_bytes, - POSIX_MADV_WILLNEED); -#endif - - uint8_t s = sink ? *sink : 0; - for (uint64_t off = 0; off < touch_bytes; off += page) { - s ^= p[off]; - } - s ^= p[touch_bytes - 1u]; - if (sink) *sink = s; - if (touched) *touched += touch_bytes; - return true; -} - -static bool metal_graph_stream_pread_range( - const ds4_model *model, - uint64_t offset, - uint64_t size, - uint64_t *read_bytes, - uint8_t *sink) { - if (!model || - model->fd < 0 || - offset > model->size || - size == 0 || - size > model->size - offset) { - return false; - } - if (offset > (uint64_t)LLONG_MAX) return false; - - const size_t chunk = 1024u * 1024u; - uint8_t *buf = xmalloc(chunk); - uint64_t pos = offset; - uint64_t rem = size; - uint8_t s = sink ? *sink : 0; - bool ok = true; - while (rem != 0) { - const size_t want = rem > (uint64_t)chunk ? chunk : (size_t)rem; - ssize_t nread; - do { - nread = pread(model->fd, buf, want, (off_t)pos); - } while (nread < 0 && errno == EINTR); - if (nread <= 0) { - ok = false; - break; - } - s ^= buf[0]; - s ^= buf[(size_t)nread - 1u]; - pos += (uint64_t)nread; - rem -= (uint64_t)nread; - if (read_bytes) { - *read_bytes = *read_bytes > UINT64_MAX - (uint64_t)nread ? - UINT64_MAX : *read_bytes + (uint64_t)nread; - } - } - if (sink) *sink = s; - free(buf); - return ok; -} - -static bool metal_graph_stream_prepare_range( - const metal_graph_stream_pagein_job *job, - uint64_t offset, - uint64_t size, - uint64_t *touched, - uint8_t *sink) { - if (!job) return false; - if (job->pread_only) { - return metal_graph_stream_pread_range(job->model, - offset, - size, - touched, - sink); - } - if (job->readahead_only) { - metal_graph_stream_readahead_range_impl(job->model, - offset, - size, - true); - if (touched) { - *touched = *touched > UINT64_MAX - size ? - UINT64_MAX : *touched + size; - } - return true; - } - if (job->madvise_only) { - return metal_graph_stream_madvise_willneed_range_impl(job->model, - offset, - size, - true, - touched); - } - return metal_graph_stream_pagein_touch_range(job->model, - offset, - size, - touched, - sink); -} - -static void *metal_graph_stream_pagein_thread_main(void *arg) { - metal_graph_stream_pagein_job *job = arg; - const double t0 = job->profile ? now_sec() : 0.0; - job->ok = true; - for (uint32_t i = 0; i < job->n_ranges; i++) { - const bool ok = metal_graph_stream_prepare_range(job, - job->ranges[i].off, - job->ranges[i].size, - &job->touched, - &job->sink); - if (!ok) { - job->ok = false; - break; - } - } - if (job->profile) { - job->thread_ms = (now_sec() - t0) * 1000.0; - } - return NULL; -} - -static void *metal_graph_stream_pagein_worker_main(void *arg) { - metal_graph_stream_pagein_worker *worker = arg; - metal_graph_stream_pagein_job *job = worker ? worker->job : NULL; - const double t0 = job && job->profile ? now_sec() : 0.0; - worker->ok = true; - if (!job || worker->stride == 0) { - worker->ok = false; - return NULL; - } - for (uint32_t i = worker->first; i < job->n_ranges; i += worker->stride) { - const bool ok = metal_graph_stream_prepare_range(job, - job->ranges[i].off, - job->ranges[i].size, - &worker->touched, - &worker->sink); - if (!ok) { - worker->ok = false; - break; - } - } - if (job->profile) { - worker->thread_ms = (now_sec() - t0) * 1000.0; - } - return NULL; -} - -static uint32_t metal_graph_stream_prefill_layer_pagein_threads(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_THREADS", - "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_THREADS"); - if (!env || !env[0]) { - env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_THREADS", - "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_THREADS"); - } - if (!env || !env[0]) return 8; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0' || v == 0) return 1; - return v > 16 ? 16u : (uint32_t)v; -} - -static uint32_t metal_graph_stream_prefill_selected_prepare_threads( - bool madvise_only) { - if (!madvise_only) return 1; - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_SELECTED_PREPARE_THREADS", - "DS4_METAL_STREAMING_PREFILL_SELECTED_PREPARE_THREADS"); - if (!env || !env[0]) { - env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_SELECTED_MADVISE_THREADS", - "DS4_METAL_STREAMING_PREFILL_SELECTED_MADVISE_THREADS"); - } - if (!env || !env[0]) return metal_graph_stream_prefill_layer_pagein_threads(); - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0' || v == 0) return 1; - return v > 16 ? 16u : (uint32_t)v; -} - -static uint32_t metal_graph_stream_prefill_selected_prepare_gap(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_SELECTED_PREPARE_GAP", - "DS4_METAL_STREAMING_PREFILL_SELECTED_PREPARE_GAP"); - if (!env || !env[0]) return 0; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0') return 0; - return v > 8 ? 8u : (uint32_t)v; -} - -static bool metal_graph_stream_prefill_layer_pagein_overlap_enabled(void) { - return !glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_NO_OVERLAP", - "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_NO_OVERLAP") && - !glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_NO_OVERLAP", - "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_NO_OVERLAP") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE_OVERLAP", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE_OVERLAP") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN_OVERLAP", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN_OVERLAP"); -} - -enum { DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD = 4 }; - -static uint32_t metal_graph_stream_prefill_layer_prepare_ahead(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_AHEAD", - "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_AHEAD"); - if (!env || !env[0]) return 1; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0' || v == 0) return 1; - if (v > DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD) { - return DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD; - } - return (uint32_t)v; -} - -static bool metal_graph_stream_prefill_selected_pagein_start( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - metal_graph_stream_pagein_job *job) { - if (!job) return false; - memset(job, 0, sizeof(*job)); - job->ok = true; - job->profile = - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_PAGEIN_PROFILE", - "DS4_METAL_STREAMING_PREFILL_SELECTED_PAGEIN_PROFILE") || - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_MADVISE_PROFILE", - "DS4_METAL_STREAMING_PREFILL_SELECTED_MADVISE_PROFILE"); - job->layer = il; - job->n_tokens = n_tokens; - - const bool madvise_only = - metal_graph_stream_prefill_selected_madvise_enabled(g); - job->madvise_only = madvise_only; - if (!metal_graph_stream_prefill_selected_pagein_enabled(g) && - !madvise_only) return true; - if (!model || !layer || !metal_graph_batch_router_selected(g) || n_tokens == 0) { - return false; - } - - const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; - if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; - int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); - - const double t_read0 = job->profile ? now_sec() : 0.0; - bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), - 0, - selected, - n_ids * sizeof(selected[0])) != 0; - if (job->profile) { - job->read_ms = (now_sec() - t_read0) * 1000.0; - } - - bool seen[DS4_MAX_EXPERT] = { false }; - if (ok) { - for (uint64_t i = 0; i < n_ids; i++) { - const int32_t expert = selected[i]; - if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: Metal streaming prefill selected page-in expert id %d is outside 0..%u at layer %u\n", - expert, - (uint32_t)DS4_N_EXPERT, - il); - ok = false; - break; - } - if (seen[expert]) continue; - seen[expert] = true; - job->unique++; - } - } - free(selected); - - metal_graph_stream_pagein_range *ranges = NULL; - uint32_t n_ranges = 0; - if (ok && job->unique != 0) { - ranges = xmalloc((size_t)DS4_N_EXPERT * 3u * sizeof(ranges[0])); - const uint32_t gap = madvise_only ? - metal_graph_stream_prefill_selected_prepare_gap() : 0; - uint32_t e = 0; - while (e < DS4_N_EXPERT) { - while (e < DS4_N_EXPERT && !seen[e]) e++; - if (e >= DS4_N_EXPERT) break; - const uint32_t first = e; - uint32_t last = e; - uint32_t skipped = 0; - e++; - while (e < DS4_N_EXPERT) { - if (seen[e]) { - last = e; - skipped = 0; - } else if (skipped < gap) { - skipped++; - } else { - break; - } - e++; - } - - const uint64_t first_id = first; - const uint64_t n_experts = (uint64_t)last - (uint64_t)first + 1u; - if (first_id > UINT64_MAX / gate_expert_bytes || - first_id > UINT64_MAX / down_expert_bytes || - n_experts > UINT64_MAX / gate_expert_bytes || - n_experts > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming prefill selected page-in offset overflow\n"); - ok = false; - break; - } - const uint64_t gate_rel = first_id * gate_expert_bytes; - const uint64_t down_rel = first_id * down_expert_bytes; - const uint64_t gate_bytes = n_experts * gate_expert_bytes; - const uint64_t down_bytes = n_experts * down_expert_bytes; - if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || - gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || - down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { - fprintf(stderr, "ds4: Metal streaming prefill selected page-in offset overflow\n"); - ok = false; - break; - } - ranges[n_ranges++] = (metal_graph_stream_pagein_range){ - layer->ffn_gate_exps->abs_offset + gate_rel, - gate_bytes, - }; - ranges[n_ranges++] = (metal_graph_stream_pagein_range){ - layer->ffn_up_exps->abs_offset + gate_rel, - gate_bytes, - }; - ranges[n_ranges++] = (metal_graph_stream_pagein_range){ - layer->ffn_down_exps->abs_offset + down_rel, - down_bytes, - }; - uint64_t run_bytes = UINT64_MAX; - if (gate_bytes <= (UINT64_MAX - down_bytes) / 2ull) { - run_bytes = gate_bytes * 2ull + down_bytes; - } - if (run_bytes == UINT64_MAX || - job->bytes > UINT64_MAX - run_bytes) { - job->bytes = UINT64_MAX; - } else { - job->bytes += run_bytes; - } - } - } - - if (!ok || n_ranges == 0) { - free(ranges); - return ok; - } - - job->model = model; - job->ranges = ranges; - job->n_ranges = n_ranges; - job->n_threads = metal_graph_stream_prefill_selected_prepare_threads(madvise_only); - if (job->n_threads <= 1) { - const int rc = pthread_create(&job->thread, - NULL, - metal_graph_stream_pagein_thread_main, - job); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill selected page-in thread failed: %s\n", - strerror(rc)); - free(ranges); - memset(job, 0, sizeof(*job)); - return false; - } - } else { - job->threads = xcalloc(job->n_threads, sizeof(job->threads[0])); - job->workers = xcalloc(job->n_threads, sizeof(job->workers[0])); - for (uint32_t t = 0; t < job->n_threads; t++) { - job->workers[t].job = job; - job->workers[t].first = t; - job->workers[t].stride = job->n_threads; - const int rc = pthread_create(&job->threads[t], - NULL, - metal_graph_stream_pagein_worker_main, - &job->workers[t]); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill selected page-in worker failed: %s\n", - strerror(rc)); - for (uint32_t j = 0; j < t; j++) { - (void)pthread_join(job->threads[j], NULL); - } - free(job->workers); - free(job->threads); - free(ranges); - memset(job, 0, sizeof(*job)); - return false; - } - } - } - job->started = true; - return true; -} - -static bool metal_graph_stream_prefill_selected_pagein_join( - metal_graph_stream_pagein_job *job) { - if (!job || !job->started) return true; - const double t0 = job->profile ? now_sec() : 0.0; - int rc = 0; - bool ok = true; - if (job->n_threads <= 1) { - rc = pthread_join(job->thread, NULL); - ok = rc == 0 && job->ok; - } else { - job->touched = 0; - job->thread_ms = 0.0; - job->sink = 0; - for (uint32_t t = 0; t < job->n_threads; t++) { - const int trc = pthread_join(job->threads[t], NULL); - if (trc != 0 && rc == 0) rc = trc; - if (trc != 0 || !job->workers[t].ok) ok = false; - if (job->touched > UINT64_MAX - job->workers[t].touched) { - job->touched = UINT64_MAX; - } else { - job->touched += job->workers[t].touched; - } - if (job->workers[t].thread_ms > job->thread_ms) { - job->thread_ms = job->workers[t].thread_ms; - } - job->sink ^= job->workers[t].sink; - } - } - const double wait_ms = job->profile ? (now_sec() - t0) * 1000.0 : 0.0; - if (job->profile) { - const char *kind = job->madvise_only ? "madvise" : "page-in"; - const char *bytes_label = job->madvise_only ? "advised" : "touched"; - fprintf(stderr, - "ds4: Metal streaming prefill selected %s layer=%u " - "tokens=%u unique=%u ranges=%u bytes=%.2f GiB " - "read=%.3f ms wait=%.3f ms thread=%.3f ms %s=%.2f GiB ok=%d\n", - kind, - job->layer, - job->n_tokens, - job->unique, - job->n_ranges, - (double)job->bytes / (1024.0 * 1024.0 * 1024.0), - job->read_ms, - wait_ms, - job->thread_ms, - bytes_label, - (double)job->touched / (1024.0 * 1024.0 * 1024.0), - ok ? 1 : 0); - } - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill selected page-in join failed: %s\n", - strerror(rc)); - } - free(job->workers); - free(job->threads); - free(job->ranges); - memset(job, 0, sizeof(*job)); - return ok; -} - -static bool metal_graph_stream_prefill_layer_pagein_start( - const ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il, - uint32_t n_tokens, - bool madvise_only, - bool pread_only, - bool readahead_only, - bool decode_only, - metal_graph_stream_pagein_job *job) { - if (!job) return false; - memset(job, 0, sizeof(*job)); - job->ok = true; - job->madvise_only = madvise_only; - job->pread_only = pread_only; - job->readahead_only = readahead_only; - job->profile = - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_PROFILE", - "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_PROFILE") || - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PREAD_PROFILE", - "DS4_METAL_STREAMING_PREFILL_LAYER_PREAD_PROFILE") || - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_MADVISE_PROFILE", - "DS4_METAL_STREAMING_PREFILL_LAYER_MADVISE_PROFILE") || - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_READAHEAD_PROFILE", - "DS4_METAL_STREAMING_PREFILL_LAYER_READAHEAD_PROFILE"); - job->layer = il; - job->n_tokens = n_tokens; - - if (pread_only) { - if (g) { - if (!metal_graph_stream_prefill_layer_pread_enabled(g)) return true; - } else if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE")) { - return true; - } - } else if (readahead_only) { - if (g) { - if (!metal_graph_stream_prefill_layer_readahead_enabled(g)) return true; - } else if (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE")) { - return true; - } - } else if (madvise_only) { - if (g) { - if (!metal_graph_stream_prefill_layer_madvise_enabled(g)) return true; - } else if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE")) { - return true; - } - } else { - if (g) { - if (!metal_graph_stream_prefill_layer_pagein_enabled(g)) return true; - } else if (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") || - glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN")) { - return true; - } - } - if (!model || !weights || il >= DS4_N_LAYER) return false; - - const uint32_t n_threads = - metal_graph_stream_prefill_layer_pagein_threads(); - ds4_model_map_span_vec spans; - const bool spans_ok = decode_only ? - weights_model_map_decode_layer_spans(weights, il, &spans) : - weights_model_map_spans(weights, il, il, false, &spans); - if (!spans_ok) return false; - metal_graph_stream_pagein_range *ranges = - xmalloc((size_t)spans.len * n_threads * sizeof(ranges[0])); - uint32_t n_ranges = 0; - const uint64_t page = (uint64_t)getpagesize(); - for (uint32_t i = 0; i < spans.len; i++) { - const uint64_t size = spans.v[i].end - spans.v[i].off; - uint64_t consumed = 0; - uint64_t chunk = size / n_threads; - if (chunk > page) chunk = (chunk / page) * page; - if (chunk == 0) chunk = size; - for (uint32_t t = 0; t < n_threads && consumed < size; t++) { - uint64_t this_size = - (t + 1u == n_threads || size - consumed <= chunk) ? - size - consumed : chunk; - ranges[n_ranges++] = (metal_graph_stream_pagein_range){ - spans.v[i].off + consumed, - this_size, - }; - consumed += this_size; - } - if (job->bytes > UINT64_MAX - size) { - job->bytes = UINT64_MAX; - } else { - job->bytes += size; - } - } - job->unique = spans.len; - free(spans.v); - - job->model = model; - job->ranges = ranges; - job->n_ranges = n_ranges; - job->n_threads = n_threads; - if (n_threads == 1) { - const int rc = pthread_create(&job->thread, - NULL, - metal_graph_stream_pagein_thread_main, - job); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill layer page-in thread failed: %s\n", - strerror(rc)); - free(ranges); - memset(job, 0, sizeof(*job)); - return false; - } - } else { - job->threads = xcalloc(n_threads, sizeof(job->threads[0])); - job->workers = xcalloc(n_threads, sizeof(job->workers[0])); - for (uint32_t t = 0; t < n_threads; t++) { - job->workers[t].job = job; - job->workers[t].first = t; - job->workers[t].stride = n_threads; - const int rc = pthread_create(&job->threads[t], - NULL, - metal_graph_stream_pagein_worker_main, - &job->workers[t]); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill layer page-in worker failed: %s\n", - strerror(rc)); - for (uint32_t j = 0; j < t; j++) { - (void)pthread_join(job->threads[j], NULL); - } - free(job->workers); - free(job->threads); - free(ranges); - memset(job, 0, sizeof(*job)); - return false; - } - } - } - job->started = true; - return true; -} - -static bool metal_graph_stream_prefill_layer_pagein_join( - metal_graph_stream_pagein_job *job) { - if (!job || !job->started) return true; - const double t0 = job->profile ? now_sec() : 0.0; - int rc = 0; - bool ok = true; - if (job->n_threads <= 1) { - rc = pthread_join(job->thread, NULL); - ok = rc == 0 && job->ok; - } else { - job->touched = 0; - job->thread_ms = 0.0; - job->sink = 0; - for (uint32_t t = 0; t < job->n_threads; t++) { - const int trc = pthread_join(job->threads[t], NULL); - if (trc != 0 && rc == 0) rc = trc; - if (trc != 0 || !job->workers[t].ok) ok = false; - if (job->touched > UINT64_MAX - job->workers[t].touched) { - job->touched = UINT64_MAX; - } else { - job->touched += job->workers[t].touched; - } - if (job->workers[t].thread_ms > job->thread_ms) { - job->thread_ms = job->workers[t].thread_ms; - } - job->sink ^= job->workers[t].sink; - } - } - const double wait_ms = job->profile ? (now_sec() - t0) * 1000.0 : 0.0; - if (job->profile) { - const char *kind = job->pread_only ? "pread" : - job->readahead_only ? "readahead" : - job->madvise_only ? "madvise" : "page-in"; - const char *bytes_label = job->pread_only ? "read" : - job->readahead_only ? "requested" : - job->madvise_only ? "advised" : "touched"; - fprintf(stderr, - "ds4: Metal streaming prefill layer %s layer=%u " - "tokens=%u threads=%u ranges=%u bytes=%.2f GiB wait=%.3f ms " - "thread=%.3f ms %s=%.2f GiB ok=%d\n", - kind, - job->layer, - job->n_tokens, - job->n_threads ? job->n_threads : 1u, - job->n_ranges, - (double)job->bytes / (1024.0 * 1024.0 * 1024.0), - wait_ms, - job->thread_ms, - bytes_label, - (double)job->touched / (1024.0 * 1024.0 * 1024.0), - ok ? 1 : 0); - } - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming prefill layer page-in join failed: %s\n", - strerror(rc)); - } - free(job->workers); - free(job->threads); - free(job->ranges); - memset(job, 0, sizeof(*job)); - return ok; -} - -typedef struct { - metal_graph_stream_pagein_job job; - uint32_t layer; - bool active; -} metal_graph_stream_prepare_slot; - -static metal_graph_stream_prepare_slot *metal_graph_stream_prepare_slot_find( - metal_graph_stream_prepare_slot *slots, - uint32_t n_slots, - uint32_t layer) { - for (uint32_t i = 0; i < n_slots; i++) { - if (slots[i].active && slots[i].layer == layer) return &slots[i]; - } - return NULL; -} - -static metal_graph_stream_prepare_slot *metal_graph_stream_prepare_slot_free( - metal_graph_stream_prepare_slot *slots, - uint32_t n_slots) { - for (uint32_t i = 0; i < n_slots; i++) { - if (!slots[i].active) return &slots[i]; - } - return NULL; -} - -static bool metal_graph_stream_prepare_start_if_needed( - const ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t layer, - uint32_t n_tokens, - bool madvise_only, - bool pread_only, - bool readahead_only, - bool decode_only, - metal_graph_stream_prepare_slot *slots, - uint32_t n_slots) { - if (layer >= DS4_N_LAYER) return true; - if (metal_graph_stream_prepare_slot_find(slots, n_slots, layer)) { - return true; - } - metal_graph_stream_prepare_slot *slot = - metal_graph_stream_prepare_slot_free(slots, n_slots); - if (!slot) { - fprintf(stderr, - "ds4: Metal streaming prefill prepare queue is full before layer %u\n", - layer); - return false; - } - memset(slot, 0, sizeof(*slot)); - slot->layer = layer; - if (!metal_graph_stream_prefill_layer_pagein_start(g, - model, - weights, - layer, - n_tokens, - madvise_only, - pread_only, - readahead_only, - decode_only, - &slot->job)) { - memset(slot, 0, sizeof(*slot)); - return false; - } - slot->active = slot->job.started; - return true; -} - -static bool metal_graph_stream_prepare_join_layer( - const ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t layer, - uint32_t n_tokens, - bool madvise_only, - bool pread_only, - bool readahead_only, - bool decode_only, - metal_graph_stream_prepare_slot *slots, - uint32_t n_slots) { - metal_graph_stream_prepare_slot *slot = - metal_graph_stream_prepare_slot_find(slots, n_slots, layer); - if (!slot) { - metal_graph_stream_pagein_job job; - memset(&job, 0, sizeof(job)); - if (!metal_graph_stream_prefill_layer_pagein_start(g, - model, - weights, - layer, - n_tokens, - madvise_only, - pread_only, - readahead_only, - decode_only, - &job)) { - return false; - } - return metal_graph_stream_prefill_layer_pagein_join(&job); - } - const bool ok = metal_graph_stream_prefill_layer_pagein_join(&slot->job); - memset(slot, 0, sizeof(*slot)); - return ok; -} - -static bool metal_graph_stream_prepare_join_all( - metal_graph_stream_prepare_slot *slots, - uint32_t n_slots) { - bool ok = true; - for (uint32_t i = 0; i < n_slots; i++) { - if (!slots[i].active) continue; - if (!metal_graph_stream_prefill_layer_pagein_join(&slots[i].job)) { - ok = false; - } - memset(&slots[i], 0, sizeof(slots[i])); - } - return ok; -} - -static void metal_graph_stream_readahead_layer( - const ds4_model *model, - const ds4_weights *weights, - uint32_t il) { - ds4_model_map_span_vec spans; - if (!weights_model_map_spans(weights, il, il, false, &spans)) return; - metal_graph_stream_readahead_spans(model, &spans); - free(spans.v); -} - -static void metal_graph_stream_readahead_layer_decode( - const ds4_model *model, - const ds4_weights *weights, - uint32_t il) { - ds4_model_map_span_vec spans; - if (!weights_model_map_decode_layer_spans(weights, il, &spans)) return; - metal_graph_stream_readahead_spans(model, &spans); - free(spans.v); -} - -static void metal_graph_stream_readahead_output( - const ds4_model *model, - const ds4_weights *weights) { - ds4_model_map_span_vec spans; - if (!weights_model_map_output_spans(weights, &spans)) return; - metal_graph_stream_readahead_spans(model, &spans); - free(spans.v); -} - -static bool metal_graph_stream_prefill_selected_readahead_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - (glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD", - "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD") || - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", - "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED")) && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD"); -} - -static bool metal_graph_stream_prefill_selected_readahead_shared_enabled( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", - "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD"); -} - -static uint32_t metal_graph_stream_prefill_selected_readahead_gap(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_PREFILL_SELECTED_READAHEAD_GAP", - "DS4_METAL_STREAMING_PREFILL_SELECTED_READAHEAD_GAP"); - if (!env || !env[0]) return 0; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0') return 0; - return v > 8 ? 8u : (uint32_t)v; -} - -static bool metal_graph_stream_readahead_selected_run( - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t first, - uint32_t last, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - uint64_t *hint_bytes) { - if (!model || !layer || first > last || last >= DS4_N_EXPERT) return false; - - const uint64_t first_id = first; - const uint64_t n_experts = (uint64_t)last - (uint64_t)first + 1u; - if (first_id > UINT64_MAX / gate_expert_bytes || - first_id > UINT64_MAX / down_expert_bytes || - n_experts > UINT64_MAX / gate_expert_bytes || - n_experts > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming prefill selected expert readahead overflow\n"); - return false; - } - - const uint64_t gate_rel = first_id * gate_expert_bytes; - const uint64_t down_rel = first_id * down_expert_bytes; - const uint64_t gate_bytes = n_experts * gate_expert_bytes; - const uint64_t down_bytes = n_experts * down_expert_bytes; - if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || - gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || - down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { - fprintf(stderr, "ds4: Metal streaming prefill selected expert readahead overflow\n"); - return false; - } - - metal_graph_stream_readahead_range_impl(model, - layer->ffn_gate_exps->abs_offset + gate_rel, - gate_bytes, - true); - metal_graph_stream_readahead_range_impl(model, - layer->ffn_up_exps->abs_offset + gate_rel, - gate_bytes, - true); - metal_graph_stream_readahead_range_impl(model, - layer->ffn_down_exps->abs_offset + down_rel, - down_bytes, - true); - if (hint_bytes) { - if (*hint_bytes > UINT64_MAX - gate_bytes || - *hint_bytes + gate_bytes > UINT64_MAX - gate_bytes || - *hint_bytes + gate_bytes * 2u > UINT64_MAX - down_bytes) { - *hint_bytes = UINT64_MAX; - } else { - *hint_bytes += gate_bytes * 2u + down_bytes; - } - } - return true; -} - -static bool metal_graph_stream_readahead_selected_experts_from_gpu( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!metal_graph_stream_prefill_selected_readahead_enabled(g)) return true; - if (!model || !layer || !g || !metal_graph_batch_router_selected(g) || n_tokens == 0) { - return false; - } - if (sizeof(int) != sizeof(int32_t) || DS4_N_EXPERT > DS4_MAX_EXPERT) { - return false; - } - - const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; - if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; - int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); - - const bool profile = - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_READAHEAD_PROFILE", - "DS4_METAL_STREAMING_PREFILL_SELECTED_READAHEAD_PROFILE"); - const double t0 = profile ? now_sec() : 0.0; - bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), - 0, - selected, - n_ids * sizeof(selected[0])) != 0; - bool seen[DS4_MAX_EXPERT] = { false }; - uint32_t unique = 0; - uint32_t ranges = 0; - uint64_t hint_bytes = 0; - const uint32_t gap = metal_graph_stream_prefill_selected_readahead_gap(); - if (ok) { - for (uint64_t i = 0; i < n_ids; i++) { - const int32_t expert = selected[i]; - if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: Metal streaming prefill selected expert id %d is outside 0..%u at layer %u\n", - expert, - (uint32_t)DS4_N_EXPERT, - il); - ok = false; - break; - } - if (seen[expert]) continue; - seen[expert] = true; - unique++; - } - } - - if (ok) { - uint32_t e = 0; - while (e < DS4_N_EXPERT) { - while (e < DS4_N_EXPERT && !seen[e]) e++; - if (e >= DS4_N_EXPERT) break; - const uint32_t first = e; - uint32_t last = e; - uint32_t skipped = 0; - e++; - while (e < DS4_N_EXPERT) { - if (seen[e]) { - last = e; - skipped = 0; - } else if (skipped < gap) { - skipped++; - } else { - break; - } - e++; - } - - if (!metal_graph_stream_readahead_selected_run(model, - layer, - first, - last, - gate_expert_bytes, - down_expert_bytes, - &hint_bytes)) { - ok = false; - break; - } - ranges++; - } - } - if (profile) { - fprintf(stderr, - "ds4: Metal streaming prefill selected readahead layer=%u " - "tokens=%u unique=%u ranges=%u gap=%u hint=%.2f GiB time=%.3f ms\n", - il, - n_tokens, - unique, - ranges, - gap, - (double)hint_bytes / (1024.0 * 1024.0 * 1024.0), - (now_sec() - t0) * 1000.0); - } - free(selected); - return ok; -} - -static bool metal_graph_stream_map_token( - const ds4_model *model, - const ds4_weights *weights) { - ds4_model_map_span_vec spans; - if (!weights_model_map_token_spans(weights, &spans)) { - fprintf(stderr, "ds4: Metal SSD streaming could not build token embedding span\n"); - return false; - } - const bool ok = metal_graph_install_model_spans(model, &spans, "token embedding"); - free(spans.v); - return ok; -} - -static bool metal_graph_stream_map_decode_static_all( - const ds4_model *model, - const ds4_weights *weights) { - ds4_model_map_span_vec spans; - if (!weights_model_map_decode_static_spans(weights, true, true, &spans)) { - fprintf(stderr, "ds4: Metal SSD streaming could not build static decode spans\n"); - return false; - } - const bool ok = metal_graph_install_model_spans(model, &spans, "static decode"); - free(spans.v); - return ok; -} - -static bool metal_graph_stream_map_layer( - const ds4_model *model, - const ds4_weights *weights, - uint32_t il) { - ds4_model_map_span_vec spans; - if (!weights_model_map_spans(weights, il, il, false, &spans)) { - fprintf(stderr, "ds4: Metal SSD streaming could not build layer %u spans\n", il); - return false; - } - const bool ok = metal_graph_install_model_spans(model, &spans, "layer"); - free(spans.v); - return ok; -} - -static bool metal_graph_stream_map_layer_decode( - const ds4_model *model, - const ds4_weights *weights, - uint32_t il) { - ds4_model_map_span_vec spans; - if (!weights_model_map_decode_layer_spans(weights, il, &spans)) { - fprintf(stderr, "ds4: Metal SSD streaming could not build decode layer %u spans\n", il); - return false; - } - const bool ok = metal_graph_install_model_spans(model, &spans, "decode layer"); - free(spans.v); - return ok; -} - -static bool metal_graph_stream_map_output( - const ds4_model *model, - const ds4_weights *weights) { - ds4_model_map_span_vec spans; - if (!weights_model_map_output_spans(weights, &spans)) { - fprintf(stderr, "ds4: Metal SSD streaming could not build output head spans\n"); - return false; - } - const bool ok = metal_graph_install_model_spans(model, &spans, "output head"); - free(spans.v); - return ok; -} - -static uint32_t metal_graph_raw_span_for_batch( - const ds4_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens) { - if (!g || g->raw_cap == 0 || n_tokens == 0) return 0; - - const uint32_t window = g->raw_window ? g->raw_window : DS4_N_SWA; - const uint32_t last_pos = pos0 + n_tokens - 1u; - uint64_t needed = (uint64_t)n_tokens; - if (window != 0) { - needed += n_tokens == 1 ? (uint64_t)window - 1u : (uint64_t)window; - } - uint64_t available = (uint64_t)last_pos + 1u; - if (needed > available) needed = available; - if (needed > g->raw_cap) needed = g->raw_cap; - return (uint32_t)needed; -} - -static uint32_t metal_graph_raw_start_for_span( - const ds4_gpu_graph *g, - uint32_t last_pos, - uint32_t n_raw) { - if (!g || g->raw_cap == 0 || n_raw == 0) return 0; - const uint32_t first_raw_pos = last_pos + 1u - n_raw; - return first_raw_pos % g->raw_cap; -} - -static uint32_t metal_graph_decode_raw_score_count( - const ds4_gpu_graph *g, - uint32_t pos, - uint32_t n_raw, - uint32_t ratio) { - if (!g || n_raw == 0) return 0; - if (ratio == 0) return n_raw > 256u ? 256u : n_raw; - - const uint32_t first_raw_pos = pos + 1u - n_raw; - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - uint32_t lo = first_raw_pos; - const uint32_t window = g->raw_window ? g->raw_window : DS4_N_SWA; - if (window != 0 && pos + 1u > window) { - const uint32_t wlo = pos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = pos < raw_last_pos ? pos : raw_last_pos; - if (hi < lo) return 0; - uint32_t raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - return raw_count; -} - -static bool metal_graph_cuda_splitkv_score_may_engage( - const ds4_gpu_graph *g, - uint32_t pos) { - if (!g) return false; - - const uint32_t min_score = metal_graph_cuda_greedy_splitkv_min_score(); - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - uint32_t visible_comp = 0; - const uint32_t n_comp = g->layer_n_comp[il]; - if (n_comp != 0) { - visible_comp = ratio == 0 ? n_comp : (pos + 1u) / ratio; - if (visible_comp > n_comp) visible_comp = n_comp; - } - const uint32_t raw_count = - metal_graph_decode_raw_score_count(g, pos, n_raw, ratio); - const uint32_t n_score = raw_count + visible_comp; - if (n_score > 1u && n_score >= min_score) return true; - } - return false; -} - -static bool metal_graph_cuda_greedy_splitkv_may_engage( - const ds4_gpu_graph *g, - uint32_t pos) { - if (!metal_graph_cuda_greedy_splitkv_requested()) return false; - return metal_graph_cuda_splitkv_score_may_engage(g, pos); -} - -/* Capture the verifier prefix after the first speculative token. - * - * Exact MTP speculation is only profitable if partial accepts are cheap. The - * target verifier computes two draft tokens together; if only the first token - * is accepted, replaying a one-token verifier throws away most of the gain. - * For compressed-attention layers the mutable frontier is just the small - * compressor state plus append counters, so we save that prefix-1 state while - * the N=2 verifier is already stepping the compressor token by token. - * - * Raw SWA rows are not captured here. This graph uses a raw ring larger than - * the 128-token logical SWA window, so writing speculative future rows does - * not evict visible raw rows. If the raw cache is ever reduced to a strict - * 128-row ring, speculative raw rows must become shadow rows and be copied - * into the ring only on commit. */ -static bool metal_graph_capture_prefix1_attn_state(ds4_gpu_graph *g, uint32_t il) { - if (!g->spec_capture_prefix1 || !g->spec_prefix1_attn_state_kv[il]) return true; - const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); - g->spec_prefix1_n_comp[il] = g->layer_n_comp[il]; - return ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_kv[il], 0, - g->layer_attn_state_kv[il], 0, bytes) != 0 && - ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_score[il], 0, - g->layer_attn_state_score[il], 0, bytes) != 0; -} - -static bool metal_graph_capture_prefix1_index_state(ds4_gpu_graph *g, uint32_t il) { - if (!g->spec_capture_prefix1 || !g->spec_prefix1_index_state_kv[il]) return true; - const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); - g->spec_prefix1_n_index_comp[il] = g->layer_n_index_comp[il]; - return ds4_gpu_tensor_copy(g->spec_prefix1_index_state_kv[il], 0, - g->layer_index_state_kv[il], 0, bytes) != 0 && - ds4_gpu_tensor_copy(g->spec_prefix1_index_state_score[il], 0, - g->layer_index_state_score[il], 0, bytes) != 0; -} - -static uint32_t metal_graph_decode_indexer_sparse_threshold(const ds4_gpu_graph *g) { - (void)g; - static int parsed = -1; - static uint32_t cached = 0; - if (parsed < 0) { - parsed = 0; -#ifndef DS4_ROCM_BUILD - const char *env = getenv("DS4_METAL_DECODE_INDEXER_SPARSE_THRESHOLD"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - while (end && isspace((unsigned char)*end)) end++; - if (end != env && end && *end == '\0' && - (v == 64ul || v == 128ul || v == 256ul || v == 512ul || - v == 1024ul || v == 2048ul || v == 4096ul)) { - cached = (uint32_t)v; - parsed = 1; - } else { - fprintf(stderr, - "ds4: invalid DS4_METAL_DECODE_INDEXER_SPARSE_THRESHOLD=%s; " - "expected 64, 128, 256, 512, 1024, 2048, or 4096\n", - env); - } - } -#endif - } - if (parsed > 0) return cached; - - /* Keep dense attention longer than the legacy 512-row window by default. - * Around the 2K frontier the sparse path's score/top-k setup dominates - * the smaller attention scan, while larger contexts benefit from sparse - * indexed attention. This threshold changes only the implementation used - * to consume the compressed rows; it must not lower the 512-row indexer - * selection defined by DS4_N_INDEXER_TOP_K. */ - return 1024u; -} - -/* ========================================================================= - * Metal Decode Release Helpers and Reference Fallbacks. - * ========================================================================= - * - * The normal generation path uses the fused helpers below. The older unfused - * kernels remain available as diagnostic reference paths selected only by the - * DS4_METAL_DISABLE_*_FUSION environment switches. - */ - -static bool metal_graph_env_flag(const char *name, int *cache) { - if (*cache == -1) { -#ifdef DS4_ROCM_BUILD - (void)name; - *cache = 0; -#else - const char *env = getenv(name); - *cache = env && env[0] && strcmp(env, "0") != 0; -#endif - } - return *cache != 0; -} - -static bool metal_graph_use_reference_hc_decode(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_HC_FUSION", &cache); -} - -static bool metal_graph_use_reference_kv_decode(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_KV_FUSION", &cache); -} - -static bool metal_graph_use_reference_qkv_norm(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_NORM_FUSION", &cache); -} - -static bool metal_graph_use_reference_qkv_pair_proj(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_PAIR_PROJ", &cache); -} - -static bool metal_graph_use_reference_compressor_pair_proj(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ", &cache); -} - -static bool metal_graph_use_reference_hc_norm_decode(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_HC_NORM_FUSION", &cache); -} - -static bool metal_graph_enable_batch_hc_norm_fusion(void) { - static int cache = -1; - if (metal_graph_use_reference_hc_norm_decode()) return false; - if (cache == -1) { - const char *disable = getenv("DS4_METAL_DISABLE_BATCH_HC_NORM_FUSION"); - if (disable && disable[0] && strcmp(disable, "0") != 0) { - cache = 0; - } else { - const char *legacy_enable = - getenv("DS4_METAL_ENABLE_BATCH_HC_NORM_FUSION"); - cache = (!legacy_enable || !legacy_enable[0] || - strcmp(legacy_enable, "0") != 0) ? 1 : 0; - } - } - return cache != 0; -} - -static bool metal_graph_use_reference_shared_down_hc(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_SHARED_DOWN_HC_FUSION", &cache); -} - -static bool metal_graph_use_reference_attn_out_hc(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_DISABLE_ATTN_OUT_HC_FUSION", &cache); -} - -static bool metal_graph_decode_hc_pre( - ds4_gpu_tensor *out, - ds4_gpu_tensor *split, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const ds4_model *model, - uint64_t scale_offset, - uint64_t base_offset) { - if (metal_graph_use_reference_hc_decode()) { - return ds4_gpu_hc_split_sinkhorn_tensor(split, - mix, - model->map, - model->size, - scale_offset, - base_offset, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0 && - ds4_gpu_hc_weighted_sum_tensor(out, - residual_hc, - split, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - - return ds4_gpu_hc_split_weighted_sum_tensor(out, - split, - mix, - residual_hc, - model->map, - model->size, - scale_offset, - base_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; -} - -static bool metal_graph_hc_norm_fusion_check_enabled(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_HC_NORM_FUSION_CHECK", &cache); -} - -static float metal_graph_hc_norm_fusion_check_tolerance(void) { - static int initialized; - static float tolerance; - if (initialized) return tolerance; - tolerance = 2.0e-4f; -#ifndef DS4_ROCM_BUILD - const char *env = getenv("DS4_METAL_HC_NORM_FUSION_CHECK_TOL"); - if (env && env[0]) { - char *end = NULL; - const float v = strtof(env, &end); - if (end != env && isfinite(v) && v > 0.0f) tolerance = v; - } -#endif - initialized = 1; - return tolerance; -} - -static bool metal_graph_check_hc_norm_fusion( - const char *label, - ds4_gpu_tensor *fused_out, - ds4_gpu_tensor *fused_norm, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const ds4_model *model, - uint64_t scale_offset, - uint64_t base_offset, - uint64_t norm_weight_offset, - uint32_t il, - uint32_t pos) { - if (!metal_graph_hc_norm_fusion_check_enabled()) return true; - if (!fused_out || !fused_norm || !mix || !residual_hc || !model) return false; - - const uint64_t n_embd = DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - ds4_gpu_tensor *ref_split = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); - ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(n_embd * sizeof(float)); - ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(n_embd * sizeof(float)); - bool ok = ref_split && ref_out && ref_norm; - - if (ok) { - ok = ds4_gpu_hc_split_sinkhorn_tensor(ref_split, - mix, - model->map, - model->size, - scale_offset, - base_offset, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0 && - ds4_gpu_hc_weighted_sum_tensor(ref_out, - residual_hc, - ref_split, - DS4_N_EMBD, - DS4_N_HC) != 0 && - ds4_gpu_rms_norm_weight_tensor(ref_norm, - ref_out, - model->map, - model->size, - norm_weight_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - } - - if (ok) ok = ds4_gpu_end_commands() != 0; - - float *fused_out_cpu = NULL; - float *ref_out_cpu = NULL; - float *fused_norm_cpu = NULL; - float *ref_norm_cpu = NULL; - if (ok) { - fused_out_cpu = xmalloc((size_t)n_embd * sizeof(float)); - ref_out_cpu = xmalloc((size_t)n_embd * sizeof(float)); - fused_norm_cpu = xmalloc((size_t)n_embd * sizeof(float)); - ref_norm_cpu = xmalloc((size_t)n_embd * sizeof(float)); - ok = ds4_gpu_tensor_read(fused_out, 0, fused_out_cpu, n_embd * sizeof(float)) != 0 && - ds4_gpu_tensor_read(ref_out, 0, ref_out_cpu, n_embd * sizeof(float)) != 0 && - ds4_gpu_tensor_read(fused_norm, 0, fused_norm_cpu, n_embd * sizeof(float)) != 0 && - ds4_gpu_tensor_read(ref_norm, 0, ref_norm_cpu, n_embd * sizeof(float)) != 0; - } - - if (ok) { - const float out_max = max_abs_diff(fused_out_cpu, ref_out_cpu, n_embd); - const float out_rms = rms_abs_diff(fused_out_cpu, ref_out_cpu, n_embd); - const float norm_max = max_abs_diff(fused_norm_cpu, ref_norm_cpu, n_embd); - const float norm_rms = rms_abs_diff(fused_norm_cpu, ref_norm_cpu, n_embd); - const float tol = metal_graph_hc_norm_fusion_check_tolerance(); - fprintf(stderr, - "ds4: Metal HC norm fusion check %s layer=%u pos=%u " - "out_max=%g out_rms=%g norm_max=%g norm_rms=%g tol=%g\n", - label ? label : "hc", - il, - pos, - out_max, - out_rms, - norm_max, - norm_rms, - tol); - if (out_max > tol || norm_max > tol) { - fprintf(stderr, - "ds4: Metal HC norm fusion check failed for %s layer=%u pos=%u\n", - label ? label : "hc", - il, - pos); - ok = false; - } - } - - free(fused_out_cpu); - free(ref_out_cpu); - free(fused_norm_cpu); - free(ref_norm_cpu); - ds4_gpu_tensor_free(ref_norm); - ds4_gpu_tensor_free(ref_out); - ds4_gpu_tensor_free(ref_split); - - const bool restart_ok = ds4_gpu_begin_commands() != 0; - return ok && restart_ok; -} - -static bool metal_graph_decode_kv_store( - ds4_gpu_tensor *kv, - ds4_gpu_tensor *raw_cache, - uint32_t raw_cap, - uint32_t raw_row) { - if (metal_graph_use_reference_kv_decode()) { - return ds4_gpu_dsv4_fp8_kv_quantize_tensor(kv, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0 && - ds4_gpu_store_raw_kv_tensor(raw_cache, kv, raw_cap, raw_row, DS4_N_HEAD_DIM) != 0; - } - - return ds4_gpu_kv_fp8_store_raw_tensor(kv, - raw_cache, - raw_cap, - raw_row, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; -} - -static uint64_t metal_graph_attn_comp_cache_row_bytes(void) { - return (uint64_t)DS4_N_HEAD_DIM * - (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); -} - -static uint32_t metal_graph_attn_comp_cache_is_f16(void) { - return DS4_GPU_ATTN_COMP_CACHE_F16 ? 1u : 0u; -} - -static bool metal_graph_store_attn_comp_stage( - ds4_gpu_graph *g, - uint32_t il, - uint32_t first_row, - uint32_t rows) { - if (!g || il >= DS4_N_LAYER) return false; - if (rows == 0) return true; - if (!g->layer_attn_comp_cache[il] || !metal_graph_attn_comp_stage(g)) return false; - if (rows > g->attn_comp_stage_cap || first_row > g->layer_comp_cap[il] || - rows > g->layer_comp_cap[il] - first_row) { - return false; - } - - const uint64_t count = (uint64_t)rows * DS4_N_HEAD_DIM; - const uint64_t dst_offset = (uint64_t)first_row * - metal_graph_attn_comp_cache_row_bytes(); - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - return ds4_gpu_tensor_copy_f32_to_f16(g->layer_attn_comp_cache[il], - dst_offset, - metal_graph_attn_comp_stage(g), - 0, - count) != 0; - } - - return ds4_gpu_tensor_copy(g->layer_attn_comp_cache[il], - dst_offset, - metal_graph_attn_comp_stage(g), - 0, - count * sizeof(float)) != 0; -} - -static ds4_gpu_tensor *metal_graph_attn_comp_update_target( - ds4_gpu_graph *g, - uint32_t il) { - return DS4_GPU_ATTN_COMP_CACHE_F16 - ? metal_graph_attn_comp_stage(g) - : g->layer_attn_comp_cache[il]; -} - -static uint32_t metal_graph_attn_comp_update_row(uint32_t row) { - return DS4_GPU_ATTN_COMP_CACHE_F16 ? 0u : row; -} - -static bool metal_graph_commit_attn_comp_stage( - ds4_gpu_graph *g, - uint32_t il, - uint32_t first_row, - uint32_t rows) { - if (!DS4_GPU_ATTN_COMP_CACHE_F16) return true; - return metal_graph_store_attn_comp_stage(g, il, first_row, rows); -} - -static ds4_gpu_tensor *metal_graph_attn_comp_row_view( - ds4_gpu_graph *g, - uint32_t il, - uint32_t row) { - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - return ds4_gpu_tensor_view(metal_graph_attn_comp_stage(g), - 0, - (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - } - return ds4_gpu_tensor_view(g->layer_attn_comp_cache[il], - (uint64_t)row * DS4_N_HEAD_DIM * sizeof(float), - (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); -} - -static ds4_gpu_tensor *metal_graph_attn_comp_prefill_target( - ds4_gpu_graph *g, - uint32_t il, - uint32_t first_row, - uint32_t rows) { - if (DS4_GPU_ATTN_COMP_CACHE_F16) return metal_graph_attn_comp_stage(g); - const uint32_t view_rows = rows ? rows : 1u; - return ds4_gpu_tensor_view(g->layer_attn_comp_cache[il], - (uint64_t)first_row * DS4_N_HEAD_DIM * sizeof(float), - (uint64_t)view_rows * DS4_N_HEAD_DIM * sizeof(float)); -} - -static void metal_graph_attn_comp_prefill_target_free(ds4_gpu_tensor *t) { - if (!DS4_GPU_ATTN_COMP_CACHE_F16) ds4_gpu_tensor_free(t); -} - -static bool metal_graph_cuda_tp_attn_cache_dup_layer_ready( - const ds4_gpu_graph *g, - uint32_t il) { - if (!g || il >= DS4_N_LAYER || !g->cuda_tp_attn_cache_dup) return false; - if (!g->placement || !g->layer_raw_cache[il] || !g->layer_raw_cache_tp[il]) { - return false; - } - const int layer_tier = g->placement[il + 1]; - if (metal_graph_cuda_tp_partner_tier(layer_tier) < 0) return false; - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0 && - (!g->layer_attn_comp_cache[il] || - !g->layer_attn_comp_cache_tp[il])) { - return false; - } - return true; -} - -static bool metal_graph_cuda_tp_attn_cache_copy_row( - ds4_gpu_tensor *dst_base, - const ds4_gpu_tensor *src_base, - uint64_t offset, - uint64_t bytes) { - if (bytes == 0) return true; - ds4_gpu_tensor *dst = ds4_gpu_tensor_view(dst_base, offset, bytes); - ds4_gpu_tensor *src = ds4_gpu_tensor_view(src_base, offset, bytes); - bool ok = dst && src && ds4_gpu_tensor_copy_xdev(dst, src, bytes) != 0; - ds4_gpu_tensor_free(src); - ds4_gpu_tensor_free(dst); - return ok; -} - -static bool metal_graph_cuda_tp_attn_cache_sync_raw_row( - ds4_gpu_graph *g, - uint32_t il, - uint32_t raw_row) { - if (!metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il)) return true; - if (raw_row >= g->raw_cap) return false; - const uint64_t row_bytes = (uint64_t)DS4_N_HEAD_DIM * sizeof(float); - return metal_graph_cuda_tp_attn_cache_copy_row( - g->layer_raw_cache_tp[il], - g->layer_raw_cache[il], - (uint64_t)raw_row * row_bytes, - row_bytes); -} - -static bool metal_graph_cuda_tp_attn_cache_sync_all(ds4_gpu_graph *g) { - if (!g || !g->cuda_tp_attn_cache_dup) return true; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (!metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il)) return false; - const uint64_t raw_bytes = - (uint64_t)g->raw_cap * DS4_N_HEAD_DIM * sizeof(float); - if (!ds4_gpu_tensor_copy_xdev(g->layer_raw_cache_tp[il], - g->layer_raw_cache[il], - raw_bytes)) { - return false; - } - const uint32_t ratio = ds4_layer_compress_ratio(il); - const uint32_t n_comp = g->layer_n_comp[il]; - if (ratio != 0 && n_comp != 0) { - const uint64_t comp_bytes = - (uint64_t)n_comp * metal_graph_attn_comp_cache_row_bytes(); - if (!ds4_gpu_tensor_copy_xdev(g->layer_attn_comp_cache_tp[il], - g->layer_attn_comp_cache[il], - comp_bytes)) { - return false; - } - } - } - return true; -} - -/* Encode one DS4 decode layer on Metal. This is the release single-token - * layer path; diagnostics reuse it so they compare exactly what generation - * runs. */ -static bool metal_graph_indexer_stage_profile_boundary( - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_comp, - double *stage_t0); -static bool metal_graph_layer_stage_profile_boundary( - const char *part, - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - double *stage_t0); -static bool metal_graph_decode_stage_profile_enabled(uint32_t il); -static bool metal_graph_matmul_plain_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok); -static bool metal_graph_matmul_dense_quant_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok); -static bool metal_graph_dense_quant_row_bytes( - const ds4_tensor *w, - uint64_t in_dim, - uint64_t *row_bytes); -static bool metal_graph_matmul_dense_quant_abs( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok); -static bool metal_graph_matmul_dense_quant_kslice( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - 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); -static bool metal_graph_attention_output_dense_quant_low( - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - uint64_t group_dim, - uint64_t rank, - uint32_t group0, - uint32_t group_cnt, - const ds4_gpu_tensor *heads); -static bool metal_graph_attention_output_dense_quant_tp( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - const ds4_tensor *out_b, - 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); -static bool metal_graph_attention_output_dense_quant_batch( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - const ds4_tensor *out_b, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens); - -static bool metal_graph_use_pro_q4_cpu_router(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_PRO_Q4_CPU_ROUTER", &cache); -} - -static bool metal_graph_use_streaming_iq2_cpu_router(void) { - return getenv("DS4_METAL_ENABLE_STREAMING_IQ2_CPU_ROUTER") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_IQ2_CPU_ROUTER") == NULL; -} - -static bool metal_graph_use_q4_selected_shared_overlap(void) { - static int cache = -1; - return metal_graph_env_flag("DS4_METAL_Q4_SELECTED_OVERLAP_SHARED", &cache); -} - -static bool metal_graph_use_cuda_selected_shared_overlap(const ds4_gpu_graph *g) { -#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) - return g && - g->ssd_streaming && - getenv("DS4_CUDA_DISABLE_STREAMING_SELECTED_SHARED_OVERLAP") == NULL; -#else - (void)g; - return false; -#endif -} - -static bool metal_graph_q4_non_streaming_opt_in_enabled(void) { - return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL; -} - -static bool metal_graph_q4_selected_paths_allowed(const ds4_gpu_graph *g) { - if (!g) return false; - if (g->ssd_streaming) return true; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) return false; - return metal_graph_q4_non_streaming_opt_in_enabled(); -} - -static bool metal_graph_use_iq2_selected_shared_overlap(const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_SHARED_OVERLAP") == NULL && - getenv("DS4_METAL_DISABLE_IQ2_SELECTED_SHARED_OVERLAP") == NULL; -} - -static bool metal_graph_use_iq2_selected_async_load(const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && -#ifndef DS4_ROCM_BUILD - getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_LOAD") == NULL; -#else - true; -#endif -} - -static bool metal_graph_use_iq2_selected_async_early_commit( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && -#ifndef DS4_ROCM_BUILD - getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_EARLY_COMMIT") == NULL; -#else - false; -#endif -} - -static bool metal_graph_use_pro_q4_expert_table_auto(const ds4_gpu_graph *g) { - if (getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") != NULL) { - return false; - } - if (!g || (!g->ssd_streaming && - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL)) { - return false; - } -#ifndef DS4_NO_GPU - return ds4_gpu_pro_q4_expert_table_auto_available() != 0; -#else - return false; -#endif -} - -static bool metal_graph_decode_cpu_router_applicable( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer) { - const bool pro_q4 = - DS4_MODEL_VARIANT == DS4_VARIANT_PRO && - metal_graph_use_pro_q4_cpu_router() && - layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q4_K; - const bool streaming_iq2 = - g && - g->ssd_streaming && - !g->quality && - metal_graph_use_streaming_iq2_cpu_router() && - layer->ffn_gate_tid2eid == NULL && - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && - DS4_N_EXPERT_USED == 6 && - DS4_N_EXPERT >= 128 && - !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && - !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && - !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", - "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); - return pro_q4 || streaming_iq2; -} - -static bool metal_graph_decode_pro_q4_expert_table_expected( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes) { - const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; - return !g->quality && - DS4_MODEL_VARIANT == DS4_VARIANT_PRO && - metal_graph_q4_selected_paths_allowed(g) && - layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && - DS4_N_EXPERT == 384 && - DS4_N_EXPERT_USED == 6 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && - !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && - (metal_graph_use_pro_q4_expert_table_auto(g) || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL) && - getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL; -} - -static bool metal_graph_decode_q4_selected_slots_expected( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes) { - if (metal_graph_decode_pro_q4_expert_table_expected(g, layer, - gate_tensor_bytes, - down_tensor_bytes)) { - return false; - } - const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; - return !g->quality && - metal_graph_q4_selected_paths_allowed(g) && - layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && - DS4_N_EXPERT_USED == 6 && - DS4_N_EXPERT >= 128 && - (g->ssd_streaming || - (gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes)) && - !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && - !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && - !glm_graph_env_present("DS4_ROCM_DISABLE_Q4_SELECTED_EXPERT_VIEWS", - "DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS"); -} - -static bool metal_graph_decode_iq2_selected_slots_expected( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer) { - return g && - g->ssd_streaming && - !g->quality && - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && - DS4_N_EXPERT_USED == 6 && - DS4_N_EXPERT >= 128 && - !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", - "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && - !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", - "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && - !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", - "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); -} - -static bool metal_graph_streaming_expert_cache_seed_layer_expected( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer) { - if (!g || - !g->ssd_streaming || - !layer || - !layer->ffn_gate_exps || - !layer->ffn_up_exps || - !layer->ffn_down_exps) { - return false; - } - if (metal_graph_decode_iq2_selected_slots_expected(g, layer)) return true; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - !g->quality && - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS && - DS4_N_EXPERT_USED != 0 && - DS4_N_EXPERT_USED <= 8 && - DS4_N_EXPERT >= 128 && - !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", - "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { - return true; - } - if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || - g->quality || - layer->ffn_gate_exps->type != layer->ffn_up_exps->type || - layer->ffn_gate_exps->type != layer->ffn_down_exps->type || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > 8 || - DS4_N_EXPERT < 128 || - glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", - "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { - return false; - } - const uint32_t type = layer->ffn_gate_exps->type; - return type == DS4_TENSOR_Q2_K || type == DS4_TENSOR_Q4_K; -} - -static bool metal_graph_decode_cuda_selected_slots_expected( - const ds4_gpu_graph *g, - const ds4_layer_weights *layer) { -#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) - if (!g || - !g->ssd_streaming || - g->quality || - !layer || - !layer->ffn_gate_exps || - !layer->ffn_up_exps || - !layer->ffn_down_exps || - DS4_N_EXPERT_USED != 6 || - DS4_N_EXPERT < 128 || - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) { - return false; - } - const bool q4 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && - layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && - getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL; - const bool iq2 = - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && - getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; - return q4 || iq2; -#else - (void)g; - (void)layer; - return false; -#endif -} - -static uint32_t metal_graph_streaming_prefill_cache_seed_k(const ds4_gpu_graph *g) { - const bool enabled = - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_CACHE_SEED", - "DS4_METAL_ENABLE_STREAMING_PREFILL_CACHE_SEED"); - if (!g || - !g->ssd_streaming || - !enabled) { - return 0; - } - - uint32_t k = 1; - const char *env = glm_graph_env_value("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_K", - "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_K"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0') { - if (v == 0) return 0; - k = v > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS ? - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS : (uint32_t)v; - } - } - return k; -} - -static bool metal_graph_streaming_prefill_cache_seed_enabled(const ds4_gpu_graph *g) { - return metal_graph_streaming_prefill_cache_seed_k(g) != 0; -} - -static bool metal_graph_streaming_expert_hotlist_enabled(const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - !g->ssd_streaming_cold && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_EXPERT_HOTLIST", - "DS4_METAL_DISABLE_STREAMING_EXPERT_HOTLIST"); -} - -static bool metal_graph_streaming_expert_hotlist_add( - uint32_t layer, - uint32_t expert, - uint32_t priority, - int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t counts[DS4_MAX_LAYER], - bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t *loaded) { - if (layer >= DS4_N_LAYER || expert >= DS4_N_EXPERT) return true; - if (layer >= DS4_MAX_LAYER || expert >= DS4_MAX_EXPERT) return true; - if (seen[layer][expert]) return true; - if (counts[layer] >= DS4_MAX_EXPERT) return false; - seen[layer][expert] = true; - if (priority == 0) priority = 1; - priorities[layer][counts[layer]] = priority; - experts[layer][counts[layer]++] = (int32_t)expert; - (*loaded)++; - return true; -} - -static bool metal_graph_streaming_expert_hotlist_load_file( - const char *path, - uint32_t max_entries, - int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t counts[DS4_MAX_LAYER], - bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t *loaded_out) { - if (!path || !path[0] || max_entries == 0 || - !experts || !priorities || !counts || !seen || !loaded_out) { - return false; - } - FILE *fp = fopen(path, "rb"); - if (!fp) { - fprintf(stderr, - "ds4: failed to open streaming expert hotlist %s: %s\n", - path, - strerror(errno)); - return false; - } - - char line[256]; - uint64_t lineno = 0; - uint32_t loaded = 0; - while (fgets(line, sizeof(line), fp)) { - if (loaded >= max_entries) break; - lineno++; - char *p = line; - while (*p && isspace((unsigned char)*p)) p++; - if (*p == '\0' || *p == '#') continue; - - errno = 0; - char *end = NULL; - unsigned long layer = strtoul(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - p = end; - while (*p && isspace((unsigned char)*p)) p++; - - errno = 0; - unsigned long expert = strtoul(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - p = end; - while (*p && isspace((unsigned char)*p)) p++; - - errno = 0; - unsigned long long hits = strtoull(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - if (hits == 0) continue; - const uint32_t priority = - hits > UINT32_MAX ? UINT32_MAX : (uint32_t)hits; - if (!metal_graph_streaming_expert_hotlist_add((uint32_t)layer, - (uint32_t)expert, - priority, - experts, - priorities, - counts, - seen, - &loaded)) { - goto bad_line; - } - continue; - -bad_line: - fprintf(stderr, - "ds4: invalid streaming expert hotlist line %" PRIu64 " in %s\n", - lineno, - path); - fclose(fp); - return false; - } - if (ferror(fp)) { - fprintf(stderr, - "ds4: failed to read streaming expert hotlist %s: %s\n", - path, - strerror(errno)); - fclose(fp); - return false; - } - fclose(fp); - - if (loaded == 0) { - fprintf(stderr, "ds4: streaming expert hotlist %s had no usable nonzero entries\n", path); - } - *loaded_out = loaded; - return true; -} - -static bool metal_graph_streaming_expert_hotlist_load_default( - uint32_t max_entries, - int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t counts[DS4_MAX_LAYER], - bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], - uint32_t *loaded_out) { - if (max_entries == 0 || !experts || !priorities || !counts || !seen || !loaded_out) { - return false; - } - const uint16_t (*hotlist)[2] = NULL; - uint32_t hotlist_count = 0; - if (g_ds4_shape.variant == DS4_VARIANT_PRO) { - hotlist = ds4_default_streaming_hotlist_pro; - hotlist_count = ds4_default_streaming_hotlist_pro_count; - } else if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { - hotlist = ds4_default_streaming_hotlist_flash; - hotlist_count = ds4_default_streaming_hotlist_flash_count; - } else if (g_ds4_shape.variant == DS4_VARIANT_GLM52) { - hotlist = ds4_default_streaming_hotlist_glm52; - hotlist_count = ds4_default_streaming_hotlist_glm52_count; - } else { - *loaded_out = 0; - return true; - } - uint32_t loaded = 0; - for (uint32_t i = 0; - i < hotlist_count && loaded < max_entries; - i++) { - if (!metal_graph_streaming_expert_hotlist_add( - hotlist[i][0], - hotlist[i][1], - max_entries - loaded, - experts, - priorities, - counts, - seen, - &loaded)) { - return false; - } - } - *loaded_out = loaded; - return true; -} - -static uint32_t metal_graph_streaming_expert_preload_count( - const ds4_gpu_graph *g, - uint32_t cache_budget) { - if (!g || cache_budget == 0) return 0; - uint32_t preload = g->streaming_preload_experts; - if (preload == 0) { - preload = cache_budget; - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_EXPERT_AUTO_PRELOAD_CAP", - "DS4_METAL_STREAMING_EXPERT_AUTO_PRELOAD_CAP"); -#ifdef DS4_ROCM_BUILD - if (g_ds4_shape.variant == DS4_VARIANT_GLM52 && - (!env || !env[0])) { - return 0; - } -#endif - /* Auto mode is a hot seed, not a request to synchronously fill the - * whole cache. Large Flash caches can otherwise spend startup doing - * thousands of preads into shared Metal buffers and trip the system - * watchdog before decode begins. ROCm GLM52 uses indexed batch prefill - * by default, which already populates the cache; explicit CLI preload - * counts and auto-preload env caps bypass that default. */ - uint32_t cap = 4096; - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0') { - cap = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v; - } - } - if (cap != 0 && preload > cap) preload = cap; - } - if (preload > cache_budget) preload = cache_budget; - const uint64_t max_possible = (uint64_t)DS4_N_LAYER * DS4_N_EXPERT; - if ((uint64_t)preload > max_possible) preload = (uint32_t)max_possible; - return preload; -} - -static bool metal_graph_decode_set_hash_selected_override( - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t token, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes, - const ds4_gpu_graph *g) { - if (!layer->ffn_gate_tid2eid) return true; - - const bool q4_selected = - metal_graph_decode_q4_selected_slots_expected(g, - layer, - gate_tensor_bytes, - down_tensor_bytes); - const bool iq2_selected = - metal_graph_decode_iq2_selected_slots_expected(g, layer); - if (!q4_selected && !iq2_selected) { - return true; - } - - int selected[DS4_MAX_EXPERT_USED]; - int32_t selected_i32[DS4_MAX_EXPERT_USED]; - layer_hash_selected_experts(selected, model, layer, (int)token); - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - selected_i32[i] = (int32_t)selected[i]; - } - if (g && g->ssd_streaming) { - if (DS4_N_EXPERT == 0 || - gate_tensor_bytes % DS4_N_EXPERT != 0 || - down_tensor_bytes % DS4_N_EXPERT != 0) { - return false; - } - const uint64_t gate_expert_bytes = gate_tensor_bytes / DS4_N_EXPERT; - const uint64_t down_expert_bytes = down_tensor_bytes / DS4_N_EXPERT; - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - if (ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - selected_i32, - DS4_N_EXPERT_USED) == 0) { - return false; - } - } - return ds4_gpu_routed_moe_set_selected_override(selected_i32, DS4_N_EXPERT_USED) != 0; -} - -static bool metal_graph_decode_cpu_router( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t token) { - const bool profile = - getenv("DS4_METAL_PRO_Q4_CPU_ROUTER_PROFILE") != NULL || - getenv("DS4_METAL_STREAMING_IQ2_CPU_ROUTER_PROFILE") != NULL; - const double t0 = profile ? now_sec() : 0.0; - if (ds4_gpu_end_commands() == 0) return false; - const double t_sync = profile ? now_sec() : 0.0; - if (ds4_gpu_tensor_read(metal_graph_ffn_norm(g), - 0, - g->cpu_router_norm, - (uint64_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])) == 0) { - return false; - } - const double t_read = profile ? now_sec() : 0.0; - - float logits[DS4_MAX_EXPERT]; - float probs[DS4_MAX_EXPERT]; - int selected[DS4_MAX_EXPERT_USED]; - int32_t selected_i32[DS4_MAX_EXPERT_USED]; - float weights[DS4_MAX_EXPERT_USED]; - - matvec_any(logits, model, layer->ffn_gate_inp, g->cpu_router_norm); - for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { - probs[i] = sqrtf(softplus_stable(logits[i])); - } - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(selected, model, layer, (int)token); - layer_hash_router_weights_from_probs(weights, probs, selected); - } else { - layer_topk_selected_experts_from_probs(selected, weights, model, layer, probs); - } - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - selected_i32[i] = (int32_t)selected[i]; - } - const double t_cpu = profile ? now_sec() : 0.0; - - if (ds4_gpu_tensor_write(metal_graph_router_logits(g), - 0, - logits, - (uint64_t)DS4_N_EXPERT * sizeof(logits[0])) == 0 || - ds4_gpu_tensor_write(metal_graph_router_probs(g), - 0, - probs, - (uint64_t)DS4_N_EXPERT * sizeof(probs[0])) == 0 || - ds4_gpu_tensor_write(metal_graph_router_selected(g), - 0, - selected_i32, - (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_i32[0])) == 0 || - ds4_gpu_tensor_write(metal_graph_router_weights(g), - 0, - weights, - (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) == 0) { - return false; - } - const double t_write = profile ? now_sec() : 0.0; - if (ds4_gpu_begin_commands() == 0) return false; - if (ds4_gpu_routed_moe_set_selected_override(selected_i32, DS4_N_EXPERT_USED) == 0) return false; - - if (profile) { - fprintf(stderr, - "ds4: Metal CPU router layer=%u gate=%s down=%s sync=%.3f ms read=%.3f ms cpu=%.3f ms write=%.3f ms total=%.3f ms\n", - il, - tensor_type_name(layer->ffn_gate_exps->type), - tensor_type_name(layer->ffn_down_exps->type), - (t_sync - t0) * 1000.0, - (t_read - t_sync) * 1000.0, - (t_cpu - t_read) * 1000.0, - (t_write - t_cpu) * 1000.0, - (t_write - t0) * 1000.0); - } - return true; -} - -static bool metal_graph_use_iq2_selected_readahead_shared_delay( - const ds4_gpu_graph *g) { - return g && - g->ssd_streaming && - getenv("DS4_METAL_ENABLE_STREAMING_SELECTED_READAHEAD_SHARED_DELAY") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_READAHEAD_SHARED_DELAY") == NULL; -} - -static bool metal_graph_decode_selected_readahead_override( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!g || !model || !layer || !metal_graph_router_selected(g) || - DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || - DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return false; - } - - const bool profile = - getenv("DS4_METAL_STREAMING_SELECTED_READAHEAD_PROFILE") != NULL; - const double t0 = profile ? now_sec() : 0.0; - if (ds4_gpu_end_commands() == 0) return false; - const double t_sync = profile ? now_sec() : 0.0; - - int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; - if (ds4_gpu_tensor_read(metal_graph_router_selected(g), - 0, - selected_ids, - (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) == 0) { - return false; - } - const double t_read = profile ? now_sec() : 0.0; - - bool seen[DS4_MAX_EXPERT] = {0}; - uint32_t unique = 0; - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: Metal streaming selected readahead expert id %d is outside 0..%u at layer %u\n", - selected_ids[i], - DS4_N_EXPERT, - il); - return false; - } - const uint32_t expert = (uint32_t)selected_ids[i]; - if (seen[expert]) continue; - seen[expert] = true; - unique++; - - const uint64_t expert_id = (uint64_t)expert; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming selected readahead offset overflow\n"); - return false; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || - gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || - down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { - fprintf(stderr, "ds4: Metal streaming selected readahead offset overflow\n"); - return false; - } - - metal_graph_stream_readahead_range_impl(model, - layer->ffn_gate_exps->abs_offset + gate_rel, - gate_expert_bytes, - true); - metal_graph_stream_readahead_range_impl(model, - layer->ffn_up_exps->abs_offset + gate_rel, - gate_expert_bytes, - true); - metal_graph_stream_readahead_range_impl(model, - layer->ffn_down_exps->abs_offset + down_rel, - down_expert_bytes, - true); - } - const double t_hint = profile ? now_sec() : 0.0; - - if (ds4_gpu_routed_moe_set_selected_override(selected_ids, - DS4_N_EXPERT_USED) == 0) { - return false; - } - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - if (ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - selected_ids, - DS4_N_EXPERT_USED) == 0) { - return false; - } - if (ds4_gpu_begin_commands() == 0) return false; - const double t_done = profile ? now_sec() : 0.0; - - if (profile) { - fprintf(stderr, - "ds4: Metal streaming selected readahead layer=%u unique=%u sync=%.3f ms read=%.3f ms hint=%.3f ms resume=%.3f ms total=%.3f ms\n", - il, - unique, - (t_sync - t0) * 1000.0, - (t_read - t_sync) * 1000.0, - (t_hint - t_read) * 1000.0, - (t_done - t_hint) * 1000.0, - (t_done - t0) * 1000.0); - } - return true; -} - -static bool metal_graph_decode_cuda_selected_load( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { -#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) - if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || - !model || - !metal_graph_router_selected(g) || - DS4_N_EXPERT == 0 || - DS4_N_EXPERT > DS4_MAX_EXPERT || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return false; - } - - const bool profile = - getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_PROFILE") != NULL; - const double t0 = profile ? now_sec() : 0.0; - - if (ds4_gpu_end_commands() == 0) return false; - const double t_sync = profile ? now_sec() : 0.0; - - int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; - bool ok = ds4_gpu_tensor_read(metal_graph_router_selected(g), - 0, - selected_ids, - (uint64_t)DS4_N_EXPERT_USED * - sizeof(selected_ids[0])) != 0; - const double t_read = profile ? now_sec() : 0.0; - - if (ok) { - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - ok = ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - selected_ids, - DS4_N_EXPERT_USED) != 0; - } - const double t_load = profile ? now_sec() : 0.0; - - if (ds4_gpu_begin_commands() == 0) ok = false; - const double t_done = profile ? now_sec() : 0.0; - - if (profile) { - fprintf(stderr, - "ds4: CUDA streaming selected load layer=%u sync=%.3f ms read=%.3f ms load=%.3f ms resume=%.3f ms total=%.3f ms\n", - il, - (t_sync - t0) * 1000.0, - (t_read - t_sync) * 1000.0, - (t_load - t_read) * 1000.0, - (t_done - t_load) * 1000.0, - (t_done - t0) * 1000.0); - } - return ok; -#else - (void)g; - (void)model; - (void)layer; - (void)il; - (void)gate_expert_bytes; - (void)down_expert_bytes; - return false; -#endif -} - -static bool metal_graph_cuda_stream_prefill_batch_selected_load( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { -#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) - if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || - !model || - !metal_graph_batch_router_selected(g) || - n_tokens <= 1 || - DS4_N_EXPERT == 0 || - DS4_N_EXPERT_USED == 0 || - getenv("DS4_CUDA_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_LOAD") != NULL) { - return true; - } - - if ((uint64_t)n_tokens > UINT64_MAX / (uint64_t)DS4_N_EXPERT_USED) { - fprintf(stderr, "ds4: CUDA streaming prefill selected-id count overflow at layer %u\n", il); - return false; - } - const uint64_t n_ids64 = (uint64_t)n_tokens * DS4_N_EXPERT_USED; - if (n_ids64 == 0 || n_ids64 > SIZE_MAX / sizeof(int32_t)) { - fprintf(stderr, "ds4: CUDA streaming prefill selected-id byte size overflow at layer %u\n", il); - return false; - } - - const bool profile = - getenv("DS4_CUDA_STREAMING_PREFILL_BATCH_SELECTED_PROFILE") != NULL; - const double t0 = profile ? now_sec() : 0.0; - - if (ds4_gpu_end_commands() == 0) return false; - const double t_sync = profile ? now_sec() : 0.0; - - int32_t *selected_ids = xmalloc((size_t)n_ids64 * sizeof(selected_ids[0])); - bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), - 0, - selected_ids, - n_ids64 * sizeof(selected_ids[0])) != 0; - const double t_read = profile ? now_sec() : 0.0; - if (ok) { - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - ok = ds4_gpu_stream_expert_cache_prepare_selected_batch( - &table, - selected_ids, - n_tokens, - DS4_N_EXPERT_USED) != 0; - } - free(selected_ids); - const double t_load = profile ? now_sec() : 0.0; - - if (ds4_gpu_begin_commands() == 0) ok = false; - const double t_done = profile ? now_sec() : 0.0; - - if (profile) { - fprintf(stderr, - "ds4: CUDA streaming prefill batch selected load layer=%u tokens=%u sync=%.3f ms read=%.3f ms load=%.3f ms resume=%.3f ms total=%.3f ms\n", - il, - n_tokens, - (t_sync - t0) * 1000.0, - (t_read - t_sync) * 1000.0, - (t_load - t_read) * 1000.0, - (t_done - t_load) * 1000.0, - (t_done - t0) * 1000.0); - } - return ok; -#else - (void)g; - (void)model; - (void)layer; - (void)il; - (void)n_tokens; - (void)gate_expert_bytes; - (void)down_expert_bytes; - return true; -#endif -} - -typedef struct metal_graph_selected_async_load { - bool active; - bool ok; - /* Selected ids remain usable for a synchronous retry if the service - * thread cannot stage the cache load without waiting on GPU work. */ - bool ids_ok; - ds4_gpu_tensor *router_selected; - const ds4_model *model; - const ds4_layer_weights *layer; - uint32_t il; - uint64_t event_value; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; - int32_t selected_ids[DS4_MAX_EXPERT_USED]; -} metal_graph_selected_async_load; - -static pthread_mutex_t g_metal_graph_selected_async_load_mutex = - PTHREAD_MUTEX_INITIALIZER; -static pthread_cond_t g_metal_graph_selected_async_load_cond = - PTHREAD_COND_INITIALIZER; -static pthread_cond_t g_metal_graph_selected_async_load_done_cond = - PTHREAD_COND_INITIALIZER; -static pthread_t g_metal_graph_selected_async_load_thread; -static bool g_metal_graph_selected_async_load_thread_started = false; -static bool g_metal_graph_selected_async_load_has_job = false; -static bool g_metal_graph_selected_async_load_done = false; -static metal_graph_selected_async_load g_metal_graph_selected_async_load_job; - -static void metal_graph_selected_async_load_run( - metal_graph_selected_async_load *job) { - job->ok = false; - - if (!job->router_selected || !job->model || !job->layer || - DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || - DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return; - } - if (job->event_value != 0) { -#ifdef DS4_ROCM_BUILD - if (ds4_gpu_tensor_read_after_selected_event( - job->router_selected, - 0, - job->selected_ids, - (uint64_t)DS4_N_EXPERT_USED * - sizeof(job->selected_ids[0]), - job->event_value, - "selected-id async expert load") == 0) { - return; - } -#else - if (ds4_gpu_wait_selected_readback_ready(job->event_value, - "selected-id async expert load") == 0) { - return; - } - if (ds4_gpu_tensor_read(job->router_selected, - 0, - job->selected_ids, - (uint64_t)DS4_N_EXPERT_USED * - sizeof(job->selected_ids[0])) == 0) { - return; - } -#endif - } - for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { - if (job->selected_ids[i] < 0 || - (uint32_t)job->selected_ids[i] >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: Metal streaming async selected expert id %d is outside 0..%u at layer %u\n", - job->selected_ids[i], - DS4_N_EXPERT, - job->il); - return; - } - } - job->ids_ok = true; - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(job->model, - job->layer, - job->il, - job->gate_expert_bytes, - job->down_expert_bytes); - if (ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - job->selected_ids, - DS4_N_EXPERT_USED) == 0) { - return; - } - - job->ok = true; -} - -static void *metal_graph_selected_async_load_worker_main(void *arg) { - (void)arg; -#ifdef __APPLE__ - /* The Metal cache paths must never wait on command buffers from this - * thread while the main thread is encoding; register it so those waits - * turn into load failures that the caller retries synchronously. */ - ds4_gpu_stream_expert_cache_note_service_thread(); -#endif - for (;;) { - pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); - while (!g_metal_graph_selected_async_load_has_job) { - pthread_cond_wait(&g_metal_graph_selected_async_load_cond, - &g_metal_graph_selected_async_load_mutex); - } - metal_graph_selected_async_load job = - g_metal_graph_selected_async_load_job; - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - - metal_graph_selected_async_load_run(&job); - - pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); - g_metal_graph_selected_async_load_job = job; - g_metal_graph_selected_async_load_has_job = false; - g_metal_graph_selected_async_load_done = true; - pthread_cond_signal(&g_metal_graph_selected_async_load_done_cond); - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - } - return NULL; -} - -static bool metal_graph_selected_async_load_ensure_worker(void) { - pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); - if (g_metal_graph_selected_async_load_thread_started) { - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - return true; - } - const int rc = pthread_create(&g_metal_graph_selected_async_load_thread, - NULL, - metal_graph_selected_async_load_worker_main, - NULL); - if (rc != 0) { - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - fprintf(stderr, - "ds4: failed to start Metal streaming async selected load worker: %s\n", - strerror(rc)); - return false; - } - g_metal_graph_selected_async_load_thread_started = true; - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - return true; -} - -static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start_tensor( - metal_graph_selected_async_load *job, - ds4_gpu_tensor *router_selected, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t event_value, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!job || !router_selected || event_value == 0) return false; - if (!metal_graph_selected_async_load_ensure_worker()) return false; - memset(job, 0, sizeof(*job)); - job->router_selected = router_selected; - job->model = model; - job->layer = layer; - job->il = il; - job->event_value = event_value; - job->gate_expert_bytes = gate_expert_bytes; - job->down_expert_bytes = down_expert_bytes; - - pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); - if (g_metal_graph_selected_async_load_has_job || - g_metal_graph_selected_async_load_done) { - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - return false; - } - g_metal_graph_selected_async_load_job = *job; - g_metal_graph_selected_async_load_job.ok = false; - g_metal_graph_selected_async_load_has_job = true; - pthread_cond_signal(&g_metal_graph_selected_async_load_cond); - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - job->active = true; - return true; -} - -static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start( - metal_graph_selected_async_load *job, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint64_t event_value, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - return metal_graph_selected_async_load_start_tensor( - job, - g ? metal_graph_router_selected(g) : NULL, - model, - layer, - il, - event_value, - gate_expert_bytes, - down_expert_bytes); -} - -static bool metal_graph_selected_async_load_finish( - metal_graph_selected_async_load *job) { - if (!job || !job->active) return false; - pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); - while (!g_metal_graph_selected_async_load_done) { - pthread_cond_wait(&g_metal_graph_selected_async_load_done_cond, - &g_metal_graph_selected_async_load_mutex); - } - *job = g_metal_graph_selected_async_load_job; - g_metal_graph_selected_async_load_done = false; - pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); - job->active = false; - if (!job->ok) return false; - return ds4_gpu_routed_moe_set_selected_override(job->selected_ids, - DS4_N_EXPERT_USED) != 0; -} - -#ifdef DS4_ROCM_BUILD -typedef struct rocm_graph_batch_selected_async_load { - bool active; - bool ok; - const ds4_gpu_tensor *selected; - const ds4_model *model; - const ds4_layer_weights *layer; - uint32_t il; - uint32_t n_tokens; - uint64_t event_value; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; - int32_t *selected_ids; -} rocm_graph_batch_selected_async_load; - -static pthread_mutex_t g_rocm_graph_batch_selected_async_load_mutex = - PTHREAD_MUTEX_INITIALIZER; -static pthread_cond_t g_rocm_graph_batch_selected_async_load_cond = - PTHREAD_COND_INITIALIZER; -static pthread_cond_t g_rocm_graph_batch_selected_async_load_done_cond = - PTHREAD_COND_INITIALIZER; -static pthread_t g_rocm_graph_batch_selected_async_load_thread; -static bool g_rocm_graph_batch_selected_async_load_thread_started = false; -static bool g_rocm_graph_batch_selected_async_load_has_job = false; -static bool g_rocm_graph_batch_selected_async_load_done = false; -static rocm_graph_batch_selected_async_load - g_rocm_graph_batch_selected_async_load_job; - -static void rocm_graph_batch_selected_async_load_run( - rocm_graph_batch_selected_async_load *job) { - job->ok = false; - if (!job->selected || !job->model || !job->layer || !job->selected_ids || - job->n_tokens <= 1 || - DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || - DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return; - } - if (DS4_N_EXPERT_USED != 0 && - job->n_tokens > UINT64_MAX / DS4_N_EXPERT_USED) { - return; - } - const uint64_t n_ids = (uint64_t)job->n_tokens * DS4_N_EXPERT_USED; - if (n_ids > SIZE_MAX / sizeof(job->selected_ids[0])) return; - if (ds4_gpu_tensor_read_after_selected_event( - job->selected, - 0, - job->selected_ids, - n_ids * sizeof(job->selected_ids[0]), - job->event_value, - "prefill selected-id async expert load") == 0) { - return; - } - for (uint64_t i = 0; i < n_ids; i++) { - if (job->selected_ids[i] < 0 || - (uint32_t)job->selected_ids[i] >= DS4_N_EXPERT) { - fprintf(stderr, - "ds4: ROCm streaming async batch selected expert id %d " - "is outside 0..%u at layer %u\n", - job->selected_ids[i], - DS4_N_EXPERT, - job->il); - return; - } - } - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(job->model, - job->layer, - job->il, - job->gate_expert_bytes, - job->down_expert_bytes); - if (ds4_gpu_stream_expert_cache_prepare_selected_batch( - &table, - job->selected_ids, - job->n_tokens, - DS4_N_EXPERT_USED) == 0) { - return; - } - job->ok = true; -} - -static void *rocm_graph_batch_selected_async_load_worker_main(void *arg) { - (void)arg; - for (;;) { - pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); - while (!g_rocm_graph_batch_selected_async_load_has_job) { - pthread_cond_wait(&g_rocm_graph_batch_selected_async_load_cond, - &g_rocm_graph_batch_selected_async_load_mutex); - } - rocm_graph_batch_selected_async_load job = - g_rocm_graph_batch_selected_async_load_job; - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - - rocm_graph_batch_selected_async_load_run(&job); - - pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); - g_rocm_graph_batch_selected_async_load_job = job; - g_rocm_graph_batch_selected_async_load_has_job = false; - g_rocm_graph_batch_selected_async_load_done = true; - pthread_cond_signal(&g_rocm_graph_batch_selected_async_load_done_cond); - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - } - return NULL; -} - -static bool rocm_graph_batch_selected_async_load_ensure_worker(void) { - pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); - if (g_rocm_graph_batch_selected_async_load_thread_started) { - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - return true; - } - const int rc = pthread_create(&g_rocm_graph_batch_selected_async_load_thread, - NULL, - rocm_graph_batch_selected_async_load_worker_main, - NULL); - if (rc != 0) { - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - fprintf(stderr, - "ds4: failed to start ROCm streaming async batch selected " - "load worker: %s\n", - strerror(rc)); - return false; - } - g_rocm_graph_batch_selected_async_load_thread_started = true; - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - return true; -} - -static bool rocm_graph_batch_selected_async_load_start( - rocm_graph_batch_selected_async_load *job, - const ds4_gpu_tensor *selected, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - uint64_t event_value, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!job || !selected || event_value == 0 || n_tokens <= 1) return false; - if (!rocm_graph_batch_selected_async_load_ensure_worker()) return false; - if (DS4_N_EXPERT_USED != 0 && - n_tokens > UINT64_MAX / DS4_N_EXPERT_USED) { - return false; - } - const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; - if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; - memset(job, 0, sizeof(*job)); - job->selected_ids = xmalloc((size_t)n_ids * sizeof(job->selected_ids[0])); - job->selected = selected; - job->model = model; - job->layer = layer; - job->il = il; - job->n_tokens = n_tokens; - job->event_value = event_value; - job->gate_expert_bytes = gate_expert_bytes; - job->down_expert_bytes = down_expert_bytes; - - pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); - if (g_rocm_graph_batch_selected_async_load_has_job || - g_rocm_graph_batch_selected_async_load_done) { - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - free(job->selected_ids); - memset(job, 0, sizeof(*job)); - return false; - } - g_rocm_graph_batch_selected_async_load_job = *job; - g_rocm_graph_batch_selected_async_load_job.ok = false; - g_rocm_graph_batch_selected_async_load_has_job = true; - pthread_cond_signal(&g_rocm_graph_batch_selected_async_load_cond); - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - job->active = true; - return true; -} - -static bool rocm_graph_batch_selected_async_load_finish( - rocm_graph_batch_selected_async_load *job) { - if (!job || !job->active) return false; - pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); - while (!g_rocm_graph_batch_selected_async_load_done) { - pthread_cond_wait(&g_rocm_graph_batch_selected_async_load_done_cond, - &g_rocm_graph_batch_selected_async_load_mutex); - } - *job = g_rocm_graph_batch_selected_async_load_job; - g_rocm_graph_batch_selected_async_load_done = false; - pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); - const bool ok = job->ok; - free(job->selected_ids); - memset(job, 0, sizeof(*job)); - return ok; -} -#endif - -static bool metal_graph_profile_router_selection( - ds4_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos) { - if (!g_expert_profile.active) return true; - if (!g || !layer || !metal_graph_router_selected(g) || !metal_graph_router_weights(g)) return false; - - if (ds4_gpu_end_commands() == 0) { - fprintf(stderr, - "ds4: failed to end Metal command batch for expert profile readback\n"); - return false; - } - - int32_t selected[DS4_MAX_EXPERT_USED] = {0}; - float weights[DS4_MAX_EXPERT_USED] = {0}; - const bool read_ok = - ds4_gpu_tensor_read(metal_graph_router_selected(g), - 0, - selected, - (uint64_t)DS4_N_EXPERT_USED * sizeof(selected[0])) != 0 && - ds4_gpu_tensor_read(metal_graph_router_weights(g), - 0, - weights, - (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) != 0; - - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, - "ds4: failed to resume Metal command batch after expert profile readback\n"); - return false; - } - if (!read_ok) { - fprintf(stderr, "ds4: failed to read Metal router tensors for expert profile\n"); - return false; - } - - ds4_expert_profile_record(il, - pos, - selected, - weights, - layer->ffn_gate_tid2eid != NULL); - return true; -} - -/* Diagnostic skip-ablation for TP profiling (DS4_TP_ABLATE=chain[,chain]): - * drops whole encode chains so their true in-situ cost shows up as a t/s - * delta. Output is semantically wrong while enabled; both ranks must set - * the same value. Chains: hcpre, router, kv, compidx. */ -static bool metal_graph_tp_ablate(const char *chain) { - static const char *env = NULL; - static int init = 0; - if (!init) { - env = getenv("DS4_TP_ABLATE"); - init = 1; - } - return env && strstr(env, chain) != NULL; -} - -static bool metal_graph_borrow_tensor_view( - ds4_gpu_tensor *view, - const ds4_gpu_tensor *base, - uint64_t offset, - uint64_t bytes) { - if (!view || !base || offset > base->bytes || bytes > base->bytes - offset) { - return false; - } - memset(view, 0, sizeof(*view)); - view->ptr = (char *)base->ptr + offset; - view->bytes = bytes; - view->owner = 0; - view->device_id = base->device_id; - return true; -} - -static bool metal_graph_cuda_tp_ep_finish_reduce( - ds4_gpu_graph *g, - int home_tier, - int partner_tier, - bool direct_return, - uint64_t return_bytes, - bool combine) { - bool ok; - if (direct_return) { - ok = ds4_gpu_tensor_wait_xdev_default( - g->routed_down_by_tier[partner_tier], home_tier) != 0; - } else { - ok = ds4_gpu_tensor_copy_xdev_default( - g->tp_peer_tmp_by_tier[home_tier], - g->routed_down_by_tier[partner_tier], - return_bytes) != 0; - } - if (!combine) return ok; - if (ok && g->cuda_tp_ep_pack_exact) { - ok = ds4_gpu_routed_moe_owned_packed_combine_tensor( - metal_graph_routed_out(g), - metal_graph_routed_down(g), - g->tp_peer_tmp_by_tier[home_tier], - metal_graph_router_selected(g), - DS4_N_EMBD, - DS4_N_EXPERT / 2u) != 0; - } else if (ok) { - ok = ds4_gpu_routed_moe_owned_slots_combine_tensor( - metal_graph_routed_out(g), - metal_graph_routed_down(g), - g->tp_peer_tmp_by_tier[home_tier], - metal_graph_router_selected(g), - DS4_N_EMBD, - DS4_N_EXPERT / 2u) != 0; - } - return ok; -} - -typedef enum { - METAL_DECODE_LAYER_FULL = 0, - METAL_DECODE_LAYER_TO_FFN, - METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN, - METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN, - METAL_DECODE_LAYER_TO_QKV, - METAL_DECODE_LAYER_FROM_QKV_TO_ATTN, - METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID, - METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN, - METAL_DECODE_LAYER_FROM_ATTN_TO_FFN, - METAL_DECODE_LAYER_TO_ROUTER, - METAL_DECODE_LAYER_TO_SHARED_MID, - METAL_DECODE_LAYER_FROM_ROUTER, -} metal_decode_layer_phase; - -static bool metal_graph_encode_decode_layer_phase( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos, - ds4_gpu_tensor *raw_cache, - uint32_t raw_cap, - uint32_t raw_row, - uint32_t n_raw, - int token, - metal_decode_layer_phase phase) { - /* switch to this layer's home tier before any Class P - * accessor reads. Single-tier (placement == NULL): no-op. */ - if (g->placement) { - const int this_tier = g->placement[il + 1]; - if (!metal_graph_set_active_tier_decode(g, this_tier)) return false; - } - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint32_t n_groups = DS4_N_OUT_GROUP; - const uint32_t group_heads = DS4_N_HEAD / n_groups; - const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; - const uint32_t rank = DS4_N_LORA_O; - const uint32_t shared_dim = (uint32_t)layer->ffn_gate_shexp->dim[1]; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - const uint64_t down_expert_bytes = routed_out_dim * down_row_bytes; - const bool compressed = ds4_layer_compress_ratio(il) != 0; - const float freq_base = layer_rope_freq_base(il); - const float freq_scale = layer_rope_freq_scale(il); - const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; - float attn_factor = 1.0f; - if (ext_factor != 0.0f && freq_scale > 0.0f) { - attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); - const int cuda_tp_home_tier = g->active_tier; - const int cuda_tp_partner_tier = g->cuda_tp_decode - ? metal_graph_cuda_tp_partner_tier(cuda_tp_home_tier) : -1; - const bool tp_split_attn = g->tp_world == 2; - const uint32_t tp_heads = tp_split_attn ? - (uint32_t)DS4_N_HEAD / 2u : (uint32_t)DS4_N_HEAD; - const uint32_t tp_head0 = tp_split_attn ? g->tp_rank * tp_heads : 0; - - bool ok = true; - const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); - double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; -#define DS4_METAL_PROFILE_DECODE_STAGE(name) do { \ - if (ok && decode_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary("decode", (name), il, pos, 1, &decode_stage_t0); \ - } \ - } while (0) - const bool tp_ablate_hcpre = metal_graph_tp_ablate("hcpre"); - if (phase != METAL_DECODE_LAYER_FROM_ROUTER) { - const bool fuse_hc_norm = - DS4_N_HC == 4 && - !metal_graph_use_reference_hc_decode() && - !metal_graph_use_reference_hc_norm_decode(); - const bool stop_before_attn = - phase == METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN || - phase == METAL_DECODE_LAYER_FROM_QKV_TO_ATTN || - phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; - const bool resume_after_qkv = - phase == METAL_DECODE_LAYER_FROM_QKV_TO_ATTN || - phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; - const bool resume_after_qa_kv_raw = - phase == METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID; - const bool resume_after_kv_store = - phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; - const bool resume_after_attn = - phase == METAL_DECODE_LAYER_FROM_ATTN_TO_FFN; - bool cuda_tp_attn_heads_active = false; - bool attn_inv_rope_done = resume_after_attn; - if (phase != METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN && - phase != METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN && - phase != METAL_DECODE_LAYER_FROM_QKV_TO_ATTN && - phase != METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID && - phase != METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN && - !resume_after_attn) { - if (ok && !tp_ablate_hcpre) { - ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_hc_mix(g), model, layer->hc_attn_fn, - hc_dim, mix_hc, metal_graph_flat_hc(g), 1); - } - if (ok && fuse_hc_norm) { - ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(metal_graph_attn_cur(g), - metal_graph_attn_norm(g), - metal_graph_hc_split(g), - metal_graph_hc_mix(g), - metal_graph_cur_hc(g), - model->map, - model->size, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset, - layer->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - if (ok) { - ok = metal_graph_check_hc_norm_fusion("attn", - metal_graph_attn_cur(g), - metal_graph_attn_norm(g), - metal_graph_hc_mix(g), - metal_graph_cur_hc(g), - model, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset, - layer->attn_norm->abs_offset, - il, - pos); - } - } else if (ok) { - ok = metal_graph_decode_hc_pre(metal_graph_attn_cur(g), - metal_graph_hc_split(g), - metal_graph_hc_mix(g), - metal_graph_cur_hc(g), - model, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset); - } - DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_pre"); - if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre_mixes", metal_graph_hc_mix(g), mix_hc, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_weights", metal_graph_hc_pre(g), DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_post_weights", metal_graph_hc_post(g), DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_comb", metal_graph_hc_comb(g), (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre", metal_graph_attn_cur(g), DS4_N_EMBD, il, pos); - } - if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_attn_norm(g), metal_graph_attn_cur(g), - model->map, model->size, - layer->attn_norm->abs_offset, - DS4_N_EMBD, DS4_RMS_EPS) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("attn_norm"); - if (ok) { - metal_graph_debug_dump_tensor("attn_norm", metal_graph_attn_norm(g), DS4_N_EMBD, il, pos); - } - if (phase == METAL_DECODE_LAYER_TO_QKV) return ok; - } - if (!resume_after_attn) { - if (!resume_after_qkv) { - bool qkv_pair_projected = resume_after_qa_kv_raw; - if (!resume_after_qa_kv_raw && ok && qkv_rms_fused && - g->cuda_qkv_pair && !metal_graph_use_reference_qkv_pair_proj()) { - qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor( - metal_graph_qr(g), - metal_graph_kv_raw(g), - model->map, - model->size, - layer->attn_q_a->abs_offset, - layer->attn_kv->abs_offset, - DS4_N_EMBD, - q_rank, - DS4_N_HEAD_DIM, - metal_graph_attn_norm(g), - 1) != 0; - } - if (!resume_after_qa_kv_raw && ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(metal_graph_qr(g), model->map, model->size, - layer->attn_q_a->abs_offset, - DS4_N_EMBD, q_rank, - metal_graph_attn_norm(g), 1) != 0; - if (ok) { - metal_graph_debug_dump_tensor("q_lora", metal_graph_qr(g), q_rank, il, pos); - } - const bool kvnorm_dump = metal_graph_debug_wants("KVnorm", il, pos); - bool kv_rope_fused = false; - if (qkv_rms_fused) { - if (!resume_after_qa_kv_raw && ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(metal_graph_kv_raw(g), model->map, model->size, - layer->attn_kv->abs_offset, - DS4_N_EMBD, DS4_N_HEAD_DIM, - metal_graph_attn_norm(g), 1) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVraw", metal_graph_kv_raw(g), DS4_N_HEAD_DIM, il, pos); - } - if (ok && g->cuda_qkv_kv_rope_fuse && !kvnorm_dump && DS4_N_HEAD_KV == 1u) { - ok = ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor( - metal_graph_qr_norm(g), - metal_graph_qr(g), - model->map, - model->size, - layer->attn_q_a_norm->abs_offset, - (uint32_t)q_rank, - metal_graph_kv(g), - metal_graph_kv_raw(g), - layer->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - 1, - DS4_N_HEAD_KV, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - kv_rope_fused = ok; - } else if (ok) { - ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(metal_graph_qr_norm(g), - metal_graph_qr(g), - model->map, - model->size, - layer->attn_q_a_norm->abs_offset, - (uint32_t)q_rank, - metal_graph_kv(g), - metal_graph_kv_raw(g), - layer->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - 1, - DS4_RMS_EPS) != 0; - } - } else { - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_qr_norm(g), metal_graph_qr(g), - model->map, model->size, - layer->attn_q_a_norm->abs_offset, - (uint32_t)q_rank, DS4_RMS_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("q_lora_norm", metal_graph_qr_norm(g), q_rank, il, pos); - } - if (qkv_rms_fused && ok && !kv_rope_fused) { - metal_graph_debug_dump_tensor("KVnorm", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); - } - /* Phase B head slice: under the real TP split this rank computes only - * its heads [tp_head0, tp_head0 + tp_heads) end to end — q_b rows, the - * per-head norm/rope, the attention core and its owned output groups. - * q and heads hold the owned half compactly at the buffer base; the - * head range lines up with the output-group split (32 heads = 4 of the - * 8 groups). */ - uint64_t tp_q_row_bytes = 0; - if (ok) ok = metal_graph_dense_quant_row_bytes(layer->attn_q_b, - q_rank, - &tp_q_row_bytes); - const uint64_t tp_q_rows_off = - (uint64_t)tp_head0 * DS4_N_HEAD_DIM * tp_q_row_bytes; - if (ok) ok = metal_graph_matmul_dense_quant_abs(metal_graph_q(g), - model, - layer->attn_q_b, - layer->attn_q_b->abs_offset + tp_q_rows_off, - q_rank, - (uint64_t)tp_heads * DS4_N_HEAD_DIM, - metal_graph_qr_norm(g), - 1); - if (ok) { - metal_graph_debug_dump_tensor("Qraw", metal_graph_q(g), q_dim, il, pos); - } - const bool decode_q_norm_debug = metal_graph_debug_wants("Qnorm", il, pos); - bool decode_q_norm_rope_fused = false; - if (ok && !decode_q_norm_debug) { - decode_q_norm_rope_fused = - ds4_gpu_head_rms_norm_rope_tail_tensor(metal_graph_q(g), - 1, - tp_heads, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } - if (!decode_q_norm_rope_fused) { - if (ok) ok = ds4_gpu_head_rms_norm_tensor(metal_graph_q(g), 1, tp_heads, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; - if (ok) { - metal_graph_debug_dump_tensor("Qnorm", metal_graph_q(g), q_dim, il, pos); - } - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_q(g), 1, tp_heads, DS4_N_HEAD_DIM, - DS4_N_ROT, pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, freq_base, freq_scale, ext_factor, attn_factor, - DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("q_path"); - if (ok) { - metal_graph_debug_dump_tensor("Qcur", metal_graph_q(g), q_dim, il, pos); - } - if (!qkv_rms_fused) { - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_kv_raw(g), - model, - layer->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - metal_graph_attn_norm(g), - 1); - if (ok) { - metal_graph_debug_dump_tensor("KVraw", metal_graph_kv_raw(g), DS4_N_HEAD_DIM, il, pos); - } - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_kv(g), metal_graph_kv_raw(g), - model->map, model->size, - layer->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVnorm", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); - } - } - const bool tp_ablate_kv = metal_graph_tp_ablate("kv"); - if (ok && !tp_ablate_kv && !kv_rope_fused) { - ok = ds4_gpu_rope_tail_tensor(metal_graph_kv(g), 1, - DS4_N_HEAD_KV, DS4_N_HEAD_DIM, - DS4_N_ROT, pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, freq_base, freq_scale, - ext_factor, attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("KVrope", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); - } - } - if (!resume_after_kv_store) { - /* The common no-debug path may fuse KV RMS with RoPE above. KV - * storage starts here after metal_graph_kv(g) contains the RoPE row. */ - if (ok) ok = metal_graph_decode_kv_store(metal_graph_kv(g), raw_cache, raw_cap, raw_row); - if (ok) ok = metal_graph_cuda_tp_attn_cache_sync_raw_row(g, il, raw_row); - DS4_METAL_PROFILE_DECODE_STAGE("kv_path"); - if (ok) { - metal_graph_debug_dump_tensor("KVcur", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); - } - } - - uint32_t n_comp = 0; - ds4_gpu_tensor *comp_cache = NULL; - ds4_gpu_tensor *comp_selected = NULL; - uint32_t n_selected = 0; - double decode_index_stage_t0 = 0.0; - const bool decode_index_stage_profile = g->decode_index_stage_profile; - if (ok && compressed) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint32_t comp_width = coff * DS4_N_HEAD_DIM; - const bool emit = ((pos + 1u) % ratio) == 0u; - if (!layer->attn_compressor_kv || !layer->attn_compressor_gate || - !layer->attn_compressor_ape || !layer->attn_compressor_norm || - layer->attn_compressor_kv->type != DS4_TENSOR_F16 || - layer->attn_compressor_gate->type != DS4_TENSOR_F16 || - layer->attn_compressor_kv->dim[0] != DS4_N_EMBD || - layer->attn_compressor_gate->dim[0] != DS4_N_EMBD || - layer->attn_compressor_kv->dim[1] != comp_width || - layer->attn_compressor_gate->dim[1] != comp_width) { - fprintf(stderr, "ds4: Metal graph compressor expects paired F16 compressor projections\n"); - ok = false; - } - if (ok && emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - } - bool comp_state_already_stored = false; - if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - const int fused_store = - ds4_gpu_matmul_f16_pair_compressor_store_tensor( - metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - model->map, - model->size, - layer->attn_compressor_kv->abs_offset, - layer->attn_compressor_gate->abs_offset, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - DS4_N_EMBD, - comp_width, - metal_graph_attn_norm(g), - ratio, - pos); - if (fused_store < 0) { - ok = false; - } else if (fused_store > 0) { - comp_state_already_stored = true; - } else { - ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - model->map, - model->size, - layer->attn_compressor_kv->abs_offset, - layer->attn_compressor_gate->abs_offset, - DS4_N_EMBD, - comp_width, - metal_graph_attn_norm(g), - 1) != 0; - } - } else { - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size, - layer->attn_compressor_kv->abs_offset, - DS4_N_EMBD, comp_width, - metal_graph_attn_norm(g), 1) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, - layer->attn_compressor_gate->abs_offset, - DS4_N_EMBD, comp_width, - metal_graph_attn_norm(g), 1) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("compressor_proj"); - const uint32_t comp_row = g->layer_n_comp[il]; - if (ok) ok = ds4_gpu_compressor_update_tensor(metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - metal_graph_attn_comp_update_target(g, il), - model->map, - model->size, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - layer->attn_compressor_norm->abs_offset, - layer->attn_compressor_norm->type, - DS4_N_HEAD_DIM, - ratio, - pos, - metal_graph_attn_comp_update_row(comp_row), - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS, - comp_state_already_stored) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("compressor_update"); - if (ok && emit) { - ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); - if (!comp_row_view) { - ok = false; - } else { - ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVcompress", comp_row_view, DS4_N_HEAD_DIM, il, pos); - } - } - ds4_gpu_tensor_free(comp_row_view); - DS4_METAL_PROFILE_DECODE_STAGE("compressor_quantize"); - if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); - DS4_METAL_PROFILE_DECODE_STAGE("compressor_commit"); - } - if (ok && emit) g->layer_n_comp[il]++; - - if (ok && ratio == 4) { - const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; - if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || - !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || - layer->indexer_compressor_kv->type != DS4_TENSOR_F16 || - layer->indexer_compressor_gate->type != DS4_TENSOR_F16 || - layer->indexer_compressor_kv->dim[0] != DS4_N_EMBD || - layer->indexer_compressor_gate->dim[0] != DS4_N_EMBD || - layer->indexer_compressor_kv->dim[1] != index_width || - layer->indexer_compressor_gate->dim[1] != index_width) { - fprintf(stderr, "ds4: Metal graph indexer compressor expects paired F16 projections\n"); - ok = false; - } - if (ok && emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - } - bool index_state_already_stored = false; - if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - const int fused_store = - ds4_gpu_matmul_f16_pair_compressor_store_tensor( - metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - model->map, - model->size, - layer->indexer_compressor_kv->abs_offset, - layer->indexer_compressor_gate->abs_offset, - layer->indexer_compressor_ape->abs_offset, - layer->indexer_compressor_ape->type, - DS4_N_EMBD, - index_width, - metal_graph_attn_norm(g), - ratio, - pos); - if (fused_store < 0) { - ok = false; - } else if (fused_store > 0) { - index_state_already_stored = true; - } else { - ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - model->map, - model->size, - layer->indexer_compressor_kv->abs_offset, - layer->indexer_compressor_gate->abs_offset, - DS4_N_EMBD, - index_width, - metal_graph_attn_norm(g), - 1) != 0; - } - } else { - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size, - layer->indexer_compressor_kv->abs_offset, - DS4_N_EMBD, index_width, - metal_graph_attn_norm(g), 1) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, - layer->indexer_compressor_gate->abs_offset, - DS4_N_EMBD, index_width, - metal_graph_attn_norm(g), 1) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_proj"); - const uint32_t index_row = g->layer_n_index_comp[il]; - if (ok) ok = ds4_gpu_compressor_update_tensor(metal_graph_comp_kv_cur(g), - metal_graph_comp_sc_cur(g), - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - g->layer_index_comp_cache[il], - model->map, - model->size, - layer->indexer_compressor_ape->abs_offset, - layer->indexer_compressor_ape->type, - layer->indexer_compressor_norm->abs_offset, - layer->indexer_compressor_norm->type, - DS4_N_INDEXER_HEAD_DIM, - ratio, - pos, - index_row, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS, - index_state_already_stored) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_update"); - if (ok && emit) { -#if defined(__APPLE__) - ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( - g->layer_index_comp_cache[il], - (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), - (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - if (!index_row_view) { - ok = false; - } else { - ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, - 1, - DS4_N_INDEXER_HEAD_DIM) != 0; - } - ds4_gpu_tensor_free(index_row_view); -#else - ds4_gpu_tensor index_row_view; - if (!metal_graph_borrow_tensor_view( - &index_row_view, - g->layer_index_comp_cache[il], - (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), - (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float))) { - ok = false; - } else { - ok = ds4_gpu_dsv4_indexer_qat_tensor(&index_row_view, - 1, - DS4_N_INDEXER_HEAD_DIM) != 0; - } -#endif - DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_qat"); - } - if (ok && emit) g->layer_n_index_comp[il]++; - const uint32_t decode_sparse_threshold = - metal_graph_decode_indexer_sparse_threshold(g); - if (ok && - g->layer_n_comp[il] > decode_sparse_threshold && - g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K) { - const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - if (!layer->indexer_attn_q_b || - !tensor_type_is_f16_or_q8_0(layer->indexer_attn_q_b->type) || - layer->indexer_attn_q_b->dim[0] != q_rank || - layer->indexer_attn_q_b->dim[1] != indexer_q_dim) { - fprintf(stderr, "ds4: Metal graph indexer q projection expects F16 or Q8_0 weights\n"); - ok = false; - } - if (ok && (!layer->indexer_proj || - layer->indexer_proj->type != DS4_TENSOR_F16 || - layer->indexer_proj->dim[0] != DS4_N_EMBD || - layer->indexer_proj->dim[1] != DS4_N_INDEXER_HEAD)) { - fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); - ok = false; - } - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_indexer_q(g), - model, - layer->indexer_attn_q_b, - q_rank, - indexer_q_dim, - metal_graph_qr_norm(g), - 1); - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_indexer_q(g), 1, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_indexer_q(g), - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_indexer_weights(g), model->map, model->size, - layer->indexer_proj->abs_offset, - DS4_N_EMBD, DS4_N_INDEXER_HEAD, - metal_graph_attn_norm(g), 1) != 0; - const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - if (ok && decode_index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary(NULL, - il, - pos, - 1, - g->layer_n_index_comp[il], - &decode_index_stage_t0); - } - if (ok) ok = ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), - metal_graph_indexer_q(g), - metal_graph_indexer_weights(g), - g->layer_index_comp_cache[il], - g->layer_n_index_comp[il], - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - index_scale) != 0; - if (ok && decode_index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("decode_score", - il, - pos, - 1, - g->layer_n_index_comp[il], - &decode_index_stage_t0); - } - if (ok) ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - metal_graph_indexer_scores(g), - g->layer_n_index_comp[il], - 1, - DS4_N_INDEXER_TOP_K) != 0; - if (ok && decode_index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("decode_topk", - il, - pos, - 1, - g->layer_n_index_comp[il], - &decode_index_stage_t0); - } - /* Decode used to materialize a dense compressed-row mask and - * call the generic gathered FlashAttention wrapper below. - * That wrapper scans every compressed row and rejects long - * contexts once raw+compressed rows exceed 8192. Ratio-4 DS4 - * attention is sparse after indexer top-k, so use the private - * indexed attention kernel instead: it scans only SWA raw rows - * plus the selected compressed rows, matching prefill and - * avoiding the long-context decode failure. */ - if (ok) { - comp_selected = metal_graph_comp_selected(g); - /* - * Contract: the indexer top-k is fixed by the model config - * and must remain the full 512 rows. Do not reduce this for - * throughput benchmarks. - * - * Why: the indexer is not just an implementation detail. It - * decides which compressed memory rows are visible to the - * attention kernel. If we keep only 128/256 rows, the later - * indexed-attention math may be perfectly computed, but it is - * computed over the wrong candidate set: rows ranked 257-512 - * are removed before softmax/PV can use them. Those rows may - * carry weak-but-necessary evidence for retrieval, name/number - * recall, or long-context disambiguation. The error is - * therefore semantic/algorithmic, not the acceptable kind of - * local numerical drift caused by a different reduction order - * or Tensor/NAX precision. - * - * Short prompt tests, first-token agreement, or even a small - * official-vector set can miss this because many prompts do - * not need the tail of the 512 selected compressed rows. The - * failure appears only when the model needs information that - * fell below the reduced cutoff. Optimizations belong inside - * the score/top-k/attention implementation while preserving - * DS4_N_INDEXER_TOP_K. - */ - n_selected = DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] - ? DS4_N_INDEXER_TOP_K - : g->layer_n_index_comp[il]; - } - } - } - - n_comp = g->layer_n_comp[il]; - comp_cache = g->layer_attn_comp_cache[il]; - } - DS4_METAL_PROFILE_DECODE_STAGE("compressor_indexer"); - - if (stop_before_attn) return ok; - if (ok) { - const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); - const bool indexed_attention = n_comp != 0 && comp_selected != NULL && n_selected != 0; - const bool cuda_tp_attn_heads_requested = g->cuda_tp_attn_heads; - const uint32_t cuda_tp_heads = DS4_N_HEAD / 2u; - const bool cuda_tp_attn_local_cache = - metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il); - cuda_tp_attn_heads_active = - cuda_tp_attn_heads_requested && - cuda_tp_partner_tier >= 0 && - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] && - (DS4_N_HEAD % 2u) == 0u && - (n_groups % 2u) == 0u && - g->q_by_tier[cuda_tp_partner_tier] && - g->heads_by_tier[cuda_tp_partner_tier] && - !metal_graph_debug_wants("kqv_out", il, pos) && - !metal_graph_debug_wants("kqv_back", il, pos); - if (cuda_tp_attn_heads_requested && !cuda_tp_attn_heads_active) { - fprintf(stderr, - "ds4: CUDA decode TP cannot split attention heads for tier %d " - "(partner=%d heads=%u peer=%d)\n", - cuda_tp_home_tier, - cuda_tp_partner_tier, - DS4_N_HEAD, - cuda_tp_partner_tier >= 0 ? - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] : 0); - ok = false; - } - if (ok && cuda_tp_attn_heads_active) { - const uint64_t tp_head_bytes = (uint64_t)cuda_tp_heads * DS4_N_HEAD_DIM * sizeof(float); - ds4_gpu_tensor q_peer_src; - ds4_gpu_tensor q_peer_dst; - ds4_gpu_tensor peer_heads_tail; - ok = metal_graph_borrow_tensor_view(&q_peer_src, - metal_graph_q(g), - tp_head_bytes, - tp_head_bytes) && - metal_graph_borrow_tensor_view(&q_peer_dst, - g->q_by_tier[cuda_tp_partner_tier], - 0, - tp_head_bytes) && - metal_graph_borrow_tensor_view(&peer_heads_tail, - g->heads_by_tier[cuda_tp_partner_tier], - tp_head_bytes, - tp_head_bytes); - if (ok) { - ok = ds4_gpu_tensor_copy_xdev(&q_peer_dst, &q_peer_src, - tp_head_bytes) != 0; - } - ds4_gpu_tensor *peer_raw_cache = cuda_tp_attn_local_cache - ? g->layer_raw_cache_tp[il] : raw_cache; - ds4_gpu_tensor *peer_comp_cache = cuda_tp_attn_local_cache - ? g->layer_attn_comp_cache_tp[il] : comp_cache; - ds4_gpu_tensor *peer_selected = comp_selected; - if (ok && indexed_attention && cuda_tp_attn_local_cache) { - peer_selected = g->comp_selected_by_tier[cuda_tp_partner_tier]; - ok = peer_selected && - ds4_gpu_tensor_copy_xdev(peer_selected, - comp_selected, - (uint64_t)n_selected * sizeof(int32_t)) != 0; - } - if (ok && !cuda_tp_attn_local_cache) { - ok = ds4_gpu_tensor_wait_xdev(raw_cache, cuda_tp_partner_tier) != 0; - } - if (ok) ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; - if (ok && indexed_attention) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - &peer_heads_tail, - model->map, - model->size, - layer->attn_sinks->abs_offset + (uint64_t)cuda_tp_heads * sizeof(float), - &q_peer_dst, - peer_raw_cache, - peer_comp_cache, - metal_graph_attn_comp_cache_is_f16(), - peer_selected, - 1, - pos, - n_raw, - raw_cap, - raw_start, - n_comp, - n_selected, - g->raw_window, - ds4_layer_compress_ratio(il), - cuda_tp_heads, - DS4_N_HEAD_DIM) != 0; - } else if (ok) { - ok = ds4_gpu_attention_decode_heads_tensor( - &peer_heads_tail, - model->map, - model->size, - layer->attn_sinks->abs_offset + (uint64_t)cuda_tp_heads * sizeof(float), - &q_peer_dst, - peer_raw_cache, - n_raw, - raw_cap, - raw_start, - n_comp ? peer_comp_cache : NULL, - metal_graph_attn_comp_cache_is_f16(), - n_comp, - NULL, - 0, - cuda_tp_heads, - DS4_N_HEAD_DIM) != 0; - } - if (ok) { - ok = ds4_gpu_rope_tail_tensor(&peer_heads_tail, - 1, cuda_tp_heads, DS4_N_HEAD_DIM, - DS4_N_ROT, pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - if (ok) ok = ds4_gpu_set_current_device(cuda_tp_home_tier) == 0; - if (ok && indexed_attention) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - metal_graph_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_q(g), - raw_cache, - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - comp_selected, - 1, - pos, - n_raw, - raw_cap, - raw_start, - n_comp, - n_selected, - g->raw_window, - ds4_layer_compress_ratio(il), - cuda_tp_heads, - DS4_N_HEAD_DIM) != 0; - } else if (ok) { - ok = ds4_gpu_attention_decode_heads_tensor(metal_graph_heads(g), - model->map, model->size, - layer->attn_sinks->abs_offset, - metal_graph_q(g), raw_cache, n_raw, - raw_cap, - raw_start, - n_comp ? comp_cache : NULL, - metal_graph_attn_comp_cache_is_f16(), - n_comp, - NULL, - 0, - cuda_tp_heads, DS4_N_HEAD_DIM) != 0; - } - if (ok) { - ok = ds4_gpu_rope_tail_tensor(metal_graph_heads(g), - 1, cuda_tp_heads, DS4_N_HEAD_DIM, - DS4_N_ROT, pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - if (ok && indexed_attention && decode_index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("decode_attention", - il, - pos, - 1, - n_comp, - &decode_index_stage_t0); - } - } else if (ok && indexed_attention) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - metal_graph_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset + (uint64_t)tp_head0 * (layer->attn_sinks->bytes / DS4_N_HEAD), - metal_graph_q(g), - raw_cache, - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - comp_selected, - 1, - pos, - n_raw, - raw_cap, - raw_start, - n_comp, - n_selected, - g->raw_window, - ds4_layer_compress_ratio(il), - tp_heads, - DS4_N_HEAD_DIM) != 0; - if (ok && decode_index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("decode_attention", - il, - pos, - 1, - n_comp, - &decode_index_stage_t0); - } - } else { - ok = ds4_gpu_attention_decode_heads_tensor(metal_graph_heads(g), - model->map, model->size, - layer->attn_sinks->abs_offset + (uint64_t)tp_head0 * (layer->attn_sinks->bytes / DS4_N_HEAD), - metal_graph_q(g), raw_cache, n_raw, - raw_cap, - raw_start, - n_comp ? comp_cache : NULL, - metal_graph_attn_comp_cache_is_f16(), - n_comp, - NULL, - 0, - tp_heads, DS4_N_HEAD_DIM) != 0; - } - } - } - if (ok && !cuda_tp_attn_heads_active && !attn_inv_rope_done) { - ok = ds4_gpu_rope_tail_tensor(metal_graph_heads(g), - 1, tp_heads, DS4_N_HEAD_DIM, - DS4_N_ROT, pos, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("attn_inv_rope"); - if (ok && !cuda_tp_attn_heads_active) { - metal_graph_debug_dump_tensor("kqv_back", metal_graph_heads(g), q_dim, il, pos); - } - ds4_gpu_tensor *cuda_tp_attn_peer = NULL; - bool cuda_tp_attn_hc_fused = false; - const bool cuda_tp_attn_requested = g->cuda_tp_attn; - const bool cuda_tp_attn = - cuda_tp_attn_requested && - !metal_graph_directional_steering_attn_enabled(g) && - cuda_tp_partner_tier >= 0 && - (n_groups % 2u) == 0u; - ds4_gpu_tensor *tp_attn_a = NULL; /* rank partials consumed directly */ - ds4_gpu_tensor *tp_attn_b = NULL; /* by the HC expand */ - const bool fuse_attn_out_hc = - !cuda_tp_attn && - g->tp_world < 2 && - layer->attn_output_a->type == DS4_TENSOR_Q8_0 && - layer->attn_output_b->type == DS4_TENSOR_Q8_0 && - !metal_graph_directional_steering_attn_enabled(g) && - !metal_graph_use_reference_attn_out_hc(); - const bool fuse_tp_attn_out_hc = - cuda_tp_attn && - !metal_graph_use_reference_attn_out_hc() && - g->cuda_tp_attn_out_hc_fuse; - if (ok && cuda_tp_attn_requested && !cuda_tp_attn) { - fprintf(stderr, - "ds4: CUDA decode TP cannot split attention output for tier %d " - "(partner=%d groups=%u)\n", - cuda_tp_home_tier, cuda_tp_partner_tier, n_groups); - ok = false; - } - if (ok && cuda_tp_attn) { - const uint32_t tp_groups = n_groups / 2u; - const uint64_t tp_heads_bytes = (uint64_t)tp_groups * group_dim * sizeof(float); - const uint64_t tp_heads_off = (uint64_t)tp_groups * group_dim * sizeof(float); - const bool cuda_tp_attn_peer_read = - !cuda_tp_attn_heads_active && - g->cuda_tp_attn_peer_read && - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; - ds4_gpu_tensor peer_heads_dst_view; - ds4_gpu_tensor peer_heads_src_view; - ds4_gpu_tensor *peer_heads_dst = NULL; - ds4_gpu_tensor *peer_heads_src = NULL; - if (cuda_tp_attn_heads_active) { - peer_heads_src = g->heads_by_tier[cuda_tp_partner_tier]; - } else if (!cuda_tp_attn_peer_read) { - ok = metal_graph_borrow_tensor_view( - &peer_heads_dst_view, - g->heads_by_tier[cuda_tp_partner_tier], - tp_heads_off, - tp_heads_bytes); - if (ok) peer_heads_dst = &peer_heads_dst_view; - } - if (ok) { - if (cuda_tp_attn_heads_active) { - peer_heads_src = g->heads_by_tier[cuda_tp_partner_tier]; - } else if (cuda_tp_attn_peer_read) { - peer_heads_src = metal_graph_heads(g); - } else { - ok = metal_graph_borrow_tensor_view(&peer_heads_src_view, - metal_graph_heads(g), - tp_heads_off, - tp_heads_bytes); - if (ok) peer_heads_src = &peer_heads_src_view; - } - } - ds4_gpu_tensor *peer_heads = cuda_tp_attn_peer_read ? - metal_graph_heads(g) : g->heads_by_tier[cuda_tp_partner_tier]; - ok = (cuda_tp_attn_heads_active || cuda_tp_attn_peer_read || peer_heads_dst) && - peer_heads_src && peer_heads && - g->attn_low_by_tier[cuda_tp_partner_tier] && - g->attn_out_by_tier[cuda_tp_partner_tier] && - g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; - if (ok && !cuda_tp_attn_heads_active && !cuda_tp_attn_peer_read) { - ok = ds4_gpu_tensor_copy_xdev(peer_heads_dst, peer_heads_src, - tp_heads_bytes) != 0; - } else if (ok && cuda_tp_attn_peer_read) { - ok = ds4_gpu_tensor_wait_xdev(peer_heads_src, cuda_tp_partner_tier) != 0; - } - if (ok) ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; - if (ok) { - ok = ds4_gpu_attention_output_q8_tp_tensor( - g->attn_out_by_tier[cuda_tp_partner_tier], - g->attn_low_by_tier[cuda_tp_partner_tier], - model->map, - model->size, - layer->attn_output_a->abs_offset, - layer->attn_output_b->abs_offset, - group_dim, - rank, - n_groups, - tp_groups, - tp_groups, - DS4_N_EMBD, - peer_heads) != 0; - } - if (ok) ok = ds4_gpu_set_current_device(cuda_tp_home_tier) == 0; - if (ok && fuse_tp_attn_out_hc) { - ok = ds4_gpu_attention_output_low_q8_tensor( - metal_graph_attn_low(g), - model->map, - model->size, - layer->attn_output_a->abs_offset, - group_dim, - rank, - tp_groups, - metal_graph_heads(g)) != 0; - } else if (ok) { - ok = ds4_gpu_attention_output_q8_tp_tensor( - metal_graph_attn_out(g), - metal_graph_attn_low(g), - model->map, - model->size, - layer->attn_output_a->abs_offset, - layer->attn_output_b->abs_offset, - group_dim, - rank, - n_groups, - 0, - tp_groups, - DS4_N_EMBD, - metal_graph_heads(g)) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy_xdev(g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - g->attn_out_by_tier[cuda_tp_partner_tier], - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - if (ok) cuda_tp_attn_peer = g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; - } - if (ok && fuse_tp_attn_out_hc) { - ok = ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( - metal_graph_after_attn_hc(g), - metal_graph_attn_out(g), - model->map, - model->size, - layer->attn_output_b->abs_offset, - (uint64_t)n_groups * rank, - DS4_N_EMBD, - 0, - (uint64_t)tp_groups * rank, - metal_graph_attn_low(g), - cuda_tp_attn_peer, - metal_graph_cur_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) cuda_tp_attn_hc_fused = true; - } - } else if (ok && fuse_attn_out_hc) { - ok = ds4_gpu_attention_output_low_q8_tensor(metal_graph_attn_low(g), - model->map, - model->size, - layer->attn_output_a->abs_offset, - group_dim, - rank, - n_groups, - metal_graph_heads(g)) != 0; - if (ok) { - ok = ds4_gpu_matmul_q8_0_hc_expand_tensor(metal_graph_after_attn_hc(g), - metal_graph_attn_out(g), - model->map, - model->size, - layer->attn_output_b->abs_offset, - (uint64_t)n_groups * rank, - DS4_N_EMBD, - metal_graph_attn_low(g), - metal_graph_cur_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - } else if (ok && g->tp_world == 2) { - /* Group-sliced attention output: this rank computes its half of the - * output groups and the matching k-window of the expand projection, - * leaving a partial block output in the gate slot. */ - const uint32_t tp_groups = n_groups / 2; - ok = metal_graph_attention_output_dense_quant_tp( - g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN], - metal_graph_attn_low(g), - g, - model, - layer->attn_output_a, - layer->attn_output_b, - group_dim, rank, - n_groups, - g->tp_rank * tp_groups, tp_groups, - DS4_N_EMBD, - metal_graph_heads(g)); - } else if (ok && layer->attn_output_a->type != DS4_TENSOR_Q8_0) { - ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? - g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); - ok = metal_graph_attention_output_dense_quant_low(metal_graph_attn_low(g), - g, - model, - layer->attn_output_a, - group_dim, - rank, - 0, - n_groups, - metal_graph_heads(g)); - if (ok) ok = metal_graph_matmul_dense_quant_tensor(attn_out_dst, - model, - layer->attn_output_b, - (uint64_t)n_groups * rank, - DS4_N_EMBD, - metal_graph_attn_low(g), - 1); - } else if (ok) { - ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? - g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); - ok = ds4_gpu_attention_output_q8_batch_tensor(attn_out_dst, - metal_graph_attn_low(g), - metal_graph_batch_group_tmp(g), - metal_graph_batch_low_tmp(g), - model->map, - model->size, - layer->attn_output_a->abs_offset, - layer->attn_output_b->abs_offset, - group_dim, rank, - n_groups, DS4_N_EMBD, - metal_graph_heads(g), 1) != 0; - } - if (ok && g->tp_world == 2) { - /* Gate ATTN: exchange the attention block output with the peer and - * rebuild the canonical sum (rank0 first, then rank1) in attn_out - * on both ranks — identical expression on both machines keeps them - * bit-exact. */ - const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN; - ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_ATTN) != 0; - if (ok) { - ds4_gpu_tensor *first = g->tp_rank == 0 ? g->tp_out[slot] : g->tp_in[slot]; - if (metal_graph_directional_steering_attn_enabled(g)) { - ds4_gpu_tensor *second = g->tp_rank == 0 ? g->tp_in[slot] : g->tp_out[slot]; - ok = ds4_gpu_add_tensor(metal_graph_attn_out(g), first, second, DS4_N_EMBD) != 0; - } else { - /* Combine folded into the HC expand below; attn_out is not - * materialized on this path. */ - tp_attn_a = first; - tp_attn_b = g->tp_rank == 0 ? g->tp_in[slot] : g->tp_out[slot]; - } - } - } - DS4_METAL_PROFILE_DECODE_STAGE("attn_output"); - if (ok) { - metal_graph_debug_dump_tensor("attn_low", metal_graph_attn_low(g), (uint64_t)n_groups * rank, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("attn_out", metal_graph_attn_out(g), DS4_N_EMBD, il, pos); - } - if (ok && metal_graph_directional_steering_attn_enabled(g)) { - ok = metal_graph_apply_directional_steering_attn(g, metal_graph_attn_out(g), il, 1); - } - if (ok && !fuse_attn_out_hc && !cuda_tp_attn_hc_fused) { - if (tp_attn_a) { - ok = ds4_gpu_hc_expand_add_tensor(metal_graph_after_attn_hc(g), tp_attn_a, tp_attn_b, - metal_graph_cur_hc(g), metal_graph_hc_post(g), metal_graph_hc_comb(g), - DS4_N_EMBD, DS4_N_HC) != 0; - } else if (cuda_tp_attn_peer) { - ok = ds4_gpu_hc_expand_add_tensor( - metal_graph_after_attn_hc(g), - metal_graph_attn_out(g), - cuda_tp_attn_peer, - metal_graph_cur_hc(g), - metal_graph_hc_post(g), - metal_graph_hc_comb(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else { - ok = ds4_gpu_hc_expand_tensor(metal_graph_after_attn_hc(g), metal_graph_attn_out(g), metal_graph_cur_hc(g), - metal_graph_hc_post(g), metal_graph_hc_comb(g), DS4_N_EMBD, DS4_N_HC) != 0; - } - } - DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_post"); - if (ok) { - metal_graph_debug_dump_tensor("hc_attn_post", metal_graph_after_attn_hc(g), hc_dim, il, pos); - } - if (ok && !tp_ablate_hcpre) { - ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_after_attn_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_hc_mix(g), model, layer->hc_ffn_fn, - hc_dim, mix_hc, metal_graph_flat_hc(g), 1); - } - if (ok && fuse_hc_norm) { - ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(metal_graph_ffn_cur(g), - metal_graph_ffn_norm(g), - metal_graph_hc_split(g), - metal_graph_hc_mix(g), - metal_graph_after_attn_hc(g), - model->map, - model->size, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset, - layer->ffn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - if (ok) { - ok = metal_graph_check_hc_norm_fusion("ffn", - metal_graph_ffn_cur(g), - metal_graph_ffn_norm(g), - metal_graph_hc_mix(g), - metal_graph_after_attn_hc(g), - model, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset, - layer->ffn_norm->abs_offset, - il, - pos); - } - } else if (ok) { - ok = metal_graph_decode_hc_pre(metal_graph_ffn_cur(g), - metal_graph_hc_split(g), - metal_graph_hc_mix(g), - metal_graph_after_attn_hc(g), - model, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset); - } - DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_pre"); - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre_mixes", metal_graph_hc_mix(g), mix_hc, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_weights", metal_graph_hc_pre(g), DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_post_weights", metal_graph_hc_post(g), DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_comb", metal_graph_hc_comb(g), (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre", metal_graph_ffn_cur(g), DS4_N_EMBD, il, pos); - } - if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_ffn_norm(g), metal_graph_ffn_cur(g), - model->map, model->size, - layer->ffn_norm->abs_offset, - DS4_N_EMBD, DS4_RMS_EPS) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("ffn_norm"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_norm", metal_graph_ffn_norm(g), DS4_N_EMBD, il, pos); - } - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t gate_expert_bytes DS4_MAYBE_UNUSED = expert_mid_dim * gate_row_bytes; - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - const uint64_t down_expert_bytes DS4_MAYBE_UNUSED = routed_out_dim * down_row_bytes; - if (ok && metal_graph_decode_cpu_router_applicable(g, layer)) { - ok = metal_graph_decode_cpu_router(g, model, layer, il, (uint32_t)token); - } else { - if (ok && !metal_graph_tp_ablate("router")) { - ok = metal_graph_matmul_plain_tensor(metal_graph_router_logits(g), model, layer->ffn_gate_inp, - DS4_N_EMBD, DS4_N_EXPERT, metal_graph_ffn_norm(g), 1); - if (ok) ok = ds4_gpu_router_select_tensor(metal_graph_router_selected(g), metal_graph_router_weights(g), metal_graph_router_probs(g), - model->map, model->size, - layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, - layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, - layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, - (uint32_t)token, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE, - 0, - 0, - layer->ffn_exp_probs_b != NULL, - layer->ffn_gate_tid2eid != NULL, - metal_graph_router_logits(g)) != 0; - } - if (ok) ok = metal_graph_decode_set_hash_selected_override(model, - layer, - il, - (uint32_t)token, - layer->ffn_gate_exps->bytes, - layer->ffn_down_exps->bytes, - g); - } - DS4_METAL_PROFILE_DECODE_STAGE("router"); - if (ok) ok = metal_graph_profile_router_selection(g, layer, il, pos); - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_logits", metal_graph_router_logits(g), DS4_N_EXPERT, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_probs", metal_graph_router_probs(g), DS4_N_EXPERT, il, pos); - metal_graph_debug_dump_i32_tensor("ffn_moe_topk", metal_graph_router_selected(g), DS4_N_EXPERT_USED, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", metal_graph_router_weights(g), DS4_N_EXPERT_USED, il, pos); - } - if (phase == METAL_DECODE_LAYER_TO_ROUTER) return ok; - } - const bool external_routed = phase == METAL_DECODE_LAYER_FROM_ROUTER; - const bool fuse_shared_gate_up = - !g->quality && - g->tp_world < 2 && - layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && - g->shared_gate_up_swiglu_fuse; - const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos); - const bool cuda_tp_shared_requested = g->cuda_tp_shared; - const bool cuda_tp_moe_requested = !external_routed && g->cuda_tp_moe; - const uint64_t shared_tp_local = shared_dim / 2u; - const uint64_t shared_tp_peer = shared_dim - shared_tp_local; - const uint64_t shared_q8_blocks = ((uint64_t)DS4_N_EMBD + 31u) / 32u; - const uint64_t shared_q8_x_bytes = shared_q8_blocks * 32u; - const uint64_t shared_q8_scale_offset = - (shared_q8_x_bytes + 15u) & ~15ull; - const uint64_t shared_q8_prequant_bytes DS4_MAYBE_UNUSED = - shared_q8_scale_offset + shared_q8_blocks * sizeof(float); - const bool cuda_tp_shared = - cuda_tp_shared_requested && - fuse_shared_gate_up && - !metal_graph_use_reference_shared_down_hc() && - cuda_tp_partner_tier >= 0 && - shared_tp_local != 0 && - shared_tp_peer != 0 && - (shared_tp_local % 32u) == 0 && - (shared_tp_peer % 32u) == 0 && - g->ffn_norm_by_tier[cuda_tp_partner_tier] && - g->shared_gate_by_tier[cuda_tp_partner_tier] && - g->shared_up_by_tier[cuda_tp_partner_tier] && - g->shared_mid_by_tier[cuda_tp_partner_tier] && - g->shared_out_by_tier[cuda_tp_partner_tier] && - g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; - const bool cuda_tp_shared_fold = - cuda_tp_shared && - g->cuda_tp_shared_fold && - !g->cuda_tp_ep && - cuda_tp_moe_requested && - !keep_ffn_out && - !metal_graph_directional_steering_ffn_enabled(g) && - !metal_graph_debug_wants("ffn_moe_out", il, pos) && - !metal_graph_debug_wants("ffn_shexp", il, pos); - const bool fuse_shared_down_hc = - g->tp_world < 2 && - layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && - !cuda_tp_shared && - !keep_ffn_out && - !metal_graph_use_reference_shared_down_hc(); - const bool cuda_tp_moe = - cuda_tp_moe_requested && - cuda_tp_partner_tier >= 0 && - (DS4_N_EXPERT_USED % 2u) == 0u && - g->ffn_norm_by_tier[cuda_tp_partner_tier] && - g->router_selected_by_tier[cuda_tp_partner_tier] && - g->router_weights_by_tier[cuda_tp_partner_tier] && - g->routed_gate_by_tier[cuda_tp_partner_tier] && - g->routed_up_by_tier[cuda_tp_partner_tier] && - g->routed_mid_by_tier[cuda_tp_partner_tier] && - g->routed_down_by_tier[cuda_tp_partner_tier] && - g->routed_out_by_tier[cuda_tp_partner_tier] && - g->tp_peer_tmp_by_tier[cuda_tp_home_tier] && - g->tp_peer_tmp_by_tier[cuda_tp_partner_tier]; - const bool cuda_tp_ep = cuda_tp_moe && g->cuda_tp_ep; - const bool cuda_tp_moe_delay_reduce = - cuda_tp_moe && - g->cuda_tp_moe_delay_reduce && - fuse_shared_down_hc && - !metal_graph_debug_wants("ffn_moe_out", il, pos); - bool cuda_tp_moe_peer_tmp = false; - bool cuda_tp_moe_peer_copy_deferred = false; - bool cuda_tp_ep_reduce_deferred = false; - bool cuda_tp_ep_fused_hc_reduce = false; - bool cuda_tp_ep_direct_return = false; - bool cuda_tp_ep_balanced_shared_mid DS4_MAYBE_UNUSED = false; - bool cuda_tp_ep_dual_prequant = false; - uint64_t cuda_tp_ep_return_bytes = 0; - bool cuda_tp_shared_fold_peer_tmp = false; - if (ok && cuda_tp_moe_requested && !cuda_tp_moe) { - fprintf(stderr, - "ds4: CUDA decode TP cannot split routed MoE for tier %d " - "(partner=%d experts=%u)\n", - cuda_tp_home_tier, cuda_tp_partner_tier, DS4_N_EXPERT_USED); - ok = false; - } - if (ok && cuda_tp_moe) { - const uint32_t tp_experts = cuda_tp_ep - ? DS4_N_EXPERT_USED : DS4_N_EXPERT_USED / 2u; - const uint64_t tp_selected_bytes = (uint64_t)tp_experts * sizeof(int32_t); - const uint64_t tp_weights_bytes = (uint64_t)tp_experts * sizeof(float); - const uint64_t peer_selected_offset = cuda_tp_ep ? 0 : tp_selected_bytes; - const uint64_t peer_weights_offset = cuda_tp_ep ? 0 : tp_weights_bytes; - ds4_gpu_tensor local_selected; - ds4_gpu_tensor local_weights; - ds4_gpu_tensor peer_selected_src; - ds4_gpu_tensor peer_weights_src; - ds4_gpu_tensor packed_peer_ffn_norm; - ds4_gpu_tensor packed_peer_selected; - ds4_gpu_tensor packed_peer_weights; - ds4_gpu_tensor direct_peer_down; - ds4_gpu_tensor *peer_down_output = NULL; - cuda_tp_ep_return_bytes = - (uint64_t)(g->cuda_tp_ep_pack_exact ? 4u : DS4_N_EXPERT_USED) * - DS4_N_EMBD * sizeof(float); - cuda_tp_ep_direct_return = - cuda_tp_ep && - metal_graph_cuda_tp_ep_direct_return_requested() && - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] && - metal_graph_borrow_tensor_view( - &direct_peer_down, - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - 0, - cuda_tp_ep_return_bytes); - if (cuda_tp_ep_direct_return) peer_down_output = &direct_peer_down; -#if !defined(__APPLE__) - /* Run the shared gate/up projection on the less-loaded EP rank. The - * two kernels use complementary predicates over the same top-k IDs; - * a partner result is ordered by the existing direct-return event. */ - cuda_tp_ep_balanced_shared_mid = - cuda_tp_ep_direct_return && - cuda_tp_moe_delay_reduce && - metal_graph_cuda_tp_ep_delay_reduce_requested() && - metal_graph_cuda_tp_ep_fused_shared_mid_requested() && - metal_graph_cuda_tp_ep_balanced_shared_mid_requested() && - fuse_shared_gate_up && - !cuda_tp_shared_requested && - !g->cuda_tp_moe_peer_read && - !g->cuda_tp_moe_peer_router && - !g->decode_stage_profile; - cuda_tp_ep_dual_prequant = - cuda_tp_ep_balanced_shared_mid && - metal_graph_cuda_tp_ep_dual_prequant_requested() && - metal_graph_shared_gate(g) && - metal_graph_shared_gate(g)->bytes >= shared_q8_prequant_bytes && - g->shared_gate_by_tier[cuda_tp_partner_tier] && - g->shared_gate_by_tier[cuda_tp_partner_tier]->bytes >= - shared_q8_prequant_bytes; -#endif - const bool cuda_tp_moe_peer_read = - g->cuda_tp_moe_peer_read && - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; - const bool cuda_tp_moe_peer_router = - !cuda_tp_moe_peer_read && - g->cuda_tp_moe_peer_router && - g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; - const ds4_gpu_tensor *peer_selected = g->router_selected_by_tier[cuda_tp_partner_tier]; - const ds4_gpu_tensor *peer_weights = g->router_weights_by_tier[cuda_tp_partner_tier]; - const ds4_gpu_tensor *peer_ffn_norm = g->ffn_norm_by_tier[cuda_tp_partner_tier]; - ok = metal_graph_borrow_tensor_view(&local_selected, - metal_graph_router_selected(g), - 0, - tp_selected_bytes) && - metal_graph_borrow_tensor_view(&local_weights, - metal_graph_router_weights(g), - 0, - tp_weights_bytes) && - metal_graph_borrow_tensor_view(&peer_selected_src, - metal_graph_router_selected(g), - peer_selected_offset, - tp_selected_bytes) && - metal_graph_borrow_tensor_view(&peer_weights_src, - metal_graph_router_weights(g), - peer_weights_offset, - tp_weights_bytes); - if (ok && cuda_tp_moe_peer_read) { - ok = ds4_gpu_tensor_wait_xdev(metal_graph_router_weights(g), - cuda_tp_partner_tier) != 0; - peer_selected = &peer_selected_src; - peer_weights = &peer_weights_src; - peer_ffn_norm = metal_graph_ffn_norm(g); - } else if (ok && cuda_tp_moe_peer_router) { - ok = ds4_gpu_tensor_copy_xdev(g->ffn_norm_by_tier[cuda_tp_partner_tier], - metal_graph_ffn_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - peer_selected = &peer_selected_src; - peer_weights = &peer_weights_src; - peer_ffn_norm = g->ffn_norm_by_tier[cuda_tp_partner_tier]; - } else if (ok && g->cuda_tp_moe_pack_handoff) { - const uint64_t norm_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - const uint64_t packed_selected_off = norm_bytes; - const uint64_t packed_weights_off = packed_selected_off + tp_selected_bytes; - const uint64_t packed_bytes = packed_weights_off + tp_weights_bytes; - ok = metal_graph_borrow_tensor_view(&packed_peer_ffn_norm, - g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], - 0, - norm_bytes) && - metal_graph_borrow_tensor_view(&packed_peer_selected, - g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], - packed_selected_off, - tp_selected_bytes) && - metal_graph_borrow_tensor_view(&packed_peer_weights, - g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], - packed_weights_off, - tp_weights_bytes); - if (ok) { - ok = ds4_gpu_moe_handoff_pack_tensor( - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - metal_graph_ffn_norm(g), - &peer_selected_src, - &peer_weights_src, - DS4_N_EMBD, - tp_experts) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy_xdev( - g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - packed_bytes) != 0; - } - peer_selected = &packed_peer_selected; - peer_weights = &packed_peer_weights; - peer_ffn_norm = &packed_peer_ffn_norm; - } else if (ok && g->cuda_tp_moe_copy3_handoff) { - ok = ds4_gpu_tensor_copy_xdev3( - g->ffn_norm_by_tier[cuda_tp_partner_tier], - metal_graph_ffn_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float), - g->router_selected_by_tier[cuda_tp_partner_tier], - &peer_selected_src, - tp_selected_bytes, - g->router_weights_by_tier[cuda_tp_partner_tier], - &peer_weights_src, - tp_weights_bytes) != 0; - } else if (ok) { - ok = ds4_gpu_tensor_copy_xdev(g->ffn_norm_by_tier[cuda_tp_partner_tier], - metal_graph_ffn_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_copy_xdev(g->router_selected_by_tier[cuda_tp_partner_tier], - &peer_selected_src, - tp_selected_bytes) != 0 && - ds4_gpu_tensor_copy_xdev(g->router_weights_by_tier[cuda_tp_partner_tier], - &peer_weights_src, - tp_weights_bytes) != 0; - } - - bool switched_to_partner = false; - if (ok) { - ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; - switched_to_partner = ok; - } - if (ok) { - if (cuda_tp_ep) { - ok = ds4_gpu_routed_moe_one_owned_tensor( - g->routed_out_by_tier[cuda_tp_partner_tier], - g->routed_gate_by_tier[cuda_tp_partner_tier], - g->routed_up_by_tier[cuda_tp_partner_tier], - g->routed_mid_by_tier[cuda_tp_partner_tier], - g->routed_down_by_tier[cuda_tp_partner_tier], - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - peer_selected, - peer_weights, - DS4_N_EXPERT, - tp_experts, - DS4_N_EXPERT / 2u, - DS4_N_EXPERT - DS4_N_EXPERT / 2u, - DS4_SWIGLU_CLAMP_EXP, - peer_ffn_norm, - peer_down_output, - g->cuda_tp_ep_pack_exact, - cuda_tp_ep_dual_prequant - ? g->shared_gate_by_tier[cuda_tp_partner_tier] - : NULL) != 0; - } else { - ok = ds4_gpu_routed_moe_one_tensor( - g->routed_out_by_tier[cuda_tp_partner_tier], - g->routed_gate_by_tier[cuda_tp_partner_tier], - g->routed_up_by_tier[cuda_tp_partner_tier], - g->routed_mid_by_tier[cuda_tp_partner_tier], - g->routed_down_by_tier[cuda_tp_partner_tier], - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - peer_selected, - peer_weights, - DS4_N_EXPERT, - tp_experts, - DS4_SWIGLU_CLAMP_EXP, - peer_ffn_norm, NULL, 0, false) != 0; - } - } -#if !defined(__APPLE__) - if (ok && cuda_tp_ep_balanced_shared_mid) { - ok = ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset, - layer->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - shared_dim, - peer_ffn_norm, - DS4_SWIGLU_CLAMP_EXP, - peer_selected, - cuda_tp_ep_dual_prequant - ? g->shared_gate_by_tier[cuda_tp_partner_tier] - : NULL, - DS4_N_EXPERT / 2u, - false) != 0; - } -#endif - if (switched_to_partner && ds4_gpu_set_current_device(cuda_tp_home_tier) != 0) { - ok = false; - } - if (ok) { - if (cuda_tp_ep) { - ok = ds4_gpu_routed_moe_one_owned_tensor( - metal_graph_routed_out(g), - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - &local_selected, - &local_weights, - DS4_N_EXPERT, - tp_experts, - 0, - DS4_N_EXPERT / 2u, - DS4_SWIGLU_CLAMP_EXP, - metal_graph_ffn_norm(g), - NULL, - false, - cuda_tp_ep_dual_prequant - ? metal_graph_shared_gate(g) - : NULL) != 0; - } else { - ok = ds4_gpu_routed_moe_one_tensor( - metal_graph_routed_out(g), - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - &local_selected, - &local_weights, - DS4_N_EXPERT, - tp_experts, - DS4_SWIGLU_CLAMP_EXP, - metal_graph_ffn_norm(g), NULL, 0, false) != 0; - } - } -#if !defined(__APPLE__) - if (ok && cuda_tp_ep_balanced_shared_mid) { - ok = ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset, - layer->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - DS4_SWIGLU_CLAMP_EXP, - &local_selected, - cuda_tp_ep_dual_prequant - ? metal_graph_shared_gate(g) - : NULL, - DS4_N_EXPERT / 2u, - true) != 0; - } -#endif - if (ok) { - if (cuda_tp_ep) { - if (cuda_tp_moe_delay_reduce && - metal_graph_cuda_tp_ep_delay_reduce_requested() && - !g->decode_stage_profile) { - cuda_tp_ep_reduce_deferred = true; - cuda_tp_ep_fused_hc_reduce = - g->cuda_tp_ep_pack_exact && - metal_graph_cuda_tp_ep_fused_hc_reduce_requested(); - } else { - ok = metal_graph_cuda_tp_ep_finish_reduce( - g, - cuda_tp_home_tier, - cuda_tp_partner_tier, - cuda_tp_ep_direct_return, - cuda_tp_ep_return_bytes, - true); - } - } else if (cuda_tp_moe_delay_reduce && - !cuda_tp_shared && !cuda_tp_shared_fold) { - /* Defer the peer routed-half copy until after the shared - * expert gate/up launch below: that work depends only on - * ffn_norm, so the home stream computes it while waiting - * for the partner instead of idling at the copy. */ - cuda_tp_moe_peer_copy_deferred = true; - } else if (cuda_tp_moe_delay_reduce) { - ok = ds4_gpu_tensor_copy_xdev( - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - g->routed_out_by_tier[cuda_tp_partner_tier], - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - cuda_tp_moe_peer_tmp = ok; - } else if (!cuda_tp_shared_fold) { - ok = ds4_gpu_add_xdev_tensor(metal_graph_routed_out(g), - metal_graph_routed_out(g), - g->routed_out_by_tier[cuda_tp_partner_tier], - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - DS4_N_EMBD) != 0; - } - } - } - /* Real TP split slices the shared expert by intermediate lanes, which - * needs the unfused gate/up/swiglu/down sequence. */ - const bool tp_split_shared = g->tp_world == 2; - const bool q4_selected_shared_overlap = - metal_graph_use_q4_selected_shared_overlap() && - metal_graph_decode_q4_selected_slots_expected(g, - layer, - layer->ffn_gate_exps->bytes, - layer->ffn_down_exps->bytes); - const bool iq2_selected_shared_overlap = - metal_graph_use_iq2_selected_shared_overlap(g) && - metal_graph_decode_iq2_selected_slots_expected(g, layer); - const bool cuda_selected_shared_overlap = - metal_graph_use_cuda_selected_shared_overlap(g) && - metal_graph_decode_cuda_selected_slots_expected(g, layer); - const bool overlap_selected_shared = - ok && - g->tp_world < 2 && - !decode_stage_profile && - !metal_graph_decode_cpu_router_applicable(g, layer) && - layer->ffn_gate_tid2eid == NULL && - getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL && - (q4_selected_shared_overlap || - iq2_selected_shared_overlap || - cuda_selected_shared_overlap); - const bool async_selected_load = - overlap_selected_shared && - ((iq2_selected_shared_overlap && - metal_graph_use_iq2_selected_async_load(g)) || - cuda_selected_shared_overlap); - const bool selected_readahead_shared_delay = - ok && - g->tp_world < 2 && - !overlap_selected_shared && - !decode_stage_profile && - metal_graph_use_iq2_selected_readahead_shared_delay(g) && - metal_graph_decode_iq2_selected_slots_expected(g, layer) && - !metal_graph_decode_cpu_router_applicable(g, layer) && - layer->ffn_gate_tid2eid == NULL && - getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; - const bool cuda_stream_selected_load = - ok && - !overlap_selected_shared && - !selected_readahead_shared_delay && - g->ssd_streaming && - metal_graph_decode_cuda_selected_slots_expected(g, layer) && - layer->ffn_gate_tid2eid == NULL && - getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; - if (cuda_stream_selected_load) { - ok = metal_graph_decode_cuda_selected_load(g, - model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - } - if (selected_readahead_shared_delay) { - if (ok) { - ok = metal_graph_decode_selected_readahead_override(g, - model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - } - if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), - metal_graph_shared_up(g), - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset, - layer->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - DS4_SWIGLU_CLAMP_EXP) != 0; - } else if (ok) { - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), - model, - layer->ffn_gate_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), - model, - layer->ffn_up_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), - shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); - if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - metal_graph_router_selected(g), metal_graph_router_weights(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), - NULL, - il, - false) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); - } - if (ok && fuse_shared_down_hc) { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), - metal_graph_shared_out(g), - model->map, - model->size, - layer->ffn_down_shexp->abs_offset, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - metal_graph_routed_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok) { - ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), - model, - layer->ffn_down_shexp, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - 1); - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); - } - if (ok && keep_ffn_out) { - ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; - } - if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), - metal_graph_ffn_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_post(g), - metal_graph_hc_comb(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), - metal_graph_routed_out(g), - metal_graph_shared_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); - } - return ok; - } - if (overlap_selected_shared) { - uint64_t selected_event = 0; - if (ok) ok = ds4_gpu_signal_selected_readback_ready(&selected_event) != 0; - metal_graph_selected_async_load async_load = {0}; - bool async_load_started = false; - const bool async_early_commit = - async_selected_load && - metal_graph_use_iq2_selected_async_early_commit(g); - if (ok && async_selected_load) { - ok = metal_graph_selected_async_load_start(&async_load, - g, - model, - layer, - il, - selected_event, - gate_expert_bytes, - down_expert_bytes); - async_load_started = ok; - } - if (ok && async_early_commit) { - ok = ds4_gpu_flush_commands() != 0; - } - if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), - metal_graph_shared_up(g), - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset, - layer->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - DS4_SWIGLU_CLAMP_EXP) != 0; - } else if (ok) { - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), - model, - layer->ffn_gate_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), - model, - layer->ffn_up_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), - shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); - if (ok && !fuse_shared_down_hc) { - ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), - model, - layer->ffn_down_shexp, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - 1); - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); - if (async_load_started) { - const bool flush_ok = ds4_gpu_flush_commands() != 0; - bool finish_ok = - metal_graph_selected_async_load_finish(&async_load); - if (!finish_ok && async_load.ids_ok) { - /* The worker read valid ids but could not stage the load - * (it is not allowed to wait on in-flight cache entries). - * This thread is, so retry the same load synchronously. */ - const ds4_gpu_stream_expert_table retry_table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - finish_ok = - ds4_gpu_stream_expert_cache_begin_selected_load( - &retry_table, - async_load.selected_ids, - DS4_N_EXPERT_USED) != 0 && - ds4_gpu_routed_moe_set_selected_override( - async_load.selected_ids, - DS4_N_EXPERT_USED) != 0; - } - ok = ok && flush_ok && finish_ok; - } else if (ok) { - ok = ds4_gpu_commit_and_wait_selected_readback(selected_event, - "selected-id shared-overlap") != 0; - } - if (ok && !async_load_started) { - int32_t selected_ids[DS4_MAX_EXPERT_USED]; - ok = ds4_gpu_tensor_read(metal_graph_router_selected(g), - 0, - selected_ids, - (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) != 0 && - ds4_gpu_routed_moe_set_selected_override(selected_ids, - DS4_N_EXPERT_USED) != 0; - if (ok) { - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - ok = ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - selected_ids, - DS4_N_EXPERT_USED) != 0; - } - } - if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - metal_graph_router_selected(g), metal_graph_router_weights(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), - NULL, - il, - false) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); - } - if (ok && fuse_shared_down_hc) { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), - metal_graph_shared_out(g), - model->map, - model->size, - layer->ffn_down_shexp->abs_offset, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - metal_graph_routed_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); - } - if (ok && keep_ffn_out) { - ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; - } - if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), - metal_graph_ffn_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_post(g), - metal_graph_hc_comb(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), - metal_graph_routed_out(g), - metal_graph_shared_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); - } - return ok; - } - /* Under the TP split the routed experts run after the shared expert so - * the sum6 kernel can fold the shared partial and write the slab slot - * directly (no separate local add). */ - const bool tp_fold_ffn = tp_split_shared && - !keep_ffn_out && - !metal_graph_directional_steering_ffn_enabled(g); - if (ok && !tp_fold_ffn && !cuda_tp_moe) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - metal_graph_router_selected(g), metal_graph_router_weights(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), - NULL, - il, - false) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), - (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); - } - if (phase == METAL_DECODE_LAYER_TO_SHARED_MID || - phase == METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID) { - return ok; - } - if (ok && tp_split_shared) { - /* Shared expert lane slice: the fused gate/up/swiglu kernel covers - * this rank's half of the intermediate (row slicing is pure offset - * math), compact at the buffer base; the down k-slice below turns - * it into a partial output. */ - const uint32_t tp_half = shared_dim / 2; - uint64_t shexp_row_bytes = 0; - ok = metal_graph_dense_quant_row_bytes(layer->ffn_gate_shexp, - DS4_N_EMBD, - &shexp_row_bytes) && - layer->ffn_gate_shexp->type == layer->ffn_up_shexp->type; - const uint64_t tp_lane_off = (uint64_t)g->tp_rank * tp_half * shexp_row_bytes; - ok = ok && (tp_half % 32u) == 0; - if (!ok) { - fprintf(stderr, "ds4: TP shared expert width %u is not sliceable\n", shared_dim); - } - if (ok && layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), - metal_graph_shared_up(g), - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset + tp_lane_off, - layer->ffn_up_shexp->abs_offset + tp_lane_off, - DS4_N_EMBD, - tp_half, - metal_graph_ffn_norm(g), - DS4_SWIGLU_CLAMP_EXP) != 0; - } else if (ok) { - ok = metal_graph_matmul_dense_quant_abs(metal_graph_shared_gate(g), - model, - layer->ffn_gate_shexp, - layer->ffn_gate_shexp->abs_offset + tp_lane_off, - DS4_N_EMBD, - tp_half, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = metal_graph_matmul_dense_quant_abs(metal_graph_shared_up(g), - model, - layer->ffn_up_shexp, - layer->ffn_up_shexp->abs_offset + tp_lane_off, - DS4_N_EMBD, - tp_half, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), - tp_half, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; - } - } else if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), - metal_graph_shared_up(g), - metal_graph_shared_mid(g), - model->map, - model->size, - layer->ffn_gate_shexp->abs_offset, - layer->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - DS4_SWIGLU_CLAMP_EXP) != 0; - } else { - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), - model, - layer->ffn_gate_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), - model, - layer->ffn_up_shexp, - DS4_N_EMBD, - shared_dim, - metal_graph_ffn_norm(g), - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), - shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); - if (ok && cuda_tp_ep_reduce_deferred) { - ok = metal_graph_cuda_tp_ep_finish_reduce( - g, - cuda_tp_home_tier, - cuda_tp_partner_tier, - cuda_tp_ep_direct_return, - cuda_tp_ep_return_bytes, - !cuda_tp_ep_fused_hc_reduce); - } - if (ok && cuda_tp_moe_peer_copy_deferred) { - ok = ds4_gpu_tensor_copy_xdev( - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - g->routed_out_by_tier[cuda_tp_partner_tier], - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - cuda_tp_moe_peer_tmp = ok; - } - if (ok && cuda_tp_shared_fold) { - bool switched_to_partner = false; - ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; - switched_to_partner = ok; - if (ok) { - ok = ds4_gpu_add_tensor( - g->routed_out_by_tier[cuda_tp_partner_tier], - g->routed_out_by_tier[cuda_tp_partner_tier], - g->shared_out_by_tier[cuda_tp_partner_tier], - DS4_N_EMBD) != 0; - } - if (switched_to_partner && - ds4_gpu_set_current_device(cuda_tp_home_tier) != 0) { - ok = false; - } - if (ok) { - ok = ds4_gpu_add_tensor(metal_graph_routed_out(g), - metal_graph_routed_out(g), - metal_graph_shared_out(g), - DS4_N_EMBD) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy_xdev( - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - g->routed_out_by_tier[cuda_tp_partner_tier], - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - cuda_tp_shared_fold_peer_tmp = ok; - } - if (ok) { - ok = cuda_tp_shared_fold_peer_tmp && - ds4_gpu_hc_expand_add_split_tensor( - metal_graph_after_ffn_hc(g), - metal_graph_routed_out(g), - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - } else if (ok && cuda_tp_shared) { - /* shared_out already contains the reduced local and partner partials. */ - } else if (ok && cuda_tp_ep_fused_hc_reduce) { - ok = ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor( - metal_graph_after_ffn_hc(g), - metal_graph_shared_out(g), - model->map, - model->size, - layer->ffn_down_shexp->abs_offset, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - metal_graph_routed_down(g), - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - metal_graph_router_selected(g), - DS4_N_EXPERT / 2u, - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok && fuse_shared_down_hc) { - if (cuda_tp_moe_peer_tmp) { - ok = ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( - metal_graph_after_ffn_hc(g), - metal_graph_shared_out(g), - model->map, - model->size, - layer->ffn_down_shexp->abs_offset, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - metal_graph_routed_out(g), - g->tp_peer_tmp_by_tier[cuda_tp_home_tier], - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor( - metal_graph_after_ffn_hc(g), - metal_graph_shared_out(g), - model->map, - model->size, - layer->ffn_down_shexp->abs_offset, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - metal_graph_routed_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - } else if (ok && tp_split_shared) { - ok = metal_graph_matmul_dense_quant_kslice(metal_graph_shared_out(g), - model, - layer->ffn_down_shexp, - shared_dim, - (uint64_t)g->tp_rank * (shared_dim / 2), - shared_dim / 2, - DS4_N_EMBD, - metal_graph_shared_mid(g), - 0); - } else if (ok) { - ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), - model, - layer->ffn_down_shexp, - shared_dim, - DS4_N_EMBD, - metal_graph_shared_mid(g), - 1); - } - DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); - if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); - } - if (ok && tp_fold_ffn) { - ok = ds4_gpu_routed_moe_one_tensor( - g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN], - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - metal_graph_router_selected(g), metal_graph_router_weights(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), - metal_graph_shared_out(g), - il, - false) != 0; - DS4_METAL_PROFILE_DECODE_STAGE("routed_moe_folded"); - } - ds4_gpu_tensor *tp_ffn_a = NULL; /* rank0/rank1 partials consumed */ - ds4_gpu_tensor *tp_ffn_b = NULL; /* directly by the HC expand */ - if (ok && g->tp_world == 2) { - /* Gate FFN: local partial = shared expert + owned routed experts. - * The HC expand below already sums two block vectors, so after the - * exchange the two rank partials feed it directly (canonical rank - * order) with no separate combine dispatch. The paths that need - * the materialized sum (ffn_out consumers) still builds it in - * routed_out. */ - const uint32_t tp_slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; - if (!tp_fold_ffn) { - ok = ds4_gpu_add_tensor(g->tp_out[tp_slot], metal_graph_shared_out(g), metal_graph_routed_out(g), - DS4_N_EMBD) != 0; - } - if (ok) ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_FFN) != 0; - if (ok) { - ds4_gpu_tensor *first = g->tp_rank == 0 ? g->tp_out[tp_slot] : g->tp_in[tp_slot]; - ds4_gpu_tensor *second = g->tp_rank == 0 ? g->tp_in[tp_slot] : g->tp_out[tp_slot]; - if (keep_ffn_out || metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_add_tensor(metal_graph_routed_out(g), first, second, DS4_N_EMBD) != 0; - } else { - tp_ffn_a = first; - tp_ffn_b = second; - } - } - } - if (ok && keep_ffn_out) { - ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(metal_graph_ffn_out(g), - g->tp_world == 2 ? g->tp_zero : metal_graph_shared_out(g), - metal_graph_routed_out(g), DS4_N_EMBD) != 0; - } - if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), - metal_graph_ffn_out(g), - metal_graph_after_attn_hc(g), - metal_graph_hc_post(g), - metal_graph_hc_comb(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok && !cuda_tp_shared_fold && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), - tp_ffn_a ? tp_ffn_a : metal_graph_routed_out(g), - tp_ffn_a ? tp_ffn_b : - (g->tp_world == 2 ? g->tp_zero : metal_graph_shared_out(g)), - metal_graph_after_attn_hc(g), - metal_graph_hc_split(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); -#undef DS4_METAL_PROFILE_DECODE_STAGE - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); - } - return ok; -} - -static bool metal_graph_encode_decode_layer( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos, - ds4_gpu_tensor *raw_cache, - uint32_t raw_cap, - uint32_t raw_row, - uint32_t n_raw, - int token) { - return metal_graph_encode_decode_layer_phase( - g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, - token, METAL_DECODE_LAYER_FULL); -} - -static bool metal_graph_output_logits_head_matmul( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - ds4_gpu_tensor *norm_full, - ds4_gpu_tensor *dst_logits, - uint32_t n_tokens, - uint64_t vocab_dim); - -/* Encode the final HC collapse, output norm, and vocab projection on Metal. */ -static bool metal_graph_encode_output_head( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint64_t vocab_dim) { - /* switch to head_tier before the output-head pipeline. - * Single-tier (placement == NULL): no-op (head_tier == 0 == active_tier). - * Note: head_tier was captured in metal_graph_alloc_raw_cap; this - * helper consults it directly (and also covers the case where the - * preceding decode layer ran on a different tier — copy_xdev ferries - * the active cur_hc across the boundary). */ - if (g->placement) { - if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; - } - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const bool output_stage_profile = g->output_stage_profile; - double output_stage_t0 = output_stage_profile ? now_sec() : 0.0; -#define DS4_METAL_PROFILE_OUTPUT_STAGE(name) do { \ - if (ok && output_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary("output", (name), DS4_N_LAYER, 0, 1, &output_stage_t0); \ - } \ - } while (0) - bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - DS4_METAL_PROFILE_OUTPUT_STAGE("hc_flat_norm"); - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), - model->map, - model->size, - weights->output_hc_fn->abs_offset, - hc_dim, - DS4_N_HC, - metal_graph_flat_hc(g), - 1) != 0; - DS4_METAL_PROFILE_OUTPUT_STAGE("hc_pre"); - if (ok) { - metal_graph_debug_dump_tensor("result_hc_pre", metal_graph_output_pre(g), DS4_N_HC, DS4_N_LAYER, 0); - } - if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), - metal_graph_output_pre(g), - model->map, - model->size, - weights->output_hc_scale->abs_offset, - weights->output_hc_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weights"); - if (ok) { - metal_graph_debug_dump_tensor("result_hc_weights", metal_graph_output_weights(g), DS4_N_HC, DS4_N_LAYER, 0); - } - bool output_sum_norm_fused = false; -#if defined(__APPLE__) - if (ok) { - output_sum_norm_fused = - ds4_gpu_hc_weighted_sum_norm_tensor( - metal_graph_output_embd(g), - metal_graph_output_norm(g), - metal_graph_cur_hc(g), - metal_graph_output_weights(g), - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_RMS_EPS) != 0; - if (!output_sum_norm_fused && - getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL) { - ok = false; - } - } -#endif - if (ok && !output_sum_norm_fused) { - ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), - metal_graph_cur_hc(g), - metal_graph_output_weights(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weighted_sum"); - if (ok) { - metal_graph_debug_dump_tensor("result_hc", metal_graph_output_embd(g), DS4_N_EMBD, DS4_N_LAYER, 0); - } - if (ok && !output_sum_norm_fused) { - ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), - metal_graph_output_embd(g), - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - } - DS4_METAL_PROFILE_OUTPUT_STAGE("output_norm"); - if (ok) { - metal_graph_debug_dump_tensor("result_norm", metal_graph_output_norm(g), DS4_N_EMBD, DS4_N_LAYER, 0); - } - if (ok && g->tp_world == 2 && g->tp_logits_half) { - /* Vocab-split: this rank computes its half of the head rows into - * its logits view; the halves are bit-identical to the full head - * (same kernel, same rows) and the worker ships its half to the - * leader after the eval. */ - const uint64_t tp_vhalf = vocab_dim / 2u; - uint64_t head_row_bytes = 0; - ok = metal_graph_dense_quant_row_bytes(weights->output, - DS4_N_EMBD, - &head_row_bytes); - if (ok) ok = metal_graph_matmul_dense_quant_abs(g->tp_logits_half, - model, - weights->output, - weights->output->abs_offset + - (uint64_t)g->tp_rank * tp_vhalf * head_row_bytes, - DS4_N_EMBD, - tp_vhalf, - metal_graph_output_norm(g), - 1); - } else if (ok && g->cuda_tp_ep && g->cuda_tp_output) { - ok = metal_graph_output_logits_head_matmul( - g, model, weights, metal_graph_output_norm(g), - metal_graph_logits(g), 1, vocab_dim); - } else if (ok) { - ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), - model, - weights->output, - DS4_N_EMBD, - vocab_dim, - metal_graph_output_norm(g), - 1); - } - if (ok) { - metal_graph_debug_dump_tensor("result_output", metal_graph_logits(g), vocab_dim, DS4_N_LAYER, 0); - } -#undef DS4_METAL_PROFILE_OUTPUT_STAGE - return ok; -} - -/* Greedy-only output head: compute one local top-1 candidate per output TP - * split and leave the full split logits on their owning tiers. This avoids - * gathering the whole vocabulary row back to the head tier when the caller only - * needs the next argmax token. */ -static bool metal_graph_encode_output_head_split_top1( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint64_t vocab_dim, - int cuda_tp_output_tiers[DS4_MAX_GPUS], - uint32_t *cuda_tp_output_ways_out) { - if (!g || !model || !weights || - !cuda_tp_output_tiers || !cuda_tp_output_ways_out || - vocab_dim > UINT32_MAX) { - return false; - } - *cuda_tp_output_ways_out = 0; - - if (g->placement) { - if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; - } - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), - metal_graph_cur_hc(g), - (uint32_t)hc_dim, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), - model->map, - model->size, - weights->output_hc_fn->abs_offset, - hc_dim, - DS4_N_HC, - metal_graph_flat_hc(g), - 1) != 0; - if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), - metal_graph_output_pre(g), - model->map, - model->size, - weights->output_hc_scale->abs_offset, - weights->output_hc_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), - metal_graph_cur_hc(g), - metal_graph_output_weights(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), - metal_graph_output_embd(g), - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - if (!ok) return false; - - const uint32_t cuda_tp_output_ways = - g->cuda_tp_output ? metal_graph_cuda_tp_output_tiers(g, cuda_tp_output_tiers) : 0; - const bool cuda_tp_output = - g->cuda_tp_output && - cuda_tp_output_ways >= 2u && - weights->output->type == DS4_TENSOR_Q8_0 && - weights->output->ndim == 2 && - weights->output->dim[0] == DS4_N_EMBD && - weights->output->dim[1] == vocab_dim && - vocab_dim >= 2; - if (!cuda_tp_output) return false; - - for (uint32_t i = 0; i < cuda_tp_output_ways; i++) { - const int t = cuda_tp_output_tiers[i]; - if (t < 0 || t >= DS4_MAX_GPUS || - !g->output_norm_by_tier[t] || - !g->logits_by_tier[t] || - !g->comp_selected_by_tier[t] || - !g->comp_mask_by_tier[t]) { - return false; - } - } - - const bool fused_top1 = metal_graph_cuda_output_fused_top1_requested(); - const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); - uint64_t split_start[DS4_MAX_GPUS] = {0}; - uint64_t split_count[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor split_logits[DS4_MAX_GPUS]; - memset(split_logits, 0, sizeof(split_logits)); - - for (uint32_t i = 0; ok && i < cuda_tp_output_ways; i++) { - const int t = cuda_tp_output_tiers[i]; - split_start[i] = (vocab_dim * (uint64_t)i) / cuda_tp_output_ways; - const uint64_t split_end = - (vocab_dim * (uint64_t)(i + 1u)) / cuda_tp_output_ways; - split_count[i] = split_end - split_start[i]; - if (split_count[i] == 0 || split_count[i] > UINT32_MAX) { - ok = false; - break; - } - if (fused_top1) { - if (t != g->head_tier) { - ok = ds4_gpu_tensor_copy_xdev( - g->output_norm_by_tier[t], - metal_graph_output_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - } - } else if (t == g->head_tier) { - ok = metal_graph_borrow_tensor_view(&split_logits[i], - metal_graph_logits(g), - split_start[i] * sizeof(float), - split_count[i] * sizeof(float)); - } else { - ok = metal_graph_borrow_tensor_view(&split_logits[i], - g->logits_by_tier[t], - 0, - split_count[i] * sizeof(float)) && - ds4_gpu_tensor_copy_xdev( - g->output_norm_by_tier[t], - metal_graph_output_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - } - } - - for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { - const int t = cuda_tp_output_tiers[i]; - ok = ds4_gpu_set_current_device(t) == 0; - if (ok && fused_top1) { - ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[t], - g->comp_mask_by_tier[t], - model->map, - model->size, - weights->output->abs_offset + - split_start[i] * row_bytes, - DS4_N_EMBD, - split_count[i], - g->output_norm_by_tier[t], - (uint32_t)split_start[i]) != 0; - } else if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[i], - model->map, - model->size, - weights->output->abs_offset + - split_start[i] * row_bytes, - DS4_N_EMBD, - split_count[i], - g->output_norm_by_tier[t], - 1) != 0; - } - if (ok && !fused_top1) { - ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[t], - g->comp_mask_by_tier[t], - &split_logits[i], - (uint32_t)split_count[i], - 1, - (uint32_t)split_start[i]) != 0; - } - } - if (ok) ok = ds4_gpu_set_current_device(g->head_tier) == 0; - if (ok && fused_top1) { - ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[g->head_tier], - g->comp_mask_by_tier[g->head_tier], - model->map, - model->size, - weights->output->abs_offset, - DS4_N_EMBD, - split_count[0], - metal_graph_output_norm(g), - (uint32_t)split_start[0]) != 0; - } else if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[0], - model->map, - model->size, - weights->output->abs_offset, - DS4_N_EMBD, - split_count[0], - metal_graph_output_norm(g), - 1) != 0; - } - if (ok && !fused_top1) { - ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[g->head_tier], - g->comp_mask_by_tier[g->head_tier], - &split_logits[0], - (uint32_t)split_count[0], - 1, - (uint32_t)split_start[0]) != 0; - } - for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { - const int t = cuda_tp_output_tiers[i]; - ds4_gpu_tensor head_id_dst; - ds4_gpu_tensor head_value_dst; - ok = metal_graph_borrow_tensor_view(&head_id_dst, - g->comp_selected_by_tier[g->head_tier], - (uint64_t)i * sizeof(uint32_t), - sizeof(uint32_t)) && - metal_graph_borrow_tensor_view(&head_value_dst, - g->comp_mask_by_tier[g->head_tier], - (uint64_t)i * sizeof(float), - sizeof(float)) && - ds4_gpu_tensor_copy_xdev3(&head_id_dst, - g->comp_selected_by_tier[t], - sizeof(uint32_t), - &head_value_dst, - g->comp_mask_by_tier[t], - sizeof(float), - NULL, - NULL, - 0) != 0; - } - if (ok) { - ok = ds4_gpu_set_current_device(g->head_tier) == 0; - *cuda_tp_output_ways_out = cuda_tp_output_ways; - } - return ok; -} - -static bool metal_graph_read_output_split_top1( - ds4_gpu_graph *g, - uint32_t output_ways, - int *top_id) { - if (!g || !top_id || output_ways == 0 || output_ways > DS4_MAX_GPUS) { - return false; - } - bool have_best = false; - uint32_t best_id = 0; - float best_value = 0.0f; - uint32_t cand_ids[DS4_MAX_GPUS] = {0}; - float cand_values[DS4_MAX_GPUS] = {0.0f}; - bool ok = ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], - 0, - cand_ids, - (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && - ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], - 0, - cand_values, - (uint64_t)output_ways * sizeof(cand_values[0])) != 0; - for (uint32_t i = 0; ok && i < output_ways; i++) { - const uint32_t cand_id = cand_ids[i]; - const float cand_value = cand_values[i]; - if (!have_best || - cand_value > best_value || - (cand_value == best_value && cand_id < best_id)) { - have_best = true; - best_id = cand_id; - best_value = cand_value; - } - } - ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; - if (ok) *top_id = (int)best_id; - return ok; -} - -/* Batched output head for speculative verification. - * - * A target verifier only needs top-1 ids for intermediate draft rows and full - * logits for the last accepted row. Running the normal one-row output head in - * a loop serializes the HC collapse, output norm, and Q8 vocab projection. For - * tiny MTP suffixes we instead process all rows together and let the GPU reduce - * each row to a top id; the CPU reads back just those ids plus the last row's - * logits needed to continue the exact target stream. */ -/* Shared vocab-head matmul: pads small batches to 8 rows for the exact-mma Q8 - * kernel and shards the vocabulary across output-TP tiers. */ -static bool metal_graph_output_logits_head_matmul( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - ds4_gpu_tensor *norm_full, - ds4_gpu_tensor *dst_logits, - uint32_t n_tokens, - uint64_t vocab_dim) { - if (!g || !model || !weights || !norm_full || n_tokens == 0 || - !dst_logits || - ds4_gpu_tensor_bytes(dst_logits) < - (uint64_t)n_tokens * vocab_dim * sizeof(float)) { - return false; - } - const uint32_t head_rows = - (n_tokens > 1 && n_tokens < 8 && - ds4_gpu_tensor_bytes(dst_logits) >= 8u * vocab_dim * sizeof(float) && - ds4_gpu_tensor_bytes(norm_full) >= - 8u * DS4_N_EMBD * sizeof(float)) ? 8u : n_tokens; - ds4_gpu_tensor *output_norm = - ds4_gpu_tensor_view(norm_full, - 0, - (uint64_t)head_rows * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *logits = - ds4_gpu_tensor_view(dst_logits, - 0, - (uint64_t)head_rows * vocab_dim * sizeof(float)); - bool ok = output_norm && logits; - if (ok && head_rows > n_tokens) { - ds4_gpu_tensor *pad = - ds4_gpu_tensor_view(norm_full, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float), - (uint64_t)(head_rows - n_tokens) * DS4_N_EMBD * - sizeof(float)); - ok = pad && - ds4_gpu_tensor_fill_f32(pad, - 0.0f, - (uint64_t)(head_rows - n_tokens) * - DS4_N_EMBD) != 0; - ds4_gpu_tensor_free(pad); - } - /* Output TP for the speculative batch, mirroring the decode head: each - * device matmuls its VRAM-resident vocab shard. Shard outputs land - * compactly in logits_by_tier[t] ([head_rows x split]) and are gathered - * into spec_logits rows. */ - int tp_tiers[DS4_MAX_GPUS] = {0}; - const uint32_t tp_ways = (ok && g->cuda_tp_output) - ? metal_graph_cuda_tp_output_tiers(g, tp_tiers) : 0; - bool tp_ok = ok && tp_ways >= 2u && - weights->output->type == DS4_TENSOR_Q8_0 && - weights->output->ndim == 2 && - weights->output->dim[0] == DS4_N_EMBD && - weights->output->dim[1] == vocab_dim && - head_rows <= DS4_DSPARK_MAX_BLOCK_SIZE && - getenv("DS4_DSPARK_VERIFY_HEAD_NO_TP") == NULL; - for (uint32_t i = 0; tp_ok && i < tp_ways; i++) { - const int t = tp_tiers[i]; - tp_ok = t >= 0 && t < DS4_MAX_GPUS && - g->logits_by_tier[t] && - ds4_gpu_tensor_bytes(g->logits_by_tier[t]) >= - (uint64_t)head_rows * - ((vocab_dim + tp_ways - 1u) / tp_ways) * - sizeof(float) && - (t == g->active_tier || - (g->batch_ffn_norm_by_tier[t] && - ds4_gpu_tensor_bytes(g->batch_ffn_norm_by_tier[t]) >= - (uint64_t)head_rows * DS4_N_EMBD * sizeof(float))); - } - if (tp_ok) { - const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); - const int home_tier = g->active_tier; - uint64_t split_start[DS4_MAX_GPUS] = {0}; - uint64_t split_count[DS4_MAX_GPUS] = {0}; - for (uint32_t i = 0; ok && i < tp_ways; i++) { - const int t = tp_tiers[i]; - split_start[i] = (vocab_dim * (uint64_t)i) / tp_ways; - const uint64_t split_end = - (vocab_dim * (uint64_t)(i + 1u)) / tp_ways; - split_count[i] = split_end - split_start[i]; - if (split_count[i] == 0) { ok = false; break; } - if (t != home_tier) { - ok = ds4_gpu_tensor_copy_xdev( - g->batch_ffn_norm_by_tier[t], - output_norm, - (uint64_t)head_rows * DS4_N_EMBD * - sizeof(float)) != 0; - } - } - for (uint32_t i = 0; ok && i < tp_ways; i++) { - const int t = tp_tiers[i]; - ok = ds4_gpu_set_current_device(t) == 0; - if (!ok) break; - ds4_gpu_tensor *shard_out = - ds4_gpu_tensor_view(g->logits_by_tier[t], - 0, - (uint64_t)head_rows * split_count[i] * - sizeof(float)); - ds4_gpu_tensor *shard_in = t == home_tier ? - NULL : - ds4_gpu_tensor_view(g->batch_ffn_norm_by_tier[t], - 0, - (uint64_t)head_rows * DS4_N_EMBD * - sizeof(float)); - ok = shard_out && (t == home_tier || shard_in) && - ds4_gpu_matmul_q8_0_tensor(shard_out, - model->map, - model->size, - weights->output->abs_offset + - split_start[i] * row_bytes, - DS4_N_EMBD, - split_count[i], - t == home_tier ? output_norm : - shard_in, - head_rows) != 0; - ds4_gpu_tensor_free(shard_in); - ds4_gpu_tensor_free(shard_out); - } - if (ok) ok = ds4_gpu_set_current_device(home_tier) == 0; - for (uint32_t i = 0; ok && i < tp_ways; i++) { - const int t = tp_tiers[i]; - for (uint32_t r = 0; ok && r < n_tokens; r++) { - ds4_gpu_tensor *dst = - ds4_gpu_tensor_view(dst_logits, - ((uint64_t)r * vocab_dim + - split_start[i]) * sizeof(float), - split_count[i] * sizeof(float)); - ds4_gpu_tensor *src = - ds4_gpu_tensor_view(g->logits_by_tier[t], - (uint64_t)r * split_count[i] * - sizeof(float), - split_count[i] * sizeof(float)); - ok = dst && src && - ds4_gpu_tensor_copy_xdev(dst, - src, - split_count[i] * - sizeof(float)) != 0; - ds4_gpu_tensor_free(src); - ds4_gpu_tensor_free(dst); - } - } - } else if (ok && !(g->cuda_tp_ep && g->cuda_tp_output)) { - ok = ds4_gpu_matmul_q8_0_tensor(logits, - model->map, - model->size, - weights->output->abs_offset, - DS4_N_EMBD, - vocab_dim, - output_norm, - head_rows) != 0; - } else if (ok) { - /* The expert-parallel cache stores only output vocabulary shards, so - * a single-device full-head fallback would access uncached weights. */ - ok = false; - } - ds4_gpu_tensor_free(logits); - ds4_gpu_tensor_free(output_norm); - return ok; -} - -static bool metal_graph_encode_output_head_batch( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t n_tokens, - uint64_t vocab_dim) { - if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - ds4_gpu_tensor *output_pre = NULL; - ds4_gpu_tensor *output_weights = NULL; - ds4_gpu_tensor *output_embd = NULL; - ds4_gpu_tensor *output_norm = NULL; - - bool ok = true; - output_pre = ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), - 0, - (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); - output_weights = ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - 0, - (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); - output_embd = ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), - 0, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - output_norm = ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), - 0, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - ok = output_pre && output_weights && output_embd && output_norm; - - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), - metal_graph_batch_cur_hc(g), - (uint32_t)hc_dim, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(output_pre, - model->map, - model->size, - weights->output_hc_fn->abs_offset, - hc_dim, - DS4_N_HC, - metal_graph_batch_flat_hc(g), - n_tokens) != 0; - if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, - output_pre, - model->map, - model->size, - weights->output_hc_scale->abs_offset, - weights->output_hc_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, - metal_graph_batch_cur_hc(g), - output_weights, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, - output_embd, - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_output_logits_head_matmul( - g, model, weights, metal_graph_batch_ffn_norm(g), - g->spec_logits, n_tokens, vocab_dim); - - ds4_gpu_tensor_free(output_norm); - ds4_gpu_tensor_free(output_embd); - ds4_gpu_tensor_free(output_weights); - ds4_gpu_tensor_free(output_pre); - return ok; -} - -static bool metal_graph_matmul_plain_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (w->type == DS4_TENSOR_F16) { - return ds4_gpu_matmul_f16_tensor(out, model->map, model->size, - w->abs_offset, in_dim, out_dim, x, n_tok) != 0; - } - if (w->type == DS4_TENSOR_F32) { - return ds4_gpu_matmul_f32_tensor(out, model->map, model->size, - w->abs_offset, in_dim, out_dim, x, n_tok) != 0; - } - if (w->type == DS4_TENSOR_Q8_0) { - return ds4_gpu_matmul_q8_0_tensor(out, model->map, model->size, - w->abs_offset, in_dim, out_dim, x, n_tok) != 0; - } - if (tensor_type_is_dense_quant(w->type)) { - return ds4_gpu_matmul_quant_tensor(out, - model->map, - model->size, - w->abs_offset, - w->type, - in_dim, - out_dim, - x, - n_tok) != 0; - } - fprintf(stderr, "ds4: Metal plain matmul does not support %s\n", tensor_type_name(w->type)); - return false; -} - -static bool metal_graph_dense_quant_row_bytes( - const ds4_tensor *w, - uint64_t in_dim, - uint64_t *row_bytes) { - if (row_bytes) *row_bytes = 0; - if (!w || !row_bytes || !tensor_type_is_dense_quant(w->type)) return false; - return tensor_nbytes(w->type, in_dim, row_bytes); -} - -static bool metal_graph_matmul_dense_quant_abs( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (!w || !tensor_type_is_dense_quant(w->type)) return false; - return ds4_gpu_matmul_quant_tensor(out, - model->map, - model->size, - weight_offset, - w->type, - in_dim, - out_dim, - x, - n_tok) != 0; -} - -static bool metal_graph_matmul_dense_quant_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (!w) return false; - return metal_graph_matmul_dense_quant_abs(out, - model, - w, - w->abs_offset, - in_dim, - out_dim, - x, - n_tok); -} - -static bool metal_graph_matmul_dense_quant_kslice( - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - 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) { - if (!w || !tensor_type_is_dense_quant(w->type)) return false; - return ds4_gpu_matmul_quant_kslice_tensor(out, - model->map, - model->size, - w->abs_offset, - w->type, - full_in_dim, - k_off, - k_cnt, - out_dim, - x, - x_elem_off) != 0; -} - -static bool metal_graph_attention_output_dense_quant_low( - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - uint64_t group_dim, - uint64_t rank, - uint32_t group0, - uint32_t group_cnt, - const ds4_gpu_tensor *heads) { - (void)g; - if (!low || !model || !out_a || !heads || - group_dim == 0 || rank == 0 || group_cnt == 0) { - return false; - } - if (out_a->type == DS4_TENSOR_Q8_0 && group0 == 0) { - return ds4_gpu_attention_output_low_q8_tensor(low, - model->map, - model->size, - out_a->abs_offset, - group_dim, - rank, - group_cnt, - heads) != 0; - } - if (out_a->type == DS4_TENSOR_Q4_K) { - return ds4_gpu_attention_output_low_q4_K_slice_tensor(low, - model->map, - model->size, - out_a->abs_offset, - group_dim, - rank, - group0, - group_cnt, - heads) != 0; - } - uint64_t row_bytes = 0; - if (!metal_graph_dense_quant_row_bytes(out_a, group_dim, &row_bytes)) return false; - const uint64_t group_weight_bytes = rank * row_bytes; - bool ok = true; - for (uint32_t i = 0; ok && i < group_cnt; i++) { - ds4_gpu_tensor *head_view = ds4_gpu_tensor_view( - heads, - (uint64_t)i * group_dim * sizeof(float), - group_dim * sizeof(float)); - ds4_gpu_tensor *low_view = ds4_gpu_tensor_view( - low, - (uint64_t)i * rank * sizeof(float), - rank * sizeof(float)); - ok = head_view && low_view && - metal_graph_matmul_dense_quant_abs(low_view, - model, - out_a, - out_a->abs_offset + - (uint64_t)(group0 + i) * group_weight_bytes, - group_dim, - rank, - head_view, - 1); - ds4_gpu_tensor_free(low_view); - ds4_gpu_tensor_free(head_view); - } - return ok; -} - -static bool metal_graph_attention_output_dense_quant_tp( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - const ds4_tensor *out_b, - 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) { - if (!out || !low || !g || !model || !out_a || !out_b || !heads || - group0 + group_cnt > n_groups_total) { - return false; - } - if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { - return ds4_gpu_attention_output_q8_tp_tensor(out, - low, - model->map, - model->size, - out_a->abs_offset, - out_b->abs_offset, - group_dim, - rank, - n_groups_total, - group0, - group_cnt, - out_dim, - heads) != 0; - } - if (!metal_graph_attention_output_dense_quant_low(low, - g, - model, - out_a, - group_dim, - rank, - group0, - group_cnt, - heads)) { - return false; - } - return metal_graph_matmul_dense_quant_kslice(out, - model, - out_b, - (uint64_t)n_groups_total * rank, - (uint64_t)group0 * rank, - (uint64_t)group_cnt * rank, - out_dim, - low, - 0); -} - -static bool metal_graph_attention_output_dense_quant_batch( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_tensor *out_a, - const ds4_tensor *out_b, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens) { - if (!out || !low || !g || !model || !out_a || !out_b || !heads || - n_groups == 0 || n_tokens == 0) { - return false; - } - if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { - return ds4_gpu_attention_output_q8_batch_tensor(out, - low, - metal_graph_batch_group_tmp(g), - metal_graph_batch_low_tmp(g), - model->map, - model->size, - out_a->abs_offset, - out_b->abs_offset, - group_dim, - rank, - n_groups, - out_dim, - heads, - n_tokens) != 0; - } - if (out_a->type == DS4_TENSOR_Q4_K && n_tokens >= 32u) { - if (ds4_gpu_attention_output_q4_K_batch_tensor(out, - low, - metal_graph_batch_group_tmp(g), - metal_graph_batch_low_tmp(g), - model->map, - model->size, - out_a->abs_offset, - out_b->abs_offset, - out_b->type, - group_dim, - rank, - n_groups, - out_dim, - heads, - n_tokens) != 0) { - return true; - } - } - - const uint64_t heads_row_elems = (uint64_t)n_groups * group_dim; - const uint64_t low_row_elems = (uint64_t)n_groups * rank; - bool ok = true; - for (uint32_t t = 0; ok && t < n_tokens; t++) { - ds4_gpu_tensor *heads_row = ds4_gpu_tensor_view( - heads, - (uint64_t)t * heads_row_elems * sizeof(float), - heads_row_elems * sizeof(float)); - ds4_gpu_tensor *low_row = ds4_gpu_tensor_view( - low, - (uint64_t)t * low_row_elems * sizeof(float), - low_row_elems * sizeof(float)); - ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( - out, - (uint64_t)t * out_dim * sizeof(float), - out_dim * sizeof(float)); - ok = heads_row && low_row && out_row && - metal_graph_attention_output_dense_quant_low(low_row, - g, - model, - out_a, - group_dim, - rank, - 0, - n_groups, - heads_row); - if (ok) ok = metal_graph_matmul_dense_quant_tensor(out_row, - model, - out_b, - low_row_elems, - out_dim, - low_row, - 1); - ds4_gpu_tensor_free(out_row); - ds4_gpu_tensor_free(low_row); - ds4_gpu_tensor_free(heads_row); - } - return ok; -} - -static bool metal_graph_matmul_q8_0_named_tensor( - const char *module, - uint32_t il, - uint32_t pos0, - ds4_gpu_tensor *out, - const ds4_model *model, - const ds4_tensor *w, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - (void)module; - (void)il; - (void)pos0; - return metal_graph_matmul_dense_quant_tensor(out, - model, - w, - in_dim, - out_dim, - x, - n_tok); -} - -static bool metal_graph_encode_output_head_mtp( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_model *mtp_model, - const ds4_mtp_weights *mtp, - uint64_t vocab_dim) { - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_output_pre(g), mtp_model, mtp->hc_head_fn, - hc_dim, DS4_N_HC, metal_graph_flat_hc(g), 1); - if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), - metal_graph_output_pre(g), - mtp_model->map, - mtp_model->size, - mtp->hc_head_scale->abs_offset, - mtp->hc_head_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), - metal_graph_cur_hc(g), - metal_graph_output_weights(g), - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), - metal_graph_output_embd(g), - mtp_model->map, - mtp_model->size, - mtp->norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), - base_model, - base_weights->output, - DS4_N_EMBD, - vocab_dim, - metal_graph_output_norm(g), - 1); - return ok; -} - -/* ========================================================================= - * Metal Diagnostic Comparisons. - * ========================================================================= - * - * These routines deliberately allocate CPU-side reference buffers and read - * Metal tensors back. They are not part of generation; command-line tests use - * them to localize drift against the C reference pipeline. - */ - -static void metal_graph_trace_layer_stages( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - const float *cpu_in_hc, - uint32_t il, - int token) { - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t shared_in_dim = layer->ffn_gate_shexp->dim[0]; - const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; - const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; - - float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); - float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); - float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); - float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); - float *cpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); - float *cpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); - float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float post[4]; - float comb[16]; - float ffn_post[4]; - float ffn_comb[16]; - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - const uint64_t shared_blocks = (shared_in_dim + 31) / 32; - int8_t *shared_xq = xmalloc((size_t)shared_blocks * 32); - float *shared_xscale = xmalloc((size_t)shared_blocks * sizeof(float)); - float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); - block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); - block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); - int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; - float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; - int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; - float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; - - hc_pre_from_state_one(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - cpu_in_hc, cpu_attn_cur, post, comb); - layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); - layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); - layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); - rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); - rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); - dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); - f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); - layer_attention_one(cpu_heads, model, layer, cpu_q, cpu_kv); - rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, true); - layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); - hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_in_hc, post, comb, DS4_N_EMBD, DS4_N_HC); - hc_pre_from_state_one(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - cpu_after_attn_hc, cpu_ffn_cur, ffn_post, ffn_comb); - rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); - quantize_q8_0_activation(cpu_ffn_norm, shared_xq, shared_xscale, shared_in_dim); - matvec_q8_0_pair_prequant(cpu_shared_gate, - cpu_shared_up, - model, - layer->ffn_gate_shexp, - layer->ffn_up_shexp, - shared_xq, - shared_xscale); - swiglu(cpu_shared_mid, cpu_shared_gate, cpu_shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP); - matvec_q8_0(cpu_shared, model, layer->ffn_down_shexp, cpu_shared_mid); - layer_routed_moe_one_prealloc(cpu_routed, - model, - layer, - cpu_ffn_norm, - il, - token, - DS4_SWIGLU_CLAMP_EXP, - routed_mid_all, - routed_xq, - routed_midq, - routed_q8_xq, - routed_q8_xscale, - routed_q8_midq, - routed_q8_midscale); - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(selected, model, layer, token); - layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); - } else { - layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); - } - for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; - hc_post_one(cpu_after_ffn_hc, cpu_ffn_out, cpu_after_attn_hc, ffn_post, ffn_comb, DS4_N_EMBD, DS4_N_HC); - - float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); - float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); - float *gpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); - float *gpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); - float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); - float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - int gpu_selected[DS4_MAX_EXPERT_USED]; - float gpu_expert_weight[DS4_MAX_EXPERT_USED]; - - bool ok = ds4_gpu_tensor_read(metal_graph_attn_cur(g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_attn_norm(g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_q(g), 0, gpu_q, q_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_kv(g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_attn_out(g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_after_attn_hc(g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_cur(g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_norm(g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_shared_gate(g), 0, gpu_shared_gate, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_shared_up(g), 0, gpu_shared_up, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_shared_mid(g), 0, gpu_shared_mid, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_shared_out(g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_router_selected(g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && - ds4_gpu_tensor_read(metal_graph_router_weights(g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && - ds4_gpu_tensor_read(metal_graph_routed_mid(g), 0, gpu_routed_mid_all, (uint64_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_routed_out(g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_out(g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_cur_hc(g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0; - - if (ok) { - fprintf(stderr, - "ds4: Metal stage layer %u attn_cur=%g/%g attn_norm=%g/%g q=%g/%g kv=%g/%g attn_out=%g/%g after_attn_hc=%g/%g ffn_cur=%g/%g ffn_norm=%g/%g shared=%g/%g router_w=%g routed=%g/%g ffn_out=%g/%g after_ffn_hc=%g/%g\n", - il, - max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), rms_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), - max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), rms_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), - max_abs_diff(cpu_q, gpu_q, q_dim), rms_abs_diff(cpu_q, gpu_q, q_dim), - max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), rms_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), - max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), rms_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), - max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), rms_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), - max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), rms_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), - max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), rms_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), - max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), - max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), - max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), - max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), rms_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), - max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), rms_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim)); - fprintf(stderr, - "ds4: Metal shared layer %u gate=%g/%g up=%g/%g mid=%g/%g down=%g/%g\n", - il, - max_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), rms_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), - max_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), rms_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), - max_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), rms_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), - max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD)); - fprintf(stderr, - "ds4: Metal routed layer %u mid=%g/%g out=%g/%g\n", - il, - max_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), - rms_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), - max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), - rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD)); - if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { - fprintf(stderr, - "ds4: Metal stage layer %u router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", - il, - selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], - gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); - } - } - - free(gpu_after_ffn_hc); - free(gpu_ffn_out); - free(gpu_routed); - free(gpu_routed_mid_all); - free(gpu_shared); - free(gpu_shared_mid); - free(gpu_shared_up); - free(gpu_shared_gate); - free(gpu_ffn_norm); - free(gpu_ffn_cur); - free(gpu_after_attn_hc); - free(gpu_attn_out); - free(gpu_kv); - free(gpu_q); - free(gpu_attn_norm); - free(gpu_attn_cur); - free(routed_q8_midscale); - free(routed_q8_midq); - free(routed_q8_xscale); - free(routed_q8_xq); - free(routed_midq); - free(routed_xq); - free(routed_mid_all); - free(shared_xscale); - free(shared_xq); - free(cpu_after_ffn_hc); - free(cpu_ffn_out); - free(cpu_routed); - free(cpu_shared); - free(cpu_shared_mid); - free(cpu_shared_up); - free(cpu_shared_gate); - free(cpu_ffn_norm); - free(cpu_ffn_cur); - free(cpu_after_attn_hc); - free(cpu_attn_out); - free(cpu_heads); - free(cpu_kv); - free(cpu_qr_norm); - free(cpu_q); - free(cpu_attn_norm); - free(cpu_attn_cur); -} - -static int metal_graph_decode_test( - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - bool quality) { - if (prompt->len <= 0) { - fprintf(stderr, "ds4: Metal graph test needs a non-empty prompt\n"); - return 1; - } - - const int token = prompt->v[0]; - const ds4_layer_weights *layer = &weights->layer[0]; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const uint64_t vocab_dim = weights->output->dim[1]; - const bool routed_q8_0 = - layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && - layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; - const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; - const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; - - float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); - float *cpu_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); - float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); - float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); - float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); - float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_ffn_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); - float *cpu_ffn_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); - float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); - float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); - float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - float *gpu_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); - float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); - int gpu_selected[DS4_MAX_EXPERT_USED]; - float gpu_expert_weight[DS4_MAX_EXPERT_USED]; - float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); - block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); - block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); - int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; - float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; - int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; - float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; - int selected[DS4_MAX_EXPERT_USED]; - float expert_weight[DS4_MAX_EXPERT_USED]; - - embed_token_f16(model, weights, token, plain); - hc_from_plain_embedding(cpu_hc, plain, DS4_N_EMBD, DS4_N_HC); - hc_pre_from_state_one(model, - layer->hc_attn_fn, - layer->hc_attn_scale, - layer->hc_attn_base, - cpu_hc, cpu_attn_cur, cpu_post, cpu_comb); - layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); - layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); - layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); - rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); - rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); - dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); - f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); - layer_attention_rows_one(cpu_heads, model, layer, cpu_q, cpu_kv, 1); - rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, true); - layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); - hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_hc, cpu_post, cpu_comb, DS4_N_EMBD, DS4_N_HC); - hc_pre_from_state_one(model, - layer->hc_ffn_fn, - layer->hc_ffn_scale, - layer->hc_ffn_base, - cpu_after_attn_hc, cpu_ffn_cur, cpu_ffn_post, cpu_ffn_comb); - rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); - layer_shared_ffn_one(cpu_shared, model, layer, cpu_ffn_norm); - layer_routed_moe_one_prealloc(cpu_routed, - model, - layer, - cpu_ffn_norm, - 0, - token, - DS4_SWIGLU_CLAMP_EXP, - routed_mid_all, - routed_xq, - routed_midq, - routed_q8_xq, - routed_q8_xscale, - routed_q8_midq, - routed_q8_midscale); - if (layer->ffn_gate_tid2eid) { - layer_hash_selected_experts(selected, model, layer, token); - layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); - } else { - layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); - } - for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; - hc_post_one(cpu_after_ffn_hc, - cpu_ffn_out, - cpu_after_attn_hc, - cpu_ffn_post, - cpu_ffn_comb, - DS4_N_EMBD, - DS4_N_HC); - output_logits_one(cpu_logits, model, weights, cpu_after_ffn_hc); - - ds4_gpu_graph g; - bool ok = metal_graph_alloc(&g, weights, layer); - g.quality = quality; - g.materialize_ffn_out = true; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = metal_graph_encode_decode_layer(&g, - model, - layer, - 0, - 0, - g.layer_raw_cache[0], - g.raw_cap, - 0, - 1, - token); - if (ok) { - /* Single-tier diagnostic: swap the active-tier slots so the head - * pipeline reads the embedded hidden state from cur_hc. */ - ds4_gpu_tensor *embedded_hc = g.cur_hc_by_tier[g.active_tier]; - g.cur_hc_by_tier[g.active_tier] = g.after_ffn_hc_by_tier[g.active_tier]; - g.after_ffn_hc_by_tier[g.active_tier] = embedded_hc; - } - if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); - if (ok) ok = ds4_gpu_end_commands() != 0; - - if (ok) { - ok = ds4_gpu_tensor_read(metal_graph_after_ffn_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_attn_cur(&g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_attn_norm(&g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_q(&g), 0, gpu_q, q_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_kv(&g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.layer_raw_cache[0], 0, gpu_raw, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_attn_out(&g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_after_attn_hc(&g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_cur(&g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_norm(&g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_shared_out(&g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_router_selected(&g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && - ds4_gpu_tensor_read(metal_graph_router_weights(&g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && - ds4_gpu_tensor_read(metal_graph_routed_out(&g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_ffn_out(&g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; - } - - if (ok) { - fprintf(stderr, - "ds4: Metal graph test layer0 diffs: embed_hc=%g hc_pre=%g attn_norm=%g q_rope=%g kv_rope=%g raw_cache=%g attn_out=%g after_attn_hc=%g ffn_cur=%g ffn_norm=%g shared=%g router_w=%g routed=%g ffn_out=%g after_ffn_hc=%g logits=%g\n", - max_abs_diff(cpu_hc, gpu_hc, hc_dim), - max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), - max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), - max_abs_diff(cpu_q, gpu_q, q_dim), - max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), - max_abs_diff(cpu_kv, gpu_raw, DS4_N_HEAD_DIM), - max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), - max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), - max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), - max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), - max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), - max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), - max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), - max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), - max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), - max_abs_diff(cpu_logits, gpu_logits, vocab_dim)); - if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { - fprintf(stderr, - "ds4: Metal graph router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", - selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], - gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); - } - print_vec_stats("metal graph q", gpu_q, q_dim); - print_vec_stats("metal graph kv", gpu_kv, DS4_N_HEAD_DIM); - print_vec_stats("metal graph routed", gpu_routed, DS4_N_EMBD); - } else { - fprintf(stderr, "ds4: Metal graph test failed while encoding first decode stages\n"); - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after graph test failure also failed\n"); - } - } - - metal_graph_free(&g); - free(routed_q8_midscale); - free(routed_q8_midq); - free(routed_q8_xscale); - free(routed_q8_xq); - free(routed_midq); - free(routed_xq); - free(routed_mid_all); - free(gpu_logits); - free(gpu_after_ffn_hc); - free(gpu_ffn_out); - free(gpu_routed); - free(gpu_shared); - free(gpu_ffn_norm); - free(gpu_ffn_cur); - free(gpu_after_attn_hc); - free(gpu_attn_out); - free(gpu_raw); - free(gpu_kv); - free(gpu_q); - free(gpu_attn_norm); - free(gpu_attn_cur); - free(gpu_hc); - free(cpu_kv); - free(cpu_q); - free(cpu_attn_out); - free(cpu_heads); - free(cpu_ffn_norm); - free(cpu_routed); - free(cpu_logits); - free(cpu_after_ffn_hc); - free(cpu_ffn_out); - free(cpu_shared); - free(cpu_ffn_comb); - free(cpu_ffn_post); - free(cpu_ffn_cur); - free(cpu_after_attn_hc); - free(cpu_qr_norm); - free(cpu_attn_norm); - free(cpu_comb); - free(cpu_post); - free(cpu_attn_cur); - free(cpu_hc); - free(plain); - return ok ? 0 : 1; -} - -static int metal_graph_first_token_full_test( - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - bool quality) { - if (prompt->len <= 0) { - fprintf(stderr, "ds4: full Metal graph test needs a non-empty prompt\n"); - return 1; - } - - const int token = prompt->v[0]; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t vocab_dim = weights->output->dim[1]; - float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); - float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); - - forward_first_token_cpu(cpu_hc, model, weights, token); - output_logits_one(cpu_logits, model, weights, cpu_hc); - - ds4_gpu_graph g; - bool ok = metal_graph_alloc(&g, weights, &weights->layer[0]); - g.quality = quality; - const bool trace_layers = getenv("DS4_METAL_GRAPH_TRACE_LAYERS") != NULL; - if (trace_layers && ok) { - g.materialize_ffn_out = true; - const bool teacher_force = getenv("DS4_METAL_GRAPH_TEACHER_FORCE") != NULL; - const char *stage_layer_env = getenv("DS4_METAL_GRAPH_TRACE_STAGE_LAYER"); - const long stage_layer = stage_layer_env ? strtol(stage_layer_env, NULL, 10) : -1; - float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); - float *cpu_cur = xmalloc((size_t)hc_dim * sizeof(float)); - float *cpu_next = xmalloc((size_t)hc_dim * sizeof(float)); - - embed_token_f16(model, weights, token, plain); - hc_from_plain_embedding(cpu_cur, plain, DS4_N_EMBD, DS4_N_HC); - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_end_commands() != 0; - - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - if (teacher_force) { - ok = ds4_gpu_tensor_write(metal_graph_cur_hc(&g), 0, cpu_cur, hc_dim * sizeof(float)) != 0; - } - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], - il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); - ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); - g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); - g.after_ffn_hc_by_tier[g.active_tier] = tmp; - if (ok) ok = ds4_gpu_end_commands() != 0; - - layer_forward_self_one(cpu_next, model, &weights->layer[il], cpu_cur, il, 0, token); - if (ok) ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0; - if (ok) { - fprintf(stderr, - "ds4: Metal full graph layer %u%s hc_max=%g hc_rms=%g\n", - il, - teacher_force ? " teacher" : "", - max_abs_diff(cpu_next, gpu_hc, hc_dim), - rms_abs_diff(cpu_next, gpu_hc, hc_dim)); - if (stage_layer == (long)il) { - metal_graph_trace_layer_stages(&g, model, &weights->layer[il], cpu_cur, il, token); - } - } - float *ctmp = cpu_cur; - cpu_cur = cpu_next; - cpu_next = ctmp; - } - - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); - if (ok) ok = ds4_gpu_end_commands() != 0; - - free(cpu_next); - free(cpu_cur); - free(plain); - } else { - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; - - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], - il, 0, g.layer_raw_cache[il], - g.raw_cap, 0, 1, token); - ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); - g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); - g.after_ffn_hc_by_tier[g.active_tier] = tmp; - } - - if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); - if (ok) ok = ds4_gpu_end_commands() != 0; - } - - if (ok) { - ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; - } - - if (ok) { - const uint64_t cpu_top = argmax_f32(cpu_logits, vocab_dim); - const uint64_t gpu_top = argmax_f32(gpu_logits, vocab_dim); - fprintf(stderr, - "ds4: Metal full first-token graph diffs: final_hc_max=%g final_hc_rms=%g logits_max=%g logits_rms=%g cpu_top=%llu gpu_top=%llu cpu_top_logit=%g gpu_top_logit=%g\n", - max_abs_diff(cpu_hc, gpu_hc, hc_dim), - rms_abs_diff(cpu_hc, gpu_hc, hc_dim), - max_abs_diff(cpu_logits, gpu_logits, vocab_dim), - rms_abs_diff(cpu_logits, gpu_logits, vocab_dim), - (unsigned long long)cpu_top, - (unsigned long long)gpu_top, - cpu_logits[cpu_top], - gpu_logits[gpu_top]); - } else { - fprintf(stderr, "ds4: Metal full first-token graph test failed\n"); - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after full graph failure also failed\n"); - } - } - - metal_graph_free(&g); - free(gpu_logits); - free(cpu_logits); - free(gpu_hc); - free(cpu_hc); - return ok ? 0 : 1; -} - -/* ========================================================================= - * Metal Release Decode and Prefill. - * ========================================================================= - * - * Everything below is the user-facing Metal backend. It uses the same layer - * encoder as diagnostics, but diagnostics are not required for normal command - * flow and their CPU reads stay outside these generation entry points. - */ - -static uint32_t metal_graph_token_split_after_layers(void) { - uint32_t split_after_layers = 4; -#ifndef DS4_ROCM_BUILD - const char *split_env = getenv("DS4_METAL_GRAPH_TOKEN_SPLIT_LAYERS"); - if (split_env && split_env[0]) { - char *end = NULL; - unsigned long v = strtoul(split_env, &end, 10); - if (end != split_env && v <= DS4_N_LAYER) split_after_layers = (uint32_t)v; - } -#endif - return split_after_layers; -} - -static int metal_graph_dspark_target_slot( - const ds4_gpu_graph *g, - uint32_t il) { - if (!g || !g->dspark_capture_enabled) return -1; - for (uint32_t i = 0; i < g->dspark_target_layer_count; i++) { - if (g->dspark_target_layers[i] == il) return (int)i; - } - return -1; -} - -static uint32_t metal_graph_dspark_capture_complete_mask( - const ds4_gpu_graph *g) { - if (!g || g->dspark_target_layer_count == 0) return 0; - return g->dspark_target_layer_count >= 32u ? - UINT32_MAX : ((1u << g->dspark_target_layer_count) - 1u); -} - -static void metal_graph_dspark_capture_note_slot(ds4_gpu_graph *g, - uint32_t slot) { - if (!g || slot >= g->dspark_target_layer_count) return; - g->dspark_capture_mask |= 1u << slot; - g->dspark_capture_valid = - g->dspark_capture_mask == metal_graph_dspark_capture_complete_mask(g); -} - -static void metal_graph_dspark_capture_row_invalidate(ds4_gpu_graph *g) { - if (!g || !g->dspark_capture_enabled) return; - g->dspark_capture_mask = 0; - g->dspark_capture_checkpoint_len = 0; - g->dspark_capture_valid = false; -} - -static void metal_graph_dspark_capture_batch_invalidate(ds4_gpu_graph *g) { - if (!g || !g->dspark_capture_enabled) return; - g->dspark_capture_batch_mask = 0; - g->dspark_capture_batch_start = 0; - g->dspark_capture_batch_tokens = 0; - g->dspark_capture_batch_valid = false; -} - -static void metal_graph_dspark_capture_invalidate(ds4_gpu_graph *g) { - metal_graph_dspark_capture_row_invalidate(g); - metal_graph_dspark_capture_batch_invalidate(g); -} - -static void metal_graph_dspark_cache_reset(ds4_gpu_graph *g) { - if (!g) return; - g->dspark_cache_start = 0; - g->dspark_cache_token_start = 0; - g->dspark_cache_len = 0; -} - -static bool metal_graph_dspark_cache_window_valid( - const ds4_gpu_graph *g, - uint32_t token_start, - uint32_t raw_start, - uint32_t len) { - if (!g || len > g->dspark_cache_cap) return false; - if (len == 0) return true; - if (g->dspark_cache_cap == 0 || - raw_start >= g->dspark_cache_cap || - token_start > UINT32_MAX - len || - raw_start != token_start % g->dspark_cache_cap) { - return false; - } - return true; -} - -static bool metal_graph_dspark_cache_current_window_valid( - const ds4_gpu_graph *g) { - if (!g) return false; - return metal_graph_dspark_cache_window_valid(g, - g->dspark_cache_token_start, - g->dspark_cache_start, - g->dspark_cache_len); -} - -static bool metal_graph_dspark_cache_set_window(ds4_gpu_graph *g, - uint32_t token_start, - uint32_t len) { - if (!g || len > g->dspark_cache_cap) return false; - const uint32_t raw_start = - len && g->dspark_cache_cap ? token_start % g->dspark_cache_cap : 0; - if (!metal_graph_dspark_cache_window_valid(g, - len ? token_start : 0, - raw_start, - len)) { - return false; - } - g->dspark_cache_start = raw_start; - g->dspark_cache_token_start = len ? token_start : 0; - g->dspark_cache_len = len; - return true; -} - -static bool metal_graph_dspark_cache_crop_to_prefix(ds4_gpu_graph *g, - uint32_t prefix_len) { - if (!g) return false; - if (g->dspark_cache_len == 0) return true; - if (!metal_graph_dspark_cache_current_window_valid(g)) return false; - - const uint32_t start = g->dspark_cache_token_start; - const uint32_t end = start + g->dspark_cache_len; - if (prefix_len <= start || prefix_len > end) { - metal_graph_dspark_cache_reset(g); - return true; - } - g->dspark_cache_len = prefix_len - start; - return true; -} - -static bool metal_graph_dspark_cache_ends_at(const ds4_gpu_graph *g, - uint32_t pos) { - if (!metal_graph_dspark_cache_current_window_valid(g)) return false; - if (g->dspark_cache_len == 0) return true; - return g->dspark_cache_token_start <= UINT32_MAX - g->dspark_cache_len && - g->dspark_cache_token_start + g->dspark_cache_len == pos; -} - -static bool metal_graph_dspark_cache_claim_appended_row(ds4_gpu_graph *g, - uint32_t pos) { - if (!g || g->dspark_cache_len == 0 || - !metal_graph_dspark_cache_ends_at(g, pos)) return false; - g->dspark_cache_len += 1u; - if (g->dspark_cache_len > g->dspark_cache_cap) { - const uint32_t excess = g->dspark_cache_len - g->dspark_cache_cap; - g->dspark_cache_token_start += excess; - g->dspark_cache_len = g->dspark_cache_cap; - g->dspark_cache_start = - g->dspark_cache_token_start % g->dspark_cache_cap; - } - return true; -} - -bool ds4_test_dspark_cache_window_crop(void) { - ds4_gpu_graph g; - memset(&g, 0, sizeof(g)); - g.dspark_cache_cap = 8; - - if (!metal_graph_dspark_cache_set_window(&g, 10, 5)) return false; - if (g.dspark_cache_token_start != 10 || - g.dspark_cache_start != 2 || - g.dspark_cache_len != 5) return false; - if (!metal_graph_dspark_cache_ends_at(&g, 15)) return false; - if (metal_graph_dspark_cache_ends_at(&g, 14)) return false; - if (metal_graph_dspark_cache_window_valid(&g, 10, 3, 5)) return false; - - if (!metal_graph_dspark_cache_crop_to_prefix(&g, 13)) return false; - if (g.dspark_cache_token_start != 10 || - g.dspark_cache_start != 2 || - g.dspark_cache_len != 3) return false; - if (!metal_graph_dspark_cache_ends_at(&g, 13)) return false; - if (!metal_graph_dspark_cache_claim_appended_row(&g, 13)) return false; - if (g.dspark_cache_token_start != 10 || - g.dspark_cache_start != 2 || - g.dspark_cache_len != 4) return false; - if (!metal_graph_dspark_cache_ends_at(&g, 14)) return false; - if (metal_graph_dspark_cache_claim_appended_row(&g, 13)) return false; - - if (!metal_graph_dspark_cache_crop_to_prefix(&g, 20)) return false; - if (g.dspark_cache_token_start != 0 || - g.dspark_cache_start != 0 || - g.dspark_cache_len != 0) return false; - if (!metal_graph_dspark_cache_ends_at(&g, 20)) return false; - - if (metal_graph_dspark_cache_set_window(&g, UINT32_MAX - 1u, 2)) { - return false; - } - return true; -} - -static void metal_graph_dspark_capture_begin(ds4_gpu_graph *g) { - metal_graph_dspark_capture_row_invalidate(g); -} - -static void metal_graph_dspark_capture_begin_prefill(ds4_gpu_graph *g) { - metal_graph_dspark_capture_invalidate(g); -} - -static bool metal_graph_dspark_capture_hc( - ds4_gpu_graph *g, - const ds4_gpu_tensor *hc, - uint32_t slot) { - if (!g || !hc || !g->dspark_target_hidden || - !g->dspark_hc_mean_weights || - slot >= g->dspark_target_layer_count) { - return false; - } - - ds4_gpu_tensor *dst = - ds4_gpu_tensor_view(g->dspark_target_hidden, - (uint64_t)slot * DS4_N_EMBD * sizeof(float), - (uint64_t)DS4_N_EMBD * sizeof(float)); - if (!dst) return false; - const bool ok = ds4_gpu_hc_weighted_sum_tensor(dst, - hc, - g->dspark_hc_mean_weights, - DS4_N_EMBD, - DS4_N_HC) != 0; - ds4_gpu_tensor_free(dst); - if (ok) metal_graph_dspark_capture_note_slot(g, slot); - return ok; -} - -static bool metal_graph_dspark_capture_batch_note_slot( - ds4_gpu_graph *g, - uint32_t slot, - uint32_t start, - uint32_t n_tokens) { - if (!g || slot >= g->dspark_target_layer_count || n_tokens == 0) { - return false; - } - if (g->dspark_capture_batch_mask == 0) { - g->dspark_capture_batch_start = start; - g->dspark_capture_batch_tokens = n_tokens; - } else if (g->dspark_capture_batch_start != start || - g->dspark_capture_batch_tokens != n_tokens) { - metal_graph_dspark_capture_batch_invalidate(g); - return false; - } - g->dspark_capture_batch_mask |= 1u << slot; - g->dspark_capture_batch_valid = - g->dspark_capture_batch_mask == - metal_graph_dspark_capture_complete_mask(g); - return true; -} - -static bool metal_graph_dspark_capture_decode_layer( - ds4_gpu_graph *g, - uint32_t il) { - const int slot = metal_graph_dspark_target_slot(g, il); - if (slot < 0) return true; - return metal_graph_dspark_capture_hc(g, metal_graph_cur_hc(g), (uint32_t)slot); -} - -static bool metal_graph_dspark_capture_prefill_layer( - ds4_gpu_graph *g, - uint32_t il, - uint32_t start, - uint32_t n_tokens) { - const int slot = metal_graph_dspark_target_slot(g, il); - if (slot < 0) return true; - if (n_tokens == 0) return false; - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - if (g->dspark_target_hidden_batch && - g->dspark_hc_mean_rows && - n_tokens <= g->prefill_cap) { - ds4_gpu_tensor *batch_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden_batch, - ((uint64_t)slot * g->prefill_cap * - DS4_N_EMBD) * sizeof(float), - (uint64_t)n_tokens * embd_bytes); - ds4_gpu_tensor *last_src = - batch_dst ? - ds4_gpu_tensor_view(batch_dst, - (uint64_t)(n_tokens - 1u) * embd_bytes, - embd_bytes) : NULL; - ds4_gpu_tensor *last_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden, - (uint64_t)slot * embd_bytes, - embd_bytes); - bool ok = batch_dst && last_src && last_dst && - ds4_gpu_hc_weighted_sum_tensor(batch_dst, - metal_graph_batch_cur_hc(g), - g->dspark_hc_mean_rows, - DS4_N_EMBD, - DS4_N_HC) != 0 && - ds4_gpu_tensor_copy(last_dst, - 0, - last_src, - 0, - embd_bytes) != 0; - ds4_gpu_tensor_free(last_dst); - ds4_gpu_tensor_free(last_src); - ds4_gpu_tensor_free(batch_dst); - if (ok) { - metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); - ok = metal_graph_dspark_capture_batch_note_slot(g, - (uint32_t)slot, - start, - n_tokens); - } - return ok; - } - - ds4_gpu_tensor *last_hc = - ds4_gpu_tensor_view(metal_graph_batch_cur_hc(g), - (uint64_t)(n_tokens - 1u) * hc_dim * sizeof(float), - hc_dim * sizeof(float)); - if (!last_hc) return false; - const bool ok = metal_graph_dspark_capture_hc(g, last_hc, (uint32_t)slot); - ds4_gpu_tensor_free(last_hc); - return ok; -} - -static bool metal_graph_dspark_capture_prefill_rows( - ds4_gpu_graph *g, - uint32_t il, - uint32_t chunk_start, - uint32_t chunk_len, - uint32_t pos0, - uint32_t n_tokens) { - const int slot = metal_graph_dspark_target_slot(g, il); - if (slot < 0) return true; - if (!g->dspark_target_hidden_batch || - !g->dspark_target_hidden || - !g->dspark_hc_mean_rows || - n_tokens == 0 || - chunk_len == 0 || - chunk_len > g->prefill_cap || - pos0 < chunk_start) { - return true; - } - const uint32_t row0 = pos0 - chunk_start; - if (row0 > chunk_len || n_tokens > chunk_len - row0) return true; - - const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - ds4_gpu_tensor *batch_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden_batch, - (((uint64_t)(uint32_t)slot * g->prefill_cap + - row0) * DS4_N_EMBD) * sizeof(float), - (uint64_t)n_tokens * embd_bytes); - bool ok = batch_dst && - ds4_gpu_hc_weighted_sum_tensor(batch_dst, - metal_graph_batch_cur_hc(g), - g->dspark_hc_mean_rows, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (!ok) fprintf(stderr, "ds4: pipeline capture rows FAIL il=%u row0=%u n=%u dst=%d\n", - il, row0, n_tokens, batch_dst != NULL); - if (ok && row0 + n_tokens == chunk_len) { - ds4_gpu_tensor *last_src = - ds4_gpu_tensor_view(batch_dst, - (uint64_t)(n_tokens - 1u) * embd_bytes, - embd_bytes); - ds4_gpu_tensor *last_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden, - (uint64_t)(uint32_t)slot * embd_bytes, - embd_bytes); - ok = last_src && last_dst && - ds4_gpu_tensor_copy(last_dst, 0, last_src, 0, embd_bytes) != 0; - ds4_gpu_tensor_free(last_dst); - ds4_gpu_tensor_free(last_src); - if (ok) { - metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); - ok = metal_graph_dspark_capture_batch_note_slot(g, - (uint32_t)slot, - chunk_start, - chunk_len); - } - } - ds4_gpu_tensor_free(batch_dst); - return ok; -} - -static bool metal_graph_dspark_capture_verified_suffix_begin( - ds4_gpu_graph *g, - uint32_t start, - uint32_t n_tokens, - bool commands_open) { - if (!g || !g->dspark_capture_enabled || - !g->dspark_target_hidden || - !g->dspark_target_hidden_batch || - start == 0 || - n_tokens == 0 || - n_tokens + 1u < n_tokens || - n_tokens + 1u > g->prefill_cap || - !g->dspark_capture_valid || - g->dspark_capture_checkpoint_len != start) { - metal_graph_dspark_capture_invalidate(g); - return false; - } - - const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - metal_graph_dspark_capture_batch_invalidate(g); - bool ok = commands_open || ds4_gpu_begin_commands() != 0; - for (uint32_t slot = 0; ok && slot < g->dspark_target_layer_count; slot++) { - ds4_gpu_tensor *dst = - ds4_gpu_tensor_view(g->dspark_target_hidden_batch, - ((uint64_t)slot * g->prefill_cap * - DS4_N_EMBD) * sizeof(float), - embd_bytes); - ds4_gpu_tensor *src = - ds4_gpu_tensor_view(g->dspark_target_hidden, - (uint64_t)slot * embd_bytes, - embd_bytes); - ok = dst && src && - ds4_gpu_tensor_copy(dst, 0, src, 0, embd_bytes) != 0; - ds4_gpu_tensor_free(src); - ds4_gpu_tensor_free(dst); - } - if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; - else if (!ok && !commands_open) (void)ds4_gpu_synchronize(); - if (!ok) { - metal_graph_dspark_capture_invalidate(g); - return false; - } - - metal_graph_dspark_capture_row_invalidate(g); - return true; -} - -static bool metal_graph_dspark_capture_verified_suffix_layer( - ds4_gpu_graph *g, - uint32_t il, - uint32_t start, - uint32_t n_tokens) { - const int slot = metal_graph_dspark_target_slot(g, il); - if (slot < 0) return true; - if (!g || !g->dspark_target_hidden_batch || - !g->dspark_target_hidden || - !g->dspark_hc_mean_rows || - start == 0 || - n_tokens == 0 || - n_tokens + 1u < n_tokens || - n_tokens + 1u > g->prefill_cap) { - return false; - } - - const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - ds4_gpu_tensor *batch_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden_batch, - (((uint64_t)(uint32_t)slot * g->prefill_cap + - 1u) * DS4_N_EMBD) * sizeof(float), - (uint64_t)n_tokens * embd_bytes); - ds4_gpu_tensor *last_src = - batch_dst ? - ds4_gpu_tensor_view(batch_dst, - (uint64_t)(n_tokens - 1u) * embd_bytes, - embd_bytes) : NULL; - ds4_gpu_tensor *last_dst = - ds4_gpu_tensor_view(g->dspark_target_hidden, - (uint64_t)(uint32_t)slot * embd_bytes, - embd_bytes); - bool ok = batch_dst && last_src && last_dst && - ds4_gpu_hc_weighted_sum_tensor(batch_dst, - metal_graph_batch_cur_hc(g), - g->dspark_hc_mean_rows, - DS4_N_EMBD, - DS4_N_HC) != 0 && - ds4_gpu_tensor_copy(last_dst, - 0, - last_src, - 0, - embd_bytes) != 0; - ds4_gpu_tensor_free(last_dst); - ds4_gpu_tensor_free(last_src); - ds4_gpu_tensor_free(batch_dst); - if (ok) { - metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); - ok = metal_graph_dspark_capture_batch_note_slot(g, - (uint32_t)slot, - start - 1u, - n_tokens + 1u); - } - return ok; -} - -/* Encode a full single-token decode step on Metal. This is the generation - * hot path: update caches, run all layers, then produce logits. */ -static bool metal_graph_encode_token_raw_swa( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - uint32_t pos, - bool need_logits, - bool allow_split_flush) { - if (g->raw_cap == 0) { - fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); - return false; - } - /* Under the vocab split both ranks materialize their logits half. */ - if (g->tp_world == 2 && g->tp_rank == 1 && - !g->tp_logits_half) need_logits = false; - const uint32_t raw_row = pos % g->raw_cap; - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); - metal_graph_dspark_capture_begin(g); - - /* write the embedded token on the embedding tier. Single- - * tier: emb_tier == 0 == active_tier; no-op. Multi-tier: switch to - * emb_tier (no cross-device copy needed — embed writes from scratch). */ - if (g->placement) { - if (!metal_graph_set_active_tier_decode(g, g->emb_tier)) return false; - } - bool ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; - - /* - * Start executing the prefix of the decode graph while the CPU is still - * encoding the rest. The split point is layer-based because this executor is - * a fixed DS4 tape, not a dynamic node graph; four layers is the measured - * point where the prefix is large enough to hide useful work without - * starving the second command buffer. - */ - const uint32_t split_after_layers = metal_graph_token_split_after_layers(); - - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - ok = metal_graph_encode_decode_layer(g, - model, - &weights->layer[il], - il, - pos, - g->layer_raw_cache[il], - g->raw_cap, - raw_row, - n_raw, - token); - ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); - g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); - g->after_ffn_hc_by_tier[g->active_tier] = tmp; - if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); - /* A TP gate uses one monotonic shared event for the whole token. A - * later command buffer may signal a higher value while the prefix is - * blocked at an earlier gate, making the transport consume a slab - * slot before its payload is ready. Keep each TP token in one command - * buffer; non-TP decode retains the encode/execute overlap. */ - if (ok && allow_split_flush && g->tp_world != 2 && - split_after_layers != 0 && il + 1u == split_after_layers) { - ok = ds4_gpu_flush_commands() != 0; - } - } - - if (ok && need_logits) { - ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - } - return ok; -} - -static ds4_gpu_tensor *metal_graph_tensor_row_view( - ds4_gpu_tensor *base, - uint32_t row, - uint64_t row_values) { - return ds4_gpu_tensor_view(base, - (uint64_t)row * row_values * sizeof(float), - row_values * sizeof(float)); -} - -/* Upload prompt token ids for kernels that need token-aware hash routing. */ -static bool metal_graph_upload_prompt_tokens( - ds4_gpu_tensor *out_tokens, - const token_vec *prompt, - uint32_t pos0, - uint32_t n_tokens) { - if (!out_tokens || pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) { - return false; - } - - int32_t *tokens = xmalloc((size_t)n_tokens * sizeof(tokens[0])); - for (uint32_t i = 0; i < n_tokens; i++) tokens[i] = prompt->v[pos0 + i]; - - const bool ok = ds4_gpu_tensor_write(out_tokens, - 0, - tokens, - (uint64_t)n_tokens * sizeof(tokens[0])) != 0; - free(tokens); - return ok; -} - -/* Rebuild ratio-4 compressor state after chunked prefill so a following decode - * token sees the same rolling compression window. */ -static bool metal_graph_refresh_ratio4_compressor_state( - ds4_gpu_graph *g, - const ds4_model *model, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_tensor *kv_weight, - const ds4_tensor *score_weight, - const ds4_tensor *ape, - uint32_t head_dim, - uint32_t width, - uint32_t pos0, - uint32_t n_tokens) { - if (n_tokens < 4) { - return true; - } - if (!g || !model || !state_kv || !state_score || !kv_weight || !score_weight || !ape || - head_dim == 0 || width == 0) { - return false; - } - - /* - * The recurrent ratio-4 state is intentionally rebuilt from the last - * four tokens using the small-batch projection kernel. The full-chunk - * projection is already available, but it uses the matrix-matrix path; - * mixing those two accumulation orders changes a few FP8 rounding - * decisions in later chunks. - */ - ds4_gpu_tensor *tail_hc = ds4_gpu_tensor_view( - metal_graph_batch_attn_norm(g), - (uint64_t)(n_tokens - 4u) * DS4_N_EMBD * sizeof(float), - 4ull * DS4_N_EMBD * sizeof(float)); - bool ok = tail_hc != NULL; - if (!ok) { - fprintf(stderr, "ds4: ratio-4 compressor tail view creation failed\n"); - } - if (ok) { -#if defined(__APPLE__) - ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), - model->map, - model->size, - kv_weight->abs_offset, - DS4_N_EMBD, - width, - tail_hc, - 4) != 0; - if (ok) { - ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), - model->map, - model->size, - score_weight->abs_offset, - DS4_N_EMBD, - width, - tail_hc, - 4) != 0; - } -#else - ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - kv_weight->abs_offset, - score_weight->abs_offset, - DS4_N_EMBD, - width, - tail_hc, - 4) != 0; -#endif - if (!ok) { - fprintf(stderr, "ds4: ratio-4 compressor tail projection failed\n"); - } - } - if (ok) { - ok = ds4_gpu_compressor_prefill_state_ratio4_tensor(state_kv, - state_score, - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - ape->abs_offset, - ape->type, - head_dim, - pos0 + n_tokens - 4u) != 0; - if (!ok) { - fprintf(stderr, "ds4: ratio-4 compressor state refresh failed\n"); - } - } - ds4_gpu_tensor_free(tail_hc); - return ok; -} - -/* CPU fallback for seeding batched HC state from token embeddings. It is still - * useful for tiny speculative verifier batches where a separate GPU embedding - * command buffer costs more than the small host write. */ -static bool metal_graph_upload_prompt_embeddings_hc_cpu( - ds4_gpu_tensor *out_hc, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t pos0, - uint32_t n_tokens) { - if (pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) return false; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t total = (uint64_t)n_tokens * hc_dim; - float *hc = xmalloc((size_t)total * sizeof(hc[0])); - float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); - - for (uint32_t t = 0; t < n_tokens; t++) { - embed_token_f16(model, weights, prompt->v[pos0 + t], plain); - float *dst = hc + (uint64_t)t * hc_dim; - for (uint32_t h = 0; h < DS4_N_HC; h++) { - memcpy(dst + (uint64_t)h * DS4_N_EMBD, - plain, - (size_t)DS4_N_EMBD * sizeof(plain[0])); - } - } - - const bool ok = ds4_gpu_tensor_write(out_hc, 0, hc, total * sizeof(hc[0])) != 0; - free(plain); - free(hc); - return ok; -} - -/* Seed the batched HC state from token ids: every HC stream starts as the same - * 4096-wide embedding. Long prefill chunks use the Metal get-rows/repeat - * kernel so the CPU does not build and upload a large [token, HC, dim] tensor. */ -static bool metal_graph_upload_prompt_embeddings_hc( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *tokens, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t pos0, - uint32_t n_tokens) { - if (pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) return false; - - uint32_t gpu_min = 512; -#ifndef DS4_ROCM_BUILD - const char *gpu_min_env = getenv("DS4_METAL_GPU_BATCH_EMBED_MIN"); - if (gpu_min_env && gpu_min_env[0]) { - char *end = NULL; - unsigned long v = strtoul(gpu_min_env, &end, 10); - if (end != gpu_min_env && v <= UINT32_MAX) gpu_min = (uint32_t)v; - } -#endif - - if (tokens && n_tokens >= gpu_min) { - return ds4_gpu_embed_tokens_hc_tensor(out_hc, - tokens, - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - n_tokens, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - - return metal_graph_upload_prompt_embeddings_hc_cpu(out_hc, - model, - weights, - prompt, - pos0, - n_tokens); -} - -static bool metal_graph_hc_rms_scale_project( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_scratch, - const ds4_model *model, - const ds4_tensor *weight, - const ds4_gpu_tensor *x, - uint64_t in_dim, - uint32_t n_tokens) { - if (!out || !norm_scratch || !model || !weight || !x || - in_dim > UINT32_MAX) { - return false; - } -#if defined(__APPLE__) - return ds4_gpu_hc_rms_scale_project_f16_tensor( - out, - norm_scratch, - model->map, - model->size, - weight->abs_offset, - (uint32_t)in_dim, - 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, - x, - n_tokens, - DS4_RMS_EPS) != 0; -#else - bool ok = ds4_gpu_rms_norm_plain_rows_tensor( - norm_scratch, - x, - (uint32_t)in_dim, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) { - ok = ds4_gpu_matmul_f16_tensor( - out, - model->map, - model->size, - weight->abs_offset, - in_dim, - 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, - norm_scratch, - n_tokens) != 0; - } - return ok; -#endif -} - -static bool metal_graph_warmup_prefill_kernels( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t n_tokens) { - static bool warmed = false; - if (g && g->ssd_streaming) return true; - if (warmed) return true; -#ifndef DS4_ROCM_BUILD - if (getenv("DS4_METAL_NO_PREFILL_KERNEL_WARMUP") != NULL) return true; -#endif - - /* - * The first batched F16 matmul can pay Metal's one-time pipeline execution - * cost. Run the same HC attention projection on scratch storage before the - * measured prefill. The output is overwritten by the real graph. - */ - if (n_tokens <= 8) return true; - - /* (B6 fix, ): warm-up uses layer-0's hc_attn_fn - * weight, which in multi-tier is resolved on placement[1]'s tier. - * Switch active_tier so the F16 matmul reads/writes the correct - * Class P scratch and resolves the weight on the right device. - * Single-tier (g->placement == NULL): no-op. */ - if (g->placement) { - if (!metal_graph_set_active_tier_batch(g, g->placement[1], n_tokens)) return false; - } - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = metal_graph_hc_rms_scale_project( - metal_graph_batch_hc_mix(g), - metal_graph_batch_flat_hc(g), - model, - weights->layer[0].hc_attn_fn, - metal_graph_batch_cur_hc(g), - hc_dim, - n_tokens); - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) { - fprintf(stderr, "ds4: Metal prefill kernel warmup failed\n"); - return false; - } - - warmed = true; - return true; -} - -/* Encode the batched prefill attention half for one layer. It mirrors the CPU - * layer-major path: HC pre/norm, Q/KV, cache/compression, prefix attention. */ -static bool metal_graph_indexer_stage_profile_boundary( - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_comp, - double *stage_t0) { - if (ds4_gpu_end_commands() == 0) return false; - const double now = now_sec(); - if (stage != NULL) { - fprintf(stderr, - "ds4: metal indexer stage layer=%u pos=%u tokens=%u comp=%u %s=%.3f ms\n", - il, - pos0, - n_tokens, - n_comp, - stage, - (now - *stage_t0) * 1000.0); - } - *stage_t0 = now; - return ds4_gpu_begin_commands() != 0; -} - -static bool metal_graph_env_value_eq(const char *v, - size_t n, - const char *literal) { - const size_t m = strlen(literal); - if (n != m) return false; - for (size_t i = 0; i < n; i++) { - if (tolower((unsigned char)v[i]) != - tolower((unsigned char)literal[i])) { - return false; - } - } - return true; -} - -static const char *metal_graph_env_trim(const char *v, size_t *len_out) { - if (!v) { - if (len_out) *len_out = 0; - return NULL; - } - while (isspace((unsigned char)*v)) v++; - size_t n = strlen(v); - while (n > 0 && isspace((unsigned char)v[n - 1])) n--; - if (len_out) *len_out = n; - return v; -} - -static bool metal_graph_profile_layer_value_match(const char *layer_env, - uint32_t il) { - size_t n = 0; - layer_env = metal_graph_env_trim(layer_env, &n); - if (!layer_env || n == 0) return true; - - char *end = NULL; - const unsigned long layer = strtoul(layer_env, &end, 10); - return end != layer_env && - (size_t)(end - layer_env) == n && - layer <= UINT32_MAX && - (uint32_t)layer == il; -} - -static bool metal_graph_stage_profile_enabled_for_layer( - const char *flag_env_name, - const char *layer_env_name, - uint32_t il) { - size_t flag_len = 0; - const char *flag = metal_graph_env_trim(getenv(flag_env_name), &flag_len); - if (!flag) return false; - - const char *layer_env = getenv(layer_env_name); - const bool has_layer_filter = layer_env && layer_env[0]; - - if (flag_len != 0) { - if (metal_graph_env_value_eq(flag, flag_len, "0") || - metal_graph_env_value_eq(flag, flag_len, "false") || - metal_graph_env_value_eq(flag, flag_len, "no") || - metal_graph_env_value_eq(flag, flag_len, "off")) { - return false; - } - if (!has_layer_filter && - !metal_graph_env_value_eq(flag, flag_len, "1") && - !metal_graph_env_value_eq(flag, flag_len, "true") && - !metal_graph_env_value_eq(flag, flag_len, "yes") && - !metal_graph_env_value_eq(flag, flag_len, "on") && - !metal_graph_env_value_eq(flag, flag_len, "all")) { - return metal_graph_profile_layer_value_match(flag, il); - } - } - - return metal_graph_profile_layer_value_match(layer_env, il); -} - -static bool metal_graph_layer_stage_profile_enabled(uint32_t il) { - return metal_graph_stage_profile_enabled_for_layer( - "DS4_ROCM_LAYER_STAGE_PROFILE", - "DS4_ROCM_LAYER_STAGE_PROFILE_LAYER", - il) || - metal_graph_stage_profile_enabled_for_layer( - "DS4_METAL_LAYER_STAGE_PROFILE", - "DS4_METAL_LAYER_STAGE_PROFILE_LAYER", - il); -} - -static bool metal_graph_decode_stage_profile_enabled(uint32_t il) { - return metal_graph_stage_profile_enabled_for_layer( - "DS4_ROCM_DECODE_STAGE_PROFILE", - "DS4_ROCM_DECODE_STAGE_PROFILE_LAYER", - il) || - metal_graph_stage_profile_enabled_for_layer( - "DS4_METAL_DECODE_STAGE_PROFILE", - "DS4_METAL_DECODE_STAGE_PROFILE_LAYER", - il); -} - -static bool metal_graph_layer_stage_profile_start(uint32_t il) { - if (!metal_graph_layer_stage_profile_enabled(il)) return true; - if (ds4_gpu_end_commands() == 0) return false; - return ds4_gpu_begin_commands() != 0; -} - -/* Optional prefill stage profiler. It intentionally ends the current Metal - * command buffer and waits, so the printed number includes encoding plus GPU - * execution for the stage just emitted. This is disabled by default because it - * adds synchronization points and changes scheduling. */ -static bool metal_graph_layer_stage_profile_boundary( - const char *part, - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - double *stage_t0) { - if (ds4_gpu_end_commands() == 0) return false; - const double now = now_sec(); - if (stage != NULL) { - fprintf(stderr, - "ds4: metal layer stage part=%s layer=%u pos=%u tokens=%u %s=%.3f ms\n", - part, - il, - pos0, - n_tokens, - stage, - (now - *stage_t0) * 1000.0); - } - *stage_t0 = now; - return ds4_gpu_begin_commands() != 0; -} - -static bool metal_graph_q_stage_profile_boundary( - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - double *stage_t0) { - if (ds4_gpu_end_commands() == 0) return false; - const double now = now_sec(); - fprintf(stderr, - "ds4: metal Q path stage layer=%u pos=%u tokens=%u %s=%.3f ms\n", - il, - pos0, - n_tokens, - stage, - (now - *stage_t0) * 1000.0); - *stage_t0 = now; - return ds4_gpu_begin_commands() != 0; -} - -static ds4_gpu_tensor *metal_graph_tensor_row_range_view( - ds4_gpu_tensor *base, - uint32_t row0, - uint32_t rows, - uint64_t row_values) { - return ds4_gpu_tensor_view(base, - (uint64_t)row0 * row_values * sizeof(float), - (uint64_t)rows * row_values * sizeof(float)); -} - -/* TP prefill threshold for row-splitting the replicated shared expert. - * Routed experts remain ownership-split at every batch size. */ -static uint32_t metal_graph_tp_prefill_split_min(void) { - static int cached = -1; - if (cached < 0) { - cached = 32; - const char *env = getenv("DS4_TP_PREFILL_SPLIT_MIN"); - if (env && env[0]) cached = atoi(env); - if (cached < 2) cached = 2; - } - return (uint32_t)cached; -} - -/* Opt-in sub-chunk gate pipelining for the TP prefill row swaps. Must be - * set on BOTH ranks (it changes the per-layer gate count; asymmetric - * settings deadlock the big gates). Default off: measured net-negative - * on the M5 Max pair, see the pipelined blocks for the numbers. */ -static bool metal_graph_tp_subgate_pipeline(void) { - static int cached = -1; - if (cached < 0) { - const char *env = getenv("DS4_TP_SUBGATE_PIPELINE"); - cached = env && env[0] && atoi(env) != 0; - } - return cached != 0; -} - -static bool metal_graph_encode_layer_attention_batch( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens) { - if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint32_t n_groups = DS4_N_OUT_GROUP; - const uint32_t group_heads = DS4_N_HEAD / n_groups; - const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; - const uint32_t rank = DS4_N_LORA_O; - const uint32_t ratio = ds4_layer_compress_ratio(il); - const bool compressed = ratio != 0; - const bool zero_prefix = pos0 == 0; - /* TP attention row split for large zero-prefix chunks: q_a and the KV - * path stay full (both ranks need every row's KV, and the compressor/ - * indexer keep updating their state from full rows), q_b onward runs on - * this rank's half of the chunk rows, and the computed row halves of - * batch_attn_out are swapped in place through one big gate per layer. - * Three chunk shapes split: full-raw (uncompressed layer at pos0 == 0, - * or a compressed layer whose chunk is too short to emit compressed - * keys, i.e. raw_prefix_tokens == n_tokens), static-mixed (compressed - * layer whose whole chunk attends through the one-shot mixed kernel - * over the full raw keys plus n_tokens / ratio compressed keys, without - * indexer top-k), and indexed (ratio-4 layer with indexer top-k, whose - * per-token score/top-k selection stays replicated while the attention - * consumption splits by rows). Every condition derives from - * pos0/n_tokens/ratio/model shape so both ranks stay in lockstep. */ - const bool tp_attn_full_raw = zero_prefix && - (ratio == 0 || (n_tokens < ratio && n_tokens <= g->raw_cap)); - const uint32_t tp_attn_n_comp = ratio != 0 ? n_tokens / ratio : 0; - const bool tp_attn_static_mixed = zero_prefix && ratio != 0 && - tp_attn_n_comp != 0 && - !(ratio == 4 && tp_attn_n_comp > DS4_N_INDEXER_TOP_K); - const bool tp_attn_indexed = zero_prefix && ratio == 4 && - tp_attn_n_comp > DS4_N_INDEXER_TOP_K; - const bool tp_row_split_attn = - g->tp_world == 2 && - g->tp_batch_rows != n_tokens && - (tp_attn_full_raw || tp_attn_static_mixed || tp_attn_indexed) && - !metal_graph_directional_steering_attn_enabled(g) && - n_tokens >= metal_graph_tp_prefill_split_min(); - const uint32_t tp_half_rows = (n_tokens + 1u) / 2u; - const uint32_t tp_row0 = (tp_row_split_attn && g->tp_rank != 0) ? tp_half_rows : 0; - const uint32_t tp_rows = tp_row_split_attn ? - (g->tp_rank == 0 ? tp_half_rows : n_tokens - tp_half_rows) : n_tokens; - const bool index_stage_profile = - glm_graph_env_present("DS4_ROCM_INDEXER_STAGE_PROFILE", - "DS4_METAL_INDEXER_STAGE_PROFILE"); - const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); - const bool q_stage_profile = - glm_graph_env_present("DS4_ROCM_Q_STAGE_PROFILE", - "DS4_METAL_Q_STAGE_PROFILE"); - double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; - double q_stage_t0 = q_stage_profile ? now_sec() : 0.0; -#define DS4_METAL_PROFILE_ATTN_STAGE(name) do { \ - if (ok && layer_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary("attn", (name), il, pos0, n_tokens, &layer_stage_t0); \ - } \ - } while (0) -#define DS4_METAL_PROFILE_Q_STAGE(name) do { \ - if (ok && q_stage_profile) { \ - ok = metal_graph_q_stage_profile_boundary((name), il, pos0, n_tokens, &q_stage_t0); \ - } \ - } while (0) - const float freq_base = layer_rope_freq_base(il); - const float freq_scale = layer_rope_freq_scale(il); - const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; - float attn_factor = 1.0f; - if (ext_factor != 0.0f && freq_scale > 0.0f) { - attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - enum { stack_count_cap = 16 }; - uint32_t comp_counts_stack[stack_count_cap]; - uint32_t index_counts_stack[stack_count_cap]; - uint32_t *comp_counts = NULL; - uint32_t *index_counts = NULL; - if (compressed) { - if (n_tokens <= stack_count_cap) { - memset(comp_counts_stack, 0, - (size_t)n_tokens * sizeof(comp_counts_stack[0])); - comp_counts = comp_counts_stack; - } else { - comp_counts = xcalloc(n_tokens, sizeof(comp_counts[0])); - } - } - if (ratio == 4) { - if (n_tokens <= stack_count_cap) { - memset(index_counts_stack, 0, - (size_t)n_tokens * sizeof(index_counts_stack[0])); - index_counts = index_counts_stack; - } else { - index_counts = xcalloc(n_tokens, sizeof(index_counts[0])); - } - } - const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); - ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( - metal_graph_batch_hc_mix(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( - metal_graph_batch_hc_split(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *attn_cur_view = ds4_gpu_tensor_view( - metal_graph_batch_attn_cur(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( - metal_graph_batch_after_attn_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); - bool ok = hc_mix_view && hc_split_view && attn_cur_view && after_attn_hc_view; - const bool fuse_hc_norm = n_tokens > 1 && - DS4_N_HC == 4 && - !metal_graph_use_reference_hc_decode() && - metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, - metal_graph_batch_flat_hc(g), - model, - layer->hc_attn_fn, - metal_graph_batch_cur_hc(g), - hc_dim, - n_tokens); - if (metal_graph_use_reference_hc_decode()) { - if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, - hc_mix_view, - model->map, - model->size, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(attn_cur_view, - metal_graph_batch_cur_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (fuse_hc_norm) { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, - metal_graph_batch_attn_norm(g), - hc_split_view, - hc_mix_view, - metal_graph_batch_cur_hc(g), - model->map, - model->size, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset, - layer->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - } else { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, - hc_split_view, - hc_mix_view, - metal_graph_batch_cur_hc(g), - model->map, - model->size, - layer->hc_attn_scale->abs_offset, - layer->hc_attn_base->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre", metal_graph_batch_attn_cur(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - DS4_METAL_PROFILE_ATTN_STAGE("hc_pre"); - if (ok && !fuse_hc_norm) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), - metal_graph_batch_attn_cur(g), - model->map, - model->size, - layer->attn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("attn_norm", metal_graph_batch_attn_norm(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - DS4_METAL_PROFILE_ATTN_STAGE("norm"); - DS4_METAL_PROFILE_Q_STAGE("pre_q"); - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_a", - il, - pos0, - metal_graph_batch_qr(g), - model, - layer->attn_q_a, - DS4_N_EMBD, - q_rank, - metal_graph_batch_attn_norm(g), - n_tokens); - if (ok) { - metal_graph_debug_dump_tensor("q_lora", metal_graph_batch_qr(g), - (uint64_t)n_tokens * q_rank, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("q_a"); - if (qkv_rms_fused) { - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", - il, - pos0, - metal_graph_batch_kv_raw(g), - model, - layer->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - metal_graph_batch_attn_norm(g), - n_tokens); - if (ok) { - metal_graph_debug_dump_tensor("KVraw", metal_graph_batch_kv_raw(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(metal_graph_batch_qr_norm(g), - metal_graph_batch_qr(g), - model->map, - model->size, - layer->attn_q_a_norm->abs_offset, - (uint32_t)q_rank, - metal_graph_batch_kv(g), - metal_graph_batch_kv_raw(g), - layer->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - n_tokens, - DS4_RMS_EPS) != 0; - } else { - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_qr_norm(g), - metal_graph_batch_qr(g), - model->map, - model->size, - layer->attn_q_a_norm->abs_offset, - (uint32_t)q_rank, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("q_lora_norm", metal_graph_batch_qr_norm(g), - (uint64_t)n_tokens * q_rank, il, pos0); - } - if (qkv_rms_fused && ok) { - metal_graph_debug_dump_tensor("KVnorm", metal_graph_batch_kv(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("q_a_norm"); - const bool q_path_debug = - metal_graph_debug_wants("Qraw", il, pos0) || - metal_graph_debug_wants("Qnorm", il, pos0); - /* Under the TP row split everything from q_b to the output projection - * runs on this rank's rows only, through row-range views of the batch - * tensors (batch_q_half is F16, so its view is built directly). */ - ds4_gpu_tensor *tp_q = tp_row_split_attn ? - metal_graph_tensor_row_range_view(metal_graph_batch_q(g), tp_row0, tp_rows, q_dim) : NULL; - ds4_gpu_tensor *tp_q_half = tp_row_split_attn ? - ds4_gpu_tensor_view(g->batch_q_half, - (uint64_t)tp_row0 * q_dim * sizeof(uint16_t), - (uint64_t)tp_rows * q_dim * sizeof(uint16_t)) : NULL; - ds4_gpu_tensor *tp_qr_norm = tp_row_split_attn ? - metal_graph_tensor_row_range_view(metal_graph_batch_qr_norm(g), tp_row0, tp_rows, q_rank) : NULL; - ds4_gpu_tensor *tp_heads = tp_row_split_attn ? - metal_graph_tensor_row_range_view(metal_graph_batch_heads(g), tp_row0, tp_rows, q_dim) : NULL; - ds4_gpu_tensor *tp_attn_out = tp_row_split_attn ? - metal_graph_tensor_row_range_view(metal_graph_batch_attn_out(g), tp_row0, tp_rows, - DS4_N_EMBD) : NULL; - if (tp_row_split_attn && - (!tp_q || !tp_q_half || !tp_qr_norm || !tp_heads || !tp_attn_out)) { - ok = false; - } - bool q_b_f16_out = false; - if (ok && !q_path_debug && layer->attn_q_b->type == DS4_TENSOR_Q8_0) { - q_b_f16_out = ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(tp_q ? tp_q : metal_graph_batch_q(g), - tp_q_half ? tp_q_half : g->batch_q_half, - model->map, - model->size, - layer->attn_q_b->abs_offset, - q_rank, - q_dim, - tp_qr_norm ? tp_qr_norm : metal_graph_batch_qr_norm(g), - tp_rows, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0 + tp_row0, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } - if (q_b_f16_out) { - DS4_METAL_PROFILE_Q_STAGE("q_b"); - DS4_METAL_PROFILE_Q_STAGE("head_norm"); - if (ok) { - metal_graph_debug_dump_tensor("Qcur", metal_graph_batch_q(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("rope"); - } else { - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_b", - il, - pos0, - tp_q ? tp_q : metal_graph_batch_q(g), - model, - layer->attn_q_b, - q_rank, - q_dim, - tp_qr_norm ? tp_qr_norm : metal_graph_batch_qr_norm(g), - tp_rows); - if (ok) { - metal_graph_debug_dump_tensor("Qraw", metal_graph_batch_q(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("q_b"); - if (ok) ok = ds4_gpu_head_rms_norm_tensor(tp_q ? tp_q : metal_graph_batch_q(g), - tp_rows, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_RMS_EPS) != 0; - if (ok) { - metal_graph_debug_dump_tensor("Qnorm", metal_graph_batch_q(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("head_norm"); - if (ok) ok = ds4_gpu_rope_tail_tensor(tp_q ? tp_q : metal_graph_batch_q(g), - tp_rows, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0 + tp_row0, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) { - metal_graph_debug_dump_tensor("Qcur", metal_graph_batch_q(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - DS4_METAL_PROFILE_Q_STAGE("rope"); - } - DS4_METAL_PROFILE_ATTN_STAGE("q_path"); - if (!qkv_rms_fused) { - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", - il, - pos0, - metal_graph_batch_kv_raw(g), - model, - layer->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - metal_graph_batch_attn_norm(g), - n_tokens); - if (ok) { - metal_graph_debug_dump_tensor("KVraw", metal_graph_batch_kv_raw(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), - metal_graph_batch_kv_raw(g), - model->map, - model->size, - layer->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVnorm", metal_graph_batch_kv(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - } - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), - n_tokens, - DS4_N_HEAD_KV, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVrope", metal_graph_batch_kv(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), - n_tokens, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVcur", metal_graph_batch_kv(g), - (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); - } - DS4_METAL_PROFILE_ATTN_STAGE("kv_path"); - /* - * Static graph order is q, kv, cpy_k(raw SWA), then attention. For a - * zero-prefix batch it is safe to store the whole batch at once: attention - * reads the contiguous batch KV, and the ring only has to end with the last - * SWA rows for later chunks/decode. For nonzero chunks the physical ring is - * sized to hold the current chunk plus the previous SWA window, while the - * attention mask still enforces the 128-token logical window. - */ - if (ok && zero_prefix) ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - metal_graph_batch_kv(g), - g->raw_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u raw KV batch store failed\n", il); - } - const bool raw_batch_attention = zero_prefix && ratio == 0; - bool batch_attention_done = false; - - if (ok && raw_batch_attention) { - if (tp_row_split_attn) { - ok = ds4_gpu_attention_prefill_raw_heads_range_tensor(tp_heads, - model->map, - model->size, - layer->attn_sinks->abs_offset, - tp_q, - metal_graph_batch_kv(g), - tp_row0, - tp_rows, - n_tokens, - g->raw_window, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } else { - ok = ds4_gpu_attention_prefill_raw_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - metal_graph_batch_kv(g), - n_tokens, - g->raw_window, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - if (ok) batch_attention_done = true; - } else if (ok && !zero_prefix && ratio == 0 && n_tokens <= g->raw_cap) { - /* - * The ubatch path stores the whole batch in the SWA cache, then runs - * one batched attention kernel with an absolute-position causal/window - * mask. This avoids mixing prefill with the different single-token - * attention path. - */ - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos0, n_tokens); - /* Nonzero prompt chunks read the SWA cache as a ring. FlashAttention - * receives a linearized window starting at raw_start, not physical row - * zero; otherwise wrapped chunks silently miss recent raw keys. */ - const uint32_t raw_start = metal_graph_raw_start_for_span(g, - pos0 + n_tokens - 1u, - n_raw); - ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - metal_graph_batch_kv(g), - g->raw_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; - if (ok) { - metal_graph_debug_dump_tensor("raw_cache", - g->layer_raw_cache[il], - (uint64_t)n_raw * DS4_N_HEAD_DIM, - il, - pos0); - } - if (ok) { - ok = ds4_gpu_attention_decode_raw_batch_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->layer_raw_cache[il], - n_tokens, - pos0, - n_raw, - g->raw_cap, - raw_start, - g->raw_window, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - if (ok) batch_attention_done = true; - } else if (ok && ratio != 0) { - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint32_t comp_width = coff * DS4_N_HEAD_DIM; - const bool have_attn_comp = layer->attn_compressor_kv && layer->attn_compressor_gate && - layer->attn_compressor_ape && layer->attn_compressor_norm; - if (!have_attn_comp) { - fprintf(stderr, "ds4: Metal layer-major prefill needs attention compressor weights\n"); - ok = false; - } - if (ok) { - ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), - model->map, - model->size, - layer->attn_compressor_kv->abs_offset, - DS4_N_EMBD, - comp_width, - metal_graph_batch_attn_norm(g), - n_tokens) != 0; - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u attention compressor KV projection failed\n", il); - } - if (ok) { - ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->attn_compressor_gate->abs_offset, - DS4_N_EMBD, - comp_width, - metal_graph_batch_attn_norm(g), - n_tokens) != 0; - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u attention compressor score projection failed\n", il); - } - } - } - if (ok) metal_graph_debug_dump_tensor("attn_comp_kv_raw", - metal_graph_batch_comp_kv(g), - (uint64_t)comp_width * n_tokens, - il, - pos0); - if (ok) metal_graph_debug_dump_tensor("attn_comp_score_raw", - metal_graph_batch_comp_sc(g), - (uint64_t)comp_width * n_tokens, - il, - pos0); - uint32_t n_comp = g->layer_n_comp[il]; - if (zero_prefix) { - n_comp = n_tokens / ratio; - if (ok && n_comp > g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal layer-major compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - } - if (ok && DS4_GPU_ATTN_COMP_CACHE_F16 && n_comp > g->attn_comp_stage_cap) { - fprintf(stderr, "ds4: Metal graph compressed KV staging capacity exceeded at layer %u\n", il); - ok = false; - } - ds4_gpu_tensor *attn_comp_target = NULL; - if (ok) { - attn_comp_target = metal_graph_attn_comp_prefill_target(g, il, 0, n_comp); - if (!attn_comp_target) { - fprintf(stderr, "ds4: gpu layer %u attention compressor target creation failed\n", il); - ok = false; - } - if (ok) ok = ds4_gpu_compressor_prefill_tensor(attn_comp_target, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - layer->attn_compressor_norm->abs_offset, - layer->attn_compressor_norm->type, - DS4_N_HEAD_DIM, - ratio, - pos0, - n_tokens, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u attention compressor prefill failed\n", il); - } - DS4_METAL_PROFILE_ATTN_STAGE("compressor_prefill"); - if (ok && n_comp != 0) { - ok = metal_graph_commit_attn_comp_stage(g, il, 0, n_comp); - } - DS4_METAL_PROFILE_ATTN_STAGE("compressor_commit"); - if (ok && ratio == 4) { - ok = metal_graph_refresh_ratio4_compressor_state(g, - model, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - layer->attn_compressor_kv, - layer->attn_compressor_gate, - layer->attn_compressor_ape, - DS4_N_HEAD_DIM, - comp_width, - pos0, - n_tokens); - } - DS4_METAL_PROFILE_ATTN_STAGE("compressor_refresh"); - } - if (ok) { - g->layer_n_comp[il] = n_comp; - for (uint32_t t = 0; t < n_tokens; t++) { - comp_counts[t] = (pos0 + t + 1u) / ratio; - } - if (n_comp != 0) { - metal_graph_debug_dump_tensor("KVcompress", - attn_comp_target, - (uint64_t)n_comp * DS4_N_HEAD_DIM, - il, - pos0); - } - metal_graph_debug_dump_tensor("attn_state_kv", - g->layer_attn_state_kv[il], - (uint64_t)comp_width * coff * ratio, - il, - pos0); - metal_graph_debug_dump_tensor("attn_state_score", - g->layer_attn_state_score[il], - (uint64_t)comp_width * coff * ratio, - il, - pos0); - } - metal_graph_attn_comp_prefill_target_free(attn_comp_target); - } else { - const bool aligned_chunk = - getenv("DS4_CUDA_NO_COMPRESSOR_PREFILL_BATCH") == NULL && - (pos0 % ratio) == 0u && (n_tokens % ratio) == 0u; - if (aligned_chunk) { - const uint32_t comp_before = g->layer_n_comp[il]; - const uint32_t comp_chunk = n_tokens / ratio; - if (comp_before + comp_chunk > g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - } - if (ok && DS4_GPU_ATTN_COMP_CACHE_F16 && comp_chunk > g->attn_comp_stage_cap) { - fprintf(stderr, "ds4: Metal graph compressed KV staging capacity exceeded at layer %u\n", il); - ok = false; - } - ds4_gpu_tensor *attn_comp_target = - ok ? metal_graph_attn_comp_prefill_target(g, il, comp_before, comp_chunk) : NULL; - if (ok && !attn_comp_target) ok = false; - if (ok && ratio == 4) { - ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( - attn_comp_target, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - layer->attn_compressor_norm->abs_offset, - layer->attn_compressor_norm->type, - DS4_N_HEAD_DIM, - pos0, - n_tokens, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } else if (ok) { - ok = ds4_gpu_compressor_prefill_tensor( - attn_comp_target, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - layer->attn_compressor_norm->abs_offset, - layer->attn_compressor_norm->type, - DS4_N_HEAD_DIM, - ratio, - pos0, - n_tokens, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } - if (ok && comp_chunk != 0) { - ok = metal_graph_commit_attn_comp_stage(g, il, comp_before, comp_chunk); - } - if (ok && ratio == 4) { - ok = metal_graph_refresh_ratio4_compressor_state(g, - model, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - layer->attn_compressor_kv, - layer->attn_compressor_gate, - layer->attn_compressor_ape, - DS4_N_HEAD_DIM, - comp_width, - pos0, - n_tokens); - } - if (ok) { - g->layer_n_comp[il] = comp_before + comp_chunk; - if (comp_counts) { - for (uint32_t t = 0; t < n_tokens; t++) { - comp_counts[t] = (pos0 + t + 1u) / ratio; - } - } - metal_graph_debug_dump_tensor("KVcompress", - attn_comp_target, - (uint64_t)comp_chunk * DS4_N_HEAD_DIM, - il, - pos0); - metal_graph_debug_dump_tensor("attn_state_kv", - g->layer_attn_state_kv[il], - (uint64_t)comp_width * coff * ratio, - il, - pos0); - metal_graph_debug_dump_tensor("attn_state_score", - g->layer_attn_state_score[il], - (uint64_t)comp_width * coff * ratio, - il, - pos0); - } - metal_graph_attn_comp_prefill_target_free(attn_comp_target); - } else { - for (uint32_t t = 0; ok && t < n_tokens; t++) { - const uint32_t pos = pos0 + t; - const bool emit = ((pos + 1u) % ratio) == 0u; - if (emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - break; - } - ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(metal_graph_batch_comp_kv(g), t, comp_width); - ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(metal_graph_batch_comp_sc(g), t, comp_width); - const uint32_t comp_row = g->layer_n_comp[il]; - ok = kv_view && sc_view && - ds4_gpu_compressor_update_tensor(kv_view, - sc_view, - g->layer_attn_state_kv[il], - g->layer_attn_state_score[il], - metal_graph_attn_comp_update_target(g, il), - model->map, - model->size, - layer->attn_compressor_ape->abs_offset, - layer->attn_compressor_ape->type, - layer->attn_compressor_norm->abs_offset, - layer->attn_compressor_norm->type, - DS4_N_HEAD_DIM, - ratio, - pos, - metal_graph_attn_comp_update_row(comp_row), - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS, - false) != 0; - if (ok && emit) { - ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); - ok = comp_row_view && - ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; - if (ok) { - metal_graph_debug_dump_tensor("KVcompress", - comp_row_view, - DS4_N_HEAD_DIM, - il, - pos); - } - ds4_gpu_tensor_free(comp_row_view); - if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); - } - if (ok && emit) g->layer_n_comp[il]++; - if (comp_counts) comp_counts[t] = g->layer_n_comp[il]; - if (ok && t == 0) ok = metal_graph_capture_prefix1_attn_state(g, il); - ds4_gpu_tensor_free(sc_view); - ds4_gpu_tensor_free(kv_view); - } - } - n_comp = g->layer_n_comp[il]; - } - DS4_METAL_PROFILE_ATTN_STAGE("compressor"); - - if (ok && ratio == 4) { - const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; - if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || - !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || - !layer->indexer_attn_q_b || !layer->indexer_proj) { - fprintf(stderr, "ds4: Metal layer-major prefill needs indexer weights\n"); - ok = false; - } - if (ok) { - ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), - model->map, - model->size, - layer->indexer_compressor_kv->abs_offset, - DS4_N_EMBD, - index_width, - metal_graph_batch_attn_norm(g), - n_tokens) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->indexer_compressor_gate->abs_offset, - DS4_N_EMBD, - index_width, - metal_graph_batch_attn_norm(g), - n_tokens) != 0; - } - if (ok) metal_graph_debug_dump_tensor("indexer_comp_kv_raw", - metal_graph_batch_comp_kv(g), - (uint64_t)index_width * n_tokens, - il, - pos0); - if (ok) metal_graph_debug_dump_tensor("indexer_comp_score_raw", - metal_graph_batch_comp_sc(g), - (uint64_t)index_width * n_tokens, - il, - pos0); - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_indexer_q(g), - model, - layer->indexer_attn_q_b, - q_rank, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - metal_graph_batch_qr_norm(g), - n_tokens); - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_indexer_q(g), - n_tokens, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - pos0, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_batch_indexer_q(g), - n_tokens * DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_indexer_weights(g), - model->map, - model->size, - layer->indexer_proj->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD, - metal_graph_batch_attn_norm(g), - n_tokens) != 0; - if (zero_prefix) { - if (ok && n_comp > g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal layer-major indexer cache capacity exceeded at layer %u\n", il); - ok = false; - } - if (ok) { - ok = ds4_gpu_compressor_prefill_tensor(g->layer_index_comp_cache[il], - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->indexer_compressor_ape->abs_offset, - layer->indexer_compressor_ape->type, - layer->indexer_compressor_norm->abs_offset, - layer->indexer_compressor_norm->type, - DS4_N_INDEXER_HEAD_DIM, - ratio, - pos0, - n_tokens, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } - if (ok && n_comp != 0) { - ok = ds4_gpu_dsv4_indexer_qat_tensor(g->layer_index_comp_cache[il], - n_comp, - DS4_N_INDEXER_HEAD_DIM) != 0; - } - if (ok) { - ok = metal_graph_refresh_ratio4_compressor_state(g, - model, - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - layer->indexer_compressor_kv, - layer->indexer_compressor_gate, - layer->indexer_compressor_ape, - DS4_N_INDEXER_HEAD_DIM, - index_width, - pos0, - n_tokens); - } - if (ok) { - g->layer_n_index_comp[il] = n_comp; - for (uint32_t t = 0; t < n_tokens; t++) { - index_counts[t] = (pos0 + t + 1u) / ratio; - } - if (n_comp != 0) { - metal_graph_debug_dump_tensor("indexer_KVcompress", - g->layer_index_comp_cache[il], - (uint64_t)n_comp * DS4_N_INDEXER_HEAD_DIM, - il, - pos0); - } - metal_graph_debug_dump_tensor("indexer_state_kv", - g->layer_index_state_kv[il], - (uint64_t)index_width * coff * ratio, - il, - pos0); - metal_graph_debug_dump_tensor("indexer_state_score", - g->layer_index_state_score[il], - (uint64_t)index_width * coff * ratio, - il, - pos0); - } - } else { - const bool aligned_chunk = - getenv("DS4_CUDA_NO_COMPRESSOR_PREFILL_BATCH") == NULL && - (pos0 % ratio) == 0u && (n_tokens % ratio) == 0u; - if (aligned_chunk) { - const uint32_t index_before = g->layer_n_index_comp[il]; - const uint32_t index_chunk = n_tokens / ratio; - if (index_before + index_chunk > g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - } - ds4_gpu_tensor *index_view = NULL; - if (ok) { - index_view = ds4_gpu_tensor_view( - g->layer_index_comp_cache[il], - (uint64_t)index_before * DS4_N_INDEXER_HEAD_DIM * sizeof(float), - (uint64_t)index_chunk * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - ok = index_view != NULL; - } - if (ok) { - ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( - index_view, - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - metal_graph_batch_comp_kv(g), - metal_graph_batch_comp_sc(g), - model->map, - model->size, - layer->indexer_compressor_ape->abs_offset, - layer->indexer_compressor_ape->type, - layer->indexer_compressor_norm->abs_offset, - layer->indexer_compressor_norm->type, - DS4_N_INDEXER_HEAD_DIM, - pos0, - n_tokens, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; - } - if (ok && index_chunk != 0) { - ok = ds4_gpu_dsv4_indexer_qat_tensor(index_view, - index_chunk, - DS4_N_INDEXER_HEAD_DIM) != 0; - } - if (ok) { - ok = metal_graph_refresh_ratio4_compressor_state(g, - model, - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - layer->indexer_compressor_kv, - layer->indexer_compressor_gate, - layer->indexer_compressor_ape, - DS4_N_INDEXER_HEAD_DIM, - index_width, - pos0, - n_tokens); - } - if (ok) { - g->layer_n_index_comp[il] = index_before + index_chunk; - if (index_counts) { - for (uint32_t t = 0; t < n_tokens; t++) { - index_counts[t] = (pos0 + t + 1u) / ratio; - } - } - metal_graph_debug_dump_tensor("indexer_KVcompress", - index_view, - (uint64_t)index_chunk * DS4_N_INDEXER_HEAD_DIM, - il, - pos0); - metal_graph_debug_dump_tensor("indexer_state_kv", - g->layer_index_state_kv[il], - (uint64_t)index_width * coff * ratio, - il, - pos0); - metal_graph_debug_dump_tensor("indexer_state_score", - g->layer_index_state_score[il], - (uint64_t)index_width * coff * ratio, - il, - pos0); - } - ds4_gpu_tensor_free(index_view); - } else { - for (uint32_t t = 0; ok && t < n_tokens; t++) { - const uint32_t pos = pos0 + t; - const bool emit = ((pos + 1u) % ratio) == 0u; - if (emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { - fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); - ok = false; - break; - } - ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(metal_graph_batch_comp_kv(g), t, index_width); - ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(metal_graph_batch_comp_sc(g), t, index_width); - const uint32_t index_row = g->layer_n_index_comp[il]; - ok = kv_view && sc_view && - ds4_gpu_compressor_update_tensor(kv_view, - sc_view, - g->layer_index_state_kv[il], - g->layer_index_state_score[il], - g->layer_index_comp_cache[il], - model->map, - model->size, - layer->indexer_compressor_ape->abs_offset, - layer->indexer_compressor_ape->type, - layer->indexer_compressor_norm->abs_offset, - layer->indexer_compressor_norm->type, - DS4_N_INDEXER_HEAD_DIM, - ratio, - pos, - index_row, - DS4_N_ROT, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS, - false) != 0; - if (ok && emit) { - ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( - g->layer_index_comp_cache[il], - (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), - (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); - if (!index_row_view) { - ok = false; - } else { - ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, - 1, - DS4_N_INDEXER_HEAD_DIM) != 0; - ds4_gpu_tensor_free(index_row_view); - } - } - if (ok && emit) g->layer_n_index_comp[il]++; - if (index_counts) index_counts[t] = g->layer_n_index_comp[il]; - if (ok && t == 0) ok = metal_graph_capture_prefix1_index_state(g, il); - ds4_gpu_tensor_free(sc_view); - ds4_gpu_tensor_free(kv_view); - } - } - } - } - if (ratio == 4) DS4_METAL_PROFILE_ATTN_STAGE("indexer_setup"); - - if (ok && !zero_prefix && n_tokens <= g->raw_cap) { - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos0, n_tokens); - /* See the raw-only branch above: batched mixed attention also - * consumes a logical raw window, linearized out of the ring. */ - const uint32_t raw_start = metal_graph_raw_start_for_span(g, - pos0 + n_tokens - 1u, - n_raw); - uint32_t use_comp_mask = 0; - bool use_indexed_comp = false; - double index_stage_t0 = 0.0; - - ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - metal_graph_batch_kv(g), - g->raw_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; - if (ok && ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K) { - const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - if (index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary(NULL, - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - ok = ds4_gpu_indexer_scores_decode_batch_tensor(metal_graph_indexer_scores(g), - metal_graph_batch_indexer_q(g), - metal_graph_batch_indexer_weights(g), - g->layer_index_comp_cache[il], - n_comp, - n_tokens, - pos0, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - ratio, - index_scale) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("score", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - if (ok) { - metal_graph_debug_dump_tensor("indexer_scores", - metal_graph_indexer_scores(g), - (uint64_t)n_comp * n_tokens, - il, - pos0); - } - if (ok) { - ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - metal_graph_indexer_scores(g), - n_comp, - n_tokens, - DS4_N_INDEXER_TOP_K) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("topk", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - if (ok) { - metal_graph_debug_dump_i32_tensor("indexer_topk", - metal_graph_comp_selected(g), - (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, - il, - pos0); - } - } - if (ok) { - use_indexed_comp = true; - } - use_comp_mask = 1; - } - if (ok) { - if (use_indexed_comp) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->layer_raw_cache[il], - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - metal_graph_comp_selected(g), - n_tokens, - pos0, - n_raw, - g->raw_cap, - raw_start, - n_comp, - DS4_N_INDEXER_TOP_K, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("attention", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - } else { - ok = ds4_gpu_attention_decode_mixed_batch_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->layer_raw_cache[il], - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - use_comp_mask ? metal_graph_comp_mask(g) : NULL, - use_comp_mask, - n_tokens, - pos0, - n_raw, - g->raw_cap, - raw_start, - n_comp, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - } - if (ok) batch_attention_done = true; - } - - const bool topk_prefill_needed = ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K; - if (ok && zero_prefix && topk_prefill_needed && n_comp != 0) { - const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - double index_stage_t0 = 0.0; - if (index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary(NULL, - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - ok = ds4_gpu_indexer_scores_prefill_tensor(metal_graph_indexer_scores(g), - metal_graph_batch_indexer_q(g), - metal_graph_batch_indexer_weights(g), - g->layer_index_comp_cache[il], - n_comp, - n_tokens, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - ratio, - index_scale) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("score", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - if (ok) { - metal_graph_debug_dump_tensor("indexer_scores", - metal_graph_indexer_scores(g), - (uint64_t)n_comp * n_tokens, - il, - pos0); - } - if (ok) { - ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - metal_graph_indexer_scores(g), - n_comp, - n_tokens, - DS4_N_INDEXER_TOP_K) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("topk", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - if (ok) { - metal_graph_debug_dump_i32_tensor("indexer_topk", - metal_graph_comp_selected(g), - (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, - il, - pos0); - } - } - if (ok && tp_row_split_attn) { - /* Score/top-k selection above ran replicated over all rows; - * only the attention consumption splits. Passing the row - * offset through pos0 and clamping n_raw to the rows this - * rank can see keeps the kernel's first_raw_pos at the - * chunk origin, so the raw ring mapping is unchanged. */ - ds4_gpu_tensor *tp_topk = metal_graph_tensor_row_range_view( - metal_graph_comp_selected(g), tp_row0, tp_rows, DS4_N_INDEXER_TOP_K); - ok = tp_topk && - ds4_gpu_attention_indexed_mixed_batch_heads_tensor(tp_heads, - model->map, - model->size, - layer->attn_sinks->abs_offset, - tp_q, - g->layer_raw_cache[il], - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - tp_topk, - tp_rows, - pos0 + tp_row0, - tp_row0 + tp_rows, - g->raw_cap, - 0, - n_comp, - DS4_N_INDEXER_TOP_K, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - ds4_gpu_tensor_free(tp_topk); - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("attention", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - } else if (ok) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->layer_raw_cache[il], - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - metal_graph_comp_selected(g), - n_tokens, - pos0, - n_tokens, - g->raw_cap, - 0, - n_comp, - DS4_N_INDEXER_TOP_K, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - if (ok && index_stage_profile) { - ok = metal_graph_indexer_stage_profile_boundary("attention", - il, - pos0, - n_tokens, - n_comp, - &index_stage_t0); - } - } - if (ok) batch_attention_done = true; - } - if (ok && zero_prefix && !topk_prefill_needed && n_comp != 0) { - if (tp_row_split_attn) { - ok = ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(tp_heads, - model->map, - model->size, - layer->attn_sinks->abs_offset, - tp_q, - metal_graph_batch_kv(g), - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - tp_row0, - tp_rows, - n_tokens, - n_comp, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } else { - ok = ds4_gpu_attention_prefill_static_mixed_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - metal_graph_batch_kv(g), - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - n_tokens, - n_comp, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - if (ok) batch_attention_done = true; - } - } - - if (ok && !raw_batch_attention && !batch_attention_done) { - uint32_t raw_prefix_tokens = 0; - if (zero_prefix && ratio != 0 && n_tokens <= g->raw_cap && comp_counts != NULL) { - while (raw_prefix_tokens < n_tokens && comp_counts[raw_prefix_tokens] == 0u) { - raw_prefix_tokens++; - } - } - - if (raw_prefix_tokens != 0) { - if (tp_row_split_attn && raw_prefix_tokens == n_tokens) { - /* tp_attn_full_raw guarantees the whole chunk stays raw - * (n_tokens < ratio), so the split covers every row. */ - ok = ds4_gpu_attention_prefill_raw_heads_range_tensor(tp_heads, - model->map, - model->size, - layer->attn_sinks->abs_offset, - tp_q, - metal_graph_batch_kv(g), - tp_row0, - tp_rows, - n_tokens, - g->raw_window, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } else { - ok = ds4_gpu_attention_prefill_raw_heads_tensor(metal_graph_batch_heads(g), - model->map, - model->size, - layer->attn_sinks->abs_offset, - metal_graph_batch_q(g), - metal_graph_batch_kv(g), - raw_prefix_tokens, - g->raw_window, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - } - if (raw_prefix_tokens < n_tokens) { - for (uint32_t t = raw_prefix_tokens; ok && t < n_tokens; t++) { - const uint32_t pos = pos0 + t; - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); - const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); - const uint32_t cur_comp = comp_counts ? comp_counts[t] : 0u; - const uint32_t cur_index = index_counts ? index_counts[t] : 0u; - uint32_t n_selected = 0; - ds4_gpu_tensor *comp_mask = NULL; - - if (ratio == 4 && cur_comp > DS4_N_INDEXER_TOP_K) { - const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - ds4_gpu_tensor *indexer_q_view = metal_graph_tensor_row_view( - metal_graph_batch_indexer_q(g), t, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); - ds4_gpu_tensor *indexer_w_view = metal_graph_tensor_row_view( - metal_graph_batch_indexer_weights(g), t, DS4_N_INDEXER_HEAD); - ok = indexer_q_view && indexer_w_view && - ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), - indexer_q_view, - indexer_w_view, - g->layer_index_comp_cache[il], - cur_index, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - index_scale) != 0 && - ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - metal_graph_indexer_scores(g), - cur_index, - 1, - DS4_N_INDEXER_TOP_K) != 0 && - ds4_gpu_dsv4_topk_mask_tensor(metal_graph_comp_mask(g), - metal_graph_comp_selected(g), - cur_index, - 1, - DS4_N_INDEXER_TOP_K) != 0; - ds4_gpu_tensor_free(indexer_w_view); - ds4_gpu_tensor_free(indexer_q_view); - if (ok) { - comp_mask = metal_graph_comp_mask(g); - n_selected = DS4_N_INDEXER_TOP_K < cur_index - ? DS4_N_INDEXER_TOP_K - : cur_index; - } - } - - ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(metal_graph_batch_q(g), t, q_dim); - ds4_gpu_tensor *kv_cache_view = metal_graph_tensor_row_view(metal_graph_batch_kv(g), t, DS4_N_HEAD_DIM); - ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(metal_graph_batch_heads(g), t, q_dim); - ok = ok && q_view && kv_cache_view && heads_view; - if (ok && !zero_prefix) { - ok = ds4_gpu_store_raw_kv_tensor(g->layer_raw_cache[il], - kv_cache_view, - g->raw_cap, - pos % g->raw_cap, - DS4_N_HEAD_DIM) != 0; - } - if (ok && comp_mask != NULL && n_selected != 0) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(heads_view, - model->map, - model->size, - layer->attn_sinks->abs_offset, - q_view, - g->layer_raw_cache[il], - g->layer_attn_comp_cache[il], - metal_graph_attn_comp_cache_is_f16(), - metal_graph_comp_selected(g), - 1, - pos, - n_raw, - g->raw_cap, - raw_start, - cur_comp, - n_selected, - g->raw_window, - ratio, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } else if (ok) { - ok = ds4_gpu_attention_decode_heads_tensor(heads_view, - model->map, - model->size, - layer->attn_sinks->abs_offset, - q_view, - g->layer_raw_cache[il], - n_raw, - g->raw_cap, - raw_start, - cur_comp ? g->layer_attn_comp_cache[il] : NULL, - metal_graph_attn_comp_cache_is_f16(), - cur_comp, - comp_mask, - n_selected, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - ds4_gpu_tensor_free(heads_view); - ds4_gpu_tensor_free(kv_cache_view); - ds4_gpu_tensor_free(q_view); - } - } - } - DS4_METAL_PROFILE_ATTN_STAGE("attention"); - - if (ok) { - metal_graph_debug_dump_tensor("kqv_out", metal_graph_batch_heads(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - if (ok) ok = ds4_gpu_rope_tail_tensor(tp_heads ? tp_heads : metal_graph_batch_heads(g), - tp_rows, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0 + tp_row0, - compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) { - metal_graph_debug_dump_tensor("kqv_back", metal_graph_batch_heads(g), - (uint64_t)n_tokens * q_dim, il, pos0); - } - DS4_METAL_PROFILE_ATTN_STAGE("inv_rope"); - const bool attn_out_debug = - metal_graph_debug_wants("attn_low", il, pos0) || - metal_graph_debug_wants("attn_out", il, pos0); - bool attn_out_f16 = false; - if (ok && - !attn_out_debug && - !tp_row_split_attn && - layer->attn_output_a->type == DS4_TENSOR_Q8_0 && - layer->attn_output_b->type == DS4_TENSOR_Q8_0 && - !metal_graph_directional_steering_attn_enabled(g)) { - attn_out_f16 = ds4_gpu_attention_output_q8_batch_f16_tensor(g->batch_q_half, - metal_graph_batch_attn_low(g), - model->map, - model->size, - layer->attn_output_a->abs_offset, - layer->attn_output_b->abs_offset, - group_dim, - rank, - n_groups, - DS4_N_EMBD, - metal_graph_batch_heads(g), - n_tokens) != 0; - } - uint64_t tp_attn_gate_seq = 0; - /* Opt-in sub-chunk gate pipelining (see metal_graph_tp_subgate_pipeline; - * measured net-negative on the M5 Max pair, kept for slower wires). - * Kernel-path constraint: the output projection picks its kernel by row - * count (direct low below 32 rows, the 64-token-tile TensorOps path at - * multiples of 64, ids-cache fallback otherwise), and the paths are not - * bit-identical per row. Parity with the single-node reference - * therefore requires every sub-call to land on the same path as the - * full-chunk call: n_tokens % 256 == 0 puts the chunk, the rank halves, - * and the quarter sub-calls all on the TensorOps path. Other sizes - * keep the proven single-gate swap. */ - const bool tp_attn_pipeline = - tp_row_split_attn && (n_tokens % 256u) == 0u && - metal_graph_tp_subgate_pipeline(); - if (!attn_out_f16) { - if (ok && tp_attn_pipeline) { - /* Sub-chunk pipelined swap: the output projection runs in two - * sub-halves of this rank's rows and each sub-half's row swap - * is kicked as soon as its rows land in batch_attn_out, so the - * first wire exchange overlaps the second sub-half's compute. - * The two kicks use opposite flag-slot parities; the wait below - * (before the HC post expand) covers both. */ - const uint32_t tp_sub1 = (tp_half_rows + 1u) / 2u; - const uint32_t tp_c1 = tp_rows < tp_sub1 ? tp_rows : tp_sub1; - const uint32_t tp_own_base = g->tp_rank == 0 ? 0u : tp_half_rows; - const uint32_t tp_peer_base = g->tp_rank == 0 ? tp_half_rows : 0u; - for (uint32_t sub = 0; ok && sub < 2u; sub++) { - const uint32_t coff = sub == 0 ? 0u : tp_c1; - const uint32_t crows = sub == 0 ? tp_c1 : tp_rows - tp_c1; - const uint32_t soff = sub == 0 ? 0u : tp_sub1; - const uint32_t srows = sub == 0 ? tp_sub1 : tp_half_rows - tp_sub1; - if (crows != 0) { - ds4_gpu_tensor *sub_heads = metal_graph_tensor_row_range_view( - metal_graph_batch_heads(g), tp_row0 + coff, crows, q_dim); - ds4_gpu_tensor *sub_out = metal_graph_tensor_row_range_view( - metal_graph_batch_attn_out(g), tp_row0 + coff, crows, DS4_N_EMBD); - ok = sub_heads && sub_out && - metal_graph_attention_output_dense_quant_batch(sub_out, - metal_graph_batch_attn_low(g), - g, - model, - layer->attn_output_a, - layer->attn_output_b, - group_dim, - rank, - n_groups, - DS4_N_EMBD, - sub_heads, - crows); - ds4_gpu_tensor_free(sub_out); - ds4_gpu_tensor_free(sub_heads); - } - if (ok && srows != 0) { - ds4_gpu_tensor *send_sub = metal_graph_tensor_row_range_view( - metal_graph_batch_attn_out(g), tp_own_base + soff, srows, DS4_N_EMBD); - ds4_gpu_tensor *recv_sub = metal_graph_tensor_row_range_view( - metal_graph_batch_attn_out(g), tp_peer_base + soff, srows, DS4_N_EMBD); - uint64_t seq = 0; - if (send_sub && recv_sub) { - seq = ds4_gpu_tp_big_gate_kick(il, n_tokens, - send_sub, recv_sub, - (uint64_t)srows * DS4_N_EMBD * sizeof(float)); - } - ok = seq != 0; - if (ok) tp_attn_gate_seq = seq; - ds4_gpu_tensor_free(recv_sub); - ds4_gpu_tensor_free(send_sub); - } - } - } else if (ok) { - ok = metal_graph_attention_output_dense_quant_batch(tp_attn_out ? tp_attn_out : metal_graph_batch_attn_out(g), - metal_graph_batch_attn_low(g), - g, - model, - layer->attn_output_a, - layer->attn_output_b, - group_dim, - rank, - n_groups, - DS4_N_EMBD, - tp_heads ? tp_heads : metal_graph_batch_heads(g), - tp_rows); - } - if (ok) { - metal_graph_debug_dump_tensor("attn_low", metal_graph_batch_attn_low(g), - (uint64_t)n_tokens * n_groups * rank, - il, - pos0); - } - if (ok) { - metal_graph_debug_dump_tensor("attn_out", metal_graph_batch_attn_out(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - } - DS4_METAL_PROFILE_ATTN_STAGE("output_proj"); - if (ok && tp_row_split_attn) { - /* Release point for the pipelined row swaps of batch_attn_out: both - * ranks reach the HC post expand with identical full tensors. */ - if (tp_attn_pipeline) { - ok = tp_attn_gate_seq != 0 && - ds4_gpu_tp_big_gate_wait(tp_attn_gate_seq) != 0; - } else { - const uint64_t half_bytes = - (uint64_t)tp_half_rows * DS4_N_EMBD * sizeof(float); - ds4_gpu_tensor *send_half = metal_graph_tensor_row_range_view( - metal_graph_batch_attn_out(g), - g->tp_rank == 0 ? 0 : tp_half_rows, tp_half_rows, DS4_N_EMBD); - ds4_gpu_tensor *recv_half = metal_graph_tensor_row_range_view( - metal_graph_batch_attn_out(g), - g->tp_rank == 0 ? tp_half_rows : 0, tp_half_rows, DS4_N_EMBD); - ok = send_half && recv_half && - ds4_gpu_tp_big_gate_encode(il, n_tokens, - send_half, recv_half, - half_bytes) != 0; - ds4_gpu_tensor_free(send_half); - ds4_gpu_tensor_free(recv_half); - } - if (!ok) fprintf(stderr, "ds4: TP prefill attention row gate failed (layer %u)\n", il); - } - if (ok && !attn_out_f16 && metal_graph_directional_steering_attn_enabled(g)) { - ok = metal_graph_apply_directional_steering_attn(g, metal_graph_batch_attn_out(g), il, n_tokens); - } - if (ok && attn_out_f16) { - ok = ds4_gpu_hc_expand_split_half_tensor(after_attn_hc_view, - g->batch_q_half, - metal_graph_batch_cur_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (ok) { - ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, - metal_graph_batch_attn_out(g), - metal_graph_batch_cur_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("hc_attn_post", metal_graph_batch_after_attn_hc(g), - (uint64_t)n_tokens * hc_dim, il, pos0); - } - DS4_METAL_PROFILE_ATTN_STAGE("hc_post"); - ds4_gpu_tensor_free(tp_attn_out); - ds4_gpu_tensor_free(tp_heads); - ds4_gpu_tensor_free(tp_qr_norm); - ds4_gpu_tensor_free(tp_q_half); - ds4_gpu_tensor_free(tp_q); - ds4_gpu_tensor_free(after_attn_hc_view); - ds4_gpu_tensor_free(attn_cur_view); - ds4_gpu_tensor_free(hc_split_view); - ds4_gpu_tensor_free(hc_mix_view); - if (index_counts != index_counts_stack) free(index_counts); - if (comp_counts != comp_counts_stack) free(comp_counts); -#undef DS4_METAL_PROFILE_ATTN_STAGE -#undef DS4_METAL_PROFILE_Q_STAGE - return ok; -} - -static bool metal_graph_encode_mixed_routed_rows( - ds4_gpu_graph *g, - ds4_decode_item *decode_items, - int decode_count, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t prefill_rows); - -/* Encode the batched prefill FFN half: HC pre/norm, shared expert, routed - * experts, sum, and HC post. A non-empty decode tail has already been - * prepared in rows [n_tokens, n_tokens + decode_count); only the routed - * expert dispatch is shared between the two arithmetic paths. */ -static bool metal_graph_encode_layer_ffn_batch( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - ds4_decode_item *decode_items, - int decode_count) { - if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; - if (decode_count < 0 || - (decode_count > 0 && - (!decode_items || (uint64_t)n_tokens + (uint32_t)decode_count > g->prefill_cap))) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; - const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; - const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; - const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; - const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - const uint64_t down_expert_bytes = routed_out_dim * down_row_bytes; - const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); - double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; -#define DS4_METAL_PROFILE_FFN_STAGE(name) do { \ - if (ok && layer_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary("ffn", (name), il, pos0, n_tokens, &layer_stage_t0); \ - } \ - } while (0) - - ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( - metal_graph_batch_hc_mix(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( - metal_graph_batch_hc_split(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *ffn_cur_view = ds4_gpu_tensor_view( - metal_graph_batch_ffn_cur(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view( - metal_graph_batch_next_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); - bool ok = hc_mix_view && hc_split_view && ffn_cur_view && next_hc_view; - const bool fuse_hc_norm = n_tokens > 1 && - DS4_N_HC == 4 && - !metal_graph_use_reference_hc_decode() && - metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, - metal_graph_batch_flat_hc(g), - model, - layer->hc_ffn_fn, - metal_graph_batch_after_attn_hc(g), - hc_dim, - n_tokens); - if (metal_graph_use_reference_hc_decode()) { - if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, - hc_mix_view, - model->map, - model->size, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(ffn_cur_view, - metal_graph_batch_after_attn_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } else if (fuse_hc_norm) { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(ffn_cur_view, - metal_graph_batch_ffn_norm(g), - hc_split_view, - hc_mix_view, - metal_graph_batch_after_attn_hc(g), - model->map, - model->size, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset, - layer->ffn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - } else { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(ffn_cur_view, - hc_split_view, - hc_mix_view, - metal_graph_batch_after_attn_hc(g), - model->map, - model->size, - layer->hc_ffn_scale->abs_offset, - layer->hc_ffn_base->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre", metal_graph_batch_ffn_cur(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - DS4_METAL_PROFILE_FFN_STAGE("hc_pre"); - if (ok && !fuse_hc_norm) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_ffn_norm(g), - metal_graph_batch_ffn_cur(g), - model->map, - model->size, - layer->ffn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_norm", metal_graph_batch_ffn_norm(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - DS4_METAL_PROFILE_FFN_STAGE("norm"); - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_router_logits(g), - model, - layer->ffn_gate_inp, - DS4_N_EMBD, - DS4_N_EXPERT, - metal_graph_batch_ffn_norm(g), - n_tokens); - - ds4_gpu_tensor *router_tokens = NULL; - if (ok) { - router_tokens = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), - (uint64_t)g->batch_token_offset * sizeof(int32_t), - (uint64_t)n_tokens * sizeof(int32_t)); - ok = router_tokens != NULL; - } - if (ok) ok = ds4_gpu_router_select_batch_tensor(metal_graph_batch_router_selected(g), - metal_graph_batch_router_weights(g), - metal_graph_batch_router_probs(g), - model->map, - model->size, - layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, - layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, - layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, - 0, - 0, - layer->ffn_exp_probs_b != NULL, - layer->ffn_gate_tid2eid != NULL, - metal_graph_batch_router_logits(g), - metal_graph_prefill_tokens(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE, - n_tokens) != 0; - ds4_gpu_tensor_free(router_tokens); - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_logits", metal_graph_batch_router_logits(g), - (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_probs", metal_graph_batch_router_probs(g), - (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); - metal_graph_debug_dump_i32_tensor("ffn_moe_topk", metal_graph_batch_router_selected(g), - (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", metal_graph_batch_router_weights(g), - (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); - } - DS4_METAL_PROFILE_FFN_STAGE("router"); - - if (ok) { - ok = metal_graph_cuda_stream_prefill_batch_selected_load(g, - model, - layer, - il, - n_tokens, - gate_expert_bytes, - down_expert_bytes); - } - -#ifdef DS4_ROCM_BUILD - rocm_graph_batch_selected_async_load rocm_batch_selected_async = {0}; - bool rocm_batch_selected_async_started = false; - const bool rocm_batch_selected_shared_overlap = - ok && - g->ssd_streaming && - !g->quality && - n_tokens > 1 && - DS4_N_EXPERT_USED == 6 && - !rocm_graph_stream_prefill_full_layer_enabled(g, layer, il, n_tokens) && - layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && - layer->ffn_down_exps->type == DS4_TENSOR_Q2_K; - if (rocm_batch_selected_shared_overlap) { - uint64_t selected_event = 0; - if (ds4_gpu_signal_selected_readback_ready(&selected_event) == 0) { - ok = false; - } else { - ok = rocm_graph_batch_selected_async_load_start( - &rocm_batch_selected_async, - metal_graph_batch_router_selected(g), - model, - layer, - il, - n_tokens, - selected_event, - gate_expert_bytes, - down_expert_bytes); - rocm_batch_selected_async_started = ok; - } - } -#endif - - const bool selected_readahead_shared = - metal_graph_stream_prefill_selected_readahead_shared_enabled(g) -#ifdef DS4_ROCM_BUILD - && !rocm_batch_selected_async_started -#endif - ; - if (ok && - metal_graph_stream_prefill_selected_readahead_enabled(g) && -#ifdef DS4_ROCM_BUILD - !rocm_batch_selected_async_started && -#endif - !selected_readahead_shared) { - if (ds4_gpu_end_commands() == 0) { - ok = false; - } else { - ok = metal_graph_stream_readahead_selected_experts_from_gpu(g, - model, - layer, - il, - n_tokens, - gate_expert_bytes, - down_expert_bytes) && - ds4_gpu_begin_commands() != 0; - } - } - - const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos0); - bool shared_down_f16 = false; - -#define DS4_METAL_TRY_SHARED_DOWN_F16() do { \ - if (ok && !tp_row_split_ffn && !keep_ffn_out && \ - !metal_graph_debug_wants("ffn_shexp", il, pos0)) { \ - shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->batch_q_half, \ - model->map, \ - model->size, \ - layer->ffn_down_shexp->abs_offset, \ - shared_dim, \ - DS4_N_EMBD, \ - metal_graph_batch_shared_mid(g), \ - n_tokens) != 0; \ - } \ - } while (0) - -#define DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT() do { \ - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_gate", \ - il, \ - pos0, \ - metal_graph_batch_shared_gate(g), \ - model, \ - layer->ffn_gate_shexp, \ - DS4_N_EMBD, \ - shared_dim, \ - tp_ffn_x ? tp_ffn_x : metal_graph_batch_ffn_norm(g), \ - tp_rows); \ - if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_up", \ - il, \ - pos0, \ - metal_graph_batch_shared_up(g), \ - model, \ - layer->ffn_up_shexp, \ - DS4_N_EMBD, \ - shared_dim, \ - tp_ffn_x ? tp_ffn_x : metal_graph_batch_ffn_norm(g), \ - tp_rows); \ - DS4_METAL_PROFILE_FFN_STAGE("shared_gate_up"); \ - if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_batch_shared_mid(g), \ - metal_graph_batch_shared_gate(g), \ - metal_graph_batch_shared_up(g), \ - (uint32_t)((uint64_t)tp_rows * shared_dim), \ - DS4_SWIGLU_CLAMP_EXP, \ - 1.0f) != 0; \ - DS4_METAL_TRY_SHARED_DOWN_F16(); \ - if (ok && !shared_down_f16) ok = metal_graph_matmul_q8_0_named_tensor("shared_down", \ - il, \ - pos0, \ - metal_graph_batch_shared_out(g), \ - model, \ - layer->ffn_down_shexp, \ - shared_dim, \ - DS4_N_EMBD, \ - metal_graph_batch_shared_mid(g), \ - tp_rows); \ - DS4_METAL_PROFILE_FFN_STAGE("shared_down"); \ - if (ok && !shared_down_f16) { \ - metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_batch_shared_out(g), \ - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); \ - } \ - } while (0) - - bool shared_done = false; - /* With 50/50 expert residency, every rank evaluates every prompt row - * against its local expert half. For large chunks the replicated shared - * expert remains row-split; its rows are folded into the local routed - * partial before the one all-reduce-style bulk exchange. */ - const bool tp_split_ffn = g->tp_world == 2; - const bool tp_row_split_ffn = - tp_split_ffn && g->tp_batch_rows != n_tokens && !keep_ffn_out && - !metal_graph_directional_steering_ffn_enabled(g) && - n_tokens >= metal_graph_tp_prefill_split_min(); - const uint32_t tp_half_rows = (n_tokens + 1u) / 2u; - const uint32_t tp_row0 = (tp_row_split_ffn && g->tp_rank != 0) ? tp_half_rows : 0; - const uint32_t tp_rows = tp_row_split_ffn ? - (g->tp_rank == 0 ? tp_half_rows : n_tokens - tp_half_rows) : n_tokens; - ds4_gpu_tensor *tp_ffn_x = tp_row_split_ffn ? - metal_graph_tensor_row_range_view(metal_graph_batch_ffn_norm(g), tp_row0, tp_rows, - DS4_N_EMBD) : NULL; - if (tp_row_split_ffn && !tp_ffn_x) ok = false; - if (ok && selected_readahead_shared) { - if (ds4_gpu_end_commands() == 0) { - ok = false; - } else { - ok = metal_graph_stream_readahead_selected_experts_from_gpu(g, - model, - layer, - il, - n_tokens, - gate_expert_bytes, - down_expert_bytes) && - ds4_gpu_begin_commands() != 0; - } - if (ok) { - DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); - shared_done = ok; - } - if (ok) { - if (ds4_gpu_end_commands() == 0) { - ok = false; - } else { - ok = ds4_gpu_begin_commands() != 0; - } - } - } - - if (ok && - !shared_done && - (metal_graph_stream_prefill_selected_pagein_enabled(g) || - metal_graph_stream_prefill_selected_madvise_enabled(g))) { - metal_graph_stream_pagein_job pagein_job; - memset(&pagein_job, 0, sizeof(pagein_job)); - bool pagein_commands_open = false; - if (ds4_gpu_end_commands() == 0) { - ok = false; - } else { - ok = metal_graph_stream_prefill_selected_pagein_start(g, - model, - layer, - il, - n_tokens, - gate_expert_bytes, - down_expert_bytes, - &pagein_job); - } - if (ok) { - if (ds4_gpu_begin_commands() == 0) { - ok = false; - } else { - pagein_commands_open = true; - } - } - if (ok) { - DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); - shared_done = ok; - } - if (pagein_commands_open) { - if (ds4_gpu_end_commands() == 0) ok = false; - } - if (!metal_graph_stream_prefill_selected_pagein_join(&pagein_job)) { - ok = false; - } - if (ok) ok = ds4_gpu_begin_commands() != 0; - } - -#ifdef DS4_ROCM_BUILD - if (rocm_batch_selected_async_started) { - if (ok && !shared_done) { - DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); - shared_done = ok; - } - const bool finish_ok = - rocm_graph_batch_selected_async_load_finish(&rocm_batch_selected_async); - ok = ok && finish_ok; - } -#endif - - const bool tp_split_batch_moe = - g->tp_batch_rows == n_tokens && n_tokens > 0 && - g->tp_world == 2 && - g->tp_batch_out && g->tp_batch_in; - const bool cuda_tp_owned_batch_moe = - g->cuda_tp_ep && g->cuda_tp_prefill_ffn; - if (ok && cuda_tp_owned_batch_moe) { - ok = metal_graph_encode_mixed_routed_rows( - g, decode_items, decode_count, model, layer, il, n_tokens); - } else if (ok && tp_split_batch_moe) { - /* Verify-block expert split: run the contiguous-half split - * single-token routed kernels per row into the slab batch-out - * rows, exchange all rows with one gate, then materialize the - * combined routed output. The add is commutative, so both ranks - * compute bit-identical sums and stay in lockstep. */ - const uint64_t vec_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - for (uint32_t r = 0; ok && r < n_tokens; r++) { - ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( - g->tp_batch_out[il], (uint64_t)r * vec_bytes, vec_bytes); - ds4_gpu_tensor *x_row = ds4_gpu_tensor_view( - metal_graph_batch_ffn_norm(g), (uint64_t)r * vec_bytes, vec_bytes); - ds4_gpu_tensor *sel_row = ds4_gpu_tensor_view( - metal_graph_batch_router_selected(g), - (uint64_t)r * DS4_N_EXPERT_USED * sizeof(int32_t), - (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); - ds4_gpu_tensor *w_row = ds4_gpu_tensor_view( - metal_graph_batch_router_weights(g), - (uint64_t)r * DS4_N_EXPERT_USED * sizeof(float), - (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); - ok = out_row && x_row && sel_row && w_row && - ds4_gpu_routed_moe_one_tensor(out_row, - metal_graph_routed_gate(g), - metal_graph_routed_up(g), - metal_graph_routed_mid(g), - metal_graph_routed_down(g), - model->map, model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - sel_row, w_row, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_SWIGLU_CLAMP_EXP, - x_row, - NULL, - il, - false) != 0; - ds4_gpu_tensor_free(w_row); - ds4_gpu_tensor_free(sel_row); - ds4_gpu_tensor_free(x_row); - ds4_gpu_tensor_free(out_row); - } - if (ok) ok = ds4_gpu_tp_batch_gate_encode(il, n_tokens) != 0; - if (ok) { - ok = ds4_gpu_add_tensor(metal_graph_batch_routed_out(g), - g->tp_batch_out[il], - g->tp_batch_in[il], - (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; - } - } else if (ok) { - ok = ds4_gpu_routed_moe_batch_tensor(metal_graph_batch_routed_out(g), - metal_graph_batch_routed_gate(g), - metal_graph_batch_routed_up(g), - metal_graph_batch_routed_mid(g), - metal_graph_batch_routed_down(g), - model->map, - model->size, - layer->ffn_gate_exps->abs_offset, - layer->ffn_up_exps->abs_offset, - layer->ffn_down_exps->abs_offset, - layer->ffn_gate_exps->type, - layer->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - (uint32_t)expert_in_dim, - (uint32_t)down_in_dim, - (uint32_t)routed_out_dim, - metal_graph_batch_router_selected(g), - metal_graph_batch_router_weights(g), - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_SWIGLU_CLAMP_EXP, - metal_graph_batch_ffn_norm(g), - il, - n_tokens, - &g->batch_routed_mid_is_f16, - false) != 0; - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_batch_routed_gate(g), - (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_batch_routed_up(g), - (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); - } - if (ok) { - const uint64_t routed_mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim; - if (g->batch_routed_mid_is_f16) { - metal_graph_debug_dump_f16_tensor("ffn_moe_weighted_swiglu", metal_graph_batch_routed_mid(g), - routed_mid_elems, il, pos0); - } else { - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_batch_routed_mid(g), - routed_mid_elems, il, pos0); - } - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_batch_routed_down(g), - (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos0); - } - if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_batch_routed_out(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - DS4_METAL_PROFILE_FFN_STAGE("routed_moe"); - if (!shared_done) { - DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); - } -#undef DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT -#undef DS4_METAL_TRY_SHARED_DOWN_F16 - - if (ok && tp_row_split_ffn) { - /* Each shared-expert row must appear exactly once in the all-reduce. - * Fold this rank's shared rows into its full-row routed partial; the - * peer does the same for the complementary rows. */ - ds4_gpu_tensor *own_rows = - metal_graph_tensor_row_range_view(metal_graph_batch_routed_out(g), tp_row0, - tp_rows, DS4_N_EMBD); - ok = own_rows && - ds4_gpu_add_tensor(own_rows, own_rows, metal_graph_batch_shared_out(g), - (uint32_t)((uint64_t)tp_rows * DS4_N_EMBD)) != 0; - ds4_gpu_tensor_free(own_rows); - } - - if (ok && tp_split_ffn && !tp_split_batch_moe) { - /* All rows contain this rank's routed-expert partial. Exchange that - * matrix in one bulk gate, then add in canonical rank order. The - * batch verify path above already performed the equivalent slab gate. */ - const uint64_t bytes = - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); - ok = metal_graph_ensure_batch_ffn_out(g) && - ds4_gpu_tp_big_gate_encode(il, n_tokens, - metal_graph_batch_routed_out(g), - metal_graph_batch_ffn_out(g), - bytes) != 0; - if (ok) { - ds4_gpu_tensor *first = g->tp_rank == 0 ? - metal_graph_batch_routed_out(g) : metal_graph_batch_ffn_out(g); - ds4_gpu_tensor *second = g->tp_rank == 0 ? - metal_graph_batch_ffn_out(g) : metal_graph_batch_routed_out(g); - ok = ds4_gpu_add_tensor(metal_graph_batch_routed_out(g), first, second, - (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; - } - if (!ok) { - fprintf(stderr, "ds4: TP prefill FFN all-reduce failed (layer %u)\n", il); - } - } - - if (ok && keep_ffn_out) { - ok = metal_graph_ensure_batch_ffn_out(g) && - ds4_gpu_add_tensor(metal_graph_batch_ffn_out(g), - metal_graph_batch_shared_out(g), - metal_graph_batch_routed_out(g), - (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; - } - if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", metal_graph_batch_ffn_out(g), - (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_batch_ffn_out(g), il, n_tokens); - } - if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, - metal_graph_batch_ffn_out(g), - metal_graph_batch_after_attn_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - else if (ok && shared_down_f16) { - ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, - metal_graph_batch_routed_out(g), - g->batch_q_half, - metal_graph_batch_after_attn_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - else if (ok && tp_row_split_ffn) { - /* Shared expert already folded into the exchanged routed rows. */ - ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, - metal_graph_batch_routed_out(g), - metal_graph_batch_after_attn_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - else if (ok) { - ok = ds4_gpu_hc_expand_add_split_tensor(next_hc_view, - metal_graph_batch_routed_out(g), - metal_graph_batch_shared_out(g), - metal_graph_batch_after_attn_hc(g), - hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - DS4_METAL_PROFILE_FFN_STAGE("hc_post"); - if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_batch_next_hc(g), - (uint64_t)n_tokens * hc_dim, il, pos0); - } - DS4_METAL_PROFILE_FFN_STAGE("hc_post"); - ds4_gpu_tensor_free(tp_ffn_x); - ds4_gpu_tensor_free(next_hc_view); - ds4_gpu_tensor_free(ffn_cur_view); - ds4_gpu_tensor_free(hc_split_view); - ds4_gpu_tensor_free(hc_mix_view); -#undef DS4_METAL_PROFILE_FFN_STAGE - return ok; -} - -/* Encode one complete layer for prefill by chaining attention and FFN batches. */ -static bool metal_graph_encode_layer_batch( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens) { - if (g->placement) { - const int this_tier = g->placement[il + 1u]; - if (!metal_graph_set_active_tier_batch(g, this_tier, n_tokens)) { - return false; - } - } - bool ok = metal_graph_layer_stage_profile_start(il); - if (ok) { - ok = metal_graph_encode_layer_attention_batch(g, model, layer, il, pos0, n_tokens); - } - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u attention batch encode failed\n", il); - } - if (ok) { - ok = metal_graph_encode_layer_ffn_batch(g, model, layer, il, pos0, - n_tokens, NULL, 0); - if (!ok) { - fprintf(stderr, "ds4: gpu layer %u ffn batch encode failed\n", il); - } - } - if (ok) { - ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); - g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); - g->batch_next_hc_by_tier[g->active_tier] = tmp; - } - return ok; -} - -static bool metal_graph_eval_token_raw_swa_streaming( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - uint32_t pos, - float *logits) { - if (g->raw_cap == 0) { - fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); - return false; - } - - const bool profile = - glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", - "DS4_METAL_GRAPH_TOKEN_PROFILE"); - const bool throttle = graph_power_throttle_enabled(g); - const double t0 = (profile || throttle) ? now_sec() : 0.0; - const uint32_t raw_row = pos % g->raw_cap; - const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); - metal_graph_dspark_capture_begin(g); - - const bool static_decode_map = metal_graph_stream_decode_static_map_enabled(); - const bool static_map_state_cache = - static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); - const bool batch_static_decode = - static_decode_map && metal_graph_stream_decode_layer_batch_enabled(g); - bool ok = true; - if (static_decode_map) { - if (!static_map_state_cache || !g->streaming_static_decode_map_current) { - ok = metal_graph_stream_map_decode_static_all(model, weights); - if (ok) g->streaming_static_decode_map_current = static_map_state_cache; - } - } else { - g->streaming_static_decode_map_current = false; - ok = metal_graph_stream_map_token(model, weights); - } - if (ok && !static_decode_map && DS4_N_LAYER > 0) { - metal_graph_stream_readahead_layer_decode(model, weights, 0); - } - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (batch_static_decode) { - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - ok = metal_graph_encode_decode_layer(g, - model, - &weights->layer[il], - il, - pos, - g->layer_raw_cache[il], - g->raw_cap, - raw_row, - n_raw, - token); - if (ok) { - ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); - g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); - g->after_ffn_hc_by_tier[g->active_tier] = tmp; - ok = metal_graph_dspark_capture_decode_layer(g, il); - } - } - if (ok && logits) { - ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - } - const double t_encoded = (profile || throttle) ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = (profile || throttle) ? now_sec() : 0.0; - if (ok && logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - const double t_read = (profile || throttle) ? now_sec() : 0.0; - if (profile) { - fprintf(stderr, - "ds4: metal SSD streaming batched token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", - pos, - (t_encoded - t0) * 1000.0, - (t_done - t_encoded) * 1000.0, - (t_read - t_done) * 1000.0, - (t_read - t0) * 1000.0, - logits != NULL); - } - if (ok && throttle) { - graph_power_note_decode_token(g, t_read - t0); - } - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after batched SSD streaming graph eval failure also failed\n"); - } - } - return ok; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - - double encode_s = 0.0; - double execute_s = 0.0; - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - const double tl0 = profile ? now_sec() : 0.0; - if (!static_decode_map && !metal_graph_stream_map_layer_decode(model, weights, il)) { - ok = false; - break; - } - if (!static_decode_map && il + 1 < DS4_N_LAYER) { - metal_graph_stream_readahead_layer_decode(model, weights, il + 1); - } else if (!static_decode_map && logits) { - metal_graph_stream_readahead_output(model, weights); - } - if (ok) ok = ds4_gpu_begin_commands() != 0; - bool encoded_layer = false; - if (ok) { - ok = metal_graph_encode_decode_layer(g, - model, - &weights->layer[il], - il, - pos, - g->layer_raw_cache[il], - g->raw_cap, - raw_row, - n_raw, - token); - encoded_layer = true; - } - if (encoded_layer) { - ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); - g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); - g->after_ffn_hc_by_tier[g->active_tier] = tmp; - if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); - } - const double tl_encoded = profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double tl_done = profile ? now_sec() : 0.0; - if (profile) { - encode_s += tl_encoded - tl0; - execute_s += tl_done - tl_encoded; - } - } - - if (ok && logits && !static_decode_map) ok = metal_graph_stream_map_output(model, weights); - const double t_head0 = profile ? now_sec() : 0.0; - if (ok && logits) ok = ds4_gpu_begin_commands() != 0; - if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - const double t_head_encoded = profile ? now_sec() : 0.0; - if (ok && logits) ok = ds4_gpu_end_commands() != 0; - const double t_done = (profile || throttle) ? now_sec() : 0.0; - if (ok && logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - const double t_read = (profile || throttle) ? now_sec() : 0.0; - - if (profile) { - if (logits) { - encode_s += t_head_encoded - t_head0; - execute_s += t_done - t_head_encoded; - } - fprintf(stderr, - "ds4: metal SSD streaming token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", - pos, - encode_s * 1000.0, - execute_s * 1000.0, - (t_read - t_done) * 1000.0, - (t_read - t0) * 1000.0, - logits != NULL); - } - if (ok) graph_power_note_decode_token(g, t_read - t0); - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after SSD streaming graph eval failure also failed\n"); - } - } - return ok; -} - -/* Execute one Metal decode token and read back logits. */ -static bool metal_graph_eval_token_raw_swa( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - uint32_t pos, - float *logits) { - if (g && g->ssd_streaming) { - return metal_graph_eval_token_raw_swa_streaming(g, model, weights, token, pos, logits); - } - - const bool profile = - glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", - "DS4_METAL_GRAPH_TOKEN_PROFILE"); - const bool throttle = graph_power_throttle_enabled(g); - const double t0 = (profile || throttle) ? now_sec() : 0.0; - - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, logits != NULL, true); - const double t_encoded = (profile || throttle) ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = (profile || throttle) ? now_sec() : 0.0; - - if (ok && logits && g->tp_world == 2 && g->tp_logits_half) { - const uint64_t tp_vhalf = (uint64_t)DS4_N_VOCAB / 2u; - const uint64_t off = (uint64_t)g->tp_rank * tp_vhalf * sizeof(float); - ok = ds4_gpu_tensor_read(metal_graph_logits(g), off, logits + g->tp_rank * tp_vhalf, - tp_vhalf * sizeof(float)) != 0; - } else if (ok && logits && !(g->tp_world == 2 && g->tp_rank == 1)) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - const double t_read = (profile || throttle) ? now_sec() : 0.0; - if (profile) { - fprintf(stderr, - "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", - pos, - (t_encoded - t0) * 1000.0, - (t_done - t_encoded) * 1000.0, - (t_read - t_done) * 1000.0, - (t_read - t0) * 1000.0, - logits != NULL); - } - if (ok) graph_power_note_decode_token(g, t_read - t0); - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after graph eval failure also failed\n"); - } - } - return ok; -} - -static bool metal_graph_streaming_decode_prefill_wide_default( - const ds4_weights *weights) { - return DS4_MODEL_VARIANT == DS4_VARIANT_FLASH && - weights && - DS4_N_LAYER > 0 && - weights->layer[0].ffn_gate_exps->type == DS4_TENSOR_Q4_K && - weights->layer[0].ffn_up_exps->type == DS4_TENSOR_Q4_K && - weights->layer[0].ffn_down_exps->type == DS4_TENSOR_Q4_K; -} - -static uint32_t metal_graph_streaming_decode_prefill_max_tokens( - const ds4_gpu_graph *g, - const ds4_weights *weights) { - (void)g; - if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_DECODE_PREFILL", - "DS4_METAL_DISABLE_STREAMING_DECODE_PREFILL")) { - return 0; - } - - const char *env = glm_graph_env_value( - "DS4_ROCM_STREAMING_DECODE_PREFILL_MAX", - "DS4_METAL_STREAMING_DECODE_PREFILL_MAX"); - if (env && env[0]) { - char *end = NULL; - const long v = strtol(env, &end, 10); - if (end != env) { - if (v <= 0) return 0; - if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; - return (uint32_t)v; - } - } - - if (DS4_MODEL_VARIANT != DS4_VARIANT_PRO && - DS4_MODEL_VARIANT != DS4_VARIANT_FLASH) { - return 0u; - } - return metal_graph_streaming_decode_prefill_wide_default(weights) ? 64u : 18u; -} - -static bool metal_graph_use_streaming_decode_prefill( - const ds4_gpu_graph *g, - const ds4_weights *weights, - uint32_t n_tokens) { - const uint32_t max_tokens = - metal_graph_streaming_decode_prefill_max_tokens(g, weights); - return g && - g->ssd_streaming && - !g->quality && - n_tokens != 0 && - max_tokens != 0 && - n_tokens <= max_tokens; -} - -static bool metal_graph_use_streaming_decode_prefill_range( - const ds4_gpu_graph *g, - const ds4_weights *weights, - uint32_t start, - uint32_t n_tokens) { - /* - * Short streamed prefill is latency-sensitive. Use the decode-style path - * by default for SSD streaming, while keeping a cold-only escape hatch for - * strict-vector tests that need canonical layer-major prefill semantics. - */ - if (start == 0) { - if (glm_graph_env_present( - "DS4_ROCM_DISABLE_STREAMING_COLD_DECODE_PREFILL", - "DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL")) { - return false; - } - } - return metal_graph_use_streaming_decode_prefill(g, weights, n_tokens); -} - -static bool metal_graph_prefill_decode_streaming_range( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - float *logits, - bool show_progress, - ds4_session_progress_fn progress, - void *progress_ud, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - ds4_session_cancel_fn cancel, - void *cancel_ud, - bool *cancelled) { - if (!metal_graph_use_streaming_decode_prefill(g, weights, n_tokens)) return false; - if (!prompt || start > (uint32_t)prompt->len || - n_tokens > (uint32_t)prompt->len - start) return false; - if (start == 0) { - ds4_gpu_stream_expert_cache_reset_route_hotness(); - } - - const bool profile = - glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", - "DS4_METAL_GRAPH_PREFILL_PROFILE"); - const double t0 = profile ? now_sec() : 0.0; - - /* - * `prefill_chunk` is not just UI progress: ds4_session_sync() wraps it to - * advance the live checkpoint, and ds4-server may save that checkpoint. - * Decode-style prefill only reads logits for the final token, so report one - * cacheable chunk at the end. `prefill_display` remains per-token UI only. - */ - if (progress) progress(progress_ud, "prefill_chunk", (int)start, prompt->len); - if (display_progress) { - display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); - } - - for (uint32_t i = 0; i < n_tokens; i++) { - if (cancel && cancel(cancel_ud)) { - if (cancelled) *cancelled = true; - return true; - } - const uint32_t pos = start + i; - const bool last = i + 1u == n_tokens; - float *token_logits = (last && logits) ? logits : NULL; - if (!metal_graph_eval_token_raw_swa(g, - model, - weights, - prompt->v[pos], - pos, - token_logits)) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after decode-style streaming prefill failure also failed\n"); - } - return false; - } - - if (last && progress && logits) { - progress(progress_ud, "prefill_chunk", (int)(pos + 1u), prompt->len); - } - if (display_progress) { - display_progress(display_progress_ud, "prefill_display", (int)(pos + 1u), prompt->len); - } - if (cancel && cancel(cancel_ud)) { - if (cancelled) *cancelled = true; - return true; - } - if (show_progress) { - fprintf(stderr, "ds4: gpu streaming prefill token %u/%u\r", - i + 1u, - n_tokens); - fflush(stderr); - } - } - if (show_progress) fputc('\n', stderr); - - if (profile) { - const double t1 = now_sec(); - fprintf(stderr, - "ds4: gpu decode-style streaming prefill start=%u tokens=%u total=%.3f ms\n", - start, - n_tokens, - (t1 - t0) * 1000.0); - } - return true; -} - -static bool metal_graph_capture_prefill_seed_router_selected( - ds4_gpu_graph *g, - uint32_t il, - uint32_t n_tokens) { - uint32_t k = metal_graph_streaming_prefill_cache_seed_k(g); - if (k == 0) return true; - if (k > n_tokens) k = n_tokens; - g->prefill_seed_tokens = k; - if (!g->prefill_seed_router_selected || !metal_graph_batch_router_selected(g) || - il >= DS4_N_LAYER || n_tokens == 0 || sizeof(int) != sizeof(int32_t)) { - return false; - } - - const uint64_t bytes = (uint64_t)k * DS4_N_EXPERT_USED * sizeof(int32_t); - const uint64_t src_off = (uint64_t)(n_tokens - k) * - DS4_N_EXPERT_USED * sizeof(int); - const uint64_t dst_off = (uint64_t)il * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * sizeof(int32_t); - return ds4_gpu_tensor_copy(g->prefill_seed_router_selected, - dst_off, - metal_graph_batch_router_selected(g), - src_off, - bytes) != 0; -} - -static bool metal_graph_seed_streaming_expert_cache_from_prefill( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; - if (!metal_graph_streaming_prefill_cache_seed_enabled(g)) return true; - if (!model || !weights || !g->prefill_seed_router_selected || seed_tokens == 0) { - return false; - } - - int32_t selected[DS4_MAX_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_MAX_EXPERT_USED]; - const uint64_t bytes = (uint64_t)DS4_N_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * sizeof(selected[0]); - if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, - 0, - selected, - bytes) == 0) { - return false; - } - - const bool profile = - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_PROFILE", - "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE"); - const double t0 = profile ? now_sec() : 0.0; - uint32_t seeded_layers = 0; - uint32_t seeded_rows = 0; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const ds4_layer_weights *layer = &weights->layer[il]; - if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { - continue; - } - - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { - fprintf(stderr, "ds4: Metal prefill expert-cache seed byte size overflow at layer %u\n", il); - return false; - } - const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; - const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - for (uint32_t row = 0; row < seed_tokens; row++) { - const size_t sel_off = ((size_t)il * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS + - row) * DS4_N_EXPERT_USED; - if (ds4_gpu_stream_expert_cache_seed_selected( - &table, - selected + sel_off, - DS4_N_EXPERT_USED) == 0) { - return false; - } - seeded_rows++; - } - seeded_layers++; - } - if (profile) { - fprintf(stderr, - "ds4: Metal streaming prefill expert-cache seed k=%u layers=%u rows=%u time=%.3f ms\n", - seed_tokens, - seeded_layers, - seeded_rows, - (now_sec() - t0) * 1000.0); - } - return true; -} - -static bool metal_graph_seed_streaming_expert_cache_from_hotlist( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - if (!metal_graph_streaming_expert_hotlist_enabled(g)) return true; - if (!model || !weights) return false; - - uint32_t cache_budget = 0; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const ds4_layer_weights *layer = &weights->layer[il]; - if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { - continue; - } - - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { - fprintf(stderr, "ds4: streaming expert hotlist budget byte size overflow at layer %u\n", il); - return false; - } - const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; - const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; - cache_budget = ds4_gpu_stream_expert_cache_budget_for_expert_size( - gate_expert_bytes, - down_expert_bytes); - break; - } - if (cache_budget == 0) return true; - const uint32_t preload_count = - metal_graph_streaming_expert_preload_count(g, cache_budget); - if (preload_count == 0) return true; - const uint32_t current_count = - ds4_gpu_stream_expert_cache_current_count(); - const char *path = glm_graph_env_value("DS4_ROCM_STREAMING_EXPERT_HOTLIST", - "DS4_METAL_STREAMING_EXPERT_HOTLIST"); - const bool from_file = path && path[0]; - const bool refresh_builtin_glm = - !from_file && g_ds4_shape.variant == DS4_VARIANT_GLM52; - const bool profile = - glm_graph_env_present("DS4_ROCM_STREAMING_EXPERT_HOTLIST_PROFILE", - "DS4_METAL_STREAMING_EXPERT_HOTLIST_PROFILE"); - if (!from_file && !refresh_builtin_glm && current_count >= preload_count) { - if (profile) { - fprintf(stderr, - "ds4: streaming expert hotlist seed skipped preload=%u current=%u\n", - preload_count, - current_count); - } - return true; - } - - int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT]; - uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT]; - uint32_t counts[DS4_MAX_LAYER]; - bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT]; - memset(experts, 0, sizeof(experts)); - memset(priorities, 0, sizeof(priorities)); - memset(counts, 0, sizeof(counts)); - memset(seen, 0, sizeof(seen)); - - uint32_t loaded = 0; - if (from_file) { - if (!metal_graph_streaming_expert_hotlist_load_file(path, - preload_count, - experts, - priorities, - counts, - seen, - &loaded)) { - return false; - } - } else if (!metal_graph_streaming_expert_hotlist_load_default(preload_count, - experts, - priorities, - counts, - seen, - &loaded)) { - return false; - } - if (loaded == 0) return true; - - const double t0 = profile ? now_sec() : 0.0; - uint32_t seeded_layers = 0; - uint32_t seeded_experts = 0; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t n = counts[il]; - if (n == 0) continue; - const ds4_layer_weights *layer = &weights->layer[il]; - if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { - continue; - } - - const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); - const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); - if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || - layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { - fprintf(stderr, "ds4: streaming expert hotlist seed byte size overflow at layer %u\n", il); - return false; - } - const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; - const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - if (ds4_gpu_stream_expert_cache_seed_experts( - &table, - experts[il], - priorities[il], - n) == 0) { - return false; - } - seeded_layers++; - seeded_experts += n; - } - if (profile) { - const char *source_name = NULL; - if (from_file) { - source_name = path; - } else if (g_ds4_shape.variant == DS4_VARIANT_GLM52) { - source_name = "built-in-glm52"; - } else if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { - source_name = "built-in-flash"; - } else if (g_ds4_shape.variant == DS4_VARIANT_PRO) { - source_name = "built-in-pro"; - } else { - source_name = "built-in"; - } - fprintf(stderr, - "ds4: streaming expert hotlist seed source=%s preload=%u loaded=%u layers=%u experts=%u time=%.3f ms\n", - source_name, - preload_count, - loaded, - seeded_layers, - seeded_experts, - (now_sec() - t0) * 1000.0); - } - return true; -} - -typedef struct { - int id0; - int id1; - float value0; - float value1; - bool valid; - bool fast_attention; -} metal_graph_top2_result; - -/* Greedy verifier helper. Speculative decoding only needs the target model's - * top token after most accepted draft rows; the full vocabulary row is needed - * once, for the final committed state that normal sampling will continue from. - * Keeping intermediate rows device-resident avoids turning verification into a - * sequence of large CPU readbacks. */ -static bool metal_graph_eval_token_raw_swa_top( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - uint32_t pos, - int *top_id, - float *logits, - bool allow_split_top1, - metal_graph_top2_result *top2, - bool force_fast_attention) { - if (!top_id) return false; - if (top2) memset(top2, 0, sizeof(*top2)); - - const bool fast_attention = - allow_split_top1 && - logits == NULL && - (force_fast_attention || metal_graph_cuda_greedy_splitkv_requested()); - if (top2) top2->fast_attention = fast_attention; - const int old_fast_attention = - ds4_gpu_set_decode_fast_attention(fast_attention ? 1 : 0); - const bool profile = getenv("DS4_METAL_GRAPH_TOKEN_PROFILE") != NULL; - const double t0 = profile ? now_sec() : 0.0; - const bool split_top1 = - allow_split_top1 && - logits == NULL && - top2 == NULL && - g->cuda_tp_output && - metal_graph_cuda_greedy_split_top1_requested(); - if (split_top1) { - int output_tiers[DS4_MAX_GPUS] = {0}; - uint32_t output_ways = 0; - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, - token, pos, false, true); - if (ok) ok = metal_graph_encode_output_head_split_top1(g, - model, - weights, - weights->output->dim[1], - output_tiers, - &output_ways); - const double t_encoded = profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = profile ? now_sec() : 0.0; - if (ok) { - bool have_best = false; - uint32_t best_id = 0; - float best_value = 0.0f; - uint32_t cand_ids[DS4_MAX_GPUS] = {0}; - float cand_values[DS4_MAX_GPUS] = {0.0f}; - ok = output_ways <= DS4_MAX_GPUS && - ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], - 0, - cand_ids, - (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && - ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], - 0, - cand_values, - (uint64_t)output_ways * sizeof(cand_values[0])) != 0; - for (uint32_t i = 0; ok && i < output_ways; i++) { - const uint32_t cand_id = cand_ids[i]; - const float cand_value = cand_values[i]; - if (ok && - (!have_best || - cand_value > best_value || - (cand_value == best_value && cand_id < best_id))) { - have_best = true; - best_id = cand_id; - best_value = cand_value; - } - } - ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; - if (ok) *top_id = (int)best_id; - } - const double t_read = profile ? now_sec() : 0.0; - if (profile) { - fprintf(stderr, - "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=0 top=1 split_top1=1\n", - pos, - (t_encoded - t0) * 1000.0, - (t_done - t_encoded) * 1000.0, - (t_read - t_done) * 1000.0, - (t_read - t0) * 1000.0); - } - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after split-top graph eval failure also failed\n"); - } - } - (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); - return ok; - } - - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, - token, pos, true, true); - if (ok) { - ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), - metal_graph_logits(g), - DS4_N_VOCAB) != 0; - } - const double t_encoded = profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = profile ? now_sec() : 0.0; - if (ok && top2) { - uint32_t ids[2] = {0, 0}; - float values[2] = {0.0f, 0.0f}; - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), - 0, - ids, - sizeof(ids)) != 0 && - ds4_gpu_tensor_read(metal_graph_comp_mask(g), - 0, - values, - sizeof(values)) != 0; - if (ok && ids[0] <= (uint32_t)INT32_MAX && ids[1] <= (uint32_t)INT32_MAX) { - top2->id0 = (int)ids[0]; - top2->id1 = (int)ids[1]; - top2->value0 = values[0]; - top2->value1 = values[1]; - top2->valid = isfinite(values[0]) && isfinite(values[1]); - *top_id = top2->id0; - } else { - ok = false; - } - } else if (ok) { - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top_id, sizeof(*top_id)) != 0; - } - if (ok && logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - const double t_read = profile ? now_sec() : 0.0; - if (profile) { - fprintf(stderr, - "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d top=1 split_top1=0\n", - pos, - (t_encoded - t0) * 1000.0, - (t_done - t_encoded) * 1000.0, - (t_read - t_done) * 1000.0, - (t_read - t0) * 1000.0, - logits != NULL); - } - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after top-only graph eval failure also failed\n"); - } - } - (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); - return ok; -} - -static bool dspark_stage0_weights_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw) { - if (!g || !dw || dw->n_stages == 0 || dw->target_layer_count == 0 || - dw->target_layer_count != g->dspark_target_layer_count || - !g->dspark_target_hidden || !g->dspark_stage0_proj || - !g->dspark_main_x) { - return false; - } - - const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; - const ds4_tensor *main_proj = stage0->main_proj; - const ds4_tensor *main_norm = stage0->main_norm; - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - return main_proj && - main_norm && - dspark_tensor_type_matches(main_proj->type, DS4_DSPARK_LAYOUT_DENSE) && - main_norm->type == DS4_TENSOR_F32 && - main_proj->ndim == 2 && - main_proj->dim[0] == in_dim && - main_proj->dim[1] == DS4_N_EMBD && - main_norm->ndim == 1 && - main_norm->dim[0] == DS4_N_EMBD; -} - -static bool metal_graph_eval_dspark_stage0( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw) { - if (!g || !dspark_model || !dw || !dspark_stage0_weights_ready(g, dw)) { - return false; - } - - const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, - dspark_model, - stage0->main_proj, - in_dim, - DS4_N_EMBD, - g->dspark_target_hidden, - 1); - } - if (ok) { - ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, - g->dspark_stage0_proj, - dspark_model->map, - dspark_model->size, - stage0->main_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - return ok; -} - -static bool dspark_stage0_batch_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw, - uint32_t n_tokens) { - if (!dspark_stage0_weights_ready(g, dw) || - n_tokens == 0 || - n_tokens > g->prefill_cap || - !g->dspark_target_hidden_batch || - !metal_graph_batch_ffn_cur(g) || - !metal_graph_batch_ffn_norm(g) || - !metal_graph_batch_cur_hc(g)) { - return false; - } - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - return ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) >= - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) >= - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_cur_hc(g)) >= - (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float) && - in_dim <= SIZE_MAX / sizeof(float); -} - -static bool metal_graph_pack_dspark_target_hidden_batch( - ds4_gpu_graph *g, - const ds4_dspark_weights *dw, - ds4_gpu_tensor *packed, - uint32_t n_tokens) { - if (!g || !dw || !packed || n_tokens == 0 || - n_tokens > g->prefill_cap || - dw->target_layer_count != g->dspark_target_layer_count) { - return false; - } - - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - const uint64_t packed_count = (uint64_t)n_tokens * in_dim; - if (packed_count == 0 || - packed_count > (uint64_t)SIZE_MAX / sizeof(float)) { - return false; - } - return ds4_gpu_pack_slot_rows_f32_tensor(packed, - g->dspark_target_hidden_batch, - n_tokens, - DS4_N_EMBD, - dw->target_layer_count, - g->prefill_cap) != 0; -} - -static bool metal_graph_eval_dspark_stage0_batch( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t n_tokens, - bool commands_open) { - if (!g || !dspark_model || !dw || - !dspark_stage0_batch_ready(g, dw, n_tokens)) { - return false; - } - - const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - const uint64_t packed_bytes = (uint64_t)n_tokens * in_dim * sizeof(float); - bool packed_owned = false; - ds4_gpu_tensor *packed = NULL; - if (g->dspark_stage0_packed && - ds4_gpu_tensor_bytes(g->dspark_stage0_packed) >= packed_bytes) { - packed = g->dspark_stage0_packed; - } else { - packed = ds4_gpu_tensor_alloc(packed_bytes); - packed_owned = true; - } - if (!packed) return false; - - bool ok = metal_graph_pack_dspark_target_hidden_batch(g, dw, packed, n_tokens); - if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = metal_graph_matmul_plain_tensor(metal_graph_batch_ffn_cur(g), - dspark_model, - stage0->main_proj, - in_dim, - DS4_N_EMBD, - packed, - n_tokens); - } - if (ok) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_ffn_norm(g), - metal_graph_batch_ffn_cur(g), - dspark_model->map, - dspark_model->size, - stage0->main_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; - if (!ok && !commands_open) (void)ds4_gpu_synchronize(); - if (packed_owned) ds4_gpu_tensor_free(packed); - return ok; -} - -static bool dspark_draft_block_ready( - const ds4_gpu_graph *g, - const ds4_weights *base_weights, - const ds4_dspark_weights *dw, - int token) { - if (!g || !base_weights || !dw || !base_weights->token_embd || - !g->dspark_draft_tokens || !g->dspark_draft_hc || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - g->dspark_block_size != dw->block_size || - !dw->has_noise_token_id) { - return false; - } - const uint32_t n_vocab = (uint32_t)base_weights->token_embd->dim[1]; - return token >= 0 && - (uint32_t)token < n_vocab && - dw->noise_token_id < n_vocab; -} - -static bool dspark_stage_input_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw) { - if (!g || !dw || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - g->dspark_block_size != dw->block_size || - !g->dspark_main_x || !g->dspark_draft_hc || - !g->dspark_target_hc || !g->dspark_stage_input_hc || - !g->dspark_position_ids) { - return false; - } - if (dw->block_size == UINT32_MAX) return false; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t rows = (uint64_t)dw->block_size + 1u; - return ds4_gpu_tensor_bytes(g->dspark_target_hc) >= - hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(g->dspark_stage_input_hc) >= - rows * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(g->dspark_position_ids) >= - rows * sizeof(int32_t); -} - -static bool dspark_stage_cache_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw) { - if (!g || !dw || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - g->dspark_cache_cap == 0 || - !metal_graph_dspark_cache_current_window_valid(g)) { - return false; - } - const uint64_t bytes = - (uint64_t)g->dspark_cache_cap * DS4_N_HEAD_DIM * sizeof(float); - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - if (!g->dspark_raw_cache[stage] || - ds4_gpu_tensor_bytes(g->dspark_raw_cache[stage]) < bytes) { - return false; - } - } - return true; -} - -static bool dspark_noncausal_attention_probe_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw) { - if (!g || !dw || - dw->n_stages == 0 || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - g->prefill_cap < dw->block_size + 1u || - !metal_graph_batch_q(g) || !metal_graph_batch_heads(g) || - !g->dspark_raw_cache[0]) { - return false; - } - const ds4_layer_weights *block = &dw->stage[0].block; - if (!block->attn_sinks) return false; - - const uint64_t rows = (uint64_t)dw->block_size + 1u; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - return ds4_gpu_tensor_bytes(metal_graph_batch_q(g)) >= - rows * q_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_heads(g)) >= - rows * q_dim * sizeof(float) && - ds4_gpu_tensor_bytes(g->dspark_raw_cache[0]) >= - rows * DS4_N_HEAD_DIM * sizeof(float); -} - -static bool metal_graph_probe_dspark_noncausal_attention( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw) { - if (!g || !dspark_model || !dspark_noncausal_attention_probe_ready(g, dw)) { - return false; - } - - const uint32_t rows = dw->block_size + 1u; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const ds4_layer_weights *block = &dw->stage[0].block; - bool ok = ds4_gpu_tensor_fill_f32(metal_graph_batch_q(g), - 0.0f, - (uint64_t)rows * q_dim) != 0 && - ds4_gpu_tensor_fill_f32(g->dspark_raw_cache[0], - 0.0f, - (uint64_t)rows * DS4_N_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = ds4_gpu_attention_noncausal_raw_batch_heads_tensor( - metal_graph_batch_heads(g), - dspark_model->map, - dspark_model->size, - block->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->dspark_raw_cache[0], - rows, - rows, - rows, - 0, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - return ok; -} - -static bool metal_graph_prepare_dspark_setup_block( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_dspark_weights *dw, - int token, - uint32_t pos) { - if (!g || !base_model || - !dspark_draft_block_ready(g, base_weights, dw, token) || - !dspark_stage_input_ready(g, dw)) { - return false; - } - if (pos > (uint32_t)INT32_MAX || - dw->block_size > (uint32_t)INT32_MAX || - pos > (uint32_t)INT32_MAX - dw->block_size) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t hc_bytes = hc_dim * sizeof(float); - int32_t positions[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; - positions[0] = (int32_t)pos; - for (uint32_t i = 0; i < dw->block_size; i++) { - positions[i + 1u] = (int32_t)(pos + i); - } - int32_t ids[DS4_DSPARK_MAX_BLOCK_SIZE]; - ids[0] = (int32_t)token; - for (uint32_t i = 1; i < dw->block_size; i++) { - ids[i] = (int32_t)dw->noise_token_id; - } - - bool ok = ds4_gpu_tensor_write(g->dspark_draft_tokens, - 0, - ids, - (uint64_t)dw->block_size * sizeof(ids[0])) != 0 && - ds4_gpu_tensor_write(g->dspark_position_ids, - 0, - positions, - ((uint64_t)dw->block_size + 1u) * - sizeof(positions[0])) != 0; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = ds4_gpu_embed_tokens_hc_tensor(g->dspark_draft_hc, - g->dspark_draft_tokens, - base_model->map, - base_model->size, - base_weights->token_embd->abs_offset, - (uint32_t)base_weights->token_embd->dim[1], - dw->block_size, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (ok) { - ok = ds4_gpu_repeat_hc_tensor(g->dspark_target_hc, - g->dspark_main_x, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, - 0, - g->dspark_target_hc, - 0, - hc_bytes) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, - hc_bytes, - g->dspark_draft_hc, - 0, - (uint64_t)dw->block_size * hc_bytes) != 0; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - return ok; -} - -static bool metal_graph_prepare_dspark_stage0_setup_block( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - int token, - uint32_t pos) { - if (!g || !base_model || !dspark_model || - !dspark_stage0_weights_ready(g, dw) || - !dspark_draft_block_ready(g, base_weights, dw, token) || - !dspark_stage_input_ready(g, dw)) { - return false; - } - if (pos > (uint32_t)INT32_MAX || - dw->block_size > (uint32_t)INT32_MAX || - pos > (uint32_t)INT32_MAX - dw->block_size) { - return false; - } - - const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t hc_bytes = hc_dim * sizeof(float); - int32_t positions[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; - positions[0] = (int32_t)pos; - for (uint32_t i = 0; i < dw->block_size; i++) { - positions[i + 1u] = (int32_t)(pos + i); - } - int32_t ids[DS4_DSPARK_MAX_BLOCK_SIZE]; - ids[0] = (int32_t)token; - for (uint32_t i = 1; i < dw->block_size; i++) { - ids[i] = (int32_t)dw->noise_token_id; - } - - /* DS4_DSPARK_PROP_PROFILE=1: break the setup block into phases to - * localize the TP-only prop_setup inflation (26ms vs 1.3ms single). */ - const bool prop_profile = getenv("DS4_DSPARK_PROP_PROFILE") != NULL; - const double pp_t0 = prop_profile ? now_sec() : 0.0; - bool ok = ds4_gpu_tensor_write(g->dspark_draft_tokens, - 0, - ids, - (uint64_t)dw->block_size * sizeof(ids[0])) != 0 && - ds4_gpu_tensor_write(g->dspark_position_ids, - 0, - positions, - ((uint64_t)dw->block_size + 1u) * - sizeof(positions[0])) != 0; - const double pp_t1 = prop_profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_begin_commands() != 0; - const double pp_t2 = prop_profile ? now_sec() : 0.0; - if (ok) { - ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, - dspark_model, - stage0->main_proj, - in_dim, - DS4_N_EMBD, - g->dspark_target_hidden, - 1); - } - if (ok) { - ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, - g->dspark_stage0_proj, - dspark_model->map, - dspark_model->size, - stage0->main_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - } - if (ok) { - ok = ds4_gpu_embed_tokens_hc_tensor(g->dspark_draft_hc, - g->dspark_draft_tokens, - base_model->map, - base_model->size, - base_weights->token_embd->abs_offset, - (uint32_t)base_weights->token_embd->dim[1], - dw->block_size, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (ok) { - ok = ds4_gpu_repeat_hc_tensor(g->dspark_target_hc, - g->dspark_main_x, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, - 0, - g->dspark_target_hc, - 0, - hc_bytes) != 0; - } - if (ok) { - ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, - hc_bytes, - g->dspark_draft_hc, - 0, - (uint64_t)dw->block_size * hc_bytes) != 0; - } - const double pp_t3 = prop_profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - if (prop_profile) { - const double pp_t4 = now_sec(); - fprintf(stderr, - "ds4: DSpark prop-setup phases: writes=%.3fms begin=%.3fms " - "encode=%.3fms end/wait=%.3fms\n", - (pp_t1 - pp_t0) * 1000.0, - (pp_t2 - pp_t1) * 1000.0, - (pp_t3 - pp_t2) * 1000.0, - (pp_t4 - pp_t3) * 1000.0); - } - if (!ok) (void)ds4_gpu_synchronize(); - return ok; -} - -static bool dspark_stage_block_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw, - uint32_t stage) { - if (!g || !dw || - stage >= dw->n_stages || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - g->prefill_cap < dw->block_size + 1u || - !g->dspark_stage_output_hc || - !dspark_stage_input_ready(g, dw) || - !dspark_stage_cache_ready(g, dw)) { - return false; - } - - const ds4_layer_weights *l = &dw->stage[stage].block; - if (!l->hc_attn_fn || !l->hc_attn_scale || !l->hc_attn_base || - !l->attn_norm || !l->attn_q_a || !l->attn_q_a_norm || - !l->attn_q_b || !l->attn_kv || !l->attn_kv_a_norm || - !l->attn_sinks || !l->attn_output_a || !l->attn_output_b || - !l->hc_ffn_fn || !l->hc_ffn_scale || !l->hc_ffn_base || - !l->ffn_norm || !l->ffn_gate_inp || !l->ffn_exp_probs_b || - !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || - !l->ffn_gate_shexp || !l->ffn_up_shexp || !l->ffn_down_shexp) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t rows = (uint64_t)dw->block_size + 1u; - const uint64_t draft = dw->block_size; - const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - const uint64_t group_dim = - (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); - - return - dspark_tensor_type_matches(l->hc_attn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && - l->hc_attn_scale->type == DS4_TENSOR_F32 && - l->hc_attn_base->type == DS4_TENSOR_F32 && - l->attn_norm->type == DS4_TENSOR_F32 && - dspark_tensor_type_matches(l->attn_q_a->type, DS4_DSPARK_LAYOUT_DENSE) && - l->attn_q_a_norm->type == DS4_TENSOR_F32 && - dspark_tensor_type_matches(l->attn_q_b->type, DS4_DSPARK_LAYOUT_DENSE) && - dspark_tensor_type_matches(l->attn_kv->type, DS4_DSPARK_LAYOUT_DENSE) && - l->attn_kv_a_norm->type == DS4_TENSOR_F32 && - l->attn_sinks->type == DS4_TENSOR_F32 && - l->attn_output_a->type == DS4_TENSOR_Q8_0 && - l->attn_output_b->type == DS4_TENSOR_Q8_0 && - dspark_tensor_type_matches(l->hc_ffn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && - l->hc_ffn_scale->type == DS4_TENSOR_F32 && - l->hc_ffn_base->type == DS4_TENSOR_F32 && - l->ffn_norm->type == DS4_TENSOR_F32 && - dspark_tensor_type_matches(l->ffn_gate_inp->type, DS4_DSPARK_LAYOUT_DENSE) && - l->ffn_exp_probs_b->type == DS4_TENSOR_F32 && - tensor_is_routed_expert_type(l->ffn_gate_exps->type) && - l->ffn_gate_exps->type == l->ffn_up_exps->type && - tensor_is_routed_expert_type(l->ffn_down_exps->type) && - l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && - l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && - l->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && - l->hc_attn_fn->ndim == 2 && - l->hc_attn_fn->dim[0] == hc_dim && - l->hc_attn_fn->dim[1] == mix_hc && - l->attn_q_a->ndim == 2 && - l->attn_q_a->dim[0] == DS4_N_EMBD && - l->attn_q_a->dim[1] == DS4_N_LORA_Q && - l->attn_q_b->ndim == 2 && - l->attn_q_b->dim[0] == DS4_N_LORA_Q && - l->attn_q_b->dim[1] == q_dim && - l->attn_kv->ndim == 2 && - l->attn_kv->dim[0] == DS4_N_EMBD && - l->attn_kv->dim[1] == DS4_N_HEAD_DIM && - l->attn_output_a->ndim == 2 && - l->attn_output_a->dim[0] == group_dim && - l->attn_output_a->dim[1] == out_low_dim && - l->attn_output_b->ndim == 2 && - l->attn_output_b->dim[0] == out_low_dim && - l->attn_output_b->dim[1] == DS4_N_EMBD && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= rows * mix_hc * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= rows * mix_hc * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= rows * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_attn_cur(g)) >= rows * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_attn_norm(g)) >= rows * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_qr(g)) >= draft * DS4_N_LORA_Q * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_qr_norm(g)) >= draft * DS4_N_LORA_Q * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_q(g)) >= draft * q_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_kv_raw(g)) >= rows * DS4_N_HEAD_DIM * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_kv(g)) >= rows * DS4_N_HEAD_DIM * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_heads(g)) >= draft * q_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_attn_out(g)) >= draft * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_after_attn_hc(g)) >= draft * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_next_hc(g)) >= draft * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_prefill_tokens(g)) >= draft * sizeof(int32_t) && - ds4_gpu_tensor_bytes(g->dspark_stage_output_hc) >= draft * hc_dim * sizeof(float); -} - -static bool dspark_stage_target_cache_seed_ready( - const ds4_gpu_graph *g, - const ds4_dspark_weights *dw, - uint32_t stage, - uint32_t n_tokens) { - if (!g || !dw || - stage >= dw->n_stages || - n_tokens == 0 || - n_tokens > g->prefill_cap || - !dspark_stage_cache_ready(g, dw) || - !metal_graph_batch_cur_hc(g) || - !metal_graph_batch_hc_mix(g) || - !metal_graph_batch_hc_split(g) || - !metal_graph_batch_flat_hc(g) || - !metal_graph_batch_attn_cur(g) || - !metal_graph_batch_attn_norm(g) || - !metal_graph_batch_kv_raw(g) || - !metal_graph_batch_kv(g)) { - return false; - } - - const ds4_layer_weights *l = &dw->stage[stage].block; - if (!l->hc_attn_fn || !l->hc_attn_scale || !l->hc_attn_base || - !l->attn_norm || !l->attn_kv || !l->attn_kv_a_norm) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - return - dspark_tensor_type_matches(l->hc_attn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && - l->hc_attn_scale->type == DS4_TENSOR_F32 && - l->hc_attn_base->type == DS4_TENSOR_F32 && - l->attn_norm->type == DS4_TENSOR_F32 && - dspark_tensor_type_matches(l->attn_kv->type, DS4_DSPARK_LAYOUT_DENSE) && - l->attn_kv_a_norm->type == DS4_TENSOR_F32 && - l->hc_attn_fn->ndim == 2 && - l->hc_attn_fn->dim[0] == hc_dim && - l->hc_attn_fn->dim[1] == mix_hc && - l->attn_kv->ndim == 2 && - l->attn_kv->dim[0] == DS4_N_EMBD && - l->attn_kv->dim[1] == DS4_N_HEAD_DIM && - ds4_gpu_tensor_bytes(metal_graph_batch_cur_hc(g)) >= - (uint64_t)n_tokens * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= - (uint64_t)n_tokens * mix_hc * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= - (uint64_t)n_tokens * mix_hc * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= - (uint64_t)n_tokens * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_attn_cur(g)) >= - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_attn_norm(g)) >= - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_kv_raw(g)) >= - (uint64_t)n_tokens * DS4_N_HEAD_DIM * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_kv(g)) >= - (uint64_t)n_tokens * DS4_N_HEAD_DIM * sizeof(float); -} - -static bool metal_graph_seed_dspark_stage_target_cache( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t stage, - uint32_t pos0, - uint32_t n_tokens, - bool commands_open) { - if (!g || !dspark_model || !dw || - !dspark_stage_target_cache_seed_ready(g, dw, stage, n_tokens) || - n_tokens > g->dspark_cache_cap) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const ds4_layer_weights *block = &dw->stage[stage].block; - const bool fuse_hc_norm = DS4_N_HC == 4 && - !metal_graph_use_reference_hc_decode() && - metal_graph_enable_batch_hc_norm_fusion(); - - ds4_gpu_tensor *hc_mix_view = - ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), - 0, - (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *hc_split_view = - ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - 0, - (uint64_t)n_tokens * mix_hc * sizeof(float)); - ds4_gpu_tensor *attn_cur_view = - ds4_gpu_tensor_view(metal_graph_batch_attn_cur(g), - 0, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - bool ok = hc_mix_view && hc_split_view && attn_cur_view; - - const float freq_base = DS4_ROPE_FREQ_BASE; - const float freq_scale = 1.0f; - const float ext_factor = 0.0f; - const float attn_factor = 1.0f; - - if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), - metal_graph_batch_cur_hc(g), - (uint32_t)hc_dim, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(hc_mix_view, - dspark_model, - block->hc_attn_fn, - hc_dim, - mix_hc, - metal_graph_batch_flat_hc(g), - n_tokens); - if (fuse_hc_norm) { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, - metal_graph_batch_attn_norm(g), - hc_split_view, - hc_mix_view, - metal_graph_batch_cur_hc(g), - dspark_model->map, - dspark_model->size, - block->hc_attn_scale->abs_offset, - block->hc_attn_base->abs_offset, - block->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - } else { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, - hc_split_view, - hc_mix_view, - metal_graph_batch_cur_hc(g), - dspark_model->map, - dspark_model->size, - block->hc_attn_scale->abs_offset, - block->hc_attn_base->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), - metal_graph_batch_attn_cur(g), - dspark_model->map, - dspark_model->size, - block->attn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_kv_raw(g), - dspark_model, - block->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - metal_graph_batch_attn_norm(g), - n_tokens); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), - metal_graph_batch_kv_raw(g), - dspark_model->map, - dspark_model->size, - block->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), - n_tokens, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos0, - 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), - n_tokens, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; - if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor(g->dspark_raw_cache[stage], - metal_graph_batch_kv(g), - g->dspark_cache_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; - if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; - - ds4_gpu_tensor_free(attn_cur_view); - ds4_gpu_tensor_free(hc_split_view); - ds4_gpu_tensor_free(hc_mix_view); - if (!ok && !commands_open) (void)ds4_gpu_synchronize(); - return ok; -} - -static bool metal_graph_seed_dspark_initial_cache_from_prefill( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t batch_start, - uint32_t n_tokens, - uint32_t *seeded_rows) { - if (seeded_rows) *seeded_rows = 0; - if (!g || !dspark_model || !dw || - n_tokens == 0 || - n_tokens > g->prefill_cap || - n_tokens > g->dspark_cache_cap || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - !dspark_stage0_batch_ready(g, dw, n_tokens) || - !dspark_stage_cache_ready(g, dw)) { - return false; - } - - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = metal_graph_eval_dspark_stage0_batch(g, - dspark_model, - dw, - n_tokens, - true); - } - if (ok) { - ok = ds4_gpu_repeat_hc_rows_tensor(metal_graph_batch_cur_hc(g), - metal_graph_batch_ffn_norm(g), - n_tokens, - DS4_N_EMBD, - DS4_N_HC) != 0; - } - for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { - ok = metal_graph_seed_dspark_stage_target_cache(g, - dspark_model, - dw, - stage, - batch_start, - n_tokens, - true); - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) { - (void)ds4_gpu_synchronize(); - return false; - } - if (!metal_graph_dspark_cache_set_window(g, batch_start, n_tokens)) { - return false; - } - if (seeded_rows) *seeded_rows = n_tokens; - return true; -} - -static bool metal_graph_encode_dspark_next_stage_draft_input_from( - ds4_gpu_graph *g, - const ds4_dspark_weights *dw, - const ds4_gpu_tensor *draft_hc) { - if (!g || !dw || !dspark_stage_input_ready(g, dw) || - !draft_hc) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t hc_bytes = hc_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(draft_hc) < - (uint64_t)dw->block_size * hc_bytes) { - return false; - } - - return ds4_gpu_tensor_copy(g->dspark_stage_input_hc, - hc_bytes, - draft_hc, - 0, - (uint64_t)dw->block_size * hc_bytes) != 0; -} - -static bool metal_graph_profile_layer_env_match(const char *env_name, uint32_t il) { - const char *layer_env = getenv(env_name); - if (!layer_env || !layer_env[0]) return true; - - char *end = NULL; - const unsigned long layer = strtoul(layer_env, &end, 10); - return end != layer_env && - *end == '\0' && - layer <= UINT32_MAX && - (uint32_t)layer == il; -} - -static bool metal_graph_dspark_stage_profile_enabled(uint32_t stage) { - return getenv("DS4_DSPARK_STAGE_PROFILE") != NULL && - metal_graph_profile_layer_env_match("DS4_DSPARK_STAGE_PROFILE_STAGE", - stage); -} - -static bool metal_graph_dspark_stage_profile_boundary( - const char *part, - uint32_t stage, - uint32_t pos, - uint32_t rows, - double *stage_t0) { - if (ds4_gpu_end_commands() == 0) return false; - const double now = now_sec(); - fprintf(stderr, - "ds4: DSpark stage profile stage=%u pos=%u rows=%u %s=%.3f ms\n", - stage, - pos, - rows, - part, - (now - *stage_t0) * 1000.0); - *stage_t0 = now; - return ds4_gpu_begin_commands() != 0; -} - -static bool metal_graph_eval_dspark_stage_block( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t stage, - uint32_t pos, - uint32_t support_len, - uint32_t raw_start, - bool prepare_next_stage_input, - bool commands_open) { - if (!g || !dspark_model || !dw || - !dspark_stage_block_ready(g, dw, stage)) { - return false; - } - - const uint32_t draft = dw->block_size; - const uint32_t rows = draft + 1u; - if (support_len > g->dspark_cache_cap || - rows > g->dspark_cache_cap - support_len || - (support_len != 0 && raw_start >= g->dspark_cache_cap)) { - return false; - } - const uint32_t visible_rows = support_len + rows; - const uint32_t attention_raw_start = - support_len ? raw_start : (pos % g->dspark_cache_cap); - const uint32_t append_pos = support_len ? - (uint32_t)(((uint64_t)raw_start + support_len) % - g->dspark_cache_cap) : - attention_raw_start; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t group_dim = - (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); - const ds4_layer_weights *block = &dw->stage[stage].block; - const bool fuse_hc_norm = DS4_N_HC == 4 && - !metal_graph_use_reference_hc_decode() && - metal_graph_enable_batch_hc_norm_fusion(); - - ds4_gpu_tensor *hc_mix_view = - ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), - 0, - (uint64_t)rows * mix_hc * sizeof(float)); - ds4_gpu_tensor *hc_split_view = - ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - 0, - (uint64_t)rows * mix_hc * sizeof(float)); - ds4_gpu_tensor *attn_cur_view = - ds4_gpu_tensor_view(metal_graph_batch_attn_cur(g), - 0, - (uint64_t)rows * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *draft_attn_norm_view = - ds4_gpu_tensor_view(metal_graph_batch_attn_norm(g), - (uint64_t)DS4_N_EMBD * sizeof(float), - (uint64_t)draft * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *draft_hc_view = - ds4_gpu_tensor_view(g->dspark_stage_input_hc, - hc_dim * sizeof(float), - (uint64_t)draft * hc_dim * sizeof(float)); - ds4_gpu_tensor *draft_hc_split_view = - ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - mix_hc * sizeof(float), - (uint64_t)draft * mix_hc * sizeof(float)); - ds4_gpu_tensor *kv_target_view = - ds4_gpu_tensor_view(metal_graph_batch_kv(g), - 0, - (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - ds4_gpu_tensor *kv_draft_view = - ds4_gpu_tensor_view(metal_graph_batch_kv(g), - (uint64_t)DS4_N_HEAD_DIM * sizeof(float), - (uint64_t)draft * DS4_N_HEAD_DIM * sizeof(float)); - ds4_gpu_tensor *after_attn_hc_view = - ds4_gpu_tensor_view(metal_graph_batch_after_attn_hc(g), - 0, - (uint64_t)draft * hc_dim * sizeof(float)); - - bool ok = hc_mix_view && hc_split_view && attn_cur_view && - draft_attn_norm_view && draft_hc_view && - draft_hc_split_view && kv_target_view && kv_draft_view && - after_attn_hc_view; - const bool saved_streaming = g->ssd_streaming; - g->ssd_streaming = false; - - const float freq_base = DS4_ROPE_FREQ_BASE; - const float freq_scale = 1.0f; - const float ext_factor = 0.0f; - const float attn_factor = 1.0f; - const bool stage_profile = - metal_graph_dspark_stage_profile_enabled(stage); - double stage_t0 = stage_profile ? now_sec() : 0.0; -#define DS4_DSPARK_PROFILE_STAGE(part_) do { \ - if (ok && stage_profile) { \ - ok = metal_graph_dspark_stage_profile_boundary((part_), \ - stage, \ - pos, \ - rows, \ - &stage_t0); \ - } \ - } while (0) - - if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), - g->dspark_stage_input_hc, - (uint32_t)hc_dim, - rows, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(hc_mix_view, - dspark_model, - block->hc_attn_fn, - hc_dim, - mix_hc, - metal_graph_batch_flat_hc(g), - rows); - DS4_DSPARK_PROFILE_STAGE("attn_hc_pre"); - if (fuse_hc_norm) { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, - metal_graph_batch_attn_norm(g), - hc_split_view, - hc_mix_view, - g->dspark_stage_input_hc, - dspark_model->map, - dspark_model->size, - block->hc_attn_scale->abs_offset, - block->hc_attn_base->abs_offset, - block->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS, - DS4_RMS_EPS) != 0; - } else { - if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, - hc_split_view, - hc_mix_view, - g->dspark_stage_input_hc, - dspark_model->map, - dspark_model->size, - block->hc_attn_scale->abs_offset, - block->hc_attn_base->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_N_HC_SINKHORN_ITER, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), - metal_graph_batch_attn_cur(g), - dspark_model->map, - dspark_model->size, - block->attn_norm->abs_offset, - DS4_N_EMBD, - rows, - DS4_RMS_EPS) != 0; - } - DS4_DSPARK_PROFILE_STAGE("attn_norm"); - - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_qr(g), - dspark_model, - block->attn_q_a, - DS4_N_EMBD, - DS4_N_LORA_Q, - draft_attn_norm_view, - draft); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_qr_norm(g), - metal_graph_batch_qr(g), - dspark_model->map, - dspark_model->size, - block->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - draft, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_q(g), - dspark_model, - block->attn_q_b, - DS4_N_LORA_Q, - q_dim, - metal_graph_batch_qr_norm(g), - draft); - if (ok) ok = ds4_gpu_head_rms_norm_tensor(metal_graph_batch_q(g), - draft, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_q(g), - draft, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_DSPARK_PROFILE_STAGE("q_path"); - - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_kv_raw(g), - dspark_model, - block->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - metal_graph_batch_attn_norm(g), - rows); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), - metal_graph_batch_kv_raw(g), - dspark_model->map, - dspark_model->size, - block->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - rows, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(kv_target_view, - 1, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(kv_draft_view, - draft, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), - rows, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; - if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor(g->dspark_raw_cache[stage], - metal_graph_batch_kv(g), - g->dspark_cache_cap, - append_pos, - rows, - DS4_N_HEAD_DIM) != 0; - DS4_DSPARK_PROFILE_STAGE("kv_path"); - - if (ok) ok = ds4_gpu_attention_noncausal_raw_batch_heads_tensor( - metal_graph_batch_heads(g), - dspark_model->map, - dspark_model->size, - block->attn_sinks->abs_offset, - metal_graph_batch_q(g), - g->dspark_raw_cache[stage], - draft, - visible_rows, - g->dspark_cache_cap, - attention_raw_start, - DS4_N_HEAD, - DS4_N_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_heads(g), - draft, - DS4_N_HEAD, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - true, - freq_base, - freq_scale, - ext_factor, - attn_factor, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_DSPARK_PROFILE_STAGE("attention"); - if (ok) ok = ds4_gpu_attention_output_q8_batch_tensor( - metal_graph_batch_attn_out(g), - metal_graph_batch_attn_low(g), - metal_graph_batch_group_tmp(g), - metal_graph_batch_low_tmp(g), - dspark_model->map, - dspark_model->size, - block->attn_output_a->abs_offset, - block->attn_output_b->abs_offset, - group_dim, - DS4_N_LORA_O, - DS4_N_OUT_GROUP, - DS4_N_EMBD, - metal_graph_batch_heads(g), - draft) != 0; - if (ok) ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, - metal_graph_batch_attn_out(g), - draft_hc_view, - draft_hc_split_view, - DS4_N_EMBD, - DS4_N_HC) != 0; - DS4_DSPARK_PROFILE_STAGE("attn_output_hc"); - - if (ok) ok = metal_graph_encode_layer_ffn_batch(g, - dspark_model, - block, - stage, - pos, - draft, - NULL, - 0); - DS4_DSPARK_PROFILE_STAGE("ffn"); - if (ok && - !prepare_next_stage_input && - getenv("DS4_DSPARK_DISABLE_FINAL_OUTPUT_ALIAS") != NULL) { - ok = ds4_gpu_tensor_copy(g->dspark_stage_output_hc, - 0, - metal_graph_batch_next_hc(g), - 0, - (uint64_t)draft * hc_dim * sizeof(float)) != 0; - } - DS4_DSPARK_PROFILE_STAGE("copy_output"); - if (ok && prepare_next_stage_input) { - ok = metal_graph_encode_dspark_next_stage_draft_input_from( - g, dw, metal_graph_batch_next_hc(g)); - } - DS4_DSPARK_PROFILE_STAGE("next_input"); - if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; - g->ssd_streaming = saved_streaming; - - ds4_gpu_tensor_free(after_attn_hc_view); - ds4_gpu_tensor_free(kv_draft_view); - ds4_gpu_tensor_free(kv_target_view); - ds4_gpu_tensor_free(draft_hc_split_view); - ds4_gpu_tensor_free(draft_hc_view); - ds4_gpu_tensor_free(draft_attn_norm_view); - ds4_gpu_tensor_free(attn_cur_view); - ds4_gpu_tensor_free(hc_split_view); - ds4_gpu_tensor_free(hc_mix_view); - if (!ok && !commands_open) (void)ds4_gpu_synchronize(); -#undef DS4_DSPARK_PROFILE_STAGE - return ok; -} - -static bool metal_graph_eval_dspark_stage_chain( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t pos, - uint32_t *completed_stages, - uint32_t *cache_start_out, - uint32_t *cache_rows_out) { - if (completed_stages) *completed_stages = 0; - if (cache_start_out) *cache_start_out = 0; - if (cache_rows_out) *cache_rows_out = 0; - if (!g || !dspark_model || !dw || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - !dspark_stage_input_ready(g, dw) || - !dspark_stage_cache_ready(g, dw) || - !metal_graph_prefill_tokens(g) || - !g->dspark_draft_tokens) { - return false; - } - - const uint32_t rows = dw->block_size + 1u; - const uint32_t support_len = g->dspark_cache_len; - const uint32_t raw_start = support_len ? g->dspark_cache_start : 0; - if (support_len > g->dspark_cache_cap || - rows > g->dspark_cache_cap - support_len || - (support_len != 0 && raw_start >= g->dspark_cache_cap) || - !metal_graph_dspark_cache_ends_at(g, pos)) { - return false; - } - if (cache_start_out) { - *cache_start_out = support_len ? raw_start : - (pos % g->dspark_cache_cap); - } - if (cache_rows_out) *cache_rows_out = support_len + rows; - - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - if (!dspark_stage_block_ready(g, dw, stage)) return false; - } - - /* The support model runs only on the coordinator. Its generic layer - * helpers share the base graph object, so temporarily disarm TP or they - * would encode expert gates that the worker can never reach. The base - * model's later verification restores and uses the normal 50/50 split. */ - const uint32_t saved_tp_world = g->tp_world; - const uint32_t saved_tp_batch_rows = g->tp_batch_rows; - g->tp_world = 0; - g->tp_batch_rows = 0; - const bool suspended_expert_sharding = saved_tp_world == 2; - if (suspended_expert_sharding) { - ds4_gpu_tp_suspend_expert_sharding(1); - } - - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) { - ok = ds4_gpu_tensor_copy(metal_graph_prefill_tokens(g), - 0, - g->dspark_draft_tokens, - 0, - (uint64_t)dw->block_size * sizeof(int32_t)) != 0; - } - for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { - const bool stage_ok = - metal_graph_eval_dspark_stage_block(g, - dspark_model, - dw, - stage, - pos, - support_len, - raw_start, - stage + 1u < dw->n_stages, - true); - if (!stage_ok) { - ok = false; - break; - } - if (completed_stages) *completed_stages = stage + 1u; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (suspended_expert_sharding) { - ds4_gpu_tp_suspend_expert_sharding(0); - } - g->tp_world = saved_tp_world; - g->tp_batch_rows = saved_tp_batch_rows; - if (!ok) { - (void)ds4_gpu_synchronize(); - return false; - } - return true; -} - -/* Keep the support KV ring aligned while the scheduler skips proposals. */ -static bool metal_graph_dspark_ring_maintain( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - uint32_t pos) { - if (!g || !dspark_model || !dw || - !g->dspark_capture_valid || - g->dspark_cache_len == 0 || - !metal_graph_dspark_cache_ends_at(g, pos) || - !dspark_stage0_weights_ready(g, dw) || - !dspark_stage_cache_ready(g, dw) || - !metal_graph_batch_kv_raw(g) || !metal_graph_batch_kv(g)) { - return false; - } - for (uint32_t stage = 0; stage < dw->n_stages; stage++) { - if (!dspark_stage_block_ready(g, dw, stage)) return false; - } - - const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; - const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; - ds4_gpu_tensor *kv_raw_view = - ds4_gpu_tensor_view(metal_graph_batch_kv_raw(g), - 0, - (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - ds4_gpu_tensor *kv_view = - ds4_gpu_tensor_view(metal_graph_batch_kv(g), - 0, - (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - bool ok = kv_raw_view && kv_view && ds4_gpu_begin_commands() != 0; - if (ok) { - ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, - dspark_model, - stage0->main_proj, - in_dim, - DS4_N_EMBD, - g->dspark_target_hidden, - 1); - } - if (ok) { - ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, - g->dspark_stage0_proj, - dspark_model->map, - dspark_model->size, - stage0->main_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - } - for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { - const ds4_layer_weights *block = &dw->stage[stage].block; - ok = metal_graph_matmul_plain_tensor(kv_raw_view, - dspark_model, - block->attn_kv, - DS4_N_EMBD, - DS4_N_HEAD_DIM, - g->dspark_main_x, - 1); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor( - kv_view, - kv_raw_view, - dspark_model->map, - dspark_model->size, - block->attn_kv_a_norm->abs_offset, - DS4_N_HEAD_DIM, - 1, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_rope_tail_tensor(kv_view, - 1, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - false, - DS4_ROPE_FREQ_BASE, - 1.0f, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(kv_view, - 1, - DS4_N_HEAD_DIM, - DS4_N_ROT) != 0; - if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor( - g->dspark_raw_cache[stage], - kv_view, - g->dspark_cache_cap, - pos, - 1, - DS4_N_HEAD_DIM) != 0; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - ds4_gpu_tensor_free(kv_view); - ds4_gpu_tensor_free(kv_raw_view); - if (ok) (void)metal_graph_dspark_cache_claim_appended_row(g, pos); - return ok; -} - -static ds4_gpu_tensor *metal_graph_dspark_final_output_hc(const ds4_gpu_graph *g) { - if (!g) return NULL; - if (getenv("DS4_DSPARK_DISABLE_FINAL_OUTPUT_ALIAS") == NULL && - metal_graph_batch_next_hc(g)) { - return metal_graph_batch_next_hc(g); - } - return g->dspark_stage_output_hc; -} - -static bool dspark_final_head_ready( - const ds4_gpu_graph *g, - const ds4_weights *base_weights, - const ds4_dspark_weights *dw) { - if (!g || !base_weights || !dw || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - !base_weights->output || - !metal_graph_dspark_final_output_hc(g) || - !metal_graph_batch_hc_mix(g) || - !metal_graph_batch_hc_split(g) || - !metal_graph_batch_flat_hc(g) || - !metal_graph_batch_ffn_cur(g) || - !metal_graph_batch_ffn_norm(g) || - !g->spec_logits) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t draft = dw->block_size; - const uint64_t vocab_dim = base_weights->output->dim[1]; - if (!final->norm || - !final->hc_head_base || - !final->hc_head_fn || - !final->hc_head_scale || - final->norm->type != DS4_TENSOR_F32 || - final->hc_head_base->type != DS4_TENSOR_F32 || - !dspark_tensor_type_matches(final->hc_head_fn->type, - DS4_DSPARK_LAYOUT_PLAIN) || - final->hc_head_scale->type != DS4_TENSOR_F32 || - !tensor_type_is_dense_quant(base_weights->output->type)) { - return false; - } - - return final->norm->ndim == 1 && - final->norm->dim[0] == DS4_N_EMBD && - final->hc_head_base->ndim == 1 && - final->hc_head_base->dim[0] == DS4_N_HC && - final->hc_head_fn->ndim == 2 && - final->hc_head_fn->dim[0] == hc_dim && - final->hc_head_fn->dim[1] == DS4_N_HC && - final->hc_head_scale->ndim == 1 && - final->hc_head_scale->dim[0] == 1 && - base_weights->output->ndim == 2 && - base_weights->output->dim[0] == DS4_N_EMBD && - vocab_dim == DS4_N_VOCAB && - ds4_gpu_tensor_bytes( - metal_graph_dspark_final_output_hc(g)) >= - draft * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= - draft * hc_dim * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= - draft * DS4_N_HC * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= - draft * DS4_N_HC * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) >= - draft * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) >= - draft * DS4_N_EMBD * sizeof(float) && - ds4_gpu_tensor_bytes(g->spec_logits) >= - draft * vocab_dim * sizeof(float); -} - -static bool metal_graph_eval_dspark_base_logits( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw) { - if (!g || !base_model || !base_weights || !dspark_model || !dw || - !dspark_final_head_ready(g, base_weights, dw)) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - const uint32_t draft = dw->block_size; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t vocab_dim = base_weights->output->dim[1]; - ds4_gpu_tensor *stage_output_hc = metal_graph_dspark_final_output_hc(g); - ds4_gpu_tensor *output_pre = - ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), - 0, - (uint64_t)draft * DS4_N_HC * sizeof(float)); - ds4_gpu_tensor *output_weights = - ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - 0, - (uint64_t)draft * DS4_N_HC * sizeof(float)); - ds4_gpu_tensor *output_embd = - ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), - 0, - (uint64_t)draft * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *output_norm = - ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), - 0, - (uint64_t)draft * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *logits = - ds4_gpu_tensor_view(g->spec_logits, - 0, - (uint64_t)draft * vocab_dim * sizeof(float)); - - bool ok = stage_output_hc && output_pre && output_weights && output_embd && - output_norm && logits; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), - stage_output_hc, - (uint32_t)hc_dim, - draft, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(output_pre, - dspark_model, - final->hc_head_fn, - hc_dim, - DS4_N_HC, - metal_graph_batch_flat_hc(g), - draft); - if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, - output_pre, - dspark_model->map, - dspark_model->size, - final->hc_head_scale->abs_offset, - final->hc_head_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, - stage_output_hc, - output_weights, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, - output_embd, - dspark_model->map, - dspark_model->size, - final->norm->abs_offset, - DS4_N_EMBD, - draft, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(logits, - base_model, - base_weights->output, - DS4_N_EMBD, - vocab_dim, - output_norm, - draft); - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - - ds4_gpu_tensor_free(logits); - ds4_gpu_tensor_free(output_norm); - ds4_gpu_tensor_free(output_embd); - ds4_gpu_tensor_free(output_weights); - ds4_gpu_tensor_free(output_pre); - return ok; -} - -static bool metal_graph_eval_dspark_final_hidden( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw) { - if (!g || !dspark_model || !dw || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - !metal_graph_dspark_final_output_hc(g) || - !metal_graph_batch_hc_mix(g) || - !metal_graph_batch_hc_split(g) || - !metal_graph_batch_flat_hc(g) || - !metal_graph_batch_ffn_cur(g) || - !metal_graph_batch_ffn_norm(g)) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - const uint32_t draft = dw->block_size; - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - if (!final->norm || - !final->hc_head_base || - !final->hc_head_fn || - !final->hc_head_scale || - final->norm->type != DS4_TENSOR_F32 || - final->hc_head_base->type != DS4_TENSOR_F32 || - !dspark_tensor_type_matches(final->hc_head_fn->type, - DS4_DSPARK_LAYOUT_PLAIN) || - final->hc_head_scale->type != DS4_TENSOR_F32 || - final->norm->ndim != 1 || - final->norm->dim[0] != DS4_N_EMBD || - final->hc_head_base->ndim != 1 || - final->hc_head_base->dim[0] != DS4_N_HC || - final->hc_head_fn->ndim != 2 || - final->hc_head_fn->dim[0] != hc_dim || - final->hc_head_fn->dim[1] != DS4_N_HC || - final->hc_head_scale->ndim != 1 || - final->hc_head_scale->dim[0] != 1 || - ds4_gpu_tensor_bytes(metal_graph_dspark_final_output_hc(g)) < - (uint64_t)draft * hc_dim * sizeof(float) || - ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) < - (uint64_t)draft * hc_dim * sizeof(float) || - ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) < - (uint64_t)draft * DS4_N_HC * sizeof(float) || - ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) < - (uint64_t)draft * DS4_N_HC * sizeof(float) || - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) < - (uint64_t)draft * DS4_N_EMBD * sizeof(float) || - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < - (uint64_t)draft * DS4_N_EMBD * sizeof(float)) { - return false; - } - - ds4_gpu_tensor *output_pre = - ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), - 0, - (uint64_t)draft * DS4_N_HC * sizeof(float)); - ds4_gpu_tensor *output_weights = - ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), - 0, - (uint64_t)draft * DS4_N_HC * sizeof(float)); - ds4_gpu_tensor *output_embd = - ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), - 0, - (uint64_t)draft * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *output_norm = - ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), - 0, - (uint64_t)draft * DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *stage_output_hc = metal_graph_dspark_final_output_hc(g); - - bool ok = stage_output_hc && output_pre && output_weights && - output_embd && output_norm; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), - stage_output_hc, - (uint32_t)hc_dim, - draft, - DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(output_pre, - dspark_model, - final->hc_head_fn, - hc_dim, - DS4_N_HC, - metal_graph_batch_flat_hc(g), - draft); - if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, - output_pre, - dspark_model->map, - dspark_model->size, - final->hc_head_scale->abs_offset, - final->hc_head_base->abs_offset, - DS4_N_HC, - DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, - stage_output_hc, - output_weights, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, - output_embd, - dspark_model->map, - dspark_model->size, - final->norm->abs_offset, - DS4_N_EMBD, - draft, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - - ds4_gpu_tensor_free(output_norm); - ds4_gpu_tensor_free(output_embd); - ds4_gpu_tensor_free(output_weights); - ds4_gpu_tensor_free(output_pre); - return ok; -} - -static bool metal_graph_eval_dspark_base_logits_from_hidden( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_dspark_weights *dw) { - if (!g || !base_model || !base_weights || !dw || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - !base_weights->output || - !tensor_type_is_dense_quant(base_weights->output->type) || - base_weights->output->ndim != 2 || - base_weights->output->dim[0] != DS4_N_EMBD || - base_weights->output->dim[1] != DS4_N_VOCAB || - !metal_graph_batch_ffn_norm(g) || - !g->spec_logits || - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < - (uint64_t)dw->block_size * DS4_N_EMBD * sizeof(float) || - ds4_gpu_tensor_bytes(g->spec_logits) < - (uint64_t)dw->block_size * DS4_N_VOCAB * sizeof(float)) { - return false; - } - - ds4_gpu_tensor *output_norm = - ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), - 0, - (uint64_t)dw->block_size * - DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *logits = - ds4_gpu_tensor_view(g->spec_logits, - 0, - (uint64_t)dw->block_size * - DS4_N_VOCAB * sizeof(float)); - bool ok = output_norm && logits; - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(logits, - base_model, - base_weights->output, - DS4_N_EMBD, - DS4_N_VOCAB, - output_norm, - dw->block_size); - if (ok) ok = ds4_gpu_end_commands() != 0; - if (!ok) (void)ds4_gpu_synchronize(); - - ds4_gpu_tensor_free(logits); - ds4_gpu_tensor_free(output_norm); - return ok; -} - -static bool dspark_markov_probe_ready( - const ds4_dspark_weights *dw) { - if (!dw || - dw->n_stages == 0 || - dw->n_stages > DS4_DSPARK_MAX_STAGES || - dw->block_size == 0 || - dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || - dw->markov_rank == 0) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - if (!final->markov_w1 || - !final->markov_w2 || - !dspark_tensor_type_matches(final->markov_w1->type, - DS4_DSPARK_LAYOUT_DENSE) || - !dspark_tensor_type_matches(final->markov_w2->type, - DS4_DSPARK_LAYOUT_DENSE)) { - return false; - } - - return final->markov_w1->ndim == 2 && - final->markov_w1->dim[0] == dw->markov_rank && - final->markov_w1->dim[1] == DS4_N_VOCAB && - final->markov_w2->ndim == 2 && - final->markov_w2->dim[0] == dw->markov_rank && - final->markov_w2->dim[1] == DS4_N_VOCAB; -} - -static bool dspark_dense_row_to_f32( - float *out, - const ds4_model *model, - const ds4_tensor *t, - uint32_t row) { - if (!out || !model || !t || t->ndim != 2 || row >= t->dim[1]) { - return false; - } - - const uint64_t width = t->dim[0]; - if (t->type == DS4_TENSOR_F32) { - const float *base = tensor_data(model, t); - memcpy(out, base + (uint64_t)row * width, width * sizeof(out[0])); - return true; - } - if (t->type == DS4_TENSOR_F16) { - const uint16_t *base = tensor_data(model, t); - const uint16_t *src = base + (uint64_t)row * width; - for (uint64_t i = 0; i < width; i++) out[i] = f16_to_f32(src[i]); - return true; - } - if (t->type == DS4_TENSOR_Q8_0) { - const uint64_t blocks = (width + 31u) / 32u; - const uint8_t *src = - (const uint8_t *)tensor_data(model, t) + - (uint64_t)row * blocks * 34u; - for (uint64_t b = 0; b < blocks; b++) { - uint16_t scale_bits; - memcpy(&scale_bits, src + b * 34u, sizeof(scale_bits)); - const float scale = f16_to_f32(scale_bits); - const int8_t *qs = (const int8_t *)(src + b * 34u + 2u); - const uint64_t i0 = b * 32u; - const uint64_t n = width - i0 < 32u ? width - i0 : 32u; - for (uint64_t i = 0; i < n; i++) { - out[i0 + i] = scale * (float)qs[i]; - } - } - return true; - } - return false; -} - -static uint32_t dspark_argmax_f32(const float *x, uint32_t n) { - uint32_t best = 0; - float best_v = x[0]; - for (uint32_t i = 1; i < n; i++) { - if (x[i] > best_v) { - best_v = x[i]; - best = i; - } - } - return best; -} - -typedef struct { - const uint8_t *data; - const int8_t *xq; - const float *xscale; - const float *logits; - uint64_t in_dim; - uint64_t blocks; - uint64_t rows_per_slot; - uint32_t best_idx[DS4_MAX_THREADS]; - float best_val[DS4_MAX_THREADS]; -} dspark_markov_q8_0_argmax_ctx; - -static void dspark_markov_q8_0_argmax_worker( - void *vctx, - uint64_t row0, - uint64_t row1) { - dspark_markov_q8_0_argmax_ctx *ctx = vctx; - uint64_t slot = ctx->rows_per_slot ? row0 / ctx->rows_per_slot : 0; - if (slot >= DS4_MAX_THREADS) slot = DS4_MAX_THREADS - 1u; - - float best_v = -FLT_MAX; - uint32_t best = (uint32_t)row0; - for (uint64_t row = row0; row < row1; row++) { - const uint8_t *wrow = ctx->data + row * ctx->blocks * 34u; - const float score = - ctx->logits[row] + - dot_q8_0_row(wrow, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks); - if (score > best_v) { - best_v = score; - best = (uint32_t)row; - } - } - - ctx->best_idx[slot] = best; - ctx->best_val[slot] = best_v; -} - -static bool dspark_markov_q8_0_argmax( - uint32_t *token_out, - const ds4_model *model, - const ds4_tensor *w, - const float *state, - const float *logits) { - if (!token_out || - !model || - !w || - !state || - !logits || - w->type != DS4_TENSOR_Q8_0 || - w->ndim != 2 || - w->dim[1] > UINT32_MAX) { - return false; - } - - const uint64_t in_dim = w->dim[0]; - const uint64_t out_dim = w->dim[1]; - const uint64_t blocks = (in_dim + 31u) / 32u; - if (out_dim == 0 || - blocks == 0 || - blocks > (uint64_t)SIZE_MAX / 32u || - blocks > (uint64_t)SIZE_MAX / sizeof(float)) { - return false; - } - - enum { DSPARK_MARKOV_ARGMAX_STACK_BLOCKS = 32 }; - int8_t xq_stack[DSPARK_MARKOV_ARGMAX_STACK_BLOCKS * 32u]; - float xscale_stack[DSPARK_MARKOV_ARGMAX_STACK_BLOCKS]; - const bool use_stack = blocks <= DSPARK_MARKOV_ARGMAX_STACK_BLOCKS; - int8_t *xq = use_stack ? xq_stack : xmalloc((size_t)blocks * 32u); - float *xscale = use_stack ? xscale_stack : - xmalloc((size_t)blocks * sizeof(xscale[0])); - quantize_q8_0_activation(state, xq, xscale, in_dim); - - ds4_threads_init(); - const uint32_t n_slots = - g_pool.n_threads == 0 ? 1u : g_pool.n_threads; - const uint64_t rows_per_slot = (out_dim + n_slots - 1u) / n_slots; - dspark_markov_q8_0_argmax_ctx ctx = { - .data = tensor_data(model, w), - .xq = xq, - .xscale = xscale, - .logits = logits, - .in_dim = in_dim, - .blocks = blocks, - .rows_per_slot = rows_per_slot, - }; - for (uint32_t i = 0; i < DS4_MAX_THREADS; i++) { - ctx.best_idx[i] = 0; - ctx.best_val[i] = -FLT_MAX; - } - - ds4_parallel_for(out_dim, dspark_markov_q8_0_argmax_worker, &ctx); - - uint32_t best = 0; - float best_v = -FLT_MAX; - for (uint32_t slot = 0; slot < n_slots && slot < DS4_MAX_THREADS; slot++) { - const uint64_t row0 = (uint64_t)slot * rows_per_slot; - if (row0 >= out_dim) break; - if (ctx.best_val[slot] > best_v) { - best_v = ctx.best_val[slot]; - best = ctx.best_idx[slot]; - } - } - - if (!use_stack) { - free(xscale); - free(xq); - } - *token_out = best; - return true; -} - -/* Exact target verification preserves correctness when this diagnostic mode - * proposes directly from the support model's base logits. */ -static bool dspark_markov_bias_disabled(void) { - static int cached = -1; - if (cached < 0) { - const char *env = getenv("DS4_DSPARK_NO_MARKOV"); - cached = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; - } - return cached == 1; -} - -static bool dspark_disable_fused_cpu_markov_argmax(void) { - static int cache = -1; - if (cache < 0) { - const char *env = getenv("DS4_DSPARK_DISABLE_FUSED_CPU_MARKOV_ARGMAX"); - cache = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; - } - return cache != 0; -} - -static bool dspark_disable_reuse_confidence0_markov(void) { - static int cache = -1; - if (cache < 0) { - const char *env = getenv("DS4_DSPARK_DISABLE_REUSE_CONFIDENCE0_MARKOV"); - cache = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; - } - return cache != 0; -} - -static bool dspark_apply_markov_greedy_probe( - float *logits, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - int first_prev_token, - float *markov_state, - float *markov_bias, - int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], - uint32_t *proposal_len) { - if (proposal_len) *proposal_len = 0; - if (!logits || - !dspark_model || - !dw || - !markov_state || - !markov_bias || - !proposal || - first_prev_token < 0 || - (uint32_t)first_prev_token >= DS4_N_VOCAB || - !dspark_markov_probe_ready(dw)) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - const bool no_bias = dspark_markov_bias_disabled(); - int32_t prev_token = first_prev_token; - for (uint32_t draft = 0; draft < dw->block_size; draft++) { - float *row = logits + (uint64_t)draft * DS4_N_VOCAB; - if (!no_bias) { - if (!dspark_dense_row_to_f32(markov_state, - dspark_model, - final->markov_w1, - (uint32_t)prev_token)) { - return false; - } - matvec_any(markov_bias, dspark_model, final->markov_w2, markov_state); - for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { - row[i] += markov_bias[i]; - } - } - const uint32_t token = dspark_argmax_f32(row, DS4_N_VOCAB); - proposal[draft] = (int32_t)token; - prev_token = (int32_t)token; - } - - if (proposal_len) *proposal_len = dw->block_size; - return true; -} - -static bool dspark_confidence_probe_ready( - const ds4_dspark_weights *dw) { - if (!dspark_markov_probe_ready(dw)) return false; - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - if (!final->confidence_proj || - !dspark_tensor_type_matches(final->confidence_proj->type, - DS4_DSPARK_LAYOUT_DENSE)) { - return false; - } - return final->confidence_proj->ndim == 2 && - final->confidence_proj->dim[0] == - (uint64_t)DS4_N_EMBD + dw->markov_rank && - final->confidence_proj->dim[1] == 1; -} - -static bool dspark_eval_confidence_probe( - float *confidence_logits, - const float *hidden_rows, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - int first_prev_token, - const int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], - float *markov_state, - float *features, - uint32_t *confidence_len) { - if (confidence_len) *confidence_len = 0; - if (!confidence_logits || - !hidden_rows || - !dspark_model || - !dw || - !proposal || - !markov_state || - !features || - first_prev_token < 0 || - (uint32_t)first_prev_token >= DS4_N_VOCAB || - !dspark_confidence_probe_ready(dw)) { - return false; - } - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - int32_t prev_token = first_prev_token; - for (uint32_t draft = 0; draft < dw->block_size; draft++) { - if (prev_token < 0 || (uint32_t)prev_token >= DS4_N_VOCAB) { - return false; - } - if (!dspark_dense_row_to_f32(markov_state, - dspark_model, - final->markov_w1, - (uint32_t)prev_token)) { - return false; - } - memcpy(features, - hidden_rows + (uint64_t)draft * DS4_N_EMBD, - (uint64_t)DS4_N_EMBD * sizeof(features[0])); - memcpy(features + DS4_N_EMBD, - markov_state, - (uint64_t)dw->markov_rank * sizeof(features[0])); - matvec_any(confidence_logits + draft, - dspark_model, - final->confidence_proj, - features); - prev_token = proposal[draft]; - } - - if (confidence_len) *confidence_len = dw->block_size; - return true; -} - -static bool dspark_apply_markov_confidence_lazy_runtime( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - int first_prev_token, - float confidence_threshold, - float *logits, - float *markov_bias, - float *features, - size_t features_cap, - int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], - uint32_t *proposal_len, - uint32_t *confidence_len, - uint32_t *confidence_prefix_len, - bool reuse_first_confidence, - float *confidence0) { - if (proposal_len) *proposal_len = 0; - if (confidence_len) *confidence_len = 0; - if (confidence_prefix_len) *confidence_prefix_len = 0; - if (confidence0 && !reuse_first_confidence) *confidence0 = 0.0f; - if (!g || - !g->spec_logits || - !metal_graph_batch_ffn_norm(g) || - !dspark_model || - !dw || - !logits || - !markov_bias || - !features || - !proposal || - confidence_threshold <= 0.0f || - first_prev_token < 0 || - (uint32_t)first_prev_token >= DS4_N_VOCAB || - (reuse_first_confidence && !confidence0) || - !dspark_markov_probe_ready(dw) || - !dspark_confidence_probe_ready(dw)) { - return false; - } - - const uint64_t logits_bytes = - (uint64_t)DS4_N_VOCAB * sizeof(float); - const uint64_t hidden_bytes = - (uint64_t)DS4_N_EMBD * sizeof(float); - const uint64_t feature_count = - (uint64_t)DS4_N_EMBD + (uint64_t)dw->markov_rank; - if (feature_count > features_cap) return false; - float *markov_state = features + DS4_N_EMBD; - bool ok = true; - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - int32_t prev_token = first_prev_token; - uint32_t produced = 0; - uint32_t confident = 0; - for (uint32_t draft = 0; ok && draft < dw->block_size; draft++) { - if (prev_token < 0 || (uint32_t)prev_token >= DS4_N_VOCAB) { - ok = false; - break; - } - float confidence_logit = 0.0f; - if (draft == 0 && reuse_first_confidence) { - confidence_logit = *confidence0; - } else { - ok = dspark_dense_row_to_f32(markov_state, - dspark_model, - final->markov_w1, - (uint32_t)prev_token); - if (!ok) break; - - ok = ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), - (uint64_t)draft * hidden_bytes, - features, - hidden_bytes) != 0; - if (!ok) break; - matvec_any(&confidence_logit, - dspark_model, - final->confidence_proj, - features); - } - if (draft == 0 && confidence0) *confidence0 = confidence_logit; - if (confidence_len) *confidence_len = draft + 1u; - if (sigmoid_stable(confidence_logit) < confidence_threshold) { - ok = true; - break; - } - - int32_t token = -1; -#ifndef __APPLE__ - /* CUDA can apply the Markov bias and argmax without reading back the - * full logits row. Metal currently falls through to the CPU path. */ - if (ok && !dspark_markov_bias_disabled() && - getenv("DS4_DSPARK_NO_GPU_MARKOV") == NULL && - g->dspark_draft_tokens && - dw->markov_rank != 0 && (dw->markov_rank & 31u) == 0 && - final->markov_w1->type == DS4_TENSOR_Q8_0 && - final->markov_w2->type == DS4_TENSOR_Q8_0) { - ds4_gpu_tensor *row_view = - ds4_gpu_tensor_view(g->spec_logits, - (uint64_t)draft * logits_bytes, - logits_bytes); - uint64_t gpu_key = 0; - bool gpu_ok = row_view && - ds4_gpu_dspark_markov_argmax_tensor( - g->dspark_draft_tokens, - row_view, - dspark_model->map, - dspark_model->size, - final->markov_w1->abs_offset, - final->markov_w2->abs_offset, - (uint32_t)prev_token, - DS4_N_VOCAB, - dw->markov_rank) != 0 && - ds4_gpu_tensor_read(g->dspark_draft_tokens, - 0, - &gpu_key, - sizeof(gpu_key)) != 0; - ds4_gpu_tensor_free(row_view); - const uint32_t gpu_token = ~(uint32_t)(gpu_key & 0xffffffffu); - if (gpu_ok && gpu_key != 0 && gpu_token < DS4_N_VOCAB) { - token = (int32_t)gpu_token; - proposal[draft] = token; - produced = draft + 1u; - confident = produced; - prev_token = token; - continue; - } - } -#endif - if (ok) { - ok = ds4_gpu_tensor_read(g->spec_logits, - (uint64_t)draft * logits_bytes, - logits, - logits_bytes) != 0; - if (ok) { - uint32_t fused_token = 0; - if (dspark_markov_bias_disabled()) { - token = (int32_t)dspark_argmax_f32(logits, DS4_N_VOCAB); - } else if (!dspark_disable_fused_cpu_markov_argmax() && - dspark_markov_q8_0_argmax(&fused_token, - dspark_model, - final->markov_w2, - markov_state, - logits)) { - token = (int32_t)fused_token; - } else { - matvec_any(markov_bias, dspark_model, final->markov_w2, markov_state); - for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { - logits[i] += markov_bias[i]; - } - token = (int32_t)dspark_argmax_f32(logits, DS4_N_VOCAB); - } - } - } - if (!ok || token < 0 || (uint32_t)token >= DS4_N_VOCAB) { - ok = false; - break; - } - proposal[draft] = token; - produced = draft + 1u; - confident = produced; - prev_token = token; - } - - if (ok) { - if (proposal_len) *proposal_len = produced; - if (confidence_prefix_len) *confidence_prefix_len = confident; - } - return ok; -} - -static bool dspark_eval_confidence0_runtime( - ds4_gpu_graph *g, - const ds4_model *dspark_model, - const ds4_dspark_weights *dw, - int first_prev_token, - float *features, - size_t features_cap, - float *confidence0) { - if (confidence0) *confidence0 = 0.0f; - if (!confidence0 || - !g || - !metal_graph_batch_ffn_norm(g) || - !dspark_model || - !dw || - !features || - first_prev_token < 0 || - (uint32_t)first_prev_token >= DS4_N_VOCAB || - !dspark_confidence_probe_ready(dw)) { - return false; - } - - const uint64_t hidden_bytes = - (uint64_t)DS4_N_EMBD * sizeof(float); - const uint64_t feature_count = - (uint64_t)DS4_N_EMBD + (uint64_t)dw->markov_rank; - if (feature_count > features_cap || - ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < hidden_bytes) { - return false; - } - - float *markov_state = features + DS4_N_EMBD; - bool ok = true; - - const ds4_dspark_stage_weights *final = - &dw->stage[dw->n_stages - 1u]; - if (ok) { - ok = dspark_dense_row_to_f32(markov_state, - dspark_model, - final->markov_w1, - (uint32_t)first_prev_token); - } - if (ok) { - ok = ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), - 0, - features, - hidden_bytes) != 0; - } - if (ok) { - matvec_any(confidence0, dspark_model, final->confidence_proj, features); - } - - return ok; -} - -static uint32_t dspark_confident_prefix_len( - const float *confidence_logits, - uint32_t confidence_len, - float threshold) { - if (!confidence_logits || confidence_len == 0 || threshold <= 0.0f) { - return confidence_len; - } - for (uint32_t i = 0; i < confidence_len; i++) { - if (sigmoid_stable(confidence_logits[i]) < threshold) return i; - } - return confidence_len; -} - -static bool metal_graph_eval_mtp_draft_from_hc( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_model *mtp_model, - const ds4_mtp_weights *mtp, - ds4_gpu_tensor *prev_hc, - ds4_gpu_tensor *out_hc, - int token, - uint32_t pos, - float *logits, - int *top_id) { - if (!mtp || !mtp->block.attn_q_a || !g->mtp_raw_cache || !prev_hc || !out_hc) return false; - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint32_t raw_row = pos % g->raw_cap; - uint32_t n_raw = g->mtp_n_raw + 1u; - if (n_raw > g->raw_window) n_raw = g->raw_window; - if (n_raw > g->raw_cap) n_raw = g->raw_cap; - - ds4_gpu_tensor *saved_cur = metal_graph_cur_hc(g); - ds4_gpu_tensor *saved_after = metal_graph_after_ffn_hc(g); - const uint32_t saved_tp_world = g->tp_world; - const uint32_t saved_tp_batch_rows = g->tp_batch_rows; - g->tp_world = 0; - g->tp_batch_rows = 0; - const bool suspended_expert_sharding = saved_tp_world == 2; - if (suspended_expert_sharding) { - ds4_gpu_tp_suspend_expert_sharding(1); - } - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(g->mtp_embed, - base_model->map, - base_model->size, - base_weights->token_embd->abs_offset, - (uint32_t)base_weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - 1) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->mtp_enorm, - g->mtp_embed, - mtp_model->map, - mtp_model->size, - mtp->enorm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_eproj, - mtp_model->map, - mtp_model->size, - mtp->e_proj->abs_offset, - DS4_N_EMBD, - DS4_N_EMBD, - g->mtp_enorm, - 1) != 0; - if (ok) ok = ds4_gpu_repeat_hc_tensor(g->mtp_eproj_hc, - g->mtp_eproj, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->mtp_hnorm_hc, - prev_hc, - mtp_model->map, - mtp_model->size, - mtp->hnorm->abs_offset, - DS4_N_EMBD, - DS4_N_HC, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_hproj_hc, - mtp_model->map, - mtp_model->size, - mtp->h_proj->abs_offset, - DS4_N_EMBD, - DS4_N_EMBD, - g->mtp_hnorm_hc, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_add_tensor(g->mtp_input_hc, - g->mtp_eproj_hc, - g->mtp_hproj_hc, - (uint32_t)hc_dim) != 0; - if (ok) { - g->cur_hc_by_tier[g->active_tier] = g->mtp_input_hc; - g->after_ffn_hc_by_tier[g->active_tier] = out_hc; - ok = metal_graph_encode_decode_layer(g, - mtp_model, - &mtp->block, - 1, - pos, - g->mtp_raw_cache, - g->raw_cap, - raw_row, - n_raw, - token); - } - if (ok) g->cur_hc_by_tier[g->active_tier] = out_hc; - if (ok) ok = metal_graph_encode_output_head_mtp(g, - base_model, - base_weights, - mtp_model, - mtp, - base_weights->output->dim[1]); - if (ok && top_id) { - ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), - metal_graph_logits(g), - DS4_N_VOCAB) != 0; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - if (suspended_expert_sharding) { - ds4_gpu_tp_suspend_expert_sharding(0); - } - g->cur_hc_by_tier[g->active_tier] = saved_cur; - g->after_ffn_hc_by_tier[g->active_tier] = saved_after; - - if (ok && logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - if (ok && top_id) { - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top_id, sizeof(*top_id)) != 0; - } - if (ok && g->mtp_n_raw < g->raw_window) g->mtp_n_raw++; - g->tp_world = saved_tp_world; - g->tp_batch_rows = saved_tp_batch_rows; - if (!ok) { - (void)ds4_gpu_synchronize(); - g->cur_hc_by_tier[g->active_tier] = saved_cur; - g->after_ffn_hc_by_tier[g->active_tier] = saved_after; - } - return ok; -} - -static bool metal_graph_eval_mtp_draft( - ds4_gpu_graph *g, - const ds4_model *base_model, - const ds4_weights *base_weights, - const ds4_model *mtp_model, - const ds4_mtp_weights *mtp, - int token, - uint32_t pos, - float *logits, - int *top_id) { - return metal_graph_eval_mtp_draft_from_hc(g, - base_model, - base_weights, - mtp_model, - mtp, - metal_graph_cur_hc(g), - g->mtp_state_hc, - token, - pos, - logits, - top_id); -} - -/* ========================================================================= - * Imatrix Collection. - * ========================================================================= - * - * The 2-bit DS4 quants care most about routed MoE experts. For expert gate - * and up matrices the matmul input is the FFN-normalized activation row. For - * expert down matrices the matmul input is the routed SwiGLU row after route - * weighting. During Metal prefill those tensors are already materialized as - * `batch_ffn_norm`, `batch_router_selected`, and `batch_routed_mid`, so the - * collector observes the exact release graph without changing inference math. - * - * The output is llama.cpp's legacy imatrix `.dat` format. Entries are packed - * by expert: one tensor entry contains `n_expert * n_columns` floats and the - * quantizer slices the vector for each expert. - */ -typedef struct { - float *gate_up_sum2; /* [active layer][active expert][hidden] */ - float *down_sum2; /* [active layer][active expert][expert FFN] */ - uint32_t gate_up_count[DS4_MAX_LAYER][DS4_MAX_EXPERT]; - uint32_t down_count[DS4_MAX_LAYER][DS4_MAX_EXPERT]; - float *ffn_norm_buf; - float *routed_mid_buf; - uint16_t *routed_mid_f16_buf; - int *selected_buf; - float *sq_tmp; - uint32_t cap_tokens; - uint64_t observed_tokens; - uint64_t observed_routes; - uint32_t chunks; - const char *dataset_path; -} ds4_imatrix_collector; - -static bool imatrix_collector_init(ds4_imatrix_collector *c, uint32_t cap_tokens, const char *dataset_path) { - memset(c, 0, sizeof(*c)); - c->cap_tokens = cap_tokens ? cap_tokens : 1u; - c->dataset_path = dataset_path; - const size_t gate_n = (size_t)DS4_N_LAYER * DS4_N_EXPERT * DS4_N_EMBD; - const size_t down_n = (size_t)DS4_N_LAYER * DS4_N_EXPERT * DS4_N_FF_EXP; - c->gate_up_sum2 = xcalloc(gate_n, sizeof(c->gate_up_sum2[0])); - c->down_sum2 = xcalloc(down_n, sizeof(c->down_sum2[0])); - c->ffn_norm_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EMBD * sizeof(c->ffn_norm_buf[0])); - c->routed_mid_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(c->routed_mid_buf[0])); - c->routed_mid_f16_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(c->routed_mid_f16_buf[0])); - c->selected_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * sizeof(c->selected_buf[0])); - c->sq_tmp = xmalloc((size_t)DS4_N_EMBD * sizeof(c->sq_tmp[0])); - return c->gate_up_sum2 && c->down_sum2 && c->ffn_norm_buf && - c->routed_mid_buf && c->routed_mid_f16_buf && c->selected_buf && c->sq_tmp; -} - -static void imatrix_collector_free(ds4_imatrix_collector *c) { - if (!c) return; - free(c->gate_up_sum2); - free(c->down_sum2); - free(c->ffn_norm_buf); - free(c->routed_mid_buf); - free(c->routed_mid_f16_buf); - free(c->selected_buf); - free(c->sq_tmp); - memset(c, 0, sizeof(*c)); -} - -static float *imatrix_gate_up_ptr(ds4_imatrix_collector *c, uint32_t il, uint32_t expert) { - return c->gate_up_sum2 + ((size_t)il * DS4_N_EXPERT + expert) * DS4_N_EMBD; -} - -static float *imatrix_down_ptr(ds4_imatrix_collector *c, uint32_t il, uint32_t expert) { - return c->down_sum2 + ((size_t)il * DS4_N_EXPERT + expert) * DS4_N_FF_EXP; -} - -static bool imatrix_collect_layer_batch( - ds4_imatrix_collector *c, - ds4_gpu_graph *g, - uint32_t il, - uint32_t n_tokens) { - if (!c || n_tokens == 0) return true; - if (n_tokens > c->cap_tokens) return false; - - const uint64_t norm_bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); - const uint64_t mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP; - const uint64_t mid_bytes = mid_elems * (g->batch_routed_mid_is_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t sel_bytes = (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int); - void *mid_dst = g->batch_routed_mid_is_f16 - ? (void *)c->routed_mid_f16_buf - : (void *)c->routed_mid_buf; - if (ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), 0, c->ffn_norm_buf, norm_bytes) == 0 || - ds4_gpu_tensor_read(metal_graph_batch_routed_mid(g), 0, mid_dst, mid_bytes) == 0 || - ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), 0, c->selected_buf, sel_bytes) == 0) - { - return false; - } - - for (uint32_t t = 0; t < n_tokens; t++) { - const float *x = c->ffn_norm_buf + (size_t)t * DS4_N_EMBD; - for (uint32_t i = 0; i < DS4_N_EMBD; i++) c->sq_tmp[i] = x[i] * x[i]; - - for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { - const int expert = c->selected_buf[(size_t)t * DS4_N_EXPERT_USED + slot]; - if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) continue; - - float *gate_up = imatrix_gate_up_ptr(c, il, (uint32_t)expert); - for (uint32_t i = 0; i < DS4_N_EMBD; i++) gate_up[i] += c->sq_tmp[i]; - c->gate_up_count[il][expert]++; - - float *down = imatrix_down_ptr(c, il, (uint32_t)expert); - const size_t mid_off = ((size_t)t * DS4_N_EXPERT_USED + slot) * DS4_N_FF_EXP; - if (g->batch_routed_mid_is_f16) { - const uint16_t *mid = c->routed_mid_f16_buf + mid_off; - for (uint32_t i = 0; i < DS4_N_FF_EXP; i++) { - const float v = f16_to_f32(mid[i]); - down[i] += v * v; - } - } else { - const float *mid = c->routed_mid_buf + mid_off; - for (uint32_t i = 0; i < DS4_N_FF_EXP; i++) down[i] += mid[i] * mid[i]; - } - c->down_count[il][expert]++; - c->observed_routes++; - } - } - c->observed_tokens += n_tokens; - c->chunks++; - return true; -} - -static void imatrix_write_i32(FILE *fp, int32_t v) { - if (fwrite(&v, sizeof(v), 1, fp) != 1) ds4_die("failed to write imatrix"); -} - -static void imatrix_write_entry( - FILE *fp, - const char *name, - const float *sum2, - const uint32_t *counts, - uint32_t n_expert, - uint32_t n_col) { - const int32_t len = (int32_t)strlen(name); - const int32_t ncall = 1; - const int32_t nval = (int32_t)((uint64_t)n_expert * n_col); - imatrix_write_i32(fp, len); - if (fwrite(name, 1, (size_t)len, fp) != (size_t)len) ds4_die("failed to write imatrix name"); - imatrix_write_i32(fp, ncall); - imatrix_write_i32(fp, nval); - - float *tmp = xmalloc((size_t)n_col * sizeof(tmp[0])); - for (uint32_t e = 0; e < n_expert; e++) { - const uint32_t count = counts[e]; - const float *src = sum2 + (size_t)e * n_col; - if (count == 0) { - for (uint32_t i = 0; i < n_col; i++) tmp[i] = 1.0f; - } else { - const float inv = 1.0f / (float)count; - for (uint32_t i = 0; i < n_col; i++) tmp[i] = src[i] * inv; - } - if (fwrite(tmp, sizeof(tmp[0]), n_col, fp) != n_col) ds4_die("failed to write imatrix values"); - } - free(tmp); -} - -static bool imatrix_collector_save( - const ds4_imatrix_collector *c, - const ds4_weights *weights, - const char *path) { - FILE *fp = fopen(path, "wb"); - if (!fp) { - fprintf(stderr, "ds4: failed to open imatrix output %s: %s\n", path, strerror(errno)); - return false; - } - - const int32_t entries = (int32_t)(DS4_N_LAYER * 3); - imatrix_write_i32(fp, entries); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const ds4_layer_weights *layer = &weights->layer[il]; - char name[256]; - snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_gate_exps->name.len, layer->ffn_gate_exps->name.ptr); - imatrix_write_entry(fp, name, - c->gate_up_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_EMBD, - c->gate_up_count[il], - DS4_N_EXPERT, - DS4_N_EMBD); - snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_up_exps->name.len, layer->ffn_up_exps->name.ptr); - imatrix_write_entry(fp, name, - c->gate_up_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_EMBD, - c->gate_up_count[il], - DS4_N_EXPERT, - DS4_N_EMBD); - snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_down_exps->name.len, layer->ffn_down_exps->name.ptr); - imatrix_write_entry(fp, name, - c->down_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_FF_EXP, - c->down_count[il], - DS4_N_EXPERT, - DS4_N_FF_EXP); - } - - const int32_t chunks = (int32_t)c->chunks; - imatrix_write_i32(fp, chunks); - const char *dataset = c->dataset_path ? c->dataset_path : ""; - const int32_t dataset_len = (int32_t)strlen(dataset); - imatrix_write_i32(fp, dataset_len); - if (dataset_len && fwrite(dataset, 1, (size_t)dataset_len, fp) != (size_t)dataset_len) { - ds4_die("failed to write imatrix dataset name"); - } - - if (fclose(fp) != 0) { - fprintf(stderr, "ds4: failed to close imatrix output %s: %s\n", path, strerror(errno)); - return false; - } - return true; -} - -static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { - memset(g->layer_n_comp, 0, sizeof(g->layer_n_comp)); - memset(g->layer_n_index_comp, 0, sizeof(g->layer_n_index_comp)); - g->mtp_n_raw = 0; - metal_graph_dspark_cache_reset(g); - metal_graph_dspark_capture_invalidate(g); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (!g->layer_raw_cache[il]) continue; - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio == 0) continue; - const uint32_t coff = ratio == 4 ? 2u : 1u; - const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; - const uint64_t attn_rows = (uint64_t)coff * ratio; - if (!metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows)) return false; - if (!metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows)) return false; - if (ratio == 4) { - const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM; - const uint64_t index_rows = (uint64_t)coff * ratio; - if (!metal_tensor_fill_f32(g->layer_index_state_kv[il], 0.0f, index_width * index_rows)) return false; - if (!metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows)) return false; - } - } - return true; -} - -/* Execute graph-backend prefill in layer-major order so intermediate - * activations stay on the GPU and cache state is built exactly once. */ -static void gpu_graph_report_prefill_display_progress( - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t start, - uint32_t n_tokens, - uint32_t layer_done, - int total) { - if (!display_progress) return; - if (layer_done > (uint32_t)DS4_N_LAYER) layer_done = (uint32_t)DS4_N_LAYER; - uint64_t done = (uint64_t)n_tokens * layer_done / (uint32_t)DS4_N_LAYER; - if (layer_done == (uint32_t)DS4_N_LAYER) done = n_tokens; - display_progress(display_progress_ud, "prefill_display", - (int)(start + (uint32_t)done), total); -} - -typedef struct { - int tier; - uint32_t first_layer; - uint32_t end_layer; -} metal_graph_prefill_stage; - -static bool metal_graph_build_prefill_stages( - const ds4_gpu_graph *g, - metal_graph_prefill_stage *stages, - uint32_t *n_stages) { - if (!g || !g->placement || !stages || !n_stages) return false; - uint32_t ns = 0; - int prev_tier = -1; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const int tier = g->placement[il + 1]; - if (tier < 0 || tier >= DS4_MAX_GPUS) return false; - if (il == 0 || tier != prev_tier) { - if (ns >= DS4_MAX_GPUS) return false; - stages[ns].tier = tier; - stages[ns].first_layer = il; - stages[ns].end_layer = il + 1u; - ns++; - prev_tier = tier; - } else { - stages[ns - 1u].end_layer = il + 1u; - } - } - *n_stages = ns; - return ns != 0; -} - -static bool metal_graph_encode_prefill_stage_batch( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const metal_graph_prefill_stage *stage, - uint32_t pos0, - uint32_t n_tokens) { - if (!g || !model || !weights || !stage || n_tokens == 0) return false; - if (!metal_graph_set_active_tier_no_copy(g, stage->tier)) return false; - for (uint32_t il = stage->first_layer; il < stage->end_layer; il++) { - if (g->placement && g->placement[il + 1] != stage->tier) return false; - if (!metal_graph_encode_layer_batch(g, - model, - &weights->layer[il], - il, - pos0, - n_tokens)) { - return false; - } - if (g->pipeline_capture_chunk_len != 0 && - !metal_graph_dspark_capture_prefill_rows( - g, il, - g->pipeline_capture_chunk_start, - g->pipeline_capture_chunk_len, - pos0, - n_tokens)) { - return false; - } - } - return true; -} - -static bool metal_graph_prefill_pipeline_stage_major( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - float *logits, - bool show_progress, - ds4_session_progress_fn display_progress, - void *display_progress_ud) { - if (!g || !model || !weights || !prompt || !g->placement || - n_tokens == 0 || n_tokens > g->prefill_cap || - start > (uint32_t)prompt->len || - n_tokens > (uint32_t)prompt->len - start) { - return false; - } - - metal_graph_prefill_stage stages[DS4_MAX_GPUS]; - uint32_t n_stages = 0; - g->pipeline_capture_chunk_start = start; - g->pipeline_capture_chunk_len = - g->dspark_capture_enabled ? n_tokens : 0; - if (!metal_graph_build_prefill_stages(g, stages, &n_stages) || n_stages < 2) { - return false; - } - if (stages[0].tier != g->emb_tier) { - return false; - } - - uint32_t mb_cap = metal_graph_cuda_prefill_pipeline_microbatch(); - if (mb_cap == 0 || mb_cap >= n_tokens) return false; - if (mb_cap > g->prefill_cap) mb_cap = g->prefill_cap; - const uint32_t n_mb = (n_tokens + mb_cap - 1u) / mb_cap; - if (n_mb < 2) return false; - - if (display_progress) - display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - bool ok = true; - const double t0 = getenv("DS4_METAL_GRAPH_PREFILL_PROFILE") ? now_sec() : 0.0; - - const bool sequential = getenv("DS4_CUDA_PREFILL_PIPELINE_SEQUENTIAL") != NULL; - const bool suppress_q8_cache = - !metal_graph_cuda_prefill_pipeline_q8_cache_requested(); - const int saved_q8_cache_suppressed = ds4_gpu_q8_cache_suppressed(); - if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(1); - ok = ds4_gpu_begin_commands() != 0; - if (sequential) { - for (uint32_t mb_i = 0; ok && mb_i < n_mb; mb_i++) { - const uint32_t mb_off = mb_i * mb_cap; - uint32_t mb_len = n_tokens - mb_off; - if (mb_len > mb_cap) mb_len = mb_cap; - const uint32_t pos0 = start + mb_off; - - for (uint32_t stage_i = 0; ok && stage_i < n_stages; stage_i++) { - g->batch_token_offset = mb_off; - if (stage_i == 0) { - ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); - ds4_gpu_tensor *tokens_view = NULL; - if (ok) { - tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), - (uint64_t)mb_off * sizeof(int32_t), - (uint64_t)mb_len * sizeof(int32_t)); - ok = tokens_view != NULL; - } - if (ok) { - ok = metal_graph_upload_prompt_embeddings_hc( - g->batch_cur_hc_by_tier[stages[0].tier], - tokens_view, - model, - weights, - prompt, - pos0, - mb_len); - } - ds4_gpu_tensor_free(tokens_view); - } - if (ok) { - ok = metal_graph_encode_prefill_stage_batch(g, - model, - weights, - &stages[stage_i], - pos0, - mb_len); - } - if (ok && stage_i + 1u < n_stages) { - ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; - ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; - if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { - ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && - ds4_gpu_synchronize() != 0; - } - ok = src && dst && - ds4_gpu_tensor_copy_xdev_ordered(dst, - src, - (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; - } - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok && display_progress) { - uint32_t done = mb_off + mb_len; - if (done > n_tokens) done = n_tokens; - display_progress(display_progress_ud, - "prefill_display", - (int)(start + done), - prompt->len); - } - if (show_progress) { - fprintf(stderr, "ds4: gpu sequential pipeline prefill microbatch %u/%u\r", - mb_i + 1u, - n_mb); - fflush(stderr); - } - if (ok && mb_i + 1u < n_mb) ok = ds4_gpu_begin_commands() != 0; - } - } else { - for (uint32_t wave = 0; ok && wave < n_mb + n_stages - 1u; wave++) { - uint32_t smax = wave < n_stages ? wave : n_stages - 1u; - for (int si = (int)smax; ok && si >= 0; si--) { - const uint32_t stage_i = (uint32_t)si; - const uint32_t mb_i = wave - stage_i; - if (mb_i >= n_mb) continue; - const uint32_t mb_off = mb_i * mb_cap; - uint32_t mb_len = n_tokens - mb_off; - if (mb_len > mb_cap) mb_len = mb_cap; - const uint32_t pos0 = start + mb_off; - - g->batch_token_offset = mb_off; - if (stage_i == 0) { - ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); - ds4_gpu_tensor *tokens_view = NULL; - if (ok) { - tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), - (uint64_t)mb_off * sizeof(int32_t), - (uint64_t)mb_len * sizeof(int32_t)); - ok = tokens_view != NULL; - } - if (ok) { - ok = metal_graph_upload_prompt_embeddings_hc( - g->batch_cur_hc_by_tier[stages[0].tier], - tokens_view, - model, - weights, - prompt, - pos0, - mb_len); - } - ds4_gpu_tensor_free(tokens_view); - } - - if (ok) { - ok = metal_graph_encode_prefill_stage_batch(g, - model, - weights, - &stages[stage_i], - pos0, - mb_len); - } - if (ok && stage_i + 1u < n_stages) { - ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; - ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; - if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { - ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && - ds4_gpu_synchronize() != 0; - } - ok = src && dst && - ds4_gpu_tensor_copy_xdev_ordered(dst, - src, - (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; - } - if (ok && display_progress && stage_i + 1u == n_stages) { - uint32_t done = mb_off + mb_len; - if (done > n_tokens) done = n_tokens; - display_progress(display_progress_ud, - "prefill_display", - (int)(start + done), - prompt->len); - } - } - if (show_progress) { - fprintf(stderr, "ds4: gpu pipeline prefill wave %u/%u\r", - wave + 1u, - n_mb + n_stages - 1u); - fflush(stderr); - } - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - } - if (show_progress) fputc('\n', stderr); - g->batch_token_offset = 0; - if (!ok) { - if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); - return false; - } - - const uint32_t final_len = n_tokens - (n_mb - 1u) * mb_cap; - const int src_tier = stages[n_stages - 1u].tier; - if (!metal_graph_set_active_tier_no_copy(g, src_tier)) { - if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); - return false; - } - - ds4_gpu_tensor *saved_cur = g->cur_hc_by_tier[src_tier]; - ds4_gpu_tensor *last_hc = NULL; - if (logits) { - last_hc = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[src_tier], - final_len - 1u, - hc_dim); - ok = last_hc != NULL; - } - if (ok && logits) { - g->cur_hc_by_tier[src_tier] = last_hc; - ok = ds4_gpu_begin_commands() != 0; - } - if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - if (ok && logits) ok = ds4_gpu_end_commands() != 0; - else if (!ok) (void)ds4_gpu_synchronize(); - g->cur_hc_by_tier[src_tier] = saved_cur; - ds4_gpu_tensor_free(last_hc); - if (g->placement && g->active_tier != src_tier) { - ok = metal_graph_set_active_tier_no_copy(g, src_tier); - } - if (ok && logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), - 0, - logits, - (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - if (ok && display_progress) - display_progress(display_progress_ud, "prefill_display", - (int)(start + n_tokens), prompt->len); - if (ok && t0 != 0.0) { - const double t1 = now_sec(); - fprintf(stderr, - "ds4: gpu pipeline prefill total tokens=%u stages=%u mb=%u total=%.3f ms\n", - n_tokens, - n_stages, - mb_cap, - (t1 - t0) * 1000.0); - } - if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); - return ok; -} - -static bool metal_graph_prefill_layer_major( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - float *logits, - bool show_progress, - ds4_imatrix_collector *imatrix, - ds4_session_progress_fn display_progress, - void *display_progress_ud) { - if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; - if (start > (uint32_t)prompt->len) return false; - if (n_tokens > (uint32_t)prompt->len - start) return false; - - if (display_progress) - display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); - - bool ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), prompt, start, n_tokens); - if (!ok) return false; - -#ifdef DS4_ROCM_BUILD - if (g->ssd_streaming && - DS4_MODEL_VARIANT == DS4_VARIANT_PRO && - n_tokens >= 1024u) { - ds4_gpu_stream_expert_cache_release_resident(); - } -#endif - - if (!metal_graph_warmup_prefill_kernels(g, model, weights, n_tokens)) return false; - if (g->placement && - !metal_graph_set_active_tier_no_copy(g, g->emb_tier)) { - return false; - } - metal_graph_dspark_capture_begin_prefill(g); - - const bool split_profile = - glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_SPLIT_PROFILE", - "DS4_METAL_GRAPH_PREFILL_SPLIT_PROFILE"); - /* - * A full long-prompt prefill can keep the GPU busy for a long time. Split - * non-tiny prefills when a frontend asked for display progress: completed - * layer command buffers are real scheduling/keepalive points, while - * callbacks emitted while encoding one huge command buffer would only be - * cosmetic. - */ - const bool throttle = graph_power_throttle_enabled(g); - const bool callback_split = display_progress != NULL && n_tokens >= 32; - const bool split_commands = g->ssd_streaming || - split_profile || throttle || callback_split || - n_tokens > 2048 || imatrix != NULL; - const bool profile = - glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", - "DS4_METAL_GRAPH_PREFILL_PROFILE") || - split_profile; - const double t0 = profile ? now_sec() : 0.0; - double encode_s = 0.0; - double execute_s = 0.0; - - const uint32_t pipeline_mb = metal_graph_cuda_prefill_pipeline_microbatch(); - if (!split_commands && - !profile && - imatrix == NULL && - metal_graph_cuda_prefill_pipeline_requested(g) && - pipeline_mb != 0 && - pipeline_mb < n_tokens) { - return metal_graph_prefill_pipeline_stage_major(g, - model, - weights, - prompt, - start, - n_tokens, - logits, - show_progress, - display_progress, - display_progress_ud); - } - - if (!split_commands) { - ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), - metal_graph_prefill_tokens(g), - model, - weights, - prompt, - start, - n_tokens); - if (ok) ok = ds4_gpu_begin_commands() != 0; - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - ok = metal_graph_encode_layer_batch(g, - model, - &weights->layer[il], - il, - start, - n_tokens); - if (!ok) { - fprintf(stderr, "ds4: gpu whole-prefill layer %u encode failed\n", il); - } - if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, - il, - start, - n_tokens); - if (show_progress) { - fprintf(stderr, "ds4: gpu prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); - fflush(stderr); - } - } - if (show_progress) fputc('\n', stderr); - if (display_progress) - display_progress(display_progress_ud, "prefill_display", - (int)(start + n_tokens), prompt->len); - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - uint32_t output_row = (uint32_t)n_tokens - 1u; - const char *output_row_env = glm_graph_env_value( - "DS4_ROCM_GRAPH_OUTPUT_ROW", - "DS4_METAL_GRAPH_OUTPUT_ROW"); - if (output_row_env && output_row_env[0]) { - char *end = NULL; - unsigned long v = strtoul(output_row_env, &end, 10); - if (end != output_row_env && v < (unsigned long)n_tokens) { - output_row = (uint32_t)v; - } - } - const int src_tier = g->active_tier; - ds4_gpu_tensor *saved_cur = g->cur_hc_by_tier[src_tier]; - ds4_gpu_tensor *last_hc = NULL; - if (ok && logits) { - last_hc = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), output_row, hc_dim); - ok = last_hc != NULL; - } - if (ok && logits) { - g->cur_hc_by_tier[src_tier] = last_hc; - ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - g->cur_hc_by_tier[src_tier] = saved_cur; - } - - const double t_encoded = profile ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = profile ? now_sec() : 0.0; - g->cur_hc_by_tier[src_tier] = saved_cur; - if (last_hc) ds4_gpu_tensor_free(last_hc); - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after whole-prefill graph failure also failed\n"); - } - return false; - } -#ifdef __APPLE__ - ds4_gpu_release_zero_prefix_prefill_mask_cache(); -#endif - - const double t_before_read = profile ? now_sec() : 0.0; - if (logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - if (profile) { - const double t_read = now_sec(); - fprintf(stderr, - "ds4: gpu graph prefill total tokens=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms\n", - n_tokens, - (t_encoded - t0) * 1000.0, - (t_done - t_encoded) * 1000.0, - (t_read - t_before_read) * 1000.0, - (t_read - t0) * 1000.0); - } - return ok; - } - - if (g->ssd_streaming) { - g->streaming_static_decode_map_current = false; - if (!metal_graph_stream_map_token(model, weights)) return false; - } - metal_graph_stream_prefill_selected_profile_reset(g); - metal_graph_stream_prepare_slot layer_prepare_slots[DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD]; - memset(layer_prepare_slots, 0, sizeof(layer_prepare_slots)); - const bool layer_pagein = - metal_graph_stream_prefill_layer_pagein_enabled(g); - const bool layer_readahead = - !layer_pagein && - metal_graph_stream_prefill_layer_readahead_enabled(g); - const bool layer_pread = - !layer_pagein && !layer_readahead && - metal_graph_stream_prefill_layer_pread_enabled(g); - const bool layer_madvise = - !layer_pagein && !layer_pread && !layer_readahead && - metal_graph_stream_prefill_layer_madvise_enabled(g); - const bool layer_prepare = - layer_pagein || layer_pread || layer_readahead || layer_madvise; - const bool layer_prepare_overlap = - layer_prepare && metal_graph_stream_prefill_layer_pagein_overlap_enabled(); - const uint32_t layer_prepare_ahead = - layer_prepare && layer_prepare_overlap ? - metal_graph_stream_prefill_layer_prepare_ahead() : 1u; - const bool batch_selected_addr = - metal_graph_stream_prefill_batch_selected_addr_enabled(g, weights, n_tokens) || - metal_graph_cuda_stream_prefill_batch_selected_addr_enabled(g, weights, n_tokens); -#ifdef DS4_ROCM_BUILD - rocm_graph_stream_layer_expert_load rocm_full_layer_load; - memset(&rocm_full_layer_load, 0, sizeof(rocm_full_layer_load)); -#endif - if (g->ssd_streaming && DS4_N_LAYER > 0) { - if (layer_prepare) { - if (!metal_graph_stream_prepare_start_if_needed(g, - model, - weights, - 0, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - batch_selected_addr, - layer_prepare_slots, - layer_prepare_ahead)) { - return false; - } - } else { - if (batch_selected_addr) { - metal_graph_stream_readahead_layer_decode(model, weights, 0); - } else { - metal_graph_stream_readahead_layer(model, weights, 0); - } - } - } -#ifdef DS4_ROCM_BUILD - if (g->ssd_streaming && DS4_N_LAYER > 0 && - !rocm_graph_stream_layer_expert_load_start_next(&rocm_full_layer_load, - g, - model, - weights, - 0, - n_tokens)) { - return false; - } -#endif - - double t_layer0 = (profile || throttle) ? now_sec() : 0.0; - ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), - metal_graph_prefill_tokens(g), - model, - weights, - prompt, - start, - n_tokens); - const double t_embed_encoded = (profile || throttle) ? now_sec() : 0.0; - const double t_embed_done = (profile || throttle) ? now_sec() : 0.0; - if (profile) { - encode_s += t_embed_encoded - t_layer0; - execute_s += t_embed_done - t_embed_encoded; - if (split_profile) { - fprintf(stderr, - "ds4: metal layer-major prefill embed encode=%.3f ms execute=%.3f ms\n", - (t_embed_encoded - t_layer0) * 1000.0, - (t_embed_done - t_embed_encoded) * 1000.0); - } - } - if (!ok) { -#ifdef DS4_ROCM_BUILD - (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); -#endif - if (layer_prepare) { - (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, - layer_prepare_ahead); - } - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after layer-major prefill embed failure also failed\n"); - } - return false; - } - - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - double layer_elapsed = 0.0; - if (layer_prepare && - !metal_graph_stream_prepare_join_layer(g, - model, - weights, - il, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - batch_selected_addr, - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - break; - } -#ifdef DS4_ROCM_BUILD - const bool rocm_full_layer_stream_prefill = - rocm_graph_stream_prefill_full_layer_enabled(g, - &weights->layer[il], - il, - n_tokens); - if (rocm_full_layer_stream_prefill && - !rocm_graph_stream_layer_expert_load_ready(&rocm_full_layer_load, - g, - model, - weights, - il, - n_tokens)) { - ok = false; - break; - } - if (rocm_full_layer_stream_prefill && - !rocm_graph_stream_layer_expert_load_start_next(&rocm_full_layer_load, - g, - model, - weights, - il + 1u, - n_tokens)) { - ok = false; - break; - } -#endif - if (g->ssd_streaming) { - g->streaming_static_decode_map_current = false; - bool decode_only_map = batch_selected_addr; -#ifdef DS4_ROCM_BUILD - decode_only_map = decode_only_map || rocm_full_layer_stream_prefill; -#endif - const bool map_ok = decode_only_map ? - metal_graph_stream_map_layer_decode(model, weights, il) : - metal_graph_stream_map_layer(model, weights, il); - if (!map_ok) { - ok = false; - break; - } - } - if (g->ssd_streaming) { - if (layer_prepare && layer_prepare_overlap) { - bool started_future = false; - for (uint32_t ahead = 1; ahead <= layer_prepare_ahead; ahead++) { - if (il + ahead >= DS4_N_LAYER) break; - started_future = true; - if (!metal_graph_stream_prepare_start_if_needed(g, - model, - weights, - il + ahead, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - batch_selected_addr, - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - break; - } - } - if (!ok) break; - if (!started_future && logits) { - metal_graph_stream_readahead_output(model, weights); - } - } else if (!layer_prepare && il + 1 < DS4_N_LAYER) { - if (batch_selected_addr) { - metal_graph_stream_readahead_layer_decode(model, weights, il + 1); - } else { - metal_graph_stream_readahead_layer(model, weights, il + 1); - } - } else if (logits) { - metal_graph_stream_readahead_output(model, weights); - } - } - if (split_profile) { - /* (B6 fix): split-profile diagnostic bypasses the - * metal_graph_encode_layer_batch wrapper that normally does - * the per-layer tier switch. Replicate the switch here so the - * diagnostic / profile mode stays multi-tier-correct. - * Single-tier (g->placement == NULL): no-op. */ - if (g->placement) { - const int this_tier = g->placement[il + 1]; - if (!metal_graph_set_active_tier_batch(g, this_tier, (uint32_t)n_tokens)) { - ok = false; - break; - } - } - const double t_attn0 = now_sec(); - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_layer_attention_batch(g, - model, - &weights->layer[il], - il, - start, - n_tokens); - if (!ok) { - fprintf(stderr, "ds4: gpu layer-major prefill layer %u attention encode failed\n", il); - } - const double t_attn_encoded = now_sec(); - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_attn_done = now_sec(); - - const double t_ffn0 = now_sec(); - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_layer_ffn_batch(g, - model, - &weights->layer[il], - il, - start, - n_tokens, - NULL, - 0); - if (!ok) { - fprintf(stderr, "ds4: gpu layer-major prefill layer %u ffn encode failed\n", il); - } - if (ok) { - ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); - g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); - g->batch_next_hc_by_tier[g->active_tier] = tmp; - } - if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, - il, - start, - n_tokens); - if (ok) ok = metal_graph_capture_prefill_seed_router_selected(g, - il, - n_tokens); - const double t_ffn_encoded = now_sec(); - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_ffn_done = now_sec(); -#ifdef DS4_ROCM_BUILD - if (ok) { - ok = rocm_graph_stream_seed_full_layer_selected(g, - model, - &weights->layer[il], - il, - n_tokens); - } -#endif - if (ok) { - ok = metal_graph_stream_prefill_selected_profile_layer( - g, - &weights->layer[il], - il, - n_tokens); - } - if (ok && imatrix) ok = imatrix_collect_layer_batch(imatrix, g, il, (uint32_t)n_tokens); - layer_elapsed = (t_attn_done - t_attn0) + (t_ffn_done - t_ffn0); - - encode_s += (t_attn_encoded - t_attn0) + (t_ffn_encoded - t_ffn0); - execute_s += (t_attn_done - t_attn_encoded) + (t_ffn_done - t_ffn_encoded); - fprintf(stderr, - "ds4: metal layer-major prefill layer %u attn encode=%.3f execute=%.3f ms ffn encode=%.3f execute=%.3f ms\n", - il, - (t_attn_encoded - t_attn0) * 1000.0, - (t_attn_done - t_attn_encoded) * 1000.0, - (t_ffn_encoded - t_ffn0) * 1000.0, - (t_ffn_done - t_ffn_encoded) * 1000.0); - } else { - const double t_chunk0 = (profile || throttle) ? now_sec() : 0.0; - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_layer_batch(g, - model, - &weights->layer[il], - il, - start, - n_tokens); - if (!ok) { - fprintf(stderr, "ds4: gpu layer-major prefill layer %u encode failed\n", il); - } - if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, - il, - start, - n_tokens); - if (ok) ok = metal_graph_capture_prefill_seed_router_selected(g, - il, - n_tokens); - const double t_encoded = (profile || throttle) ? now_sec() : 0.0; - if (ok) ok = ds4_gpu_end_commands() != 0; - const double t_done = (profile || throttle) ? now_sec() : 0.0; -#ifdef DS4_ROCM_BUILD - if (ok) { - ok = rocm_graph_stream_seed_full_layer_selected(g, - model, - &weights->layer[il], - il, - n_tokens); - } -#endif - if (ok) { - ok = metal_graph_stream_prefill_selected_profile_layer( - g, - &weights->layer[il], - il, - n_tokens); - } - if (ok && imatrix) ok = imatrix_collect_layer_batch(imatrix, g, il, (uint32_t)n_tokens); - layer_elapsed = t_done - t_chunk0; - if (profile) { - encode_s += t_encoded - t_chunk0; - execute_s += t_done - t_encoded; - fprintf(stderr, - "ds4: gpu layer-major prefill layer %u encode=%.3f ms execute=%.3f ms\n", - il, - (t_encoded - t_chunk0) * 1000.0, - (t_done - t_encoded) * 1000.0); - } - } - if (ok && - g->ssd_streaming && - layer_prepare && - !layer_prepare_overlap) { - if (il + 1 < DS4_N_LAYER) { - if (!metal_graph_stream_prepare_start_if_needed(g, - model, - weights, - il + 1, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - batch_selected_addr, - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - } - } else if (logits) { - metal_graph_stream_readahead_output(model, weights); - } - } - if (!ok) { -#ifdef DS4_ROCM_BUILD - (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); -#endif - if (layer_prepare) { - (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, - layer_prepare_ahead); - } - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after layer-major prefill failure also failed\n"); - } - return false; - } - graph_power_note_prefill_layer(g, il, layer_elapsed); - gpu_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - start, - n_tokens, - il + 1, - prompt->len); - if (show_progress) { - fprintf(stderr, "ds4: gpu prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); - fflush(stderr); - } - } - if (!ok) { -#ifdef DS4_ROCM_BUILD - (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); -#endif - if (layer_prepare) { - (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, - layer_prepare_ahead); - } - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after layer-major prefill failure also failed\n"); - } - return false; - } -#ifdef __APPLE__ - /* Zero-prefix masks are shared across the 43 per-layer command batches, - * then become dead weight. Release them before the output head and later - * replay/decode chunks so the prefill win does not add residency pressure. */ - ds4_gpu_release_zero_prefix_prefill_mask_cache(); -#endif - if (show_progress) fputc('\n', stderr); - metal_graph_stream_prefill_selected_profile_summary(g); -#ifdef DS4_ROCM_BUILD - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); - if (g->ssd_streaming) ds4_gpu_release_q8_f16_cache(); -#endif - if (!metal_graph_seed_streaming_expert_cache_from_hotlist(g, model, weights)) { - return false; - } - if (!metal_graph_seed_streaming_expert_cache_from_prefill(g, model, weights)) { - return false; - } - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - uint32_t output_row = (uint32_t)n_tokens - 1u; - const char *output_row_env = glm_graph_env_value( - "DS4_ROCM_GRAPH_OUTPUT_ROW", - "DS4_METAL_GRAPH_OUTPUT_ROW"); - if (output_row_env && output_row_env[0]) { - char *end = NULL; - unsigned long v = strtoul(output_row_env, &end, 10); - if (end != output_row_env && v < (unsigned long)n_tokens) { - output_row = (uint32_t)v; - } - } - ds4_gpu_tensor *saved_cur = metal_graph_cur_hc(g); - ds4_gpu_tensor *last_hc = NULL; - - const double t_head0 = profile ? now_sec() : 0.0; - if (logits) { - last_hc = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), - output_row, - hc_dim); - ok = last_hc != NULL; - } - if (ok && logits && g->ssd_streaming) { - const bool static_decode_map = - metal_graph_stream_decode_static_map_enabled(); - const bool static_map_state_cache = - static_decode_map && - metal_graph_stream_decode_static_map_state_cache_enabled(); - g->streaming_static_decode_map_current = false; - if (static_map_state_cache) { - ok = metal_graph_stream_map_decode_static_all(model, weights); - if (ok) g->streaming_static_decode_map_current = true; - } else { - ok = metal_graph_stream_map_output(model, weights); - } - } - if (ok && logits) { - g->cur_hc_by_tier[g->active_tier] = last_hc; - ok = ds4_gpu_begin_commands() != 0; - } - if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - const double t_head_encoded = profile ? now_sec() : 0.0; - if (ok && logits) ok = ds4_gpu_end_commands() != 0; - const double t_head_done = profile ? now_sec() : 0.0; - g->cur_hc_by_tier[g->active_tier] = saved_cur; - if (last_hc) ds4_gpu_tensor_free(last_hc); - if (!ok) return false; - - const double t_before_read = profile ? now_sec() : 0.0; - if (logits) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - if (profile) { - const double t_read = now_sec(); - encode_s += t_head_encoded - t_head0; - execute_s += t_head_done - t_head_encoded; - if (split_profile) { - fprintf(stderr, - "ds4: gpu layer-major prefill head encode=%.3f ms execute=%.3f ms\n", - (t_head_encoded - t_head0) * 1000.0, - (t_head_done - t_head_encoded) * 1000.0); - } - fprintf(stderr, - "ds4: gpu layer-major prefill total tokens=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms\n", - n_tokens, - encode_s * 1000.0, - execute_s * 1000.0, - (t_read - t_before_read) * 1000.0, - (t_read - t0) * 1000.0); - } - return ok; -} - -static bool metal_graph_prefill_raw_swa( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - int n_tokens, - float *logits, - bool show_progress, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - ds4_session_cancel_fn cancel, - void *cancel_ud, - bool *cancelled) { - if (n_tokens <= 0 || n_tokens > prompt->len) return false; - if ((uint32_t)n_tokens > g->prefill_cap) return false; - if (metal_graph_use_streaming_decode_prefill_range(g, weights, 0, - (uint32_t)n_tokens)) { - return metal_graph_prefill_decode_streaming_range(g, - model, - weights, - prompt, - 0, - (uint32_t)n_tokens, - logits, - show_progress, - NULL, - NULL, - display_progress, - display_progress_ud, - cancel, - cancel_ud, - cancelled); - } - /* The layer-major fallback below may submit the whole short prefill as one - * Metal command buffer. Once that command is in flight there is no useful - * safe prefix to expose: by the time cancellation can be observed again, - * the prompt has already been fully read and the KV is valid. Let the - * caller observe the pending interrupt at generation time instead. */ - (void)cancel; - (void)cancel_ud; - (void)cancelled; - return metal_graph_prefill_layer_major(g, - model, - weights, - prompt, - 0, - (uint32_t)n_tokens, - logits, - show_progress, - NULL, - display_progress, - display_progress_ud); -} - -/* Prefill a contiguous token range in fixed-size chunks. - * - * The common case starts at token zero, but server sessions also use this to - * extend an existing KV cache with a long suffix. Resumed chunks are aligned - * to the same absolute prefill-cap boundaries used by a cold full prompt, so - * compression windows and row finalization follow the same schedule after the - * cached prefix. - */ -static bool metal_graph_prefill_chunked_range( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - float *logits, - bool show_progress, - ds4_session_progress_fn progress, - void *progress_ud, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - ds4_imatrix_collector *imatrix, - ds4_session_cancel_fn cancel, - void *cancel_ud, - bool *cancelled) { - if (n_tokens == 0 || g->prefill_cap == 0) return false; - if (start > (uint32_t)prompt->len) return false; - if (n_tokens > (uint32_t)prompt->len - start) return false; - if (g->ssd_streaming && start == 0) { - ds4_gpu_stream_expert_cache_reset_route_hotness(); - } - if (!imatrix && - metal_graph_use_streaming_decode_prefill_range(g, weights, - start, n_tokens)) { - return metal_graph_prefill_decode_streaming_range(g, - model, - weights, - prompt, - start, - n_tokens, - logits, - show_progress, - progress, - progress_ud, - display_progress, - display_progress_ud, - cancel, - cancel_ud, - cancelled); - } - - uint32_t chunk_cap = g->prefill_cap; - if (start != 0 && chunk_cap > g->raw_cap) chunk_cap = g->raw_cap; - if (chunk_cap == 0) return false; - - const bool profile = - glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", - "DS4_METAL_GRAPH_PREFILL_PROFILE"); - const double t0 = profile ? now_sec() : 0.0; - const uint32_t end = start + n_tokens; - - if (progress) { - progress(progress_ud, "prefill_chunk", (int)start, prompt->len); - } - if (display_progress) { - display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); - } - - for (uint32_t pos0 = start; pos0 < end; ) { - if (cancel && cancel(cancel_ud)) { - if (cancelled) *cancelled = true; - return true; - } - const uint32_t remaining = end - pos0; - uint32_t local_cap = chunk_cap; - if (start != 0 && g->prefill_cap != 0) { - const uint32_t mod = pos0 % g->prefill_cap; - if (mod != 0) { - const uint32_t to_boundary = g->prefill_cap - mod; - if (to_boundary < local_cap) local_cap = to_boundary; - } - } - const uint32_t chunk = remaining < local_cap ? remaining : local_cap; - const uint32_t chunk_end = pos0 + chunk; - float *chunk_logits = (progress || chunk_end == end) ? logits : NULL; - bool ok = metal_graph_prefill_layer_major(g, - model, - weights, - prompt, - pos0, - chunk, - chunk_logits, - show_progress, - imatrix, - display_progress, - display_progress_ud); - if (!ok) { - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after chunked prefill failure also failed\n"); - } - return false; - } - if (progress) { - progress(progress_ud, "prefill_chunk", (int)chunk_end, prompt->len); - } - if (display_progress) { - display_progress(display_progress_ud, "prefill_display", (int)chunk_end, prompt->len); - } - if (cancel && cancel(cancel_ud)) { - if (cancelled) *cancelled = true; - return true; - } - pos0 = chunk_end; - } - if (show_progress) fputc('\n', stderr); - if (profile) { - const double t_read = now_sec(); - fprintf(stderr, - "ds4: gpu chunked prefill start=%u tokens=%u chunk=%u total=%.3f ms\n", - start, - n_tokens, - chunk_cap, - (t_read - t0) * 1000.0); - } - return true; -} - -/* Long prompts are prefetched in fixed-size chunks. Chunks bound transient - * attention buffers while preserving the same final KV/cache state. */ -static bool metal_graph_prefill_chunked( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - int n_tokens, - float *logits, - bool show_progress, - ds4_session_progress_fn progress, - void *progress_ud, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - ds4_session_cancel_fn cancel, - void *cancel_ud, - bool *cancelled) { - if (n_tokens <= 0) return false; - return metal_graph_prefill_chunked_range(g, - model, - weights, - prompt, - 0, - (uint32_t)n_tokens, - logits, - show_progress, - progress, - progress_ud, - display_progress, - display_progress_ud, - NULL, - cancel, - cancel_ud, - cancelled); -} - -typedef struct ds4_verify_suffix_timing { - double upload_ms; - double layer_ms; - double head_ms; - double read_ms; - bool fused_head; -} ds4_verify_suffix_timing; - -static bool metal_graph_dspark_verify_selected_profile_enabled(void) { - return getenv("DS4_DSPARK_VERIFY_SELECTED_PROFILE") != NULL && - getenv("DS4_DSPARK_DISABLE_VERIFY_SELECTED_PROFILE") == NULL; -} - -/* Layer-major speculative target verifier for tiny MTP suffixes. - * - * This is the first production-shaped verifier attempt: unlike repeated decode - * it runs the target model layer-by-layer for the whole speculative suffix, and - * unlike the diagnostic path it does not read back full logits for every row. - * The verifier returns the row top-1 ids needed for acceptance. The caller - * then reads exactly one logits row: the row that becomes the new continuation - * state. It still reuses the existing batch layer kernels, so it is not yet - * the final hand-written N=2/N=4 decode microbatch, but it exercises the right - * verifier contract and removes the obvious diagnostic overheads first. */ -static bool metal_graph_verify_suffix_tops_impl( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - bool capture_prefix1, - bool capture_dspark_hidden, - int *row_tops, - float *row_logits, - ds4_verify_suffix_timing *timing) { - if (timing) memset(timing, 0, sizeof(*timing)); - if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; - if (start > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - start) return false; - const uint32_t top_rows = n_tokens > 1 ? n_tokens - 1 : 0; - if (top_rows && !row_tops) return false; - - const double upload_t0 = timing ? now_sec() : 0.0; - bool ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), prompt, start, n_tokens); - if (ok) ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), - metal_graph_prefill_tokens(g), - model, - weights, - prompt, - start, - n_tokens); - if (!ok) return false; - - const bool saved_capture = g->spec_capture_prefix1; - g->spec_capture_prefix1 = capture_prefix1 && n_tokens == 2; - const char *split_head_env = getenv("DS4_DSPARK_VERIFY_SPLIT_HEAD"); - const bool fuse_head = - !split_head_env || !split_head_env[0] || - strcmp(split_head_env, "0") == 0; - if (timing) timing->fused_head = fuse_head; - if (timing) timing->upload_ms += (now_sec() - upload_t0) * 1000.0; - - const bool selected_profile = - metal_graph_dspark_verify_selected_profile_enabled(); - if (selected_profile) { - metal_graph_stream_prefill_selected_profile_reset(g); - } - - /* Under TP, verify every speculative block against the two resident - * expert halves. Both ranks encode identically, so one batch gate per - * layer reconstructs the routed result while preserving their KV state. */ - g->tp_batch_rows = (g->tp_world == 2 && - g->tp_batch_out != NULL && g->tp_batch_in != NULL && - n_tokens <= (uint32_t)DS4_TP_BATCH_MAX_ROWS) - ? n_tokens : 0; - const double layer_t0 = timing ? now_sec() : 0.0; - ok = ds4_gpu_begin_commands() != 0; - const bool dspark_capture_active = - ok && - capture_dspark_hidden && - metal_graph_dspark_capture_verified_suffix_begin(g, - start, - n_tokens, - true); - static int verify_profile_left = -1; - if (verify_profile_left < 0) { - verify_profile_left = - getenv("DS4_DSPARK_VERIFY_PROFILE") != NULL ? 1 : 0; - } - const bool verify_profile = verify_profile_left > 0 && ok; - if (verify_profile) verify_profile_left--; - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - if (verify_profile) { - ok = ds4_gpu_end_commands() != 0; - if (ok) (void)ds4_gpu_synchronize(); - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (!ok) break; - } - ok = metal_graph_encode_layer_batch(g, - model, - &weights->layer[il], - il, - start, - n_tokens); - if (ok && dspark_capture_active) { - ok = metal_graph_dspark_capture_verified_suffix_layer(g, - il, - start, - n_tokens); - } - if (ok && selected_profile) { - ok = ds4_gpu_end_commands() != 0 && - metal_graph_selected_profile_layer_impl( - g, - &weights->layer[il], - il, - n_tokens, - "DSpark verifier selected profile") && - ds4_gpu_begin_commands() != 0; - } - } - g->tp_batch_rows = 0; - if (ok && fuse_head) { - ok = metal_graph_encode_output_head_batch(g, - model, - weights, - n_tokens, - weights->output->dim[1]); - } - if (ok && fuse_head) { - if (top_rows == 1) { - ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), - g->spec_logits, - DS4_N_VOCAB) != 0; - } else if (top_rows) { - ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - g->spec_logits, - DS4_N_VOCAB, - top_rows, - 1) != 0; - } - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - g->spec_capture_prefix1 = saved_capture; - if (!ok && dspark_capture_active) { - metal_graph_dspark_capture_invalidate(g); - } - if (timing) timing->layer_ms += (now_sec() - layer_t0) * 1000.0; - if (!ok) return false; - if (selected_profile) { - metal_graph_selected_profile_summary_impl( - g, - "DSpark verifier selected profile"); - } - - if (!fuse_head) { - const double head_t0 = timing ? now_sec() : 0.0; - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_output_head_batch(g, - model, - weights, - n_tokens, - weights->output->dim[1]); - if (ok) { - if (top_rows == 1) { - /* Common K=2 verify case: top_k=1 over n_vocab → use the dedicated - * argmax kernel (single-block tree-reduce) instead of the legacy - * indexer_topk_kernel's single-thread O(n_vocab * top_k) fall-through. */ - ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), - g->spec_logits, - DS4_N_VOCAB) != 0; - } else if (top_rows) { - /* top-1 of each of the top_rows rows: n_tokens=top_rows, top_k=1. - * The order is transposed vs the indexer-score callers; a swap - * silently scores row 0's runner-ups instead of each row. */ - ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - g->spec_logits, - DS4_N_VOCAB, - top_rows, - 1) != 0; - } - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (timing) timing->head_ms += (now_sec() - head_t0) * 1000.0; - } - const double read_t0 = timing ? now_sec() : 0.0; - if (ok && top_rows) { - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), - 0, - row_tops, - (uint64_t)top_rows * sizeof(row_tops[0])) != 0; - if (ok && getenv("DS4_DSPARK_VERIFY_TOPS_CHECK") != NULL) { - float *chk = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); - for (uint32_t r = 0; r < top_rows; r++) { - if (ds4_gpu_tensor_read(g->spec_logits, - (uint64_t)r * DS4_N_VOCAB * - sizeof(float), - chk, - (uint64_t)DS4_N_VOCAB * - sizeof(float)) == 0) break; - uint32_t am = 0; - for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { - if (chk[i] > chk[am]) am = i; - } - fprintf(stderr, - "ds4: verify tops-check row=%u gpu_top=%d cpu_argmax=%u " - "cpu_max=%.3f\n", - r, row_tops[r], am, chk[am]); - } - free(chk); - } - } - if (ok && row_logits) { - ok = ds4_gpu_tensor_read(g->spec_logits, - 0, - row_logits, - (uint64_t)n_tokens * DS4_N_VOCAB * sizeof(row_logits[0])) != 0; - } - if (timing) timing->read_ms += (now_sec() - read_t0) * 1000.0; - return ok; -} - -/* The verify block keeps the GPU genuinely busy, so the TP DVFS - * keep-alive is a pure parasite for its duration — pause it. */ -static bool metal_graph_verify_suffix_tops( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - uint32_t start, - uint32_t n_tokens, - bool capture_prefix1, - bool capture_dspark_hidden, - int *row_tops, - float *row_logits, - ds4_verify_suffix_timing *timing) { - ds4_gpu_tp_keepalive_pause(1); - const bool ok = metal_graph_verify_suffix_tops_impl(g, model, weights, - prompt, start, - n_tokens, - capture_prefix1, - capture_dspark_hidden, - row_tops, row_logits, - timing); - ds4_gpu_tp_keepalive_pause(0); - return ok; -} - -static bool metal_graph_read_spec_logits_row(ds4_gpu_graph *g, uint32_t row, float *logits) { - if (!g || !g->spec_logits || !logits || row >= g->prefill_cap) return false; - const uint64_t row_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); - return ds4_gpu_tensor_read(g->spec_logits, - (uint64_t)row * row_bytes, - logits, - row_bytes) != 0; -} - -/* Exact N=2 target verifier for MTP. - * - * The generic batch prefill path is fast, but it is not a safe substitute for - * autoregressive decode: small row-wise differences in HC/MoE/output kernels - * are enough to flip future greedy tokens. This verifier keeps the exact - * decode kernels and cache update order, but encodes the two proposed tokens - * layer-by-layer in one command stream. It returns the exact target top after - * token0, and exact logits after token1. */ -static bool metal_graph_verify_decode2_exact( - ds4_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token0, - int token1, - uint32_t start, - int *top0, - int *top1, - float *logits0, - float *logits1) { - if (!g || !top0 || (!top1 && !logits1) || g->raw_cap == 0) return false; - - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t hc_bytes = hc_dim * sizeof(float); - ds4_gpu_tensor *cur0_by_tier[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor *cur1_by_tier[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor *next0_by_tier[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor *next1_by_tier[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor *saved_cur_by_tier[DS4_MAX_GPUS] = {0}; - ds4_gpu_tensor *saved_after_by_tier[DS4_MAX_GPUS] = {0}; - const int saved_active_tier = g->active_tier; - const bool saved_capture = g->spec_capture_prefix1; - - bool ok = true; - for (int t = 0; t < DS4_MAX_GPUS; t++) { - saved_cur_by_tier[t] = g->cur_hc_by_tier[t]; - saved_after_by_tier[t] = g->after_ffn_hc_by_tier[t]; - if (!g->batch_cur_hc_by_tier[t] && !g->batch_next_hc_by_tier[t]) continue; - if (!g->batch_cur_hc_by_tier[t] || !g->batch_next_hc_by_tier[t]) { - ok = false; - break; - } - cur0_by_tier[t] = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[t], 0, hc_dim); - cur1_by_tier[t] = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[t], 1, hc_dim); - next0_by_tier[t] = metal_graph_tensor_row_view(g->batch_next_hc_by_tier[t], 0, hc_dim); - next1_by_tier[t] = metal_graph_tensor_row_view(g->batch_next_hc_by_tier[t], 1, hc_dim); - if (!cur0_by_tier[t] || !cur1_by_tier[t] || - !next0_by_tier[t] || !next1_by_tier[t]) { - ok = false; - break; - } - } - - int cur_tier = g->emb_tier; - if (cur_tier < 0 || cur_tier >= DS4_MAX_GPUS || - !cur0_by_tier[cur_tier] || !cur1_by_tier[cur_tier]) { - ok = false; - } - if (ok) ok = metal_graph_set_active_tier_no_copy(g, cur_tier); - if (ok) ok = ds4_gpu_embed_token_hc_tensor(cur0_by_tier[cur_tier], - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token0, - DS4_N_EMBD, - DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(cur1_by_tier[cur_tier], - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token1, - DS4_N_EMBD, - DS4_N_HC) != 0; - - g->spec_capture_prefix1 = true; - if (ok) ok = ds4_gpu_begin_commands() != 0; - for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { - const uint32_t pos0 = start; - const uint32_t pos1 = start + 1u; - const int this_tier = g->placement ? g->placement[il + 1] : cur_tier; - if (this_tier < 0 || this_tier >= DS4_MAX_GPUS || - !cur0_by_tier[this_tier] || !cur1_by_tier[this_tier]) { - ok = false; - break; - } - if (this_tier != cur_tier) { - ok = ds4_gpu_tensor_copy_xdev(cur0_by_tier[this_tier], - cur0_by_tier[cur_tier], - hc_bytes) != 0 && - ds4_gpu_tensor_copy_xdev(cur1_by_tier[this_tier], - cur1_by_tier[cur_tier], - hc_bytes) != 0; - if (!ok) break; - cur_tier = this_tier; - } - ok = metal_graph_set_active_tier_no_copy(g, this_tier); - if (!ok) break; - - g->cur_hc_by_tier[this_tier] = cur0_by_tier[this_tier]; - g->after_ffn_hc_by_tier[this_tier] = next0_by_tier[this_tier]; - ok = metal_graph_encode_decode_layer(g, - model, - &weights->layer[il], - il, - pos0, - g->layer_raw_cache[il], - g->raw_cap, - pos0 % g->raw_cap, - metal_graph_raw_span_for_batch(g, pos0, 1), - token0); - if (!ok) break; - ok = metal_graph_capture_prefix1_attn_state(g, il) && - metal_graph_capture_prefix1_index_state(g, il); - if (!ok) break; - - g->cur_hc_by_tier[this_tier] = cur1_by_tier[this_tier]; - g->after_ffn_hc_by_tier[this_tier] = next1_by_tier[this_tier]; - ok = metal_graph_encode_decode_layer(g, - model, - &weights->layer[il], - il, - pos1, - g->layer_raw_cache[il], - g->raw_cap, - pos1 % g->raw_cap, - metal_graph_raw_span_for_batch(g, pos1, 1), - token1); - if (!ok) break; - - ds4_gpu_tensor *tmp = cur0_by_tier[this_tier]; - cur0_by_tier[this_tier] = next0_by_tier[this_tier]; - next0_by_tier[this_tier] = tmp; - tmp = cur1_by_tier[this_tier]; - cur1_by_tier[this_tier] = next1_by_tier[this_tier]; - next1_by_tier[this_tier] = tmp; - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - g->spec_capture_prefix1 = saved_capture; - - if (ok) { - ok = metal_graph_set_active_tier_no_copy(g, cur_tier); - } - if (ok) { - const bool split_top1 = - logits0 == NULL && - g->cuda_tp_output && - metal_graph_cuda_verify_decode2_split_top1_requested(); - uint32_t output_ways = 0; - g->cur_hc_by_tier[cur_tier] = cur0_by_tier[cur_tier]; - ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - if (ok) ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), - metal_graph_logits(g), - DS4_N_VOCAB) != 0; - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok && split_top1) { - ok = metal_graph_read_output_split_top1(g, output_ways, top0); - } else if (ok) { - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top0, sizeof(*top0)) != 0; - } - if (ok && logits0) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), - 0, - logits0, - (uint64_t)DS4_N_VOCAB * sizeof(logits0[0])) != 0; - } - } - - if (ok) { - ok = metal_graph_set_active_tier_no_copy(g, cur_tier); - } - if (ok) { - const bool split_top1 = - logits1 == NULL && - top1 != NULL && - g->cuda_tp_output && - metal_graph_cuda_verify_decode2_split_top1_requested(); - int output_tiers[DS4_MAX_GPUS] = {0}; - uint32_t output_ways = 0; - g->cur_hc_by_tier[cur_tier] = cur1_by_tier[cur_tier]; - ok = ds4_gpu_begin_commands() != 0; - if (ok && split_top1) { - ok = metal_graph_encode_output_head_split_top1(g, - model, - weights, - weights->output->dim[1], - output_tiers, - &output_ways); - } else if (ok) { - ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - if (ok && top1) { - ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), - metal_graph_logits(g), - DS4_N_VOCAB, - 1, - 1) != 0; - } - } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok && split_top1) { - ok = metal_graph_read_output_split_top1(g, output_ways, top1); - } else if (ok && top1) { - ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top1, sizeof(*top1)) != 0; - } - if (ok) { - if (logits1) { - ok = ds4_gpu_tensor_read(metal_graph_logits(g), - 0, - logits1, - (uint64_t)DS4_N_VOCAB * sizeof(logits1[0])) != 0; - } - } - } - g->spec_capture_prefix1 = saved_capture; - for (int t = 0; t < DS4_MAX_GPUS; t++) { - g->cur_hc_by_tier[t] = saved_cur_by_tier[t]; - g->after_ffn_hc_by_tier[t] = saved_after_by_tier[t]; - } - if (g->placement) { - if (saved_active_tier >= 0) { - (void)metal_graph_set_active_tier_no_copy(g, saved_active_tier); - } else { - g->active_tier = saved_active_tier; - } - } - for (int t = 0; t < DS4_MAX_GPUS; t++) { - ds4_gpu_tensor_free(next1_by_tier[t]); - ds4_gpu_tensor_free(next0_by_tier[t]); - ds4_gpu_tensor_free(cur1_by_tier[t]); - ds4_gpu_tensor_free(cur0_by_tier[t]); - } - return ok; -} - -/* Pick a raw SWA cache size for Metal. During batched prefill it must cover - * the previous window plus the current ubatch. */ -static uint32_t metal_graph_raw_cap_for_context(int ctx_size, uint32_t prefill_cap) { - uint32_t raw_window = DS4_N_SWA; - if (raw_window > (uint32_t)ctx_size) raw_window = (uint32_t)ctx_size; - if (raw_window == 0) raw_window = 1; - - /* - * During batched prefill the SWA cache must hold the current ubatch plus - * the previous logical window. The cache is padded to a 256-row multiple - * so the physical row order and FlashAttention block grouping match the - * model path we compare against. - */ - uint64_t wanted = (uint64_t)raw_window + prefill_cap; - if (wanted > (uint32_t)ctx_size) wanted = (uint32_t)ctx_size; - if (wanted == 0) wanted = 1; - wanted = align_up(wanted, 256u); - if (wanted > 8192u) wanted = 8192u; - uint32_t raw_cap = (uint32_t)wanted; - if (raw_cap < raw_window) raw_cap = raw_window; - -#ifndef DS4_ROCM_BUILD - const char *env = getenv("DS4_METAL_GRAPH_RAW_CAP"); - if (env && env[0]) { - char *endp = NULL; - const long v = strtol(env, &endp, 10); - if (endp != env && v > 0) { - raw_cap = (uint32_t)v; - if (raw_cap > (uint32_t)ctx_size) raw_cap = (uint32_t)ctx_size; - if (raw_cap > 8192u) raw_cap = 8192u; - if (raw_cap < raw_window) raw_cap = raw_window; - } - } -#endif - - return raw_cap; -} - -/* Choose the prefill ubatch size. Whole-batch is fastest for normal prompts. - * Long Flash prompts default to 4096-token chunks; PRO defaults to 8192. */ -static uint32_t metal_graph_prefill_cap_for_prompt(int prompt_len, - uint32_t prefill_chunk) { - return ds4_prefill_cap_for_prompt(prompt_len, prefill_chunk); -} - -/* When a server request shares a large prefix with the live checkpoint, extend - * the KV cache with batched prefill instead of single-token decode. On an M3 - * Max, prefill is faster from 2-token suffixes upward; keep the default at 4 - * as a conservative crossover. The env knob remains useful for retuning. */ -static uint32_t metal_graph_resume_prefill_min_tokens(void) { -#ifndef DS4_ROCM_BUILD - const char *env = getenv("DS4_METAL_RESUME_PREFILL_MIN"); - if (env && env[0]) { - char *endp = NULL; - const long v = strtol(env, &endp, 10); - if (endp != env) { - if (v <= 0) return UINT32_MAX; - return (uint32_t)v; - } - } -#endif - return 4u; -} - -static uint32_t glm_graph_resume_prefill_min_tokens(void) { -#ifndef DS4_ROCM_BUILD - const char *env = getenv("DS4_GLM_RESUME_PREFILL_MIN"); - if (env && env[0]) { - char *endp = NULL; - const long v = strtol(env, &endp, 10); - if (endp != env) { - if (v <= 0) return UINT32_MAX; - return (uint32_t)v; - } - } -#endif - return 4u; -} - -#define DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT 4096u -#define DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT 8192u -#define DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT 2048u -#define DS4_GLM_METAL_DISPLAY_PROGRESS_LAYER_TOKENS 32u -#define DS4_GLM_METAL_SMALL_PREFILL_STAGE_SYNC_TOKENS 0u -#define DS4_GLM_METAL_LONG_CONTEXT_THRESHOLD 65536u -#define DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT 4096u -#define DS4_GLM_METAL_INDEXED_PREFILL_CHUNK_TOKENS 4096u -#define DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB 256u - -static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, - bool ssd_streaming); -static uint32_t glm_graph_indexed_prefill_chunk_tokens( - uint32_t full_attention_cap, - uint32_t compact_cap); -static uint32_t glm_graph_indexed_prefill_score_tokens( - uint32_t indexed_prefill_cap, - uint32_t compact_cap); - -static uint64_t glm_graph_compact_cache_elem_bytes(void) { - return DS4_GPU_GLM_COMPACT_CACHE_F16 ? sizeof(uint16_t) : sizeof(float); -} - -static uint32_t glm_graph_compact_cache_is_f16(void) { - return DS4_GPU_GLM_COMPACT_CACHE_F16 ? 1u : 0u; -} - -static bool glm_graph_expanded_kv_cache_enabled(bool ssd_streaming) { - (void)ssd_streaming; - return false; -} - -static bool glm_graph_layer_uses_full_indexer(uint32_t il) { - if (il < DS4_N_LEADING_DENSE) return true; - return il >= 6u && ((il - 6u) % 4u) == 0u; -} - -static uint32_t glm_graph_normal_layer_count(void) { - if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT || DS4_N_LAYER > DS4_MAX_LAYER) { - return 0; - } - return DS4_N_LAYER - DS4_N_NEXTN_PREDICT; -} - -static uint32_t glm_graph_full_indexer_layer_count_range(uint32_t layer_start, - uint32_t layer_end) { - if (layer_start > layer_end) return 0; - uint32_t n = 0; - for (uint32_t il = layer_start; il <= layer_end; il++) { - if (glm_graph_layer_uses_full_indexer(il)) n++; - } - return n; -} - -static uint64_t glm_graph_full_kv_cache_elem_bytes(void) { - return sizeof(uint16_t); -} - -static uint32_t glm_graph_indexer_top_k_limit(void) { - return DS4_N_INDEXER_TOP_K; -} - -static uint32_t glm_tp_head_split_min(void) { - static int cached = -1; - if (cached < 0) { - cached = 64; - const char *env = getenv("DS4_GLM_TP_HEAD_SPLIT_MIN"); - if (env && env[0]) cached = atoi(env); - if (cached < 0) cached = 0; - } - return (uint32_t)cached; -} - -/* Correctness isolation: dump a hidden row (pre-output-norm), overwriting - * on each call — run with -n 0 so the file ends as the final prompt row. */ -static void glm_debug_dump_hidden_row(const ds4_gpu_tensor *t, uint32_t row) { - const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); - if (!path || !path[0] || !t) return; - float *buf = malloc((size_t)DS4_N_EMBD * sizeof(float)); - if (!buf) return; - if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, - (uint64_t)row * DS4_N_EMBD * sizeof(float), - buf, - (uint64_t)DS4_N_EMBD * sizeof(float))) { - FILE *f = fopen(path, "wb"); - if (f) { - fwrite(buf, sizeof(float), (size_t)DS4_N_EMBD, f); - fclose(f); - } - } - free(buf); -} - -/* Layer bisect: which layer's output hidden to dump (-1 = final/off, - * -2 = every layer, one file per layer). */ -static int glm_debug_hidden_dump_layer(void) { - const char *v = getenv("DS4_GLM_HIDDEN_DUMP_LAYER"); - if (!v || !v[0]) return -1; - if (strcmp(v, "all") == 0) return -2; - return atoi(v); -} - -static bool glm_debug_hidden_dump_layer_match(uint32_t il) { - const int dl = glm_debug_hidden_dump_layer(); - return dl == -2 || dl == (int)il; -} - -static void glm_debug_dump_raw_layer(const ds4_gpu_tensor *t, - const char *tag, - uint64_t bytes, - uint32_t il, - int pos) { - const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); - if (!path || !path[0] || !t) return; - char full[1024]; - if (pos >= 0) - snprintf(full, sizeof(full), "%s.%s.L%02u.T%02u", path, tag, il, pos); - else - snprintf(full, sizeof(full), "%s.%s.L%02u", path, tag, il); - void *buf = malloc(bytes); - if (!buf) return; - if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, 0, buf, bytes)) { - FILE *f = fopen(full, "wb"); - if (f) { fwrite(buf, 1, bytes, f); fclose(f); } - } - free(buf); -} - -static void glm_debug_dump_hidden_layer(const ds4_gpu_tensor *t, - uint32_t row, - uint32_t il, - uint32_t pos) { - const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); - if (!path || !path[0] || !t) return; - char full[1024]; - snprintf(full, sizeof(full), "%s.L%02u.T%02u", path, il, pos); - float *buf = malloc((size_t)DS4_N_EMBD * sizeof(float)); - if (!buf) return; - if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, - (uint64_t)row * DS4_N_EMBD * sizeof(float), - buf, - (uint64_t)DS4_N_EMBD * sizeof(float))) { - FILE *f = fopen(full, "wb"); - if (f) { - fwrite(buf, sizeof(float), (size_t)DS4_N_EMBD, f); - fclose(f); - } - } - free(buf); -} - -/* Correctness isolation: dump the post-prefill logits vector once. */ -static void glm_debug_dump_prefill_logits(const float *logits) { - static int dumped; - const char *path = getenv("DS4_GLM_LOGIT_DUMP"); - if (!path || !path[0] || dumped || !logits) return; - FILE *f = fopen(path, "wb"); - if (!f) return; - fwrite(logits, sizeof(float), (size_t)DS4_N_VOCAB, f); - fclose(f); - dumped = 1; - fprintf(stderr, "ds4: prefill logits dumped to %s\n", path); -} - -static bool glm_graph_indexed_prefill_trace_enabled(void) { - return false; -} - -static bool glm_graph_indexed_prefill_trace_all(void) { - return false; -} - -static uint32_t glm_graph_indexed_prefill_trace_slow_ms(void) { - return 100u; -} - -static uint32_t glm_graph_indexed_prefill_drain_interval(void) { - return 16u; -} - -static bool glm_graph_full_prefill_trace_enabled(void) { - return false; -} - -static bool glm_graph_full_prefill_trace_all(void) { - return false; -} - -static uint32_t glm_graph_full_prefill_trace_slow_ms(void) { - return 100u; -} - -static uint32_t glm_graph_full_prefill_drain_interval(void) { - return 16u; -} - -static void glm_graph_full_prefill_tracef(const char *fmt, ...) { - va_list ap; - va_start(ap, fmt); - fprintf(stderr, "ds4: GLM full prefill trace "); - vfprintf(stderr, fmt, ap); - fputc('\n', stderr); - fflush(stderr); - va_end(ap); -} - -static void glm_graph_indexed_prefill_tracef(const char *fmt, ...) { - va_list ap; - va_start(ap, fmt); - fprintf(stderr, "ds4: GLM indexed prefill trace "); - vfprintf(stderr, fmt, ap); - fputc('\n', stderr); - fflush(stderr); - va_end(ap); -} - -static uint32_t glm_graph_compact_cache_initial_cap( - uint32_t ctx_size, - uint32_t full_attention_cap) { - if (ctx_size == 0) return 0; - if (ctx_size <= full_attention_cap) return ctx_size; - - uint32_t cap = ctx_size; - if (cap == 0) cap = full_attention_cap ? full_attention_cap : 1u; - if (cap > ctx_size) cap = ctx_size; - return cap; -} - -static uint64_t glm_graph_compact_cache_bytes_for_cap( - uint32_t normal_layers, - uint32_t indexer_layers, - uint32_t compact_cap) { - if (compact_cap == 0) return 0; - const uint64_t elem = glm_graph_compact_cache_elem_bytes(); - uint64_t total = - (uint64_t)normal_layers * - compact_cap * - ((uint64_t)DS4_N_KV_LORA + DS4_N_ROT) * - elem; - total += - (uint64_t)indexer_layers * - compact_cap * - DS4_N_INDEXER_HEAD_DIM * - elem; - return total; -} - -static uint64_t glm_graph_indexed_scratch_bytes_for_cap( - uint32_t full_attention_cap, - uint32_t compact_cap) { - if (compact_cap == 0) return 0; - const uint64_t indexed_rows = - glm_graph_indexed_prefill_chunk_tokens(full_attention_cap, compact_cap); - const uint64_t indexed_score_rows = - glm_graph_indexed_prefill_score_tokens((uint32_t)indexed_rows, - compact_cap); - const uint64_t indexer_q_elems = - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - const uint64_t qk_low_elems = - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA; - uint64_t bytes = (uint64_t)compact_cap * sizeof(float); - const uint64_t indexer_top_k = glm_graph_indexer_top_k_limit(); - bytes += indexed_score_rows * (uint64_t)compact_cap * sizeof(float); - bytes += indexed_rows * indexer_q_elems * sizeof(float); - bytes += indexed_rows * DS4_N_INDEXER_HEAD * sizeof(float); - bytes += indexed_rows * indexer_top_k * sizeof(uint32_t); /* batch_indexer_selected */ - bytes += indexed_rows * qk_low_elems * sizeof(float); /* batch_qk_low */ - bytes += indexed_rows * qk_low_elems * sizeof(float); /* batch_attn_lora */ - return bytes; -} - -static uint64_t glm_graph_workspace_add_bytes( - uint64_t total, - uint64_t count, - uint64_t elem_bytes) { - return ds4_add_sat_u64(total, ds4_mul_sat_u64(count, elem_bytes)); -} - -static uint32_t glm_graph_indexed_decode_split_blocks(void); - -static uint64_t glm_graph_workspace_bytes_for_cap( - uint32_t full_attention_cap, - uint32_t compact_cap, - bool ssd_streaming) { - const bool expanded_kv = - glm_graph_expanded_kv_cache_enabled(ssd_streaming); - const uint64_t indexed_rows = - compact_cap != 0 ? - glm_graph_indexed_prefill_chunk_tokens(full_attention_cap, - compact_cap) : - 0; - const uint64_t batch_rows = - expanded_kv || indexed_rows == 0 ? full_attention_cap : indexed_rows; - const uint64_t indexer_top_k = glm_graph_indexer_top_k_limit(); - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; - const uint64_t q_nope = - DS4_N_KEY_MLA > DS4_N_ROT ? (uint64_t)DS4_N_KEY_MLA - DS4_N_ROT : 0; - const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; - const uint64_t kv_raw_dim = (uint64_t)DS4_N_KV_LORA + DS4_N_ROT; - uint64_t dense_hidden_max = - DS4_N_FF_DENSE > DS4_N_FF_EXP ? DS4_N_FF_DENSE : DS4_N_FF_EXP; - if (dense_hidden_max == 0) dense_hidden_max = DS4_N_FF_EXP; - const uint64_t sparse_mid_elems = - (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; - const uint64_t ffn_mid_elems = - dense_hidden_max > sparse_mid_elems ? - dense_hidden_max : - sparse_mid_elems; - const uint64_t qk_low_elems = - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA; - const uint64_t split_attn_blocks = - glm_graph_indexed_decode_split_blocks(); - - uint64_t bytes = - glm_graph_indexed_scratch_bytes_for_cap(full_attention_cap, - compact_cap); - - bytes = glm_graph_workspace_add_bytes(bytes, 3u, DS4_N_EMBD * sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, 2u, DS4_N_LORA_Q * sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, q_dim, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_INDEXER_HEAD_DIM, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - (uint64_t)DS4_N_INDEXER_HEAD * - DS4_N_INDEXER_HEAD_DIM, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_INDEXER_HEAD, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, indexer_top_k, sizeof(uint32_t)); - bytes = glm_graph_workspace_add_bytes(bytes, qk_low_elems, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - split_attn_blocks * qk_low_elems, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - split_attn_blocks * - DS4_N_HEAD * 2u, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, kv_raw_dim, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_KV_LORA, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - DS4_N_HEAD * q_nope, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, 2u * heads_dim, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, 6u, DS4_N_EMBD * sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, 2u * dense_hidden_max, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, ffn_mid_elems, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT * 2u, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT_USED, sizeof(int32_t)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT_USED, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_VOCAB, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - (uint64_t)DS4_N_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED, - sizeof(int32_t)); - - bytes = glm_graph_workspace_add_bytes(bytes, batch_rows, sizeof(int32_t)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_EXPERT * 2u, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_EXPERT_USED, - sizeof(int32_t)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_EXPERT_USED, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_EMBD * 7u, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_LORA_Q * 2u, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * q_dim, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_INDEXER_HEAD_DIM, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * kv_raw_dim, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_KV_LORA, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_HEAD * q_nope, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * heads_dim * 2u, sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * dense_hidden_max * 2u, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * DS4_N_FF_EXP, - sizeof(float)); - bytes = glm_graph_workspace_add_bytes(bytes, - batch_rows * ffn_mid_elems, - sizeof(float)); - if (indexed_rows != 0) { - bytes = glm_graph_workspace_add_bytes(bytes, - indexed_rows * qk_low_elems, - sizeof(float)); - } - return bytes; -} - -static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap_slice( - uint32_t ctx, - uint32_t work_ctx, - uint32_t compact_cap, - bool ssd_streaming, - uint32_t layer_start, - uint32_t layer_end) { - ds4_context_memory m = {0}; - const uint32_t normal_layers = glm_graph_normal_layer_count(); - if (normal_layers == 0 || layer_start >= normal_layers || - layer_end < layer_start) { - return m; - } - if (layer_end >= normal_layers) layer_end = normal_layers - 1u; - const uint32_t layer_count = layer_end - layer_start + 1u; - if (compact_cap > ctx) compact_cap = ctx; - - const bool expanded_kv = glm_graph_expanded_kv_cache_enabled(ssd_streaming); - const uint32_t indexed_rows = - compact_cap != 0 ? - glm_graph_indexed_prefill_chunk_tokens(work_ctx, compact_cap) : - 0; - const uint32_t batch_rows = - expanded_kv || indexed_rows == 0 ? work_ctx : indexed_rows; - - m.prefill_cap = batch_rows; - m.raw_cap = expanded_kv ? work_ctx : 0; - if (expanded_kv) { - m.raw_bytes = (uint64_t)layer_count * - work_ctx * - ((uint64_t)DS4_N_HEAD * (DS4_N_KEY_MLA + DS4_N_VALUE_MLA)) * - glm_graph_full_kv_cache_elem_bytes(); - } - m.scratch_bytes = - glm_graph_workspace_bytes_for_cap(work_ctx, - compact_cap, - ssd_streaming); - if (compact_cap != 0) { - m.comp_cap = compact_cap; - m.compressed_bytes = - glm_graph_compact_cache_bytes_for_cap( - layer_count, - glm_graph_full_indexer_layer_count_range(layer_start, - layer_end), - compact_cap); - } - m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; - return m; -} - -static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( - uint32_t ctx, - uint32_t work_ctx, - uint32_t compact_cap, - bool ssd_streaming) { - const uint32_t normal_layers = glm_graph_normal_layer_count(); - if (normal_layers == 0) { - const ds4_context_memory empty = {0}; - return empty; - } - return glm_graph_context_memory_estimate_for_compact_cap_slice( - ctx, - work_ctx, - compact_cap, - ssd_streaming, - 0, - normal_layers - 1u); -} - -ds4_context_memory ds4_context_memory_estimate_with_prefill_mode( - ds4_backend backend, - int ctx_size, - uint32_t prefill_chunk, - bool ssd_streaming) { - ds4_context_memory m = {0}; - uint32_t ctx = ctx_size > 0 ? (uint32_t)ctx_size : 1u; - - if (ds4_backend_uses_graph(backend)) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx, ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx, work_ctx); - m = glm_graph_context_memory_estimate_for_compact_cap(ctx, - work_ctx, - compact_cap, - ssd_streaming); - return m; - } - m.prefill_cap = metal_graph_prefill_cap_for_prompt((int)ctx, - prefill_chunk); - m.raw_cap = metal_graph_raw_cap_for_context((int)ctx, m.prefill_cap); - - uint32_t min_ratio = UINT32_MAX; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; - } - if (min_ratio == UINT32_MAX) min_ratio = ctx; - m.comp_cap = ctx / min_ratio + 2u; - if (m.comp_cap < 2u) m.comp_cap = 2u; - - m.raw_bytes = (uint64_t)DS4_N_LAYER * - m.raw_cap * - DS4_N_HEAD_DIM * - sizeof(float); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio == 0) continue; - const uint32_t layer_comp_cap = ctx / ratio + 2u; - m.compressed_bytes += (uint64_t)layer_comp_cap * - DS4_N_HEAD_DIM * - (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); - if (ratio == 4) { - m.compressed_bytes += (uint64_t)layer_comp_cap * - DS4_N_INDEXER_HEAD_DIM * - sizeof(float); - } - } - uint64_t attn_stage_cap = (uint64_t)(m.prefill_cap / min_ratio + 2u); - if (attn_stage_cap < 2u) attn_stage_cap = 2u; - m.scratch_bytes = 2ull * - m.comp_cap * - m.prefill_cap * - sizeof(float) + - attn_stage_cap * DS4_N_HEAD_DIM * sizeof(float); - } else { - m.raw_cap = ds4_default_raw_cap(ctx); - m.raw_bytes = (uint64_t)DS4_N_LAYER * - m.raw_cap * - DS4_N_HEAD_DIM * - sizeof(float); - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t ratio = ds4_layer_compress_ratio(il); - if (ratio == 0) continue; - const uint32_t comp_cap = ctx / ratio + 2u; - if (ratio == 4) m.comp_cap = comp_cap; - m.compressed_bytes += (uint64_t)comp_cap * - DS4_N_HEAD_DIM * - sizeof(float); - if (ratio == 4) { - m.compressed_bytes += (uint64_t)comp_cap * - DS4_N_INDEXER_HEAD_DIM * - sizeof(float); - } - } - if (m.comp_cap == 0) m.comp_cap = ctx / 4u + 2u; - m.scratch_bytes = ((uint64_t)(m.raw_cap + m.comp_cap) * sizeof(float)) + - ((uint64_t)m.comp_cap * sizeof(float)) + - ((uint64_t)m.comp_cap * sizeof(bool)); - } - - m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; - return m; -} - -ds4_context_memory ds4_context_memory_estimate_with_prefill( - ds4_backend backend, - int ctx_size, - uint32_t prefill_chunk) { - return ds4_context_memory_estimate_with_prefill_mode(backend, - ctx_size, - prefill_chunk, - false); -} - -ds4_context_memory ds4_context_memory_estimate(ds4_backend backend, - int ctx_size) { - return ds4_context_memory_estimate_with_prefill(backend, ctx_size, 0); -} - -static int metal_graph_prompt_logits_test( - const ds4_model *model, - const ds4_weights *weights, - const token_vec *prompt, - int ctx_size) { - int n_test = prompt->len; - const char *n_test_env = getenv("DS4_METAL_GRAPH_PROMPT_TOKENS"); - if (n_test_env && n_test_env[0]) { - char *endp = NULL; - const long v = strtol(n_test_env, &endp, 10); - if (endp != n_test_env && v > 0 && v <= prompt->len) n_test = (int)v; - } - - if (n_test <= 0 || n_test > ctx_size) { - fprintf(stderr, "ds4: Metal graph prompt test needs 1..%d prompt tokens\n", ctx_size); - return 1; - } - - const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, (uint32_t)n_test); - - ds4_gpu_graph g; - /* diagnostic single-tier callsite; placement=NULL. */ - bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], - raw_cap, (uint32_t)ctx_size, - (uint32_t)n_test, false, NULL, false, NULL); - if (!ok) { - metal_graph_free(&g); - fprintf(stderr, "ds4: failed to initialize Metal graph prompt test runtime\n"); - return 1; - } - const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; - if (memory_report) ds4_gpu_print_memory_report("after graph alloc"); - - ds4_kv_cache cpu_cache; - kv_cache_init(&cpu_cache, (uint32_t)ctx_size, raw_cap); - float *cpu_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); - float *gpu_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); - float *oracle_logits = NULL; - - const char *oracle_path = getenv("DS4_ORACLE_LOGITS"); - if (oracle_path && oracle_path[0]) { - oracle_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); - if (!read_f32_binary_file(oracle_path, oracle_logits, DS4_N_VOCAB)) { - free(oracle_logits); - oracle_logits = NULL; - } - } - - for (int t = 0; t < n_test; t++) { - const bool last = t == n_test - 1; - forward_token_raw_swa_cpu(last ? cpu_logits : NULL, - model, - weights, - &cpu_cache, - prompt->v[t], - (uint32_t)t); - } - ok = metal_graph_prefill_raw_swa(&g, model, weights, prompt, n_test, - gpu_logits, true, NULL, NULL, - NULL, NULL, NULL); - if (memory_report) ds4_gpu_print_memory_report("after prompt graph"); - - if (ok) { - const char *dump_gpu = getenv("DS4_METAL_GRAPH_DUMP_LOGITS"); - if (dump_gpu && dump_gpu[0]) { - if (write_f32_binary_file(dump_gpu, gpu_logits, DS4_N_VOCAB)) { - fprintf(stderr, "ds4: wrote Metal graph logits to %s\n", dump_gpu); - } - } - const char *dump_cpu = getenv("DS4_CPU_DUMP_LOGITS"); - if (dump_cpu && dump_cpu[0]) { - if (write_f32_binary_file(dump_cpu, cpu_logits, DS4_N_VOCAB)) { - fprintf(stderr, "ds4: wrote CPU logits to %s\n", dump_cpu); - } - } - if (getenv("DS4_METAL_GRAPH_TRACE_CACHE") != NULL || - getenv("DS4_METAL_GRAPH_TRACE_COMP") != NULL) { - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - const uint32_t n_raw = cpu_cache.layer[il].n_raw; - if (n_raw != 0) { - const uint64_t raw_phys_n = (uint64_t)raw_cap * DS4_N_HEAD_DIM; - const uint64_t raw_logical_n = (uint64_t)n_raw * DS4_N_HEAD_DIM; - const uint32_t raw_start = n_raw < raw_cap ? 0u : ((uint32_t)n_test % raw_cap); - float *gpu_raw_phys = xmalloc((size_t)raw_phys_n * sizeof(float)); - float *gpu_raw_logical = xmalloc((size_t)raw_logical_n * sizeof(float)); - if (ds4_gpu_tensor_read(g.layer_raw_cache[il], 0, gpu_raw_phys, raw_phys_n * sizeof(float)) != 0) { - for (uint32_t r = 0; r < n_raw; r++) { - const uint32_t phys = (raw_start + r) % raw_cap; - memcpy(gpu_raw_logical + (uint64_t)r * DS4_N_HEAD_DIM, - gpu_raw_phys + (uint64_t)phys * DS4_N_HEAD_DIM, - (size_t)DS4_N_HEAD_DIM * sizeof(float)); - } - fprintf(stderr, - "ds4: cache trace layer %u raw_n=%u raw_start=%u raw_max=%g raw_rms=%g\n", - il, n_raw, raw_start, - max_abs_diff(cpu_cache.layer[il].raw_kv, gpu_raw_logical, raw_logical_n), - rms_abs_diff(cpu_cache.layer[il].raw_kv, gpu_raw_logical, raw_logical_n)); - } - free(gpu_raw_logical); - free(gpu_raw_phys); - } - - const uint32_t n_comp = cpu_cache.layer[il].n_comp; - if (n_comp == 0) continue; - const uint64_t n = (uint64_t)n_comp * DS4_N_HEAD_DIM; - float *gpu_comp = xmalloc((size_t)n * sizeof(float)); - bool comp_read = false; - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - uint16_t *gpu_comp_h = xmalloc((size_t)n * sizeof(uint16_t)); - if (ds4_gpu_tensor_read(g.layer_attn_comp_cache[il], 0, - gpu_comp_h, n * sizeof(uint16_t)) != 0) { - for (uint64_t i = 0; i < n; i++) gpu_comp[i] = f16_to_f32(gpu_comp_h[i]); - comp_read = true; - } - free(gpu_comp_h); - } else { - comp_read = ds4_gpu_tensor_read(g.layer_attn_comp_cache[il], 0, - gpu_comp, n * sizeof(float)) != 0; - } - if (comp_read) { - fprintf(stderr, - "ds4: comp trace layer %u n=%u attn_max=%g attn_rms=%g\n", - il, n_comp, - max_abs_diff(cpu_cache.layer[il].attn_comp_kv, gpu_comp, n), - rms_abs_diff(cpu_cache.layer[il].attn_comp_kv, gpu_comp, n)); - } - free(gpu_comp); - - const uint32_t n_index = cpu_cache.layer[il].n_index_comp; - if (n_index != 0 && g.layer_index_comp_cache[il]) { - const uint64_t ni = (uint64_t)n_index * DS4_N_INDEXER_HEAD_DIM; - float *gpu_index = xmalloc((size_t)ni * sizeof(float)); - if (ds4_gpu_tensor_read(g.layer_index_comp_cache[il], 0, gpu_index, ni * sizeof(float)) != 0) { - fprintf(stderr, - "ds4: comp trace layer %u n=%u index_max=%g index_rms=%g\n", - il, n_index, - max_abs_diff(cpu_cache.layer[il].index_comp_kv, gpu_index, ni), - rms_abs_diff(cpu_cache.layer[il].index_comp_kv, gpu_index, ni)); - } - free(gpu_index); - } - } - } - const uint64_t cpu_top = argmax_f32(cpu_logits, DS4_N_VOCAB); - const uint64_t gpu_top = argmax_f32(gpu_logits, DS4_N_VOCAB); - fprintf(stderr, - "ds4: Metal prompt graph logits: tokens=%d logits_max=%g logits_rms=%g cpu_top=%llu gpu_top=%llu cpu_top_logit=%g gpu_top_logit=%g\n", - n_test, - max_abs_diff(cpu_logits, gpu_logits, DS4_N_VOCAB), - rms_abs_diff(cpu_logits, gpu_logits, DS4_N_VOCAB), - (unsigned long long)cpu_top, - (unsigned long long)gpu_top, - cpu_logits[cpu_top], - gpu_logits[gpu_top]); - if (oracle_logits) { - const uint64_t oracle_top = argmax_f32(oracle_logits, DS4_N_VOCAB); - fprintf(stderr, - "ds4: oracle logits: tokens=%d oracle_top=%llu oracle_top_logit=%g cpu_max=%g cpu_rms=%g metal_max=%g metal_rms=%g\n", - n_test, - (unsigned long long)oracle_top, - oracle_logits[oracle_top], - max_abs_diff(cpu_logits, oracle_logits, DS4_N_VOCAB), - rms_abs_diff(cpu_logits, oracle_logits, DS4_N_VOCAB), - max_abs_diff(gpu_logits, oracle_logits, DS4_N_VOCAB), - rms_abs_diff(gpu_logits, oracle_logits, DS4_N_VOCAB)); - } - } else { - fprintf(stderr, "ds4: Metal prompt graph logits test failed\n"); - if (ds4_gpu_synchronize() == 0) { - fprintf(stderr, "ds4: Metal synchronize after prompt graph failure also failed\n"); - } - } - - free(gpu_logits); - free(cpu_logits); - free(oracle_logits); - kv_cache_free(&cpu_cache); - metal_graph_free(&g); - return ok ? 0 : 1; -} - -#endif - -typedef struct ds4_vocab ds4_vocab; - -static void embed_prompt( - const ds4_model * model, - const ds4_weights * weights, - const token_vec * tokens, - uint32_t n_embd, - float * out) { - for (int i = 0; i < tokens->len; i++) { - embed_token_any(model, weights, tokens->v[i], out + (uint64_t)i * n_embd); - } -} - -/* ========================================================================= - * Tokenizer and Chat Prompt Encoding. - * ========================================================================= - * - * DeepSeek V4 Flash stores a GPT-2 style byte-level BPE tokenizer in GGUF. - * The implementation below is intentionally small. It loads token strings - * and merge ranks from the mmaped file, builds two open-addressed hash tables, - * and applies BPE to user text. Chat special tokens are inserted directly by - * ID; user text goes through BPE. - */ - -typedef struct { - ds4_str key; - int value; - bool used; -} str_i32_entry; - -typedef struct { - str_i32_entry *entry; - uint64_t cap; - uint64_t used; -} str_i32_table; - -static uint64_t next_pow2(uint64_t n) { - uint64_t p = 1; - while (p < n) p <<= 1; - return p; -} - -static void table_init(str_i32_table *t, uint64_t expected) { - t->cap = next_pow2(expected * 2 + 16); - t->used = 0; - t->entry = xcalloc((size_t)t->cap, sizeof(t->entry[0])); -} - -static void table_free(str_i32_table *t) { - free(t->entry); - memset(t, 0, sizeof(*t)); -} - -static void table_put(str_i32_table *t, ds4_str key, int value) { - uint64_t mask = t->cap - 1; - uint64_t i = hash_bytes(key.ptr, key.len) & mask; - - while (t->entry[i].used) { - if (ds4_str_eq(t->entry[i].key, key)) { - t->entry[i].value = value; - return; - } - i = (i + 1) & mask; - } - - t->entry[i].used = true; - t->entry[i].key = key; - t->entry[i].value = value; - t->used++; -} - -static bool table_get(const str_i32_table *t, const char *ptr, uint64_t len, int *value) { - if (t->cap == 0) return false; - - uint64_t mask = t->cap - 1; - uint64_t i = hash_bytes(ptr, len) & mask; - - while (t->entry[i].used) { - ds4_str key = t->entry[i].key; - if (key.len == len && memcmp(key.ptr, ptr, len) == 0) { - *value = t->entry[i].value; - return true; - } - i = (i + 1) & mask; - } - return false; -} - -static void token_vec_push(token_vec *tv, int token) { - if (tv->len == tv->cap) { - tv->cap = tv->cap ? tv->cap * 2 : 64; - tv->v = xrealloc(tv->v, (size_t)tv->cap * sizeof(tv->v[0])); - } - tv->v[tv->len++] = token; -} - -static void token_vec_free(token_vec *tv) { - free(tv->v); - memset(tv, 0, sizeof(*tv)); -} - -void ds4_tokens_push(ds4_tokens *tv, int token) { - token_vec_push(tv, token); -} - -void ds4_tokens_free(ds4_tokens *tv) { - token_vec_free(tv); -} - -void ds4_tokens_copy(ds4_tokens *dst, const ds4_tokens *src) { - dst->len = 0; - for (int i = 0; i < src->len; i++) token_vec_push(dst, src->v[i]); -} - -bool ds4_tokens_starts_with(const ds4_tokens *tokens, const ds4_tokens *prefix) { - if (prefix->len > tokens->len) return false; - for (int i = 0; i < prefix->len; i++) { - if (tokens->v[i] != prefix->v[i]) return false; - } - return true; -} - -struct ds4_vocab { - ds4_str *token; - int n_vocab; - int bos_id; - int eos_id; - int system_id; - int user_id; - int assistant_id; - int observation_id; - int sop_id; - int think_start_id; - int think_end_id; - int tool_call_start_id; - int tool_call_end_id; - int tool_response_start_id; - int tool_response_end_id; - int arg_key_start_id; - int arg_key_end_id; - int arg_value_start_id; - int arg_value_end_id; - int dsml_id; - str_i32_table token_to_id; - str_i32_table merge_rank; -}; - -/* Engine-side tensor-parallel state. The transport context is owned by the - * frontend (CLI leader or ds4_tp_worker_run); the engine owns the GPU slab, - * the per-slot views and the gate machinery lifetime. */ -typedef struct { - struct ds4_tp *ctx; - ds4_gpu_tensor *slab; - ds4_gpu_tensor **out_views; - ds4_gpu_tensor **in_views; - ds4_gpu_tensor **batch_out_views; /* [layer] verify-block row partials */ - ds4_gpu_tensor **batch_in_views; - ds4_gpu_tensor *zero_vec; - uint64_t eval_seq; /* leader: mirrored eval counter */ - uint64_t next_session_id; /* leader: stable worker-session handle */ - int rank; - bool vocab_split; /* DS4-only: logits halves cross the wire */ - bool active; -} ds4_engine_tp_state; - -struct ds4_engine { - ds4_model model; - ds4_model mtp_model; - ds4_vocab vocab; - ds4_weights weights; - ds4_mtp_weights mtp_weights; - ds4_dspark_weights dspark_weights; - ds4_backend backend; - ds4_support_kind support_kind; - int dspark_exec_tier; - uint32_t support_stages; - int mtp_draft_tokens; - float mtp_margin; - float dspark_confidence_threshold; - char *directional_steering_file; - float *directional_steering_dirs; - float directional_steering_attn_scale; - float directional_steering_ffn_scale; - int power_percent; - uint32_t prefill_chunk; - uint32_t ssd_streaming_cache_experts; - uint64_t ssd_streaming_cache_bytes; - uint64_t ssd_streaming_prefill_headroom_bytes; - uint64_t ssd_streaming_full_layer_bytes; - uint32_t ssd_streaming_full_layers; - uint32_t ssd_streaming_preload_experts; - uint64_t startup_model_span_bytes; - ds4_ssd_memory_lock simulated_memory; - bool quality; - bool glm_mtp; - bool glm_mtp_timing; - bool dspark; - bool dspark_strict; - bool cuda_tensor_parallel; - bool glm_tp_token_prefill; - bool ssd_streaming; - bool ssd_streaming_cold; - bool ssd_streaming_full_layers_set; - ds4_distributed_options distributed; - ds4_engine_tp_state tp; - bool metal_ready; - bool mtp_ready; - bool share_session_prefill_workspace; -#ifndef DS4_NO_GPU - bool shared_prefill_workspace_ready; - ds4_gpu_graph shared_prefill_workspace; -#endif - - /* Wave-2 multi-GPU placement scaffolding: optional multi-GPU placement - * state. Zero-initialized for every existing caller (gpu_cfg == NULL) - * via xcalloc, so the single-tier path observes identical engine - * state to pre-multi-GPU CLI main. multi_tier == 1 is the gate for all - * new code paths. */ - ds4_gpu_config gpu_cfg; - int placement[DS4_MAX_LAYER + 2]; - int n_placement_entries; - int multi_tier; - - /* Max-context hint copied from - * ds4_engine_options.placement_ctx_hint. Used by - * engine_compute_entry_bytes for per-layer KV estimation. - * Zero / negative = legacy 4096 fallback (single-tier paths and any - * caller that doesn't set the option observe the prior behavior). */ - int placement_ctx_hint; -}; - -static uint64_t ds4_engine_dynamic_expert_cache_bytes( - const ds4_engine *e) { - if (!e || !e->ssd_streaming) return 0; - if (e->ssd_streaming_cache_bytes != 0) { - return e->ssd_streaming_cache_bytes; - } - if (e->ssd_streaming_cache_experts == 0) return 0; - - uint64_t per_expert_bytes = 0; - if (!ds4_streaming_routed_expert_bytes(&e->weights, - &per_expert_bytes)) { - return 0; - } - if (e->ssd_streaming_cache_experts > UINT64_MAX / per_expert_bytes) { - return UINT64_MAX; - } - return (uint64_t)e->ssd_streaming_cache_experts * per_expert_bytes; -} - -static uint64_t ds4_engine_streaming_transient_guard_bytes( - const ds4_engine *e) { - if (!e || !e->ssd_streaming) return 0; - uint64_t total = ds4_engine_dynamic_expert_cache_bytes(e); - total = ds4_add_sat_u64(total, e->ssd_streaming_full_layer_bytes); - total = ds4_add_sat_u64(total, e->ssd_streaming_prefill_headroom_bytes); - return total; -} - -static void ds4_engine_print_startup_memory( - const ds4_engine *e, - int ctx_size) { - if (!e || ctx_size <= 0) return; - - ds4_context_memory mem; -#ifndef DS4_NO_GPU - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - e->distributed.role != DS4_DISTRIBUTED_NONE && - e->distributed.layers.set) { - const uint32_t normal_layers = glm_graph_normal_layer_count(); - const uint32_t layer_end = e->distributed.layers.has_output ? - (normal_layers ? normal_layers - 1u : 0u) : - e->distributed.layers.end; - const uint32_t ctx = (uint32_t)ctx_size; - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx, e->ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx, work_ctx); - mem = glm_graph_context_memory_estimate_for_compact_cap_slice( - ctx, - work_ctx, - compact_cap, - e->ssd_streaming, - e->distributed.layers.start, - layer_end); - } else { -#endif - mem = ds4_context_memory_estimate_with_prefill_mode(e->backend, - ctx_size, - e->prefill_chunk, - e->ssd_streaming); -#ifndef DS4_NO_GPU - } -#endif - const uint64_t kv_bytes = - ds4_add_sat_u64(mem.raw_bytes, mem.compressed_bytes); - const uint64_t dynamic_expert_cache_bytes = - ds4_engine_dynamic_expert_cache_bytes(e); - const uint64_t expert_reserved_bytes = - e->ssd_streaming_prefill_headroom_bytes; - uint64_t total = kv_bytes; - total = ds4_add_sat_u64(total, mem.scratch_bytes); - total = ds4_add_sat_u64(total, e->startup_model_span_bytes); - total = ds4_add_sat_u64(total, dynamic_expert_cache_bytes); - total = ds4_add_sat_u64(total, e->ssd_streaming_full_layer_bytes); - total = ds4_add_sat_u64(total, expert_reserved_bytes); - - const bool color = ds4_log_is_tty(stderr); - const char *green = color ? "\x1b[32m" : ""; - const char *bright_green = color ? "\x1b[1;32m" : ""; - const char *reset = color ? "\x1b[0m" : ""; - - fprintf(stderr, - "%sds4: memory: KV %.2f GiB (raw %.2f + compressed %.2f) " - "+ buffers %.2f GiB + resident model %.2f GiB", - green, - ds4_bytes_to_gib(kv_bytes), - ds4_bytes_to_gib(mem.raw_bytes), - ds4_bytes_to_gib(mem.compressed_bytes), - ds4_bytes_to_gib(mem.scratch_bytes), - ds4_bytes_to_gib(e->startup_model_span_bytes)); - if (e->ssd_streaming_full_layer_bytes != 0) { - fprintf(stderr, - " + full-layer experts %.2f GiB", - ds4_bytes_to_gib(e->ssd_streaming_full_layer_bytes)); - } - if (dynamic_expert_cache_bytes != 0) { - fprintf(stderr, - " + expert cache %.2f GiB", - ds4_bytes_to_gib(dynamic_expert_cache_bytes)); - } - if (expert_reserved_bytes != 0) { - fprintf(stderr, - " + prefill expert reserve %.2f GiB", - ds4_bytes_to_gib(expert_reserved_bytes)); - } - fprintf(stderr, - " = %s%.2f GiB planned%s\n", - bright_green, - ds4_bytes_to_gib(total), - reset); - - fprintf(stderr, - "%sds4: memory detail: ctx=%d prefill_cap=%u raw_kv_rows=%u " - "compressed_kv_rows=%u backend=%s%s\n", - green, - ctx_size, - mem.prefill_cap, - mem.raw_cap, - mem.comp_cap, - ds4_backend_name(e->backend), - reset); -} - -static bool cpu_directional_steering_enabled( - const float *dirs, - float scale) { - return dirs && scale != 0.0f; -} - -static void cpu_directional_steering_project_rows( - float *x, - const float *dirs, - uint32_t il, - uint32_t rows, - float scale) { - if (!cpu_directional_steering_enabled(dirs, scale) || !x || rows == 0) return; - - const float *dir = dirs + (uint64_t)il * DS4_N_EMBD; - for (uint32_t row = 0; row < rows; row++) { - float *xr = x + (uint64_t)row * DS4_N_EMBD; - float dot = 0.0f; - for (uint32_t i = 0; i < DS4_N_EMBD; i++) { - dot += xr[i] * dir[i]; - } - const float coeff = scale * dot; - for (uint32_t i = 0; i < DS4_N_EMBD; i++) { - xr[i] -= coeff * dir[i]; - } - } -} - -static bool cpu_load_directional_steering(ds4_engine *e) { - if (!e || - (e->directional_steering_attn_scale == 0.0f && - e->directional_steering_ffn_scale == 0.0f)) { - return true; - } - - const char *path = e->directional_steering_file; - if (!path || !path[0]) { - fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); - return false; - } - - const uint64_t n = (uint64_t)DS4_N_LAYER * DS4_N_EMBD; - e->directional_steering_dirs = xmalloc((size_t)n * sizeof(e->directional_steering_dirs[0])); - if (!read_f32_binary_file(path, e->directional_steering_dirs, n)) { - free(e->directional_steering_dirs); - e->directional_steering_dirs = NULL; - fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); - return false; - } - fprintf(stderr, "ds4: CPU directional steering enabled: %s attn=%g ffn=%g\n", - path, - (double)e->directional_steering_attn_scale, - (double)e->directional_steering_ffn_scale); - return true; -} - -static void utf8_put(char **p, uint32_t cp) { - if (cp <= 0x7f) { - *(*p)++ = (char)cp; - } else if (cp <= 0x7ff) { - *(*p)++ = (char)(0xc0 | (cp >> 6)); - *(*p)++ = (char)(0x80 | (cp & 0x3f)); - } else if (cp <= 0xffff) { - *(*p)++ = (char)(0xe0 | (cp >> 12)); - *(*p)++ = (char)(0x80 | ((cp >> 6) & 0x3f)); - *(*p)++ = (char)(0x80 | (cp & 0x3f)); - } else { - *(*p)++ = (char)(0xf0 | (cp >> 18)); - *(*p)++ = (char)(0x80 | ((cp >> 12) & 0x3f)); - *(*p)++ = (char)(0x80 | ((cp >> 6) & 0x3f)); - *(*p)++ = (char)(0x80 | (cp & 0x3f)); - } -} - -static uint32_t gpt2_byte_to_codepoint(uint8_t b) { - if ((b >= 33 && b <= 126) || (b >= 161 && b <= 172) || (b >= 174)) { - return b; - } - - uint32_t n = 0; - for (uint32_t x = 0; x < 256; x++) { - if ((x >= 33 && x <= 126) || (x >= 161 && x <= 172) || (x >= 174)) { - continue; - } - if (x == b) return 256 + n; - n++; - } - return b; -} - -/* GPT-2 byte-level BPE first maps raw bytes to printable Unicode codepoints - * so merges can operate on UTF-8 strings without losing byte identity. */ -static char *byte_encode(ds4_str in, uint64_t *out_len) { - char *out = xmalloc((size_t)in.len * 4 + 1); - char *p = out; - - for (uint64_t i = 0; i < in.len; i++) { - utf8_put(&p, gpt2_byte_to_codepoint((uint8_t)in.ptr[i])); - } - *p = '\0'; - *out_len = (uint64_t)(p - out); - return out; -} - -static int utf8_len_from_first_byte(uint8_t c) { - if (c < 0x80) return 1; - if ((c & 0xe0) == 0xc0) return 2; - if ((c & 0xf0) == 0xe0) return 3; - if ((c & 0xf8) == 0xf0) return 4; - return 1; -} - -typedef struct { - char *ptr; - uint64_t len; -} owned_str; - -static owned_str owned_copy(const char *ptr, uint64_t len) { - owned_str s; - s.ptr = xmalloc((size_t)len); - memcpy(s.ptr, ptr, (size_t)len); - s.len = len; - return s; -} - -/* Look up the merge rank for two adjacent BPE symbols. */ -static int bpe_rank(const ds4_vocab *vocab, const owned_str *a, const owned_str *b) { - uint64_t len = a->len + 1 + b->len; - char stack[512]; - char *buf = len <= sizeof(stack) ? stack : xmalloc((size_t)len); - - memcpy(buf, a->ptr, (size_t)a->len); - buf[a->len] = ' '; - memcpy(buf + a->len + 1, b->ptr, (size_t)b->len); - - int rank = -1; - table_get(&vocab->merge_rank, buf, len, &rank); - - if (buf != stack) free(buf); - return rank; -} - -/* Apply byte-level BPE to one regex-like pre-tokenized piece and emit token ids. */ -static void bpe_emit_piece(const ds4_vocab *vocab, ds4_str raw_piece, token_vec *out) { - uint64_t encoded_len = 0; - char *encoded = byte_encode(raw_piece, &encoded_len); - - int n_sym = 0; - int cap_sym = 32; - owned_str *sym = xcalloc((size_t)cap_sym, sizeof(sym[0])); - - for (uint64_t off = 0; off < encoded_len;) { - int n = utf8_len_from_first_byte((uint8_t)encoded[off]); - if (off + (uint64_t)n > encoded_len) n = 1; - if (n_sym == cap_sym) { - cap_sym *= 2; - sym = xrealloc(sym, (size_t)cap_sym * sizeof(sym[0])); - } - sym[n_sym++] = owned_copy(encoded + off, (uint64_t)n); - off += (uint64_t)n; - } - - for (;;) { - int best_i = -1; - int best_rank = INT32_MAX; - - for (int i = 0; i + 1 < n_sym; i++) { - int rank = bpe_rank(vocab, &sym[i], &sym[i + 1]); - if (rank >= 0 && rank < best_rank) { - best_rank = rank; - best_i = i; - } - } - - if (best_i < 0) break; - - owned_str merged; - merged.len = sym[best_i].len + sym[best_i + 1].len; - merged.ptr = xmalloc((size_t)merged.len); - memcpy(merged.ptr, sym[best_i].ptr, (size_t)sym[best_i].len); - memcpy(merged.ptr + sym[best_i].len, sym[best_i + 1].ptr, (size_t)sym[best_i + 1].len); - - free(sym[best_i].ptr); - free(sym[best_i + 1].ptr); - sym[best_i] = merged; - - for (int j = best_i + 1; j + 1 < n_sym; j++) { - sym[j] = sym[j + 1]; - } - n_sym--; - } - - for (int i = 0; i < n_sym; i++) { - int token = -1; - if (table_get(&vocab->token_to_id, sym[i].ptr, sym[i].len, &token)) { - token_vec_push(out, token); - } else { - for (uint64_t j = 0; j < sym[i].len; j++) { - if (table_get(&vocab->token_to_id, sym[i].ptr + j, 1, &token)) { - token_vec_push(out, token); - } - } - } - free(sym[i].ptr); - } - - free(sym); - free(encoded); -} - -static uint64_t next_utf8_char(const char *s, uint64_t len, uint64_t pos) { - int n = utf8_len_from_first_byte((uint8_t)s[pos]); - if (pos + (uint64_t)n > len) n = 1; - return pos + (uint64_t)n; -} - -static bool ascii_alpha(uint8_t c) { - return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z'); -} - -static bool ascii_digit(uint8_t c) { - return c >= '0' && c <= '9'; -} - -static bool ascii_space(uint8_t c) { - return c == ' ' || c == '\t' || c == '\n' || c == '\r' || - c == '\v' || c == '\f'; -} - -static bool ascii_newline(uint8_t c) { - return c == '\n' || c == '\r'; -} - -static bool joyai_ascii_punct_symbol(uint8_t c) { - return (c >= '!' && c <= '/') || - (c >= ':' && c <= '@') || - (c >= '[' && c <= '`') || - (c >= '{' && c <= '~'); -} - -static bool utf8_is_cjk_hira_kata(uint32_t cp) { - return (cp >= 0x4e00 && cp <= 0x9fa5) || - (cp >= 0x3040 && cp <= 0x309f) || - (cp >= 0x30a0 && cp <= 0x30ff); -} - -static uint32_t utf8_peek_one(const char *s, uint64_t len, uint64_t pos, uint64_t *next) { - const uint8_t c0 = (uint8_t)s[pos]; - int n = utf8_len_from_first_byte(c0); - if (pos + (uint64_t)n > len) n = 1; - *next = pos + (uint64_t)n; - - if (n == 1) return c0; - if (n == 2) { - return ((uint32_t)(c0 & 0x1f) << 6) | - ((uint32_t)((uint8_t)s[pos + 1] & 0x3f)); - } - if (n == 3) { - return ((uint32_t)(c0 & 0x0f) << 12) | - ((uint32_t)((uint8_t)s[pos + 1] & 0x3f) << 6) | - ((uint32_t)((uint8_t)s[pos + 2] & 0x3f)); - } - return ((uint32_t)(c0 & 0x07) << 18) | - ((uint32_t)((uint8_t)s[pos + 1] & 0x3f) << 12) | - ((uint32_t)((uint8_t)s[pos + 2] & 0x3f) << 6) | - ((uint32_t)((uint8_t)s[pos + 3] & 0x3f)); -} - -static bool joyai_letter_like_at(const char *s, uint64_t len, uint64_t pos) { - (void)len; - uint8_t c = (uint8_t)s[pos]; - if (c < 128) return ascii_alpha(c); - - /* - * The JoyAI tokenizer maps Unicode letters into a collapsed regex alphabet before - * applying the JoyAI pre-tokenizer. The prompts we care about are mostly - * ASCII, but treating non-ASCII non-control bytes as letters preserves the - * useful behavior for ordinary UTF-8 text such as Italian accents. CJK and - * kana are isolated by the JoyAI pre-tokenizer before the generic letter - * rule, below. - */ - return true; -} - -static uint64_t joyai_consume_letters(const char *s, uint64_t len, uint64_t pos) { - while (pos < len && joyai_letter_like_at(s, len, pos)) { - pos = next_utf8_char(s, len, pos); - } - return pos; -} - -static bool joyai_cjk_at(const char *s, uint64_t len, uint64_t pos) { - if ((uint8_t)s[pos] < 128) return false; - uint64_t next = pos; - uint32_t cp = utf8_peek_one(s, len, pos, &next); - return utf8_is_cjk_hira_kata(cp); -} - -typedef struct { - uint32_t cp; - uint64_t next; - bool valid; - bool is_letter; - bool is_number; - bool is_whitespace; -} glm4_char_info; - -static bool glm4_unicode_whitespace(uint32_t cp) { - if (cp < 128) return ascii_space((uint8_t)cp); - return cp == 0x0085 || - cp == 0x00a0 || - cp == 0x1680 || - (cp >= 0x2000 && cp <= 0x200a) || - cp == 0x2028 || - cp == 0x2029 || - cp == 0x202f || - cp == 0x205f || - cp == 0x3000; -} - -static bool glm4_unicode_number(uint32_t cp) { - if (cp < 128) return ascii_digit((uint8_t)cp); - return (cp >= 0x0660 && cp <= 0x0669) || - (cp >= 0x06f0 && cp <= 0x06f9) || - (cp >= 0x07c0 && cp <= 0x07c9) || - (cp >= 0x0966 && cp <= 0x096f) || - (cp >= 0x09e6 && cp <= 0x09ef) || - (cp >= 0x0a66 && cp <= 0x0a6f) || - (cp >= 0x0ae6 && cp <= 0x0aef) || - (cp >= 0x0b66 && cp <= 0x0b6f) || - (cp >= 0x0be6 && cp <= 0x0bef) || - (cp >= 0x0c66 && cp <= 0x0c6f) || - (cp >= 0x0ce6 && cp <= 0x0cef) || - (cp >= 0x0d66 && cp <= 0x0d6f) || - (cp >= 0x0de6 && cp <= 0x0def) || - (cp >= 0x0e50 && cp <= 0x0e59) || - (cp >= 0x0ed0 && cp <= 0x0ed9) || - (cp >= 0x0f20 && cp <= 0x0f29) || - (cp >= 0x1040 && cp <= 0x1049) || - (cp >= 0x1090 && cp <= 0x1099) || - (cp >= 0x17e0 && cp <= 0x17e9) || - (cp >= 0x1810 && cp <= 0x1819) || - (cp >= 0xff10 && cp <= 0xff19); -} - -static bool glm4_unicode_punct_symbol(uint32_t cp) { - if (cp < 128) return joyai_ascii_punct_symbol((uint8_t)cp); - return (cp >= 0x00a1 && cp <= 0x00a9) || - (cp >= 0x00ab && cp <= 0x00ac) || - (cp >= 0x00ae && cp <= 0x00b1) || - cp == 0x00b4 || - (cp >= 0x00b6 && cp <= 0x00b8) || - cp == 0x00bb || - cp == 0x00bf || - cp == 0x00d7 || - cp == 0x00f7 || - (cp >= 0x02c2 && cp <= 0x02df) || - (cp >= 0x02e5 && cp <= 0x02eb) || - (cp >= 0x02ed && cp <= 0x02ff) || - (cp >= 0x0375 && cp <= 0x037e) || - (cp >= 0x0384 && cp <= 0x0385) || - cp == 0x0387 || - (cp >= 0x055a && cp <= 0x055f) || - (cp >= 0x0589 && cp <= 0x058a) || - (cp >= 0x05be && cp <= 0x05c0) || - cp == 0x05c3 || - (cp >= 0x05c6 && cp <= 0x05c7) || - (cp >= 0x0609 && cp <= 0x060a) || - (cp >= 0x060c && cp <= 0x060d) || - cp == 0x061b || - (cp >= 0x061e && cp <= 0x061f) || - cp == 0x066a || - cp == 0x066d || - cp == 0x06d4 || - (cp >= 0x2000 && cp <= 0x206f) || - (cp >= 0x20a0 && cp <= 0x20cf) || - (cp >= 0x2100 && cp <= 0x214f) || - (cp >= 0x2190 && cp <= 0x23ff) || - (cp >= 0x2460 && cp <= 0x24ff) || - (cp >= 0x2500 && cp <= 0x2775) || - (cp >= 0x2794 && cp <= 0x2bff) || - (cp >= 0x2e00 && cp <= 0x2e7f) || - (cp >= 0x3000 && cp <= 0x303f) || - (cp >= 0xfd3e && cp <= 0xfd3f) || - (cp >= 0xfe10 && cp <= 0xfe6f) || - (cp >= 0xff01 && cp <= 0xff0f) || - (cp >= 0xff1a && cp <= 0xff20) || - (cp >= 0xff3b && cp <= 0xff40) || - (cp >= 0xff5b && cp <= 0xff65) || - (cp >= 0x1f000 && cp <= 0x1faff); -} - -static glm4_char_info glm4_char_at(const char *s, uint64_t len, uint64_t pos) { - glm4_char_info info; - memset(&info, 0, sizeof(info)); - if (pos >= len) return info; - - info.valid = true; - info.cp = utf8_peek_one(s, len, pos, &info.next); - info.is_whitespace = glm4_unicode_whitespace(info.cp); - info.is_number = glm4_unicode_number(info.cp); - if (info.cp < 128) { - info.is_letter = ascii_alpha((uint8_t)info.cp); - } else { - info.is_letter = - !info.is_whitespace && - !info.is_number && - !glm4_unicode_punct_symbol(info.cp); - } - return info; -} - -static uint32_t ascii_tolower_cp(uint32_t cp) { - if (cp >= 'A' && cp <= 'Z') return cp + ('a' - 'A'); - return cp; -} - -/* ChatGLM4/GLM pre-tokenization. GLM GGUFs use tokenizer.ggml.pre="glm4", - * which shares the llama3-style split shape used by llama.cpp's CHATGLM4 path. */ -static void bpe_tokenize_text_glm4(const ds4_vocab *vocab, const char *text, token_vec *out) { - const uint64_t len = strlen(text); - uint64_t pos = 0; - - while (pos < len) { - uint64_t start = pos; - glm4_char_info cur = glm4_char_at(text, len, pos); - - if (!cur.valid) break; - - if (cur.cp == '\'' && cur.next < len) { - glm4_char_info next = glm4_char_at(text, len, cur.next); - uint32_t n1 = ascii_tolower_cp(next.cp); - if (n1 == 's' || n1 == 't' || n1 == 'm' || n1 == 'd') { - pos = next.next; - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - if (next.valid && next.next < len) { - glm4_char_info next2 = glm4_char_at(text, len, next.next); - uint32_t n2 = ascii_tolower_cp(next2.cp); - if ((n1 == 'r' && n2 == 'e') || - (n1 == 'v' && n2 == 'e') || - (n1 == 'l' && n2 == 'l')) { - pos = next2.next; - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - } - } - - if (!(cur.cp == '\r' || cur.cp == '\n' || cur.is_number)) { - glm4_char_info next = glm4_char_at(text, len, cur.next); - if (cur.is_letter || next.is_letter) { - pos = cur.next; - while (pos < len) { - glm4_char_info scan = glm4_char_at(text, len, pos); - if (!scan.valid || !scan.is_letter) break; - pos = scan.next; - } - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - } - - if (cur.is_number) { - int ndigits = 0; - while (pos < len && ndigits < 3) { - glm4_char_info scan = glm4_char_at(text, len, pos); - if (!scan.valid || !scan.is_number) break; - pos = scan.next; - ndigits++; - } - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - - glm4_char_info punct = cur; - uint64_t punct_pos = pos; - if (cur.cp == ' ') { - punct_pos = cur.next; - punct = glm4_char_at(text, len, punct_pos); - } - if (punct.valid && - !punct.is_whitespace && - !punct.is_letter && - !punct.is_number) { - pos = punct_pos; - while (pos < len) { - glm4_char_info scan = glm4_char_at(text, len, pos); - if (!scan.valid || - scan.is_whitespace || - scan.is_letter || - scan.is_number) { - break; - } - pos = scan.next; - } - while (pos < len) { - glm4_char_info scan = glm4_char_at(text, len, pos); - if (!scan.valid || !(scan.cp == '\r' || scan.cp == '\n')) break; - pos = scan.next; - } - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - - if (cur.is_whitespace) { - uint64_t p = pos; - uint64_t last_newline_end = 0; - uint64_t last_ws_start = pos; - int nspace = 0; - while (p < len) { - glm4_char_info scan = glm4_char_at(text, len, p); - if (!scan.valid || !scan.is_whitespace) break; - last_ws_start = p; - if (scan.cp == '\r' || scan.cp == '\n') last_newline_end = scan.next; - p = scan.next; - nspace++; - } - if (last_newline_end) { - pos = last_newline_end; - } else if (nspace > 1 && p < len) { - pos = last_ws_start; - } else { - pos = p; - } - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - continue; - } - - pos = cur.next; - if (pos == start) pos = next_utf8_char(text, len, pos); - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - } -} - -/* - * DeepSeek V4 Flash declares tokenizer.ggml.pre = "joyai-llm". The split - * below mirrors the JoyAI BPE pre-tokenizer for the cases this model - * uses in normal text and source-code prompts: - * - * \p{N}{1,3} - * [CJK/Hiragana/Katakana]+ - * [P/S][A-Za-z]+ - * [^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+ - * ?[\p{P}\p{S}]+[\r\n]* - * \s*[\r\n]+ - * \s+(?!\S) - * \s+ - * - * The punctuation rule intentionally keeps trailing newlines in the same BPE - * word (for example ">;\n"). Splitting those newlines separately changes the - * token stream for code prompts and produces wrong long-context logits. - */ -static void bpe_tokenize_text(const ds4_vocab *vocab, const char *text, token_vec *out) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - bpe_tokenize_text_glm4(vocab, text, out); - return; - } - - const uint64_t len = strlen(text); - uint64_t pos = 0; - - while (pos < len) { - uint64_t start = pos; - uint8_t c = (uint8_t)text[pos]; - - if (ascii_digit(c)) { - int ndigits = 0; - while (pos < len && ascii_digit((uint8_t)text[pos]) && ndigits < 3) { - pos++; - ndigits++; - } - } else if (joyai_cjk_at(text, len, pos)) { - do { - pos = next_utf8_char(text, len, pos); - } while (pos < len && joyai_cjk_at(text, len, pos)); - } else if (joyai_ascii_punct_symbol(c) && - pos + 1 < len && - ascii_alpha((uint8_t)text[pos + 1])) { - pos++; - while (pos < len && ascii_alpha((uint8_t)text[pos])) pos++; - } else if (joyai_letter_like_at(text, len, pos)) { - pos = joyai_consume_letters(text, len, pos); - } else if (!ascii_newline(c) && - !joyai_ascii_punct_symbol(c) && - pos + 1 < len && - joyai_letter_like_at(text, len, pos + 1)) { - pos++; - pos = joyai_consume_letters(text, len, pos); - } else if (c == ' ' && - pos + 1 < len && - joyai_ascii_punct_symbol((uint8_t)text[pos + 1])) { - pos++; - while (pos < len && joyai_ascii_punct_symbol((uint8_t)text[pos])) pos++; - while (pos < len && ascii_newline((uint8_t)text[pos])) pos++; - } else if (joyai_ascii_punct_symbol(c)) { - while (pos < len && joyai_ascii_punct_symbol((uint8_t)text[pos])) pos++; - while (pos < len && ascii_newline((uint8_t)text[pos])) pos++; - } else if (ascii_space(c)) { - uint64_t p = pos; - uint64_t last_newline_end = 0; - while (p < len && ascii_space((uint8_t)text[p])) { - uint8_t sc = (uint8_t)text[p++]; - if (ascii_newline(sc)) last_newline_end = p; - } - if (last_newline_end) { - pos = last_newline_end; - } else if (p < len && p > pos + 1 && - (joyai_letter_like_at(text, len, p) || - joyai_ascii_punct_symbol((uint8_t)text[p]))) { - /* - * JoyAI lets a single leading space join the following word or - * punctuation run. For " int", the pre-tokenizer therefore emits - * " " then " int", not " " then "int". - */ - pos = p - 1; - } else { - pos = p; - } - } else { - pos = next_utf8_char(text, len, pos); - } - - if (pos == start) pos = next_utf8_char(text, len, pos); - bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); - } -} - -static int vocab_lookup(const ds4_vocab *vocab, const char *text) { - int token = -1; - if (!table_get(&vocab->token_to_id, text, strlen(text), &token)) { - fprintf(stderr, "ds4: required tokenizer token is missing: %s\n", text); - exit(1); - } - return token; -} - -static int vocab_lookup_optional(const ds4_vocab *vocab, const char *text) { - int token = -1; - if (!table_get(&vocab->token_to_id, text, strlen(text), &token)) return -1; - return token; -} - -/* Load token strings, special token ids, and merge ranks from GGUF metadata. */ - -static void vocab_load(ds4_vocab *vocab, const ds4_model *model) { - memset(vocab, 0, sizeof(*vocab)); - - ds4_array_ref tokens; - ds4_array_ref merges; - if (!model_get_array(model, "tokenizer.ggml.tokens", &tokens) || - tokens.type != GGUF_VALUE_STRING || - tokens.len > INT32_MAX) { - ds4_die("GGUF tokenizer token table is missing or invalid"); - } - if (!model_get_array(model, "tokenizer.ggml.merges", &merges) || - merges.type != GGUF_VALUE_STRING) { - ds4_die("GGUF tokenizer merge table is missing or invalid"); - } - - vocab->n_vocab = (int)tokens.len; - vocab->token = xcalloc((size_t)vocab->n_vocab, sizeof(vocab->token[0])); - table_init(&vocab->token_to_id, tokens.len); - - ds4_cursor c = cursor_at(model, tokens.data_pos); - for (int i = 0; i < vocab->n_vocab; i++) { - if (!cursor_string(&c, &vocab->token[i])) ds4_die(c.error); - table_put(&vocab->token_to_id, vocab->token[i], i); - } - - table_init(&vocab->merge_rank, merges.len); - c = cursor_at(model, merges.data_pos); - for (uint64_t i = 0; i < merges.len; i++) { - ds4_str merge; - if (!cursor_string(&c, &merge)) ds4_die(c.error); - table_put(&vocab->merge_rank, merge, (int)i); - } - - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - if (!model_get_token_id(model, "tokenizer.ggml.bos_token_id", &vocab->bos_id)) { - vocab->bos_id = vocab_lookup_optional(vocab, ""); - } - if (!model_get_token_id(model, "tokenizer.ggml.eos_token_id", &vocab->eos_id)) { - vocab->eos_id = vocab_lookup_optional(vocab, "<|endoftext|>"); - } - vocab->system_id = vocab_lookup_optional(vocab, "<|system|>"); - vocab->user_id = vocab_lookup_optional(vocab, "<|user|>"); - vocab->assistant_id = vocab_lookup_optional(vocab, "<|assistant|>"); - vocab->observation_id = vocab_lookup_optional(vocab, "<|observation|>"); - vocab->sop_id = vocab_lookup_optional(vocab, ""); - vocab->think_start_id = vocab_lookup_optional(vocab, ""); - vocab->think_end_id = vocab_lookup_optional(vocab, ""); - vocab->tool_call_start_id = vocab_lookup_optional(vocab, ""); - vocab->tool_call_end_id = vocab_lookup_optional(vocab, ""); - vocab->tool_response_start_id = vocab_lookup_optional(vocab, ""); - vocab->tool_response_end_id = vocab_lookup_optional(vocab, ""); - vocab->arg_key_start_id = vocab_lookup_optional(vocab, ""); - vocab->arg_key_end_id = vocab_lookup_optional(vocab, ""); - vocab->arg_value_start_id = vocab_lookup_optional(vocab, ""); - vocab->arg_value_end_id = vocab_lookup_optional(vocab, ""); - vocab->dsml_id = -1; - return; - } - - vocab->bos_id = vocab_lookup(vocab, "<｜begin▁of▁sentence｜>"); - vocab->eos_id = vocab_lookup(vocab, "<｜end▁of▁sentence｜>"); - vocab->system_id = -1; - vocab->user_id = vocab_lookup(vocab, "<｜User｜>"); - vocab->assistant_id = vocab_lookup(vocab, "<｜Assistant｜>"); - vocab->observation_id = -1; - vocab->sop_id = -1; - vocab->think_start_id = vocab_lookup(vocab, ""); - vocab->think_end_id = vocab_lookup(vocab, ""); - vocab->tool_call_start_id = -1; - vocab->tool_call_end_id = -1; - vocab->tool_response_start_id = -1; - vocab->tool_response_end_id = -1; - vocab->arg_key_start_id = -1; - vocab->arg_key_end_id = -1; - vocab->arg_value_start_id = -1; - vocab->arg_value_end_id = -1; - vocab->dsml_id = vocab_lookup(vocab, "｜DSML｜"); -} - -static void vocab_free(ds4_vocab *vocab) { - free(vocab->token); - table_free(&vocab->token_to_id); - table_free(&vocab->merge_rank); - memset(vocab, 0, sizeof(*vocab)); -} - -/* Build the DS4 chat prompt: BOS, optional system text, user prompt, assistant - * marker, and either or depending on the requested mode. Max - * thinking is only a prompt prefix: the model still enters through . */ -static void chat_push_bos_sequence(const ds4_vocab *vocab, token_vec *out) { - token_vec_push(out, vocab->bos_id); - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && vocab->sop_id >= 0) - token_vec_push(out, vocab->sop_id); -} - -const char *ds4_glm_reasoning_effort_text(ds4_think_mode mode) { - switch (mode) { - case DS4_THINK_HIGH: return "Reasoning Effort: High"; - case DS4_THINK_MAX: return "Reasoning Effort: Max"; - case DS4_THINK_NONE: return NULL; - } - return NULL; -} - -static void chat_push_think_prefix(const ds4_vocab *vocab, - ds4_think_mode think_mode, - token_vec *out) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - const char *effort = ds4_glm_reasoning_effort_text(think_mode); - if (effort) { - token_vec_push(out, vocab->system_id); - bpe_tokenize_text(vocab, effort, out); - } - } else if (think_mode == DS4_THINK_MAX) { - bpe_tokenize_text(vocab, DS4_REASONING_EFFORT_MAX_PREFIX, out); - } -} - -static void encode_chat_prompt( - const ds4_vocab *vocab, - const char *system, - const char *prompt, - ds4_think_mode think_mode, - token_vec *out) { - const bool need_think_start = - ds4_think_mode_enabled(think_mode) || - DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; - if (vocab->bos_id < 0 || - vocab->user_id < 0 || - vocab->assistant_id < 0 || - vocab->think_end_id < 0 || - (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && vocab->system_id < 0) || - (need_think_start && vocab->think_start_id < 0)) { - ds4_die("this tokenizer does not provide the DeepSeek chat markers; use raw prompt tokenization"); - } - - chat_push_bos_sequence(vocab, out); - chat_push_think_prefix(vocab, think_mode, out); - if (system && system[0]) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) - token_vec_push(out, vocab->system_id); - bpe_tokenize_text(vocab, system, out); - } - token_vec_push(out, vocab->user_id); - bpe_tokenize_text(vocab, prompt, out); - token_vec_push(out, vocab->assistant_id); - if (ds4_think_mode_enabled(think_mode)) { - token_vec_push(out, vocab->think_start_id); - } else if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - token_vec_push(out, vocab->think_start_id); - token_vec_push(out, vocab->think_end_id); - } else { - token_vec_push(out, vocab->think_end_id); - } -} - -void ds4_tokenize_text(ds4_engine *e, const char *text, ds4_tokens *out) { - bpe_tokenize_text(&e->vocab, text ? text : "", out); -} - -static bool special_token_at(const ds4_vocab *vocab, const char *p, int *token, size_t *len) { - struct special { - const char *text; - int token; - } specials[] = { - {"<｜begin▁of▁sentence｜>", vocab->bos_id}, - {"<｜end▁of▁sentence｜>", vocab->eos_id}, - {"[gMASK]", vocab->bos_id}, - {"", vocab->sop_id}, - {"<|system|>", vocab->system_id}, - {"<｜User｜>", vocab->user_id}, - {"<｜Assistant｜>", vocab->assistant_id}, - {"<|user|>", vocab->user_id}, - {"<|assistant|>", vocab->assistant_id}, - {"<|observation|>", vocab->observation_id}, - {"", vocab->think_start_id}, - {"", vocab->think_end_id}, - {"", vocab->tool_call_start_id}, - {"", vocab->tool_call_end_id}, - {"", vocab->tool_response_start_id}, - {"", vocab->tool_response_end_id}, - {"", vocab->arg_key_start_id}, - {"", vocab->arg_key_end_id}, - {"", vocab->arg_value_start_id}, - {"", vocab->arg_value_end_id}, - {"｜DSML｜", vocab->dsml_id}, - }; - - for (size_t i = 0; i < sizeof(specials) / sizeof(specials[0]); i++) { - if (specials[i].token < 0) continue; - size_t n = strlen(specials[i].text); - if (!strncmp(p, specials[i].text, n)) { - *token = specials[i].token; - *len = n; - return true; - } - } - return false; -} - -static void tokenize_span(const ds4_vocab *vocab, const char *p, size_t n, token_vec *out) { - if (!n) return; - char *tmp = xmalloc(n + 1); - memcpy(tmp, p, n); - tmp[n] = '\0'; - bpe_tokenize_text(vocab, tmp, out); - free(tmp); -} - - - - -static void tokenize_rendered_chat_vocab(const ds4_vocab *vocab, const char *text, - token_vec *out) { - if (!text) text = ""; - - const char *span = text; - const char *p = text; - while (*p) { - int token = -1; - size_t len = 0; - if (special_token_at(vocab, p, &token, &len)) { - tokenize_span(vocab, span, (size_t)(p - span), out); - token_vec_push(out, token); - p += len; - span = p; - continue; - } - p++; - } - tokenize_span(vocab, span, (size_t)(p - span), out); -} - -void ds4_tokenize_rendered_chat(ds4_engine *e, const char *text, ds4_tokens *out) { - tokenize_rendered_chat_vocab(&e->vocab, text, out); -} - -void ds4_chat_begin(ds4_engine *e, ds4_tokens *tokens) { - chat_push_bos_sequence(&e->vocab, tokens); -} - -void ds4_encode_chat_prompt( - ds4_engine *e, - const char *system, - const char *prompt, - ds4_think_mode think_mode, - ds4_tokens *out) { - encode_chat_prompt(&e->vocab, system, prompt ? prompt : "", think_mode, out); -} - -void ds4_chat_append_max_effort_prefix(ds4_engine *e, ds4_tokens *tokens) { - bpe_tokenize_text(&e->vocab, DS4_REASONING_EFFORT_MAX_PREFIX, tokens); -} - -static void bpe_tokenize_wrapped_payload_text(ds4_vocab *vocab, const char *content, - const char *end, token_vec *out) { - /* Tool output is plain data inside the model-family wrapper. - * Preserve literal '<', '>' and '&' so shell output and file snippets stay - * intact, but escape the exact closing sentinel so a malicious or accidental - * tool payload cannot terminate the wrapper early. */ - const size_t endlen = strlen(end); - const char *span = content ? content : ""; - const char *p = span; - while (*p) { - if (!strncmp(p, end, endlen)) { - tokenize_span(vocab, span, (size_t)(p - span), out); - bpe_tokenize_text(vocab, "<", out); - p++; - span = p; - } else { - p++; - } - } - tokenize_span(vocab, span, (size_t)(p - span), out); -} - -static void bpe_tokenize_tool_result_text(ds4_vocab *vocab, const char *content, token_vec *out) { - bpe_tokenize_wrapped_payload_text(vocab, content, "", out); -} - -static void bpe_tokenize_tool_response_text(ds4_vocab *vocab, const char *content, token_vec *out) { - bpe_tokenize_wrapped_payload_text(vocab, content, "", out); -} - -void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role, const char *content) { - ds4_vocab *vocab = &e->vocab; - if (!role) role = "user"; - if (!content) content = ""; - - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - if (!strcmp(role, "system") || !strcmp(role, "developer")) { - if (vocab->system_id >= 0) token_vec_push(tokens, vocab->system_id); - tokenize_rendered_chat_vocab(vocab, content, tokens); - } else if (!strcmp(role, "assistant")) { - token_vec_push(tokens, vocab->assistant_id); - if (strncmp(content, "", 7) != 0 && - strncmp(content, "", 8) != 0) { - token_vec_push(tokens, vocab->think_start_id); - token_vec_push(tokens, vocab->think_end_id); - } - tokenize_rendered_chat_vocab(vocab, content, tokens); - } else if (!strcmp(role, "tool") || !strcmp(role, "function")) { - if (vocab->observation_id >= 0) token_vec_push(tokens, vocab->observation_id); - tokenize_rendered_chat_vocab(vocab, "", tokens); - bpe_tokenize_tool_response_text(vocab, content, tokens); - tokenize_rendered_chat_vocab(vocab, "", tokens); - } else { - token_vec_push(tokens, vocab->user_id); - bpe_tokenize_text(vocab, content, tokens); - } - return; - } - - if (!strcmp(role, "system") || !strcmp(role, "developer")) { - bpe_tokenize_text(vocab, content, tokens); - } else if (!strcmp(role, "assistant")) { - token_vec_push(tokens, vocab->assistant_id); - if (strncmp(content, "", 7) != 0 && strncmp(content, "", 8) != 0) { - token_vec_push(tokens, vocab->think_end_id); - } - bpe_tokenize_text(vocab, content, tokens); - } else if (!strcmp(role, "tool") || !strcmp(role, "function")) { - token_vec_push(tokens, vocab->user_id); - bpe_tokenize_text(vocab, "", tokens); - bpe_tokenize_tool_result_text(vocab, content, tokens); - bpe_tokenize_text(vocab, "", tokens); - } else { - token_vec_push(tokens, vocab->user_id); - bpe_tokenize_text(vocab, content, tokens); - } -} - - -void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode) { - token_vec_push(tokens, e->vocab.assistant_id); - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && - !ds4_think_mode_enabled(think_mode)) { - token_vec_push(tokens, e->vocab.think_start_id); - token_vec_push(tokens, e->vocab.think_end_id); - return; - } - token_vec_push(tokens, ds4_think_mode_enabled(think_mode) ? - e->vocab.think_start_id : e->vocab.think_end_id); -} - -static void dump_tokens_fp(FILE *fp, const ds4_vocab *vocab, const token_vec *tokens) { - fprintf(fp, "["); - for (int i = 0; i < tokens->len; i++) { - if (i) fprintf(fp, ", "); - fprintf(fp, "%d", tokens->v[i]); - } - fprintf(fp, "]\n"); - - for (int i = 0; i < tokens->len; i++) { - int id = tokens->v[i]; - if (id >= 0 && id < vocab->n_vocab) { - fprintf(fp, "%6d %.*s\n", id, (int)vocab->token[id].len, vocab->token[id].ptr); - } - } -} - -static void dump_tokens(const ds4_vocab *vocab, const token_vec *tokens) { - dump_tokens_fp(stdout, vocab, tokens); -} - -static uint32_t utf8_decode_one(const char *s, uint64_t len, uint64_t *pos) { - const uint8_t c = (uint8_t)s[*pos]; - if (c < 0x80 || *pos + 1 >= len) { - (*pos)++; - return c; - } - if ((c & 0xe0) == 0xc0 && *pos + 1 < len) { - uint32_t cp = ((uint32_t)(c & 0x1f) << 6) | ((uint8_t)s[*pos + 1] & 0x3f); - *pos += 2; - return cp; - } - if ((c & 0xf0) == 0xe0 && *pos + 2 < len) { - uint32_t cp = ((uint32_t)(c & 0x0f) << 12) | - ((uint32_t)((uint8_t)s[*pos + 1] & 0x3f) << 6) | - ((uint8_t)s[*pos + 2] & 0x3f); - *pos += 3; - return cp; - } - if ((c & 0xf8) == 0xf0 && *pos + 3 < len) { - uint32_t cp = ((uint32_t)(c & 0x07) << 18) | - ((uint32_t)((uint8_t)s[*pos + 1] & 0x3f) << 12) | - ((uint32_t)((uint8_t)s[*pos + 2] & 0x3f) << 6) | - ((uint8_t)s[*pos + 3] & 0x3f); - *pos += 4; - return cp; - } - (*pos)++; - return c; -} - -static int gpt2_codepoint_to_byte(uint32_t cp) { - if ((cp >= 33 && cp <= 126) || (cp >= 161 && cp <= 172) || (cp >= 174 && cp <= 255)) { - return (int)cp; - } - - uint32_t n = 0; - for (uint32_t b = 0; b < 256; b++) { - if ((b >= 33 && b <= 126) || (b >= 161 && b <= 172) || (b >= 174)) { - continue; - } - if (cp == 256 + n) return (int)b; - n++; - } - return -1; -} - -static bool vocab_token_is_literal_special(ds4_str s) { - const unsigned char bar[] = {0xef, 0xbd, 0x9c}; /* U+FF5C fullwidth vertical bar. */ - if (s.len < sizeof(bar)) return false; - for (uint64_t i = 0; i + sizeof(bar) <= s.len; i++) { - if (!memcmp(s.ptr + i, bar, sizeof(bar))) return true; - } - return false; -} - -char *ds4_token_text(ds4_engine *e, int token, size_t *len) { - ds4_vocab *vocab = &e->vocab; - if (token < 0 || token >= vocab->n_vocab) { - if (len) *len = 0; - char *out = xmalloc(1); - out[0] = '\0'; - return out; - } - - ds4_str s = vocab->token[token]; - char *out = xmalloc((size_t)s.len + 1); - if (vocab_token_is_literal_special(s)) { - memcpy(out, s.ptr, (size_t)s.len); - out[s.len] = '\0'; - if (len) *len = (size_t)s.len; - return out; - } - - size_t n = 0; - uint64_t pos = 0; - while (pos < s.len) { - uint32_t cp = utf8_decode_one(s.ptr, s.len, &pos); - int b = gpt2_codepoint_to_byte(cp); - if (b >= 0) out[n++] = (char)b; - } - out[n] = '\0'; - if (len) *len = n; - return out; -} - -static bool vocab_token_is_generation_stop(const ds4_vocab *vocab, int token) { - if (!vocab || token < 0) return false; - if (token == vocab->eos_id) return true; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { - return (vocab->system_id >= 0 && token == vocab->system_id) || - (vocab->user_id >= 0 && token == vocab->user_id) || - (vocab->assistant_id >= 0 && token == vocab->assistant_id) || - (vocab->observation_id >= 0 && token == vocab->observation_id); - } - return false; -} - -int ds4_token_eos(ds4_engine *e) { - return e->vocab.eos_id; -} - -bool ds4_token_is_stop(ds4_engine *e, int token) { - return e ? vocab_token_is_generation_stop(&e->vocab, token) : false; -} - -bool ds4_token_is_thinking_control(ds4_engine *e, int token) { - if (!e || token < 0) return false; - return (e->vocab.think_start_id >= 0 && - token == e->vocab.think_start_id) || - (e->vocab.think_end_id >= 0 && - token == e->vocab.think_end_id); -} - -bool ds4_token_is_stop_for_think_mode( - ds4_engine *e, - int token, - ds4_think_mode mode) { - if (ds4_token_is_stop(e, token)) return true; - /* In no-thinking mode the prompt already supplied the protocol close tag. - * If the model emits another thinking tag, do not print or feed it back: - * it is a control marker, not assistant content. */ - if (!ds4_think_mode_enabled(mode) && - ds4_token_is_thinking_control(e, token)) { - return true; - } - return false; -} - -int ds4_token_user(ds4_engine *e) { - return e->vocab.user_id; -} - -int ds4_token_assistant(ds4_engine *e) { - return e->vocab.assistant_id; -} - -static int sample_argmax(const float *logits, uint32_t n_vocab) { - int best = 0; - float best_v = DS4_NEG_INF; - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - if (v > best_v) { - best_v = v; - best = (int)i; - } - } - return best; -} - -static DS4_MAYBE_UNUSED void logits_top2(const float *logits, uint32_t n_vocab, - int *top0, float *logit0, - int *top1, float *logit1) { - int b0 = -1, b1 = -1; - float v0 = DS4_NEG_INF, v1 = DS4_NEG_INF; - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - if (v > v0) { - b1 = b0; v1 = v0; - b0 = (int)i; v0 = v; - } else if (v > v1) { - b1 = (int)i; v1 = v; - } - } - if (top0) *top0 = b0; - if (logit0) *logit0 = v0; - if (top1) *top1 = b1; - if (logit1) *logit1 = v1; -} - -static uint64_t sample_rng_next(uint64_t *state) { - uint64_t x = *state; - if (x == 0) x = 0x9e3779b97f4a7c15ULL; - x ^= x >> 12; - x ^= x << 25; - x ^= x >> 27; - *state = x; - return x * 0x2545f4914f6cdd1dULL; -} - -static float sample_rng_f32(uint64_t *state) { - const uint64_t x = sample_rng_next(state); - return (float)((x >> 40) & 0xffffffu) / 16777216.0f; -} - -typedef struct { - int id; - float logit; - float prob; -} sample_candidate; - -static int sample_candidate_cmp_desc(const void *a, const void *b) { - const sample_candidate *ca = a; - const sample_candidate *cb = b; - const int logit_order = - (cb->logit > ca->logit) - (cb->logit < ca->logit); - if (logit_order != 0) return logit_order; - return (ca->id > cb->id) - (ca->id < cb->id); -} - -static bool sample_candidate_gt(sample_candidate a, sample_candidate b) { - if (a.logit != b.logit) return a.logit > b.logit; - return a.id < b.id; -} - -static void sample_heap_sift_up(sample_candidate *heap, uint32_t idx) { - while (idx > 0) { - const uint32_t parent = (idx - 1u) / 2u; - if (!sample_candidate_gt(heap[parent], heap[idx])) break; - sample_candidate tmp = heap[parent]; - heap[parent] = heap[idx]; - heap[idx] = tmp; - idx = parent; - } -} - -static void sample_heap_sift_down(sample_candidate *heap, uint32_t n, uint32_t idx) { - for (;;) { - const uint32_t left = idx * 2u + 1u; - const uint32_t right = left + 1u; - uint32_t smallest = idx; - if (left < n && sample_candidate_gt(heap[smallest], heap[left])) { - smallest = left; - } - if (right < n && sample_candidate_gt(heap[smallest], heap[right])) { - smallest = right; - } - if (smallest == idx) break; - sample_candidate tmp = heap[idx]; - heap[idx] = heap[smallest]; - heap[smallest] = tmp; - idx = smallest; - } -} - -static bool sample_fast_top_p( - const float *logits, - uint32_t n_vocab, - uint32_t finite, - float max_logit, - int best, - float temperature, - float top_p, - float min_p, - uint64_t *rng, - int *token_out) { - enum { SAMPLE_FAST_TOP_P_CAP = 512 }; - if (!logits || !rng || !token_out || finite == 0) return false; - if (finite > SAMPLE_FAST_TOP_P_CAP && top_p >= 0.999f) return false; - - const uint32_t cap = finite < SAMPLE_FAST_TOP_P_CAP ? - finite : (uint32_t)SAMPLE_FAST_TOP_P_CAP; - sample_candidate heap[SAMPLE_FAST_TOP_P_CAP]; - uint32_t n = 0; - float sum = 0.0f; - float heap_sum = 0.0f; - - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - if (!isfinite(v)) continue; - const float p = expf((v - max_logit) / temperature); - sum += p; - sample_candidate cand = {.id = (int)i, .logit = v, .prob = p}; - if (n < cap) { - heap[n] = cand; - heap_sum += p; - sample_heap_sift_up(heap, n); - n++; - } else if (sample_candidate_gt(cand, heap[0])) { - heap_sum -= heap[0].prob; - heap[0] = cand; - heap_sum += p; - sample_heap_sift_down(heap, n, 0); - } - } - if (sum <= 0.0f || !isfinite(sum)) { - *token_out = best; - return true; - } - - if (n < finite && heap_sum < top_p * sum) { - return false; - } - - qsort(heap, n, sizeof(heap[0]), sample_candidate_cmp_desc); - const float min_prob = (heap[0].prob / sum) * (min_p > 0.0f ? min_p : 0.0f); - const float min_prob_raw = heap[0].prob * (min_p > 0.0f ? min_p : 0.0f); - float filtered_sum = 0.0f; - uint32_t filtered = 0; - bool stopped_by_min_p = false; - for (uint32_t i = 0; i < n; i++) { - const float p = heap[i].prob / sum; - if (i > 0 && p < min_prob) { - stopped_by_min_p = true; - break; - } - filtered_sum += heap[i].prob; - filtered++; - if (filtered_sum / sum >= top_p) break; - } - if (n < finite && - stopped_by_min_p && - min_p > 0.0f && - heap[n - 1u].prob >= min_prob_raw) { - return false; - } - if (filtered == 0) { - *token_out = best; - return true; - } - - float r = sample_rng_f32(rng) * filtered_sum; - for (uint32_t i = 0; i < filtered; i++) { - r -= heap[i].prob; - if (r <= 0.0f) { - *token_out = heap[i].id; - return true; - } - } - *token_out = heap[filtered - 1u].id; - return true; -} - -static int sample_full_vocab( - const float *logits, - uint32_t n_vocab, - float temperature, - float top_p, - float min_p, - uint64_t *rng, - float *prob_scratch) { - float max_logit = DS4_NEG_INF; - int best = 0; - uint32_t finite = 0; - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - if (!isfinite(v)) continue; - finite++; - if (v > max_logit) { - max_logit = v; - best = (int)i; - } - } - if (finite == 0) return sample_argmax(logits, n_vocab); + ds4_select_shape_from_metadata(n_layer, + n_embd, + n_vocab, + n_head, + n_head_kv, + n_head_dim, + n_value_dim, + n_rot, + n_lora_q, + n_lora_o, + n_out_group, + n_expert, + n_expert_used, + n_ff_exp, + n_expert_shared, + n_hash_layer, + n_swa, + n_indexer_head, + n_indexer_head_dim, + n_indexer_top_k, + n_hc, + n_hc_sinkhorn_iter); - int fast_token = best; - if (top_p < 1.0f && - sample_fast_top_p(logits, - n_vocab, - finite, - max_logit, - best, - temperature, - top_p, - min_p, - rng, - &fast_token)) { - return fast_token; - } + config_expect_u32("embedding_length", n_embd, DS4_N_EMBD); + config_expect_u32("vocab_size", n_vocab, DS4_N_VOCAB); + config_expect_u32("attention.head_count", n_head, DS4_N_HEAD); + config_expect_u32("attention.key_length", n_head_dim, DS4_N_HEAD_DIM); + config_expect_u32("attention.head_count_kv", n_head_kv, DS4_N_HEAD_KV); + config_expect_u32("attention.value_length", n_value_dim, DS4_N_VALUE_DIM); + config_expect_u32("rope.dimension_count", n_rot, DS4_N_ROT); + config_expect_u32("attention.output_group_count", n_out_group, DS4_N_OUT_GROUP); + config_expect_u32("attention.q_lora_rank", n_lora_q, DS4_N_LORA_Q); + config_expect_u32("attention.output_lora_rank", n_lora_o, DS4_N_LORA_O); + config_expect_u32("expert_count", n_expert, DS4_N_EXPERT); + config_expect_u32("expert_used_count", n_expert_used, DS4_N_EXPERT_USED); + config_expect_u32("expert_feed_forward_length", n_ff_exp, DS4_N_FF_EXP); + config_expect_u32("expert_shared_count", n_expert_shared, DS4_N_EXPERT_SHARED); + config_expect_u32("hash_layer_count", n_hash_layer, DS4_N_HASH_LAYER); + config_expect_u32("expert_group_count", n_expert_groups, 0); + config_expect_u32("expert_group_used_count", n_group_used, 0); - if (top_p >= 1.0f) { - float sum = 0.0f; - const float min_rel = min_p > 0.0f ? min_p : 0.0f; - if (min_rel > 1.0f) return best; + config_expect_u32("attention.sliding_window", n_swa, DS4_N_SWA); + config_expect_u32("attention.indexer.head_count", n_indexer_head, DS4_N_INDEXER_HEAD); + config_expect_u32("attention.indexer.key_length", n_indexer_head_dim, DS4_N_INDEXER_HEAD_DIM); + config_expect_u32("attention.indexer.top_k", n_indexer_top_k, DS4_N_INDEXER_TOP_K); + config_expect_u32("hyper_connection.count", n_hc, DS4_N_HC); + config_expect_u32("hyper_connection.sinkhorn_iterations", n_hc_sinkhorn_iter, DS4_N_HC_SINKHORN_ITER); - /* Find a conservative log-space rejection boundary using the same - * expf implementation as the probability path. Values below this - * boundary are guaranteed to fail min-p, avoiding an expf for the - * overwhelming majority of a large vocabulary. Near-boundary values - * still take the ordinary expf comparison. */ - float reject_scaled = DS4_NEG_INF; - bool have_reject_scaled = false; - if (min_rel > 0.0f && isfinite(min_rel)) { - float cutoff = logf(min_rel); - for (int i = 0; i < 8 && isfinite(cutoff); i++) { - cutoff = nextafterf(cutoff, -FLT_MAX); - if (expf(cutoff) < min_rel) { - reject_scaled = cutoff; - have_reject_scaled = true; - break; - } - } - } + config_validate_fixed_shape(n_layer); + validate_compress_ratio_metadata(m); - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - prob_scratch[i] = -1.0f; - if (!isfinite(v)) continue; - const float scaled = (v - max_logit) / temperature; - if (have_reject_scaled && scaled <= reject_scaled) continue; - const float p = expf(scaled); - if (p < min_rel) continue; - prob_scratch[i] = p; - sum += p; - } - if (sum <= 0.0f || !isfinite(sum)) return best; - float r = sample_rng_f32(rng) * sum; - for (uint32_t i = 0; i < n_vocab; i++) { - const float p = prob_scratch[i]; - if (p < 0.0f) continue; - r -= p; - if (r <= 0.0f) return (int)i; - } - return best; + validate_swiglu_clamp_metadata(m); + + uint64_t rope_orig_ctx = DS4_ROPE_ORIG_CTX; + model_get_u64_compat(m, "deepseek4.rope.scaling.original_context_length", &rope_orig_ctx); + if (rope_orig_ctx != DS4_ROPE_ORIG_CTX) { + fprintf(stderr, "ds4: expected rope.scaling.original_context_length=%" PRIu64 + " for %s, got %" PRIu64 "\n", + (uint64_t)DS4_ROPE_ORIG_CTX, DS4_MODEL_SHAPE_NAME, rope_orig_ctx); + exit(1); } + const float rope_freq_base = required_f32(m, "deepseek4.rope.freq_base"); + config_expect_f32("rope.freq_base", rope_freq_base, DS4_ROPE_FREQ_BASE); + float rope_scale_factor = DS4_ROPE_SCALE_FACTOR; + model_get_f32_compat(m, "deepseek4.rope.scaling.factor", &rope_scale_factor); + config_expect_f32("rope.scaling.factor", rope_scale_factor, DS4_ROPE_SCALE_FACTOR); + float rope_yarn_beta_fast = DS4_ROPE_YARN_BETA_FAST; + model_get_f32_compat(m, "deepseek4.rope.scaling.yarn_beta_fast", &rope_yarn_beta_fast); + config_expect_f32("rope.scaling.yarn_beta_fast", rope_yarn_beta_fast, DS4_ROPE_YARN_BETA_FAST); + float rope_yarn_beta_slow = DS4_ROPE_YARN_BETA_SLOW; + model_get_f32_compat(m, "deepseek4.rope.scaling.yarn_beta_slow", &rope_yarn_beta_slow); + config_expect_f32("rope.scaling.yarn_beta_slow", rope_yarn_beta_slow, DS4_ROPE_YARN_BETA_SLOW); + const float compress_rope_freq_base = required_f32(m, "deepseek4.attention.compress_rope_freq_base"); + config_expect_f32("attention.compress_rope_freq_base", compress_rope_freq_base, DS4_COMPRESS_ROPE_FREQ_BASE); + const float expert_weight_scale = required_f32(m, "deepseek4.expert_weights_scale"); + config_expect_f32("expert_weights_scale", expert_weight_scale, DS4_EXPERT_WEIGHT_SCALE); + const float rms_eps = required_f32(m, "deepseek4.attention.layer_norm_rms_epsilon"); + config_expect_f32("attention.layer_norm_rms_epsilon", rms_eps, DS4_RMS_EPS); + const float hc_eps = required_f32(m, "deepseek4.hyper_connection.epsilon"); + config_expect_f32("hyper_connection.epsilon", hc_eps, DS4_HC_EPS); + const bool expert_weight_norm = required_bool(m, "deepseek4.expert_weights_norm"); + config_expect_bool("expert_weights_norm", expert_weight_norm, true); +} - uint32_t n = 0; - float sum = 0.0f; - sample_candidate *cand = NULL; - if (min_p > 0.0f && min_p <= 1.0f) { - /* The later min-p comparison is equivalent to - * exp((logit-max)/temperature) >= min_p; its normalization cancels. - * Still compute the full softmax sum in the original order, then sort - * only candidates that can survive. This preserves the nucleus mass - * and RNG semantics while avoiding a full-vocabulary qsort. */ - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - prob_scratch[i] = -1.0f; - if (!isfinite(v)) continue; - const float p = expf((v - max_logit) / temperature); - prob_scratch[i] = p; - sum += p; - } - if (sum <= 0.0f || !isfinite(sum)) return best; +static void config_validate_glm_dsa_model(const ds4_model *m) { + g_ds4_shape = DS4_SHAPE_GLM52; + memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); - const float min_prob = (1.0f / sum) * min_p; - for (uint32_t i = 0; i < n_vocab; i++) { - const float p = prob_scratch[i]; - if (p < 0.0f || p / sum < min_prob) continue; - n++; - } - if (n == 0) return best; - cand = xmalloc((size_t)n * sizeof(cand[0])); - uint32_t out = 0; - for (uint32_t i = 0; i < n_vocab; i++) { - const float p = prob_scratch[i]; - if (p < 0.0f || p / sum < min_prob) continue; - cand[out++] = (sample_candidate){ - .id = (int)i, .logit = logits[i], .prob = p - }; - } - } else { - cand = xmalloc((size_t)finite * sizeof(cand[0])); - for (uint32_t i = 0; i < n_vocab; i++) { - const float v = logits[i]; - if (!isfinite(v)) continue; - const float p = expf((v - max_logit) / temperature); - cand[n++] = (sample_candidate){.id = (int)i, .logit = v, .prob = p}; - sum += p; - } - } - if (sum <= 0.0f || !isfinite(sum)) { - free(cand); - return best; - } + const uint32_t n_layer = required_u32(m, "glm-dsa.block_count"); + const uint64_t n_ctx = required_u64_compat(m, "glm-dsa.context_length"); + const uint32_t n_embd = required_u32(m, "glm-dsa.embedding_length"); + const uint32_t n_vocab = required_u32(m, "glm-dsa.vocab_size"); + const uint32_t n_ff_dense = required_u32(m, "glm-dsa.feed_forward_length"); + const uint32_t n_head = required_u32(m, "glm-dsa.attention.head_count"); + const uint32_t n_head_kv = required_u32(m, "glm-dsa.attention.head_count_kv"); + const uint32_t n_head_dim = required_u32(m, "glm-dsa.attention.key_length"); + const uint32_t n_value_dim = required_u32(m, "glm-dsa.attention.value_length"); + const uint32_t n_rot = required_u32(m, "glm-dsa.rope.dimension_count"); + const uint32_t n_lora_q = required_u32(m, "glm-dsa.attention.q_lora_rank"); + const uint32_t n_kv_lora = required_u32(m, "glm-dsa.attention.kv_lora_rank"); + const uint32_t n_key_mla = required_u32(m, "glm-dsa.attention.key_length_mla"); + const uint32_t n_value_mla = required_u32(m, "glm-dsa.attention.value_length_mla"); + const uint32_t n_expert = required_u32(m, "glm-dsa.expert_count"); + const uint32_t n_expert_used = required_u32(m, "glm-dsa.expert_used_count"); + const uint32_t n_ff_exp = required_u32(m, "glm-dsa.expert_feed_forward_length"); + const uint32_t n_expert_shared = required_u32(m, "glm-dsa.expert_shared_count"); + const uint32_t n_expert_group = required_u32(m, "glm-dsa.expert_group_count"); + const uint32_t n_expert_group_used = required_u32(m, "glm-dsa.expert_group_used_count"); + const uint32_t expert_gating_func = required_u32(m, "glm-dsa.expert_gating_func"); + const uint32_t n_leading_dense = required_u32(m, "glm-dsa.leading_dense_block_count"); + const uint32_t n_nextn = required_u32(m, "glm-dsa.nextn_predict_layers"); + const uint32_t n_indexer_head = required_u32(m, "glm-dsa.attention.indexer.head_count"); + const uint32_t n_indexer_head_dim = required_u32(m, "glm-dsa.attention.indexer.key_length"); + const uint32_t n_indexer_top_k = required_u32(m, "glm-dsa.attention.indexer.top_k"); - qsort(cand, n, sizeof(cand[0]), sample_candidate_cmp_desc); - const float min_prob = (cand[0].prob / sum) * (min_p > 0.0f ? min_p : 0.0f); - float filtered_sum = 0.0f; - uint32_t filtered = 0; - for (uint32_t i = 0; i < n; i++) { - const float p = cand[i].prob / sum; - if (i > 0 && p < min_prob) break; - filtered_sum += cand[i].prob; - filtered++; - if (filtered_sum / sum >= top_p) break; - } - if (filtered == 0) { - free(cand); - return best; - } + config_expect_u32("block_count", n_layer, DS4_N_LAYER); + config_expect_u64("context_length", n_ctx, DS4_ROPE_ORIG_CTX); + config_expect_u32("embedding_length", n_embd, DS4_N_EMBD); + config_expect_u32("vocab_size", n_vocab, DS4_N_VOCAB); + config_expect_u32("feed_forward_length", n_ff_dense, DS4_N_FF_DENSE); + config_expect_u32("attention.head_count", n_head, DS4_N_HEAD); + config_expect_u32("attention.head_count_kv", n_head_kv, DS4_N_HEAD_KV); + config_expect_u32("attention.key_length", n_head_dim, DS4_N_HEAD_DIM); + config_expect_u32("attention.value_length", n_value_dim, DS4_N_VALUE_DIM); + config_expect_u32("rope.dimension_count", n_rot, DS4_N_ROT); + config_expect_u32("attention.q_lora_rank", n_lora_q, DS4_N_LORA_Q); + config_expect_u32("attention.kv_lora_rank", n_kv_lora, DS4_N_KV_LORA); + config_expect_u32("attention.key_length_mla", n_key_mla, DS4_N_KEY_MLA); + config_expect_u32("attention.value_length_mla", n_value_mla, DS4_N_VALUE_MLA); + config_expect_u32("expert_count", n_expert, DS4_N_EXPERT); + config_expect_u32("expert_used_count", n_expert_used, DS4_N_EXPERT_USED); + config_expect_u32("expert_feed_forward_length", n_ff_exp, DS4_N_FF_EXP); + config_expect_u32("expert_shared_count", n_expert_shared, DS4_N_EXPERT_SHARED); + config_expect_u32("expert_group_count", n_expert_group, 1); + config_expect_u32("expert_group_used_count", n_expert_group_used, 1); + config_expect_u32("expert_gating_func", expert_gating_func, 2); + config_expect_u32("leading_dense_block_count", n_leading_dense, DS4_N_LEADING_DENSE); + config_expect_u32("nextn_predict_layers", n_nextn, DS4_N_NEXTN_PREDICT); + config_expect_u32("attention.indexer.head_count", n_indexer_head, DS4_N_INDEXER_HEAD); + config_expect_u32("attention.indexer.key_length", n_indexer_head_dim, DS4_N_INDEXER_HEAD_DIM); + config_expect_u32("attention.indexer.top_k", n_indexer_top_k, DS4_N_INDEXER_TOP_K); - float r = sample_rng_f32(rng) * filtered_sum; - for (uint32_t i = 0; i < filtered; i++) { - r -= cand[i].prob; - if (r <= 0.0f) { - const int id = cand[i].id; - free(cand); - return id; - } - } - const int id = cand[filtered - 1].id; - free(cand); - return id; + const float rope_freq_base = required_f32(m, "glm-dsa.rope.freq_base"); + config_expect_f32("rope.freq_base", rope_freq_base, DS4_ROPE_FREQ_BASE); + const float rms_eps = required_f32(m, "glm-dsa.attention.layer_norm_rms_epsilon"); + config_expect_f32("attention.layer_norm_rms_epsilon", rms_eps, DS4_RMS_EPS); + const float expert_weight_scale = required_f32(m, "glm-dsa.expert_weights_scale"); + config_expect_f32("expert_weights_scale", expert_weight_scale, DS4_EXPERT_WEIGHT_SCALE); + const bool expert_weight_norm = required_bool(m, "glm-dsa.expert_weights_norm"); + config_expect_bool("expert_weights_norm", expert_weight_norm, true); } -static int sample_top_p_min_p( - const float *logits, - uint32_t n_vocab, - float temperature, - int top_k, - float top_p, - float min_p, - uint64_t *rng, - float *prob_scratch) { - if (temperature <= 0.0f) return sample_argmax(logits, n_vocab); - if (top_p <= 0.0f || top_p > 1.0f) top_p = 1.0f; - if (min_p < 0.0f) min_p = 0.0f; - if (top_k <= 0) { - const bool owned_scratch = prob_scratch == NULL; - if (owned_scratch) { - prob_scratch = xmalloc((size_t)n_vocab * sizeof(prob_scratch[0])); - } - const int token = sample_full_vocab(logits, n_vocab, temperature, - top_p, min_p, rng, prob_scratch); - if (owned_scratch) free(prob_scratch); - return token; +static void config_validate_model(const ds4_model *m) { + ds4_str arch = {0}; + if (model_get_string(m, "general.architecture", &arch) && + ds4_streq(arch, "glm-dsa")) { + config_validate_glm_dsa_model(m); + return; } - if (top_k > 1024) top_k = 1024; - if ((uint32_t)top_k > n_vocab) top_k = (int)n_vocab; + config_validate_deepseek4_model(m); +} - int ids[1024]; - float vals[1024]; - int n = 0; - for (uint32_t i = 0; i < n_vocab; i++) { - float v = logits[i]; - if (!isfinite(v)) continue; - if (n == top_k && v <= vals[n - 1]) continue; - int j = n < top_k ? n++ : n - 1; - while (j > 0 && vals[j - 1] < v) { - vals[j] = vals[j - 1]; - ids[j] = ids[j - 1]; - j--; +static void weights_bind_output( + ds4_weights *w, + const ds4_model *m, + bool required, + bool optional) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (required) { + w->output_norm = required_tensor(m, "output_norm.weight"); + w->output = required_tensor(m, "output.weight"); + } else if (optional) { + w->output_norm = model_find_tensor(m, "output_norm.weight"); + w->output = model_find_tensor(m, "output.weight"); } - vals[j] = v; - ids[j] = (int)i; + } else if (required) { + w->output_hc_base = required_tensor(m, "output_hc_base.weight"); + w->output_hc_fn = required_tensor(m, "output_hc_fn.weight"); + w->output_hc_scale = required_tensor(m, "output_hc_scale.weight"); + w->output_norm = required_tensor(m, "output_norm.weight"); + w->output = required_tensor(m, "output.weight"); + } else if (optional) { + w->output_hc_base = model_find_tensor(m, "output_hc_base.weight"); + w->output_hc_fn = model_find_tensor(m, "output_hc_fn.weight"); + w->output_hc_scale = model_find_tensor(m, "output_hc_scale.weight"); + w->output_norm = model_find_tensor(m, "output_norm.weight"); + w->output = model_find_tensor(m, "output.weight"); } - if (n == 0) return sample_argmax(logits, n_vocab); - float probs[1024]; - const float max_logit = vals[0]; - float sum = 0.0f; - for (int i = 0; i < n; i++) { - probs[i] = expf((vals[i] - max_logit) / temperature); - sum += probs[i]; + if (optional && + weights_have_partial_output_head(w) && + !weights_have_output_head(w)) { + ds4_die("partial output head in GGUF"); } - if (sum <= 0.0f || !isfinite(sum)) return ids[0]; +} - const float min_prob = (probs[0] / sum) * min_p; - float filtered_sum = 0.0f; - int filtered = 0; - for (int i = 0; i < n; i++) { - float p = probs[i] / sum; - if (i > 0 && p < min_prob) break; - filtered_sum += probs[i]; - filtered++; - if (filtered_sum / sum >= top_p) break; +static void weights_bind_glm_dsa_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { + l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); + l->attn_q_a = required_tensorf(m, "blk.%u.attn_q_a.weight", il); + l->attn_q_a_norm = required_tensorf(m, "blk.%u.attn_q_a_norm.weight", il); + l->attn_q_b = required_tensorf(m, "blk.%u.attn_q_b.weight", il); + l->attn_kv_a_mqa = required_tensorf(m, "blk.%u.attn_kv_a_mqa.weight", il); + l->attn_kv_a_norm = required_tensorf(m, "blk.%u.attn_kv_a_norm.weight", il); + l->attn_k_b = required_tensorf(m, "blk.%u.attn_k_b.weight", il); + l->attn_v_b = required_tensorf(m, "blk.%u.attn_v_b.weight", il); + l->attn_output = required_tensorf(m, "blk.%u.attn_output.weight", il); + l->indexer_attn_q_b = required_tensorf(m, "blk.%u.indexer.attn_q_b.weight", il); + l->indexer_attn_k = required_tensorf(m, "blk.%u.indexer.attn_k.weight", il); + l->indexer_k_norm = required_tensorf(m, "blk.%u.indexer.k_norm.weight", il); + l->indexer_k_norm_b = required_tensorf(m, "blk.%u.indexer.k_norm.bias", il); + l->indexer_proj = required_tensorf(m, "blk.%u.indexer.proj.weight", il); + l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); + + if (il < DS4_N_LEADING_DENSE) { + l->ffn_gate = required_tensorf(m, "blk.%u.ffn_gate.weight", il); + l->ffn_up = required_tensorf(m, "blk.%u.ffn_up.weight", il); + l->ffn_down = required_tensorf(m, "blk.%u.ffn_down.weight", il); + } else { + l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); + l->ffn_exp_probs_b = required_tensorf(m, "blk.%u.exp_probs_b.bias", il); + l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); + l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); + l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); + l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); + l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); + l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); } - if (filtered <= 0) return ids[0]; - float r = sample_rng_f32(rng) * filtered_sum; - for (int i = 0; i < filtered; i++) { - r -= probs[i]; - if (r <= 0.0f) return ids[i]; + if (DS4_N_NEXTN_PREDICT != 0 && + il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { + l->nextn_eh_proj = required_tensorf(m, "blk.%u.nextn.eh_proj.weight", il); + l->nextn_enorm = required_tensorf(m, "blk.%u.nextn.enorm.weight", il); + l->nextn_hnorm = required_tensorf(m, "blk.%u.nextn.hnorm.weight", il); + l->nextn_shared_head_norm = + required_tensorf(m, "blk.%u.nextn.shared_head_norm.weight", il); } - return ids[filtered - 1]; } -#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) { - if (!logits || !rng || n_vocab == 0) return -1; - return sample_top_p_min_p(logits, n_vocab, temperature, top_k, - top_p, min_p, rng, prob_scratch); -} -#endif +static void weights_bind_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + weights_bind_glm_dsa_layer(l, m, il); + return; + } -static void print_top_logits( - FILE * fp, - const char * label, - const ds4_vocab * vocab, - const float * logits, - uint32_t n_vocab, - int k) { - int best[16]; - if (k > 16) k = 16; - for (int i = 0; i < k; i++) best[i] = -1; + const uint32_t compress_ratio = ds4_layer_compress_ratio(il); - for (uint32_t i = 0; i < n_vocab; i++) { - for (int j = 0; j < k; j++) { - if (best[j] < 0 || logits[i] > logits[best[j]]) { - for (int l = k - 1; l > j; l--) best[l] = best[l - 1]; - best[j] = (int)i; - break; - } - } + l->hc_attn_fn = required_tensorf(m, "blk.%u.hc_attn_fn.weight", il); + l->hc_attn_scale = required_tensorf(m, "blk.%u.hc_attn_scale.weight", il); + l->hc_attn_base = required_tensorf(m, "blk.%u.hc_attn_base.weight", il); + l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); + l->attn_q_a = required_tensorf(m, "blk.%u.attn_q_a.weight", il); + l->attn_q_a_norm = required_tensorf(m, "blk.%u.attn_q_a_norm.weight", il); + l->attn_q_b = required_tensorf(m, "blk.%u.attn_q_b.weight", il); + l->attn_kv = required_tensorf(m, "blk.%u.attn_kv.weight", il); + l->attn_kv_a_norm = required_tensorf(m, "blk.%u.attn_kv_a_norm.weight", il); + l->attn_sinks = required_tensorf(m, "blk.%u.attn_sinks.weight", il); + l->attn_output_a = required_tensorf(m, "blk.%u.attn_output_a.weight", il); + l->attn_output_b = required_tensorf(m, "blk.%u.attn_output_b.weight", il); + if (compress_ratio != 0) { + l->attn_compressor_ape = required_tensorf(m, "blk.%u.attn_compressor_ape.weight", il); + l->attn_compressor_kv = required_tensorf(m, "blk.%u.attn_compressor_kv.weight", il); + l->attn_compressor_gate = required_tensorf(m, "blk.%u.attn_compressor_gate.weight", il); + l->attn_compressor_norm = required_tensorf(m, "blk.%u.attn_compressor_norm.weight", il); + } + if (compress_ratio == 4) { + l->indexer_attn_q_b = required_tensorf(m, "blk.%u.indexer.attn_q_b.weight", il); + l->indexer_proj = required_tensorf(m, "blk.%u.indexer.proj.weight", il); + l->indexer_compressor_ape = required_tensorf(m, "blk.%u.indexer_compressor_ape.weight", il); + l->indexer_compressor_kv = required_tensorf(m, "blk.%u.indexer_compressor_kv.weight", il); + l->indexer_compressor_gate = required_tensorf(m, "blk.%u.indexer_compressor_gate.weight", il); + l->indexer_compressor_norm = required_tensorf(m, "blk.%u.indexer_compressor_norm.weight", il); } + l->hc_ffn_fn = required_tensorf(m, "blk.%u.hc_ffn_fn.weight", il); + l->hc_ffn_scale = required_tensorf(m, "blk.%u.hc_ffn_scale.weight", il); + l->hc_ffn_base = required_tensorf(m, "blk.%u.hc_ffn_base.weight", il); + l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); + l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); + l->ffn_exp_probs_b = tensor_by_namef(m, "blk.%u.exp_probs_b.bias", il); + l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); + l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); + l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); + l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); + l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); + l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); - fprintf(fp, "ds4: top logits %s:\n", label); - for (int i = 0; i < k && best[i] >= 0; i++) { - const int id = best[i]; - fprintf(fp, " %2d %7d % .9g ", i, id, logits[id]); - if (id >= 0 && id < vocab->n_vocab) { - fprintf(fp, "%.*s", (int)vocab->token[id].len, vocab->token[id].ptr); - } - fputc('\n', fp); + if (il < DS4_N_HASH_LAYER) { + l->ffn_gate_tid2eid = required_tensorf(m, "blk.%u.ffn_gate_tid2eid.weight", il); } } -/* CPU generation entry point. It runs layer-major prefill once, then decodes - * one token at a time using the persistent KV cache and scratch arena. */ -static int generate_raw_swa_cpu( - const ds4_model * model, - const ds4_vocab * vocab, - const ds4_weights * weights, - const token_vec * prompt, - int n_predict, - int ctx_size, - const float * directional_steering_dirs, - float directional_steering_attn, - float directional_steering_ffn, - ds4_token_emit_fn emit, - ds4_generation_done_fn done, - void * emit_ud, - ds4_session_progress_fn progress, - void * progress_ud) { - (void)progress; - (void)progress_ud; - fprintf(stderr, "ds4: using CPU generation with layer-major prefill\n"); - - ds4_kv_cache cache; - kv_cache_init(&cache, (uint32_t)ctx_size, 0); - ds4_cpu_decode_scratch decode_scratch; - cpu_decode_scratch_init(&decode_scratch, (uint32_t)ctx_size); +/* Bind tensor names once into the fixed DS4 layer layout. This is the point + * where stringly GGUF metadata becomes direct model-specific pointers. */ +static void weights_bind( + ds4_weights *w, + const ds4_model *m, + bool load_slice, + uint32_t load_layer_start, + uint32_t load_layer_end, + bool require_output, + bool optional_output) { + memset(w, 0, sizeof(*w)); - float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); - int pos = prompt->len; - const bool trace_top = getenv("DS4_TRACE_TOP") != NULL; - const double t_prefill0 = now_sec(); + uint32_t executable_layers = DS4_N_LAYER; + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + DS4_N_LAYER > DS4_N_NEXTN_PREDICT) { + executable_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + } + uint32_t start = 0; + uint32_t end = executable_layers - 1u; + bool require_token_embd = true; + if (load_slice) { + if (load_layer_start >= executable_layers) ds4_die("invalid model load layer slice"); + start = load_layer_start; + end = load_layer_end == UINT32_MAX ? executable_layers - 1u : load_layer_end; + if (end >= executable_layers || end < start) ds4_die("invalid model load layer slice"); + require_token_embd = start == 0; + } else { + require_output = true; + optional_output = false; + } - if (prompt->len <= 0 || prompt->len > ctx_size) { - fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); - free(logits); - cpu_decode_scratch_free(&decode_scratch); - kv_cache_free(&cache); - return 1; + if (require_token_embd) { + w->token_embd = required_tensor(m, "token_embd.weight"); + } else { + w->token_embd = model_find_tensor(m, "token_embd.weight"); } + weights_bind_output(w, m, require_output, optional_output); - prefill_layer_major_cpu(logits, model, weights, &cache, prompt, - directional_steering_dirs, - directional_steering_attn, - directional_steering_ffn); - - const double t_prefill1 = now_sec(); - fprintf(stderr, "ds4: prefill %d/%d done\n", prompt->len, prompt->len); - const char *dump_prefill_logits = getenv("DS4_CPU_DUMP_PREFILL_LOGITS"); - if (dump_prefill_logits && dump_prefill_logits[0]) { - if (!write_f32_binary_file(dump_prefill_logits, logits, DS4_N_VOCAB)) { - free(logits); - cpu_decode_scratch_free(&decode_scratch); - kv_cache_free(&cache); - return 1; + for (uint32_t il = start; il <= end; il++) { + weights_bind_layer(&w->layer[il], m, il); + } + /* GLM nextn/MTP block(s): excluded from the executable pass but bound + * so the drafter can run them. Only when the full model is loaded. */ + if (!load_slice && + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + start == 0 && end == executable_layers - 1u) { + for (uint32_t il = executable_layers; il < DS4_N_LAYER; il++) { + weights_bind_layer(&w->layer[il], m, il); } - fprintf(stderr, "ds4: wrote CPU prefill logits to %s\n", dump_prefill_logits); } - int n_generated = 0; - int n_decode_eval = 0; - const bool token_timing = getenv("DS4_TOKEN_TIMING") != NULL; - const double t_decode0 = now_sec(); - for (int i = 0; i < n_predict && pos < ctx_size; i++) { - if (trace_top) { - char label[64]; - snprintf(label, sizeof(label), "step %d", i); - print_top_logits(stderr, label, vocab, logits, DS4_N_VOCAB, 10); - } + weights_validate_layout(w, start, end, require_token_embd, require_output); +} - int token = sample_argmax(logits, DS4_N_VOCAB); - if (vocab_token_is_generation_stop(vocab, token)) break; +typedef struct { + uint64_t off; + uint64_t end; + bool isolate; +} ds4_model_map_span; - if (emit) emit(emit_ud, token); - n_generated++; +typedef struct { + ds4_model_map_span *v; + uint32_t len; + uint32_t cap; + uint64_t max_tensor_bytes; +} ds4_model_map_span_vec; - if (i == n_predict - 1 || pos + 1 >= ctx_size) { - pos++; - break; - } +static void model_map_span_include_tensor( + const ds4_tensor *t, + uint64_t *lo, + uint64_t *hi, + uint64_t *max_tensor_bytes) { + if (!t || t->bytes == 0) return; + const uint64_t end = t->abs_offset + t->bytes; + if (*lo == UINT64_MAX || t->abs_offset < *lo) *lo = t->abs_offset; + if (end > *hi) *hi = end; + if (t->bytes > *max_tensor_bytes) *max_tensor_bytes = t->bytes; +} - const double t_eval0 = token_timing ? now_sec() : 0.0; - /* The CPU decode step is expected to reuse buffers from - * cpu_decode_scratch. Keep the allocation guard tightly scoped to the - * decode math itself; sampling, token emission, tracing, and callbacks - * may allocate small temporary strings without invalidating that - * guarantee. */ - ds4_alloc_guard_begin("CPU token decode"); - forward_token_raw_swa_cpu_decode_scratch(logits, model, weights, &cache, token, (uint32_t)pos, - directional_steering_dirs, - directional_steering_attn, - directional_steering_ffn, - &decode_scratch); - ds4_alloc_guard_end(); - if (token_timing) { - const double t_eval1 = now_sec(); - fprintf(stderr, "ds4: decode eval %d took %.3f ms\n", n_decode_eval + 1, (t_eval1 - t_eval0) * 1000.0); +static void model_map_span_vec_append(ds4_model_map_span_vec *spans, uint64_t lo, uint64_t hi, bool isolate) { + if (!spans || lo == UINT64_MAX || hi <= lo) return; + if (spans->len == spans->cap) { + uint32_t new_cap = spans->cap ? spans->cap * 2u : 16u; + spans->v = xrealloc(spans->v, (size_t)new_cap * sizeof(spans->v[0])); + spans->cap = new_cap; + } + spans->v[spans->len++] = (ds4_model_map_span){lo, hi, isolate}; +} + +static uint32_t model_map_q4_pro_group_views(void) { + uint32_t views = 1; + const char *env = getenv("DS4_METAL_Q4_PRO_MAP_GROUPS"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0' && v > 0 && v <= 384 && (384u % (uint32_t)v) == 0) { + views = (uint32_t)v; } - n_decode_eval++; - pos++; } - const double t_decode1 = now_sec(); - if (done) done(emit_ud); + return views; +} - const double prefill_s = t_prefill1 - t_prefill0; - const double decode_s = t_decode1 - t_decode0; - ds4_log(stderr, - DS4_LOG_TIMING, - "ds4: prefill: %.2f t/s, generation: %.2f t/s\n", - prefill_s > 0.0 ? (double)prompt->len / prefill_s : 0.0, - decode_s > 0.0 ? (double)n_generated / decode_s : 0.0); +static void model_map_span_vec_include_one(ds4_model_map_span_vec *spans, const ds4_tensor *t) { + if (!t || t->bytes == 0) return; + const uint64_t q4_isolated_min_bytes = 2ull * 1024ull * 1024ull * 1024ull; + const uint32_t q4_pro_group_views = model_map_q4_pro_group_views(); + if (t->type == DS4_TENSOR_Q4_K && + t->ndim == 3 && + t->dim[2] == 384 && + t->bytes >= q4_isolated_min_bytes && + (t->bytes % q4_pro_group_views) == 0) + { + /* + * PRO Q4 routed expert tensors are too large to hide inside broad + * layer spans. Isolate them so the default selected-expert path does + * not stack large aliases on top of layer-sized model views. Optional + * group splits are enabled by DS4_METAL_Q4_PRO_MAP_GROUPS for Metal + * experiments that bind stable grouped views. + */ + const uint64_t group_bytes = t->bytes / q4_pro_group_views; + if (group_bytes > spans->max_tensor_bytes) spans->max_tensor_bytes = group_bytes; + for (uint32_t i = 0; i < q4_pro_group_views; i++) { + const uint64_t lo = t->abs_offset + (uint64_t)i * group_bytes; + model_map_span_vec_append(spans, lo, lo + group_bytes, true); + } + return; + } - free(logits); - cpu_decode_scratch_free(&decode_scratch); - kv_cache_free(&cache); - return 0; + uint64_t lo = UINT64_MAX, hi = 0; + model_map_span_include_tensor(t, &lo, &hi, &spans->max_tensor_bytes); + const bool isolate = t->type == DS4_TENSOR_Q4_K && + t->bytes >= q4_isolated_min_bytes; + model_map_span_vec_append(spans, lo, hi, isolate); } -#ifndef DS4_NO_GPU -typedef struct { - uint32_t ctx_size; - uint32_t ctx_cap; - uint32_t normal_layers; - uint32_t layer_start; - uint32_t layer_end; - uint32_t layer_count; - uint64_t q_dim; - uint64_t q_nope; - uint64_t heads_dim; - uint64_t kv_raw_dim; - uint64_t dense_hidden_max; - uint64_t ffn_mid_elems; - - ds4_gpu_tensor *cur; - ds4_gpu_tensor *next; - ds4_gpu_tensor *attn_norm; - ds4_gpu_tensor *q_rank; - ds4_gpu_tensor *q_rank_norm; - ds4_gpu_tensor *q; - ds4_gpu_tensor *kv_raw; - ds4_gpu_tensor *kv_norm; - ds4_gpu_tensor *k_nope; - ds4_gpu_tensor *value; - ds4_gpu_tensor *heads; - ds4_gpu_tensor *attn_out; - ds4_gpu_tensor *after_attn; - ds4_gpu_tensor *ffn_norm; - ds4_gpu_tensor *ffn_gate; - ds4_gpu_tensor *ffn_up; - ds4_gpu_tensor *ffn_mid; - ds4_gpu_tensor *routed_gate; - ds4_gpu_tensor *routed_up; - ds4_gpu_tensor *routed_down; - ds4_gpu_tensor *ffn_out; - ds4_gpu_tensor *ffn_sum; - ds4_gpu_tensor *router_logits; - ds4_gpu_tensor *router_probs; - ds4_gpu_tensor *router_selected; - ds4_gpu_tensor *router_weights; - ds4_gpu_tensor *output_norm; - ds4_gpu_tensor *logits; - ds4_gpu_tensor *batch_router_logits; - ds4_gpu_tensor *batch_router_probs; - ds4_gpu_tensor *batch_router_selected; - ds4_gpu_tensor *batch_router_weights; - ds4_gpu_tensor *prefill_seed_router_selected; - uint32_t prefill_seed_tokens; - bool prefill_seed_layer_captured[DS4_MAX_LAYER]; - - ds4_gpu_tensor *prefill_tokens; - ds4_gpu_tensor *batch_cur; - ds4_gpu_tensor *batch_next; - ds4_gpu_tensor *batch_attn_norm; - ds4_gpu_tensor *batch_q_rank; - ds4_gpu_tensor *batch_q_rank_norm; - ds4_gpu_tensor *batch_q; - ds4_gpu_tensor *batch_kv_raw; - ds4_gpu_tensor *batch_kv_norm; - ds4_gpu_tensor *batch_k_nope; - ds4_gpu_tensor *batch_value; - ds4_gpu_tensor *batch_heads; - ds4_gpu_tensor *batch_attn_out; - ds4_gpu_tensor *batch_after_attn; - ds4_gpu_tensor *batch_ffn_norm; - ds4_gpu_tensor *batch_ffn_gate; - ds4_gpu_tensor *batch_ffn_up; - ds4_gpu_tensor *batch_shared_mid; - ds4_gpu_tensor *batch_ffn_mid; - ds4_gpu_tensor *batch_routed_gate; - ds4_gpu_tensor *batch_routed_up; - ds4_gpu_tensor *batch_routed_down; - ds4_gpu_tensor *batch_ffn_out; - bool batch_routed_mid_is_f16; - - uint32_t compact_cache_cap; - uint32_t indexed_prefill_cap; - uint32_t indexed_prefill_score_cap; - uint32_t indexer_full_layers; - ds4_gpu_tensor *indexer_k; - ds4_gpu_tensor *indexer_q; - ds4_gpu_tensor *indexer_weights; - ds4_gpu_tensor *indexer_scores; - ds4_gpu_tensor *indexer_selected; - ds4_gpu_tensor *qk_low; - ds4_gpu_tensor *attn_partial_lora; - ds4_gpu_tensor *attn_partial_ms; - ds4_gpu_tensor *batch_indexer_k; - ds4_gpu_tensor *batch_indexer_q; - ds4_gpu_tensor *batch_indexer_weights; - ds4_gpu_tensor *batch_indexer_scores; - ds4_gpu_tensor *batch_indexer_selected; - ds4_gpu_tensor *batch_qk_low; - ds4_gpu_tensor *batch_attn_lora; - ds4_gpu_tensor *layer_kv_lora_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_k_rope_cache[DS4_MAX_LAYER]; - /* GLM MTP (nextn block) drafting: private compact caches for the nextn - * layer (slot = absolute position; only [mtp_min_pos..pos] is ever - * selected) plus small scratch. Allocated lazily on first draft. */ - ds4_gpu_tensor *mtp_kv_lora_cache; - ds4_gpu_tensor *mtp_k_rope_cache; - ds4_gpu_tensor *mtp_concat; - ds4_gpu_tensor *mtp_selected; - float *mtp_logits_host; - int mtp_ready; - ds4_gpu_tensor *layer_indexer_key_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_key_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_value_cache[DS4_MAX_LAYER]; - bool full_kv_cache; - bool has_token_embd; - bool has_output_head; - bool quality; - bool ssd_streaming; - bool ssd_streaming_cold; - bool generic_routed_moe; - bool streaming_static_decode_map_current; - /* Tensor parallelism (50/50 expert sharding): tp_world 2 means - * this rank computes only its contiguous half of the routed experts - * and exchanges the 24KB routed-FFN partial at one gate per sparse - * layer. Views alias the engine's TP slab slots [layer*2 + FFN]. */ - uint32_t tp_world; - uint32_t tp_rank; - ds4_gpu_tensor **tp_out; - ds4_gpu_tensor **tp_in; - /* Prefill batch gate bounce buffers (shared storage; grow on demand). */ - ds4_gpu_tensor *tp_bounce_out; - ds4_gpu_tensor *tp_bounce_in; - /* CUDA multi-tier placement and device-local decode scratch mirrors. */ - const int *placement; -#define DS4_GLM_WS_SLOTS 29 - ds4_gpu_tensor *ws_mirror[DS4_MAX_GPUS][DS4_GLM_WS_SLOTS]; - ds4_gpu_tensor *ws_orig[DS4_GLM_WS_SLOTS]; - int ws_ready; - int ws_tier; -#define DS4_GLM_VERIFY_WS_SLOTS 28 - ds4_gpu_tensor *verify_ws_mirror[DS4_MAX_GPUS][DS4_GLM_VERIFY_WS_SLOTS]; - ds4_gpu_tensor *verify_ws_orig[DS4_GLM_VERIFY_WS_SLOTS]; - int verify_ws_ready; - int verify_ws_tier; -} ds4_glm_gpu_graph; - -static uint32_t glm_graph_model_context_limit(void) { - if (DS4_ROPE_ORIG_CTX > UINT32_MAX) return UINT32_MAX; - return (uint32_t)DS4_ROPE_ORIG_CTX; -} - -static double glm_graph_bytes_to_gib(uint64_t bytes) { - return (double)bytes / (1024.0 * 1024.0 * 1024.0); -} - -static uint64_t glm_graph_saturating_add_u64(uint64_t a, uint64_t b) { - return a > UINT64_MAX - b ? UINT64_MAX : a + b; +static void model_map_span_vec_include_layer(ds4_model_map_span_vec *spans, const ds4_layer_weights *l) { +#define DS4_INCLUDE_TENSOR(t_) model_map_span_vec_include_one(spans, (t_)) + DS4_INCLUDE_TENSOR(l->hc_attn_fn); + DS4_INCLUDE_TENSOR(l->hc_attn_scale); + DS4_INCLUDE_TENSOR(l->hc_attn_base); + DS4_INCLUDE_TENSOR(l->attn_norm); + DS4_INCLUDE_TENSOR(l->attn_q_a); + DS4_INCLUDE_TENSOR(l->attn_q_a_norm); + DS4_INCLUDE_TENSOR(l->attn_q_b); + DS4_INCLUDE_TENSOR(l->attn_kv); + DS4_INCLUDE_TENSOR(l->attn_kv_a_mqa); + DS4_INCLUDE_TENSOR(l->attn_kv_a_norm); + DS4_INCLUDE_TENSOR(l->attn_k_b); + DS4_INCLUDE_TENSOR(l->attn_v_b); + DS4_INCLUDE_TENSOR(l->attn_sinks); + DS4_INCLUDE_TENSOR(l->attn_output); + DS4_INCLUDE_TENSOR(l->attn_output_a); + DS4_INCLUDE_TENSOR(l->attn_output_b); + DS4_INCLUDE_TENSOR(l->attn_compressor_ape); + DS4_INCLUDE_TENSOR(l->attn_compressor_kv); + DS4_INCLUDE_TENSOR(l->attn_compressor_gate); + DS4_INCLUDE_TENSOR(l->attn_compressor_norm); + DS4_INCLUDE_TENSOR(l->indexer_attn_q_b); + DS4_INCLUDE_TENSOR(l->indexer_attn_k); + DS4_INCLUDE_TENSOR(l->indexer_k_norm); + DS4_INCLUDE_TENSOR(l->indexer_k_norm_b); + DS4_INCLUDE_TENSOR(l->indexer_proj); + DS4_INCLUDE_TENSOR(l->indexer_compressor_ape); + DS4_INCLUDE_TENSOR(l->indexer_compressor_kv); + DS4_INCLUDE_TENSOR(l->indexer_compressor_gate); + DS4_INCLUDE_TENSOR(l->indexer_compressor_norm); + DS4_INCLUDE_TENSOR(l->hc_ffn_fn); + DS4_INCLUDE_TENSOR(l->hc_ffn_scale); + DS4_INCLUDE_TENSOR(l->hc_ffn_base); + DS4_INCLUDE_TENSOR(l->ffn_norm); + DS4_INCLUDE_TENSOR(l->ffn_gate_tid2eid); + DS4_INCLUDE_TENSOR(l->ffn_gate); + DS4_INCLUDE_TENSOR(l->ffn_up); + DS4_INCLUDE_TENSOR(l->ffn_down); + DS4_INCLUDE_TENSOR(l->ffn_gate_inp); + DS4_INCLUDE_TENSOR(l->ffn_exp_probs_b); + DS4_INCLUDE_TENSOR(l->ffn_gate_exps); + DS4_INCLUDE_TENSOR(l->ffn_up_exps); + DS4_INCLUDE_TENSOR(l->ffn_down_exps); + DS4_INCLUDE_TENSOR(l->ffn_gate_shexp); + DS4_INCLUDE_TENSOR(l->ffn_up_shexp); + DS4_INCLUDE_TENSOR(l->ffn_down_shexp); + DS4_INCLUDE_TENSOR(l->nextn_eh_proj); + DS4_INCLUDE_TENSOR(l->nextn_enorm); + DS4_INCLUDE_TENSOR(l->nextn_hnorm); + DS4_INCLUDE_TENSOR(l->nextn_shared_head_norm); +#undef DS4_INCLUDE_TENSOR } -static bool glm_graph_env_disabled(const char *name) { - const char *env = getenv(name); - if (!env || !env[0]) return false; - return strcmp(env, "0") == 0 || - strcasecmp(env, "false") == 0 || - strcasecmp(env, "off") == 0 || - strcasecmp(env, "no") == 0; -} - -static double glm_graph_env_double( - const char *name, - double fallback, - double min_value, - double max_value) { - const char *env = getenv(name); - if (!env || !env[0]) return fallback; - char *end = NULL; - errno = 0; - const double v = strtod(env, &end); - if (end == env || errno != 0 || !isfinite(v)) return fallback; - if (v < min_value) return min_value; - if (v > max_value) return max_value; - return v; +static void model_map_span_vec_include_layer_decode_static(ds4_model_map_span_vec *spans, const ds4_layer_weights *l) { +#define DS4_INCLUDE_TENSOR(t_) model_map_span_vec_include_one(spans, (t_)) + DS4_INCLUDE_TENSOR(l->hc_attn_fn); + DS4_INCLUDE_TENSOR(l->hc_attn_scale); + DS4_INCLUDE_TENSOR(l->hc_attn_base); + DS4_INCLUDE_TENSOR(l->attn_norm); + DS4_INCLUDE_TENSOR(l->attn_q_a); + DS4_INCLUDE_TENSOR(l->attn_q_a_norm); + DS4_INCLUDE_TENSOR(l->attn_q_b); + DS4_INCLUDE_TENSOR(l->attn_kv); + DS4_INCLUDE_TENSOR(l->attn_kv_a_mqa); + DS4_INCLUDE_TENSOR(l->attn_kv_a_norm); + DS4_INCLUDE_TENSOR(l->attn_k_b); + DS4_INCLUDE_TENSOR(l->attn_v_b); + DS4_INCLUDE_TENSOR(l->attn_sinks); + DS4_INCLUDE_TENSOR(l->attn_output); + DS4_INCLUDE_TENSOR(l->attn_output_a); + DS4_INCLUDE_TENSOR(l->attn_output_b); + DS4_INCLUDE_TENSOR(l->attn_compressor_ape); + DS4_INCLUDE_TENSOR(l->attn_compressor_kv); + DS4_INCLUDE_TENSOR(l->attn_compressor_gate); + DS4_INCLUDE_TENSOR(l->attn_compressor_norm); + DS4_INCLUDE_TENSOR(l->indexer_attn_q_b); + DS4_INCLUDE_TENSOR(l->indexer_attn_k); + DS4_INCLUDE_TENSOR(l->indexer_k_norm); + DS4_INCLUDE_TENSOR(l->indexer_k_norm_b); + DS4_INCLUDE_TENSOR(l->indexer_proj); + DS4_INCLUDE_TENSOR(l->indexer_compressor_ape); + DS4_INCLUDE_TENSOR(l->indexer_compressor_kv); + DS4_INCLUDE_TENSOR(l->indexer_compressor_gate); + DS4_INCLUDE_TENSOR(l->indexer_compressor_norm); + DS4_INCLUDE_TENSOR(l->hc_ffn_fn); + DS4_INCLUDE_TENSOR(l->hc_ffn_scale); + DS4_INCLUDE_TENSOR(l->hc_ffn_base); + DS4_INCLUDE_TENSOR(l->ffn_norm); + DS4_INCLUDE_TENSOR(l->ffn_gate_tid2eid); + DS4_INCLUDE_TENSOR(l->ffn_gate); + DS4_INCLUDE_TENSOR(l->ffn_up); + DS4_INCLUDE_TENSOR(l->ffn_down); + DS4_INCLUDE_TENSOR(l->ffn_gate_inp); + DS4_INCLUDE_TENSOR(l->ffn_exp_probs_b); + DS4_INCLUDE_TENSOR(l->ffn_gate_shexp); + DS4_INCLUDE_TENSOR(l->ffn_up_shexp); + DS4_INCLUDE_TENSOR(l->ffn_down_shexp); + DS4_INCLUDE_TENSOR(l->nextn_eh_proj); + DS4_INCLUDE_TENSOR(l->nextn_enorm); + DS4_INCLUDE_TENSOR(l->nextn_hnorm); + DS4_INCLUDE_TENSOR(l->nextn_shared_head_norm); +#undef DS4_INCLUDE_TENSOR } -static uint64_t glm_graph_host_memory_bytes(void) { -#if defined(__APPLE__) - uint64_t mem = 0; - size_t len = sizeof(mem); - if (sysctlbyname("hw.memsize", &mem, &len, NULL, 0) != 0) return 0; - return mem; -#else - return 0; -#endif +static bool glm_stream_resident_decode_layer_supported( + const ds4_layer_weights *l, + uint32_t il) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || + !l || + il < DS4_N_LEADING_DENSE || + !l->ffn_gate_exps || + !l->ffn_up_exps || + !l->ffn_down_exps) { + return false; + } + if (l->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + l->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS || + l->ffn_down_exps->type == DS4_TENSOR_Q2_K)) { + return true; + } + return l->ffn_gate_exps->type == l->ffn_up_exps->type && + l->ffn_gate_exps->type == l->ffn_down_exps->type && + (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || + l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); } -static uint64_t glm_graph_streaming_active_model_bytes( - const ds4_weights *weights) { - if (!weights) return 0; +static uint32_t g_glm_streaming_full_resident_start; +static uint32_t g_glm_streaming_full_resident_layers; - uint64_t max_bytes = 0; - ds4_model_map_span_vec spans; +static bool glm_stream_resident_decode_layer_enabled( + const ds4_layer_weights *l, + uint32_t il) { + if (!glm_stream_resident_decode_layer_supported(l, il)) return false; + return g_glm_streaming_full_resident_layers != 0 && + il >= g_glm_streaming_full_resident_start && + il - g_glm_streaming_full_resident_start < + g_glm_streaming_full_resident_layers; +} - if (weights_layer_has_required(&weights->layer[0], 0) && - weights_model_map_token_spans(weights, &spans)) { - max_bytes = model_map_span_vec_total_bytes(&spans); - free(spans.v); - } - if (weights_have_output_head(weights) && - weights_model_map_output_spans(weights, &spans)) { - const uint64_t bytes = model_map_span_vec_total_bytes(&spans); - if (bytes > max_bytes) max_bytes = bytes; - free(spans.v); - } - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (!weights_model_map_spans(weights, il, il, false, &spans)) { - continue; - } - const uint64_t bytes = model_map_span_vec_total_bytes(&spans); - if (bytes > max_bytes) max_bytes = bytes; - free(spans.v); +static bool glm_stream_expert_cache_addr_layout_supported( + const ds4_weights *w, + const ds4_layer_weights *l, + uint32_t il) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || + !w || + !l || + il >= DS4_N_LAYER || + il < DS4_N_LEADING_DENSE || + !l->ffn_gate_exps || + !l->ffn_up_exps || + !l->ffn_down_exps || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > 8 || + DS4_N_EXPERT < 128 || + glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", + "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { + return false; } + if (!weights_streaming_layer_experts_uniform(w, il)) return false; - return max_bytes; + if (l->ffn_gate_exps->type != l->ffn_up_exps->type) return false; + const bool q2_addr = + l->ffn_gate_exps->type == DS4_TENSOR_Q2_K && + l->ffn_down_exps->type == DS4_TENSOR_Q2_K; + const bool q4_addr = + l->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + l->ffn_down_exps->type == DS4_TENSOR_Q4_K; + return q2_addr || q4_addr; } -/* TP shard bytes: dense weights plus this rank's routed-expert range. - * Zero when not sharding. Set during engine open, before the GLM memory - * guard runs. */ -static uint64_t g_tp_shard_model_bytes; +static DS4_MAYBE_UNUSED bool glm_stream_expert_cache_addr_supported( + const ds4_weights *w, + const ds4_layer_weights *l, + uint32_t il) { + if (!glm_stream_expert_cache_addr_layout_supported(w, l, il)) { + return false; + } -/* A user-raised iogpu.wired_limit_mb is an explicit GPU budget grant; - * prefer it over the fraction/reserve heuristics. */ -static uint64_t glm_graph_wired_limit_bytes(void) { -#ifdef __APPLE__ - int64_t mb = 0; - size_t len = sizeof(mb); - if (sysctlbyname("iogpu.wired_limit_mb", &mb, &len, NULL, 0) != 0) return 0; - if (mb <= 0) return 0; - return (uint64_t)mb * 1024ull * 1024ull; +#ifdef DS4_NO_GPU + return false; #else - return 0; + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!streaming_layer_gate_down_expert_bytes(l, + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + return ds4_gpu_stream_expert_cache_budget_for_expert_size( + gate_expert_bytes, + down_expert_bytes) >= DS4_N_EXPERT_USED; #endif } -static uint64_t glm_graph_model_bytes_for_guard( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - bool load_slice, - uint32_t layer_start, - uint32_t layer_end, - bool include_token, - bool include_output) { - if (!model) return 0; - /* Under TP, the sharded map bytes are authoritative regardless of - * how the caller frames the request (TP excludes real layer slicing, - * so any slice request here is the session's full-range accounting). */ - if (!ssd_streaming && g_tp_shard_model_bytes != 0) { - return g_tp_shard_model_bytes; - } - if (load_slice && weights) { - ds4_model_map_span_vec spans; - bool ok = false; - if (ssd_streaming) { - ok = weights_model_map_decode_static_slice_spans(weights, - layer_start, - layer_end, - include_token, - include_output, - &spans); - } else { - ok = weights_model_map_spans(weights, - layer_start, - layer_end, - include_output, - &spans); - } - if (ok) { - const uint64_t bytes = model_map_span_vec_total_bytes(&spans); - free(spans.v); - if (bytes != 0) return bytes; - } - } - if (!ssd_streaming) { - if (g_tp_shard_model_bytes != 0) return g_tp_shard_model_bytes; - return model->size; +static bool glm_stream_selected_expert_cache_supported( + const ds4_layer_weights *l, + uint32_t il) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || + !l || + il < DS4_N_LEADING_DENSE || + !l->ffn_gate_exps || + !l->ffn_up_exps || + !l->ffn_down_exps || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > 8 || + DS4_N_EXPERT < 128 || + glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") || + glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") || + glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", + "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { + return false; } - const uint64_t active_bytes = glm_graph_streaming_active_model_bytes(weights); - return active_bytes != 0 ? active_bytes : model->size; -} -static double glm_graph_memory_guard_default_reserve_gib( - uint64_t budget_base, - uint64_t model_bytes) { - const double base_gib = glm_graph_bytes_to_gib(budget_base); - const double model_gib = glm_graph_bytes_to_gib(model_bytes); - if (base_gib >= 480.0 && - base_gib <= 640.0 && - model_gib >= base_gib * 0.80) { - return 24.0; + if (l->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || + l->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS) { + return false; } - return 32.0; -} -static bool glm_graph_memory_guard_for_compact_cap( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - bool load_slice, - uint32_t layer_start, - uint32_t layer_end, - bool include_token, - bool include_output, - uint32_t ctx_size, - uint32_t compact_cap, - uint64_t transient_extra_bytes, - const char *phase) { - if (!model || glm_graph_env_disabled("DS4_GLM_MEMORY_GUARD")) return true; - - const uint64_t host_bytes = glm_graph_host_memory_bytes(); - uint64_t budget_base = host_bytes; - if (budget_base == 0) { - budget_base = ds4_gpu_recommended_working_set_size(); - } - if (budget_base == 0) return true; - const uint64_t wired_limit = glm_graph_wired_limit_bytes(); - - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const ds4_context_memory mem = load_slice ? - glm_graph_context_memory_estimate_for_compact_cap_slice( - ctx_size, - work_ctx, - compact_cap, - ssd_streaming, - layer_start, - layer_end) : - glm_graph_context_memory_estimate_for_compact_cap( - ctx_size, - work_ctx, - compact_cap, - ssd_streaming); - const uint64_t graph_bytes = mem.total_bytes; - const uint64_t model_bytes = - glm_graph_model_bytes_for_guard(model, - weights, - ssd_streaming, - load_slice, - layer_start, - layer_end, - include_token, - include_output); - uint64_t required = glm_graph_saturating_add_u64(model_bytes, graph_bytes); - required = glm_graph_saturating_add_u64(required, transient_extra_bytes); - - const double fraction = - glm_graph_env_double("DS4_GLM_MEMORY_GUARD_FRACTION", 0.99, 0.50, 1.00); - double default_reserve_gib = - glm_graph_memory_guard_default_reserve_gib(budget_base, model_bytes); -#ifdef DS4_ROCM_BUILD - if (load_slice && !ssd_streaming) { - /* The original fixed reserve protects Metal's shared host/GPU heap. - * A resident ROCm layer slice already accounts its exact model spans - * and owned graph state above. Keep proportional backend headroom for - * driver and temporary allocations without rejecting viable UMA - * slices merely because the heap is smaller than a high-memory Mac. */ - double rocm_reserve_gib = glm_graph_bytes_to_gib(budget_base) / 16.0; - if (rocm_reserve_gib < 8.0) rocm_reserve_gib = 8.0; - if (rocm_reserve_gib < default_reserve_gib) { - default_reserve_gib = rocm_reserve_gib; - } - } -#endif - const double reserve_gib = - glm_graph_env_double("DS4_GLM_MEMORY_GUARD_RESERVE_GB", - default_reserve_gib, - 0.0, - 1024.0); - const uint64_t fraction_budget = (uint64_t)((double)budget_base * fraction); - const uint64_t reserve_bytes = - (uint64_t)(reserve_gib * 1024.0 * 1024.0 * 1024.0); - const uint64_t reserve_budget = - reserve_bytes >= budget_base ? 0 : budget_base - reserve_bytes; - uint64_t budget = fraction_budget; - if (reserve_bytes != 0 && reserve_budget < budget) budget = reserve_budget; - if (wired_limit != 0) { - /* An explicitly raised iogpu.wired_limit_mb is the user granting - * the GPU that much wired memory; it overrides the heuristics - * (keep a small margin for non-model GPU allocations). */ - const uint64_t margin = 2ull * 1024ull * 1024ull * 1024ull; - const uint64_t wired_budget = - wired_limit > margin ? wired_limit - margin : wired_limit; - if (wired_budget > budget) budget = wired_budget; - } - - if (required <= budget) { - const char *report = getenv("DS4_GLM_MEMORY_GUARD_REPORT"); - if (report && report[0]) { - fprintf(stderr, - "ds4: GLM memory guard ctx=%u compact_cap=%u required=%.2f GiB " - "budget=%.2f GiB (model %.2f GiB, graph %.2f GiB, transient %.2f GiB)\n", - ctx_size, - mem.comp_cap, - glm_graph_bytes_to_gib(required), - glm_graph_bytes_to_gib(budget), - glm_graph_bytes_to_gib(model_bytes), - glm_graph_bytes_to_gib(graph_bytes), - glm_graph_bytes_to_gib(transient_extra_bytes)); - if (ssd_streaming && model_bytes != model->size) { - fprintf(stderr, - "ds4: GLM streaming guard uses active model span %.2f GiB " - "(full GGUF %.2f GiB)\n", - glm_graph_bytes_to_gib(model_bytes), - glm_graph_bytes_to_gib(model->size)); - } else if (load_slice && model_bytes != model->size) { - fprintf(stderr, - "ds4: GLM memory guard uses sliced model span %.2f GiB " - "(full GGUF %.2f GiB)\n", - glm_graph_bytes_to_gib(model_bytes), - glm_graph_bytes_to_gib(model->size)); - } - } - return true; + if (l->ffn_down_exps->type == DS4_TENSOR_Q2_K) { + return !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", + "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); } - - fprintf(stderr, - "ds4: GLM memory guard refused ctx=%u compact_cap=%u %s\n", - ctx_size, - mem.comp_cap, - phase ? phase : "before Metal graph allocation"); - if (ssd_streaming && model_bytes != model->size) { - fprintf(stderr, - "ds4: streamed active model map: %.2f GiB " - "(full GGUF %.2f GiB)\n", - glm_graph_bytes_to_gib(model_bytes), - glm_graph_bytes_to_gib(model->size)); - } else if (load_slice && model_bytes != model->size) { - fprintf(stderr, - "ds4: sliced model map: %.2f GiB " - "(full GGUF %.2f GiB)\n", - glm_graph_bytes_to_gib(model_bytes), - glm_graph_bytes_to_gib(model->size)); - } else { - fprintf(stderr, - "ds4: model map: %.2f GiB\n", - glm_graph_bytes_to_gib(model_bytes)); + if (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS) { + return !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_STREAM_ADDR_TABLE", + "DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE"); } - fprintf(stderr, - "ds4: graph cache/scratch: %.2f GiB " - "(full KV %.2f GiB, compact DSA %.2f GiB, scratch %.2f GiB)\n", - glm_graph_bytes_to_gib(graph_bytes), - glm_graph_bytes_to_gib(mem.raw_bytes), - glm_graph_bytes_to_gib(mem.compressed_bytes), - glm_graph_bytes_to_gib(mem.scratch_bytes)); - fprintf(stderr, - "ds4: required model+graph: %.2f GiB; guard budget: %.2f GiB " - "(base %.2f GiB, fraction %.2f, reserve %.2f GiB, transient %.2f GiB)\n", - glm_graph_bytes_to_gib(required), - glm_graph_bytes_to_gib(budget), - glm_graph_bytes_to_gib(budget_base), - fraction, - reserve_gib, - glm_graph_bytes_to_gib(transient_extra_bytes)); - fprintf(stderr, - "ds4: set DS4_GLM_MEMORY_GUARD=0 to bypass, use a smaller --ctx, " - "or use SSD streaming\n"); return false; } -static bool glm_graph_memory_guard( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - uint32_t ctx_size) { - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); - return glm_graph_memory_guard_for_compact_cap( - model, - weights, - ssd_streaming, - false, - 0, - 0, - true, - true, - ctx_size, - compact_cap, - 0, - "before GLM graph allocation"); -} - -static bool glm_graph_memory_guard_with_transient( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - uint32_t ctx_size, - uint64_t transient_extra_bytes, - const char *phase) { - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); - return glm_graph_memory_guard_for_compact_cap( - model, - weights, - ssd_streaming, - false, - 0, - 0, - true, - true, - ctx_size, - compact_cap, - transient_extra_bytes, - phase); -} - -static bool glm_graph_memory_guard_slice( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - uint32_t layer_start, - uint32_t layer_end, - bool include_token, - bool include_output, - uint32_t ctx_size) { - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); - return glm_graph_memory_guard_for_compact_cap( - model, - weights, - ssd_streaming, - true, - layer_start, - layer_end, - include_token, - include_output, - ctx_size, - compact_cap, - 0, - "before GLM graph allocation"); -} - -static bool glm_graph_memory_guard_slice_with_transient( - const ds4_model *model, - const ds4_weights *weights, - bool ssd_streaming, - uint32_t layer_start, - uint32_t layer_end, - bool include_token, - bool include_output, - uint32_t ctx_size, - uint64_t transient_extra_bytes, - const char *phase) { - const uint32_t work_ctx = - glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const uint32_t compact_cap = - glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); - return glm_graph_memory_guard_for_compact_cap( - model, - weights, - ssd_streaming, - true, - layer_start, - layer_end, - include_token, - include_output, - ctx_size, - compact_cap, - transient_extra_bytes, - phase); -} - -static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, - bool ssd_streaming) { - uint32_t cap = ssd_streaming ? - DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT : - DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT; - if (ctx_size >= DS4_GLM_METAL_LONG_CONTEXT_THRESHOLD && - cap > DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT) { - cap = DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT; - } - if (ctx_size > 0 && cap > ctx_size) cap = ctx_size; - if (cap == 0) cap = 1; - return cap; -} - -static uint32_t glm_graph_full_prefill_layer_flush_interval( - uint32_t n_tokens, - uint32_t command_rows, - bool logits_requested) { - /* Tiny logits-bearing passes (MTP verify, short prefills) must NOT - * flush per layer: 76 command-buffer round-trips cost ~35ms while the - * whole pass is ~70ms of GPU work. Real prefill chunks keep the - * interactive per-layer flush. */ - return (n_tokens > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || - command_rows > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || - (logits_requested && n_tokens > 8u)) ? 1u : 0u; -} - -static uint32_t glm_graph_prefill_progress_flush_interval( - uint32_t layer_flush_interval, - uint32_t n_tokens, - ds4_session_progress_fn display_progress, - uint32_t work_total) { - if (layer_flush_interval != 0) return layer_flush_interval; - (void)n_tokens; - (void)display_progress; - (void)work_total; - return 0; +static bool glm_stream_decode_experts_are_streamed( + const ds4_weights *w, + const ds4_layer_weights *l, + uint32_t il) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) return false; + return glm_stream_expert_cache_addr_layout_supported(w, l, il) || + glm_stream_selected_expert_cache_supported(l, il); } -static void glm_graph_report_prefill_display_progress( - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t absolute_base, - uint32_t work_done_base, - uint32_t n_tokens, - uint32_t layer_done, - uint32_t normal_layers, - uint32_t work_total, - bool allow_complete) { - if (!display_progress || work_total == 0) return; - - uint64_t chunk_done = 0; - if (normal_layers == 0 || layer_done >= normal_layers) { - chunk_done = n_tokens; - } else { - chunk_done = (uint64_t)n_tokens * (uint64_t)layer_done / - (uint64_t)normal_layers; - } - - uint64_t done = (uint64_t)work_done_base + chunk_done; - if (done > (uint64_t)work_total) done = work_total; - if (!allow_complete && done >= (uint64_t)work_total) { - done = work_total > 0 ? (uint64_t)work_total - 1u : 0u; +/* + * Decode-time spans for one layer. The static set excludes routed expert + * tensors only when the streaming expert-cache path can really serve them. + * Boosted layers, mixed GLM quant layouts such as Q4 gate/up plus Q5 down, or + * undersized expert caches fall back to direct model-range reads. Include + * those expert tensors so cache-hit prefill extension and decode are covered. + */ +static void model_map_span_vec_include_layer_decode( + ds4_model_map_span_vec *spans, + const ds4_weights *w, + uint32_t il) { + const ds4_layer_weights *l = &w->layer[il]; + model_map_span_vec_include_layer_decode_static(spans, l); + if (!weights_streaming_layer_experts_uniform(w, il) || + (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + !glm_stream_decode_experts_are_streamed(w, l, il)) || + glm_stream_resident_decode_layer_enabled(l, il)) { + model_map_span_vec_include_one(spans, l->ffn_gate_exps); + model_map_span_vec_include_one(spans, l->ffn_up_exps); + model_map_span_vec_include_one(spans, l->ffn_down_exps); } - display_progress(display_progress_ud, - "prefill_display", - (int)((uint64_t)absolute_base + done), - (int)((uint64_t)absolute_base + (uint64_t)work_total)); } -static bool glm_graph_small_prefill_stage_sync( - uint32_t n_tokens, - bool logits_requested) { - return logits_requested && - n_tokens > 0 && - n_tokens <= DS4_GLM_METAL_SMALL_PREFILL_STAGE_SYNC_TOKENS; +static void model_map_span_vec_include_output(ds4_model_map_span_vec *spans, const ds4_weights *w) { + model_map_span_vec_include_one(spans, w->output_hc_base); + model_map_span_vec_include_one(spans, w->output_hc_fn); + model_map_span_vec_include_one(spans, w->output_hc_scale); + model_map_span_vec_include_one(spans, w->output_norm); + model_map_span_vec_include_one(spans, w->output); } -static uint32_t glm_graph_indexed_decode_split_min_block_rows(void) { - return 32u; +static int model_map_span_cmp(const void *a, const void *b) { + const ds4_model_map_span *sa = a; + const ds4_model_map_span *sb = b; + if (sa->off < sb->off) return -1; + if (sa->off > sb->off) return 1; + if (sa->end < sb->end) return -1; + if (sa->end > sb->end) return 1; + return 0; } -static uint32_t glm_graph_indexed_decode_split_blocks(void) { - const uint32_t block_rows = glm_graph_indexed_decode_split_min_block_rows(); - const uint32_t top_k = glm_graph_indexer_top_k_limit(); - return (top_k + block_rows - 1u) / block_rows; -} +static bool model_map_span_vec_finish(ds4_model_map_span_vec *spans) { + if (!spans || spans->len == 0 || spans->max_tensor_bytes == 0) return false; -static uint32_t glm_graph_indexed_decode_split_block_rows_for(uint32_t n_selected) { - return n_selected <= 1024u ? 32u : 128u; + qsort(spans->v, spans->len, sizeof(spans->v[0]), model_map_span_cmp); + uint32_t out = 0; + for (uint32_t i = 0; i < spans->len; i++) { + if (out == 0 || + spans->v[i].off > spans->v[out - 1u].end || + spans->v[i].isolate || + spans->v[out - 1u].isolate) { + spans->v[out++] = spans->v[i]; + } else if (spans->v[i].end > spans->v[out - 1u].end) { + spans->v[out - 1u].end = spans->v[i].end; + } + } + spans->len = out; + return spans->len != 0; } -static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) { - const uint32_t block_rows = glm_graph_indexed_decode_split_block_rows_for(n_selected); - const uint32_t needed_blocks = - block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; - return n_selected > 512u && - block_rows > 0 && - needed_blocks > 0 && - needed_blocks <= glm_graph_indexed_decode_split_blocks() && - glm_graph_indexed_decode_split_blocks() <= 64u && - (DS4_N_HEAD % 8u) == 0 && - DS4_N_KV_LORA == 512u && - DS4_N_ROT == 64u && - glm_graph_compact_cache_is_f16(); -} +static DS4_MAYBE_UNUSED bool weights_model_map_spans( + const ds4_weights *w, + uint32_t layer_start, + uint32_t layer_end, + bool include_output, + ds4_model_map_span_vec *spans) { + if (!w || !spans) return false; + if (layer_start >= DS4_N_LAYER) return false; + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; -static bool glm_graph_prefill_stage_sync_boundary(void) { - if (ds4_gpu_end_commands() == 0) return false; - return ds4_gpu_begin_commands() != 0; + memset(spans, 0, sizeof(*spans)); + if (layer_start == 0) model_map_span_vec_include_one(spans, w->token_embd); + for (uint32_t il = layer_start; il <= layer_end; il++) { + model_map_span_vec_include_layer(spans, &w->layer[il]); + } + if (include_output) model_map_span_vec_include_output(spans, w); + return model_map_span_vec_finish(spans); } -static bool glm_graph_indexed_prefill_attention_boundary(void) { -#ifdef DS4_ROCM_BUILD - /* - * ROCm launches in this path are ordered on the default stream. The Metal - * backend still needs the encoder flush, but on ROCm it is a full-device - * synchronize and stalls every indexed-prefill layer. - */ - return true; -#else - return ds4_gpu_flush_encoder() != 0; -#endif -} +static const uint8_t *tensor_expert_bytes( + const ds4_model *m, + const ds4_tensor *w, + uint32_t expert, + uint64_t *in_dim, + uint64_t *out_dim, + uint64_t *row_bytes); -static DS4_MAYBE_UNUSED bool glm_graph_env_truthy(const char *env) { - return env && - env[0] && - strcmp(env, "0") != 0 && - strcasecmp(env, "false") != 0 && - strcasecmp(env, "off") != 0 && - strcasecmp(env, "no") != 0; +/* TP sharding keeps full layers but restricts every routed-expert blob to + * one contiguous rank range (rank 0 owns the lower expert ids, matching + * ds4_tp_owns_expert in metal/moe.metal). */ +static DS4_MAYBE_UNUSED bool weights_model_map_sharded_spans( + const ds4_weights *w, + const ds4_model *m, + int rank, + ds4_model_map_span_vec *spans) { + if (!w || !m || !spans || (rank != 0 && rank != 1)) return false; + memset(spans, 0, sizeof(*spans)); + model_map_span_vec_include_one(spans, w->token_embd); + for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { + const ds4_layer_weights *l = &w->layer[il]; + /* Dense/attention/router tensors only — the plain decode include + * would map the full expert blobs in non-streaming mode. */ + model_map_span_vec_include_layer_decode_static(spans, l); + const ds4_tensor *exps[3] = { l->ffn_gate_exps, l->ffn_up_exps, + l->ffn_down_exps }; + for (int t = 0; t < 3; t++) { + const ds4_tensor *x = exps[t]; + if (!x || x->ndim != 3 || x->dim[2] < 2) continue; + uint64_t in_dim = 0, out_dim = 0, row_bytes = 0; + (void)tensor_expert_bytes(m, x, 0, &in_dim, &out_dim, &row_bytes); + const uint64_t expert_bytes = out_dim * row_bytes; + const uint64_t low_experts = x->dim[2] / 2; + const uint64_t first_expert = rank == 1 ? low_experts : 0; + const uint64_t owned_experts = rank == 1 ? + x->dim[2] - low_experts : low_experts; + const uint64_t owned_bytes = owned_experts * expert_bytes; + const uint64_t lo = x->abs_offset + first_expert * expert_bytes; + /* Kernels index experts from the blob base, so the owned range + * must sit in one contiguous view. Rank 1 takes any remainder. */ + model_map_span_vec_append(spans, lo, lo + owned_bytes, true); + if (owned_bytes > spans->max_tensor_bytes) { + spans->max_tensor_bytes = owned_bytes; + } + } + } + model_map_span_vec_include_output(spans, w); + return model_map_span_vec_finish(spans); } -static bool glm_graph_streaming_prefill_sync_each_layer( - bool full_layer_prefill) { -#ifdef DS4_ROCM_BUILD - /* - * ROCm command boundaries are full device synchronizes. Compact streaming - * prefill can keep queued default-stream work alive across layer mappings: - * streamed model-range eviction synchronizes before freeing ranges, while - * selected-expert cache reuse/eviction is protected by reuse events. The - * full-layer expert cache is only double-buffered, so keep its old boundary. - */ - if (full_layer_prefill) return true; - const char *env = glm_graph_env_value( - "DS4_ROCM_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER", - "DS4_METAL_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER"); - if (!env) env = getenv("DS4_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER"); - return glm_graph_env_truthy(env); -#else - (void)full_layer_prefill; - return true; -#endif +static DS4_MAYBE_UNUSED bool weights_model_map_decode_layer_spans( + const ds4_weights *w, + uint32_t il, + ds4_model_map_span_vec *spans) { + if (!w || !spans || il >= DS4_N_LAYER) return false; + memset(spans, 0, sizeof(*spans)); + model_map_span_vec_include_layer_decode(spans, w, il); + return model_map_span_vec_finish(spans); } -static bool glm_graph_indexed_prefill_batch_available( - const ds4_glm_gpu_graph *g) { - return g && - g->compact_cache_cap != 0 && - g->indexed_prefill_cap != 0 && - g->indexed_prefill_score_cap != 0 && - g->batch_indexer_q && - g->batch_indexer_weights && - g->batch_indexer_scores && - g->batch_indexer_selected && - g->batch_qk_low && - g->batch_attn_lora; -} - -static bool glm_graph_indexed_prefill_batch_ready( - const ds4_glm_gpu_graph *g, - uint32_t pos) { - return glm_graph_indexed_prefill_batch_available(g) && - (!g->full_kv_cache || pos >= g->ctx_cap); -} - -static uint32_t glm_graph_limit_indexed_prefill_chunk( - uint32_t pos, - uint32_t chunk) { - const uint32_t top_k = glm_graph_indexer_top_k_limit(); - if (pos < top_k) { - const uint32_t bridge = top_k - pos; - if (bridge != 0 && chunk > bridge) chunk = bridge; - } - return chunk; -} - -static uint32_t glm_graph_indexed_prefill_chunk_tokens( - uint32_t full_attention_cap, - uint32_t compact_cap) { - (void)full_attention_cap; - uint32_t chunk = DS4_GLM_METAL_INDEXED_PREFILL_CHUNK_TOKENS; - if (compact_cap > 0 && chunk > compact_cap) chunk = compact_cap; - if (chunk == 0) chunk = 1; - return chunk; -} - -static uint32_t glm_graph_indexed_prefill_score_tokens( - uint32_t indexed_prefill_cap, - uint32_t compact_cap) { - if (indexed_prefill_cap == 0 || compact_cap == 0) return 0; - const uint32_t scratch_mb = DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB; - const uint64_t budget_bytes = (uint64_t)scratch_mb * 1024ull * 1024ull; - uint64_t budget_rows = budget_bytes / ((uint64_t)compact_cap * sizeof(float)); - if (budget_rows == 0) budget_rows = 1; - if (budget_rows > indexed_prefill_cap) budget_rows = indexed_prefill_cap; - if (budget_rows > UINT32_MAX) budget_rows = UINT32_MAX; - return (uint32_t)budget_rows; -} - -static bool glm_graph_context_request(int ctx_size, uint32_t *ctx_out) { - if (!ctx_out || ctx_size <= 0) return false; - const uint32_t model_ctx = glm_graph_model_context_limit(); - if ((uint64_t)(uint32_t)ctx_size > (uint64_t)model_ctx) { - fprintf(stderr, - "ds4: GLM context %d exceeds model context %u\n", - ctx_size, - model_ctx); - return false; +static DS4_MAYBE_UNUSED bool weights_model_map_decode_static_spans( + const ds4_weights *w, + bool include_token, + bool include_output, + ds4_model_map_span_vec *spans) { + if (!w || !spans) return false; + memset(spans, 0, sizeof(*spans)); + if (include_token) model_map_span_vec_include_one(spans, w->token_embd); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + model_map_span_vec_include_layer_decode(spans, w, il); } - *ctx_out = (uint32_t)ctx_size; - return true; + if (include_output) model_map_span_vec_include_output(spans, w); + return model_map_span_vec_finish(spans); } -static bool glm_graph_span_fits_context( - const ds4_glm_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens) { - return g && n_tokens > 0 && pos0 < g->ctx_size && n_tokens <= g->ctx_size - pos0; -} +static DS4_MAYBE_UNUSED bool weights_model_map_decode_static_slice_spans( + const ds4_weights *w, + uint32_t layer_start, + uint32_t layer_end, + bool include_token, + bool include_output, + ds4_model_map_span_vec *spans) { + if (!w || !spans) return false; + if (layer_start >= DS4_N_LAYER) return false; + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; -static bool glm_graph_span_fits_full_attention( - const ds4_glm_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens) { - return g && n_tokens > 0 && pos0 < g->ctx_cap && n_tokens <= g->ctx_cap - pos0; + memset(spans, 0, sizeof(*spans)); + if (include_token) model_map_span_vec_include_one(spans, w->token_embd); + for (uint32_t il = layer_start; il <= layer_end; il++) { + model_map_span_vec_include_layer_decode(spans, w, il); + } + if (include_output) model_map_span_vec_include_output(spans, w); + return model_map_span_vec_finish(spans); } -static void glm_graph_log_full_attention_limit( - const ds4_glm_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens) { - const uint32_t end = pos0 + n_tokens; - fprintf(stderr, - "ds4: GLM Metal full-attention work cap is %u tokens; " - "requested span [%u,%u) in ctx %u needs compact indexed attention\n", - g ? g->ctx_cap : 0, - pos0, - end, - g ? g->ctx_size : 0); +static DS4_MAYBE_UNUSED uint64_t model_map_span_vec_total_bytes( + const ds4_model_map_span_vec *spans) { + if (!spans) return 0; + uint64_t total = 0; + for (uint32_t i = 0; i < spans->len; i++) { + const uint64_t bytes = spans->v[i].end - spans->v[i].off; + if (total > UINT64_MAX - bytes) return UINT64_MAX; + total += bytes; + } + return total; } -static bool glm_graph_tensor_layout( - const ds4_tensor *t, - uint32_t type, - uint32_t ndim, - uint64_t dim0, - uint64_t dim1, - uint64_t dim2) { - if (!t || t->type != type || t->ndim != ndim) return false; - if (ndim > 0 && t->dim[0] != dim0) return false; - if (ndim > 1 && t->dim[1] != dim1) return false; - if (ndim > 2 && t->dim[2] != dim2) return false; - return true; -} +static DS4_MAYBE_UNUSED bool weights_streaming_non_routed_bytes( + const ds4_weights *w, + uint64_t *bytes_out) { + if (bytes_out) *bytes_out = 0; + if (!w || !bytes_out) return false; -static bool glm_graph_dense_tensor_layout( - const ds4_tensor *t, - uint32_t ndim, - uint64_t dim0, - uint64_t dim1, - uint64_t dim2) { - if (!t || !tensor_type_is_glm_dense_quant(t->type) || t->ndim != ndim) return false; - if (ndim > 0 && t->dim[0] != dim0) return false; - if (ndim > 1 && t->dim[1] != dim1) return false; - if (ndim > 2 && t->dim[2] != dim2) return false; + ds4_model_map_span_vec spans; + const bool include_token = + weights_layer_has_required(&w->layer[0], 0); + if (!weights_model_map_decode_static_spans(w, + include_token, + weights_have_output_head(w), + &spans)) { + return false; + } + *bytes_out = model_map_span_vec_total_bytes(&spans); + free(spans.v); return true; } -static bool glm_graph_layer_uses_generic_routed_moe( - const ds4_layer_weights *l) { - return l && - l->ffn_gate_exps && - l->ffn_up_exps && - l->ffn_down_exps && - l->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS; -} - -static bool glm_graph_stream_map_token( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - if (!g || !g->ssd_streaming) return true; - g->streaming_static_decode_map_current = false; - return metal_graph_stream_map_token(model, weights); -} - -static bool glm_graph_stream_layer_expert_cache_supported( - const ds4_weights *weights, - const ds4_layer_weights *l, - uint32_t il) { - if (!weights || !l) return false; - if (il < DS4_N_LEADING_DENSE) return true; - return glm_stream_decode_experts_are_streamed(weights, l, il); +static DS4_MAYBE_UNUSED bool weights_model_map_token_spans( + const ds4_weights *w, + ds4_model_map_span_vec *spans) { + if (!w || !spans) return false; + memset(spans, 0, sizeof(*spans)); + model_map_span_vec_include_one(spans, w->token_embd); + return model_map_span_vec_finish(spans); } -static bool glm_graph_stream_prefill_expert_addr_supported( - const ds4_weights *weights, - const ds4_layer_weights *l, - uint32_t il, - uint32_t n_tokens) { - if (il < DS4_N_LEADING_DENSE) return true; - if (n_tokens <= 1) return false; -#ifdef DS4_ROCM_BUILD - /* - * ROCm selected-address batch prefill has pointer kernels for the - * IQ2-gate/Q2-down generic path and the uniform Q2_K GLM path. Q4_K still - * maps the full layer until matching pointer kernels exist. - */ - if (glm_stream_selected_expert_cache_supported(l, il)) return true; - return l && - l->ffn_gate_exps && - l->ffn_up_exps && - l->ffn_down_exps && - l->ffn_gate_exps->type == DS4_TENSOR_Q2_K && - l->ffn_up_exps->type == DS4_TENSOR_Q2_K && - l->ffn_down_exps->type == DS4_TENSOR_Q2_K && - glm_stream_expert_cache_addr_layout_supported(weights, l, il); -#else - return glm_stream_expert_cache_addr_supported(weights, l, il); -#endif +static DS4_MAYBE_UNUSED bool weights_model_map_output_spans( + const ds4_weights *w, + ds4_model_map_span_vec *spans) { + if (!w || !spans) return false; + memset(spans, 0, sizeof(*spans)); + model_map_span_vec_include_output(spans, w); + return model_map_span_vec_finish(spans); } -static bool rocm_graph_glm_stream_prefill_full_layer_enabled( - const ds4_glm_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens); - -static bool glm_graph_stream_map_decode_layer( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il) { - if (!g || !g->ssd_streaming) return true; - g->streaming_static_decode_map_current = false; - if (glm_graph_env_present("DS4_ROCM_GLM_STREAMING_DECODE_FULL_LAYER_MAP", - "DS4_METAL_GLM_STREAMING_DECODE_FULL_LAYER_MAP") || - getenv("DS4_GLM_STREAMING_DECODE_FULL_LAYER_MAP") != NULL) { - return metal_graph_stream_map_layer(model, weights, il); - } - if (weights && il < DS4_N_LAYER && - glm_stream_resident_decode_layer_enabled(&weights->layer[il], il)) { - return metal_graph_stream_map_layer(model, weights, il); - } - if (weights && il < DS4_N_LAYER && - glm_graph_stream_layer_expert_cache_supported(weights, - &weights->layer[il], - il)) { - return metal_graph_stream_map_layer_decode(model, weights, il); - } - return metal_graph_stream_map_layer(model, weights, il); -} +static void mtp_weights_bind(ds4_mtp_weights *w, const ds4_model *m) { + memset(w, 0, sizeof(*w)); -static bool glm_graph_stream_map_prefill_layer( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il, - uint32_t n_tokens, - bool full_layer_prefill) { - if (!g || !g->ssd_streaming) return true; - g->streaming_static_decode_map_current = false; - if (full_layer_prefill) { - const char *trace = glm_graph_env_value("DS4_ROCM_STREAMING_MAP_TRACE", - "DS4_METAL_STREAMING_MAP_TRACE"); - if (trace && trace[0] && strcmp(trace, "0") != 0) { - fprintf(stderr, - "ds4: GLM SSD prefill map layer=%u tokens=%u mode=full-prefill\n", - il, - n_tokens); - } -#ifdef DS4_ROCM_BUILD - if (weights && - il < DS4_N_LAYER && - rocm_graph_glm_stream_prefill_full_layer_enabled(g, - &weights->layer[il], - il, - n_tokens)) { - return metal_graph_stream_map_layer_decode(model, weights, il); - } -#endif - return metal_graph_stream_map_layer(model, weights, il); - } - const bool addr_supported = - weights && il < DS4_N_LAYER && - glm_graph_stream_prefill_expert_addr_supported(weights, - &weights->layer[il], - il, - n_tokens); - const char *trace = glm_graph_env_value("DS4_ROCM_STREAMING_MAP_TRACE", - "DS4_METAL_STREAMING_MAP_TRACE"); - if (trace && trace[0] && strcmp(trace, "0") != 0) { - fprintf(stderr, - "ds4: GLM SSD prefill map layer=%u tokens=%u mode=%s\n", - il, - n_tokens, - addr_supported ? "decode-expert-cache" : "full-layer"); - } - if (addr_supported) { - return metal_graph_stream_map_layer_decode(model, weights, il); - } - return metal_graph_stream_map_layer(model, weights, il); -} + w->hc_head_base = required_tensor(m, "mtp.0.hc_head_base.weight"); + w->hc_head_fn = required_tensor(m, "mtp.0.hc_head_fn.weight"); + w->hc_head_scale = required_tensor(m, "mtp.0.hc_head_scale.weight"); + w->e_proj = required_tensor(m, "mtp.0.e_proj.weight"); + w->h_proj = required_tensor(m, "mtp.0.h_proj.weight"); + w->enorm = required_tensor(m, "mtp.0.enorm.weight"); + w->hnorm = required_tensor(m, "mtp.0.hnorm.weight"); + w->norm = required_tensor(m, "mtp.0.norm.weight"); -#ifdef DS4_ROCM_BUILD -enum { DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 1024 }; -#else -enum { DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 64 }; -#endif + ds4_layer_weights *l = &w->block; + l->hc_attn_fn = required_tensor(m, "mtp.0.hc_attn_fn.weight"); + l->hc_attn_scale = required_tensor(m, "mtp.0.hc_attn_scale.weight"); + l->hc_attn_base = required_tensor(m, "mtp.0.hc_attn_base.weight"); + l->attn_norm = required_tensor(m, "mtp.0.attn_norm.weight"); + l->attn_q_a = required_tensor(m, "mtp.0.attn_q_a.weight"); + l->attn_q_a_norm = required_tensor(m, "mtp.0.attn_q_a_norm.weight"); + l->attn_q_b = required_tensor(m, "mtp.0.attn_q_b.weight"); + l->attn_kv = required_tensor(m, "mtp.0.attn_kv.weight"); + l->attn_kv_a_norm = required_tensor(m, "mtp.0.attn_kv_a_norm.weight"); + l->attn_sinks = required_tensor(m, "mtp.0.attn_sinks.weight"); + l->attn_output_a = required_tensor(m, "mtp.0.attn_output_a.weight"); + l->attn_output_b = required_tensor(m, "mtp.0.attn_output_b.weight"); + l->hc_ffn_fn = required_tensor(m, "mtp.0.hc_ffn_fn.weight"); + l->hc_ffn_scale = required_tensor(m, "mtp.0.hc_ffn_scale.weight"); + l->hc_ffn_base = required_tensor(m, "mtp.0.hc_ffn_base.weight"); + l->ffn_norm = required_tensor(m, "mtp.0.ffn_norm.weight"); + l->ffn_gate_inp = required_tensor(m, "mtp.0.ffn_gate_inp.weight"); + l->ffn_exp_probs_b = required_tensor(m, "mtp.0.exp_probs_b.bias"); + l->ffn_gate_exps = required_tensor(m, "mtp.0.ffn_gate_exps.weight"); + l->ffn_up_exps = required_tensor(m, "mtp.0.ffn_up_exps.weight"); + l->ffn_down_exps = required_tensor(m, "mtp.0.ffn_down_exps.weight"); + l->ffn_gate_shexp = required_tensor(m, "mtp.0.ffn_gate_shexp.weight"); + l->ffn_up_shexp = required_tensor(m, "mtp.0.ffn_up_shexp.weight"); + l->ffn_down_shexp = required_tensor(m, "mtp.0.ffn_down_shexp.weight"); -static uint32_t glm_graph_stream_prefill_full_layer_min_tokens(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_GLM_STREAMING_PREFILL_FULL_LAYER_MIN_TOKENS", - "DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER_MIN_TOKENS"); - if (!env) return DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS; - char *end = NULL; - errno = 0; - unsigned long v = strtoul(env, &end, 10); - if (end == env || errno != 0 || v == 0 || v > UINT32_MAX) { - return DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS; - } - return (uint32_t)v; + mtp_weights_validate_layout(w); } -static bool glm_graph_stream_prefill_full_layer_enabled( - const ds4_glm_gpu_graph *g, - uint32_t n_tokens) { - if (!g || !g->ssd_streaming) return false; - if (glm_graph_env_present( - "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER", - "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER")) { - return false; - } - if (glm_graph_env_present("DS4_ROCM_GLM_STREAMING_PREFILL_FULL_LAYER", - "DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER")) { - return true; +static ds4_tensor *dspark_bind_tensor( + ds4_dspark_weights *dw, + const ds4_model *m, + uint32_t stage, + const char *suffix, + bool required) { + ds4_tensor *t = tensor_by_mtp_stage_suffix(m, stage, suffix); + if (t) { + dw->present_tensors++; + } else if (required) { + dw->missing_tensors++; } - return n_tokens >= glm_graph_stream_prefill_full_layer_min_tokens(); + return t; } -static bool glm_graph_stream_prefill_full_layer_prepare_enabled( - const ds4_glm_gpu_graph *g, - bool full_layer_prefill) { - return g && - g->ssd_streaming && - full_layer_prefill && - !glm_graph_env_present( - "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER_PREPARE", - "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER_PREPARE"); +static void dspark_bind_block( + ds4_dspark_weights *dw, + ds4_layer_weights *l, + const ds4_model *m, + uint32_t stage) { + l->hc_attn_fn = dspark_bind_tensor(dw, m, stage, "hc_attn_fn.weight", true); + l->hc_attn_scale = dspark_bind_tensor(dw, m, stage, "hc_attn_scale.weight", true); + l->hc_attn_base = dspark_bind_tensor(dw, m, stage, "hc_attn_base.weight", true); + l->attn_norm = dspark_bind_tensor(dw, m, stage, "attn_norm.weight", true); + l->attn_q_a = dspark_bind_tensor(dw, m, stage, "attn_q_a.weight", true); + l->attn_q_a_norm = dspark_bind_tensor(dw, m, stage, "attn_q_a_norm.weight", true); + l->attn_q_b = dspark_bind_tensor(dw, m, stage, "attn_q_b.weight", true); + l->attn_kv = dspark_bind_tensor(dw, m, stage, "attn_kv.weight", true); + l->attn_kv_a_norm = dspark_bind_tensor(dw, m, stage, "attn_kv_a_norm.weight", true); + l->attn_sinks = dspark_bind_tensor(dw, m, stage, "attn_sinks.weight", true); + l->attn_output_a = dspark_bind_tensor(dw, m, stage, "attn_output_a.weight", true); + l->attn_output_b = dspark_bind_tensor(dw, m, stage, "attn_output_b.weight", true); + l->hc_ffn_fn = dspark_bind_tensor(dw, m, stage, "hc_ffn_fn.weight", true); + l->hc_ffn_scale = dspark_bind_tensor(dw, m, stage, "hc_ffn_scale.weight", true); + l->hc_ffn_base = dspark_bind_tensor(dw, m, stage, "hc_ffn_base.weight", true); + l->ffn_norm = dspark_bind_tensor(dw, m, stage, "ffn_norm.weight", true); + l->ffn_gate_inp = dspark_bind_tensor(dw, m, stage, "ffn_gate_inp.weight", true); + l->ffn_exp_probs_b = dspark_bind_tensor(dw, m, stage, "exp_probs_b.bias", true); + l->ffn_gate_exps = dspark_bind_tensor(dw, m, stage, "ffn_gate_exps.weight", true); + l->ffn_up_exps = dspark_bind_tensor(dw, m, stage, "ffn_up_exps.weight", true); + l->ffn_down_exps = dspark_bind_tensor(dw, m, stage, "ffn_down_exps.weight", true); + l->ffn_gate_shexp = dspark_bind_tensor(dw, m, stage, "ffn_gate_shexp.weight", true); + l->ffn_up_shexp = dspark_bind_tensor(dw, m, stage, "ffn_up_shexp.weight", true); + l->ffn_down_shexp = dspark_bind_tensor(dw, m, stage, "ffn_down_shexp.weight", true); } -#ifdef DS4_ROCM_BUILD -static bool rocm_graph_glm_stream_prefill_full_layer_enabled( - const ds4_glm_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens) { - return glm_graph_stream_prefill_full_layer_enabled(g, n_tokens) && - layer && - glm_stream_resident_decode_layer_supported(layer, il); -} - -static bool rocm_graph_glm_stream_layer_expert_load_start_next( - rocm_graph_stream_layer_expert_load *job, - const ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t first_il, - uint32_t last_il, - uint32_t n_tokens) { - if (!job || !model || !weights || first_il > last_il) return true; - if (job->active) return true; - for (uint32_t il = first_il; il <= last_il && il < DS4_N_LAYER; il++) { - const ds4_layer_weights *layer = &weights->layer[il]; - if (!rocm_graph_glm_stream_prefill_full_layer_enabled(g, - layer, - il, - n_tokens)) { - continue; - } - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!rocm_graph_stream_layer_expert_bytes(layer, - &gate_expert_bytes, - &down_expert_bytes)) { - return false; +static void dspark_weights_bind_optional( + ds4_dspark_weights *dw, + const ds4_model *m, + const ds4_dspark_summary *summary) { + memset(dw, 0, sizeof(*dw)); + if (!m || !summary) return; + + dw->n_stages = summary->stages < DS4_DSPARK_MAX_STAGES ? + summary->stages : DS4_DSPARK_MAX_STAGES; + dw->block_size = summary->block_size; + dw->markov_rank = summary->markov_rank; + dw->noise_token_id = summary->noise_token_id; + dw->target_layer_count = summary->target_layer_count; + dw->has_block_size = summary->has_block_size; + dw->has_markov_rank = summary->has_markov_rank; + dw->has_noise_token_id = summary->has_noise_token_id; + dw->has_target_layers = summary->has_target_layers; + memcpy(dw->target_layers, + summary->target_layers, + (size_t)dw->target_layer_count * sizeof(dw->target_layers[0])); + if (summary->stages > DS4_DSPARK_MAX_STAGES) dw->missing_tensors++; + + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + ds4_dspark_stage_weights *sw = &dw->stage[stage]; + dspark_bind_block(dw, &sw->block, m, stage); + if (stage == 0) { + sw->main_proj = dspark_bind_tensor(dw, m, stage, "main_proj.weight", true); + sw->main_norm = dspark_bind_tensor(dw, m, stage, "main_norm.weight", true); } - return rocm_graph_stream_layer_expert_load_start(job, - model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); } - return true; -} -static bool rocm_graph_glm_stream_layer_expert_load_ready( - rocm_graph_stream_layer_expert_load *job, - const ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t il, - uint32_t n_tokens) { - if (!model || !weights || il >= DS4_N_LAYER) return false; - const ds4_layer_weights *layer = &weights->layer[il]; - if (!rocm_graph_glm_stream_prefill_full_layer_enabled(g, - layer, - il, - n_tokens)) { - return true; - } - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!rocm_graph_stream_layer_expert_bytes(layer, - &gate_expert_bytes, - &down_expert_bytes)) { - return false; - } - if (job && job->active) { - if (job->il != il) { - fprintf(stderr, - "ds4: GLM ROCm streaming full-layer expert load expected " - "layer %u but pending job is layer %u\n", - il, - job->il); - return false; - } - return rocm_graph_stream_layer_expert_load_join(job); + if (dw->n_stages != 0) { + const uint32_t final_stage = dw->n_stages - 1u; + ds4_dspark_stage_weights *sw = &dw->stage[final_stage]; + sw->norm = dspark_bind_tensor(dw, m, final_stage, "norm.weight", true); + sw->hc_head_base = + dspark_bind_tensor(dw, m, final_stage, "hc_head_base.weight", true); + sw->hc_head_fn = + dspark_bind_tensor(dw, m, final_stage, "hc_head_fn.weight", true); + sw->hc_head_scale = + dspark_bind_tensor(dw, m, final_stage, "hc_head_scale.weight", true); + sw->markov_w1 = + dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w1.weight", true); + sw->markov_w2 = + dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w2.weight", true); + sw->confidence_proj = + dspark_bind_tensor(dw, m, final_stage, "confidence_head.proj.weight", true); } - return rocm_graph_stream_layer_expert_load_sync(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); -} -#else -static bool rocm_graph_glm_stream_prefill_full_layer_enabled( - const ds4_glm_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens) { - (void)g; - (void)layer; - (void)il; - (void)n_tokens; - return false; -} -#endif -static bool glm_graph_stream_map_output( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - if (!g || !g->ssd_streaming) return true; - g->streaming_static_decode_map_current = false; - return metal_graph_stream_map_output(model, weights); + dspark_weights_validate_layout(dw); } -static bool glm_graph_validate_expert_layout( - const ds4_model *model, - const ds4_tensor *gate, - const ds4_tensor *up, - const ds4_tensor *down, - uint64_t *gate_row_bytes, - uint64_t *up_row_bytes, - uint64_t *down_row_bytes) { - if (!gate || !up || !down) return false; - if (!glm_graph_gate_pair_type_supported(gate->type, up->type) || - !glm_graph_down_type_supported(down->type) || - gate->ndim != 3 || up->ndim != 3 || down->ndim != 3 || - gate->dim[0] != DS4_N_EMBD || - gate->dim[1] != DS4_N_FF_EXP || - gate->dim[2] != DS4_N_EXPERT || - up->dim[0] != DS4_N_EMBD || - up->dim[1] != DS4_N_FF_EXP || - up->dim[2] != DS4_N_EXPERT || - down->dim[0] != DS4_N_FF_EXP || - down->dim[1] != DS4_N_EMBD || - down->dim[2] != DS4_N_EXPERT) { - return false; - } +#include "kernels/cpu_matmul.inc" +#include "models/deepseek/cpu.inc" +#include "models/glm/cpu.inc" - uint64_t gate_in = 0, gate_out = 0; - uint64_t up_in = 0, up_out = 0; - uint64_t down_in = 0, down_out = 0; - (void)tensor_expert_bytes(model, gate, 0, &gate_in, &gate_out, gate_row_bytes); - (void)tensor_expert_bytes(model, up, 0, &up_in, &up_out, up_row_bytes); - (void)tensor_expert_bytes(model, down, 0, &down_in, &down_out, down_row_bytes); - return gate_in == DS4_N_EMBD && - up_in == DS4_N_EMBD && - down_in == DS4_N_FF_EXP && - gate_out == DS4_N_FF_EXP && - up_out == DS4_N_FF_EXP && - down_out == DS4_N_EMBD; -} +#ifndef DS4_NO_GPU +static int sample_argmax(const float *logits, uint32_t n_vocab); -static bool glm_graph_validate_layer_layout( - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint64_t q_dim, - uint64_t q_nope, - uint64_t heads_dim, - uint64_t *kv_raw_dim_out, - uint64_t *dense_hidden_max) { - if (!l) return false; - const uint64_t kv_raw_dim = l->attn_kv_a_mqa ? l->attn_kv_a_mqa->dim[1] : 0; - const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - if (!glm_graph_tensor_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0) || - !glm_graph_dense_tensor_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0) || - !glm_graph_tensor_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0) || - !glm_graph_dense_tensor_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0) || - !l->attn_kv_a_mqa || - !tensor_type_is_glm_dense_quant(l->attn_kv_a_mqa->type) || - l->attn_kv_a_mqa->ndim != 2 || - l->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || - kv_raw_dim < (uint64_t)DS4_N_KV_LORA + DS4_N_ROT || - !glm_graph_tensor_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_KV_LORA, 0, 0) || - !glm_graph_dense_tensor_layout(l->attn_k_b, 3, q_nope, DS4_N_KV_LORA, DS4_N_HEAD) || - !glm_graph_dense_tensor_layout(l->attn_v_b, 3, DS4_N_KV_LORA, DS4_N_VALUE_MLA, DS4_N_HEAD) || - !glm_graph_dense_tensor_layout(l->attn_output, 2, heads_dim, DS4_N_EMBD, 0) || - !glm_graph_dense_tensor_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, indexer_q_dim, 0) || - !glm_graph_dense_tensor_layout(l->indexer_attn_k, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 0) || - !glm_graph_tensor_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0) || - !glm_graph_tensor_layout(l->indexer_k_norm_b, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0) || - !glm_graph_tensor_layout(l->indexer_proj, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0) || - !glm_graph_tensor_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0)) { - fprintf(stderr, "ds4: GLM Metal graph found unexpected attention layout in layer %u\n", il); - return false; - } - if (kv_raw_dim > *kv_raw_dim_out) *kv_raw_dim_out = kv_raw_dim; +/* ========================================================================= + * Metal Reference Comparison Helpers. + * ========================================================================= + * + * These small scalar helpers are used only by diagnostics that compare the C + * reference path with the Metal executor. + */ - if (il < DS4_N_LEADING_DENSE) { - const uint64_t hidden = l->ffn_gate ? l->ffn_gate->dim[1] : 0; - if (!l->ffn_gate || - !glm_graph_dense_tensor_layout(l->ffn_gate, 2, DS4_N_EMBD, hidden, 0) || - !glm_graph_dense_tensor_layout(l->ffn_up, 2, DS4_N_EMBD, hidden, 0) || - !glm_graph_dense_tensor_layout(l->ffn_down, 2, hidden, DS4_N_EMBD, 0)) { - fprintf(stderr, "ds4: GLM Metal graph found unexpected dense FFN layout in layer %u\n", il); - return false; - } - if (hidden > *dense_hidden_max) *dense_hidden_max = hidden; - } else { - uint64_t gate_row_bytes = 0, up_row_bytes = 0, down_row_bytes = 0; - if (!glm_graph_tensor_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0) || - !glm_graph_tensor_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0) || - !glm_graph_validate_expert_layout(model, - l->ffn_gate_exps, - l->ffn_up_exps, - l->ffn_down_exps, - &gate_row_bytes, - &up_row_bytes, - &down_row_bytes) || - !glm_graph_dense_tensor_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0) || - !glm_graph_dense_tensor_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0) || - !glm_graph_dense_tensor_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0)) { - fprintf(stderr, "ds4: GLM Metal graph found unexpected sparse FFN layout in layer %u\n", il); - return false; - } - (void)gate_row_bytes; - (void)up_row_bytes; - (void)down_row_bytes; +static float max_abs_diff(const float *a, const float *b, uint64_t n) { + float max_diff = 0.0f; + for (uint64_t i = 0; i < n; i++) { + const float diff = fabsf(a[i] - b[i]); + if (diff > max_diff) max_diff = diff; } - return true; + return max_diff; } -static bool glm_graph_validate_layout( - const ds4_model *model, - const ds4_weights *weights, - ds4_glm_gpu_graph *g, - uint32_t layer_start, - uint32_t layer_end, - bool require_token_embd, - bool require_output_head) { - if (!model || !weights || !g) return false; - const uint32_t normal_layers = glm_graph_normal_layer_count(); - if (normal_layers == 0 || DS4_N_ROT >= DS4_N_KEY_MLA) { - fprintf(stderr, "ds4: GLM Metal graph found unsupported layer/key dimensions\n"); - return false; - } - if (layer_end == UINT32_MAX) layer_end = normal_layers - 1u; - if (layer_start > layer_end || layer_end >= normal_layers) { - fprintf(stderr, - "ds4: GLM Metal graph found invalid layer slice %u:%u for %u normal layers\n", - layer_start, - layer_end, - normal_layers); - return false; - } - - g->has_token_embd = weights->token_embd != NULL; - g->has_output_head = weights_have_output_head(weights); - if (require_token_embd && !g->has_token_embd) { - fprintf(stderr, "ds4: GLM Metal graph layer slice requires token embeddings\n"); - return false; - } - if (g->has_token_embd && - !glm_graph_dense_tensor_layout(weights->token_embd, 2, - DS4_N_EMBD, DS4_N_VOCAB, 0)) { - fprintf(stderr, "ds4: GLM Metal graph found unexpected token embedding layout\n"); - return false; - } - if (require_output_head && !g->has_output_head) { - fprintf(stderr, "ds4: GLM Metal graph layer slice requires the output head\n"); - return false; - } - if (weights_have_partial_output_head(weights) && !g->has_output_head) { - fprintf(stderr, "ds4: GLM Metal graph found partial output head\n"); - return false; - } - if (g->has_output_head && - (!glm_graph_tensor_layout(weights->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0) || - !glm_graph_dense_tensor_layout(weights->output, 2, - DS4_N_EMBD, DS4_N_VOCAB, 0))) { - fprintf(stderr, "ds4: GLM Metal graph found unexpected output head layout\n"); - return false; +static float rms_abs_diff(const float *a, const float *b, uint64_t n) { + double ss = 0.0; + for (uint64_t i = 0; i < n; i++) { + const double d = (double)a[i] - (double)b[i]; + ss += d * d; } + return n ? (float)sqrt(ss / (double)n) : 0.0f; +} - g->normal_layers = normal_layers; - g->layer_start = layer_start; - g->layer_end = layer_end; - g->layer_count = layer_end - layer_start + 1u; - g->q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; - g->q_nope = (uint64_t)DS4_N_KEY_MLA - DS4_N_ROT; - g->heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; - g->dense_hidden_max = DS4_N_FF_EXP; - g->kv_raw_dim = 0; - for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { - if (!glm_graph_validate_layer_layout(model, - &weights->layer[il], - il, - g->q_dim, - g->q_nope, - g->heads_dim, - &g->kv_raw_dim, - &g->dense_hidden_max)) { - return false; - } - if (glm_graph_layer_uses_generic_routed_moe(&weights->layer[il])) { - g->generic_routed_moe = true; - } +static uint64_t argmax_f32(const float *x, uint64_t n) { + uint64_t best = 0; + for (uint64_t i = 1; i < n; i++) { + if (x[i] > x[best]) best = i; } - const uint64_t sparse_mid_elems = (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; - g->ffn_mid_elems = - g->dense_hidden_max > sparse_mid_elems ? g->dense_hidden_max : sparse_mid_elems; - return g->layer_count > 0 && - g->kv_raw_dim >= (uint64_t)DS4_N_KV_LORA + DS4_N_ROT; + return best; } -/* Per-tier decode scratch keeps layer kernels off peer-mapped work buffers. */ -#define DS4_GLM_WS_FIELDS(X) \ - X(cur) X(next) X(attn_norm) X(q_rank) X(q_rank_norm) X(q) X(kv_raw) \ - X(kv_norm) X(k_nope) X(value) X(heads) X(attn_out) X(after_attn) \ - X(ffn_norm) X(ffn_gate) X(ffn_up) X(ffn_mid) X(ffn_out) X(ffn_sum) \ - X(router_logits) X(router_probs) X(router_selected) X(router_weights) \ - X(indexer_k) X(indexer_q) X(indexer_weights) X(indexer_scores) \ - X(indexer_selected) X(qk_low) +#endif -static ds4_gpu_tensor **glm_graph_ws_slot(ds4_glm_gpu_graph *g, int i) { - int n = 0; -#define DS4_GLM_WS_SLOT_CASE(field) if (n++ == i) return &g->field; - DS4_GLM_WS_FIELDS(DS4_GLM_WS_SLOT_CASE) -#undef DS4_GLM_WS_SLOT_CASE - return NULL; -} +static void print_vec_stats(const char *name, const float *x, uint64_t n) { + float minv = DS4_POS_INF; + float maxv = DS4_NEG_INF; + double ss = 0.0; -static void glm_graph_ws_free(ds4_glm_gpu_graph *g) { - if (!g || !g->ws_ready) return; - for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); - if (slot) *slot = g->ws_orig[i]; - } - for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { - for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { - ds4_gpu_tensor_free(g->ws_mirror[tier][i]); - g->ws_mirror[tier][i] = NULL; - } + for (uint64_t i = 0; i < n; i++) { + const float v = x[i]; + if (v < minv) minv = v; + if (v > maxv) maxv = v; + ss += (double)v * v; } - g->ws_ready = 0; - g->ws_tier = -1; -} -static void glm_graph_ws_init(ds4_glm_gpu_graph *g) { - g->ws_ready = 0; - g->ws_tier = -1; - if (!g->placement) return; - bool used[DS4_MAX_GPUS] = { false }; - if (g->placement[0] >= 0 && g->placement[0] < DS4_MAX_GPUS) { - used[g->placement[0]] = true; - } - for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { - const int tier = g->placement[il + 1u]; - if (tier >= 0 && tier < DS4_MAX_GPUS) used[tier] = true; - } - for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); - if (!slot) return; - g->ws_orig[i] = *slot; - } - for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { - if (!used[tier]) continue; - for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { - const ds4_gpu_tensor *orig = g->ws_orig[i]; - if (!orig) continue; - g->ws_mirror[tier][i] = - ds4_gpu_tensor_alloc_ptr_on(tier, ds4_gpu_tensor_bytes(orig)); - if (!g->ws_mirror[tier][i]) { - fprintf(stderr, - "ds4: GLM per-tier working-set alloc failed (tier %d); " - "falling back to base buffers\n", - tier); - g->ws_ready = 1; - glm_graph_ws_free(g); - return; - } - } - } - g->ws_ready = 1; + printf("%s: min=%g max=%g rms=%g\n", + name, minv, maxv, sqrt(ss / (double)n)); } -static bool glm_graph_ws_switch(ds4_glm_gpu_graph *g, - int tier, - bool carry_hidden) { - if (!g->placement) return true; - if (tier < 0 || tier >= DS4_MAX_GPUS || - ds4_gpu_set_current_device_fenced(tier) != 0) { - return false; - } - if (!g->ws_ready) return true; - if (g->ws_tier == tier) return true; - ds4_gpu_tensor *old_cur = g->cur; - for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); - if (slot && g->ws_mirror[tier][i]) { - *slot = g->ws_mirror[tier][i]; - } - } - if (carry_hidden && old_cur && g->cur && g->cur != old_cur) { - const uint64_t bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - if (ds4_gpu_tensor_copy_async(g->cur, old_cur, bytes) == 0) { - return false; - } +#ifndef DS4_NO_GPU +#include "models/deepseek/graph.inc" +#endif + +typedef struct ds4_vocab ds4_vocab; + +static void embed_prompt( + const ds4_model * model, + const ds4_weights * weights, + const token_vec * tokens, + uint32_t n_embd, + float * out) { + for (int i = 0; i < tokens->len; i++) { + embed_token_any(model, weights, tokens->v[i], out + (uint64_t)i * n_embd); } - g->ws_tier = tier; - return true; } -#define DS4_GLM_VERIFY_WS_FIELDS(X) \ - X(batch_cur) X(batch_next) X(batch_attn_norm) X(batch_q_rank) \ - X(batch_q_rank_norm) X(batch_q) X(batch_indexer_k) X(batch_kv_raw) \ - X(batch_kv_norm) X(batch_qk_low) X(batch_attn_lora) \ - X(batch_indexer_selected) X(batch_heads) X(batch_attn_out) \ - X(batch_after_attn) X(batch_ffn_norm) X(batch_ffn_gate) X(batch_ffn_up) \ - X(batch_shared_mid) X(batch_ffn_mid) X(batch_routed_gate) \ - X(batch_routed_up) X(batch_routed_down) X(batch_ffn_out) \ - X(batch_router_logits) X(batch_router_probs) X(batch_router_selected) \ - X(batch_router_weights) +/* ========================================================================= + * Tokenizer and Chat Prompt Encoding. + * ========================================================================= + * + * DeepSeek V4 Flash stores a GPT-2 style byte-level BPE tokenizer in GGUF. + * The implementation below is intentionally small. It loads token strings + * and merge ranks from the mmaped file, builds two open-addressed hash tables, + * and applies BPE to user text. Chat special tokens are inserted directly by + * ID; user text goes through BPE. + */ + +typedef struct { + ds4_str key; + int value; + bool used; +} str_i32_entry; -static ds4_gpu_tensor **glm_graph_verify_ws_slot(ds4_glm_gpu_graph *g, - int i) { - int n = 0; -#define DS4_GLM_VERIFY_WS_SLOT_CASE(field) if (n++ == i) return &g->field; - DS4_GLM_VERIFY_WS_FIELDS(DS4_GLM_VERIFY_WS_SLOT_CASE) -#undef DS4_GLM_VERIFY_WS_SLOT_CASE - return NULL; -} +typedef struct { + str_i32_entry *entry; + uint64_t cap; + uint64_t used; +} str_i32_table; -static void glm_graph_verify_ws_restore(ds4_glm_gpu_graph *g) { - if (!g || !g->verify_ws_ready) return; - for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); - if (slot) *slot = g->verify_ws_orig[i]; - } - g->verify_ws_tier = -1; +static uint64_t next_pow2(uint64_t n) { + uint64_t p = 1; + while (p < n) p <<= 1; + return p; } -static void glm_graph_verify_ws_free(ds4_glm_gpu_graph *g) { - if (!g || !g->verify_ws_ready) return; - glm_graph_verify_ws_restore(g); - for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { - for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { - ds4_gpu_tensor_free(g->verify_ws_mirror[tier][i]); - g->verify_ws_mirror[tier][i] = NULL; - } - } - memset(g->verify_ws_orig, 0, sizeof(g->verify_ws_orig)); - g->verify_ws_ready = 0; +static void table_init(str_i32_table *t, uint64_t expected) { + t->cap = next_pow2(expected * 2 + 16); + t->used = 0; + t->entry = xcalloc((size_t)t->cap, sizeof(t->entry[0])); } -static bool glm_graph_verify_ws_init(ds4_glm_gpu_graph *g) { - if (!g) return false; - if (g->verify_ws_ready) return true; - if (!g->placement) { - g->verify_ws_ready = 1; - g->verify_ws_tier = -1; - return true; - } - bool used[DS4_MAX_GPUS] = { false }; - for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { - const int tier = g->placement[il + 1u]; - if (tier < 0 || tier >= DS4_MAX_GPUS) return false; - used[tier] = true; - } - for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); - if (!slot) return false; - g->verify_ws_orig[i] = *slot; - } - for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { - if (!used[tier]) continue; - for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { - const ds4_gpu_tensor *orig = g->verify_ws_orig[i]; - if (!orig) continue; - const uint32_t cap = (i >= 9 && i <= 11) ? - g->indexed_prefill_cap : g->ctx_cap; - const uint64_t orig_bytes = ds4_gpu_tensor_bytes(orig); - if (cap == 0 || orig_bytes % cap != 0 || - orig_bytes / cap > UINT64_MAX / 2u) { - g->verify_ws_ready = 1; - glm_graph_verify_ws_free(g); - return false; - } - const uint64_t bytes = (orig_bytes / cap) * 2u; - g->verify_ws_mirror[tier][i] = - ds4_gpu_tensor_alloc_ptr_on(tier, bytes); - if (!g->verify_ws_mirror[tier][i]) { - g->verify_ws_ready = 1; - glm_graph_verify_ws_free(g); - return false; - } - } - } - g->verify_ws_ready = 1; - g->verify_ws_tier = -1; - return true; +static void table_free(str_i32_table *t) { + free(t->entry); + memset(t, 0, sizeof(*t)); } -static bool glm_graph_verify_ws_switch(ds4_glm_gpu_graph *g, - int tier, - bool carry_hidden, - uint32_t n_rows) { - if (!g || n_rows == 0 || n_rows > 2) return false; - if (!g->placement) return true; - if (!glm_graph_verify_ws_init(g) || - tier < 0 || tier >= DS4_MAX_GPUS || - ds4_gpu_set_current_device_fenced(tier) != 0) { - return false; - } - if (g->verify_ws_tier == tier) return true; - ds4_gpu_tensor *old_cur = g->batch_cur; - for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { - ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); - if (slot && g->verify_ws_mirror[tier][i]) { - *slot = g->verify_ws_mirror[tier][i]; - } - } - if (carry_hidden && old_cur && g->batch_cur && old_cur != g->batch_cur) { - const uint64_t bytes = - (uint64_t)n_rows * DS4_N_EMBD * sizeof(float); - if (ds4_gpu_tensor_copy_async(g->batch_cur, old_cur, bytes) == 0) { - return false; +static void table_put(str_i32_table *t, ds4_str key, int value) { + uint64_t mask = t->cap - 1; + uint64_t i = hash_bytes(key.ptr, key.len) & mask; + + while (t->entry[i].used) { + if (ds4_str_eq(t->entry[i].key, key)) { + t->entry[i].value = value; + return; } + i = (i + 1) & mask; } - g->verify_ws_tier = tier; - return true; -} -#undef DS4_GLM_VERIFY_WS_FIELDS - -static void glm_graph_free(ds4_glm_gpu_graph *g) { - glm_graph_verify_ws_free(g); - glm_graph_ws_free(g); - if (!g) return; - ds4_gpu_tensor_free(g->mtp_kv_lora_cache); - ds4_gpu_tensor_free(g->mtp_k_rope_cache); - ds4_gpu_tensor_free(g->mtp_concat); - ds4_gpu_tensor_free(g->mtp_selected); - free(g->mtp_logits_host); - g->mtp_kv_lora_cache = NULL; - g->mtp_k_rope_cache = NULL; - g->mtp_concat = NULL; - g->mtp_selected = NULL; - g->mtp_logits_host = NULL; - g->mtp_ready = 0; - for (uint32_t il = 0; il < DS4_MAX_LAYER; il++) { - ds4_gpu_tensor_free(g->layer_indexer_key_cache[il]); - ds4_gpu_tensor_free(g->layer_k_rope_cache[il]); - ds4_gpu_tensor_free(g->layer_kv_lora_cache[il]); - ds4_gpu_tensor_free(g->layer_value_cache[il]); - ds4_gpu_tensor_free(g->layer_key_cache[il]); - } - ds4_gpu_tensor_free(g->logits); - ds4_gpu_tensor_free(g->batch_router_weights); - ds4_gpu_tensor_free(g->prefill_seed_router_selected); - ds4_gpu_tensor_free(g->batch_router_selected); - ds4_gpu_tensor_free(g->batch_router_probs); - ds4_gpu_tensor_free(g->batch_router_logits); - ds4_gpu_tensor_free(g->batch_routed_down); - ds4_gpu_tensor_free(g->batch_routed_up); - ds4_gpu_tensor_free(g->batch_routed_gate); - ds4_gpu_tensor_free(g->batch_ffn_out); - ds4_gpu_tensor_free(g->batch_ffn_mid); - ds4_gpu_tensor_free(g->batch_shared_mid); - ds4_gpu_tensor_free(g->batch_ffn_up); - ds4_gpu_tensor_free(g->batch_ffn_gate); - ds4_gpu_tensor_free(g->batch_ffn_norm); - ds4_gpu_tensor_free(g->batch_after_attn); - ds4_gpu_tensor_free(g->batch_attn_out); - ds4_gpu_tensor_free(g->batch_heads); - ds4_gpu_tensor_free(g->batch_value); - ds4_gpu_tensor_free(g->batch_k_nope); - ds4_gpu_tensor_free(g->batch_kv_norm); - ds4_gpu_tensor_free(g->batch_kv_raw); - ds4_gpu_tensor_free(g->batch_attn_lora); - ds4_gpu_tensor_free(g->batch_qk_low); - ds4_gpu_tensor_free(g->batch_indexer_selected); - ds4_gpu_tensor_free(g->batch_indexer_scores); - ds4_gpu_tensor_free(g->batch_indexer_weights); - ds4_gpu_tensor_free(g->batch_indexer_q); - ds4_gpu_tensor_free(g->batch_indexer_k); - ds4_gpu_tensor_free(g->batch_q); - ds4_gpu_tensor_free(g->batch_q_rank_norm); - ds4_gpu_tensor_free(g->batch_q_rank); - ds4_gpu_tensor_free(g->batch_attn_norm); - ds4_gpu_tensor_free(g->batch_next); - ds4_gpu_tensor_free(g->batch_cur); - ds4_gpu_tensor_free(g->prefill_tokens); - ds4_gpu_tensor_free(g->output_norm); - ds4_gpu_tensor_free(g->router_weights); - ds4_gpu_tensor_free(g->router_selected); - ds4_gpu_tensor_free(g->router_probs); - ds4_gpu_tensor_free(g->router_logits); - ds4_gpu_tensor_free(g->ffn_sum); - ds4_gpu_tensor_free(g->ffn_out); - ds4_gpu_tensor_free(g->routed_down); - ds4_gpu_tensor_free(g->routed_up); - ds4_gpu_tensor_free(g->tp_bounce_out); - ds4_gpu_tensor_free(g->tp_bounce_in); - ds4_gpu_tensor_free(g->routed_gate); - ds4_gpu_tensor_free(g->ffn_mid); - ds4_gpu_tensor_free(g->ffn_up); - ds4_gpu_tensor_free(g->ffn_gate); - ds4_gpu_tensor_free(g->ffn_norm); - ds4_gpu_tensor_free(g->after_attn); - ds4_gpu_tensor_free(g->attn_out); - ds4_gpu_tensor_free(g->heads); - ds4_gpu_tensor_free(g->value); - ds4_gpu_tensor_free(g->k_nope); - ds4_gpu_tensor_free(g->kv_norm); - ds4_gpu_tensor_free(g->kv_raw); - ds4_gpu_tensor_free(g->attn_partial_ms); - ds4_gpu_tensor_free(g->attn_partial_lora); - ds4_gpu_tensor_free(g->qk_low); - ds4_gpu_tensor_free(g->indexer_selected); - ds4_gpu_tensor_free(g->indexer_scores); - ds4_gpu_tensor_free(g->indexer_weights); - ds4_gpu_tensor_free(g->indexer_q); - ds4_gpu_tensor_free(g->indexer_k); - ds4_gpu_tensor_free(g->q); - ds4_gpu_tensor_free(g->q_rank_norm); - ds4_gpu_tensor_free(g->q_rank); - ds4_gpu_tensor_free(g->attn_norm); - ds4_gpu_tensor_free(g->next); - ds4_gpu_tensor_free(g->cur); - memset(g, 0, sizeof(*g)); -} - -static bool glm_graph_ensure_compact_cache( - const ds4_glm_gpu_graph *g, - uint32_t needed_rows) { - if (!g || needed_rows == 0) return false; - if (needed_rows <= g->ctx_cap && g->compact_cache_cap == 0) return true; - if (needed_rows <= g->compact_cache_cap) return true; - fprintf(stderr, - "ds4: GLM compact DSA cache capacity %u is smaller than required row %u " - "(ctx=%u, full_cap=%u)\n", - g->compact_cache_cap, - needed_rows, - g->ctx_size, - g->ctx_cap); - return false; + + t->entry[i].used = true; + t->entry[i].key = key; + t->entry[i].value = value; + t->used++; } -static bool glm_graph_warm_compact_indexer_store( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t warm_pos) { - if (!g || !model || !weights) return false; - if (g->compact_cache_cap == 0 || g->indexer_full_layers == 0) return true; - if (!g->indexer_k) return false; - if (warm_pos >= g->compact_cache_cap) warm_pos = g->compact_cache_cap - 1u; - - if (ds4_gpu_tensor_fill_f32(g->indexer_k, - 0.0f, - DS4_N_INDEXER_HEAD_DIM) == 0) { - return false; - } +static bool table_get(const str_i32_table *t, const char *ptr, uint64_t len, int *value) { + if (t->cap == 0) return false; - const bool profile = false; - const double t0 = 0.0; - bool ok = ds4_gpu_begin_commands() != 0; - uint32_t warmed = 0; - for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { - if (!glm_graph_layer_uses_full_indexer(il)) continue; - const ds4_layer_weights *l = &weights->layer[il]; - if (!g->layer_indexer_key_cache[il] || - !l->indexer_k_norm || - !l->indexer_k_norm_b) { - ok = false; - break; + uint64_t mask = t->cap - 1; + uint64_t i = hash_bytes(ptr, len) & mask; + + while (t->entry[i].used) { + ds4_str key = t->entry[i].key; + if (key.len == len && memcmp(key.ptr, ptr, len) == 0) { + *value = t->entry[i].value; + return true; } - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - ok = ds4_gpu_glm_store_indexer_k_tensor( - g->layer_indexer_key_cache[il], - g->indexer_k, - model->map, - model->size, - l->indexer_k_norm->abs_offset, - l->indexer_k_norm_b->abs_offset, - warm_pos, - 1, - g->compact_cache_cap, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - 0, - 1.0e-6f, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - glm_graph_compact_cache_is_f16()) != 0; - if (ok) warmed++; + i = (i + 1) & mask; } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); + return false; +} - if (profile) { - fprintf(stderr, - "ds4: GLM compact indexer warmup pos=%u layers=%u %.3f ms\n", - warm_pos, - warmed, - (now_sec() - t0) * 1000.0); +static void token_vec_push(token_vec *tv, int token) { + if (tv->len == tv->cap) { + tv->cap = tv->cap ? tv->cap * 2 : 64; + tv->v = xrealloc(tv->v, (size_t)tv->cap * sizeof(tv->v[0])); } - return ok; + tv->v[tv->len++] = token; } -static bool glm_graph_alloc_slice( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int ctx_size, - bool ssd_streaming, - bool ssd_streaming_cold, - uint64_t streaming_transient_guard_bytes, - uint32_t layer_start, - uint32_t layer_end, - bool require_token_embd, - bool require_output_head) { - if (!g || !model || !weights || ctx_size <= 0) return false; - const int *placement = g->placement; - memset(g, 0, sizeof(*g)); - g->placement = placement; - g->ssd_streaming = ssd_streaming; - g->ssd_streaming_cold = ssd_streaming_cold; - - if (!glm_graph_context_request(ctx_size, &g->ctx_size)) return false; - if (!glm_graph_memory_guard_slice_with_transient( - model, - weights, - g->ssd_streaming, - layer_start, - layer_end, - require_token_embd, - require_output_head, - g->ctx_size, - streaming_transient_guard_bytes, - "before GLM graph allocation")) { - return false; - } - if (!glm_graph_validate_layout(model, - weights, - g, - layer_start, - layer_end, - require_token_embd, - require_output_head)) { - return false; - } - g->ctx_cap = glm_graph_full_attention_cap(g->ctx_size, - g->ssd_streaming); - g->full_kv_cache = glm_graph_expanded_kv_cache_enabled(g->ssd_streaming); - g->compact_cache_cap = - glm_graph_compact_cache_initial_cap(g->ctx_size, g->ctx_cap); - g->indexed_prefill_cap = - g->compact_cache_cap != 0 ? - glm_graph_indexed_prefill_chunk_tokens(g->ctx_cap, g->compact_cache_cap) : - 0; - g->indexed_prefill_score_cap = - glm_graph_indexed_prefill_score_tokens(g->indexed_prefill_cap, - g->compact_cache_cap); - g->indexer_full_layers = - glm_graph_full_indexer_layer_count_range(g->layer_start, - g->layer_end); - if (g->ctx_size > g->ctx_cap) { - fprintf(stderr, - "ds4: GLM Metal session ctx=%u (model max=%u); " - "full-attention prefill/work cap=%u; compact indexed decode is used beyond the cap\n", - g->ctx_size, - glm_graph_model_context_limit(), - g->ctx_cap); - } +static void token_vec_free(token_vec *tv) { + free(tv->v); + memset(tv, 0, sizeof(*tv)); +} - const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - const uint64_t q_rank_bytes = (uint64_t)DS4_N_LORA_Q * sizeof(float); - const uint64_t q_bytes = g->q_dim * sizeof(float); - const uint64_t kv_raw_bytes = g->kv_raw_dim * sizeof(float); - const uint64_t kv_norm_bytes = (uint64_t)DS4_N_KV_LORA * sizeof(float); - const uint64_t k_nope_bytes = (uint64_t)DS4_N_HEAD * g->q_nope * sizeof(float); - const uint64_t heads_bytes = g->heads_dim * sizeof(float); - const uint64_t indexer_k_bytes = (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float); - const uint64_t indexer_q_bytes = - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float); - const uint64_t indexer_weights_bytes = - (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); - const uint64_t indexer_work_cap = - g->compact_cache_cap != 0 ? g->compact_cache_cap : g->ctx_cap; - const uint64_t indexer_scores_bytes = indexer_work_cap * sizeof(float); - const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); - const uint64_t indexer_selected_bytes = - (uint64_t)indexer_top_k * sizeof(uint32_t); - const uint64_t qk_low_bytes = - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float); - const uint32_t split_attn_blocks = glm_graph_indexed_decode_split_blocks(); - const uint64_t attn_partial_lora_bytes = - (uint64_t)split_attn_blocks * qk_low_bytes; - const uint64_t attn_partial_ms_bytes = - (uint64_t)split_attn_blocks * DS4_N_HEAD * 2u * sizeof(float); - const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); - const uint64_t full_kv_elem_bytes = glm_graph_full_kv_cache_elem_bytes(); - const uint64_t key_cache_bytes = g->full_kv_cache ? - (uint64_t)g->ctx_cap * g->q_dim * full_kv_elem_bytes : 0; - const uint64_t value_cache_bytes = g->full_kv_cache ? - (uint64_t)g->ctx_cap * g->heads_dim * full_kv_elem_bytes : 0; - const uint64_t compact_kv_lora_bytes = - (uint64_t)g->compact_cache_cap * DS4_N_KV_LORA * - glm_graph_compact_cache_elem_bytes(); - const uint64_t compact_k_rope_bytes = - (uint64_t)g->compact_cache_cap * DS4_N_ROT * - glm_graph_compact_cache_elem_bytes(); - const uint64_t compact_indexer_key_bytes = - (uint64_t)g->compact_cache_cap * DS4_N_INDEXER_HEAD_DIM * - glm_graph_compact_cache_elem_bytes(); - const uint64_t batch_rows = - g->full_kv_cache || g->indexed_prefill_cap == 0 ? - g->ctx_cap : - g->indexed_prefill_cap; - const uint64_t indexed_batch_rows = g->indexed_prefill_cap; - const uint64_t indexed_score_rows = g->indexed_prefill_score_cap; - const uint64_t batch_indexer_q_bytes = indexed_batch_rows * indexer_q_bytes; - const uint64_t batch_indexer_weights_bytes = indexed_batch_rows * indexer_weights_bytes; - const uint64_t batch_indexer_scores_bytes = - indexed_score_rows * indexer_work_cap * sizeof(float); - const uint64_t batch_indexer_selected_bytes = - indexed_batch_rows * indexer_top_k * sizeof(uint32_t); - const uint64_t batch_qk_low_bytes = indexed_batch_rows * qk_low_bytes; - const uint64_t batch_attn_lora_bytes = - indexed_batch_rows * (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float); - const uint64_t routed_mid_bytes = - (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float); - const uint64_t routed_down_bytes = - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); - const double cache_gib = - (double)(g->layer_count * (key_cache_bytes + value_cache_bytes)) / - (1024.0 * 1024.0 * 1024.0); - if (g->full_kv_cache) { - fprintf(stderr, - "ds4: GLM graph allocating full-attention KV cache: work_ctx=%u layers=%u:%u (%u) %s %.2f GiB\n", - g->ctx_cap, - g->layer_start, - g->layer_end, - g->layer_count, - "f16", - cache_gib); - } else { - fprintf(stderr, - "ds4: GLM graph using compact DSA KV only; expanded full-attention KV cache is skipped\n"); - } -#ifdef DS4_ROCM_BUILD - if (glm_graph_env_truthy( - getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE"))) { - fprintf(stderr, - "ds4: ROCm GLM one-token layer slices use the optimized token graph\n"); - } -#endif - if (g->compact_cache_cap != 0) { - const uint64_t compact_kv_total = - (uint64_t)g->layer_count * (compact_kv_lora_bytes + compact_k_rope_bytes); - const uint64_t compact_indexer_total = - (uint64_t)g->indexer_full_layers * compact_indexer_key_bytes; - const double compact_gib = - (double)(compact_kv_total + compact_indexer_total) / - (1024.0 * 1024.0 * 1024.0); - fprintf(stderr, - "ds4: GLM graph allocating compact DSA cache: rows=%u logical_ctx=%u kv_layers=%u indexer_layers=%u %s %.2f GiB\n", - g->compact_cache_cap, - g->ctx_size, - g->layer_count, - g->indexer_full_layers, - glm_graph_compact_cache_is_f16() ? "f16" : "f32", - compact_gib); - fprintf(stderr, - "ds4: GLM compact indexed prefill chunk=%u score_rows=%u score_scratch=%.2f MiB\n", - g->indexed_prefill_cap, - g->indexed_prefill_score_cap, - (double)batch_indexer_scores_bytes / (1024.0 * 1024.0)); - } +void ds4_tokens_push(ds4_tokens *tv, int token) { + token_vec_push(tv, token); +} - bool ok = true; -#define DS4_GLM_GRAPH_ALLOC_TENSOR(var, bytes_) \ - do { \ - (var) = ds4_gpu_tensor_alloc((bytes_)); \ - if (!(var)) { \ - fprintf(stderr, "ds4: GLM Metal graph could not allocate %s\n", #var); \ - ok = false; \ - } \ - } while (0) +void ds4_tokens_free(ds4_tokens *tv) { + token_vec_free(tv); +} - DS4_GLM_GRAPH_ALLOC_TENSOR(g->cur, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->next, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_norm, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->q_rank, q_rank_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->q_rank_norm, q_rank_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->q, q_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_k, indexer_k_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_q, indexer_q_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_weights, indexer_weights_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_scores, indexer_scores_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_selected, indexer_selected_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->qk_low, qk_low_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_lora, attn_partial_lora_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_ms, attn_partial_ms_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_raw, kv_raw_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_norm, kv_norm_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->k_nope, k_nope_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->value, heads_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->heads, heads_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_out, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->after_attn, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_norm, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_gate, g->dense_hidden_max * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_up, g->dense_hidden_max * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_mid, g->ffn_mid_elems * sizeof(float)); - if (g->generic_routed_moe) { - DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_gate, routed_mid_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_up, routed_mid_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_down, routed_down_bytes); - } - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_out, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_sum, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_logits, (uint64_t)DS4_N_EXPERT * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_probs, (uint64_t)DS4_N_EXPERT * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->output_norm, emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->logits, logits_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_logits, batch_rows * DS4_N_EXPERT * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_probs, batch_rows * DS4_N_EXPERT * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_selected, batch_rows * DS4_N_EXPERT_USED * sizeof(int32_t)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_weights, batch_rows * DS4_N_EXPERT_USED * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->prefill_seed_router_selected, - (uint64_t)DS4_N_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * - sizeof(int32_t)); - - DS4_GLM_GRAPH_ALLOC_TENSOR(g->prefill_tokens, batch_rows * sizeof(int32_t)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_cur, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_next, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_norm, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q_rank, batch_rows * q_rank_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q_rank_norm, batch_rows * q_rank_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q, batch_rows * q_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_k, batch_rows * indexer_k_bytes); - if (g->compact_cache_cap != 0 && g->indexed_prefill_cap != 0) { - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_q, batch_indexer_q_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_weights, batch_indexer_weights_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_scores, batch_indexer_scores_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_selected, batch_indexer_selected_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_qk_low, batch_qk_low_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_lora, batch_attn_lora_bytes); - } - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_kv_raw, batch_rows * kv_raw_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_kv_norm, batch_rows * kv_norm_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_k_nope, batch_rows * k_nope_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_value, batch_rows * heads_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_heads, batch_rows * heads_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_out, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_after_attn, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_norm, batch_rows * emb_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_gate, batch_rows * g->dense_hidden_max * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_up, batch_rows * g->dense_hidden_max * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_shared_mid, batch_rows * DS4_N_FF_EXP * sizeof(float)); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_mid, batch_rows * g->ffn_mid_elems * sizeof(float)); - if (g->generic_routed_moe) { - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_gate, batch_rows * routed_mid_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_up, batch_rows * routed_mid_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_down, batch_rows * routed_down_bytes); - } - DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_out, batch_rows * emb_bytes); - - for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { - int cache_tier = 0; - if (g->placement && g->placement[il + 1u] >= 0 && - g->placement[il + 1u] < DS4_MAX_GPUS) { - cache_tier = g->placement[il + 1u]; - } -#define DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(var, bytes_) \ - do { \ - (var) = ds4_gpu_tensor_alloc_ptr_on(cache_tier, (bytes_)); \ - if (!(var)) { \ - fprintf(stderr, "ds4: GLM graph could not allocate %s on tier %d\n", \ - #var, cache_tier); \ - ok = false; \ - } \ - } while (0) - if (g->full_kv_cache) { - DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_key_cache[il], key_cache_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_value_cache[il], value_cache_bytes); - } - if (g->compact_cache_cap != 0) { - DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_kv_lora_cache[il], compact_kv_lora_bytes); - DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_k_rope_cache[il], compact_k_rope_bytes); - if (glm_graph_layer_uses_full_indexer(il)) { - DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_indexer_key_cache[il], - compact_indexer_key_bytes); - } - } -#undef DS4_GLM_GRAPH_ALLOC_TENSOR_TIER - } -#undef DS4_GLM_GRAPH_ALLOC_TENSOR +void ds4_tokens_copy(ds4_tokens *dst, const ds4_tokens *src) { + dst->len = 0; + for (int i = 0; i < src->len; i++) token_vec_push(dst, src->v[i]); +} - if (!ok) { - glm_graph_free(g); - return false; +bool ds4_tokens_starts_with(const ds4_tokens *tokens, const ds4_tokens *prefix) { + if (prefix->len > tokens->len) return false; + for (int i = 0; i < prefix->len; i++) { + if (tokens->v[i] != prefix->v[i]) return false; } - glm_graph_ws_init(g); return true; } -static bool glm_graph_alloc( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int ctx_size, - bool ssd_streaming, - bool ssd_streaming_cold) { - const uint32_t normal_layers = glm_graph_normal_layer_count(); - if (normal_layers == 0) { - fprintf(stderr, "ds4: GLM Metal graph found no normal transformer layers\n"); - return false; - } - return glm_graph_alloc_slice(g, - model, - weights, - ctx_size, - ssd_streaming, - ssd_streaming_cold, - 0, - 0, - normal_layers - 1u, - true, - true); -} +struct ds4_vocab { + ds4_str *token; + int n_vocab; + int bos_id; + int eos_id; + int system_id; + int user_id; + int assistant_id; + int observation_id; + int sop_id; + int think_start_id; + int think_end_id; + int tool_call_start_id; + int tool_call_end_id; + int tool_response_start_id; + int tool_response_end_id; + int arg_key_start_id; + int arg_key_end_id; + int arg_value_start_id; + int arg_value_end_id; + int dsml_id; + str_i32_table token_to_id; + str_i32_table merge_rank; +}; + +/* Engine-side tensor-parallel state. The transport context is owned by the + * frontend (CLI leader or ds4_tp_worker_run); the engine owns the GPU slab, + * the per-slot views and the gate machinery lifetime. */ +typedef struct { + struct ds4_tp *ctx; + ds4_gpu_tensor *slab; + ds4_gpu_tensor **out_views; + ds4_gpu_tensor **in_views; + ds4_gpu_tensor **batch_out_views; /* [layer] verify-block row partials */ + ds4_gpu_tensor **batch_in_views; + ds4_gpu_tensor *zero_vec; + uint64_t eval_seq; /* leader: mirrored eval counter */ + uint64_t next_session_id; /* leader: stable worker-session handle */ + int rank; + bool vocab_split; /* DS4-only: logits halves cross the wire */ + bool active; +} ds4_engine_tp_state; + +struct ds4_engine { + const ds4_model_provider_v1 *provider; + ds4_model model; + ds4_model mtp_model; + ds4_vocab vocab; + ds4_weights weights; + ds4_mtp_weights mtp_weights; + ds4_dspark_weights dspark_weights; + ds4_backend backend; + ds4_support_kind support_kind; + int dspark_exec_tier; + uint32_t support_stages; + int mtp_draft_tokens; + float mtp_margin; + float dspark_confidence_threshold; + char *directional_steering_file; + float *directional_steering_dirs; + float directional_steering_attn_scale; + float directional_steering_ffn_scale; + int power_percent; + uint32_t prefill_chunk; + uint32_t ssd_streaming_cache_experts; + uint64_t ssd_streaming_cache_bytes; + uint64_t ssd_streaming_prefill_headroom_bytes; + uint64_t ssd_streaming_full_layer_bytes; + uint32_t ssd_streaming_full_layers; + uint32_t ssd_streaming_preload_experts; + uint64_t startup_model_span_bytes; + ds4_ssd_memory_lock simulated_memory; + bool quality; + bool glm_mtp; + bool glm_mtp_timing; + bool dspark; + bool dspark_strict; + bool cuda_tensor_parallel; + bool glm_tp_token_prefill; + bool ssd_streaming; + bool ssd_streaming_cold; + bool ssd_streaming_full_layers_set; + ds4_distributed_options distributed; + ds4_engine_tp_state tp; + bool metal_ready; + bool mtp_ready; + bool share_session_prefill_workspace; +#ifndef DS4_NO_GPU + bool shared_prefill_workspace_ready; + ds4_gpu_graph shared_prefill_workspace; +#endif -static uint32_t glm_graph_weight_type_for_offset( - const ds4_model *model, - uint64_t weight_offset); - -static int glm_graph_matmul_q8_0_decode_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - bool ssd_streaming) { - if (!model) return 0; - const uint32_t weight_type = glm_graph_weight_type_for_offset(model, weight_offset); - if (ssd_streaming) { - return ds4_gpu_matmul_quant_tensor(out, - model->map, - model->size, - weight_offset, - weight_type, - in_dim, - out_dim, - x, - 1); - } - return ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(out, - model->map, - model->size, - weight_offset, - weight_type, - in_dim, - out_dim, - x, - 1); -} + /* Wave-2 multi-GPU placement scaffolding: optional multi-GPU placement + * state. Zero-initialized for every existing caller (gpu_cfg == NULL) + * via xcalloc, so the single-tier path observes identical engine + * state to pre-multi-GPU CLI main. multi_tier == 1 is the gate for all + * new code paths. */ + ds4_gpu_config gpu_cfg; + int placement[DS4_MAX_LAYER + 2]; + int n_placement_entries; + int multi_tier; -static bool glm_graph_q8_decode_profile_enabled(uint32_t il, const char *label) { - (void)il; - (void)label; - return false; -} + /* Max-context hint copied from + * ds4_engine_options.placement_ctx_hint. Used by + * engine_compute_entry_bytes for per-layer KV estimation. + * Zero / negative = legacy 4096 fallback (single-tier paths and any + * caller that doesn't set the option observe the prior behavior). */ + int placement_ctx_hint; +}; -static uint32_t glm_graph_weight_type_for_offset( - const ds4_model *model, - uint64_t weight_offset) { - if (!model || !model->tensors) return DS4_TENSOR_Q8_0; - for (uint64_t i = 0; i < model->n_tensors; i++) { - const ds4_tensor *t = &model->tensors[i]; - if (t->abs_offset == weight_offset) return t->type; +static uint64_t ds4_engine_dynamic_expert_cache_bytes( + const ds4_engine *e) { + if (!e || !e->ssd_streaming) return 0; + if (e->ssd_streaming_cache_bytes != 0) { + return e->ssd_streaming_cache_bytes; } - return DS4_TENSOR_Q8_0; -} + if (e->ssd_streaming_cache_experts == 0) return 0; -static bool glm_graph_weights_are_q8_0( - const ds4_model *model, - uint64_t offset_a, - uint64_t offset_b) { - return glm_graph_weight_type_for_offset(model, offset_a) == DS4_TENSOR_Q8_0 && - glm_graph_weight_type_for_offset(model, offset_b) == DS4_TENSOR_Q8_0; -} - -static int glm_graph_matmul_q8_0_decode_profiled_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint32_t il, - uint32_t pos, - const char *label, - bool ssd_streaming) { - const bool profile = glm_graph_q8_decode_profile_enabled(il, label); - if (profile) { - if (ds4_gpu_end_commands() == 0) return 0; - if (ds4_gpu_begin_commands() == 0) return 0; - } - const double t0 = profile ? now_sec() : 0.0; - int ok = glm_graph_matmul_q8_0_decode_tensor(out, - model, - weight_offset, - in_dim, - out_dim, - x, - ssd_streaming); - if (profile) { - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM Q8 decode profile layer=%u pos=%u label=%s in=%llu out=%llu %.3f ms\n", - il, - pos, - label ? label : "?", - (unsigned long long)in_dim, - (unsigned long long)out_dim, - (now - t0) * 1000.0); - if (ok) ok = ds4_gpu_begin_commands() != 0; + uint64_t per_expert_bytes = 0; + if (!ds4_streaming_routed_expert_bytes(&e->weights, + &per_expert_bytes)) { + return 0; } - return ok; -} - -static bool glm_graph_encode_output_head_from( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const ds4_gpu_tensor *hidden) { - bool ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, - hidden, - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->logits, - model, - weights->output->abs_offset, - DS4_N_EMBD, - DS4_N_VOCAB, - g->output_norm, - g->ssd_streaming) != 0; - return ok; -} - -static bool glm_graph_encode_output_head( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - return glm_graph_encode_output_head_from(g, model, weights, g->cur); -} - -static bool glm_graph_forward_output_head( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const ds4_gpu_tensor *hidden, - float *logits_out) { - if (!g || !model || !weights || !hidden || !logits_out) return false; - bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = glm_graph_encode_output_head_from(g, model, weights, hidden); - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok && glm_debug_hidden_dump_layer() < 0) - glm_debug_dump_hidden_row(hidden, 0); - if (ok) { - ok = ds4_gpu_tensor_read(g->logits, - 0, - logits_out, - (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + if (e->ssd_streaming_cache_experts > UINT64_MAX / per_expert_bytes) { + return UINT64_MAX; } - return ok; + return (uint64_t)e->ssd_streaming_cache_experts * per_expert_bytes; } -static bool glm_graph_profile_stage( - bool enabled, - const char *part, - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - double *stage_t0) { - if (!enabled) return true; - if (!stage_t0) return false; - return metal_graph_layer_stage_profile_boundary(part, stage, il, pos0, n_tokens, stage_t0); +static uint64_t ds4_engine_streaming_transient_guard_bytes( + const ds4_engine *e) { + if (!e || !e->ssd_streaming) return 0; + uint64_t total = ds4_engine_dynamic_expert_cache_bytes(e); + total = ds4_add_sat_u64(total, e->ssd_streaming_full_layer_bytes); + total = ds4_add_sat_u64(total, e->ssd_streaming_prefill_headroom_bytes); + return total; } -static bool glm_graph_profile_router_selection( - ds4_glm_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos) { - if (!g_expert_profile.active) return true; - if (!g || !layer || !g->router_selected || !g->router_weights) return false; - - if (ds4_gpu_end_commands() == 0) { - fprintf(stderr, - "ds4: failed to end GLM Metal command batch for expert profile readback\n"); - return false; - } - - int32_t selected[DS4_MAX_EXPERT_USED] = {0}; - float weights[DS4_MAX_EXPERT_USED] = {0}; - const bool read_ok = - ds4_gpu_tensor_read(g->router_selected, - 0, - selected, - (uint64_t)DS4_N_EXPERT_USED * sizeof(selected[0])) != 0 && - ds4_gpu_tensor_read(g->router_weights, - 0, - weights, - (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) != 0; - - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, - "ds4: failed to resume GLM Metal command batch after expert profile readback\n"); - return false; - } - if (!read_ok) { - fprintf(stderr, "ds4: failed to read GLM Metal router tensors for expert profile\n"); - return false; - } - - ds4_expert_profile_record(il, - pos, - selected, - weights, - layer->ffn_gate_tid2eid != NULL); - return true; -} +static void ds4_engine_print_startup_memory( + const ds4_engine *e, + int ctx_size) { + if (!e || ctx_size <= 0) return; -static bool glm_graph_profile_router_selection_batch( - ds4_glm_gpu_graph *g, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens) { - if (!g_expert_profile.active) return true; - if (!g || !layer || !g->batch_router_selected || - !g->batch_router_weights || n_tokens == 0) { - return false; + ds4_context_memory mem; +#ifndef DS4_NO_GPU + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + const uint32_t layer_end = e->distributed.layers.has_output ? + (normal_layers ? normal_layers - 1u : 0u) : + e->distributed.layers.end; + const uint32_t ctx = (uint32_t)ctx_size; + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx, e->ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx, work_ctx); + mem = glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx, + work_ctx, + compact_cap, + e->ssd_streaming, + e->distributed.layers.start, + layer_end); + } else { +#endif + mem = ds4_context_memory_estimate_with_prefill_mode(e->backend, + ctx_size, + e->prefill_chunk, + e->ssd_streaming); +#ifndef DS4_NO_GPU } +#endif + const uint64_t kv_bytes = + ds4_add_sat_u64(mem.raw_bytes, mem.compressed_bytes); + const uint64_t dynamic_expert_cache_bytes = + ds4_engine_dynamic_expert_cache_bytes(e); + const uint64_t expert_reserved_bytes = + e->ssd_streaming_prefill_headroom_bytes; + uint64_t total = kv_bytes; + total = ds4_add_sat_u64(total, mem.scratch_bytes); + total = ds4_add_sat_u64(total, e->startup_model_span_bytes); + total = ds4_add_sat_u64(total, dynamic_expert_cache_bytes); + total = ds4_add_sat_u64(total, e->ssd_streaming_full_layer_bytes); + total = ds4_add_sat_u64(total, expert_reserved_bytes); - const size_t selected_count = (size_t)n_tokens * DS4_N_EXPERT_USED; - if (n_tokens != 0 && selected_count / n_tokens != DS4_N_EXPERT_USED) { - return false; - } - if (selected_count > SIZE_MAX / sizeof(int32_t) || - selected_count > SIZE_MAX / sizeof(float)) { - return false; - } + const bool color = ds4_log_is_tty(stderr); + const char *green = color ? "\x1b[32m" : ""; + const char *bright_green = color ? "\x1b[1;32m" : ""; + const char *reset = color ? "\x1b[0m" : ""; - if (ds4_gpu_end_commands() == 0) { + fprintf(stderr, + "%sds4: memory: KV %.2f GiB (raw %.2f + compressed %.2f) " + "+ buffers %.2f GiB + resident model %.2f GiB", + green, + ds4_bytes_to_gib(kv_bytes), + ds4_bytes_to_gib(mem.raw_bytes), + ds4_bytes_to_gib(mem.compressed_bytes), + ds4_bytes_to_gib(mem.scratch_bytes), + ds4_bytes_to_gib(e->startup_model_span_bytes)); + if (e->ssd_streaming_full_layer_bytes != 0) { fprintf(stderr, - "ds4: failed to end GLM Metal command batch for batch expert profile readback\n"); - return false; + " + full-layer experts %.2f GiB", + ds4_bytes_to_gib(e->ssd_streaming_full_layer_bytes)); } - - int32_t *selected = xmalloc(selected_count * sizeof(selected[0])); - float *weights = xmalloc(selected_count * sizeof(weights[0])); - const bool read_ok = - ds4_gpu_tensor_read(g->batch_router_selected, - 0, - selected, - (uint64_t)selected_count * sizeof(selected[0])) != 0 && - ds4_gpu_tensor_read(g->batch_router_weights, - 0, - weights, - (uint64_t)selected_count * sizeof(weights[0])) != 0; - - if (ds4_gpu_begin_commands() == 0) { - free(weights); - free(selected); + if (dynamic_expert_cache_bytes != 0) { fprintf(stderr, - "ds4: failed to resume GLM Metal command batch after batch expert profile readback\n"); - return false; - } - if (!read_ok) { - free(weights); - free(selected); - fprintf(stderr, "ds4: failed to read GLM Metal batch router tensors for expert profile\n"); - return false; - } - - for (uint32_t t = 0; t < n_tokens; t++) { - const size_t off = (size_t)t * DS4_N_EXPERT_USED; - ds4_expert_profile_record(il, - pos0 + t, - selected + off, - weights + off, - layer->ffn_gate_tid2eid != NULL); - } - - free(weights); - free(selected); - return true; -} - -static bool glm_graph_prefill_stage_boundary( - bool stage_profile, - bool stage_sync, - const char *part, - const char *stage, - uint32_t il, - uint32_t pos0, - uint32_t n_tokens, - double *stage_t0) { - if (stage_profile) { - return glm_graph_profile_stage(true, part, stage, il, pos0, n_tokens, stage_t0); + " + expert cache %.2f GiB", + ds4_bytes_to_gib(dynamic_expert_cache_bytes)); } - if (stage_sync) return glm_graph_prefill_stage_sync_boundary(); - return true; -} - -static int glm_graph_routed_moe_one_dispatch( - const ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *x, - bool force_resident) { - if (!g || !model || !l) return 0; - /* Under the TP expert split only the ownership-aware kernels may run: - * the generic mul_mv_id family and the GLM q2_K resident pair/down. - * Anything else would silently compute the full expert set. */ - if (g->tp_world == 2 && - !glm_graph_layer_uses_generic_routed_moe(l) && - l->ffn_gate_exps->type != DS4_TENSOR_Q2_K) { + if (expert_reserved_bytes != 0) { fprintf(stderr, - "ds4: GLM TP split lacks ownership-aware kernels for expert type %u (layer %u)\n", - l->ffn_gate_exps->type, il); - return 0; - } - if (glm_graph_layer_uses_generic_routed_moe(l)) { - if (!g->routed_gate || !g->routed_up || !g->routed_down || - l->ffn_gate_exps->type != l->ffn_up_exps->type) { - if (getenv("DS4_GLM_TP_DEBUG")) { - fprintf(stderr, - "ds4: glm dispatch guard: gate=%p up=%p down=%p types=%u/%u\n", - (void *)g->routed_gate, (void *)g->routed_up, - (void *)g->routed_down, - l->ffn_gate_exps->type, l->ffn_up_exps->type); - } - return 0; - } - return ds4_gpu_routed_moe_one_tensor(out, - g->routed_gate, - g->routed_up, - mid, - g->routed_down, - model->map, - model->size, - l->ffn_gate_exps->abs_offset, - l->ffn_up_exps->abs_offset, - l->ffn_down_exps->abs_offset, - l->ffn_gate_exps->type, - l->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - DS4_N_EMBD, - DS4_N_FF_EXP, - DS4_N_EMBD, - selected, - weights, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - 0.0f, - x, - NULL, - il, - force_resident); - } - - return ds4_gpu_glm_routed_moe_one_tensor(out, - mid, - model->map, - model->size, - l->ffn_gate_exps->abs_offset, - l->ffn_up_exps->abs_offset, - l->ffn_down_exps->abs_offset, - l->ffn_gate_exps->type, - l->ffn_up_exps->type, - l->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - DS4_N_EMBD, - DS4_N_FF_EXP, - DS4_N_EMBD, - selected, - weights, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - il, - x, - force_resident); -} - -/* Post-compute visibility for GLM TP debugging: the combine stashes the - * selected-ids contents pointer; by exchange time the router kernels have - * completed, so the service thread sees final ids. */ -static const int32_t *g_glm_tp_debug_ids DS4_MAYBE_UNUSED; - -/* After the TP ownership-split batch routed MoE, exchange the - * per-token routed partial rows with the peer through shared bounce - * buffers (one gate per sparse layer per chunk) and rebuild the full - * routed output with a commutative add. */ -/* The routed batch dispatch writes its local partial DIRECTLY into the - * shared bounce buffer (graph scratch may be private/untracked on M5, so - * a blit from it is not reliable); ensure capacity before dispatching. */ -static bool glm_graph_tp_batch_bounce_ready(ds4_glm_gpu_graph *g, - uint32_t n_tokens) { - const uint64_t bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); - if (!g->tp_bounce_out || ds4_gpu_tensor_bytes(g->tp_bounce_out) < bytes) { - ds4_gpu_tensor_free(g->tp_bounce_out); - ds4_gpu_tensor_free(g->tp_bounce_in); - g->tp_bounce_out = ds4_gpu_tensor_alloc(bytes); - g->tp_bounce_in = ds4_gpu_tensor_alloc(bytes); - } - return g->tp_bounce_out && g->tp_bounce_in; -} - -static bool glm_graph_tp_batch_ffn_combine( - ds4_glm_gpu_graph *g, - uint32_t il, - ds4_gpu_tensor *ffn_out, - uint32_t n_tokens) { - const uint64_t bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); - if (!g->tp_bounce_out || !g->tp_bounce_in) return false; - if (getenv("DS4_GLM_ABLATE_COMBINE")) { - /* Timing probe: local half only, no exchange (garbage output; - * both ranks must set the env or the gates desync). */ - return ds4_gpu_add_tensor(ffn_out, - g->tp_bounce_out, - g->tp_bounce_out, - (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; - } - if (!ds4_gpu_tp_big_gate_encode(il, n_tokens, - g->tp_bounce_out, g->tp_bounce_in, - bytes)) { - return false; - } - return ds4_gpu_add_tensor(ffn_out, - g->tp_bounce_out, - g->tp_bounce_in, - (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; -} - -static int glm_graph_routed_moe_batch_dispatch( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride, - bool force_resident, - bool direct_scalar_q4) { - if (!g || !model || !l) return 0; - g->batch_routed_mid_is_f16 = false; - - if (glm_graph_layer_uses_generic_routed_moe(l)) { - if (!g->batch_routed_gate || !g->batch_routed_up || !g->batch_routed_down || - l->ffn_gate_exps->type != l->ffn_up_exps->type) { - return 0; - } - return ds4_gpu_routed_moe_batch_tensor(out, - g->batch_routed_gate, - g->batch_routed_up, - mid, - g->batch_routed_down, - model->map, - model->size, - l->ffn_gate_exps->abs_offset, - l->ffn_up_exps->abs_offset, - l->ffn_down_exps->abs_offset, - l->ffn_gate_exps->type, - l->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - DS4_N_EMBD, - DS4_N_FF_EXP, - DS4_N_EMBD, - selected, - weights, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - 0.0f, - x, - il, - n_tokens, - &g->batch_routed_mid_is_f16, - force_resident); - } - - if (direct_scalar_q4) { - return ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( - out, - mid, - model->map, - model->size, - l->ffn_gate_exps->abs_offset, - l->ffn_up_exps->abs_offset, - l->ffn_down_exps->abs_offset, - l->ffn_gate_exps->type, - l->ffn_up_exps->type, - l->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - DS4_N_EMBD, - DS4_N_FF_EXP, - DS4_N_EMBD, - selected, - weights, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - il, - x, - n_tokens, - mid_token_stride); + " + prefill expert reserve %.2f GiB", + ds4_bytes_to_gib(expert_reserved_bytes)); } + fprintf(stderr, + " = %s%.2f GiB planned%s\n", + bright_green, + ds4_bytes_to_gib(total), + reset); - return ds4_gpu_glm_routed_moe_batch_tensor( - out, - mid, - model->map, - model->size, - l->ffn_gate_exps->abs_offset, - l->ffn_up_exps->abs_offset, - l->ffn_down_exps->abs_offset, - l->ffn_gate_exps->type, - l->ffn_up_exps->type, - l->ffn_down_exps->type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - DS4_N_EMBD, - DS4_N_FF_EXP, - DS4_N_EMBD, - selected, - weights, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - il, - x, - n_tokens, - mid_token_stride, - force_resident); -} - -static bool glm_graph_disable_add3_residual(void); - -static bool glm_graph_use_streaming_selected_async_load( - const ds4_glm_gpu_graph *g) { - if (!g || !g->ssd_streaming) return false; - if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_SELECTED_ASYNC_LOAD", - "DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_LOAD") || - glm_graph_env_present("DS4_ROCM_DISABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD", - "DS4_METAL_DISABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD")) { - return false; - } -#ifdef DS4_ROCM_BUILD - return true; -#else - return getenv("DS4_METAL_ENABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD") != NULL; -#endif + fprintf(stderr, + "%sds4: memory detail: ctx=%d prefill_cap=%u raw_kv_rows=%u " + "compressed_kv_rows=%u backend=%s%s\n", + green, + ctx_size, + mem.prefill_cap, + mem.raw_cap, + mem.comp_cap, + ds4_backend_name(e->backend), + reset); } -typedef struct glm_graph_streaming_async_profile { - uint64_t async_calls; - uint64_t sync_calls; - double async_total_ms; - double sync_total_ms; - double async_signal_start_ms; - double async_flush_router_ms; - double async_shared_ms; - double async_flush_shared_ms; - double async_finish_ms; - double async_routed_ms; - double async_post_ms; - double sync_early_load_ms; - double sync_shared_ms; - double sync_routed_ms; - double sync_post_ms; -} glm_graph_streaming_async_profile; - -static glm_graph_streaming_async_profile g_glm_streaming_async_profile; -static bool g_glm_streaming_async_profile_registered; - -static bool glm_graph_streaming_async_profile_enabled(void) { - return glm_graph_env_present("DS4_ROCM_GLM_STREAMING_ASYNC_PROFILE", - "DS4_METAL_GLM_STREAMING_ASYNC_PROFILE"); -} - -static void glm_graph_streaming_async_profile_print(void) { - const glm_graph_streaming_async_profile *p = - &g_glm_streaming_async_profile; - if (p->async_calls == 0 && p->sync_calls == 0) return; - const double async_calls = p->async_calls ? (double)p->async_calls : 1.0; - const double sync_calls = p->sync_calls ? (double)p->sync_calls : 1.0; - fprintf(stderr, - "ds4: GLM streaming async profile async_calls=%llu " - "total=%.3f ms avg=%.3f ms signal_start=%.3f ms " - "flush_router=%.3f ms shared=%.3f ms flush_shared=%.3f ms " - "finish=%.3f ms routed=%.3f ms post=%.3f ms\n", - (unsigned long long)p->async_calls, - p->async_total_ms, - p->async_total_ms / async_calls, - p->async_signal_start_ms, - p->async_flush_router_ms, - p->async_shared_ms, - p->async_flush_shared_ms, - p->async_finish_ms, - p->async_routed_ms, - p->async_post_ms); - fprintf(stderr, - "ds4: GLM streaming sync profile calls=%llu " - "total=%.3f ms avg=%.3f ms early_load=%.3f ms " - "shared=%.3f ms routed=%.3f ms post=%.3f ms\n", - (unsigned long long)p->sync_calls, - p->sync_total_ms, - p->sync_total_ms / sync_calls, - p->sync_early_load_ms, - p->sync_shared_ms, - p->sync_routed_ms, - p->sync_post_ms); -} - -static void glm_graph_streaming_async_profile_register(void) { - if (g_glm_streaming_async_profile_registered) return; - if (!glm_graph_streaming_async_profile_enabled()) return; - atexit(glm_graph_streaming_async_profile_print); - g_glm_streaming_async_profile_registered = true; -} - -static double glm_graph_streaming_async_profile_ms(void) { - return now_sec() * 1000.0; -} - -/* Timing-only skip-ablation for the GLM decode layer (comma list in - * DS4_GLM_DECODE_ABLATE): the skipped stage's output buffer keeps stale - * contents, so the run produces garbage text but every remaining dispatch - * (and every TP gate) still executes. Whole-token time deltas against a - * baseline run are the only reliable per-stage cost measurement — the - * stage profiler's per-stage command-buffer splits inflate small stages. */ -#define DS4_GLM_ABLATE_ATTN_OUT (1u << 0) -#define DS4_GLM_ABLATE_ATTN_CORE (1u << 1) -#define DS4_GLM_ABLATE_QPATH (1u << 2) -#define DS4_GLM_ABLATE_INDEXER (1u << 3) -#define DS4_GLM_ABLATE_ROUTED (1u << 4) -#define DS4_GLM_ABLATE_SHARED (1u << 5) -#define DS4_GLM_ABLATE_QKLOW (1u << 6) - -static uint32_t glm_decode_ablate_mask(void) { - static int cached = -1; - if (cached < 0) { - uint32_t mask = 0; - const char *env = getenv("DS4_GLM_DECODE_ABLATE"); - if (env) { - if (strstr(env, "attn_out")) mask |= DS4_GLM_ABLATE_ATTN_OUT; - if (strstr(env, "attn_core")) mask |= DS4_GLM_ABLATE_ATTN_CORE; - if (strstr(env, "qpath")) mask |= DS4_GLM_ABLATE_QPATH; - if (strstr(env, "indexer")) mask |= DS4_GLM_ABLATE_INDEXER; - if (strstr(env, "routed")) mask |= DS4_GLM_ABLATE_ROUTED; - if (strstr(env, "shared")) mask |= DS4_GLM_ABLATE_SHARED; - if (strstr(env, "qklow")) mask |= DS4_GLM_ABLATE_QKLOW; - if (mask) { - fprintf(stderr, "ds4: GLM decode ablation active (mask 0x%x) — output is garbage, timing only\n", mask); - } - } - cached = (int)mask; - } - return (uint32_t)cached; +static bool cpu_directional_steering_enabled( + const float *dirs, + float scale) { + return dirs && scale != 0.0f; } -static bool glm_graph_encode_shared_swiglu_one( - ds4_gpu_tensor *mid, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos, - const ds4_gpu_tensor *x, - bool ssd_streaming, - bool stage_profile, - double *stage_t0) { - if (!mid || !gate || !up || !model || !l || !x || - !l->ffn_gate_shexp || !l->ffn_up_shexp) { - return false; - } - - bool ok = true; - if (glm_graph_weights_are_q8_0(model, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset)) { - ok = ds4_gpu_shared_mid_swiglu_q8_0_tensor( - mid, - model->map, - model->size, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - x, - 0.0f) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "shared_gate_up_swiglu", - il, - pos, - 1, - stage_t0); - return ok; - } - - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(gate, - model, - l->ffn_gate_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - x, - il, - pos, - "shared_gate", - ssd_streaming) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(up, - model, - l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - x, - il, - pos, - "shared_up", - ssd_streaming) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "shared_gate_up", - il, - pos, - 1, - stage_t0); - if (ok) ok = ds4_gpu_swiglu_tensor(mid, - gate, - up, - DS4_N_FF_EXP, - 0.0f, - 1.0f) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "shared_swiglu", - il, - pos, - 1, - stage_t0); - return ok; -} +static void cpu_directional_steering_project_rows( + float *x, + const float *dirs, + uint32_t il, + uint32_t rows, + float scale) { + if (!cpu_directional_steering_enabled(dirs, scale) || !x || rows == 0) return; -static bool glm_graph_encode_sparse_ffn_one( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos, - const ds4_gpu_tensor *ffn_norm, - const ds4_gpu_tensor *after_attn, - ds4_gpu_tensor *next, - ds4_gpu_tensor *ffn_gate, - ds4_gpu_tensor *ffn_up, - ds4_gpu_tensor *ffn_mid, - ds4_gpu_tensor *ffn_out, - ds4_gpu_tensor *ffn_sum, - ds4_gpu_tensor *tmp, - bool stage_profile, - double *stage_t0) { - uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; - uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; - uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; - (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); - (void)gate_in; - (void)up_in; - (void)down_in; - - bool ok = ds4_gpu_matmul_f32_tensor(g->router_logits, - model->map, - model->size, - l->ffn_gate_inp->abs_offset, - DS4_N_EMBD, - DS4_N_EXPERT, - ffn_norm, - 1) != 0; - if (ok) ok = ds4_gpu_glm_router_select_tensor(g->router_selected, - g->router_weights, - g->router_probs, - model->map, - model->size, - l->ffn_exp_probs_b->abs_offset, - g->router_logits, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "router", - il, - pos, - 1, - stage_t0); - if (ok) ok = glm_graph_profile_router_selection(g, l, il, pos); - const bool resident_decode_layer = - g->ssd_streaming && glm_stream_resident_decode_layer_enabled(l, il); - const bool generic_streaming_selected_cache = - g->ssd_streaming && - !resident_decode_layer && - glm_graph_layer_uses_generic_routed_moe(l); - const bool uniform_streaming_selected_cache = - g->ssd_streaming && - !resident_decode_layer && - l->ffn_gate_exps->type == l->ffn_up_exps->type && - l->ffn_gate_exps->type == l->ffn_down_exps->type && - (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || - l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); - const bool streaming_selected_cache = - generic_streaming_selected_cache || - uniform_streaming_selected_cache; - const bool shared_first = streaming_selected_cache; - metal_graph_selected_async_load async_load = {0}; - bool async_load_started = false; - const bool async_profile = - streaming_selected_cache && - glm_graph_streaming_async_profile_enabled(); - if (async_profile) glm_graph_streaming_async_profile_register(); - const double stream_total_t0 = - async_profile ? glm_graph_streaming_async_profile_ms() : 0.0; - double stream_t0 = stream_total_t0; - bool async_path_profiled = false; - if (ok && streaming_selected_cache) { - const ds4_gpu_stream_expert_table table = { - .model_map = model->map, - .model_size = model->size, - .layer = il, - .n_total_expert = DS4_N_EXPERT, - .gate_offset = l->ffn_gate_exps->abs_offset, - .up_offset = l->ffn_up_exps->abs_offset, - .down_offset = l->ffn_down_exps->abs_offset, - .gate_expert_bytes = gate_out * gate_row_bytes, - .down_expert_bytes = down_out * down_row_bytes, - }; - const bool async_selected_load = -#ifdef DS4_ROCM_BUILD - streaming_selected_cache && -#else - glm_graph_layer_uses_generic_routed_moe(l) && -#endif - glm_graph_use_streaming_selected_async_load(g); - async_path_profiled = false; - uint64_t selected_event = 0; - if (async_selected_load) { - if (ds4_gpu_signal_selected_readback_ready(&selected_event) != 0) { - async_load_started = metal_graph_selected_async_load_start_tensor( - &async_load, - g->router_selected, - model, - l, - il, - selected_event, - gate_out * gate_row_bytes, - down_out * down_row_bytes); - async_path_profiled = async_profile && async_load_started; - } - if (async_profile) { - g_glm_streaming_async_profile.async_signal_start_ms += - glm_graph_streaming_async_profile_ms() - stream_t0; - stream_t0 = glm_graph_streaming_async_profile_ms(); - } -#ifndef DS4_ROCM_BUILD - if (ok && async_load_started) { - ok = ds4_gpu_flush_commands() != 0; - } -#endif - if (async_profile) { - g_glm_streaming_async_profile.async_flush_router_ms += - glm_graph_streaming_async_profile_ms() - stream_t0; - stream_t0 = glm_graph_streaming_async_profile_ms(); - } - } - if (!async_load_started) { - if (async_selected_load && selected_event != 0) { - ok = ds4_gpu_wait_selected_readback_ready( - selected_event, - "selected-id sync expert load fallback") != 0; - } - if (ok) { - ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( - &table, - g->router_selected, - DS4_N_EXPERT_USED) != 0; - } - if (async_profile) { - g_glm_streaming_async_profile.sync_early_load_ms += - glm_graph_streaming_async_profile_ms() - stream_t0; - stream_t0 = glm_graph_streaming_async_profile_ms(); - } - } - } - if (ok && shared_first) { - ok = glm_graph_encode_shared_swiglu_one(ffn_mid, - ffn_gate, - ffn_up, - model, - l, - il, - pos, - ffn_norm, - g->ssd_streaming, - stage_profile, - stage_t0); - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - ffn_mid, - il, - pos, - "shared_down", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "shared_down", - il, - pos, - 1, - stage_t0); - } - if (async_profile) { - const double now_ms = glm_graph_streaming_async_profile_ms(); - if (async_path_profiled) { - g_glm_streaming_async_profile.async_shared_ms += - now_ms - stream_t0; - } else { - g_glm_streaming_async_profile.sync_shared_ms += - now_ms - stream_t0; + const float *dir = dirs + (uint64_t)il * DS4_N_EMBD; + for (uint32_t row = 0; row < rows; row++) { + float *xr = x + (uint64_t)row * DS4_N_EMBD; + float dot = 0.0f; + for (uint32_t i = 0; i < DS4_N_EMBD; i++) { + dot += xr[i] * dir[i]; } - stream_t0 = now_ms; - } - if (async_load_started) { - bool flush_ok = true; -#ifndef DS4_ROCM_BUILD - flush_ok = ds4_gpu_flush_commands() != 0; -#endif - if (async_profile) { - g_glm_streaming_async_profile.async_flush_shared_ms += - glm_graph_streaming_async_profile_ms() - stream_t0; - stream_t0 = glm_graph_streaming_async_profile_ms(); - } - const bool finish_ok = metal_graph_selected_async_load_finish(&async_load); - ok = ok && flush_ok && finish_ok; - if (async_profile) { - g_glm_streaming_async_profile.async_finish_ms += - glm_graph_streaming_async_profile_ms() - stream_t0; - stream_t0 = glm_graph_streaming_async_profile_ms(); - } - } - /* 50/50 TP: this rank's routed partial goes straight into the - * slab out slot, the gate exchanges it with the peer's half, and the - * commutative add rebuilds the full routed output on both ranks - * bit-identically. The shared expert and everything else stay - * replicated, so no other exchange is needed. */ - const bool tp_split_ffn = g->tp_world == 2 && g->tp_out && g->tp_in; - const uint32_t tp_ffn_slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; - ds4_gpu_tensor *routed_dst = tp_split_ffn ? g->tp_out[tp_ffn_slot] : ffn_out; - if (ok && tp_split_ffn && g->ssd_streaming) { - fprintf(stderr, "ds4: GLM tensor parallelism requires resident weights\n"); - ok = false; - } - if (!ok && tp_split_ffn && getenv("DS4_GLM_TP_DEBUG")) { - fprintf(stderr, "ds4: glm sparse ffn: failed before routed dispatch (layer %u)\n", il); - } - if (ok && !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { - ok = glm_graph_routed_moe_one_dispatch( - g, - model, - l, - il, - routed_dst, - ffn_mid, - gate_out * gate_row_bytes, - gate_row_bytes, - up_out * up_row_bytes, - up_row_bytes, - down_out * down_row_bytes, - down_row_bytes, - g->router_selected, - g->router_weights, - ffn_norm, - resident_decode_layer) != 0; - } - if (ok && tp_split_ffn) { - ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_FFN) != 0; - if (ok) ok = ds4_gpu_add_tensor(ffn_out, - g->tp_out[tp_ffn_slot], - g->tp_in[tp_ffn_slot], - DS4_N_EMBD) != 0; - if (!ok) fprintf(stderr, "ds4: GLM TP gate/combine failed (layer %u)\n", il); - } else if (!ok && tp_split_ffn) { - fprintf(stderr, "ds4: GLM TP routed dispatch failed before the gate (layer %u)\n", il); - } - if (async_profile) { - const double now_ms = glm_graph_streaming_async_profile_ms(); - if (async_path_profiled) { - g_glm_streaming_async_profile.async_routed_ms += - now_ms - stream_t0; - } else { - g_glm_streaming_async_profile.sync_routed_ms += - now_ms - stream_t0; - } - stream_t0 = now_ms; - } - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "routed_moe", - il, - pos, - 1, - stage_t0); - if (ok && !shared_first && - !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_SHARED)) { - ok = glm_graph_encode_shared_swiglu_one(ffn_mid, - ffn_gate, - ffn_up, - model, - l, - il, - pos, - ffn_norm, - g->ssd_streaming, - stage_profile, - stage_t0); - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - ffn_mid, - il, - pos, - "shared_down", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "shared_down", - il, - pos, - 1, - stage_t0); - } - if (ok && !glm_graph_disable_add3_residual()) { - ok = ds4_gpu_add3_tensor(next, - after_attn, - ffn_out, - ffn_sum, - DS4_N_EMBD) != 0; - } else if (ok) { - ok = ds4_gpu_add_tensor(tmp, - ffn_out, - ffn_sum, - DS4_N_EMBD) != 0; - if (ok) ok = ds4_gpu_add_tensor(next, - after_attn, - tmp, - DS4_N_EMBD) != 0; - } - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "residual", - il, - pos, - 1, - stage_t0); - if (async_profile) { - const double now_ms = glm_graph_streaming_async_profile_ms(); - if (async_path_profiled) { - g_glm_streaming_async_profile.async_calls++; - g_glm_streaming_async_profile.async_post_ms += now_ms - stream_t0; - g_glm_streaming_async_profile.async_total_ms += - now_ms - stream_total_t0; - } else { - g_glm_streaming_async_profile.sync_calls++; - g_glm_streaming_async_profile.sync_post_ms += now_ms - stream_t0; - g_glm_streaming_async_profile.sync_total_ms += - now_ms - stream_total_t0; + const float coeff = scale * dot; + for (uint32_t i = 0; i < DS4_N_EMBD; i++) { + xr[i] -= coeff * dir[i]; } } - return ok; } -static bool glm_graph_encode_ffn_one_normed_from( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos, - const ds4_gpu_tensor *ffn_norm, - const ds4_gpu_tensor *after_attn, - ds4_gpu_tensor *next, - ds4_gpu_tensor *ffn_gate, - ds4_gpu_tensor *ffn_up, - ds4_gpu_tensor *ffn_mid, - ds4_gpu_tensor *ffn_out, - ds4_gpu_tensor *ffn_sum, - ds4_gpu_tensor *tmp, - bool stage_profile, - double *stage_t0) { - if (!g || !model || !l || !ffn_norm || !after_attn || !next || - !ffn_gate || !ffn_up || !ffn_mid || !ffn_out || - !ffn_sum || !tmp) { - return false; +static bool cpu_load_directional_steering(ds4_engine *e) { + if (!e || + (e->directional_steering_attn_scale == 0.0f && + e->directional_steering_ffn_scale == 0.0f)) { + return true; } - if (il < DS4_N_LEADING_DENSE) { - const uint64_t hidden = l->ffn_gate->dim[1]; - const bool can_fuse_gate_up = - glm_graph_weights_are_q8_0(model, - l->ffn_gate->abs_offset, - l->ffn_up->abs_offset); - const bool fused_gate_up = can_fuse_gate_up && - ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( - ffn_gate, - ffn_up, - ffn_mid, - model->map, - model->size, - l->ffn_gate->abs_offset, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - ffn_norm, - 0.0f) != 0; - bool ok = fused_gate_up; - if (fused_gate_up) { - ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "dense_gate_up_swiglu", - il, - pos, - 1, - stage_t0); - } else { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_gate, - model, - l->ffn_gate->abs_offset, - DS4_N_EMBD, - hidden, - ffn_norm, - il, - pos, - "dense_gate", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_up, - model, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - ffn_norm, - il, - pos, - "dense_up", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "dense_gate_up", - il, - pos, - 1, - stage_t0); - if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, - ffn_gate, - ffn_up, - (uint32_t)hidden, - 0.0f, - 1.0f) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "dense_swiglu", - il, - pos, - 1, - stage_t0); - } - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_out, - model, - l->ffn_down->abs_offset, - hidden, - DS4_N_EMBD, - ffn_mid, - il, - pos, - "dense_down", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "dense_down", - il, - pos, - 1, - stage_t0); - if (ok) ok = ds4_gpu_add_tensor(next, - after_attn, - ffn_out, - DS4_N_EMBD) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "residual", - il, - pos, - 1, - stage_t0); - return ok; + const char *path = e->directional_steering_file; + if (!path || !path[0]) { + fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); + return false; } - return glm_graph_encode_sparse_ffn_one(g, - model, - l, - il, - pos, - ffn_norm, - after_attn, - next, - ffn_gate, - ffn_up, - ffn_mid, - ffn_out, - ffn_sum, - tmp, - stage_profile, - stage_t0); -} - -static bool glm_graph_encode_ffn_one_from( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos, - const ds4_gpu_tensor *after_attn, - ds4_gpu_tensor *next, - ds4_gpu_tensor *ffn_norm, - ds4_gpu_tensor *ffn_gate, - ds4_gpu_tensor *ffn_up, - ds4_gpu_tensor *ffn_mid, - ds4_gpu_tensor *ffn_out, - ds4_gpu_tensor *ffn_sum, - ds4_gpu_tensor *tmp, - bool stage_profile, - double *stage_t0) { - if (!g || !model || !l || !after_attn || !next || - !ffn_norm || !ffn_gate || !ffn_up || !ffn_mid || !ffn_out || - !ffn_sum || !tmp) { + const uint64_t n = (uint64_t)DS4_N_LAYER * DS4_N_EMBD; + e->directional_steering_dirs = xmalloc((size_t)n * sizeof(e->directional_steering_dirs[0])); + if (!read_f32_binary_file(path, e->directional_steering_dirs, n)) { + free(e->directional_steering_dirs); + e->directional_steering_dirs = NULL; + fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); return false; } - - bool ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, - after_attn, - model->map, - model->size, - l->ffn_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_profile_stage(stage_profile, - "glm_decode_ffn", - "ffn_norm", - il, - pos, - 1, - stage_t0); - if (!ok) return false; - - return glm_graph_encode_ffn_one_normed_from(g, - model, - l, - il, - pos, - ffn_norm, - after_attn, - next, - ffn_gate, - ffn_up, - ffn_mid, - ffn_out, - ffn_sum, - tmp, - stage_profile, - stage_t0); + fprintf(stderr, "ds4: CPU directional steering enabled: %s attn=%g ffn=%g\n", + path, + (double)e->directional_steering_attn_scale, + (double)e->directional_steering_ffn_scale); + return true; } -static ds4_gpu_tensor *glm_graph_tensor_row_view_strided( - ds4_gpu_tensor *base, - uint32_t row, - uint64_t stride_values, - uint64_t row_values) { - return ds4_gpu_tensor_view(base, - (uint64_t)row * stride_values * sizeof(float), - row_values * sizeof(float)); +static void utf8_put(char **p, uint32_t cp) { + if (cp <= 0x7f) { + *(*p)++ = (char)cp; + } else if (cp <= 0x7ff) { + *(*p)++ = (char)(0xc0 | (cp >> 6)); + *(*p)++ = (char)(0x80 | (cp & 0x3f)); + } else if (cp <= 0xffff) { + *(*p)++ = (char)(0xe0 | (cp >> 12)); + *(*p)++ = (char)(0x80 | ((cp >> 6) & 0x3f)); + *(*p)++ = (char)(0x80 | (cp & 0x3f)); + } else { + *(*p)++ = (char)(0xf0 | (cp >> 18)); + *(*p)++ = (char)(0x80 | ((cp >> 12) & 0x3f)); + *(*p)++ = (char)(0x80 | ((cp >> 6) & 0x3f)); + *(*p)++ = (char)(0x80 | (cp & 0x3f)); + } } -static uint32_t glm_graph_q8_stripe_tokens(void) { - return 2048u; -} +static uint32_t gpt2_byte_to_codepoint(uint8_t b) { + if ((b >= 33 && b <= 126) || (b >= 161 && b <= 172) || (b >= 174)) { + return b; + } -static bool glm_graph_flash_attention_prefill_enabled(void) { - return getenv("DS4_GLM_DISABLE_FLASH_PREFILL") == NULL; + uint32_t n = 0; + for (uint32_t x = 0; x < 256; x++) { + if ((x >= 33 && x <= 126) || (x >= 161 && x <= 172) || (x >= 174)) { + continue; + } + if (x == b) return 256 + n; + n++; + } + return b; } -static uint32_t glm_graph_flash_attention_prefill_min_tokens(void) { - return 24u; -} +/* GPT-2 byte-level BPE first maps raw bytes to printable Unicode codepoints + * so merges can operate on UTF-8 strings without losing byte identity. */ +static char *byte_encode(ds4_str in, uint64_t *out_len) { + char *out = xmalloc((size_t)in.len * 4 + 1); + char *p = out; -static bool glm_graph_use_flash_attention_prefill(uint32_t n_tokens) { - return glm_graph_flash_attention_prefill_enabled() && - n_tokens >= glm_graph_flash_attention_prefill_min_tokens(); + for (uint64_t i = 0; i < in.len; i++) { + utf8_put(&p, gpt2_byte_to_codepoint((uint8_t)in.ptr[i])); + } + *p = '\0'; + *out_len = (uint64_t)(p - out); + return out; } -static bool glm_graph_use_flash_attention_staged_kv( - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len) { - return pos0 == 0 && - n_tokens == cache_len; +static int utf8_len_from_first_byte(uint8_t c) { + if (c < 0x80) return 1; + if ((c & 0xe0) == 0xc0) return 2; + if ((c & 0xf0) == 0xe0) return 3; + if ((c & 0xf8) == 0xf0) return 4; + return 1; } -static bool glm_graph_force_indexed_decode(void) { - return false; -} +typedef struct { + char *ptr; + uint64_t len; +} owned_str; -static bool glm_graph_disable_indexed_decode(void) { - return false; +static owned_str owned_copy(const char *ptr, uint64_t len) { + owned_str s; + s.ptr = xmalloc((size_t)len); + memcpy(s.ptr, ptr, (size_t)len); + s.len = len; + return s; } -static bool glm_graph_decode_uses_indexed_attention(const ds4_glm_gpu_graph *g, - uint32_t pos, - const float *logits_out) { - return g && g->compact_cache_cap != 0 && - (!g->full_kv_cache || - pos >= g->ctx_cap || - glm_graph_force_indexed_decode() || - (logits_out != NULL && !glm_graph_disable_indexed_decode())); -} +/* Look up the merge rank for two adjacent BPE symbols. */ +static int bpe_rank(const ds4_vocab *vocab, const owned_str *a, const owned_str *b) { + uint64_t len = a->len + 1 + b->len; + char stack[512]; + char *buf = len <= sizeof(stack) ? stack : xmalloc((size_t)len); -static bool glm_graph_decode_updates_dense_cache(const ds4_glm_gpu_graph *g, - uint32_t pos, - const float *logits_out) { - return g && pos < g->ctx_cap && - !glm_graph_decode_uses_indexed_attention(g, pos, logits_out); -} + memcpy(buf, a->ptr, (size_t)a->len); + buf[a->len] = ' '; + memcpy(buf + a->len + 1, b->ptr, (size_t)b->len); -static bool glm_graph_indexed_prefill_scalar_kernels(void) { - return false; -} + int rank = -1; + table_get(&vocab->merge_rank, buf, len, &rank); -static bool glm_graph_indexed_prefill_scalar_indexer(void) { - return false; + if (buf != stack) free(buf); + return rank; } -static bool glm_graph_indexed_prefill_batch_indexer(void) { - return true; -} +/* Apply byte-level BPE to one regex-like pre-tokenized piece and emit token ids. */ +static void bpe_emit_piece(const ds4_vocab *vocab, ds4_str raw_piece, token_vec *out) { + uint64_t encoded_len = 0; + char *encoded = byte_encode(raw_piece, &encoded_len); -static bool glm_graph_indexed_prefill_scalar_attn(void) { - return false; -} + int n_sym = 0; + int cap_sym = 32; + owned_str *sym = xcalloc((size_t)cap_sym, sizeof(sym[0])); -static bool glm_graph_indexed_prefill_batch_qk_low(void) { - return true; -} + for (uint64_t off = 0; off < encoded_len;) { + int n = utf8_len_from_first_byte((uint8_t)encoded[off]); + if (off + (uint64_t)n > encoded_len) n = 1; + if (n_sym == cap_sym) { + cap_sym *= 2; + sym = xrealloc(sym, (size_t)cap_sym * sizeof(sym[0])); + } + sym[n_sym++] = owned_copy(encoded + off, (uint64_t)n); + off += (uint64_t)n; + } -static bool glm_graph_indexed_prefill_batch_attn_kernel(void) { - return true; -} + for (;;) { + int best_i = -1; + int best_rank = INT32_MAX; + + for (int i = 0; i + 1 < n_sym; i++) { + int rank = bpe_rank(vocab, &sym[i], &sym[i + 1]); + if (rank >= 0 && rank < best_rank) { + best_rank = rank; + best_i = i; + } + } + + if (best_i < 0) break; + + owned_str merged; + merged.len = sym[best_i].len + sym[best_i + 1].len; + merged.ptr = xmalloc((size_t)merged.len); + memcpy(merged.ptr, sym[best_i].ptr, (size_t)sym[best_i].len); + memcpy(merged.ptr + sym[best_i].len, sym[best_i + 1].ptr, (size_t)sym[best_i + 1].len); -static uint32_t glm_graph_indexed_prefill_batch_attn_slice_tokens(void) { - return 2048u; -} + free(sym[best_i].ptr); + free(sym[best_i + 1].ptr); + sym[best_i] = merged; -static bool glm_graph_indexer_qat(void) { - return false; -} + for (int j = best_i + 1; j + 1 < n_sym; j++) { + sym[j] = sym[j + 1]; + } + n_sym--; + } -static bool glm_graph_indexed_prefill_batch_ffn(void) { - return true; -} + for (int i = 0; i < n_sym; i++) { + int token = -1; + if (table_get(&vocab->token_to_id, sym[i].ptr, sym[i].len, &token)) { + token_vec_push(out, token); + } else { + for (uint64_t j = 0; j < sym[i].len; j++) { + if (table_get(&vocab->token_to_id, sym[i].ptr + j, 1, &token)) { + token_vec_push(out, token); + } + } + } + free(sym[i].ptr); + } -static bool glm_graph_indexed_prefill_batch_ffn_norm(void) { - return true; + free(sym); + free(encoded); } -static bool glm_graph_indexed_prefill_batch_routed_moe(void) { - return true; +static uint64_t next_utf8_char(const char *s, uint64_t len, uint64_t pos) { + int n = utf8_len_from_first_byte((uint8_t)s[pos]); + if (pos + (uint64_t)n > len) n = 1; + return pos + (uint64_t)n; } -static bool glm_graph_indexed_prefill_batch_router_select(void) { - return true; +static bool ascii_alpha(uint8_t c) { + return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z'); } -static bool glm_graph_indexed_prefill_batch_residual(void) { - return true; +static bool ascii_digit(uint8_t c) { + return c >= '0' && c <= '9'; } -static bool glm_graph_indexed_prefill_batch_f32_rows(void) { - return true; +static bool ascii_space(uint8_t c) { + return c == ' ' || c == '\t' || c == '\n' || c == '\r' || + c == '\v' || c == '\f'; } -static bool glm_graph_indexed_prefill_batch_q8_rows(void) { - return true; +static bool ascii_newline(uint8_t c) { + return c == '\n' || c == '\r'; } -static bool glm_graph_indexed_prefill_batch_shared_expert(void) { - return true; +static bool joyai_ascii_punct_symbol(uint8_t c) { + return (c >= '!' && c <= '/') || + (c >= ':' && c <= '@') || + (c >= '[' && c <= '`') || + (c >= '{' && c <= '~'); } -static bool glm_graph_matmul_q8_0_tensor( - ds4_gpu_tensor *out, - const ds4_model *model, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint32_t n_tokens) { - if (!out || !model || !x || n_tokens == 0) return false; - - const uint32_t q8_stripe_tokens = glm_graph_q8_stripe_tokens(); - if (n_tokens <= q8_stripe_tokens) { - return ds4_gpu_matmul_quant_tensor(out, - model->map, - model->size, - weight_offset, - glm_graph_weight_type_for_offset(model, weight_offset), - in_dim, - out_dim, - x, - n_tokens) != 0; - } - if (in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float)) { - return false; - } - - uint32_t done = 0; - while (done < n_tokens) { - uint32_t chunk = n_tokens - done; - if (chunk > q8_stripe_tokens) chunk = q8_stripe_tokens; - /* - * The Q8 prefill TensorOps path needs token counts divisible by 32. - * For a final chunk like 1736 rows, keep the aligned 1728 rows on that - * path and leave only the tiny tail to the small-batch kernel. Avoid - * splitting small batches where the extra launch would dominate. - */ - if (chunk >= 256u && chunk == n_tokens - done) { - const uint32_t tail = chunk & 31u; - if (tail != 0u && tail <= 16u) chunk -= tail; - } - - ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( - x, - (uint64_t)done * in_dim * sizeof(float), - (uint64_t)chunk * in_dim * sizeof(float)); - ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( - out, - (uint64_t)done * out_dim * sizeof(float), - (uint64_t)chunk * out_dim * sizeof(float)); - const bool ok = x_view && out_view && - ds4_gpu_matmul_quant_tensor(out_view, - model->map, - model->size, - weight_offset, - glm_graph_weight_type_for_offset(model, weight_offset), - in_dim, - out_dim, - x_view, - chunk) != 0; - ds4_gpu_tensor_free(out_view); - ds4_gpu_tensor_free(x_view); - if (!ok) return false; - done += chunk; - } - return true; +static bool utf8_is_cjk_hira_kata(uint32_t cp) { + return (cp >= 0x4e00 && cp <= 0x9fa5) || + (cp >= 0x3040 && cp <= 0x309f) || + (cp >= 0x30a0 && cp <= 0x30ff); } -static bool glm_graph_matmul_q8_0_rows_scalar( - ds4_gpu_tensor *out, - const ds4_model *model, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint32_t n_tokens) { - if (!out || !model || !x || n_tokens == 0) return false; - if (in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float)) { - return false; - } - - if (glm_graph_indexed_prefill_batch_q8_rows() && - ds4_gpu_matmul_quant_rows_scalar_tensor(out, - model->map, - model->size, - weight_offset, - glm_graph_weight_type_for_offset(model, weight_offset), - in_dim, - out_dim, - x, - n_tokens) != 0) { - return true; - } - - for (uint32_t t = 0; t < n_tokens; t++) { - ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( - (ds4_gpu_tensor *)x, - (uint64_t)t * in_dim * sizeof(float), - in_dim * sizeof(float)); - ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( - out, - (uint64_t)t * out_dim * sizeof(float), - out_dim * sizeof(float)); - const bool ok = x_view && out_view && - ds4_gpu_matmul_quant_tensor(out_view, - model->map, - model->size, - weight_offset, - glm_graph_weight_type_for_offset(model, weight_offset), - in_dim, - out_dim, - x_view, - 1) != 0; - ds4_gpu_tensor_free(out_view); - ds4_gpu_tensor_free(x_view); - if (!ok) return false; - } - return true; -} +static uint32_t utf8_peek_one(const char *s, uint64_t len, uint64_t pos, uint64_t *next) { + const uint8_t c0 = (uint8_t)s[pos]; + int n = utf8_len_from_first_byte(c0); + if (pos + (uint64_t)n > len) n = 1; + *next = pos + (uint64_t)n; -static bool glm_graph_matmul_f32_rows_scalar( - ds4_gpu_tensor *out, - const ds4_model *model, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint32_t n_tokens) { - if (!out || !model || !x || n_tokens == 0) return false; - if (in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float)) { - return false; + if (n == 1) return c0; + if (n == 2) { + return ((uint32_t)(c0 & 0x1f) << 6) | + ((uint32_t)((uint8_t)s[pos + 1] & 0x3f)); } - - if (glm_graph_indexed_prefill_batch_f32_rows()) { - return ds4_gpu_matmul_f32_tensor(out, - model->map, - model->size, - weight_offset, - in_dim, - out_dim, - x, - n_tokens) != 0; - } - - for (uint32_t t = 0; t < n_tokens; t++) { - ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( - (ds4_gpu_tensor *)x, - (uint64_t)t * in_dim * sizeof(float), - in_dim * sizeof(float)); - ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( - out, - (uint64_t)t * out_dim * sizeof(float), - out_dim * sizeof(float)); - const bool ok = x_view && out_view && - ds4_gpu_matmul_f32_tensor(out_view, - model->map, - model->size, - weight_offset, - in_dim, - out_dim, - x_view, - 1) != 0; - ds4_gpu_tensor_free(out_view); - ds4_gpu_tensor_free(x_view); - if (!ok) return false; + if (n == 3) { + return ((uint32_t)(c0 & 0x0f) << 12) | + ((uint32_t)((uint8_t)s[pos + 1] & 0x3f) << 6) | + ((uint32_t)((uint8_t)s[pos + 2] & 0x3f)); } - return true; + return ((uint32_t)(c0 & 0x07) << 18) | + ((uint32_t)((uint8_t)s[pos + 1] & 0x3f) << 12) | + ((uint32_t)((uint8_t)s[pos + 2] & 0x3f) << 6) | + ((uint32_t)((uint8_t)s[pos + 3] & 0x3f)); } -static bool glm_graph_shared_gate_up_swiglu_q8_0_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const ds4_model *model, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - float clamp) { - if (!gate || !up || !mid || !model || !x || n_tokens == 0) return false; - if (!glm_graph_weights_are_q8_0(model, gate_offset, up_offset)) return false; - - const uint32_t q8_stripe_tokens = glm_graph_q8_stripe_tokens(); - if (n_tokens == 1) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(gate, - up, - mid, - model->map, - model->size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp) != 0; - } - if (n_tokens <= q8_stripe_tokens) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate, - up, - mid, - model->map, - model->size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - n_tokens, - clamp) != 0; - } - if (in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float)) { - return false; - } +static bool joyai_letter_like_at(const char *s, uint64_t len, uint64_t pos) { + (void)len; + uint8_t c = (uint8_t)s[pos]; + if (c < 128) return ascii_alpha(c); - uint32_t done = 0; - while (done < n_tokens) { - uint32_t chunk = n_tokens - done; - if (chunk > q8_stripe_tokens) chunk = q8_stripe_tokens; - - ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( - x, - (uint64_t)done * in_dim * sizeof(float), - (uint64_t)chunk * in_dim * sizeof(float)); - ds4_gpu_tensor *gate_view = ds4_gpu_tensor_view( - gate, - (uint64_t)done * out_dim * sizeof(float), - (uint64_t)chunk * out_dim * sizeof(float)); - ds4_gpu_tensor *up_view = ds4_gpu_tensor_view( - up, - (uint64_t)done * out_dim * sizeof(float), - (uint64_t)chunk * out_dim * sizeof(float)); - ds4_gpu_tensor *mid_view = ds4_gpu_tensor_view( - mid, - (uint64_t)done * out_dim * sizeof(float), - (uint64_t)chunk * out_dim * sizeof(float)); - const bool ok = x_view && gate_view && up_view && mid_view && - ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate_view, - up_view, - mid_view, - model->map, - model->size, - gate_offset, - up_offset, - in_dim, - out_dim, - x_view, - chunk, - clamp) != 0; - ds4_gpu_tensor_free(mid_view); - ds4_gpu_tensor_free(up_view); - ds4_gpu_tensor_free(gate_view); - ds4_gpu_tensor_free(x_view); - if (!ok) return false; - done += chunk; - } + /* + * The JoyAI tokenizer maps Unicode letters into a collapsed regex alphabet before + * applying the JoyAI pre-tokenizer. The prompts we care about are mostly + * ASCII, but treating non-ASCII non-control bytes as letters preserves the + * useful behavior for ordinary UTF-8 text such as Italian accents. CJK and + * kana are isolated by the JoyAI pre-tokenizer before the generic letter + * rule, below. + */ return true; } -static bool glm_graph_indexed_prefill_grouped_moe_default( - const ds4_glm_gpu_graph *g) { - return g && - !g->quality; -} - -static uint32_t glm_graph_streaming_prefill_cache_seed_k( - const ds4_glm_gpu_graph *g) { - const bool enabled = - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_CACHE_SEED", - "DS4_METAL_ENABLE_STREAMING_PREFILL_CACHE_SEED"); - if (!g || - !g->ssd_streaming || - !enabled) { - return 0; - } - - uint32_t k = 1; - const char *env = glm_graph_env_value("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_K", - "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_K"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0') { - if (v == 0) return 0; - k = v > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS ? - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS : (uint32_t)v; - } +static uint64_t joyai_consume_letters(const char *s, uint64_t len, uint64_t pos) { + while (pos < len && joyai_letter_like_at(s, len, pos)) { + pos = next_utf8_char(s, len, pos); } - return k; -} - -static bool glm_graph_streaming_prefill_cache_seed_enabled( - const ds4_glm_gpu_graph *g) { - return glm_graph_streaming_prefill_cache_seed_k(g) != 0; + return pos; } -static void glm_graph_reset_prefill_seed_capture(ds4_glm_gpu_graph *g) { - if (!g) return; - g->prefill_seed_tokens = 0; - memset(g->prefill_seed_layer_captured, - 0, - sizeof(g->prefill_seed_layer_captured)); +static bool joyai_cjk_at(const char *s, uint64_t len, uint64_t pos) { + if ((uint8_t)s[pos] < 128) return false; + uint64_t next = pos; + uint32_t cp = utf8_peek_one(s, len, pos, &next); + return utf8_is_cjk_hira_kata(cp); } -static bool glm_graph_streaming_expert_cache_seed_layer_expected( - const ds4_glm_gpu_graph *g, - const ds4_weights *weights, - const ds4_layer_weights *layer, - uint32_t il) { - if (!g || - !g->ssd_streaming || - g->quality || - !weights || - !layer || - !layer->ffn_gate_exps || - !layer->ffn_up_exps || - !layer->ffn_down_exps) { - return false; - } - if (glm_stream_resident_decode_layer_enabled(layer, il)) return false; - return glm_stream_selected_expert_cache_supported(layer, il) || - glm_stream_expert_cache_addr_layout_supported(weights, layer, il); -} +typedef struct { + uint32_t cp; + uint64_t next; + bool valid; + bool is_letter; + bool is_number; + bool is_whitespace; +} glm4_char_info; -static bool glm_graph_capture_prefill_seed_router_selected( - ds4_glm_gpu_graph *g, - uint32_t il, - uint32_t n_tokens) { - uint32_t k = glm_graph_streaming_prefill_cache_seed_k(g); - if (k == 0) return true; - if (!g->prefill_seed_router_selected || - !g->batch_router_selected || - il >= DS4_N_LAYER || - il >= DS4_MAX_LAYER || - n_tokens == 0 || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return false; - } - if (k > n_tokens) k = n_tokens; - - const uint64_t bytes = (uint64_t)k * DS4_N_EXPERT_USED * sizeof(int32_t); - const uint64_t src_off = (uint64_t)(n_tokens - k) * - DS4_N_EXPERT_USED * sizeof(int32_t); - const uint64_t dst_off = (uint64_t)il * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * sizeof(int32_t); - if (ds4_gpu_tensor_copy(g->prefill_seed_router_selected, - dst_off, - g->batch_router_selected, - src_off, - bytes) == 0) { - return false; - } - g->prefill_seed_tokens = k; - g->prefill_seed_layer_captured[il] = true; - return true; +static bool glm4_unicode_whitespace(uint32_t cp) { + if (cp < 128) return ascii_space((uint8_t)cp); + return cp == 0x0085 || + cp == 0x00a0 || + cp == 0x1680 || + (cp >= 0x2000 && cp <= 0x200a) || + cp == 0x2028 || + cp == 0x2029 || + cp == 0x202f || + cp == 0x205f || + cp == 0x3000; } -static bool glm_graph_seed_streaming_expert_cache_from_prefill( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights) { - if (!glm_graph_streaming_prefill_cache_seed_enabled(g)) return true; - const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; - if (seed_tokens == 0) return true; - if (!g || - !model || - !weights || - !g->prefill_seed_router_selected || - seed_tokens > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS || - DS4_N_LAYER > DS4_MAX_LAYER || - DS4_N_EXPERT_USED == 0 || - DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { - return false; - } - bool any_captured = false; - for (uint32_t il = g->layer_start; - il <= g->layer_end && il < DS4_N_LAYER && il < DS4_MAX_LAYER; - il++) { - if (g->prefill_seed_layer_captured[il]) { - any_captured = true; - break; - } - } - if (!any_captured) return true; - - int32_t selected[DS4_MAX_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_MAX_EXPERT_USED]; - const uint64_t bytes = (uint64_t)DS4_N_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * sizeof(selected[0]); - if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, - 0, - selected, - bytes) == 0) { - return false; - } - - const bool profile = - glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_PROFILE", - "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE"); - const double t0 = profile ? now_sec() : 0.0; - uint32_t seeded_layers = 0; - uint32_t seeded_rows = 0; - for (uint32_t il = g->layer_start; - il <= g->layer_end && il < DS4_N_LAYER; - il++) { - if (il >= DS4_MAX_LAYER || !g->prefill_seed_layer_captured[il]) { - continue; - } - const ds4_layer_weights *layer = &weights->layer[il]; - if (!glm_graph_streaming_expert_cache_seed_layer_expected(g, - weights, - layer, - il)) { - continue; - } - - uint64_t gate_expert_bytes = 0; - uint64_t down_expert_bytes = 0; - if (!streaming_layer_gate_down_expert_bytes(layer, - &gate_expert_bytes, - &down_expert_bytes)) { - fprintf(stderr, - "ds4: GLM prefill expert-cache seed byte size overflow at layer %u\n", - il); - return false; - } - - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - for (uint32_t row = 0; row < seed_tokens; row++) { - const size_t sel_off = ((size_t)il * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS + - row) * DS4_N_EXPERT_USED; - if (ds4_gpu_stream_expert_cache_seed_selected( - &table, - selected + sel_off, - DS4_N_EXPERT_USED) == 0) { - return false; - } - seeded_rows++; - } - seeded_layers++; - } - if (profile) { - fprintf(stderr, - "ds4: GLM streaming prefill expert-cache seed k=%u layers=%u rows=%u time=%.3f ms\n", - seed_tokens, - seeded_layers, - seeded_rows, - (now_sec() - t0) * 1000.0); - } - return true; +static bool glm4_unicode_number(uint32_t cp) { + if (cp < 128) return ascii_digit((uint8_t)cp); + return (cp >= 0x0660 && cp <= 0x0669) || + (cp >= 0x06f0 && cp <= 0x06f9) || + (cp >= 0x07c0 && cp <= 0x07c9) || + (cp >= 0x0966 && cp <= 0x096f) || + (cp >= 0x09e6 && cp <= 0x09ef) || + (cp >= 0x0a66 && cp <= 0x0a6f) || + (cp >= 0x0ae6 && cp <= 0x0aef) || + (cp >= 0x0b66 && cp <= 0x0b6f) || + (cp >= 0x0be6 && cp <= 0x0bef) || + (cp >= 0x0c66 && cp <= 0x0c6f) || + (cp >= 0x0ce6 && cp <= 0x0cef) || + (cp >= 0x0d66 && cp <= 0x0d6f) || + (cp >= 0x0de6 && cp <= 0x0def) || + (cp >= 0x0e50 && cp <= 0x0e59) || + (cp >= 0x0ed0 && cp <= 0x0ed9) || + (cp >= 0x0f20 && cp <= 0x0f29) || + (cp >= 0x1040 && cp <= 0x1049) || + (cp >= 0x1090 && cp <= 0x1099) || + (cp >= 0x17e0 && cp <= 0x17e9) || + (cp >= 0x1810 && cp <= 0x1819) || + (cp >= 0xff10 && cp <= 0xff19); } -static bool glm_graph_seed_streaming_expert_cache_from_full_layer( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const ds4_layer_weights *layer, - uint32_t il, - uint32_t n_tokens, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - bool full_layer_prefill) { -#ifdef DS4_ROCM_BUILD - uint32_t seed_tokens = glm_graph_streaming_prefill_cache_seed_k(g); - if (seed_tokens == 0) return true; - if (!full_layer_prefill || - !model || - !weights || - !layer || - !g || - !g->batch_router_selected || - n_tokens == 0 || - il >= DS4_N_LAYER || - il >= DS4_MAX_LAYER || - gate_expert_bytes == 0 || - down_expert_bytes == 0 || - !g->prefill_seed_layer_captured[il] || - !glm_graph_streaming_expert_cache_seed_layer_expected(g, - weights, - layer, - il)) { - return true; - } - if (seed_tokens > n_tokens) seed_tokens = n_tokens; - const ds4_gpu_stream_expert_table table = - graph_stream_expert_table_make(model, - layer, - il, - gate_expert_bytes, - down_expert_bytes); - if (ds4_gpu_stream_expert_cache_seed_from_layer_selected( - &table, - g->batch_router_selected, - n_tokens, - seed_tokens, - DS4_N_EXPERT_USED) != 0) { - g->prefill_seed_layer_captured[il] = false; - return true; - } - - static bool warned = false; - if (!warned) { - fprintf(stderr, - "ds4: GLM ROCm full-layer prefill expert-cache seed skipped; " - "falling back to end-of-prefill selected seed\n"); - warned = true; - } - return true; -#else - (void)g; - (void)model; - (void)weights; - (void)layer; - (void)il; - (void)n_tokens; - (void)gate_expert_bytes; - (void)down_expert_bytes; - (void)full_layer_prefill; - return true; -#endif +static bool glm4_unicode_punct_symbol(uint32_t cp) { + if (cp < 128) return joyai_ascii_punct_symbol((uint8_t)cp); + return (cp >= 0x00a1 && cp <= 0x00a9) || + (cp >= 0x00ab && cp <= 0x00ac) || + (cp >= 0x00ae && cp <= 0x00b1) || + cp == 0x00b4 || + (cp >= 0x00b6 && cp <= 0x00b8) || + cp == 0x00bb || + cp == 0x00bf || + cp == 0x00d7 || + cp == 0x00f7 || + (cp >= 0x02c2 && cp <= 0x02df) || + (cp >= 0x02e5 && cp <= 0x02eb) || + (cp >= 0x02ed && cp <= 0x02ff) || + (cp >= 0x0375 && cp <= 0x037e) || + (cp >= 0x0384 && cp <= 0x0385) || + cp == 0x0387 || + (cp >= 0x055a && cp <= 0x055f) || + (cp >= 0x0589 && cp <= 0x058a) || + (cp >= 0x05be && cp <= 0x05c0) || + cp == 0x05c3 || + (cp >= 0x05c6 && cp <= 0x05c7) || + (cp >= 0x0609 && cp <= 0x060a) || + (cp >= 0x060c && cp <= 0x060d) || + cp == 0x061b || + (cp >= 0x061e && cp <= 0x061f) || + cp == 0x066a || + cp == 0x066d || + cp == 0x06d4 || + (cp >= 0x2000 && cp <= 0x206f) || + (cp >= 0x20a0 && cp <= 0x20cf) || + (cp >= 0x2100 && cp <= 0x214f) || + (cp >= 0x2190 && cp <= 0x23ff) || + (cp >= 0x2460 && cp <= 0x24ff) || + (cp >= 0x2500 && cp <= 0x2775) || + (cp >= 0x2794 && cp <= 0x2bff) || + (cp >= 0x2e00 && cp <= 0x2e7f) || + (cp >= 0x3000 && cp <= 0x303f) || + (cp >= 0xfd3e && cp <= 0xfd3f) || + (cp >= 0xfe10 && cp <= 0xfe6f) || + (cp >= 0xff01 && cp <= 0xff0f) || + (cp >= 0xff1a && cp <= 0xff20) || + (cp >= 0xff3b && cp <= 0xff40) || + (cp >= 0xff5b && cp <= 0xff65) || + (cp >= 0x1f000 && cp <= 0x1faff); } -static bool glm_graph_disable_add3_residual(void); - -static bool glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos0, - const ds4_gpu_tensor *after_attn, - ds4_gpu_tensor *next, - uint32_t n_tokens, - bool stage_profile, - bool stage_sync, - double *stage_t0) { - if (!g || !model || !l || !after_attn || !next || - !g->batch_ffn_norm || - !g->batch_router_logits || - !g->batch_router_probs || - !g->batch_router_selected || - !g->batch_router_weights || - !g->batch_ffn_out || - !g->batch_ffn_mid || - n_tokens <= 1 || - il < DS4_N_LEADING_DENSE || - g->ffn_mid_elems > UINT32_MAX) { - return false; - } - - uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; - uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; - uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; - (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); - (void)gate_in; - (void)up_in; - (void)down_in; +static glm4_char_info glm4_char_at(const char *s, uint64_t len, uint64_t pos) { + glm4_char_info info; + memset(&info, 0, sizeof(info)); + if (pos >= len) return info; - bool ok = glm_graph_matmul_f32_rows_scalar(g->batch_router_logits, - model, - l->ffn_gate_inp->abs_offset, - DS4_N_EMBD, - DS4_N_EXPERT, - g->batch_ffn_norm, - n_tokens); - const bool use_batch_router_select = - glm_graph_indexed_prefill_batch_router_select(); - if (ok && use_batch_router_select) { - ok = ds4_gpu_glm_router_select_batch_tensor(g->batch_router_selected, - g->batch_router_weights, - g->batch_router_probs, - model->map, - model->size, - l->ffn_exp_probs_b->abs_offset, - g->batch_router_logits, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE, - n_tokens) != 0; - } - for (uint32_t t = 0; ok && !use_batch_router_select && t < n_tokens; t++) { - ds4_gpu_tensor *logits_view = - glm_graph_tensor_row_view_strided(g->batch_router_logits, - t, - DS4_N_EXPERT, - DS4_N_EXPERT); - ds4_gpu_tensor *probs_view = - glm_graph_tensor_row_view_strided(g->batch_router_probs, - t, - DS4_N_EXPERT, - DS4_N_EXPERT); - ds4_gpu_tensor *selected_view = - ds4_gpu_tensor_view(g->batch_router_selected, - (uint64_t)t * DS4_N_EXPERT_USED * sizeof(int32_t), - (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); - ds4_gpu_tensor *weights_view = - glm_graph_tensor_row_view_strided(g->batch_router_weights, - t, - DS4_N_EXPERT_USED, - DS4_N_EXPERT_USED); - ok = logits_view && probs_view && selected_view && weights_view; - if (ok) { - ok = ds4_gpu_glm_router_select_tensor(selected_view, - weights_view, - probs_view, - model->map, - model->size, - l->ffn_exp_probs_b->abs_offset, - logits_view, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE) != 0; - } - ds4_gpu_tensor_free(weights_view); - ds4_gpu_tensor_free(selected_view); - ds4_gpu_tensor_free(probs_view); - ds4_gpu_tensor_free(logits_view); - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_indexed_ffn", - "router", - il, - pos0, - n_tokens, - stage_t0); - if (ok) ok = glm_graph_profile_router_selection_batch(g, - l, - il, - pos0, - n_tokens); - if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, - il, - n_tokens); - metal_graph_debug_dump_tensor("glm_indexed_router_logits", - g->batch_router_logits, - (uint64_t)n_tokens * DS4_N_EXPERT, - il, - pos0); - metal_graph_debug_dump_i32_tensor("glm_indexed_router_selected", - g->batch_router_selected, - (uint64_t)n_tokens * DS4_N_EXPERT_USED, - il, - pos0); - metal_graph_debug_dump_tensor("glm_indexed_router_weights", - g->batch_router_weights, - (uint64_t)n_tokens * DS4_N_EXPERT_USED, - il, - pos0); - - const bool tp_batch_split_ffn2 = g->tp_world == 2; - if (ok && tp_batch_split_ffn2) { - ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); - } - if (ok) { - const bool use_grouped_moe = - glm_graph_indexed_prefill_grouped_moe_default(g); - ok = glm_graph_routed_moe_batch_dispatch( - g, - model, - l, - il, - tp_batch_split_ffn2 ? g->tp_bounce_out : g->batch_ffn_out, - g->batch_ffn_mid, - gate_out * gate_row_bytes, - gate_row_bytes, - up_out * up_row_bytes, - up_row_bytes, - down_out * down_row_bytes, - down_row_bytes, - g->batch_router_selected, - g->batch_router_weights, - g->batch_ffn_norm, - n_tokens, - (uint32_t)g->ffn_mid_elems, - false, - !use_grouped_moe) != 0; - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_indexed_ffn", - "routed_moe", - il, - pos0, - n_tokens, - stage_t0); - metal_graph_debug_dump_tensor("glm_indexed_routed_out", - g->batch_ffn_out, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - - const bool use_batch_residual = - glm_graph_indexed_prefill_batch_residual(); - const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; - if (use_batch_residual && residual_elems > UINT32_MAX) return false; - - bool shared_expert_done = false; - if (ok && - use_batch_residual && - glm_graph_indexed_prefill_batch_shared_expert() && - g->batch_ffn_gate && - g->batch_ffn_up && - g->batch_shared_mid && - glm_graph_weights_are_q8_0(model, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset) && - ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( - g->batch_ffn_gate, - g->batch_ffn_up, - g->batch_shared_mid, - model->map, - model->size, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - g->batch_ffn_norm, - n_tokens, - 0.0f) != 0) { - shared_expert_done = - glm_graph_matmul_q8_0_rows_scalar(g->batch_attn_out, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - g->batch_shared_mid, - n_tokens); + info.valid = true; + info.cp = utf8_peek_one(s, len, pos, &info.next); + info.is_whitespace = glm4_unicode_whitespace(info.cp); + info.is_number = glm4_unicode_number(info.cp); + if (info.cp < 128) { + info.is_letter = ascii_alpha((uint8_t)info.cp); + } else { + info.is_letter = + !info.is_whitespace && + !info.is_number && + !glm4_unicode_punct_symbol(info.cp); } - - for (uint32_t t = 0; ok && !shared_expert_done && t < n_tokens; t++) { - ds4_gpu_tensor *ffn_norm_view = - glm_graph_tensor_row_view_strided(g->batch_ffn_norm, - t, - DS4_N_EMBD, - DS4_N_EMBD); - ds4_gpu_tensor *shared_out_view = use_batch_residual ? - glm_graph_tensor_row_view_strided(g->batch_attn_out, - t, - DS4_N_EMBD, - DS4_N_EMBD) : - NULL; - ds4_gpu_tensor *after_attn_view = !use_batch_residual ? - glm_graph_tensor_row_view_strided((ds4_gpu_tensor *)after_attn, - t, - DS4_N_EMBD, - DS4_N_EMBD) : - NULL; - ds4_gpu_tensor *routed_out_view = !use_batch_residual ? - glm_graph_tensor_row_view_strided(g->batch_ffn_out, - t, - DS4_N_EMBD, - DS4_N_EMBD) : - NULL; - ds4_gpu_tensor *next_view = !use_batch_residual ? - glm_graph_tensor_row_view_strided(next, - t, - DS4_N_EMBD, - DS4_N_EMBD) : - NULL; - ok = ffn_norm_view && - (use_batch_residual ? (shared_out_view != NULL) : - (after_attn_view && routed_out_view && next_view)); - if (ok && glm_graph_weights_are_q8_0(model, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset)) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( - g->ffn_gate, - g->ffn_up, - g->ffn_mid, - model->map, - model->size, - l->ffn_gate_shexp->abs_offset, - l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - ffn_norm_view, - 0.0f) != 0; - } else if (ok) { - ok = glm_graph_matmul_q8_0_tensor(g->ffn_gate, - model, - l->ffn_gate_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - ffn_norm_view, - 1); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->ffn_up, - model, - l->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - ffn_norm_view, - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(g->ffn_mid, - g->ffn_gate, - g->ffn_up, - DS4_N_FF_EXP, - 0.0f, - 1.0f) != 0; - } - if (ok) ok = glm_graph_matmul_q8_0_tensor(use_batch_residual ? - shared_out_view : - g->ffn_sum, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - g->ffn_mid, - 1); - if (ok && !use_batch_residual) { - ok = ds4_gpu_add_tensor(g->attn_out, - routed_out_view, - g->ffn_sum, - DS4_N_EMBD) != 0; - } - if (ok && !use_batch_residual) { - ok = ds4_gpu_add_tensor(next_view, - after_attn_view, - g->attn_out, - DS4_N_EMBD) != 0; - } - - ds4_gpu_tensor_free(next_view); - ds4_gpu_tensor_free(routed_out_view); - ds4_gpu_tensor_free(shared_out_view); - ds4_gpu_tensor_free(ffn_norm_view); - ds4_gpu_tensor_free(after_attn_view); - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_indexed_ffn", - "shared_expert", - il, - pos0, - n_tokens, - stage_t0); - if (ok && use_batch_residual) { - if (!glm_graph_disable_add3_residual()) { - ok = ds4_gpu_add3_tensor(next, - after_attn, - g->batch_ffn_out, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - } else { - ok = ds4_gpu_add_tensor(g->batch_heads, - g->batch_ffn_out, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - if (ok) ok = ds4_gpu_add_tensor(next, - after_attn, - g->batch_heads, - (uint32_t)residual_elems) != 0; - } - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_indexed_ffn", - "residual", - il, - pos0, - n_tokens, - stage_t0); - metal_graph_debug_dump_tensor("glm_indexed_next", - next, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - - (void)pos0; - return ok; + return info; } -static bool glm_graph_upload_tokens( - ds4_gpu_tensor *out_tokens, - const int *tokens, - uint32_t n_tokens) { - if (!out_tokens || !tokens || n_tokens == 0) return false; - - int32_t *ids = xmalloc((size_t)n_tokens * sizeof(ids[0])); - for (uint32_t i = 0; i < n_tokens; i++) ids[i] = (int32_t)tokens[i]; - const bool ok = ds4_gpu_tensor_write(out_tokens, - 0, - ids, - (uint64_t)n_tokens * sizeof(ids[0])) != 0; - free(ids); - return ok; +static uint32_t ascii_tolower_cp(uint32_t cp) { + if (cp >= 'A' && cp <= 'Z') return cp + ('a' - 'A'); + return cp; } -static bool glm_graph_disable_add3_residual(void) { - return false; -} +/* ChatGLM4/GLM pre-tokenization. GLM GGUFs use tokenizer.ggml.pre="glm4", + * which shares the llama3-style split shape used by llama.cpp's CHATGLM4 path. */ +static void bpe_tokenize_text_glm4(const ds4_vocab *vocab, const char *text, token_vec *out) { + const uint64_t len = strlen(text); + uint64_t pos = 0; -static bool glm_graph_encode_ffn_batch( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const ds4_layer_weights *l, - uint32_t il, - uint32_t pos0, - ds4_gpu_tensor *after_attn, - ds4_gpu_tensor *next, - uint32_t n_tokens, - bool full_layer_prefill, - bool stage_profile, - bool stage_sync, - double *stage_t0) { - if (!g || !model || !weights || !l || !after_attn || !next || n_tokens == 0) return false; - - bool ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, - after_attn, - model->map, - model->size, - l->ffn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "ffn_norm", - il, - pos0, - n_tokens, - stage_t0); - if (ok) { - metal_graph_debug_dump_tensor("glm_ffn_norm", - g->batch_ffn_norm, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - } - if (!ok) return false; + while (pos < len) { + uint64_t start = pos; + glm4_char_info cur = glm4_char_at(text, len, pos); - if (il < DS4_N_LEADING_DENSE) { - const uint64_t hidden = l->ffn_gate->dim[1]; - if (hidden == 0 || hidden > UINT32_MAX / n_tokens) return false; - const uint32_t mid_elems = (uint32_t)(hidden * n_tokens); - const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; - if (residual_elems > UINT32_MAX) return false; - - const bool fused_gate_up = glm_graph_shared_gate_up_swiglu_q8_0_tensor( - g->batch_ffn_gate, - g->batch_ffn_up, - g->batch_ffn_mid, - model, - l->ffn_gate->abs_offset, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - g->batch_ffn_norm, - n_tokens, - 0.0f); - if (fused_gate_up) { - ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "dense_gate_up_swiglu", - il, - pos0, - n_tokens, - stage_t0); - } else { - ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_gate, - model, - l->ffn_gate->abs_offset, - DS4_N_EMBD, - hidden, - g->batch_ffn_norm, - n_tokens); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, - model, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - g->batch_ffn_norm, - n_tokens); - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "dense_gate_up", - il, - pos0, - n_tokens, - stage_t0); - if (ok) ok = ds4_gpu_swiglu_tensor(g->batch_ffn_mid, - g->batch_ffn_gate, - g->batch_ffn_up, - mid_elems, - 0.0f, - 1.0f) != 0; - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "dense_swiglu", - il, - pos0, - n_tokens, - stage_t0); - } - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_out, - model, - l->ffn_down->abs_offset, - hidden, - DS4_N_EMBD, - g->batch_ffn_mid, - n_tokens); - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "dense_down", - il, - pos0, - n_tokens, - stage_t0); - if (ok) ok = ds4_gpu_add_tensor(next, - after_attn, - g->batch_ffn_out, - (uint32_t)residual_elems) != 0; - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "residual", - il, - pos0, - n_tokens, - stage_t0); - return ok; - } + if (!cur.valid) break; - if (g->ffn_mid_elems > UINT32_MAX || - g->dense_hidden_max < DS4_N_FF_EXP || - (uint64_t)n_tokens > UINT32_MAX / DS4_N_FF_EXP) { - return false; - } - const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; - if (residual_elems > UINT32_MAX) return false; - - uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; - uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; - uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; - (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); - (void)gate_in; - (void)up_in; - (void)down_in; - - ok = ds4_gpu_matmul_f32_tensor(g->batch_router_logits, - model->map, - model->size, - l->ffn_gate_inp->abs_offset, - DS4_N_EMBD, - DS4_N_EXPERT, - g->batch_ffn_norm, - n_tokens) != 0; - if (ok) ok = ds4_gpu_glm_router_select_batch_tensor(g->batch_router_selected, - g->batch_router_weights, - g->batch_router_probs, - model->map, - model->size, - l->ffn_exp_probs_b->abs_offset, - g->batch_router_logits, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE, - n_tokens) != 0; - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "router", - il, - pos0, - n_tokens, - stage_t0); - if (ok) { - metal_graph_debug_dump_tensor("glm_ffn_router_logits", - g->batch_router_logits, - (uint64_t)n_tokens * DS4_N_EXPERT, - il, - pos0); - metal_graph_debug_dump_tensor("glm_ffn_router_probs", - g->batch_router_probs, - (uint64_t)n_tokens * DS4_N_EXPERT, - il, - pos0); - metal_graph_debug_dump_i32_tensor("glm_ffn_router_selected", - g->batch_router_selected, - (uint64_t)n_tokens * DS4_N_EXPERT_USED, - il, - pos0); - metal_graph_debug_dump_tensor("glm_ffn_router_weights", - g->batch_router_weights, - (uint64_t)n_tokens * DS4_N_EXPERT_USED, - il, - pos0); - } - if (ok) ok = glm_graph_profile_router_selection_batch(g, - l, - il, - pos0, - n_tokens); - const bool tp_batch_split_ffn = g->tp_world == 2; - if (ok && tp_batch_split_ffn) { - ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); - } - if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, - il, - n_tokens); - if (ok) ok = glm_graph_seed_streaming_expert_cache_from_full_layer( - g, - model, - weights, - l, - il, - n_tokens, - gate_out * gate_row_bytes, - down_out * down_row_bytes, - full_layer_prefill); - bool shared_done = false; -#define DS4_GLM_ENCODE_FFN_BATCH_SHARED() do { \ - if (ok) { \ - const bool fused_shared = glm_graph_shared_gate_up_swiglu_q8_0_tensor( \ - g->batch_ffn_gate, \ - g->batch_ffn_up, \ - g->batch_shared_mid, \ - model, \ - l->ffn_gate_shexp->abs_offset, \ - l->ffn_up_shexp->abs_offset, \ - DS4_N_EMBD, \ - DS4_N_FF_EXP, \ - g->batch_ffn_norm, \ - n_tokens, \ - 0.0f); \ - if (fused_shared) { \ - ok = glm_graph_prefill_stage_boundary(stage_profile, \ - stage_sync, \ - "glm_ffn", \ - "shared_gate_up_swiglu", \ - il, \ - pos0, \ - n_tokens, \ - stage_t0); \ - } else { \ - ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_gate, \ - model, \ - l->ffn_gate_shexp->abs_offset, \ - DS4_N_EMBD, \ - DS4_N_FF_EXP, \ - g->batch_ffn_norm, \ - n_tokens); \ - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, \ - model, \ - l->ffn_up_shexp->abs_offset, \ - DS4_N_EMBD, \ - DS4_N_FF_EXP, \ - g->batch_ffn_norm, \ - n_tokens); \ - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ - stage_sync, \ - "glm_ffn", \ - "shared_gate_up", \ - il, \ - pos0, \ - n_tokens, \ - stage_t0); \ - if (ok) ok = ds4_gpu_swiglu_tensor(g->batch_shared_mid, \ - g->batch_ffn_gate, \ - g->batch_ffn_up, \ - n_tokens * DS4_N_FF_EXP, \ - 0.0f, \ - 1.0f) != 0; \ - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ - stage_sync, \ - "glm_ffn", \ - "shared_swiglu", \ - il, \ - pos0, \ - n_tokens, \ - stage_t0); \ - } \ - } \ - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, \ - model, \ - l->ffn_down_shexp->abs_offset, \ - DS4_N_FF_EXP, \ - DS4_N_EMBD, \ - g->batch_shared_mid, \ - n_tokens); \ - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ - stage_sync, \ - "glm_ffn", \ - "shared_down", \ - il, \ - pos0, \ - n_tokens, \ - stage_t0); \ - if (ok) shared_done = true; \ - } while (0) -#ifdef DS4_ROCM_BUILD - rocm_graph_batch_selected_async_load rocm_batch_selected_async = {0}; - bool rocm_batch_selected_async_started = false; - const bool rocm_batch_selected_shared_overlap = - ok && - g->ssd_streaming && - !g->quality && - n_tokens > 1 && - !full_layer_prefill && - !glm_graph_env_present( - "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_SELECTED_ASYNC_LOAD", - "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_SELECTED_ASYNC_LOAD") && - !glm_graph_env_present( - "DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_ASYNC_LOAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_ASYNC_LOAD") && - glm_graph_stream_prefill_expert_addr_supported(weights, l, il, n_tokens); - if (rocm_batch_selected_shared_overlap) { - uint64_t selected_event = 0; - if (ds4_gpu_signal_selected_readback_ready(&selected_event) != 0) { - rocm_batch_selected_async_started = - rocm_graph_batch_selected_async_load_start( - &rocm_batch_selected_async, - g->batch_router_selected, - model, - l, - il, - n_tokens, - selected_event, - gate_out * gate_row_bytes, - down_out * down_row_bytes); + if (cur.cp == '\'' && cur.next < len) { + glm4_char_info next = glm4_char_at(text, len, cur.next); + uint32_t n1 = ascii_tolower_cp(next.cp); + if (n1 == 's' || n1 == 't' || n1 == 'm' || n1 == 'd') { + pos = next.next; + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; + } + if (next.valid && next.next < len) { + glm4_char_info next2 = glm4_char_at(text, len, next.next); + uint32_t n2 = ascii_tolower_cp(next2.cp); + if ((n1 == 'r' && n2 == 'e') || + (n1 == 'v' && n2 == 'e') || + (n1 == 'l' && n2 == 'l')) { + pos = next2.next; + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; + } + } } - } - if (rocm_batch_selected_async_started) { - DS4_GLM_ENCODE_FFN_BATCH_SHARED(); - const bool finish_ok = - rocm_graph_batch_selected_async_load_finish(&rocm_batch_selected_async); - if (!finish_ok) rocm_batch_selected_async_started = false; - } -#endif - if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { /* ablate: keep the gate */ } else - if (ok) ok = glm_graph_routed_moe_batch_dispatch( - g, - model, - l, - il, - tp_batch_split_ffn ? g->tp_bounce_out : g->batch_ffn_out, - g->batch_ffn_mid, - gate_out * gate_row_bytes, - gate_row_bytes, - up_out * up_row_bytes, - up_row_bytes, - down_out * down_row_bytes, - down_row_bytes, - g->batch_router_selected, - g->batch_router_weights, - g->batch_ffn_norm, - n_tokens, - (uint32_t)g->ffn_mid_elems, - full_layer_prefill, - false) != 0; - if (ok && g->tp_world == 2) { - ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_ffn_out, n_tokens); - if (!ok) fprintf(stderr, "ds4: GLM TP batch gate failed (layer %u)\n", il); - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "routed_moe", - il, - pos0, - n_tokens, - stage_t0); - if (ok) { - metal_graph_debug_dump_tensor("glm_ffn_routed_out", - g->batch_ffn_out, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - } - if (ok && !shared_done) DS4_GLM_ENCODE_FFN_BATCH_SHARED(); -#undef DS4_GLM_ENCODE_FFN_BATCH_SHARED - if (ok) { - metal_graph_debug_dump_tensor("glm_ffn_shared_out", - g->batch_attn_out, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - } - if (ok && !glm_graph_disable_add3_residual()) { - ok = ds4_gpu_add3_tensor(next, - after_attn, - g->batch_ffn_out, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - } else if (ok) { - ok = ds4_gpu_add_tensor(g->batch_heads, - g->batch_ffn_out, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - if (ok) ok = ds4_gpu_add_tensor(next, - after_attn, - g->batch_heads, - (uint32_t)residual_elems) != 0; - } - if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, - stage_sync, - "glm_ffn", - "residual", - il, - pos0, - n_tokens, - stage_t0); - if (ok) { - metal_graph_debug_dump_tensor("glm_ffn_next", - next, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - } - return ok; -} + if (!(cur.cp == '\r' || cur.cp == '\n' || cur.is_number)) { + glm4_char_info next = glm4_char_at(text, len, cur.next); + if (cur.is_letter || next.is_letter) { + pos = cur.next; + while (pos < len) { + glm4_char_info scan = glm4_char_at(text, len, pos); + if (!scan.valid || !scan.is_letter) break; + pos = scan.next; + } + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; + } + } -static bool glm_graph_begin_commands_if_needed(void); - -/* ------------------------------------------------------------------------ - * GLM MTP (nextn block) drafting. - * - * blk.(N_LAYER-1) is GLM 5.2's multi-token-prediction block: a full - * attention+MoE layer fed with eh_proj(concat(enorm(embed(token[p+1])), - * hnorm(h[p]))), predicting token[p+2] through the shared output head. - * It keeps a private compact KV cache (slot = absolute position; only - * positions >= mtp_min_pos are ever selected, so the unwritten prompt - * range is never read). Under TP the routed experts are combined over - * the BIG-gate exchange, never the decode row gate, so the RDMA row-gate - * schedule stays intact. - * --------------------------------------------------------------------- */ -static bool glm_graph_mtp_ensure(ds4_glm_gpu_graph *g) { - if (g->mtp_ready) return true; - if (g->compact_cache_cap == 0) return false; - const uint64_t elem = glm_graph_compact_cache_elem_bytes(); - const uint64_t kv_bytes = (uint64_t)g->compact_cache_cap * DS4_N_KV_LORA * elem; - const uint64_t rope_bytes = (uint64_t)g->compact_cache_cap * DS4_N_ROT * elem; - g->mtp_kv_lora_cache = ds4_gpu_tensor_alloc(kv_bytes); - g->mtp_k_rope_cache = ds4_gpu_tensor_alloc(rope_bytes); - g->mtp_concat = ds4_gpu_tensor_alloc(2ull * DS4_N_EMBD * sizeof(float)); - g->mtp_selected = ds4_gpu_tensor_alloc((uint64_t)g->compact_cache_cap * sizeof(int32_t)); - g->mtp_logits_host = malloc((size_t)DS4_N_VOCAB * sizeof(float)); - if (!g->mtp_kv_lora_cache || !g->mtp_k_rope_cache || !g->mtp_concat || - !g->mtp_selected || !g->mtp_logits_host) { - return false; - } - g->mtp_ready = 1; - return true; -} + if (cur.is_number) { + int ndigits = 0; + while (pos < len && ndigits < 3) { + glm4_char_info scan = glm4_char_at(text, len, pos); + if (!scan.valid || !scan.is_number) break; + pos = scan.next; + ndigits++; + } + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; + } -/* One MTP step at (absolute) position pos: consumes the main model's last - * hidden h[pos] (expected in g->cur, pre output-norm) and next_token - * (= token[pos+1]), writes the nextn KV at slot pos, and returns the - * drafted token[pos+2] by greedy argmax. Clobbers the decode scratch. */ -static bool glm_graph_mtp_step( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int next_token, - uint32_t pos, - uint32_t min_pos, - int *draft_out) { - if (!g || !model || !weights || !draft_out) return false; - if (DS4_N_NEXTN_PREDICT == 0) return false; - if (pos >= g->compact_cache_cap || min_pos > pos) { - fprintf(stderr, "ds4: glm mtp: pos %u/min %u out of range (cap %u)\n", - pos, min_pos, g->compact_cache_cap); - return false; - } - const uint32_t il = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; - const ds4_layer_weights *l = &weights->layer[il]; - if (!l->nextn_eh_proj || !l->nextn_enorm || !l->nextn_hnorm || - !l->nextn_shared_head_norm || !l->ffn_gate_exps) { - fprintf(stderr, "ds4: glm mtp: nextn weights missing at layer %u\n", il); - return false; - } - const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - const uint32_t n_selected = pos - min_pos + 1u; - bool input_ready = false; - - if (g->placement) { - const int embedding_tier = g->placement[0]; - const int mtp_tier = g->placement[il + 1u]; - bool handoff_ok = glm_graph_ws_switch(g, embedding_tier, true); - if (handoff_ok) handoff_ok = glm_graph_begin_commands_if_needed(); - if (handoff_ok) { - handoff_ok = ds4_gpu_embed_token_quant_tensor( - g->next, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - (uint32_t)next_token, - DS4_N_EMBD) != 0; + glm4_char_info punct = cur; + uint64_t punct_pos = pos; + if (cur.cp == ' ') { + punct_pos = cur.next; + punct = glm4_char_at(text, len, punct_pos); } - if (handoff_ok) handoff_ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - ds4_gpu_tensor *embedded_next = handoff_ok ? g->next : NULL; - if (handoff_ok) handoff_ok = glm_graph_ws_switch(g, mtp_tier, true); - if (handoff_ok) { - handoff_ok = ds4_gpu_tensor_copy_async( - g->next, embedded_next, - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - } - if (!handoff_ok) { - fprintf(stderr, "ds4: glm mtp: multi-tier input handoff failed\n"); - return false; + if (punct.valid && + !punct.is_whitespace && + !punct.is_letter && + !punct.is_number) { + pos = punct_pos; + while (pos < len) { + glm4_char_info scan = glm4_char_at(text, len, pos); + if (!scan.valid || + scan.is_whitespace || + scan.is_letter || + scan.is_number) { + break; + } + pos = scan.next; + } + while (pos < len) { + glm4_char_info scan = glm4_char_at(text, len, pos); + if (!scan.valid || !(scan.cp == '\r' || scan.cp == '\n')) break; + pos = scan.next; + } + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; + } + + if (cur.is_whitespace) { + uint64_t p = pos; + uint64_t last_newline_end = 0; + uint64_t last_ws_start = pos; + int nspace = 0; + while (p < len) { + glm4_char_info scan = glm4_char_at(text, len, p); + if (!scan.valid || !scan.is_whitespace) break; + last_ws_start = p; + if (scan.cp == '\r' || scan.cp == '\n') last_newline_end = scan.next; + p = scan.next; + nspace++; + } + if (last_newline_end) { + pos = last_newline_end; + } else if (nspace > 1 && p < len) { + pos = last_ws_start; + } else { + pos = p; + } + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); + continue; } - input_ready = true; - } - if (!glm_graph_mtp_ensure(g)) { - fprintf(stderr, "ds4: glm mtp: ensure failed (cap %u)\n", g->compact_cache_cap); - return false; - } - /* Draft attention window: absolute cache slots [min_pos..pos]. */ - { - int32_t *sel = malloc((size_t)n_selected * sizeof(int32_t)); - if (!sel) return false; - for (uint32_t i = 0; i < n_selected; i++) sel[i] = (int32_t)(min_pos + i); - const int wr = ds4_gpu_tensor_write(g->mtp_selected, 0, sel, - (uint64_t)n_selected * sizeof(int32_t)); - free(sel); - if (!wr) { - fprintf(stderr, "ds4: glm mtp: selected write failed (%u)\n", n_selected); - return false; - } + pos = cur.next; + if (pos == start) pos = next_utf8_char(text, len, pos); + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); } +} - ds4_gpu_tensor *enorm_view = - ds4_gpu_tensor_view(g->mtp_concat, 0, (uint64_t)DS4_N_EMBD * sizeof(float)); - ds4_gpu_tensor *hnorm_view = - ds4_gpu_tensor_view(g->mtp_concat, - (uint64_t)DS4_N_EMBD * sizeof(float), - (uint64_t)DS4_N_EMBD * sizeof(float)); - if (!enorm_view || !hnorm_view) { - ds4_gpu_tensor_free(enorm_view); - ds4_gpu_tensor_free(hnorm_view); - fprintf(stderr, "ds4: glm mtp: concat views failed\n"); - return false; +/* + * DeepSeek V4 Flash declares tokenizer.ggml.pre = "joyai-llm". The split + * below mirrors the JoyAI BPE pre-tokenizer for the cases this model + * uses in normal text and source-code prompts: + * + * \p{N}{1,3} + * [CJK/Hiragana/Katakana]+ + * [P/S][A-Za-z]+ + * [^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+ + * ?[\p{P}\p{S}]+[\r\n]* + * \s*[\r\n]+ + * \s+(?!\S) + * \s+ + * + * The punctuation rule intentionally keeps trailing newlines in the same BPE + * word (for example ">;\n"). Splitting those newlines separately changes the + * token stream for code prompts and produces wrong long-context logits. + */ +static void bpe_tokenize_text(const ds4_vocab *vocab, const char *text, token_vec *out) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + bpe_tokenize_text_glm4(vocab, text, out); + return; } - const char *mtp_stage = "begin"; -#define DS4_GLM_MTP_STAGE(name_) do { if (ok) mtp_stage = (name_); } while (0) - bool ok = glm_graph_begin_commands_if_needed(); - /* MTP input: concat(enorm(embed(next_token)), hnorm(h)) -> eh_proj. */ - if (ok && !input_ready) { - ok = ds4_gpu_embed_token_quant_tensor(g->next, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - (uint32_t)next_token, - DS4_N_EMBD) != 0; - } - DS4_GLM_MTP_STAGE("enorm"); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(enorm_view, - g->next, - model->map, - model->size, - l->nextn_enorm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("hnorm"); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(hnorm_view, - g->cur, - model->map, - model->size, - l->nextn_hnorm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("eh_proj"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->cur, - model, - l->nextn_eh_proj->abs_offset, - 2ull * DS4_N_EMBD, - DS4_N_EMBD, - g->mtp_concat, - false); - /* nextn attention (full causal over the MTP window, no indexer). */ - DS4_GLM_MTP_STAGE("attn_norm"); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->attn_norm, - g->cur, - model->map, - model->size, - l->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("q_a"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->attn_norm, - false); - DS4_GLM_MTP_STAGE("q_a_norm"); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->q_rank_norm, - g->q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("q_b"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->q_rank_norm, - false); - DS4_GLM_MTP_STAGE("rope"); - if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->q, - 1, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_ROT, - pos, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_GLM_MTP_STAGE("kv_a"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - false); - DS4_GLM_MTP_STAGE("kv_norm"); - if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, - g->kv_raw, - model->map, - model->size, - l->attn_kv_a_norm->abs_offset, - 1, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("kv_store"); - if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(g->mtp_kv_lora_cache, - g->mtp_k_rope_cache, - g->kv_norm, - g->kv_raw, - pos, - 1, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16()) != 0; - DS4_GLM_MTP_STAGE("qk_low"); - if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor(g->qk_low, - g->q, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_KEY_MLA) != 0; - DS4_GLM_MTP_STAGE("attention"); - if (ok) ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor(g->heads, - g->q, - g->qk_low, - g->mtp_kv_lora_cache, - g->mtp_k_rope_cache, - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - g->mtp_selected, - n_selected, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_GLM_MTP_STAGE("attn_out"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->heads, - false); - DS4_GLM_MTP_STAGE("ffn_norm"); - if (ok) ok = ds4_gpu_add_rms_norm_weight_tensor(g->ffn_norm, - g->after_attn, - g->cur, - g->attn_out, - model->map, - model->size, - l->ffn_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - /* nextn sparse FFN: router + split routed experts (BIG-gate combine - * under TP) + shared expert. */ - DS4_GLM_MTP_STAGE("router"); - if (ok) ok = ds4_gpu_matmul_f32_tensor(g->router_logits, - model->map, - model->size, - l->ffn_gate_inp->abs_offset, - DS4_N_EMBD, - DS4_N_EXPERT, - g->ffn_norm, - 1) != 0; - DS4_GLM_MTP_STAGE("router_select"); - if (ok) ok = ds4_gpu_glm_router_select_tensor(g->router_selected, - g->router_weights, - g->router_probs, - model->map, - model->size, - l->ffn_exp_probs_b->abs_offset, - g->router_logits, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE) != 0; - DS4_GLM_MTP_STAGE("routed"); - if (ok) { - uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; - uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; - uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; - (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); - (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); -#ifdef DS4_ROCM_BUILD - if (g->ssd_streaming) { - const ds4_gpu_stream_expert_table table = { - .model_map = model->map, - .model_size = model->size, - .layer = il, - .n_total_expert = DS4_N_EXPERT, - .gate_offset = l->ffn_gate_exps->abs_offset, - .up_offset = l->ffn_up_exps->abs_offset, - .down_offset = l->ffn_down_exps->abs_offset, - .gate_expert_bytes = gate_out * gate_row_bytes, - .down_expert_bytes = down_out * down_row_bytes, - }; - ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( - &table, g->router_selected, DS4_N_EXPERT_USED) != 0; - } -#endif - const bool tp_split = g->tp_world == 2 && g->tp_out && g->tp_in; - ds4_gpu_tensor *routed_dst = g->ffn_out; - if (tp_split) { - ok = glm_graph_tp_batch_bounce_ready(g, 1); - routed_dst = g->tp_bounce_out; - } - if (ok) ok = glm_graph_routed_moe_one_dispatch(g, - model, - l, - il, - routed_dst, - g->ffn_mid, - gate_out * gate_row_bytes, - gate_row_bytes, - up_out * up_row_bytes, - up_row_bytes, - down_out * down_row_bytes, - down_row_bytes, - g->router_selected, - g->router_weights, - g->ffn_norm, - false) != 0; - if (ok && tp_split) { - ok = glm_graph_tp_batch_ffn_combine(g, il, g->ffn_out, 1); - } - } - DS4_GLM_MTP_STAGE("shared"); - if (ok) ok = glm_graph_encode_shared_swiglu_one(g->ffn_mid, - g->ffn_gate, - g->ffn_up, - model, - l, - il, - pos, - g->ffn_norm, - false, - false, - NULL); - DS4_GLM_MTP_STAGE("shared_down"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->ffn_sum, - model, - l->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - g->ffn_mid, - false); - DS4_GLM_MTP_STAGE("residual"); - if (ok) ok = ds4_gpu_add3_tensor(g->next, - g->after_attn, - g->ffn_out, - g->ffn_sum, - DS4_N_EMBD) != 0; - /* Shared output head behind the nextn head norm. */ - DS4_GLM_MTP_STAGE("head_norm"); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, - g->next, - model->map, - model->size, - l->nextn_shared_head_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_MTP_STAGE("head"); - if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->logits, - model, - weights->output->abs_offset, - DS4_N_EMBD, - DS4_N_VOCAB, - g->output_norm, - false); - DS4_GLM_MTP_STAGE("end"); - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok) { - ok = ds4_gpu_tensor_read(g->logits, - 0, - g->mtp_logits_host, - (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - } - ds4_gpu_tensor_free(enorm_view); - ds4_gpu_tensor_free(hnorm_view); - if (!ok) { - fprintf(stderr, "ds4: glm mtp step failed at stage '%s' (pos %u)\n", - mtp_stage, pos); - return false; - } -#undef DS4_GLM_MTP_STAGE - int best = 0; - float best_v = g->mtp_logits_host[0]; - for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { - if (g->mtp_logits_host[i] > best_v) { - best_v = g->mtp_logits_host[i]; - best = (int)i; + const uint64_t len = strlen(text); + uint64_t pos = 0; + + while (pos < len) { + uint64_t start = pos; + uint8_t c = (uint8_t)text[pos]; + + if (ascii_digit(c)) { + int ndigits = 0; + while (pos < len && ascii_digit((uint8_t)text[pos]) && ndigits < 3) { + pos++; + ndigits++; + } + } else if (joyai_cjk_at(text, len, pos)) { + do { + pos = next_utf8_char(text, len, pos); + } while (pos < len && joyai_cjk_at(text, len, pos)); + } else if (joyai_ascii_punct_symbol(c) && + pos + 1 < len && + ascii_alpha((uint8_t)text[pos + 1])) { + pos++; + while (pos < len && ascii_alpha((uint8_t)text[pos])) pos++; + } else if (joyai_letter_like_at(text, len, pos)) { + pos = joyai_consume_letters(text, len, pos); + } else if (!ascii_newline(c) && + !joyai_ascii_punct_symbol(c) && + pos + 1 < len && + joyai_letter_like_at(text, len, pos + 1)) { + pos++; + pos = joyai_consume_letters(text, len, pos); + } else if (c == ' ' && + pos + 1 < len && + joyai_ascii_punct_symbol((uint8_t)text[pos + 1])) { + pos++; + while (pos < len && joyai_ascii_punct_symbol((uint8_t)text[pos])) pos++; + while (pos < len && ascii_newline((uint8_t)text[pos])) pos++; + } else if (joyai_ascii_punct_symbol(c)) { + while (pos < len && joyai_ascii_punct_symbol((uint8_t)text[pos])) pos++; + while (pos < len && ascii_newline((uint8_t)text[pos])) pos++; + } else if (ascii_space(c)) { + uint64_t p = pos; + uint64_t last_newline_end = 0; + while (p < len && ascii_space((uint8_t)text[p])) { + uint8_t sc = (uint8_t)text[p++]; + if (ascii_newline(sc)) last_newline_end = p; + } + if (last_newline_end) { + pos = last_newline_end; + } else if (p < len && p > pos + 1 && + (joyai_letter_like_at(text, len, p) || + joyai_ascii_punct_symbol((uint8_t)text[p]))) { + /* + * JoyAI lets a single leading space join the following word or + * punctuation run. For " int", the pre-tokenizer therefore emits + * " " then " int", not " " then "int". + */ + pos = p - 1; + } else { + pos = p; + } + } else { + pos = next_utf8_char(text, len, pos); } + + if (pos == start) pos = next_utf8_char(text, len, pos); + bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); } - *draft_out = best; - return true; } -static bool glm_graph_forward_token( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - const float *input_hc, - uint32_t pos, - float *output_hc, - float *logits_out, - bool defer_completion); - - -/* Decode-style verify pass for tiny row counts (MTP): the indexed batch - * fn measures ~1.4ms/layer at n=2 (gate-profile: gpu-wait 1.22ms/layer) - * while decode does the same math in 0.79ms. This pass mirrors the - * decode encoders at n rows over the batch scratch, attends causally - * over the compact caches (valid while pos+n fits the indexer window), - * and reuses the batch FFN encoder (routed split + big-gate combine). - * KV/indexer caches are updated exactly like the indexed batch path. */ -static bool glm_graph_verify_rows( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - uint32_t pos, - uint32_t n, - float *output_hc, - float *logits_out) { - if (!g || !model || !weights || !tokens || n == 0 || - g->compact_cache_cap == 0 || - pos + n > g->compact_cache_cap || - !g->batch_cur || !g->batch_next || !g->prefill_tokens) { - return false; - } - const uint32_t executable = glm_graph_normal_layer_count(); - ds4_gpu_tensor *cur = g->batch_cur; - ds4_gpu_tensor *nxt = g->batch_next; - if (!ds4_gpu_tensor_write(g->prefill_tokens, 0, tokens, - (uint64_t)n * sizeof(int32_t))) { - return false; - } - if (g->placement && - !glm_graph_verify_ws_switch(g, g->placement[0], false, n)) { - glm_graph_verify_ws_restore(g); - return false; - } - cur = g->batch_cur; - nxt = g->batch_next; - bool ok = glm_graph_begin_commands_if_needed(); - if (ok) ok = ds4_gpu_embed_tokens_quant_tensor(cur, - g->prefill_tokens, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - n, - DS4_N_EMBD) != 0; - for (uint32_t il = 0; ok && il < executable; il++) { - if (g->placement) { - g->batch_cur = cur; - g->batch_next = nxt; - ok = glm_graph_verify_ws_switch(g, - g->placement[il + 1u], - il != 0u, - n); - if (ok) { - ok = glm_graph_ws_switch(g, - g->placement[il + 1u], - false); - } - cur = g->batch_cur; - nxt = g->batch_next; - } - const ds4_layer_weights *l = &weights->layer[il]; - const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, - cur, - model->map, - model->size, - l->attn_norm->abs_offset, - DS4_N_EMBD, - n, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->batch_attn_norm, - n); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, - g->batch_q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - n, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->batch_q_rank_norm, - n); - if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->batch_q, - n, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_ROT, - pos, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok && glm_graph_layer_uses_full_indexer(il)) { - ok = glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, - model, - l->indexer_attn_k->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD_DIM, - cur, - n); - if (ok) ok = ds4_gpu_glm_store_indexer_k_tensor( - g->layer_indexer_key_cache[il], - g->batch_indexer_k, - model->map, - model->size, - l->indexer_k_norm->abs_offset, - l->indexer_k_norm_b->abs_offset, - pos, - n, - g->compact_cache_cap, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - 0, - 1.0e-6f, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - glm_graph_compact_cache_is_f16()) != 0; - } - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->batch_attn_norm, - n); - if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, - g->batch_kv_raw, - model->map, - model->size, - l->attn_kv_a_norm->abs_offset, - n, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - g->batch_kv_norm, - g->batch_kv_raw, - pos, - n, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16()) != 0; - if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_batch_tensor(g->batch_qk_low, - g->batch_q, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - n, - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_KEY_MLA) != 0; - if (ok) ok = ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( - g->batch_attn_lora, - g->batch_q, - g->batch_qk_low, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - n, - pos, - pos + n, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_glm_value_project_typed_batch_heads_tensor( - g->batch_heads, - g->batch_attn_lora, - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - n, - DS4_N_HEAD, - DS4_N_KV_LORA, - DS4_N_VALUE_MLA) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->batch_heads, - n); - if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, - cur, - g->batch_attn_out, - (uint32_t)((uint64_t)n * DS4_N_EMBD)) != 0; - if (ok) ok = glm_graph_encode_ffn_batch(g, - model, - weights, - l, - il, - pos, - g->batch_after_attn, - nxt, - n, - false, - false, - false, - NULL); - if (ok) { - ds4_gpu_tensor *tmp = cur; - cur = nxt; - nxt = tmp; - } +static int vocab_lookup(const ds4_vocab *vocab, const char *text) { + int token = -1; + if (!table_get(&vocab->token_to_id, text, strlen(text), &token)) { + fprintf(stderr, "ds4: required tokenizer token is missing: %s\n", text); + exit(1); } - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - if (ok && output_hc) { - ok = ds4_gpu_tensor_read(cur, - 0, - output_hc, - (uint64_t)n * DS4_N_EMBD * sizeof(float)) != 0; - } - if (ok && logits_out) { - ds4_gpu_tensor *last = glm_graph_tensor_row_view_strided(cur, - n - 1u, - DS4_N_EMBD, - DS4_N_EMBD); - ok = last != NULL && - glm_graph_forward_output_head(g, model, weights, last, logits_out); - ds4_gpu_tensor_free(last); - } - glm_graph_verify_ws_restore(g); - return ok; + return token; +} + +static int vocab_lookup_optional(const ds4_vocab *vocab, const char *text) { + int token = -1; + if (!table_get(&vocab->token_to_id, text, strlen(text), &token)) return -1; + return token; } -static bool glm_graph_forward_tokens( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - const float *input_hc, - uint32_t pos0, - uint32_t n_tokens, - float *output_hc, - float *logits_out, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t display_absolute_base, - uint32_t work_done_base, - uint32_t work_total) { - if (!g || !model || !weights || !tokens || - n_tokens == 0 || - g->layer_count == 0 || - !glm_graph_span_fits_context(g, pos0, n_tokens)) { - return false; - } - if (!glm_graph_span_fits_full_attention(g, pos0, n_tokens)) { - glm_graph_log_full_attention_limit(g, pos0, n_tokens); - return false; - } - if (!g->full_kv_cache) { - fprintf(stderr, - "ds4: GLM full-attention prefill was requested without an expanded KV cache\n"); - return false; +/* Load token strings, special token ids, and merge ranks from GGUF metadata. */ + +static void vocab_load(ds4_vocab *vocab, const ds4_model *model) { + memset(vocab, 0, sizeof(*vocab)); + + ds4_array_ref tokens; + ds4_array_ref merges; + if (!model_get_array(model, "tokenizer.ggml.tokens", &tokens) || + tokens.type != GGUF_VALUE_STRING || + tokens.len > INT32_MAX) { + ds4_die("GGUF tokenizer token table is missing or invalid"); } - for (uint32_t i = 0; i < n_tokens; i++) { - if (tokens[i] < 0 || tokens[i] >= (int)DS4_N_VOCAB) return false; - } - if (!input_hc && !g->has_token_embd) return false; - if (logits_out && !g->has_output_head) return false; - glm_graph_reset_prefill_seed_capture(g); - const uint32_t n_rows = pos0 + n_tokens; - const bool trace = glm_graph_full_prefill_trace_enabled(); - const bool trace_all = trace && glm_graph_full_prefill_trace_all(); - const double trace_slow_ms = trace ? - (double)glm_graph_full_prefill_trace_slow_ms() : 0.0; - const double trace_chunk_t0 = trace ? now_sec() : 0.0; - if (trace) { - glm_graph_full_prefill_tracef( - "chunk begin pos=%u tokens=%u rows=%u compact_cap=%u work_base=%u work_total=%u", - pos0, - n_tokens, - n_rows, - g->compact_cache_cap, - work_done_base, - work_total); - } - if (g->compact_cache_cap != 0) { - const double trace_cache_t0 = trace ? now_sec() : 0.0; - if (!glm_graph_ensure_compact_cache(g, n_rows)) { - if (trace) { - glm_graph_full_prefill_tracef( - "ensure_cache failed pos=%u tokens=%u rows=%u compact_cap=%u", - pos0, - n_tokens, - n_rows, - g->compact_cache_cap); - } - return false; - } - if (trace) { - const double ms = (now_sec() - trace_cache_t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms) { - glm_graph_full_prefill_tracef( - "ensure_cache done pos=%u tokens=%u rows=%u compact_cap=%u %.3f ms", - pos0, - n_tokens, - n_rows, - g->compact_cache_cap, - ms); - } - } + if (!model_get_array(model, "tokenizer.ggml.merges", &merges) || + merges.type != GGUF_VALUE_STRING) { + ds4_die("GGUF tokenizer merge table is missing or invalid"); } - const double trace_upload_t0 = trace ? now_sec() : 0.0; - bool ok = glm_graph_upload_tokens(g->prefill_tokens, tokens, n_tokens); - if (trace) { - const double ms = (now_sec() - trace_upload_t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_full_prefill_tracef( - "upload_tokens %s pos=%u tokens=%u %.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } - } - ds4_gpu_tensor *cur = g->batch_cur; - ds4_gpu_tensor *next = g->batch_next; - ds4_gpu_tensor *last_hidden = NULL; + vocab->n_vocab = (int)tokens.len; + vocab->token = xcalloc((size_t)vocab->n_vocab, sizeof(vocab->token[0])); + table_init(&vocab->token_to_id, tokens.len); - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - 0, - g->layer_count, - work_total, - true); - - const bool stage_sync = - glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); - const uint32_t layer_flush_interval = stage_sync ? 0u : - glm_graph_full_prefill_layer_flush_interval(n_tokens, - n_rows, - logits_out != NULL); - const uint32_t progress_flush_interval = - glm_graph_prefill_progress_flush_interval(layer_flush_interval, - n_tokens, - display_progress, - work_total); - const bool progress_requested = display_progress && work_total > 0; - const uint32_t drain_interval = - (progress_requested && progress_flush_interval != 0) ? - glm_graph_full_prefill_drain_interval() : 0u; - metal_graph_stream_prepare_slot layer_prepare_slots[DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD]; - memset(layer_prepare_slots, 0, sizeof(layer_prepare_slots)); - const bool full_layer_prefill = - glm_graph_stream_prefill_full_layer_enabled(g, n_tokens); - ds4_gpu_set_glm_streaming_prefill_full_layer(full_layer_prefill); - const bool streaming_prefill_sync_each_layer = - !g->ssd_streaming || - glm_graph_streaming_prefill_sync_each_layer(full_layer_prefill); -#ifdef DS4_ROCM_BUILD - rocm_graph_stream_layer_expert_load rocm_full_layer_load; - memset(&rocm_full_layer_load, 0, sizeof(rocm_full_layer_load)); -#endif - const bool full_layer_prepare_base = - glm_graph_stream_prefill_full_layer_prepare_enabled(g, - full_layer_prefill); - const bool layer_pagein = - full_layer_prepare_base && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN"); - const bool layer_readahead = - full_layer_prepare_base && - !layer_pagein && - glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); - const bool layer_pread = - full_layer_prepare_base && - !layer_pagein && - !layer_readahead && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); - const bool layer_madvise = - full_layer_prepare_base && - !layer_pagein && - !layer_pread && - !layer_readahead && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") && - !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", - "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE"); - const bool layer_prepare = - layer_pagein || layer_pread || layer_readahead || layer_madvise; - const bool layer_prepare_overlap = - layer_prepare && - metal_graph_stream_prefill_layer_pagein_overlap_enabled(); - const bool full_layer_flush_intermediate = full_layer_prefill; - const uint32_t layer_prepare_ahead = - layer_prepare && layer_prepare_overlap ? - metal_graph_stream_prefill_layer_prepare_ahead() : 1u; - if (trace) { - glm_graph_full_prefill_tracef( - "mode pos=%u tokens=%u stage_sync=%u layer_flush_interval=%u progress_flush_interval=%u drain_interval=%u", - pos0, - n_tokens, - stage_sync ? 1u : 0u, - layer_flush_interval, - progress_flush_interval, - drain_interval); - } - if (ok && layer_prepare && g->layer_count > 0 && - !metal_graph_stream_prepare_start_if_needed(NULL, - model, - weights, - g->layer_start, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - full_layer_prefill && - rocm_graph_glm_stream_prefill_full_layer_enabled( - g, - &weights->layer[g->layer_start], - g->layer_start, - n_tokens), - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; + ds4_cursor c = cursor_at(model, tokens.data_pos); + for (int i = 0; i < vocab->n_vocab; i++) { + if (!cursor_string(&c, &vocab->token[i])) ds4_die(c.error); + table_put(&vocab->token_to_id, vocab->token[i], i); } -#ifdef DS4_ROCM_BUILD - if (ok && - full_layer_prefill && - !rocm_graph_glm_stream_layer_expert_load_start_next( - &rocm_full_layer_load, - g, - model, - weights, - g->layer_start, - g->layer_end, - n_tokens)) { - ok = false; + + table_init(&vocab->merge_rank, merges.len); + c = cursor_at(model, merges.data_pos); + for (uint64_t i = 0; i < merges.len; i++) { + ds4_str merge; + if (!cursor_string(&c, &merge)) ds4_die(c.error); + table_put(&vocab->merge_rank, merge, (int)i); } -#endif - if (ok) { - const double t0 = trace ? now_sec() : 0.0; - if (input_hc) { - ok = ds4_gpu_tensor_write(cur, - 0, - input_hc, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; - } else { - ok = glm_graph_stream_map_token(g, model, weights); - } - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (trace) { - const double ms = (now_sec() - t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_full_prefill_tracef( - "begin_commands%s %s pos=%u tokens=%u %.3f ms", - input_hc ? "_from_hidden" : "", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } + + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (!model_get_token_id(model, "tokenizer.ggml.bos_token_id", &vocab->bos_id)) { + vocab->bos_id = vocab_lookup_optional(vocab, ""); } - } - if (ok && !input_hc) { - const double t0 = trace ? now_sec() : 0.0; - ok = ds4_gpu_embed_tokens_quant_tensor(cur, - g->prefill_tokens, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - n_tokens, - DS4_N_EMBD) != 0; - if (trace) { - const double ms = (now_sec() - t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_full_prefill_tracef( - "embed %s pos=%u tokens=%u %.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } + if (!model_get_token_id(model, "tokenizer.ggml.eos_token_id", &vocab->eos_id)) { + vocab->eos_id = vocab_lookup_optional(vocab, "<|endoftext|>"); } + vocab->system_id = vocab_lookup_optional(vocab, "<|system|>"); + vocab->user_id = vocab_lookup_optional(vocab, "<|user|>"); + vocab->assistant_id = vocab_lookup_optional(vocab, "<|assistant|>"); + vocab->observation_id = vocab_lookup_optional(vocab, "<|observation|>"); + vocab->sop_id = vocab_lookup_optional(vocab, ""); + vocab->think_start_id = vocab_lookup_optional(vocab, ""); + vocab->think_end_id = vocab_lookup_optional(vocab, ""); + vocab->tool_call_start_id = vocab_lookup_optional(vocab, ""); + vocab->tool_call_end_id = vocab_lookup_optional(vocab, ""); + vocab->tool_response_start_id = vocab_lookup_optional(vocab, ""); + vocab->tool_response_end_id = vocab_lookup_optional(vocab, ""); + vocab->arg_key_start_id = vocab_lookup_optional(vocab, ""); + vocab->arg_key_end_id = vocab_lookup_optional(vocab, ""); + vocab->arg_value_start_id = vocab_lookup_optional(vocab, ""); + vocab->arg_value_end_id = vocab_lookup_optional(vocab, ""); + vocab->dsml_id = -1; + return; } - if (ok && g->ssd_streaming && streaming_prefill_sync_each_layer) { - ok = ds4_gpu_end_commands() != 0; - } -#define DS4_GLM_PROFILE_PREFILL_STAGE(part_, name_) do { \ - if (ok && layer_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos0, n_tokens, &layer_stage_t0); \ - } else if (ok && stage_sync) { \ - ok = glm_graph_prefill_stage_sync_boundary(); \ - } \ - } while (0) - for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { - const uint32_t slice_layer_done = il - g->layer_start + 1u; - if (layer_prepare && - !metal_graph_stream_prepare_join_layer(NULL, - model, - weights, - il, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - full_layer_prefill && - rocm_graph_glm_stream_prefill_full_layer_enabled( - g, - &weights->layer[il], - il, - n_tokens), - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - break; - } -#ifdef DS4_ROCM_BUILD - if (full_layer_prefill && - !rocm_graph_glm_stream_layer_expert_load_ready( - &rocm_full_layer_load, - g, - model, - weights, - il, - n_tokens)) { - ok = false; - break; - } - if (full_layer_prefill && - !rocm_graph_glm_stream_layer_expert_load_start_next( - &rocm_full_layer_load, - g, - model, - weights, - il + 1u, - g->layer_end, - n_tokens)) { - ok = false; - break; - } -#endif - if (g->ssd_streaming) { - ok = glm_graph_stream_map_prefill_layer(g, - model, - weights, - il, - n_tokens, - full_layer_prefill); - if (ok && layer_prepare && layer_prepare_overlap) { - bool started_future = false; - for (uint32_t ahead = 1; ahead <= layer_prepare_ahead; ahead++) { - if (il + ahead > g->layer_end) break; - started_future = true; - if (!metal_graph_stream_prepare_start_if_needed(NULL, - model, - weights, - il + ahead, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - full_layer_prefill && - rocm_graph_glm_stream_prefill_full_layer_enabled( - g, - &weights->layer[il + ahead], - il + ahead, - n_tokens), - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - break; - } - } - if (ok && !started_future && logits_out) { - metal_graph_stream_readahead_output(model, weights); - } - } - if (ok && !ds4_gpu_commands_active()) { - ok = ds4_gpu_begin_commands() != 0; - } - } - const ds4_layer_weights *l = &weights->layer[il]; - const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - const uint32_t cache_len = pos0 + n_tokens; - const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; - const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); - double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; - const double trace_layer_t0 = trace ? now_sec() : 0.0; - bool trace_layer_flushed = false; - if (trace && trace_all) { - glm_graph_full_prefill_tracef( - "layer begin layer=%u pos=%u tokens=%u rows=%u", - il, - pos0, - n_tokens, - n_rows); - } - if (residual_elems > UINT32_MAX) { - ok = false; - break; - } - if (layer_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_attn", - NULL, - il, - pos0, - n_tokens, - &layer_stage_t0); - } - - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, - cur, - model->map, - model->size, - l->attn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attn_norm"); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->batch_attn_norm, - n_tokens); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, - g->batch_q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->batch_q_rank_norm, - n_tokens); - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, - n_tokens, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_ROT, - pos0, - 0, - false, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "q_path"); - if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il)) { - ok = glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, - model, - l->indexer_attn_k->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD_DIM, - cur, - n_tokens); - if (ok) { - ok = ds4_gpu_glm_store_indexer_k_tensor( - g->layer_indexer_key_cache[il], - g->batch_indexer_k, - model->map, - model->size, - l->indexer_k_norm->abs_offset, - l->indexer_k_norm_b->abs_offset, - pos0, - n_tokens, - g->compact_cache_cap, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - 0, - 1.0e-6f, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - glm_graph_compact_cache_is_f16()) != 0; - } - } - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "indexer_k"); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->batch_attn_norm, - n_tokens); - if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, - g->batch_kv_raw, - model->map, - model->size, - l->attn_kv_a_norm->abs_offset, - n_tokens, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_RMS_EPS) != 0; - if (ok && g->compact_cache_cap != 0) { - ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - g->batch_kv_norm, - g->batch_kv_raw, - pos0, - n_tokens, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16()) != 0; - } - if (ok) ok = ds4_gpu_glm_k_b_project_typed_tensor(g->batch_k_nope, - g->batch_kv_norm, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - n_tokens, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_HEAD) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_value, - model, - l->attn_v_b->abs_offset, - DS4_N_KV_LORA, - g->heads_dim, - g->batch_kv_norm, - n_tokens); - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_path"); - const bool flash_requested = glm_graph_use_flash_attention_prefill(n_tokens); - const bool use_staged_flash_kv = - flash_requested && glm_graph_use_flash_attention_staged_kv(pos0, n_tokens, cache_len); - const bool use_flash_attn = flash_requested; - if (ok) { - if (use_staged_flash_kv) { - ok = ds4_gpu_glm_build_kv_cache_flash_tensor(g->layer_key_cache[il], - g->layer_value_cache[il], - g->batch_kv_raw, - g->batch_k_nope, - g->batch_value, - pos0, - n_tokens, - g->ctx_cap, - DS4_N_HEAD, - kv_raw_dim, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - true) != 0; - } else { - ok = ds4_gpu_glm_build_kv_cache_tensor(g->layer_key_cache[il], - g->layer_value_cache[il], - g->batch_kv_raw, - g->batch_k_nope, - g->batch_value, - pos0, - n_tokens, - g->ctx_cap, - DS4_N_HEAD, - kv_raw_dim, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - true) != 0; - } - } - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_cache"); - if (ok) { - if (use_flash_attn) { - if (use_staged_flash_kv) { - ok = ds4_gpu_glm_attention_flash_staged_tensor(g->batch_heads, - g->batch_q, - g->layer_key_cache[il], - g->layer_value_cache[il], - pos0, - n_tokens, - cache_len, - g->ctx_cap, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_VALUE_MLA, - true) != 0; - } else { - ok = ds4_gpu_glm_attention_flash_tensor(g->batch_heads, - g->batch_q, - g->layer_key_cache[il], - g->layer_value_cache[il], - pos0, - n_tokens, - cache_len, - g->ctx_cap, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_VALUE_MLA, - true) != 0; - } - } else { - ok = ds4_gpu_glm_attention_full_tensor(g->batch_heads, - g->batch_q, - g->layer_key_cache[il], - g->layer_value_cache[il], - pos0, - n_tokens, - cache_len, - g->ctx_cap, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_VALUE_MLA, - true) != 0; - } - } - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attention"); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->batch_heads, - n_tokens); - if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, - cur, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attn_output"); - if (ok) ok = glm_graph_encode_ffn_batch(g, - model, - weights, - l, - il, - pos0, - g->batch_after_attn, - next, - n_tokens, - full_layer_prefill, - layer_stage_profile, - stage_sync, - layer_stage_profile ? &layer_stage_t0 : NULL); - if (ok) { - ds4_gpu_tensor *tmp = cur; - cur = next; - next = tmp; - } - if (ok && glm_debug_hidden_dump_layer_match(il)) { - ok = ds4_gpu_end_commands() != 0; - if (ok) { - for (uint32_t r = 0; r < n_tokens; r++) - glm_debug_dump_hidden_layer(cur, r, il, pos0 + r); - glm_debug_dump_raw_layer(g->batch_router_selected, "sel", - (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int32_t), - il, -1); - glm_debug_dump_raw_layer(g->batch_router_weights, "selw", - (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(float), - il, -1); - ok = ds4_gpu_begin_commands() != 0; - } - } - if (ok && - !g->ssd_streaming && - progress_flush_interval != 0 && - (il < g->layer_end || progress_requested) && - (slice_layer_done % progress_flush_interval) == 0) { - const uint32_t work_done = - work_done_base + (uint32_t)(((uint64_t)n_tokens * slice_layer_done) / g->layer_count); - const bool drain_now = - drain_interval != 0 && - il < g->layer_end && - (slice_layer_done % drain_interval) == 0; - const char *command_action = drain_now ? "drain" : "flush"; - const double trace_command_t0 = trace ? now_sec() : 0.0; - if (trace && (trace_all || drain_now)) { - glm_graph_full_prefill_tracef( - "layer %s begin layer=%u pos=%u tokens=%u work=%u/%u", - command_action, - il, - pos0, - n_tokens, - work_done, - work_total); - } - if (drain_now) { - ok = ds4_gpu_end_commands() != 0; - if (ok) ok = ds4_gpu_begin_commands() != 0; - } else { - ok = ds4_gpu_flush_commands() != 0; - } - if (trace) { - const double trace_command_done = now_sec(); - const double command_ms = (trace_command_done - trace_command_t0) * 1000.0; - const double layer_ms = (trace_command_done - trace_layer_t0) * 1000.0; - trace_layer_flushed = true; - if (trace_all || drain_now || - command_ms >= trace_slow_ms || layer_ms >= trace_slow_ms || !ok) { - glm_graph_full_prefill_tracef( - "layer %s %s layer=%u pos=%u tokens=%u command=%.3f ms layer_total=%.3f ms work=%u/%u", - command_action, - ok ? "done" : "failed", - il, - pos0, - n_tokens, - command_ms, - layer_ms, - work_done, - work_total); - } - } - if (ok) { - const bool progress_completed = - drain_interval == 0 || drain_now; - if (progress_completed) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - slice_layer_done, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); - } - } - } - if (g->ssd_streaming) { - if (streaming_prefill_sync_each_layer) { - if (ok && full_layer_flush_intermediate && - il < g->layer_end) { - ok = ds4_gpu_flush_commands() != 0; - } else if (ok) { - ok = ds4_gpu_end_commands() != 0; - } - else (void)ds4_gpu_synchronize(); - } - else if (!ok) (void)ds4_gpu_synchronize(); - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - slice_layer_done, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); - } - } - if (ok && - g->ssd_streaming && - layer_prepare && - !layer_prepare_overlap) { - if (il < g->layer_end) { - if (!metal_graph_stream_prepare_start_if_needed(NULL, - model, - weights, - il + 1u, - n_tokens, - layer_madvise, - layer_pread, - layer_readahead, - full_layer_prefill && - rocm_graph_glm_stream_prefill_full_layer_enabled( - g, - &weights->layer[il + 1u], - il + 1u, - n_tokens), - layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - } - } else if (logits_out) { - metal_graph_stream_readahead_output(model, weights); - } - } - if (trace && ok) { - const double layer_ms = (now_sec() - trace_layer_t0) * 1000.0; - if (trace_all || (!trace_layer_flushed && layer_ms >= trace_slow_ms)) { - glm_graph_full_prefill_tracef( - "layer end layer=%u pos=%u tokens=%u flushed=%u layer_total=%.3f ms", - il, - pos0, - n_tokens, - trace_layer_flushed ? 1u : 0u, - layer_ms); - } - } + + vocab->bos_id = vocab_lookup(vocab, "<｜begin▁of▁sentence｜>"); + vocab->eos_id = vocab_lookup(vocab, "<｜end▁of▁sentence｜>"); + vocab->system_id = -1; + vocab->user_id = vocab_lookup(vocab, "<｜User｜>"); + vocab->assistant_id = vocab_lookup(vocab, "<｜Assistant｜>"); + vocab->observation_id = -1; + vocab->sop_id = -1; + vocab->think_start_id = vocab_lookup(vocab, ""); + vocab->think_end_id = vocab_lookup(vocab, ""); + vocab->tool_call_start_id = -1; + vocab->tool_call_end_id = -1; + vocab->tool_response_start_id = -1; + vocab->tool_response_end_id = -1; + vocab->arg_key_start_id = -1; + vocab->arg_key_end_id = -1; + vocab->arg_value_start_id = -1; + vocab->arg_value_end_id = -1; + vocab->dsml_id = vocab_lookup(vocab, "｜DSML｜"); +} + +static void vocab_free(ds4_vocab *vocab) { + free(vocab->token); + table_free(&vocab->token_to_id); + table_free(&vocab->merge_rank); + memset(vocab, 0, sizeof(*vocab)); +} + +/* Build the DS4 chat prompt: BOS, optional system text, user prompt, assistant + * marker, and either or depending on the requested mode. Max + * thinking is only a prompt prefix: the model still enters through . */ +static void chat_push_bos_sequence(const ds4_vocab *vocab, token_vec *out) { + token_vec_push(out, vocab->bos_id); + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && vocab->sop_id >= 0) + token_vec_push(out, vocab->sop_id); +} + +const char *ds4_glm_reasoning_effort_text(ds4_think_mode mode) { + switch (mode) { + case DS4_THINK_HIGH: return "Reasoning Effort: High"; + case DS4_THINK_MAX: return "Reasoning Effort: Max"; + case DS4_THINK_NONE: return NULL; } -#undef DS4_GLM_PROFILE_PREFILL_STAGE - if (ok && !g->ssd_streaming) { - const double trace_end_t0 = trace ? now_sec() : 0.0; - if (trace) { - glm_graph_full_prefill_tracef( - "chunk end_commands begin pos=%u tokens=%u", - pos0, - n_tokens); - } - ok = ds4_gpu_end_commands() != 0; - if (trace) { - const double end_ms = (now_sec() - trace_end_t0) * 1000.0; - const double chunk_ms = (now_sec() - trace_chunk_t0) * 1000.0; - glm_graph_full_prefill_tracef( - "chunk end_commands %s pos=%u tokens=%u end=%.3f ms chunk_total=%.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - end_ms, - chunk_ms); - } - } else if (!ok) { - if (trace) { - glm_graph_full_prefill_tracef( - "chunk failed before end pos=%u tokens=%u elapsed=%.3f ms", - pos0, - n_tokens, - (now_sec() - trace_chunk_t0) * 1000.0); - } -#ifdef DS4_ROCM_BUILD - (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); - if (full_layer_prefill) { - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); - } -#endif - if (layer_prepare) { - (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, - layer_prepare_ahead); + return NULL; +} + +static void chat_push_think_prefix(const ds4_vocab *vocab, + ds4_think_mode think_mode, + token_vec *out) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + const char *effort = ds4_glm_reasoning_effort_text(think_mode); + if (effort) { + token_vec_push(out, vocab->system_id); + bpe_tokenize_text(vocab, effort, out); } - (void)ds4_gpu_synchronize(); - } - if (ok && layer_prepare && - !metal_graph_stream_prepare_join_all(layer_prepare_slots, - layer_prepare_ahead)) { - ok = false; - } -#ifdef DS4_ROCM_BUILD - if (!rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load)) { - ok = false; + } else if (think_mode == DS4_THINK_MAX) { + bpe_tokenize_text(vocab, DS4_REASONING_EFFORT_MAX_PREFIX, out); } - if (full_layer_prefill) { - (void)ds4_gpu_stream_expert_cache_release_layer_cache(); +} + +static void encode_chat_prompt( + const ds4_vocab *vocab, + const char *system, + const char *prompt, + ds4_think_mode think_mode, + token_vec *out) { + const bool need_think_start = + ds4_think_mode_enabled(think_mode) || + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; + if (vocab->bos_id < 0 || + vocab->user_id < 0 || + vocab->assistant_id < 0 || + vocab->think_end_id < 0 || + (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && vocab->system_id < 0) || + (need_think_start && vocab->think_start_id < 0)) { + ds4_die("this tokenizer does not provide the DeepSeek chat markers; use raw prompt tokenization"); } -#endif - if (ok && - g->ssd_streaming && - !streaming_prefill_sync_each_layer && - !output_hc && - !logits_out) { - ok = ds4_gpu_end_commands() != 0; + + chat_push_bos_sequence(vocab, out); + chat_push_think_prefix(vocab, think_mode, out); + if (system && system[0]) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) + token_vec_push(out, vocab->system_id); + bpe_tokenize_text(vocab, system, out); } - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - g->layer_count, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); + token_vec_push(out, vocab->user_id); + bpe_tokenize_text(vocab, prompt, out); + token_vec_push(out, vocab->assistant_id); + if (ds4_think_mode_enabled(think_mode)) { + token_vec_push(out, vocab->think_start_id); + } else if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + token_vec_push(out, vocab->think_start_id); + token_vec_push(out, vocab->think_end_id); + } else { + token_vec_push(out, vocab->think_end_id); } - if (ok && output_hc) { - ok = ds4_gpu_tensor_read(cur, - 0, - output_hc, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; - } - if (ok && logits_out) { - ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, - model, - weights); - } - if (ok && logits_out) { - last_hidden = glm_graph_tensor_row_view_strided(cur, - n_tokens - 1u, - DS4_N_EMBD, - DS4_N_EMBD); - ok = last_hidden != NULL; - if (ok && g->ssd_streaming) ok = glm_graph_stream_map_output(g, model, weights); - if (ok) ok = glm_graph_forward_output_head(g, model, weights, last_hidden, logits_out); - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - g->layer_count, - g->layer_count, - work_total, - true); - } - } - ds4_gpu_tensor_free(last_hidden); - ds4_gpu_set_glm_streaming_prefill_full_layer(false); - return ok; } -static uint32_t glm_graph_prefill_chunk_tokens(uint32_t full_attention_cap) { - return full_attention_cap ? full_attention_cap : DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT; +void ds4_tokenize_text(ds4_engine *e, const char *text, ds4_tokens *out) { + bpe_tokenize_text(&e->vocab, text ? text : "", out); } -static bool glm_graph_forward_indexed_tokens( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - const float *input_hc, - uint32_t pos0, - uint32_t n_tokens, - float *output_hc, - float *logits_out, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t display_absolute_base, - uint32_t work_done_base, - uint32_t work_total) { - if (!g || !model || !weights || !tokens || - g->compact_cache_cap == 0 || - g->indexed_prefill_cap == 0 || - g->indexed_prefill_score_cap == 0 || - !g->batch_indexer_q || - !g->batch_indexer_weights || - !g->batch_indexer_scores || - !g->batch_indexer_selected || - !g->batch_qk_low || - n_tokens == 0 || - g->layer_count == 0 || - n_tokens > g->indexed_prefill_cap || - !glm_graph_span_fits_context(g, pos0, n_tokens)) { - return false; - } - const uint32_t n_rows = pos0 + n_tokens; - const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); - if (pos0 < indexer_top_k && n_rows > indexer_top_k) { - return false; - } - const uint32_t indexed_selected_count = - n_rows <= indexer_top_k ? n_rows : indexer_top_k; - const bool use_causal_range_select = n_rows <= indexer_top_k; - const bool trace = glm_graph_indexed_prefill_trace_enabled(); - const bool trace_all = trace && glm_graph_indexed_prefill_trace_all(); - const double trace_slow_ms = trace ? - (double)glm_graph_indexed_prefill_trace_slow_ms() : 0.0; - const double trace_chunk_t0 = trace ? now_sec() : 0.0; - if (trace) { - glm_graph_indexed_prefill_tracef( - "chunk begin pos=%u tokens=%u rows=%u selected=%u compact_cap=%u score_cap=%u work_base=%u work_total=%u", - pos0, - n_tokens, - n_rows, - indexed_selected_count, - g->compact_cache_cap, - g->indexed_prefill_score_cap, - work_done_base, - work_total); - } - const double trace_cache_t0 = trace ? now_sec() : 0.0; - if (!glm_graph_ensure_compact_cache(g, n_rows)) { - if (trace) { - glm_graph_indexed_prefill_tracef( - "ensure_cache failed pos=%u tokens=%u rows=%u compact_cap=%u", - pos0, - n_tokens, - n_rows, - g->compact_cache_cap); - } - return false; - } - if (trace) { - const double ms = (now_sec() - trace_cache_t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms) { - glm_graph_indexed_prefill_tracef( - "ensure_cache done pos=%u tokens=%u rows=%u compact_cap=%u %.3f ms", - pos0, - n_tokens, - n_rows, - g->compact_cache_cap, - ms); - } - } - for (uint32_t i = 0; i < n_tokens; i++) { - if (tokens[i] < 0 || tokens[i] >= (int)DS4_N_VOCAB) return false; - } - if (!input_hc && !g->has_token_embd) return false; - if (logits_out && !g->has_output_head) return false; - glm_graph_reset_prefill_seed_capture(g); - - const double trace_upload_t0 = trace ? now_sec() : 0.0; - bool ok = glm_graph_upload_tokens(g->prefill_tokens, tokens, n_tokens); - if (trace) { - const double ms = (now_sec() - trace_upload_t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_indexed_prefill_tracef( - "upload_tokens %s pos=%u tokens=%u %.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } - } - ds4_gpu_tensor *cur = g->batch_cur; - ds4_gpu_tensor *next = g->batch_next; - ds4_gpu_tensor *last_hidden = NULL; - - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - 0, - g->layer_count, - work_total, - true); - - const bool use_all_scalar_kernels = - n_tokens == 1u && glm_graph_indexed_prefill_scalar_kernels(); - const bool use_scalar_indexer = - use_all_scalar_kernels || - glm_graph_indexed_prefill_scalar_indexer() || - !glm_graph_indexed_prefill_batch_indexer(); - const bool force_scalar_attn = - use_all_scalar_kernels || glm_graph_indexed_prefill_scalar_attn(); - const bool use_batch_qk_low = - !force_scalar_attn && glm_graph_indexed_prefill_batch_qk_low(); - const bool use_batch_attn_kernel = - !force_scalar_attn && glm_graph_indexed_prefill_batch_attn_kernel(); - const bool use_split_value_proj = - use_batch_attn_kernel && - g->batch_attn_lora; - /* Tensor-parallel attention head split: each rank computes half the - * heads in the qk-low / attention-lora / value-project kernels, the - * unowned half of batch_heads stays zero, and the full-width attn - * output projection yields partials combined over the big-gate - * exchange (same commutative add as the routed-FFN combine). Only the - * split-value-proj batch chain has head ownership. */ - const bool tp_attn_head_split = - g->tp_world == 2 && - use_batch_attn_kernel && - use_split_value_proj && - (DS4_N_HEAD % 16u) == 0u && - n_tokens >= glm_tp_head_split_min(); /* small batches replicate; - * the floor is env-tunable for correctness - * isolation (DS4_GLM_TP_HEAD_SPLIT_MIN). */ - const bool use_batch_q_rank_proj = true; - const bool use_batch_q_proj = true; - const bool use_batch_indexer_k_proj = true; - const bool use_batch_kv_proj = true; - const bool use_batch_indexer_q_proj = true; - const bool use_batch_indexer_weights_proj = true; - const bool use_batch_attn_out_proj = true; - const bool use_batch_ffn = glm_graph_indexed_prefill_batch_ffn(); - const bool stage_sync = - glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); - const uint32_t layer_flush_interval = stage_sync ? 0u : - glm_graph_full_prefill_layer_flush_interval(n_tokens, - n_tokens, - logits_out != NULL); - const uint32_t progress_flush_interval = - glm_graph_prefill_progress_flush_interval(layer_flush_interval, - n_tokens, - display_progress, - work_total); - const uint32_t drain_interval = - progress_flush_interval != 0 ? glm_graph_indexed_prefill_drain_interval() : 0u; - const bool progress_requested = display_progress && work_total > 0; - ds4_gpu_set_glm_streaming_prefill_full_layer(false); - const bool streaming_prefill_sync_each_layer = - !g->ssd_streaming || - glm_graph_streaming_prefill_sync_each_layer(false); - - if (trace) { - glm_graph_indexed_prefill_tracef( - "mode pos=%u tokens=%u scalar_indexer=%u batch_qk_low=%u batch_attn=%u split_value=%u batch_ffn=%u progress_flush_interval=%u drain_interval=%u", - pos0, - n_tokens, - use_scalar_indexer ? 1u : 0u, - use_batch_qk_low ? 1u : 0u, - use_batch_attn_kernel ? 1u : 0u, - use_split_value_proj ? 1u : 0u, - use_batch_ffn ? 1u : 0u, - progress_flush_interval, - drain_interval); - } +static bool special_token_at(const ds4_vocab *vocab, const char *p, int *token, size_t *len) { + struct special { + const char *text; + int token; + } specials[] = { + {"<｜begin▁of▁sentence｜>", vocab->bos_id}, + {"<｜end▁of▁sentence｜>", vocab->eos_id}, + {"[gMASK]", vocab->bos_id}, + {"", vocab->sop_id}, + {"<|system|>", vocab->system_id}, + {"<｜User｜>", vocab->user_id}, + {"<｜Assistant｜>", vocab->assistant_id}, + {"<|user|>", vocab->user_id}, + {"<|assistant|>", vocab->assistant_id}, + {"<|observation|>", vocab->observation_id}, + {"", vocab->think_start_id}, + {"", vocab->think_end_id}, + {"", vocab->tool_call_start_id}, + {"", vocab->tool_call_end_id}, + {"", vocab->tool_response_start_id}, + {"", vocab->tool_response_end_id}, + {"", vocab->arg_key_start_id}, + {"", vocab->arg_key_end_id}, + {"", vocab->arg_value_start_id}, + {"", vocab->arg_value_end_id}, + {"｜DSML｜", vocab->dsml_id}, + }; - if (ok) { - const double t0 = trace ? now_sec() : 0.0; - if (input_hc) { - ok = ds4_gpu_tensor_write(cur, - 0, - input_hc, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; - } else { - ok = glm_graph_stream_map_token(g, model, weights); - } - if (ok) ok = ds4_gpu_begin_commands() != 0; - if (trace) { - const double ms = (now_sec() - t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_indexed_prefill_tracef( - "begin_commands%s %s pos=%u tokens=%u %.3f ms", - input_hc ? "_from_hidden" : "", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } + for (size_t i = 0; i < sizeof(specials) / sizeof(specials[0]); i++) { + if (specials[i].token < 0) continue; + size_t n = strlen(specials[i].text); + if (!strncmp(p, specials[i].text, n)) { + *token = specials[i].token; + *len = n; + return true; } } - if (ok && !input_hc) { - const double t0 = trace ? now_sec() : 0.0; - ok = ds4_gpu_embed_tokens_quant_tensor(cur, - g->prefill_tokens, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - n_tokens, - DS4_N_EMBD) != 0; - if (trace) { - const double ms = (now_sec() - t0) * 1000.0; - if (trace_all || ms >= trace_slow_ms || !ok) { - glm_graph_indexed_prefill_tracef( - "embed %s pos=%u tokens=%u %.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - ms); - } + return false; +} + +static void tokenize_span(const ds4_vocab *vocab, const char *p, size_t n, token_vec *out) { + if (!n) return; + char *tmp = xmalloc(n + 1); + memcpy(tmp, p, n); + tmp[n] = '\0'; + bpe_tokenize_text(vocab, tmp, out); + free(tmp); +} + + + + +static void tokenize_rendered_chat_vocab(const ds4_vocab *vocab, const char *text, + token_vec *out) { + if (!text) text = ""; + + const char *span = text; + const char *p = text; + while (*p) { + int token = -1; + size_t len = 0; + if (special_token_at(vocab, p, &token, &len)) { + tokenize_span(vocab, span, (size_t)(p - span), out); + token_vec_push(out, token); + p += len; + span = p; + continue; } + p++; } - if (ok && g->ssd_streaming && streaming_prefill_sync_each_layer) { - ok = ds4_gpu_end_commands() != 0; - } - -#define DS4_GLM_PROFILE_INDEXED_STAGE(part_, name_) do { \ - if (ok && trace) { \ - const double _trace_stage_now = now_sec(); \ - const double _trace_stage_ms = (_trace_stage_now - trace_stage_t0) * 1000.0; \ - if (trace_all || _trace_stage_ms >= trace_slow_ms) { \ - glm_graph_indexed_prefill_tracef( \ - "stage layer=%u pos=%u tokens=%u %s.%s encode %.3f ms", \ - il, \ - pos0, \ - n_tokens, \ - (part_), \ - (name_), \ - _trace_stage_ms); \ - } \ - trace_stage_t0 = _trace_stage_now; \ - } \ - if (ok && layer_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos0, n_tokens, &layer_stage_t0); \ - } else if (ok && stage_sync) { \ - ok = glm_graph_prefill_stage_sync_boundary(); \ - } \ - } while (0) - ds4_gpu_tensor *last_indexer_selected = NULL; - uint32_t last_indexer_selected_count = 0; - if (ok && tp_attn_head_split) { - /* The unowned head range of batch_heads must be exactly zero so the - * full-width attn-output matmul produces partial sums. Owned heads - * are rewritten every layer, so one fill per chunk suffices. */ - ok = ds4_gpu_tensor_fill_f32(g->batch_heads, 0.0f, - (uint64_t)n_tokens * g->heads_dim) != 0; - } - ds4_gpu_tp_set_attn_head_split(tp_attn_head_split ? 1 : 0); - for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { - const uint32_t slice_layer_done = il - g->layer_start + 1u; - if (g->ssd_streaming) { - ok = glm_graph_stream_map_prefill_layer(g, - model, - weights, - il, - n_tokens, - false); - if (ok) ok = ds4_gpu_begin_commands() != 0; - } - const ds4_layer_weights *l = &weights->layer[il]; - const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; - const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); - double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; - double trace_stage_t0 = trace ? now_sec() : 0.0; - const double trace_layer_t0 = trace_stage_t0; - const bool trace_full_indexer = glm_graph_layer_uses_full_indexer(il); - bool trace_layer_flushed = false; - if (trace && (trace_all || trace_full_indexer)) { - glm_graph_indexed_prefill_tracef( - "layer begin layer=%u pos=%u tokens=%u rows=%u selected=%u full_indexer=%u", - il, - pos0, - n_tokens, - n_rows, - indexed_selected_count, - trace_full_indexer ? 1u : 0u); - } - if (residual_elems > UINT32_MAX) { - ok = false; - break; - } - if (layer_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", - NULL, - il, - pos0, - n_tokens, - &layer_stage_t0); - } - - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, - cur, - model->map, - model->size, - l->attn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attn_norm"); - if (ok) { - if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_QPATH)) { /* ablate */ } else - ok = (use_batch_q_rank_proj ? - glm_graph_matmul_q8_0_tensor(g->batch_q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->batch_attn_norm, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(g->batch_q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->batch_attn_norm, - n_tokens)); - } - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, - g->batch_q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) { - ok = (use_batch_q_proj ? - glm_graph_matmul_q8_0_tensor(g->batch_q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->batch_q_rank_norm, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(g->batch_q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->batch_q_rank_norm, - n_tokens)); - } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, - n_tokens, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_ROT, - pos0, - 0, - false, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "q_path"); + tokenize_span(vocab, span, (size_t)(p - span), out); +} - if (ok && glm_graph_layer_uses_full_indexer(il)) { - ok = (use_batch_indexer_k_proj ? - glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, - model, - l->indexer_attn_k->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD_DIM, - cur, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(g->batch_indexer_k, - model, - l->indexer_attn_k->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD_DIM, - cur, - n_tokens)); - if (ok) { - ok = ds4_gpu_glm_store_indexer_k_tensor( - g->layer_indexer_key_cache[il], - g->batch_indexer_k, - model->map, - model->size, - l->indexer_k_norm->abs_offset, - l->indexer_k_norm_b->abs_offset, - pos0, - n_tokens, - g->compact_cache_cap, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - 0, - 1.0e-6f, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - glm_graph_compact_cache_is_f16()) != 0; - } - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_k"); +void ds4_tokenize_rendered_chat(ds4_engine *e, const char *text, ds4_tokens *out) { + tokenize_rendered_chat_vocab(&e->vocab, text, out); +} - if (ok) { - ok = (use_batch_kv_proj ? - glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->batch_attn_norm, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(g->batch_kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->batch_attn_norm, - n_tokens)); - } - if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, - g->batch_kv_raw, - model->map, - model->size, - l->attn_kv_a_norm->abs_offset, - n_tokens, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_RMS_EPS) != 0; - if (ok) { - ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - g->batch_kv_norm, - g->batch_kv_raw, - pos0, - n_tokens, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16()) != 0; - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "kv_path"); - - if (ok && glm_graph_layer_uses_full_indexer(il)) { - if (ok && !use_causal_range_select) { - ok = (use_batch_indexer_q_proj ? - glm_graph_matmul_q8_0_tensor(g->batch_indexer_q, - model, - l->indexer_attn_q_b->abs_offset, - DS4_N_LORA_Q, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - g->batch_q_rank_norm, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(g->batch_indexer_q, - model, - l->indexer_attn_q_b->abs_offset, - DS4_N_LORA_Q, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - g->batch_q_rank_norm, - n_tokens)); - if (ok) ok = ds4_gpu_glm_indexer_rope_tail_tensor(g->batch_indexer_q, - n_tokens, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - pos0, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok && glm_graph_indexer_qat()) { - ok = ds4_gpu_dsv4_indexer_qat_tensor(g->batch_indexer_q, - n_tokens * DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM) != 0; - } - if (ok) { - ok = (use_batch_indexer_weights_proj ? - ds4_gpu_matmul_f32_tensor(g->batch_indexer_weights, - model->map, - model->size, - l->indexer_proj->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD, - cur, - n_tokens) != 0 : - glm_graph_matmul_f32_rows_scalar(g->batch_indexer_weights, - model, - l->indexer_proj->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD, - cur, - n_tokens)); - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_q_weights"); - } - if (ok) { - if (use_causal_range_select) { - ok = ds4_gpu_glm_fill_selected_range_batch_tensor( - g->batch_indexer_selected, - n_tokens, - pos0, - indexed_selected_count, - g->compact_cache_cap) != 0; - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_range"); - } else if (use_scalar_indexer) { - const float indexer_scale = - 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - for (uint32_t t = 0; ok && t < n_tokens; t++) { - const uint32_t visible = pos0 + t + 1u; - ds4_gpu_tensor *indexer_q_view = - glm_graph_tensor_row_view_strided( - g->batch_indexer_q, - t, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); - ds4_gpu_tensor *indexer_weights_view = - glm_graph_tensor_row_view_strided(g->batch_indexer_weights, - t, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD); - ds4_gpu_tensor *scores_view = - ds4_gpu_tensor_view(g->batch_indexer_scores, - 0, - (uint64_t)visible * sizeof(float)); - ds4_gpu_tensor *selected_view = - ds4_gpu_tensor_view(g->batch_indexer_selected, - (uint64_t)t * indexed_selected_count * sizeof(uint32_t), - (uint64_t)indexed_selected_count * sizeof(uint32_t)); - ok = indexer_q_view && indexer_weights_view && scores_view && selected_view; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create indexer row views at layer %u token %u\n", il, t); - if (ok) { - int rc = ds4_gpu_glm_indexer_score_one_tensor( - scores_view, - indexer_q_view, - indexer_weights_view, - g->layer_indexer_key_cache[il], - visible, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - indexer_scale, - glm_graph_compact_cache_is_f16()); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill indexer scores failed at layer %u token %u\n", il, t); - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_score_scalar"); - if (ok) { - int rc = ds4_gpu_indexer_topk_tensor(selected_view, - scores_view, - visible, - 1, - indexed_selected_count); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill topk failed at layer %u token %u\n", il, t); - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_topk_scalar"); - ds4_gpu_tensor_free(selected_view); - ds4_gpu_tensor_free(scores_view); - ds4_gpu_tensor_free(indexer_weights_view); - ds4_gpu_tensor_free(indexer_q_view); - } - } else { - const float indexer_scale = - 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - const uint32_t score_cap = - g->indexed_prefill_score_cap != 0 ? - g->indexed_prefill_score_cap : - g->indexed_prefill_cap; - const uint64_t indexer_q_row_bytes = - (uint64_t)DS4_N_INDEXER_HEAD * - DS4_N_INDEXER_HEAD_DIM * - sizeof(float); - const uint64_t indexer_weights_row_bytes = - (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); - const uint64_t selected_row_bytes = - (uint64_t)indexed_selected_count * sizeof(uint32_t); - for (uint32_t t0 = 0; ok && t0 < n_tokens; ) { - uint32_t slice = n_tokens - t0; - if (slice > score_cap) slice = score_cap; - if (slice == 0) { - ok = false; - break; - } +void ds4_chat_begin(ds4_engine *e, ds4_tokens *tokens) { + chat_push_bos_sequence(&e->vocab, tokens); +} - ds4_gpu_tensor *indexer_q_view = - ds4_gpu_tensor_view(g->batch_indexer_q, - (uint64_t)t0 * indexer_q_row_bytes, - (uint64_t)slice * indexer_q_row_bytes); - ds4_gpu_tensor *indexer_weights_view = - ds4_gpu_tensor_view(g->batch_indexer_weights, - (uint64_t)t0 * indexer_weights_row_bytes, - (uint64_t)slice * indexer_weights_row_bytes); - ds4_gpu_tensor *selected_view = - ds4_gpu_tensor_view(g->batch_indexer_selected, - (uint64_t)t0 * selected_row_bytes, - (uint64_t)slice * selected_row_bytes); - ok = indexer_q_view && indexer_weights_view && selected_view; - if (!ok) { - fprintf(stderr, - "ds4: GLM indexed prefill failed to create indexer score slice views at layer %u token %u\n", - il, - t0); - } - if (ok) { - ok = ds4_gpu_glm_indexer_scores_batch_tensor( - g->batch_indexer_scores, - indexer_q_view, - indexer_weights_view, - g->layer_indexer_key_cache[il], - n_rows, - slice, - pos0 + t0, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - indexer_scale, - glm_graph_compact_cache_is_f16()) != 0; - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_score"); - if (ok) { - ok = ds4_gpu_indexer_topk_tensor(selected_view, - g->batch_indexer_scores, - n_rows, - slice, - indexed_selected_count) != 0; - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_topk"); - ds4_gpu_tensor_free(selected_view); - ds4_gpu_tensor_free(indexer_weights_view); - ds4_gpu_tensor_free(indexer_q_view); - t0 += slice; - } - } - } - if (ok) { - last_indexer_selected = g->batch_indexer_selected; - last_indexer_selected_count = indexed_selected_count; - } - } else if (ok && (!last_indexer_selected || last_indexer_selected_count == 0)) { - ok = false; - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_select"); - metal_graph_debug_dump_tensor("glm_indexed_q", - g->batch_q, - (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KEY_MLA, - il, - pos0); - if (use_batch_qk_low) { - if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_CORE)) { /* ablate */ } else - if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_batch_tensor(g->batch_qk_low, - g->batch_q, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - n_tokens, - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_KEY_MLA) != 0; - } else { - for (uint32_t t = 0; ok && t < n_tokens; t++) { - ds4_gpu_tensor *q_view = - glm_graph_tensor_row_view_strided(g->batch_q, - t, - g->q_dim, - g->q_dim); - ds4_gpu_tensor *qk_low_view = - glm_graph_tensor_row_view_strided( - g->batch_qk_low, - t, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); - ok = q_view && qk_low_view; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create qk-low row views at layer %u token %u\n", il, t); - if (ok) { - int rc = ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low_view, - q_view, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_KEY_MLA); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill qk-low failed at layer %u token %u\n", il, t); - } - ds4_gpu_tensor_free(qk_low_view); - ds4_gpu_tensor_free(q_view); - } - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "qk_low"); - metal_graph_debug_dump_tensor("glm_indexed_qk_low", - g->batch_qk_low, - (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KV_LORA, - il, - pos0); - if (ok && use_batch_attn_kernel) ok = glm_graph_indexed_prefill_attention_boundary(); - - if (use_batch_attn_kernel) { - const uint32_t attn_slice_cap = - glm_graph_indexed_prefill_batch_attn_slice_tokens(); - for (uint32_t t0 = 0; ok && t0 < n_tokens; ) { - uint32_t slice = n_tokens - t0; - if (slice > attn_slice_cap) slice = attn_slice_cap; - - ds4_gpu_tensor *q_view = - ds4_gpu_tensor_view(g->batch_q, - (uint64_t)t0 * g->q_dim * sizeof(float), - (uint64_t)slice * g->q_dim * sizeof(float)); - ds4_gpu_tensor *qk_low_view = - ds4_gpu_tensor_view(g->batch_qk_low, - (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), - (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)); - ds4_gpu_tensor *heads_view = - ds4_gpu_tensor_view(g->batch_heads, - (uint64_t)t0 * g->heads_dim * sizeof(float), - (uint64_t)slice * g->heads_dim * sizeof(float)); - ds4_gpu_tensor *attn_lora_view = use_split_value_proj ? - ds4_gpu_tensor_view(g->batch_attn_lora, - (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), - (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)) : - NULL; - ds4_gpu_tensor *selected_view = - ds4_gpu_tensor_view(last_indexer_selected, - (uint64_t)t0 * last_indexer_selected_count * sizeof(uint32_t), - (uint64_t)slice * last_indexer_selected_count * sizeof(uint32_t)); - ok = q_view && qk_low_view && heads_view && selected_view && - (!use_split_value_proj || attn_lora_view); - if (!ok) { - fprintf(stderr, "ds4: GLM sliced indexed prefill failed to create attention views at layer %u token %u\n", il, t0); - } - if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_CORE)) { /* ablate */ } else if (ok && use_split_value_proj) { - int rc = 0; - if (use_causal_range_select) { - rc = ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( - attn_lora_view, - q_view, - qk_low_view, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - slice, - pos0 + t0, - last_indexer_selected_count, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW); - } else { - rc = ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( - attn_lora_view, - q_view, - qk_low_view, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - selected_view, - slice, - last_indexer_selected_count, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW); - } - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill attention-lora failed at layer %u token %u\n", il, t0); - if (ok && layer_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", - "attention_lora", - il, - pos0 + t0, - slice, - &layer_stage_t0); - } - if (ok) { - rc = ds4_gpu_glm_value_project_typed_batch_heads_tensor( - heads_view, - attn_lora_view, - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - slice, - DS4_N_HEAD, - DS4_N_KV_LORA, - DS4_N_VALUE_MLA); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill value project failed at layer %u token %u\n", il, t0); - } - if (ok && layer_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", - "value_project", - il, - pos0 + t0, - slice, - &layer_stage_t0); - } - } else if (ok) { - int rc = ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads_view, - q_view, - qk_low_view, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - selected_view, - slice, - last_indexer_selected_count, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill indexed attention failed at layer %u token %u\n", il, t0); - if (ok && layer_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", - "attention_fused", - il, - pos0 + t0, - slice, - &layer_stage_t0); - } - } - ds4_gpu_tensor_free(selected_view); - ds4_gpu_tensor_free(attn_lora_view); - ds4_gpu_tensor_free(heads_view); - ds4_gpu_tensor_free(qk_low_view); - ds4_gpu_tensor_free(q_view); - t0 += slice; - } - if (ok) ok = glm_graph_indexed_prefill_attention_boundary(); - } else { - for (uint32_t t = 0; ok && t < n_tokens; t++) { - ds4_gpu_tensor *q_view = - glm_graph_tensor_row_view_strided(g->batch_q, - t, - g->q_dim, - g->q_dim); - ds4_gpu_tensor *qk_low_view = - glm_graph_tensor_row_view_strided( - g->batch_qk_low, - t, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); - ds4_gpu_tensor *heads_view = - glm_graph_tensor_row_view_strided(g->batch_heads, - t, - g->heads_dim, - g->heads_dim); - ds4_gpu_tensor *selected_view = - ds4_gpu_tensor_view(last_indexer_selected, - (uint64_t)t * last_indexer_selected_count * sizeof(uint32_t), - (uint64_t)last_indexer_selected_count * sizeof(uint32_t)); - ok = q_view && qk_low_view && heads_view && selected_view; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create attention row views at layer %u token %u\n", il, t); - if (ok) { - int rc = ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads_view, - q_view, - qk_low_view, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - selected_view, - last_indexer_selected_count, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW); - ok = rc != 0; - if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill indexed attention failed at layer %u token %u\n", il, t); - } - ds4_gpu_tensor_free(selected_view); - ds4_gpu_tensor_free(heads_view); - ds4_gpu_tensor_free(qk_low_view); - ds4_gpu_tensor_free(q_view); - } - } - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attention"); - metal_graph_debug_dump_tensor("glm_indexed_heads", - g->batch_heads, - (uint64_t)n_tokens * g->heads_dim, - il, - pos0); - if (ok && tp_attn_head_split) { - ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); - } - if (ok) { - /* Under the head split the projection input has zeros in the - * unowned head columns, so the result is this rank's partial; - * it must land in the shared bounce tensor for the exchange. */ - ds4_gpu_tensor *attn_out_dst = - tp_attn_head_split ? g->tp_bounce_out : g->batch_attn_out; - if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_OUT)) { /* ablate */ } else - ok = (use_batch_attn_out_proj ? - glm_graph_matmul_q8_0_tensor(attn_out_dst, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->batch_heads, - n_tokens) : - glm_graph_matmul_q8_0_rows_scalar(attn_out_dst, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->batch_heads, - n_tokens)); - } - if (ok && tp_attn_head_split) { - ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_attn_out, n_tokens); - } - if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, - cur, - g->batch_attn_out, - (uint32_t)residual_elems) != 0; - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attn_output"); - metal_graph_debug_dump_tensor("glm_indexed_after_attn", - g->batch_after_attn, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - metal_graph_debug_dump_tensor("glm_indexed_attn_out", - g->batch_attn_out, - (uint64_t)n_tokens * DS4_N_EMBD, - il, - pos0); - if (ok && use_batch_ffn) { - ok = glm_graph_encode_ffn_batch(g, - model, - weights, - l, - il, - pos0, - g->batch_after_attn, - next, - n_tokens, - false, - layer_stage_profile, - stage_sync, - layer_stage_profile ? &layer_stage_t0 : NULL); - } else if (ok) { - const bool use_batch_ffn_norm = - n_tokens > 1 && glm_graph_indexed_prefill_batch_ffn_norm(); - if (use_batch_ffn_norm) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, - g->batch_after_attn, - model->map, - model->size, - l->ffn_norm->abs_offset, - DS4_N_EMBD, - n_tokens, - DS4_RMS_EPS) != 0; - } - if (use_batch_ffn_norm) { - DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_ffn", "ffn_norm"); - } - if (ok && - use_batch_ffn_norm && - il >= DS4_N_LEADING_DENSE && - glm_graph_indexed_prefill_batch_routed_moe()) { - ok = glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( - g, - model, - l, - il, - pos0, - g->batch_after_attn, - next, - n_tokens, - layer_stage_profile, - stage_sync, - layer_stage_profile ? &layer_stage_t0 : NULL); - } else for (uint32_t t = 0; ok && t < n_tokens; t++) { - ds4_gpu_tensor *after_attn_view = - glm_graph_tensor_row_view_strided(g->batch_after_attn, - t, - DS4_N_EMBD, - DS4_N_EMBD); - ds4_gpu_tensor *ffn_norm_view = use_batch_ffn_norm ? - glm_graph_tensor_row_view_strided(g->batch_ffn_norm, - t, - DS4_N_EMBD, - DS4_N_EMBD) : - NULL; - ds4_gpu_tensor *next_view = - glm_graph_tensor_row_view_strided(next, - t, - DS4_N_EMBD, - DS4_N_EMBD); - ok = after_attn_view && next_view && - (!use_batch_ffn_norm || ffn_norm_view); - if (ok && use_batch_ffn_norm) { - ok = glm_graph_encode_ffn_one_normed_from(g, - model, - l, - il, - pos0 + t, - ffn_norm_view, - after_attn_view, - next_view, - g->ffn_gate, - g->ffn_up, - g->ffn_mid, - g->ffn_out, - g->ffn_sum, - g->attn_out, - false, - NULL); - } else if (ok) { - ok = glm_graph_encode_ffn_one_from(g, - model, - l, - il, - pos0 + t, - after_attn_view, - next_view, - g->ffn_norm, - g->ffn_gate, - g->ffn_up, - g->ffn_mid, - g->ffn_out, - g->ffn_sum, - g->attn_out, - false, - NULL); - } - ds4_gpu_tensor_free(next_view); - ds4_gpu_tensor_free(ffn_norm_view); - ds4_gpu_tensor_free(after_attn_view); - } - } - if (ok) { - ds4_gpu_tensor *tmp = cur; - cur = next; - next = tmp; - } - if (ok && glm_debug_hidden_dump_layer_match(il)) { - ok = ds4_gpu_end_commands() != 0; - if (ok) { - for (uint32_t r = 0; r < n_tokens; r++) - glm_debug_dump_hidden_layer(cur, r, il, pos0 + r); - glm_debug_dump_raw_layer(g->batch_router_selected, "sel", - (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int32_t), - il, -1); - glm_debug_dump_raw_layer(g->batch_router_weights, "selw", - (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(float), - il, -1); - ok = ds4_gpu_begin_commands() != 0; - } +void ds4_encode_chat_prompt( + ds4_engine *e, + const char *system, + const char *prompt, + ds4_think_mode think_mode, + ds4_tokens *out) { + encode_chat_prompt(&e->vocab, system, prompt ? prompt : "", think_mode, out); +} + +void ds4_chat_append_max_effort_prefix(ds4_engine *e, ds4_tokens *tokens) { + bpe_tokenize_text(&e->vocab, DS4_REASONING_EFFORT_MAX_PREFIX, tokens); +} + +static void bpe_tokenize_wrapped_payload_text(ds4_vocab *vocab, const char *content, + const char *end, token_vec *out) { + /* Tool output is plain data inside the model-family wrapper. + * Preserve literal '<', '>' and '&' so shell output and file snippets stay + * intact, but escape the exact closing sentinel so a malicious or accidental + * tool payload cannot terminate the wrapper early. */ + const size_t endlen = strlen(end); + const char *span = content ? content : ""; + const char *p = span; + while (*p) { + if (!strncmp(p, end, endlen)) { + tokenize_span(vocab, span, (size_t)(p - span), out); + bpe_tokenize_text(vocab, "<", out); + p++; + span = p; + } else { + p++; } - if (ok && - !g->ssd_streaming && - progress_flush_interval != 0 && - (il < g->layer_end || progress_requested) && - (slice_layer_done % progress_flush_interval) == 0) { - const uint32_t work_done = - work_done_base + (uint32_t)(((uint64_t)n_tokens * slice_layer_done) / g->layer_count); - const bool drain_now = - drain_interval != 0 && - il < g->layer_end && - (slice_layer_done % drain_interval) == 0; - const char *command_action = drain_now ? "drain" : "flush"; - const double trace_command_t0 = trace ? now_sec() : 0.0; - if (trace && (trace_all || trace_full_indexer || drain_now)) { - glm_graph_indexed_prefill_tracef( - "layer %s begin layer=%u pos=%u tokens=%u work=%u/%u", - command_action, - il, - pos0, - n_tokens, - work_done, - work_total); - } - if (drain_now) { - ok = ds4_gpu_end_commands() != 0; - if (ok) ok = ds4_gpu_begin_commands() != 0; - } else { - ok = ds4_gpu_flush_commands() != 0; - } - if (trace) { - const double trace_command_done = now_sec(); - const double command_ms = (trace_command_done - trace_command_t0) * 1000.0; - const double layer_ms = (trace_command_done - trace_layer_t0) * 1000.0; - trace_layer_flushed = true; - if (trace_all || trace_full_indexer || drain_now || - command_ms >= trace_slow_ms || layer_ms >= trace_slow_ms || !ok) { - glm_graph_indexed_prefill_tracef( - "layer %s %s layer=%u pos=%u tokens=%u command=%.3f ms layer_total=%.3f ms work=%u/%u", - command_action, - ok ? "done" : "failed", - il, - pos0, - n_tokens, - command_ms, - layer_ms, - work_done, - work_total); - } - } - if (ok) { - const bool progress_completed = - drain_interval == 0 || drain_now; - if (progress_completed) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - slice_layer_done, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); - } + } + tokenize_span(vocab, span, (size_t)(p - span), out); +} + +static void bpe_tokenize_tool_result_text(ds4_vocab *vocab, const char *content, token_vec *out) { + bpe_tokenize_wrapped_payload_text(vocab, content, "", out); +} + +static void bpe_tokenize_tool_response_text(ds4_vocab *vocab, const char *content, token_vec *out) { + bpe_tokenize_wrapped_payload_text(vocab, content, "", out); +} + +void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role, const char *content) { + ds4_vocab *vocab = &e->vocab; + if (!role) role = "user"; + if (!content) content = ""; + + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (!strcmp(role, "system") || !strcmp(role, "developer")) { + if (vocab->system_id >= 0) token_vec_push(tokens, vocab->system_id); + tokenize_rendered_chat_vocab(vocab, content, tokens); + } else if (!strcmp(role, "assistant")) { + token_vec_push(tokens, vocab->assistant_id); + if (strncmp(content, "", 7) != 0 && + strncmp(content, "", 8) != 0) { + token_vec_push(tokens, vocab->think_start_id); + token_vec_push(tokens, vocab->think_end_id); } + tokenize_rendered_chat_vocab(vocab, content, tokens); + } else if (!strcmp(role, "tool") || !strcmp(role, "function")) { + if (vocab->observation_id >= 0) token_vec_push(tokens, vocab->observation_id); + tokenize_rendered_chat_vocab(vocab, "", tokens); + bpe_tokenize_tool_response_text(vocab, content, tokens); + tokenize_rendered_chat_vocab(vocab, "", tokens); + } else { + token_vec_push(tokens, vocab->user_id); + bpe_tokenize_text(vocab, content, tokens); } - if (g->ssd_streaming) { - if (streaming_prefill_sync_each_layer) { - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - } else if (!ok) { - (void)ds4_gpu_synchronize(); - } - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - slice_layer_done, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); - } + return; + } + + if (!strcmp(role, "system") || !strcmp(role, "developer")) { + bpe_tokenize_text(vocab, content, tokens); + } else if (!strcmp(role, "assistant")) { + token_vec_push(tokens, vocab->assistant_id); + if (strncmp(content, "", 7) != 0 && strncmp(content, "", 8) != 0) { + token_vec_push(tokens, vocab->think_end_id); } - if (trace && ok) { - const double layer_ms = (now_sec() - trace_layer_t0) * 1000.0; - if (trace_all || (!trace_layer_flushed && layer_ms >= trace_slow_ms)) { - glm_graph_indexed_prefill_tracef( - "layer end layer=%u pos=%u tokens=%u flushed=%u layer_total=%.3f ms", - il, - pos0, - n_tokens, - trace_layer_flushed ? 1u : 0u, - layer_ms); - } + bpe_tokenize_text(vocab, content, tokens); + } else if (!strcmp(role, "tool") || !strcmp(role, "function")) { + token_vec_push(tokens, vocab->user_id); + bpe_tokenize_text(vocab, "", tokens); + bpe_tokenize_tool_result_text(vocab, content, tokens); + bpe_tokenize_text(vocab, "", tokens); + } else { + token_vec_push(tokens, vocab->user_id); + bpe_tokenize_text(vocab, content, tokens); + } +} + + +void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode) { + token_vec_push(tokens, e->vocab.assistant_id); + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + !ds4_think_mode_enabled(think_mode)) { + token_vec_push(tokens, e->vocab.think_start_id); + token_vec_push(tokens, e->vocab.think_end_id); + return; + } + token_vec_push(tokens, ds4_think_mode_enabled(think_mode) ? + e->vocab.think_start_id : e->vocab.think_end_id); +} + +static void dump_tokens_fp(FILE *fp, const ds4_vocab *vocab, const token_vec *tokens) { + fprintf(fp, "["); + for (int i = 0; i < tokens->len; i++) { + if (i) fprintf(fp, ", "); + fprintf(fp, "%d", tokens->v[i]); + } + fprintf(fp, "]\n"); + + for (int i = 0; i < tokens->len; i++) { + int id = tokens->v[i]; + if (id >= 0 && id < vocab->n_vocab) { + fprintf(fp, "%6d %.*s\n", id, (int)vocab->token[id].len, vocab->token[id].ptr); } } - ds4_gpu_tp_set_attn_head_split(0); -#undef DS4_GLM_PROFILE_INDEXED_STAGE - if (ok && !g->ssd_streaming) { - const double trace_end_t0 = trace ? now_sec() : 0.0; - if (trace) { - glm_graph_indexed_prefill_tracef( - "chunk end_commands begin pos=%u tokens=%u", - pos0, - n_tokens); - } - ok = ds4_gpu_end_commands() != 0; - if (trace) { - const double end_ms = (now_sec() - trace_end_t0) * 1000.0; - const double chunk_ms = (now_sec() - trace_chunk_t0) * 1000.0; - glm_graph_indexed_prefill_tracef( - "chunk end_commands %s pos=%u tokens=%u end=%.3f ms chunk_total=%.3f ms", - ok ? "done" : "failed", - pos0, - n_tokens, - end_ms, - chunk_ms); - } - } else if (!ok) { - if (trace) { - glm_graph_indexed_prefill_tracef( - "chunk failed before end pos=%u tokens=%u elapsed=%.3f ms", - pos0, - n_tokens, - (now_sec() - trace_chunk_t0) * 1000.0); - } - (void)ds4_gpu_synchronize(); - } - if (ok && - g->ssd_streaming && - !streaming_prefill_sync_each_layer && - !output_hc && - !logits_out) { - ok = ds4_gpu_end_commands() != 0; +} + +static void dump_tokens(const ds4_vocab *vocab, const token_vec *tokens) { + dump_tokens_fp(stdout, vocab, tokens); +} + +static uint32_t utf8_decode_one(const char *s, uint64_t len, uint64_t *pos) { + const uint8_t c = (uint8_t)s[*pos]; + if (c < 0x80 || *pos + 1 >= len) { + (*pos)++; + return c; } - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - g->layer_count, - g->layer_count, - work_total, - logits_out == NULL && output_hc == NULL); + if ((c & 0xe0) == 0xc0 && *pos + 1 < len) { + uint32_t cp = ((uint32_t)(c & 0x1f) << 6) | ((uint8_t)s[*pos + 1] & 0x3f); + *pos += 2; + return cp; } - if (ok && output_hc) { - ok = ds4_gpu_tensor_read(cur, - 0, - output_hc, - (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; - } - if (ok && logits_out) { - ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, - model, - weights); - } - if (ok && logits_out) { - last_hidden = glm_graph_tensor_row_view_strided(cur, - n_tokens - 1u, - DS4_N_EMBD, - DS4_N_EMBD); - ok = last_hidden != NULL; - if (ok && g->ssd_streaming) ok = glm_graph_stream_map_output(g, model, weights); - if (ok) ok = glm_graph_forward_output_head(g, model, weights, last_hidden, logits_out); - if (ok) { - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base, - n_tokens, - g->layer_count, - g->layer_count, - work_total, - true); - } - } - ds4_gpu_tensor_free(last_hidden); - ds4_gpu_set_glm_streaming_prefill_full_layer(false); - return ok; + if ((c & 0xf0) == 0xe0 && *pos + 2 < len) { + uint32_t cp = ((uint32_t)(c & 0x0f) << 12) | + ((uint32_t)((uint8_t)s[*pos + 1] & 0x3f) << 6) | + ((uint8_t)s[*pos + 2] & 0x3f); + *pos += 3; + return cp; + } + if ((c & 0xf8) == 0xf0 && *pos + 3 < len) { + uint32_t cp = ((uint32_t)(c & 0x07) << 18) | + ((uint32_t)((uint8_t)s[*pos + 1] & 0x3f) << 12) | + ((uint32_t)((uint8_t)s[*pos + 2] & 0x3f) << 6) | + ((uint8_t)s[*pos + 3] & 0x3f); + *pos += 4; + return cp; + } + (*pos)++; + return c; } -static bool glm_graph_use_streaming_token_prefill( - const ds4_glm_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens); -static uint32_t glm_graph_streaming_token_prefill_max_tokens(void); -static bool glm_graph_prefill_token_major( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - uint32_t pos0, - uint32_t n_tokens, - float *logits_out, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t display_absolute_base, - uint32_t work_done_base, - uint32_t work_total); - -static bool glm_graph_prefill_range( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - uint32_t pos0, - uint32_t n_tokens, - float *logits_out, - ds4_session_progress_fn progress, - void *progress_ud, - uint32_t progress_total) { - if (n_tokens == 0) return true; - if (!glm_graph_span_fits_context(g, pos0, n_tokens)) return false; - const uint32_t chunk_max = glm_graph_prefill_chunk_tokens(g->ctx_cap); - uint32_t done = 0; - while (done < n_tokens) { - const uint32_t pos = pos0 + done; - if (!g->full_kv_cache) { - const uint32_t remaining = n_tokens - done; - uint32_t chunk = 1; - if (glm_graph_use_streaming_token_prefill(g, pos, remaining)) { - chunk = remaining; - const uint32_t token_prefill_max = - glm_graph_streaming_token_prefill_max_tokens(); - if (token_prefill_max != 0 && chunk > token_prefill_max) { - chunk = token_prefill_max; - } - float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; - if (!glm_graph_prefill_token_major(g, - model, - weights, - tokens + done, - pos, - chunk, - dst_logits, - progress, - progress_ud, - pos0, - done, - n_tokens)) { - return false; - } - } else if (glm_graph_indexed_prefill_batch_ready(g, pos)) { - chunk = remaining; - if (chunk > g->indexed_prefill_cap) chunk = g->indexed_prefill_cap; - chunk = glm_graph_limit_indexed_prefill_chunk(pos, chunk); - if (chunk == 0) chunk = 1; - float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; - if (!glm_graph_forward_indexed_tokens(g, - model, - weights, - tokens + done, - NULL, - pos, - chunk, - NULL, - dst_logits, - progress, - progress_ud, - pos0, - done, - n_tokens)) { - return false; - } - } else { - float *dst_logits = (done + 1u == n_tokens) ? logits_out : NULL; - if (!glm_graph_forward_token(g, - model, - weights, - tokens[done], - NULL, - pos, - NULL, - dst_logits, - false)) { - return false; - } - } - done += chunk; - if (progress) { - const uint32_t current = pos0 + done; - progress(progress_ud, - "prefill_chunk", - current, - progress_total ? progress_total : pos0 + n_tokens); - } +static int gpt2_codepoint_to_byte(uint32_t cp) { + if ((cp >= 33 && cp <= 126) || (cp >= 161 && cp <= 172) || (cp >= 174 && cp <= 255)) { + return (int)cp; + } + + uint32_t n = 0; + for (uint32_t b = 0; b < 256; b++) { + if ((b >= 33 && b <= 126) || (b >= 161 && b <= 172) || (b >= 174)) { continue; } - if (pos >= g->ctx_cap) { - if (g->compact_cache_cap == 0) { - glm_graph_log_full_attention_limit(g, pos, n_tokens - done); - return false; - } - while (done < n_tokens) { - const uint32_t cur_pos = pos0 + done; - const bool use_indexed_batch = - glm_graph_indexed_prefill_batch_ready(g, cur_pos); - if (use_indexed_batch) { - uint32_t chunk = n_tokens - done; - if (chunk > g->indexed_prefill_cap) chunk = g->indexed_prefill_cap; - chunk = glm_graph_limit_indexed_prefill_chunk(cur_pos, chunk); - float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; - if (!glm_graph_forward_indexed_tokens(g, - model, - weights, - tokens + done, - NULL, - cur_pos, - chunk, - NULL, - dst_logits, - progress, - progress_ud, - pos0, - done, - n_tokens)) { - return false; - } - done += chunk; - } else { - float *dst_logits = (done + 1u == n_tokens) ? logits_out : NULL; - if (!glm_graph_forward_token(g, - model, - weights, - tokens[done], - NULL, - cur_pos, - NULL, - dst_logits, - false)) { - return false; - } - done++; - } - if (progress) { - const uint32_t current = pos0 + done; - progress(progress_ud, - "prefill_chunk", - current, - progress_total ? progress_total : pos0 + n_tokens); - } - } - return true; - } - uint32_t chunk = n_tokens - done; - const uint32_t full_remaining = g->ctx_cap - pos; - if (chunk > full_remaining) chunk = full_remaining; - if (chunk > chunk_max) chunk = chunk_max; - float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; - if (glm_graph_use_streaming_token_prefill(g, pos, chunk)) { - if (!glm_graph_prefill_token_major(g, - model, - weights, - tokens + done, - pos, - chunk, - dst_logits, - progress, - progress_ud, - pos0, - done, - n_tokens)) { - return false; - } - } else if (!glm_graph_forward_tokens(g, - model, - weights, - tokens + done, - NULL, - pos, - chunk, - NULL, - dst_logits, - progress, - progress_ud, - pos0, - done, - n_tokens)) { - return false; - } - done += chunk; - if (progress) { - const uint32_t current = pos0 + done; - progress(progress_ud, - "prefill_chunk", - current, - progress_total ? progress_total : pos0 + n_tokens); - } + if (cp == 256 + n) return (int)b; + n++; + } + return -1; +} + +static bool vocab_token_is_literal_special(ds4_str s) { + const unsigned char bar[] = {0xef, 0xbd, 0x9c}; /* U+FF5C fullwidth vertical bar. */ + if (s.len < sizeof(bar)) return false; + for (uint64_t i = 0; i + sizeof(bar) <= s.len; i++) { + if (!memcmp(s.ptr + i, bar, sizeof(bar))) return true; + } + return false; +} + +char *ds4_token_text(ds4_engine *e, int token, size_t *len) { + ds4_vocab *vocab = &e->vocab; + if (token < 0 || token >= vocab->n_vocab) { + if (len) *len = 0; + char *out = xmalloc(1); + out[0] = '\0'; + return out; + } + + ds4_str s = vocab->token[token]; + char *out = xmalloc((size_t)s.len + 1); + if (vocab_token_is_literal_special(s)) { + memcpy(out, s.ptr, (size_t)s.len); + out[s.len] = '\0'; + if (len) *len = (size_t)s.len; + return out; + } + + size_t n = 0; + uint64_t pos = 0; + while (pos < s.len) { + uint32_t cp = utf8_decode_one(s.ptr, s.len, &pos); + int b = gpt2_codepoint_to_byte(cp); + if (b >= 0) out[n++] = (char)b; + } + out[n] = '\0'; + if (len) *len = n; + return out; +} + +static bool vocab_token_is_generation_stop(const ds4_vocab *vocab, int token) { + if (!vocab || token < 0) return false; + if (token == vocab->eos_id) return true; + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + return (vocab->system_id >= 0 && token == vocab->system_id) || + (vocab->user_id >= 0 && token == vocab->user_id) || + (vocab->assistant_id >= 0 && token == vocab->assistant_id) || + (vocab->observation_id >= 0 && token == vocab->observation_id); } - return true; + return false; } -/* - * For very short GLM SSD-streaming prefills, Metal still benefits from the - * token-major path because it reuses the normal decode graph and warms the - * decode expert cache. On ROCm/Strix Halo the indexed batch prefill is faster - * now that streamed batch routing and expert cache seeding are implemented, so - * ROCm defaults to canonical batch prefill unless the env override below opts - * token-major prefill back in. - */ -enum { DS4_GLM_STREAM_PREFILL_TOKEN_MAJOR_MAX_TOKENS = 64 }; +int ds4_token_eos(ds4_engine *e) { + return e->vocab.eos_id; +} -static uint32_t glm_graph_streaming_token_prefill_default_max_tokens(void) { -#ifdef DS4_ROCM_BUILD - return 0; -#else - return DS4_GLM_STREAM_PREFILL_TOKEN_MAJOR_MAX_TOKENS; -#endif +bool ds4_token_is_stop(ds4_engine *e, int token) { + return e ? vocab_token_is_generation_stop(&e->vocab, token) : false; } -static uint32_t glm_graph_streaming_token_prefill_max_tokens(void) { - const char *env = glm_graph_env_value( - "DS4_ROCM_GLM_STREAMING_TOKEN_PREFILL_MAX", - "DS4_METAL_GLM_STREAMING_TOKEN_PREFILL_MAX"); - if (!env || !env[0]) env = getenv("DS4_GLM_STREAMING_TOKEN_PREFILL_MAX"); - const uint32_t default_max = - glm_graph_streaming_token_prefill_default_max_tokens(); - if (!env || !env[0]) return default_max; - char *end = NULL; - errno = 0; - unsigned long v = strtoul(env, &end, 10); - if (end == env || errno != 0 || v > UINT32_MAX) { - return default_max; - } - return (uint32_t)v; +bool ds4_token_is_thinking_control(ds4_engine *e, int token) { + if (!e || token < 0) return false; + return (e->vocab.think_start_id >= 0 && + token == e->vocab.think_start_id) || + (e->vocab.think_end_id >= 0 && + token == e->vocab.think_end_id); } -static bool glm_graph_use_streaming_token_prefill( - const ds4_glm_gpu_graph *g, - uint32_t pos0, - uint32_t n_tokens) { - if (!g || !g->ssd_streaming || g->quality || n_tokens == 0) return false; - if (getenv("DS4_GLM_DISABLE_STREAMING_TOKEN_PREFILL") != NULL || - glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_TOKEN_PREFILL", - "DS4_METAL_GLM_DISABLE_STREAMING_TOKEN_PREFILL")) { - return false; +bool ds4_token_is_stop_for_think_mode( + ds4_engine *e, + int token, + ds4_think_mode mode) { + if (ds4_token_is_stop(e, token)) return true; + /* In no-thinking mode the prompt already supplied the protocol close tag. + * If the model emits another thinking tag, do not print or feed it back: + * it is a control marker, not assistant content. */ + if (!ds4_think_mode_enabled(mode) && + ds4_token_is_thinking_control(e, token)) { + return true; } - if (!glm_graph_span_fits_full_attention(g, pos0, n_tokens)) return false; - const uint32_t max_tokens = glm_graph_streaming_token_prefill_max_tokens(); - return max_tokens != 0 && n_tokens <= max_tokens; + return false; } -static bool glm_graph_prefill_token_major( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - const int *tokens, - uint32_t pos0, - uint32_t n_tokens, - float *logits_out, - ds4_session_progress_fn display_progress, - void *display_progress_ud, - uint32_t display_absolute_base, - uint32_t work_done_base, - uint32_t work_total) { - if (!g || !model || !weights || !tokens || n_tokens == 0) return false; - for (uint32_t i = 0; i < n_tokens; i++) { - const bool last = i + 1u == n_tokens; - float *dst_logits = (last && logits_out) ? logits_out : NULL; - if (!glm_graph_forward_token(g, - model, - weights, - tokens[i], - NULL, - pos0 + i, - NULL, - dst_logits, - false)) { - return false; +int ds4_token_user(ds4_engine *e) { + return e->vocab.user_id; +} + +int ds4_token_assistant(ds4_engine *e) { + return e->vocab.assistant_id; +} + +static int sample_argmax(const float *logits, uint32_t n_vocab) { + int best = 0; + float best_v = DS4_NEG_INF; + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + if (v > best_v) { + best_v = v; + best = (int)i; } - glm_graph_report_prefill_display_progress(display_progress, - display_progress_ud, - display_absolute_base, - work_done_base + i + 1u, - n_tokens, - 0, - 1, - work_total, - false); } - return true; + return best; } -static bool glm_graph_maybe_warm_compact_indexer_after_prefill( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - uint32_t next_pos) { - if (!g || !model || !weights) return false; - if (g->compact_cache_cap == 0 || - g->indexer_full_layers == 0 || - next_pos < g->ctx_cap) { - return true; +static DS4_MAYBE_UNUSED void logits_top2(const float *logits, uint32_t n_vocab, + int *top0, float *logit0, + int *top1, float *logit1) { + int b0 = -1, b1 = -1; + float v0 = DS4_NEG_INF, v1 = DS4_NEG_INF; + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + if (v > v0) { + b1 = b0; v1 = v0; + b0 = (int)i; v0 = v; + } else if (v > v1) { + b1 = (int)i; v1 = v; + } } - if (next_pos >= g->ctx_size) return true; - return glm_graph_warm_compact_indexer_store(g, model, weights, next_pos); + if (top0) *top0 = b0; + if (logit0) *logit0 = v0; + if (top1) *top1 = b1; + if (logit1) *logit1 = v1; } -static bool glm_graph_begin_commands_if_needed(void) { - return ds4_gpu_commands_active() || ds4_gpu_begin_commands() != 0; +static uint64_t sample_rng_next(uint64_t *state) { + uint64_t x = *state; + if (x == 0) x = 0x9e3779b97f4a7c15ULL; + x ^= x >> 12; + x ^= x << 25; + x ^= x >> 27; + *state = x; + return x * 0x2545f4914f6cdd1dULL; } -static bool glm_graph_end_commands_if_active(void) { - return !ds4_gpu_commands_active() || ds4_gpu_end_commands() != 0; +static float sample_rng_f32(uint64_t *state) { + const uint64_t x = sample_rng_next(state); + return (float)((x >> 40) & 0xffffffu) / 16777216.0f; } -static bool glm_graph_streaming_decode_sync_each_layer(void) { -#ifdef DS4_ROCM_BUILD - const char *env = glm_graph_env_value( - "DS4_ROCM_GLM_STREAMING_DECODE_SYNC_EACH_LAYER", - "DS4_METAL_GLM_STREAMING_DECODE_SYNC_EACH_LAYER"); - if (!env) env = getenv("DS4_GLM_STREAMING_DECODE_SYNC_EACH_LAYER"); - return glm_graph_env_truthy(env); -#else - return true; -#endif -} +typedef struct { + int id; + float logit; + float prob; +} sample_candidate; -static bool glm_graph_forward_token( - ds4_glm_gpu_graph *g, - const ds4_model *model, - const ds4_weights *weights, - int token, - const float *input_hc, - uint32_t pos, - float *output_hc, - float *logits_out, - bool defer_completion) { -#define DS4_GLM_FT_FAIL(why) do { \ - if (getenv("DS4_GLM_TP_DEBUG")) \ - fprintf(stderr, "ds4: glm forward_token fail pos=%u: %s\n", pos, why); \ - } while (0) - if (!g || !model || !weights || - token < 0 || token >= (int)DS4_N_VOCAB || - g->layer_count == 0 || - pos >= g->ctx_size || - (defer_completion && - (g->ssd_streaming || !ds4_gpu_commands_active()))) { - DS4_GLM_FT_FAIL("arg guard"); - return false; - } - if (!input_hc && !g->has_token_embd) { DS4_GLM_FT_FAIL("no token embd"); return false; } - if (logits_out && !g->has_output_head) { DS4_GLM_FT_FAIL("no output head"); return false; } - const bool use_indexed_attention = - glm_graph_decode_uses_indexed_attention(g, pos, logits_out); - uint32_t decode_layer_flush_interval = 0; - if (logits_out != NULL) { - decode_layer_flush_interval = use_indexed_attention ? 4u : 32u; - const char *dfi = getenv("DS4_GLM_DECODE_FLUSH_INTERVAL"); - if (dfi && dfi[0]) { - int v = atoi(dfi); - decode_layer_flush_interval = v <= 0 ? 0u : (uint32_t)v; - } - if (decode_layer_flush_interval > g->layer_count) { - decode_layer_flush_interval = g->layer_count; - } - if (defer_completion) decode_layer_flush_interval = 0; - } - if (pos >= g->ctx_cap && !use_indexed_attention) { - glm_graph_log_full_attention_limit(g, pos, 1); - DS4_GLM_FT_FAIL("full attention limit"); - return false; - } - if (g->compact_cache_cap != 0 && - !glm_graph_ensure_compact_cache(g, pos + 1u)) { - DS4_GLM_FT_FAIL("compact cache ensure"); - return false; - } +static int sample_candidate_cmp_desc(const void *a, const void *b) { + const sample_candidate *ca = a; + const sample_candidate *cb = b; + const int logit_order = + (cb->logit > ca->logit) - (cb->logit < ca->logit); + if (logit_order != 0) return logit_order; + return (ca->id > cb->id) - (ca->id < cb->id); +} - const bool decode_output_profile = false; - const bool merge_indexed_output = - logits_out != NULL && use_indexed_attention && !decode_output_profile; - double decode_output_stage_t0 = decode_output_profile ? now_sec() : 0.0; - const bool decode_flush_profile = false; - uint32_t decode_flush_layer0 = 0; - double decode_flush_stage_t0 = decode_flush_profile ? now_sec() : 0.0; +static bool sample_candidate_gt(sample_candidate a, sample_candidate b) { + if (a.logit != b.logit) return a.logit > b.logit; + return a.id < b.id; +} - const bool static_decode_map = - !input_hc && - g->has_token_embd && - g->ssd_streaming && - metal_graph_stream_decode_static_map_enabled(); - const bool static_map_state_cache = - static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); - const bool streaming_decode_sync_each_layer = - g->ssd_streaming && - !static_decode_map && - glm_graph_streaming_decode_sync_each_layer(); - bool ok = true; - if (input_hc) { - ok = ds4_gpu_tensor_write(g->cur, - 0, - input_hc, - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - } else if (static_decode_map) { - if (!static_map_state_cache || !g->streaming_static_decode_map_current) { - ok = metal_graph_stream_map_decode_static_all(model, weights); - if (ok) g->streaming_static_decode_map_current = static_map_state_cache; - } - } else { - ok = glm_graph_stream_map_token(g, model, weights); - } - if (ok) ok = glm_graph_begin_commands_if_needed(); - if (ok && !input_hc) { - if (g->placement) { - ok = glm_graph_ws_switch(g, g->placement[0], false); - } - } - if (ok && !input_hc) { - ok = ds4_gpu_embed_token_quant_tensor(g->cur, - model->map, - model->size, - weights->token_embd->abs_offset, - weights->token_embd->type, - DS4_N_VOCAB, - (uint32_t)token, - DS4_N_EMBD) != 0; - } - if (ok && streaming_decode_sync_each_layer) { - ok = ds4_gpu_end_commands() != 0; - } - const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); - ds4_gpu_tensor *last_indexer_selected = NULL; - uint32_t last_indexer_selected_count = 0; -#define DS4_GLM_PROFILE_DECODE_STAGE(part_, name_) do { \ - if (ok && decode_stage_profile) { \ - ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos, 1, &decode_stage_t0); \ - } \ - } while (0) - uint32_t glm_ft_fail_il = UINT32_MAX; - for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { - if (g->placement) { - ok = glm_graph_ws_switch(g, g->placement[il + 1u], true); - if (!ok) break; - } - glm_ft_fail_il = il; - const uint32_t slice_layer_done = il - g->layer_start + 1u; - if (g->ssd_streaming) { - if (!static_decode_map) { - ok = glm_graph_stream_map_decode_layer(g, model, weights, il); - } - if (ok) ok = glm_graph_begin_commands_if_needed(); - } - const ds4_layer_weights *l = &weights->layer[il]; - const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; - const float rope_base = layer_rope_freq_base(il); - const float rope_scale = layer_rope_freq_scale(il); - const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); - double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; - if (decode_stage_profile) { - ok = metal_graph_layer_stage_profile_boundary("glm_decode_attn", - NULL, - il, - pos, - 1, - &decode_stage_t0); - } - - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->attn_norm, - g->cur, - model->map, - model->size, - l->attn_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_norm"); - const uint32_t decode_ablate = glm_decode_ablate_mask(); - if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q_rank, - model, - l->attn_q_a->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - g->attn_norm, - il, - pos, - "attn_q_a", - g->ssd_streaming) != 0; - } - const bool fuse_qkv_norm_store = use_indexed_attention && - !decode_stage_profile && - g->compact_cache_cap != 0; - const bool fuse_qkv_norm = !decode_stage_profile && !fuse_qkv_norm_store; - if (ok && fuse_qkv_norm_store) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_store", - g->ssd_streaming) != 0; - if (ok) { - ok = ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( - g->q_rank_norm, - g->q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - g->kv_raw, - l->attn_kv_a_norm->abs_offset, - pos, - 1, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16(), - DS4_RMS_EPS) != 0; - } - } else if (ok && fuse_qkv_norm) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_norm", - g->ssd_streaming) != 0; - if (ok) { - ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->q_rank_norm, - g->q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - g->kv_norm, - g->kv_raw, - l->attn_kv_a_norm->abs_offset, - DS4_N_KV_LORA, - 1, - DS4_RMS_EPS) != 0; - } - } else if (ok) { - ok = ds4_gpu_rms_norm_weight_tensor(g->q_rank_norm, - g->q_rank, - model->map, - model->size, - l->attn_q_a_norm->abs_offset, - DS4_N_LORA_Q, - DS4_RMS_EPS) != 0; - } - if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->q_rank_norm, - il, - pos, - "attn_q_b", - g->ssd_streaming) != 0; - if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->q, - 1, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_ROT, - pos, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "q_path"); - if (ok) metal_graph_debug_dump_tensor("glm_decode_q", - g->q, - (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA, - il, - pos); - if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il) && - !(decode_ablate & DS4_GLM_ABLATE_INDEXER)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->indexer_k, - model, - l->indexer_attn_k->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD_DIM, - g->cur, - il, - pos, - "indexer_k", - g->ssd_streaming) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k_proj"); - if (ok) { - ok = ds4_gpu_glm_store_indexer_k_tensor( - g->layer_indexer_key_cache[il], - g->indexer_k, - model->map, - model->size, - l->indexer_k_norm->abs_offset, - l->indexer_k_norm_b->abs_offset, - pos, - 1, - g->compact_cache_cap, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - 0, - 1.0e-6f, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - glm_graph_compact_cache_is_f16()) != 0; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k_store"); - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k"); - if (ok && !fuse_qkv_norm && !fuse_qkv_norm_store) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a", - g->ssd_streaming) != 0; - if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, - g->kv_raw, - model->map, - model->size, - l->attn_kv_a_norm->abs_offset, - 1, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_RMS_EPS) != 0; - } - if (ok && g->compact_cache_cap != 0 && !fuse_qkv_norm_store) { - ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - g->kv_norm, - g->kv_raw, - pos, - 1, - g->compact_cache_cap, - kv_raw_dim, - DS4_N_KV_LORA, - DS4_N_ROT, - glm_graph_compact_cache_is_f16()) != 0; - } - if (use_indexed_attention) { - if (ok && glm_graph_layer_uses_full_indexer(il)) { - const uint32_t visible = pos + 1u; - if (ok && visible <= indexer_top_k) { - ok = ds4_gpu_glm_fill_selected_range_tensor(g->indexer_selected, - visible) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_fill"); - last_indexer_selected_count = visible; - } else if (ok && (decode_ablate & DS4_GLM_ABLATE_INDEXER)) { - /* Ablation: valid selected ids without the score/topk - * chain, so downstream attention timing stays real. */ - ok = ds4_gpu_glm_fill_selected_range_tensor(g->indexer_selected, - indexer_top_k) != 0; - last_indexer_selected_count = indexer_top_k; - } else if (ok) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->indexer_q, - model, - l->indexer_attn_q_b->abs_offset, - DS4_N_LORA_Q, - (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - g->q_rank_norm, - il, - pos, - "indexer_q", - g->ssd_streaming) != 0; - if (ok) ok = ds4_gpu_glm_indexer_rope_tail_tensor(g->indexer_q, - 1, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - DS4_N_ROT, - pos, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok && glm_graph_indexer_qat()) { - ok = ds4_gpu_dsv4_indexer_qat_tensor(g->indexer_q, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM) != 0; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_q"); - if (ok) ok = ds4_gpu_matmul_f32_tensor(g->indexer_weights, - model->map, - model->size, - l->indexer_proj->abs_offset, - DS4_N_EMBD, - DS4_N_INDEXER_HEAD, - g->cur, - 1) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_weights"); - const float indexer_scale = - 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); - ok = ds4_gpu_glm_indexer_score_one_tensor(g->indexer_scores, - g->indexer_q, - g->indexer_weights, - g->layer_indexer_key_cache[il], - visible, - DS4_N_INDEXER_HEAD, - DS4_N_INDEXER_HEAD_DIM, - indexer_scale, - glm_graph_compact_cache_is_f16()) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_scores"); - if (ok) ok = ds4_gpu_indexer_topk_tensor(g->indexer_selected, - g->indexer_scores, - visible, - 1, - indexer_top_k) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_topk"); - last_indexer_selected_count = indexer_top_k; - } - if (ok) last_indexer_selected = g->indexer_selected; - } else if (ok && (!last_indexer_selected || last_indexer_selected_count == 0)) { - ok = false; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_select"); - if (ok && !(decode_ablate & (DS4_GLM_ABLATE_ATTN_CORE | DS4_GLM_ABLATE_QKLOW))) { - ok = ds4_gpu_glm_qk_lowrank_typed_tensor(g->qk_low, - g->q, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_KEY_MLA) != 0; - if (ok) metal_graph_debug_dump_tensor("glm_decode_qk_low", - g->qk_low, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, - il, - pos); - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_path"); - if (ok && (decode_ablate & DS4_GLM_ABLATE_ATTN_CORE)) { - /* Skip the indexed attention kernels; zero heads so the - * rest of the layer stays finite (timing-only). */ - ok = ds4_gpu_tensor_fill_f32(g->heads, 0.0f, - (uint64_t)g->heads_dim) != 0; - } else if (ok && glm_graph_indexed_decode_split_group8_available(last_indexer_selected_count)) { - const uint32_t split_block_rows = - glm_graph_indexed_decode_split_block_rows_for(last_indexer_selected_count); - const uint32_t split_blocks = - (last_indexer_selected_count + split_block_rows - 1u) / split_block_rows; - ok = ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(g->heads, - g->attn_partial_lora, - g->attn_partial_ms, - g->q, - g->qk_low, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - last_indexer_selected, - last_indexer_selected_count, - true, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - split_block_rows, - split_blocks, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } else if (ok) { - ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor(g->heads, - g->q, - g->qk_low, - g->layer_kv_lora_cache[il], - g->layer_k_rope_cache[il], - model->map, - model->size, - l->attn_v_b->abs_offset, - l->attn_v_b->type, - last_indexer_selected, - last_indexer_selected_count, - g->compact_cache_cap, - glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW) != 0; - } - } else { - if (ok) ok = ds4_gpu_glm_k_b_project_typed_tensor(g->k_nope, - g->kv_norm, - model->map, - model->size, - l->attn_k_b->abs_offset, - l->attn_k_b->type, - 1, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_HEAD) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->value, - model, - l->attn_v_b->abs_offset, - DS4_N_KV_LORA, - g->heads_dim, - g->kv_norm, - il, - pos, - "attn_v_b", - g->ssd_streaming) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_path"); - if (ok) ok = ds4_gpu_glm_build_kv_cache_tensor(g->layer_key_cache[il], - g->layer_value_cache[il], - g->kv_raw, - g->k_nope, - g->value, - pos, - 1, - g->ctx_cap, - DS4_N_HEAD, - kv_raw_dim, - DS4_N_KV_LORA, - (uint32_t)g->q_nope, - DS4_N_ROT, - DS4_N_VALUE_MLA, - 0, - rope_base, - rope_scale, - 0.0f, - 1.0f, - DS4_ROPE_YARN_BETA_FAST, - DS4_ROPE_YARN_BETA_SLOW, - true) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_cache"); - if (ok) ok = ds4_gpu_glm_attention_full_tensor(g->heads, - g->q, - g->layer_key_cache[il], - g->layer_value_cache[il], - pos, - 1, - pos + 1u, - g->ctx_cap, - DS4_N_HEAD, - DS4_N_KEY_MLA, - DS4_N_VALUE_MLA, - true) != 0; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attention"); - if (ok) metal_graph_debug_dump_tensor("glm_decode_heads", - g->heads, - g->heads_dim, - il, - pos); - if (ok && !(decode_ablate & DS4_GLM_ABLATE_ATTN_OUT)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->heads, - il, - pos, - "attn_o", - g->ssd_streaming) != 0; - } - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_output"); - if (ok) ok = ds4_gpu_add_rms_norm_weight_tensor(g->ffn_norm, - g->after_attn, - g->cur, - g->attn_out, - model->map, - model->size, - l->ffn_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; - DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_ffn", "ffn_norm"); - if (ok) ok = glm_graph_encode_ffn_one_normed_from(g, - model, - l, - il, - pos, - g->ffn_norm, - g->after_attn, - g->next, - g->ffn_gate, - g->ffn_up, - g->ffn_mid, - g->ffn_out, - g->ffn_sum, - g->attn_out, - decode_stage_profile, - decode_stage_profile ? &decode_stage_t0 : NULL); - if (ok) { - ds4_gpu_tensor *tmp = g->cur; - g->cur = g->next; - g->next = tmp; - } - if (ok && glm_debug_hidden_dump_layer_match(il)) { - ok = ds4_gpu_end_commands() != 0; - if (ok) { - glm_debug_dump_hidden_layer(g->cur, 0, il, pos); - glm_debug_dump_raw_layer(g->router_selected, "sel", - (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t), - il, (int)pos); - glm_debug_dump_raw_layer(g->router_weights, "selw", - (uint64_t)DS4_N_EXPERT_USED * sizeof(float), - il, (int)pos); - ok = ds4_gpu_begin_commands() != 0; - } - } - if (ok && - !g->ssd_streaming && - decode_layer_flush_interval != 0 && - il < g->layer_end && - (slice_layer_done % decode_layer_flush_interval) == 0) { - if (decode_flush_profile) { - ok = ds4_gpu_flush_commands() != 0; - if (ok) ok = ds4_gpu_synchronize() != 0; - if (ok) { - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM decode layer flush pos=%u layers=%u..%u %.3f ms\n", - pos, - decode_flush_layer0, - il, - (now - decode_flush_stage_t0) * 1000.0); - decode_flush_layer0 = il + 1u; - decode_flush_stage_t0 = now; - ok = ds4_gpu_begin_commands() != 0; - } - } else { - ok = ds4_gpu_flush_commands() != 0; - } - } - if (streaming_decode_sync_each_layer) { - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - } +static void sample_heap_sift_up(sample_candidate *heap, uint32_t idx) { + while (idx > 0) { + const uint32_t parent = (idx - 1u) / 2u; + if (!sample_candidate_gt(heap[parent], heap[idx])) break; + sample_candidate tmp = heap[parent]; + heap[parent] = heap[idx]; + heap[idx] = tmp; + idx = parent; } -#undef DS4_GLM_PROFILE_DECODE_STAGE - if (ok && (merge_indexed_output || - (defer_completion && logits_out != NULL))) { - if (g->ssd_streaming) { - if (!static_decode_map) { - ok = glm_graph_stream_map_output(g, model, weights); - } - if (ok) ok = glm_graph_begin_commands_if_needed(); +} + +static void sample_heap_sift_down(sample_candidate *heap, uint32_t n, uint32_t idx) { + for (;;) { + const uint32_t left = idx * 2u + 1u; + const uint32_t right = left + 1u; + uint32_t smallest = idx; + if (left < n && sample_candidate_gt(heap[smallest], heap[left])) { + smallest = left; } - ok = glm_graph_encode_output_head(g, model, weights); - if (g->ssd_streaming) { - if (ok) ok = glm_graph_end_commands_if_active(); - else (void)ds4_gpu_synchronize(); + if (right < n && sample_candidate_gt(heap[smallest], heap[right])) { + smallest = right; } + if (smallest == idx) break; + sample_candidate tmp = heap[idx]; + heap[idx] = heap[smallest]; + heap[smallest] = tmp; + idx = smallest; } - if (!g->ssd_streaming && !defer_completion) { - if (ok) ok = ds4_gpu_end_commands() != 0; - else (void)ds4_gpu_synchronize(); - } else if (!ok) { - (void)ds4_gpu_synchronize(); - } - if (decode_output_profile) { - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM decode output profile pos=%u layers=%.3f ms\n", - pos, - (now - decode_output_stage_t0) * 1000.0); - decode_output_stage_t0 = now; - } - if (ok && output_hc && !defer_completion) { - ok = ds4_gpu_tensor_read(g->cur, - 0, - output_hc, - (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; - } - if (ok && logits_out && !defer_completion) { - if (use_indexed_attention) { - if (!merge_indexed_output) { - if (g->ssd_streaming && !static_decode_map) { - ok = glm_graph_stream_map_output(g, model, weights); - } - if (ok) ok = glm_graph_begin_commands_if_needed(); - if (ok) ok = glm_graph_encode_output_head(g, model, weights); - if (ok) ok = glm_graph_end_commands_if_active(); - else (void)ds4_gpu_synchronize(); - if (decode_output_profile) { - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM decode output profile pos=%u output_head=%.3f ms\n", - pos, - (now - decode_output_stage_t0) * 1000.0); - decode_output_stage_t0 = now; - } - } - if (ok) { - if (glm_debug_hidden_dump_layer() < 0) - glm_debug_dump_hidden_row(g->cur, 0); - ok = ds4_gpu_tensor_read(g->logits, - 0, - logits_out, - (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; - if (decode_output_profile) { - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM decode output profile pos=%u logits_read=%.3f ms\n", - pos, - (now - decode_output_stage_t0) * 1000.0); - } - } - } else { - if (g->ssd_streaming && !static_decode_map) { - ok = glm_graph_stream_map_output(g, model, weights); - } - if (ok && ds4_gpu_commands_active()) { - ok = ds4_gpu_end_commands() != 0; - } - if (ok) ok = glm_graph_forward_output_head(g, model, weights, g->cur, logits_out); - if (decode_output_profile) { - const double now = now_sec(); - fprintf(stderr, - "ds4: GLM decode output profile pos=%u fallback_output=%.3f ms\n", - pos, - (now - decode_output_stage_t0) * 1000.0); - } +} + +static bool sample_fast_top_p( + const float *logits, + uint32_t n_vocab, + uint32_t finite, + float max_logit, + int best, + float temperature, + float top_p, + float min_p, + uint64_t *rng, + int *token_out) { + enum { SAMPLE_FAST_TOP_P_CAP = 512 }; + if (!logits || !rng || !token_out || finite == 0) return false; + if (finite > SAMPLE_FAST_TOP_P_CAP && top_p >= 0.999f) return false; + + const uint32_t cap = finite < SAMPLE_FAST_TOP_P_CAP ? + finite : (uint32_t)SAMPLE_FAST_TOP_P_CAP; + sample_candidate heap[SAMPLE_FAST_TOP_P_CAP]; + uint32_t n = 0; + float sum = 0.0f; + float heap_sum = 0.0f; + + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + if (!isfinite(v)) continue; + const float p = expf((v - max_logit) / temperature); + sum += p; + sample_candidate cand = {.id = (int)i, .logit = v, .prob = p}; + if (n < cap) { + heap[n] = cand; + heap_sum += p; + sample_heap_sift_up(heap, n); + n++; + } else if (sample_candidate_gt(cand, heap[0])) { + heap_sum -= heap[0].prob; + heap[0] = cand; + heap_sum += p; + sample_heap_sift_down(heap, n, 0); } } - if (ok && - !logits_out && - !output_hc && - g->ssd_streaming && - !streaming_decode_sync_each_layer) { - ok = ds4_gpu_end_commands() != 0; - } else if (ok && !logits_out && g->ssd_streaming) { - ok = glm_graph_end_commands_if_active(); - } else if (!ok) { - (void)ds4_gpu_synchronize(); - } - if (!ok && getenv("DS4_GLM_TP_DEBUG")) { - fprintf(stderr, - "ds4: glm forward_token fail pos=%u around layer %u\n", - pos, glm_ft_fail_il); + if (sum <= 0.0f || !isfinite(sum)) { + *token_out = best; + return true; } - (void)glm_ft_fail_il; - return ok; -#undef DS4_GLM_FT_FAIL -} -static int glm_metal_first_token_logits( - const ds4_model *model, - const ds4_weights *weights, - int token, - float *logits_out) { - if (!model || !weights || !logits_out) return 1; - if (token < 0 || token >= (int)DS4_N_VOCAB) { - fprintf(stderr, "ds4: GLM token %d is outside vocab\n", token); - return 1; + if (n < finite && heap_sum < top_p * sum) { + return false; } - if (!weights->token_embd || weights->token_embd->type != DS4_TENSOR_Q8_0 || - !weights->output_norm || weights->output_norm->type != DS4_TENSOR_F32 || - !weights->output || weights->output->type != DS4_TENSOR_Q8_0 || - weights->output_norm->dim[0] != DS4_N_EMBD || - weights->output->dim[0] != DS4_N_EMBD || - weights->output->dim[1] != DS4_N_VOCAB) { - fprintf(stderr, "ds4: GLM Metal first-token path found unexpected embedding/output layout\n"); - return 1; + + qsort(heap, n, sizeof(heap[0]), sample_candidate_cmp_desc); + const float min_prob = (heap[0].prob / sum) * (min_p > 0.0f ? min_p : 0.0f); + const float min_prob_raw = heap[0].prob * (min_p > 0.0f ? min_p : 0.0f); + float filtered_sum = 0.0f; + uint32_t filtered = 0; + bool stopped_by_min_p = false; + for (uint32_t i = 0; i < n; i++) { + const float p = heap[i].prob / sum; + if (i > 0 && p < min_prob) { + stopped_by_min_p = true; + break; + } + filtered_sum += heap[i].prob; + filtered++; + if (filtered_sum / sum >= top_p) break; } - if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) { - fprintf(stderr, "ds4: GLM Metal first-token path has no normal transformer layers\n"); - return 1; + if (n < finite && + stopped_by_min_p && + min_p > 0.0f && + heap[n - 1u].prob >= min_prob_raw) { + return false; } - - const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; - const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; - uint64_t kv_raw_dim = 0; - uint64_t dense_hidden_max = DS4_N_FF_EXP; - bool generic_routed_moe = false; - for (uint32_t il = 0; il < normal_layers; il++) { - const ds4_layer_weights *l = &weights->layer[il]; - if (l->attn_kv_a_mqa && l->attn_kv_a_mqa->dim[1] > kv_raw_dim) { - kv_raw_dim = l->attn_kv_a_mqa->dim[1]; - } - if (il < DS4_N_LEADING_DENSE && l->ffn_gate && - l->ffn_gate->dim[1] > dense_hidden_max) { - dense_hidden_max = l->ffn_gate->dim[1]; - } - if (glm_graph_layer_uses_generic_routed_moe(l)) generic_routed_moe = true; - } - if (kv_raw_dim < DS4_N_KV_LORA) { - fprintf(stderr, "ds4: GLM Metal first-token path found no valid KV projection\n"); - return 1; + if (filtered == 0) { + *token_out = best; + return true; } - const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); - const uint64_t sparse_mid_elems = (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; - const uint64_t ffn_mid_elems = - dense_hidden_max > sparse_mid_elems ? dense_hidden_max : sparse_mid_elems; - const uint64_t routed_mid_bytes = - (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float); - const uint64_t routed_down_bytes = - (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); - const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); - - ds4_gpu_tensor *cur = NULL; - ds4_gpu_tensor *attn_norm = NULL; - ds4_gpu_tensor *kv_raw = NULL; - ds4_gpu_tensor *kv_norm = NULL; - ds4_gpu_tensor *heads = NULL; - ds4_gpu_tensor *attn_out = NULL; - ds4_gpu_tensor *after_attn = NULL; - ds4_gpu_tensor *ffn_norm = NULL; - ds4_gpu_tensor *ffn_gate = NULL; - ds4_gpu_tensor *ffn_up = NULL; - ds4_gpu_tensor *ffn_mid = NULL; - ds4_gpu_tensor *routed_gate = NULL; - ds4_gpu_tensor *routed_up = NULL; - ds4_gpu_tensor *routed_down = NULL; - ds4_gpu_tensor *ffn_out = NULL; - ds4_gpu_tensor *ffn_sum = NULL; - ds4_gpu_tensor *next = NULL; - ds4_gpu_tensor *router_logits = NULL; - ds4_gpu_tensor *router_probs = NULL; - ds4_gpu_tensor *router_selected = NULL; - ds4_gpu_tensor *router_weights = NULL; - ds4_gpu_tensor *logits = NULL; - - int ok = 1; -#define DS4_GLM_FIRST_ALLOC_TENSOR(var, bytes_) \ - do { \ - (var) = ds4_gpu_tensor_alloc((bytes_)); \ - if (!(var)) { \ - fprintf(stderr, "ds4: GLM Metal first-token path could not allocate %s\n", #var); \ - ok = 0; \ - } \ - } while (0) - - DS4_GLM_FIRST_ALLOC_TENSOR(cur, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(attn_norm, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(kv_raw, kv_raw_dim * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(kv_norm, (uint64_t)DS4_N_KV_LORA * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(heads, heads_dim * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(attn_out, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(after_attn, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_norm, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_gate, dense_hidden_max * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_up, dense_hidden_max * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_mid, ffn_mid_elems * sizeof(float)); - if (generic_routed_moe) { - DS4_GLM_FIRST_ALLOC_TENSOR(routed_gate, routed_mid_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(routed_up, routed_mid_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(routed_down, routed_down_bytes); - } - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_out, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(ffn_sum, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(next, emb_bytes); - DS4_GLM_FIRST_ALLOC_TENSOR(router_logits, (uint64_t)DS4_N_EXPERT * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(router_probs, (uint64_t)DS4_N_EXPERT * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); - DS4_GLM_FIRST_ALLOC_TENSOR(router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); - DS4_GLM_FIRST_ALLOC_TENSOR(logits, logits_bytes); -#undef DS4_GLM_FIRST_ALLOC_TENSOR - - if (ok) { - ok = ds4_gpu_embed_token_q8_0_tensor(cur, - model->map, - model->size, - weights->token_embd->abs_offset, - DS4_N_VOCAB, - (uint32_t)token, - DS4_N_EMBD); + float r = sample_rng_f32(rng) * filtered_sum; + for (uint32_t i = 0; i < filtered; i++) { + r -= heap[i].prob; + if (r <= 0.0f) { + *token_out = heap[i].id; + return true; + } } - for (uint32_t il = 0; ok && il < normal_layers; il++) { - const ds4_layer_weights *gl = &weights->layer[il]; - const uint64_t gl_kv_raw_dim = gl->attn_kv_a_mqa ? gl->attn_kv_a_mqa->dim[1] : 0; - if (!gl->attn_norm || - !gl->attn_kv_a_mqa || - !gl->attn_kv_a_norm || - !gl->attn_v_b || - !gl->attn_output || - !gl->ffn_norm || - gl_kv_raw_dim < DS4_N_KV_LORA || - gl_kv_raw_dim > kv_raw_dim || - gl->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || - gl->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || - gl->attn_v_b->type != DS4_TENSOR_Q8_0 || - gl->attn_v_b->dim[0] != DS4_N_KV_LORA || - gl->attn_v_b->dim[1] != DS4_N_VALUE_MLA || - gl->attn_v_b->dim[2] != DS4_N_HEAD || - gl->attn_output->type != DS4_TENSOR_Q8_0 || - gl->attn_output->dim[0] != heads_dim || - gl->attn_output->dim[1] != DS4_N_EMBD) { - fprintf(stderr, - "ds4: GLM Metal first-token path found unexpected attention layout in layer %u\n", - il); - ok = 0; - break; + *token_out = heap[filtered - 1u].id; + return true; +} + +static int sample_full_vocab( + const float *logits, + uint32_t n_vocab, + float temperature, + float top_p, + float min_p, + uint64_t *rng, + float *prob_scratch) { + float max_logit = DS4_NEG_INF; + int best = 0; + uint32_t finite = 0; + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + if (!isfinite(v)) continue; + finite++; + if (v > max_logit) { + max_logit = v; + best = (int)i; } + } + if (finite == 0) return sample_argmax(logits, n_vocab); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(attn_norm, cur, - model->map, model->size, - gl->attn_norm->abs_offset, - DS4_N_EMBD, DS4_RMS_EPS); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(kv_raw, - model->map, - model->size, - gl->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - gl_kv_raw_dim, - attn_norm, - 1); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(kv_norm, kv_raw, - model->map, model->size, - gl->attn_kv_a_norm->abs_offset, - DS4_N_KV_LORA, DS4_RMS_EPS); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(heads, - model->map, - model->size, - gl->attn_v_b->abs_offset, - DS4_N_KV_LORA, - heads_dim, - kv_norm, - 1); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(attn_out, - model->map, - model->size, - gl->attn_output->abs_offset, - heads_dim, - DS4_N_EMBD, - heads, - 1); - if (ok) ok = ds4_gpu_add_tensor(after_attn, cur, attn_out, DS4_N_EMBD); - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, after_attn, - model->map, model->size, - gl->ffn_norm->abs_offset, - DS4_N_EMBD, DS4_RMS_EPS); - if (il < DS4_N_LEADING_DENSE) { - const uint64_t gl_ffn_hidden = gl->ffn_gate ? gl->ffn_gate->dim[1] : 0; - if (!gl->ffn_gate || - !gl->ffn_up || - !gl->ffn_down || - gl->ffn_gate->type != DS4_TENSOR_Q8_0 || - gl->ffn_up->type != DS4_TENSOR_Q8_0 || - gl->ffn_down->type != DS4_TENSOR_Q8_0 || - gl->ffn_gate->dim[0] != DS4_N_EMBD || - gl->ffn_up->dim[0] != DS4_N_EMBD || - gl->ffn_up->dim[1] != gl_ffn_hidden || - gl->ffn_down->dim[0] != gl_ffn_hidden || - gl->ffn_down->dim[1] != DS4_N_EMBD || - gl_ffn_hidden > dense_hidden_max) { - fprintf(stderr, - "ds4: GLM Metal first-token path found unexpected dense FFN layout in layer %u\n", - il); - ok = 0; - break; - } - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_gate, - model->map, - model->size, - gl->ffn_gate->abs_offset, - DS4_N_EMBD, - gl_ffn_hidden, - ffn_norm, - 1); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_up, - model->map, - model->size, - gl->ffn_up->abs_offset, - DS4_N_EMBD, - gl_ffn_hidden, - ffn_norm, - 1); - if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, ffn_gate, ffn_up, - (uint32_t)gl_ffn_hidden, 0.0f, 1.0f); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_out, - model->map, - model->size, - gl->ffn_down->abs_offset, - gl_ffn_hidden, - DS4_N_EMBD, - ffn_mid, - 1); - if (ok) ok = ds4_gpu_add_tensor(next, after_attn, ffn_out, DS4_N_EMBD); - } else { - const uint32_t gl_gate_type = gl->ffn_gate_exps ? gl->ffn_gate_exps->type : 0; - const uint32_t gl_up_type = gl->ffn_up_exps ? gl->ffn_up_exps->type : 0; - const bool gl_gate_pair_supported = - glm_graph_gate_pair_type_supported(gl_gate_type, gl_up_type); - uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; - uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; - uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + int fast_token = best; + if (top_p < 1.0f && + sample_fast_top_p(logits, + n_vocab, + finite, + max_logit, + best, + temperature, + top_p, + min_p, + rng, + &fast_token)) { + return fast_token; + } - if (!gl->ffn_gate_inp || - !gl->ffn_exp_probs_b || - !gl->ffn_gate_exps || - !gl->ffn_up_exps || - !gl->ffn_down_exps || - !gl->ffn_gate_shexp || - !gl->ffn_up_shexp || - !gl->ffn_down_shexp || - gl->ffn_gate_inp->type != DS4_TENSOR_F32 || - gl->ffn_gate_inp->dim[0] != DS4_N_EMBD || - gl->ffn_gate_inp->dim[1] != DS4_N_EXPERT || - gl->ffn_exp_probs_b->type != DS4_TENSOR_F32 || - gl->ffn_exp_probs_b->dim[0] != DS4_N_EXPERT || - !gl_gate_pair_supported || - !glm_graph_down_type_supported(gl->ffn_down_exps->type) || - gl->ffn_gate_shexp->type != DS4_TENSOR_Q8_0 || - gl->ffn_up_shexp->type != DS4_TENSOR_Q8_0 || - gl->ffn_down_shexp->type != DS4_TENSOR_Q8_0 || - gl->ffn_gate_shexp->dim[0] != DS4_N_EMBD || - gl->ffn_gate_shexp->dim[1] != DS4_N_FF_EXP || - gl->ffn_up_shexp->dim[0] != DS4_N_EMBD || - gl->ffn_up_shexp->dim[1] != DS4_N_FF_EXP || - gl->ffn_down_shexp->dim[0] != DS4_N_FF_EXP || - gl->ffn_down_shexp->dim[1] != DS4_N_EMBD || - sparse_mid_elems > ffn_mid_elems) { - fprintf(stderr, - "ds4: GLM Metal first-token path found unexpected sparse FFN layout in layer %u\n", - il); - ok = 0; - break; - } + if (top_p >= 1.0f) { + float sum = 0.0f; + const float min_rel = min_p > 0.0f ? min_p : 0.0f; + if (min_rel > 1.0f) return best; - (void)tensor_expert_bytes(model, gl->ffn_gate_exps, 0, - &gate_in, &gate_out, &gate_row_bytes); - (void)tensor_expert_bytes(model, gl->ffn_up_exps, 0, - &up_in, &up_out, &up_row_bytes); - (void)tensor_expert_bytes(model, gl->ffn_down_exps, 0, - &down_in, &down_out, &down_row_bytes); - if (gate_in != DS4_N_EMBD || - up_in != DS4_N_EMBD || - down_in != DS4_N_FF_EXP || - gate_out != DS4_N_FF_EXP || - up_out != DS4_N_FF_EXP || - down_out != DS4_N_EMBD) { - fprintf(stderr, - "ds4: GLM Metal first-token path found unexpected expert strides in layer %u\n", - il); - ok = 0; - break; + /* Find a conservative log-space rejection boundary using the same + * expf implementation as the probability path. Values below this + * boundary are guaranteed to fail min-p, avoiding an expf for the + * overwhelming majority of a large vocabulary. Near-boundary values + * still take the ordinary expf comparison. */ + float reject_scaled = DS4_NEG_INF; + bool have_reject_scaled = false; + if (min_rel > 0.0f && isfinite(min_rel)) { + float cutoff = logf(min_rel); + for (int i = 0; i < 8 && isfinite(cutoff); i++) { + cutoff = nextafterf(cutoff, -FLT_MAX); + if (expf(cutoff) < min_rel) { + reject_scaled = cutoff; + have_reject_scaled = true; + break; + } } + } - if (ok) ok = ds4_gpu_matmul_f32_tensor(router_logits, - model->map, - model->size, - gl->ffn_gate_inp->abs_offset, - DS4_N_EMBD, - DS4_N_EXPERT, - ffn_norm, - 1); - if (ok) ok = ds4_gpu_glm_router_select_tensor(router_selected, - router_weights, - router_probs, - model->map, - model->size, - gl->ffn_exp_probs_b->abs_offset, - router_logits, - DS4_N_EXPERT, - DS4_N_EXPERT_USED, - DS4_EXPERT_WEIGHT_SCALE); - if (ok) { - const ds4_gpu_stream_expert_table table = { - .model_map = model->map, - .model_size = model->size, - .layer = il, - .n_total_expert = DS4_N_EXPERT, - .gate_offset = gl->ffn_gate_exps->abs_offset, - .up_offset = gl->ffn_up_exps->abs_offset, - .down_offset = gl->ffn_down_exps->abs_offset, - .gate_expert_bytes = gate_out * gate_row_bytes, - .down_expert_bytes = down_out * down_row_bytes, - }; - ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( - &table, - router_selected, - DS4_N_EXPERT_USED) != 0; - } - ds4_glm_gpu_graph route_g = { - .routed_gate = routed_gate, - .routed_up = routed_up, - .routed_down = routed_down, - .ssd_streaming = false, - }; - if (ok) ok = glm_graph_routed_moe_one_dispatch( - &route_g, - model, - gl, - il, - ffn_out, - ffn_mid, - gate_out * gate_row_bytes, - gate_row_bytes, - up_out * up_row_bytes, - up_row_bytes, - down_out * down_row_bytes, - down_row_bytes, - router_selected, - router_weights, - ffn_norm, - false); - if (ok) ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( - ffn_gate, - ffn_up, - ffn_mid, - model->map, - model->size, - gl->ffn_gate_shexp->abs_offset, - gl->ffn_up_shexp->abs_offset, - DS4_N_EMBD, - DS4_N_FF_EXP, - ffn_norm, - 0.0f); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_sum, - model->map, - model->size, - gl->ffn_down_shexp->abs_offset, - DS4_N_FF_EXP, - DS4_N_EMBD, - ffn_mid, - 1); - if (ok) ok = ds4_gpu_add_tensor(attn_out, ffn_out, ffn_sum, DS4_N_EMBD); - if (ok) ok = ds4_gpu_add_tensor(next, after_attn, attn_out, DS4_N_EMBD); + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + prob_scratch[i] = -1.0f; + if (!isfinite(v)) continue; + const float scaled = (v - max_logit) / temperature; + if (have_reject_scaled && scaled <= reject_scaled) continue; + const float p = expf(scaled); + if (p < min_rel) continue; + prob_scratch[i] = p; + sum += p; + } + if (sum <= 0.0f || !isfinite(sum)) return best; + float r = sample_rng_f32(rng) * sum; + for (uint32_t i = 0; i < n_vocab; i++) { + const float p = prob_scratch[i]; + if (p < 0.0f) continue; + r -= p; + if (r <= 0.0f) return (int)i; } + return best; + } - if (ok) { - ds4_gpu_tensor *tmp = cur; - cur = next; - next = tmp; + uint32_t n = 0; + float sum = 0.0f; + sample_candidate *cand = NULL; + if (min_p > 0.0f && min_p <= 1.0f) { + /* The later min-p comparison is equivalent to + * exp((logit-max)/temperature) >= min_p; its normalization cancels. + * Still compute the full softmax sum in the original order, then sort + * only candidates that can survive. This preserves the nucleus mass + * and RNG semantics while avoiding a full-vocabulary qsort. */ + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + prob_scratch[i] = -1.0f; + if (!isfinite(v)) continue; + const float p = expf((v - max_logit) / temperature); + prob_scratch[i] = p; + sum += p; + } + if (sum <= 0.0f || !isfinite(sum)) return best; + + const float min_prob = (1.0f / sum) * min_p; + for (uint32_t i = 0; i < n_vocab; i++) { + const float p = prob_scratch[i]; + if (p < 0.0f || p / sum < min_prob) continue; + n++; + } + if (n == 0) return best; + cand = xmalloc((size_t)n * sizeof(cand[0])); + uint32_t out = 0; + for (uint32_t i = 0; i < n_vocab; i++) { + const float p = prob_scratch[i]; + if (p < 0.0f || p / sum < min_prob) continue; + cand[out++] = (sample_candidate){ + .id = (int)i, .logit = logits[i], .prob = p + }; + } + } else { + cand = xmalloc((size_t)finite * sizeof(cand[0])); + for (uint32_t i = 0; i < n_vocab; i++) { + const float v = logits[i]; + if (!isfinite(v)) continue; + const float p = expf((v - max_logit) / temperature); + cand[n++] = (sample_candidate){.id = (int)i, .logit = v, .prob = p}; + sum += p; } } - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, cur, - model->map, model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, DS4_RMS_EPS); - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(logits, - model->map, - model->size, - weights->output->abs_offset, - DS4_N_EMBD, - DS4_N_VOCAB, - ffn_norm, - 1); - if (ok) ok = ds4_gpu_tensor_read(logits, 0, logits_out, logits_bytes) != 0; + if (sum <= 0.0f || !isfinite(sum)) { + free(cand); + return best; + } + + qsort(cand, n, sizeof(cand[0]), sample_candidate_cmp_desc); + const float min_prob = (cand[0].prob / sum) * (min_p > 0.0f ? min_p : 0.0f); + float filtered_sum = 0.0f; + uint32_t filtered = 0; + for (uint32_t i = 0; i < n; i++) { + const float p = cand[i].prob / sum; + if (i > 0 && p < min_prob) break; + filtered_sum += cand[i].prob; + filtered++; + if (filtered_sum / sum >= top_p) break; + } + if (filtered == 0) { + free(cand); + return best; + } - ds4_gpu_tensor_free(router_weights); - ds4_gpu_tensor_free(router_selected); - ds4_gpu_tensor_free(router_probs); - ds4_gpu_tensor_free(router_logits); - ds4_gpu_tensor_free(logits); - ds4_gpu_tensor_free(next); - ds4_gpu_tensor_free(ffn_sum); - ds4_gpu_tensor_free(ffn_out); - ds4_gpu_tensor_free(routed_down); - ds4_gpu_tensor_free(routed_up); - ds4_gpu_tensor_free(routed_gate); - ds4_gpu_tensor_free(ffn_mid); - ds4_gpu_tensor_free(ffn_up); - ds4_gpu_tensor_free(ffn_gate); - ds4_gpu_tensor_free(ffn_norm); - ds4_gpu_tensor_free(after_attn); - ds4_gpu_tensor_free(attn_out); - ds4_gpu_tensor_free(heads); - ds4_gpu_tensor_free(kv_norm); - ds4_gpu_tensor_free(kv_raw); - ds4_gpu_tensor_free(attn_norm); - ds4_gpu_tensor_free(cur); - return ok ? 0 : 1; + float r = sample_rng_f32(rng) * filtered_sum; + for (uint32_t i = 0; i < filtered; i++) { + r -= cand[i].prob; + if (r <= 0.0f) { + const int id = cand[i].id; + free(cand); + return id; + } + } + const int id = cand[filtered - 1].id; + free(cand); + return id; } -static DS4_MAYBE_UNUSED int generate_glm_metal_first_token( - const ds4_model * model, - const ds4_vocab * vocab, - const ds4_weights * weights, - const token_vec * prompt, - int n_predict, - int ctx_size, - ds4_token_emit_fn emit, - ds4_generation_done_fn done, - void * emit_ud) { - fprintf(stderr, "ds4: using GLM Metal first-token generation path\n"); - - if (prompt->len != 1 || prompt->len > ctx_size) { - fprintf(stderr, - "ds4: GLM Metal generation currently supports exactly one prompt token; " - "multi-token prefill needs the GLM KV/DSA graph\n"); - return 1; +static int sample_top_p_min_p( + const float *logits, + uint32_t n_vocab, + float temperature, + int top_k, + float top_p, + float min_p, + uint64_t *rng, + float *prob_scratch) { + if (temperature <= 0.0f) return sample_argmax(logits, n_vocab); + if (top_p <= 0.0f || top_p > 1.0f) top_p = 1.0f; + if (min_p < 0.0f) min_p = 0.0f; + if (top_k <= 0) { + const bool owned_scratch = prob_scratch == NULL; + if (owned_scratch) { + prob_scratch = xmalloc((size_t)n_vocab * sizeof(prob_scratch[0])); + } + const int token = sample_full_vocab(logits, n_vocab, temperature, + top_p, min_p, rng, prob_scratch); + if (owned_scratch) free(prob_scratch); + return token; } - if (n_predict <= 0) { - if (done) done(emit_ud); - return 0; + if (top_k > 1024) top_k = 1024; + if ((uint32_t)top_k > n_vocab) top_k = (int)n_vocab; + + int ids[1024]; + float vals[1024]; + int n = 0; + for (uint32_t i = 0; i < n_vocab; i++) { + float v = logits[i]; + if (!isfinite(v)) continue; + if (n == top_k && v <= vals[n - 1]) continue; + int j = n < top_k ? n++ : n - 1; + while (j > 0 && vals[j - 1] < v) { + vals[j] = vals[j - 1]; + ids[j] = ids[j - 1]; + j--; + } + vals[j] = v; + ids[j] = (int)i; } - if (n_predict > 1) { - fprintf(stderr, - "ds4: GLM Metal generation currently emits only the first generated token; " - "stopping after one token\n"); + if (n == 0) return sample_argmax(logits, n_vocab); + + float probs[1024]; + const float max_logit = vals[0]; + float sum = 0.0f; + for (int i = 0; i < n; i++) { + probs[i] = expf((vals[i] - max_logit) / temperature); + sum += probs[i]; } + if (sum <= 0.0f || !isfinite(sum)) return ids[0]; - float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); - const double t0 = now_sec(); - const int rc = glm_metal_first_token_logits(model, weights, prompt->v[0], logits); - const double t1 = now_sec(); - if (rc != 0) { - free(logits); - return 1; + const float min_prob = (probs[0] / sum) * min_p; + float filtered_sum = 0.0f; + int filtered = 0; + for (int i = 0; i < n; i++) { + float p = probs[i] / sum; + if (i > 0 && p < min_prob) break; + filtered_sum += probs[i]; + filtered++; + if (filtered_sum / sum >= top_p) break; } + if (filtered <= 0) return ids[0]; - if (getenv("DS4_TRACE_TOP") != NULL) { - print_top_logits(stderr, "GLM first-token", vocab, logits, DS4_N_VOCAB, 10); + float r = sample_rng_f32(rng) * filtered_sum; + for (int i = 0; i < filtered; i++) { + r -= probs[i]; + if (r <= 0.0f) return ids[i]; } - const int token = sample_argmax(logits, DS4_N_VOCAB); - if (!vocab_token_is_generation_stop(vocab, token) && emit) emit(emit_ud, token); - if (done) done(emit_ud); + return ids[filtered - 1]; +} - const double eval_s = t1 - t0; - ds4_log(stderr, - DS4_LOG_TIMING, - "ds4: GLM first-token eval: %.2f t/s\n", - eval_s > 0.0 ? 1.0 / eval_s : 0.0); +#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) { + if (!logits || !rng || n_vocab == 0) return -1; + return sample_top_p_min_p(logits, n_vocab, temperature, top_k, + top_p, min_p, rng, prob_scratch); +} +#endif - free(logits); - return 0; +static void print_top_logits( + FILE * fp, + const char * label, + const ds4_vocab * vocab, + const float * logits, + uint32_t n_vocab, + int k) { + int best[16]; + if (k > 16) k = 16; + for (int i = 0; i < k; i++) best[i] = -1; + + for (uint32_t i = 0; i < n_vocab; i++) { + for (int j = 0; j < k; j++) { + if (best[j] < 0 || logits[i] > logits[best[j]]) { + for (int l = k - 1; l > j; l--) best[l] = best[l - 1]; + best[j] = (int)i; + break; + } + } + } + + fprintf(fp, "ds4: top logits %s:\n", label); + for (int i = 0; i < k && best[i] >= 0; i++) { + const int id = best[i]; + fprintf(fp, " %2d %7d % .9g ", i, id, logits[id]); + if (id >= 0 && id < vocab->n_vocab) { + fprintf(fp, "%.*s", (int)vocab->token[id].len, vocab->token[id].ptr); + } + fputc('\n', fp); + } } -static int generate_glm_metal_argmax( +/* CPU generation entry point. It runs layer-major prefill once, then decodes + * one token at a time using the persistent KV cache and scratch arena. */ +static int generate_raw_swa_cpu( const ds4_model * model, const ds4_vocab * vocab, const ds4_weights * weights, const token_vec * prompt, int n_predict, int ctx_size, - bool quality, - bool ssd_streaming, - bool ssd_streaming_cold, - uint32_t ssd_streaming_preload_experts, - uint64_t ssd_streaming_cache_bytes, - uint64_t ssd_streaming_prefill_headroom_bytes, + const float * directional_steering_dirs, + float directional_steering_attn, + float directional_steering_ffn, ds4_token_emit_fn emit, ds4_generation_done_fn done, void * emit_ud, ds4_session_progress_fn progress, void * progress_ud) { - fprintf(stderr, "ds4: using GLM full-attention argmax generation path\n"); - - if (!prompt || prompt->len <= 0 || prompt->len > ctx_size) { - fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); - return 1; - } - if (n_predict <= 0) { - if (done) done(emit_ud); - return 0; - } + (void)progress; + (void)progress_ud; + fprintf(stderr, "ds4: using CPU generation with layer-major prefill\n"); - ds4_glm_gpu_graph g = {0}; - if (!glm_graph_alloc(&g, - model, - weights, - ctx_size, - ssd_streaming, - ssd_streaming_cold)) { - fprintf(stderr, "ds4: failed to allocate GLM graph runtime\n"); - return 1; - } - g.quality = quality; - if ((uint32_t)prompt->len >= g.ctx_size) { - fprintf(stderr, - "ds4: prompt length %d leaves no GLM context room (ctx %u)\n", - prompt->len, - g.ctx_size); - glm_graph_free(&g); - return 1; - } - const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; - if (memory_report) ds4_gpu_print_memory_report("after GLM graph alloc"); + ds4_kv_cache cache; + kv_cache_init(&cache, (uint32_t)ctx_size, 0); + ds4_cpu_decode_scratch decode_scratch; + cpu_decode_scratch_init(&decode_scratch, (uint32_t)ctx_size); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); - bool ok = true; - const bool seed_before_prefill = - ssd_streaming && - !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL", - "DS4_METAL_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL"); + int pos = prompt->len; + const bool trace_top = getenv("DS4_TRACE_TOP") != NULL; const double t_prefill0 = now_sec(); - if (seed_before_prefill) { - ds4_gpu_graph seed_graph; - memset(&seed_graph, 0, sizeof(seed_graph)); - seed_graph.quality = quality; - seed_graph.ssd_streaming = ssd_streaming; - seed_graph.ssd_streaming_cold = ssd_streaming_cold; - seed_graph.streaming_preload_experts = ssd_streaming_preload_experts; - ok = metal_graph_seed_streaming_expert_cache_from_hotlist(&seed_graph, - model, - weights); - } - if (ok) { - ok = glm_graph_prefill_range(&g, - model, - weights, - prompt->v, - 0, - (uint32_t)prompt->len, - logits, - progress, - progress_ud, - (uint32_t)prompt->len); - } - const double t_prefill1 = now_sec(); - if (memory_report) ds4_gpu_print_memory_report("after GLM prefill"); - if (!ok) { - fprintf(stderr, "ds4: GLM prefill failed\n"); + + if (prompt->len <= 0 || prompt->len > ctx_size) { + fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); free(logits); - glm_graph_free(&g); + cpu_decode_scratch_free(&decode_scratch); + kv_cache_free(&cache); return 1; } -#ifdef DS4_ROCM_BUILD - /* - * Decode is SSD-read bound, so the prefill expert headroom is worth more - * as extra dynamic cache once prefill is done. Opt out with - * DS4_ROCM_GLM_STREAMING_GROW_CACHE_AFTER_PREFILL=0. - */ - const char *grow_cache_env = - getenv("DS4_ROCM_GLM_STREAMING_GROW_CACHE_AFTER_PREFILL"); - if (ssd_streaming && - ssd_streaming_cache_bytes != 0 && - ssd_streaming_prefill_headroom_bytes != 0 && - (grow_cache_env == NULL || glm_graph_env_truthy(grow_cache_env))) { - uint64_t budget_bytes = 0; - uint64_t per_expert_bytes = 0; - if (ssd_streaming_cache_bytes <= - UINT64_MAX - ssd_streaming_prefill_headroom_bytes) { - budget_bytes = - ssd_streaming_cache_bytes + ssd_streaming_prefill_headroom_bytes; - } - const uint32_t grown_budget = - ds4_streaming_cache_experts_for_byte_budget(weights, - budget_bytes, - &per_expert_bytes); - const uint32_t current_budget = - ds4_gpu_stream_expert_cache_configured_count(); - if (grown_budget > current_budget) { - ds4_gpu_set_streaming_expert_cache_budget(grown_budget); - fprintf(stderr, - "ds4: ROCm GLM streaming expert cache grew after prefill: " - "%u -> %u experts (%.2f GiB)\n", - current_budget, - grown_budget, - (double)((uint64_t)grown_budget * per_expert_bytes) / - 1073741824.0); + + prefill_layer_major_cpu(logits, model, weights, &cache, prompt, + directional_steering_dirs, + directional_steering_attn, + directional_steering_ffn); + + const double t_prefill1 = now_sec(); + fprintf(stderr, "ds4: prefill %d/%d done\n", prompt->len, prompt->len); + const char *dump_prefill_logits = getenv("DS4_CPU_DUMP_PREFILL_LOGITS"); + if (dump_prefill_logits && dump_prefill_logits[0]) { + if (!write_f32_binary_file(dump_prefill_logits, logits, DS4_N_VOCAB)) { + free(logits); + cpu_decode_scratch_free(&decode_scratch); + kv_cache_free(&cache); + return 1; } + fprintf(stderr, "ds4: wrote CPU prefill logits to %s\n", dump_prefill_logits); } -#else - (void)ssd_streaming_cache_bytes; - (void)ssd_streaming_prefill_headroom_bytes; -#endif int n_generated = 0; int n_decode_eval = 0; - uint32_t pos = (uint32_t)prompt->len; const bool token_timing = getenv("DS4_TOKEN_TIMING") != NULL; const double t_decode0 = now_sec(); - for (int i = 0; i < n_predict && pos < g.ctx_size; i++) { - if (getenv("DS4_TRACE_TOP") != NULL) { + for (int i = 0; i < n_predict && pos < ctx_size; i++) { + if (trace_top) { char label[64]; - snprintf(label, sizeof(label), "GLM step %d", i); + snprintf(label, sizeof(label), "step %d", i); print_top_logits(stderr, label, vocab, logits, DS4_N_VOCAB, 10); } - const int token = sample_argmax(logits, DS4_N_VOCAB); + + int token = sample_argmax(logits, DS4_N_VOCAB); if (vocab_token_is_generation_stop(vocab, token)) break; + if (emit) emit(emit_ud, token); n_generated++; - if (i == n_predict - 1 || pos + 1u >= g.ctx_size) { + if (i == n_predict - 1 || pos + 1 >= ctx_size) { pos++; break; } const double t_eval0 = token_timing ? now_sec() : 0.0; - ok = glm_graph_forward_token(&g, model, weights, token, NULL, pos, - NULL, logits, false); - if (!ok) { - fprintf(stderr, "ds4: GLM decode failed at position %u\n", pos); - free(logits); - glm_graph_free(&g); - return 1; - } + /* The CPU decode step is expected to reuse buffers from + * cpu_decode_scratch. Keep the allocation guard tightly scoped to the + * decode math itself; sampling, token emission, tracing, and callbacks + * may allocate small temporary strings without invalidating that + * guarantee. */ + ds4_alloc_guard_begin("CPU token decode"); + forward_token_raw_swa_cpu_decode_scratch(logits, model, weights, &cache, token, (uint32_t)pos, + directional_steering_dirs, + directional_steering_attn, + directional_steering_ffn, + &decode_scratch); + ds4_alloc_guard_end(); if (token_timing) { const double t_eval1 = now_sec(); - fprintf(stderr, - "ds4: GLM decode eval %d took %.3f ms\n", - n_decode_eval + 1, - (t_eval1 - t_eval0) * 1000.0); + fprintf(stderr, "ds4: decode eval %d took %.3f ms\n", n_decode_eval + 1, (t_eval1 - t_eval0) * 1000.0); } n_decode_eval++; pos++; @@ -46782,16 +7697,19 @@ static int generate_glm_metal_argmax( const double decode_s = t_decode1 - t_decode0; ds4_log(stderr, DS4_LOG_TIMING, - "ds4: GLM prefill: %.2f t/s, generation: %.2f t/s\n", + "ds4: prefill: %.2f t/s, generation: %.2f t/s\n", prefill_s > 0.0 ? (double)prompt->len / prefill_s : 0.0, decode_s > 0.0 ? (double)n_generated / decode_s : 0.0); - if (memory_report) ds4_gpu_print_memory_report("before GLM graph free"); free(logits); - glm_graph_free(&g); + cpu_decode_scratch_free(&decode_scratch); + kv_cache_free(&cache); return 0; } +#ifndef DS4_NO_GPU +#include "models/glm/graph.inc" + /* Metal generation entry point. The model runs as one local whole-graph * pipeline: graph prefill followed by graph decode steps. Streaming PRO may * use decode-style prefill for short prompts. */ @@ -47208,9 +8126,8 @@ static size_t engine_per_layer_kv_bytes_planner(uint32_t il, /* Per-used-tier Class-P graph overhead estimate. Mirrors the * `*_by_tier[t]` allocations in - * metal_graph_alloc_raw_cap (ds4.c:10664-10686 + 10760-10800 + 10806-10816 - * for head extras + 10844 for prefill_tokens + 10852-10892 for batch - * chunked-prefill scratch). + * metal_graph_alloc_raw_cap() in models/deepseek/graph.inc, including decode, + * FFN, head, prompt-token, and chunked-prefill scratch. * * The runtime loop replicates an entire set of Class-P kernel-scratch * buffers on EVERY used tier. The packer's budget math (entry_bytes) only @@ -47222,11 +8139,11 @@ static size_t engine_per_layer_kv_bytes_planner(uint32_t il, * still reserve the overhead, so the packer cannot accept a layout that * would later OOM). * - * Head-tier extras (output_pre/weights/embd/norm/logits at ds4.c:10806-10816) - * and prefill_tokens (ds4.c:10844 — emb_tier only) are charged to ALL tiers - * conservatively: only the head_tier / emb_tier actually pays at runtime, but - * since the pre-subtract is a single scalar applied to every device, charging - * to all is the simplest correct posture. The over-charge is a few MB total. + * Head-tier extras and the embedding-tier prefill_tokens allocation are + * charged to ALL tiers conservatively: only the head_tier / emb_tier actually + * pays at runtime, but since the pre-subtract is a single scalar applied to + * every device, charging to all is the simplest correct posture. The + * over-charge is a few MB total. * * Uses g_ds4_shape compile-time constants (DS4_N_*) and ctx-derived caps * from engine_planner_prefill_cap. The runtime computes some dims from @@ -47317,10 +8234,10 @@ static size_t engine_per_tier_graph_overhead_bytes(const ds4_engine *e) { const uint32_t prefill_cap = engine_planner_prefill_cap(est_ctx, e ? e->prefill_chunk : 0); - /* comp_cap and attn_comp_stage_cap: same formula as runtime line 10597. + /* comp_cap and attn_comp_stage_cap: same formula as the runtime allocation. * If no layer has a compression ratio set (test path, where * g_ds4_compress_ratios is zero-init), min_ratio falls back to ctx, - * matching runtime line 10596. */ + * matching metal_graph_alloc_raw_cap(). */ uint32_t min_ratio = UINT32_MAX; for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); @@ -47340,7 +8257,7 @@ static size_t engine_per_tier_graph_overhead_bytes(const ds4_engine *e) { size_t total = 0; - /* === Class P decode HC scratch (mirrors ds4.c:10664-10686). The + /* === Class P decode HC scratch. The * hc_pre/hc_post/hc_comb buffers are VIEWS of hc_split and are NOT * counted (they would double-count). === */ total += hc_dim * sizeof(float); /* cur_hc_by_tier */ @@ -47356,7 +8273,7 @@ static size_t engine_per_tier_graph_overhead_bytes(const ds4_engine *e) { total += (uint64_t)DS4_N_HEAD_DIM * sizeof(float); /* kv_raw_by_tier */ total += (uint64_t)DS4_N_HEAD_DIM * sizeof(float); /* kv_by_tier */ - /* === Class P FFN / routed-expert state (mirrors ds4.c:10760-10800). === */ + /* === Class P FFN / routed-expert state. === */ total += comp_width_max * sizeof(float); /* comp_kv_cur_by_tier */ total += comp_width_max * sizeof(float); /* comp_sc_cur_by_tier */ if (DS4_PLANNER_ATTN_COMP_CACHE_F16) { @@ -47435,12 +8352,11 @@ static size_t engine_per_tier_graph_overhead_bytes(const ds4_engine *e) { total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_routed_out_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_ffn_out_by_tier */ - /* === Class E embedding-tier prefill_tokens (mirrors ds4.c:10844). + /* === Class E embedding-tier prefill_tokens. * Charged to ALL tiers conservatively. Negligible (pc * int32). === */ total += pc * sizeof(int32_t); /* prefill_tokens_by_tier */ - /* === Head-tier-only extras (mirrors ds4.c:10806-10816). Charged - * conservatively to EVERY tier. === */ + /* === Head-tier-only extras. Charged conservatively to EVERY tier. === */ total += (uint64_t)DS4_N_HC * sizeof(float); /* output_pre_by_tier */ total += (uint64_t)DS4_N_HC * sizeof(float); /* output_weights_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* output_embd_by_tier */ @@ -48471,7 +9387,8 @@ static void ds4_session_dspark_capture_note_checkpoint(ds4_session *s) { } static bool ds4_session_is_glm(const ds4_session *s) { - return s && s->engine && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; + return s && s->engine && + s->engine->provider == ds4_glm_model_provider(); } #ifndef DS4_NO_GPU @@ -48567,11 +9484,26 @@ static bool ds4_layer_payload_range_valid(uint32_t layer_start, uint32_t layer_e uint64_t ds4_session_layer_payload_bytes(ds4_session *s, uint32_t layer_start, uint32_t layer_end) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + return 0; + } + return s->engine->provider->session_layer_payload_bytes( + s, layer_start, layer_end); +} + +static uint64_t ds4_builtin_session_layer_payload_bytes( + ds4_session *s, + uint32_t layer_start, + uint32_t layer_end, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !s->checkpoint_valid || !ds4_layer_payload_range_valid(layer_start, layer_end)) return 0; if (ds4_session_is_cpu(s)) return 0; - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU (void)layer_start; (void)layer_end; @@ -48635,9 +9567,44 @@ uint64_t ds4_session_layer_payload_bytes(ds4_session *s, #endif } +uint64_t ds4_deepseek_session_layer_payload_bytes( + ds4_session *session, + uint32_t layer_start, + uint32_t layer_end) { + return ds4_builtin_session_layer_payload_bytes( + session, layer_start, layer_end, false); +} + +uint64_t ds4_glm_session_layer_payload_bytes( + ds4_session *session, + uint32_t layer_start, + uint32_t layer_end) { + return ds4_builtin_session_layer_payload_bytes( + session, layer_start, layer_end, true); +} + int ds4_session_save_layer_payload(ds4_session *s, FILE *fp, uint32_t layer_start, uint32_t layer_end, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + payload_set_err(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_save_layer_payload( + s, fp, layer_start, layer_end, err, errlen); +} + +static int ds4_builtin_session_save_layer_payload( + ds4_session *s, + FILE *fp, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !fp || !s->checkpoint_valid || !ds4_layer_payload_range_valid(layer_start, layer_end)) { payload_set_err(err, errlen, "invalid session layer payload save"); @@ -48647,7 +9614,7 @@ int ds4_session_save_layer_payload(ds4_session *s, FILE *fp, payload_set_err(err, errlen, "distributed layer payloads require the graph backend"); return 1; } - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; @@ -48870,11 +9837,56 @@ int ds4_session_save_layer_payload(ds4_session *s, FILE *fp, #endif } +int ds4_deepseek_session_save_layer_payload( + ds4_session *session, + FILE *file, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen) { + return ds4_builtin_session_save_layer_payload( + session, file, layer_start, layer_end, err, errlen, false); +} + +int ds4_glm_session_save_layer_payload( + ds4_session *session, + FILE *file, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen) { + return ds4_builtin_session_save_layer_payload( + session, file, layer_start, layer_end, err, errlen, true); +} + int ds4_session_load_layer_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, const int *tokens, uint32_t n_tokens, uint32_t layer_start, uint32_t layer_end, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + payload_set_err(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_load_layer_payload( + s, fp, payload_bytes, tokens, n_tokens, + layer_start, layer_end, err, errlen); +} + +static int ds4_builtin_session_load_layer_payload( + ds4_session *s, + FILE *fp, + uint64_t payload_bytes, + const int *tokens, + uint32_t n_tokens, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !fp || !tokens || !ds4_layer_payload_range_valid(layer_start, layer_end)) { payload_set_err(err, errlen, "invalid session layer payload load"); @@ -48884,7 +9896,7 @@ int ds4_session_load_layer_payload(ds4_session *s, FILE *fp, payload_set_err(err, errlen, "distributed layer payloads require the graph backend"); return 1; } - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU (void)payload_bytes; (void)n_tokens; @@ -49316,6 +10328,36 @@ int ds4_session_load_layer_payload(ds4_session *s, FILE *fp, #endif } +int ds4_deepseek_session_load_layer_payload( + ds4_session *session, + FILE *file, + uint64_t payload_bytes, + const int *tokens, + uint32_t token_count, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen) { + return ds4_builtin_session_load_layer_payload( + session, file, payload_bytes, tokens, token_count, + layer_start, layer_end, err, errlen, false); +} + +int ds4_glm_session_load_layer_payload( + ds4_session *session, + FILE *file, + uint64_t payload_bytes, + const int *tokens, + uint32_t token_count, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen) { + return ds4_builtin_session_load_layer_payload( + session, file, payload_bytes, tokens, token_count, + layer_start, layer_end, err, errlen, true); +} + int ds4_engine_routed_quant_bits(ds4_engine *e) { if (!e) return 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { @@ -49492,6 +10534,16 @@ static void session_greedy_splitkv_reset(ds4_session *s) { #endif uint64_t ds4_session_payload_bytes(ds4_session *s) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + return 0; + } + return s->engine->provider->session_payload_bytes(s); +} + +static uint64_t ds4_builtin_session_payload_bytes(ds4_session *s, bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !s->checkpoint_valid) return 0; if (s->distributed) return 0; if (ds4_session_is_cpu(s)) { @@ -49503,7 +10555,7 @@ uint64_t ds4_session_payload_bytes(ds4_session *s) { bytes += session_cpu_payload_live_tensor_bytes(s); return bytes; } - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU return 0; #else @@ -49536,6 +10588,14 @@ uint64_t ds4_session_payload_bytes(ds4_session *s) { #endif } +uint64_t ds4_deepseek_session_payload_bytes(ds4_session *session) { + return ds4_builtin_session_payload_bytes(session, false); +} + +uint64_t ds4_glm_session_payload_bytes(ds4_session *session) { + return ds4_builtin_session_payload_bytes(session, true); +} + int ds4_session_write_staged_payload(const ds4_session_payload_file *payload, FILE *fp, char *err, size_t errlen) { if (!payload || !payload->path || !fp) { @@ -49619,6 +10679,21 @@ int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out, } int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + payload_set_err(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_save_payload(s, fp, err, errlen); +} + +static int ds4_builtin_session_save_payload(ds4_session *s, + FILE *fp, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !fp || !s->checkpoint_valid) { payload_set_err(err, errlen, "session has no valid checkpoint to save"); return 1; @@ -49626,7 +10701,7 @@ int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) if (s->distributed) { return ds4_dist_session_save_payload(s->distributed, s, fp, err, errlen); } - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; @@ -49927,7 +11002,40 @@ int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) #endif } +int ds4_deepseek_session_save_payload(ds4_session *session, + FILE *file, + char *err, + size_t errlen) { + return ds4_builtin_session_save_payload( + session, file, err, errlen, false); +} + +int ds4_glm_session_save_payload(ds4_session *session, + FILE *file, + char *err, + size_t errlen) { + return ds4_builtin_session_save_payload( + session, file, err, errlen, true); +} + int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + payload_set_err(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_load_payload( + s, fp, payload_bytes, err, errlen); +} + +static int ds4_builtin_session_load_payload(ds4_session *s, + FILE *fp, + uint64_t payload_bytes, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !fp) { payload_set_err(err, errlen, "invalid session payload load"); return 1; @@ -49944,7 +11052,7 @@ int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, c payload_set_err(err, errlen, "unsupported session payload version"); return 1; } - if (ds4_session_is_glm(s)) { + if (glm) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; @@ -50492,6 +11600,24 @@ int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, c #endif } +int ds4_deepseek_session_load_payload(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + char *err, + size_t errlen) { + return ds4_builtin_session_load_payload( + session, file, payload_bytes, err, errlen, false); +} + +int ds4_glm_session_load_payload(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + char *err, + size_t errlen) { + return ds4_builtin_session_load_payload( + session, file, payload_bytes, err, errlen, true); +} + int ds4_session_save_snapshot(ds4_session *s, ds4_session_snapshot *snap, char *err, size_t errlen) { if (!s || !snap) { payload_set_err(err, errlen, "invalid session snapshot save"); @@ -55561,6 +16687,16 @@ static int ds4_engine_open_internal(ds4_engine **out, model_open(&e->model, opt->model_path, graph_backend, !opt->inspect_only); if (opt->warm_weights) model_warm_weights(&e->model); config_validate_model(&e->model); + e->provider = + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA + ? ds4_glm_model_provider() + : ds4_deepseek_model_provider(); + if (!ds4_model_provider_valid(e->provider)) { + fprintf(stderr, "ds4: model provider is missing or incompatible\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } if (load_slice && load_layer_end == UINT32_MAX) { const uint32_t normal_layers = ds4_model_normal_layer_count(); if (normal_layers == 0) { @@ -56771,9 +17907,20 @@ static int ds4_session_tp_register(ds4_session *s) { } int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { + if (!out || !e || ctx_size <= 0 || + !ds4_model_provider_valid(e->provider)) { + return 1; + } + return e->provider->session_create(out, e, ctx_size); +} + +static int ds4_builtin_session_create(ds4_session **out, + ds4_engine *e, + int ctx_size, + bool glm) { if (!out || !e || ctx_size <= 0) return 1; if (e->backend == DS4_BACKEND_CPU) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (glm) { fprintf(stderr, "ds4: GLM sessions currently require a graph backend\n"); return 1; } @@ -56805,7 +17952,7 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { ds4_session *s = xcalloc(1, sizeof(*s)); s->engine = e; s->ctx_size = ctx_size; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (glm) { const uint32_t normal_layers = glm_graph_normal_layer_count(); uint32_t layer_start = 0; uint32_t layer_end = normal_layers ? normal_layers - 1u : 0; @@ -57066,6 +18213,34 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { #endif } +int ds4_deepseek_session_create(ds4_session **out, + ds4_engine *engine, + int context_size) { + return ds4_builtin_session_create(out, engine, context_size, false); +} + +int ds4_glm_session_create(ds4_session **out, + ds4_engine *engine, + int context_size) { + return ds4_builtin_session_create(out, engine, context_size, true); +} + +void ds4_deepseek_session_destroy(ds4_session *session) { +#ifndef DS4_NO_GPU + metal_graph_free(&session->graph); +#else + (void)session; +#endif +} + +void ds4_glm_session_destroy(ds4_session *session) { +#ifndef DS4_NO_GPU + glm_graph_free(&session->glm_graph); +#else + (void)session; +#endif +} + void ds4_session_free(ds4_session *s) { if (!s) return; if (ds4_session_tp_leader(s) && s->tp_session_id != 0 && @@ -57090,11 +18265,7 @@ void ds4_session_free(ds4_session *s) { } #ifndef DS4_NO_GPU else { - if (ds4_session_is_glm(s)) { - glm_graph_free(&s->glm_graph); - } else { - metal_graph_free(&s->graph); - } + s->engine->provider->session_destroy(s); } #endif token_vec_free(&s->checkpoint); @@ -57178,6 +18349,20 @@ void ds4_session_report_progress(ds4_session *s, const char *event, int current, } int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (errlen) snprintf(err, errlen, "missing layer-slice session or model provider"); + return 1; + } + return s->engine->provider->session_layer_slice_reset(s, err, errlen); +} + +static int ds4_builtin_session_layer_slice_reset(ds4_session *s, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s) { if (errlen) snprintf(err, errlen, "missing layer-slice session"); return 1; @@ -57191,7 +18376,7 @@ int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen) { if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else - if (ds4_session_is_glm(s)) { + if (glm) { s->checkpoint.len = 0; s->checkpoint_valid = false; s->mtp_draft_valid = false; @@ -57208,12 +18393,43 @@ int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen) { #endif } +int ds4_deepseek_session_layer_slice_reset(ds4_session *session, + char *err, + size_t errlen) { + return ds4_builtin_session_layer_slice_reset(session, err, errlen, false); +} + +int ds4_glm_session_layer_slice_reset(ds4_session *session, + char *err, + size_t errlen) { + return ds4_builtin_session_layer_slice_reset(session, err, errlen, true); +} + int ds4_session_eval_output_head_from_hc(ds4_session *s, const float *hidden_hc, uint32_t n_tokens, float *logits, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (errlen) snprintf(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_eval_output_head( + s, hidden_hc, n_tokens, logits, err, errlen); +} + +static int ds4_builtin_session_eval_output_head( + ds4_session *s, + const float *hidden_hc, + uint32_t n_tokens, + float *logits, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !s->engine || !hidden_hc || n_tokens == 0 || !logits) { if (errlen) snprintf(err, errlen, "invalid output-head hidden-state input"); return 1; @@ -57236,7 +18452,7 @@ int ds4_session_eval_output_head_from_hc(ds4_session *s, if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else - if (ds4_session_is_glm(s)) { + if (glm) { ds4_glm_gpu_graph *gg = &s->glm_graph; bool ok = ds4_gpu_tensor_write(gg->cur, 0, @@ -57286,6 +18502,28 @@ int ds4_session_eval_output_head_from_hc(ds4_session *s, #endif } +int ds4_deepseek_session_eval_output_head( + ds4_session *session, + const float *hidden_state, + uint32_t token_count, + float *logits, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_output_head( + session, hidden_state, token_count, logits, err, errlen, false); +} + +int ds4_glm_session_eval_output_head( + ds4_session *session, + const float *hidden_state, + uint32_t token_count, + float *logits, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_output_head( + session, hidden_state, token_count, logits, err, errlen, true); +} + static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, char *err, size_t errlen); @@ -57507,6 +18745,32 @@ int ds4_session_eval_layer_slice(ds4_session *s, float *logits, char *err, size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (errlen) snprintf(err, errlen, "missing layer-slice session or model provider"); + return 1; + } + return s->engine->provider->session_eval_layer_slice( + s, tokens, n_tokens, pos0, layer_start, layer_end, + input_hc, output_hc, output_logits, logits, err, errlen); +} + +static int ds4_builtin_session_eval_layer_slice( + ds4_session *s, + const int *tokens, + uint32_t n_tokens, + uint32_t pos0, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hc, + float *output_hc, + bool output_logits, + float *logits, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !s->engine) { if (errlen) snprintf(err, errlen, "missing layer-slice session"); return 1; @@ -57561,7 +18825,7 @@ int ds4_session_eval_layer_slice(ds4_session *s, s->checkpoint_valid = false; return 1; #else - if (ds4_session_is_glm(s)) { + if (glm) { ds4_engine *e = s->engine; ds4_glm_gpu_graph *g = &s->glm_graph; if (!s->glm_graph_ready) { @@ -57988,6 +19252,44 @@ int ds4_session_eval_layer_slice(ds4_session *s, #endif } +int ds4_deepseek_session_eval_layer_slice( + ds4_session *session, + const int *tokens, + uint32_t token_count, + uint32_t position, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hidden_state, + float *output_hidden_state, + bool output_logits, + float *logits, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_layer_slice( + session, tokens, token_count, position, layer_start, layer_end, + input_hidden_state, output_hidden_state, output_logits, logits, + err, errlen, false); +} + +int ds4_glm_session_eval_layer_slice( + ds4_session *session, + const int *tokens, + uint32_t token_count, + uint32_t position, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hidden_state, + float *output_hidden_state, + bool output_logits, + float *logits, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_layer_slice( + session, tokens, token_count, position, layer_start, layer_end, + input_hidden_state, output_hidden_state, output_logits, logits, + err, errlen, true); +} + #ifndef DS4_NO_GPU typedef struct { ds4_session *session; @@ -58090,7 +19392,14 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t return rc; } -static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen) { +static int ds4_builtin_session_sync(ds4_session *s, + const ds4_tokens *prompt, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s || !prompt) { snprintf(err, errlen, "missing session or prompt"); return 1; @@ -58179,7 +19488,7 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c ds4_engine *e = s->engine; const char *backend_name = ds4_backend_name(e->backend); (void)backend_name; (void)e; - if (ds4_session_is_glm(s)) { + if (glm) { /* Debug: truncate the prompt so the dumped prefill logits line up * with the CPU first-token reference (DS4_GLM_LOGIT_DUMP). */ ds4_tokens glm_trunc_prompt; @@ -58903,6 +20212,31 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c #endif } +int ds4_deepseek_session_sync(ds4_session *session, + const ds4_tokens *prompt, + char *err, + size_t errlen) { + return ds4_builtin_session_sync(session, prompt, err, errlen, false); +} + +int ds4_glm_session_sync(ds4_session *session, + const ds4_tokens *prompt, + char *err, + size_t errlen) { + return ds4_builtin_session_sync(session, prompt, err, errlen, true); +} + +static int ds4_session_sync_internal(ds4_session *s, + const ds4_tokens *prompt, + char *err, + size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (errlen) snprintf(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_sync(s, prompt, err, errlen); +} + /* Return true when canonicalization would replace already-sampled tokens. * * A DS4 session checkpoint is more than a token vector: the backend state also @@ -59718,8 +21052,15 @@ static void ds4_session_prepare_support_draft(ds4_session *s, } #endif -static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, - char *err, size_t errlen) { +static int ds4_builtin_session_eval(ds4_session *s, + int token, + bool probe_mtp, + char *err, + size_t errlen, + bool glm) { +#ifdef DS4_NO_GPU + (void)glm; +#endif if (!s) return 1; if (s->distributed) { if (!s->checkpoint_valid) { @@ -59762,7 +21103,7 @@ static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, return 1; #else ds4_engine *e = s->engine; - if (ds4_session_is_glm(s)) { + if (glm) { /* TP worker under GLM MTP: run the full speculative cycle off the * mirrored EVAL frame so drafts, verify batches, and gate traffic * stay in lockstep with the leader's cycle. */ @@ -59880,6 +21221,48 @@ static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, #endif } +int ds4_deepseek_session_eval(ds4_session *session, + int token, + bool probe_support_model, + char *err, + size_t errlen) { + return ds4_builtin_session_eval(session, + token, + probe_support_model, + err, + errlen, + false); +} + +int ds4_glm_session_eval(ds4_session *session, + int token, + bool probe_support_model, + char *err, + size_t errlen) { + return ds4_builtin_session_eval(session, + token, + probe_support_model, + err, + errlen, + true); +} + +static int ds4_session_eval_internal(ds4_session *s, + int token, + bool probe_mtp, + char *err, + size_t errlen) { + if (!s || !s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (errlen) snprintf(err, errlen, "missing session or model provider"); + return 1; + } + return s->engine->provider->session_eval(s, + token, + probe_mtp, + err, + errlen); +} + /* TP-aware eval: mirrors the eval to the worker, runs it locally with the * requested draft-probe flag, then merges the vocab-split logits halves. * Both ds4_session_eval and the speculative driver funnel through here so @@ -60747,6 +22130,21 @@ int ds4_sessions_eval_batch(ds4_decode_item *items, int count, } } + if (!ds4_model_provider_valid(e->provider)) { + if (err && errlen) snprintf(err, errlen, "missing model provider"); + return 1; + } + return e->provider->sessions_eval_batch(items, count, err, errlen); +} + +int ds4_builtin_sessions_eval_batch(ds4_decode_item *items, + int count, + char *err, + size_t errlen) { + ds4_engine *e = items[0].session->engine; +#ifdef DS4_NO_GPU + (void)e; +#endif #ifndef DS4_NO_GPU if (e->backend == DS4_BACKEND_CUDA) { return ds4_sessions_eval_batch_cuda(items, count, err, errlen); @@ -60816,6 +22214,23 @@ int ds4_sessions_eval_batch_with_prefill( } } + const ds4_model_provider_v1 *provider = + prefill_session->engine->provider; + if (!ds4_model_provider_valid(provider)) { + if (err && errlen) snprintf(err, errlen, "missing model provider"); + return 1; + } + return provider->sessions_eval_batch_with_prefill( + items, count, prefill_session, prefill_prompt, err, errlen); +} + +int ds4_builtin_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) { #ifndef DS4_NO_GPU if (prefill_session->engine->backend == DS4_BACKEND_CUDA) { return ds4_sessions_eval_batch_with_prefill_cuda( @@ -64131,6 +25546,26 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, int *accepted, int accepted_cap, char *err, size_t errlen) { if (!s || max_tokens <= 0 || accepted_cap <= 0) return 0; + if (!s->engine || !ds4_model_provider_valid(s->engine->provider)) { + if (err && errlen) snprintf(err, errlen, "missing model provider"); + return -1; + } + return s->engine->provider->session_eval_speculative( + s, first_token, max_tokens, eos_token, + accepted, accepted_cap, err, errlen); +} + +static int ds4_builtin_session_eval_speculative( + ds4_session *s, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen, + bool glm) { + if (!s || max_tokens <= 0 || accepted_cap <= 0) return 0; if (s->distributed) { if (!accepted) return 0; if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; @@ -64145,7 +25580,7 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, accepted[0] = first_token; return 1; } - if (ds4_session_is_glm(s)) { + if (glm) { (void)max_tokens; (void)eos_token; if (!accepted || accepted_cap <= 0) return 0; @@ -64182,7 +25617,7 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, return -1; #else ds4_engine *e = s->engine; - if (ds4_session_is_glm(s) && ds4_engine_glm_mtp_spec_enabled(e)) { + if (glm && ds4_engine_glm_mtp_spec_enabled(e)) { int cycle_cap = accepted_cap; if (cycle_cap > max_tokens) cycle_cap = max_tokens; return ds4_session_glm_spec_cycle(s, first_token, @@ -64848,6 +26283,58 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, #endif } +int ds4_deepseek_session_eval_speculative( + ds4_session *session, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_speculative( + session, first_token, max_tokens, eos_token, + accepted, accepted_cap, err, errlen, false); +} + +int ds4_glm_session_eval_speculative( + ds4_session *session, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen) { + return ds4_builtin_session_eval_speculative( + session, first_token, max_tokens, eos_token, + accepted, accepted_cap, err, errlen, true); +} + +void ds4_deepseek_session_invalidate(ds4_session *session) { + ds4_session_dspark_capture_invalidate(session); +} + +void ds4_glm_session_invalidate(ds4_session *session) { + ds4_session_dspark_capture_invalidate(session); +#ifndef DS4_NO_GPU + ds4_session_glm_reset_dense_cache(session); +#endif +} + +void ds4_deepseek_session_rewind(ds4_session *session, int position) { + (void)position; + ds4_session_dspark_capture_invalidate(session); +} + +void ds4_glm_session_rewind(ds4_session *session, int position) { + (void)position; + ds4_session_dspark_capture_invalidate(session); +#ifndef DS4_NO_GPU + ds4_session_glm_cap_dense_cache(session); +#endif +} + void ds4_session_invalidate(ds4_session *s) { if (!s) return; if (ds4_session_tp_leader(s) && @@ -64857,13 +26344,13 @@ void ds4_session_invalidate(ds4_session *s) { s->checkpoint_valid = false; s->checkpoint.len = 0; s->mtp_draft_valid = false; - ds4_session_dspark_capture_invalidate(s); -#ifndef DS4_NO_GPU - ds4_session_glm_reset_dense_cache(s); -#endif + if (s->engine && ds4_model_provider_valid(s->engine->provider)) { + s->engine->provider->session_invalidate(s); + } } void ds4_session_rewind(ds4_session *s, int pos) { + if (!s) return; if (ds4_session_tp_leader(s) && !ds4_tp_failed(s->engine->tp.ctx)) { (void)ds4_tp_send_rewind(s->engine->tp.ctx, s->tp_session_id, pos); @@ -64872,10 +26359,9 @@ void ds4_session_rewind(ds4_session *s, int pos) { if (pos > s->checkpoint.len) pos = s->checkpoint.len; s->checkpoint.len = pos; s->mtp_draft_valid = false; - ds4_session_dspark_capture_invalidate(s); -#ifndef DS4_NO_GPU - ds4_session_glm_cap_dense_cache(s); -#endif + if (s->engine && ds4_model_provider_valid(s->engine->provider)) { + s->engine->provider->session_rewind(s, pos); + } } int ds4_session_pos(ds4_session *s) { diff --git a/ds4_cuda.cu b/ds4_cuda.cu index aaa4df1134..956bb44d06 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -234,27433 +234,11 @@ static int cuda_q4_mma_tile16_shmem_ok(int which_down); static void routed_moe_decode_graph_destroy_one(int logical_tier); -/* ========================================================================= - * Multi-GPU plumbing (device-aware CUDA). - * ========================================================================= */ - -static_assert(DS4_MAX_GPUS == 16, "DS4_MAX_GPUS stack tables sized for 16"); - -ds4_gpu_ctx g_gpu[DS4_MAX_GPUS]; -int g_n_gpus = 0; -int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS]; - -/* Per-pair pinned-host bounce buffers, indexed [src][dst]. Lazily grown - * to the largest copy seen for that pair. Each pair is its own allocation - * so concurrent fan-out copies from a single source GPU to multiple - * destinations cannot race for staging memory. */ -static void *g_xdev_bounce[DS4_MAX_GPUS][DS4_MAX_GPUS]; -static size_t g_xdev_bounce_bytes[DS4_MAX_GPUS][DS4_MAX_GPUS]; - -/* Internal helper: resolve a tensor's device index. -1 (untagged) is - * treated as device 0 for legacy callers. */ -static inline int ds4_tensor_device_idx(const ds4_gpu_tensor *t) { - if (!t) return 0; - int d = t->device_id; - if (d < 0) return 0; - return d; -} - -/* Debug/override flags are read once per CUDA init. The hot decode path calls - * the xdev helpers many times per token, so they must not re-enter getenv(). */ -static void cuda_xdev_env_refresh(void) { - g_xdev_sync_debug = getenv("DS4_CUDA_SYNC_XDEV") != NULL; - g_xdev_force_cuda_peer = getenv("DS4_FORCE_CUDA_PEER") != NULL; - g_xdev_force_host_bounce = getenv("DS4_FORCE_HOST_BOUNCE") != NULL; -} - -static void cuda_decode_dispatch_env_refresh(void) { - g_cuda_disable_qkv_rms_fused = getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL; - g_cuda_no_window_attention = getenv("DS4_CUDA_NO_WINDOW_ATTENTION") != NULL; - g_cuda_decode_heads8_online = getenv("DS4_CUDA_DECODE_HEADS8_ONLINE") != NULL; - g_cuda_decode_score4 = getenv("DS4_CUDA_DECODE_SCORE4") != NULL; - g_cuda_decode_score8 = getenv("DS4_CUDA_DECODE_SCORE8") != NULL; - g_cuda_no_decode_value512 = getenv("DS4_CUDA_NO_DECODE_VALUE512") != NULL; - g_cuda_no_top1 = getenv("DS4_CUDA_NO_TOP1") != NULL; - g_cuda_end_stream_sync = getenv("DS4_CUDA_END_STREAM_SYNC") != NULL; - g_cuda_no_setdevice_cache = getenv("DS4_CUDA_NO_SETDEVICE_CACHE") != NULL; - g_cuda_exact_score_split_graph = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_GRAPH") != NULL; - g_cuda_exact_score_split_ldg = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_LDG") != NULL; - g_cuda_exact_score_split_vec4 = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4") != NULL; - g_cuda_exact_score_split_vec4_plain = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4_PLAIN") != NULL; - g_cuda_exact_score_split_dim2 = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_DIM2") != NULL && - getenv("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DIM2") == NULL; - g_cuda_exact_score_split_fuse_inv_rope = - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_FUSE_INV_ROPE") != NULL; - g_cuda_moe_decode_graph = getenv("DS4_CUDA_MOE_DECODE_GRAPH") != NULL; -} - -/* WITH_DEVICE(d) { ... } scope macro. - * - * Save the calling thread's current CUDA device, switch to device `d`, - * run the body exactly once, then restore the previous device. If the - * required CUDA calls fail, the body still runs (we don't have a clean - * way to early-exit a containing function from a macro), but the next - * CUDA call inside the body will surface the error naturally. - * - * Implementation: a for-loop with two synthetic variables. Iter 0 runs - * the body; on iter 1, the iteration step restores the previous device - * via cudaSetDevice and sets _wd_first = 0 so the loop exits. The - * single-statement-body restriction of for-loops is removed by the - * required `{ ... }` block in the call site. - */ -#define WITH_DEVICE(d) \ - for (int _wd_prev = -1, _wd_first = 1; \ - _wd_first; \ - _wd_first = 0, \ - (_wd_prev >= 0 ? (void)cudaSetDevice(_wd_prev) : (void)0)) \ - if (cudaGetDevice(&_wd_prev) != cudaSuccess) { /* leave */ } else \ - if (cudaSetDevice(d) != cudaSuccess) { /* leave */ } else - -/* ========================================================================= - * Per-device selective model cache (selective model cache). - * - * The public API in ds4_gpu.h declares ds4_tensor_range and the - * device_cache_tensors / lookup_cache entry points. ds4_cuda.cu does NOT - * include ds4_gpu.h historically (a pre-existing project convention), so - * we redeclare the struct here with the same layout the header uses. - * The implementation links by C linkage; struct compatibility is by - * field layout. */ -typedef struct { - uint64_t source_offset; - uint64_t bytes; - int target_device; -} ds4_tensor_range; - -struct cuda_device_cache { - void *base; /* device-side slab base */ - size_t bytes; - int present; -}; -static cuda_device_cache g_dev_cache[DS4_MAX_GPUS]; - -struct cache_range_entry { - uint64_t source_offset; - uint64_t bytes; - int device_id; - void *device_ptr; -}; -static std::vector g_cache_ranges; - -struct cuda_model_range { - const void *host_base; - uint64_t offset; - uint64_t bytes; - char *device_ptr; - void *registered_base; - char *registered_device_base; - uint64_t registered_bytes; - int host_registered; - int arena_allocated; -}; - -struct cuda_model_arena { - char *device_ptr; - uint64_t bytes; - uint64_t used; -}; - -struct cuda_q8_f16_range { - const void *host_base; - uint64_t offset; - uint64_t weight_bytes; - uint64_t in_dim; - uint64_t out_dim; - __half *device_ptr; - int device_id; /* physical CUDA device id; 0 in single-tier */ -}; - -struct cuda_q8_f32_range { - const void *host_base; - uint64_t offset; - uint64_t weight_bytes; - uint64_t in_dim; - uint64_t out_dim; - float *device_ptr; - int device_id; /* physical CUDA device id; 0 in single-tier */ -}; - -static std::vector g_model_ranges; -static std::vector g_model_arenas; -static std::unordered_map g_model_range_by_offset; -static std::vector g_q8_f16_ranges; -static std::unordered_map g_q8_f16_by_offset; -static std::vector g_q8_f32_ranges; -static std::unordered_map g_q8_f32_by_offset; -static uint64_t g_model_range_bytes; -static uint64_t g_q8_f16_bytes; -static uint64_t g_q8_f32_bytes; -static int g_q8_cache_suppressed; -static int g_q8_f16_disabled_after_oom; -static int g_q8_f16_budget_notice_printed; -static uint64_t g_model_load_progress_next; -static double g_model_load_progress_last; -static int g_model_load_progress_started; -static int g_model_load_progress_tty; -static void *g_cuda_tmp; -static uint64_t g_cuda_tmp_bytes; -static void *g_model_stage_raw[4]; -static void *g_model_stage[4]; -static cudaEvent_t g_model_stage_event[4]; -static uint64_t g_model_stage_bytes; -static void *g_stream_selected_stage_raw[4]; -static void *g_stream_selected_stage[4]; -static cudaEvent_t g_stream_selected_stage_event[4]; -static uint64_t g_stream_selected_stage_bytes; -static cudaStream_t g_stream_selected_upload_stream; - -static int cuda_ok(cudaError_t err, const char *what); -static const char *cuda_model_range_ptr_from_fd( - const void *model_map, - uint64_t offset, - uint64_t bytes, - const char *what); - -/* Forward declaration: defined later in this file. The resolver wrapper - * below uses it for multi-tier dispatch. */ -extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset, - uint64_t bytes, - int expected_device, - void **out_device_ptr); -__global__ static void dequant_q8_0_to_f16_kernel( - __half *out, - const unsigned char *w, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks); -__global__ static void dequant_q8_0_to_f32_kernel( - float *out, - const unsigned char *w, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks); - -static void *cuda_tmp_alloc(uint64_t bytes, const char *what) { - if (bytes == 0) return NULL; - if (g_cuda_tmp_bytes >= bytes) return g_cuda_tmp; - if (g_cuda_tmp) { - (void)cudaFree(g_cuda_tmp); - g_cuda_tmp = NULL; - g_cuda_tmp_bytes = 0; - } - void *ptr = NULL; - cudaError_t err = cudaMalloc(&ptr, (size_t)bytes); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA temp alloc failed for %s (%.2f MiB): %s\n", - what ? what : "scratch", (double)bytes / 1048576.0, cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - g_cuda_tmp = ptr; - g_cuda_tmp_bytes = bytes; - return g_cuda_tmp; -} - -/* Per-tier scratch accessor. - * - * Behavior: - * - At g_n_gpus <= 1 (single-tier), delegates to cuda_tmp_alloc which - * manages the legacy g_cuda_tmp slab. This guarantees byte-identical - * behavior to pre-task code for the gpu_cfg == NULL case. - * - For multi-tier, grows the per-device g_gpu[logical_tier].scratch - * slab on the corresponding physical device. Cleanup is already - * handled by ds4_gpu_cleanup (which walks g_gpu[i].scratch). - * - * The legacy g_cuda_tmp slab is untouched: init-time / preload callers - * still use cuda_tmp_alloc directly. No aliasing between g_cuda_tmp - * and g_gpu[0].scratch — they are independently owned and freed. - * - * Added for multi-GPU execution (multi-GPU execution), step A3 of the - * spec (sub-area 2). */ -static void *cuda_tmp_alloc_on(int logical_tier, uint64_t bytes, const char *what) { - if (bytes == 0) return NULL; - if (g_n_gpus <= 1) { - return cuda_tmp_alloc(bytes, what); - } - if (logical_tier < 0 || logical_tier >= g_n_gpus) { - fprintf(stderr, "ds4: cuda_tmp_alloc_on: bad tier %d (n_gpus=%d, what=%s)\n", - logical_tier, g_n_gpus, what ? what : "?"); - return NULL; - } - ds4_gpu_ctx *ctx = &g_gpu[logical_tier]; - if (ctx->scratch_bytes >= bytes) return ctx->scratch; - int prev = -1; - cudaError_t derr = cudaGetDevice(&prev); - if (derr != cudaSuccess) { - fprintf(stderr, - "ds4: cudaGetDevice failed before scratch alloc on tier %d (dev=%d, what=%s): %s\n", - logical_tier, ctx->device_id, what ? what : "scratch", - cudaGetErrorString(derr)); - (void)cudaGetLastError(); - return NULL; - } - derr = cudaSetDevice(ctx->device_id); - if (derr != cudaSuccess) { - fprintf(stderr, - "ds4: cudaSetDevice(%d) failed before scratch alloc on tier %d (what=%s): %s\n", - ctx->device_id, logical_tier, what ? what : "scratch", - cudaGetErrorString(derr)); - (void)cudaGetLastError(); - if (prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - if (ctx->scratch) { - (void)cudaFree(ctx->scratch); - ctx->scratch = NULL; - ctx->scratch_bytes = 0; - } - void *p = NULL; - cudaError_t err = cudaMalloc(&p, (size_t)bytes); - if (prev >= 0) (void)cudaSetDevice(prev); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA scratch alloc on tier %d (dev=%d) failed for %s (%.2f MiB): %s\n", - logical_tier, ctx->device_id, what ? what : "scratch", - (double)bytes / 1048576.0, cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - ctx->scratch = p; - ctx->scratch_bytes = (size_t)bytes; - return p; -} - -static int cuda_attention_score_buffer_fits(uint32_t n_comp) { - return n_comp <= DS4_CUDA_ATTENTION_SCORE_CAP - DS4_CUDA_ATTENTION_RAW_SCORE_CAP; -} - -static const char *cuda_model_ptr(const void *model_map, uint64_t offset) { - if (model_map == g_model_host_base && g_model_device_base) return g_model_device_base + offset; - return (const char *)model_map + offset; -} - -static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset, uint64_t bytes, const char *what) { - if (bytes == 0) return cuda_model_ptr(model_map, offset); - if (g_model_device_owned || g_model_registered) return cuda_model_ptr(model_map, offset); - if (g_model_hmm_direct && - getenv("DS4_CUDA_WEIGHT_CACHE") == NULL && - getenv("DS4_CUDA_WEIGHT_PRELOAD") == NULL) { - return cuda_model_ptr(model_map, offset); - } - const char *direct_env = getenv("DS4_CUDA_DIRECT_MODEL"); - if (direct_env && direct_env[0]) return cuda_model_ptr(model_map, offset); - - const uint64_t end = offset + bytes; - auto exact = g_model_range_by_offset.find(offset); - if (exact != g_model_range_by_offset.end()) { - const cuda_model_range &r = g_model_ranges[exact->second]; - if (r.host_base == model_map && end >= offset && bytes <= r.bytes) return r.device_ptr; - } - for (const cuda_model_range &r : g_model_ranges) { - if (r.host_base == model_map && offset >= r.offset && end >= offset && end <= r.offset + r.bytes) { - return r.device_ptr + (offset - r.offset); - } - if (r.host_base == model_map && r.host_registered && r.registered_base && r.registered_device_base) { - const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); - const uintptr_t h1 = h0 + bytes; - const uintptr_t r0 = (uintptr_t)r.registered_base; - const uintptr_t r1 = r0 + r.registered_bytes; - if (h1 >= h0 && h0 >= r0 && h1 <= r1) return r.registered_device_base + (h0 - r0); - } - } - - if (getenv("DS4_CUDA_NO_FD_CACHE") == NULL) { - const char *fd_ptr = cuda_model_range_ptr_from_fd(model_map, offset, bytes, what); - if (fd_ptr) return fd_ptr; - } - - cudaError_t err = cudaSuccess; - if (g_model_range_mapping_supported) { - const long page_sz_l = sysconf(_SC_PAGESIZE); - const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; - const uintptr_t host_addr = (uintptr_t)((const char *)model_map + offset); - const uintptr_t reg_addr = host_addr & ~(uintptr_t)(page_sz - 1u); - const uint64_t reg_delta = (uint64_t)(host_addr - reg_addr); - const uint64_t reg_bytes = (reg_delta + bytes + page_sz - 1u) & ~(page_sz - 1u); - void *reg_dev = NULL; - err = cudaHostRegister((void *)reg_addr, - (size_t)reg_bytes, - cudaHostRegisterMapped | cudaHostRegisterReadOnly); - if (err == cudaSuccess) { - err = cudaHostGetDevicePointer(®_dev, (void *)reg_addr, 0); - if (err == cudaSuccess && reg_dev) { - char *dev_ptr = (char *)reg_dev + reg_delta; - g_model_ranges.push_back({model_map, offset, bytes, dev_ptr, (void *)reg_addr, (char *)reg_dev, reg_bytes, 1, 0}); - g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA mapped %s %.2f MiB\n", - what ? what : "weights", - (double)bytes / 1048576.0); - } - return dev_ptr; - } - fprintf(stderr, "ds4: CUDA model range map pointer failed for %s: %s\n", - what ? what : "weights", cudaGetErrorString(err)); - (void)cudaHostUnregister((void *)reg_addr); - (void)cudaGetLastError(); - } else { - if (err == cudaErrorNotSupported || err == cudaErrorInvalidValue) g_model_range_mapping_supported = 0; - (void)cudaGetLastError(); - } - } - - void *dev = NULL; - err = cudaMalloc(&dev, (size_t)bytes); - if (err != cudaSuccess) { - (void)cudaGetLastError(); - fprintf(stderr, "ds4: CUDA model range alloc failed for %s (%.2f MiB): %s\n", - what ? what : "weights", (double)bytes / 1048576.0, cudaGetErrorString(err)); - return NULL; - } - - const char *src = (const char *)model_map + offset; - const uint64_t chunk = 64ull * 1024ull * 1024ull; - for (uint64_t done = 0; done < bytes; done += chunk) { - uint64_t n = bytes - done < chunk ? bytes - done : chunk; - err = cudaMemcpy((char *)dev + done, src + done, (size_t)n, cudaMemcpyHostToDevice); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f/%.2f MiB: %s\n", - what ? what : "weights", - (double)done / 1048576.0, - (double)bytes / 1048576.0, - cudaGetErrorString(err)); - (void)cudaFree(dev); - (void)cudaGetLastError(); - return NULL; - } - } - g_model_ranges.push_back({model_map, offset, bytes, (char *)dev, NULL, NULL, 0, 0, 0}); - g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; - g_model_range_bytes += bytes; - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA cached %s %.2f MiB (total %.2f GiB)\n", - what ? what : "weights", - (double)bytes / 1048576.0, - (double)g_model_range_bytes / 1073741824.0); - } - return (const char *)dev; -} - -/* Per-tier cuBLAS handle. Used by kernel-dispatch wrappers; returns the - * cuBLAS handle for the logical tier. The wrapper is expected to have - * cudaSetDevice'd to that tier's physical device already (kernels and - * cuBLAS calls ride the default stream and are naturally serialized). - * - * Added for multi-GPU execution (multi-GPU execution), sub-area 1. */ -static inline cublasHandle_t cuda_cublas_for_tier(int logical_tier) { - if (g_n_gpus <= 1) { - return (cublasHandle_t)g_gpu[0].cublas; - } - /* The executing device is authoritative: GLM per-layer switching runs - * generic launchers whose out tensors live on device 0 while the layer - * executes elsewhere. On DS4 paths the current device always equals the - * requested tier, so this is behavior-preserving there. */ - int cur_dev = -1; - if (cudaGetDevice(&cur_dev) == cudaSuccess) { - for (int t = 0; t < g_n_gpus; t++) { - if (g_gpu[t].device_id == cur_dev) { - return (cublasHandle_t)g_gpu[t].cublas; - } - } - } - if (logical_tier < 0 || logical_tier >= g_n_gpus) { - return (cublasHandle_t)g_gpu[0].cublas; - } - return (cublasHandle_t)g_gpu[logical_tier].cublas; -} - -/* Multi-tier-aware weight pointer resolver. - * - * Used by kernel-dispatch wrappers in the per-layer execution path to - * obtain a device pointer for a weight slice on the layer's logical - * tier. Behavior: - * - * - When g_n_gpus <= 1 (single-tier engine), delegates to the existing - * cuda_model_range_ptr path. This short-circuit guarantees byte- - * identical behavior to pre-multi-tier code for the gpu_cfg == NULL - * case. - * - * - When g_n_gpus >= 2 (multi-tier engine), translates the logical - * tier index to the corresponding physical CUDA device id via - * g_gpu[logical_tier].device_id and looks up the slice strictly - * in the per-device selective cache via ds4_gpu_lookup_cache_strict. - * On miss, logs a diagnostic and returns NULL (no host-pointer - * fallback — a miss here is a placement/install bug). - * - * The caller is responsible for cudaSetDevice'ing to the right physical - * device before launching the kernel that consumes the returned pointer. - * Wrappers in this file thread `int logical_tier` from the dispatch - * caller; single-tier callers pass 0, which hits the short-circuit. - * - * Added for multi-GPU execution (multi-GPU execution), sub-area 3 of the - * spec. */ -/* Optional second (support) model map for speculative decoding. The strict - * multi-tier cache is keyed by source offset only, so support tensors are - * installed and resolved with a large disjoint offset bias. */ -static const void *g_support_host_base = NULL; -static uint64_t g_support_host_size = 0; -static uint64_t g_support_offset_bias = 0; - -extern "C" uint64_t ds4_gpu_tier_free_vram(int logical_tier) { - if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0; - int prev = -1; - if (cudaGetDevice(&prev) != cudaSuccess) prev = -1; - if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) return 0; - size_t free_b = 0, total_b = 0; - uint64_t out = 0; - if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) out = (uint64_t)free_b; - if (prev >= 0) (void)cudaSetDevice(prev); - return out; -} - -extern "C" int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias) { - if (!map || size == 0 || bias == 0) return 0; - g_support_host_base = map; - g_support_host_size = size; - g_support_offset_bias = bias; - return 1; -} - -static const char *cuda_resolve_weight_ptr(const void *model_map, - uint64_t offset, - uint64_t bytes, - int logical_tier, - const char *label) { - if (g_n_gpus <= 1) { - return cuda_model_range_ptr(model_map, offset, bytes, label); - } - if (g_support_host_base && model_map == g_support_host_base) { - offset += g_support_offset_bias; - } - if (logical_tier < 0 || logical_tier >= g_n_gpus) { - fprintf(stderr, - "ds4: cuda_resolve_weight_ptr: bad tier %d (n_gpus=%d, label=%s)\n", - logical_tier, g_n_gpus, label ? label : "?"); - return NULL; - } - const int physical_device = g_gpu[logical_tier].device_id; - void *dev_ptr = NULL; - if (ds4_gpu_lookup_cache_strict(offset, bytes, physical_device, &dev_ptr) - && dev_ptr) { - return (const char *)dev_ptr; - } - /* GLM multi-tier: generic launchers resolve by the OUT tensor's tier, - * but the executing device (set per layer) is where the weights were - * cached. Retry with the current device before declaring a miss; - * DS4 paths never reach this (out tier == current device). */ - int cur_dev = -1; - if (cudaGetDevice(&cur_dev) == cudaSuccess && - cur_dev != physical_device && - ds4_gpu_lookup_cache_strict(offset, bytes, cur_dev, &dev_ptr) && - dev_ptr) { - return (const char *)dev_ptr; - } - fprintf(stderr, - "ds4: selective-cache miss for offset=%llu bytes=%llu on " - "logical_tier=%d (physical_device=%d, current_device=%d, " - "label=%s); this is a placement/cache-install bug\n", - (unsigned long long)offset, (unsigned long long)bytes, - logical_tier, physical_device, cur_dev, label ? label : "?"); - return NULL; -} - -static int cuda_model_range_is_cached(const void *model_map, uint64_t offset, uint64_t bytes) { - if (bytes == 0) return 1; - if (g_model_device_owned || g_model_registered) return 1; - - const uint64_t end = offset + bytes; - if (end < offset) return 0; - for (const cuda_model_range &r : g_model_ranges) { - if (r.host_base == model_map && - offset >= r.offset && - end <= r.offset + r.bytes) { - return 1; - } - if (r.host_base == model_map && - r.host_registered && - r.registered_base && - r.registered_device_base) { - const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); - const uintptr_t h1 = h0 + bytes; - const uintptr_t r0 = (uintptr_t)r.registered_base; - const uintptr_t r1 = r0 + r.registered_bytes; - if (h1 >= h0 && h0 >= r0 && h1 <= r1) return 1; - } - } - return 0; -} - -static void cuda_q8_f16_cache_release_all(void) { - for (const cuda_q8_f16_range &r : g_q8_f16_ranges) { - (void)cudaFree(r.device_ptr); - } - g_q8_f16_ranges.clear(); - g_q8_f16_by_offset.clear(); - g_q8_f16_bytes = 0; -} - -static uint64_t cuda_parse_mib_env(const char *name, int *present) { - const char *env = getenv(name); - if (present) *present = 0; - if (!env || !env[0]) return 0; - char *end = NULL; - unsigned long long v = strtoull(env, &end, 10); - if (end == env || *end != '\0') return 0; - if (present) *present = 1; - if (v > UINT64_MAX / 1048576ull) return UINT64_MAX; - return (uint64_t)v * 1048576ull; -} - -static uint32_t cuda_parse_u32_env_clamped(const char *name, uint32_t fallback, - uint32_t min_value, uint32_t max_value, - int *present) { - const char *env = getenv(name); - if (present) *present = 0; - if (!env || !env[0]) return fallback; - errno = 0; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (errno != 0 || end == env || *end != '\0') return fallback; - if (present) *present = 1; - if (v < min_value) return min_value; - if (v > max_value) return max_value; - return (uint32_t)v; -} - -static int cuda_env_flag_enabled(const char *name, int fallback) { - const char *env = getenv(name); - if (!env || !env[0]) return fallback; - return strcmp(env, "0") != 0; -} - -extern "C" int ds4_gpu_set_decode_fast_attention(int enabled) { - const int old = g_decode_fast_attention; - g_decode_fast_attention = enabled != 0; - return old; -} - -extern "C" int ds4_gpu_set_decode_score_vec4(int enabled) { - const int old = g_decode_score_vec4; - g_decode_score_vec4 = enabled != 0; - return old; -} - -static bool cuda_splitkv_decode_requested(void) { - if (cuda_env_flag_enabled("DS4_CUDA_NO_SPLITKV_DECODE", 0)) return false; - return g_decode_fast_attention || - cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0); -} - -static uint64_t cuda_q8_f16_cache_limit_bytes(void) { - int present = 0; - const uint64_t limit = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_MB", &present); - return present ? limit : UINT64_MAX; -} - -static uint64_t cuda_q8_f16_cache_reserve_bytes(uint64_t total_bytes) { - int present = 0; - const uint64_t reserve = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", &present); - if (present) return reserve; - - if (total_bytes >= 112ull * 1024ull * 1024ull * 1024ull) { - return 512ull * 1048576ull; - } - - /* High-VRAM cards (>= 40 GiB, e.g. 48 GiB RTX 6000 Ada): use a small - * reserve so the selective Q8->F16 cache can actually engage at tight - * budgets (e.g. --gpu-vram 47,47, where the 81 GB model leaves only ~1.3 - * GiB free and the old 4 GiB floor rejected every cache allocation, - * forcing the scalar DP4A prefill kernel). - * - * NOTE: this 768 MiB value is a *bounded cache-growth guard*, not a hard - * guarantee that live free VRAM stays >= 768 MiB. cuda_q8_f16_cache_has_budget - * only blocks a *cache* allocation when free - request < reserve at that - * moment; allocations made outside cache accounting (cuda_tmp_alloc_on - * activation/prequant buffers, cuBLAS internal workspaces) can still dip - * below it. 768 MiB is chosen to leave headroom above the ~0.5 GiB - * memory-safety floor for those out-of-cache allocations; actual minimum - * free VRAM is verified by measurement, and the disable-after-failure path - * degrades gracefully if cuBLAS/alloc ever fails under pressure. Set - * DS4_CUDA_Q8_F16_CACHE_RESERVE_MB=4096 to restore the prior behavior. */ - if (total_bytes >= 40ull * 1024ull * 1024ull * 1024ull) { - const uint64_t hi_min_reserve = 768ull * 1048576ull; - const uint64_t hi_pct_reserve = total_bytes / 100u; /* 1% */ - return hi_pct_reserve > hi_min_reserve ? hi_pct_reserve : hi_min_reserve; - } - - /* Smaller cards (< 40 GiB): keep the conservative reserve. The expanded - * Q8->F16 cache is only an acceleration path; on a small card a sub-GiB - * reserve would be a large fraction of total VRAM, so keep enough free for - * cuBLAS workspaces, transient graph buffers, and driver bookkeeping. */ - const uint64_t min_reserve = 4096ull * 1048576ull; - const uint64_t pct_reserve = total_bytes / 20u; /* 5% */ - return pct_reserve > min_reserve ? pct_reserve : min_reserve; -} - -static void cuda_q8_f16_cache_budget_notice( - const char *reason, - uint64_t request_bytes, - uint64_t free_bytes, - uint64_t total_bytes, - uint64_t reserve_bytes, - uint64_t limit_bytes) { - if (g_q8_f16_budget_notice_printed && getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") == NULL) return; - g_q8_f16_budget_notice_printed = 1; - if (limit_bytes != UINT64_MAX && free_bytes == 0 && total_bytes == 0 && reserve_bytes == 0) { - fprintf(stderr, - "ds4: CUDA q8 fp16 cache %s; using q8 kernels " - "(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB)\n", - reason, - (double)request_bytes / 1048576.0, - (double)g_q8_f16_bytes / 1073741824.0, - (double)limit_bytes / 1073741824.0); - } else if (limit_bytes == UINT64_MAX) { - fprintf(stderr, - "ds4: CUDA q8 fp16 cache %s; using q8 kernels " - "(request=%.2f MiB cached=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n", - reason, - (double)request_bytes / 1048576.0, - (double)g_q8_f16_bytes / 1073741824.0, - (double)free_bytes / 1073741824.0, - (double)reserve_bytes / 1073741824.0, - (double)total_bytes / 1073741824.0); - } else { - fprintf(stderr, - "ds4: CUDA q8 fp16 cache %s; using q8 kernels " - "(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n", - reason, - (double)request_bytes / 1048576.0, - (double)g_q8_f16_bytes / 1073741824.0, - (double)limit_bytes / 1073741824.0, - (double)free_bytes / 1073741824.0, - (double)reserve_bytes / 1073741824.0, - (double)total_bytes / 1073741824.0); - } -} - -static int cuda_q8_f16_cache_has_budget(uint64_t request_bytes, const char *label) { - (void)label; - const uint64_t limit = cuda_q8_f16_cache_limit_bytes(); - if (limit == 0) return 0; - if (g_q8_f16_bytes > limit || request_bytes > limit - g_q8_f16_bytes) { - cuda_q8_f16_cache_budget_notice("limit reached", request_bytes, 0, 0, 0, limit); - return 0; - } - - size_t free_b = 0; - size_t total_b = 0; - cudaError_t err = cudaMemGetInfo(&free_b, &total_b); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA q8 fp16 cache memory query failed: %s; using q8 kernels\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - - const uint64_t free_bytes = (uint64_t)free_b; - const uint64_t total_bytes = (uint64_t)total_b; - const uint64_t reserve_bytes = cuda_q8_f16_cache_reserve_bytes(total_bytes); - if (request_bytes > free_bytes || - free_bytes - request_bytes < reserve_bytes) { - cuda_q8_f16_cache_budget_notice("budget exhausted", request_bytes, - free_bytes, total_bytes, - reserve_bytes, limit); - return 0; - } - return 1; -} - -static void cuda_q8_f16_cache_disable_after_failure(const char *what, uint64_t request_bytes) { - if (!g_q8_f16_disabled_after_oom) { - fprintf(stderr, - "ds4: CUDA q8 fp16 cache disabled after %s " - "(request=%.2f MiB cached=%.2f GiB); using q8 kernels\n", - what ? what : "allocation failure", - (double)request_bytes / 1048576.0, - (double)g_q8_f16_bytes / 1073741824.0); - } - g_q8_f16_disabled_after_oom = 1; - if (!g_q8_f16_ranges.empty()) { - (void)cudaDeviceSynchronize(); - cuda_q8_f16_cache_release_all(); - } - (void)cudaGetLastError(); -} - -static int cuda_q8_f16_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { - if (g_quality_mode) return 0; - if (g_q8_cache_suppressed) return 0; - if (g_q8_f16_disabled_after_oom) return 0; - if (getenv("DS4_CUDA_NO_Q8_F16_CACHE") != NULL) return 0; - if (cuda_q8_f16_cache_limit_bytes() == 0) return 0; - if (getenv("DS4_CUDA_Q8_F16_ALL") != NULL) return 1; - if (!label) return 0; - if (strstr(label, "attn_output_a") != NULL || - strstr(label, "attn_output_b") != NULL || - strstr(label, "attention_output_a") != NULL || - strstr(label, "attention_output_b") != NULL) { - return getenv("DS4_CUDA_NO_ATTENTION_OUTPUT_F16_CACHE") == NULL; - } - if (strstr(label, "attn_q_b") != NULL) { - return getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL; - } - if (strstr(label, "ffn_gate_shexp") != NULL || - strstr(label, "ffn_up_shexp") != NULL || - strstr(label, "ffn_down_shexp") != NULL) { - return 1; - } - return (in_dim == 4096u && out_dim == 2048u) || - (in_dim == 2048u && out_dim == 4096u) || - (in_dim == 4096u && out_dim == 1024u) || - (in_dim == 4096u && out_dim == 512u) || - (getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL && - in_dim == 1024u && out_dim == 32768u); -} - -static int cuda_q8_label_is_attention_output(const char *label) { - return label && - (strstr(label, "attn_output_a") != NULL || - strstr(label, "attn_output_b") != NULL || - strstr(label, "attention_output_a") != NULL || - strstr(label, "attention_output_b") != NULL); -} - -static int cuda_q8_use_dp4a(void) { - return getenv("DS4_CUDA_NO_Q8_DP4A") == NULL; -} - -static unsigned cuda_q8_exact_threads(uint64_t blocks) { - if (blocks <= 64u) return 64u; - if (blocks <= 128u) return 128u; - return 256u; -} - -static int cuda_q8_f16_preload_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { - if (cuda_q8_label_is_attention_output(label) && - getenv("DS4_CUDA_ATTENTION_OUTPUT_PRELOAD") == NULL && - getenv("DS4_CUDA_Q8_F16_ALL") == NULL) { - return 0; - } - return cuda_q8_f16_cache_allowed(label, in_dim, out_dim); -} - -static int cuda_q8_f32_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) { - if (g_q8_cache_suppressed) return 0; - if (getenv("DS4_CUDA_NO_Q8_F32_CACHE") != NULL) return 0; - if (getenv("DS4_CUDA_Q8_F32_ALL") != NULL) return 1; - if (label && strstr(label, "attn_q_b") != NULL) { - return getenv("DS4_CUDA_ATTN_Q_B_F32_CACHE") != NULL; - } - return getenv("DS4_CUDA_Q8_F32_LARGE") != NULL && - in_dim == 1024u && out_dim == 32768u; -} - -/* Look up a per-device dequantized fp16 slice of the Q8_0 weight at - * (model_map, offset, weight_bytes, in_dim, out_dim). expected_device is a - * PHYSICAL CUDA device id (0 in single-tier; g_gpu[logical_tier].device_id in - * multi-tier). On hit returns the cached pointer for that device. On miss - * cudaSetDevice's to expected_device, allocates + dequants there, stamps the - * new entry with device_id == expected_device, restores the previous device, - * and returns the new pointer. - * - * Single-tier (g_n_gpus <= 1) uses the offset-keyed map for a fast path — - * legacy entries were stamped device_id=0, so the map remains authoritative. - * Multi-tier linear-scans the ranges vector filtering on device_id (the same - * offset may now legitimately map to multiple entries, one per device). - */ -static const __half *cuda_q8_f16_ptr( - const void *model_map, - uint64_t offset, - uint64_t weight_bytes, - uint64_t in_dim, - uint64_t out_dim, - int expected_device, - const char *label) { - if (g_n_gpus <= 1) { - auto exact = g_q8_f16_by_offset.find(offset); - if (exact != g_q8_f16_by_offset.end()) { - const cuda_q8_f16_range &r = g_q8_f16_ranges[exact->second]; - if (r.host_base == model_map && r.weight_bytes == weight_bytes && - r.in_dim == in_dim && r.out_dim == out_dim) { - return r.device_ptr; - } - } - } else { - for (const cuda_q8_f16_range &r : g_q8_f16_ranges) { - if (r.host_base == model_map && - r.offset == offset && - r.weight_bytes == weight_bytes && - r.in_dim == in_dim && - r.out_dim == out_dim && - r.device_id == expected_device) { - return r.device_ptr; - } - } - } - if (!cuda_q8_f16_cache_allowed(label, in_dim, out_dim)) return NULL; - - /* Source Q8 bytes: - * - Single-tier (g_n_gpus <= 1): cuda_model_range_ptr — preserves the - * legacy behavior (FD cache, host-register, or cudaMalloc-and-copy). - * - Multi-tier (g_n_gpus > 1): the per-device selective cache must - * already contain the weight on expected_device. Use the strict - * lookup; on miss this is a placement bug and we hard-fail. - */ - const char *q8; - if (g_n_gpus <= 1) { - q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, "q8_0"); - } else { - void *strict_ptr = NULL; - if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) || - !strict_ptr) { - fprintf(stderr, - "ds4: q8 fp16 cache miss: source bytes not in selective cache for " - "offset=%llu bytes=%llu device=%d (label=%s); placement bug\n", - (unsigned long long)offset, (unsigned long long)weight_bytes, - expected_device, label ? label : "?"); - return NULL; - } - q8 = (const char *)strict_ptr; - } - if (!q8) return NULL; - - if (in_dim != 0 && out_dim > UINT64_MAX / in_dim / sizeof(__half)) return NULL; - const uint64_t out_bytes = in_dim * out_dim * sizeof(__half); - if (!cuda_q8_f16_cache_has_budget(out_bytes, label)) return NULL; - - int prev = -1; - if (g_n_gpus > 1) { - cudaError_t derr = cudaGetDevice(&prev); - if (derr != cudaSuccess) { - fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp16 alloc on device %d: %s\n", - expected_device, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - return NULL; - } - derr = cudaSetDevice(expected_device); - if (derr != cudaSuccess) { - fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp16 alloc: %s\n", - expected_device, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - if (prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - } - __half *dev = NULL; - cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA q8 fp16 cache alloc failed on device %d (%.2f MiB): %s\n", - expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err)); - cuda_q8_f16_cache_disable_after_failure("allocation failure", out_bytes); - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - const uint64_t blocks = (in_dim + 31) / 32; - const uint64_t n = in_dim * out_dim; - dequant_q8_0_to_f16_kernel<<<(n + 255) / 256, 256>>>(dev, - (const unsigned char *)q8, - in_dim, - out_dim, - blocks); - if (!cuda_ok(cudaGetLastError(), "q8 fp16 dequant launch")) { - (void)cudaFree(dev); - cuda_q8_f16_cache_disable_after_failure("dequant launch failure", out_bytes); - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - g_q8_f16_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device}); - if (g_n_gpus <= 1) { - g_q8_f16_by_offset[offset] = g_q8_f16_ranges.size() - 1u; - } - g_q8_f16_bytes += out_bytes; - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA cached q8 fp16 %.2f MiB on device %d (total %.2f GiB)\n", - (double)out_bytes / 1048576.0, expected_device, - (double)g_q8_f16_bytes / 1073741824.0); - } - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return dev; -} - -/* Per-device dequantized fp32 cache. Same conventions as cuda_q8_f16_ptr. */ -static float *cuda_q8_f32_ptr( - const void *model_map, - uint64_t offset, - uint64_t weight_bytes, - uint64_t in_dim, - uint64_t out_dim, - int expected_device, - const char *label) { - if (g_n_gpus <= 1) { - auto exact = g_q8_f32_by_offset.find(offset); - if (exact != g_q8_f32_by_offset.end()) { - const cuda_q8_f32_range &r = g_q8_f32_ranges[exact->second]; - if (r.host_base == model_map && r.weight_bytes == weight_bytes && - r.in_dim == in_dim && r.out_dim == out_dim) { - return r.device_ptr; - } - } - } else { - for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { - if (r.host_base == model_map && - r.offset == offset && - r.weight_bytes == weight_bytes && - r.in_dim == in_dim && - r.out_dim == out_dim && - r.device_id == expected_device) { - return r.device_ptr; - } - } - } - if (!cuda_q8_f32_cache_allowed(label, in_dim, out_dim)) return NULL; - - /* Source Q8 bytes: legacy path in single-tier; strict per-device lookup - * in multi-tier (same rationale as cuda_q8_f16_ptr). */ - const char *q8; - if (g_n_gpus <= 1) { - q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, label ? label : "q8_0"); - } else { - void *strict_ptr = NULL; - if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) || - !strict_ptr) { - fprintf(stderr, - "ds4: q8 fp32 cache miss: source bytes not in selective cache for " - "offset=%llu bytes=%llu device=%d (label=%s); placement bug\n", - (unsigned long long)offset, (unsigned long long)weight_bytes, - expected_device, label ? label : "?"); - return NULL; - } - q8 = (const char *)strict_ptr; - } - if (!q8) return NULL; - - const uint64_t out_bytes = in_dim * out_dim * sizeof(float); - int prev = -1; - if (g_n_gpus > 1) { - cudaError_t derr = cudaGetDevice(&prev); - if (derr != cudaSuccess) { - fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp32 alloc on device %d: %s\n", - expected_device, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - return NULL; - } - derr = cudaSetDevice(expected_device); - if (derr != cudaSuccess) { - fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp32 alloc: %s\n", - expected_device, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - if (prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - } - float *dev = NULL; - cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA q8 fp32 cache alloc failed on device %d (%.2f MiB): %s\n", - expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err)); - (void)cudaGetLastError(); - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - const uint64_t blocks = (in_dim + 31) / 32; - const uint64_t n = in_dim * out_dim; - dequant_q8_0_to_f32_kernel<<<(n + 255) / 256, 256>>>(dev, - (const unsigned char *)q8, - in_dim, - out_dim, - blocks); - if (!cuda_ok(cudaGetLastError(), "q8 fp32 dequant launch")) { - (void)cudaFree(dev); - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return NULL; - } - g_q8_f32_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device}); - if (g_n_gpus <= 1) { - g_q8_f32_by_offset[offset] = g_q8_f32_ranges.size() - 1u; - } - g_q8_f32_bytes += out_bytes; - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA cached q8 fp32 %.2f MiB on device %d (total %.2f GiB)\n", - (double)out_bytes / 1048576.0, expected_device, - (double)g_q8_f32_bytes / 1073741824.0); - } - if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev); - return dev; -} - -static int cuda_ok(cudaError_t err, const char *what) { - if (err == cudaSuccess) return 1; - fprintf(stderr, "ds4: CUDA %s failed: %s\n", what, cudaGetErrorString(err)); - return 0; -} - -static double cuda_wall_sec(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (double)ts.tv_sec + (double)ts.tv_nsec * 1.0e-9; -} - -static int cuda_model_load_progress_enabled(void) { - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") != NULL) return 0; - return 1; -} - -static void cuda_model_load_progress_reset(void) { - g_model_load_progress_next = 0; - g_model_load_progress_last = 0.0; - g_model_load_progress_started = 0; - g_model_load_progress_tty = 0; -} - -static void cuda_model_load_progress_note(uint64_t cached_bytes) { - if (!cuda_model_load_progress_enabled()) return; - - const double now = cuda_wall_sec(); - if (!g_model_load_progress_started) { - g_model_load_progress_started = 1; - g_model_load_progress_tty = isatty(STDERR_FILENO) != 0; - g_model_load_progress_next = (g_model_load_progress_tty ? 2ull : 16ull) * - 1024ull * 1024ull * 1024ull; - g_model_load_progress_last = now; - if (g_model_load_progress_tty) { - fprintf(stderr, "ds4: CUDA loading model tensors into device cache: 0.00 GiB"); - } else { - fprintf(stderr, "ds4: CUDA loading model tensors into device cache\n"); - } - } - - if (cached_bytes < g_model_load_progress_next && - now - g_model_load_progress_last < (g_model_load_progress_tty ? 2.0 : 10.0)) { - return; - } - - if (g_model_load_progress_tty) { - fprintf(stderr, "\rds4: CUDA loading model tensors into device cache: %.2f GiB", - (double)cached_bytes / 1073741824.0); - } else { - fprintf(stderr, "ds4: CUDA loading model tensors %.2f GiB cached\n", - (double)cached_bytes / 1073741824.0); - } - fflush(stderr); - g_model_load_progress_last = now; - const uint64_t step = (g_model_load_progress_tty ? 2ull : 16ull) * - 1024ull * 1024ull * 1024ull; - while (g_model_load_progress_next <= cached_bytes) { - g_model_load_progress_next += step; - } -} - -static int cuda_model_prefetch_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) { - if (!model_map || map_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0; - if (getenv("DS4_CUDA_NO_MODEL_PREFETCH") != NULL || - getenv("DS4_CUDA_COPY_MODEL") != NULL || - getenv("DS4_CUDA_WEIGHT_CACHE") != NULL || - getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) { - return 0; - } - - int device = 0; - if (cudaGetDevice(&device) != cudaSuccess) { - (void)cudaGetLastError(); - return 0; - } - - int pageable = 0; - cudaError_t err = cudaDeviceGetAttribute(&pageable, cudaDevAttrPageableMemoryAccess, device); - if (err != cudaSuccess || !pageable) { - (void)cudaGetLastError(); - return 0; - } -#if CUDART_VERSION >= 13000 - cudaMemLocation loc; - memset(&loc, 0, sizeof(loc)); - loc.type = cudaMemLocationTypeDevice; - loc.id = device; -#else - int loc = device; -#endif - - const long page_sz_l = sysconf(_SC_PAGESIZE); - const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; - const uintptr_t host_addr = (uintptr_t)((const char *)model_map + map_offset); - const uintptr_t pre_addr = host_addr & ~(uintptr_t)(page_sz - 1u); - const uint64_t pre_delta = (uint64_t)(host_addr - pre_addr); - const uint64_t pre_bytes = (pre_delta + map_size + page_sz - 1u) & ~(page_sz - 1u); - void *pre_ptr = (void *)pre_addr; - - const double t0 = cuda_wall_sec(); - err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetReadMostly, loc); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model read-mostly advise skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetPreferredLocation, loc); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model preferred-location advise skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - - if (!g_model_prefetch_stream) { - err = cudaStreamCreateWithFlags(&g_model_prefetch_stream, cudaStreamNonBlocking); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model prefetch stream creation skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - } - -#if CUDART_VERSION >= 13000 - err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, 0, g_model_prefetch_stream); -#else - err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, g_model_prefetch_stream); -#endif - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model prefetch skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - if (getenv("DS4_CUDA_MODEL_PREFETCH_SYNC") != NULL) { - err = cudaStreamSynchronize(g_model_prefetch_stream); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model prefetch sync failed: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - } - const double t1 = cuda_wall_sec(); - fprintf(stderr, - "ds4: CUDA ATS/HMM prefetch queued %.2f GiB of model tensors in %.3fs\n", - (double)map_size / 1073741824.0, - t1 - t0); - g_model_hmm_direct = 1; - return 1; -} - -static uint64_t cuda_model_copy_chunk_bytes(void) { - uint64_t mb = 64; - const char *env = getenv("DS4_CUDA_MODEL_COPY_CHUNK_MB"); - if (env && env[0]) { - char *end = NULL; - unsigned long long v = strtoull(env, &end, 10); - if (end != env && v > 0) mb = (uint64_t)v; - } - if (mb < 16) mb = 16; - if (mb > 4096) mb = 4096; - return mb * 1048576ull; -} - -static void cuda_model_discard_source_pages(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes) { -#if defined(POSIX_MADV_DONTNEED) - if (getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || !model_map || bytes == 0 || offset > model_size) return; - if (bytes > model_size - offset) bytes = model_size - offset; - const long page_sz_l = sysconf(_SC_PAGESIZE); - const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u; - const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset); - const uintptr_t h1 = h0 + bytes; - const uintptr_t p0 = h0 & ~(uintptr_t)(page_sz - 1u); - const uintptr_t p1 = (h1 + page_sz - 1u) & ~(uintptr_t)(page_sz - 1u); - if (p1 > p0) (void)posix_madvise((void *)p0, (size_t)(p1 - p0), POSIX_MADV_DONTNEED); -#else - (void)model_map; - (void)model_size; - (void)offset; - (void)bytes; -#endif -} - -static void cuda_model_drop_file_pages(uint64_t offset, uint64_t bytes) { -#if defined(POSIX_FADV_DONTNEED) - if (g_model_fd < 0 || getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || bytes == 0) return; - (void)posix_fadvise(g_model_fd, (off_t)offset, (off_t)bytes, POSIX_FADV_DONTNEED); -#else - (void)offset; - (void)bytes; -#endif -} - -static uint64_t cuda_round_down(uint64_t v, uint64_t align) { - if (align <= 1) return v; - return (v / align) * align; -} - -static uint64_t cuda_round_up(uint64_t v, uint64_t align) { - if (align <= 1) return v; - const uint64_t rem = v % align; - return rem == 0 ? v : v + (align - rem); -} - -static void *cuda_align_ptr(void *ptr, uint64_t align) { - if (align <= 1) return ptr; - uintptr_t p = (uintptr_t)ptr; - uintptr_t a = (uintptr_t)align; - return (void *)(((p + a - 1u) / a) * a); -} - -static int cuda_model_stage_pool_alloc(uint64_t bytes) { - if (g_model_stage_bytes >= bytes) return 1; - for (size_t i = 0; i < 4; i++) { - if (g_model_stage_event[i]) { - (void)cudaEventDestroy(g_model_stage_event[i]); - g_model_stage_event[i] = NULL; - } - if (g_model_stage_raw[i]) { - (void)cudaFreeHost(g_model_stage_raw[i]); - g_model_stage_raw[i] = NULL; - g_model_stage[i] = NULL; - } - } - g_model_stage_bytes = 0; - if (!g_model_upload_stream) { - cudaError_t err = cudaStreamCreateWithFlags(&g_model_upload_stream, cudaStreamNonBlocking); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model upload stream creation failed: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - } - for (size_t i = 0; i < 4; i++) { - cudaError_t err = cudaMallocHost(&g_model_stage_raw[i], (size_t)bytes); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - g_model_stage[i] = cuda_align_ptr(g_model_stage_raw[i], g_model_direct_align); - err = cudaEventCreateWithFlags(&g_model_stage_event[i], cudaEventDisableTiming); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model staging event creation failed: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - } - g_model_stage_bytes = bytes; - return 1; -} - -static int cuda_pread_full(int fd, void *buf, uint64_t bytes, uint64_t offset) { - uint64_t done = 0; - while (done < bytes) { - const size_t n_req = (bytes - done > (uint64_t)SSIZE_MAX) ? (size_t)SSIZE_MAX : (size_t)(bytes - done); - ssize_t n = pread(fd, (char *)buf + done, n_req, (off_t)(offset + done)); - if (n < 0) { - if (errno == EINTR) continue; - return 0; - } - if (n == 0) return 0; - done += (uint64_t)n; - } - return 1; -} - -static int cuda_model_stage_read(void *stage, uint64_t stage_bytes, - uint64_t offset, uint64_t bytes, - const char **payload) { - *payload = (const char *)stage; -#if defined(__linux__) && defined(O_DIRECT) - if (g_model_direct_fd >= 0 && g_model_direct_align > 1 && g_model_file_size != 0) { - const uint64_t aligned_off = cuda_round_down(offset, g_model_direct_align); - const uint64_t delta = offset - aligned_off; - uint64_t read_size = cuda_round_up(delta + bytes, g_model_direct_align); - if (aligned_off <= g_model_file_size && - read_size <= stage_bytes && - read_size <= g_model_file_size - aligned_off) { - const int saved_errno = errno; - errno = 0; - if (cuda_pread_full(g_model_direct_fd, stage, read_size, aligned_off)) { - *payload = (const char *)stage + delta; - errno = saved_errno; - return 1; - } - const int direct_errno = errno; - if (direct_errno == EINVAL || direct_errno == EFAULT || direct_errno == ENOTSUP || direct_errno == EOPNOTSUPP) { - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA direct model read disabled: %s\n", strerror(direct_errno)); - } - (void)close(g_model_direct_fd); - g_model_direct_fd = -1; - g_model_direct_align = 1; - } - errno = direct_errno; - } - } -#else - (void)stage_bytes; -#endif - return cuda_pread_full(g_model_fd, stage, bytes, offset); -} - -static void cuda_stream_selected_stage_release(void) { - for (size_t i = 0; i < 4; i++) { - if (g_stream_selected_stage_event[i]) { - (void)cudaEventDestroy(g_stream_selected_stage_event[i]); - g_stream_selected_stage_event[i] = NULL; - } - if (g_stream_selected_stage_raw[i]) { - (void)cudaFreeHost(g_stream_selected_stage_raw[i]); - g_stream_selected_stage_raw[i] = NULL; - g_stream_selected_stage[i] = NULL; - } - } - g_stream_selected_stage_bytes = 0; - if (g_stream_selected_upload_stream) { - (void)cudaStreamDestroy(g_stream_selected_upload_stream); - g_stream_selected_upload_stream = NULL; - } -} - -static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { - if (g_stream_selected_stage_bytes >= bytes) return 1; - cuda_stream_selected_stage_release(); - cudaError_t err = cudaStreamCreateWithFlags( - &g_stream_selected_upload_stream, cudaStreamNonBlocking); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected upload stream creation failed: %s\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - for (size_t i = 0; i < 4; i++) { - err = cudaMallocHost(&g_stream_selected_stage_raw[i], (size_t)bytes); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected staging allocation failed: %s\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - cuda_stream_selected_stage_release(); - return 0; - } - g_stream_selected_stage[i] = cuda_align_ptr( - g_stream_selected_stage_raw[i], g_model_direct_align); - err = cudaEventCreateWithFlags(&g_stream_selected_stage_event[i], - cudaEventDisableTiming); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected staging event creation failed: %s\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - cuda_stream_selected_stage_release(); - return 0; - } - } - g_stream_selected_stage_bytes = bytes; - return 1; -} - -static int cuda_model_copy_to_device_streamed( - char *dst, - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t bytes, - const char *what) { - if (!dst || !model_map || offset > model_size || - bytes > model_size - offset) { - return 0; - } - if (bytes == 0) return 1; - if (g_model_fd < 0 || - (g_model_fd_host_base != NULL && model_map != g_model_fd_host_base)) { - return cuda_ok(cudaMemcpy(dst, - (const char *)model_map + offset, - (size_t)bytes, - cudaMemcpyHostToDevice), - what ? what : "stream selected expert copy"); - } - - const uint64_t chunk = cuda_model_copy_chunk_bytes(); - const uint64_t stage_bytes = - chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); - if (!cuda_stream_selected_stage_pool_alloc(stage_bytes)) return 0; - - uint64_t copied = 0; - uint64_t chunk_idx = 0; - while (copied < bytes) { - const uint64_t n = bytes - copied < chunk ? bytes - copied : chunk; - const uint64_t bi = chunk_idx % 4u; - cudaError_t err; - if (chunk_idx >= 4u) { - err = cudaEventSynchronize(g_stream_selected_stage_event[bi]); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected staging wait failed for %s: %s\n", - what ? what : "expert", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - } - const char *payload = NULL; - if (!cuda_model_stage_read(g_stream_selected_stage[bi], - g_stream_selected_stage_bytes, - offset + copied, n, &payload)) { - fprintf(stderr, - "ds4: CUDA streaming selected read failed for %s at %.2f MiB: %s\n", - what ? what : "expert", (double)copied / 1048576.0, - strerror(errno)); - return 0; - } - err = cudaMemcpyAsync(dst + copied, payload, (size_t)n, - cudaMemcpyHostToDevice, - g_stream_selected_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected copy failed for %s at %.2f MiB: %s\n", - what ? what : "expert", (double)copied / 1048576.0, - cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - err = cudaEventRecord(g_stream_selected_stage_event[bi], - g_stream_selected_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected staging record failed for %s: %s\n", - what ? what : "expert", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - cuda_model_drop_file_pages(offset + copied, n); - cuda_model_discard_source_pages(model_map, model_size, - offset + copied, n); - copied += n; - chunk_idx++; - } - - const cudaError_t err = - cudaStreamSynchronize(g_stream_selected_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: CUDA streaming selected upload sync failed for %s: %s\n", - what ? what : "expert", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - return 1; -} - -static uint64_t cuda_model_cache_limit_bytes(void) { - uint64_t gb = 0; - const char *env = getenv("DS4_CUDA_WEIGHT_CACHE_LIMIT_GB"); - if (env && env[0]) { - char *end = NULL; - unsigned long long v = strtoull(env, &end, 10); - if (end != env) gb = (uint64_t)v; - } - if (gb == 0) return UINT64_MAX; - return gb * 1073741824ull; -} - -static uint64_t cuda_model_arena_chunk_bytes(uint64_t need) { - uint64_t mb = 1792; - const char *env = getenv("DS4_CUDA_WEIGHT_ARENA_CHUNK_MB"); - if (env && env[0]) { - char *end = NULL; - unsigned long long v = strtoull(env, &end, 10); - if (end != env && v > 0) mb = (uint64_t)v; - } - if (mb < 256) mb = 256; - if (mb > 8192) mb = 8192; - uint64_t bytes = mb * 1048576ull; - if (bytes < need) { - const uint64_t align = 256ull * 1048576ull; - bytes = (need + align - 1u) & ~(align - 1u); - } - return bytes; -} - -static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) { - if (bytes == 0) return NULL; - if (g_model_cache_full) return NULL; - const uint64_t align = 256u; - const uint64_t aligned = (bytes + align - 1u) & ~(align - 1u); - - for (cuda_model_arena &a : g_model_arenas) { - const uint64_t used = (a.used + align - 1u) & ~(align - 1u); - if (used <= a.bytes && aligned <= a.bytes - used) { - char *ptr = a.device_ptr + used; - a.used = used + aligned; - return ptr; - } - } - - const uint64_t limit = cuda_model_cache_limit_bytes(); - if (g_model_range_bytes > limit || aligned > limit - g_model_range_bytes) return NULL; - - const uint64_t chunk = cuda_model_arena_chunk_bytes(aligned); - void *dev = NULL; - cudaError_t err = cudaMalloc(&dev, (size_t)chunk); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model arena alloc failed for %s (%.2f MiB chunk): %s\n", - what ? what : "weights", - (double)chunk / 1048576.0, - cudaGetErrorString(err)); - (void)cudaGetLastError(); - g_model_cache_full = 1; - return NULL; - } - g_model_arenas.push_back({(char *)dev, chunk, aligned}); - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - uint64_t arena_bytes = 0; - for (const cuda_model_arena &a : g_model_arenas) arena_bytes += a.bytes; - fprintf(stderr, "ds4: CUDA model arena allocated %.2f MiB (arenas %.2f GiB)\n", - (double)chunk / 1048576.0, - (double)arena_bytes / 1073741824.0); - } - return (char *)dev; -} - -static const char *cuda_model_range_ptr_from_fd( - const void *model_map, - uint64_t offset, - uint64_t bytes, - const char *what) { - if (g_model_fd < 0 || bytes == 0) return NULL; - if (g_model_fd_host_base != NULL && model_map != g_model_fd_host_base) return NULL; - const uint64_t limit = cuda_model_cache_limit_bytes(); - if (g_model_range_bytes > limit || bytes > limit - g_model_range_bytes) { - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA direct %s %.2f MiB (cache budget %.2f GiB exhausted)\n", - what ? what : "weights", - (double)bytes / 1048576.0, - (double)limit / 1073741824.0); - } - return cuda_model_ptr(model_map, offset); - } - - char *dev = cuda_model_arena_alloc(bytes, what); - if (!dev) { - if (getenv("DS4_CUDA_STRICT_WEIGHT_CACHE") != NULL) return NULL; - return cuda_model_ptr(model_map, offset); - } - cudaError_t err = cudaSuccess; - - const uint64_t chunk = cuda_model_copy_chunk_bytes(); - const uint64_t stage_bytes = chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); - if (!cuda_model_stage_pool_alloc(stage_bytes)) return NULL; - - uint64_t copied = 0; - uint64_t chunk_idx = 0; - while (copied < bytes) { - const uint64_t n = (bytes - copied < chunk) ? (bytes - copied) : chunk; - const uint64_t bi = chunk_idx % 4u; - if (chunk_idx >= 4u) { - err = cudaEventSynchronize(g_model_stage_event[bi]); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model staging wait failed for %s: %s\n", - what ? what : "weights", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - } - const char *payload = NULL; - if (!cuda_model_stage_read(g_model_stage[bi], g_model_stage_bytes, - offset + copied, n, &payload)) { - fprintf(stderr, "ds4: CUDA model range read failed for %s at %.2f MiB: %s\n", - what ? what : "weights", - (double)copied / 1048576.0, - strerror(errno)); - return NULL; - } - err = cudaMemcpyAsync(dev + copied, payload, (size_t)n, - cudaMemcpyHostToDevice, g_model_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f MiB: %s\n", - what ? what : "weights", - (double)copied / 1048576.0, - cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - err = cudaEventRecord(g_model_stage_event[bi], g_model_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model staging record failed for %s: %s\n", - what ? what : "weights", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - cuda_model_drop_file_pages(offset + copied, n); - cuda_model_discard_source_pages(model_map, g_model_registered_size, offset + copied, n); - copied += n; - cuda_model_load_progress_note(g_model_range_bytes + copied); - chunk_idx++; - } - err = cudaStreamSynchronize(g_model_upload_stream); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model range upload sync failed for %s: %s\n", - what ? what : "weights", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return NULL; - } - - g_model_ranges.push_back({model_map, offset, bytes, dev, NULL, NULL, 0, 0, 1}); - g_model_range_by_offset[offset] = g_model_ranges.size() - 1u; - g_model_range_bytes += bytes; - cuda_model_load_progress_note(g_model_range_bytes); - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA fd-cached %s %.2f MiB (total %.2f GiB)\n", - what ? what : "weights", - (double)bytes / 1048576.0, - (double)g_model_range_bytes / 1073741824.0); - } - return (const char *)dev; -} - -static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) { - if (!model_map || model_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0; - if (getenv("DS4_CUDA_NO_MODEL_COPY") != NULL || - getenv("DS4_CUDA_DIRECT_MODEL") != NULL || - getenv("DS4_CUDA_WEIGHT_CACHE") != NULL || - getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) { - return 0; - } - if (g_model_device_owned || g_model_registered) return 1; - - void *dev = NULL; - const double t0 = cuda_wall_sec(); - cudaError_t err = cudaMalloc(&dev, (size_t)model_size); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - - fprintf(stderr, "ds4: CUDA chunk-copying %.2f GiB model image\n", - (double)model_size / 1073741824.0); - - const uint64_t chunk = cuda_model_copy_chunk_bytes(); - void *stage = NULL; - err = cudaMallocHost(&stage, (size_t)chunk); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err)); - (void)cudaFree(dev); - (void)cudaGetLastError(); - return 0; - } - - if (map_offset > 0) { - uint64_t copied_header = 0; - while (copied_header < map_offset) { - const uint64_t n = (map_offset - copied_header < chunk) ? (map_offset - copied_header) : chunk; - memcpy(stage, (const char *)model_map + copied_header, (size_t)n); - err = cudaMemcpy((char *)dev + copied_header, stage, (size_t)n, cudaMemcpyHostToDevice); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model header copy failed: %s\n", cudaGetErrorString(err)); - (void)cudaFreeHost(stage); - (void)cudaFree(dev); - (void)cudaGetLastError(); - return 0; - } - copied_header += n; - } - } - - uint64_t copied = 0; - double last_report = t0; - while (copied < map_size) { - const uint64_t n = (map_size - copied < chunk) ? (map_size - copied) : chunk; - const uint64_t off = map_offset + copied; - memcpy(stage, (const char *)model_map + off, (size_t)n); - err = cudaMemcpy((char *)dev + off, stage, (size_t)n, cudaMemcpyHostToDevice); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA model chunk copy failed at %.2f GiB: %s\n", - (double)copied / 1073741824.0, cudaGetErrorString(err)); - (void)cudaFreeHost(stage); - (void)cudaFree(dev); - (void)cudaGetLastError(); - return 0; - } - cuda_model_discard_source_pages(model_map, model_size, off, n); - copied += n; - const double now = cuda_wall_sec(); - if (getenv("DS4_CUDA_MODEL_COPY_VERBOSE") != NULL && now - last_report >= 2.0) { - fprintf(stderr, "ds4: CUDA model chunk copy %.2f/%.2f GiB\n", - (double)copied / 1073741824.0, - (double)map_size / 1073741824.0); - last_report = now; - } - } - - (void)cudaFreeHost(stage); - g_model_device_base = (const char *)dev; - g_model_device_owned = 1; - g_model_hmm_direct = 0; - const double t1 = cuda_wall_sec(); - fprintf(stderr, - "ds4: CUDA model chunk copy complete in %.3fs (%.2f GiB tensors)\n", - t1 - t0, - (double)map_size / 1073741824.0); - return 1; -} - -static void cuda_model_range_release_all(void) { - for (const cuda_model_range &r : g_model_ranges) { - if (r.host_registered && r.registered_base) { - (void)cudaHostUnregister(r.registered_base); - } else if (r.device_ptr && !r.arena_allocated) { - (void)cudaFree(r.device_ptr); - } - } - for (const cuda_model_arena &a : g_model_arenas) { - if (a.device_ptr) (void)cudaFree(a.device_ptr); - } - g_model_arenas.clear(); - g_model_ranges.clear(); - g_model_range_by_offset.clear(); - g_model_range_bytes = 0; - cuda_model_load_progress_reset(); -} - -static int cublas_ok(cublasStatus_t st, const char *what) { - if (st == CUBLAS_STATUS_SUCCESS) return 1; - fprintf(stderr, "ds4: cuBLAS %s failed: status %d\n", what, (int)st); - return 0; -} - -extern "C" int ds4_gpu_init_multi(const ds4_gpu_config *cfg) { - if (!cfg || cfg->n_gpus < 1 || cfg->n_gpus > DS4_MAX_GPUS) return 0; - cuda_xdev_env_refresh(); - cuda_decode_dispatch_env_refresh(); - g_current_logical_tier = -1; - - /* g_n_gpus is published incrementally so ds4_gpu_cleanup() can unwind - * partial state on failure. We publish `i + 1` BEFORE allocating any - * resources for context `i`, so even if (e.g.) stream creation - * succeeds but event creation fails, cleanup still walks device `i` - * and destroys the stream. ds4_gpu_cleanup is null-safe per field — - * partial state is OK. */ - for (int i = 0; i < cfg->n_gpus; i++) { - ds4_gpu_ctx *c = &g_gpu[i]; - c->device_id = cfg->device_indices[i]; - if (c->device_id < 0) return 0; - /* Publish the in-progress device id so cleanup can target it on - * any later failure. cudaSetDevice is also required before - * cleanup's cudaEventDestroy / cudaStreamDestroy / cublasDestroy - * calls hit the right context. */ - g_n_gpus = i + 1; - if (!cuda_ok(cudaSetDevice(c->device_id), "init set device")) return 0; - cudaDeviceProp prop; - if (cudaGetDeviceProperties(&prop, c->device_id) == cudaSuccess) { - fprintf(stderr, "ds4: CUDA backend initialized on %s (sm_%d%d) dev=%d\n", - prop.name, prop.major, prop.minor, c->device_id); - } - /* Per-device stream. */ - cudaStream_t s = NULL; - if (!cuda_ok(cudaStreamCreate(&s), "init stream")) return 0; - c->stream = (void *)s; - /* Per-device boundary event (reusable, no timing). */ - cudaEvent_t ev = NULL; - if (!cuda_ok(cudaEventCreateWithFlags(&ev, cudaEventDisableTiming), - "init event")) return 0; - c->boundary_event = (void *)ev; - /* Per-device cuBLAS handle. */ - cublasHandle_t h = NULL; - if (!cublas_ok(cublasCreate(&h), "init cublas")) return 0; - c->cublas = (void *)h; - const cublasMath_t math_mode = - (g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL) - ? CUBLAS_DEFAULT_MATH - : CUBLAS_TF32_TENSOR_OP_MATH; - (void)cublasSetMathMode(h, math_mode); - c->cublas_ready = 1; - c->budget_bytes = cfg->vram_bytes[i]; - c->used_bytes = 0; - c->scratch = NULL; - c->scratch_bytes = 0; - } - - /* NxN peer-access matrix. - * - * Driver semantics: cudaDeviceCanAccessPeer + cudaDeviceEnablePeerAccess - * can both succeed even on hardware/drivers where cudaMemcpyPeerAsync - * silently delivers wrong data (notably RTX 6000 Ada under recent - * NVIDIA drivers, per the v0 design doc). The corruption is - * non-deterministic and can affect either or both directions of a - * pair. To guard against this, we run a multi-size, multi-iteration - * validation at init (see the loop below): write distinct known - * patterns, peer-copy them to the destination, read back, and only - * set peer_ok[i][j] if every probe round-trips byte-perfect. A - * single small probe is not sufficient — it can pass while realistic - * activation-sized copies still corrupt. On any failure the entry - * stays at 0 and cross-device copies fall back to the pinned-host - * bounce path automatically. */ - for (int i = 0; i < g_n_gpus; i++) { - for (int j = 0; j < g_n_gpus; j++) { - if (i == j) { g_gpu_peer_ok[i][j] = 1; continue; } - int can = 0; - (void)cudaDeviceCanAccessPeer(&can, g_gpu[i].device_id, - g_gpu[j].device_id); - if (!can) { g_gpu_peer_ok[i][j] = 0; continue; } - (void)cudaSetDevice(g_gpu[i].device_id); - cudaError_t e = cudaDeviceEnablePeerAccess(g_gpu[j].device_id, 0); - int enabled = (e == cudaSuccess || - e == cudaErrorPeerAccessAlreadyEnabled); - (void)cudaGetLastError(); - if (!enabled) { g_gpu_peer_ok[i][j] = 0; continue; } - - /* Runtime validation: peer copies on RTX 6000 Ada under recent - * NVIDIA drivers silently corrupt at realistic sizes even though - * the API returns success. cudaDeviceCanAccessPeer and - * cudaDeviceEnablePeerAccess can both report success while - * cudaMemcpyPeer delivers wrong data non-deterministically. - * Probe with multiple sizes and iterations; ALL must round-trip - * byte-perfect or we disable peer for this pair and silently fall - * back to the pinned-host bounce path. */ - static const size_t kValidateSizes[] = { - 4u * 1024u, - 256u * 1024u, - 1u * 1024u * 1024u, - 16u * 1024u * 1024u, - }; - const int kValidateIters = 4; - const int kNValidateSizes = (int)(sizeof(kValidateSizes) / - sizeof(kValidateSizes[0])); - const size_t kMaxValidate = kValidateSizes[kNValidateSizes - 1]; - - unsigned char *vh_src = (unsigned char *)malloc(kMaxValidate); - unsigned char *vh_dst = (unsigned char *)malloc(kMaxValidate); - if (!vh_src || !vh_dst) { - free(vh_src); free(vh_dst); - g_gpu_peer_ok[i][j] = 0; continue; - } - - void *src_dev = NULL; void *dst_dev = NULL; - (void)cudaSetDevice(g_gpu[i].device_id); - if (cudaMalloc(&src_dev, kMaxValidate) != cudaSuccess) { - (void)cudaGetLastError(); - free(vh_src); free(vh_dst); - g_gpu_peer_ok[i][j] = 0; continue; - } - (void)cudaSetDevice(g_gpu[j].device_id); - if (cudaMalloc(&dst_dev, kMaxValidate) != cudaSuccess) { - (void)cudaGetLastError(); - (void)cudaSetDevice(g_gpu[i].device_id); - (void)cudaFree(src_dev); - free(vh_src); free(vh_dst); - g_gpu_peer_ok[i][j] = 0; continue; - } - - int peer_validated = 1; - size_t failed_bytes = 0; - int failed_iter = -1; - for (int s_idx = 0; - s_idx < kNValidateSizes && peer_validated; - s_idx++) { - size_t n = kValidateSizes[s_idx]; - for (int it = 0; it < kValidateIters && peer_validated; it++) { - for (size_t k = 0; k < n; k++) { - vh_src[k] = (unsigned char) - ((k * 31u + (size_t)it * 17u + - (size_t)s_idx * 53u + 11u) & 0xffu); - } - (void)cudaSetDevice(g_gpu[i].device_id); - if (cudaMemcpy(src_dev, vh_src, n, - cudaMemcpyHostToDevice) != cudaSuccess) { - peer_validated = 0; failed_bytes = n; failed_iter = it; - break; - } - cudaError_t pc = cudaMemcpyPeer( - dst_dev, g_gpu[j].device_id, - src_dev, g_gpu[i].device_id, n); - if (pc != cudaSuccess) { - peer_validated = 0; failed_bytes = n; failed_iter = it; - break; - } - (void)cudaSetDevice(g_gpu[j].device_id); - if (cudaMemcpy(vh_dst, dst_dev, n, - cudaMemcpyDeviceToHost) != cudaSuccess) { - peer_validated = 0; failed_bytes = n; failed_iter = it; - break; - } - if (memcmp(vh_src, vh_dst, n) != 0) { - peer_validated = 0; failed_bytes = n; failed_iter = it; - break; - } - } - } - - (void)cudaSetDevice(g_gpu[j].device_id); - (void)cudaFree(dst_dev); - (void)cudaSetDevice(g_gpu[i].device_id); - (void)cudaFree(src_dev); - free(vh_src); - free(vh_dst); - - g_gpu_peer_ok[i][j] = peer_validated; - if (peer_validated) { - fprintf(stderr, - "ds4: peer access %d->%d validated across %d sizes x %d" - " iterations (max %zu MiB)\n", - g_gpu[i].device_id, g_gpu[j].device_id, - kNValidateSizes, kValidateIters, - kMaxValidate / (1024u * 1024u)); - } else { - fprintf(stderr, - "ds4: peer access %d->%d FAILED validation at" - " size=%zu iter=%d; falling back to pinned-host bounce\n", - g_gpu[i].device_id, g_gpu[j].device_id, - failed_bytes, failed_iter); - } - } - } - - g_cublas_ready = 1; - return 1; -} - -extern "C" int ds4_gpu_init(void) { - ds4_gpu_config cfg; - memset(&cfg, 0, sizeof(cfg)); - cfg.device_indices[0] = 0; - cfg.n_gpus = 1; - return ds4_gpu_init_multi(&cfg); -} - -extern "C" void ds4_gpu_cleanup(void) { - (void)cudaDeviceSynchronize(); - g_current_logical_tier = -1; - - /* Multi-GPU teardown: events, streams, cublas handles, scratch - * slabs, per-pair bounce buffers. */ - for (int i = 0; i < g_n_gpus; i++) { - ds4_gpu_ctx *c = &g_gpu[i]; - (void)cudaSetDevice(c->device_id); - attention_decode_score_split_graph_destroy_one(i); - routed_moe_decode_graph_destroy_one(i); - if (c->boundary_event) { - (void)cudaEventDestroy((cudaEvent_t)c->boundary_event); - c->boundary_event = NULL; - } - if (c->stream) { - (void)cudaStreamDestroy((cudaStream_t)c->stream); - c->stream = NULL; - } - if (c->cublas) { - (void)cublasDestroy((cublasHandle_t)c->cublas); - c->cublas = NULL; - c->cublas_ready = 0; - } - if (c->scratch) { - (void)cudaFree(c->scratch); - c->scratch = NULL; - c->scratch_bytes = 0; - } - } - for (int i = 0; i < DS4_MAX_GPUS; i++) { - for (int j = 0; j < DS4_MAX_GPUS; j++) { - if (g_xdev_bounce[i][j]) { - (void)cudaFreeHost(g_xdev_bounce[i][j]); - g_xdev_bounce[i][j] = NULL; - g_xdev_bounce_bytes[i][j] = 0; - } - } - } - cuda_stream_selected_cache_release(); - cuda_stream_selected_stage_release(); - g_n_gpus = 0; - g_cublas_ready = 0; - - /* Per-device selective cache teardown (selective model cache). */ - for (int d = 0; d < DS4_MAX_GPUS; d++) { - if (!g_dev_cache[d].present) continue; - int prev = -1; - (void)cudaGetDevice(&prev); - (void)cudaSetDevice(d); - if (g_dev_cache[d].base) (void)cudaFree(g_dev_cache[d].base); - g_dev_cache[d].base = NULL; - g_dev_cache[d].bytes = 0; - g_dev_cache[d].present = 0; - if (prev >= 0) (void)cudaSetDevice(prev); - } - g_cache_ranges.clear(); - - /* Continue with legacy global teardown below. */ - - cuda_model_range_release_all(); - cuda_q8_f16_cache_release_all(); - g_q8_f16_disabled_after_oom = 0; - g_q8_f16_budget_notice_printed = 0; - for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { - (void)cudaFree(r.device_ptr); - } - g_q8_f32_ranges.clear(); - g_q8_f32_by_offset.clear(); - g_q8_f32_bytes = 0; - if (g_cuda_tmp) { - (void)cudaFree(g_cuda_tmp); - g_cuda_tmp = NULL; - g_cuda_tmp_bytes = 0; - } - for (size_t i = 0; i < 4; i++) { - if (g_model_stage_event[i]) { - (void)cudaEventDestroy(g_model_stage_event[i]); - g_model_stage_event[i] = NULL; - } - if (g_model_stage_raw[i]) { - (void)cudaFreeHost(g_model_stage_raw[i]); - g_model_stage_raw[i] = NULL; - g_model_stage[i] = NULL; - } - } - g_model_stage_bytes = 0; - if (g_model_upload_stream) { - (void)cudaStreamDestroy(g_model_upload_stream); - g_model_upload_stream = NULL; - } - if (g_model_device_owned && g_model_device_base) { - (void)cudaFree((void *)g_model_device_base); - } - if (g_model_registered && g_model_host_base) { - (void)cudaHostUnregister((void *)g_model_host_base); - } - g_model_host_base = NULL; - g_model_device_base = NULL; - g_model_registered_size = 0; - g_model_registered = 0; - g_model_device_owned = 0; - g_model_range_mapping_supported = 1; - g_model_hmm_direct = 0; - g_model_fd = -1; - if (g_model_direct_fd >= 0) { - (void)close(g_model_direct_fd); - g_model_direct_fd = -1; - } - g_model_direct_align = 1; - g_model_file_size = 0; - g_model_cache_full = 0; - if (g_model_prefetch_stream) { - (void)cudaStreamDestroy(g_model_prefetch_stream); - g_model_prefetch_stream = NULL; - } -} - -__global__ static void fill_f32_kernel(float *x, uint64_t n, float v); - -extern "C" int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, - uint64_t bytes) { - if (!t) return 1; - if (device_id < 0 || device_id >= g_n_gpus) return 2; - if (bytes == 0) bytes = 1; - int ok = 0; - WITH_DEVICE(g_gpu[device_id].device_id) { - ok = cuda_ok(cudaMalloc(&t->ptr, (size_t)bytes), "tensor alloc"); - } - if (!ok) return 3; - t->bytes = bytes; - t->owner = 1; - t->device_id = device_id; - g_gpu[device_id].used_bytes += bytes; - return 0; -} - -/* Async D2D copy queued on the destination device's default stream — - * ordering against the producer comes from the caller's fence; the CPU - * does not block (unlike ds4_gpu_tensor_copy's sync cudaMemcpy). */ -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 cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes, - cudaMemcpyDeviceToDevice, 0), - "tensor copy async"); -} - -extern "C" void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t) { - if (!t) return; - int d = ds4_tensor_device_idx(t); - if (t->owner && t->ptr) { - WITH_DEVICE(g_gpu[d].device_id) { - (void)cudaFree(t->ptr); - } - } - t->ptr = NULL; - t->bytes = 0; - t->owner = 0; -} - -extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { - ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); - if (!t) return NULL; - if (ds4_gpu_tensor_alloc_on(t, 0, bytes) != 0) { - free(t); - return NULL; - } - return t; -} - -extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { - if (bytes == 0) bytes = 1; - ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); - if (!t) return NULL; - int ok = 0; - /* Managed memory is not device-bound, but we record device 0 so that - * subsequent ds4_gpu_tensor_free pairs with WITH_DEVICE(0) safely. */ - WITH_DEVICE(g_gpu[0].device_id) { - ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes), - "managed tensor alloc"); - } - if (!ok) { free(t); return NULL; } - t->bytes = bytes; - t->owner = 1; - t->device_id = 0; - return t; -} - -/* Heap-allocated tensor on a specific logical tier. - * - * Mirrors the legacy ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *) - * with an explicit tier argument. Internally calls - * ds4_gpu_tensor_alloc_on on a freshly malloc'd struct. - * - * The legacy ds4_gpu_tensor_alloc(bytes) (above) calls - * ds4_gpu_tensor_alloc_on(t, 0, bytes); ds4_gpu_tensor_alloc_ptr_on(0, - * bytes) is byte-equivalent. Single-tier callers MAY remain on the - * legacy 1-arg helper; new multi-tier callers in ds4.c use _ptr_on. */ -extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes) { - if (tier < 0 || tier >= g_n_gpus) { - fprintf(stderr, - "ds4: ds4_gpu_tensor_alloc_ptr_on: bad tier %d (n_gpus=%d)\n", - tier, g_n_gpus); - return NULL; - } - ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); - if (!t) return NULL; - if (ds4_gpu_tensor_alloc_on(t, tier, bytes) != 0) { - free(t); - return NULL; - } - return t; -} - -/* Heap-allocated managed-memory tensor on a specific logical tier. - * Differs from ds4_gpu_tensor_alloc_managed only in stamping - * tier instead of 0. Used by the per-layer KV cache when tier !=0. - * - * Managed-memory paging behavior: cudaMallocManaged pages between - * devices on first-touch. In a single-tier pipeline the page lives on - * tier 0; in a multi-tier pipeline the layer's kernels run on the - * layer's tier so the page lives there after first-touch and stays - * unless another device touches it. Stamping tier matches the home - * device for free-time accounting. */ -extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes) { - if (tier < 0 || tier >= g_n_gpus) { - fprintf(stderr, - "ds4: ds4_gpu_tensor_alloc_managed_on: bad tier %d (n_gpus=%d)\n", - tier, g_n_gpus); - return NULL; - } - if (bytes == 0) bytes = 1; - ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); - if (!t) return NULL; - int ok = 0; - /* Run the cudaMallocManaged call under the home tier's device so the - * first-touch home matches the stamped device_id; the page itself can - * migrate freely under managed-memory semantics. */ - WITH_DEVICE(g_gpu[tier].device_id) { - ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes), - "managed tensor alloc (tier)"); - } - if (!ok) { free(t); return NULL; } - t->bytes = bytes; - t->owner = 1; - t->device_id = tier; - return t; -} - -extern "C" int ds4_gpu_tensor_device(const ds4_gpu_tensor *t) { - return t ? t->device_id : -1; -} - -static uint64_t cuda_managed_kv_reserve_bytes(uint64_t total_bytes) { - const uint64_t min_reserve = 8ull * 1073741824ull; - const uint64_t max_reserve = 40ull * 1073741824ull; - uint64_t reserve = total_bytes / 4u; - if (reserve < min_reserve) reserve = min_reserve; - if (reserve > max_reserve) reserve = max_reserve; - return reserve; -} - -extern "C" int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) { - if (kv_cache_bytes == 0) return 0; - - /* Very large KV caches are where device-only cudaMalloc() can make a - * unified-memory machine unresponsive. Managed memory restores the old - * demand-paged behavior for this one long-lived allocation class only. */ - const uint64_t huge_kv = 8ull * 1073741824ull; - if (kv_cache_bytes >= huge_kv) return 1; - - const uint64_t large_context = 8ull * 1073741824ull; - if (context_bytes < large_context) return 0; - - size_t free_b = 0; - size_t total_b = 0; - cudaError_t err = cudaMemGetInfo(&free_b, &total_b); - if (err != cudaSuccess) { - (void)cudaGetLastError(); - return 0; - } - - const uint64_t free_bytes = (uint64_t)free_b; - const uint64_t total_bytes = (uint64_t)total_b; - const uint64_t reserve_bytes = cuda_managed_kv_reserve_bytes(total_bytes); - if (context_bytes > free_bytes) return 1; - return free_bytes - context_bytes < reserve_bytes; -} - -extern "C" ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) { - if (!base || offset > base->bytes || bytes > base->bytes - offset) return NULL; - ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t)); - if (!t) return NULL; - t->ptr = (char *)base->ptr + offset; - t->bytes = bytes; - t->owner = 0; - t->device_id = base->device_id; /* inherit owning device */ - return t; -} - -extern "C" void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) { - if (!tensor) return; - int d = ds4_tensor_device_idx(tensor); - if (tensor->owner && tensor->ptr) { - WITH_DEVICE(g_gpu[d].device_id) { - (void)cudaFree(tensor->ptr); - } - } - free(tensor); -} - -extern "C" uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) { - return tensor ? tensor->bytes : 0; -} - -extern "C" void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) { - if (!tensor) return NULL; - /* Full-device sync preserves legacy semantics. */ - (void)cudaDeviceSynchronize(); - return tensor->ptr; -} - -extern "C" int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) { - if (!tensor || count > tensor->bytes / sizeof(float)) return 0; - if (count == 0) return 1; - int d = ds4_tensor_device_idx(tensor); - int ok = 0; - WITH_DEVICE(g_gpu[d].device_id) { - fill_f32_kernel<<<(count + 255u) / 256u, 256>>>((float *)tensor->ptr, count, value); - ok = cuda_ok(cudaGetLastError(), "tensor fill f32 launch"); - } - return ok; -} - -extern "C" int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) { - if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0; - int d = ds4_tensor_device_idx(tensor); - int ok = 0; - WITH_DEVICE(g_gpu[d].device_id) { - ok = cuda_ok(cudaMemcpy((char *)tensor->ptr + offset, data, (size_t)bytes, - cudaMemcpyHostToDevice), - "tensor write"); - } - return ok; -} - -extern "C" int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) { - if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0; - int d = ds4_tensor_device_idx(tensor); - int ok = 0; - WITH_DEVICE(g_gpu[d].device_id) { - ok = cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset, (size_t)bytes, - cudaMemcpyDeviceToHost), - "tensor read"); - } - return ok; -} - -extern "C" int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, - const ds4_gpu_tensor *src, uint64_t src_offset, - uint64_t bytes) { - if (!dst || !src || dst_offset > dst->bytes || src_offset > src->bytes || - bytes > dst->bytes - dst_offset || bytes > src->bytes - src_offset) { - return 0; - } - if (bytes == 0) return 1; - /* Same-device fast path; for cross-device, callers should use - * ds4_gpu_tensor_copy_xdev. We still tolerate cross-device callers - * here by routing to D2D copy on the destination's device. */ - int d = ds4_tensor_device_idx(dst); - int ok = 0; - WITH_DEVICE(g_gpu[d].device_id) { - ok = cuda_ok(cudaMemcpy((char *)dst->ptr + dst_offset, - (const char *)src->ptr + src_offset, - (size_t)bytes, - cudaMemcpyDeviceToDevice), - "tensor copy"); - } - return ok; -} - -__global__ static void moe_handoff_pack_kernel( - unsigned char *packed, - const float *ffn_norm, - const int32_t *selected, - const float *weights, - uint32_t n_embd, - uint32_t n_expert) { - const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; - float *packed_norm = (float *)packed; - int32_t *packed_selected = (int32_t *)(packed + (uint64_t)n_embd * sizeof(float)); - float *packed_weights = (float *)(packed + (uint64_t)n_embd * sizeof(float) + - (uint64_t)n_expert * sizeof(int32_t)); - if (i < n_embd) packed_norm[i] = ffn_norm[i]; - if (i < n_expert) { - packed_selected[i] = selected[i]; - packed_weights[i] = weights[i]; - } -} - -extern "C" int ds4_gpu_moe_handoff_pack_tensor( - ds4_gpu_tensor *packed, - const ds4_gpu_tensor *ffn_norm, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_embd, - uint32_t n_expert) { - if (!packed || !ffn_norm || !selected || !weights || - n_embd == 0 || n_expert == 0) { - return 0; - } - const uint64_t bytes = (uint64_t)n_embd * sizeof(float) + - (uint64_t)n_expert * sizeof(int32_t) + - (uint64_t)n_expert * sizeof(float); - if (packed->bytes < bytes || - ffn_norm->bytes < (uint64_t)n_embd * sizeof(float) || - selected->bytes < (uint64_t)n_expert * sizeof(int32_t) || - weights->bytes < (uint64_t)n_expert * sizeof(float)) { - return 0; - } - const uint32_t n = n_embd > n_expert ? n_embd : n_expert; - moe_handoff_pack_kernel<<<(n + 255u) / 256u, 256>>>( - (unsigned char *)packed->ptr, - (const float *)ffn_norm->ptr, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - n_embd, - n_expert); - return cuda_ok(cudaGetLastError(), "moe handoff pack launch"); -} - -/* Cross-device copy primitive. Path selection (highest priority first): - * DS4_FORCE_HOST_BOUNCE=1 -> always pinned-host bounce - * DS4_FORCE_CUDA_PEER=1 -> always cudaMemcpyPeerAsync (manual-testing - * override; bypasses g_gpu_peer_ok) - * otherwise -> peer if validation passed at init, else - * pinned-host bounce. */ -static int ds4_gpu_tensor_copy_xdev_impl(ds4_gpu_tensor *dst, - const ds4_gpu_tensor *src, - uint64_t bytes, - bool order_dst_before_write) { - if (!dst || !src) return 0; - if (bytes == 0) return 1; - if (bytes > dst->bytes || bytes > src->bytes) return 0; - int sd = ds4_tensor_device_idx(src); - int dd = ds4_tensor_device_idx(dst); - - /* Same-device fast path. */ - if (sd == dd) { - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, bytes, - cudaMemcpyDeviceToDevice, s), - "xdev same-device copy"); - if (ok && g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s), "xdev same-device sync"); - } - } - return ok; - } - - int peer = g_gpu_peer_ok[sd][dd]; - if (g_xdev_force_cuda_peer) peer = 1; - if (g_xdev_force_host_bounce) peer = 0; - - if (peer) { - int ok = 0; - if (order_dst_before_write) { - WITH_DEVICE(g_gpu[dd].device_id) { - cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; - cudaEvent_t e2 = (cudaEvent_t)g_gpu[dd].boundary_event; - ok = cuda_ok(cudaEventRecord(e2, s2), "peer dst-ready event record"); - } - if (!ok) return 0; - } else { - ok = 1; - } - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; - if (order_dst_before_write) { - ok = cuda_ok(cudaStreamWaitEvent(s, (cudaEvent_t)g_gpu[dd].boundary_event, 0), - "peer src wait dst-ready"); - } - if (ok) ok = cuda_ok(cudaMemcpyPeerAsync( - dst->ptr, g_gpu[dd].device_id, - src->ptr, g_gpu[sd].device_id, - bytes, s), - "peer copy"); - if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer event record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; - (void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0); - if (g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s2), "peer dst sync"); - } - } - return ok; - } - - /* Per-pair pinned-host bounce buffer. */ - if (g_xdev_bounce_bytes[sd][dd] < bytes) { - if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]); - if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes), - "bounce alloc")) return 0; - g_xdev_bounce_bytes[sd][dd] = bytes; - } - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; - ok = cuda_ok(cudaMemcpyAsync(g_xdev_bounce[sd][dd], src->ptr, bytes, - cudaMemcpyDeviceToHost, s), - "bounce d2h"); - if (ok) ok = cuda_ok(cudaEventRecord(e, s), "bounce event record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; - (void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0); - ok = cuda_ok(cudaMemcpyAsync(dst->ptr, g_xdev_bounce[sd][dd], bytes, - cudaMemcpyHostToDevice, s2), - "bounce h2d"); - if (ok) ok = cuda_ok(cudaStreamSynchronize(s2), "bounce dst sync"); - } - return ok; -} - -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_xdev_impl(dst, src, bytes, false); -} - -static int ds4_gpu_tensor_copy_xdev_default_impl(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 == 0u) return 1; - const int sd = ds4_tensor_device_idx(src); - const int dd = ds4_tensor_device_idx(dst); - if (sd == dd) { - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes, - cudaMemcpyDeviceToDevice, 0), - "default-stream same-device copy"); - } - return ok; - } - - int peer = g_gpu_peer_ok[sd][dd]; - if (g_xdev_force_cuda_peer) peer = 1; - if (g_xdev_force_host_bounce) peer = 0; - if (peer) { - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - ok = cuda_ok(cudaMemcpyPeerAsync( - dst->ptr, g_gpu[dd].device_id, - src->ptr, g_gpu[sd].device_id, - (size_t)bytes, 0), - "default-stream peer copy"); - if (ok) { - ok = cuda_ok(cudaEventRecord( - (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "default-stream peer event record"); - } - } - if (ok) { - WITH_DEVICE(g_gpu[dd].device_id) { - ok = cuda_ok(cudaStreamWaitEvent( - 0, - (cudaEvent_t)g_gpu[sd].boundary_event, - 0), - "default-stream peer destination wait"); - } - } - return ok; - } - - if (g_xdev_bounce_bytes[sd][dd] < bytes) { - if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]); - if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes), - "default-stream bounce alloc")) return 0; - g_xdev_bounce_bytes[sd][dd] = bytes; - } - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - ok = cuda_ok(cudaMemcpy(g_xdev_bounce[sd][dd], src->ptr, (size_t)bytes, - cudaMemcpyDeviceToHost), - "default-stream bounce d2h"); - } - if (ok) { - WITH_DEVICE(g_gpu[dd].device_id) { - ok = cuda_ok(cudaMemcpy(dst->ptr, g_xdev_bounce[sd][dd], - (size_t)bytes, cudaMemcpyHostToDevice), - "default-stream bounce h2d"); - } - } - return ok; -} - -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_default_impl(dst, src, bytes); -} - -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) { - ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2}; - const ds4_gpu_tensor *srcs[3] = {src0, src1, src2}; - const uint64_t sizes[3] = {bytes0, bytes1, bytes2}; - int sd = -1; - int dd = -1; - for (int i = 0; i < 3; i++) { - if (sizes[i] == 0u) continue; - if (!dsts[i] || !srcs[i] || sizes[i] > dsts[i]->bytes || - sizes[i] > srcs[i]->bytes) { - return 0; - } - const int this_sd = ds4_tensor_device_idx(srcs[i]); - const int this_dd = ds4_tensor_device_idx(dsts[i]); - if (sd < 0) { - sd = this_sd; - dd = this_dd; - } else if (sd != this_sd || dd != this_dd) { - return 0; - } - } - if (sd < 0) return 1; - if (sd == dd) { - int ok = 1; - WITH_DEVICE(g_gpu[sd].device_id) { - for (int i = 0; ok && i < 3; i++) { - if (sizes[i] == 0u) continue; - ok = cuda_ok(cudaMemcpyAsync( - dsts[i]->ptr, srcs[i]->ptr, (size_t)sizes[i], - cudaMemcpyDeviceToDevice, 0), - "grouped default same-device copy"); - } - } - return ok; - } - - int peer = g_gpu_peer_ok[dd][sd]; - if (g_xdev_force_cuda_peer) peer = 1; - if (g_xdev_force_host_bounce) peer = 0; - if (!peer) { - for (int i = 0; i < 3; i++) { - if (sizes[i] != 0u && - !ds4_gpu_tensor_copy_xdev_default_impl( - dsts[i], srcs[i], sizes[i])) { - return 0; - } - } - return 1; - } - - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - ok = cuda_ok(cudaEventRecord( - (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "grouped default source-ready record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - ok = cuda_ok(cudaStreamWaitEvent( - 0, (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "grouped default destination wait"); - for (int i = 0; ok && i < 3; i++) { - if (sizes[i] == 0u) continue; - ok = cuda_ok(cudaMemcpyPeerAsync( - dsts[i]->ptr, g_gpu[dd].device_id, - srcs[i]->ptr, g_gpu[sd].device_id, - (size_t)sizes[i], 0), - "grouped destination-stream peer copy"); - } - } - return ok; -} - -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) { - ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2}; - const ds4_gpu_tensor *srcs[3] = {src0, src1, src2}; - uint64_t bytes[3] = {bytes0, bytes1, bytes2}; - int first = -1; - for (int i = 0; i < 3; i++) { - if (bytes[i] == 0) continue; - if (!dsts[i] || !srcs[i] || - bytes[i] > dsts[i]->bytes || bytes[i] > srcs[i]->bytes) { - return 0; - } - if (first < 0) first = i; - } - if (first < 0) return 1; - - const int sd = ds4_tensor_device_idx(srcs[first]); - const int dd = ds4_tensor_device_idx(dsts[first]); - for (int i = first + 1; i < 3; i++) { - if (bytes[i] == 0) continue; - if (ds4_tensor_device_idx(srcs[i]) != sd || - ds4_tensor_device_idx(dsts[i]) != dd) { - int ok = 1; - for (int j = 0; ok && j < 3; j++) { - if (bytes[j] == 0) continue; - ok = ds4_gpu_tensor_copy_xdev(dsts[j], srcs[j], bytes[j]); - } - return ok; - } - } - - if (sd == dd) { - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - ok = 1; - for (int i = 0; ok && i < 3; i++) { - if (bytes[i] == 0) continue; - ok = cuda_ok(cudaMemcpyAsync(dsts[i]->ptr, srcs[i]->ptr, bytes[i], - cudaMemcpyDeviceToDevice, s), - "xdev3 same-device copy"); - } - if (ok && g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s), "xdev3 same-device sync"); - } - } - return ok; - } - - int peer = g_gpu_peer_ok[sd][dd]; - if (g_xdev_force_cuda_peer) peer = 1; - if (g_xdev_force_host_bounce) peer = 0; - if (!peer) { - int ok = 1; - for (int i = 0; ok && i < 3; i++) { - if (bytes[i] == 0) continue; - ok = ds4_gpu_tensor_copy_xdev(dsts[i], srcs[i], bytes[i]); - } - return ok; - } - - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; - ok = 1; - for (int i = 0; ok && i < 3; i++) { - if (bytes[i] == 0) continue; - ok = cuda_ok(cudaMemcpyPeerAsync( - dsts[i]->ptr, g_gpu[dd].device_id, - srcs[i]->ptr, g_gpu[sd].device_id, - bytes[i], s), - "peer copy3"); - } - if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer copy3 event record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; - ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "peer copy3 dst wait"); - if (ok && g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s2), "peer copy3 dst sync"); - } - } - return ok; -} - -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_impl(dst, src, bytes, true); -} - -extern "C" int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) { - if (!src) return 0; - if (dst_tier < 0 || dst_tier >= g_n_gpus) return 0; - int sd = ds4_tensor_device_idx(src); - int dd = dst_tier; - if (sd == dd) return 1; - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[sd].stream; - cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event; - ok = cuda_ok(cudaEventRecord(e, s), "xdev wait source event record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream; - ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "xdev wait destination wait"); - if (ok && g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s2), "xdev wait dst sync"); - } - } - return ok; -} - -extern "C" int ds4_gpu_tensor_wait_xdev_default( - const ds4_gpu_tensor *src, - int dst_tier) { - if (!src || dst_tier < 0 || dst_tier >= g_n_gpus) return 0; - const int sd = ds4_tensor_device_idx(src); - const int dd = dst_tier; - if (sd == dd) return 1; - int ok = 0; - WITH_DEVICE(g_gpu[sd].device_id) { - ok = cuda_ok(cudaEventRecord( - (cudaEvent_t)g_gpu[sd].boundary_event, 0), - "default xdev wait source event record"); - } - if (!ok) return 0; - WITH_DEVICE(g_gpu[dd].device_id) { - ok = cuda_ok(cudaStreamWaitEvent( - 0, - (cudaEvent_t)g_gpu[sd].boundary_event, - 0), - "default xdev wait destination wait"); - } - return ok; -} - -extern "C" int ds4_gpu_q8_cache_suppressed(void) { - return g_q8_cache_suppressed; -} - -extern "C" void ds4_gpu_set_q8_cache_suppressed(int suppressed) { - g_q8_cache_suppressed = suppressed ? 1 : 0; -} - -__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap); - -extern "C" int ds4_gpu_pack_slot_rows_f32_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *slots, - uint32_t n_rows, - uint32_t width, - uint32_t n_slots, - uint32_t slot_cap) { - uint64_t slot_rows = 0; - uint64_t slot_elems = 0; - uint64_t out_rows = 0; - uint64_t out_elems = 0; - if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 || - slot_cap == 0 || n_rows > slot_cap || - (uint64_t)n_slots > UINT64_MAX / slot_cap || - (slot_rows = (uint64_t)n_slots * slot_cap) > UINT64_MAX / width || - (slot_elems = slot_rows * width) > UINT64_MAX / sizeof(float) || - (uint64_t)n_rows > UINT64_MAX / n_slots || - (out_rows = (uint64_t)n_rows * n_slots) > UINT64_MAX / width || - (out_elems = out_rows * width) > UINT64_MAX / sizeof(float) || - slots->bytes < slot_elems * sizeof(float) || - out->bytes < out_elems * sizeof(float)) { - return 0; - } - const uint64_t blocks = (out_elems + 255u) / 256u; - if (blocks > UINT32_MAX) return 0; - pack_slot_rows_f32_kernel<<<(unsigned)blocks, 256>>>( - (float *)out->ptr, - (const float *)slots->ptr, - n_rows, - width, - n_slots, - slot_cap); - return cuda_ok(cudaGetLastError(), "pack_slot_rows_f32 launch"); -} - -extern "C" int ds4_gpu_begin_commands(void) { return 1; } -extern "C" int ds4_gpu_flush_commands(void) { return cuda_ok(cudaDeviceSynchronize(), "flush"); } -extern "C" int ds4_gpu_end_commands(void) { - if (g_cuda_end_stream_sync) { - return cuda_ok(cudaStreamSynchronize(0), "end commands stream"); - } - return cuda_ok(cudaDeviceSynchronize(), "end commands"); -} -extern "C" int ds4_gpu_synchronize(void) { return cuda_ok(cudaDeviceSynchronize(), "synchronize"); } - -extern "C" int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) { - if (!model_map || model_size == 0) return 0; - if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1; - cuda_stream_selected_cache_release(); - cuda_model_range_release_all(); - cuda_q8_f16_cache_release_all(); - g_q8_f16_disabled_after_oom = 0; - g_q8_f16_budget_notice_printed = 0; - for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { - (void)cudaFree(r.device_ptr); - } - g_q8_f32_ranges.clear(); - g_q8_f32_by_offset.clear(); - g_q8_f32_bytes = 0; - if (g_model_device_owned && g_model_device_base) { - (void)cudaFree((void *)g_model_device_base); - g_model_device_owned = 0; - } - if (g_model_registered && g_model_host_base) { - (void)cudaHostUnregister((void *)g_model_host_base); - g_model_registered = 0; - } - g_model_host_base = model_map; - g_model_device_base = (const char *)model_map; - g_model_registered_size = model_size; - g_model_range_mapping_supported = 1; - g_model_hmm_direct = 0; - g_model_cache_full = 0; - if (g_model_fd >= 0 && g_model_fd_host_base == NULL) { - g_model_fd_host_base = model_map; - } - - const char *copy_env = getenv("DS4_CUDA_COPY_MODEL"); - if (copy_env && copy_env[0]) { - void *dev = NULL; - const double t0 = clock() / (double)CLOCKS_PER_SEC; - cudaError_t err = cudaMalloc(&dev, (size_t)model_size); - if (err == cudaSuccess) { - fprintf(stderr, "ds4: CUDA copying %.2f GiB model to device memory\n", - (double)model_size / 1073741824.0); - err = cudaMemcpy(dev, model_map, (size_t)model_size, cudaMemcpyHostToDevice); - if (err == cudaSuccess) { - g_model_device_base = (const char *)dev; - g_model_device_owned = 1; - const double t1 = clock() / (double)CLOCKS_PER_SEC; - fprintf(stderr, "ds4: CUDA model copy complete in %.3fs\n", t1 - t0); - return 1; - } - fprintf(stderr, "ds4: CUDA model copy failed: %s\n", cudaGetErrorString(err)); - (void)cudaFree(dev); - (void)cudaGetLastError(); - } else { - fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - } - - cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size, - cudaHostRegisterMapped | cudaHostRegisterReadOnly); - if (err == cudaSuccess) { - void *dev = NULL; - err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0); - if (err == cudaSuccess && dev) { - g_model_device_base = (const char *)dev; - g_model_registered = 1; - fprintf(stderr, "ds4: CUDA registered %.2f GiB model mapping for device access\n", - (double)model_size / 1073741824.0); - } else { - fprintf(stderr, "ds4: CUDA host registration pointer lookup failed: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - } else { - fprintf(stderr, "ds4: CUDA host registration skipped: %s\n", cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - return 1; -} - -extern "C" int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) { - (void)max_tensor_bytes; - if (!ds4_gpu_register_model_map_no_copy(model_map, model_size)) return 0; - if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL && - !cuda_model_copy_chunked(model_map, model_size, map_offset, map_size)) { - (void)cuda_model_prefetch_range(model_map, model_size, map_offset, map_size); - } - return 1; -} - -/* Register the mmap'd host model pointer for selective-cache lookups WITHOUT - * triggering any device-side copy. Used by multi-GPU placement scaffolding's - * multi-tier path so DS4_CUDA_COPY_MODEL cannot reintroduce a full-model - * copy that defeats the per-device selective cache. - * - * This is the no-copy subset of ds4_gpu_set_model_map: same bookkeeping - * for the host pointer plus cudaHostRegister, but skipping the - * DS4_CUDA_COPY_MODEL branch that allocates and copies the entire model. */ -extern "C" int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size) { - if (!model_map || model_size == 0) return 0; - if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1; - - cuda_stream_selected_cache_release(); - cuda_model_range_release_all(); - cuda_q8_f16_cache_release_all(); - g_q8_f16_disabled_after_oom = 0; - g_q8_f16_budget_notice_printed = 0; - for (const cuda_q8_f32_range &r : g_q8_f32_ranges) { - (void)cudaFree(r.device_ptr); - } - g_q8_f32_ranges.clear(); - g_q8_f32_by_offset.clear(); - g_q8_f32_bytes = 0; - if (g_model_device_owned && g_model_device_base) { - (void)cudaFree((void *)g_model_device_base); - g_model_device_owned = 0; - } - if (g_model_registered && g_model_host_base) { - (void)cudaHostUnregister((void *)g_model_host_base); - g_model_registered = 0; - } - g_model_host_base = model_map; - g_model_device_base = (const char *)model_map; - g_model_registered_size = model_size; - g_model_range_mapping_supported = 1; - g_model_hmm_direct = 0; - g_model_cache_full = 0; - if (g_model_fd >= 0 && g_model_fd_host_base == NULL) { - g_model_fd_host_base = model_map; - } - - /* No DS4_CUDA_COPY_MODEL branch — that is the entire point. */ - - cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size, - cudaHostRegisterMapped | cudaHostRegisterReadOnly); - if (err == cudaSuccess) { - void *dev = NULL; - err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0); - if (err == cudaSuccess && dev) { - g_model_device_base = (const char *)dev; - g_model_registered = 1; - fprintf(stderr, - "ds4: CUDA (no-copy) registered %.2f GiB model mapping for multi-tier selective cache\n", - (double)model_size / 1073741824.0); - } else { - fprintf(stderr, - "ds4: CUDA (no-copy) host registration pointer lookup failed: %s\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - } else { - fprintf(stderr, - "ds4: CUDA (no-copy) host registration skipped: %s\n", - cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - return 1; -} - -/* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1). - * Maps to the physical CUDA device id stored in g_gpu[].device_id. - * Added for multi-GPU placement scaffolding (multi-GPU CLI); first executed by - * multi-GPU execution (follow-up). */ -extern "C" int ds4_gpu_set_current_device(int logical_tier) { - if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1; - if (!g_cuda_no_setdevice_cache && g_current_logical_tier == logical_tier) { - return 0; - } - if (cudaSetDevice(g_gpu[logical_tier].device_id) == cudaSuccess) { - g_current_logical_tier = logical_tier; - return 0; - } - g_current_logical_tier = -1; - return -1; -} - -/* Fenced device switch for sequential cross-device pipelines (GLM - * per-layer placement): work queued on the next device's default stream - * waits for everything queued so far on the previous device's default - * stream. Async — no host sync. Falls back to a plain switch when the - * device does not change. */ -extern "C" int ds4_gpu_set_current_device_fenced(int logical_tier) { - if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1; - static cudaEvent_t fence_ev[DS4_MAX_GPUS]; - /* Resolve the ACTUAL current device: WITH_DEVICE blocks and direct - * cudaSetDevice calls can leave g_current_logical_tier stale, and a - * false "already there" here strands work on the wrong device. */ - int cur_dev = -1; - (void)cudaGetDevice(&cur_dev); - int prev = -1; - for (int t = 0; t < g_n_gpus; t++) { - if (g_gpu[t].device_id == cur_dev) { prev = t; break; } - } - if (getenv("DS4_GLM_FENCE_TRACE")) { - fprintf(stderr, "ds4: fenced switch %d -> %d\n", prev, logical_tier); - } - if (prev == logical_tier) { - g_current_logical_tier = logical_tier; - return 0; - } - if (prev >= 0 && prev < g_n_gpus && prev != logical_tier) { - if (cudaSetDevice(g_gpu[prev].device_id) != cudaSuccess) return -1; - if (!fence_ev[prev] && - cudaEventCreateWithFlags(&fence_ev[prev], - cudaEventDisableTiming) != cudaSuccess) { - fence_ev[prev] = NULL; - } - if (fence_ev[prev]) { - (void)cudaEventRecord(fence_ev[prev], 0); - } - if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) { - g_current_logical_tier = -1; - return -1; - } - g_current_logical_tier = logical_tier; - if (fence_ev[prev]) { - (void)cudaStreamWaitEvent(0, fence_ev[prev], 0); - } - return 0; - } - return ds4_gpu_set_current_device(logical_tier); -} - -/* ========================================================================= - * Per-device selective model cache (selective model cache). - * - * ds4_gpu_device_cache_tensors copies the listed source ranges from the - * host mmap onto device_id's selective slab and appends sorted lookup - * entries. The legacy chunked-copy machinery (cuda_model_range_*) is - * NOT disturbed — it continues to drive all existing callers. New - * lookups fall back to it when no selective entry covers the range. - * - * Caller-context preference for overlap: when the same source range is - * cached on multiple devices, ds4_gpu_lookup_cache returns the entry - * whose device matches cudaGetDevice(). - * ========================================================================= */ - -extern "C" int ds4_gpu_device_cache_tensors(int device_id, - const ds4_tensor_range *ranges, - int n_ranges) { - if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1; - if (n_ranges < 0 || (!ranges && n_ranges > 0)) return 2; - if (n_ranges == 0) return 0; - - if (!g_model_host_base || g_model_registered_size == 0) return 3; - - /* Validate ranges against the mmap'd model bounds; reject ranges - * that overflow or extend past the mapped region. Done in a - * separate pass before any allocation so a bad input doesn't - * partially grow the slab. */ - uint64_t want_bytes = 0; - for (int i = 0; i < n_ranges; i++) { - if (ranges[i].target_device != device_id) continue; - const uint64_t off = ranges[i].source_offset; - const uint64_t nb = ranges[i].bytes; - /* Overflow-safe upper bound: off + nb must not exceed model - * size, and the sum must not wrap. */ - if (nb == 0) continue; - if (off > g_model_registered_size) return 8; - if (nb > g_model_registered_size - off) return 9; - /* Accumulate into want_bytes with overflow check. */ - if (want_bytes > UINT64_MAX - nb) return 10; - want_bytes += nb; - } - if (want_bytes == 0) return 0; - - cuda_device_cache &c = g_dev_cache[device_id]; - - int prev_device = -1; - if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1; - if (cudaSetDevice(device_id) != cudaSuccess) return 4; - - /* Allocate or grow the slab via cudaMalloc + d2d copy. */ - void *new_base = NULL; - size_t new_bytes = c.bytes + want_bytes; - - /* Refuse cleanly before cudaMalloc if the device clearly cannot hold - * the slab. The multi-tier packer reserves per-tier runtime scratch - * before placing tensors, but it cannot predict the cudaMalloc - * allocator's overhead (alignment, fragmentation after CUDA context - * init, default driver-side reservations). On a borderline budget - * that overhead pushes a "fits-by-packer-math" layout past the actual - * free pool and the cudaMalloc below OOMs after the engine already - * committed to the layout — same silent-late-OOM failure mode the - * upfront refusal path was added to eliminate. Catch it here too. */ - { - size_t free_b = 0, total_b = 0; - if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { - /* free_b already excludes the existing slab (it's still - * allocated), so the additional cudaMalloc only needs - * new_bytes free — not new_bytes + c.bytes. The old slab is - * freed AFTER the d2d copy succeeds. 2 GiB safety covers what - * the engine will allocate AFTER the cache slab in the same - * session_create: per-tier graph scratch (the planner can't - * predict its cumulative cudaMalloc alignment overhead), - * cuBLAS workspace beyond the 64 MiB the packer already - * reserves, and driver-side allocator slack. Without this - * headroom a borderline budget that fits the slab itself can - * still OOM at the per-tier tensor allocations a few moments - * later — same silent-late-OOM failure mode, one layer up. */ - const size_t safety = (size_t)2ull * 1024ull * 1024ull * 1024ull; - const size_t need = new_bytes + safety; - if (need > free_b) { - fprintf(stderr, - "ds4: device cache slab needs %.2f GiB on device %d " - "but only %.2f GiB free (slab=%.2f GiB + %.2f GiB safety). " - "Lower --gpu-vram / --ctx-max, or use --gpu-vram auto on " - "a host with more free VRAM. Refusing upfront to avoid " - "late OOM at cudaMalloc.\n", - (double)need / 1073741824.0, - device_id, - (double)free_b / 1073741824.0, - (double)new_bytes / 1073741824.0, - (double)safety / 1073741824.0); - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 5; - } - } - /* If cudaMemGetInfo itself failed, fall through; cudaMalloc's own - * error path still catches the late case, just with a less helpful - * message. */ - } - - if (!cuda_ok(cudaMalloc(&new_base, new_bytes), "device cache alloc")) { - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 5; - } - if (c.present && c.bytes > 0) { - cudaError_t e = cudaMemcpy(new_base, c.base, c.bytes, - cudaMemcpyDeviceToDevice); - if (e != cudaSuccess) { - cuda_ok(e, "device cache grow d2d"); - (void)cudaFree(new_base); - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 6; - } - /* Re-base existing entries on this device. */ - char *old_base = (char *)c.base; - char *grown = (char *)new_base; - for (size_t k = 0; k < g_cache_ranges.size(); k++) { - if (g_cache_ranges[k].device_id == device_id) { - g_cache_ranges[k].device_ptr = - grown + ((char *)g_cache_ranges[k].device_ptr - old_base); - } - } - (void)cudaFree(c.base); - } - c.base = new_base; - c.bytes = new_bytes; - c.present = 1; - - /* Copy ranges and append entries. */ - const char *host_base = (const char *)g_model_host_base; - size_t write_off = c.bytes - want_bytes; - for (int i = 0; i < n_ranges; i++) { - if (ranges[i].target_device != device_id) continue; - char *dev_ptr = (char *)c.base + write_off; - cudaError_t e = cudaMemcpy(dev_ptr, - host_base + ranges[i].source_offset, - (size_t)ranges[i].bytes, - cudaMemcpyHostToDevice); - if (e != cudaSuccess) { - cuda_ok(e, "device cache range h2d"); - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 7; - } - cache_range_entry ent; - ent.source_offset = ranges[i].source_offset; - ent.bytes = ranges[i].bytes; - ent.device_id = device_id; - ent.device_ptr = dev_ptr; - g_cache_ranges.push_back(ent); - write_off += ranges[i].bytes; - } - - /* Keep sorted by source_offset for binary-search lookup. */ - std::sort(g_cache_ranges.begin(), g_cache_ranges.end(), - [](const cache_range_entry &a, const cache_range_entry &b) { - if (a.source_offset != b.source_offset) - return a.source_offset < b.source_offset; - return a.device_id < b.device_id; - }); - - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 0; -} - -/* Install support-model tensor ranges into device_id's strict cache, - * copying from the registered support map and keying entries at - * source_offset + bias. Standalone slab (does not touch the main cache - * slab growth path). */ -extern "C" int ds4_gpu_device_cache_support_tensors(int device_id, - int entry_device_id, - const ds4_tensor_range *ranges, - int n_ranges, - int from_main_map) { - if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1; - if (entry_device_id < 0 || entry_device_id >= DS4_MAX_GPUS) return 1; - if (n_ranges <= 0 || !ranges) return 2; - const char *src_base; - uint64_t src_size; - uint64_t key_bias; - if (from_main_map) { - /* Auxiliary main-model ranges (e.g. the embedding bucket for the - * DSpark executor tier): standalone slab, unbiased offsets. */ - src_base = (const char *)g_model_host_base; - src_size = g_model_registered_size; - key_bias = 0; - } else { - src_base = (const char *)g_support_host_base; - src_size = g_support_host_size; - key_bias = g_support_offset_bias; - if (key_bias == 0) return 3; - } - if (!src_base || src_size == 0) return 3; - uint64_t want = 0; - for (int i = 0; i < n_ranges; i++) { - const uint64_t off = ranges[i].source_offset; - const uint64_t nb = ranges[i].bytes; - if (nb == 0) continue; - if (off > src_size || nb > src_size - off) return 8; - if (want > UINT64_MAX - nb) return 9; - want += nb; - } - if (want == 0) return 0; - int prev_device = -1; - if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1; - if (cudaSetDevice(device_id) != cudaSuccess) return 4; - void *base = NULL; - if (!cuda_ok(cudaMalloc(&base, (size_t)want), "support cache alloc")) { - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 5; - } - const char *host_base = src_base; - size_t write_off = 0; - for (int i = 0; i < n_ranges; i++) { - if (ranges[i].bytes == 0) continue; - char *dev_ptr = (char *)base + write_off; - cudaError_t e = cudaMemcpy(dev_ptr, - host_base + ranges[i].source_offset, - (size_t)ranges[i].bytes, - cudaMemcpyHostToDevice); - if (e != cudaSuccess) { - cuda_ok(e, "support cache range h2d"); - (void)cudaFree(base); - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 7; - } - cache_range_entry ent; - ent.source_offset = ranges[i].source_offset + key_bias; - ent.bytes = ranges[i].bytes; - /* Entries can claim a different (executor) device than the one the - * slab physically lives on: strict lookups filter by entry device, - * and peer access lets the executor's kernels dereference the - * spilled pointer directly. */ - ent.device_id = entry_device_id; - ent.device_ptr = dev_ptr; - g_cache_ranges.push_back(ent); - write_off += ranges[i].bytes; - } - std::sort(g_cache_ranges.begin(), g_cache_ranges.end(), - [](const cache_range_entry &a, const cache_range_entry &b) { - if (a.source_offset != b.source_offset) - return a.source_offset < b.source_offset; - return a.device_id < b.device_id; - }); - if (getenv("DS4_DSPARK_VERIFY_CACHE") != NULL) { - /* Read back every installed range and compare with the host copy. */ - int bad = 0; - write_off = 0; - for (int i = 0; i < n_ranges; i++) { - if (ranges[i].bytes == 0) continue; - char *dev_ptr = (char *)base + write_off; - std::vector tmp((size_t)ranges[i].bytes); - if (cudaMemcpy(tmp.data(), dev_ptr, (size_t)ranges[i].bytes, - cudaMemcpyDeviceToHost) != cudaSuccess || - memcmp(tmp.data(), host_base + ranges[i].source_offset, - (size_t)ranges[i].bytes) != 0) { - fprintf(stderr, - "ds4: support cache VERIFY MISMATCH offset=%llu bytes=%llu dev=%d\n", - (unsigned long long)ranges[i].source_offset, - (unsigned long long)ranges[i].bytes, device_id); - bad++; - } - write_off += ranges[i].bytes; - } - fprintf(stderr, "ds4: support cache verify dev=%d ranges=%d bad=%d\n", - device_id, n_ranges, bad); - } - if (prev_device >= 0) (void)cudaSetDevice(prev_device); - return 0; -} - -extern "C" int ds4_gpu_lookup_cache(uint64_t source_offset, uint64_t bytes, - int *out_device_id, void **out_device_ptr) { - int active_device = -1; - (void)cudaGetDevice(&active_device); - - if (!g_cache_ranges.empty()) { - /* upper_bound: first entry with source_offset > query. - * Candidates are at strictly earlier positions; scan all of - * them rather than breaking on the first non-covering entry, - * because the table allows overlap across devices. */ - auto it = std::upper_bound( - g_cache_ranges.begin(), g_cache_ranges.end(), - source_offset, - [](uint64_t off, const cache_range_entry &e) { - return off < e.source_offset; - }); - const cache_range_entry *match_any = NULL; - const cache_range_entry *match_pref = NULL; - while (it != g_cache_ranges.begin()) { - --it; - /* Overflow-safe coverage check: - * 1. source_offset >= it->source_offset - * 2. bytes <= it->bytes - (source_offset - it->source_offset) - * The second form computes only the remaining capacity inside - * the entry, so neither side can overflow even with bytes == - * UINT64_MAX. */ - if (source_offset >= it->source_offset) { - uint64_t into = source_offset - it->source_offset; - if (into <= it->bytes && bytes <= it->bytes - into) { - if (it->device_id == active_device) { - match_pref = &*it; - break; - } - if (!match_any) match_any = &*it; - } - } - /* Do NOT break on non-covering: an earlier entry may still - * cover if its bytes extend far enough. */ - } - const cache_range_entry *m = match_pref ? match_pref : match_any; - if (m) { - if (out_device_id) *out_device_id = m->device_id; - if (out_device_ptr) { - *out_device_ptr = - (char *)m->device_ptr + (source_offset - m->source_offset); - } - return 1; - } - } - - /* Legacy chunk-aware fallback (device 0 only). */ - const char *p = cuda_model_range_ptr_from_fd(g_model_host_base, - source_offset, bytes, - "lookup_cache"); - if (p) { - if (out_device_id) *out_device_id = 0; - if (out_device_ptr) *out_device_ptr = (void *)p; - return 1; - } - return 0; -} - -extern "C" int ds4_gpu_lookup_cache_device(uint64_t source_offset, uint64_t bytes) { - int d = -1; - if (!ds4_gpu_lookup_cache(source_offset, bytes, &d, NULL)) return -1; - return d; -} - -/* Strict per-device selective-cache lookup. - * - * Returns 1 only if a covering entry exists whose device_id matches the - * caller-supplied expected_device. Otherwise returns 0 with *out_device_ptr - * untouched. Unlike ds4_gpu_lookup_cache, this variant performs NO host- - * pointer fallback (no FD-cache, no model_range_ptr_from_fd) and NO - * different-device fallback. It is the canonical lookup for multi-tier - * dispatch where consuming a different device's pointer would be a - * correctness bug. expected_device is a PHYSICAL CUDA device id (the - * value stored in g_gpu[logical_tier].device_id, not the logical tier - * index). 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. Added for - * multi-GPU execution (multi-GPU execution). */ -extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset, - uint64_t bytes, - int expected_device, - void **out_device_ptr) { - if (g_cache_ranges.empty()) return 0; - - auto it = std::upper_bound( - g_cache_ranges.begin(), g_cache_ranges.end(), - source_offset, - [](uint64_t off, const cache_range_entry &e) { - return off < e.source_offset; - }); - while (it != g_cache_ranges.begin()) { - --it; - if (source_offset < it->source_offset) { - /* Should not happen given upper_bound semantics, but defensive. */ - continue; - } - uint64_t into = source_offset - it->source_offset; - if (into > it->bytes) continue; - if (bytes > it->bytes - into) continue; - if (it->device_id != expected_device) continue; - if (out_device_ptr) { - *out_device_ptr = - (char *)it->device_ptr + (source_offset - it->source_offset); - } - return 1; - } - return 0; -} - -extern "C" int ds4_gpu_set_model_fd(int fd) { - g_model_fd = fd; - g_model_fd_host_base = g_model_host_base; - g_model_file_size = 0; - if (g_model_direct_fd >= 0) { - (void)close(g_model_direct_fd); - g_model_direct_fd = -1; - } - g_model_direct_align = 1; - if (fd >= 0) { - struct stat st; - if (fstat(fd, &st) == 0 && st.st_size > 0) { - g_model_file_size = (uint64_t)st.st_size; - if (st.st_blksize > 1) g_model_direct_align = (uint64_t)st.st_blksize; - } -#if defined(__linux__) && defined(O_DIRECT) - if (getenv("DS4_CUDA_NO_DIRECT_IO") == NULL) { - char proc_path[64]; - snprintf(proc_path, sizeof(proc_path), "/proc/self/fd/%d", fd); - int direct_fd = open(proc_path, O_RDONLY | O_DIRECT); - if (direct_fd >= 0) { - g_model_direct_fd = direct_fd; - if (g_model_direct_align < 512) g_model_direct_align = 512; - if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA model direct I/O enabled (align=%llu)\n", - (unsigned long long)g_model_direct_align); - } - } else if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { - fprintf(stderr, "ds4: CUDA model direct I/O unavailable: %s\n", strerror(errno)); - } - } -#endif - } - return 1; -} - -extern "C" int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label) { - if (!model_map || bytes == 0) return 1; - if (offset > model_size || bytes > model_size - offset) return 0; - if (!cuda_model_range_ptr(model_map, offset, bytes, label ? label : "model_tensor")) return 0; - return cuda_model_range_is_cached(model_map, offset, bytes); -} - -extern "C" int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label) { - if (!model_map || bytes == 0) return 1; - if (offset > model_size || bytes > model_size - offset) return 0; - static int optional_q8_preload_disabled = 0; - if (optional_q8_preload_disabled) return 1; - const char *cache_label = label ? label : "q8_0"; - /* Preload runs before any multi-tier dispatch. The cache entries it creates - * are device-0 by construction; multi-tier callers in kernel wrappers will - * miss the linear scan (device_id filter) and allocate fresh per-device - * copies the first time they're consulted. */ - if (getenv("DS4_CUDA_Q8_F32_PRELOAD") != NULL && - cuda_q8_f32_cache_allowed(cache_label, in_dim, out_dim)) { - if (cuda_q8_f32_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1; - optional_q8_preload_disabled = 1; - return 1; - } - if (!cuda_q8_f16_preload_allowed(cache_label, in_dim, out_dim)) return 1; - if (cuda_q8_f16_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1; - optional_q8_preload_disabled = 1; - return 1; -} - -extern "C" void ds4_gpu_print_memory_report(const char *label) { - size_t free_b = 0, total_b = 0; - (void)cudaMemGetInfo(&free_b, &total_b); - fprintf(stderr, "ds4: CUDA memory report %s: free %.2f MiB total %.2f MiB\n", - label ? label : "", (double)free_b / 1048576.0, (double)total_b / 1048576.0); -} - -extern "C" void ds4_gpu_set_quality(bool quality) { - g_quality_mode = quality ? 1 : 0; - const cublasMath_t math_mode = - (g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL) - ? CUBLAS_DEFAULT_MATH - : CUBLAS_TF32_TENSOR_OP_MATH; - /* Walk every initialized per-tier handle. Single-tier (g_n_gpus == 1) - * walks exactly one entry. On any device-switch failure, - * skip the tier and continue — the function is void and the math-mode - * setting is advisory, but log so misconfiguration is visible. */ - for (int i = 0; i < g_n_gpus; i++) { - if (!g_gpu[i].cublas_ready || !g_gpu[i].cublas) continue; - int prev = -1; - cudaError_t derr = cudaGetDevice(&prev); - if (derr != cudaSuccess) { - fprintf(stderr, - "ds4: ds4_gpu_set_quality: cudaGetDevice failed before tier %d " - "(dev=%d): %s; skipping\n", - i, g_gpu[i].device_id, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - continue; - } - derr = cudaSetDevice(g_gpu[i].device_id); - if (derr != cudaSuccess) { - fprintf(stderr, - "ds4: ds4_gpu_set_quality: cudaSetDevice(%d) failed for tier %d: " - "%s; skipping\n", - g_gpu[i].device_id, i, cudaGetErrorString(derr)); - (void)cudaGetLastError(); - if (prev >= 0) (void)cudaSetDevice(prev); - continue; - } - cublasStatus_t st = cublasSetMathMode((cublasHandle_t)g_gpu[i].cublas, math_mode); - if (st != CUBLAS_STATUS_SUCCESS) { - fprintf(stderr, - "ds4: ds4_gpu_set_quality: cublasSetMathMode failed on tier %d " - "(dev=%d): status %d\n", - i, g_gpu[i].device_id, (int)st); - } - if (prev >= 0) (void)cudaSetDevice(prev); - } -} - -__global__ static void embed_token_hc_kernel(float *out, const unsigned short *w, uint32_t token, uint32_t n_embd, uint32_t n_hc) { - uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; - uint32_t n = n_embd * n_hc; - if (i >= n) return; - uint32_t e = i % n_embd; - out[i] = __half2float(reinterpret_cast(w)[(uint64_t)token * n_embd + e]); -} - -__global__ static void embed_tokens_hc_kernel( - float *out, - const int32_t *tokens, - const __half *w, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd, - uint32_t n_hc) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; - if (gid >= n) return; - uint32_t d = gid % n_embd; - uint64_t tmp = gid / n_embd; - uint32_t t = tmp / n_hc; - int32_t tok_i = tokens[t]; - uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i; - if (tok >= n_vocab) tok = 0; - out[gid] = __half2float(w[(uint64_t)tok * n_embd + d]); -} - -__global__ static void matmul_f16_kernel( - float *out, - const __half *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok) return; - - float sum = 0.0f; - const __half *wr = w + row * in_dim; - const float *xr = x + tok * in_dim; - for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) { - sum += __half2float(wr[i]) * xr[i]; - } - - __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(); - } - if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; -} - -__global__ static void matmul_f16_serial_kernel( - float *out, - const __half *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok || threadIdx.x != 0) return; - - float sum = 0.0f; - const __half *wr = w + row * in_dim; - const float *xr = x + tok * in_dim; - for (uint64_t i = 0; i < in_dim; i++) { - sum += __half2float(wr[i]) * xr[i]; - } - out[tok * out_dim + row] = sum; -} - -__global__ static void matmul_f16_ordered_chunks_kernel( - float *out, - const __half *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok) return; - - __shared__ float partial[32]; - const uint32_t tid = threadIdx.x; - float sum = 0.0f; - const uint64_t chunk = (in_dim + 31u) / 32u; - const uint64_t k0 = (uint64_t)tid * chunk; - uint64_t k1 = k0 + chunk; - if (k1 > in_dim) k1 = in_dim; - const __half *wr = w + row * in_dim; - const float *xr = x + tok * in_dim; - for (uint64_t i = k0; i < k1; i++) { - sum += __half2float(wr[i]) * xr[i]; - } - partial[tid] = sum; - __syncthreads(); - if (tid == 0) { - float total = 0.0f; - for (uint32_t i = 0; i < 32u; i++) total += partial[i]; - out[tok * out_dim + row] = total; - } -} - -__global__ static void matmul_f16_small_out_hx_ordered_chunks_kernel( - float *out, - const __half *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim) { - uint64_t row = (uint64_t)blockIdx.x; - if (row >= out_dim) return; - - __shared__ float partial[32]; - const uint32_t tid = threadIdx.x; - float sum = 0.0f; - const uint64_t chunk = (in_dim + 31u) / 32u; - const uint64_t k0 = (uint64_t)tid * chunk; - uint64_t k1 = k0 + chunk; - if (k1 > in_dim) k1 = in_dim; - const __half *wr = w + row * in_dim; - for (uint64_t i = k0; i < k1; i++) { - const float xv = __half2float(__float2half(x[i])); - sum += __half2float(wr[i]) * xv; - } - partial[tid] = sum; - __syncthreads(); - if (tid == 0) { - float total = 0.0f; - for (uint32_t i = 0; i < 32u; i++) total += partial[i]; - out[row] = total; - } -} - -__global__ static void matmul_f16_small_out_batch_kernel( - float *out, - const __half *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - const uint64_t tok = (uint64_t)blockIdx.x; - const uint32_t tid = threadIdx.x; - if (tok >= n_tok || out_dim > 32u || blockDim.x != 256u) return; - - float acc[32]; - #pragma unroll - for (uint32_t r = 0; r < 32u; r++) acc[r] = 0.0f; - - const float *xr = x + tok * in_dim; - for (uint64_t i = tid; i < in_dim; i += 256u) { - const float xv = xr[i]; - #pragma unroll - for (uint32_t r = 0; r < 32u; r++) { - if (r < out_dim) { - acc[r] += __half2float(w[(uint64_t)r * in_dim + i]) * xv; - } - } - } - - __shared__ float partial[32 * 256]; - #pragma unroll - for (uint32_t r = 0; r < 32u; r++) { - if (r < out_dim) partial[r * 256u + tid] = acc[r]; - } - __syncthreads(); - - for (uint32_t stride = 128u; stride > 0u; stride >>= 1u) { - if (tid < stride) { - #pragma unroll - for (uint32_t r = 0; r < 32u; r++) { - if (r < out_dim) { - partial[r * 256u + tid] += partial[r * 256u + tid + stride]; - } - } - } - __syncthreads(); - } - - if (tid == 0) { - #pragma unroll - for (uint32_t r = 0; r < 32u; r++) { - if (r < out_dim) out[tok * out_dim + r] = partial[r * 256u]; - } - } -} - -__global__ static void matmul_f16_pair_ordered_chunks_kernel( - float *out0, - float *out1, - const __half *w0, - const __half *w1, - const float *x, - uint64_t in_dim, - uint64_t out0_dim, - uint64_t out1_dim) { - uint64_t row = (uint64_t)blockIdx.x; - if (row >= out0_dim && row >= out1_dim) return; - - __shared__ float partial0[32]; - __shared__ float partial1[32]; - const uint32_t tid = threadIdx.x; - float sum0 = 0.0f; - float sum1 = 0.0f; - const uint64_t chunk = (in_dim + 31u) / 32u; - const uint64_t k0 = (uint64_t)tid * chunk; - uint64_t k1 = k0 + chunk; - if (k1 > in_dim) k1 = in_dim; - const __half *wr0 = row < out0_dim ? w0 + row * in_dim : w0; - const __half *wr1 = row < out1_dim ? w1 + row * in_dim : w1; - for (uint64_t i = k0; i < k1; i++) { - const float xv = x[i]; - if (row < out0_dim) sum0 += __half2float(wr0[i]) * xv; - if (row < out1_dim) sum1 += __half2float(wr1[i]) * xv; - } - partial0[tid] = sum0; - partial1[tid] = sum1; - __syncthreads(); - if (tid == 0) { - float total0 = 0.0f; - float total1 = 0.0f; - for (uint32_t i = 0; i < 32u; i++) { - total0 += partial0[i]; - total1 += partial1[i]; - } - if (row < out0_dim) out0[row] = total0; - if (row < out1_dim) out1[row] = total1; - } -} - -__global__ static void matmul_f32_kernel( - float *out, - const float *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok) return; - - float sum = 0.0f; - const float *wr = w + row * in_dim; - const float *xr = x + tok * in_dim; - for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) { - sum += wr[i] * xr[i]; - } - - __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(); - } - if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; -} - -__global__ static void repeat_hc_kernel(float *out, const float *row, uint32_t n_embd, uint32_t n_hc) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_embd * n_hc; - if (i >= n) return; - out[i] = row[i % n_embd]; -} - -__global__ static void repeat_hc_rows_kernel(float *out, const float *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; - if (i >= n) return; - - uint64_t hc_row = (uint64_t)n_hc * n_embd; - uint64_t tok = i / hc_row; - uint64_t embd = i % n_embd; - out[i] = rows[tok * n_embd + embd]; -} - -__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_rows * n_slots * width; - if (i >= n) return; - - uint64_t col = i % width; - uint64_t slot = (i / width) % n_slots; - uint64_t row = i / ((uint64_t)n_slots * width); - out[i] = slots[((slot * slot_cap) + row) * width + col]; -} - -__global__ static void f32_to_f16_kernel(__half *out, const float *x, uint64_t n) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) out[i] = __float2half(x[i]); -} - -__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 warp_max_f32(float v) { - for (int offset = 16; offset > 0; offset >>= 1) { - v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, offset)); - } - return v; -} - -__device__ static float dot4_f32(float4 a, float4 b) { - return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w; -} - -__device__ __forceinline__ static int32_t load_i8x4_i32_aligned(const int8_t *p) { - return *(const int32_t *)p; -} - -__device__ __forceinline__ static 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__ static 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__ static 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; -} - -__global__ static DS4_CUDA_UNUSED void matmul_q8_0_kernel( - float *out, - const unsigned char *w, - const float *x, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok) return; - const uint64_t blocks = (in_dim + 31) / 32; - const unsigned char *wr = w + row * blocks * 34; - const float *xr = x + tok * in_dim; - float acc = 0.0f; - - for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { - uint64_t i0 = b * 32; - uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32; - float amax = 0.0f; - for (uint64_t i = 0; i < bn; i++) amax = fmaxf(amax, fabsf(xr[i0 + i])); - float d = amax / 127.0f; - float id = d != 0.0f ? 1.0f / d : 0.0f; - const __half *scale_h = (const __half *)(wr + b * 34); - const int8_t *qs = (const int8_t *)(wr + b * 34 + 2); - int dot = 0; - for (uint64_t i = 0; i < bn; i++) { - int q = (int)lrintf(xr[i0 + i] * id); - q = q > 127 ? 127 : (q < -128 ? -128 : q); - dot += (int)qs[i] * q; - } - acc += __half2float(*scale_h) * d * (float)dot; - } - - __shared__ float partial[256]; - partial[threadIdx.x] = acc; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; -} - -__global__ static void quantize_q8_0_f32_kernel( - int8_t *xq, - float *xscale, - const float *x, - uint64_t in_dim, - uint64_t blocks) { - uint64_t b = blockIdx.x; - uint64_t tok = blockIdx.y; - if (b >= blocks) return; - uint64_t i0 = b * 32; - uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32; - const float *xr = x + tok * in_dim + i0; - - float a = 0.0f; - if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]); - __shared__ float vals[32]; - vals[threadIdx.x] = a; - __syncthreads(); - for (uint32_t stride = 16; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) vals[threadIdx.x] = fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]); - __syncthreads(); - } - const float d = vals[0] / 127.0f; - const float id = d != 0.0f ? 1.0f / d : 0.0f; - if (threadIdx.x == 0) xscale[tok * blocks + b] = d; - int8_t *dst = xq + (tok * blocks + b) * 32; - if (threadIdx.x < bn) { - int v = (int)lrintf(xr[threadIdx.x] * id); - v = v > 127 ? 127 : (v < -128 ? -128 : v); - dst[threadIdx.x] = (int8_t)v; - } else { - dst[threadIdx.x] = 0; - } -} - -__global__ static void quantize_q8_0_group_slice_rows_kernel( - int8_t *xq, - float *xscale, - const float *x, - uint64_t group_dim, - uint64_t blocks, - uint32_t n_groups_total, - uint32_t group0, - uint32_t group_cnt) { - const uint64_t b = blockIdx.x; - const uint64_t packed_row = blockIdx.y; - if (b >= blocks) return; - const uint64_t token = packed_row / group_cnt; - const uint64_t group = group0 + packed_row - token * group_cnt; - const uint64_t i0 = b * 32u; - const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u; - const float *xr = x + - (token * n_groups_total + group) * group_dim + i0; - - float a = 0.0f; - if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]); - __shared__ float vals[32]; - vals[threadIdx.x] = a; - __syncthreads(); - for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - vals[threadIdx.x] = - fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]); - } - __syncthreads(); - } - const float d = vals[0] / 127.0f; - const float id = d != 0.0f ? 1.0f / d : 0.0f; - if (threadIdx.x == 0u) xscale[packed_row * blocks + b] = d; - int8_t *dst = xq + (packed_row * blocks + b) * 32u; - if (threadIdx.x < bn) { - int v = (int)lrintf(xr[threadIdx.x] * id); - v = v > 127 ? 127 : (v < -128 ? -128 : v); - dst[threadIdx.x] = (int8_t)v; - } else { - dst[threadIdx.x] = 0; - } -} - -__global__ static void matmul_q8_0_preq_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - int use_dp4a) { - uint64_t row = (uint64_t)blockIdx.x; - uint64_t tok = (uint64_t)blockIdx.y; - if (row >= out_dim || tok >= n_tok) return; - const unsigned char *wr = w + row * blocks * 34; - const int8_t *xqr = xq + tok * blocks * 32; - const float *xsr = xscale + tok * blocks; - float acc = 0.0f; - for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { - uint64_t i0 = b * 32; - 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 = xqr + b * 32; - int dot = dot_i8_block(qs, xqb, bn, use_dp4a); - acc += __half2float(*scale_h) * xsr[b] * (float)dot; - } - __shared__ float partial[256]; - partial[threadIdx.x] = acc; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; -} - -__global__ static void matmul_q8_0_preq_warp8_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks, - int use_dp4a) { - uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); - const uint64_t tok = (uint64_t)blockIdx.y; - uint32_t lane = threadIdx.x & 31u; - if (row >= out_dim) return; - const unsigned char *wr = w + row * blocks * 34; - const int8_t *xqr = xq + tok * blocks * 32u; - const float *xsr = xscale + tok * blocks; - float acc = 0.0f; - 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 __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) out[tok * out_dim + row] = acc; -} - -__device__ __forceinline__ static uint32_t q8_top1_float_ordered_key(float v) { - const uint32_t u = __float_as_uint(v); - return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); -} - -__device__ __forceinline__ static uint64_t q8_top1_pack_key(float v, uint32_t idx) { - return ((uint64_t)q8_top1_float_ordered_key(v) << 32u) | - (uint64_t)(0xffffffffu - idx); -} - -__device__ __forceinline__ static float q8_top1_unpack_value(uint32_t ordered) { - const uint32_t u = (ordered & 0x80000000u) - ? (ordered ^ 0x80000000u) - : ~ordered; - return __uint_as_float(u); -} - -__global__ static void matmul_q8_0_top1_preq_warp8_kernel( - unsigned long long *best_key, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks, - uint32_t index_offset, - int use_dp4a) { - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t lane = threadIdx.x & 31u; - const uint64_t row = (uint64_t)blockIdx.x * 8u + warp; - const bool valid = row < out_dim; - float acc = 0.0f; - - if (valid) { - const unsigned char *wr = w + row * blocks * 34; - 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 __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); - - __shared__ unsigned long long keys[8]; - if (lane == 0u) { - keys[warp] = valid - ? (unsigned long long)q8_top1_pack_key(acc, index_offset + (uint32_t)row) - : 0ull; - } - __syncthreads(); - - if (threadIdx.x == 0u) { - unsigned long long block_best = keys[0]; - #pragma unroll - for (uint32_t i = 1u; i < 8u; i++) { - if (keys[i] > block_best) block_best = keys[i]; - } - (void)atomicMax(best_key, block_best); - } -} - -__global__ static void matmul_q8_0_top1_unpack_kernel( - uint32_t *selected, - float *values, - const unsigned long long *best_key) { - if (threadIdx.x != 0u || blockIdx.x != 0u) return; - const uint64_t key = (uint64_t)best_key[0]; - const uint32_t ordered = (uint32_t)(key >> 32u); - const uint32_t idx = 0xffffffffu - (uint32_t)key; - selected[0] = idx; - values[0] = q8_top1_unpack_value(ordered); -} - -__global__ static void matmul_q8_0_kslice_preq_warp8_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t slice_dim, - uint64_t out_dim, - uint64_t full_blocks, - uint64_t block_start, - uint64_t slice_blocks, - int use_dp4a) { - const uint64_t row = - (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); - const uint64_t tok = blockIdx.y; - const uint32_t lane = threadIdx.x & 31u; - if (row >= out_dim) return; - out += tok * out_dim; - xq += tok * slice_blocks * 32u; - xscale += tok * slice_blocks; - const unsigned char *wr = - w + row * full_blocks * 34u + block_start * 34u; - float acc = 0.0f; - for (uint64_t b = lane; b < slice_blocks; b += 32u) { - uint64_t i0 = b * 32u; - uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u; - const __half *scale_h = (const __half *)(wr + b * 34u); - const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); - const int8_t *xqb = xq + b * 32u; - 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) out[row] = acc; -} - -__global__ static void matmul_q8_0_pair_preq_warp8_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); - const uint64_t tok = (uint64_t)blockIdx.y; - 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; - const int8_t *xqr = xq + tok * blocks * 32u; - const float *xsr = xscale + tok * blocks; - 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 = xqr + b * 32; - const float xs = xsr[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[tok * out0_dim + row] = acc0; - if (row < out1_dim) out1[tok * out1_dim + row] = acc1; - } -} - -__global__ static void shared_mid_q8_0_preq_warp8_exact_kernel( - float *mid, - const unsigned char *gate_w, - const unsigned char *up_w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks, - float clamp, - const int32_t *selected, - uint32_t expert_split, - bool home_rank, - 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; - if (selected) { - /* Complementary predicates select exactly one writer; ties stay on - * the home rank to avoid an unnecessary peer store. */ - uint32_t home_count = 0u; - uint32_t peer_count = 0u; - #pragma unroll - for (uint32_t i = 0; i < 6u; i++) { - const int32_t expert = selected[i]; - if (expert >= 0 && (uint32_t)expert < expert_split) { - home_count++; - } else if (expert >= 0 && - (uint32_t)expert < 2u * expert_split) { - peer_count++; - } - } - const bool assigned = home_rank - ? home_count <= peer_count : peer_count < home_count; - if (!assigned) return; - } - const unsigned char *gate_row = gate_w + row * blocks * 34u; - const unsigned char *up_row = up_w + row * blocks * 34u; - float gate = 0.0f; - float up = 0.0f; - for (uint64_t b = lane; b < blocks; b += 32u) { - const uint64_t i0 = b * 32u; - const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; - const int8_t *xqb = xq + b * 32u; - const float xs = xscale[b]; - const unsigned char *gb = gate_row + b * 34u; - const unsigned char *ub = up_row + b * 34u; - gate += __half2float(*(const __half *)gb) * xs * - (float)dot_i8_block((const int8_t *)(gb + 2u), xqb, bn, - use_dp4a); - up += __half2float(*(const __half *)ub) * xs * - (float)dot_i8_block((const int8_t *)(ub + 2u), xqb, bn, - use_dp4a); - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - gate = fminf(gate, clamp); - up = fminf(fmaxf(up, -clamp), clamp); - } - const float silu = gate / (1.0f + expf(-gate)); - mid[row] = silu * up * 1.0f; - } -} - -__global__ static void matmul_q8_0_pair_preq_batch_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 n_tok, - uint64_t blocks, - int use_dp4a) { - const uint64_t row = (uint64_t)blockIdx.x; - const uint64_t tok = (uint64_t)blockIdx.y; - if (tok >= n_tok) return; - const int has0 = row < out0_dim; - const int has1 = row < out1_dim; - if (!has0 && !has1) return; - - const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL; - const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL; - const int8_t *xqr = xq + tok * blocks * 32u; - const float *xsr = xscale + tok * blocks; - float acc0 = 0.0f; - float acc1 = 0.0f; - - for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { - const uint64_t i0 = b * 32u; - const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; - const int8_t *xqb = xqr + b * 32u; - const float xs = xsr[b]; - if (has0) { - const __half *scale_h = (const __half *)(wr0 + b * 34u); - const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u); - const int dot = dot_i8_block(qs, xqb, bn, use_dp4a); - acc0 += __half2float(*scale_h) * xs * (float)dot; - } - if (has1) { - const __half *scale_h = (const __half *)(wr1 + b * 34u); - const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u); - const int dot = dot_i8_block(qs, xqb, bn, use_dp4a); - acc1 += __half2float(*scale_h) * xs * (float)dot; - } - } - - __shared__ float partial0[256]; - __shared__ float partial1[256]; - partial0[threadIdx.x] = acc0; - partial1[threadIdx.x] = acc1; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - partial0[threadIdx.x] += partial0[threadIdx.x + stride]; - partial1[threadIdx.x] += partial1[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - if (has0) out0[tok * out0_dim + row] = partial0[0]; - if (has1) out1[tok * out1_dim + row] = partial1[0]; - } -} - -__global__ static void matmul_q8_0_pair_preq_batch_tok2_exact_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 n_tok, - uint64_t blocks, - int use_dp4a) { - const uint64_t row = (uint64_t)blockIdx.x; - const uint64_t tok0 = (uint64_t)blockIdx.y * 2u; - if (tok0 >= n_tok) return; - const int has0 = row < out0_dim; - const int has1 = row < out1_dim; - if (!has0 && !has1) return; - const int valid1 = tok0 + 1u < n_tok; - - const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL; - const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL; - const int8_t *xqr0 = xq + tok0 * blocks * 32u; - const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0; - const float *xsr0 = xscale + tok0 * blocks; - const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; - float acc00 = 0.0f; - float acc01 = 0.0f; - float acc10 = 0.0f; - float acc11 = 0.0f; - - for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { - const uint64_t i0 = b * 32u; - const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; - const int8_t *xqb0 = xqr0 + b * 32u; - const int8_t *xqb1 = xqr1 + b * 32u; - const float xs0 = xsr0[b]; - const float xs1 = valid1 ? xsr1[b] : 0.0f; - if (has0) { - const __half *scale_h = (const __half *)(wr0 + b * 34u); - const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u); - const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - int dot1 = 0; - if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - const float ws = __half2float(*scale_h); - acc00 += ws * xs0 * (float)dot0; - if (valid1) acc01 += ws * xs1 * (float)dot1; - } - if (has1) { - const __half *scale_h = (const __half *)(wr1 + b * 34u); - const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u); - const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - int dot1 = 0; - if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - const float ws = __half2float(*scale_h); - acc10 += ws * xs0 * (float)dot0; - if (valid1) acc11 += ws * xs1 * (float)dot1; - } - } - - __shared__ float partial00[256]; - __shared__ float partial01[256]; - __shared__ float partial10[256]; - __shared__ float partial11[256]; - partial00[threadIdx.x] = acc00; - partial01[threadIdx.x] = acc01; - partial10[threadIdx.x] = acc10; - partial11[threadIdx.x] = acc11; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - partial00[threadIdx.x] += partial00[threadIdx.x + stride]; - partial01[threadIdx.x] += partial01[threadIdx.x + stride]; - partial10[threadIdx.x] += partial10[threadIdx.x + stride]; - partial11[threadIdx.x] += partial11[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - if (has0) { - out0[tok0 * out0_dim + row] = partial00[0]; - if (valid1) out0[(tok0 + 1u) * out0_dim + row] = partial01[0]; - } - if (has1) { - out1[tok0 * out1_dim + row] = partial10[0]; - if (valid1) out1[(tok0 + 1u) * out1_dim + row] = partial11[0]; - } - } -} - -__device__ static float moe_owned_packed_combine_row( - const float *home_slots, - const float *peer_packed, - const int32_t *selected, - uint32_t row, - uint32_t out_dim, - uint32_t expert_split); - -__global__ static void matmul_q8_0_hc_expand_preq_warp8_kernel( - float *out_hc, - float *block_out, - const float *block_add, - const float *block_add2, - const float *owned_home_slots, - const float *owned_peer_packed, - const int32_t *owned_selected, - 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 has_add2, - int has_owned_slots, - uint32_t owned_expert_split, - 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_owned_slots) { - const float routed = moe_owned_packed_combine_row( - owned_home_slots, - owned_peer_packed, - owned_selected, - d, - (uint32_t)out_dim, - owned_expert_split); - block_v = __fadd_rn(block_v, routed); - } else if (has_add) { - float add_v = block_add[d]; - if (has_add2) add_v += block_add2[d]; - block_v += add_v; - } - 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; - } - } -} - -__global__ static void matmul_q8_0_kslice_hc_expand_add_preq_warp8_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 slice_dim, - uint64_t out_dim, - uint64_t full_blocks, - uint64_t block_start, - uint64_t slice_blocks, - uint32_t n_embd, - uint32_t n_hc, - 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 * full_blocks * 34u + block_start * 34u; - float acc = 0.0f; - for (uint64_t b = lane; b < slice_blocks; b += 32u) { - const uint64_t i0 = b * 32u; - const uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u; - const __half *scale_h = (const __half *)(wr + b * 34u); - const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); - const int8_t *xqb = xq + b * 32u; - const 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; - const float block_v = acc + 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; - } - } -} - -__global__ static void matmul_q8_0_preq_batch_warp8_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - 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; - if (row >= out_dim || tok >= n_tok) return; - - const unsigned char *wr = w + row * blocks * 34; - const int8_t *xqr = xq + tok * blocks * 32; - const float *xsr = xscale + tok * blocks; - 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 = 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) out[tok * out_dim + row] = acc; -} - -__global__ static void matmul_q8_0_preq_batch_warp8_tok2_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks, - 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 * 34u; - const int8_t *xqr0 = xq; - const int8_t *xqr1 = xq + blocks * 32u; - const float *xsr0 = xscale; - const float *xsr1 = xscale + blocks; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint64_t b = lane; b < blocks; b += 32u) { - const uint64_t i0 = b * 32u; - const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; - const __half *scale_h = (const __half *)(wr + b * 34u); - const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); - const int8_t *xqb0 = xqr0 + b * 32u; - const int8_t *xqb1 = xqr1 + b * 32u; - int dot0 = 0; - int dot1 = 0; - if (use_dp4a && bn == 32u) { -#pragma unroll - for (uint32_t i = 0; i < 32u; i += 4u) { - const int32_t w4 = load_i8x4_i32_unaligned(qs + i); - dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); - dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); - } - } else { - dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - } - const float ws = __half2float(*scale_h); - acc0 += ws * xsr0[b] * (float)dot0; - acc1 += ws * xsr1[b] * (float)dot1; - } - acc0 = warp_sum_f32(acc0); - acc1 = warp_sum_f32(acc1); - if (lane == 0) { - out[row] = acc0; - out[out_dim + row] = acc1; - } -} - -__global__ static void matmul_q8_0_preq_batch_warp8_tok4_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - int use_dp4a) { - const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); - const uint64_t tok0 = (uint64_t)blockIdx.y * 4u; - const uint32_t lane = threadIdx.x & 31u; - if (row >= out_dim || tok0 >= n_tok) return; - - const unsigned char *wr = w + row * blocks * 34; - const int8_t *xqr0 = xq + tok0 * blocks * 32; - const int8_t *xqr1 = xqr0 + blocks * 32; - const int8_t *xqr2 = xqr1 + blocks * 32; - const int8_t *xqr3 = xqr2 + blocks * 32; - const float *xsr0 = xscale + tok0 * blocks; - const float *xsr1 = xsr0 + blocks; - const float *xsr2 = xsr1 + blocks; - const float *xsr3 = xsr2 + blocks; - const int valid1 = tok0 + 1u < n_tok; - const int valid2 = tok0 + 2u < n_tok; - const int valid3 = tok0 + 3u < n_tok; - - float acc0 = 0.0f; - float acc1 = 0.0f; - float acc2 = 0.0f; - float acc3 = 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 *xqb0 = xqr0 + b * 32; - const int8_t *xqb1 = xqr1 + b * 32; - const int8_t *xqb2 = xqr2 + b * 32; - const int8_t *xqb3 = xqr3 + b * 32; - int dot0 = 0; - int dot1 = 0; - int dot2 = 0; - int dot3 = 0; - if (use_dp4a && bn == 32u) { -#pragma unroll - for (uint32_t i = 0; i < 32u; i += 4u) { - const int32_t w4 = load_i8x4_i32_unaligned(qs + i); - dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); - if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); - if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2); - if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3); - } - } else { - dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a); - if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a); - } - const float ws = __half2float(*scale_h); - acc0 += ws * xsr0[b] * (float)dot0; - if (valid1) acc1 += ws * xsr1[b] * (float)dot1; - if (valid2) acc2 += ws * xsr2[b] * (float)dot2; - if (valid3) acc3 += ws * xsr3[b] * (float)dot3; - } - acc0 = warp_sum_f32(acc0); - acc1 = warp_sum_f32(acc1); - acc2 = warp_sum_f32(acc2); - acc3 = warp_sum_f32(acc3); - if (lane == 0) { - out[tok0 * out_dim + row] = acc0; - if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1; - if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2; - if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3; - } -} - -__global__ static void matmul_q8_0_preq_batch_warp8_tok8_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - int use_dp4a) { - const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); - const uint64_t tok0 = (uint64_t)blockIdx.y * 8u; - const uint32_t lane = threadIdx.x & 31u; - if (row >= out_dim || tok0 >= n_tok) return; - - const unsigned char *wr = w + row * blocks * 34; - const uint64_t xq_stride = blocks * 32u; - const int8_t *xqr0 = xq + tok0 * xq_stride; - const int valid1 = tok0 + 1u < n_tok; - const int valid2 = tok0 + 2u < n_tok; - const int valid3 = tok0 + 3u < n_tok; - const int valid4 = tok0 + 4u < n_tok; - const int valid5 = tok0 + 5u < n_tok; - const int valid6 = tok0 + 6u < n_tok; - const int valid7 = tok0 + 7u < n_tok; - const int8_t *xqr1 = valid1 ? xqr0 + xq_stride : xqr0; - const int8_t *xqr2 = valid2 ? xqr1 + xq_stride : xqr0; - const int8_t *xqr3 = valid3 ? xqr2 + xq_stride : xqr0; - const int8_t *xqr4 = valid4 ? xqr3 + xq_stride : xqr0; - const int8_t *xqr5 = valid5 ? xqr4 + xq_stride : xqr0; - const int8_t *xqr6 = valid6 ? xqr5 + xq_stride : xqr0; - const int8_t *xqr7 = valid7 ? xqr6 + xq_stride : xqr0; - const float *xsr0 = xscale + tok0 * blocks; - const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; - const float *xsr2 = valid2 ? xsr1 + blocks : xsr0; - const float *xsr3 = valid3 ? xsr2 + blocks : xsr0; - const float *xsr4 = valid4 ? xsr3 + blocks : xsr0; - const float *xsr5 = valid5 ? xsr4 + blocks : xsr0; - const float *xsr6 = valid6 ? xsr5 + blocks : xsr0; - const float *xsr7 = valid7 ? xsr6 + blocks : xsr0; - - float acc0 = 0.0f; - float acc1 = 0.0f; - float acc2 = 0.0f; - float acc3 = 0.0f; - float acc4 = 0.0f; - float acc5 = 0.0f; - float acc6 = 0.0f; - float acc7 = 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 *xqb0 = xqr0 + b * 32; - const int8_t *xqb1 = xqr1 + b * 32; - const int8_t *xqb2 = xqr2 + b * 32; - const int8_t *xqb3 = xqr3 + b * 32; - const int8_t *xqb4 = xqr4 + b * 32; - const int8_t *xqb5 = xqr5 + b * 32; - const int8_t *xqb6 = xqr6 + b * 32; - const int8_t *xqb7 = xqr7 + b * 32; - int dot0 = 0; - int dot1 = 0; - int dot2 = 0; - int dot3 = 0; - int dot4 = 0; - int dot5 = 0; - int dot6 = 0; - int dot7 = 0; - if (use_dp4a && bn == 32u) { -#pragma unroll - for (uint32_t i = 0; i < 32u; i += 4u) { - const int32_t w4 = load_i8x4_i32_unaligned(qs + i); - dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); - if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); - if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2); - if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3); - if (valid4) dot4 = __dp4a(w4, load_i8x4_i32_aligned(xqb4 + i), dot4); - if (valid5) dot5 = __dp4a(w4, load_i8x4_i32_aligned(xqb5 + i), dot5); - if (valid6) dot6 = __dp4a(w4, load_i8x4_i32_aligned(xqb6 + i), dot6); - if (valid7) dot7 = __dp4a(w4, load_i8x4_i32_aligned(xqb7 + i), dot7); - } - } else { - dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a); - if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a); - if (valid4) dot4 = dot_i8_block(qs, xqb4, bn, use_dp4a); - if (valid5) dot5 = dot_i8_block(qs, xqb5, bn, use_dp4a); - if (valid6) dot6 = dot_i8_block(qs, xqb6, bn, use_dp4a); - if (valid7) dot7 = dot_i8_block(qs, xqb7, bn, use_dp4a); - } - const float ws = __half2float(*scale_h); - acc0 += ws * xsr0[b] * (float)dot0; - if (valid1) acc1 += ws * xsr1[b] * (float)dot1; - if (valid2) acc2 += ws * xsr2[b] * (float)dot2; - if (valid3) acc3 += ws * xsr3[b] * (float)dot3; - if (valid4) acc4 += ws * xsr4[b] * (float)dot4; - if (valid5) acc5 += ws * xsr5[b] * (float)dot5; - if (valid6) acc6 += ws * xsr6[b] * (float)dot6; - if (valid7) acc7 += ws * xsr7[b] * (float)dot7; - } - acc0 = warp_sum_f32(acc0); - acc1 = warp_sum_f32(acc1); - acc2 = warp_sum_f32(acc2); - acc3 = warp_sum_f32(acc3); - acc4 = warp_sum_f32(acc4); - acc5 = warp_sum_f32(acc5); - acc6 = warp_sum_f32(acc6); - acc7 = warp_sum_f32(acc7); - if (lane == 0) { - out[tok0 * out_dim + row] = acc0; - if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1; - if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2; - if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3; - if (valid4) out[(tok0 + 4u) * out_dim + row] = acc4; - if (valid5) out[(tok0 + 5u) * out_dim + row] = acc5; - if (valid6) out[(tok0 + 6u) * out_dim + row] = acc6; - if (valid7) out[(tok0 + 7u) * out_dim + row] = acc7; - } -} - -__global__ static void matmul_q8_0_preq_batch_tok2_exact_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - int use_dp4a) { - const uint64_t row = (uint64_t)blockIdx.x; - const uint64_t tok0 = (uint64_t)blockIdx.y * 2u; - if (row >= out_dim || tok0 >= n_tok) return; - const int valid1 = tok0 + 1u < n_tok; - const unsigned char *wr = w + row * blocks * 34u; - const int8_t *xqr0 = xq + tok0 * blocks * 32u; - const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0; - const float *xsr0 = xscale + tok0 * blocks; - const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; - float acc0 = 0.0f; - float acc1 = 0.0f; - - for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { - const uint64_t i0 = b * 32u; - const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; - const __half *scale_h = (const __half *)(wr + b * 34u); - const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); - const int8_t *xqb0 = xqr0 + b * 32u; - const int8_t *xqb1 = xqr1 + b * 32u; - const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); - int dot1 = 0; - if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); - const float ws = __half2float(*scale_h); - acc0 += ws * xsr0[b] * (float)dot0; - if (valid1) acc1 += ws * xsr1[b] * (float)dot1; - } - - __shared__ float partial0[256]; - __shared__ float partial1[256]; - partial0[threadIdx.x] = acc0; - partial1[threadIdx.x] = acc1; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - partial0[threadIdx.x] += partial0[threadIdx.x + stride]; - partial1[threadIdx.x] += partial1[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - out[tok0 * out_dim + row] = partial0[0]; - if (valid1) out[(tok0 + 1u) * out_dim + row] = partial1[0]; - } -} - - -/* ---- INT8 tensor-core exact Q8_0 batch matmul -------------------------- - * Bit-identical replacement for the exact tok2/warp8-family batched Q8_0 - * kernels. Each output element's reduction is the reference's strided - * halving tree over T slots (T = reduction width: 32 for the warp kernels, - * cuda_q8_exact_threads(blocks) for the exact kernels; slots >= blocks hold - * +0.0f). The kernel decomposes that tree as: 32 streams at stride T/32 - * whose 32 terms per outer step j combine via an adjacent-pairwise static - * register stack taken in bit-reversed stream order (== the top five strided - * tree levels), plus per-(j&3) sequential accumulators and a fixed tail for - * the remaining levels. Fuzz-verified bitwise against both reference - * kernels across shapes, including blocks < T and ragged out_dim/n_tok. - * Rollback: DS4_CUDA_NO_Q8_MMA=1. */ -__device__ __forceinline__ static uint32_t ldu32_unaligned(const uint8_t *p) { - const uintptr_t addr = (uintptr_t)p; - const uint32_t *base = (const uint32_t *)(addr & ~(uintptr_t)3); - const uint32_t lo = base[0]; - const uint32_t hi = base[1]; - return __funnelshift_r(lo, hi, (uint32_t)(addr & 3u) * 8u); -} - -__device__ __forceinline__ static void mma_m16n8k32_s8( - int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3, - uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, - uint32_t b0, uint32_t b1) { -#if __CUDA_ARCH__ >= 800 - asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" - : "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3) - : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1)); -#else - (void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3; -#endif -} - -__device__ __forceinline__ static uint32_t bitrev5(uint32_t i) { - return ((i & 1u) << 4) | ((i & 2u) << 2) | (i & 4u) | ((i & 8u) >> 2) | ((i & 16u) >> 4); -} - -template -__global__ static void matmul_q8_0_mma_exact_kernel( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - uint64_t a_stride_blocks, /* activation row stride in blocks (>= blocks) */ - uint64_t out_stride) { /* output token stride in floats (>= out_dim) */ - extern __shared__ unsigned char q8mma_sh[]; - __half *sh_ws = (__half *)q8mma_sh; /* 64 rows x blocks */ - float *sh_xs = (float *)(q8mma_sh + 64u * blocks * 2u); /* 16 toks x blocks */ - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint64_t row_base = (uint64_t)blockIdx.x * 64u; - const uint64_t tok_base = (uint64_t)blockIdx.y * 16u; - - /* stage weight scales (64 rows) and activation scales (16 tokens) */ - for (uint32_t idx = threadIdx.x; idx < 64u * (uint32_t)blocks; idx += blockDim.x) { - const uint32_t rl = idx / (uint32_t)blocks; - const uint32_t b = idx - rl * (uint32_t)blocks; - uint64_t row = row_base + rl; - if (row >= out_dim) row = out_dim - 1u; - sh_ws[idx] = *(const __half *)(w + row * blocks * 34u + (uint64_t)b * 34u); - } - for (uint32_t idx = threadIdx.x; idx < 16u * (uint32_t)blocks; idx += blockDim.x) { - const uint32_t tl = idx / (uint32_t)blocks; - const uint32_t b = idx - tl * (uint32_t)blocks; - const uint64_t tok = tok_base + tl; - sh_xs[idx] = tok < n_tok ? xscale[tok * a_stride_blocks + b] : 0.0f; - } - __syncthreads(); - - const uint64_t row0 = row_base + (uint64_t)warp * 8u; - /* thread's C elements: rows n0,n0+1; tokens mt0, mt0+8 */ - const uint32_t n0 = (lane & 3u) * 2u; - const uint32_t mt0 = lane >> 2u; - const uint64_t tokA = tok_base + mt0; - const uint64_t tokB = tok_base + mt0 + 8u; - /* A source rows for loads (fragment layout): rows lane>>2 and (lane>>2)+8 */ - const uint64_t a_tok_lo = tok_base + (lane >> 2u); - const uint64_t a_tok_hi = a_tok_lo + 8u; - const int8_t *aq_lo = xq + (a_tok_lo < n_tok ? a_tok_lo : 0u) * a_stride_blocks * 32u; - const int8_t *aq_hi = xq + (a_tok_hi < n_tok ? a_tok_hi : 0u) * a_stride_blocks * 32u; - const bool a_lo_ok = a_tok_lo < n_tok; - const bool a_hi_ok = a_tok_hi < n_tok; - /* B source row for loads: row lane>>2 within the warp tile */ - uint64_t b_row = row0 + (lane >> 2u); - if (b_row >= out_dim) b_row = out_dim - 1u; - const unsigned char *b_wr = w + b_row * blocks * 34u; - - /* per-element (4) x per-(j&3) accumulators */ - float acc00 = 0.0f, acc01 = 0.0f, acc02 = 0.0f, acc03 = 0.0f; - float acc10 = 0.0f, acc11 = 0.0f, acc12 = 0.0f, acc13 = 0.0f; - float acc20 = 0.0f, acc21 = 0.0f, acc22 = 0.0f, acc23 = 0.0f; - float acc30 = 0.0f, acc31 = 0.0f, acc32 = 0.0f, acc33 = 0.0f; - - const uint32_t stride = T / 32u; - const uint32_t rl_ws0 = warp * 8u + n0; /* local row for ws of element cols */ - const uint32_t tl_xsA = mt0; /* local token rows for xs */ - const uint32_t tl_xsB = mt0 + 8u; - - for (uint32_t j = 0; j < stride; j++) { - /* adjacent-pairwise static stack over 32 terms in bitrev5 m order */ - float s0e0 = 0, s1e0 = 0, s2e0 = 0, s3e0 = 0, s4e0 = 0; - float s0e1 = 0, s1e1 = 0, s2e1 = 0, s3e1 = 0, s4e1 = 0; - float s0e2 = 0, s1e2 = 0, s2e2 = 0, s3e2 = 0, s4e2 = 0; - float s0e3 = 0, s1e3 = 0, s2e3 = 0, s3e3 = 0, s4e3 = 0; -#pragma unroll - for (uint32_t i = 0; i < 32u; i++) { - const uint32_t m = bitrev5(i); - const uint32_t s = j + m * stride; /* slot index */ - float t0 = 0.0f, t1 = 0.0f, t2 = 0.0f, t3 = 0.0f; - /* slot s sums blocks {s + k*T} sequentially (multi-term when - * blocks > T, exactly like the per-lane strided walk). */ - for (uint32_t b = s; b < blocks; b += T) { - const uint32_t koff = (lane & 3u) * 4u; - const int8_t *ablk_lo = aq_lo + b * 32u; - const int8_t *ablk_hi = aq_hi + b * 32u; - const uint32_t a0 = a_lo_ok ? *(const uint32_t *)(ablk_lo + koff) : 0u; - const uint32_t a1 = a_hi_ok ? *(const uint32_t *)(ablk_hi + koff) : 0u; - const uint32_t a2 = a_lo_ok ? *(const uint32_t *)(ablk_lo + 16u + koff) : 0u; - const uint32_t a3 = a_hi_ok ? *(const uint32_t *)(ablk_hi + 16u + koff) : 0u; - const uint8_t *bq = (const uint8_t *)(b_wr + (uint64_t)b * 34u + 2u); - const uint32_t b0 = ldu32_unaligned(bq + koff); - const uint32_t b1 = ldu32_unaligned(bq + 16u + koff); - int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; - mma_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); - /* term = ws * xs * dot, same expression as reference */ - const float ws0 = __half2float(sh_ws[rl_ws0 * (uint32_t)blocks + b]); - const float ws1 = __half2float(sh_ws[(rl_ws0 + 1u) * (uint32_t)blocks + b]); - const float xsA = sh_xs[tl_xsA * (uint32_t)blocks + b]; - const float xsB = sh_xs[tl_xsB * (uint32_t)blocks + b]; - t0 += ws0 * xsA * (float)c0; - t1 += ws1 * xsA * (float)c1; - t2 += ws0 * xsB * (float)c2; - t3 += ws1 * xsB * (float)c3; - } - /* static adjacent stack push (compile-time resolved) */ - if ((i & 1u) == 0u) { s0e0 = t0; s0e1 = t1; s0e2 = t2; s0e3 = t3; } - else { - t0 = s0e0 + t0; t1 = s0e1 + t1; t2 = s0e2 + t2; t3 = s0e3 + t3; - if ((i & 2u) == 0u) { s1e0 = t0; s1e1 = t1; s1e2 = t2; s1e3 = t3; } - else { - t0 = s1e0 + t0; t1 = s1e1 + t1; t2 = s1e2 + t2; t3 = s1e3 + t3; - if ((i & 4u) == 0u) { s2e0 = t0; s2e1 = t1; s2e2 = t2; s2e3 = t3; } - else { - t0 = s2e0 + t0; t1 = s2e1 + t1; t2 = s2e2 + t2; t3 = s2e3 + t3; - if ((i & 8u) == 0u) { s3e0 = t0; s3e1 = t1; s3e2 = t2; s3e3 = t3; } - else { - t0 = s3e0 + t0; t1 = s3e1 + t1; t2 = s3e2 + t2; t3 = s3e3 + t3; - if ((i & 16u) == 0u) { s4e0 = t0; s4e1 = t1; s4e2 = t2; s4e3 = t3; } - else { - t0 = s4e0 + t0; t1 = s4e1 + t1; t2 = s4e2 + t2; t3 = s4e3 + t3; - /* i == 31: t is the finished x_j */ - switch (j & 3u) { - case 0u: acc00 += t0; acc10 += t1; acc20 += t2; acc30 += t3; break; - case 1u: acc01 += t0; acc11 += t1; acc21 += t2; acc31 += t3; break; - case 2u: acc02 += t0; acc12 += t1; acc22 += t2; acc32 += t3; break; - default: acc03 += t0; acc13 += t1; acc23 += t2; acc33 += t3; break; - } - } - } - } - } - } - } - } - /* tail combine per T */ - float r0, r1, r2, r3; - if (T == 32u) { - r0 = acc00; r1 = acc10; r2 = acc20; r3 = acc30; - } else if (T == 64u) { - r0 = acc00 + acc01; r1 = acc10 + acc11; r2 = acc20 + acc21; r3 = acc30 + acc31; - } else { - r0 = (acc00 + acc02) + (acc01 + acc03); - r1 = (acc10 + acc12) + (acc11 + acc13); - r2 = (acc20 + acc22) + (acc21 + acc23); - r3 = (acc30 + acc32) + (acc31 + acc33); - } - /* writes */ - const uint64_t rowa = row0 + n0; - const uint64_t rowb = rowa + 1u; - if (tokA < n_tok) { - if (rowa < out_dim) out[tokA * out_stride + rowa] = r0; - if (rowb < out_dim) out[tokA * out_stride + rowb] = r1; - } - if (tokB < n_tok) { - if (rowa < out_dim) out[tokB * out_stride + rowa] = r2; - if (rowb < out_dim) out[tokB * out_stride + rowb] = r3; - } -} - -static int cuda_q4_mma_ok(void); -static int cuda_q8_mma_attr_ready[DS4_MAX_GPUS][4]; -static int cuda_q8_mma_try_launch( - float *out, - const unsigned char *w, - const int8_t *xq, - const float *xscale, - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t blocks, - uint64_t a_stride_blocks, - uint64_t out_stride, - uint32_t T) { - static int disabled = -1; - if (disabled < 0) disabled = getenv("DS4_CUDA_NO_Q8_MMA") != NULL ? 1 : 0; - if (disabled || !cuda_q4_mma_ok()) return 0; - if ((in_dim & 31u) != 0u || blocks > 256u || n_tok < 8u) return 0; - if (((uintptr_t)w & 1u) || ((uintptr_t)xq & 3u) || ((uintptr_t)xscale & 3u)) return 0; - const size_t shmem = (size_t)(64u * blocks * 2u + 16u * blocks * 4u); - int dev = 0; - cudaGetDevice(&dev); - if (dev < 0 || dev >= DS4_MAX_GPUS) return 0; - const int ti = T == 32u ? 0 : (T == 64u ? 1 : (T == 128u ? 2 : 3)); - dim3 grid(((unsigned)out_dim + 63u) / 64u, ((unsigned)n_tok + 15u) / 16u, 1); -#define DS4_Q8_MMA_LAUNCH(TT) \ - do { \ - if (!cuda_q8_mma_attr_ready[dev][ti]) { \ - cudaFuncAttributes fn_attr; \ - if (cudaFuncGetAttributes(&fn_attr, matmul_q8_0_mma_exact_kernel) != cudaSuccess || \ - fn_attr.binaryVersion < 80) { \ - disabled = 1; \ - return 0; \ - } \ - if (cudaFuncSetAttribute(matmul_q8_0_mma_exact_kernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - (int)(64u * 256u * 2u + 16u * 256u * 4u)) != cudaSuccess) { \ - disabled = 1; \ - return 0; \ - } \ - cuda_q8_mma_attr_ready[dev][ti] = 1; \ - } \ - matmul_q8_0_mma_exact_kernel<<>>( \ - out, w, xq, xscale, in_dim, out_dim, n_tok, blocks, \ - a_stride_blocks, out_stride); \ - } while (0) - if (T == 32u) DS4_Q8_MMA_LAUNCH(32u); - else if (T == 64u) DS4_Q8_MMA_LAUNCH(64u); - else if (T == 128u) DS4_Q8_MMA_LAUNCH(128u); - else DS4_Q8_MMA_LAUNCH(256u); -#undef DS4_Q8_MMA_LAUNCH - return cuda_ok(cudaGetLastError(), "matmul_q8_0 mma launch") ? 1 : -1; -} - - -__global__ static void dequant_q8_0_to_f16_kernel( - __half *out, - const unsigned char *w, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = in_dim * out_dim; - if (gid >= n) return; - uint64_t row = gid / in_dim; - uint64_t i = gid - row * in_dim; - uint64_t b = i / 32; - uint64_t j = i - b * 32; - const unsigned char *blk = w + (row * blocks + b) * 34; - const __half scale = *(const __half *)blk; - const int8_t q = *(const int8_t *)(blk + 2 + j); - out[gid] = __hmul(scale, __float2half((float)q)); -} - -__global__ static void dequant_q8_0_to_f32_kernel( - float *out, - const unsigned char *w, - uint64_t in_dim, - uint64_t out_dim, - uint64_t blocks) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = in_dim * out_dim; - if (gid >= n) return; - uint64_t row = gid / in_dim; - uint64_t i = gid - row * in_dim; - uint64_t b = i / 32; - uint64_t j = i - b * 32; - const unsigned char *blk = w + (row * blocks + b) * 34; - const float scale = __half2float(*(const __half *)blk); - const int8_t q = *(const int8_t *)(blk + 2 + j); - out[gid] = scale * (float)q; -} - -__global__ static void grouped_q8_0_a_preq_warp8_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__ static void grouped_q8_0_a_preq_warp8_tok2_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 uint32_t tid_in_tok = threadIdx.x & 255u; - const uint64_t row = (uint64_t)blockIdx.x * 8u + (tid_in_tok >> 5u); - const uint64_t tok = (uint64_t)blockIdx.y * 2u + (threadIdx.x >> 8u); - const uint32_t lane = threadIdx.x & 31u; - const uint64_t low_dim = (uint64_t)n_groups * rank; - - float acc = 0.0f; - if (row < low_dim && tok < n_tokens) { - 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 * 34u; - const uint64_t xrow = tok * (uint64_t)n_groups + group; - const int8_t *xqr = xq + xrow * blocks * 32u; - const float *xsr = xscale + xrow * blocks; - - for (uint64_t b = lane; b < blocks; b += 32u) { - const uint64_t i0 = b * 32u; - const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u; - const __half *scale_h = (const __half *)(wr + b * 34u); - const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); - const int8_t *xqb = xqr + b * 32u; - const 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 && row < low_dim && tok < n_tokens) { - low[tok * low_dim + row] = acc; - } -} - -__global__ static void rms_norm_plain_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { - uint32_t row = blockIdx.x; - if (row >= rows) return; - const float *xr = x + (uint64_t)row * n; - float *orow = out + (uint64_t)row * n; - float sum = 0.0f; - for (uint32_t i = threadIdx.x; i < n; 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)n + eps); - for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { - orow[i] = xr[i] * scale; - } -} - -/* Latency-optimized RMS norm for the common n==4096 decode shape: one global - * read pass with register-batched loads, same per-thread accumulation order - * and shared-memory tree as rms_norm_plain_kernel (bit-identical, fuzz - * checked). */ -__global__ static void rms_norm_plain_fast4096_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { - uint32_t row = blockIdx.x; - if (row >= rows) return; - const float *xr = x + (uint64_t)row * n; - float *orow = out + (uint64_t)row * n; - float v[16]; -#pragma unroll - for (uint32_t j = 0; j < 16u; j++) v[j] = xr[threadIdx.x + j * 256u]; - float sum = 0.0f; -#pragma unroll - for (uint32_t j = 0; j < 16u; j++) sum += v[j] * v[j]; - __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)n + eps); -#pragma unroll - for (uint32_t j = 0; j < 16u; j++) orow[threadIdx.x + j * 256u] = v[j] * scale; -} - -/* Batched-load RMS norm for larger rows (n multiple of 2048, e.g. the 16384 - * HC-concatenated decode rows). Two passes like the reference kernel, but - * eight independent loads are issued per accumulation group; the per-thread - * accumulation order (ascending i with stride 256) is unchanged, so results - * are bit-identical. */ -__global__ static void rms_norm_plain_batch8_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { - uint32_t row = blockIdx.x; - if (row >= rows) return; - const float *xr = x + (uint64_t)row * n; - float *orow = out + (uint64_t)row * n; - float sum = 0.0f; -#pragma unroll 1 - for (uint32_t i = threadIdx.x; i < n; i += 2048u) { - const float v0 = xr[i]; - const float v1 = xr[i + 256u]; - const float v2 = xr[i + 512u]; - const float v3 = xr[i + 768u]; - const float v4 = xr[i + 1024u]; - const float v5 = xr[i + 1280u]; - const float v6 = xr[i + 1536u]; - const float v7 = xr[i + 1792u]; - sum += v0 * v0; - sum += v1 * v1; - sum += v2 * v2; - sum += v3 * v3; - sum += v4 * v4; - sum += v5 * v5; - sum += v6 * v6; - sum += v7 * v7; - } - __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)n + eps); -#pragma unroll 1 - for (uint32_t i = threadIdx.x; i < n; i += 2048u) { - const float v0 = xr[i]; - const float v1 = xr[i + 256u]; - const float v2 = xr[i + 512u]; - const float v3 = xr[i + 768u]; - const float v4 = xr[i + 1024u]; - const float v5 = xr[i + 1280u]; - const float v6 = xr[i + 1536u]; - const float v7 = xr[i + 1792u]; - orow[i] = v0 * scale; - orow[i + 256u] = v1 * scale; - orow[i + 512u] = v2 * scale; - orow[i + 768u] = v3 * scale; - orow[i + 1024u] = v4 * scale; - orow[i + 1280u] = v5 * scale; - orow[i + 1536u] = v6 * scale; - orow[i + 1792u] = v7 * scale; - } -} - -__global__ static void rms_norm_weight_kernel(float *out, const float *x, const float *w, uint32_t n, uint32_t rows, float eps) { - uint32_t row = blockIdx.x; - if (row >= rows) return; - const float *xr = x + (uint64_t)row * n; - float *orow = out + (uint64_t)row * n; - float sum = 0.0f; - for (uint32_t i = threadIdx.x; i < n; 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)n + eps); - for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { - orow[i] = xr[i] * scale * w[i]; - } -} - -__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; -} - -__device__ static float rope_yarn_ramp_dev(float low, float high, int i0); - -__global__ static void dsv4_qkv_rms_norm_rows_kv_rope_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, - uint32_t kv_n_head, - uint32_t kv_head_dim, - uint32_t n_rot, - uint32_t pos0, - 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) { - 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); - if (which == 0u) { - for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { - orow[i] = xr[i] * scale * w[i]; - } - return; - } - - const uint32_t n_nope = kv_head_dim - n_rot; - for (uint32_t h = 0; h < kv_n_head; h++) { - const uint32_t head_base = h * kv_head_dim; - for (uint32_t d = threadIdx.x; d < n_nope; d += blockDim.x) { - const uint32_t i = head_base + d; - orow[i] = xr[i] * scale * w[i]; - } - } - - 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); - } - const uint32_t pairs_per_head = n_rot / 2u; - const uint32_t total_pairs = kv_n_head * pairs_per_head; - for (uint32_t p = threadIdx.x; p < total_pairs; p += blockDim.x) { - const uint32_t h = p / pairs_per_head; - const uint32_t pair = p - h * pairs_per_head; - const uint32_t d = n_nope + pair * 2u; - const uint32_t i0 = h * kv_head_dim + d; - const uint32_t i = pair * 2u; - float theta_extrap = (float)(pos0 + row) * 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; - const float x0 = xr[i0] * scale * w[i0]; - const float x1 = xr[i0 + 1u] * scale * w[i0 + 1u]; - orow[i0] = x0 * c - x1 * s; - orow[i0 + 1u] = x0 * s + x1 * c; - } -} - -__global__ static void head_rms_norm_rope_tail_kernel( - float *x, - uint32_t n_tok, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos0, - 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; - } - - 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) * 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; - } -} - -__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)); -} - -__global__ static void rope_tail_kernel( - float *x, - uint32_t n_tok, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos0, - 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; - - 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 rope_tail_decode_rows_kernel( - float *x, - cuda_attention_decode_row_table rows, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - 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 gid = blockIdx.x * blockDim.x + threadIdx.x; - const uint32_t pairs = n_rows * n_head * (n_rot / 2u); - if (gid >= pairs) return; - const uint32_t pair = gid % (n_rot / 2u); - const uint32_t tmp = gid / (n_rot / 2u); - const uint32_t h = tmp % n_head; - const uint32_t row = tmp / n_head; - const uint32_t n_nope = head_dim - n_rot; - const uint32_t i = pair * 2u; - - float corr0 = 0.0f, corr1 = 0.0f; - if (ext_factor != 0.0f) { - const 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 - 1u), corr1); - } - - const float theta_extrap = (float)rows.row[row].pos * - powf(freq_base, -((float)i) / (float)n_rot); - const float theta_interp = freq_scale * theta_extrap; - float theta = theta_interp; - float mscale = attn_factor; - if (ext_factor != 0.0f) { - const 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); - } - const float c = cosf(theta) * mscale; - float s = sinf(theta) * mscale; - if (inverse) s = -s; - - float *tail = x + ((uint64_t)row * n_head + h) * head_dim + n_nope; - const float x0 = tail[i]; - const float x1 = tail[i + 1u]; - tail[i] = x0 * c - x1 * s; - tail[i + 1u] = x0 * s + x1 * c; -} - -__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 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 float dsv4_e2m1fn_dequant_dev(float x) { - float sign = x < 0.0f ? -1.0f : 1.0f; - float ax = fminf(fabsf(x), 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 sign * dsv4_e2m1fn_value_dev(best); -} - -__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]; -} - -__device__ static float rope_yarn_ramp_cpu_equiv_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 DS4_CUDA_UNUSED void rope_tail_one_dev(float *x, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint32_t n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) { - uint32_t n_nope = head_dim - n_rot; - float corr0 = 0.0f, corr1 = 0.0f; - if (ext_factor != 0.0f) { - float denom = 2.0f * logf(freq_base); - corr0 = fmaxf(0.0f, floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom)); - corr1 = fminf((float)(n_rot - 1), ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom)); - } - for (uint32_t i = 0; i < n_rot; i += 2) { - float theta_extrap = (float)pos * 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 mix = rope_yarn_ramp_cpu_equiv_dev(corr0, corr1, (int)i) * ext_factor; - theta = theta_interp * (1.0f - mix) + theta_extrap * mix; - mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - float c = cosf(theta) * mscale; - float s = sinf(theta) * mscale; - float x0 = x[n_nope + i]; - float x1 = x[n_nope + i + 1]; - x[n_nope + i] = x0 * c - x1 * s; - x[n_nope + i + 1] = x0 * s + x1 * c; - } -} - -__device__ static void fp8_kv_quantize_row( - float *xr, - uint32_t head_dim, - uint32_t n_rot, - float *scratch) { - uint32_t tid = threadIdx.x; - uint32_t n_nope = head_dim - n_rot; - 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 q = dsv4_e4m3fn_dequant_dev(fminf(448.0f, fmaxf(-448.0f, v / scale))) * scale; - xr[off + tid] = q; - } - __syncthreads(); - } -} - -__global__ static void fp8_kv_quantize_kernel( - float *x, - uint32_t n_tok, - uint32_t head_dim, - uint32_t n_rot) { - uint32_t row = blockIdx.x; - if (row >= n_tok) return; - __shared__ float scratch[64]; - fp8_kv_quantize_row( - x + (uint64_t)row * head_dim, head_dim, n_rot, scratch); -} - -__global__ static void fp8_kv_quantize_store_rows_kernel( - float *x, - cuda_attention_decode_row_table rows, - uint32_t n_rows, - uint32_t head_dim, - uint32_t n_rot) { - const uint32_t row = blockIdx.x; - if (row >= n_rows) return; - __shared__ float scratch[64]; - float *xr = x + (uint64_t)row * head_dim; - fp8_kv_quantize_row(xr, head_dim, n_rot, scratch); - - const ds4_gpu_attention_decode_row dsc = rows.row[row]; - float *raw = (float *)(uintptr_t)dsc.raw_kv; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - raw[(uint64_t)dsc.raw_start * head_dim + d] = - __half2float(__float2half(xr[d])); - } -} - -__global__ static void indexer_hadamard_fp4_kernel(float *x, uint32_t n_rows, uint32_t head_dim) { - 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))); - xr[tid] = dsv4_e2m1fn_dequant_dev(fminf(6.0f, fmaxf(-6.0f, v / scale))) * scale; -} - -__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) { - 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 = (pos0 + t) % raw_cap; - raw[(uint64_t)row * head_dim + d] = __half2float(__float2half(kv[(uint64_t)t * head_dim + d])); -} - -__global__ static void attention_prefill_raw_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - uint32_t n_tokens, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || h >= n_head) return; - uint32_t raw_count = t + 1 < window ? t + 1 : window; - uint32_t raw_start = t + 1 - raw_count; - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - __shared__ float scores[256]; - __shared__ float partial[128]; - __shared__ float max_s; - __shared__ float denom; - float scale = rsqrtf((float)head_dim); - float local_max = sinks[h]; - __syncthreads(); - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - const float *kv = raw_kv + (uint64_t)(raw_start + r) * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; - scores[r] = dot * scale; - local_max = fmaxf(local_max, scores[r]); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - if (threadIdx.x == 0) { - float den = expf(sinks[h] - max_s); - for (uint32_t r = 0; r < raw_count; r++) { - scores[r] = expf(scores[r] - max_s); - den += scores[r]; - } - denom = den; - } - __syncthreads(); - float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r]; - } - oh[d] = acc / denom; - } -} - -__global__ static void attention_prefill_mixed_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || h >= n_head) return; - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - uint32_t raw_start = (window != 0 && t + 1u > window) ? t + 1u - window : 0u; - uint32_t raw_count = t + 1u - raw_start; - uint32_t visible_comp = (t + 1u) / ratio; - if (visible_comp > n_comp) visible_comp = n_comp; - __shared__ float scores[512]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - float scale = rsqrtf((float)head_dim); - float local_max = sinks[h]; - uint32_t n_score = raw_count + visible_comp; - - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - const float *kvrow = raw_kv + (uint64_t)(raw_start + r) * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - scores[r] = dot * scale; - local_max = fmaxf(local_max, scores[r]); - } - for (uint32_t c = threadIdx.x; c < visible_comp; c += blockDim.x) { - float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - float s = -INFINITY; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)c * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale + add; - } - scores[raw_count + c] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r]; - for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; - oh[d] = acc / denom; - } -} - -__global__ static void attention_prefill_raw_softmax_kernel( - float *scores, - const float *sinks, - uint32_t n_tokens, - uint32_t window, - uint32_t n_keys) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens) return; - float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - float local_max = sinks[h]; - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { - bool valid = k <= t && (window == 0 || t - k < window); - float s = valid ? row[k] : -INFINITY; - row[k] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { - float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f; - row[k] = p; - den_local += p; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom; -} - -__global__ static void attention_prefill_mixed_softmax_kernel( - float *scores, - const float *sinks, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_keys) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || ratio == 0) return; - float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - float local_max = sinks[h]; - const uint32_t visible_comp = (t + 1u) / ratio; - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { - float s = -INFINITY; - if (k < n_tokens) { - if (k <= t && (window == 0 || t - k < window)) s = row[k]; - } else { - uint32_t c = k - n_tokens; - if (c < n_comp && c < visible_comp) { - float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - if (add > -1.0e20f) s = row[k] + add; - } - } - row[k] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { - float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f; - row[k] = p; - den_local += p; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom; -} - -__global__ static void attention_prefill_pack_mixed_kv_kernel( - float *dst, - const float *raw_kv, - const float *comp_kv, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t head_dim) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)(n_tokens + n_comp) * head_dim; - if (gid >= n) return; - uint32_t d = gid % head_dim; - uint32_t r = gid / head_dim; - dst[gid] = r < n_tokens ? raw_kv[(uint64_t)r * head_dim + d] - : comp_kv[(uint64_t)(r - n_tokens) * head_dim + d]; -} - -__global__ static void attention_prefill_unpack_heads_kernel( - float *heads, - const float *tmp, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * n_head * head_dim; - if (gid >= n) return; - uint32_t d = gid % head_dim; - uint64_t q = gid / head_dim; - uint32_t h = q % n_head; - uint32_t t = q / n_head; - heads[gid] = tmp[((uint64_t)h * n_tokens + t) * head_dim + d]; -} - -__global__ static void attention_pack_group_heads_f16_kernel( - __half *dst, - const float *heads, - uint32_t n_tokens, - uint32_t n_groups, - uint32_t group_dim) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_groups * n_tokens * group_dim; - if (gid >= n) return; - uint32_t d = gid % group_dim; - uint64_t q = gid / group_dim; - uint32_t t = q % n_tokens; - uint32_t g = q / n_tokens; - dst[gid] = __float2half(heads[((uint64_t)t * n_groups + g) * group_dim + d]); -} - -__global__ static void attention_unpack_group_low_kernel( - float *low, - const float *tmp, - uint32_t n_tokens, - uint32_t n_groups, - uint32_t rank) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_groups * n_tokens * rank; - if (gid >= n) return; - uint32_t r = gid % rank; - uint64_t q = gid / rank; - uint32_t t = q % n_tokens; - uint32_t g = q / n_tokens; - uint32_t low_dim = n_groups * rank; - low[(uint64_t)t * low_dim + (uint64_t)g * rank + r] = tmp[gid]; -} - -__global__ static void attention_decode_mixed_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t score_lanes_single) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || h >= n_head) return; - const bool single_all = (n_tokens == 1u && ratio == 0u); - uint32_t qpos = pos0 + t; - uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; - __shared__ uint32_t raw_rows[256]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; - const bool score_thread = threadIdx.x < score_threads; - float scale = rsqrtf((float)head_dim); - if (threadIdx.x == 0) { - raw_count = 0; - raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - if (score_thread) { - for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - } - __syncthreads(); - uint32_t n_score = raw_count + visible_comp; - float local_max = sinks[h]; - if (score_thread) { - if (visible_comp == 0 || (n_tokens == 1u && score_lanes_single == 0u)) { - for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { - const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - scores[r] = dot * scale; - local_max = fmaxf(local_max, scores[r]); - } - for (uint32_t c = threadIdx.x; c < visible_comp; c += score_threads) { - float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - float s = -INFINITY; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)c * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale + add; - } - scores[raw_count + c] = s; - local_max = fmaxf(local_max, s); - } - } else if (n_tokens == 1u && score_lanes_single == 4u) { - uint32_t qlane = threadIdx.x & 3u; - uint32_t qgroup = threadIdx.x >> 2u; - for (uint32_t row0 = 0; row0 < n_score; row0 += 64u) { - uint32_t row = row0 + qgroup; - if (row < n_score) { - float add = 0.0f; - const float *kvrow = NULL; - if (row < raw_count) { - kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim; - } else { - uint32_t c = row - raw_count; - add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim; - } - float s = -INFINITY; - if (kvrow) { - float dot = 0.0f; - for (uint32_t d = qlane; d < head_dim; d += 4u) dot += qh[d] * kvrow[d]; - const uint32_t mask = 0xfu << (threadIdx.x & 28u); - dot += __shfl_down_sync(mask, dot, 2, 4); - dot += __shfl_down_sync(mask, dot, 1, 4); - s = dot * scale + add; - } - if (qlane == 0) scores[row] = s; - } - } - __syncthreads(); - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - local_max = fmaxf(local_max, scores[i]); - } - } else { - uint32_t qlane = threadIdx.x & 7u; - uint32_t qgroup = threadIdx.x >> 3u; - for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { - uint32_t row = row0 + qgroup; - if (row < n_score) { - float add = 0.0f; - const float *kvrow = NULL; - if (row < raw_count) { - kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim; - } else { - uint32_t c = row - raw_count; - add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim; - } - float s = -INFINITY; - if (kvrow) { - float dot = 0.0f; - for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; - const uint32_t mask = 0xffu << (threadIdx.x & 24u); - for (uint32_t off = 4u; off > 0u; off >>= 1u) { - dot += __shfl_down_sync(mask, dot, off, 8); - } - s = dot * scale + add; - } - if (qlane == 0) scores[row] = s; - } - } - __syncthreads(); - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - local_max = fmaxf(local_max, scores[i]); - } - } - } - if (score_thread) partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - } - if (score_thread) partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; - if (head_dim == 512u && blockDim.x >= 512u) { - uint32_t d = threadIdx.x; - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc += kv[d] * s; - } - for (uint32_t c = 0; c < visible_comp; c++) { - float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc += kv[d] * s; - } - oh[d] = acc / denom; - } else if (head_dim == 512u && blockDim.x == 256u) { - uint32_t d0 = threadIdx.x; - uint32_t d1 = d0 + 256u; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - for (uint32_t c = 0; c < visible_comp; c++) { - float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - oh[d0] = acc0 / denom; - oh[d1] = acc1 / denom; - } else { - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; - for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; - oh[d] = acc / denom; - } - } -} - -__global__ static void attention_decode_score_split_scores_kernel( - float *score_out, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - const uint32_t h = blockIdx.y; - const uint32_t j = blockIdx.z; - if (h >= n_head || j >= S) return; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return; - - const uint32_t qbase = n_score / S; - const uint32_t rem = n_score % S; - const uint32_t g0 = j * qbase + (j < rem ? j : rem); - const uint32_t cnt = qbase + (j < rem ? 1u : 0u); - const uint32_t g1 = g0 + cnt; - const float *qh = q + (uint64_t)h * head_dim; - float *row_scores = score_out + (uint64_t)h * n_score; - const float scale = rsqrtf((float)head_dim); - - for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { - float s = -INFINITY; - if (g < raw_count) { - const uint32_t raw_row = - (raw_start + raw_first_idx + g) % raw_cap; - const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale; - } else { - const uint32_t cidx = g - raw_count; - const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale + add; - } - } - row_scores[g] = s; - } -} - -__device__ __forceinline__ float ds4_dot_scalar_ldg( - const float *a, - const float *b, - uint32_t n) { - float dot = 0.0f; - for (uint32_t d = 0; d < n; d++) dot += __ldg(a + d) * __ldg(b + d); - return dot; -} - -__global__ static void attention_decode_score_split_scores_ldg_kernel( - float *score_out, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - const uint32_t h = blockIdx.y; - const uint32_t j = blockIdx.z; - if (h >= n_head || j >= S) return; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return; - - const uint32_t qbase = n_score / S; - const uint32_t rem = n_score % S; - const uint32_t g0 = j * qbase + (j < rem ? j : rem); - const uint32_t cnt = qbase + (j < rem ? 1u : 0u); - const uint32_t g1 = g0 + cnt; - const float *qh = q + (uint64_t)h * head_dim; - float *row_scores = score_out + (uint64_t)h * n_score; - const float scale = rsqrtf((float)head_dim); - - for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { - float s = -INFINITY; - if (g < raw_count) { - const uint32_t raw_row = - (raw_start + raw_first_idx + g) % raw_cap; - const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; - const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim); - s = dot * scale; - } else { - const uint32_t cidx = g - raw_count; - const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; - const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim); - s = dot * scale + add; - } - } - row_scores[g] = s; - } -} - -/* Head-tiled exact score kernel for head_dim==512. - * - * The reference score kernel assigns one (head, row-chunk) per block and lets - * every thread walk one KV row with a scalar sequential dot. Because MQA - * shares the same KV rows across all 64 heads, that reference layout re-reads - * every KV row once per head, and the per-thread row walk is fully - * uncoalesced (threads stride 2KB apart), which multiplies L2 traffic again. - * - * This kernel keeps the per-score arithmetic bit-identical (same ascending-d - * scalar accumulation `dot += q[d] * kv[d]`, same `dot * scale [+ add]` - * epilogue, same masked-row/raw-window classification) but stages a 16-row KV - * tile and a 16-head Q tile in shared memory with coalesced global loads, so - * each KV row is read from L2 once per 16 heads instead of once per head. - * Scores are independent outputs, so retiling the (head, row) space cannot - * change any output bit as long as each individual dot keeps its order. */ -#define DS4_SCORE_TILE_HEADS 16u -#define DS4_SCORE_TILE_ROWS 16u -#define DS4_SCORE_TILE_STRIDE 516u /* 512 + 4 floats: 16B-aligned rows, banks shifted by 4 */ - -__global__ static void attention_decode_score_split_scores_tile512_kernel( - float *score_out, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return; - - extern __shared__ float score_tile_shared[]; - float *sh_q = score_tile_shared; /* 16 x 516 */ - float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE; /* 16 x 516 */ - __shared__ float sh_add[DS4_SCORE_TILE_ROWS]; - - const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS; - const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS; - if (g_base >= n_score || h_base >= n_head) return; - - /* Cooperative Q tile load: 16 heads x 512 floats, float4 coalesced. */ - { - const float4 *q4 = (const float4 *)(q + (uint64_t)h_base * 512u); - const uint32_t tile_heads = - n_head - h_base < DS4_SCORE_TILE_HEADS ? n_head - h_base : DS4_SCORE_TILE_HEADS; - for (uint32_t idx = threadIdx.x; idx < tile_heads * 128u; idx += blockDim.x) { - const uint32_t hh = idx >> 7u; /* head within tile */ - const uint32_t dd = idx & 127u; /* float4 within row */ - const float4 v = q4[hh * 128u + dd]; - float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u; - dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; - } - } - /* Row classification + mask staging (thread per row). */ - if (threadIdx.x < DS4_SCORE_TILE_ROWS) { - const uint32_t g = g_base + threadIdx.x; - float add = -INFINITY; - if (g < n_score) { - if (g < raw_count) { - add = 0.0f; /* raw rows are always visible */ - } else { - const uint32_t cidx = g - raw_count; - add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; - } - } - sh_add[threadIdx.x] = add; - } - __syncthreads(); - /* Cooperative KV tile load: two rows at a time, float4 coalesced. - * Masked rows (add <= -1e20) are skipped; their scores never read KV. */ - { - const uint32_t rows_per_pass = blockDim.x >> 7u; /* 128 threads per row */ - const uint32_t rr0 = threadIdx.x >> 7u; - const uint32_t dd = threadIdx.x & 127u; - for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) { - const uint32_t g = g_base + r; - if (g >= n_score) continue; - const bool visible = g < raw_count || sh_add[r] > -1.0e20f; - if (!visible) continue; - const float4 *src; - if (g < raw_count) { - const uint32_t raw_row = (raw_start + raw_first_idx + g) % raw_cap; - src = (const float4 *)(raw_kv + (uint64_t)raw_row * 512u); - } else { - const uint32_t cidx = g - raw_count; - src = (const float4 *)(comp_kv + (uint64_t)cidx * 512u); - } - const float4 v = src[dd]; - float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u; - dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; - } - } - __syncthreads(); - - /* One score per thread: r = tid&15 (consecutive threads, coalesced score - * writes), h = tid>>4. The dot keeps the reference kernel's exact scalar - * ascending-d accumulation. */ - const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u); - const uint32_t h = h_base + (threadIdx.x >> 4u); - const uint32_t g = g_base + r; - if (h >= n_head || g >= n_score) return; - const float scale = rsqrtf((float)head_dim); - float *row_scores = score_out + (uint64_t)h * n_score; - const float *qh = sh_q + (uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE; - const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE; - float s = -INFINITY; - const bool need_dot = g < raw_count || sh_add[r] > -1.0e20f; - if (need_dot) { - /* The reference kernel's runtime-trip loop compiles to one sequential - * FFMA chain. Keep exactly that accumulation order here: batched loads - * for latency hiding, but a single explicit ascending fma chain. */ - float dot = 0.0f; -#pragma unroll 1 - for (uint32_t d = 0; d < 512u; d += 8u) { - const float a0 = qh[d + 0u], a1 = qh[d + 1u]; - const float a2 = qh[d + 2u], a3 = qh[d + 3u]; - const float a4 = qh[d + 4u], a5 = qh[d + 5u]; - const float a6 = qh[d + 6u], a7 = qh[d + 7u]; - const float b0 = kvrow[d + 0u], b1 = kvrow[d + 1u]; - const float b2 = kvrow[d + 2u], b3 = kvrow[d + 3u]; - const float b4 = kvrow[d + 4u], b5 = kvrow[d + 5u]; - const float b6 = kvrow[d + 6u], b7 = kvrow[d + 7u]; - dot = __fmaf_rn(a0, b0, dot); - dot = __fmaf_rn(a1, b1, dot); - dot = __fmaf_rn(a2, b2, dot); - dot = __fmaf_rn(a3, b3, dot); - dot = __fmaf_rn(a4, b4, dot); - dot = __fmaf_rn(a5, b5, dot); - dot = __fmaf_rn(a6, b6, dot); - dot = __fmaf_rn(a7, b7, dot); - } - if (g < raw_count) { - s = dot * scale; - } else { - /* The reference expression `dot * scale + add` contracts to one - * FFMA; keep that exact contraction explicit. */ - s = __fmaf_rn(dot, scale, sh_add[r]); - } - } - row_scores[g] = s; -} - -/* Multi-session form of the exact tiled score kernel. Each z-slice selects a - * private KV table entry, while every individual score keeps the same scalar - * ascending-d FMA chain as the one-session kernel. */ -__global__ static void attention_decode_score_split_scores_tile512_rows_kernel( - float *score_out, - const float *q, - cuda_attention_decode_row_table rows, - uint32_t n_rows, - uint32_t score_stride, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t row = blockIdx.z; - if (row >= n_rows) return; - const ds4_gpu_attention_decode_row dsc = rows.row[row]; - if (dsc.indexed) return; - - const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; - const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; - const bool single_all = dsc.ratio == 0u; - const uint32_t qpos = dsc.pos; - const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; - uint32_t visible_comp = single_all - ? dsc.n_comp - : (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u); - if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; - - uint32_t raw_count = 0u; - uint32_t raw_first_idx = 0u; - if (dsc.n_raw != 0u) { - const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; - if (single_all) { - raw_count = dsc.n_raw > 256u ? 256u : dsc.n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (dsc.window != 0u && qpos + 1u > dsc.window) { - const uint32_t wlo = qpos + 1u - dsc.window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - - extern __shared__ float score_tile_shared[]; - float *sh_q = score_tile_shared; - float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE; - - const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS; - const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS; - if (g_base >= n_score || h_base >= n_head) return; - - { - const float4 *q4 = (const float4 *)( - q + ((uint64_t)row * n_head + h_base) * head_dim); - const uint32_t tile_heads = - n_head - h_base < DS4_SCORE_TILE_HEADS - ? n_head - h_base : DS4_SCORE_TILE_HEADS; - for (uint32_t idx = threadIdx.x; - idx < tile_heads * 128u; - idx += blockDim.x) { - const uint32_t hh = idx >> 7u; - const uint32_t dd = idx & 127u; - const float4 v = q4[hh * 128u + dd]; - float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u; - dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; - } - } - __syncthreads(); - { - const uint32_t rows_per_pass = blockDim.x >> 7u; - const uint32_t rr0 = threadIdx.x >> 7u; - const uint32_t dd = threadIdx.x & 127u; - for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) { - const uint32_t g = g_base + r; - if (g >= n_score) continue; - const float4 *src; - if (g < raw_count) { - const uint32_t raw_row = - (dsc.raw_start + raw_first_idx + g) % dsc.raw_cap; - src = (const float4 *)(raw_kv + (uint64_t)raw_row * head_dim); - } else { - src = (const float4 *)(comp_kv + - (uint64_t)(g - raw_count) * head_dim); - } - const float4 v = src[dd]; - float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u; - dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; - } - } - __syncthreads(); - - const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u); - const uint32_t h = h_base + (threadIdx.x >> 4u); - const uint32_t g = g_base + r; - if (h >= n_head || g >= n_score) return; - const float scale = rsqrtf((float)head_dim); - float *row_scores = score_out + - ((uint64_t)row * n_head + h) * score_stride; - const float *qh = sh_q + - (uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE; - const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE; - float dot = 0.0f; -#pragma unroll 1 - for (uint32_t dd = 0; dd < 512u; dd += 8u) { - const float a0 = qh[dd + 0u], a1 = qh[dd + 1u]; - const float a2 = qh[dd + 2u], a3 = qh[dd + 3u]; - const float a4 = qh[dd + 4u], a5 = qh[dd + 5u]; - const float a6 = qh[dd + 6u], a7 = qh[dd + 7u]; - const float b0 = kvrow[dd + 0u], b1 = kvrow[dd + 1u]; - const float b2 = kvrow[dd + 2u], b3 = kvrow[dd + 3u]; - const float b4 = kvrow[dd + 4u], b5 = kvrow[dd + 5u]; - const float b6 = kvrow[dd + 6u], b7 = kvrow[dd + 7u]; - dot = __fmaf_rn(a0, b0, dot); - dot = __fmaf_rn(a1, b1, dot); - dot = __fmaf_rn(a2, b2, dot); - dot = __fmaf_rn(a3, b3, dot); - dot = __fmaf_rn(a4, b4, dot); - dot = __fmaf_rn(a5, b5, dot); - dot = __fmaf_rn(a6, b6, dot); - dot = __fmaf_rn(a7, b7, dot); - } - row_scores[g] = g < raw_count - ? dot * scale - : __fmaf_rn(dot, scale, 0.0f); -} - -__device__ __forceinline__ float ds4_dot512_float4_ordered( - const float *a, - const float *b) { - const float4 *a4 = (const float4 *)a; - const float4 *b4 = (const float4 *)b; - float dot = 0.0f; -#pragma unroll 1 - for (uint32_t i = 0; i < 128u; i++) { - const float4 av = a4[i]; - const float4 bv = b4[i]; - dot = __fadd_rn(dot, __fmul_rn(av.x, bv.x)); - dot = __fadd_rn(dot, __fmul_rn(av.y, bv.y)); - dot = __fadd_rn(dot, __fmul_rn(av.z, bv.z)); - dot = __fadd_rn(dot, __fmul_rn(av.w, bv.w)); - } - return dot; -} - -__device__ __forceinline__ float ds4_dot512_float4_plain( - const float *a, - const float *b) { - const float4 *a4 = (const float4 *)a; - const float4 *b4 = (const float4 *)b; - float dot = 0.0f; -#pragma unroll 1 - for (uint32_t i = 0; i < 128u; i++) { - const float4 av = a4[i]; - const float4 bv = b4[i]; - dot += av.x * bv.x; - dot += av.y * bv.y; - dot += av.z * bv.z; - dot += av.w * bv.w; - } - return dot; -} - -__global__ static void attention_decode_score_split_scores_vec4_kernel( - float *score_out, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t S) { - const uint32_t h = blockIdx.y; - const uint32_t j = blockIdx.z; - if (h >= n_head || j >= S) return; - const uint32_t head_dim = 512u; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return; - - const uint32_t qbase = n_score / S; - const uint32_t rem = n_score % S; - const uint32_t g0 = j * qbase + (j < rem ? j : rem); - const uint32_t cnt = qbase + (j < rem ? 1u : 0u); - const uint32_t g1 = g0 + cnt; - const float *qh = q + (uint64_t)h * head_dim; - float *row_scores = score_out + (uint64_t)h * n_score; - const float scale = rsqrtf((float)head_dim); - - for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { - float s = -INFINITY; - if (g < raw_count) { - const uint32_t raw_row = - (raw_start + raw_first_idx + g) % raw_cap; - const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; - const float dot = ds4_dot512_float4_ordered(qh, kvrow); - s = dot * scale; - } else { - const uint32_t cidx = g - raw_count; - const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; - const float dot = ds4_dot512_float4_ordered(qh, kvrow); - s = dot * scale + add; - } - } - row_scores[g] = s; - } -} - -__global__ static void attention_decode_score_split_scores_vec4_plain_kernel( - float *score_out, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t S) { - const uint32_t h = blockIdx.y; - const uint32_t j = blockIdx.z; - if (h >= n_head || j >= S) return; - const uint32_t head_dim = 512u; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return; - - const uint32_t qbase = n_score / S; - const uint32_t rem = n_score % S; - const uint32_t g0 = j * qbase + (j < rem ? j : rem); - const uint32_t cnt = qbase + (j < rem ? 1u : 0u); - const uint32_t g1 = g0 + cnt; - const float *qh = q + (uint64_t)h * head_dim; - float *row_scores = score_out + (uint64_t)h * n_score; - const float scale = rsqrtf((float)head_dim); - - for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { - float s = -INFINITY; - if (g < raw_count) { - const uint32_t raw_row = - (raw_start + raw_first_idx + g) % raw_cap; - const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; - const float dot = ds4_dot512_float4_plain(qh, kvrow); - s = dot * scale; - } else { - const uint32_t cidx = g - raw_count; - const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; - const float dot = ds4_dot512_float4_plain(qh, kvrow); - s = dot * scale + add; - } - } - row_scores[g] = s; - } -} - -__global__ static void attention_decode_score_split_finalize_kernel( - float *heads, - const float *sinks, - const float *score_in, - const float *raw_kv, - const float *comp_kv, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t h = blockIdx.y; - if (h >= n_head) return; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; - __shared__ uint32_t raw_rows[256]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count_s; - __shared__ uint32_t raw_first_idx_s; - - const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; - const bool score_thread = threadIdx.x < score_threads; - if (threadIdx.x == 0) { - raw_count_s = 0; - raw_first_idx_s = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count_s = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx_s = lo - first_raw_pos; - raw_count_s = hi - lo + 1u; - if (raw_count_s > 256u) raw_count_s = 256u; - } - } - } - } - __syncthreads(); - const uint32_t raw_count = raw_count_s; - const uint32_t raw_first_idx = raw_first_idx_s; - if (score_thread) { - for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - } - __syncthreads(); - const uint32_t n_score = raw_count + visible_comp; - const float *row_scores = score_in + (uint64_t)h * n_score; - float local_max = sinks[h]; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - const float s = row_scores[i]; - scores[i] = s; - local_max = fmaxf(local_max, s); - } - } - if (score_thread) partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = - fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - } - if (score_thread) partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - float *oh = heads + (uint64_t)h * head_dim; - if (head_dim == 512u && blockDim.x >= 512u) { - const uint32_t d = threadIdx.x; - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - const float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc += kv[d] * s; - } - for (uint32_t c = 0; c < visible_comp; c++) { - const float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc += kv[d] * s; - } - oh[d] = acc / denom; - } else if (head_dim == 512u && blockDim.x == 256u) { - const uint32_t d0 = threadIdx.x; - const uint32_t d1 = d0 + 256u; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - const float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - for (uint32_t c = 0; c < visible_comp; c++) { - const float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - oh[d0] = acc0 / denom; - oh[d1] = acc1 / denom; - } else { - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; - } - for (uint32_t c = 0; c < visible_comp; c++) { - acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; - } - oh[d] = acc / denom; - } - } -} - -__global__ static void attention_decode_score_split_finalize_rows_kernel( - float *heads, - const float *sinks, - const float *score_in, - cuda_attention_decode_row_table rows, - uint32_t n_rows, - uint32_t score_stride, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t row = blockIdx.x; - const uint32_t h = blockIdx.y; - if (row >= n_rows || h >= n_head) return; - const ds4_gpu_attention_decode_row dsc = rows.row[row]; - if (dsc.indexed) return; - const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; - const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; - const bool single_all = dsc.ratio == 0u; - const uint32_t qpos = dsc.pos; - const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; - uint32_t visible_comp = single_all - ? dsc.n_comp - : (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u); - if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; - - __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; - __shared__ uint32_t raw_rows[256]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count_s; - __shared__ uint32_t raw_first_idx_s; - - const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; - const bool score_thread = threadIdx.x < score_threads; - if (threadIdx.x == 0u) { - raw_count_s = 0u; - raw_first_idx_s = 0u; - if (dsc.n_raw != 0u) { - const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; - if (single_all) { - raw_count_s = dsc.n_raw > 256u ? 256u : dsc.n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (dsc.window != 0u && qpos + 1u > dsc.window) { - const uint32_t wlo = qpos + 1u - dsc.window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx_s = lo - first_raw_pos; - raw_count_s = hi - lo + 1u; - if (raw_count_s > 256u) raw_count_s = 256u; - } - } - } - } - __syncthreads(); - const uint32_t raw_count = raw_count_s; - const uint32_t raw_first_idx = raw_first_idx_s; - if (score_thread) { - for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { - raw_rows[r] = - (dsc.raw_start + raw_first_idx + r) % dsc.raw_cap; - } - } - __syncthreads(); - const uint32_t n_score = raw_count + visible_comp; - const float *row_scores = score_in + - ((uint64_t)row * n_head + h) * score_stride; - float local_max = sinks[h]; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - const float s = row_scores[i]; - scores[i] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - } - __syncthreads(); - for (uint32_t stride = score_threads >> 1u; - stride > 0u; - stride >>= 1u) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = - fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0u) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - partial[threadIdx.x] = den_local; - } - __syncthreads(); - for (uint32_t stride = score_threads >> 1u; - stride > 0u; - stride >>= 1u) { - if (threadIdx.x < stride) { - partial[threadIdx.x] += partial[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0u) { - denom = partial[0] + expf(sinks[h] - max_s); - } - __syncthreads(); - - float *oh = heads + ((uint64_t)row * n_head + h) * head_dim; - if (head_dim == 512u && blockDim.x >= 512u) { - const uint32_t dim = threadIdx.x; - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc += kv[dim] * scores[r]; - } - for (uint32_t c = 0; c < visible_comp; c++) { - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc += kv[dim] * scores[raw_count + c]; - } - oh[dim] = acc / denom; - } else { - for (uint32_t dim = threadIdx.x; - dim < head_dim; - dim += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] * - scores[r]; - } - for (uint32_t c = 0; c < visible_comp; c++) { - acc += comp_kv[(uint64_t)c * head_dim + dim] * - scores[raw_count + c]; - } - oh[dim] = acc / denom; - } - } -} - -__global__ static void attention_decode_score_split_finalize_dim2_kernel( - float *heads, - const float *sinks, - const float *score_in, - const float *raw_kv, - const float *comp_kv, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t dim_half = blockIdx.x; - const uint32_t h = blockIdx.y; - if (h >= n_head || head_dim != 512u || dim_half >= 2u) return; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - - __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; - __shared__ uint32_t raw_rows[256]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count_s; - __shared__ uint32_t raw_first_idx_s; - - const uint32_t score_threads = 256u; - if (threadIdx.x == 0) { - raw_count_s = 0; - raw_first_idx_s = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count_s = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx_s = lo - first_raw_pos; - raw_count_s = hi - lo + 1u; - if (raw_count_s > 256u) raw_count_s = 256u; - } - } - } - } - __syncthreads(); - const uint32_t raw_count = raw_count_s; - const uint32_t raw_first_idx = raw_first_idx_s; - for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - __syncthreads(); - - const uint32_t n_score = raw_count + visible_comp; - const float *row_scores = score_in + (uint64_t)h * n_score; - float local_max = sinks[h]; - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - const float s = row_scores[i]; - scores[i] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = - fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - - float den_local = 0.0f; - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - - const uint32_t d = dim_half * 256u + threadIdx.x; - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - const float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc += kv[d] * s; - } - for (uint32_t c = 0; c < visible_comp; c++) { - const float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)c * head_dim; - acc += kv[d] * s; - } - heads[(uint64_t)h * head_dim + d] = acc / denom; -} - -__global__ static void attention_decode_global_softmax_kernel( - float *score_inout, - float *denom_out, - const float *sinks, - uint32_t n_score, - uint32_t n_head) { - const uint32_t h = blockIdx.x; - if (h >= n_head || n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return; - __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom_s; - const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; - const bool score_thread = threadIdx.x < score_threads; - float *row_scores = score_inout + (uint64_t)h * n_score; - - float local_max = sinks[h]; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - const float s = row_scores[i]; - scores[i] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - } - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = - fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - - float den_local = 0.0f; - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - const float e = expf(scores[i] - max_s); - scores[i] = e; - den_local += e; - } - partial[threadIdx.x] = den_local; - } - __syncthreads(); - for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom_s = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - - if (score_thread) { - for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { - row_scores[i] = scores[i]; - } - } - if (threadIdx.x == 0) denom_out[h] = denom_s; -} - -__global__ static void attention_decode_split_value_kernel( - float *partials, - const float *score_exp, - const float *raw_kv, - const float *comp_kv, - uint32_t raw_count, - uint32_t raw_first_idx, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_score, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - const uint32_t h = blockIdx.y; - const uint32_t j = blockIdx.z; - if (h >= n_head || j >= S || n_score == 0u) return; - const uint32_t qbase = n_score / S; - const uint32_t rem = n_score % S; - const uint32_t g0 = j * qbase + (j < rem ? j : rem); - const uint32_t cnt = qbase + (j < rem ? 1u : 0u); - const uint32_t g1 = g0 + cnt; - const float *row_scores = score_exp + (uint64_t)h * n_score; - float *pout = partials + ((uint64_t)h * S + j) * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t g = g0; g < g1; g++) { - const float s = row_scores[g]; - if (g < raw_count) { - const uint32_t raw_row = - (raw_start + raw_first_idx + g) % raw_cap; - acc += raw_kv[(uint64_t)raw_row * head_dim + d] * s; - } else { - const uint32_t c = g - raw_count; - acc += comp_kv[(uint64_t)c * head_dim + d] * s; - } - } - pout[d] = acc; - } -} - -__global__ static void attention_decode_split_value_combine_kernel( - float *heads, - const float *partials, - const float *denom, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - const uint32_t h = blockIdx.y; - if (h >= n_head) return; - const float *base = partials + (uint64_t)h * S * head_dim; - const float den = denom[h]; - float *oh = heads + (uint64_t)h * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t j = 0; j < S; j++) { - acc += base[(uint64_t)j * head_dim + d]; - } - oh[d] = acc / den; - } -} - -typedef struct { - uint32_t n_rot; - uint32_t pos0; - uint32_t n_ctx_orig; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -} cuda_attention_inv_rope_params; - -static void attention_decode_score_split_graph_destroy_one(int logical_tier) { - if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return; - cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier]; - if (c->exec) (void)cudaGraphExecDestroy(c->exec); - if (c->graph) (void)cudaGraphDestroy(c->graph); - memset(c, 0, sizeof(*c)); -} - -static int attention_decode_score_split_graph_launch( - int logical_tier, - float *heads, - const float *sinks, - float *scores, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t final_threads, - uint32_t S, - const cuda_attention_inv_rope_params *inv_rope) { - if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0; - cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier]; - const bool graph_inv_rope = - inv_rope && - head_dim == 512u && - inv_rope->n_rot != 0u && - inv_rope->n_rot <= head_dim && - (inv_rope->n_rot & 1u) == 0u; - const bool shape_match = - c->valid && - c->n_head == n_head && - c->head_dim == head_dim && - c->S == S && - c->final_threads == final_threads && - c->fuses_inv_rope == (graph_inv_rope ? 1 : 0) && - (!graph_inv_rope || c->n_rot == inv_rope->n_rot); - if (c->valid && !shape_match) { - attention_decode_score_split_graph_destroy_one(logical_tier); - c = &g_score_split_graph[logical_tier]; - } - - dim3 score_grid(1, n_head, S); - dim3 final_grid(1, n_head, 1); - dim3 score_block(256, 1, 1); - dim3 final_block(final_threads, 1, 1); - - void *score_args[] = { - &scores, &q, &raw_kv, &comp_kv, &comp_mask, &use_comp_mask, - &pos0, &n_raw, &raw_cap, &raw_start, &n_comp, &window, &ratio, - &n_head, &head_dim, &S - }; - cudaKernelNodeParams score_params; - memset(&score_params, 0, sizeof(score_params)); - score_params.func = (void *)attention_decode_score_split_scores_kernel; - score_params.gridDim = score_grid; - score_params.blockDim = score_block; - score_params.sharedMemBytes = 0; - score_params.kernelParams = score_args; - score_params.extra = NULL; - - void *final_args[] = { - &heads, &sinks, &scores, &raw_kv, &comp_kv, &pos0, &n_raw, - &raw_cap, &raw_start, &n_comp, &window, &ratio, &n_head, &head_dim - }; - cudaKernelNodeParams final_params; - memset(&final_params, 0, sizeof(final_params)); - final_params.func = (void *)attention_decode_score_split_finalize_kernel; - final_params.gridDim = final_grid; - final_params.blockDim = final_block; - final_params.sharedMemBytes = 0; - final_params.kernelParams = final_args; - final_params.extra = NULL; - - uint32_t rope_n_tok = 1u; - uint32_t rope_pos_stride = 1u; - int rope_inverse = 1; - uint32_t rope_n_rot = graph_inv_rope ? inv_rope->n_rot : 0u; - uint32_t rope_pos0 = graph_inv_rope ? inv_rope->pos0 : 0u; - uint32_t rope_n_ctx_orig = graph_inv_rope ? inv_rope->n_ctx_orig : 0u; - float rope_freq_base = graph_inv_rope ? inv_rope->freq_base : 0.0f; - float rope_freq_scale = graph_inv_rope ? inv_rope->freq_scale : 0.0f; - float rope_ext_factor = graph_inv_rope ? inv_rope->ext_factor : 0.0f; - float rope_attn_factor = graph_inv_rope ? inv_rope->attn_factor : 0.0f; - float rope_beta_fast = graph_inv_rope ? inv_rope->beta_fast : 0.0f; - float rope_beta_slow = graph_inv_rope ? inv_rope->beta_slow : 0.0f; - void *rope_args[] = { - &heads, &rope_n_tok, &n_head, &head_dim, &rope_n_rot, - &rope_pos0, &rope_pos_stride, &rope_n_ctx_orig, &rope_inverse, - &rope_freq_base, &rope_freq_scale, &rope_ext_factor, - &rope_attn_factor, &rope_beta_fast, &rope_beta_slow - }; - cudaKernelNodeParams rope_params; - memset(&rope_params, 0, sizeof(rope_params)); - if (graph_inv_rope) { - const uint32_t pairs = n_head * (rope_n_rot / 2u); - rope_params.func = (void *)rope_tail_kernel; - rope_params.gridDim = dim3((pairs + 255u) / 256u, 1, 1); - rope_params.blockDim = dim3(256, 1, 1); - rope_params.sharedMemBytes = 0; - rope_params.kernelParams = rope_args; - rope_params.extra = NULL; - } - - if (!c->valid) { - cudaError_t err = cudaGraphCreate(&c->graph, 0); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: attention score-split graph create failed: %s\n", - cudaGetErrorString(err)); - attention_decode_score_split_graph_destroy_one(logical_tier); - return -1; - } - err = cudaGraphAddKernelNode(&c->score_node, c->graph, NULL, 0, - &score_params); - if (err == cudaSuccess) { - err = cudaGraphAddKernelNode(&c->final_node, c->graph, - &c->score_node, 1, &final_params); - } - if (err == cudaSuccess && graph_inv_rope) { - err = cudaGraphAddKernelNode(&c->rope_node, c->graph, - &c->final_node, 1, &rope_params); - } - if (err == cudaSuccess) { - err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0); - } - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: attention score-split graph instantiate failed: %s\n", - cudaGetErrorString(err)); - attention_decode_score_split_graph_destroy_one(logical_tier); - return -1; - } - c->n_head = n_head; - c->head_dim = head_dim; - c->S = S; - c->final_threads = final_threads; - c->n_rot = graph_inv_rope ? inv_rope->n_rot : 0u; - c->fuses_inv_rope = graph_inv_rope ? 1 : 0; - c->valid = 1; - } else { - cudaError_t err = - cudaGraphExecKernelNodeSetParams(c->exec, c->score_node, - &score_params); - if (err == cudaSuccess) { - err = cudaGraphExecKernelNodeSetParams(c->exec, c->final_node, - &final_params); - } - if (err == cudaSuccess && graph_inv_rope) { - err = cudaGraphExecKernelNodeSetParams(c->exec, c->rope_node, - &rope_params); - } - if (err != cudaSuccess) { - fprintf(stderr, - "ds4: attention score-split graph update failed: %s\n", - cudaGetErrorString(err)); - attention_decode_score_split_graph_destroy_one(logical_tier); - return -1; - } - } - - cudaError_t err = cudaGraphLaunch(c->exec, 0); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: attention score-split graph launch failed: %s\n", - cudaGetErrorString(err)); - attention_decode_score_split_graph_destroy_one(logical_tier); - return -1; - } - return 1; -} - -static int attention_decode_score_split_launch( - int logical_tier, - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t final_threads, - const cuda_attention_inv_rope_params *inv_rope) { - if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0)) return 0; - const int explicit_exact = - cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 0); - if (!cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1)) return 0; - if (!explicit_exact && cuda_splitkv_decode_requested()) return 0; - if (g_cuda_decode_score4 || g_cuda_decode_score8) return 0; - if (head_dim == 0u || n_head == 0u) return 0; - const bool single_all = (ratio == 0u); - const uint32_t qpos = pos0; - const uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - uint32_t raw_count = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - const uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return 0; - /* With the head-tiled score kernel the exact score-split path beats the - * one-block mixed kernel even for short score counts, so the gate that - * used to protect short contexts (512) now defaults to 1. */ - const uint32_t min_score = cuda_parse_u32_env_clamped( - "DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE", 1u, 0u, - DS4_CUDA_ATTENTION_SCORE_CAP, NULL); - if (n_score < min_score) return 0; - uint32_t chunk = cuda_parse_u32_env_clamped( - "DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK", DS4_CUDA_SPLITKV_CHUNK, - 1u, DS4_CUDA_ATTENTION_SCORE_CAP, NULL); - uint32_t s_floor = cuda_parse_u32_env_clamped( - "DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR", 6u, - 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); - uint32_t s_max = cuda_parse_u32_env_clamped( - "DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX", DS4_CUDA_SPLITKV_S_MAX, - 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); - int exact_present = 0; - uint32_t S = cuda_parse_u32_env_clamped( - "DS4_CUDA_EXACT_SCORE_SPLIT_S", 0u, 1u, - DS4_CUDA_SPLITKV_S_MAX, &exact_present); - if (!exact_present) { - S = (n_score + chunk - 1u) / chunk; - if (S < s_floor) S = s_floor < n_score ? s_floor : n_score; - if (S > s_max) S = s_max; - } - if (S > n_score) S = n_score; - if (S <= 1u) return 0; - const bool graph_inv_rope = - g_cuda_exact_score_split_fuse_inv_rope && - inv_rope && - head_dim == 512u && - final_threads >= 512u && - inv_rope->n_rot != 0u && - inv_rope->n_rot <= 512u && - (inv_rope->n_rot & 1u) == 0u; - - const uint64_t score_count = (uint64_t)n_head * n_score; - float *scores = (float *)cuda_tmp_alloc_on(logical_tier, - score_count * sizeof(float), - "attention exact score split"); - if (!scores) return 0; - const bool use_ldg_scores = g_cuda_exact_score_split_ldg; - const bool use_vec4_plain_scores = - !use_ldg_scores && - g_cuda_exact_score_split_vec4_plain && - head_dim == 512u; - const bool use_vec4_scores = - !use_ldg_scores && - !use_vec4_plain_scores && - (g_cuda_exact_score_split_vec4 || g_decode_score_vec4) && - head_dim == 512u; - const bool use_dim2_finalize = - g_cuda_exact_score_split_dim2 && - head_dim == 512u && - final_threads >= 512u && - !graph_inv_rope; - if ((g_cuda_exact_score_split_graph || graph_inv_rope) && - !use_dim2_finalize && - !use_ldg_scores && - !use_vec4_plain_scores && - !use_vec4_scores) - { - int rc = attention_decode_score_split_graph_launch( - logical_tier, heads, sinks, scores, q, raw_kv, comp_kv, comp_mask, - use_comp_mask, pos0, n_raw, raw_cap, raw_start, n_comp, window, - ratio, n_head, head_dim, final_threads, S, - graph_inv_rope ? inv_rope : NULL); - if (rc == 1) return 1; - if (rc < 0) return -1; - } - if (graph_inv_rope) return 0; - static int score_tile_disabled = -1; - if (score_tile_disabled < 0) { - score_tile_disabled = getenv("DS4_CUDA_NO_SCORE_TILE") != NULL ? 1 : 0; - } - if (!score_tile_disabled && - head_dim == 512u && - !use_ldg_scores && - !use_vec4_plain_scores && - !use_vec4_scores) { - /* cudaFuncSetAttribute() applies to the current device only, so opt in - * to >48KB dynamic shared memory once per device. */ - static int tile_shmem_ready[DS4_MAX_GPUS] = {0}; - const size_t tile_shmem = - (size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) * - DS4_SCORE_TILE_STRIDE * sizeof(float); - int tile_dev = 0; - cudaGetDevice(&tile_dev); - if (tile_dev >= 0 && tile_dev < DS4_MAX_GPUS && - !tile_shmem_ready[tile_dev]) { - if (!cuda_ok(cudaFuncSetAttribute( - attention_decode_score_split_scores_tile512_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - (int)tile_shmem), - "attention score tile shared-memory opt-in")) { - score_tile_disabled = 1; - } - tile_shmem_ready[tile_dev] = 1; - } - if (score_tile_disabled) { - return 0; /* retry via the generic path on the next call */ - } - dim3 tile_grid((n_score + DS4_SCORE_TILE_ROWS - 1u) / DS4_SCORE_TILE_ROWS, - (n_head + DS4_SCORE_TILE_HEADS - 1u) / DS4_SCORE_TILE_HEADS, - 1); - attention_decode_score_split_scores_tile512_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), "attention exact score split tile launch")) return -1; - } else { - dim3 score_grid(1, n_head, S); - if (use_ldg_scores) { - attention_decode_score_split_scores_ldg_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim, S); - } else if (use_vec4_plain_scores) { - attention_decode_score_split_scores_vec4_plain_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, S); - } else if (use_vec4_scores) { - attention_decode_score_split_scores_vec4_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, S); - } else { - attention_decode_score_split_scores_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim, S); - } - if (!cuda_ok(cudaGetLastError(), "attention exact score split scores launch")) return -1; - } - if (use_dim2_finalize) { - dim3 final_grid(2, n_head, 1); - attention_decode_score_split_finalize_dim2_kernel<<>>( - heads, sinks, scores, raw_kv, comp_kv, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), "attention exact score split dim2 finalize launch")) return -1; - } else { - dim3 final_grid(1, n_head, 1); - attention_decode_score_split_finalize_kernel<<>>( - heads, sinks, scores, raw_kv, comp_kv, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), "attention exact score split finalize launch")) return -1; - } - return 1; -} - -/* ---- perf-02 split-KV / flash-decode (opt-in, default OFF) ---------------- - * - * attention_decode_splitkv_kernel computes a partial online-softmax over a - * contiguous chunk of the flattened logical row set [0, n_score) used by - * attention_decode_mixed_kernel (raw rows first, then compressed rows, same - * ascending ordering). Each block handles (t = blockIdx.x, h = blockIdx.y, - * chunk = blockIdx.z) and writes a partial (m_j, l_j, acc_j[head_dim]) WITHOUT - * the sink term. attention_decode_splitkv_combine_kernel merges the S partials - * per (t,h), folds the sink once, and writes the final normalized head output. - * - * The math is the standard flash-attention online-softmax rescale and is - * algebraically identical to attention_decode_mixed_kernel; it is NOT - * guaranteed bit-identical in FP32 (different expf inputs + add/mul grouping), - * hence default-OFF behind DS4_CUDA_SPLITKV_DECODE and the S==1 dispatch to the - * old kernel as the bit-exact anchor (handled in the launch helper). - * - * Partials scratch layout (per logical tier), contiguous floats: - * stride = head_dim + 2 - * base(t,h,j) = ((t*n_head + h)*S + j) * stride - * [0] = m_j (chunk running max; -INF if empty/all-masked) - * [1] = l_j (chunk denominator sum exp(s - m_j)) - * [2 .. 2+head_dim) = acc_j[head_dim] (chunk weighted value sum) - */ -__global__ static void attention_decode_splitkv_kernel( - float *partials, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - uint32_t j = blockIdx.z; - if (t >= n_tokens || h >= n_head || j >= S) return; - const bool single_all = (n_tokens == 1u && ratio == 0u); - uint32_t qpos = pos0 + t; - uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - /* scores buffer holds only this chunk's rows. The launch helper guarantees - * cnt <= DS4_CUDA_SPLITKV_SCORE_CAP, including env-tuned split counts. */ - __shared__ float scores[DS4_CUDA_SPLITKV_SCORE_CAP]; - __shared__ uint32_t raw_rows[256]; - __shared__ float partial[256]; - __shared__ float m_s; - __shared__ float l_s; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - float scale = rsqrtf((float)head_dim); - if (threadIdx.x == 0) { - raw_count = 0; - raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - uint32_t n_score = raw_count + visible_comp; - /* even split of [0, n_score) across S chunks: first (n_score % S) chunks - * get base+1, identical deterministic partition for every block. */ - uint32_t qbase = n_score / S; - uint32_t rem = n_score % S; - uint32_t g0 = j * qbase + (j < rem ? j : rem); - uint32_t cnt = qbase + (j < rem ? 1u : 0u); - uint32_t g1 = g0 + cnt; /* exclusive end of this chunk */ - /* Map raw rows that fall in this chunk into shared raw_rows[]. The chunk's - * raw portion is [raw_lo, raw_hi). cnt <= CHUNK and raw rows <= 256, so - * the slice fits raw_rows[256]. */ - uint32_t raw_lo = g0 < raw_count ? g0 : raw_count; - uint32_t raw_hi = g1 < raw_count ? g1 : raw_count; - for (uint32_t r = raw_lo + threadIdx.x; r < raw_hi; r += blockDim.x) { - raw_rows[r - raw_lo] = (raw_start + raw_first_idx + r) % raw_cap; - } - __syncthreads(); - float *pout = partials + (((uint64_t)t * n_head + h) * S + j) * (head_dim + 2u); - /* Pass 1: scores for this chunk's rows into shared scores[0..cnt). */ - float local_max = -INFINITY; - for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) { - uint32_t g = g0 + i; - float s; - if (g < raw_count) { - const float *kvrow = raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale; - } else { - uint32_t c = g - raw_count; - float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; - s = -INFINITY; - if (add > -1.0e20f) { - const float *kvrow = comp_kv + (uint64_t)c * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - s = dot * scale + add; - } - } - scores[i] = s; - local_max = fmaxf(local_max, s); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) m_s = partial[0]; - __syncthreads(); - float chunk_max = m_s; - /* All-masked / empty-chunk guard: never evaluate exp(-INF - -INF) -> NaN. - * Write zero partial (m=-INF, l=0, acc=0) and return. */ - if (!isfinite(chunk_max)) { - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) pout[2u + d] = 0.0f; - if (threadIdx.x == 0) { pout[0] = -INFINITY; pout[1] = 0.0f; } - return; - } - /* Pass 2: exponentiate in place and reduce denominator. */ - float den_local = 0.0f; - for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) { - float e = expf(scores[i] - chunk_max); - scores[i] = e; - den_local += e; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) l_s = partial[0]; - __syncthreads(); - /* Pass 3: weighted value accumulation over this chunk's rows (ascending g), - * preserving raw-then-comp ordering to match the reference accumulation. */ - if (head_dim == 512u && blockDim.x == 256u) { - uint32_t d0 = threadIdx.x; - uint32_t d1 = d0 + 256u; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t i = 0; i < cnt; i++) { - uint32_t g = g0 + i; - float s = scores[i]; - const float *kv = (g < raw_count) - ? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim - : comp_kv + (uint64_t)(g - raw_count) * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - pout[2u + d0] = acc0; - pout[2u + d1] = acc1; - } else { - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t i = 0; i < cnt; i++) { - uint32_t g = g0 + i; - float s = scores[i]; - const float *kv = (g < raw_count) - ? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim - : comp_kv + (uint64_t)(g - raw_count) * head_dim; - acc += kv[d] * s; - } - pout[2u + d] = acc; - } - } - if (threadIdx.x == 0) { - pout[0] = chunk_max; - pout[1] = l_s; - } -} - -__global__ static void attention_decode_splitkv_combine_kernel( - float *heads, - const float *sinks, - const float *partials, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t S) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || h >= n_head) return; - const float *base = partials + (((uint64_t)t * n_head + h) * S) * (head_dim + 2u); - uint32_t stride = head_dim + 2u; - __shared__ float M_s; - __shared__ float L_s; - if (threadIdx.x == 0) { - /* Global max M = max(sink, max_j m_j); sink placed first to match the - * reference (sink seeds local_max). */ - float M = sinks[h]; - for (uint32_t jj = 0; jj < S; jj++) { - float m_j = base[(uint64_t)jj * stride]; - M = fmaxf(M, m_j); /* -INF partials never raise M */ - } - M_s = M; - /* L = Σ_j exp(m_j - M) * l_j + exp(sink - M); sink term added last to - * mirror the reference's denom = Σ scores + expf(sink - max). Chunks - * with l_j == 0 / m_j == -INF contribute exactly 0 (guarded to avoid - * exp(-INF - finite) * 0 edge cases). */ - float L = 0.0f; - for (uint32_t jj = 0; jj < S; jj++) { - float m_j = base[(uint64_t)jj * stride]; - float l_j = base[(uint64_t)jj * stride + 1u]; - if (l_j != 0.0f && isfinite(m_j)) L += expf(m_j - M) * l_j; - } - L += expf(sinks[h] - M); - L_s = L; - } - __syncthreads(); - float M = M_s; - float L = L_s; - float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float A = 0.0f; - for (uint32_t jj = 0; jj < S; jj++) { - float m_j = base[(uint64_t)jj * stride]; - float l_j = base[(uint64_t)jj * stride + 1u]; - if (l_j != 0.0f && isfinite(m_j)) { - A += expf(m_j - M) * base[(uint64_t)jj * stride + 2u + d]; - } - } - oh[d] = A / L; - } -} - -__device__ __forceinline__ void attention_compact_topk_stable( - uint32_t *comp_rows, - uint32_t *comp_count, - uint32_t *warp_offsets, - const int32_t *topk, - uint32_t top_k, - uint32_t visible_comp) { - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t n_warp = blockDim.x >> 5u; - if (threadIdx.x == 0u) *comp_count = 0u; - __syncthreads(); - - for (uint32_t base = 0u; base < 512u; base += blockDim.x) { - const uint32_t i = base + threadIdx.x; - const int32_t c = i < top_k ? topk[i] : -1; - const bool valid = c >= 0 && (uint32_t)c < visible_comp; - const uint32_t mask = __ballot_sync(0xffffffffu, valid); - if (lane == 0u) warp_offsets[warp] = __popc(mask); - __syncthreads(); - if (threadIdx.x == 0u) { - uint32_t out = *comp_count; - for (uint32_t w = 0u; w < n_warp; w++) { - const uint32_t count = warp_offsets[w]; - warp_offsets[w] = out; - out += count; - } - *comp_count = out; - } - __syncthreads(); - if (valid) { - const uint32_t lanes_before = lane == 0u - ? 0u : ((1u << lane) - 1u); - const uint32_t slot = warp_offsets[warp] + - __popc(mask & lanes_before); - if (slot < 512u) comp_rows[slot] = (uint32_t)c; - } - __syncthreads(); - } -} - -__global__ static void attention_indexed_mixed_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const int32_t *topk, - uint32_t n_tokens, - uint32_t pos0, - 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 n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t h = blockIdx.y; - if (t >= n_tokens || h >= n_head) return; - uint32_t qpos = pos0 + t; - uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t visible_comp = n_comp; - if (ratio != 0) { - visible_comp = (qpos + 1u) / ratio; - if (visible_comp > n_comp) visible_comp = n_comp; - } - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - __shared__ float scores[768]; - __shared__ uint32_t raw_rows[256]; - __shared__ uint32_t comp_rows[512]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - __shared__ uint32_t comp_count; - __shared__ uint32_t comp_warp_offsets[8]; - float scale = rsqrtf((float)head_dim); - if (threadIdx.x == 0) { - raw_count = 0; - raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - attention_compact_topk_stable( - comp_rows, &comp_count, comp_warp_offsets, - topk + (uint64_t)t * top_k, top_k, visible_comp); - uint32_t n_score = raw_count + comp_count; - float local_max = sinks[h]; - if (comp_count == 0) { - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; - scores[r] = dot * scale; - local_max = fmaxf(local_max, scores[r]); - } - } else { - uint32_t qlane = threadIdx.x & 7u; - uint32_t qgroup = threadIdx.x >> 3u; - for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { - uint32_t row = row0 + qgroup; - if (row < n_score) { - const float *kvrow = row < raw_count - ? raw_kv + (uint64_t)raw_rows[row] * head_dim - : comp_kv + (uint64_t)comp_rows[row - raw_count] * head_dim; - float dot = 0.0f; - for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; - const uint32_t mask = 0xffu << (threadIdx.x & 24u); - for (uint32_t off = 4u; off > 0u; off >>= 1u) { - dot += __shfl_down_sync(mask, dot, off, 8); - } - if (qlane == 0) scores[row] = dot * scale; - } - } - __syncthreads(); - for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { - local_max = fmaxf(local_max, scores[i]); - } - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; - if (head_dim == 512u && blockDim.x == 256u) { - uint32_t d0 = threadIdx.x; - uint32_t d1 = d0 + 256u; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - for (uint32_t c = 0; c < comp_count; c++) { - float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - oh[d0] = acc0 / denom; - oh[d1] = acc1 / denom; - } else { - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; - for (uint32_t s = 0; s < comp_count; s++) acc += comp_kv[(uint64_t)comp_rows[s] * head_dim + d] * scores[raw_count + s]; - oh[d] = acc / denom; - } - } -} - -__global__ static void attention_indexed_mixed_decode_rows_kernel( - float *heads, - const float *sinks, - const float *q, - cuda_attention_decode_row_table rows, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t row = blockIdx.x; - const uint32_t h = blockIdx.y; - if (row >= n_rows || h >= n_head) return; - const ds4_gpu_attention_decode_row dsc = rows.row[row]; - if (!dsc.indexed) return; - const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; - const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; - const int32_t *topk = (const int32_t *)(uintptr_t)dsc.topk; - const uint32_t qpos = dsc.pos; - const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; - uint32_t visible_comp = dsc.n_comp; - if (dsc.ratio != 0u) { - visible_comp = (qpos + 1u) / dsc.ratio; - if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; - } - const float *qh = q + ((uint64_t)row * n_head + h) * head_dim; - __shared__ float scores[768]; - __shared__ uint32_t raw_rows[256]; - __shared__ uint32_t comp_rows[512]; - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - __shared__ uint32_t comp_count; - __shared__ uint32_t comp_warp_offsets[8]; - const float scale = rsqrtf((float)head_dim); - if (threadIdx.x == 0u) { - raw_count = 0u; - raw_first_idx = 0u; - if (dsc.n_raw != 0u) { - const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; - if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (dsc.window != 0u && qpos + 1u > dsc.window) { - const uint32_t wlo = qpos + 1u - dsc.window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - raw_rows[r] = - (dsc.raw_start + raw_first_idx + r) % dsc.raw_cap; - } - attention_compact_topk_stable( - comp_rows, &comp_count, comp_warp_offsets, - topk, dsc.top_k, visible_comp); - const uint32_t n_score = raw_count + comp_count; - float local_max = sinks[h]; - if (comp_count == 0u) { - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; - float dot = 0.0f; - for (uint32_t dim = 0; dim < head_dim; dim++) { - dot += qh[dim] * kvrow[dim]; - } - scores[r] = dot * scale; - local_max = fmaxf(local_max, scores[r]); - } - } else { - const uint32_t qlane = threadIdx.x & 7u; - const uint32_t qgroup = threadIdx.x >> 3u; - for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { - const uint32_t score_row = row0 + qgroup; - if (score_row < n_score) { - const float *kvrow = score_row < raw_count - ? raw_kv + (uint64_t)raw_rows[score_row] * head_dim - : comp_kv + - (uint64_t)comp_rows[score_row - raw_count] * head_dim; - float dot = 0.0f; - for (uint32_t dim = qlane; dim < head_dim; dim += 8u) { - dot += qh[dim] * kvrow[dim]; - } - const uint32_t mask = 0xffu << (threadIdx.x & 24u); - for (uint32_t off = 4u; off > 0u; off >>= 1u) { - dot += __shfl_down_sync(mask, dot, off, 8); - } - if (qlane == 0u) scores[score_row] = dot * scale; - } - } - __syncthreads(); - for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { - local_max = fmaxf(local_max, scores[i]); - } - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; - stride > 0u; - stride >>= 1u) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = - fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0u) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { - scores[i] = expf(scores[i] - max_s); - den_local += scores[i]; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; - stride > 0u; - stride >>= 1u) { - if (threadIdx.x < stride) { - partial[threadIdx.x] += partial[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0u) { - denom = partial[0] + expf(sinks[h] - max_s); - } - __syncthreads(); - float *oh = heads + ((uint64_t)row * n_head + h) * head_dim; - if (head_dim == 512u && blockDim.x == 256u) { - const uint32_t d0 = threadIdx.x; - const uint32_t d1 = d0 + 256u; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - const float s = scores[r]; - const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - for (uint32_t c = 0; c < comp_count; c++) { - const float s = scores[raw_count + c]; - const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim; - acc0 += kv[d0] * s; - acc1 += kv[d1] * s; - } - oh[d0] = acc0 / denom; - oh[d1] = acc1 / denom; - } else { - for (uint32_t dim = threadIdx.x; - dim < head_dim; - dim += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < raw_count; r++) { - acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] * - scores[r]; - } - for (uint32_t c = 0; c < comp_count; c++) { - acc += comp_kv[(uint64_t)comp_rows[c] * head_dim + dim] * - scores[raw_count + c]; - } - oh[dim] = acc / denom; - } - } -} - -__global__ static void attention_indexed_mixed_heads8_rb4_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const int32_t *topk, - uint32_t n_tokens, - uint32_t pos0, - 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 n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t head_group = blockIdx.y; - if (t >= n_tokens || head_dim != 512u) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t head = head_group * 8u + warp; - const bool valid_head = head < n_head; - - __shared__ uint32_t raw_rows[256]; - __shared__ uint32_t comp_rows[512]; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - __shared__ uint32_t comp_count; - __shared__ float4 kv_shared[4 * 128]; - __shared__ float scores[8 * 768]; - - uint32_t qpos = pos0 + t; - uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t visible_comp = n_comp; - if (ratio != 0) { - visible_comp = (qpos + 1u) / ratio; - if (visible_comp > n_comp) visible_comp = n_comp; - } - - if (threadIdx.x == 0) { - raw_count = 0; - raw_first_idx = 0; - comp_count = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - if (threadIdx.x == 0) { - for (uint32_t i = 0; i < top_k && comp_count < 512u; i++) { - int32_t c = topk[(uint64_t)t * top_k + i]; - if (c >= 0 && (uint32_t)c < visible_comp) comp_rows[comp_count++] = (uint32_t)c; - } - } - __syncthreads(); - - const uint32_t n_score = raw_count + comp_count; - const float scale = rsqrtf((float)head_dim); - const float4 *q4 = valid_head - ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) - : NULL; - float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 q1 = q0, q2 = q0, q3 = q0; - if (valid_head) { - q0 = q4[lane + 0u]; - q1 = q4[lane + 32u]; - q2 = q4[lane + 64u]; - q3 = q4[lane + 96u]; - } - - for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { - const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; - for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { - const uint32_t rr = off >> 7u; - const uint32_t c4 = off & 127u; - const uint32_t sr = row0 + rr; - const float4 *src = sr < raw_count - ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) - : (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim); - kv_shared[off] = src[c4]; - } - __syncthreads(); - if (valid_head) { - for (uint32_t rr = 0; rr < nr; rr++) { - const float4 *kv4 = kv_shared + rr * 128u; - float dot = dot4_f32(q0, kv4[lane + 0u]) + - dot4_f32(q1, kv4[lane + 32u]) + - dot4_f32(q2, kv4[lane + 64u]) + - dot4_f32(q3, kv4[lane + 96u]); - dot = warp_sum_f32(dot); - if (lane == 0) scores[warp * 768u + row0 + rr] = dot * scale; - } - } - __syncthreads(); - } - - float max_s = valid_head ? sinks[head] : -INFINITY; - if (valid_head) { - const float *score_row = scores + warp * 768u; - for (uint32_t i = lane; i < n_score; i += 32u) max_s = fmaxf(max_s, score_row[i]); - max_s = warp_max_f32(max_s); - max_s = __shfl_sync(0xffffffffu, max_s, 0); - } - float den = 0.0f; - if (valid_head) { - float *score_row = scores + warp * 768u; - for (uint32_t i = lane; i < n_score; i += 32u) { - float p = expf(score_row[i] - max_s); - score_row[i] = p; - den += p; - } - den = warp_sum_f32(den); - den += expf(sinks[head] - max_s); - den = __shfl_sync(0xffffffffu, den, 0); - } - - float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 o1 = o0, o2 = o0, o3 = o0; - for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { - const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; - for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { - const uint32_t rr = off >> 7u; - const uint32_t c4 = off & 127u; - const uint32_t sr = row0 + rr; - const float4 *src = sr < raw_count - ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) - : (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim); - kv_shared[off] = src[c4]; - } - __syncthreads(); - if (valid_head) { - const float *score_row = scores + warp * 768u; - for (uint32_t rr = 0; rr < nr; rr++) { - const float p = den == 0.0f ? 0.0f : score_row[row0 + rr] / den; - const float4 *kv4 = kv_shared + rr * 128u; - float4 k0 = kv4[lane + 0u]; - float4 k1 = kv4[lane + 32u]; - float4 k2 = kv4[lane + 64u]; - float4 k3 = kv4[lane + 96u]; - o0.x += k0.x * p; o0.y += k0.y * p; o0.z += k0.z * p; o0.w += k0.w * p; - o1.x += k1.x * p; o1.y += k1.y * p; o1.z += k1.z * p; o1.w += k1.w * p; - o2.x += k2.x * p; o2.y += k2.y * p; o2.z += k2.z * p; o2.w += k2.w * p; - o3.x += k3.x * p; o3.y += k3.y * p; o3.z += k3.z * p; o3.w += k3.w * p; - } - } - __syncthreads(); - } - if (valid_head) { - float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); - out4[lane + 0u] = o0; - out4[lane + 32u] = o1; - out4[lane + 64u] = o2; - out4[lane + 96u] = o3; - } -} - -template -__global__ static void attention_indexed_mixed_heads8_online_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const int32_t *topk, - uint32_t n_tokens, - uint32_t pos0, - 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 n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t head_group = blockIdx.y; - if (t >= n_tokens || head_dim != 512u) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t head = head_group * HEADS_PER_GROUP + warp; - const bool valid_head = head < n_head; - - __shared__ uint32_t raw_rows[256]; - __shared__ uint32_t raw_count; - __shared__ uint32_t raw_first_idx; - __shared__ float4 kv_shared[ROWS_PER_STAGE * 128]; - - uint32_t qpos = pos0 + t; - uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t visible_comp = n_comp; - if (ratio != 0) { - visible_comp = (qpos + 1u) / ratio; - if (visible_comp > n_comp) visible_comp = n_comp; - } - - if (threadIdx.x == 0) { - raw_count = 0; - raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - } - __syncthreads(); - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - __syncthreads(); - - uint32_t comp_count = top_k < visible_comp ? top_k : visible_comp; - if (comp_count > 512u) comp_count = 512u; - const uint32_t n_score = raw_count + comp_count; - const float scale = rsqrtf((float)head_dim); - const float4 *q4 = valid_head - ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) - : NULL; - float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 q1 = q0, q2 = q0, q3 = q0; - if (valid_head) { - q0 = q4[lane + 0u]; - q1 = q4[lane + 32u]; - q2 = q4[lane + 64u]; - q3 = q4[lane + 96u]; - } - - float max_s = -INFINITY; - float sum_s = 0.0f; - float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 o1 = o0, o2 = o0, o3 = o0; - - for (uint32_t row0 = 0; row0 < n_score; row0 += ROWS_PER_STAGE) { - const uint32_t nr = n_score - row0 < ROWS_PER_STAGE ? n_score - row0 : ROWS_PER_STAGE; - for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { - const uint32_t rr = off >> 7u; - const uint32_t c4 = off & 127u; - const uint32_t sr = row0 + rr; - const uint32_t comp_idx = sr < raw_count - ? 0u - : (uint32_t)topk[(uint64_t)t * top_k + (sr - raw_count)]; - const float4 *src = sr < raw_count - ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) - : (const float4 *)(comp_kv + (uint64_t)comp_idx * head_dim); - kv_shared[off] = src[c4]; - } - __syncthreads(); - if (valid_head) { - for (uint32_t rr = 0; rr < nr; rr++) { - const float4 *kv4 = kv_shared + rr * 128u; - float4 k0 = kv4[lane + 0u]; - float4 k1 = kv4[lane + 32u]; - float4 k2 = kv4[lane + 64u]; - float4 k3 = kv4[lane + 96u]; - float score = dot4_f32(q0, k0) + - dot4_f32(q1, k1) + - dot4_f32(q2, k2) + - dot4_f32(q3, k3); - score = warp_sum_f32(score) * scale; - score = __shfl_sync(0xffffffffu, score, 0); - - const float new_m = fmaxf(max_s, score); - const float old_scale = expf(max_s - new_m); - const float row_scale = expf(score - new_m); - sum_s = sum_s * old_scale + row_scale; - o0.x = o0.x * old_scale + k0.x * row_scale; - o0.y = o0.y * old_scale + k0.y * row_scale; - o0.z = o0.z * old_scale + k0.z * row_scale; - o0.w = o0.w * old_scale + k0.w * row_scale; - o1.x = o1.x * old_scale + k1.x * row_scale; - o1.y = o1.y * old_scale + k1.y * row_scale; - o1.z = o1.z * old_scale + k1.z * row_scale; - o1.w = o1.w * old_scale + k1.w * row_scale; - o2.x = o2.x * old_scale + k2.x * row_scale; - o2.y = o2.y * old_scale + k2.y * row_scale; - o2.z = o2.z * old_scale + k2.z * row_scale; - o2.w = o2.w * old_scale + k2.w * row_scale; - o3.x = o3.x * old_scale + k3.x * row_scale; - o3.y = o3.y * old_scale + k3.y * row_scale; - o3.z = o3.z * old_scale + k3.z * row_scale; - o3.w = o3.w * old_scale + k3.w * row_scale; - max_s = new_m; - } - } - __syncthreads(); - } - - if (valid_head) { - const float sink = sinks[head]; - const float new_m = fmaxf(max_s, sink); - const float old_scale = expf(max_s - new_m); - const float sink_scale = expf(sink - new_m); - sum_s = sum_s * old_scale + sink_scale; - o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; - o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; - o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; - o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; - - const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; - o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; - o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; - o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; - o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; - float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); - out4[lane + 0u] = o0; - out4[lane + 32u] = o1; - out4[lane + 64u] = o2; - out4[lane + 96u] = o3; - } -} - -__global__ static void attention_static_mixed_heads8_online_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t head_group = blockIdx.y; - if (t >= n_tokens || head_dim != 512u) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t head = head_group * 8u + warp; - const bool valid_head = head < n_head; - - __shared__ float4 kv_shared[4 * 128]; - - const uint32_t raw_count = window != 0u && t + 1u > window ? window : t + 1u; - const uint32_t raw_start = t + 1u - raw_count; - uint32_t comp_count = 0; - if (n_comp != 0u && ratio != 0u) { - comp_count = (t + 1u) / ratio; - if (comp_count > n_comp) comp_count = n_comp; - } - const uint32_t n_score = raw_count + comp_count; - const float scale = rsqrtf((float)head_dim); - const float4 *q4 = valid_head - ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) - : NULL; - float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 q1 = q0, q2 = q0, q3 = q0; - if (valid_head) { - q0 = q4[lane + 0u]; - q1 = q4[lane + 32u]; - q2 = q4[lane + 64u]; - q3 = q4[lane + 96u]; - } - - float max_s = -INFINITY; - float sum_s = 0.0f; - float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 o1 = o0, o2 = o0, o3 = o0; - - for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { - const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; - for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { - const uint32_t rr = off >> 7u; - const uint32_t c4 = off & 127u; - const uint32_t sr = row0 + rr; - const float4 *src = sr < raw_count - ? (const float4 *)(raw_kv + (uint64_t)(raw_start + sr) * head_dim) - : (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim); - kv_shared[off] = src[c4]; - } - __syncthreads(); - if (valid_head) { - for (uint32_t rr = 0; rr < nr; rr++) { - const float4 *kv4 = kv_shared + rr * 128u; - float4 k0 = kv4[lane + 0u]; - float4 k1 = kv4[lane + 32u]; - float4 k2 = kv4[lane + 64u]; - float4 k3 = kv4[lane + 96u]; - float score = dot4_f32(q0, k0) + - dot4_f32(q1, k1) + - dot4_f32(q2, k2) + - dot4_f32(q3, k3); - score = warp_sum_f32(score) * scale; - score = __shfl_sync(0xffffffffu, score, 0); - - const float new_m = fmaxf(max_s, score); - const float old_scale = expf(max_s - new_m); - const float row_scale = expf(score - new_m); - sum_s = sum_s * old_scale + row_scale; - o0.x = o0.x * old_scale + k0.x * row_scale; - o0.y = o0.y * old_scale + k0.y * row_scale; - o0.z = o0.z * old_scale + k0.z * row_scale; - o0.w = o0.w * old_scale + k0.w * row_scale; - o1.x = o1.x * old_scale + k1.x * row_scale; - o1.y = o1.y * old_scale + k1.y * row_scale; - o1.z = o1.z * old_scale + k1.z * row_scale; - o1.w = o1.w * old_scale + k1.w * row_scale; - o2.x = o2.x * old_scale + k2.x * row_scale; - o2.y = o2.y * old_scale + k2.y * row_scale; - o2.z = o2.z * old_scale + k2.z * row_scale; - o2.w = o2.w * old_scale + k2.w * row_scale; - o3.x = o3.x * old_scale + k3.x * row_scale; - o3.y = o3.y * old_scale + k3.y * row_scale; - o3.z = o3.z * old_scale + k3.z * row_scale; - o3.w = o3.w * old_scale + k3.w * row_scale; - max_s = new_m; - } - } - __syncthreads(); - } - - if (valid_head) { - const float sink = sinks[head]; - const float new_m = fmaxf(max_s, sink); - const float old_scale = expf(max_s - new_m); - const float sink_scale = expf(sink - new_m); - sum_s = sum_s * old_scale + sink_scale; - o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; - o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; - o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; - o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; - - const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; - o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; - o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; - o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; - o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; - float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); - out4[lane + 0u] = o0; - out4[lane + 32u] = o1; - out4[lane + 64u] = o2; - out4[lane + 96u] = o3; - } -} - -__global__ static void attention_decode_mixed_heads8_online_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - uint32_t t = blockIdx.x; - uint32_t head_group = blockIdx.y; - if (t >= n_tokens || head_dim != 512u) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t head = head_group * 8u + warp; - const bool valid_head = head < n_head; - - __shared__ uint32_t raw_rows[256]; - __shared__ uint32_t raw_count_s; - __shared__ uint32_t raw_first_idx_s; - __shared__ float4 kv_shared[4 * 128]; - - const uint32_t qpos = pos0 + t; - const uint32_t first_raw_pos = pos0 + n_tokens - n_raw; - uint32_t comp_count = 0; - if (n_comp != 0u) { - if (n_tokens == 1u && ratio == 0u) { - comp_count = n_comp; - } else if (ratio != 0u) { - comp_count = (qpos + 1u) / ratio; - if (comp_count > n_comp) comp_count = n_comp; - } - } - if (threadIdx.x == 0) { - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0u) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0u && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - raw_count_s = raw_count; - raw_first_idx_s = raw_first_idx; - } - __syncthreads(); - const uint32_t raw_count = raw_count_s; - const uint32_t raw_first_idx = raw_first_idx_s; - for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { - raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; - } - __syncthreads(); - - const uint32_t n_score = raw_count + comp_count; - const float scale = rsqrtf((float)head_dim); - const float4 *q4 = valid_head - ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) - : NULL; - float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 q1 = q0, q2 = q0, q3 = q0; - if (valid_head) { - q0 = q4[lane + 0u]; - q1 = q4[lane + 32u]; - q2 = q4[lane + 64u]; - q3 = q4[lane + 96u]; - } - - float max_s = -INFINITY; - float sum_s = 0.0f; - float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - float4 o1 = o0, o2 = o0, o3 = o0; - - for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { - const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; - for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { - const uint32_t rr = off >> 7u; - const uint32_t c4 = off & 127u; - const uint32_t sr = row0 + rr; - const float4 *src = sr < raw_count - ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) - : (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim); - kv_shared[off] = src[c4]; - } - __syncthreads(); - if (valid_head) { - for (uint32_t rr = 0; rr < nr; rr++) { - const float4 *kv4 = kv_shared + rr * 128u; - float4 k0 = kv4[lane + 0u]; - float4 k1 = kv4[lane + 32u]; - float4 k2 = kv4[lane + 64u]; - float4 k3 = kv4[lane + 96u]; - float score = dot4_f32(q0, k0) + - dot4_f32(q1, k1) + - dot4_f32(q2, k2) + - dot4_f32(q3, k3); - score = warp_sum_f32(score) * scale; - score = __shfl_sync(0xffffffffu, score, 0); - - const float new_m = fmaxf(max_s, score); - const float old_scale = expf(max_s - new_m); - const float row_scale = expf(score - new_m); - sum_s = sum_s * old_scale + row_scale; - o0.x = o0.x * old_scale + k0.x * row_scale; - o0.y = o0.y * old_scale + k0.y * row_scale; - o0.z = o0.z * old_scale + k0.z * row_scale; - o0.w = o0.w * old_scale + k0.w * row_scale; - o1.x = o1.x * old_scale + k1.x * row_scale; - o1.y = o1.y * old_scale + k1.y * row_scale; - o1.z = o1.z * old_scale + k1.z * row_scale; - o1.w = o1.w * old_scale + k1.w * row_scale; - o2.x = o2.x * old_scale + k2.x * row_scale; - o2.y = o2.y * old_scale + k2.y * row_scale; - o2.z = o2.z * old_scale + k2.z * row_scale; - o2.w = o2.w * old_scale + k2.w * row_scale; - o3.x = o3.x * old_scale + k3.x * row_scale; - o3.y = o3.y * old_scale + k3.y * row_scale; - o3.z = o3.z * old_scale + k3.z * row_scale; - o3.w = o3.w * old_scale + k3.w * row_scale; - max_s = new_m; - } - } - __syncthreads(); - } - - if (valid_head) { - const float sink = sinks[head]; - const float new_m = fmaxf(max_s, sink); - const float old_scale = expf(max_s - new_m); - const float sink_scale = expf(sink - new_m); - sum_s = sum_s * old_scale + sink_scale; - o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; - o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; - o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; - o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; - - const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; - o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; - o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; - o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; - o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; - float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); - out4[lane + 0u] = o0; - out4[lane + 32u] = o1; - out4[lane + 64u] = o2; - out4[lane + 96u] = o3; - } -} - -__device__ static void hc4_split_one(float *out, const float *mix, const float *scale, const float *base, uint32_t sinkhorn_iters, float epsv) { - const float pre_scale = scale[0]; - const float post_scale = scale[1]; - const float comb_scale = scale[2]; - for (int i = 0; i < 4; i++) { - float z = mix[i] * pre_scale + base[i]; - out[i] = 1.0f / (1.0f + expf(-z)) + epsv; - } - for (int i = 0; i < 4; i++) { - float z = mix[4 + i] * post_scale + base[4 + i]; - out[4 + i] = 2.0f / (1.0f + expf(-z)); - } - float c[16]; - for (int r = 0; r < 4; r++) { - float m = -INFINITY; - for (int col = 0; col < 4; col++) { - float v = mix[8 + r * 4 + col] * comb_scale + base[8 + r * 4 + col]; - c[r * 4 + col] = v; - m = fmaxf(m, v); - } - float s = 0.0f; - for (int col = 0; col < 4; col++) { - float v = expf(c[r * 4 + col] - m); - c[r * 4 + col] = v; - s += v; - } - for (int col = 0; col < 4; col++) c[r * 4 + col] = c[r * 4 + col] / s + epsv; - } - for (int col = 0; col < 4; col++) { - float s = epsv; - for (int r = 0; r < 4; r++) s += c[r * 4 + col]; - for (int r = 0; r < 4; r++) c[r * 4 + col] /= s; - } - for (uint32_t iter = 1; iter < sinkhorn_iters; iter++) { - for (int r = 0; r < 4; r++) { - float s = epsv; - for (int col = 0; col < 4; col++) s += c[r * 4 + col]; - for (int col = 0; col < 4; col++) c[r * 4 + col] /= s; - } - for (int col = 0; col < 4; col++) { - float s = epsv; - for (int r = 0; r < 4; r++) s += c[r * 4 + col]; - for (int r = 0; r < 4; r++) c[r * 4 + col] /= s; - } - } - for (int i = 0; i < 16; i++) out[8 + i] = c[i]; -} - -__global__ static void hc_split_sinkhorn_kernel(float *out, const float *mix, const float *scale, const float *base, uint32_t n_rows, uint32_t sinkhorn_iters, float epsv) { - uint32_t row = blockIdx.x * blockDim.x + threadIdx.x; - if (row >= n_rows) return; - hc4_split_one(out + (uint64_t)row * 24, mix + (uint64_t)row * 24, scale, base, sinkhorn_iters, epsv); -} - -__global__ static void hc_weighted_sum_kernel(float *out, const float *x, const float *w, uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens, uint32_t weight_stride_f32) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_embd * n_tokens; - if (gid >= n) return; - uint32_t d = gid % n_embd; - uint32_t t = gid / n_embd; - float acc = 0.0f; - for (uint32_t h = 0; h < n_hc; h++) { - acc += x[(uint64_t)t * n_hc * n_embd + (uint64_t)h * n_embd + d] * - w[(uint64_t)t * weight_stride_f32 + h]; - } - out[(uint64_t)t * n_embd + d] = acc; -} - -__global__ static void hc_expand_kernel( - float *out_hc, - const float *block_out, - const float *block_add, - const float *block_add2, - const float *residual_hc, - const float *post, - const float *comb, - uint32_t n_embd, - uint32_t n_hc, - uint32_t n_tokens, - uint32_t post_stride, - uint32_t comb_stride, - int has_add, - int has_add2) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - if (gid >= n_elem) return; - uint32_t d = gid % n_embd; - uint64_t tmp = gid / n_embd; - uint32_t dst_hc = tmp % n_hc; - uint32_t t = tmp / n_hc; - - float block_v = block_out[(uint64_t)t * n_embd + d]; - if (has_add) { - float add_v = block_add[(uint64_t)t * n_embd + d]; - if (has_add2) add_v += block_add2[(uint64_t)t * n_embd + d]; - block_v += add_v; - } - float acc = block_v * post[(uint64_t)t * post_stride + dst_hc]; - for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) { - float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc]; - float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d]; - acc += comb_v * res_v; - } - out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc; -} - -__global__ static void hc_split_weighted_sum_fused_kernel( - float *out, - float *split, - const float *mix, - const float *residual_hc, - const float *scale, - const float *base, - uint32_t n_embd, - uint32_t n_hc, - uint32_t n_rows, - uint32_t sinkhorn_iters, - float epsv) { - uint32_t t = blockIdx.x; - uint32_t d = threadIdx.x; - if (t >= n_rows || n_hc != 4) return; - const uint32_t mix_hc = 24; - float *sp = split + (uint64_t)t * mix_hc; - if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv); - __syncthreads(); - for (uint32_t col = d; col < n_embd; col += blockDim.x) { - float acc = 0.0f; - for (uint32_t h = 0; h < 4; h++) { - acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h]; - } - out[(uint64_t)t * n_embd + col] = acc; - } -} - -__global__ static void hc_split_weighted_sum_norm_fused_kernel( - float *out, - float *norm_out, - float *split, - const float *mix, - const float *residual_hc, - const float *scale, - const float *base, - const float *norm_w, - uint32_t n_embd, - uint32_t n_hc, - uint32_t n_rows, - uint32_t sinkhorn_iters, - float epsv, - float norm_eps) { - const uint32_t t = blockIdx.x; - const uint32_t d = threadIdx.x; - if (t >= n_rows || n_hc != 4) return; - const uint32_t mix_hc = 24; - float *sp = split + (uint64_t)t * mix_hc; - if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv); - __syncthreads(); - - float sum = 0.0f; - for (uint32_t col = d; col < n_embd; col += blockDim.x) { - float acc = 0.0f; - for (uint32_t h = 0; h < 4; h++) { - acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h]; - } - out[(uint64_t)t * n_embd + col] = acc; - sum += acc * acc; - } - - __shared__ float partial[256]; - partial[d] = sum; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (d < stride) partial[d] += partial[d + stride]; - __syncthreads(); - } - const float norm_scale = rsqrtf(partial[0] / (float)n_embd + norm_eps); - for (uint32_t col = d; col < n_embd; col += blockDim.x) { - const float v = out[(uint64_t)t * n_embd + col]; - norm_out[(uint64_t)t * n_embd + col] = v * norm_scale * norm_w[col]; - } -} - -__global__ static void output_hc_weights_kernel( - float *out, - const float *pre, - const float *scale, - const float *base, - uint32_t n_hc, - uint32_t n_tokens, - float epsv) { - uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; - uint32_t n = n_tokens * n_hc; - if (gid >= n) return; - uint32_t h = gid % n_hc; - float z = pre[gid] * scale[0] + base[h]; - out[gid] = 1.0f / (1.0f + expf(-z)) + epsv; -} - -__global__ static void fill_f32_kernel(float *x, uint64_t n, float v) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) x[i] = v; -} - -__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) { - 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); -} - -__global__ static void compressor_set_rows_kernel( - float *state_kv, - float *state_score, - const float *kv, - const float *sc, - const void *model_map, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t width, - uint32_t ratio, - uint32_t pos0, - uint32_t src0, - uint32_t dst0, - uint32_t rows) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)rows * width; - if (gid >= n) return; - uint32_t r = gid / width; - uint32_t j = gid - (uint64_t)r * width; - uint32_t src = src0 + r; - uint32_t dst = dst0 + r; - uint32_t phase = (pos0 + src) % ratio; - state_kv[(uint64_t)dst * width + j] = kv[(uint64_t)src * width + j]; - state_score[(uint64_t)dst * width + j] = - sc[(uint64_t)src * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)phase * width + j); -} - -__global__ static void compressor_prefill_pool_kernel( - float *comp, - const float *kv, - const float *sc, - const float *state_kv, - const 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_comp, - uint32_t replay) { - uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; - uint32_t c = blockIdx.y; - if (d >= head_dim || c >= n_comp) return; - uint32_t coff = ratio == 4u ? 2u : 1u; - uint32_t width = coff * head_dim; - float vals[128]; - float scores[128]; - float max_s = -INFINITY; - uint32_t n_cand = 0; - if (ratio == 4u) { - if (replay && c == 0) { - for (uint32_t r = 0; r < 4; r++) { - vals[n_cand] = state_kv[(uint64_t)r * width + d]; - scores[n_cand] = state_score[(uint64_t)r * width + d]; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } else if (c > 0) { - uint32_t base = (c - 1u) * ratio; - for (uint32_t r = 0; r < 4; r++) { - uint32_t t = base + r; - float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d); - vals[n_cand] = kv[(uint64_t)t * width + d]; - scores[n_cand] = sc[(uint64_t)t * width + d] + ape; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } - uint32_t base = c * ratio; - for (uint32_t r = 0; r < 4; r++) { - uint32_t t = base + r; - float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + head_dim + d); - vals[n_cand] = kv[(uint64_t)t * width + head_dim + d]; - scores[n_cand] = sc[(uint64_t)t * width + head_dim + d] + ape; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } else { - uint32_t base = c * ratio; - for (uint32_t r = 0; r < ratio; r++) { - uint32_t t = base + r; - float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d); - vals[n_cand] = kv[(uint64_t)t * width + d]; - scores[n_cand] = sc[(uint64_t)t * width + d] + ape; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } - float den = 0.0f, acc = 0.0f; - for (uint32_t i = 0; i < n_cand; i++) { - float w = expf(scores[i] - max_s); - den += w; - acc += vals[i] * w; - } - comp[(uint64_t)c * head_dim + d] = den != 0.0f ? acc / den : 0.0f; -} - -__global__ static void compressor_update_pool_kernel( - float *row, - const float *state_kv, - const float *state_score, - uint32_t head_dim, - uint32_t ratio) { - uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; - if (d >= head_dim) return; - uint32_t coff = ratio == 4u ? 2u : 1u; - uint32_t width = coff * head_dim; - float vals[128]; - float scores[128]; - float max_s = -INFINITY; - uint32_t n_cand = 0; - if (ratio == 4u) { - for (uint32_t r = 0; r < 4; r++) { - vals[n_cand] = state_kv[(uint64_t)r * width + d]; - scores[n_cand] = state_score[(uint64_t)r * width + d]; - max_s = fmaxf(max_s, scores[n_cand++]); - } - for (uint32_t r = 0; r < 4; r++) { - vals[n_cand] = state_kv[(uint64_t)(ratio + r) * width + head_dim + d]; - scores[n_cand] = state_score[(uint64_t)(ratio + r) * width + head_dim + d]; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } else { - for (uint32_t r = 0; r < ratio; r++) { - vals[n_cand] = state_kv[(uint64_t)r * width + d]; - scores[n_cand] = state_score[(uint64_t)r * width + d]; - max_s = fmaxf(max_s, scores[n_cand++]); - } - } - float den = 0.0f, acc = 0.0f; - for (uint32_t i = 0; i < n_cand; i++) { - float w = expf(scores[i] - max_s); - den += w; - acc += vals[i] * w; - } - row[d] = den != 0.0f ? acc / den : 0.0f; -} - -__global__ static void compressor_shift_ratio4_kernel(float *state_kv, float *state_score, uint32_t width) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t half = 4ull * width; - if (i >= half) return; - float v = state_kv[half + i]; - float s = state_score[half + i]; - state_kv[i] = v; - state_score[i] = s; - state_kv[half + i] = v; - state_score[half + i] = 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)); -} - -__global__ static void router_select_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) { - uint32_t t = blockIdx.x; - if (t >= n_tokens || threadIdx.x != 0) return; - const float *log = logits + (uint64_t)t * 256; - float *prob = probs + (uint64_t)t * 256; - int32_t *sel = selected + (uint64_t)t * 6; - float *w = weights + (uint64_t)t * 6; - - for (int i = 0; i < 256; i++) prob[i] = sqrtf(softplus_dev(log[i])); - - if (hash_mode) { - int32_t tok = tokens ? tokens[t] : token_scalar; - if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0; - const int32_t *row = hash + (uint64_t)tok * 6; - for (int i = 0; i < 6; i++) sel[i] = row[i]; - } else { - for (int i = 0; i < 6; i++) sel[i] = -1; - for (int i = 0; i < 256; i++) { - float score = prob[i] + (has_bias ? bias[i] : 0.0f); - for (int j = 0; j < 6; j++) { - if (sel[j] < 0 || score > prob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) { - for (int k = 5; k > j; k--) sel[k] = sel[k - 1]; - sel[j] = i; - break; - } - } - } - } - - float sum = 0.0f; - for (int i = 0; i < 6; i++) { - int e = sel[i]; - float v = (e >= 0 && e < 256) ? prob[e] : 0.0f; - w[i] = v; - sum += v; - } - sum = fmaxf(sum, 6.103515625e-5f); - for (int i = 0; i < 6; i++) w[i] = w[i] / sum * 1.5f; -} - -__global__ static void router_select_parallel_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) { - uint32_t t = blockIdx.x; - uint32_t i = threadIdx.x; - if (t >= n_tokens || i >= 256u) return; - const float *log = logits + (uint64_t)t * 256; - float *prob = probs + (uint64_t)t * 256; - int32_t *sel = selected + (uint64_t)t * 6; - float *w = weights + (uint64_t)t * 6; - __shared__ float sprob[256]; - - const float p = sqrtf(softplus_dev(log[i])); - sprob[i] = p; - prob[i] = p; - __syncthreads(); - - if (i != 0) return; - if (hash_mode) { - int32_t tok = tokens ? tokens[t] : token_scalar; - if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0; - const int32_t *row = hash + (uint64_t)tok * 6; - for (int j = 0; j < 6; j++) sel[j] = row[j]; - } else { - for (int j = 0; j < 6; j++) sel[j] = -1; - for (int e = 0; e < 256; e++) { - float score = sprob[e] + (has_bias ? bias[e] : 0.0f); - for (int j = 0; j < 6; j++) { - if (sel[j] < 0 || score > sprob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) { - for (int k = 5; k > j; k--) sel[k] = sel[k - 1]; - sel[j] = e; - break; - } - } - } - } - - float sum = 0.0f; - for (int j = 0; j < 6; j++) { - int e = sel[j]; - float v = (e >= 0 && e < 256) ? sprob[e] : 0.0f; - w[j] = v; - sum += v; - } - sum = fmaxf(sum, 6.103515625e-5f); - for (int j = 0; j < 6; j++) w[j] = w[j] / sum * 1.5f; -} - -__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 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] : 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; - } -} - -__global__ static void swiglu_kernel(float *out, const float *gate, const float *up, uint32_t n, float clamp, float weight) { - uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - float g = gate[i]; - float u = up[i]; - if (clamp > 1.0e-6f) { - g = fminf(g, clamp); - u = fminf(fmaxf(u, -clamp), clamp); - } - float s = g / (1.0f + expf(-g)); - out[i] = s * u * weight; -} - -__global__ static void add_kernel(float *out, const float *a, const float *b, uint32_t n) { - uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - out[i] = a[i] + b[i]; -} - -__global__ static void directional_steering_project_kernel( - float *x, - const float *directions, - uint32_t layer, - uint32_t width, - uint32_t rows, - float scale) { - const uint32_t row = blockIdx.x; - if (row >= rows || width == 0) return; - - float *xr = x + (uint64_t)row * width; - const float *dir = directions + (uint64_t)layer * width; - float sum = 0.0f; - for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) { - sum += xr[i] * dir[i]; - } - - __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 coeff = scale * partial[0]; - for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) { - xr[i] -= coeff * dir[i]; - } -} - -__global__ static void zero_kernel(float *out, uint64_t n) { - uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) out[i] = 0.0f; -} - -__global__ static void indexer_scores_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - int causal) { - uint32_t c = blockIdx.x; - uint32_t t = blockIdx.y; - if (c >= n_comp || t >= n_tokens) return; - if (causal) { - uint32_t n_visible = (pos0 + t + 1u) / ratio; - if (c >= n_visible) { - if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = -INFINITY; - return; - } - } - float total = 0.0f; - for (uint32_t h = 0; h < n_head; h++) { - const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; - const float *kh = index_comp + (uint64_t)c * head_dim; - float dot = 0.0f; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) dot += qh[d] * kh[d]; - __shared__ float partial[256]; - partial[threadIdx.x] = dot; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - total += fmaxf(partial[0], 0.0f) * weights[(uint64_t)t * n_head + h]; - __syncthreads(); - } - if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = total * scale; -} - -__global__ static void indexer_score_one_direct_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t pos0, - uint32_t ratio, - float scale, - int causal) { - const uint32_t c = blockIdx.x; - const uint32_t tid = threadIdx.x; - const uint32_t lane = tid & 31u; - const uint32_t warp = tid >> 5u; - if (c >= n_comp || tid >= 128u) return; - if (causal) { - const uint32_t visible = ratio ? (pos0 + 1u) / ratio : n_comp; - if (c >= visible) { - if (tid == 0) scores[c] = -INFINITY; - return; - } - } - - __shared__ float krow[128]; - __shared__ float partial[4]; - if (tid < 128u) krow[tid] = index_comp[(uint64_t)c * 128u + tid]; - __syncthreads(); - - float total = 0.0f; - for (uint32_t h0 = 0; h0 < 64u; h0 += 4u) { - const uint32_t h = h0 + warp; - const float4 qv = ((const float4 *)(q + (uint64_t)h * 128u))[lane]; - const float4 kv = ((const float4 *)krow)[lane]; - float dot = qv.x * kv.x + qv.y * kv.y + qv.z * kv.z + qv.w * kv.w; - dot = warp_sum_f32(dot); - if (lane == 0) partial[warp] = fmaxf(dot, 0.0f) * weights[h] * scale; - __syncthreads(); - if (tid == 0) total += partial[0] + partial[1] + partial[2] + partial[3]; - __syncthreads(); - } - if (tid == 0) scores[c] = total; -} - -__global__ static void indexer_scores_wmma_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - int causal) { -#if __CUDA_ARCH__ >= 700 - namespace wmma = nvcuda::wmma; - const uint32_t tile_c = blockIdx.x * 16u; - const uint32_t tile_t = blockIdx.y * 16u; - const uint32_t tid = threadIdx.x; - if (tid >= 32u || head_dim != 128u) return; - - if (causal) { - const uint32_t last_token = min(tile_t + 16u, n_tokens); - const uint32_t max_visible = last_token > tile_t - ? min((pos0 + last_token) / ratio, n_comp) - : 0u; - if (tile_c >= max_visible) { - for (uint32_t i = tid; i < 16u * 16u; i += 32u) { - const uint32_t r = i >> 4u; - const uint32_t c = i & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + c; - if (token < n_tokens && comp < n_comp) { - scores[(uint64_t)token * n_comp + comp] = -INFINITY; - } - } - return; - } - } - - __shared__ __half a_sh[16 * 128]; - __shared__ __half b_sh[16 * 128]; - __shared__ float c_sh[16 * 16]; - __shared__ float acc_sh[16 * 16]; - - for (uint32_t i = tid; i < 16u * 16u; i += 32u) acc_sh[i] = 0.0f; - for (uint32_t i = tid; i < 16u * 128u; i += 32u) { - const uint32_t c = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t comp = tile_c + c; - float v = 0.0f; - if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; - b_sh[d + c * 128u] = __float2half(v); - } - __syncthreads(); - - for (uint32_t h = 0; h < n_head; h++) { - for (uint32_t i = tid; i < 16u * 128u; i += 32u) { - const uint32_t r = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t token = tile_t + r; - float v = 0.0f; - if (token < n_tokens) { - v = q[((uint64_t)token * n_head + h) * head_dim + d]; - } - a_sh[i] = __float2half(v); - } - __syncthreads(); - - wmma::fragment a_frag; - wmma::fragment b_frag; - wmma::fragment c_frag; - wmma::fill_fragment(c_frag, 0.0f); - for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { - wmma::load_matrix_sync(a_frag, a_sh + k0, 128); - wmma::load_matrix_sync(b_frag, b_sh + k0, 128); - wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); - } - wmma::store_matrix_sync(c_sh, c_frag, 16, wmma::mem_row_major); - __syncthreads(); - - for (uint32_t i = tid; i < 16u * 16u; i += 32u) { - const uint32_t r = i >> 4u; - const uint32_t token = tile_t + r; - if (token < n_tokens) { - const float w = weights[(uint64_t)token * n_head + h]; - acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; - } - } - __syncthreads(); - } - - for (uint32_t i = tid; i < 16u * 16u; i += 32u) { - const uint32_t r = i >> 4u; - const uint32_t c = i & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + c; - if (token < n_tokens && comp < n_comp) { - float out = acc_sh[i] * scale; - if (causal) { - const uint32_t visible = (pos0 + token + 1u) / ratio; - if (comp >= visible) out = -INFINITY; - } - scores[(uint64_t)token * n_comp + comp] = out; - } - } -#endif -} - -__global__ static void indexer_scores_wmma32_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - int causal) { -#if __CUDA_ARCH__ >= 700 - namespace wmma = nvcuda::wmma; - const uint32_t tile_c = blockIdx.x * 32u; - const uint32_t tile_t = blockIdx.y * 16u; - const uint32_t tid = threadIdx.x; - const uint32_t warp = tid >> 5u; - if (tid >= 64u || head_dim != 128u) return; - - if (causal) { - const uint32_t last_token = min(tile_t + 16u, n_tokens); - const uint32_t max_visible = last_token > tile_t - ? min((pos0 + last_token) / ratio, n_comp) - : 0u; - if (tile_c >= max_visible) { - for (uint32_t i = tid; i < 16u * 32u; i += 64u) { - const uint32_t r = i >> 5u; - const uint32_t c = i & 31u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + c; - if (token < n_tokens && comp < n_comp) { - scores[(uint64_t)token * n_comp + comp] = -INFINITY; - } - } - return; - } - } - - __shared__ __half a_sh[16 * 128]; - __shared__ __half b_sh[32 * 128]; - __shared__ float c_sh[2 * 16 * 16]; - __shared__ float acc_sh[2 * 16 * 16]; - - for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) acc_sh[i] = 0.0f; - for (uint32_t i = tid; i < 32u * 128u; i += 64u) { - const uint32_t c = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t comp = tile_c + c; - float v = 0.0f; - if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; - b_sh[d + c * 128u] = __float2half(v); - } - __syncthreads(); - - for (uint32_t h = 0; h < n_head; h++) { - for (uint32_t i = tid; i < 16u * 128u; i += 64u) { - const uint32_t r = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t token = tile_t + r; - float v = 0.0f; - if (token < n_tokens) { - v = q[((uint64_t)token * n_head + h) * head_dim + d]; - } - a_sh[i] = __float2half(v); - } - __syncthreads(); - - wmma::fragment a_frag; - wmma::fragment b_frag; - wmma::fragment c_frag; - wmma::fill_fragment(c_frag, 0.0f); - const uint32_t col0 = warp * 16u; - for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { - wmma::load_matrix_sync(a_frag, a_sh + k0, 128); - wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); - wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); - } - wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); - __syncthreads(); - - for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - const float w = weights[(uint64_t)token * n_head + h]; - acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; - } - } - __syncthreads(); - } - - for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - float out = acc_sh[i] * scale; - if (causal) { - const uint32_t visible = (pos0 + token + 1u) / ratio; - if (comp >= visible) out = -INFINITY; - } - scores[(uint64_t)token * n_comp + comp] = out; - } - } -#endif -} - -__global__ static void indexer_scores_wmma64_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - int causal) { -#if __CUDA_ARCH__ >= 700 - namespace wmma = nvcuda::wmma; - const uint32_t tile_c = blockIdx.x * 64u; - const uint32_t tile_t = blockIdx.y * 16u; - const uint32_t tid = threadIdx.x; - const uint32_t warp = tid >> 5u; - if (tid >= 128u || head_dim != 128u) return; - - if (causal) { - const uint32_t last_token = min(tile_t + 16u, n_tokens); - const uint32_t max_visible = last_token > tile_t - ? min((pos0 + last_token) / ratio, n_comp) - : 0u; - if (tile_c >= max_visible) { - for (uint32_t i = tid; i < 16u * 64u; i += 128u) { - const uint32_t r = i >> 6u; - const uint32_t c = i & 63u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + c; - if (token < n_tokens && comp < n_comp) { - scores[(uint64_t)token * n_comp + comp] = -INFINITY; - } - } - return; - } - } - - __shared__ __half a_sh[16 * 128]; - __shared__ __half b_sh[64 * 128]; - __shared__ float c_sh[4 * 16 * 16]; - __shared__ float acc_sh[4 * 16 * 16]; - - for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) acc_sh[i] = 0.0f; - for (uint32_t i = tid; i < 64u * 128u; i += 128u) { - const uint32_t c = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t comp = tile_c + c; - float v = 0.0f; - if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; - b_sh[d + c * 128u] = __float2half(v); - } - __syncthreads(); - - for (uint32_t h = 0; h < n_head; h++) { - for (uint32_t i = tid; i < 16u * 128u; i += 128u) { - const uint32_t r = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t token = tile_t + r; - float v = 0.0f; - if (token < n_tokens) { - v = q[((uint64_t)token * n_head + h) * head_dim + d]; - } - a_sh[i] = __float2half(v); - } - __syncthreads(); - - wmma::fragment a_frag; - wmma::fragment b_frag; - wmma::fragment c_frag; - wmma::fill_fragment(c_frag, 0.0f); - const uint32_t col0 = warp * 16u; - for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { - wmma::load_matrix_sync(a_frag, a_sh + k0, 128); - wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); - wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); - } - wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); - __syncthreads(); - - for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - const float w = weights[(uint64_t)token * n_head + h]; - acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; - } - } - __syncthreads(); - } - - for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - float out = acc_sh[i] * scale; - if (causal) { - const uint32_t visible = (pos0 + token + 1u) / ratio; - if (comp >= visible) out = -INFINITY; - } - scores[(uint64_t)token * n_comp + comp] = out; - } - } -#endif -} - -__global__ static void indexer_scores_wmma128_kernel( - float *scores, - const float *q, - const float *weights, - const float *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - int causal) { -#if __CUDA_ARCH__ >= 700 - namespace wmma = nvcuda::wmma; - const uint32_t tile_c = blockIdx.x * 128u; - const uint32_t tile_t = blockIdx.y * 16u; - const uint32_t tid = threadIdx.x; - const uint32_t warp = tid >> 5u; - if (tid >= 256u || head_dim != 128u) return; - - if (causal) { - const uint32_t last_token = min(tile_t + 16u, n_tokens); - const uint32_t max_visible = last_token > tile_t - ? min((pos0 + last_token) / ratio, n_comp) - : 0u; - if (tile_c >= max_visible) { - for (uint32_t i = tid; i < 16u * 128u; i += 256u) { - const uint32_t r = i >> 7u; - const uint32_t c = i & 127u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + c; - if (token < n_tokens && comp < n_comp) { - scores[(uint64_t)token * n_comp + comp] = -INFINITY; - } - } - return; - } - } - - __shared__ __half a_sh[16 * 128]; - __shared__ __half b_sh[128 * 128]; - __shared__ float c_sh[8 * 16 * 16]; - - float acc[8]; -#pragma unroll - for (uint32_t i = 0; i < 8u; i++) acc[i] = 0.0f; - - for (uint32_t i = tid; i < 128u * 128u; i += 256u) { - const uint32_t c = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t comp = tile_c + c; - float v = 0.0f; - if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; - b_sh[d + c * 128u] = __float2half(v); - } - __syncthreads(); - - for (uint32_t h = 0; h < n_head; h++) { - for (uint32_t i = tid; i < 16u * 128u; i += 256u) { - const uint32_t r = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t token = tile_t + r; - float v = 0.0f; - if (token < n_tokens) { - v = q[((uint64_t)token * n_head + h) * head_dim + d]; - } - a_sh[i] = __float2half(v); - } - __syncthreads(); - - wmma::fragment a_frag; - wmma::fragment b_frag; - wmma::fragment c_frag; - wmma::fill_fragment(c_frag, 0.0f); - const uint32_t col0 = warp * 16u; - for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { - wmma::load_matrix_sync(a_frag, a_sh + k0, 128); - wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); - wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); - } - wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); - __syncthreads(); - - const uint32_t local0 = tid & 255u; - const uint32_t token0 = tile_t + (local0 >> 4u); - const float w0 = token0 < n_tokens ? weights[(uint64_t)token0 * n_head + h] : 0.0f; - uint32_t slot = 0; - for (uint32_t i = tid; i < 8u * 16u * 16u; i += 256u, slot++) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - acc[slot] += fmaxf(c_sh[i], 0.0f) * w0; - } - } - __syncthreads(); - } - - uint32_t slot = 0; - for (uint32_t i = tid; i < 8u * 16u * 16u; i += 256u, slot++) { - const uint32_t wtile = i >> 8u; - const uint32_t local = i & 255u; - const uint32_t r = local >> 4u; - const uint32_t c = local & 15u; - const uint32_t token = tile_t + r; - const uint32_t comp = tile_c + wtile * 16u + c; - if (token < n_tokens && comp < n_comp) { - float out = acc[slot] * scale; - if (causal) { - const uint32_t visible = (pos0 + token + 1u) / ratio; - if (comp >= visible) out = -INFINITY; - } - scores[(uint64_t)token * n_comp + comp] = out; - } - } -#endif -} - -__global__ static void indexer_topk_kernel(uint32_t *selected, const float *scores, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) { - uint32_t t = blockIdx.x; - if (t >= n_tokens || threadIdx.x != 0) return; - const float *row = scores + (uint64_t)t * n_comp; - uint32_t *sel = selected + (uint64_t)t * top_k; - for (uint32_t k = 0; k < top_k; k++) sel[k] = 0; - for (uint32_t c = 0; c < n_comp; c++) { - float v = row[c]; - for (uint32_t k = 0; k < top_k; k++) { - if ((k >= c) || v > row[sel[k]]) { - for (uint32_t j = top_k - 1; j > k; j--) sel[j] = sel[j - 1]; - sel[k] = c; - break; - } - } - } -} - -__device__ __forceinline__ static bool topk_score_better(float av, uint32_t ai, float bv, uint32_t bi) { - return av > bv || (av == bv && ai < bi); -} - -__device__ __forceinline__ static void top2_insert_candidate( - float v, - uint32_t i, - float *v0, - uint32_t *i0, - float *v1, - uint32_t *i1) { - if (i == *i0 || i == *i1) return; - if (topk_score_better(v, i, *v0, *i0)) { - *v1 = *v0; - *i1 = *i0; - *v0 = v; - *i0 = i; - } else if (topk_score_better(v, i, *v1, *i1)) { - *v1 = v; - *i1 = i; - } -} - -/* DSpark markov chain step: out = argmax_i(logits[i] + dot(w2[i], w1[prev])) - * over the vocab, entirely on-device (logits row stays resident; the chain - * loop only reads back 4 bytes per draft). w1/w2 are q8_0 with 272-byte rows - * (8 blocks of 32). Single block; ~35 MB w2 read per step. */ -__global__ static void dspark_markov_argmax_kernel( - unsigned long long *out_key, - const float *logits, - const unsigned char *w1_row, - const unsigned char *w2, - uint32_t vocab, - uint32_t rank_blocks) { - __shared__ float state[256]; - const uint32_t tid = threadIdx.x; - if (tid < rank_blocks * 32u) { - const uint32_t b = tid >> 5, k = tid & 31u; - const unsigned char *blk = w1_row + (uint64_t)b * 34u; - const float d = __half2float(*(const __half *)blk); - state[tid] = d * (float)((const int8_t *)(blk + 2))[k]; - } - __syncthreads(); - - float best_v = -INFINITY; - uint32_t best_i = 0; - for (uint32_t i = blockIdx.x * blockDim.x + tid; i < vocab; - i += gridDim.x * blockDim.x) { - const unsigned char *row = w2 + (uint64_t)i * rank_blocks * 34u; - float acc = 0.0f; - for (uint32_t b = 0; b < rank_blocks; b++) { - const unsigned char *blk = row + (uint64_t)b * 34u; - const float d = __half2float(*(const __half *)blk); - const int8_t *q = (const int8_t *)(blk + 2); - float s = 0.0f; - #pragma unroll - for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * state[b * 32u + k]; - acc += d * s; - } - const float v = logits[i] + acc; - if (topk_score_better(v, i, best_v, best_i)) { - best_v = v; - best_i = i; - } - } - - __shared__ float vals[256]; - __shared__ uint32_t idxs[256]; - vals[tid] = best_v; - idxs[tid] = best_i; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0u; stride >>= 1u) { - if (tid < stride) { - if (topk_score_better(vals[tid + stride], idxs[tid + stride], - vals[tid], idxs[tid])) { - vals[tid] = vals[tid + stride]; - idxs[tid] = idxs[tid + stride]; - } - } - __syncthreads(); - } - if (tid == 0u) { - /* Monotonic float key; ~idx in the low bits makes ties resolve to - * the smaller index under atomicMax (matches topk_score_better). */ - const unsigned int f = __float_as_uint(vals[0]); - const unsigned int fkey = (f & 0x80000000u) ? ~f : (f | 0x80000000u); - const unsigned long long key = - ((unsigned long long)fkey << 32) | (unsigned int)(~idxs[0]); - atomicMax(out_key, key); - } -} - -__global__ static void indexer_top1_kernel( - uint32_t *selected, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens) { - const uint32_t t = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= 1024u) return; - - const float *row = scores + (uint64_t)t * n_comp; - float best_v = -INFINITY; - uint32_t best_i = 0; - for (uint32_t i = tid; i < n_comp; i += 1024u) { - const float v = row[i]; - if (topk_score_better(v, i, best_v, best_i)) { - best_v = v; - best_i = i; - } - } - - __shared__ float vals[1024]; - __shared__ uint32_t idxs[1024]; - vals[tid] = best_v; - idxs[tid] = best_i; - __syncthreads(); - - for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { - if (tid < stride) { - const float ov = vals[tid + stride]; - const uint32_t oi = idxs[tid + stride]; - if (topk_score_better(ov, oi, vals[tid], idxs[tid])) { - vals[tid] = ov; - idxs[tid] = oi; - } - } - __syncthreads(); - } - - if (tid == 0u) selected[t] = idxs[0]; -} - -__global__ static void indexer_top1_value_kernel( - uint32_t *selected, - float *values, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t index_offset) { - const uint32_t t = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= 1024u) return; - - const float *row = scores + (uint64_t)t * n_comp; - float best_v = -INFINITY; - uint32_t best_i = 0; - for (uint32_t i = tid; i < n_comp; i += 1024u) { - const float v = row[i]; - const uint32_t gi = index_offset + i; - const uint32_t best_gi = index_offset + best_i; - if (topk_score_better(v, gi, best_v, best_gi)) { - best_v = v; - best_i = i; - } - } - - __shared__ float vals[1024]; - __shared__ uint32_t idxs[1024]; - vals[tid] = best_v; - idxs[tid] = best_i; - __syncthreads(); - - for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { - if (tid < stride) { - const float ov = vals[tid + stride]; - const uint32_t oi = idxs[tid + stride]; - const uint32_t ogi = index_offset + oi; - const uint32_t gi = index_offset + idxs[tid]; - if (topk_score_better(ov, ogi, vals[tid], gi)) { - vals[tid] = ov; - idxs[tid] = oi; - } - } - __syncthreads(); - } - - if (tid == 0u) { - selected[t] = index_offset + idxs[0]; - values[t] = vals[0]; - } -} - -__global__ static void indexer_top2_value_kernel( - uint32_t *selected, - float *values, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t index_offset) { - const uint32_t t = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= 1024u) return; - - const float *row = scores + (uint64_t)t * n_comp; - float best0_v = -INFINITY; - float best1_v = -INFINITY; - uint32_t best0_i = UINT32_MAX; - uint32_t best1_i = UINT32_MAX; - for (uint32_t i = tid; i < n_comp; i += 1024u) { - const uint32_t gi = index_offset + i; - top2_insert_candidate(row[i], gi, - &best0_v, &best0_i, - &best1_v, &best1_i); - } - - __shared__ float vals0[1024]; - __shared__ float vals1[1024]; - __shared__ uint32_t idxs0[1024]; - __shared__ uint32_t idxs1[1024]; - vals0[tid] = best0_v; - vals1[tid] = best1_v; - idxs0[tid] = best0_i; - idxs1[tid] = best1_i; - __syncthreads(); - - for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { - if (tid < stride) { - top2_insert_candidate(vals0[tid + stride], idxs0[tid + stride], - &vals0[tid], &idxs0[tid], - &vals1[tid], &idxs1[tid]); - top2_insert_candidate(vals1[tid + stride], idxs1[tid + stride], - &vals0[tid], &idxs0[tid], - &vals1[tid], &idxs1[tid]); - } - __syncthreads(); - } - - if (tid == 0u) { - selected[(uint64_t)t * 2u + 0u] = idxs0[0]; - selected[(uint64_t)t * 2u + 1u] = idxs1[0]; - values[(uint64_t)t * 2u + 0u] = vals0[0]; - values[(uint64_t)t * 2u + 1u] = vals1[0]; - } -} - -__device__ __forceinline__ static uint32_t topk_float_ordered_key(float v) { - const uint32_t u = __float_as_uint(v); - return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); -} - -__device__ __forceinline__ static uint64_t topk_pack_key(float v, uint32_t idx) { - return ((uint64_t)topk_float_ordered_key(v) << 32u) | (uint64_t)(0xffffffffu - idx); -} - -__global__ static void indexer_topk_8192_cub_kernel( - uint32_t *selected, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - constexpr uint32_t BLOCK_THREADS = 512u; - constexpr uint32_t ITEMS_PER_THREAD = 16u; - using BlockSort = cub::BlockRadixSort; - extern __shared__ __align__(16) unsigned char sort_smem[]; - typename BlockSort::TempStorage &sort_storage = - *reinterpret_cast(sort_smem); - - const uint32_t t = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= BLOCK_THREADS) return; - - const float *row = scores + (uint64_t)t * n_comp; - uint64_t keys[ITEMS_PER_THREAD]; -#pragma unroll - for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) { - const uint32_t i = tid * ITEMS_PER_THREAD + item; - if (i < n_comp) { - keys[item] = topk_pack_key(row[i], i); - } else { - keys[item] = topk_pack_key(-INFINITY, UINT32_MAX); - } - } - - BlockSort(sort_storage).SortDescending(keys); - -#pragma unroll - for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) { - const uint32_t i = tid * ITEMS_PER_THREAD + item; - if (i < top_k) { - selected[(uint64_t)t * top_k + i] = 0xffffffffu - (uint32_t)keys[item]; - } - } -} - -__global__ static void indexer_topk_1024_kernel( - uint32_t *selected, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - uint32_t t = blockIdx.x; - uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= 1024u) return; - __shared__ float vals[1024]; - __shared__ uint32_t idxs[1024]; - - const float *row = scores + (uint64_t)t * n_comp; - if (tid < n_comp) { - vals[tid] = row[tid]; - idxs[tid] = tid; - } else { - vals[tid] = -INFINITY; - idxs[tid] = UINT32_MAX; - } - __syncthreads(); - - for (uint32_t k = 2u; k <= 1024u; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - uint32_t other = tid ^ j; - if (other > tid && other < 1024u) { - const float av = vals[tid]; - const float bv = vals[other]; - const uint32_t ai = idxs[tid]; - const uint32_t bi = idxs[other]; - const bool desc_half = (tid & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[tid] = bv; - idxs[tid] = bi; - vals[other] = av; - idxs[other] = ai; - } - } - __syncthreads(); - } - } - - if (tid < top_k) selected[(uint64_t)t * top_k + tid] = idxs[tid]; -} - -template -__global__ static void indexer_topk_pow2_kernel( - uint32_t *selected, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - uint32_t t = blockIdx.x; - uint32_t tid = threadIdx.x; - if (t >= n_tokens) return; - __shared__ float vals[SORT_N]; - __shared__ uint32_t idxs[SORT_N]; - - const float *row = scores + (uint64_t)t * n_comp; - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - if (i < n_comp) { - vals[i] = row[i]; - idxs[i] = i; - } else { - vals[i] = -INFINITY; - idxs[i] = UINT32_MAX; - } - } - __syncthreads(); - - for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t other = i ^ j; - if (other > i && other < SORT_N) { - const float av = vals[i]; - const float bv = vals[other]; - const uint32_t ai = idxs[i]; - const uint32_t bi = idxs[other]; - const bool desc_half = (i & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[i] = bv; - idxs[i] = bi; - vals[other] = av; - idxs[other] = ai; - } - } - } - __syncthreads(); - } - } - - for (uint32_t i = tid; i < top_k; i += blockDim.x) { - selected[(uint64_t)t * top_k + i] = idxs[i]; - } -} - -template -__global__ static void indexer_topk_pow2_u16_kernel( - uint32_t *selected, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - uint32_t t = blockIdx.x; - uint32_t tid = threadIdx.x; - if (t >= n_tokens) return; - __shared__ float vals[SORT_N]; - __shared__ uint16_t idxs[SORT_N]; - - const float *row = scores + (uint64_t)t * n_comp; - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - if (i < n_comp) { - vals[i] = row[i]; - idxs[i] = (uint16_t)i; - } else { - vals[i] = -INFINITY; - idxs[i] = UINT16_MAX; - } - } - __syncthreads(); - - for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t other = i ^ j; - if (other > i && other < SORT_N) { - const float av = vals[i]; - const float bv = vals[other]; - const uint32_t ai = idxs[i]; - const uint32_t bi = idxs[other]; - const bool desc_half = (i & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[i] = bv; - idxs[i] = (uint16_t)bi; - vals[other] = av; - idxs[other] = (uint16_t)ai; - } - } - } - __syncthreads(); - } - } - - for (uint32_t i = tid; i < top_k; i += blockDim.x) { - selected[(uint64_t)t * top_k + i] = idxs[i]; - } -} - -template -__global__ static void indexer_topk_chunk_pow2_kernel( - uint32_t *candidates, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k, - uint32_t candidate_stride) { - uint32_t t = blockIdx.x; - uint32_t chunk = blockIdx.y; - uint32_t tid = threadIdx.x; - if (t >= n_tokens) return; - - const uint32_t chunk_start = chunk * SORT_N; - if (chunk_start >= n_comp) return; - const uint32_t chunk_n = n_comp - chunk_start < SORT_N ? n_comp - chunk_start : SORT_N; - __shared__ float vals[SORT_N]; - __shared__ uint32_t idxs[SORT_N]; - - const float *row = scores + (uint64_t)t * n_comp; - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - if (i < chunk_n) { - vals[i] = row[chunk_start + i]; - idxs[i] = chunk_start + i; - } else { - vals[i] = -INFINITY; - idxs[i] = UINT32_MAX; - } - } - __syncthreads(); - - for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t other = i ^ j; - if (other > i && other < SORT_N) { - const float av = vals[i]; - const float bv = vals[other]; - const uint32_t ai = idxs[i]; - const uint32_t bi = idxs[other]; - const bool desc_half = (i & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[i] = bv; - idxs[i] = bi; - vals[other] = av; - idxs[other] = ai; - } - } - } - __syncthreads(); - } - } - - uint32_t *out = candidates + (uint64_t)t * candidate_stride + chunk * top_k; - for (uint32_t i = tid; i < top_k; i += blockDim.x) { - out[i] = idxs[i]; - } -} - -template -__global__ static void indexer_topk_merge_pow2_kernel( - uint32_t *selected, - const uint32_t *candidates, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k, - uint32_t candidate_count, - uint32_t candidate_stride) { - uint32_t t = blockIdx.x; - uint32_t tid = threadIdx.x; - if (t >= n_tokens) return; - __shared__ float vals[SORT_N]; - __shared__ uint32_t idxs[SORT_N]; - - const float *row = scores + (uint64_t)t * n_comp; - const uint32_t *cand = candidates + (uint64_t)t * candidate_stride; - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t idx = UINT32_MAX; - float v = -INFINITY; - if (i < candidate_count) { - idx = cand[i]; - if (idx < n_comp) v = row[idx]; - } - vals[i] = v; - idxs[i] = idx; - } - __syncthreads(); - - for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t other = i ^ j; - if (other > i && other < SORT_N) { - const float av = vals[i]; - const float bv = vals[other]; - const uint32_t ai = idxs[i]; - const uint32_t bi = idxs[other]; - const bool desc_half = (i & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[i] = bv; - idxs[i] = bi; - vals[other] = av; - idxs[other] = ai; - } - } - } - __syncthreads(); - } - } - - for (uint32_t i = tid; i < top_k; i += blockDim.x) { - selected[(uint64_t)t * top_k + i] = idxs[i]; - } -} - -template -__global__ static void indexer_topk_tree_merge_pow2_kernel( - uint32_t *out, - const uint32_t *candidates, - const float *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k, - uint32_t n_sets, - uint32_t merge_group, - uint32_t candidate_stride, - uint32_t out_stride) { - uint32_t t = blockIdx.x; - uint32_t group = blockIdx.y; - uint32_t tid = threadIdx.x; - if (t >= n_tokens) return; - - const uint32_t set0 = group * merge_group; - if (set0 >= n_sets) return; - uint32_t set_count = n_sets - set0; - if (set_count > merge_group) set_count = merge_group; - const uint32_t candidate_count = set_count * top_k; - - __shared__ float vals[SORT_N]; - __shared__ uint32_t idxs[SORT_N]; - - const float *row = scores + (uint64_t)t * n_comp; - const uint32_t *cand = candidates + (uint64_t)t * candidate_stride + set0 * top_k; - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t idx = UINT32_MAX; - float v = -INFINITY; - if (i < candidate_count) { - idx = cand[i]; - if (idx < n_comp) v = row[idx]; - } - vals[i] = v; - idxs[i] = idx; - } - __syncthreads(); - - for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { - uint32_t other = i ^ j; - if (other > i && other < SORT_N) { - const float av = vals[i]; - const float bv = vals[other]; - const uint32_t ai = idxs[i]; - const uint32_t bi = idxs[other]; - const bool desc_half = (i & k) == 0u; - const bool swap = desc_half - ? topk_score_better(bv, bi, av, ai) - : topk_score_better(av, ai, bv, bi); - if (swap) { - vals[i] = bv; - idxs[i] = bi; - vals[other] = av; - idxs[other] = ai; - } - } - } - __syncthreads(); - } - } - - uint32_t *dst = out + (uint64_t)t * out_stride + group * top_k; - for (uint32_t i = tid; i < top_k; i += blockDim.x) { - dst[i] = idxs[i]; - } -} - -__global__ static void indexed_topk_sort_512_asc_kernel( - int32_t *dst, - const int32_t *src, - uint32_t n_tokens) { - const uint32_t t = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (t >= n_tokens || tid >= 512u) return; - __shared__ int32_t rows[512]; - - const int32_t *src_row = src + (uint64_t)t * 512u; - int32_t *dst_row = dst + (uint64_t)t * 512u; - rows[tid] = src_row[tid]; - __syncthreads(); - - for (uint32_t k = 2u; k <= 512u; k <<= 1u) { - for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { - const uint32_t other = tid ^ j; - if (other > tid && other < 512u) { - const int32_t a = rows[tid]; - const int32_t b = rows[other]; - const bool up = (tid & k) == 0u; - if ((up && a > b) || (!up && a < b)) { - rows[tid] = b; - rows[other] = a; - } - } - __syncthreads(); - } - } - - dst_row[tid] = rows[tid]; -} - -__global__ static void topk_mask_kernel(float *mask, const uint32_t *topk, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * n_comp; - if (gid >= n) return; - uint32_t t = gid / n_comp; - uint32_t c = gid - (uint64_t)t * n_comp; - float v = -INFINITY; - for (uint32_t k = 0; k < top_k; k++) { - if (topk[(uint64_t)t * top_k + k] == c) { - v = 0.0f; - break; - } - } - mask[gid] = v; -} - -extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n_vocab, uint32_t token, uint32_t n_embd, uint32_t n_hc) { - (void)n_vocab; - if (!out_hc || !model_map || weight_offset >= model_size) return 0; - uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0; - const int logical_tier = ds4_tensor_device_idx(out_hc); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "token_embd"); - if (!wptr) return 0; - uint32_t n = n_embd * n_hc; - embed_token_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "embed token launch"); -} - -extern "C" int ds4_gpu_embed_tokens_hc_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *tokens_t, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd, - uint32_t n_hc) { - if (!out_hc || !tokens_t || !model_map || - weight_offset > model_size || - (uint64_t)n_vocab * n_embd * sizeof(uint16_t) > model_size - weight_offset || - tokens_t->bytes < (uint64_t)n_tokens * sizeof(int32_t) || - out_hc->bytes < (uint64_t)n_tokens * n_hc * n_embd * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out_hc); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, - (uint64_t)n_vocab * n_embd * sizeof(uint16_t), - logical_tier, - "token_embd"); - if (!wptr) return 0; - uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; - embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>( - (float *)out_hc->ptr, - (const int32_t *)tokens_t->ptr, - (const __half *)wptr, - n_vocab, n_tokens, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "embed tokens launch"); -} - -static int indexer_scores_launch( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale, - uint32_t causal) { - if (!scores || !q || !weights || !index_comp || - n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) || - index_comp->bytes < (uint64_t)n_comp * head_dim * sizeof(float) || - scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float)) { - return 0; - } - if (causal && ratio == 0) return 0; - if (n_tokens == 1u && head_dim == 128u && n_head == 64u && - getenv("DS4_CUDA_NO_INDEXER_DIRECT_ONE") == NULL) { - indexer_score_one_direct_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, pos0, ratio, - scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer score one direct launch"); - } - if (!g_quality_mode && head_dim == 128u && n_head == 64u && - getenv("DS4_CUDA_NO_INDEXER_WMMA") == NULL) { - if (getenv("DS4_CUDA_NO_INDEXER_WMMA128") == NULL) { - dim3 grid((n_comp + 127u) / 128u, (n_tokens + 15u) / 16u, 1); - indexer_scores_wmma128_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, n_tokens, pos0, n_head, - head_dim, ratio, scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer scores wmma128 launch"); - } else if (getenv("DS4_CUDA_NO_INDEXER_WMMA64") == NULL) { - dim3 grid((n_comp + 63u) / 64u, (n_tokens + 15u) / 16u, 1); - indexer_scores_wmma64_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, n_tokens, pos0, n_head, - head_dim, ratio, scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer scores wmma64 launch"); - } else if (getenv("DS4_CUDA_NO_INDEXER_WMMA32") == NULL) { - dim3 grid((n_comp + 31u) / 32u, (n_tokens + 15u) / 16u, 1); - indexer_scores_wmma32_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, n_tokens, pos0, n_head, - head_dim, ratio, scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer scores wmma32 launch"); - } else { - dim3 grid((n_comp + 15u) / 16u, (n_tokens + 15u) / 16u, 1); - indexer_scores_wmma_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, n_tokens, pos0, n_head, - head_dim, ratio, scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer scores wmma launch"); - } - } - dim3 grid(n_comp, n_tokens, 1); - indexer_scores_kernel<<>>((float *)scores->ptr, - (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)index_comp->ptr, - n_comp, n_tokens, pos0, n_head, - head_dim, ratio, scale, causal ? 1 : 0); - return cuda_ok(cudaGetLastError(), "indexer scores launch"); -} - -extern "C" int ds4_gpu_indexer_score_one_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_head, - uint32_t head_dim, - float scale) { - return indexer_scores_launch(scores, q, weights, index_comp, n_comp, 1, 0, - n_head, head_dim, 1, scale, 0); -} - -extern "C" int ds4_gpu_indexer_scores_prefill_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale) { - return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, 0, - n_head, head_dim, ratio, scale, 1); -} - -extern "C" int ds4_gpu_indexer_scores_decode_batch_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale) { - return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, pos0, - n_head, head_dim, ratio, scale, 1); -} - -extern "C" int ds4_gpu_dspark_markov_argmax_tensor( - ds4_gpu_tensor *out_idx, - const ds4_gpu_tensor *logits_row, - const void *model_map, - uint64_t model_size, - uint64_t w1_offset, - uint64_t w2_offset, - uint32_t prev_token, - uint32_t vocab, - uint32_t rank) { - if (!out_idx || !logits_row || !model_map || vocab == 0 || - rank == 0 || (rank & 31u) != 0u || rank > 256u || - out_idx->bytes < sizeof(unsigned long long) || - logits_row->bytes < (uint64_t)vocab * sizeof(float)) { - return 0; - } - const uint32_t rank_blocks = rank / 32u; - const uint64_t row_bytes = (uint64_t)rank_blocks * 34u; - if (w1_offset > model_size || - (uint64_t)prev_token * row_bytes + row_bytes > model_size - w1_offset || - w2_offset > model_size || - (uint64_t)vocab * row_bytes > model_size - w2_offset) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(logits_row); - const unsigned char *w1_row = (const unsigned char *)cuda_resolve_weight_ptr( - model_map, w1_offset + (uint64_t)prev_token * row_bytes, - row_bytes, logical_tier, "markov_w1_row"); - const unsigned char *w2 = (const unsigned char *)cuda_resolve_weight_ptr( - model_map, w2_offset, (uint64_t)vocab * row_bytes, - logical_tier, "markov_w2"); - if (!w1_row || !w2) return 0; - int dev_save = 0; - if (cudaGetDevice(&dev_save) != cudaSuccess) return 0; - if (logical_tier != dev_save && cudaSetDevice(logical_tier) != cudaSuccess) { - return 0; - } - int rc = cudaMemsetAsync(out_idx->ptr, 0, - sizeof(unsigned long long)) == cudaSuccess; - if (rc) { - dspark_markov_argmax_kernel<<<128, 256>>>( - (unsigned long long *)out_idx->ptr, - (const float *)logits_row->ptr, - w1_row, w2, vocab, rank_blocks); - rc = cuda_ok(cudaGetLastError(), "dspark markov argmax launch"); - } - if (logical_tier != dev_save) (void)cudaSetDevice(dev_save); - return rc; -} - -extern "C" int ds4_gpu_indexer_topk_tensor( - ds4_gpu_tensor *selected, - const ds4_gpu_tensor *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || - top_k > n_comp || - scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) { - return 0; - } - if (top_k == 1u && !g_cuda_no_top1) { - indexer_top1_kernel<<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, - n_tokens); - return cuda_ok(cudaGetLastError(), "indexer top1 launch"); - } - if (top_k == 2048u && n_comp <= 4096u && - getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) { - indexer_topk_pow2_kernel<4096><<>>( - (uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 2048-wide launch"); - } - if (top_k == 2048u && n_comp > 4096u && - getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) { - const uint32_t chunk_n = 4096u; - const uint32_t merge_group = 2u; - const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n; - const uint32_t candidate_stride = n_chunks * top_k; - uint32_t n_sets = n_chunks; - uint64_t scratch_u32_per_token = candidate_stride; - while (n_sets > merge_group) { - n_sets = (n_sets + merge_group - 1u) / merge_group; - scratch_u32_per_token += (uint64_t)n_sets * top_k; - } - if (scratch_u32_per_token > - UINT64_MAX / n_tokens / sizeof(uint32_t)) { - return 0; - } - int exec_tier = ds4_tensor_device_idx(selected); - int current_device = -1; - if (cudaGetDevice(¤t_device) == cudaSuccess) { - for (int t = 0; t < g_n_gpus; t++) { - if (g_gpu[t].device_id == current_device) { - exec_tier = t; - break; - } - } - } - const uint64_t tmp_bytes = - (uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t); - uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on( - exec_tier, tmp_bytes, "indexer topk 2048-wide tree"); - if (!scratch) return 0; - - uint32_t *cur = scratch; - n_sets = n_chunks; - uint32_t cur_stride = candidate_stride; - dim3 grid_chunks(n_tokens, n_chunks, 1); - indexer_topk_chunk_pow2_kernel<4096><<>>( - cur, (const float *)scores->ptr, - n_comp, n_tokens, top_k, candidate_stride); - if (!cuda_ok(cudaGetLastError(), - "indexer topk 2048-wide chunk launch")) { - return 0; - } - - while (n_sets > merge_group) { - const uint32_t next_sets = - (n_sets + merge_group - 1u) / merge_group; - const uint32_t next_stride = next_sets * top_k; - uint32_t *next = cur + (uint64_t)n_tokens * cur_stride; - dim3 grid_merge(n_tokens, next_sets, 1); - indexer_topk_tree_merge_pow2_kernel<4096><<>>( - next, cur, (const float *)scores->ptr, - n_comp, n_tokens, top_k, n_sets, merge_group, - cur_stride, next_stride); - if (!cuda_ok(cudaGetLastError(), - "indexer topk 2048-wide merge launch")) { - return 0; - } - cur = next; - n_sets = next_sets; - cur_stride = next_stride; - } - - indexer_topk_merge_pow2_kernel<4096><<>>( - (uint32_t *)selected->ptr, - cur, (const float *)scores->ptr, - n_comp, n_tokens, top_k, n_sets * top_k, cur_stride); - return cuda_ok(cudaGetLastError(), - "indexer topk 2048-wide final launch"); - } - if (top_k == 512u && n_comp <= 1024u && - getenv("DS4_CUDA_NO_TOPK1024") == NULL) { - indexer_topk_1024_kernel<<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 1024 launch"); - } - if (top_k == 512u && n_comp <= 2048u && - getenv("DS4_CUDA_NO_TOPK2048") == NULL) { - indexer_topk_pow2_kernel<2048><<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 2048 launch"); - } - if (top_k == 512u && n_comp <= 4096u && - getenv("DS4_CUDA_NO_TOPK2048") == NULL) { - if (n_comp == 4096u) { - using TopkCubSort = cub::BlockRadixSort; - const int smem = (int)sizeof(typename TopkCubSort::TempStorage); - int dev = 0; - int max_optin_smem = 0; - cudaError_t attr_err = cudaGetDevice(&dev); - if (attr_err == cudaSuccess) { - attr_err = cudaDeviceGetAttribute(&max_optin_smem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, - dev); - } - if (attr_err == cudaSuccess && max_optin_smem >= smem) { - attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - if (attr_err == cudaSuccess) { - indexer_topk_8192_cub_kernel<<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 4096 cub launch"); - } - } - } - indexer_topk_pow2_kernel<4096><<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 4096 launch"); - } - if (top_k == 512u && n_comp <= 8192u && - getenv("DS4_CUDA_NO_TOPK2048") == NULL && - getenv("DS4_CUDA_NO_TOPK8192") == NULL) { - if (n_comp > 4096u) { - using TopkCubSort = cub::BlockRadixSort; - const int smem = (int)sizeof(typename TopkCubSort::TempStorage); - int dev = 0; - int max_optin_smem = 0; - cudaError_t attr_err = cudaGetDevice(&dev); - if (attr_err == cudaSuccess) { - attr_err = cudaDeviceGetAttribute(&max_optin_smem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, - dev); - } - if (attr_err == cudaSuccess && max_optin_smem >= smem) { - attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - if (attr_err == cudaSuccess) { - indexer_topk_8192_cub_kernel<<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 8192 cub launch"); - } - } - } - indexer_topk_pow2_u16_kernel<8192><<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk 8192 launch"); - } - if (top_k == 512u && getenv("DS4_CUDA_NO_TOPK2048") == NULL && - getenv("DS4_CUDA_NO_TOPK_CHUNKED") == NULL) { - const uint32_t chunk_n = 4096u; - const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n; - const uint32_t candidate_stride = n_chunks * top_k; - uint32_t n_sets = n_chunks; - uint64_t scratch_u32_per_token = candidate_stride; - while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) { - n_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP; - scratch_u32_per_token += (uint64_t)n_sets * top_k; - } - if (scratch_u32_per_token > UINT64_MAX / n_tokens / sizeof(uint32_t)) return 0; - const uint64_t tmp_bytes = (uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t); - const int logical_tier = ds4_tensor_device_idx(selected); - uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "indexer topk tree"); - if (!scratch) return 0; - - uint32_t *cur = scratch; - n_sets = n_chunks; - uint32_t cur_stride = candidate_stride; - dim3 grid_chunks(n_tokens, n_chunks, 1); - indexer_topk_chunk_pow2_kernel<4096><<>>(cur, - (const float *)scores->ptr, - n_comp, - n_tokens, - top_k, - candidate_stride); - if (!cuda_ok(cudaGetLastError(), "indexer topk chunk launch")) return 0; - - while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) { - const uint32_t next_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP; - const uint32_t next_stride = next_sets * top_k; - uint32_t *next = cur + (uint64_t)n_tokens * cur_stride; - dim3 grid_merge(n_tokens, next_sets, 1); - indexer_topk_tree_merge_pow2_kernel<4096><<>>( - next, - cur, - (const float *)scores->ptr, - n_comp, - n_tokens, - top_k, - n_sets, - DS4_CUDA_TOPK_MERGE_GROUP, - cur_stride, - next_stride); - if (!cuda_ok(cudaGetLastError(), "indexer topk tree merge launch")) return 0; - cur = next; - n_sets = next_sets; - cur_stride = next_stride; - } - - indexer_topk_merge_pow2_kernel<4096><<>>((uint32_t *)selected->ptr, - cur, - (const float *)scores->ptr, - n_comp, - n_tokens, - top_k, - n_sets * top_k, - cur_stride); - return cuda_ok(cudaGetLastError(), "indexer topk tree final launch"); - } - indexer_topk_kernel<<>>((uint32_t *)selected->ptr, - (const float *)scores->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "indexer topk launch"); -} - -extern "C" int ds4_gpu_indexer_top1_value_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *values, - const ds4_gpu_tensor *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t index_offset) { - if (!selected || !values || !scores || n_comp == 0 || n_tokens == 0 || - scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * sizeof(uint32_t) || - values->bytes < (uint64_t)n_tokens * sizeof(float)) { - return 0; - } - indexer_top1_value_kernel<<>>((uint32_t *)selected->ptr, - (float *)values->ptr, - (const float *)scores->ptr, - n_comp, - n_tokens, - index_offset); - return cuda_ok(cudaGetLastError(), "indexer top1 value launch"); -} - -extern "C" int ds4_gpu_indexer_top2_value_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *values, - const ds4_gpu_tensor *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t index_offset) { - if (!selected || !values || !scores || n_comp < 2u || n_tokens == 0 || - scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * 2u * sizeof(uint32_t) || - values->bytes < (uint64_t)n_tokens * 2u * sizeof(float)) { - return 0; - } - indexer_top2_value_kernel<<>>((uint32_t *)selected->ptr, - (float *)values->ptr, - (const float *)scores->ptr, - n_comp, - n_tokens, - index_offset); - return cuda_ok(cudaGetLastError(), "indexer top2 value launch"); -} - -extern "C" int ds4_gpu_dsv4_topk_mask_tensor( - ds4_gpu_tensor *mask, - const ds4_gpu_tensor *topk, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0 || - mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || - topk->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) { - return 0; - } - uint64_t n = (uint64_t)n_tokens * n_comp; - uint64_t nk = (uint64_t)n_tokens * top_k; - uint64_t blocks = ((n > nk ? n : nk) + 255) / 256; - topk_mask_kernel<<>>((float *)mask->ptr, - (const uint32_t *)topk->ptr, - n_comp, n_tokens, top_k); - return cuda_ok(cudaGetLastError(), "topk mask launch"); -} -/* GLM opt-in: batched q8_0 matmuls with blocks > 32 may run as a - * streaming dequant-to-f16 GEMM (exact-q8 native kernels only cover - * blocks <= 32). Never enabled on DS4 paths, keeping them byte-stable. */ -static int g_q8_dequant_gemm_enabled = 0; -extern "C" void ds4_gpu_enable_q8_dequant_gemm(void) { - g_q8_dequant_gemm_enabled = 1; -} - -__global__ static void q8_0_dequant_f16_kernel( - __half *out, - const unsigned char *w, - uint64_t total_blocks, - uint32_t blocks_per_row, - uint32_t in_dim) { - /* Two threads per q8_0 block; each converts 16 values with half2 - * stores so a warp writes 512B contiguously per block pair. */ - const uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - const uint64_t b = tid >> 1; - if (b >= total_blocks) return; - const uint32_t half_idx = (uint32_t)tid & 1u; - const unsigned char *blk = w + b * 34u; - const float d = __half2float(*(const __half *)blk); - const int8_t *q = (const int8_t *)(blk + 2) + half_idx * 16u; - const uint64_t row = b / blocks_per_row; - const uint32_t col = (uint32_t)(b - row * blocks_per_row) * 32u + - half_idx * 16u; - __half2 *dst = (__half2 *)(out + row * in_dim + col); - #pragma unroll - for (int k = 0; k < 8; k++) { - dst[k] = __floats2half2_rn(d * (float)q[2 * k], - d * (float)q[2 * k + 1]); - } -} - -static int cuda_matmul_q8_0_tensor_labeled(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_tok, const char *label) { - if (!out || !x || !model_map) return 0; - uint64_t blocks = (in_dim + 31) / 32; - if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0; - uint64_t weight_bytes = out_dim * blocks * 34; - if (weight_bytes > model_size - weight_offset) return 0; - if (x->bytes < n_tok * in_dim * sizeof(float) || - out->bytes < n_tok * out_dim * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const int physical_device = - (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) - ? g_gpu[logical_tier].device_id : 0; - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "q8_0"); - if (!wptr) return 0; - if (g_cublas_ready && n_tok > 1) { - const float *w_f32 = cuda_q8_f32_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label); - if (w_f32) { - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w_f32, - (int)in_dim, - (const float *)x->ptr, - (int)in_dim, - &beta, - (float *)out->ptr, - (int)out_dim); - return cublas_ok(st, "q8 fp32 matmul"); - } - const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label); - if (w_f16) { - const uint64_t xh_count = n_tok * in_dim; - __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "q8 f16 gemm activations"); - if (!xh) return 0; - f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count); - if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0; - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w_f16, - CUDA_R_16F, - (int)in_dim, - xh, - CUDA_R_16F, - (int)in_dim, - &beta, - out->ptr, - CUDA_R_32F, - (int)out_dim, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - if (st == CUBLAS_STATUS_SUCCESS) return 1; - fprintf(stderr, "ds4: cuBLAS q8 f16 matmul failed: status %d\n", (int)st); - cuda_q8_f16_cache_disable_after_failure("cuBLAS f16 matmul failure", - in_dim * out_dim * sizeof(__half)); - /* The F16 expansion cache is only an optimization. If cuBLAS - * rejects the cached path under memory pressure, retry the same - * operation through the native Q8 kernels below. */ - } - } - if (g_q8_dequant_gemm_enabled && g_cublas_ready && - n_tok >= 128u && blocks > 32u && (in_dim & 31u) == 0u) { - /* Streaming dequant + f16 GEMM: the exact-q8 batched kernels only - * cover blocks <= 32 (DS4 TP shard widths); the per-token fallback - * re-reads the full weight per token (~30x the bytes at GLM dims). - * Scratch layout: [w_f16][x_f16] in one arena grab. */ - const uint64_t wh_bytes = in_dim * out_dim * sizeof(__half); - const uint64_t xh_off = (wh_bytes + 255u) & ~255ull; - const uint64_t oo_off = - (xh_off + n_tok * in_dim * sizeof(__half) + 255u) & ~255ull; - const uint64_t gemm_tmp = oo_off + n_tok * out_dim * sizeof(float); - /* Scratch must live on the EXECUTING device: logical_tier is the - * out tensor's tier (0 for GLM graph buffers), and a GEMM reading - * its staged weights across PCIe costs ~20ms instead of ~0.1ms. */ - int exec_tier = logical_tier; - { - int cur_dev = -1; - if (cudaGetDevice(&cur_dev) == cudaSuccess) { - for (int t = 0; t < g_n_gpus; t++) { - if (g_gpu[t].device_id == cur_dev) { exec_tier = t; break; } - } - } - } - void *tmp16 = cuda_tmp_alloc_on(exec_tier, gemm_tmp, "q8 dequant gemm"); - if (tmp16) { - __half *wh = (__half *)tmp16; - __half *xh = (__half *)((char *)tmp16 + xh_off); - /* GEMM into device-local scratch, then one bulk D2D to the - * (possibly peer-mapped) out tensor: scattered peer stores - * from GEMM kernels run at <1GB/s over PCIe. */ - float *olocal = (float *)((char *)tmp16 + oo_off); - const uint64_t total_blocks = out_dim * blocks; - q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256>>>( - wh, reinterpret_cast(wptr), - total_blocks, (uint32_t)blocks, (uint32_t)in_dim); - const uint64_t xh_count = n_tok * in_dim; - f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>( - xh, (const float *)x->ptr, xh_count); - if (cuda_ok(cudaGetLastError(), "q8 dequant gemm staging")) { - const int gemm_trace = getenv("DS4_GLM_GEMM_TRACE") != NULL; - cudaEvent_t ev0, ev1, ev2; - if (gemm_trace) { - cudaEventCreate(&ev0); cudaEventCreate(&ev1); cudaEventCreate(&ev2); - cudaEventRecord(ev0); - } - const float alpha = 1.0f; - const float beta = 0.0f; - if (gemm_trace) cudaEventRecord(ev1); - cublasStatus_t st = cublasGemmEx( - cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, CUBLAS_OP_N, - (int)out_dim, (int)n_tok, (int)in_dim, - &alpha, - wh, CUDA_R_16F, (int)in_dim, - xh, CUDA_R_16F, (int)in_dim, - &beta, - olocal, CUDA_R_32F, (int)out_dim, - CUDA_R_32F, CUBLAS_GEMM_DEFAULT); - if (st == CUBLAS_STATUS_SUCCESS) { - if (!cuda_ok(cudaMemcpyAsync(out->ptr, olocal, - n_tok * out_dim * sizeof(float), - cudaMemcpyDeviceToDevice, 0), - "q8 dequant gemm out copy")) { - st = CUBLAS_STATUS_INTERNAL_ERROR; - } - } - if (gemm_trace) { - cudaEventRecord(ev2); - cudaEventSynchronize(ev2); - float stage_ms = 0, gemm_ms = 0; - cudaEventElapsedTime(&stage_ms, ev0, ev1); - cudaEventElapsedTime(&gemm_ms, ev1, ev2); - fprintf(stderr, - "ds4: gemm trace in=%llu out=%llu n=%llu stage(before)=%.2f gemm=%.2f ms\n", - (unsigned long long)in_dim, (unsigned long long)out_dim, - (unsigned long long)n_tok, stage_ms, gemm_ms); - cudaEventDestroy(ev0); cudaEventDestroy(ev1); cudaEventDestroy(ev2); - } - if (st == CUBLAS_STATUS_SUCCESS) return 1; - fprintf(stderr, - "ds4: q8 dequant gemm failed: status %d; using native path\n", - (int)st); - } - } - } - const uint64_t xq_bytes = n_tok * blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1); - quantize_q8_0_f32_kernel<<>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0; - if (n_tok == 1) { - matmul_q8_0_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 warp launch"); - } - const bool force_decode_warp = - n_tok == 2u && g_glm_mtp_verify_mode; - if (n_tok > 1u && !force_decode_warp) { - /* T matches the reduction width of whichever reference kernel would - * have run: warp tree (32) for blocks <= 32, exact-thread tree - * otherwise. */ - const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks); - const int mma_rc = cuda_q8_mma_try_launch( - (float *)out->ptr, reinterpret_cast(wptr), - xq, xscale, in_dim, out_dim, n_tok, blocks, blocks, out_dim, mma_T); - if (mma_rc) return mma_rc > 0; - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL && - blocks <= 32u && - n_tok >= 8u) { - dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); - matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok8 warp launch"); - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL && - blocks <= 32u && - n_tok >= 4u) { - dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); - matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok4 warp launch"); - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - (blocks <= 32u || force_decode_warp)) { - if (force_decode_warp && - getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) { - matmul_q8_0_preq_batch_warp8_tok2_kernel - <<<((unsigned)out_dim + 7u) / 8u, 256>>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), - "matmul_q8_0 batch tok2 warp launch"); - } - dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1); - matmul_q8_0_preq_batch_warp8_kernel<<>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch warp launch"); - } - const unsigned exact_threads = cuda_q8_exact_threads(blocks); - if (getenv("DS4_CUDA_NO_Q8_BATCH_EXACT_TOK2") == NULL && - n_tok >= 2u) { - dim3 bgrid((unsigned)out_dim, ((unsigned)n_tok + 1u) / 2u, 1); - matmul_q8_0_preq_batch_tok2_exact_kernel<<>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 exact tok2 launch"); - } - dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); - matmul_q8_0_preq_kernel<<>>((float *)out->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, out_dim, n_tok, blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 launch"); -} - -extern "C" int ds4_gpu_matmul_q8_0_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_tok) { - return cuda_matmul_q8_0_tensor_labeled(out, model_map, model_size, weight_offset, - in_dim, out_dim, x, n_tok, "q8_0"); -} - -extern "C" int ds4_gpu_matmul_q8_0_top1_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *values, - 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 index_offset) { - if (!selected || !values || !x || !model_map || - in_dim == 0 || out_dim == 0 || out_dim > UINT32_MAX) { - return 0; - } - const uint64_t blocks = (in_dim + 31u) / 32u; - if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34u)) { - return 0; - } - const uint64_t weight_bytes = out_dim * blocks * 34u; - if (weight_bytes > model_size - weight_offset || - x->bytes < in_dim * sizeof(float) || - selected->bytes < sizeof(uint32_t) || - values->bytes < sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(selected); - const char *wptr = cuda_resolve_weight_ptr(model_map, - weight_offset, - weight_bytes, - logical_tier, - "q8_0_top1"); - if (!wptr) return 0; - - const uint64_t xq_bytes = blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t key_offset = - (scale_offset + blocks * sizeof(float) + 7u) & ~7ull; - const uint64_t tmp_bytes = key_offset + sizeof(unsigned long long); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 top1 prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - unsigned long long *best_key = - (unsigned long long *)((char *)tmp + key_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - - if (!cuda_ok(cudaMemsetAsync(best_key, 0, sizeof(*best_key)), - "matmul_q8_0_top1 clear")) { - return 0; - } - quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>( - xq, - xscale, - (const float *)x->ptr, - in_dim, - blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 quantize launch")) return 0; - matmul_q8_0_top1_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( - best_key, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - blocks, - index_offset, - use_dp4a); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 launch")) return 0; - matmul_q8_0_top1_unpack_kernel<<<1, 1>>>( - (uint32_t *)selected->ptr, - (float *)values->ptr, - best_key); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 unpack launch"); -} - -extern "C" int ds4_gpu_matmul_q8_0_kslice_rows_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, - uint64_t in_start, - uint64_t in_count, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || - in_count == 0 || n_tok == 0 || n_tok > 65535u) return 0; - if ((in_start % 32u) != 0 || (in_count % 32u) != 0 || - in_start > in_dim || in_count > in_dim - in_start) return 0; - const uint64_t full_blocks = (in_dim + 31u) / 32u; - const uint64_t block_start = in_start / 32u; - const uint64_t slice_blocks = in_count / 32u; - if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0; - const uint64_t weight_bytes = out_dim * full_blocks * 34u; - if (in_count > UINT64_MAX / n_tok || out_dim > UINT64_MAX / n_tok) { - return 0; - } - if (weight_bytes > model_size - weight_offset || - x->bytes < n_tok * in_count * sizeof(float) || - out->bytes < n_tok * out_dim * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const unsigned char *wptr = reinterpret_cast( - cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, - logical_tier, "q8_0_kslice")); - if (!wptr) return 0; - - const uint64_t xq_bytes = n_tok * slice_blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = - scale_offset + n_tok * slice_blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - const dim3 qgrid((unsigned)slice_blocks, (unsigned)n_tok, 1u); - quantize_q8_0_f32_kernel<<>>( - xq, - xscale, - (const float *)x->ptr, - in_count, - slice_blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice quantize launch")) return 0; - const dim3 grid(((unsigned)out_dim + 7u) / 8u, - (unsigned)n_tok, 1u); - matmul_q8_0_kslice_preq_warp8_kernel<<>>( - (float *)out->ptr, - wptr, - xq, - xscale, - in_count, - out_dim, - full_blocks, - block_start, - slice_blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice launch"); -} - -extern "C" int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *block_out, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint64_t in_dim, - uint64_t out_dim, - uint64_t in_start, - uint64_t in_count, - const ds4_gpu_tensor *x, - const ds4_gpu_tensor *block_add, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (!out_hc || !block_out || !x || !block_add || !residual_hc || !split || - !model_map || in_dim == 0 || out_dim == 0 || in_count == 0 || - n_embd == 0 || n_hc == 0 || out_dim != (uint64_t)n_embd) { - return 0; - } - if ((in_start % 32u) != 0 || (in_count % 32u) != 0 || - in_start > in_dim || in_count > in_dim - in_start) return 0; - const uint64_t full_blocks = (in_dim + 31u) / 32u; - const uint64_t block_start = in_start / 32u; - const uint64_t slice_blocks = in_count / 32u; - if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0; - const uint64_t weight_bytes = out_dim * full_blocks * 34u; - const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t split_bytes = - (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float); - if (weight_bytes > model_size - weight_offset || - x->bytes < in_count * sizeof(float) || - block_out->bytes < out_dim * sizeof(float) || - block_add->bytes < out_dim * sizeof(float) || - residual_hc->bytes < hc_bytes || - split->bytes < split_bytes || - out_hc->bytes < hc_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out_hc); - const unsigned char *wptr = reinterpret_cast( - cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, - logical_tier, "q8_0_kslice_hc_expand_add")); - if (!wptr) return 0; - - const uint64_t xq_bytes = slice_blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + slice_blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice hc expand prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - quantize_q8_0_f32_kernel<<<(unsigned)slice_blocks, 32>>>( - xq, - xscale, - (const float *)x->ptr, - in_count, - slice_blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add quantize launch")) return 0; - matmul_q8_0_kslice_hc_expand_add_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( - (float *)out_hc->ptr, - (float *)block_out->ptr, - (const float *)block_add->ptr, - (const float *)residual_hc->ptr, - (const float *)split->ptr, - wptr, - xq, - xscale, - in_count, - out_dim, - full_blocks, - block_start, - slice_blocks, - n_embd, - n_hc, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add launch"); -} - -extern "C" int ds4_gpu_matmul_q8_0_pair_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, - uint64_t n_tok) { - if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out0_dim == 0 || out1_dim == 0 || n_tok == 0) { - return 0; - } - const uint64_t blocks = (in_dim + 31) / 32; - if (weight0_offset > model_size || weight1_offset > model_size || - out0_dim > UINT64_MAX / (blocks * 34) || - out1_dim > UINT64_MAX / (blocks * 34)) { - return 0; - } - const uint64_t weight0_bytes = out0_dim * blocks * 34; - const uint64_t weight1_bytes = out1_dim * blocks * 34; - if (weight0_bytes > model_size - weight0_offset || - weight1_bytes > model_size - weight1_offset || - x->bytes < in_dim * sizeof(float) || - out0->bytes < out0_dim * sizeof(float) || - out1->bytes < out1_dim * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out0); - const char *w0 = cuda_resolve_weight_ptr(model_map, weight0_offset, weight0_bytes, logical_tier, "q8_0_pair0"); - const char *w1 = cuda_resolve_weight_ptr(model_map, weight1_offset, weight1_bytes, logical_tier, "q8_0_pair1"); - if (!w0 || !w1) return 0; - - const bool force_decode_warp = - n_tok == 2u && g_glm_mtp_verify_mode; - if (n_tok != 1 && !force_decode_warp && !g_q8_cache_suppressed && - getenv("DS4_CUDA_Q8_PAIR_BATCH") == NULL) { - return cuda_matmul_q8_0_tensor_labeled(out0, model_map, model_size, weight0_offset, - in_dim, out0_dim, x, n_tok, "q8_0_pair0") && - cuda_matmul_q8_0_tensor_labeled(out1, model_map, model_size, weight1_offset, - in_dim, out1_dim, x, n_tok, "q8_0_pair1"); - } - - const uint64_t xq_bytes = n_tok * blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 pair prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1); - quantize_q8_0_f32_kernel<<>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) return 0; - if (n_tok != 1) { - if (force_decode_warp && - getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) { - matmul_q8_0_preq_batch_warp8_tok2_kernel - <<<((unsigned)out0_dim + 7u) / 8u, 256>>>( - (float *)out0->ptr, - reinterpret_cast(w0), - xq, xscale, in_dim, out0_dim, blocks, use_dp4a); - if (!cuda_ok(cudaGetLastError(), - "matmul_q8_0 pair0 tok2 warp launch")) { - return 0; - } - matmul_q8_0_preq_batch_warp8_tok2_kernel - <<<((unsigned)out1_dim + 7u) / 8u, 256>>>( - (float *)out1->ptr, - reinterpret_cast(w1), - xq, xscale, in_dim, out1_dim, blocks, use_dp4a); - return cuda_ok(cudaGetLastError(), - "matmul_q8_0 pair1 tok2 warp launch"); - } - const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks); - int mma_rc = cuda_q8_mma_try_launch( - (float *)out0->ptr, reinterpret_cast(w0), - xq, xscale, in_dim, out0_dim, n_tok, blocks, blocks, out0_dim, mma_T); - if (mma_rc < 0) return 0; - if (mma_rc > 0) { - mma_rc = cuda_q8_mma_try_launch( - (float *)out1->ptr, reinterpret_cast(w1), - xq, xscale, in_dim, out1_dim, n_tok, blocks, blocks, out1_dim, mma_T); - if (mma_rc > 0) return 1; - return 0; - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL && - blocks <= 32u && - n_tok >= 8u) { - dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); - matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( - (float *)out0->ptr, - reinterpret_cast(w0), - xq, - xscale, - in_dim, - out0_dim, - n_tok, - blocks, - use_dp4a); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok8 launch")) return 0; - dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1); - matmul_q8_0_preq_batch_warp8_tok8_kernel<<>>( - (float *)out1->ptr, - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out1_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok8 launch"); - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL && - blocks <= 32u && - n_tok >= 4u) { - dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); - matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( - (float *)out0->ptr, - reinterpret_cast(w0), - xq, - xscale, - in_dim, - out0_dim, - n_tok, - blocks, - use_dp4a); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok4 launch")) return 0; - dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1); - matmul_q8_0_preq_batch_warp8_tok4_kernel<<>>( - (float *)out1->ptr, - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out1_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok4 launch"); - } - if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL && - blocks <= 32u) { - dim3 grid0(((unsigned)out0_dim + 7u) / 8u, (unsigned)n_tok, 1); - matmul_q8_0_preq_batch_warp8_kernel<<>>( - (float *)out0->ptr, - reinterpret_cast(w0), - xq, - xscale, - in_dim, - out0_dim, - n_tok, - blocks, - use_dp4a); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch warp launch")) return 0; - dim3 grid1(((unsigned)out1_dim + 7u) / 8u, (unsigned)n_tok, 1); - matmul_q8_0_preq_batch_warp8_kernel<<>>( - (float *)out1->ptr, - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out1_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch warp launch"); - } - if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT") == NULL) { - const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; - const unsigned exact_threads = cuda_q8_exact_threads(blocks); - if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT_TOK2") == NULL && - n_tok >= 2u) { - dim3 grid((unsigned)max_out_dim, ((unsigned)n_tok + 1u) / 2u, 1); - matmul_q8_0_pair_preq_batch_tok2_exact_kernel<<>>( - (float *)out0->ptr, - (float *)out1->ptr, - reinterpret_cast(w0), - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out0_dim, - out1_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair exact tok2 launch"); - } - dim3 grid((unsigned)max_out_dim, (unsigned)n_tok, 1); - matmul_q8_0_pair_preq_batch_kernel<<>>( - (float *)out0->ptr, - (float *)out1->ptr, - reinterpret_cast(w0), - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out0_dim, - out1_dim, - n_tok, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair batch exact launch"); - } - const unsigned exact_threads = cuda_q8_exact_threads(blocks); - dim3 grid0((unsigned)out0_dim, (unsigned)n_tok, 1); - matmul_q8_0_preq_kernel<<>>((float *)out0->ptr, - reinterpret_cast(w0), - xq, - xscale, - in_dim, out0_dim, n_tok, blocks, - use_dp4a); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch launch")) return 0; - dim3 grid1((unsigned)out1_dim, (unsigned)n_tok, 1); - matmul_q8_0_preq_kernel<<>>((float *)out1->ptr, - reinterpret_cast(w1), - xq, - xscale, - in_dim, out1_dim, n_tok, blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch launch"); - } - const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; - matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>( - (float *)out0->ptr, - (float *)out1->ptr, - reinterpret_cast(w0), - reinterpret_cast(w1), - xq, - xscale, - in_dim, - out0_dim, - out1_dim, - blocks, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch"); -} - -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) { - if (!out || !x || !model_map || in_dim == 0u || out_dim == 0u || - n_rows == 0u || - x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || - out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) { - return 0; - } - const uint64_t blocks = (in_dim + 31u) / 32u; - if (weight_offset > model_size || - out_dim > UINT64_MAX / (blocks * 34u)) { - return 0; - } - const uint64_t weight_bytes = out_dim * blocks * 34u; - if (weight_bytes > model_size - weight_offset) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - if (logical_tier < 0 || logical_tier >= g_n_gpus || - ds4_tensor_device_idx(x) != logical_tier) { - return 0; - } - const char *wptr = cuda_resolve_weight_ptr( - model_map, weight_offset, weight_bytes, logical_tier, - "q8_0 decode rows exact"); - if (!wptr) return 0; - - const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = - scale_offset + (uint64_t)n_rows * blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on( - logical_tier, tmp_bytes, "q8_0 decode rows exact prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - dim3 qgrid((unsigned)blocks, n_rows, 1u); - quantize_q8_0_f32_kernel<<>>( - xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), - "q8_0 decode rows exact quantize launch")) { - return 0; - } - dim3 grid(((unsigned)out_dim + 7u) / 8u, n_rows, 1u); - matmul_q8_0_preq_warp8_kernel<<>>( - (float *)out->ptr, - reinterpret_cast(wptr), - xq, xscale, in_dim, out_dim, blocks, cuda_q8_use_dp4a()); - return cuda_ok(cudaGetLastError(), - "q8_0 decode rows exact warp launch"); -} - -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) { - if (!out0 || !out1 || !x || !model_map || in_dim == 0u || - out0_dim == 0u || out1_dim == 0u || n_rows == 0u || - x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || - out0->bytes < (uint64_t)n_rows * out0_dim * sizeof(float) || - out1->bytes < (uint64_t)n_rows * out1_dim * sizeof(float)) { - return 0; - } - const uint64_t blocks = (in_dim + 31u) / 32u; - if (weight0_offset > model_size || weight1_offset > model_size || - out0_dim > UINT64_MAX / (blocks * 34u) || - out1_dim > UINT64_MAX / (blocks * 34u)) { - return 0; - } - const uint64_t weight0_bytes = out0_dim * blocks * 34u; - const uint64_t weight1_bytes = out1_dim * blocks * 34u; - if (weight0_bytes > model_size - weight0_offset || - weight1_bytes > model_size - weight1_offset) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out0); - if (logical_tier < 0 || logical_tier >= g_n_gpus || - ds4_tensor_device_idx(out1) != logical_tier || - ds4_tensor_device_idx(x) != logical_tier) { - return 0; - } - const char *w0 = cuda_resolve_weight_ptr( - model_map, weight0_offset, weight0_bytes, logical_tier, - "q8_0 pair decode rows exact gate"); - const char *w1 = cuda_resolve_weight_ptr( - model_map, weight1_offset, weight1_bytes, logical_tier, - "q8_0 pair decode rows exact up"); - if (!w0 || !w1) return 0; - - const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = - scale_offset + (uint64_t)n_rows * blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on( - logical_tier, tmp_bytes, "q8_0 pair decode rows exact prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - dim3 qgrid((unsigned)blocks, n_rows, 1u); - quantize_q8_0_f32_kernel<<>>( - xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), - "q8_0 pair decode rows exact quantize launch")) { - return 0; - } - const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; - dim3 grid(((unsigned)max_out + 7u) / 8u, n_rows, 1u); - matmul_q8_0_pair_preq_warp8_kernel<<>>( - (float *)out0->ptr, - (float *)out1->ptr, - reinterpret_cast(w0), - reinterpret_cast(w1), - xq, xscale, in_dim, out0_dim, out1_dim, blocks, - cuda_q8_use_dp4a()); - return cuda_ok(cudaGetLastError(), - "q8_0 pair decode rows exact warp launch"); -} - -static int cuda_matmul_q8_0_hc_expand_tensor_labeled( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *block_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, - const ds4_gpu_tensor *block_add, - const ds4_gpu_tensor *block_add2, - const ds4_gpu_tensor *owned_home_slots, - const ds4_gpu_tensor *owned_peer_packed, - const ds4_gpu_tensor *owned_selected, - uint32_t owned_expert_split, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc, - const char *label) { - if (!out_hc || !block_out || !x || !residual_hc || !split || !model_map || - in_dim == 0 || out_dim == 0 || n_embd == 0 || n_hc == 0 || - out_dim != (uint64_t)n_embd) { - return 0; - } - const uint64_t blocks = (in_dim + 31) / 32; - if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0; - const uint64_t weight_bytes = out_dim * blocks * 34; - const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t split_bytes = (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float); - if (weight_bytes > model_size - weight_offset || - x->bytes < in_dim * sizeof(float) || - block_out->bytes < out_dim * sizeof(float) || - residual_hc->bytes < hc_bytes || - split->bytes < split_bytes || - out_hc->bytes < hc_bytes || - (block_add && block_add->bytes < out_dim * sizeof(float)) || - (block_add2 && block_add2->bytes < out_dim * sizeof(float)) || - ((owned_home_slots || owned_peer_packed || owned_selected) && - (!owned_home_slots || !owned_peer_packed || !owned_selected || - owned_expert_split == 0u || - owned_home_slots->bytes < 6u * out_dim * sizeof(float) || - owned_peer_packed->bytes < 4u * out_dim * sizeof(float) || - owned_selected->bytes < 6u * sizeof(int32_t)))) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out_hc); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, label ? label : "q8_0_hc_expand"); - if (!wptr) return 0; - - const uint64_t xq_bytes = blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 hc expand prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) return 0; - matmul_q8_0_hc_expand_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( - (float *)out_hc->ptr, - (float *)block_out->ptr, - block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr, - block_add2 ? (const float *)block_add2->ptr : (const float *)block_out->ptr, - owned_home_slots ? (const float *)owned_home_slots->ptr : NULL, - owned_peer_packed ? (const float *)owned_peer_packed->ptr : NULL, - owned_selected ? (const int32_t *)owned_selected->ptr : NULL, - (const float *)residual_hc->ptr, - (const float *)split->ptr, - reinterpret_cast(wptr), - xq, - xscale, - in_dim, - out_dim, - n_embd, - n_hc, - blocks, - block_add ? 1 : 0, - block_add2 ? 1 : 0, - owned_home_slots ? 1 : 0, - owned_expert_split, - use_dp4a); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand launch"); -} - -extern "C" int ds4_gpu_matmul_f16_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_tok) { - if (!out || !x || !model_map) return 0; - if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0; - uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t); - if (weight_bytes > model_size - weight_offset) return 0; - if (x->bytes < n_tok * in_dim * sizeof(float) || - out->bytes < n_tok * out_dim * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f16"); - if (!wptr) return 0; - const __half *w = (const __half *)wptr; - const int serial_f16 = getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL; - const int router_shape = in_dim == 4096u && out_dim == 256u && n_tok == 1u; - const int serial_router = - !serial_f16 && - router_shape && - getenv("DS4_CUDA_SERIAL_ROUTER") != NULL; - const int ordered_router = - !serial_f16 && - !serial_router && - n_tok == 1u && - getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL; - const int small_out_one_token = - !serial_f16 && - !serial_router && - !g_quality_mode && - n_tok == 1u && - out_dim <= 32u && - in_dim >= 8192u && - getenv("DS4_CUDA_F16_SMALL_OUT") != NULL && - getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL && - getenv("DS4_CUDA_NO_F16_SMALL_OUT") == NULL; - if (small_out_one_token) { - matmul_f16_small_out_hx_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>( - (float *)out->ptr, - w, - (const float *)x->ptr, - in_dim, - out_dim); - return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_hx_ordered_chunks launch"); - } - const int small_out_batch = - !serial_f16 && - !serial_router && - n_tok > 1u && - out_dim <= 32u && - in_dim >= 4096u && - (g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) && - getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL; - if (small_out_batch) { - matmul_f16_small_out_batch_kernel<<<(unsigned)n_tok, 256>>>( - (float *)out->ptr, - w, - (const float *)x->ptr, - in_dim, - out_dim, - n_tok); - return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_batch launch"); - } - const int cublas_one_token = - n_tok == 1u && - getenv("DS4_CUDA_NO_F16_CUBLAS_ONE") == NULL && - (!g_quality_mode || getenv("DS4_CUDA_F16_CUBLAS_ONE") != NULL); - const int cublas_batch = - n_tok > 1u && getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL; - if (!serial_f16 && g_cublas_ready && (cublas_batch || cublas_one_token)) { - const uint64_t xh_count = n_tok * in_dim; - __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "f16 gemm activations"); - if (!xh) return 0; - f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count); - if (!cuda_ok(cudaGetLastError(), "f16 activation convert launch")) return 0; - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w, - CUDA_R_16F, - (int)in_dim, - xh, - CUDA_R_16F, - (int)in_dim, - &beta, - out->ptr, - CUDA_R_32F, - (int)out_dim, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - return cublas_ok(st, "f16 matmul"); - } - dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); - if (serial_f16 || serial_router) { - matmul_f16_serial_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); - return cuda_ok(cudaGetLastError(), serial_router ? "matmul_f16_router_serial launch" : "matmul_f16_serial launch"); - } - if (ordered_router) { - matmul_f16_ordered_chunks_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); - return cuda_ok(cudaGetLastError(), "matmul_f16_ordered_chunks launch"); - } - matmul_f16_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); - return cuda_ok(cudaGetLastError(), "matmul_f16 launch"); -} - -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) { - const uint64_t in_dim = 4096u; - const uint64_t out_dim = 256u; - if (!out || !x || !model_map || n_rows == 0u || - weight_offset > model_size) { - return 0; - } - const uint64_t weight_bytes = in_dim * out_dim * sizeof(uint16_t); - if (weight_bytes > model_size - weight_offset || - x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) || - out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) { - return 0; - } - if (n_rows == 1u) { - return ds4_gpu_matmul_f16_tensor( - out, model_map, model_size, weight_offset, - in_dim, out_dim, x, 1); - } - const int logical_tier = ds4_tensor_device_idx(out); - if (ds4_tensor_device_idx(x) != logical_tier || !g_cublas_ready) return 0; - const __half *w = (const __half *)cuda_resolve_weight_ptr( - model_map, weight_offset, weight_bytes, logical_tier, - "f16_router_rows_exact"); - if (!w) return 0; - - const uint64_t xh_count = (uint64_t)n_rows * in_dim; - __half *xh = (__half *)cuda_tmp_alloc_on( - logical_tier, xh_count * sizeof(__half), - "f16 exact router batch activations"); - if (!xh) return 0; - f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>( - xh, (const float *)x->ptr, xh_count); - if (!cuda_ok(cudaGetLastError(), - "f16 exact router activation convert launch")) { - return 0; - } - const float alpha = 1.0f; - const float beta = 0.0f; - /* Larger batchCount values let cuBLAS select a different reduction and - * change logits. Four-row calls match the one-row decode bit for bit on - * this projection, while still replacing most per-session launches. */ - uint32_t row = 0; - for (; row + 4u <= n_rows; row += 4u) { - cublasStatus_t st = cublasGemmStridedBatchedEx( - cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - 1, - (int)in_dim, - &alpha, - w, - CUDA_R_16F, - (int)in_dim, - 0, - xh + (uint64_t)row * in_dim, - CUDA_R_16F, - (int)in_dim, - (long long int)in_dim, - &beta, - (float *)out->ptr + (uint64_t)row * out_dim, - CUDA_R_32F, - (int)out_dim, - (long long int)out_dim, - 4, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - if (!cublas_ok(st, "f16 exact router row batch")) return 0; - } - for (; row < n_rows; row++) { - ds4_gpu_tensor out_row = *out; - ds4_gpu_tensor x_row = *x; - out_row.ptr = (float *)out->ptr + (uint64_t)row * out_dim; - out_row.bytes = out_dim * sizeof(float); - x_row.ptr = (float *)x->ptr + (uint64_t)row * in_dim; - x_row.bytes = in_dim * sizeof(float); - if (!ds4_gpu_matmul_f16_tensor( - &out_row, model_map, model_size, weight_offset, - in_dim, out_dim, &x_row, 1)) { - return 0; - } - } - return 1; -} - -extern "C" int ds4_gpu_matmul_f16_pair_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 out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) { - return 0; - } - if (getenv("DS4_CUDA_NO_F16_PAIR_MATMUL") != NULL || - getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL || - getenv("DS4_CUDA_SERIAL_ROUTER") != NULL || - getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") != NULL) { - return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, - in_dim, out_dim, x, n_tok) && - ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, - in_dim, out_dim, x, n_tok); - } - if (weight0_offset > model_size || weight1_offset > model_size || - out_dim > UINT64_MAX / in_dim || - n_tok > UINT64_MAX / in_dim || - n_tok > UINT64_MAX / out_dim) { - return 0; - } - const uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (weight_bytes > model_size - weight0_offset || - weight_bytes > model_size - weight1_offset || - x->bytes < x_bytes || - out0->bytes < out_bytes || - out1->bytes < out_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out0); - if (ds4_tensor_device_idx(out1) != logical_tier) { - return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, - in_dim, out_dim, x, n_tok) && - ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, - in_dim, out_dim, x, n_tok); - } - const __half *w0 = (const __half *)cuda_resolve_weight_ptr(model_map, weight0_offset, weight_bytes, logical_tier, "f16_pair0"); - const __half *w1 = (const __half *)cuda_resolve_weight_ptr(model_map, weight1_offset, weight_bytes, logical_tier, "f16_pair1"); - if (!w0 || !w1) return 0; - if (n_tok > 1) { - const bool small_out_batch_requested = - out_dim <= 32u && - in_dim >= 4096u && - (g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) && - getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL; - if (!small_out_batch_requested && - g_cublas_ready && - getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL) { - const uint64_t xh_count = n_tok * in_dim; - __half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, - xh_count * sizeof(__half), - "f16 pair gemm activations"); - if (!xh) return 0; - f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>( - xh, (const float *)x->ptr, xh_count); - if (!cuda_ok(cudaGetLastError(), "f16 pair activation convert launch")) return 0; - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w0, - CUDA_R_16F, - (int)in_dim, - xh, - CUDA_R_16F, - (int)in_dim, - &beta, - out0->ptr, - CUDA_R_32F, - (int)out_dim, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - if (!cublas_ok(st, "f16 pair matmul0")) return 0; - st = cublasGemmEx(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w1, - CUDA_R_16F, - (int)in_dim, - xh, - CUDA_R_16F, - (int)in_dim, - &beta, - out1->ptr, - CUDA_R_32F, - (int)out_dim, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - return cublas_ok(st, "f16 pair matmul1"); - } - return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset, - in_dim, out_dim, x, n_tok) && - ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset, - in_dim, out_dim, x, n_tok); - } - matmul_f16_pair_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>( - (float *)out0->ptr, - (float *)out1->ptr, - w0, - w1, - (const float *)x->ptr, - in_dim, - out_dim, - out_dim); - return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch"); -} - -extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor( - ds4_gpu_tensor *out_kv, - ds4_gpu_tensor *out_score, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const void *model_map, - uint64_t model_size, - uint64_t weight_kv_offset, - uint64_t weight_score_offset, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t in_dim, - uint32_t width, - const ds4_gpu_tensor *x, - uint32_t ratio, - uint32_t pos) { - (void)out_kv; - (void)out_score; - (void)state_kv; - (void)state_score; - (void)model_map; - (void)model_size; - (void)weight_kv_offset; - (void)weight_score_offset; - (void)ape_offset; - (void)ape_type; - (void)in_dim; - (void)width; - (void)x; - (void)ratio; - (void)pos; - return 0; -} - -extern "C" int ds4_gpu_matmul_f32_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_tok) { - if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) return 0; - if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0; - uint64_t weight_elems = out_dim * in_dim; - if (weight_elems > UINT64_MAX / sizeof(float)) return 0; - uint64_t weight_bytes = weight_elems * sizeof(float); - if (weight_bytes > model_size - weight_offset) return 0; - if (x->bytes < n_tok * in_dim * sizeof(float) || - out->bytes < n_tok * out_dim * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f32"); - if (!wptr) return 0; - const float *w = (const float *)wptr; - if (g_cublas_ready && n_tok > 1) { - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)out_dim, - (int)n_tok, - (int)in_dim, - &alpha, - w, - (int)in_dim, - (const float *)x->ptr, - (int)in_dim, - &beta, - (float *)out->ptr, - (int)out_dim); - return cublas_ok(st, "f32 matmul"); - } - dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1); - matmul_f32_kernel<<>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok); - return cuda_ok(cudaGetLastError(), "matmul_f32 launch"); -} - -extern "C" int ds4_gpu_repeat_hc_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *row, uint32_t n_embd, uint32_t n_hc) { - if (!out || !row || n_embd == 0 || n_hc == 0 || - row->bytes < (uint64_t)n_embd * sizeof(float) || - out->bytes < (uint64_t)n_embd * n_hc * sizeof(float)) { - return 0; - } - uint64_t n = (uint64_t)n_embd * n_hc; - repeat_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)row->ptr, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "repeat_hc launch"); -} - - -/* Non-causal batch attention over a raw KV ring for the DSpark draft block. - * Every query row attends over all n_raw visible rows plus the per-head sink, - * with the same exact one-block max/denominator/value accumulation order as - * the reference decode attention (scores in shared, sequential value pass). */ -__global__ static void attention_noncausal_raw_batch_heads_kernel( - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - uint32_t n_tokens, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t tok = blockIdx.x; - const uint32_t h = blockIdx.y; - if (tok >= n_tokens || h >= n_head) return; - extern __shared__ float sh_scores[]; /* n_raw floats */ - const float *qh = q + ((uint64_t)tok * n_head + h) * head_dim; - const float scale = rsqrtf((float)head_dim); - for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { - const uint32_t row = (raw_start + r) % raw_cap; - const float *kv = raw_kv + (uint64_t)row * head_dim; - float dot = 0.0f; - for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; - sh_scores[r] = dot * scale; - } - __syncthreads(); - __shared__ float partial[256]; - __shared__ float max_s; - __shared__ float denom; - float local_max = sinks[h]; - for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { - local_max = fmaxf(local_max, sh_scores[r]); - } - partial[threadIdx.x] = local_max; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); - } - __syncthreads(); - } - if (threadIdx.x == 0) max_s = partial[0]; - __syncthreads(); - float den_local = 0.0f; - for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { - sh_scores[r] = expf(sh_scores[r] - max_s); - den_local += sh_scores[r]; - } - partial[threadIdx.x] = den_local; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); - __syncthreads(); - float *oh = heads + ((uint64_t)tok * n_head + h) * head_dim; - for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { - float acc = 0.0f; - for (uint32_t r = 0; r < n_raw; r++) { - const uint32_t row = (raw_start + r) % raw_cap; - acc += raw_kv[(uint64_t)row * head_dim + d] * sh_scores[r]; - } - oh[d] = acc / denom; - } -} - -extern "C" int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_head, - uint32_t head_dim) { - if (!heads || !q || !raw_kv || !model_map || - n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || - raw_start >= raw_cap || n_head == 0 || head_dim == 0 || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, - "dspark_attn_sinks"); - if (!sinks) return 0; - const size_t shmem = (size_t)n_raw * sizeof(float); - if (shmem > 32768) return 0; /* draft blocks are tiny; guard anyway */ - dim3 grid(n_tokens, n_head, 1); - attention_noncausal_raw_batch_heads_kernel<<>>( - (float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_tokens, n_raw, raw_cap, raw_start, n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), "attention noncausal raw batch heads launch")) return 0; - static int verify_left = -1; - if (verify_left < 0) { - verify_left = getenv("DS4_DSPARK_VERIFY_NONCAUSAL") != NULL ? 3 : 0; - } - if (verify_left > 0) { - verify_left--; - (void)cudaDeviceSynchronize(); - const uint64_t qn = (uint64_t)n_tokens * n_head * head_dim; - const uint64_t kn = (uint64_t)raw_cap * head_dim; - std::vector hq(qn), hkv(kn), hout(qn), hsink(n_head); - (void)cudaMemcpy(hq.data(), q->ptr, qn * 4, cudaMemcpyDeviceToHost); - (void)cudaMemcpy(hkv.data(), raw_kv->ptr, kn * 4, cudaMemcpyDeviceToHost); - (void)cudaMemcpy(hout.data(), heads->ptr, qn * 4, cudaMemcpyDeviceToHost); - (void)cudaMemcpy(hsink.data(), sinks, (uint64_t)n_head * 4, cudaMemcpyDeviceToHost); - double max_abs = 0.0, max_rel = 0.0; - const double scale = 1.0 / sqrt((double)head_dim); - for (uint32_t t = 0; t < n_tokens; t++) { - for (uint32_t h = 0; h < n_head; h++) { - std::vector sc(n_raw); - double mx = (double)hsink[h]; - for (uint32_t r = 0; r < n_raw; r++) { - const uint32_t row = (raw_start + r) % raw_cap; - double dot = 0.0; - for (uint32_t d = 0; d < head_dim; d++) { - dot += (double)hq[((uint64_t)t * n_head + h) * head_dim + d] * - (double)hkv[(uint64_t)row * head_dim + d]; - } - sc[r] = dot * scale; - if (sc[r] > mx) mx = sc[r]; - } - double den = exp((double)hsink[h] - mx); - for (uint32_t r = 0; r < n_raw; r++) den += exp(sc[r] - mx); - for (uint32_t d = 0; d < head_dim; d++) { - double acc = 0.0; - for (uint32_t r = 0; r < n_raw; r++) { - const uint32_t row = (raw_start + r) % raw_cap; - acc += exp(sc[r] - mx) * (double)hkv[(uint64_t)row * head_dim + d]; - } - const double ref = acc / den; - const double got = (double)hout[((uint64_t)t * n_head + h) * head_dim + d]; - const double ad = fabs(ref - got); - if (ad > max_abs) max_abs = ad; - if (fabs(ref) > 1e-3 && ad / fabs(ref) > max_rel) max_rel = ad / fabs(ref); - } - } - } - fprintf(stderr, - "ds4: DSpark noncausal verify n_tok=%u n_raw=%u start=%u cap=%u " - "max_abs=%.3e max_rel=%.3e\n", - n_tokens, n_raw, raw_start, raw_cap, max_abs, max_rel); - } - return 1; -} - -extern "C" int ds4_gpu_repeat_hc_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) { - uint64_t rows_elems = 0; - uint64_t out_elems = 0; - if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0 || - (uint64_t)n_tokens > UINT64_MAX / n_embd || - (rows_elems = (uint64_t)n_tokens * n_embd) > UINT64_MAX / n_hc || - (out_elems = rows_elems * n_hc) > UINT64_MAX / sizeof(float) || - rows_elems > UINT64_MAX / sizeof(float) || - rows->bytes < rows_elems * sizeof(float) || - out->bytes < out_elems * sizeof(float)) { - return 0; - } - const uint64_t blocks = (out_elems + 255u) / 256u; - if (blocks > UINT32_MAX) return 0; - repeat_hc_rows_kernel<<<(unsigned)blocks, 256>>>((float *)out->ptr, (const float *)rows->ptr, n_tokens, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "repeat_hc_rows launch"); -} - -extern "C" int ds4_gpu_rms_norm_plain_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, float eps) { - if (!out || !x || out->bytes < (uint64_t)n * sizeof(float) || - x->bytes < (uint64_t)n * sizeof(float)) return 0; - if (n == 4096u) { - rms_norm_plain_fast4096_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); - } else if ((n & 2047u) == 0u) { - rms_norm_plain_batch8_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); - } else { - rms_norm_plain_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps); - } - return cuda_ok(cudaGetLastError(), "rms_norm_plain launch"); -} -extern "C" int ds4_gpu_rms_norm_plain_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, uint32_t rows, float eps) { - if (!out || !x || out->bytes < (uint64_t)n * rows * sizeof(float) || - x->bytes < (uint64_t)n * rows * sizeof(float)) return 0; - if (n == 4096u) { - rms_norm_plain_fast4096_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); - } else if ((n & 2047u) == 0u) { - rms_norm_plain_batch8_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); - } else { - rms_norm_plain_kernel<<>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps); - } - return cuda_ok(cudaGetLastError(), "rms_norm_plain launch"); -} -extern "C" int ds4_gpu_rms_norm_weight_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, float eps) { - if (!out || !x || !model_map || weight_offset > model_size || - model_size - weight_offset < (uint64_t)n * sizeof(float) || - out->bytes < (uint64_t)n * sizeof(float) || - x->bytes < (uint64_t)n * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight"); - if (!wptr) return 0; - const float *w = (const float *)wptr; - rms_norm_weight_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, w, n, 1, eps); - return cuda_ok(cudaGetLastError(), "rms_norm_weight launch"); -} -extern "C" int ds4_gpu_rms_norm_weight_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, uint32_t rows, float eps) { - if (!out || !x || !model_map || weight_offset > model_size || - model_size - weight_offset < (uint64_t)n * sizeof(float) || - out->bytes < (uint64_t)n * rows * sizeof(float) || - x->bytes < (uint64_t)n * rows * sizeof(float)) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight"); - if (!wptr) return 0; - const float *w = (const float *)wptr; - rms_norm_weight_kernel<<>>((float *)out->ptr, (const float *)x->ptr, w, n, rows, eps); - return cuda_ok(cudaGetLastError(), "rms_norm_weight launch"); -} -extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_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, - float eps) { - if (!g_cuda_disable_qkv_rms_fused) { - if (!q_out || !q || !kv_out || !kv || !model_map || - q_weight_offset > model_size || - kv_weight_offset > model_size || - model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) || - model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) || - q_out->bytes < (uint64_t)q_n * rows * sizeof(float) || - q->bytes < (uint64_t)q_n * rows * sizeof(float) || - kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) || - kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(q_out); - const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map, - q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight"); - const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map, - kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight"); - if (!q_w || !kv_w) return 0; - dim3 grid(rows, 2u, 1u); - dsv4_qkv_rms_norm_rows_kernel<<>>( - (float *)q_out->ptr, - (const float *)q->ptr, - q_w, - q_n, - (float *)kv_out->ptr, - (const float *)kv->ptr, - kv_w, - kv_n, - rows, - eps); - return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm rows launch"); - } - return ds4_gpu_rms_norm_weight_rows_tensor(q_out, q, model_map, model_size, - q_weight_offset, q_n, rows, eps) && - ds4_gpu_rms_norm_weight_rows_tensor(kv_out, kv, model_map, model_size, - kv_weight_offset, kv_n, rows, eps); -} - -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) { - if (g_cuda_disable_qkv_rms_fused) return 0; - if (!q_out || !q || !kv_out || !kv || !model_map || - q_weight_offset > model_size || - kv_weight_offset > model_size || - kv_n_head == 0 || kv_head_dim == 0 || - n_rot > kv_head_dim || (n_rot & 1u) || - kv_n != kv_n_head * kv_head_dim || - model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) || - model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) || - q_out->bytes < (uint64_t)q_n * rows * sizeof(float) || - q->bytes < (uint64_t)q_n * rows * sizeof(float) || - kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) || - kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(q_out); - const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map, - q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight"); - const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map, - kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight"); - if (!q_w || !kv_w) return 0; - dim3 grid(rows, 2u, 1u); - dsv4_qkv_rms_norm_rows_kv_rope_kernel<<>>( - (float *)q_out->ptr, - (const float *)q->ptr, - q_w, - q_n, - (float *)kv_out->ptr, - (const float *)kv->ptr, - kv_w, - kv_n, - rows, - kv_n_head, - kv_head_dim, - n_rot, - pos0, - n_ctx_orig, - inverse ? 1 : 0, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow, - eps); - return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm kv rope launch"); -} - -extern "C" int ds4_gpu_head_rms_norm_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) { - if (!x || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; - head_rms_norm_kernel<<>>((float *)x->ptr, n_tok, n_head, head_dim, eps); - return cuda_ok(cudaGetLastError(), "head_rms_norm launch"); -} -extern "C" int ds4_gpu_head_rms_norm_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t 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) { - if (!x || n_rot > head_dim || (n_rot & 1u) || - x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; - head_rms_norm_rope_tail_kernel<<>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, eps); - return cuda_ok(cudaGetLastError(), "head_rms_norm_rope_tail launch"); -} -extern "C" int ds4_gpu_dsv4_fp8_kv_quantize_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) { - if (!x || n_rot > head_dim || x->bytes < (uint64_t)n_tok * head_dim * sizeof(float)) return 0; - fp8_kv_quantize_kernel<<>>((float *)x->ptr, n_tok, head_dim, n_rot); - return cuda_ok(cudaGetLastError(), "fp8_kv_quantize launch"); -} -extern "C" int ds4_gpu_dsv4_indexer_qat_tensor(ds4_gpu_tensor *x, uint32_t n_rows, uint32_t head_dim) { - if (!x || n_rows == 0 || head_dim != 128u || - x->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) { - return 0; - } - indexer_hadamard_fp4_kernel<<>>((float *)x->ptr, n_rows, head_dim); - return cuda_ok(cudaGetLastError(), "indexer_hadamard_fp4 launch"); -} -extern "C" int ds4_gpu_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t 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) { - if (!x || n_rot > head_dim || (n_rot & 1) || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; - uint32_t pairs = n_tok * n_head * (n_rot / 2); - rope_tail_kernel<<<(pairs + 255) / 256, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, 1, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - return cuda_ok(cudaGetLastError(), "rope_tail launch"); -} -extern "C" int ds4_gpu_rope_tail_decode_rows_tensor( - ds4_gpu_tensor *x, - const ds4_gpu_attention_decode_row *rows, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - 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) { - if (!x || !rows || n_rows == 0u || - n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u || - n_rot == 0u || n_rot > head_dim || (n_rot & 1u) != 0u || - x->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) { - return 0; - } - cuda_attention_decode_row_table table; - memset(&table, 0, sizeof(table)); - for (uint32_t i = 0; i < n_rows; i++) table.row[i].pos = rows[i].pos; - const uint32_t pairs = n_rows * n_head * (n_rot / 2u); - rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>( - (float *)x->ptr, table, n_rows, n_head, head_dim, n_rot, - n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - return cuda_ok(cudaGetLastError(), "rope tail decode rows launch"); -} -extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim); -extern "C" int ds4_gpu_kv_fp8_store_raw_tensor( - ds4_gpu_tensor *kv, - ds4_gpu_tensor *raw_cache, - uint32_t raw_cap, - uint32_t raw_row, - uint32_t head_dim, - uint32_t n_rot) { - if (!kv || !raw_cache || raw_cap == 0u || n_rot > head_dim || - kv->device_id != raw_cache->device_id || - kv->bytes < (uint64_t)head_dim * sizeof(float) || - raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) { - return 0; - } - cuda_attention_decode_row_table table; - memset(&table, 0, sizeof(table)); - table.row[0].raw_kv = (uint64_t)(uintptr_t)raw_cache->ptr; - table.row[0].raw_cap = raw_cap; - table.row[0].raw_start = raw_row % raw_cap; - fp8_kv_quantize_store_rows_kernel<<<1, 64>>>( - (float *)kv->ptr, table, 1u, head_dim, n_rot); - return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store launch"); -} - -extern "C" int ds4_gpu_kv_fp8_store_raw_decode_rows_tensor( - ds4_gpu_tensor *kv, - ds4_gpu_tensor *const *raw_caches, - const uint32_t *raw_caps, - const uint32_t *raw_rows, - uint32_t n_rows, - uint32_t head_dim, - uint32_t n_rot) { - if (!kv || !raw_caches || !raw_caps || !raw_rows || n_rows == 0u || - n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_rot > head_dim || - kv->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) { - return 0; - } - cuda_attention_decode_row_table table; - memset(&table, 0, sizeof(table)); - for (uint32_t i = 0; i < n_rows; i++) { - const ds4_gpu_tensor *raw = raw_caches[i]; - if (!raw || raw_caps[i] == 0u || raw_rows[i] >= raw_caps[i] || - raw->device_id != kv->device_id || - raw->bytes < (uint64_t)raw_caps[i] * head_dim * sizeof(float)) { - return 0; - } - table.row[i].raw_kv = (uint64_t)(uintptr_t)raw->ptr; - table.row[i].raw_cap = raw_caps[i]; - table.row[i].raw_start = raw_rows[i]; - } - fp8_kv_quantize_store_rows_kernel<<>>( - (float *)kv->ptr, table, n_rows, head_dim, n_rot); - return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store rows launch"); -} -extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim) { - if (!raw_cache || !kv || raw_cap == 0 || - raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - kv->bytes < (uint64_t)head_dim * sizeof(float)) return 0; - store_raw_kv_batch_kernel<<<(head_dim + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, row, 1, head_dim); - return cuda_ok(cudaGetLastError(), "store_raw_kv launch"); -} -extern "C" int ds4_gpu_store_raw_kv_batch_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) { - if (!raw_cache || !kv || raw_cap == 0 || - raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float)) return 0; - uint64_t n = (uint64_t)n_tokens * head_dim; - store_raw_kv_batch_kernel<<<(n + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, pos0, n_tokens, head_dim); - return cuda_ok(cudaGetLastError(), "store_raw_kv_batch launch"); -} -extern "C" int ds4_gpu_compressor_store_batch_tensor( - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos0, - uint32_t n_tokens) { - if (!kv || !sc || !state_kv || !state_score || !model_map || - head_dim == 0 || ratio == 0 || n_tokens == 0 || - (ape_type != 0u && ape_type != 1u)) { - return 0; - } - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - kv->bytes < kv_bytes || sc->bytes < kv_bytes || - state_kv->bytes < state_bytes || state_score->bytes < state_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(state_kv); - const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); - if (!ape) return 0; - uint64_t n = (uint64_t)n_tokens * width; - compressor_store_kernel<<<(n + 255) / 256, 256>>>( - (const float *)kv->ptr, - (const float *)sc->ptr, - (float *)state_kv->ptr, - (float *)state_score->ptr, - ape, - 0, - ape_type, - head_dim, - ratio, - pos0, - n_tokens); - return cuda_ok(cudaGetLastError(), "compressor store launch"); -} - -extern "C" int ds4_gpu_compressor_update_tensor( - const ds4_gpu_tensor *kv_cur, - const ds4_gpu_tensor *sc_cur, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - ds4_gpu_tensor *comp_cache, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos, - uint32_t comp_row, - uint32_t n_rot, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps, - bool state_already_stored) { - if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || - !model_map || head_dim == 0 || ratio == 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { - return 0; - } - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset || - kv_cur->bytes < kv_bytes || sc_cur->bytes < kv_bytes || - state_kv->bytes < state_bytes || state_score->bytes < state_bytes || - (emit && comp_cache->bytes < comp_bytes)) { - return 0; - } - if (!state_already_stored) { - if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score, - model_map, model_size, ape_offset, ape_type, - head_dim, ratio, pos, 1)) { - return 0; - } - } - if (!emit) return 1; - ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( - comp_cache, - (uint64_t)comp_row * head_dim * sizeof(float), - (uint64_t)head_dim * sizeof(float)); - if (!comp_row_view) return 0; - compressor_update_pool_kernel<<<(head_dim + 255) / 256, 256>>>( - (float *)comp_row_view->ptr, - (const float *)state_kv->ptr, - (const float *)state_score->ptr, - head_dim, - ratio); - int ok = cuda_ok(cudaGetLastError(), "compressor update pool launch"); - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view, - model_map, model_size, norm_offset, - head_dim, 1, rms_eps); - if (ok) ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, 1, head_dim, n_rot, - pos + 1u - ratio, n_ctx_orig, false, - freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow); - ds4_gpu_tensor_free(comp_row_view); - if (ok && ratio == 4u) { - uint64_t half = 4ull * width; - compressor_shift_ratio4_kernel<<<(half + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, width); - ok = cuda_ok(cudaGetLastError(), "compressor ratio4 shift launch"); - } - return ok; -} -extern "C" int ds4_gpu_compressor_prefill_tensor( - ds4_gpu_tensor *comp_cache, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_rot, - uint32_t n_ctx_orig, - bool quantize_fp8, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps) { - if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || - head_dim == 0 || ratio == 0 || n_tokens == 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { - return 0; - } - - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint32_t n_comp = n_tokens / ratio; - const uint32_t cutoff = n_comp * ratio; - const uint32_t rem = n_tokens - cutoff; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset || - kv->bytes < kv_bytes || sc->bytes < kv_bytes || - state_kv->bytes < state_bytes || state_score->bytes < state_bytes || - (n_comp && comp_cache->bytes < comp_bytes)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(state_kv); - const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); - if (!ape) return 0; - - uint64_t state_n = (uint64_t)state_rows * width; - if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), - "compressor state kv zero")) return 0; - fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); - if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0; - - if (ratio == 4u) { - if (cutoff >= ratio) { - uint32_t prev_start = cutoff - ratio; - uint64_t n = (uint64_t)ratio * width; - compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, - (const float *)kv->ptr, (const float *)sc->ptr, - ape, 0, ape_type, width, ratio, pos0, - prev_start, 0, ratio); - if (!cuda_ok(cudaGetLastError(), "compressor prefill prev state launch")) return 0; - } - if (rem != 0) { - uint64_t n = (uint64_t)rem * width; - compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, - (const float *)kv->ptr, (const float *)sc->ptr, - ape, 0, ape_type, width, ratio, pos0, - cutoff, ratio, rem); - if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0; - } - } else if (rem != 0) { - uint64_t n = (uint64_t)rem * width; - compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, - (const float *)kv->ptr, (const float *)sc->ptr, - ape, 0, ape_type, width, ratio, pos0, - cutoff, 0, rem); - if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0; - } - if (n_comp != 0) { - dim3 grid((head_dim + 255) / 256, n_comp, 1); - compressor_prefill_pool_kernel<<>>( - (float *)comp_cache->ptr, - (const float *)kv->ptr, - (const float *)sc->ptr, - (const float *)state_kv->ptr, - (const float *)state_score->ptr, - ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 0); - if (!cuda_ok(cudaGetLastError(), "compressor prefill pool launch")) return 0; - if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache, - model_map, model_size, norm_offset, - head_dim, n_comp, rms_eps)) return 0; - if (n_rot != 0) { - const uint32_t pairs = n_comp * (n_rot / 2u); - rope_tail_kernel<<<(pairs + 255) / 256, 256>>>( - (float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot, - pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - if (!cuda_ok(cudaGetLastError(), "compressor prefill rope launch")) return 0; - } - if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0; - } - return 1; -} -extern "C" int ds4_gpu_compressor_prefill_ratio4_replay_tensor( - ds4_gpu_tensor *comp_cache, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_rot, - uint32_t n_ctx_orig, - bool quantize_fp8, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps) { - if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || - head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || norm_type != 0u) { - return 0; - } - - const uint32_t ratio = 4u; - const uint32_t width = 2u * head_dim; - const uint32_t state_rows = 8u; - const uint32_t n_comp = n_tokens / ratio; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset || - kv->bytes < kv_bytes || sc->bytes < kv_bytes || - state_kv->bytes < state_bytes || state_score->bytes < state_bytes || - comp_cache->bytes < comp_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(comp_cache); - const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); - if (!ape) return 0; - dim3 grid((head_dim + 255) / 256, n_comp, 1); - compressor_prefill_pool_kernel<<>>( - (float *)comp_cache->ptr, - (const float *)kv->ptr, - (const float *)sc->ptr, - (const float *)state_kv->ptr, - (const float *)state_score->ptr, - ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 1); - if (!cuda_ok(cudaGetLastError(), "compressor replay pool launch")) return 0; - if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache, - model_map, model_size, norm_offset, - head_dim, n_comp, rms_eps)) return 0; - if (n_rot != 0) { - const uint32_t pairs = n_comp * (n_rot / 2u); - rope_tail_kernel<<<(pairs + 255) / 256, 256>>>( - (float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot, - pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - if (!cuda_ok(cudaGetLastError(), "compressor replay rope launch")) return 0; - } - if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0; - - uint64_t state_n = (uint64_t)state_rows * width; - if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), - "compressor replay state kv zero")) return 0; - fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); - if (!cuda_ok(cudaGetLastError(), "compressor replay state score fill launch")) return 0; - uint32_t prev_start = n_tokens - ratio; - uint64_t n = (uint64_t)ratio * width; - compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, - (const float *)kv->ptr, (const float *)sc->ptr, - ape, 0, ape_type, width, ratio, pos0, - prev_start, 0, ratio); - return cuda_ok(cudaGetLastError(), "compressor replay state launch"); -} -extern "C" int ds4_gpu_compressor_prefill_state_ratio4_tensor( - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv_tail, - const ds4_gpu_tensor *sc_tail, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t head_dim, - uint32_t pos0) { - if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map || - head_dim == 0 || (ape_type != 0u && ape_type != 1u)) { - return 0; - } - const uint32_t ratio = 4u; - const uint32_t width = 2u * head_dim; - const uint32_t state_rows = 8u; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape; - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - kv_tail->bytes < tail_bytes || sc_tail->bytes < tail_bytes || - state_kv->bytes < state_bytes || state_score->bytes < state_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(state_kv); - const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape"); - if (!ape) return 0; - uint64_t state_n = (uint64_t)state_rows * width; - if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))), - "compressor state kv zero")) return 0; - fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY); - if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0; - uint64_t n = (uint64_t)ratio * width; - compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>( - (float *)state_kv->ptr, (float *)state_score->ptr, - (const float *)kv_tail->ptr, (const float *)sc_tail->ptr, - ape, 0, ape_type, width, ratio, pos0, - 0, 0, ratio); - return cuda_ok(cudaGetLastError(), "compressor state set launch"); -} - -/* perf-02 split-KV / flash-decode launch helper (opt-in, default OFF). - * Returns 1 if the split path handled the launch, 0 if the caller should fall - * through to the existing attention_decode_mixed_kernel path. - * - * Engages only for the single-token decode shape (n_tokens==1) and only when - * DS4_CUDA_SPLITKV_DECODE is set. S==1 is NOT handled here: the caller dispatches - * the old kernel as the bit-exact anchor when S would be 1. - */ -static int attention_decode_splitkv_launch( - int logical_tier, - float *heads, - const float *sinks, - const float *q, - const float *raw_kv, - const float *comp_kv, - const float *comp_mask, - uint32_t use_comp_mask, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - /* n_tokens is fixed at 1 for the split path; compute the EXACT logical row - * count the kernel will use (raw_count + visible_comp) so S is sized to the - * real work. raw_count MUST apply the same window logic as the kernel / - * reference, otherwise a true-S==1 case (e.g. ratio=1, window=1, n_raw>=2, - * n_comp=0) could be over-estimated to S>1 and engage split-KV instead of - * the bit-exact old-kernel anchor. The count is head-independent. */ - const bool single_all = (ratio == 0u); - uint32_t qpos = pos0; /* t==0, n_tokens==1 */ - uint32_t first_raw_pos = pos0 + 1u - n_raw; - uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); - if (visible_comp > n_comp) visible_comp = n_comp; - uint32_t raw_count = 0; - uint32_t raw_first_idx = 0; - if (n_raw != 0) { - const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; - if (single_all) { - raw_count = n_raw > 256u ? 256u : n_raw; - } else if (qpos >= first_raw_pos) { - uint32_t lo = first_raw_pos; - if (window != 0 && qpos + 1u > window) { - const uint32_t wlo = qpos + 1u - window; - if (wlo > lo) lo = wlo; - } - const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; - if (hi >= lo) { - raw_first_idx = lo - first_raw_pos; - raw_count = hi - lo + 1u; - if (raw_count > 256u) raw_count = 256u; - } - } - } - uint32_t n_score = raw_count + visible_comp; - if (n_score == 0u) return 0; /* nothing to do; let old path handle it */ - const int manual_splitkv = cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0); - const uint32_t scoped_min_score = - (g_decode_fast_attention && !manual_splitkv) ? 512u : 0u; - uint32_t min_score = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_MIN_SCORE", - scoped_min_score, 0u, - DS4_CUDA_ATTENTION_SCORE_CAP, - NULL); - if (n_score < min_score) return 0; - /* S = clamp(ceil(n_score / CHUNK), 1, S_MAX); raise to S_FLOOR for short - * context to fill more SMs, but never exceed n_score (no empty chunks). - * Optional tuning knobs are guarded by min_needed so every block's chunk - * still fits the fixed shared score buffer. */ - const uint32_t split_cap = DS4_CUDA_SPLITKV_SCORE_CAP; - const uint32_t min_needed = (n_score + split_cap - 1u) / split_cap; - uint32_t chunk = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_CHUNK", - DS4_CUDA_SPLITKV_CHUNK, - 1u, split_cap, NULL); - uint32_t s_floor = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_FLOOR", - DS4_CUDA_SPLITKV_S_FLOOR, - 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); - uint32_t s_max = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_MAX", - DS4_CUDA_SPLITKV_S_MAX, - 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); - int exact_present = 0; - uint32_t S = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S", - 0u, 1u, DS4_CUDA_SPLITKV_S_MAX, - &exact_present); - if (!exact_present) { - S = (n_score + chunk - 1u) / chunk; - if (S < s_floor) S = s_floor < n_score ? s_floor : n_score; - if (S > s_max) S = s_max; - } - if (S < min_needed) S = min_needed; - if (S > n_score) S = n_score; - if (S <= 1u) return 0; /* S==1: caller uses the old kernel anchor */ - if (cuda_env_flag_enabled("DS4_CUDA_SPLITKV_GLOBAL_SOFTMAX", 0)) { - const uint64_t score_count = (uint64_t)n_head * n_score; - const uint64_t score_bytes = score_count * sizeof(float); - const uint64_t denom_offset = (score_bytes + 255u) & ~255ull; - const uint64_t denom_bytes = (uint64_t)n_head * sizeof(float); - const uint64_t partial_offset = (denom_offset + denom_bytes + 255u) & ~255ull; - const uint64_t partial_bytes = (uint64_t)n_head * S * head_dim * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, - partial_offset + partial_bytes, - "attention splitkv global softmax"); - if (!tmp) return 0; - float *scores = (float *)tmp; - float *denom = (float *)((char *)tmp + denom_offset); - float *partials = (float *)((char *)tmp + partial_offset); - dim3 score_grid(1, n_head, S); - attention_decode_score_split_scores_kernel<<>>( - scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim, S); - if (!cuda_ok(cudaGetLastError(), "attention splitkv global score launch")) return -1; - attention_decode_global_softmax_kernel<<>>( - scores, denom, sinks, n_score, n_head); - if (!cuda_ok(cudaGetLastError(), "attention splitkv global softmax launch")) return -1; - dim3 value_grid(1, n_head, S); - attention_decode_split_value_kernel<<>>( - partials, scores, raw_kv, comp_kv, raw_count, raw_first_idx, - raw_cap, raw_start, n_score, n_head, head_dim, S); - if (!cuda_ok(cudaGetLastError(), "attention splitkv global value launch")) return -1; - dim3 combine_grid(1, n_head, 1); - attention_decode_split_value_combine_kernel<<>>( - heads, partials, denom, n_head, head_dim, S); - if (!cuda_ok(cudaGetLastError(), "attention splitkv global combine launch")) return -1; - return 1; - } - /* Partials scratch: n_head * S * (head_dim + 2) floats (n_tokens==1). */ - uint64_t stride = (uint64_t)head_dim + 2u; - uint64_t count = (uint64_t)n_head * S * stride; - float *partials = (float *)cuda_tmp_alloc_on(logical_tier, count * sizeof(float), - "attention splitkv partials"); - if (!partials) return 0; - dim3 split_grid(1, n_head, S); - attention_decode_splitkv_kernel<<>>(partials, - q, - raw_kv, - comp_kv, - comp_mask, - use_comp_mask, - 1, pos0, n_raw, raw_cap, raw_start, - n_comp, window, ratio, n_head, head_dim, S); - if (!cuda_ok(cudaGetLastError(), "attention splitkv partial launch")) return -1; - dim3 combine_grid(1, n_head, 1); - attention_decode_splitkv_combine_kernel<<>>(heads, - sinks, - partials, - 1, n_head, head_dim, S); - if (!cuda_ok(cudaGetLastError(), "attention splitkv combine launch")) return -1; - return 1; -} - -extern "C" int ds4_gpu_attention_decode_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_comp, - const ds4_gpu_tensor *comp_mask, - uint32_t use_mask, - uint32_t n_head, - uint32_t head_dim) { - if (comp_kv_f16 || - !heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw || - raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || - (use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float))) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); - if (!sinks) return 0; - if (!cuda_attention_score_buffer_fits(n_comp)) { - if (!use_mask && head_dim == 512u && - !g_cuda_no_window_attention) { - const uint32_t synthetic_pos0 = n_raw - 1u; - dim3 online_grid(1, (n_head + 7u) / 8u, 1); - attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - 1, - synthetic_pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - 0, - 0, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention decode online launch"); - } - fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp); - return 0; - } - if (!use_mask && head_dim == 512u && - g_cuda_decode_heads8_online && - !g_cuda_no_window_attention) { - const uint32_t synthetic_pos0 = n_raw - 1u; - dim3 online_grid(1, (n_head + 7u) / 8u, 1); - attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - 1, - synthetic_pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - 0, - 0, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention decode heads8 online launch"); - } - const uint32_t score_lanes = - g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); - const uint32_t threads = - head_dim == 512u && score_lanes == 0u && - !g_cuda_no_decode_value512 ? 512u : 256u; - int score_split_rc = attention_decode_score_split_launch( - logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, - 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, - n_head, head_dim, threads, NULL); - if (score_split_rc == 1) { - return cuda_ok(cudaGetLastError(), "attention exact score split launch"); - } - if (score_split_rc < 0) return 0; - /* perf-02 split-KV opt-in (default OFF). n_tokens==1 here by construction. - * S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */ - if (cuda_splitkv_decode_requested()) { - int rc = attention_decode_splitkv_launch( - logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, - 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, n_head, head_dim); - if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv launch"); - if (rc < 0) return 0; - } - dim3 grid(1, n_head, 1); - attention_decode_mixed_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_mask ? (const float *)comp_mask->ptr : NULL, - use_mask, - 1, 0, n_raw, raw_cap, raw_start, n_comp, - 0, 0, n_head, head_dim, - score_lanes); - return cuda_ok(cudaGetLastError(), "attention decode launch"); -} - -extern "C" int ds4_gpu_attention_decode_heads_rope_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_comp, - const ds4_gpu_tensor *comp_mask, - uint32_t use_mask, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - int *fused_inv_rope) { - if (fused_inv_rope) *fused_inv_rope = 0; - if (!g_cuda_exact_score_split_fuse_inv_rope || - n_rot == 0u || n_rot > head_dim || (n_rot & 1u) || - head_dim != 512u) { - return ds4_gpu_attention_decode_heads_tensor( - heads, model_map, model_size, sinks_offset, q, raw_kv, - n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, - comp_mask, use_mask, n_head, head_dim); - } - if (!use_mask && g_cuda_decode_heads8_online && !g_cuda_no_window_attention) { - return ds4_gpu_attention_decode_heads_tensor( - heads, model_map, model_size, sinks_offset, q, raw_kv, - n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, - comp_mask, use_mask, n_head, head_dim); - } - if (comp_kv_f16 || - !heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw || - raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || - (use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float)) || - !cuda_attention_score_buffer_fits(n_comp)) { - return ds4_gpu_attention_decode_heads_tensor( - heads, model_map, model_size, sinks_offset, q, raw_kv, - n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, - comp_mask, use_mask, n_head, head_dim); - } - const uint32_t score_lanes = - g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); - const uint32_t threads = - score_lanes == 0u && !g_cuda_no_decode_value512 ? 512u : 256u; - if (threads < 512u) { - return ds4_gpu_attention_decode_heads_tensor( - heads, model_map, model_size, sinks_offset, q, raw_kv, - n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, - comp_mask, use_mask, n_head, head_dim); - } - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), - logical_tier, "attn_sinks"); - if (!sinks) return 0; - cuda_attention_inv_rope_params rope; - rope.n_rot = n_rot; - rope.pos0 = pos0; - rope.n_ctx_orig = n_ctx_orig; - rope.freq_base = freq_base; - rope.freq_scale = freq_scale; - rope.ext_factor = ext_factor; - rope.attn_factor = attn_factor; - rope.beta_fast = beta_fast; - rope.beta_slow = beta_slow; - int score_split_rc = attention_decode_score_split_launch( - logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_mask ? (const float *)comp_mask->ptr : NULL, use_mask, - 0, n_raw, raw_cap, raw_start, n_comp, 0, 0, - n_head, head_dim, threads, &rope); - if (score_split_rc == 1) { - if (fused_inv_rope) *fused_inv_rope = 1; - return cuda_ok(cudaGetLastError(), - "attention exact score split fused inv rope launch"); - } - if (score_split_rc < 0) return 0; - return ds4_gpu_attention_decode_heads_tensor( - heads, model_map, model_size, sinks_offset, q, raw_kv, - n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp, - comp_mask, use_mask, n_head, head_dim); -} - -extern "C" int ds4_gpu_attention_decode_rows_rope_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_attention_decode_row *rows, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (!heads || !q || !rows || !model_map || n_rows < 2u || - n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u || - head_dim != 512u || n_rot == 0u || n_rot > head_dim || - (n_rot & 1u) != 0u || sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(heads); - if (logical_tier < 0 || logical_tier >= g_n_gpus || - ds4_tensor_device_idx(q) != logical_tier) { - return 0; - } - - /* This first grouped path mirrors the promoted default decode exactly. - * Alternative score kernels and graph/split-KV experiments retain the - * one-session dispatcher until they gain equivalent row-table variants. */ - if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0) || - !cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1) || - cuda_splitkv_decode_requested() || - g_cuda_decode_heads8_online || g_cuda_decode_score4 || - g_cuda_decode_score8 || g_cuda_no_decode_value512 || - g_cuda_exact_score_split_graph || g_cuda_exact_score_split_ldg || - g_cuda_exact_score_split_vec4 || - g_cuda_exact_score_split_vec4_plain || - g_cuda_exact_score_split_dim2 || - g_cuda_exact_score_split_fuse_inv_rope || - getenv("DS4_CUDA_NO_SCORE_TILE") != NULL || - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE") != NULL || - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK") != NULL || - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR") != NULL || - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX") != NULL || - getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S") != NULL) { - return 0; - } - - cuda_attention_decode_row_table table; - memset(&table, 0, sizeof(table)); - uint32_t max_dense_score = 0u; - bool have_dense = false; - bool have_indexed = false; - for (uint32_t i = 0; i < n_rows; i++) { - const ds4_gpu_attention_decode_row r = rows[i]; - if (r.raw_kv == 0u || r.n_raw == 0u || r.raw_cap < r.n_raw || - r.raw_start >= r.raw_cap || (r.n_comp != 0u && r.comp_kv == 0u)) { - return 0; - } - if (r.indexed) { - if (r.comp_kv == 0u || r.topk == 0u || r.n_comp == 0u || - r.top_k == 0u || r.top_k > 512u || r.ratio == 0u) { - return 0; - } - have_indexed = true; - } else { - const uint32_t raw_count = r.n_raw > 256u ? 256u : r.n_raw; - const uint32_t n_score = raw_count + r.n_comp; - /* n_score==1 takes the legacy one-block kernel and is not a - * score-split shape. Decode after any nonempty prompt is >1. */ - if (n_score <= 1u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) { - return 0; - } - if (n_score > max_dense_score) max_dense_score = n_score; - have_dense = true; - } - table.row[i] = r; - } - - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), - logical_tier, "attn_sinks_rows"); - if (!sinks) return 0; - - if (have_dense) { - if ((uint64_t)n_rows > UINT64_MAX / n_head || - (uint64_t)n_rows * n_head > UINT64_MAX / max_dense_score) { - return 0; - } - const uint64_t score_count = - (uint64_t)n_rows * n_head * max_dense_score; - float *scores = (float *)cuda_tmp_alloc_on( - logical_tier, score_count * sizeof(float), - "attention exact decode rows"); - if (!scores) return 0; - - const size_t tile_shmem = - (size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) * - DS4_SCORE_TILE_STRIDE * sizeof(float); - static int tile_shmem_ready[DS4_MAX_GPUS] = {0}; - int physical_device = 0; - if (cudaGetDevice(&physical_device) != cudaSuccess || - physical_device < 0 || physical_device >= DS4_MAX_GPUS) { - return 0; - } - if (!tile_shmem_ready[physical_device]) { - if (!cuda_ok(cudaFuncSetAttribute( - attention_decode_score_split_scores_tile512_rows_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - (int)tile_shmem), - "attention score rows shared-memory opt-in")) { - return 0; - } - tile_shmem_ready[physical_device] = 1; - } - dim3 score_grid( - (max_dense_score + DS4_SCORE_TILE_ROWS - 1u) / - DS4_SCORE_TILE_ROWS, - (n_head + DS4_SCORE_TILE_HEADS - 1u) / - DS4_SCORE_TILE_HEADS, - n_rows); - attention_decode_score_split_scores_tile512_rows_kernel - <<>>( - scores, (const float *)q->ptr, table, n_rows, - max_dense_score, n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), - "attention exact score rows launch")) { - return 0; - } - dim3 final_grid(n_rows, n_head, 1u); - attention_decode_score_split_finalize_rows_kernel - <<>>( - (float *)heads->ptr, sinks, scores, table, n_rows, - max_dense_score, n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), - "attention exact finalize rows launch")) { - return 0; - } - } - if (have_indexed) { - dim3 indexed_grid(n_rows, n_head, 1u); - attention_indexed_mixed_decode_rows_kernel<<>>( - (float *)heads->ptr, sinks, (const float *)q->ptr, table, - n_rows, n_head, head_dim); - if (!cuda_ok(cudaGetLastError(), - "attention indexed decode rows launch")) { - return 0; - } - } - - const uint32_t pairs = n_rows * n_head * (n_rot / 2u); - rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>( - (float *)heads->ptr, table, n_rows, n_head, head_dim, n_rot, - n_ctx_orig, 1, freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow); - return cuda_ok(cudaGetLastError(), - "attention decode rows inverse rope launch"); -} - -extern "C" int ds4_gpu_attention_prefill_raw_heads_tensor(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, uint64_t sinks_offset, const ds4_gpu_tensor *q, const ds4_gpu_tensor *raw_kv, uint32_t n_tokens, uint32_t window, uint32_t n_head, uint32_t head_dim) { - if (!heads || !q || !raw_kv || !model_map || sinks_offset > model_size || - model_size - sinks_offset < (uint64_t)n_head * sizeof(float) || - heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || - window > 256) return 0; - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); - if (!sinks) return 0; - if (n_tokens > 1 && head_dim == 512 && - getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL && - (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { - dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); - attention_static_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - (const float *)raw_kv->ptr, - n_tokens, - 0, - window, - 1, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention raw window launch"); - } - if (g_cublas_ready && n_tokens > 1 && head_dim == 512 && - getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) { - const uint32_t n_keys = n_tokens; - const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys; - const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim; - const uint64_t score_bytes = score_count * sizeof(float); - const uint64_t out_offset = (score_bytes + 255u) & ~255ull; - const uint64_t tmp_bytes = out_offset + out_count * sizeof(float); - float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention raw cublas"); - if (!tmp) return 0; - float *scores = tmp; - float *out_tmp = (float *)((char *)tmp + out_offset); - const float alpha = rsqrtf((float)head_dim); - const float beta = 0.0f; - cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)n_keys, - (int)n_tokens, - (int)head_dim, - &alpha, - (const float *)raw_kv->ptr, - (int)head_dim, - 0, - (const float *)q->ptr, - (int)(n_head * head_dim), - (long long)head_dim, - &beta, - scores, - (int)n_keys, - (long long)n_keys * n_tokens, - (int)n_head); - if (!cublas_ok(st, "attention raw score gemm")) return 0; - dim3 sgrid(n_tokens, n_head, 1); - attention_prefill_raw_softmax_kernel<<>>(scores, sinks, n_tokens, window, n_keys); - if (!cuda_ok(cudaGetLastError(), "attention raw softmax launch")) return 0; - const float one = 1.0f; - st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_N, - CUBLAS_OP_N, - (int)head_dim, - (int)n_tokens, - (int)n_keys, - &one, - (const float *)raw_kv->ptr, - (int)head_dim, - 0, - scores, - (int)n_keys, - (long long)n_keys * n_tokens, - &beta, - out_tmp, - (int)head_dim, - (long long)head_dim * n_tokens, - (int)n_head); - if (!cublas_ok(st, "attention raw value gemm")) return 0; - uint64_t n = (uint64_t)n_tokens * n_head * head_dim; - attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr, - out_tmp, - n_tokens, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention raw unpack launch"); - } - dim3 grid(n_tokens, n_head, 1); - attention_prefill_raw_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_tokens, window, n_head, head_dim); - return cuda_ok(cudaGetLastError(), "attention_prefill_raw launch"); -} -static int attention_decode_batch_launch( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (comp_kv_f16 || - !heads || !q || !raw_kv || !model_map || n_tokens == 0 || - n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || - (n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || - (use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) { - return 0; - } - if (n_comp != 0 && ratio == 0) return 0; - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); - if (!sinks) return 0; - if (!cuda_attention_score_buffer_fits(n_comp)) { - if (!use_comp_mask && head_dim == 512u && - !g_cuda_no_window_attention) { - dim3 online_grid(n_tokens, (n_head + 7u) / 8u, 1); - attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention decode online launch"); - } - fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp); - return 0; - } - if (!use_comp_mask && n_tokens > 1 && head_dim == 512 && - !g_cuda_no_window_attention && - (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { - dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); - attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention decode window launch"); - } - if (!use_comp_mask && n_tokens == 1u && head_dim == 512 && - g_cuda_decode_heads8_online && - !g_cuda_no_window_attention) { - dim3 grid(1, (n_head + 7u) / 8u, 1); - attention_decode_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention decode heads8 online batch launch"); - } - const uint32_t score_lanes = - g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u); - const uint32_t threads = - n_tokens == 1u && head_dim == 512u && score_lanes == 0u && - !g_cuda_no_decode_value512 ? 512u : 256u; - if (n_tokens == 1u) { - int score_split_rc = attention_decode_score_split_launch( - logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, - n_head, head_dim, threads, NULL); - if (score_split_rc == 1) { - return cuda_ok(cudaGetLastError(), "attention exact score split batch launch"); - } - if (score_split_rc < 0) return 0; - } - /* perf-02 split-KV opt-in (default OFF). Single-token decode only; multi- - * token batch shapes already fill the grid and fall through unchanged. - * S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */ - if (n_tokens == 1u && cuda_splitkv_decode_requested()) { - int rc = attention_decode_splitkv_launch( - logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask, - pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, n_head, head_dim); - if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv batch launch"); - if (rc < 0) return 0; - } - dim3 grid(n_tokens, n_head, 1); - attention_decode_mixed_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_comp_mask ? (const float *)comp_mask->ptr : NULL, - use_comp_mask, n_tokens, pos0, n_raw, raw_cap, - raw_start, n_comp, window, ratio, n_head, head_dim, - score_lanes); - return cuda_ok(cudaGetLastError(), "attention decode batch launch"); -} - -extern "C" int ds4_gpu_attention_decode_raw_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset, - q, raw_kv, NULL, 0, NULL, 0, n_tokens, pos0, - n_raw, raw_cap, raw_start, 0, window, 1, - n_head, head_dim); -} - -extern "C" int ds4_gpu_attention_decode_mixed_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (comp_kv_f16) return 0; - return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset, - q, raw_kv, comp_kv, comp_kv_f16, comp_mask, use_comp_mask, - n_tokens, pos0, n_raw, raw_cap, raw_start, - n_comp, window, ratio, n_head, head_dim); -} - -extern "C" int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *topk, - uint32_t n_tokens, - uint32_t pos0, - 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 n_head, - uint32_t head_dim) { - if (comp_kv_f16 || - !heads || !q || !raw_kv || !comp_kv || !topk || !model_map || - n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || - n_comp == 0 || top_k == 0 || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) || - comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float) || - topk->bytes < (uint64_t)n_tokens * top_k * sizeof(int32_t)) { - return 0; - } - if (top_k > 512u) return 0; - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); - if (!sinks) return 0; - const int32_t *topk_ptr = (const int32_t *)topk->ptr; - if (n_tokens > 1u && top_k == 512u && - getenv("DS4_CUDA_NO_INDEXED_TOPK_SORT") == NULL) { - const uint64_t sort_bytes = (uint64_t)n_tokens * top_k * sizeof(int32_t); - int32_t *sorted = (int32_t *)cuda_tmp_alloc_on(logical_tier, sort_bytes, "indexed attention topk sort"); - if (!sorted) return 0; - indexed_topk_sort_512_asc_kernel<<>>(sorted, topk_ptr, n_tokens); - if (!cuda_ok(cudaGetLastError(), "indexed attention topk sort launch")) return 0; - topk_ptr = sorted; - } - if (n_tokens > 1 && head_dim == 512 && top_k <= 512u && - getenv("DS4_CUDA_NO_INDEXED_HEADS8") == NULL) { - if (getenv("DS4_CUDA_INDEXED_TWOPASS") == NULL) { - dim3 grid(n_tokens, (n_head + 15u) / 16u, 1); - attention_indexed_mixed_heads8_online_kernel<8, 16><<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - (const float *)comp_kv->ptr, - topk_ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - top_k, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention indexed online launch"); - } - dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); - attention_indexed_mixed_heads8_rb4_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - (const float *)comp_kv->ptr, - topk_ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - top_k, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention indexed heads8 launch"); - } - dim3 grid(n_tokens, n_head, 1); - attention_indexed_mixed_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - (const float *)comp_kv->ptr, - topk_ptr, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - top_k, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention indexed mixed launch"); -} - -static int attention_prefill_mixed_launch( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || ratio == 0 || - (n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) || - sinks_offset > model_size || - (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || - heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || - (n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) || - (use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(heads); - const float *sinks = (const float *)cuda_resolve_weight_ptr( - model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks"); - if (!sinks) return 0; - if (!use_comp_mask && n_tokens > 1 && head_dim == 512 && - getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL && - (getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) { - dim3 grid(n_tokens, (n_head + 7u) / 8u, 1); - attention_static_mixed_heads8_online_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention mixed window launch"); - } - if (g_cublas_ready && n_tokens > 1 && head_dim == 512 && - getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) { - const uint32_t n_keys = n_tokens + n_comp; - const uint64_t kv_count = (uint64_t)n_keys * head_dim; - const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys; - const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim; - const uint64_t kv_bytes = kv_count * sizeof(float); - const uint64_t score_offset = (kv_bytes + 255u) & ~255ull; - const uint64_t score_bytes = score_count * sizeof(float); - const uint64_t out_offset = score_offset + ((score_bytes + 255u) & ~255ull); - const uint64_t tmp_bytes = out_offset + out_count * sizeof(float); - float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention mixed cublas"); - if (!tmp) return 0; - float *kv = tmp; - float *scores = (float *)((char *)tmp + score_offset); - float *out_tmp = (float *)((char *)tmp + out_offset); - attention_prefill_pack_mixed_kv_kernel<<<(kv_count + 255) / 256, 256>>>( - kv, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - n_tokens, - n_comp, - head_dim); - if (!cuda_ok(cudaGetLastError(), "attention mixed kv pack launch")) return 0; - const float alpha = rsqrtf((float)head_dim); - const float beta = 0.0f; - cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)n_keys, - (int)n_tokens, - (int)head_dim, - &alpha, - kv, - (int)head_dim, - 0, - (const float *)q->ptr, - (int)(n_head * head_dim), - (long long)head_dim, - &beta, - scores, - (int)n_keys, - (long long)n_keys * n_tokens, - (int)n_head); - if (!cublas_ok(st, "attention mixed score gemm")) return 0; - dim3 sgrid(n_tokens, n_head, 1); - attention_prefill_mixed_softmax_kernel<<>>( - scores, - sinks, - use_comp_mask ? (const float *)comp_mask->ptr : NULL, - use_comp_mask, - n_tokens, - n_comp, - window, - ratio, - n_keys); - if (!cuda_ok(cudaGetLastError(), "attention mixed softmax launch")) return 0; - const float one = 1.0f; - st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_N, - CUBLAS_OP_N, - (int)head_dim, - (int)n_tokens, - (int)n_keys, - &one, - kv, - (int)head_dim, - 0, - scores, - (int)n_keys, - (long long)n_keys * n_tokens, - &beta, - out_tmp, - (int)head_dim, - (long long)head_dim * n_tokens, - (int)n_head); - if (!cublas_ok(st, "attention mixed value gemm")) return 0; - uint64_t n = (uint64_t)n_tokens * n_head * head_dim; - attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr, - out_tmp, - n_tokens, - n_head, - head_dim); - return cuda_ok(cudaGetLastError(), "attention mixed unpack launch"); - } - dim3 grid(n_tokens, n_head, 1); - attention_prefill_mixed_kernel<<>>((float *)heads->ptr, - sinks, - (const float *)q->ptr, - (const float *)raw_kv->ptr, - n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr, - use_comp_mask ? (const float *)comp_mask->ptr : NULL, - use_comp_mask, n_tokens, n_comp, window, ratio, - n_head, head_dim); - return cuda_ok(cudaGetLastError(), "attention prefill mixed launch"); -} - -extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (comp_kv_f16) return 0; - return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset, - q, raw_kv, comp_kv, NULL, 0, n_tokens, - n_comp, window, ratio, n_head, head_dim); -} - -extern "C" int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (comp_kv_f16) return 0; - return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset, - q, raw_kv, comp_kv, comp_mask, 1, n_tokens, - n_comp, window, ratio, n_head, head_dim); -} -extern "C" int ds4_gpu_attention_output_q8_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_tensor *group_tmp, - ds4_gpu_tensor *low_tmp, - 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, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens) { - (void)group_tmp; - (void)low_tmp; - if (!out || !low || !heads || !model_map || - group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0) { - return 0; - } - const uint64_t low_dim = (uint64_t)n_groups * rank; - const uint64_t blocks_a = (group_dim + 31) / 32; - const uint64_t blocks_b = (low_dim + 31) / 32; - const uint64_t out_a_bytes = (uint64_t)n_groups * rank * blocks_a * 34; - const uint64_t out_b_bytes = out_dim * blocks_b * 34; - if (out_a_offset > model_size || out_b_offset > model_size || - out_a_bytes > model_size - out_a_offset || - out_b_bytes > model_size - out_b_offset || - heads->bytes < (uint64_t)n_tokens * n_groups * group_dim * sizeof(float) || - low->bytes < (uint64_t)n_tokens * low_dim * sizeof(float) || - out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out); - const int physical_device = - (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) - ? g_gpu[logical_tier].device_id : 0; - const unsigned char *out_a = reinterpret_cast( - cuda_resolve_weight_ptr(model_map, out_a_offset, out_a_bytes, logical_tier, "attn_out_a")); - const unsigned char *out_b = reinterpret_cast( - cuda_resolve_weight_ptr(model_map, out_b_offset, out_b_bytes, logical_tier, "attn_out_b")); - if (!out_a || !out_b) return 0; - - const uint32_t profile = getenv("DS4_CUDA_ATTN_OUTPUT_PROFILE") != NULL; - cudaEvent_t prof_ev[3] = {NULL, NULL, NULL}; - if (profile) { - for (uint32_t i = 0; i < 3u; i++) { - if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) { - for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]); - memset(prof_ev, 0, sizeof(prof_ev)); - break; - } - } - if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0); - } - - const __half *out_a_f16 = NULL; - uint32_t out_a_cublas_min_tokens = 2u; - const char *out_a_min_env = getenv("DS4_CUDA_ATTENTION_OUTPUT_A_CUBLAS_MIN"); - if (out_a_min_env && out_a_min_env[0]) { - char *endp = NULL; - long v = strtol(out_a_min_env, &endp, 10); - if (endp != out_a_min_env && v > 1 && v < 4096) out_a_cublas_min_tokens = (uint32_t)v; - } - if (!g_quality_mode && - g_cublas_ready && - n_tokens >= out_a_cublas_min_tokens && - getenv("DS4_CUDA_NO_CUBLAS_ATTENTION_OUTPUT_A") == NULL) { - out_a_f16 = cuda_q8_f16_ptr(model_map, out_a_offset, out_a_bytes, group_dim, low_dim, physical_device, "attn_output_a"); - } - if (out_a_f16) { - const uint64_t heads_h_count = (uint64_t)n_groups * n_tokens * group_dim; - const uint64_t low_tmp_count = (uint64_t)n_groups * n_tokens * rank; - const uint64_t heads_h_bytes = heads_h_count * sizeof(__half); - const uint64_t low_tmp_offset = (heads_h_bytes + 255u) & ~255ull; - const uint64_t tmp_bytes = low_tmp_offset + low_tmp_count * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a cublas"); - if (!tmp) return 0; - __half *heads_h = (__half *)tmp; - float *low_packed = (float *)((char *)tmp + low_tmp_offset); - attention_pack_group_heads_f16_kernel<<<(heads_h_count + 255) / 256, 256>>>( - heads_h, - (const float *)heads->ptr, - n_tokens, - n_groups, - group_dim); - if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a pack launch")) return 0; - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasGemmStridedBatchedEx(cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)rank, - (int)n_tokens, - (int)group_dim, - &alpha, - out_a_f16, - CUDA_R_16F, - (int)group_dim, - (long long)rank * group_dim, - heads_h, - CUDA_R_16F, - (int)group_dim, - (long long)n_tokens * group_dim, - &beta, - low_packed, - CUDA_R_32F, - (int)rank, - (long long)rank * n_tokens, - (int)n_groups, - CUDA_R_32F, - CUBLAS_GEMM_DEFAULT); - if (!cublas_ok(st, "attention output a gemm")) return 0; - attention_unpack_group_low_kernel<<<(low_tmp_count + 255) / 256, 256>>>( - (float *)low->ptr, - low_packed, - n_tokens, - n_groups, - rank); - if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a unpack launch")) return 0; - } else { - const uint64_t x_rows = (uint64_t)n_tokens * n_groups; - const uint64_t xq_bytes = x_rows * blocks_a * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a q8 prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1); - quantize_q8_0_f32_kernel<<>>(xq, - xscale, - (const float *)heads->ptr, - group_dim, - blocks_a); - if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a prequant launch")) return 0; - int grouped_mma_done = 0; - if (n_tokens >= 8u) { - /* One mma launch per group: T=32 matches the warp tree of the - * grouped reference kernels (multi-term slots for blocks > 32). */ - grouped_mma_done = 1; - for (uint32_t g = 0; g < n_groups && grouped_mma_done == 1; g++) { - const int rc = cuda_q8_mma_try_launch( - (float *)low->ptr + (uint64_t)g * rank, - reinterpret_cast(out_a) + - (uint64_t)g * rank * blocks_a * 34u, - xq + (uint64_t)g * blocks_a * 32u, - xscale + (uint64_t)g * blocks_a, - group_dim, rank, n_tokens, blocks_a, - (uint64_t)n_groups * blocks_a, low_dim, 32u); - if (rc < 0) return 0; - if (rc == 0) grouped_mma_done = 0; - } - } - if (grouped_mma_done) { - /* handled */ - } else if (getenv("DS4_CUDA_NO_ATTN_A_TOK2") == NULL && n_tokens >= 2u) { - dim3 grid_a(((unsigned)low_dim + 7u) / 8u, ((unsigned)n_tokens + 1u) / 2u, 1); - grouped_q8_0_a_preq_warp8_tok2_kernel<<>>((float *)low->ptr, - out_a, - xq, - xscale, - group_dim, - rank, - n_groups, - n_tokens, - blocks_a, - use_dp4a); - } else { - dim3 grid_a(((unsigned)low_dim + 7u) / 8u, (unsigned)n_tokens, 1); - grouped_q8_0_a_preq_warp8_kernel<<>>((float *)low->ptr, - out_a, - xq, - xscale, - group_dim, - rank, - n_groups, - n_tokens, - blocks_a, - use_dp4a); - } - if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a preq launch")) return 0; - } - - if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0); - (void)out_b; - int ok = cuda_matmul_q8_0_tensor_labeled(out, - model_map, - model_size, - out_b_offset, - low_dim, - out_dim, - low, - n_tokens, - "attn_output_b"); - if (prof_ev[2]) { - (void)cudaEventRecord(prof_ev[2], 0); - if (cudaEventSynchronize(prof_ev[2]) == cudaSuccess) { - float ms_a = 0.0f, ms_b = 0.0f, ms_total = 0.0f; - (void)cudaEventElapsedTime(&ms_a, prof_ev[0], prof_ev[1]); - (void)cudaEventElapsedTime(&ms_b, prof_ev[1], prof_ev[2]); - (void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[2]); - fprintf(stderr, - "ds4: CUDA attention output profile tokens=%u groups=%u group_dim=%llu rank=%llu low=%llu out=%llu A=%.3f B=%.3f total=%.3f ms\n", - n_tokens, - n_groups, - (unsigned long long)group_dim, - (unsigned long long)rank, - (unsigned long long)low_dim, - (unsigned long long)out_dim, - ms_a, - ms_b, - ms_total); - } - for (uint32_t i = 0; i < 3u; i++) (void)cudaEventDestroy(prof_ev[i]); - } - return ok; -} -extern "C" int ds4_gpu_attention_output_low_q8_rows_exact_tensor( - ds4_gpu_tensor *low, - const void *model_map, - uint64_t model_size, - uint64_t out_a_offset, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups_total, - uint32_t group0, - uint32_t group_cnt, - const ds4_gpu_tensor *heads, - uint32_t n_rows) { - if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || - n_groups_total == 0 || group_cnt == 0 || - group0 > n_groups_total || group_cnt > n_groups_total - group0 || - n_rows == 0 || (uint64_t)n_rows * group_cnt > 65535u) { - return 0; - } - const uint64_t low_dim = (uint64_t)group_cnt * rank; - const uint64_t blocks_a = (group_dim + 31) / 32; - const uint64_t row_a_bytes = blocks_a * 34u; - const uint64_t a_offset = - out_a_offset + (uint64_t)group0 * rank * row_a_bytes; - const uint64_t out_a_bytes = low_dim * row_a_bytes; - if (a_offset < out_a_offset || a_offset > model_size || - out_a_bytes > model_size - a_offset || - heads->bytes < (uint64_t)n_rows * n_groups_total * group_dim * - sizeof(float) || - low->bytes < (uint64_t)n_rows * low_dim * sizeof(float)) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(low); - const unsigned char *out_a = reinterpret_cast( - cuda_resolve_weight_ptr(model_map, a_offset, out_a_bytes, - logical_tier, "attn_out_a_rows")); - if (!out_a) return 0; - - const uint64_t x_rows = (uint64_t)n_rows * group_cnt; - const uint64_t xq_bytes = x_rows * blocks_a * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output low q8 prequant"); - if (!tmp) return 0; - int8_t *xq = (int8_t *)tmp; - float *xscale = (float *)((char *)tmp + scale_offset); - const int use_dp4a = cuda_q8_use_dp4a(); - dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1); - quantize_q8_0_group_slice_rows_kernel<<>>( - xq, - xscale, - (const float *)heads->ptr, - group_dim, - blocks_a, - n_groups_total, - group0, - group_cnt); - if (!cuda_ok(cudaGetLastError(), - "attention_output_low_q8 rows prequant launch")) return 0; - dim3 grid_a(((unsigned)low_dim + 7u) / 8u, n_rows, 1u); - grouped_q8_0_a_preq_warp8_kernel<<>>((float *)low->ptr, - out_a, - xq, - xscale, - group_dim, - rank, - group_cnt, - n_rows, - blocks_a, - use_dp4a); - return cuda_ok(cudaGetLastError(), - "attention_output_low_q8 rows launch"); -} - -extern "C" int ds4_gpu_attention_output_low_q8_tensor( - ds4_gpu_tensor *low, - const void *model_map, - uint64_t model_size, - uint64_t out_a_offset, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups, - const ds4_gpu_tensor *heads) { - return ds4_gpu_attention_output_low_q8_rows_exact_tensor( - low, model_map, model_size, out_a_offset, group_dim, rank, - n_groups, 0u, n_groups, heads, 1u); -} - -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) { - if (!out || !low || !heads || !model_map || - group_dim == 0 || rank == 0 || n_groups_total == 0 || - group_cnt == 0 || group0 > n_groups_total || - group_cnt > n_groups_total - group0 || out_dim == 0) { - return 0; - } - const uint64_t blocks_a = (group_dim + 31u) / 32u; - const uint64_t row_a_bytes = blocks_a * 34u; - const uint64_t low_dim_total = (uint64_t)n_groups_total * rank; - const uint64_t k_off = (uint64_t)group0 * rank; - const uint64_t k_cnt = (uint64_t)group_cnt * rank; - if ((k_off % 32u) != 0 || (k_cnt % 32u) != 0) return 0; - if (heads->bytes < (uint64_t)(group0 + group_cnt) * group_dim * sizeof(float) || - low->bytes < k_cnt * sizeof(float) || - out->bytes < out_dim * sizeof(float)) { - return 0; - } - - ds4_gpu_tensor heads_slice = *heads; - heads_slice.ptr = (char *)heads->ptr + (uint64_t)group0 * group_dim * sizeof(float); - heads_slice.bytes = (uint64_t)group_cnt * group_dim * sizeof(float); - heads_slice.owner = 0; - - const uint64_t a_off = out_a_offset + (uint64_t)group0 * rank * row_a_bytes; - return ds4_gpu_attention_output_low_q8_tensor(low, - model_map, - model_size, - a_off, - group_dim, - rank, - group_cnt, - &heads_slice) && - ds4_gpu_matmul_q8_0_kslice_rows_tensor(out, - model_map, - model_size, - out_b_offset, - low_dim_total, - out_dim, - k_off, - k_cnt, - low, - 1); -} -extern "C" int ds4_gpu_swiglu_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *gate, const ds4_gpu_tensor *up, uint32_t n, float clamp, float weight) { - if (!out || !gate || !up || - out->bytes < (uint64_t)n * sizeof(float) || - gate->bytes < (uint64_t)n * sizeof(float) || - up->bytes < (uint64_t)n * sizeof(float)) return 0; - swiglu_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)gate->ptr, (const float *)up->ptr, n, clamp, weight); - return cuda_ok(cudaGetLastError(), "swiglu launch"); -} -extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - if (getenv("DS4_CUDA_DISABLE_SHARED_GATE_UP_PAIR") == NULL) { - return ds4_gpu_matmul_q8_0_pair_tensor(gate, up, - model_map, model_size, - gate_offset, up_offset, - in_dim, out_dim, out_dim, - x, 1) && - ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f); - } - return ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size, - gate_offset, in_dim, out_dim, x, 1) && - ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size, - up_offset, in_dim, out_dim, x, 1) && - ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f); -} - -extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *prequant, - uint32_t expert_split, - bool home_rank) { - if (!mid || !x || !model_map || in_dim == 0u || out_dim == 0u || - x->bytes < in_dim * sizeof(float) || - mid->bytes < out_dim * sizeof(float) || - (selected && (selected->bytes < 6u * sizeof(int32_t) || - expert_split == 0u))) { - return 0; - } - const uint64_t blocks = (in_dim + 31u) / 32u; - if (gate_offset > model_size || up_offset > model_size || - out_dim > UINT64_MAX / (blocks * 34u)) { - return 0; - } - const uint64_t weight_bytes = out_dim * blocks * 34u; - if (weight_bytes > model_size - gate_offset || - weight_bytes > model_size - up_offset) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(x); - if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0; - if (selected && ds4_tensor_device_idx(selected) != logical_tier) return 0; - if (prequant && ds4_tensor_device_idx(prequant) != logical_tier) return 0; - const int mid_tier = ds4_tensor_device_idx(mid); - if (mid_tier != logical_tier && !g_gpu_peer_ok[logical_tier][mid_tier]) { - return 0; - } - const char *gate_w = cuda_resolve_weight_ptr( - model_map, gate_offset, weight_bytes, logical_tier, - "shared_mid_gate_exact"); - const char *up_w = cuda_resolve_weight_ptr( - model_map, up_offset, weight_bytes, logical_tier, - "shared_mid_up_exact"); - if (!gate_w || !up_w) return 0; - - const uint64_t xq_bytes = blocks * 32u; - const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; - const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); - int8_t *xq; - float *xscale; - if (prequant) { - if (prequant->bytes < tmp_bytes) return 0; - xq = (int8_t *)prequant->ptr; - xscale = (float *)((char *)prequant->ptr + scale_offset); - } else { - void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, - "shared mid q8 exact prequant"); - if (!tmp) return 0; - xq = (int8_t *)tmp; - xscale = (float *)((char *)tmp + scale_offset); - quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>( - xq, xscale, (const float *)x->ptr, in_dim, blocks); - if (!cuda_ok(cudaGetLastError(), - "shared mid q8 exact quantize launch")) { - return 0; - } - } - shared_mid_q8_0_preq_warp8_exact_kernel<<< - ((unsigned)out_dim + 7u) / 8u, 256>>>( - (float *)mid->ptr, - (const unsigned char *)gate_w, - (const unsigned char *)up_w, - xq, - xscale, - in_dim, - out_dim, - blocks, - clamp, - selected ? (const int32_t *)selected->ptr : NULL, - expert_split, - home_rank, - cuda_q8_use_dp4a()); - return cuda_ok(cudaGetLastError(), "shared mid q8 exact launch"); -} -extern "C" int ds4_gpu_add_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, uint32_t n) { - if (!out || !a || !b || - out->bytes < (uint64_t)n * sizeof(float) || - a->bytes < (uint64_t)n * sizeof(float) || - b->bytes < (uint64_t)n * sizeof(float)) return 0; - add_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)a->ptr, (const float *)b->ptr, n); - return cuda_ok(cudaGetLastError(), "add launch"); -} - -extern "C" 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) { - if (!out || !local || !remote || - out->bytes < (uint64_t)n * sizeof(float) || - local->bytes < (uint64_t)n * sizeof(float) || - remote->bytes < (uint64_t)n * sizeof(float)) return 0; - if (n == 0) return 1; - - const int od = ds4_tensor_device_idx(out); - const int ld = ds4_tensor_device_idx(local); - const int rd = ds4_tensor_device_idx(remote); - if (od != ld) return 0; - - const ds4_gpu_tensor *rhs = remote; - if (rd != od) { - if (!remote_tmp || - remote_tmp->bytes < (uint64_t)n * sizeof(float) || - ds4_tensor_device_idx(remote_tmp) != od) return 0; - if (!ds4_gpu_tensor_copy_xdev(remote_tmp, remote, - (uint64_t)n * sizeof(float))) return 0; - rhs = remote_tmp; - } - - int ok = 0; - WITH_DEVICE(g_gpu[od].device_id) { - cudaStream_t s = (cudaStream_t)g_gpu[od].stream; - add_kernel<<<(n + 255u) / 256u, 256, 0, s>>>( - (float *)out->ptr, - (const float *)local->ptr, - (const float *)rhs->ptr, - n); - ok = cuda_ok(cudaGetLastError(), "xdev add launch"); - cudaEvent_t e = (cudaEvent_t)g_gpu[od].boundary_event; - if (ok) ok = cuda_ok(cudaEventRecord(e, s), "xdev add event record"); - if (ok) ok = cuda_ok(cudaStreamWaitEvent(0, e, 0), "xdev add default wait"); - if (ok && g_xdev_sync_debug) { - ok = cuda_ok(cudaStreamSynchronize(s), "xdev add sync"); - } - } - return ok; -} -extern "C" int ds4_gpu_directional_steering_project_tensor( - ds4_gpu_tensor *x, - const ds4_gpu_tensor *directions, - uint32_t layer, - uint32_t width, - uint32_t rows, - float scale) { - if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0; - const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float); - const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float); - if (x->bytes < x_bytes || directions->bytes < dir_bytes) return 0; - - uint32_t nth = 256u; - while (nth > width && nth > 1u) nth >>= 1; - directional_steering_project_kernel<<>>( - (float *)x->ptr, - (const float *)directions->ptr, - layer, - width, - rows, - scale); - return cuda_ok(cudaGetLastError(), "directional steering launch"); -} -extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t token, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits) { - if (!selected || !weights || !probs || !logits || !model_map || n_expert_groups > 1u || n_group_used > 0u) return 0; - if (n_expert != 256u || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0; - int32_t tok = (int32_t)token; - int ok = 1; - const float *bias = NULL; - const int32_t *hash = NULL; - const int logical_tier = ds4_tensor_device_idx(selected); - if (ok && has_bias && !hash_mode) { - if (bias_offset > model_size || model_size - bias_offset < 256u * sizeof(float)) ok = 0; - else bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, 256u * sizeof(float), logical_tier, "router_bias"); - if (!bias) ok = 0; - } - if (ok && hash_mode) { - const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t); - if (hash_offset > model_size || hash_bytes > model_size - hash_offset) ok = 0; - else hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash"); - if (!hash) ok = 0; - } - if (ok) { - if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL && - getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { - dim3 block(32, 4, 1); - router_select_warp_topk_kernel<<<1, block>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, - bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, - has_bias && !hash_mode, hash_mode); - } else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { - router_select_parallel_kernel<<<1, 256>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, - bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, - has_bias && !hash_mode, hash_mode); - } else { - router_select_kernel<<<1, 1>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr, - bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1, - has_bias && !hash_mode, hash_mode); - } - ok = cuda_ok(cudaGetLastError(), "router_select launch"); - } - return ok; -} -extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits, const ds4_gpu_tensor *tokens, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_tokens) { - if (n_expert != 256u || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0; - if (!selected || !weights || !probs || !logits || !tokens || !model_map || n_tokens == 0 || - n_expert_groups > 1u || n_group_used > 0u || - logits->bytes < (uint64_t)n_tokens * 256u * sizeof(float) || - probs->bytes < (uint64_t)n_tokens * 256u * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * 6u * sizeof(int32_t) || - weights->bytes < (uint64_t)n_tokens * 6u * sizeof(float)) { - return 0; - } - const float *bias = NULL; - const int32_t *hash = NULL; - const int logical_tier = ds4_tensor_device_idx(selected); - if (has_bias && !hash_mode) { - if (bias_offset > model_size || model_size - bias_offset < 256u * sizeof(float)) return 0; - bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, 256u * sizeof(float), logical_tier, "router_bias"); - if (!bias) return 0; - } - if (hash_mode) { - const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t); - if (hash_offset > model_size || hash_bytes > model_size - hash_offset) return 0; - hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash"); - if (!hash) return 0; - } - if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL && - getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { - dim3 block(32, 4, 1); - router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>((int32_t *)selected->ptr, - (float *)weights->ptr, - (float *)probs->ptr, - bias, - hash, - (const float *)logits->ptr, - (const int32_t *)tokens->ptr, - 0, - hash_rows, - n_tokens, - has_bias && !hash_mode, - hash_mode); - } else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) { - router_select_parallel_kernel<<>>((int32_t *)selected->ptr, - (float *)weights->ptr, - (float *)probs->ptr, - bias, - hash, - (const float *)logits->ptr, - (const int32_t *)tokens->ptr, - 0, - hash_rows, - n_tokens, - has_bias && !hash_mode, - hash_mode); - } else { - router_select_kernel<<>>((int32_t *)selected->ptr, - (float *)weights->ptr, - (float *)probs->ptr, - bias, - hash, - (const float *)logits->ptr, - (const int32_t *)tokens->ptr, - 0, - hash_rows, - n_tokens, - has_bias && !hash_mode, - hash_mode); - } - return cuda_ok(cudaGetLastError(), "router_select launch"); -} - -__device__ static float dev_f16_to_f32(uint16_t v) { - return __half2float(*reinterpret_cast(&v)); -} - -__device__ __forceinline__ static uint32_t dev_unpack_iq2_signs(uint32_t v) { - const uint32_t p = __popc(v) & 1u; - const uint32_t s = v ^ (p << 7u); - return s * 0x01010101u; -} - -__device__ __forceinline__ static int32_t dev_iq2_dp4a_8(uint64_t grid, uint32_t sign, const int8_t *q8, int32_t acc) { - const uint32_t signs = dev_unpack_iq2_signs(sign); - const int32_t sm0 = __vcmpne4(signs & 0x08040201u, 0); - const int32_t sm1 = __vcmpne4(signs & 0x80402010u, 0); - const int32_t g0 = __vsub4((int32_t)(uint32_t)grid ^ sm0, sm0); - const int32_t g1 = __vsub4((int32_t)(uint32_t)(grid >> 32) ^ sm1, sm1); - acc = __dp4a(g0, *(const int32_t *)(q8 + 0), acc); - acc = __dp4a(g1, *(const int32_t *)(q8 + 4), acc); - return acc; -} - -__device__ static int32_t dev_dot_q2_16(const uint8_t *q2, const int8_t *q8, int shift) { - int32_t sum = 0; - #pragma unroll - for (uint32_t i = 0; i < 16; i += 4) { - const int32_t v = (*(const int32_t *)(q2 + i) >> shift) & 0x03030303; - sum = __dp4a(v, *(const int32_t *)(q8 + i), sum); - } - return sum; -} - -__device__ static int32_t dev_dot_iq2_pair_16(uint8_t grid0, uint32_t sign0, uint8_t grid1, uint32_t sign1, const int8_t *q8) { - int32_t sum = 0; - sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid0], cuda_ksigns_iq2xs[sign0], q8, sum); - sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid1], cuda_ksigns_iq2xs[sign1], q8 + 8, sum); - return sum; -} - -__device__ __forceinline__ static void dev_iq2_i8x8_lut( - const uint64_t *grid, - const uint8_t *signs, - uint8_t grid_idx, - uint32_t sign_idx, - int32_t *w0, - int32_t *w1) { - const uint32_t s = dev_unpack_iq2_signs(signs[sign_idx]); - const int32_t sm0 = __vcmpne4(s & 0x08040201u, 0); - const int32_t sm1 = __vcmpne4(s & 0x80402010u, 0); - const uint64_t g = grid[grid_idx]; - *w0 = __vsub4((int32_t)(uint32_t)g ^ sm0, sm0); - *w1 = __vsub4((int32_t)(uint32_t)(g >> 32) ^ sm1, sm1); -} - -__device__ static float dev_dot_iq2_xxs_q8_K_block_lut( - const cuda_block_iq2_xxs *x, - const cuda_block_q8_K *y, - const uint64_t *grid, - const uint8_t *signs) { - const float xd = dev_f16_to_f32(x->d); - const uint16_t *q2 = x->qs; - const int8_t *q8 = y->qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u); - int32_t w[8]; - dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]); - int32_t sumi = 0; - sumi = __dp4a(w[0], *(const int32_t *)(q8 + ib32 * 32u + 0), sumi); - sumi = __dp4a(w[1], *(const int32_t *)(q8 + ib32 * 32u + 4), sumi); - sumi = __dp4a(w[2], *(const int32_t *)(q8 + ib32 * 32u + 8), sumi); - sumi = __dp4a(w[3], *(const int32_t *)(q8 + ib32 * 32u + 12), sumi); - sumi = __dp4a(w[4], *(const int32_t *)(q8 + ib32 * 32u + 16), sumi); - sumi = __dp4a(w[5], *(const int32_t *)(q8 + ib32 * 32u + 20), sumi); - sumi = __dp4a(w[6], *(const int32_t *)(q8 + ib32 * 32u + 24), sumi); - sumi = __dp4a(w[7], *(const int32_t *)(q8 + ib32 * 32u + 28), sumi); - bsum += sumi * ls; - } - return 0.125f * xd * y->d * (float)bsum; -} - -__device__ static float dev_dot_iq2_xxs_q8_K_block(const cuda_block_iq2_xxs *x, const cuda_block_q8_K *y) { - const float d = dev_f16_to_f32(x->d) * y->d; - const uint16_t *q2 = x->qs; - const int8_t *q8 = y->qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const uint32_t ls = 2u * (aux1 >> 28) + 1u; - const uint8_t a0 = (uint8_t)(aux0 & 0xffu); - const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); - const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); - const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); - int32_t sumi = 0; - sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8); - q8 += 16; - sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8); - q8 += 16; - bsum += sumi * (int32_t)ls; - } - return 0.125f * d * (float)bsum; -} - -__device__ static void dev_dot_iq2_xxs_q8_K_block8_deq_lut( - const cuda_block_iq2_xxs *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - const cuda_block_q8_K *y4, - const cuda_block_q8_K *y5, - const cuda_block_q8_K *y6, - const cuda_block_q8_K *y7, - uint32_t n, - float acc[8], - const uint64_t *grid, - const uint8_t *signs) { - const float xd = dev_f16_to_f32(x->d); - const uint16_t *q2 = x->qs; - int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const int8_t *q8[8] = { - y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL, - y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL, - }; - for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u); - int32_t w[8]; - dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]); - dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]); - for (uint32_t p = 0; p < n; p++) { - const int8_t *q = q8[p] + ib32 * 32; - int32_t sumi = 0; - sumi = __dp4a(w[0], *(const int32_t *)(q + 0), sumi); - sumi = __dp4a(w[1], *(const int32_t *)(q + 4), sumi); - sumi = __dp4a(w[2], *(const int32_t *)(q + 8), sumi); - sumi = __dp4a(w[3], *(const int32_t *)(q + 12), sumi); - sumi = __dp4a(w[4], *(const int32_t *)(q + 16), sumi); - sumi = __dp4a(w[5], *(const int32_t *)(q + 20), sumi); - sumi = __dp4a(w[6], *(const int32_t *)(q + 24), sumi); - sumi = __dp4a(w[7], *(const int32_t *)(q + 28), sumi); - bsum[p] += sumi * ls; - } - } - const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; - for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; -} - -__device__ static void dev_dot_iq2_xxs_q8_K_block4( - const cuda_block_iq2_xxs *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - uint32_t n, - float acc[4]) { - const float xd = dev_f16_to_f32(x->d); - const uint16_t *q2 = x->qs; - int32_t bsum[4] = {0, 0, 0, 0}; - const int8_t *q8[4] = { - y0 ? y0->qs : NULL, - y1 ? y1->qs : NULL, - y2 ? y2->qs : NULL, - y3 ? y3->qs : NULL, - }; - for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const uint32_t ls = 2u * (aux1 >> 28) + 1u; - const uint8_t a0 = (uint8_t)(aux0 & 0xffu); - const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); - const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); - const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); - for (uint32_t p = 0; p < n; p++) { - int32_t sumi = 0; - sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32); - sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16); - bsum[p] += sumi * (int32_t)ls; - } - } - const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 }; - for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; -} - -__device__ static DS4_CUDA_UNUSED void dev_dot_iq2_xxs_q8_K_block8( - const cuda_block_iq2_xxs *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - const cuda_block_q8_K *y4, - const cuda_block_q8_K *y5, - const cuda_block_q8_K *y6, - const cuda_block_q8_K *y7, - uint32_t n, - float acc[8]) { - const float xd = dev_f16_to_f32(x->d); - const uint16_t *q2 = x->qs; - int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const int8_t *q8[8] = { - y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL, - y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL, - }; - for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const uint32_t ls = 2u * (aux1 >> 28) + 1u; - const uint8_t a0 = (uint8_t)(aux0 & 0xffu); - const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); - const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); - const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); - for (uint32_t p = 0; p < n; p++) { - int32_t sumi = 0; - sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32); - sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16); - bsum[p] += sumi * (int32_t)ls; - } - } - const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; - for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; -} - -__device__ static void dev_q4_K_get_scale_min( - uint32_t j, - const uint8_t *scales, - uint8_t *d_out, - uint8_t *m_out) { - if (j < 4u) { - *d_out = scales[j] & 63u; - *m_out = scales[j + 4u] & 63u; - } else { - *d_out = (scales[j + 4u] & 0x0fu) | ((scales[j - 4u] >> 6u) << 4u); - *m_out = (scales[j + 4u] >> 4u) | ((scales[j] >> 6u) << 4u); - } -} - -__device__ __forceinline__ static int32_t dev_dot_q4_32(const uint8_t *qs, const int8_t *q8, int shift) { - int32_t sum = 0; - #pragma unroll - for (uint32_t i = 0; i < 32u; i += 4u) { - const int32_t v = (*(const int32_t *)(qs + i) >> shift) & 0x0f0f0f0f; - sum = __dp4a(v, *(const int32_t *)(q8 + i), sum); - } - return sum; -} - -__device__ static float dev_dot_q4_K_q8_K_block(const cuda_block_q4_K *x, const cuda_block_q8_K *y) { - const float xd = dev_f16_to_f32(x->d); - const float xmin = dev_f16_to_f32(x->dmin); - int isum = 0; - int summs = 0; - #pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - uint8_t sc, m; - dev_q4_K_get_scale_min(j, x->scales, &sc, &m); - summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]); - const uint32_t byte_off = (j >> 1u) * 32u; - const int shift = (j & 1u) ? 4 : 0; - isum += (int)sc * dev_dot_q4_32(x->qs + byte_off, y->qs + j * 32u, shift); - } - return y->d * xd * (float)isum - y->d * xmin * (float)summs; -} - -/* Vector-load variant of dev_dot_q4_K_q8_K_block: loads the whole 144-byte - * Q4_K block with nine 16B loads (requires a 16B-aligned tensor base; block - * stride 144 and row strides are 16B multiples), then computes the exact same - * integer sums and float finish. Same values in the same order, so results - * are bit-identical; the wide loads just improve DRAM/memory-level - * parallelism for the bandwidth-bound decode matvecs. */ -__device__ __forceinline__ static void dev_dot_q4_K_q8_K_block_vec( - const cuda_block_q4_K *x, - const cuda_block_q8_K *y, - float *out_acc) { - const uint4 hdr = *(const uint4 *)x; /* d, dmin, scales[12] */ - uint4 qv[8]; -#pragma unroll - for (uint32_t i = 0; i < 8u; i++) qv[i] = ((const uint4 *)(x->qs))[i]; - const uint16_t xd_u = (uint16_t)(hdr.x & 0xffffu); - const uint16_t xmin_u = (uint16_t)(hdr.x >> 16u); - uint8_t scales[12]; - scales[0] = (uint8_t)(hdr.y); - scales[1] = (uint8_t)(hdr.y >> 8); - scales[2] = (uint8_t)(hdr.y >> 16); - scales[3] = (uint8_t)(hdr.y >> 24); - scales[4] = (uint8_t)(hdr.z); - scales[5] = (uint8_t)(hdr.z >> 8); - scales[6] = (uint8_t)(hdr.z >> 16); - scales[7] = (uint8_t)(hdr.z >> 24); - scales[8] = (uint8_t)(hdr.w); - scales[9] = (uint8_t)(hdr.w >> 8); - scales[10] = (uint8_t)(hdr.w >> 16); - scales[11] = (uint8_t)(hdr.w >> 24); - const float xd = dev_f16_to_f32(xd_u); - const float xmin = dev_f16_to_f32(xmin_u); - int isum = 0; - int summs = 0; - const int32_t *qw = (const int32_t *)qv; -#pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - uint8_t sc, m; - dev_q4_K_get_scale_min(j, scales, &sc, &m); - summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]); - const uint32_t word_off = (j >> 1u) * 8u; - const int shift = (j & 1u) ? 4 : 0; - int32_t sum = 0; -#pragma unroll - for (uint32_t i = 0; i < 8u; i++) { - const int32_t v = (qw[word_off + i] >> shift) & 0x0f0f0f0f; - sum = __dp4a(v, *(const int32_t *)(y->qs + j * 32u + i * 4u), sum); - } - isum += (int)sc * sum; - } - *out_acc += y->d * xd * (float)isum - y->d * xmin * (float)summs; -} - -__device__ static void dev_dot_q4_K_q8_K_block8( - const cuda_block_q4_K *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - const cuda_block_q8_K *y4, - const cuda_block_q8_K *y5, - const cuda_block_q8_K *y6, - const cuda_block_q8_K *y7, - uint32_t n, - float acc[8]) { - const float xd = dev_f16_to_f32(x->d); - const float xmin = dev_f16_to_f32(x->dmin); - const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; - int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - - #pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - uint8_t sc, m; - dev_q4_K_get_scale_min(j, x->scales, &sc, &m); - const uint32_t byte_off = (j >> 1u) * 32u; - const int shift = (j & 1u) ? 4 : 0; - for (uint32_t p = 0; p < n; p++) { - summs[p] += (int)m * (int)(ys[p]->bsums[2u * j] + ys[p]->bsums[2u * j + 1u]); - isum[p] += (int)sc * dev_dot_q4_32(x->qs + byte_off, ys[p]->qs + j * 32u, shift); - } - } - - for (uint32_t p = 0; p < n; p++) { - acc[p] += ys[p]->d * xd * (float)isum[p] - ys[p]->d * xmin * (float)summs[p]; - } -} - -__device__ static float dev_dot_q2_K_q8_K_block(const cuda_block_q2_K *x, const cuda_block_q8_K *y) { - const uint8_t *q2 = x->qs; - const int8_t *q8 = y->qs; - const uint8_t *sc = x->scales; - int summs = 0; - for (int j = 0; j < 16; j++) summs += y->bsums[j] * (sc[j] >> 4); - const float dall = y->d * dev_f16_to_f32(x->d); - const float dmin = y->d * dev_f16_to_f32(x->dmin); - int isum = 0; - int is = 0; - for (int k = 0; k < CUDA_QK_K / 128; k++) { - int shift = 0; - for (int j = 0; j < 4; j++) { - int d = sc[is++] & 0x0f; - isum += d * dev_dot_q2_16(q2, q8, shift); - d = sc[is++] & 0x0f; - isum += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); - shift += 2; - q8 += 32; - } - q2 += 32; - } - return dall * (float)isum - dmin * (float)summs; -} - -__device__ static void dev_dot_q2_K_q8_K_block4( - const cuda_block_q2_K *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - uint32_t n, - float acc[4]) { - const uint8_t *sc = x->scales; - const float xd = dev_f16_to_f32(x->d); - const float xmin = dev_f16_to_f32(x->dmin); - const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 }; - int isum[4] = {0, 0, 0, 0}; - int summs[4] = {0, 0, 0, 0}; - for (uint32_t p = 0; p < n; p++) { - for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4); - } - for (uint32_t p = 0; p < n; p++) { - const uint8_t *q2 = x->qs; - const int8_t *q8 = ys[p]->qs; - int is = 0; - for (int k = 0; k < CUDA_QK_K / 128; k++) { - int shift = 0; - for (int j = 0; j < 4; j++) { - int d = sc[is++] & 0x0f; - isum[p] += d * dev_dot_q2_16(q2, q8, shift); - d = sc[is++] & 0x0f; - isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); - shift += 2; - q8 += 32; - } - q2 += 32; - } - } - for (uint32_t p = 0; p < n; p++) { - const float yd = ys[p]->d; - acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p]; - } -} - -__device__ static void dev_dot_q2_K_q8_K_block8( - const cuda_block_q2_K *x, - const cuda_block_q8_K *y0, - const cuda_block_q8_K *y1, - const cuda_block_q8_K *y2, - const cuda_block_q8_K *y3, - const cuda_block_q8_K *y4, - const cuda_block_q8_K *y5, - const cuda_block_q8_K *y6, - const cuda_block_q8_K *y7, - uint32_t n, - float acc[8]) { - const uint8_t *sc = x->scales; - const float xd = dev_f16_to_f32(x->d); - const float xmin = dev_f16_to_f32(x->dmin); - const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; - int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - for (uint32_t p = 0; p < n; p++) { - for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4); - } - for (uint32_t p = 0; p < n; p++) { - const uint8_t *q2 = x->qs; - const int8_t *q8 = ys[p]->qs; - int is = 0; - for (int k = 0; k < CUDA_QK_K / 128; k++) { - int shift = 0; - for (int j = 0; j < 4; j++) { - int d = sc[is++] & 0x0f; - isum[p] += d * dev_dot_q2_16(q2, q8, shift); - d = sc[is++] & 0x0f; - isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); - shift += 2; - q8 += 32; - } - q2 += 32; - } - } - for (uint32_t p = 0; p < n; p++) { - const float yd = ys[p]->d; - acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p]; - } -} - -__device__ static float half_warp_sum_f32(float v, uint32_t lane16) { - uint32_t mask = 0xffffu << (threadIdx.x & 16u); - for (int offset = 8; offset > 0; offset >>= 1) { - v += __shfl_down_sync(mask, v, offset, 16); - } - (void)lane16; - return v; -} - -__device__ static float quarter_warp_sum_f32(float v, uint32_t lane8) { - uint32_t mask = 0xffu << (threadIdx.x & 24u); - for (int offset = 4; offset > 0; offset >>= 1) { - v += __shfl_down_sync(mask, v, offset, 8); - } - (void)lane8; - return v; -} - -__global__ static void q8_K_quantize_kernel(cuda_block_q8_K *out, const float *x, uint32_t in_dim, uint32_t n_rows) { - uint32_t b = blockIdx.x; - uint32_t row = blockIdx.y; - if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; - const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; - cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b; - __shared__ float abs_part[256]; - __shared__ float val_part[256]; - __shared__ float maxv_s; - __shared__ float iscale_s; - uint32_t tid = threadIdx.x; - float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; - abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; - val_part[tid] = v; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { - abs_part[tid] = abs_part[tid + stride]; - val_part[tid] = val_part[tid + stride]; - } - __syncthreads(); - } - float amax = abs_part[0]; - if (amax == 0.0f) { - if (tid == 0) yb->d = 0.0f; - if (tid < CUDA_QK_K) yb->qs[tid] = 0; - if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; - return; - } - if (tid == 0) { - maxv_s = val_part[0]; - iscale_s = -127.0f / maxv_s; - } - __syncthreads(); - if (tid < CUDA_QK_K) { - int qv = (int)lrintf(iscale_s * xr[tid]); - if (qv > 127) qv = 127; - if (qv < -128) qv = -128; - yb->qs[tid] = (int8_t)qv; - } - __syncthreads(); - if (tid < CUDA_QK_K / 16) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; - yb->bsums[tid] = (int16_t)sum; - } - if (tid == 0) yb->d = 1.0f / iscale_s; -} - -/* Decode-only dual quantizer. The Q8_0 half mirrors - * quantize_q8_0_f32_kernel's 32-thread reduction and expression order, while - * the Q8_K half remains byte-for-byte the ordinary routed-MoE quantizer. */ -__global__ static void q8_K_q8_0_quantize_kernel( - cuda_block_q8_K *out, - int8_t *q8_0, - float *q8_0_scale, - const float *x, - uint32_t in_dim, - uint32_t n_rows) { - const uint32_t b = blockIdx.x; - const uint32_t row = blockIdx.y; - if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; - const float *xr = x + (uint64_t)row * in_dim + - (uint64_t)b * CUDA_QK_K; - cuda_block_q8_K *yb = out + - (uint64_t)row * (in_dim / CUDA_QK_K) + b; - __shared__ float abs_part[256]; - __shared__ float val_part[256]; - __shared__ float maxv_s; - __shared__ float iscale_s; - const uint32_t tid = threadIdx.x; - const uint32_t lane = tid & 31u; - const uint32_t warp = tid >> 5u; - const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; - - abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; - __syncthreads(); - for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { - if (lane < stride) { - abs_part[tid] = fmaxf(abs_part[tid], abs_part[tid + stride]); - } - __syncthreads(); - } - const uint32_t q8_blocks = in_dim / 32u; - const uint32_t q8_block = b * 8u + warp; - const float d = abs_part[warp * 32u] / 127.0f; - const float id = d != 0.0f ? 1.0f / d : 0.0f; - if (lane == 0u) { - q8_0_scale[(uint64_t)row * q8_blocks + q8_block] = d; - } - int qv = (int)lrintf(v * id); - qv = qv > 127 ? 127 : (qv < -128 ? -128 : qv); - q8_0[((uint64_t)row * q8_blocks + q8_block) * 32u + lane] = - (int8_t)qv; - __syncthreads(); - - abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; - val_part[tid] = v; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { - abs_part[tid] = abs_part[tid + stride]; - val_part[tid] = val_part[tid + stride]; - } - __syncthreads(); - } - const float amax = abs_part[0]; - if (amax == 0.0f) { - if (tid == 0) yb->d = 0.0f; - if (tid < CUDA_QK_K) yb->qs[tid] = 0; - if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; - return; - } - if (tid == 0) { - maxv_s = val_part[0]; - iscale_s = -127.0f / maxv_s; - } - __syncthreads(); - if (tid < CUDA_QK_K) { - int kv = (int)lrintf(iscale_s * xr[tid]); - if (kv > 127) kv = 127; - if (kv < -128) kv = -128; - yb->qs[tid] = (int8_t)kv; - } - __syncthreads(); - if (tid < CUDA_QK_K / 16) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; - yb->bsums[tid] = (int16_t)sum; - } - if (tid == 0) yb->d = 1.0f / iscale_s; -} - -__device__ __forceinline__ static bool moe_owned_local_expert( - int32_t expert, - uint32_t expert_base, - uint32_t expert_count, - uint32_t *local_expert) { - if (expert < 0) return false; - const uint32_t e = (uint32_t)expert; - if (e < expert_base || e - expert_base >= expert_count) return false; - if (local_expert) *local_expert = e - expert_base; - return true; -} - -/* Quantize only selected slots owned by this expert-parallel rank. Rows keep - * their original slot index so the final rank-local reduction can visit slots - * in canonical order without compaction or a host synchronization. */ -__global__ static void q8_K_quantize_owned_kernel( - cuda_block_q8_K *out, - const float *x, - const int32_t *selected, - uint32_t in_dim, - uint32_t n_rows, - uint32_t expert_base, - uint32_t expert_count) { - const uint32_t b = blockIdx.x; - const uint32_t row = blockIdx.y; - if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; - if (!moe_owned_local_expert(selected[row], expert_base, expert_count, NULL)) return; - - const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; - cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b; - __shared__ float abs_part[256]; - __shared__ float val_part[256]; - __shared__ float maxv_s; - __shared__ float iscale_s; - const uint32_t tid = threadIdx.x; - const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; - abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; - val_part[tid] = v; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { - abs_part[tid] = abs_part[tid + stride]; - val_part[tid] = val_part[tid + stride]; - } - __syncthreads(); - } - const float amax = abs_part[0]; - if (amax == 0.0f) { - if (tid == 0) yb->d = 0.0f; - if (tid < CUDA_QK_K) yb->qs[tid] = 0; - if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; - return; - } - if (tid == 0) { - maxv_s = val_part[0]; - iscale_s = -127.0f / maxv_s; - } - __syncthreads(); - if (tid < CUDA_QK_K) { - int qv = (int)lrintf(iscale_s * xr[tid]); - if (qv > 127) qv = 127; - if (qv < -128) qv = -128; - yb->qs[tid] = (int8_t)qv; - } - __syncthreads(); - if (tid < CUDA_QK_K / 16) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; - yb->bsums[tid] = (int16_t)sum; - } - if (tid == 0) yb->d = 1.0f / iscale_s; -} - -__global__ static void moe_filter_owned_pairs_kernel( - int32_t *selected, - float *weights, - uint64_t pair_count, - uint32_t n_total_expert, - uint32_t expert_base, - uint32_t expert_count) { - const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (pair >= pair_count) return; - const int32_t expert_i = selected[pair]; - if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert && - (uint32_t)expert_i >= expert_base && - (uint32_t)expert_i - expert_base < expert_count) { - selected[pair] = expert_i - (int32_t)expert_base; - } else { - selected[pair] = -1; - weights[pair] = 0.0f; - } -} - -__global__ static void q8_K_quantize_sidecar_kernel( - cuda_block_q8_K *out, - const float *x, - const float *amax_sidecar, - uint32_t in_dim, - uint32_t n_rows) { - uint32_t b = blockIdx.x; - uint32_t row = blockIdx.y; - const uint32_t blocks = in_dim / CUDA_QK_K; - if (row >= n_rows || b >= blocks) return; - const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; - const float *sc = amax_sidecar + ((uint64_t)row * blocks + b) * 32u; - cuda_block_q8_K *yb = out + (uint64_t)row * blocks + b; - __shared__ float abs_part[32]; - __shared__ float val_part[32]; - __shared__ float iscale_s; - const uint32_t tid = threadIdx.x; - if (tid < 32u) { - const float v = sc[tid]; - abs_part[tid] = fabsf(v); - val_part[tid] = v; - } - __syncthreads(); - for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { - if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { - abs_part[tid] = abs_part[tid + stride]; - val_part[tid] = val_part[tid + stride]; - } - __syncthreads(); - } - const float amax = abs_part[0]; - if (amax == 0.0f) { - if (tid == 0u) yb->d = 0.0f; - if (tid < CUDA_QK_K) yb->qs[tid] = 0; - if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0; - return; - } - if (tid == 0u) iscale_s = -127.0f / val_part[0]; - __syncthreads(); - if (tid < CUDA_QK_K) { - int qv = (int)lrintf(iscale_s * xr[tid]); - if (qv > 127) qv = 127; - if (qv < -128) qv = -128; - yb->qs[tid] = (int8_t)qv; - } - __syncthreads(); - if (tid < CUDA_QK_K / 16u) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i]; - yb->bsums[tid] = (int16_t)sum; - } - if (tid == 0u) yb->d = 1.0f / iscale_s; -} - -__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t row = blockIdx.x; - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = threadIdx.x; b < xq_blocks; b += blockDim.x) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - __shared__ float partial_gate[256]; - __shared__ float partial_up[256]; - partial_gate[threadIdx.x] = gate; - partial_up[threadIdx.x] = up; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride]; - partial_up[threadIdx.x] += partial_up[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - gate = partial_gate[0]; - up = partial_up[0]; - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_warp8_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t lane = threadIdx.x & 31u; - uint32_t warp = threadIdx.x >> 5u; - uint32_t row = blockIdx.x * 8u + warp; - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_hwarp16_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t lane = threadIdx.x & 15u; - uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 16u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = half_warp_sum_f32(gate, lane); - up = half_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -// perf-04: launch-geometry tuning for the routed-MoE gate/up decode kernels -// (moe_gate_up_mid_qwarp32 / _decode_lut_qwarp32 / _decode_q4K_qwarp32). Each -// block processes MOE_DECODE_ROW_TILES tiles of 32 rows (row_lane in [0,32)). -// The historical value was 4 (128 rows/block -> ~96 blocks, occupancy ~16%, -// "grid too small to fill the device"). Lowering it issues correspondingly more -// blocks (e.g. 1 tile -> 32 rows/block -> ~4x more blocks -> ~384) to fill the -// SMs. The per-row arithmetic is identical regardless of this value, so output -// is bit-identical; only the qgrid.x divisor must match MOE_DECODE_ROWS_PER_BLOCK. -#ifndef MOE_DECODE_ROW_TILES -#define MOE_DECODE_ROW_TILES 1u -#endif -#define MOE_DECODE_ROWS_PER_BLOCK (32u * MOE_DECODE_ROW_TILES) - -__global__ static void moe_gate_up_mid_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t pair = blockIdx.y; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { - uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } - } -} - -__global__ static void moe_gate_up_mid_decode_lut_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t pair = blockIdx.y; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - __shared__ uint64_t s_iq2_grid[256]; - __shared__ uint8_t s_iq2_signs[128]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - xqb = sxq; - } - for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { - uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs); - up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } - } -} - -__global__ static void moe_gate_up_mid_decode_lut_owned_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t expert_base, - uint32_t expert_count, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t pair = blockIdx.y; - uint32_t expert = 0u; - if (!moe_owned_local_expert(selected[pair], expert_base, expert_count, - &expert)) return; - const cuda_block_q8_K *xqb = xq; - __shared__ cuda_block_q8_K sxq[16]; - __shared__ uint64_t s_iq2_grid[256]; - __shared__ uint8_t s_iq2_signs[128]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - xqb = sxq; - } - for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { - uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs); - up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair]; - } - } -} - -__global__ static void moe_count_sorted_pairs_kernel( - uint32_t *counts, - const int32_t *selected, - uint32_t pair_count, - uint32_t n_total_expert) { - uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); - if (pair >= pair_count) return; - int32_t expert_i = selected[pair]; - if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return; - atomicAdd(counts + (uint32_t)expert_i, 1u); -} - -__global__ static void moe_prefix_sorted_pairs_kernel( - uint32_t *offsets, - uint32_t *cursors, - const uint32_t *counts, - uint32_t n_total_expert) { - if (threadIdx.x == 0) { - uint32_t sum = 0; - for (uint32_t e = 0; e < n_total_expert; e++) { - offsets[e] = sum; - cursors[e] = sum; - sum += counts[e]; - } - offsets[n_total_expert] = sum; - } -} - -__global__ static void moe_scatter_sorted_pairs_kernel( - uint32_t *sorted_pairs, - uint32_t *cursors, - const int32_t *selected, - uint32_t pair_count, - uint32_t n_total_expert) { - uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); - if (pair >= pair_count) return; - int32_t expert_i = selected[pair]; - if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return; - uint32_t pos = atomicAdd(cursors + (uint32_t)expert_i, 1u); - sorted_pairs[pos] = pair; -} - -__global__ static void moe_build_expert_tile_offsets_kernel( - uint32_t *tile_offsets, - uint32_t *tile_total, - const uint32_t *counts, - uint32_t block_m, - uint32_t n_total_expert) { - if (threadIdx.x == 0) { - uint32_t sum = 0; - for (uint32_t e = 0; e < n_total_expert; e++) { - tile_offsets[e] = sum; - sum += (counts[e] + block_m - 1u) / block_m; - } - tile_offsets[n_total_expert] = sum; - *tile_total = sum; - } -} - -__global__ static void moe_build_expert_tiles_kernel( - uint32_t *tile_experts, - uint32_t *tile_starts, - const uint32_t *tile_offsets, - const uint32_t *counts, - uint32_t block_m, - uint32_t n_total_expert) { - uint32_t e = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); - if (e >= n_total_expert) return; - uint32_t ntiles = (counts[e] + block_m - 1u) / block_m; - uint32_t off = tile_offsets[e]; - for (uint32_t t = 0; t < ntiles; t++) { - tile_experts[off + t] = e; - tile_starts[off + t] = t * block_m; - } -} - -/* Decode-sized routed batches spend more host time launching metadata kernels - * than doing the <= 96 pair / 128 expert setup. Build both tile lists in one - * deterministic block; the expensive expert kernels remain unchanged. */ -__global__ static void moe_prepare_sorted_tiles_small_kernel( - uint32_t *counts, - uint32_t *offsets, - uint32_t *cursors, - uint32_t *sorted_pairs, - uint32_t *tile_offsets, - uint32_t *tile_total, - uint32_t *tile_experts, - uint32_t *tile_starts, - uint32_t *tile16_offsets, - uint32_t *tile16_total, - uint32_t *tile16_experts, - uint32_t *tile16_starts, - const int32_t *selected, - uint32_t pair_count, - uint32_t n_total_expert, - uint32_t block_m, - bool build_tile16) { - if (blockIdx.x != 0) return; - const uint32_t tid = threadIdx.x; - __shared__ uint32_t local_counts[128]; - __shared__ int32_t local_selected[96]; - - for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { - local_counts[e] = 0u; - } - for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { - local_selected[pair] = selected[pair]; - } - __syncthreads(); - - for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { - const int32_t expert_i = local_selected[pair]; - if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) { - atomicAdd(local_counts + (uint32_t)expert_i, 1u); - } - } - __syncthreads(); - - for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { - counts[e] = local_counts[e]; - } - if (tid == 0u) { - uint32_t pair_sum = 0u; - uint32_t tile_sum = 0u; - uint32_t tile16_sum = 0u; - for (uint32_t e = 0; e < n_total_expert; e++) { - const uint32_t count = local_counts[e]; - offsets[e] = pair_sum; - pair_sum += count; - cursors[e] = pair_sum; - tile_offsets[e] = tile_sum; - tile_sum += (count + block_m - 1u) / block_m; - if (build_tile16) { - tile16_offsets[e] = tile16_sum; - tile16_sum += (count + 15u) / 16u; - } - } - offsets[n_total_expert] = pair_sum; - tile_offsets[n_total_expert] = tile_sum; - *tile_total = tile_sum; - if (build_tile16) { - tile16_offsets[n_total_expert] = tile16_sum; - *tile16_total = tile16_sum; - } - } - __syncthreads(); - - for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { - const int32_t expert_i = local_selected[pair]; - if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) { - uint32_t rank = 0u; - for (uint32_t prev = 0; prev < pair; prev++) { - rank += local_selected[prev] == expert_i; - } - sorted_pairs[offsets[(uint32_t)expert_i] + rank] = pair; - } - } - for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { - const uint32_t count = local_counts[e]; - const uint32_t ntiles = (count + block_m - 1u) / block_m; - const uint32_t tile_off = tile_offsets[e]; - for (uint32_t t = 0; t < ntiles; t++) { - tile_experts[tile_off + t] = e; - tile_starts[tile_off + t] = t * block_m; - } - if (build_tile16) { - const uint32_t ntiles16 = (count + 15u) / 16u; - const uint32_t tile16_off = tile16_offsets[e]; - for (uint32_t t = 0; t < ntiles16; t++) { - tile16_experts[tile16_off + t] = e; - tile16_starts[tile16_off + t] = t * 16u; - } - } - } -} - -__global__ static void moe_gate_up_mid_sorted_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t pair = sorted_pairs[blockIdx.y]; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_expert_tile8_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t group = threadIdx.x >> 3u; - uint32_t lane = threadIdx.x & 7u; - uint32_t pair_slot = group & 7u; - uint32_t row_lane = group >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_pair = tile_starts[tile] + pair_slot; - if (local_pair >= counts[expert]) return; - uint32_t sorted_idx = offsets[expert] + local_pair; - uint32_t pair = sorted_pairs[sorted_idx]; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - - for (uint32_t rr = 0; rr < 2u; rr++) { - uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } - } -} - -__global__ static void moe_gate_up_mid_expert_tile4_row32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[4][16]; - uint32_t pair[4] = {0, 0, 0, 0}; - uint32_t tok[4] = {0, 0, 0, 0}; - uint32_t slot[4] = {0, 0, 0, 0}; - const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 4u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - tok[np] = pair[np] / n_expert; - slot[np] = pair[np] - tok[np] * n_expert; - xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; - } - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { - uint32_t p = i / xq_blocks; - uint32_t b = i - p * xq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - if (row >= expert_mid_dim) return; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate[4] = {0.0f, 0.0f, 0.0f, 0.0f}; - float up[4] = {0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - dev_dot_iq2_xxs_q8_K_block4(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, gate); - dev_dot_iq2_xxs_q8_K_block4(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, up); - } - for (uint32_t p = 0; p < np; p++) { - gate[p] = quarter_warp_sum_f32(gate[p], lane); - up[p] = quarter_warp_sum_f32(up[p], lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate[p] > clamp) gate[p] = clamp; - if (up[p] > clamp) up[p] = clamp; - if (up[p] < -clamp) up[p] = -clamp; - } - const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate[p]; - up_out[off] = up[p]; - } - mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; - } - } -} - -__global__ static void moe_gate_up_mid_expert_tile8_row32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][16]; - __shared__ uint64_t s_iq2_grid[256]; - __shared__ uint8_t s_iq2_signs[128]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - tok[np] = pair[np] / n_expert; - slot[np] = pair[np] - tok[np] * n_expert; - xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; - } - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { - uint32_t p = i / xq_blocks; - uint32_t b = i - p * xq_blocks; - sxq[p][b] = xqb[p][b]; - } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - if (row >= expert_mid_dim) return; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, - s_iq2_grid, s_iq2_signs); - dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, - s_iq2_grid, s_iq2_signs); - } - for (uint32_t p = 0; p < np; p++) { - gate[p] = quarter_warp_sum_f32(gate[p], lane); - up[p] = quarter_warp_sum_f32(up[p], lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate[p] > clamp) gate[p] = clamp; - if (up[p] > clamp) up[p] = clamp; - if (up[p] < -clamp) up[p] = -clamp; - } - const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate[p]; - up_out[off] = up[p]; - } - mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; - } - } -} - -__global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][16]; - __shared__ uint64_t s_iq2_grid[256]; - __shared__ uint8_t s_iq2_signs[128]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - tok[np] = pair[np] / n_expert; - slot[np] = pair[np] - tok[np] * n_expert; - xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; - } - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { - uint32_t p = i / xq_blocks; - uint32_t b = i - p * xq_blocks; - sxq[p][b] = xqb[p][b]; - } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < 64u; rr++) { - uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, - s_iq2_grid, s_iq2_signs); - dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, - s_iq2_grid, s_iq2_signs); - } - for (uint32_t p = 0; p < np; p++) { - gate[p] = quarter_warp_sum_f32(gate[p], lane); - up[p] = quarter_warp_sum_f32(up[p], lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate[p] > clamp) gate[p] = clamp; - if (up[p] > clamp) up[p] = clamp; - if (up[p] < -clamp) up[p] = -clamp; - } - const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate[p]; - up_out[off] = up[p]; - } - mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; - } - } - } -} - -template -__global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][16]; - __shared__ uint64_t s_iq2_grid[256]; - __shared__ uint8_t s_iq2_signs[128]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - tok[np] = pair[np] / n_expert; - slot[np] = pair[np] - tok[np] * n_expert; - xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; - } - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { - uint32_t p = i / xq_blocks; - uint32_t b = i - p * xq_blocks; - sxq[p][b] = xqb[p][b]; - } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { - uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, - s_iq2_grid, s_iq2_signs); - dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, - s_iq2_grid, s_iq2_signs); - } - for (uint32_t p = 0; p < np; p++) { - gate[p] = quarter_warp_sum_f32(gate[p], lane); - up[p] = quarter_warp_sum_f32(up[p], lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate[p] > clamp) gate[p] = clamp; - if (up[p] > clamp) up[p] = clamp; - if (up[p] < -clamp) up[p] = -clamp; - } - const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate[p]; - up_out[off] = up[p]; - } - mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; - } - } - } -} - -__global__ static void moe_gate_up_mid_sorted_p2_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t pair_count, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t pair_lane = (threadIdx.x >> 3u) & 1u; - uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); - uint32_t sorted_idx = blockIdx.y * 2u + pair_lane; - if (row >= expert_mid_dim || sorted_idx >= pair_count) return; - uint32_t pair = sorted_pairs[sorted_idx]; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); - up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static DS4_CUDA_UNUSED void moe_down_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t row = blockIdx.x; - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = threadIdx.x; b < midq_blocks; b += blockDim.x) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - __shared__ float partial[256]; - partial[threadIdx.x] = acc; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0]; -} - -__global__ static DS4_CUDA_UNUSED void moe_down_warp8_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t lane = threadIdx.x & 31u; - uint32_t warp = threadIdx.x >> 5u; - uint32_t row = blockIdx.x * 8u + warp; - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 32u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = warp_sum_f32(acc); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -__global__ static DS4_CUDA_UNUSED void moe_down_hwarp16_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t lane = threadIdx.x & 15u; - uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 16u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = half_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -__global__ static void moe_down_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -__global__ static void moe_gate_up_mid_decode_q4K_qwarp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t pair = blockIdx.y; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { - uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 15u; - uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 16u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = half_warp_sum_f32(gate, lane); - up = half_warp_sum_f32(up, lane); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 15u; - uint32_t group = threadIdx.x >> 4u; - uint32_t pair = blockIdx.y; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - if (group >= 8u) return; - uint32_t row = blockIdx.x * 8u + group; - if (row >= expert_mid_dim) return; - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 16u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = half_warp_sum_f32(gate, lane); - up = half_warp_sum_f32(up, lane); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_warp32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 31u; - uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_kernel( - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t lane = threadIdx.x & 31u; - uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - /* Word-wise cooperative staging copy (same bytes, all lanes busy). */ - const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - uint32_t *dst = (uint32_t *)sxq; - const uint32_t *srcw = (const uint32_t *)xqb; - for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i]; - __syncthreads(); - xqb = sxq; - } - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base | - gate_row_bytes | gate_expert_bytes) & 15u) == 0u); - if (vec_ok) { - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate); - dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up); - } - } else { - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel( - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t expert_base, - uint32_t expert_count, - float clamp) { - uint32_t lane = threadIdx.x & 31u; - uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t expert = 0u; - if (!moe_owned_local_expert(selected[pair], expert_base, expert_count, - &expert)) return; - const cuda_block_q8_K *xqb = xq; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - uint32_t *dst = (uint32_t *)sxq; - const uint32_t *srcw = (const uint32_t *)xqb; - for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i]; - __syncthreads(); - xqb = sxq; - } - const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base | - gate_row_bytes | gate_expert_bytes) & 15u) == 0u); - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - if (vec_ok) { - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate); - dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up); - } - } else { - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair]; - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel( - float *mid_out, - float *amax_sidecar, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t row = blockIdx.x * 8u + warp; - const uint32_t pair = blockIdx.y; - const uint32_t tok = pair / n_expert; - const uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - __shared__ float tile_vals[8]; - __shared__ float tile_abs[8]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - - float midv = 0.0f; - const bool valid = row < expert_mid_dim; - if (valid) { - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - midv = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - mid_out[off] = midv; - } - } - if (lane == 0u) { - tile_vals[warp] = midv; - tile_abs[warp] = valid ? fabsf(midv) : 0.0f; - } - __syncthreads(); - if (threadIdx.x == 0u) { - float best_abs = tile_abs[0]; - float best_val = tile_vals[0]; - #pragma unroll - for (uint32_t i = 1u; i < 8u; i++) { - if (tile_abs[i] > best_abs) { - best_abs = tile_abs[i]; - best_val = tile_vals[i]; - } - } - const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; - const uint32_t qblock = blockIdx.x / 32u; - const uint32_t tile = blockIdx.x & 31u; - if (qblock < midq_blocks) { - amax_sidecar[((uint64_t)pair * midq_blocks + qblock) * 32u + tile] = best_val; - } - } -} - -__global__ static void moe_gate_up_mid_decode_q4K_warp32_row16_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t lane = threadIdx.x & 31u; - uint32_t warp = threadIdx.x >> 5u; - uint32_t row = blockIdx.x * 16u + warp; - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = warp_sum_f32(gate); - up = warp_sum_f32(up); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_gate_up_midq_decode_q4K_qwarp32_kernel( - float *mid_out, - cuda_block_q8_K *midq, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row_lane = threadIdx.x >> 3u; - const uint32_t qblock = blockIdx.x; - const uint32_t pair = blockIdx.y; - const uint32_t tok = pair / n_expert; - const uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const uint32_t expert = (uint32_t)expert_i; - const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; - __shared__ cuda_block_q8_K sxq[16]; - __shared__ float vals[CUDA_QK_K]; - __shared__ float abs_part[CUDA_QK_K]; - __shared__ float val_part[CUDA_QK_K]; - __shared__ float iscale_s; - - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - __syncthreads(); - xqb = sxq; - } - - const float w = weights[(uint64_t)tok * n_expert + slot]; - #pragma unroll - for (uint32_t rr = 0; rr < 8u; rr++) { - const uint32_t row_in_block = row_lane + rr * 32u; - const uint32_t row = qblock * CUDA_QK_K + row_in_block; - float midv = 0.0f; - if (row < expert_mid_dim) { - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); - up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); - } - gate = quarter_warp_sum_f32(gate, lane); - up = quarter_warp_sum_f32(up, lane); - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - midv = (gate / (1.0f + expf(-gate))) * up * w; - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - mid_out[off] = midv; - } - } - if (lane == 0u) vals[row_in_block] = midv; - } - __syncthreads(); - - cuda_block_q8_K *yb = midq + (uint64_t)pair * (expert_mid_dim / CUDA_QK_K) + qblock; - const uint32_t tid = threadIdx.x; - const float v = vals[tid]; - abs_part[tid] = fabsf(v); - val_part[tid] = v; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { - abs_part[tid] = abs_part[tid + stride]; - val_part[tid] = val_part[tid + stride]; - } - __syncthreads(); - } - const float amax = abs_part[0]; - if (amax == 0.0f) { - if (tid == 0u) yb->d = 0.0f; - if (tid < CUDA_QK_K) yb->qs[tid] = 0; - if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0; - return; - } - if (tid == 0u) { - iscale_s = -127.0f / val_part[0]; - } - __syncthreads(); - int qv = (int)lrintf(iscale_s * v); - if (qv > 127) qv = 127; - if (qv < -128) qv = -128; - yb->qs[tid] = (int8_t)qv; - __syncthreads(); - if (tid < CUDA_QK_K / 16u) { - int sum = 0; - for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i]; - yb->bsums[tid] = (int16_t)sum; - } - if (tid == 0u) yb->d = 1.0f / iscale_s; -} - -template -__global__ static void moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][16]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - tok[np] = pair[np] / n_expert; - slot[np] = pair[np] - tok[np] * n_expert; - xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; - } - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { - uint32_t p = i / xq_blocks; - uint32_t b = i - p * xq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { - uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; - if (row >= expert_mid_dim) continue; - const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 8u) { - dev_dot_q4_K_q8_K_block8(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate); - dev_dot_q4_K_q8_K_block8(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up); - } - for (uint32_t p = 0; p < np; p++) { - gate[p] = quarter_warp_sum_f32(gate[p], lane); - up[p] = quarter_warp_sum_f32(up[p], lane); - if (lane == 0) { - if (clamp > 1.0e-6f) { - if (gate[p] > clamp) gate[p] = clamp; - if (up[p] > clamp) up[p] = clamp; - if (up[p] < -clamp) up[p] = -clamp; - } - const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate[p]; - up_out[off] = up[p]; - } - mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; - } - } - } -} - -__global__ static void moe_down_sum6_qwarp32_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - if (row >= out_dim) return; - float total = 0.0f; - #pragma unroll - for (uint32_t slot = 0; slot < 6u; slot++) { - int32_t expert_i = selected[slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) total += acc; - } - if (lane == 0) out[row] = total; -} - -__global__ static void moe_down_owned_slots_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t expert_base, - uint32_t expert_count) { - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - const uint32_t slot = blockIdx.y; - if (row >= out_dim || slot >= 6u) return; - uint32_t expert = 0; - if (!moe_owned_local_expert(selected[slot], expert_base, - expert_count, &expert)) { - return; - } - const cuda_block_q2_K *wr = - (const cuda_block_q2_K *)(down_base + - (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc; -} - -/* Map one of two packed operands for a three-slot reduction group. The only - * multi-slot operand is the peer-owned prefix (slots 0+1 within the group), - * which can be pre-added exactly because the reference reduction starts from - * +0. Every other peer slot remains a distinct operand in original order. */ -__device__ __forceinline__ static int moe_owned_packed_component( - const int32_t *selected, - uint32_t group, - uint32_t component, - uint32_t expert_base, - uint32_t expert_count, - bool *prefix_pair) { - const uint32_t slot0 = group * 3u; - uint32_t mask = 0u; - #pragma unroll - for (uint32_t i = 0; i < 3u; i++) { - if (moe_owned_local_expert(selected[slot0 + i], expert_base, - expert_count, NULL)) { - mask |= 1u << i; - } - } - *prefix_pair = false; - if ((mask & 3u) == 3u) { - if (component == 0u) { - *prefix_pair = true; - return (int)slot0; - } - return (mask & 4u) != 0u ? (int)(slot0 + 2u) : -1; - } - uint32_t ordinal = 0u; - #pragma unroll - for (uint32_t i = 0; i < 3u; i++) { - if ((mask & (1u << i)) == 0u) continue; - if (ordinal++ == component) return (int)(slot0 + i); - } - return -1; -} - -__global__ static void moe_down_owned_packed_qwarp32_kernel( - float *packed_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t expert_base, - uint32_t expert_count) { - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - const uint32_t packed_slot = blockIdx.y; - if (row >= out_dim || packed_slot >= 4u) return; - bool prefix_pair = false; - const int first_slot = moe_owned_packed_component( - selected, packed_slot / 2u, packed_slot & 1u, - expert_base, expert_count, &prefix_pair); - if (first_slot < 0) { - if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f; - return; - } - - float packed = 0.0f; - const uint32_t n_slots = prefix_pair ? 2u : 1u; - #pragma unroll - for (uint32_t i = 0; i < 2u; i++) { - if (i >= n_slots) break; - const uint32_t slot = (uint32_t)first_slot + i; - uint32_t expert = 0; - if (!moe_owned_local_expert(selected[slot], expert_base, - expert_count, &expert)) { - continue; - } - const cuda_block_q2_K *wr = - (const cuda_block_q2_K *)(down_base + - (uint64_t)expert * down_expert_bytes + - (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0u) { - packed = prefix_pair ? __fadd_rn(packed, acc) : acc; - } - } - if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed; -} - -__global__ static void moe_down_sum3_qwarp32_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - if (row >= out_dim) return; - float total = 0.0f; - #pragma unroll - for (uint32_t slot = 0; slot < 3u; slot++) { - int32_t expert_i = selected[slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) total += acc; - } - if (lane == 0) out[row] = total; -} - -__global__ static void moe_down_q4K_sum6_qwarp32_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - if (row >= out_dim) return; - const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u); - float total = 0.0f; - #pragma unroll - for (uint32_t slot = 0; slot < 6u; slot++) { - int32_t expert_i = selected[slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - if (vec_ok) { - for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); - } else { - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) total += acc; - } - if (lane == 0) out[row] = total; -} - -__global__ static void moe_down_q4K_owned_slots_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t expert_base, - uint32_t expert_count) { - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - const uint32_t slot = blockIdx.y; - if (row >= out_dim || slot >= 6u) return; - uint32_t expert = 0; - if (!moe_owned_local_expert(selected[slot], expert_base, - expert_count, &expert)) { - return; - } - const cuda_block_q4_K *wr = - (const cuda_block_q4_K *)(down_base + - (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | - down_expert_bytes) & 15u) == 0u); - float acc = 0.0f; - if (vec_ok) { - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); - } - } else { - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - } - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc; -} - -__global__ static void moe_down_q4K_owned_packed_qwarp32_kernel( - float *packed_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t expert_base, - uint32_t expert_count) { - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - const uint32_t packed_slot = blockIdx.y; - if (row >= out_dim || packed_slot >= 4u) return; - bool prefix_pair = false; - const int first_slot = moe_owned_packed_component( - selected, packed_slot / 2u, packed_slot & 1u, - expert_base, expert_count, &prefix_pair); - if (first_slot < 0) { - if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f; - return; - } - - const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | - down_expert_bytes) & 15u) == 0u); - float packed = 0.0f; - const uint32_t n_slots = prefix_pair ? 2u : 1u; - #pragma unroll - for (uint32_t i = 0; i < 2u; i++) { - if (i >= n_slots) break; - const uint32_t slot = (uint32_t)first_slot + i; - uint32_t expert = 0; - if (!moe_owned_local_expert(selected[slot], expert_base, - expert_count, &expert)) { - continue; - } - const cuda_block_q4_K *wr = - (const cuda_block_q4_K *)(down_base + - (uint64_t)expert * down_expert_bytes + - (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - if (vec_ok) { - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); - } - } else { - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - } - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0u) { - packed = prefix_pair ? __fadd_rn(packed, acc) : acc; - } - } - if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed; -} - -__global__ static void moe_owned_slots_combine_fixed3_kernel( - float *out, - const float *home_slots, - const float *peer_slots, - const int32_t *selected, - uint32_t out_dim, - uint32_t expert_split) { - const uint32_t col = - (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); - const uint32_t row = blockIdx.y; - if (col >= out_dim) return; - out += (uint64_t)row * out_dim; - home_slots += (uint64_t)row * 6u * out_dim; - peer_slots += (uint64_t)row * 6u * out_dim; - selected += (uint64_t)row * 6u; - float slotv[6]; - #pragma unroll - for (uint32_t slot = 0; slot < 6u; slot++) { - const int32_t expert = selected[slot]; - if (expert < 0 || - (uint32_t)expert >= 2u * expert_split) { - slotv[slot] = 0.0f; - } else { - const bool on_home = (uint32_t)expert < expert_split; - const float *src = on_home ? home_slots : peer_slots; - slotv[slot] = src[(uint64_t)slot * out_dim + col]; - } - } - float home = __fadd_rn(0.0f, slotv[0]); - home = __fadd_rn(home, slotv[1]); - home = __fadd_rn(home, slotv[2]); - float peer = __fadd_rn(0.0f, slotv[3]); - peer = __fadd_rn(peer, slotv[4]); - peer = __fadd_rn(peer, slotv[5]); - out[col] = __fadd_rn(home, peer); -} - -__device__ static float moe_owned_packed_combine_row( - const float *home_slots, - const float *peer_packed, - const int32_t *selected, - uint32_t row, - uint32_t out_dim, - uint32_t expert_split) { - float groups[2]; - #pragma unroll - for (uint32_t group = 0; group < 2u; group++) { - const uint32_t slot0 = group * 3u; - uint32_t peer_mask = 0u; - uint32_t valid_mask = 0u; - #pragma unroll - for (uint32_t i = 0; i < 3u; i++) { - const int32_t expert = selected[slot0 + i]; - if (expert >= 0 && (uint32_t)expert < 2u * expert_split) { - valid_mask |= 1u << i; - } - if (expert >= 0 && (uint32_t)expert >= expert_split && - (uint32_t)expert < 2u * expert_split) { - peer_mask |= 1u << i; - } - } - const float *packed = peer_packed + - (uint64_t)group * 2u * out_dim + row; - float acc; - if ((peer_mask & 3u) == 3u) { - /* packed[0] is already (+0 + slot0) + slot1. */ - acc = packed[0]; - float slot2 = 0.0f; - if ((peer_mask & 4u) != 0u) { - slot2 = packed[out_dim]; - } else if ((valid_mask & 4u) != 0u) { - slot2 = home_slots[(uint64_t)(slot0 + 2u) * out_dim + row]; - } - acc = __fadd_rn(acc, slot2); - } else { - acc = 0.0f; - uint32_t peer_operand = 0u; - #pragma unroll - for (uint32_t i = 0; i < 3u; i++) { - float value; - if ((peer_mask & (1u << i)) != 0u) { - value = packed[(uint64_t)peer_operand * out_dim]; - peer_operand++; - } else if ((valid_mask & (1u << i)) != 0u) { - value = home_slots[(uint64_t)(slot0 + i) * out_dim + row]; - } else { - value = 0.0f; - } - acc = __fadd_rn(acc, value); - } - } - groups[group] = acc; - } - return __fadd_rn(groups[0], groups[1]); -} - -__global__ static void moe_owned_packed_combine_fixed3_kernel( - float *out, - const float *home_slots, - const float *peer_packed, - const int32_t *selected, - uint32_t out_dim, - uint32_t expert_split) { - const uint32_t row = - (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); - if (row >= out_dim) return; - out[row] = moe_owned_packed_combine_row( - home_slots, peer_packed, selected, row, out_dim, expert_split); -} - -__global__ static void moe_down_q4K_sum3_qwarp32_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - if (row >= out_dim) return; - const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u); - float total = 0.0f; - #pragma unroll - for (uint32_t slot = 0; slot < 3u; slot++) { - int32_t expert_i = selected[slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - float acc = 0.0f; - if (vec_ok) { - for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); - } else { - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - } - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) total += acc; - } - if (lane == 0) out[row] = total; -} - -__global__ static void moe_down_q4K_sum3_slotwarp_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim) { - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t slot = lane >> 3u; - const uint32_t qlane = lane & 7u; - const uint32_t row = blockIdx.x * 8u + warp; - if (row >= out_dim) return; - - float acc = 0.0f; - if (slot < 3u) { - int32_t expert_i = selected[slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q4_K *wr = - (const cuda_block_q4_K *)(down_base + - (uint64_t)(uint32_t)expert_i * down_expert_bytes + - (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; - for (uint32_t b = qlane; b < midq_blocks; b += 8u) { - acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - } - acc = quarter_warp_sum_f32(acc, qlane); - } - - const float s1 = __shfl_sync(0xffffffffu, acc, 8); - const float s2 = __shfl_sync(0xffffffffu, acc, 16); - if (lane == 0u) { - const float s0 = acc; - out[row] = (s0 + s1) + s2; - } -} - -static void routed_moe_decode_graph_destroy_one(int logical_tier) { - if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return; - cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier]; - if (c->exec) (void)cudaGraphExecDestroy(c->exec); - if (c->graph) (void)cudaGraphDestroy(c->graph); - memset(c, 0, sizeof(*c)); -} - -static int routed_moe_decode_q4_graph_launch( - int logical_tier, - float *out, - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_w, - const char *up_w, - const char *down_w, - cuda_block_q8_K *xq, - cuda_block_q8_K *midq, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp, - const float *x) { - if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0; - if (n_expert != 3u && n_expert != 6u) return 0; - uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; - uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; - if (xq_blocks == 0u || midq_blocks == 0u) return 0; - - cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier]; - const bool shape_match = - c->valid && - c->n_expert == n_expert && - c->expert_in_dim == expert_in_dim && - c->expert_mid_dim == expert_mid_dim && - c->out_dim == out_dim; - if (c->valid && !shape_match) { - routed_moe_decode_graph_destroy_one(logical_tier); - c = &g_moe_decode_graph[logical_tier]; - } - - uint32_t x_rows = 1u; - uint32_t mid_rows = n_expert; - dim3 xq_grid(xq_blocks, 1, 1); - dim3 gate_grid((expert_mid_dim + 7u) / 8u, n_expert, 1); - dim3 midq_grid(midq_blocks, n_expert, 1); - dim3 down_grid((out_dim + 31u) / 32u, 1, 1); - dim3 block(256, 1, 1); - - void *xq_args[] = { &xq, &x, &expert_in_dim, &x_rows }; - cudaKernelNodeParams xq_params; - memset(&xq_params, 0, sizeof(xq_params)); - xq_params.func = (void *)q8_K_quantize_kernel; - xq_params.gridDim = xq_grid; - xq_params.blockDim = block; - xq_params.kernelParams = xq_args; - - void *gate_args[] = { - &gate_out, &up_out, &mid_out, &gate_w, &up_w, &xq, &selected, - &weights, &gate_expert_bytes, &gate_row_bytes, &xq_blocks, - &expert_mid_dim, &n_expert, &write_aux, &clamp - }; - cudaKernelNodeParams gate_params; - memset(&gate_params, 0, sizeof(gate_params)); - gate_params.func = (void *)moe_gate_up_mid_decode_q4K_warp32_kernel; - gate_params.gridDim = gate_grid; - gate_params.blockDim = block; - gate_params.kernelParams = gate_args; - - void *midq_args[] = { &midq, &mid_out, &expert_mid_dim, &mid_rows }; - cudaKernelNodeParams midq_params; - memset(&midq_params, 0, sizeof(midq_params)); - midq_params.func = (void *)q8_K_quantize_kernel; - midq_params.gridDim = midq_grid; - midq_params.blockDim = block; - midq_params.kernelParams = midq_args; - - void *down_args[] = { - &out, &down_w, &midq, &selected, &down_expert_bytes, - &down_row_bytes, &midq_blocks, &out_dim - }; - cudaKernelNodeParams down_params; - memset(&down_params, 0, sizeof(down_params)); - down_params.func = n_expert == 6u - ? (void *)moe_down_q4K_sum6_qwarp32_kernel - : (void *)moe_down_q4K_sum3_qwarp32_kernel; - down_params.gridDim = down_grid; - down_params.blockDim = block; - down_params.kernelParams = down_args; - - if (!c->valid) { - cudaError_t err = cudaGraphCreate(&c->graph, 0); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: routed MoE decode graph create failed: %s\n", - cudaGetErrorString(err)); - routed_moe_decode_graph_destroy_one(logical_tier); - return -1; - } - err = cudaGraphAddKernelNode(&c->xq_node, c->graph, NULL, 0, - &xq_params); - if (err == cudaSuccess) { - err = cudaGraphAddKernelNode(&c->gate_node, c->graph, - &c->xq_node, 1, &gate_params); - } - if (err == cudaSuccess) { - err = cudaGraphAddKernelNode(&c->midq_node, c->graph, - &c->gate_node, 1, &midq_params); - } - if (err == cudaSuccess) { - err = cudaGraphAddKernelNode(&c->down_node, c->graph, - &c->midq_node, 1, &down_params); - } - if (err == cudaSuccess) { - err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0); - } - if (err != cudaSuccess) { - fprintf(stderr, "ds4: routed MoE decode graph instantiate failed: %s\n", - cudaGetErrorString(err)); - routed_moe_decode_graph_destroy_one(logical_tier); - return -1; - } - c->n_expert = n_expert; - c->expert_in_dim = expert_in_dim; - c->expert_mid_dim = expert_mid_dim; - c->out_dim = out_dim; - c->valid = 1; - } else { - cudaError_t err = - cudaGraphExecKernelNodeSetParams(c->exec, c->xq_node, - &xq_params); - if (err == cudaSuccess) { - err = cudaGraphExecKernelNodeSetParams(c->exec, c->gate_node, - &gate_params); - } - if (err == cudaSuccess) { - err = cudaGraphExecKernelNodeSetParams(c->exec, c->midq_node, - &midq_params); - } - if (err == cudaSuccess) { - err = cudaGraphExecKernelNodeSetParams(c->exec, c->down_node, - &down_params); - } - if (err != cudaSuccess) { - fprintf(stderr, "ds4: routed MoE decode graph update failed: %s\n", - cudaGetErrorString(err)); - routed_moe_decode_graph_destroy_one(logical_tier); - return -1; - } - } - - cudaError_t err = cudaGraphLaunch(c->exec, 0); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: routed MoE decode graph launch failed: %s\n", - cudaGetErrorString(err)); - routed_moe_decode_graph_destroy_one(logical_tier); - return -1; - } - return 1; -} - -/* Q4_K prefill (n_tokens > 1) down kernel. Mirrors moe_down_qwarp32_kernel - * geometry exactly; only the weight block type and dot helper differ. The - * pair = blockIdx.y indexing means the same grid shape (out_dim/32, n_tokens*n_expert) - * used by the IQ2 path applies here. The downstream moe_sum_kernel is - * weight-type-agnostic and sums these per-pair outputs into the final output. */ -__global__ static void moe_down_q4K_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -template -__global__ static void moe_down_q4K_expert_tile8_rowspan_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][8]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { - uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; - if (row >= out_dim) continue; - const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q4_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc); - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } -} - - -/* INT8 tensor-core (m8n8k16) exact MoE prefill tile kernels. - * - * Each warp computes an 8-token x 8-row tile. The Q4_K x Q8_K superblock dot - * keeps its integer sums (order-invariant, exact) but computes the 32-wide - * group dots on tensor cores; every output element keeps 8 float slot - * accumulators (slot[b & 7] += term_b, b ascending) and reduces them with the - * exact quarter_warp_sum_f32 grouping, so results are bit-identical to the - * scalar expert-tile kernels (fuzz-verified). Requires sm_75+, 16B-aligned - * expert tensors, and the staged activation-block counts (<=16 gate/up, - * <=8 down). Rollback: DS4_CUDA_MOE_NO_Q4_MMA=1. */ -__device__ __forceinline__ static void mma_m8n8k16_s8(int32_t &c0, int32_t &c1, uint32_t a, uint32_t b) { -#if __CUDA_ARCH__ >= 750 - asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%0,%1};" - : "+r"(c0), "+r"(c1) : "r"(a), "r"(b)); -#else - (void)a; (void)b; (void)c0; (void)c1; -#endif -} - -template -__global__ static void moe_gate_up_mid_q4K_tile8_mma_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][16]; - __shared__ uint32_t s_pair[8]; - __shared__ uint32_t s_tok[8]; - __shared__ uint32_t s_slot[8]; - __shared__ uint32_t s_np; - if (threadIdx.x == 0) { - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - uint32_t pr = sorted_pairs[offsets[expert] + local_pair]; - s_pair[np] = pr; - s_tok[np] = pr / n_expert; - s_slot[np] = pr - s_tok[np] * n_expert; - } - s_np = np; - } - __syncthreads(); - const uint32_t np = s_np; - if (xq_blocks <= 16u) { - for (uint32_t i = threadIdx.x; i < np * xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); i += blockDim.x) { - const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - uint32_t p = i / words_per_tok; - uint32_t w = i - p * words_per_tok; - ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w]; - } - __syncthreads(); - } - const uint32_t mtok = lane >> 2u; /* token row of this thread's C elems */ - const uint32_t n0 = (lane & 3u) * 2u; /* first C column (weight row) */ - /* 8 warps x 8 rows = 64 rows per pass */ - for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { - const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; - if (row0 >= expert_mid_dim) continue; - const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; - const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; - /* per-element slot accumulators (2 elements x 8 slots) */ - float sg0[8] = {0,0,0,0,0,0,0,0}, sg1[8] = {0,0,0,0,0,0,0,0}; - float su0[8] = {0,0,0,0,0,0,0,0}, su1[8] = {0,0,0,0,0,0,0,0}; - for (uint32_t b = 0; b < xq_blocks; b++) { - /* headers for this thread's two C columns */ - const uint4 ghdr0 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)n0 * gate_row_bytes) + b); - const uint4 ghdr1 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b); - const uint4 uhdr0 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)n0 * gate_row_bytes) + b); - const uint4 uhdr1 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b); - /* B-fragment source rows for loads: n_load = lane>>2. - * Batch all global loads for this superblock upfront so the - * memory system sees independent requests instead of a - * load->mma dependency chain. */ - const uint32_t *gqw = (const uint32_t *)(((const cuda_block_q4_K *)(grow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs); - const uint32_t *uqw = (const uint32_t *)(((const cuda_block_q4_K *)(urow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs); - const int8_t *aqs = sxq[mtok][b].qs; - uint32_t gw8[8], uw8[8]; -#pragma unroll - for (uint32_t k = 0; k < 8u; k++) { - gw8[k] = gqw[k * 4u + (lane & 3u)]; - uw8[k] = uqw[k * 4u + (lane & 3u)]; - } - int gi0 = 0, gi1 = 0, ui0 = 0, ui1 = 0; - int gs0 = 0, gs1 = 0, us0 = 0, us1 = 0; -#pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - const int shift = (j & 1u) ? 4 : 0; - /* dot32 via two chained k16 mmas, per matrix */ - int32_t gc0 = 0, gc1 = 0, uc0 = 0, uc1 = 0; -#pragma unroll - for (uint32_t h = 0; h < 2u; h++) { - const uint32_t koff = h * 16u + (lane & 3u) * 4u; - const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff); - const uint32_t gw = (gw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; - const uint32_t uw = (uw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; - mma_m8n8k16_s8(gc0, gc1, a, gw); - mma_m8n8k16_s8(uc0, uc1, a, uw); - } - /* integer scale application for this thread's two columns */ - uint8_t sc, m; - dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr0.y, &sc, &m); - gi0 += (int)sc * gc0; - const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u]; - gs0 += (int)m * bs; - dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr1.y, &sc, &m); - gi1 += (int)sc * gc1; - gs1 += (int)m * bs; - dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr0.y, &sc, &m); - ui0 += (int)sc * uc0; - us0 += (int)m * bs; - dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr1.y, &sc, &m); - ui1 += (int)sc * uc1; - us1 += (int)m * bs; - } - /* float finish, exact dev_dot_q4_K_q8_K_block8 expression */ - const float yd = sxq[mtok][b].d; - const uint32_t sl = b & 7u; - sg0[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr0.x & 0xffffu)) * (float)gi0 - - yd * dev_f16_to_f32((uint16_t)(ghdr0.x >> 16u)) * (float)gs0; - sg1[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr1.x & 0xffffu)) * (float)gi1 - - yd * dev_f16_to_f32((uint16_t)(ghdr1.x >> 16u)) * (float)gs1; - su0[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr0.x & 0xffffu)) * (float)ui0 - - yd * dev_f16_to_f32((uint16_t)(uhdr0.x >> 16u)) * (float)us0; - su1[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr1.x & 0xffffu)) * (float)ui1 - - yd * dev_f16_to_f32((uint16_t)(uhdr1.x >> 16u)) * (float)us1; - } - /* quarter_warp_sum_f32 order: ((s0+s4)+(s2+s6)) + ((s1+s5)+(s3+s7)) */ - const uint32_t p = mtok; - if (p < np) { - const uint32_t rowa = row0 + n0; - const uint32_t rowb = row0 + n0 + 1u; - float gate2[2], up2[2]; - { - float a0 = sg0[0] + sg0[4], a1 = sg0[1] + sg0[5], a2 = sg0[2] + sg0[6], a3 = sg0[3] + sg0[7]; - gate2[0] = (a0 + a2) + (a1 + a3); - a0 = sg1[0] + sg1[4]; a1 = sg1[1] + sg1[5]; a2 = sg1[2] + sg1[6]; a3 = sg1[3] + sg1[7]; - gate2[1] = (a0 + a2) + (a1 + a3); - a0 = su0[0] + su0[4]; a1 = su0[1] + su0[5]; a2 = su0[2] + su0[6]; a3 = su0[3] + su0[7]; - up2[0] = (a0 + a2) + (a1 + a3); - a0 = su1[0] + su1[4]; a1 = su1[1] + su1[5]; a2 = su1[2] + su1[6]; a3 = su1[3] + su1[7]; - up2[1] = (a0 + a2) + (a1 + a3); - } -#pragma unroll - for (uint32_t e = 0; e < 2u; e++) { - const uint32_t row = e ? rowb : rowa; - if (row >= expert_mid_dim) continue; - float gate = gate2[e]; - float up = up2[e]; - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]]; - } - } - } -} - -template -__global__ static void moe_down_q4K_tile8_mma_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][8]; - __shared__ uint32_t s_pair[8]; - __shared__ uint32_t s_np; - if (threadIdx.x == 0) { - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - s_pair[np] = sorted_pairs[offsets[expert] + local_pair]; - } - s_np = np; - } - __syncthreads(); - const uint32_t np = s_np; - if (midq_blocks <= 8u) { - const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { - uint32_t p = i / words_per_tok; - uint32_t w = i - p * words_per_tok; - ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w]; - } - __syncthreads(); - } - const uint32_t mtok = lane >> 2u; - const uint32_t n0 = (lane & 3u) * 2u; - for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { - const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; - if (row0 >= out_dim) continue; - const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes; - float s0[8] = {0,0,0,0,0,0,0,0}, s1[8] = {0,0,0,0,0,0,0,0}; - for (uint32_t b = 0; b < midq_blocks; b++) { - const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b); - const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b); - const uint32_t *wqw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs); - const int8_t *aqs = sxq[mtok][b].qs; - uint32_t w8[8]; -#pragma unroll - for (uint32_t k = 0; k < 8u; k++) w8[k] = wqw[k * 4u + (lane & 3u)]; - int i0 = 0, i1 = 0, m0 = 0, m1 = 0; -#pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - const int shift = (j & 1u) ? 4 : 0; - int32_t c0 = 0, c1 = 0; -#pragma unroll - for (uint32_t h = 0; h < 2u; h++) { - const uint32_t koff = h * 16u + (lane & 3u) * 4u; - const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff); - const uint32_t w = (w8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; - mma_m8n8k16_s8(c0, c1, a, w); - } - uint8_t sc, m; - const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u]; - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc, &m); - i0 += (int)sc * c0; - m0 += (int)m * bs; - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc, &m); - i1 += (int)sc * c1; - m1 += (int)m * bs; - } - const float yd = sxq[mtok][b].d; - const uint32_t sl = b & 7u; - s0[sl] += yd * dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)) * (float)i0 - - yd * dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)) * (float)m0; - s1[sl] += yd * dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)) * (float)i1 - - yd * dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)) * (float)m1; - } - const uint32_t p = mtok; - if (p < np) { - float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; - const float r0 = (a0 + a2) + (a1 + a3); - a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; - const float r1 = (a0 + a2) + (a1 + a3); - if (row0 + n0 < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0] = r0; - if (row0 + n0 + 1u < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0 + 1u] = r1; - } - } -} - -/* 16-pair MoE expert tile kernels on sm_80+ m16n8k32 INT8 tensor cores. - * - * Same per-output math and reduction order as the 8-pair expert tile - * kernels (slot[b & 7] += term_b with b ascending, then the exact - * quarter_warp_sum_f32 grouping), so results are bit-identical; grouping - * 16 pairs per tile just halves how often each expert's weights are - * streamed from DRAM. Gate and up run as two passes over the superblocks - * to keep register pressure at the 8-pair kernel's level. */ - -__device__ __forceinline__ static void mma16_m16n8k32_s8( - int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3, - uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, - uint32_t b0, uint32_t b1) { -#if __CUDA_ARCH__ >= 800 - asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" - : "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3) - : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1)); -#else - (void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3; -#endif -} - -/* One matrix pass over all superblocks for this thread's 4 C elements - * (tokens mtokA/mtokB x rows n0/n0+1). Returns the quarter-tree-reduced - * values in r[4] with the exact reference ordering. */ -__device__ __forceinline__ static void moe_tile16_mma_pass( - const char *wrow, /* row0 base of this matrix */ - uint64_t row_bytes, - const cuda_block_q8_K (*sxq)[16], - uint32_t xq_blocks, - uint32_t lane, - float r[4]) { - const uint32_t mtokA = lane >> 2u; - const uint32_t mtokB = mtokA + 8u; - const uint32_t n0 = (lane & 3u) * 2u; - float s0[8] = {0,0,0,0,0,0,0,0}; - float s1[8] = {0,0,0,0,0,0,0,0}; - float s2[8] = {0,0,0,0,0,0,0,0}; - float s3[8] = {0,0,0,0,0,0,0,0}; - for (uint32_t b = 0; b < xq_blocks; b++) { - const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * row_bytes) + b); - const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * row_bytes) + b); - const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * row_bytes) + b)->qs); - uint32_t w8[8]; -#pragma unroll - for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)]; - const int8_t *aqsA = sxq[mtokA][b].qs; - const int8_t *aqsB = sxq[mtokB][b].qs; - int i0 = 0, i1 = 0, i2 = 0, i3 = 0; - int m0 = 0, m1 = 0, m2 = 0, m3 = 0; -#pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - const int shift = (j & 1u) ? 4 : 0; - const uint32_t koff = (lane & 3u) * 4u; - const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff); - const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff); - const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff); - const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff); - const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu; - const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu; - int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; - mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); - uint8_t sc0, sm0, sc1, sm1; - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0); - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1); - const int bsA = (int)sxq[mtokA][b].bsums[2u * j] + (int)sxq[mtokA][b].bsums[2u * j + 1u]; - const int bsB = (int)sxq[mtokB][b].bsums[2u * j] + (int)sxq[mtokB][b].bsums[2u * j + 1u]; - i0 += (int)sc0 * c0; - i1 += (int)sc1 * c1; - i2 += (int)sc0 * c2; - i3 += (int)sc1 * c3; - m0 += (int)sm0 * bsA; - m1 += (int)sm1 * bsA; - m2 += (int)sm0 * bsB; - m3 += (int)sm1 * bsB; - } - const float ydA = sxq[mtokA][b].d; - const float ydB = sxq[mtokB][b].d; - const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)); - const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)); - const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)); - const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)); - const uint32_t sl = b & 7u; - s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0; - s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1; - s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2; - s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3; - } - { - float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; - r[0] = (a0 + a2) + (a1 + a3); - a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; - r[1] = (a0 + a2) + (a1 + a3); - a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7]; - r[2] = (a0 + a2) + (a1 + a3); - a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7]; - r[3] = (a0 + a2) + (a1 + a3); - } -} - -template -__global__ static void moe_gate_up_mid_q4K_tile16_mma_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t write_aux, - float clamp) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - extern __shared__ unsigned char t16_sh[]; - cuda_block_q8_K (*sxq)[16] = (cuda_block_q8_K (*)[16])t16_sh; /* [16][16] */ - __shared__ uint32_t s_pair[16]; - __shared__ uint32_t s_tok[16]; - __shared__ uint32_t s_slot[16]; - __shared__ uint32_t s_np; - if (threadIdx.x == 0) { - uint32_t np = 0; - for (; np < 16u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - uint32_t pr = sorted_pairs[offsets[expert] + local_pair]; - s_pair[np] = pr; - s_tok[np] = pr / n_expert; - s_slot[np] = pr - s_tok[np] * n_expert; - } - s_np = np; - } - __syncthreads(); - const uint32_t np = s_np; - if (xq_blocks <= 16u) { - const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { - uint32_t p = i / words_per_tok; - uint32_t w = i - p * words_per_tok; - ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w]; - } - /* zero-fill missing pairs so the A fragments are defined */ - const uint32_t total_words = 16u * words_per_tok; - for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) { - ((uint32_t *)t16_sh)[i] = 0u; - } - __syncthreads(); - } - const uint32_t mtokA = lane >> 2u; - const uint32_t mtokB = mtokA + 8u; - const uint32_t n0 = (lane & 3u) * 2u; - for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { - const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; - if (row0 >= expert_mid_dim) continue; - const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; - const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; - float gr[4], ur[4]; - moe_tile16_mma_pass(grow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, gr); - moe_tile16_mma_pass(urow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, ur); -#pragma unroll - for (uint32_t e = 0; e < 4u; e++) { - const uint32_t p = (e < 2u) ? mtokA : mtokB; - const uint32_t row = row0 + n0 + (e & 1u); - if (p >= np || row >= expert_mid_dim) continue; - float gate = gr[e]; - float up = ur[e]; - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row; - if (write_aux) { - gate_out[off] = gate; - up_out[off] = up; - } - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * - weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]]; - } - } -} - -template -__global__ static void moe_down_q4K_tile16_mma_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t warp = threadIdx.x >> 5u; - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - extern __shared__ unsigned char t16_sh[]; - __shared__ uint32_t s_pair[16]; - __shared__ uint32_t s_np; - if (threadIdx.x == 0) { - uint32_t np = 0; - for (; np < 16u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - s_pair[np] = sorted_pairs[offsets[expert] + local_pair]; - } - s_np = np; - } - __syncthreads(); - const uint32_t np = s_np; - if (midq_blocks <= 16u) { - const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); - for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { - uint32_t p = i / words_per_tok; - uint32_t w = i - p * words_per_tok; - ((uint32_t *)t16_sh)[i] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w]; - } - const uint32_t total_words = 16u * words_per_tok; - for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) { - ((uint32_t *)t16_sh)[i] = 0u; - } - __syncthreads(); - } - const uint32_t mtokA = lane >> 2u; - const uint32_t mtokB = mtokA + 8u; - const uint32_t n0 = (lane & 3u) * 2u; - for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { - const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; - if (row0 >= out_dim) continue; - const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes; - float s0[8] = {0,0,0,0,0,0,0,0}; - float s1[8] = {0,0,0,0,0,0,0,0}; - float s2[8] = {0,0,0,0,0,0,0,0}; - float s3[8] = {0,0,0,0,0,0,0,0}; - for (uint32_t b = 0; b < midq_blocks; b++) { - const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b); - const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b); - const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs); - uint32_t w8[8]; -#pragma unroll - for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)]; - /* activation rows: midq_blocks stride within the staged region */ - const int8_t *aqsA = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b)->qs; - const int8_t *aqsB = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b)->qs; - const cuda_block_q8_K *blkA = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b; - const cuda_block_q8_K *blkB = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b; - int i0 = 0, i1 = 0, i2 = 0, i3 = 0; - int m0 = 0, m1 = 0, m2 = 0, m3 = 0; -#pragma unroll - for (uint32_t j = 0; j < 8u; j++) { - const int shift = (j & 1u) ? 4 : 0; - const uint32_t koff = (lane & 3u) * 4u; - const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff); - const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff); - const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff); - const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff); - const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu; - const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu; - int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; - mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); - uint8_t sc0, sm0, sc1, sm1; - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0); - dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1); - const int bsA = (int)blkA->bsums[2u * j] + (int)blkA->bsums[2u * j + 1u]; - const int bsB = (int)blkB->bsums[2u * j] + (int)blkB->bsums[2u * j + 1u]; - i0 += (int)sc0 * c0; - i1 += (int)sc1 * c1; - i2 += (int)sc0 * c2; - i3 += (int)sc1 * c3; - m0 += (int)sm0 * bsA; - m1 += (int)sm1 * bsA; - m2 += (int)sm0 * bsB; - m3 += (int)sm1 * bsB; - } - const float ydA = blkA->d; - const float ydB = blkB->d; - const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)); - const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)); - const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)); - const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)); - const uint32_t sl = b & 7u; - s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0; - s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1; - s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2; - s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3; - } - float rr4[4]; - { - float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; - rr4[0] = (a0 + a2) + (a1 + a3); - a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; - rr4[1] = (a0 + a2) + (a1 + a3); - a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7]; - rr4[2] = (a0 + a2) + (a1 + a3); - a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7]; - rr4[3] = (a0 + a2) + (a1 + a3); - } -#pragma unroll - for (uint32_t e = 0; e < 4u; e++) { - const uint32_t p = (e < 2u) ? mtokA : mtokB; - const uint32_t row = row0 + n0 + (e & 1u); - if (p >= np || row >= out_dim) continue; - down_out[(uint64_t)s_pair[p] * out_dim + row] = rr4[e]; - } - } -} - -static int cuda_q4_mma_tile16_shmem_ok(int which_down) { - /* Opt the tile16 kernels into >48KB dynamic shared memory, per device. */ - static int ready[DS4_MAX_GPUS][2]; - static int failed = 0; - if (failed) return 0; - int dev = 0; - cudaGetDevice(&dev); - if (dev < 0 || dev >= DS4_MAX_GPUS) return 0; - if (ready[dev][which_down]) return 1; - cudaFuncAttributes fn_attr; - cudaError_t err = which_down - ? cudaFuncGetAttributes(&fn_attr, moe_down_q4K_tile16_mma_kernel<512>) - : cudaFuncGetAttributes(&fn_attr, moe_gate_up_mid_q4K_tile16_mma_kernel<512>); - if (err != cudaSuccess || fn_attr.binaryVersion < 80) { - failed = 1; - return 0; - } - const int bytes = (int)(16u * 16u * sizeof(cuda_block_q8_K)); - if (which_down) { - err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<512>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - if (err == cudaSuccess) - err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<1024>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - if (err == cudaSuccess) - err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<2048>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - } else { - err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<512>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - if (err == cudaSuccess) - err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<1024>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - if (err == cudaSuccess) - err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<2048>, - cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); - } - if (err != cudaSuccess) { - failed = 1; - return 0; - } - ready[dev][which_down] = 1; - return 1; -} - - - - -__global__ static void moe_down_sorted_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t pair = sorted_pairs[blockIdx.y]; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -__global__ static DS4_CUDA_UNUSED void moe_down_expert_tile8_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t group = threadIdx.x >> 3u; - uint32_t lane = threadIdx.x & 7u; - uint32_t pair_slot = group & 7u; - uint32_t row_lane = group >> 3u; - uint32_t expert = tile_experts[tile]; - uint32_t local_pair = tile_starts[tile] + pair_slot; - if (local_pair >= counts[expert]) return; - uint32_t sorted_idx = offsets[expert] + local_pair; - uint32_t pair = sorted_pairs[sorted_idx]; - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - - for (uint32_t rr = 0; rr < 2u; rr++) { - uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u; - if (row >= out_dim) continue; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; - } -} - -__global__ static void moe_down_expert_tile4_row32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t atomic_out) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[4][8]; - uint32_t pair[4] = {0, 0, 0, 0}; - const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 4u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - if (row >= out_dim) return; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q2_K_q8_K_block4(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, acc); - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) { - if (atomic_out) { - uint32_t tok = pair[p] / n_expert; - atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); - } else { - down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } - } -} - -__global__ static void moe_down_expert_tile8_row32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t atomic_out) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t expert = tile_experts[tile]; - uint32_t local_start = tile_starts[tile]; - __shared__ cuda_block_q8_K sxq[8][8]; - uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - uint32_t np = 0; - for (; np < 8u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - if (row >= out_dim) return; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc); - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) { - if (atomic_out) { - uint32_t tok = pair[p] / n_expert; - atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); - } else { - down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } - } -} - -__global__ static void moe_down_expert_tile16_row32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t atomic_out) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t local_start = tile_starts[tile]; - if (local_start & 8u) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - uint32_t expert = tile_experts[tile]; - __shared__ cuda_block_q8_K sxq[16][8]; - uint32_t pair[16] = {0}; - const cuda_block_q8_K *xqb[16] = {NULL}; - uint32_t np = 0; - for (; np < 16u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - if (row >= out_dim) return; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[16] = {0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); - if (np > 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, - xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, - xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, - xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); - } - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) { - if (atomic_out) { - uint32_t tok = pair[p] / n_expert; - atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); - } else { - down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } - } -} - -__global__ static void moe_down_expert_tile16_row2048_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t atomic_out) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t local_start = tile_starts[tile]; - if (local_start & 8u) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - __shared__ cuda_block_q8_K sxq[16][8]; - uint32_t pair[16] = {0}; - const cuda_block_q8_K *xqb[16] = {NULL}; - uint32_t np = 0; - for (; np < 16u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < 64u; rr++) { - uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u; - if (row >= out_dim) continue; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[16] = {0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); - if (np > 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, - xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, - xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, - xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); - } - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) { - if (atomic_out) { - uint32_t tok = pair[p] / n_expert; - atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); - } else { - down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } - } - } -} - -template -__global__ static void moe_down_expert_tile16_rowspan_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const uint32_t *offsets, - const uint32_t *counts, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t atomic_out) { - uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - uint32_t local_start = tile_starts[tile]; - if (local_start & 8u) return; - uint32_t lane = threadIdx.x & 7u; - uint32_t row_lane = threadIdx.x >> 3u; - uint32_t expert = tile_experts[tile]; - __shared__ cuda_block_q8_K sxq[16][8]; - uint32_t pair[16] = {0}; - const cuda_block_q8_K *xqb[16] = {NULL}; - uint32_t np = 0; - for (; np < 16u; np++) { - uint32_t local_pair = local_start + np; - if (local_pair >= counts[expert]) break; - pair[np] = sorted_pairs[offsets[expert] + local_pair]; - xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; - } - if (midq_blocks <= 8u) { - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - uint32_t p = i / midq_blocks; - uint32_t b = i - p * midq_blocks; - sxq[p][b] = xqb[p][b]; - } - __syncthreads(); - for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; - } - for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { - uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; - if (row >= out_dim) continue; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); - float acc[16] = {0.0f}; - for (uint32_t b = lane; b < midq_blocks; b += 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, - xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, - xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, - xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); - if (np > 8u) { - dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, - xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, - xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, - xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); - } - } - for (uint32_t p = 0; p < np; p++) { - acc[p] = quarter_warp_sum_f32(acc[p], lane); - if (lane == 0) { - if (atomic_out) { - uint32_t tok = pair[p] / n_expert; - atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); - } else { - down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; - } - } - } - } -} - -__global__ static void moe_down_sorted_p2_qwarp32_kernel( - float *down_out, - const char *down_base, - const cuda_block_q8_K *midq, - const uint32_t *sorted_pairs, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t pair_count) { - uint32_t lane = threadIdx.x & 7u; - uint32_t pair_lane = (threadIdx.x >> 3u) & 1u; - uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); - uint32_t sorted_idx = blockIdx.y * 2u + pair_lane; - if (row >= out_dim || sorted_idx >= pair_count) return; - uint32_t pair = sorted_pairs[sorted_idx]; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; - float acc = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); - acc = quarter_warp_sum_f32(acc, lane); - if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; -} - -__global__ static void moe_sum_kernel(float *out, const float *down, uint32_t out_dim, uint32_t n_expert, uint32_t n_tokens) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * out_dim; - if (gid >= n) return; - uint32_t tok = gid / out_dim; - uint32_t row = gid - (uint64_t)tok * out_dim; - float acc = 0.0f; - for (uint32_t e = 0; e < n_expert; e++) acc += down[((uint64_t)tok * n_expert + e) * out_dim + row]; - out[gid] = acc; -} - -__global__ static void moe_sum_owned_kernel( - float *out, - const float *down, - const int32_t *selected, - uint32_t out_dim, - uint32_t n_expert, - uint32_t n_tokens) { - const uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - const uint64_t n = (uint64_t)n_tokens * out_dim; - if (gid >= n) return; - const uint32_t tok = (uint32_t)(gid / out_dim); - const uint32_t row = (uint32_t)(gid - (uint64_t)tok * out_dim); - float acc = 0.0f; - #pragma unroll - for (uint32_t slot = 0; slot < 6u; slot++) { - if (slot >= n_expert) break; - const uint64_t pair = (uint64_t)tok * n_expert + slot; - const float value = selected[pair] >= 0 - ? down[pair * out_dim + row] : 0.0f; - acc = __fadd_rn(acc, value); - } - out[gid] = acc; -} - -__device__ static float dev_iq2_xxs_dot_f32(const cuda_block_iq2_xxs *row, const float *x, uint32_t nb) { - float acc = 0.0f; - for (uint32_t b = 0; b < nb; b++) { - const cuda_block_iq2_xxs *xb = row + b; - const float d = dev_f16_to_f32(xb->d); - const uint16_t *q2 = xb->qs; - const float *xf = x + (uint64_t)b * CUDA_QK_K; - for (uint32_t ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { - const uint32_t aux_g = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); - const uint32_t aux_s = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); - q2 += 4; - const float dl = d * (0.5f + (float)(aux_s >> 28)) * 0.25f; - const uint8_t grids[4] = { - (uint8_t)(aux_g & 0xffu), - (uint8_t)((aux_g >> 8) & 0xffu), - (uint8_t)((aux_g >> 16) & 0xffu), - (uint8_t)((aux_g >> 24) & 0xffu), - }; - for (uint32_t half = 0; half < 2; half++) { - for (uint32_t g = 0; g < 2; g++) { - const uint32_t gi = half * 2 + g; - const uint64_t grid = cuda_iq2xxs_grid[grids[gi]]; - const uint8_t signs = cuda_ksigns_iq2xs[(aux_s >> (14u * half + 7u * g)) & 127u]; - for (uint32_t i = 0; i < 8; i++) { - float w = (float)((grid >> (8u * i)) & 0xffu); - if (signs & (1u << i)) w = -w; - acc += dl * w * xf[ib32 * 32u + half * 16u + g * 8u + i]; - } - } - } - } - } - return acc; -} - -__device__ static float dev_q2_K_dot_f32(const cuda_block_q2_K *row, const float *x, uint32_t nb) { - float acc = 0.0f; - for (uint32_t b = 0; b < nb; b++) { - const cuda_block_q2_K *xb = row + b; - const float d = dev_f16_to_f32(xb->d); - const float dmin = dev_f16_to_f32(xb->dmin); - for (uint32_t il = 0; il < 16; il++) { - const uint32_t chunk = il / 8u; - const uint32_t pair = il & 1u; - const uint32_t shift = ((il / 2u) & 3u) * 2u; - const uint8_t sc = xb->scales[il]; - const float dl = d * (float)(sc & 0x0fu); - const float ml = dmin * (float)(sc >> 4); - const uint8_t *q = xb->qs + 32u * chunk + 16u * pair; - const float *xf = x + (uint64_t)b * CUDA_QK_K + chunk * 128u + ((il % 8u) / 2u) * 32u + pair * 16u; - for (uint32_t i = 0; i < 16; i++) { - const float w = dl * (float)((q[i] >> shift) & 3u) - ml; - acc += w * xf[i]; - } - } - } - return acc; -} - -__global__ static void moe_gate_up_mid_f32_kernel( - float *gate_out, - float *up_out, - float *mid_out, - const char *gate_base, - const char *up_base, - const float *x, - const int32_t *selected, - const float *weights, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t n_expert, - float clamp) { - uint32_t row = blockIdx.x; - uint32_t pair = blockIdx.y; - if (row >= expert_mid_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - uint32_t expert = (uint32_t)expert_i; - const uint32_t nb = expert_in_dim / CUDA_QK_K; - const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); - const float *xr = x + (uint64_t)tok * expert_in_dim; - float gate = 0.0f; - float up = 0.0f; - for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) { - gate += dev_iq2_xxs_dot_f32(gr + b, xr + (uint64_t)b * CUDA_QK_K, 1); - up += dev_iq2_xxs_dot_f32(ur + b, xr + (uint64_t)b * CUDA_QK_K, 1); - } - __shared__ float partial_gate[256]; - __shared__ float partial_up[256]; - partial_gate[threadIdx.x] = gate; - partial_up[threadIdx.x] = up; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride]; - partial_up[threadIdx.x] += partial_up[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - gate = partial_gate[0]; - up = partial_up[0]; - if (clamp > 1.0e-6f) { - if (gate > clamp) gate = clamp; - if (up > clamp) up = clamp; - if (up < -clamp) up = -clamp; - } - const uint64_t off = (uint64_t)pair * expert_mid_dim + row; - gate_out[off] = gate; - up_out[off] = up; - mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; - } -} - -__global__ static void moe_down_f32_kernel( - float *down_out, - const char *down_base, - const float *mid, - const int32_t *selected, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_mid_dim, - uint32_t out_dim, - uint32_t n_expert) { - uint32_t row = blockIdx.x; - uint32_t pair = blockIdx.y; - if (row >= out_dim) return; - uint32_t tok = pair / n_expert; - uint32_t slot = pair - tok * n_expert; - int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; - if (expert_i < 0) expert_i = 0; - const uint32_t nb = expert_mid_dim / CUDA_QK_K; - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); - const float *xr = mid + (uint64_t)pair * expert_mid_dim; - float acc = 0.0f; - for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) acc += dev_q2_K_dot_f32(wr + b, xr + (uint64_t)b * CUDA_QK_K, 1); - __shared__ float partial[256]; - partial[threadIdx.x] = acc; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; - __syncthreads(); - } - if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0]; -} - -static int routed_moe_launch( - ds4_gpu_tensor *out, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - ds4_gpu_tensor *down, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - float clamp, - const ds4_gpu_tensor *x, - uint32_t layer_index, - uint32_t n_tokens, - int allow_streaming, - int owned_filtered) { - if (!out || !gate || !up || !mid || !down || !model_map || !selected || !weights || !x || - n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || - expert_in_dim % CUDA_QK_K != 0 || expert_mid_dim % CUDA_QK_K != 0 || - gate_offset > model_size || up_offset > model_size || down_offset > model_size || - x->bytes < (uint64_t)n_tokens * expert_in_dim * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * n_expert * sizeof(int32_t) || - weights->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) || - gate->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || - up->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || - mid->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || - down->bytes < (uint64_t)n_tokens * n_expert * out_dim * sizeof(float) || - out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) { - return 0; - } - const int q4k_path = (gate_type == 12u && down_type == 12u); - if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; - /* Q4_K routed-MoE dispatch: - * n_tokens == 1 and n_expert == 6: - * use_direct_down_sum + moe_gate_up_mid_decode_q4K_qwarp32 - * + moe_down_q4K_sum6_qwarp32. - * n_tokens == 1 and n_expert == 3: - * use the same direct path with moe_down_q4K_sum3_qwarp32. - * Decode TP relies on this for splitting the six selected - * experts into two groups. - * n_tokens == 1 and other n_expert: - * use the same per-pair gate/up kernel plus the generic - * q4K down + sum path. - * n_tokens > 1: default sorted-pairs expert-tile path groups token/expert - * pairs by expert and uses Q4_K tile8 gate/up + down kernels - * (`DS4_CUDA_MOE_NO_Q4_SORTED=1` restores the older - * token-indexed decode-style prefill kernels). */ - const uint64_t gate_bytes = (uint64_t)n_total_expert * gate_expert_bytes; - const uint64_t down_bytes = (uint64_t)n_total_expert * down_expert_bytes; - if (gate_bytes > model_size - gate_offset || - gate_bytes > model_size - up_offset || - down_bytes > model_size - down_offset) { - return 0; - } - const uint64_t required_slot_count = (uint64_t)n_tokens * n_expert; - const int logical_tier = ds4_tensor_device_idx(out); - const int use_stream_selected_cache = - allow_streaming && - g_ssd_streaming_mode && - g_stream_selected_cache.valid && - g_stream_selected_cache.logical_tier == logical_tier && - g_stream_selected_cache.model_map == model_map && - g_stream_selected_cache.layer == layer_index && - g_stream_selected_cache.n_total_expert == n_total_expert && - g_stream_selected_cache.slot_count >= required_slot_count && - g_stream_selected_cache.gate_offset == gate_offset && - g_stream_selected_cache.up_offset == up_offset && - g_stream_selected_cache.down_offset == down_offset && - g_stream_selected_cache.gate_expert_bytes == gate_expert_bytes && - g_stream_selected_cache.down_expert_bytes == down_expert_bytes && - g_stream_selected_cache.gate_ptr && - g_stream_selected_cache.up_ptr && - g_stream_selected_cache.down_ptr && - g_stream_selected_cache.slot_selected_tensor.ptr && - g_stream_selected_cache.slot_selected_tensor.bytes >= - required_slot_count * sizeof(int32_t); - if (g_ssd_streaming_mode && allow_streaming && - !use_stream_selected_cache) { - fprintf(stderr, - "ds4: CUDA streaming selected experts are unavailable for layer %u\n", - layer_index); - return 0; - } - if (use_stream_selected_cache) { - selected = &g_stream_selected_cache.slot_selected_tensor; - } - const char *gate_w = use_stream_selected_cache ? - g_stream_selected_cache.gate_ptr : - cuda_resolve_weight_ptr(model_map, gate_offset, gate_bytes, - logical_tier, "moe_gate"); - const char *up_w = use_stream_selected_cache ? - g_stream_selected_cache.up_ptr : - cuda_resolve_weight_ptr(model_map, up_offset, gate_bytes, - logical_tier, "moe_up"); - const char *down_w = use_stream_selected_cache ? - g_stream_selected_cache.down_ptr : - cuda_resolve_weight_ptr(model_map, down_offset, down_bytes, - logical_tier, "moe_down"); - if (!gate_w || !up_w || !down_w) return 0; - - int ok = 1; - const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; - const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; - const uint64_t xq_count = (uint64_t)n_tokens * xq_blocks; - const uint64_t midq_count = (uint64_t)n_tokens * n_expert * midq_blocks; - const uint64_t xq_bytes = xq_count * sizeof(cuda_block_q8_K); - const uint64_t midq_bytes = midq_count * sizeof(cuda_block_q8_K); - if (down->bytes >= xq_bytes && gate->bytes >= midq_bytes) { - cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; - cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; - const uint32_t profile_moe = getenv("DS4_CUDA_MOE_PROFILE") != NULL; - cudaEvent_t prof_ev[7] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL}; - if (profile_moe) { - for (uint32_t i = 0; i < 7u; i++) { - if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) { - for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]); - memset(prof_ev, 0, sizeof(prof_ev)); - break; - } - } - if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0); - } - const uint32_t pair_count = n_tokens * n_expert; - const uint32_t use_q4_sorted_pairs = - q4k_path && n_tokens > 1u && - (owned_filtered || - (getenv("DS4_CUDA_MOE_NO_Q4_SORTED") == NULL && - getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL && - getenv("DS4_CUDA_MOE_TILE4") == NULL)); - const uint32_t use_sorted_pairs = - n_tokens > 1u && - (owned_filtered || !q4k_path || use_q4_sorted_pairs); - const uint32_t use_expert_tiles = - use_sorted_pairs && - (owned_filtered || getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL); - /* Small batches (DSpark stage chain / verify, n<=8) leave most of an - * 8-slot expert tile empty (1-2 rows per expert): tile4 halves the - * wasted dot-slots and measures ~2x faster there. Large prefill - * keeps tile8. Env overrides both ways. */ - const uint32_t expert_tile_m = - getenv("DS4_CUDA_MOE_TILE4") ? 4u : - (getenv("DS4_CUDA_MOE_TILE8") ? 8u : - (n_tokens <= 8u ? 4u : 8u)); - const uint32_t write_gate_up = getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL; - const uint32_t use_p2_sorted = - use_sorted_pairs && !owned_filtered && - getenv("DS4_CUDA_MOE_NO_P2") == NULL; - const uint32_t use_atomic_down = !q4k_path && use_expert_tiles && - (getenv("DS4_CUDA_MOE_ATOMIC_DOWN") != NULL || - (n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_ATOMIC_DOWN") == NULL)); - const uint32_t use_owned_sparse_buffers = owned_filtered && - getenv("DS4_CUDA_MOE_NO_OWNED_SPARSE_BUFFERS") == NULL; - const uint32_t use_gate_row2048 = use_expert_tiles && expert_tile_m == 8u && - (getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL || - getenv("DS4_CUDA_MOE_GATE_ROW256") != NULL || - getenv("DS4_CUDA_MOE_GATE_ROW128") != NULL || - (n_tokens >= 128u && - getenv("DS4_CUDA_MOE_NO_GATE_ROW2048") == NULL && - getenv("DS4_CUDA_MOE_NO_GATE_ROW256") == NULL && - getenv("DS4_CUDA_MOE_NO_GATE_ROW128") == NULL)); - const uint32_t use_q4_mma_tiles16 = q4k_path && use_expert_tiles && - expert_tile_m == 8u && cuda_q4_mma_ok() && - getenv("DS4_CUDA_MOE_NO_Q4_MMA_TILE16") == NULL; - const uint32_t use_down_tile16 = !q4k_path && use_atomic_down && expert_tile_m == 8u && - n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_DOWN_TILE16") == NULL; - const uint32_t use_small_sorted_prep = - owned_filtered && q4k_path && n_tokens <= 16u && pair_count <= 96u && - n_total_expert <= 128u && use_sorted_pairs && use_expert_tiles && - getenv("DS4_CUDA_MOE_NO_SMALL_SORTED_PREP") == NULL; - const uint32_t use_q4_down_rowspan = q4k_path && use_expert_tiles && expert_tile_m == 8u && - n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_Q4_DOWN_ROWSPAN") == NULL; - const uint32_t use_decode_lut_gate = - n_tokens == 1u && xq_blocks <= 16u && - getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") == NULL; - const uint32_t gate_row_span = - getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL ? 2048u : - getenv("DS4_CUDA_MOE_GATE_ROW1024") != NULL ? 1024u : 512u; - const uint32_t down_row_span = - getenv("DS4_CUDA_MOE_DOWN_ROW512") != NULL ? 512u : - getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL ? 2048u : - getenv("DS4_CUDA_MOE_DOWN_ROW1024") != NULL ? 1024u : 512u; - const uint32_t use_down_row2048 = !q4k_path && use_atomic_down && expert_tile_m == 8u && - (getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL || - getenv("DS4_CUDA_MOE_DOWN_ROW256") != NULL || - getenv("DS4_CUDA_MOE_DOWN_ROW128") != NULL || - getenv("DS4_CUDA_MOE_DOWN_ROW64") != NULL || - (use_down_tile16 && - getenv("DS4_CUDA_MOE_NO_DOWN_ROW2048") == NULL && - getenv("DS4_CUDA_MOE_NO_DOWN_ROW256") == NULL && - getenv("DS4_CUDA_MOE_NO_DOWN_ROW128") == NULL && - getenv("DS4_CUDA_MOE_NO_DOWN_ROW64") == NULL)); - const uint32_t use_direct_down_sum = - n_tokens == 1u && (n_expert == 6u || n_expert == 3u) && - getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6") == NULL; - const uint32_t use_direct_midq = - q4k_path && use_direct_down_sum && !write_gate_up && - getenv("DS4_CUDA_MOE_DIRECT_MIDQ") != NULL && - getenv("DS4_CUDA_MOE_NO_DIRECT_MIDQ") == NULL; - const uint32_t use_q4_gate_h16r8 = - q4k_path && !use_direct_midq && - getenv("DS4_CUDA_MOE_Q4_GATE_H16R8") != NULL && - getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16R8") == NULL; - const uint32_t use_q4_gate_h16 = - q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && - getenv("DS4_CUDA_MOE_Q4_GATE_H16") != NULL && - getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16") == NULL; - const uint32_t use_q4_gate_w32r16 = - q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 && - getenv("DS4_CUDA_MOE_Q4_GATE_W32R16") != NULL && - getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32R16") == NULL; - const uint32_t use_q4_gate_w32 = - q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 && - !use_q4_gate_w32r16 && - getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32") == NULL; - const uint32_t use_q4_gate_w32_noaux = - use_q4_gate_w32 && !write_gate_up && - getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32_NOAUX") == NULL; - const uint32_t use_q4_down_slot3 = - q4k_path && use_direct_down_sum && n_expert == 3u && - getenv("DS4_CUDA_MOE_Q4_DOWN_SLOT3") != NULL && - getenv("DS4_CUDA_MOE_NO_Q4_DOWN_SLOT3") == NULL; - const uint32_t use_q4_midq_sidecar = - q4k_path && use_direct_down_sum && use_q4_gate_w32_noaux && - !use_direct_midq && !write_gate_up && - (expert_mid_dim % CUDA_QK_K) == 0u && - getenv("DS4_CUDA_MOE_MIDQ_SIDECAR") != NULL && - getenv("DS4_CUDA_MOE_NO_MIDQ_SIDECAR") == NULL; - float *midq_sidecar = use_q4_midq_sidecar ? (float *)up->ptr : NULL; - if (g_cuda_moe_decode_graph && - !owned_filtered && - !profile_moe && - q4k_path && - n_tokens == 1u && - use_direct_down_sum && - use_q4_gate_w32 && - !use_q4_gate_w32r16 && - !use_q4_down_slot3 && - !use_direct_midq && - !use_q4_midq_sidecar && - (n_expert == 3u || n_expert == 6u)) { - int grc = routed_moe_decode_q4_graph_launch( - logical_tier, - (float *)out->ptr, - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - down_w, - xq, - midq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - n_expert, - write_gate_up, - clamp, - (const float *)x->ptr); - if (grc == 1) return 1; - if (grc < 0) return 0; - } - uint32_t *sorted_pairs = NULL; - uint32_t *sorted_offsets = NULL; - uint32_t *sorted_counts = NULL; - uint32_t *tile_total = NULL; - uint32_t *tile_experts = NULL; - uint32_t *tile_starts = NULL; - uint32_t *tile16_total = NULL; - uint32_t *tile16_experts = NULL; - uint32_t *tile16_starts = NULL; - uint32_t tile_capacity = 0; - uint32_t tile16_capacity = 0; - dim3 xq_grid(xq_blocks, n_tokens, 1); - q8_K_quantize_kernel<<>>(xq, (const float *)x->ptr, expert_in_dim, n_tokens); - ok = cuda_ok(cudaGetLastError(), "routed_moe x quantize launch"); - if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0); - if (ok && use_sorted_pairs) { - const uint64_t counts_bytes = (uint64_t)n_total_expert * sizeof(uint32_t); - const uint64_t offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t); - const uint64_t cursors_bytes = (uint64_t)n_total_expert * sizeof(uint32_t); - const uint64_t sorted_bytes = (uint64_t)pair_count * sizeof(uint32_t); - tile_capacity = (pair_count + expert_tile_m - 1u) / expert_tile_m + n_total_expert; - tile16_capacity = (use_down_tile16 || use_q4_mma_tiles16) ? ((pair_count + 15u) / 16u + n_total_expert) : 0u; - const uint64_t tile_offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t); - const uint64_t tile_total_bytes = sizeof(uint32_t); - const uint64_t tile_experts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t); - const uint64_t tile_starts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t); - const uint64_t tile16_offsets_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? (((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t)) : 0u; - const uint64_t tile16_total_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? sizeof(uint32_t) : 0u; - const uint64_t tile16_experts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t); - const uint64_t tile16_starts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t); - const uint64_t tile_offsets_off = counts_bytes + offsets_bytes + cursors_bytes + sorted_bytes; - const uint64_t tile_total_off = tile_offsets_off + tile_offsets_bytes; - const uint64_t tile_experts_off = tile_total_off + tile_total_bytes; - const uint64_t tile_starts_off = tile_experts_off + tile_experts_bytes; - const uint64_t tile16_offsets_off = tile_starts_off + tile_starts_bytes; - const uint64_t tile16_total_off = tile16_offsets_off + tile16_offsets_bytes; - const uint64_t tile16_experts_off = tile16_total_off + tile16_total_bytes; - const uint64_t tile16_starts_off = tile16_experts_off + tile16_experts_bytes; - const uint64_t scratch_bytes = tile16_starts_off + tile16_starts_bytes; - uint8_t *scratch = (uint8_t *)cuda_tmp_alloc_on(logical_tier, scratch_bytes, - "routed_moe sorted pairs"); - if (!scratch) { - ok = 0; - } else { - uint32_t *counts = (uint32_t *)scratch; - uint32_t *offsets = (uint32_t *)(scratch + counts_bytes); - uint32_t *cursors = (uint32_t *)(scratch + counts_bytes + offsets_bytes); - sorted_pairs = (uint32_t *)(scratch + counts_bytes + offsets_bytes + cursors_bytes); - sorted_offsets = offsets; - sorted_counts = counts; - uint32_t *tile_offsets = (uint32_t *)(scratch + tile_offsets_off); - tile_total = (uint32_t *)(scratch + tile_total_off); - tile_experts = (uint32_t *)(scratch + tile_experts_off); - tile_starts = (uint32_t *)(scratch + tile_starts_off); - uint32_t *tile16_offsets = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_offsets_off) : NULL; - tile16_total = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_total_off) : NULL; - tile16_experts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_experts_off) : NULL; - tile16_starts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_starts_off) : NULL; - if (use_small_sorted_prep) { - moe_prepare_sorted_tiles_small_kernel<<<1, 128>>>( - counts, offsets, cursors, sorted_pairs, - tile_offsets, tile_total, tile_experts, tile_starts, - tile16_offsets, tile16_total, tile16_experts, tile16_starts, - (const int32_t *)selected->ptr, pair_count, n_total_expert, - expert_tile_m, use_down_tile16 || use_q4_mma_tiles16); - ok = cuda_ok(cudaGetLastError(), - "routed_moe small sorted setup launch"); - } else { - ok = cuda_ok(cudaMemset(counts, 0, counts_bytes), - "routed_moe sorted counts clear"); - } - if (ok && !use_small_sorted_prep) { - moe_count_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( - counts, - (const int32_t *)selected->ptr, - pair_count, - n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe sorted count launch"); - } - if (ok && !use_small_sorted_prep) { - moe_prefix_sorted_pairs_kernel<<<1, 1>>>(offsets, cursors, counts, n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe sorted prefix launch"); - } - if (ok && !use_small_sorted_prep) { - moe_scatter_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( - sorted_pairs, - cursors, - (const int32_t *)selected->ptr, - pair_count, - n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe sorted scatter launch"); - } - if (ok && use_expert_tiles && !use_small_sorted_prep) { - moe_build_expert_tile_offsets_kernel<<<1, 1>>>(tile_offsets, tile_total, counts, expert_tile_m, n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile offsets launch"); - } - if (ok && use_expert_tiles && !use_small_sorted_prep) { - moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256>>>( - tile_experts, tile_starts, tile_offsets, counts, expert_tile_m, n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe expert tiles launch"); - } - if (ok && use_expert_tiles && !use_small_sorted_prep && - (use_down_tile16 || use_q4_mma_tiles16)) { - moe_build_expert_tile_offsets_kernel<<<1, 1>>>(tile16_offsets, tile16_total, counts, 16u, n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 offsets launch"); - } - if (ok && use_expert_tiles && !use_small_sorted_prep && - (use_down_tile16 || use_q4_mma_tiles16)) { - moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256>>>( - tile16_experts, tile16_starts, tile16_offsets, counts, 16u, n_total_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 launch"); - } - } - } - if (prof_ev[2]) (void)cudaEventRecord(prof_ev[2], 0); - if (ok && owned_filtered && use_sorted_pairs && - !use_owned_sparse_buffers) { - const uint64_t mid_bytes = - (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float); - ok = cuda_ok(cudaMemset(mid->ptr, 0, (size_t)mid_bytes), - "owned routed_moe mid clear"); - } - if (ok) { - dim3 mgrid((expert_mid_dim + 31u) / 32u, n_tokens * n_expert, 1); - if (ok && sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && tile_total && tile_experts && tile_starts) { - if (q4k_path) { - const int use_q4_mma = cuda_q4_mma_ok() && - ((((uintptr_t)gate_w | (uintptr_t)up_w | - gate_row_bytes | gate_expert_bytes) & 15u) == 0u) && - xq_blocks <= 16u && (expert_mid_dim & 7u) == 0u; - const int use_q4_mma_t16 = use_q4_mma && use_q4_mma_tiles16 && - tile16_total && tile16_experts && tile16_starts && - xq_blocks == 16u && cuda_q4_mma_tile16_shmem_ok(0); - if (use_q4_mma_t16 && use_gate_row2048) { - const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert); - const size_t t16sh = 16u * 16u * sizeof(cuda_block_q8_K); - if (gate_row_span == 512u) { - dim3 tgrid((expert_mid_dim + 511u) / 512u, t16cap, 1); - moe_gate_up_mid_q4K_tile16_mma_kernel<512><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else if (gate_row_span == 1024u) { - dim3 tgrid((expert_mid_dim + 1023u) / 1024u, t16cap, 1); - moe_gate_up_mid_q4K_tile16_mma_kernel<1024><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else { - dim3 tgrid((expert_mid_dim + 2047u) / 2048u, t16cap, 1); - moe_gate_up_mid_q4K_tile16_mma_kernel<2048><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else if (use_q4_mma && use_gate_row2048) { - if (gate_row_span == 512u) { - dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); - moe_gate_up_mid_q4K_tile8_mma_kernel<512><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else if (gate_row_span == 1024u) { - dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); - moe_gate_up_mid_q4K_tile8_mma_kernel<1024><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else { - dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); - moe_gate_up_mid_q4K_tile8_mma_kernel<2048><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else if (use_gate_row2048) { - if (gate_row_span == 512u) { - dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); - moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<512><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else if (gate_row_span == 1024u) { - dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); - moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<1024><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else { - dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); - moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<2048><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else { - dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); - moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<32><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else if (use_gate_row2048) { - if (gate_row_span == 512u) { - dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); - moe_gate_up_mid_expert_tile8_rowspan_kernel<512><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else if (gate_row_span == 1024u) { - dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); - moe_gate_up_mid_expert_tile8_rowspan_kernel<1024><<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else { - dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); - moe_gate_up_mid_expert_tile8_row2048_kernel<<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else if (expert_tile_m == 8u) { - dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); - moe_gate_up_mid_expert_tile8_row32_kernel<<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } else { - dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); - moe_gate_up_mid_expert_tile4_row32_kernel<<>>( - (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, - gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, - tile_total, tile_experts, tile_starts, (const float *)weights->ptr, - gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, - write_gate_up, clamp); - } - } else if (ok && sorted_pairs && use_p2_sorted) { - dim3 p2_mgrid((expert_mid_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1); - moe_gate_up_mid_sorted_p2_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - sorted_pairs, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - pair_count, - clamp); - } else if (ok && sorted_pairs) { - moe_gate_up_mid_sorted_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - sorted_pairs, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - clamp); - } else if (ok) { - dim3 qgrid((expert_mid_dim + MOE_DECODE_ROWS_PER_BLOCK - 1u) / MOE_DECODE_ROWS_PER_BLOCK, n_tokens * n_expert, 1); - if (q4k_path) { - /* Q4_K gate/up: the decode kernel is token-indexed via - * pair = blockIdx.y; tok = pair / n_expert, so the same - * launch covers both n_tokens == 1 (decode) and n_tokens > 1 - * (prefill). q4k_path is steered here by use_sorted_pairs = 0 - * cascading the IQ2 sorted/expert-tile branches off. */ - if (use_direct_midq) { - dim3 mqgrid(midq_blocks, n_tokens * n_expert, 1); - moe_gate_up_midq_decode_q4K_qwarp32_kernel<<>>( - (float *)mid->ptr, - midq, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - clamp); - } else if (use_q4_gate_h16r8) { - dim3 h8grid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1); - moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } else if (use_q4_gate_w32r16) { - dim3 w16grid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1); - moe_gate_up_mid_decode_q4K_warp32_row16_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } else if (use_q4_gate_w32) { - dim3 wgrid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1); - if (use_q4_gate_w32_noaux) { - if (use_q4_midq_sidecar) { - moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel<<>>( - (float *)mid->ptr, - midq_sidecar, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - clamp); - } else { - moe_gate_up_mid_decode_q4K_warp32_noaux_kernel<<>>( - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - clamp); - } - } else { - moe_gate_up_mid_decode_q4K_warp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } - } else if (use_q4_gate_h16) { - dim3 hgrid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1); - moe_gate_up_mid_decode_q4K_hwarp16_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } else { - moe_gate_up_mid_decode_q4K_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } - } else if (use_decode_lut_gate) { - moe_gate_up_mid_decode_lut_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - write_gate_up, - clamp); - } else { - moe_gate_up_mid_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - n_expert, - clamp); - } - } - ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch"); - } - if (prof_ev[3]) (void)cudaEventRecord(prof_ev[3], 0); - if (ok && !use_direct_midq) { - dim3 midq_grid(midq_blocks, n_tokens * n_expert, 1); - if (use_q4_midq_sidecar) { - q8_K_quantize_sidecar_kernel<<>>( - midq, - (const float *)mid->ptr, - midq_sidecar, - expert_mid_dim, - n_tokens * n_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe mid sidecar quantize launch"); - } else if (use_owned_sparse_buffers) { - q8_K_quantize_owned_kernel<<>>( - midq, - (const float *)mid->ptr, - (const int32_t *)selected->ptr, - expert_mid_dim, - n_tokens * n_expert, - 0u, - n_total_expert); - ok = cuda_ok(cudaGetLastError(), - "owned routed_moe active mid quantize launch"); - } else { - q8_K_quantize_kernel<<>>(midq, (const float *)mid->ptr, expert_mid_dim, n_tokens * n_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe mid quantize launch"); - } - } - if (prof_ev[4]) (void)cudaEventRecord(prof_ev[4], 0); - if (ok && owned_filtered && use_sorted_pairs && !use_atomic_down && - !use_owned_sparse_buffers) { - const uint64_t down_clear_bytes = - (uint64_t)n_tokens * n_expert * out_dim * sizeof(float); - ok = cuda_ok(cudaMemset(down->ptr, 0, (size_t)down_clear_bytes), - "owned routed_moe down clear"); - } - if (ok) { - dim3 dgrid((out_dim + 31u) / 32u, n_tokens * n_expert, 1); - uint32_t *down_tile_total = tile_total; - uint32_t *down_tile_experts = tile_experts; - uint32_t *down_tile_starts = tile_starts; - uint32_t down_tile_capacity = tile_capacity; - if (use_down_tile16 && tile16_total && tile16_experts && tile16_starts) { - down_tile_total = tile16_total; - down_tile_experts = tile16_experts; - down_tile_starts = tile16_starts; - down_tile_capacity = tile16_capacity; - } - if (use_direct_down_sum) { - dim3 sgrid((out_dim + 31u) / 32u, 1, 1); - if (q4k_path) { - if (n_expert == 6u) { - moe_down_q4K_sum6_qwarp32_kernel<<>>( - (float *)out->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim); - } else { - if (use_q4_down_slot3) { - dim3 swgrid((out_dim + 7u) / 8u, 1, 1); - moe_down_q4K_sum3_slotwarp_kernel<<>>( - (float *)out->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim); - } else { - moe_down_q4K_sum3_qwarp32_kernel<<>>( - (float *)out->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim); - } - } - } else { - if (n_expert == 6u) { - moe_down_sum6_qwarp32_kernel<<>>( - (float *)out->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim); - } else { - moe_down_sum3_qwarp32_kernel<<>>( - (float *)out->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim); - } - } - } else if (use_atomic_down) { - uint64_t n = (uint64_t)n_tokens * out_dim; - zero_kernel<<<(n + 255u) / 256u, 256>>>((float *)out->ptr, n); - ok = cuda_ok(cudaGetLastError(), "routed_moe atomic zero launch"); - } - if (use_direct_down_sum) { - /* The direct decode kernel writes the final token row. */ - } else if (sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && - down_tile_total && down_tile_experts && down_tile_starts) { - if (q4k_path) { - const int use_q4_down_mma = cuda_q4_mma_ok() && - ((((uintptr_t)down_w | down_row_bytes | down_expert_bytes) & 15u) == 0u) && - midq_blocks <= 8u && (out_dim & 7u) == 0u; - const int use_q4_down_t16 = use_q4_down_mma && use_q4_mma_tiles16 && - tile16_total && tile16_experts && tile16_starts && - midq_blocks <= 16u && cuda_q4_mma_tile16_shmem_ok(1); - if (use_q4_down_t16 && use_q4_down_rowspan) { - const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert); - const size_t dt16sh = 16u * (size_t)midq_blocks * sizeof(cuda_block_q8_K); - if (down_row_span == 512u) { - dim3 tgrid((out_dim + 511u) / 512u, t16cap, 1); - moe_down_q4K_tile16_mma_kernel<512><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else if (down_row_span == 1024u) { - dim3 tgrid((out_dim + 1023u) / 1024u, t16cap, 1); - moe_down_q4K_tile16_mma_kernel<1024><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else { - dim3 tgrid((out_dim + 2047u) / 2048u, t16cap, 1); - moe_down_q4K_tile16_mma_kernel<2048><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } - } else if (use_q4_down_mma && use_q4_down_rowspan) { - if (down_row_span == 512u) { - dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); - moe_down_q4K_tile8_mma_kernel<512><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else if (down_row_span == 1024u) { - dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); - moe_down_q4K_tile8_mma_kernel<1024><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else { - dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); - moe_down_q4K_tile8_mma_kernel<2048><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } - } else if (use_q4_down_rowspan) { - if (down_row_span == 512u) { - dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); - moe_down_q4K_expert_tile8_rowspan_kernel<512><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else if (down_row_span == 1024u) { - dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); - moe_down_q4K_expert_tile8_rowspan_kernel<1024><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } else { - dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); - moe_down_q4K_expert_tile8_rowspan_kernel<2048><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } - } else { - dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); - moe_down_q4K_expert_tile8_rowspan_kernel<32><<>>( - (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert); - } - } else if (use_down_row2048) { - if (down_row_span == 512u) { - dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); - moe_down_expert_tile16_rowspan_kernel<512><<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } else if (down_row_span == 1024u) { - dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); - moe_down_expert_tile16_rowspan_kernel<1024><<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } else { - dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); - moe_down_expert_tile16_row2048_kernel<<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } - } else if (use_down_tile16) { - dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); - moe_down_expert_tile16_row32_kernel<<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } else if (expert_tile_m == 8u) { - dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); - moe_down_expert_tile8_row32_kernel<<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } else { - dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); - moe_down_expert_tile4_row32_kernel<<>>( - use_atomic_down ? (float *)out->ptr : (float *)down->ptr, - down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, - down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, use_atomic_down); - } - } else if (sorted_pairs && use_p2_sorted) { - dim3 p2_dgrid((out_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1); - moe_down_sorted_p2_qwarp32_kernel<<>>( - (float *)down->ptr, - down_w, - midq, - sorted_pairs, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - n_expert, - pair_count); - } else if (sorted_pairs) { - moe_down_sorted_qwarp32_kernel<<>>( - (float *)down->ptr, - down_w, - midq, - sorted_pairs, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - n_expert); - } else if (q4k_path) { - /* Q4_K prefill down. New kernel mirrors moe_down_qwarp32_kernel - * grid/geometry, swapping the weight block type to cuda_block_q4_K - * and the dot helper to dev_dot_q4_K_q8_K_block. Writes per-pair - * outputs into down->ptr; moe_sum_kernel below sums them across - * experts into out->ptr. */ - moe_down_q4K_qwarp32_kernel<<>>( - (float *)down->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - n_expert); - } else { - moe_down_qwarp32_kernel<<>>( - (float *)down->ptr, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - n_expert); - } - ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); - } - if (prof_ev[5]) (void)cudaEventRecord(prof_ev[5], 0); - if (ok && !use_atomic_down && !use_direct_down_sum) { - uint64_t n = (uint64_t)n_tokens * out_dim; - if (use_owned_sparse_buffers) { - moe_sum_owned_kernel<<<(n + 255) / 256, 256>>>( - (float *)out->ptr, - (const float *)down->ptr, - (const int32_t *)selected->ptr, - out_dim, - n_expert, - n_tokens); - } else { - moe_sum_kernel<<<(n + 255) / 256, 256>>>( - (float *)out->ptr, - (const float *)down->ptr, - out_dim, - n_expert, - n_tokens); - } - ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch"); - } - if (prof_ev[6]) { - (void)cudaEventRecord(prof_ev[6], 0); - if (cudaEventSynchronize(prof_ev[6]) == cudaSuccess) { - float ms_xq = 0.0f, ms_sort = 0.0f, ms_gate = 0.0f, ms_midq = 0.0f, ms_down = 0.0f, ms_sum = 0.0f, ms_total = 0.0f; - (void)cudaEventElapsedTime(&ms_xq, prof_ev[0], prof_ev[1]); - (void)cudaEventElapsedTime(&ms_sort, prof_ev[1], prof_ev[2]); - (void)cudaEventElapsedTime(&ms_gate, prof_ev[2], prof_ev[3]); - (void)cudaEventElapsedTime(&ms_midq, prof_ev[3], prof_ev[4]); - (void)cudaEventElapsedTime(&ms_down, prof_ev[4], prof_ev[5]); - (void)cudaEventElapsedTime(&ms_sum, prof_ev[5], prof_ev[6]); - (void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[6]); - fprintf(stderr, - "ds4: CUDA MoE profile tokens=%u pairs=%u xq=%.3f sort=%.3f gateup=%.3f midq=%.3f down=%.3f sum=%.3f total=%.3f ms\n", - n_tokens, pair_count, ms_xq, ms_sort, ms_gate, ms_midq, ms_down, ms_sum, ms_total); - } - for (uint32_t i = 0; i < 7u; i++) (void)cudaEventDestroy(prof_ev[i]); - } - return ok; - } - - if (ok) { - dim3 mgrid(expert_mid_dim, n_tokens * n_expert, 1); - moe_gate_up_mid_f32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - (const float *)x->ptr, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - expert_in_dim, - expert_mid_dim, - n_expert, - clamp); - ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch"); - } - if (ok) { - dim3 dgrid(out_dim, n_tokens * n_expert, 1); - moe_down_f32_kernel<<>>( - (float *)down->ptr, - down_w, - (const float *)mid->ptr, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - expert_mid_dim, - out_dim, - n_expert); - ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); - } - if (ok) { - uint64_t n = (uint64_t)n_tokens * out_dim; - moe_sum_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)down->ptr, out_dim, n_expert, n_tokens); - ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch"); - } - return ok; -} - -extern "C" int ds4_gpu_routed_moe_one_owned_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - ds4_gpu_tensor *down, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t resident_expert_base, - uint32_t resident_expert_count, - float clamp, - const ds4_gpu_tensor *x, - ds4_gpu_tensor *down_output, - bool pack_fixed3, - ds4_gpu_tensor *shared_prequant) { - if (!out || !gate || !up || !mid || !down || !model_map || - !selected || !weights || !x || n_expert != 6u || - n_total_expert == 0u || resident_expert_count == 0u || - gate_expert_bytes == 0u || gate_row_bytes == 0u || - down_expert_bytes == 0u || down_row_bytes == 0u || - resident_expert_base >= n_total_expert || - resident_expert_count > n_total_expert - resident_expert_base || - expert_in_dim % CUDA_QK_K != 0u || - expert_mid_dim % CUDA_QK_K != 0u || - selected->bytes < 6u * sizeof(int32_t) || - weights->bytes < 6u * sizeof(float) || - x->bytes < (uint64_t)expert_in_dim * sizeof(float) || - mid->bytes < 6ull * expert_mid_dim * sizeof(float) || - out->bytes < (uint64_t)out_dim * sizeof(float)) { - return 0; - } - if (pack_fixed3 && resident_expert_base == 0u) return 0; - const bool q4k_path = gate_type == 12u && down_type == 12u; - if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; - if (q4k_path && getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL) { - fprintf(stderr, "ds4: CUDA owned Q4 decode does not support gate/up auxiliary output\n"); - return 0; - } - if (!q4k_path && getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") != NULL) { - fprintf(stderr, "ds4: CUDA owned IQ2 decode requires the LUT gate path\n"); - return 0; - } - const bool write_aux = - !q4k_path && getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL; - - if (resident_expert_base > UINT64_MAX / gate_expert_bytes || - resident_expert_count > UINT64_MAX / gate_expert_bytes || - resident_expert_base > UINT64_MAX / down_expert_bytes || - resident_expert_count > UINT64_MAX / down_expert_bytes) { - return 0; - } - const uint64_t gate_shift = (uint64_t)resident_expert_base * gate_expert_bytes; - const uint64_t down_shift = (uint64_t)resident_expert_base * down_expert_bytes; - const uint64_t gate_bytes = (uint64_t)resident_expert_count * gate_expert_bytes; - const uint64_t down_bytes = (uint64_t)resident_expert_count * down_expert_bytes; - if (gate_offset > model_size || gate_shift > model_size - gate_offset || - gate_bytes > model_size - gate_offset - gate_shift || - up_offset > model_size || gate_shift > model_size - up_offset || - gate_bytes > model_size - up_offset - gate_shift || - down_offset > model_size || down_shift > model_size - down_offset || - down_bytes > model_size - down_offset - down_shift) { - return 0; - } - - const int logical_tier = ds4_tensor_device_idx(out); - const char *gate_w = (const char *)cuda_resolve_weight_ptr( - model_map, gate_offset + gate_shift, gate_bytes, - logical_tier, "moe_owned_gate"); - const char *up_w = (const char *)cuda_resolve_weight_ptr( - model_map, up_offset + gate_shift, gate_bytes, - logical_tier, "moe_owned_up"); - const char *down_w = (const char *)cuda_resolve_weight_ptr( - model_map, down_offset + down_shift, down_bytes, - logical_tier, "moe_owned_down"); - if (!gate_w || !up_w || !down_w) return 0; - - const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; - const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; - const uint64_t xq_bytes = (uint64_t)xq_blocks * sizeof(cuda_block_q8_K); - const uint64_t midq_bytes = 6ull * midq_blocks * sizeof(cuda_block_q8_K); - const uint64_t down_output_bytes = - (uint64_t)(pack_fixed3 ? 4u : 6u) * out_dim * sizeof(float); - const uint64_t aux_bytes = 6ull * expert_mid_dim * sizeof(float); - const uint64_t shared_q8_blocks = expert_in_dim / 32u; - const uint64_t shared_q8_bytes = shared_q8_blocks * 32u; - const uint64_t shared_scale_offset = - (shared_q8_bytes + 15u) & ~15ull; - const uint64_t shared_prequant_bytes = - shared_scale_offset + shared_q8_blocks * sizeof(float); - if (down->bytes < xq_bytes || down->bytes < down_output_bytes || - (down_output && down_output->bytes < down_output_bytes) || - gate->bytes < midq_bytes || - (shared_prequant && - (shared_prequant->bytes < shared_prequant_bytes || - ds4_tensor_device_idx(shared_prequant) != logical_tier)) || - (write_aux && (gate->bytes < aux_bytes || up->bytes < aux_bytes))) { - return 0; - } - float *down_dst = (float *)(down_output ? down_output->ptr : down->ptr); - cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; - cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; - - dim3 xq_grid(xq_blocks, 1, 1); - if (shared_prequant) { - int8_t *shared_xq = (int8_t *)shared_prequant->ptr; - float *shared_scale = (float *)((char *)shared_prequant->ptr + - shared_scale_offset); - q8_K_q8_0_quantize_kernel<<>>( - xq, shared_xq, shared_scale, (const float *)x->ptr, - expert_in_dim, 1u); - } else { - q8_K_quantize_kernel<<>>( - xq, (const float *)x->ptr, expert_in_dim, 1u); - } - if (!cuda_ok(cudaGetLastError(), "owned routed_moe x quantize launch")) return 0; - - if (q4k_path) { - dim3 gate_grid((expert_mid_dim + 7u) / 8u, 6u, 1u); - moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel<<>>( - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - 6u, - resident_expert_base, - resident_expert_count, - clamp); - } else { - dim3 gate_grid((expert_mid_dim + 31u) / 32u, 6u, 1u); - moe_gate_up_mid_decode_lut_owned_qwarp32_kernel<<>>( - (float *)gate->ptr, - (float *)up->ptr, - (float *)mid->ptr, - gate_w, - up_w, - xq, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - gate_expert_bytes, - gate_row_bytes, - xq_blocks, - expert_mid_dim, - 6u, - resident_expert_base, - resident_expert_count, - write_aux, - clamp); - } - if (!cuda_ok(cudaGetLastError(), "owned routed_moe gate/up launch")) return 0; - - dim3 midq_grid(midq_blocks, 6u, 1u); - q8_K_quantize_owned_kernel<<>>( - midq, - (const float *)mid->ptr, - (const int32_t *)selected->ptr, - expert_mid_dim, - 6u, - resident_expert_base, - resident_expert_count); - if (!cuda_ok(cudaGetLastError(), "owned routed_moe mid quantize launch")) return 0; - - dim3 down_grid((out_dim + 31u) / 32u, pack_fixed3 ? 4u : 6u, 1u); - if (q4k_path && pack_fixed3) { - moe_down_q4K_owned_packed_qwarp32_kernel<<>>( - down_dst, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - resident_expert_base, - resident_expert_count); - } else if (q4k_path) { - moe_down_q4K_owned_slots_qwarp32_kernel<<>>( - down_dst, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - resident_expert_base, - resident_expert_count); - } else if (pack_fixed3) { - moe_down_owned_packed_qwarp32_kernel<<>>( - down_dst, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - resident_expert_base, - resident_expert_count); - } else { - moe_down_owned_slots_qwarp32_kernel<<>>( - down_dst, - down_w, - midq, - (const int32_t *)selected->ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - resident_expert_base, - resident_expert_count); - } - return cuda_ok(cudaGetLastError(), "owned routed_moe down launch"); -} - -extern "C" int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *home_slots, - const ds4_gpu_tensor *peer_slots, - const ds4_gpu_tensor *selected, - uint32_t out_dim, - uint32_t expert_split, - uint32_t rows) { - if (!out || !home_slots || !peer_slots || !selected || out_dim == 0u || - rows == 0u || rows > 65535u) { - return 0; - } - const uint64_t row_elems = (uint64_t)rows * out_dim; - if (row_elems > UINT64_MAX / (6u * sizeof(float))) return 0; - const uint64_t out_bytes = row_elems * sizeof(float); - const uint64_t slots_bytes = row_elems * 6u * sizeof(float); - const uint64_t selected_bytes = (uint64_t)rows * 6u * sizeof(int32_t); - if (out->bytes < out_bytes || home_slots->bytes < slots_bytes || - peer_slots->bytes < slots_bytes || selected->bytes < selected_bytes) { - return 0; - } - const dim3 grid((out_dim + 255u) / 256u, rows, 1u); - moe_owned_slots_combine_fixed3_kernel<<>>( - (float *)out->ptr, - (const float *)home_slots->ptr, - (const float *)peer_slots->ptr, - (const int32_t *)selected->ptr, - out_dim, - expert_split); - return cuda_ok(cudaGetLastError(), - "owned routed_moe slot rows combine launch"); -} - -extern "C" int ds4_gpu_routed_moe_owned_slots_combine_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *home_slots, - const ds4_gpu_tensor *peer_slots, - const ds4_gpu_tensor *selected, - uint32_t out_dim, - uint32_t expert_split) { - return ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( - out, home_slots, peer_slots, selected, - out_dim, expert_split, 1u); -} - -extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *home_slots, - const ds4_gpu_tensor *peer_packed, - const ds4_gpu_tensor *selected, - uint32_t out_dim, - uint32_t expert_split) { - const uint64_t home_bytes = 6ull * out_dim * sizeof(float); - const uint64_t peer_bytes = 4ull * out_dim * sizeof(float); - if (!out || !home_slots || !peer_packed || !selected || out_dim == 0u || - out->bytes < (uint64_t)out_dim * sizeof(float) || - home_slots->bytes < home_bytes || peer_packed->bytes < peer_bytes || - selected->bytes < 6u * sizeof(int32_t)) { - return 0; - } - moe_owned_packed_combine_fixed3_kernel<<< - (out_dim + 255u) / 256u, 256>>>( - (float *)out->ptr, - (const float *)home_slots->ptr, - (const float *)peer_packed->ptr, - (const int32_t *)selected->ptr, - out_dim, - expert_split); - return cuda_ok(cudaGetLastError(), - "owned routed_moe packed combine launch"); -} - -extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, - const ds4_gpu_tensor *add_in, - uint32_t layer_index, - bool force_resident) { - if (add_in) { - if (!ds4_gpu_add_tensor(out, out, add_in, - (uint32_t)(out->bytes / sizeof(float)))) return 0; - } - return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, - gate_offset, up_offset, down_offset, - gate_type, down_type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - expert_in_dim, expert_mid_dim, out_dim, - selected, weights, n_total_expert, n_expert, clamp, x, - layer_index, 1, force_resident ? 0 : 1, 0); -} -extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) { - (void)force_resident; - if (mid_is_f16) *mid_is_f16 = false; - return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, - gate_offset, up_offset, down_offset, - gate_type, down_type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - expert_in_dim, expert_mid_dim, out_dim, - selected, weights, n_total_expert, n_expert, clamp, x, - layer_index, n_tokens, 1, 0); -} - -extern "C" int ds4_gpu_routed_moe_batch_owned_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - ds4_gpu_tensor *down, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t resident_expert_base, - uint32_t resident_expert_count, - float clamp, - const ds4_gpu_tensor *x, - uint32_t layer_index, - uint32_t n_tokens, - bool *mid_is_f16) { - if (mid_is_f16) *mid_is_f16 = false; - if (!selected || !weights || n_tokens == 0u || n_expert == 0u || - n_total_expert == 0u || resident_expert_count == 0u || - gate_expert_bytes == 0u || gate_row_bytes == 0u || - down_expert_bytes == 0u || down_row_bytes == 0u || - resident_expert_base >= n_total_expert || - resident_expert_count > n_total_expert - resident_expert_base || - resident_expert_base > UINT64_MAX / gate_expert_bytes || - resident_expert_base > UINT64_MAX / down_expert_bytes) { - return 0; - } - const uint64_t pair_count = (uint64_t)n_tokens * n_expert; - if (pair_count > UINT32_MAX || - selected->bytes < pair_count * sizeof(int32_t) || - weights->bytes < pair_count * sizeof(float)) { - return 0; - } - const uint64_t gate_shift = - (uint64_t)resident_expert_base * gate_expert_bytes; - const uint64_t down_shift = - (uint64_t)resident_expert_base * down_expert_bytes; - if (gate_offset > model_size || gate_shift > model_size - gate_offset || - up_offset > model_size || gate_shift > model_size - up_offset || - down_offset > model_size || down_shift > model_size - down_offset) { - return 0; - } - moe_filter_owned_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( - (int32_t *)selected->ptr, - (float *)weights->ptr, - pair_count, - n_total_expert, - resident_expert_base, - resident_expert_count); - if (!cuda_ok(cudaGetLastError(), "owned routed_moe pair filter launch")) return 0; - - return routed_moe_launch( - out, gate, up, mid, down, model_map, model_size, - gate_offset + gate_shift, - up_offset + gate_shift, - down_offset + down_shift, - gate_type, down_type, - gate_expert_bytes, gate_row_bytes, - down_expert_bytes, down_row_bytes, - expert_in_dim, expert_mid_dim, out_dim, - selected, weights, resident_expert_count, n_expert, - clamp, x, layer_index, n_tokens, 0, 1); -} -extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) { - if (!out || !mix || !model_map || n_hc != 4) return 0; - const uint64_t mix_bytes = 24ull * sizeof(float); - if (scale_offset > model_size || model_size - scale_offset < 3ull * sizeof(float) || - base_offset > model_size || model_size - base_offset < mix_bytes || - mix->bytes < mix_bytes || out->bytes < mix_bytes) return 0; - const int logical_tier = ds4_tensor_device_idx(out); - const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale"); - const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base"); - if (!scale || !base) return 0; - uint32_t n_rows = (uint32_t)(mix->bytes / mix_bytes); - if (out->bytes / mix_bytes < n_rows) n_rows = (uint32_t)(out->bytes / mix_bytes); - hc_split_sinkhorn_kernel<<<(n_rows + 255) / 256, 256>>>( - (float *)out->ptr, (const float *)mix->ptr, - scale, - base, - n_rows, sinkhorn_iters, eps); - return cuda_ok(cudaGetLastError(), "hc_split_sinkhorn launch"); -} -extern "C" int ds4_gpu_hc_weighted_sum_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *weights, uint32_t n_embd, uint32_t n_hc) { - if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float))); - hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>( - (float *)out->ptr, (const float *)residual_hc->ptr, (const float *)weights->ptr, - n_embd, n_hc, n_tokens, n_hc); - return cuda_ok(cudaGetLastError(), "hc_weighted_sum launch"); -} -extern "C" int ds4_gpu_hc_weighted_sum_split_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { - if (!out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float))); - uint32_t stride = (uint32_t)(2u * n_hc + n_hc * n_hc); - hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>( - (float *)out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr, - n_embd, n_hc, n_tokens, stride); - return cuda_ok(cudaGetLastError(), "hc_weighted_sum_split launch"); -} -extern "C" int ds4_gpu_hc_split_weighted_sum_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *split, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint32_t n_embd, - uint32_t n_hc, - uint32_t sinkhorn_iters, - float eps) { - if (!out || !split || !mix || !residual_hc || !model_map || - n_embd == 0 || n_hc != 4) { - return 0; - } - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t mix_bytes = mix_hc * sizeof(float); - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 || - scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset || - base_offset > model_size || mix_bytes > model_size - base_offset) { - return 0; - } - uint64_t n_rows = out->bytes / out_row_bytes; - if (mix->bytes < n_rows * mix_bytes || - split->bytes < n_rows * mix_bytes || - residual_hc->bytes < n_rows * residual_row_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out); - const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale"); - const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base"); - if (!scale || !base) return 0; - hc_split_weighted_sum_fused_kernel<<<(uint32_t)n_rows, 256>>>( - (float *)out->ptr, - (float *)split->ptr, - (const float *)mix->ptr, - (const float *)residual_hc->ptr, - scale, - base, - n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps); - return cuda_ok(cudaGetLastError(), "hc split weighted sum launch"); -} -extern "C" int ds4_gpu_hc_split_weighted_sum_norm_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_out, - ds4_gpu_tensor *split, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint64_t norm_weight_offset, - uint32_t n_embd, - uint32_t n_hc, - uint32_t sinkhorn_iters, - float eps, - float norm_eps) { - if (getenv("DS4_CUDA_DISABLE_HC_SPLIT_NORM_FUSED") == NULL) { - if (!out || !norm_out || !split || !mix || !residual_hc || !model_map || - n_embd == 0 || n_hc != 4) { - return 0; - } - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t mix_bytes = mix_hc * sizeof(float); - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 || - norm_out->bytes < out->bytes || - scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset || - base_offset > model_size || mix_bytes > model_size - base_offset || - norm_weight_offset > model_size || - (uint64_t)n_embd * sizeof(float) > model_size - norm_weight_offset) { - return 0; - } - uint64_t n_rows = out->bytes / out_row_bytes; - if (n_rows == 1) { - if (mix->bytes < n_rows * mix_bytes || - split->bytes < n_rows * mix_bytes || - residual_hc->bytes < n_rows * residual_row_bytes) { - return 0; - } - const int logical_tier = ds4_tensor_device_idx(out); - const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, - 3ull * sizeof(float), logical_tier, "hc_scale"); - const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, - mix_bytes, logical_tier, "hc_base"); - const float *norm_w = (const float *)cuda_resolve_weight_ptr(model_map, norm_weight_offset, - (uint64_t)n_embd * sizeof(float), logical_tier, "hc_norm_weight"); - if (!scale || !base || !norm_w) return 0; - hc_split_weighted_sum_norm_fused_kernel<<<(uint32_t)n_rows, 256>>>( - (float *)out->ptr, - (float *)norm_out->ptr, - (float *)split->ptr, - (const float *)mix->ptr, - (const float *)residual_hc->ptr, - scale, - base, - norm_w, - n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps, norm_eps); - return cuda_ok(cudaGetLastError(), "hc split weighted sum norm launch"); - } - } - /* Multi-row fallback: norm EVERY row (rms_norm_weight_tensor is the - * single-row entry and would leave rows 1..n-1 of norm_out untouched). */ - if (!out || n_embd == 0) return 0; - return ds4_gpu_hc_split_weighted_sum_tensor(out, split, mix, residual_hc, - model_map, model_size, - scale_offset, base_offset, - n_embd, n_hc, - sinkhorn_iters, eps) && - ds4_gpu_rms_norm_weight_rows_tensor( - norm_out, out, model_map, model_size, - norm_weight_offset, n_embd, - (uint32_t)(out->bytes / - ((uint64_t)n_embd * sizeof(float))), - norm_eps); -} -extern "C" int ds4_gpu_output_hc_weights_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *pre, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint32_t n_hc, - float eps) { - if (!out || !pre || !model_map || n_hc == 0) return 0; - const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float); - if (row_bytes == 0 || out->bytes < row_bytes || out->bytes % row_bytes != 0 || - pre->bytes < out->bytes || - scale_offset > model_size || sizeof(float) > model_size - scale_offset || - base_offset > model_size || row_bytes > model_size - base_offset) { - return 0; - } - const uint64_t n_tokens = out->bytes / row_bytes; - const int logical_tier = ds4_tensor_device_idx(out); - const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, sizeof(float), logical_tier, "output_hc_scale"); - const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, row_bytes, logical_tier, "output_hc_base"); - if (!scale || !base) return 0; - uint64_t n = n_tokens * n_hc; - output_hc_weights_kernel<<<(n + 255) / 256, 256>>>( - (float *)out->ptr, - (const float *)pre->ptr, - scale, - base, - n_hc, - (uint32_t)n_tokens, - eps); - return cuda_ok(cudaGetLastError(), "output hc weights launch"); -} -extern "C" int ds4_gpu_hc_expand_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) { - if (!out_hc || !block_out || !residual_hc || !post || !comb || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, - (const float *)block_out->ptr, - (const float *)block_out->ptr, - (const float *)block_out->ptr, - (const float *)residual_hc->ptr, - (const float *)post->ptr, - (const float *)comb->ptr, - n_embd, n_hc, n_tokens, - n_hc, n_hc * n_hc, 0, 0); - return cuda_ok(cudaGetLastError(), "hc_expand launch"); -} -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) { - if (!out_hc || !block_out || !block_add || !residual_hc || !post || !comb || - n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, - (const float *)block_out->ptr, - (const float *)block_add->ptr, - (const float *)block_out->ptr, - (const float *)residual_hc->ptr, - (const float *)post->ptr, - (const float *)comb->ptr, - n_embd, n_hc, n_tokens, - n_hc, n_hc * n_hc, 1, 0); - return cuda_ok(cudaGetLastError(), "hc_expand_add launch"); -} -extern "C" int ds4_gpu_hc_expand_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { - if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); - uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - const float *base = (const float *)split->ptr; - hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, - (const float *)block_out->ptr, - (const float *)block_out->ptr, - (const float *)block_out->ptr, - (const float *)residual_hc->ptr, - base + n_hc, - base + 2u * n_hc, - n_embd, n_hc, n_tokens, - mix_hc, mix_hc, 0, 0); - return cuda_ok(cudaGetLastError(), "hc_expand_split launch"); -} -extern "C" int ds4_gpu_hc_expand_add_split_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 *split, uint32_t n_embd, uint32_t n_hc) { - if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); - uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - const float *base = (const float *)split->ptr; - hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, - (const float *)block_out->ptr, - (const float *)block_add->ptr, - (const float *)block_out->ptr, - (const float *)residual_hc->ptr, - base + n_hc, - base + 2u * n_hc, - n_embd, n_hc, n_tokens, - mix_hc, mix_hc, 1, 0); - return cuda_ok(cudaGetLastError(), "hc_expand_add_split launch"); -} - -extern "C" int ds4_gpu_hc_expand_add2_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *block_add2, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { - if (!out_hc || !block_out || !block_add || !block_add2 || - !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); - uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; - uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; - const float *base = (const float *)split->ptr; - hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, - (const float *)block_out->ptr, - (const float *)block_add->ptr, - (const float *)block_add2->ptr, - (const float *)residual_hc->ptr, - base + n_hc, - base + 2u * n_hc, - n_embd, n_hc, n_tokens, - mix_hc, mix_hc, 1, 1); - return cuda_ok(cudaGetLastError(), "hc_expand_add2_split launch"); -} - -extern "C" int ds4_gpu_shared_down_hc_expand_q8_0_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *shared_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 *shared_mid, - const ds4_gpu_tensor *routed_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { - return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out, - model_map, model_size, - weight_offset, - in_dim, out_dim, - shared_mid, - routed_out, - NULL, - NULL, NULL, NULL, 0, - residual_hc, - split, - n_embd, n_hc, - "shared_down_hc_expand"); - } - return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size, - weight_offset, in_dim, out_dim, - shared_mid, 1) && - ds4_gpu_hc_expand_add_split_tensor(out_hc, shared_out, routed_out, - residual_hc, split, n_embd, n_hc); -} - -extern "C" int ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *shared_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 *shared_mid, - const ds4_gpu_tensor *routed_out, - const ds4_gpu_tensor *routed_add, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { - return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out, - model_map, model_size, - weight_offset, - in_dim, out_dim, - shared_mid, - routed_out, - routed_add, - NULL, NULL, NULL, 0, - residual_hc, - split, - n_embd, n_hc, - "shared_down_hc_expand_add"); - } - return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size, - weight_offset, in_dim, out_dim, - shared_mid, 1) && - ds4_gpu_hc_expand_add2_split_tensor(out_hc, shared_out, routed_out, - routed_add, residual_hc, split, - n_embd, n_hc); -} - -extern "C" int ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *shared_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 *shared_mid, - const ds4_gpu_tensor *home_slots, - const ds4_gpu_tensor *peer_packed, - const ds4_gpu_tensor *selected, - uint32_t expert_split, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") != NULL) return 0; - return cuda_matmul_q8_0_hc_expand_tensor_labeled( - out_hc, - shared_out, - model_map, - model_size, - weight_offset, - in_dim, - out_dim, - shared_mid, - NULL, - NULL, - home_slots, - peer_packed, - selected, - expert_split, - residual_hc, - split, - n_embd, - n_hc, - "shared_down_hc_expand_owned"); -} - -extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *block_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, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { - return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, block_out, - model_map, model_size, - weight_offset, - in_dim, out_dim, - x, - NULL, - NULL, - NULL, NULL, NULL, 0, - residual_hc, - split, - n_embd, n_hc, - "q8_hc_expand"); - } - return ds4_gpu_matmul_q8_0_tensor(block_out, model_map, model_size, - weight_offset, in_dim, out_dim, x, 1) && - ds4_gpu_hc_expand_split_tensor(out_hc, block_out, residual_hc, - split, n_embd, n_hc); -} - -/* --gpu-vram auto probe. Defined here (in the .cu unit) so the - * C-side parser (ds4_gpu_args.c) does not need to include - * . Returns 0 on success, nonzero on error - * (errbuf populated). See ds4_gpu_args.h. - * - * Side-effect-light: changes cudaSetDevice during probing; callers - * that care about the active device should reset it themselves - * before continuing. (The mgpu init path resets it anyway.) */ -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; - cudaError_t rc = cudaGetDeviceCount(&visible); - if (rc != cudaSuccess || visible <= 0) { - if (errbuf && errbuflen) { - snprintf(errbuf, errbuflen, - "cudaGetDeviceCount failed: %s", - rc == cudaSuccess ? "no devices" : cudaGetErrorString(rc)); - } - return 1; - } - /* Build the device list: either the explicit filter or 0..visible-1. */ - int devs[DS4_MAX_GPUS]; - int n_dev = 0; - if (device_filter && filter_len > 0) { - if (filter_len > DS4_MAX_GPUS) { - if (errbuf && errbuflen) { - snprintf(errbuf, errbuflen, - "--gpu-devices filter has %d entries (max %d)", - filter_len, DS4_MAX_GPUS); - } - return 1; - } - for (int i = 0; i < filter_len; i++) { - int d = device_filter[i]; - if (d < 0 || d >= visible) { - if (errbuf && errbuflen) { - snprintf(errbuf, errbuflen, - "--gpu-devices: device %d not in 0..%d", - d, visible - 1); - } - return 1; - } - devs[n_dev++] = d; - } - } else { - int cap = visible < DS4_MAX_GPUS ? visible : DS4_MAX_GPUS; - for (int i = 0; i < cap; i++) devs[n_dev++] = i; - } - out->n_gpus = n_dev; - out->safety_margin_bytes = safety_margin_bytes; - for (int i = 0; i < n_dev; i++) { - int d = devs[i]; - rc = cudaSetDevice(d); - if (rc != cudaSuccess) { - if (errbuf && errbuflen) { - snprintf(errbuf, errbuflen, - "cudaSetDevice(%d) failed: %s", - d, cudaGetErrorString(rc)); - } - return 1; - } - size_t free_b = 0, total_b = 0; - rc = cudaMemGetInfo(&free_b, &total_b); - if (rc != cudaSuccess) { - if (errbuf && errbuflen) { - snprintf(errbuf, errbuflen, - "cudaMemGetInfo on device %d failed: %s", - d, cudaGetErrorString(rc)); - } - return 1; - } - /* Auto-mode reserve. Auto-probe is the only place we override - * the user's stated budget, so this is where the conservative- - * on-the-user's-behalf reserve belongs. The - * engine path (engine_classify_multi_tier) still subtracts the - * user-supplied safety_margin_bytes + the cuBLAS workspace from - * whatever budget we hand back; that math is unchanged and - * applies on top of the reserve we trim here. - * - * Reserve = max(2 GiB, 5 % of free). Why these numbers: - * - 2 GiB floor covers runtime scratch / Q8 dequant caches / - * MTP optional state on small GPUs (8-12 GB cards) where - * 5 % is < 1 GiB and not enough headroom. - * - 5 % of free scales the reserve up on larger cards where - * workspace + KV growth needs proportionally more room. - * Explicit --gpu-vram 47,37 budgets do not go through this - * probe and are unaffected. */ - const size_t reserve_floor = (size_t)2ull * 1024ull * 1024ull * 1024ull; - const size_t reserve_pct = free_b / 20u; - const size_t reserve = reserve_floor > reserve_pct ? reserve_floor : reserve_pct; - const size_t budget = free_b > reserve ? (free_b - reserve) : 0; - (void)safety_margin_bytes; - out->device_indices[i] = d; - out->vram_bytes[i] = budget; - } - return 0; -} - -typedef struct ds4_gpu_stream_expert_table { - const void *model_map; - uint64_t model_size; - uint32_t layer; - uint32_t n_total_expert; - uint64_t gate_offset; - uint64_t up_offset; - uint64_t down_offset; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; -} ds4_gpu_stream_expert_table; - -static int cuda_stream_selected_ensure_bytes( - char **ptr, uint64_t *capacity, uint64_t bytes, const char *label) { - if (*ptr && *capacity >= bytes) return 1; - if (*ptr) { - (void)cudaFree(*ptr); - *ptr = NULL; - *capacity = 0; - } - if (bytes == 0 || bytes > (uint64_t)SIZE_MAX) return 0; - cudaError_t err = cudaMalloc((void **)ptr, (size_t)bytes); - if (err != cudaSuccess) { - fprintf(stderr, "ds4: CUDA streaming %s allocation failed for %.2f MiB: %s\n", - label, (double)bytes / 1048576.0, cudaGetErrorString(err)); - (void)cudaGetLastError(); - return 0; - } - *capacity = bytes; - return 1; -} - -static int cuda_stream_selected_ensure_i32(uint64_t count) { - if (count == 0 || count > UINT64_MAX / sizeof(int32_t)) return 0; - const uint64_t bytes = count * sizeof(int32_t); - return cuda_stream_selected_ensure_bytes( - (char **)&g_stream_selected_cache.slot_selected_ptr, - &g_stream_selected_cache.slot_selected_capacity, - bytes, - "selected-id remap"); -} - -static int cuda_stream_selected_ranges_valid( - const ds4_gpu_stream_expert_table *table) { - if (!table || !table->model_map || table->model_size == 0 || - table->n_total_expert == 0 || table->gate_expert_bytes == 0 || - table->down_expert_bytes == 0) { - return 0; - } - if ((uint64_t)table->n_total_expert > - UINT64_MAX / table->gate_expert_bytes || - (uint64_t)table->n_total_expert > - UINT64_MAX / table->down_expert_bytes) { - return 0; - } - const uint64_t gate_bytes = - (uint64_t)table->n_total_expert * table->gate_expert_bytes; - const uint64_t down_bytes = - (uint64_t)table->n_total_expert * table->down_expert_bytes; - return table->gate_offset <= table->model_size && - gate_bytes <= table->model_size - table->gate_offset && - table->up_offset <= table->model_size && - gate_bytes <= table->model_size - table->up_offset && - table->down_offset <= table->model_size && - down_bytes <= table->model_size - table->down_offset; -} - -static int cuda_stream_selected_cache_begin_load( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t slot_count) { - cuda_stream_selected_cache_invalidate(); - if (!g_ssd_streaming_mode) return 1; - if (!cuda_stream_selected_ranges_valid(table) || !selected_ids || - slot_count == 0) { - return 0; - } - if (g_n_gpus != 1) { - fprintf(stderr, - "ds4: CUDA SSD streaming requires single-GPU placement\n"); - return 0; - } - - std::vector expert_to_slot; - std::vector compact_ids; - std::vector slot_ids; - try { - expert_to_slot.assign(table->n_total_expert, -1); - compact_ids.reserve(slot_count < table->n_total_expert ? - slot_count : table->n_total_expert); - slot_ids.resize(slot_count); - } catch (...) { - return 0; - } - for (uint32_t i = 0; i < slot_count; i++) { - const int32_t expert = selected_ids[i]; - if (expert < 0 || (uint32_t)expert >= table->n_total_expert) { - fprintf(stderr, - "ds4: CUDA streaming expert id %d is outside 0..%u at layer %u\n", - expert, table->n_total_expert, table->layer); - return 0; - } - int32_t compact = expert_to_slot[(uint32_t)expert]; - if (compact < 0) { - compact = (int32_t)compact_ids.size(); - expert_to_slot[(uint32_t)expert] = compact; - compact_ids.push_back(expert); - } - slot_ids[i] = compact; - } - if (compact_ids.empty() || compact_ids.size() > UINT32_MAX) return 0; - const uint64_t compact_count = compact_ids.size(); - if (compact_count > UINT64_MAX / table->gate_expert_bytes || - compact_count > UINT64_MAX / table->down_expert_bytes) { - return 0; - } - const uint64_t gate_bytes = compact_count * table->gate_expert_bytes; - const uint64_t down_bytes = compact_count * table->down_expert_bytes; - const int logical_tier = 0; - if (g_stream_selected_cache.logical_tier != logical_tier && - (g_stream_selected_cache.gate_ptr || - g_stream_selected_cache.up_ptr || - g_stream_selected_cache.down_ptr || - g_stream_selected_cache.slot_selected_ptr)) { - cuda_stream_selected_cache_release(); - } - if (ds4_gpu_set_current_device(logical_tier) != 0 || - !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.gate_ptr, - &g_stream_selected_cache.gate_capacity, - gate_bytes, "gate experts") || - !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.up_ptr, - &g_stream_selected_cache.up_capacity, - gate_bytes, "up experts") || - !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.down_ptr, - &g_stream_selected_cache.down_capacity, - down_bytes, "down experts") || - !cuda_stream_selected_ensure_i32(slot_count)) { - cuda_stream_selected_cache_invalidate(); - return 0; - } - - for (uint32_t i = 0; i < compact_ids.size(); i++) { - const uint64_t expert = (uint32_t)compact_ids[i]; - const uint64_t gate_src = - table->gate_offset + expert * table->gate_expert_bytes; - const uint64_t up_src = - table->up_offset + expert * table->gate_expert_bytes; - const uint64_t down_src = - table->down_offset + expert * table->down_expert_bytes; - const uint64_t gate_dst = (uint64_t)i * table->gate_expert_bytes; - const uint64_t down_dst = (uint64_t)i * table->down_expert_bytes; - if (!cuda_model_copy_to_device_streamed( - g_stream_selected_cache.gate_ptr + gate_dst, - table->model_map, table->model_size, - gate_src, table->gate_expert_bytes, - "stream gate expert copy") || - !cuda_model_copy_to_device_streamed( - g_stream_selected_cache.up_ptr + gate_dst, - table->model_map, table->model_size, - up_src, table->gate_expert_bytes, - "stream up expert copy") || - !cuda_model_copy_to_device_streamed( - g_stream_selected_cache.down_ptr + down_dst, - table->model_map, table->model_size, - down_src, table->down_expert_bytes, - "stream down expert copy")) { - cuda_stream_selected_cache_invalidate(); - return 0; - } - } - if (!cuda_ok(cudaMemcpy(g_stream_selected_cache.slot_selected_ptr, - slot_ids.data(), - (size_t)slot_count * sizeof(int32_t), - cudaMemcpyHostToDevice), - "stream selected-id remap copy")) { - cuda_stream_selected_cache_invalidate(); - return 0; - } - - g_stream_selected_cache.logical_tier = logical_tier; - g_stream_selected_cache.model_map = table->model_map; - g_stream_selected_cache.layer = table->layer; - g_stream_selected_cache.n_total_expert = table->n_total_expert; - g_stream_selected_cache.slot_count = slot_count; - g_stream_selected_cache.compact_count = (uint32_t)compact_count; - g_stream_selected_cache.gate_offset = table->gate_offset; - g_stream_selected_cache.up_offset = table->up_offset; - g_stream_selected_cache.down_offset = table->down_offset; - g_stream_selected_cache.gate_expert_bytes = table->gate_expert_bytes; - g_stream_selected_cache.down_expert_bytes = table->down_expert_bytes; - g_stream_selected_cache.slot_selected_tensor.ptr = - g_stream_selected_cache.slot_selected_ptr; - g_stream_selected_cache.slot_selected_tensor.bytes = - (uint64_t)slot_count * sizeof(int32_t); - g_stream_selected_cache.slot_selected_tensor.owner = 0; - g_stream_selected_cache.slot_selected_tensor.device_id = logical_tier; - g_stream_selected_cache.valid = 1; - return 1; -} - -__device__ __forceinline__ static float glm_rope_yarn_corr_factor_dev( - int n_dims, int n_ctx_orig, float n_rot, float base) { - return n_dims * logf(n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) / - (2.0f * logf(base)); -} -__device__ __forceinline__ static float glm_rope_yarn_ramp_dev( - float low, float high, int i0) { - const float y = (i0 / 2 - low) / fmaxf(0.001f, high - low); - return 1.0f - fminf(1.0f, fmaxf(0.0f, y)); -} -__device__ __forceinline__ static void glm_rope_yarn_dev( - float theta_extrap, float freq_scale, const float corr_dims[2], - int i0, float ext_factor, float mscale, - float *cos_theta, float *sin_theta) { - float theta_interp = freq_scale * theta_extrap; - float theta = theta_interp; - if (ext_factor != 0.0f) { - float ramp_mix = glm_rope_yarn_ramp_dev(corr_dims[0], corr_dims[1], i0) * - ext_factor; - theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; - mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - *cos_theta = cosf(theta) * mscale; - *sin_theta = sinf(theta) * mscale; -} - -static int cuda_current_tier(void) { - int dev = 0; - if (cudaGetDevice(&dev) != cudaSuccess) return 0; - return dev; -} - -/* ===== GLM 5.2 stubs (to be implemented; fail loudly) ===== */ -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-parameter" -__global__ static void add3_kernel(float *out, const float *a, - const float *b, const float *c, - uint32_t n) { - uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) out[i] = a[i] + b[i] + c[i]; -} - -extern "C" int ds4_gpu_add3_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *a, - const ds4_gpu_tensor *b, - const ds4_gpu_tensor *c, - uint32_t n) { - if (!out || !a || !b || !c || n == 0 || - out->bytes < (uint64_t)n * sizeof(float) || - a->bytes < (uint64_t)n * sizeof(float) || - b->bytes < (uint64_t)n * sizeof(float) || - c->bytes < (uint64_t)n * sizeof(float)) { - return 0; - } - add3_kernel<<<(n + 255) / 256, 256>>>( - (float *)out->ptr, (const float *)a->ptr, - (const float *)b->ptr, (const float *)c->ptr, n); - return cuda_ok(cudaGetLastError(), "add3 launch"); -} - -/* Fused decode residual: sum_out = a + b; norm_out = rmsnorm(sum) * w. - * Single row, one block (two-pass over n with a shared reduction). */ -__global__ static void glm_add_rms_norm_weight_kernel( - float *norm_out, - float *sum_out, - const float *a, - const float *b, - const float *w, - uint32_t n, - float eps) { - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - __shared__ float sh[32]; - float sumsq = 0.0f; - for (uint32_t i = tid; i < n; i += nth) { - const float v = a[i] + b[i]; - sum_out[i] = v; - sumsq += v * v; - } - for (int off = 16; off > 0; off >>= 1) { - sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off); - } - if ((tid & 31u) == 0u) sh[tid >> 5] = sumsq; - __syncthreads(); - if (tid < 32u) { - sumsq = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) { - sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off); - } - if (tid == 0u) sh[0] = sumsq; - } - __syncthreads(); - const float scale = rsqrtf(sh[0] / (float)n + eps); - for (uint32_t i = tid; i < n; i += nth) { - norm_out[i] = (sum_out[i] * scale) * w[i]; - } -} - -extern "C" int ds4_gpu_add_rms_norm_weight_tensor( - ds4_gpu_tensor *norm_out, - ds4_gpu_tensor *sum_out, - const ds4_gpu_tensor *a, - const ds4_gpu_tensor *b, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n, - float eps) { - if (!norm_out || !sum_out || !a || !b || !model_map || n == 0 || - norm_out->bytes < (uint64_t)n * sizeof(float) || - sum_out->bytes < (uint64_t)n * sizeof(float) || - a->bytes < (uint64_t)n * sizeof(float) || - b->bytes < (uint64_t)n * sizeof(float) || - weight_offset > model_size || - (uint64_t)n * sizeof(float) > model_size - weight_offset) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const float *w = (const float *)cuda_resolve_weight_ptr( - model_map, weight_offset, (uint64_t)n * sizeof(float), - logical_tier, "rms_weight"); - if (!w) return 0; - glm_add_rms_norm_weight_kernel<<<1, 1024>>>( - (float *)norm_out->ptr, (float *)sum_out->ptr, - (const float *)a->ptr, (const float *)b->ptr, w, n, eps); - return cuda_ok(cudaGetLastError(), "add rms norm weight"); -} - -extern "C" bool ds4_gpu_commands_active(void) { - return false; -} - -__global__ static void glm_embed_token_q8_0_kernel( - float *out, - const unsigned char *w, - uint32_t token, - uint32_t n_embd) { - uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; - if (d >= n_embd) return; - const uint64_t row_blocks = n_embd / 32u; - const unsigned char *blk = - w + ((uint64_t)token * row_blocks + (d >> 5)) * 34u; - const float scale = __half2float(*(const __half *)blk); - out[d] = scale * (float)((const int8_t *)(blk + 2))[d & 31u]; -} - -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 (!out || !model_map || n_embd == 0 || (n_embd & 31u) != 0u || - token >= n_vocab) { - return 0; - } - if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ - fprintf(stderr, "ds4: embed_token_quant: unsupported type %u\n", - weight_type); - return 0; - } - const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; - if (weight_offset > model_size || - (uint64_t)n_vocab * row_bytes > model_size - weight_offset || - out->bytes < (uint64_t)n_embd * sizeof(float)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr( - model_map, weight_offset, (uint64_t)n_vocab * row_bytes, - logical_tier, "glm_token_embd"); - if (!w) return 0; - glm_embed_token_q8_0_kernel<<<(n_embd + 255) / 256, 256>>>( - (float *)out->ptr, w, token, n_embd); - return cuda_ok(cudaGetLastError(), "glm embed token launch"); -} - -__global__ static void glm_embed_tokens_q8_0_kernel( - float *out, - const int32_t *tokens, - const unsigned char *w, - uint32_t n_tokens, - uint32_t n_embd) { - uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - uint64_t n = (uint64_t)n_tokens * n_embd; - if (gid >= n) return; - uint32_t t = gid / n_embd; - uint32_t d = gid - (uint64_t)t * n_embd; - int32_t tok = tokens[t]; - const uint64_t row_blocks = n_embd / 32u; - const unsigned char *blk = w + ((uint64_t)tok * row_blocks + (d >> 5)) * 34u; - const float scale = __half2float(*(const __half *)blk); - out[gid] = scale * (float)((const int8_t *)(blk + 2))[d & 31u]; -} - -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 (!out || !tokens || !model_map || n_tokens == 0 || n_embd == 0 || - (n_embd & 31u) != 0u) { - return 0; - } - if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ - fprintf(stderr, "ds4: embed_tokens_quant: unsupported type %u\n", - weight_type); - return 0; - } - const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; - if (weight_offset > model_size || - (uint64_t)n_vocab * row_bytes > model_size - weight_offset || - out->bytes < (uint64_t)n_tokens * n_embd * sizeof(float) || - tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr( - model_map, weight_offset, (uint64_t)n_vocab * row_bytes, - logical_tier, "glm_token_embd"); - if (!w) return 0; - uint64_t n = (uint64_t)n_tokens * n_embd; - glm_embed_tokens_q8_0_kernel<<<(n + 255) / 256, 256>>>( - (float *)out->ptr, - (const int32_t *)tokens->ptr, - w, n_tokens, n_embd); - return cuda_ok(cudaGetLastError(), "glm embed tokens launch"); -} - -extern "C" int ds4_gpu_flush_encoder(void) { - /* Metal encoder flush: CUDA kernels are already queued in stream - * order, nothing to split. */ - return 1; -} - -extern "C" int ds4_gpu_glm_attention_flash_staged_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_flash_staged_tensor\n"); - return 0; -} - -extern "C" int ds4_gpu_glm_attention_flash_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_flash_tensor\n"); - return 0; -} - -extern "C" int ds4_gpu_glm_attention_full_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_full_tensor\n"); - return 0; -} - -template -__device__ __forceinline__ static float2 glm_cache_rope_pair_f16_dev( - const CT *rope_cache, uint64_t rope_base, uint32_t r, - uint32_t row, uint32_t qk_rope, float freq_base, float freq_scale, - float ext_factor, float attn_factor, const float corr_dims[2]) { - const float theta_base = (float)row; - const float inv_ndims = -1.0f / (float)qk_rope; - const float theta = theta_base * powf(freq_base, inv_ndims * (float)r); - float ct, st; - glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)r, - ext_factor, attn_factor, &ct, &st); - const float x0 = (float)rope_cache[rope_base + r]; - const float x1 = (float)rope_cache[rope_base + r + 1u]; - return make_float2(x0 * ct - x1 * st, x0 * st + x1 * ct); -} - -/* Scalar-correct MLA attention: one warp per head, grid - * (ceil(n_head/8), n_tokens). The row loop handles either a contiguous - * causal range or an explicit selected-row list and mirrors the Metal online - * softmax so numerics stay comparable. */ -template -__global__ static void glm_attention_lora_causal_kernel( - float *lora_out, - const float *q, - const float *qk_low, - const CT *kv_lora_cache, - const CT *k_rope_cache, - const uint32_t *selected, - uint32_t cache_cap, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_selected, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_rope, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float scale) { - const uint32_t token = blockIdx.y; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t head = blockIdx.x * 8u + warp; - if (token >= n_tokens || head >= n_head || kv_lora_dim != 512u || - qk_rope != 64u) { - return; - } - const uint32_t visible = selected_rows - ? n_selected : min(n_selected, pos0 + token + 1u); - if (visible == 0u) return; - - const uint32_t qk_dim = qk_nope + qk_rope; - const float *qh = q + (uint64_t)token * n_head * qk_dim + - (uint64_t)head * qk_dim; - const float4 *low4 = (const float4 *)(qk_low + - (uint64_t)token * n_head * kv_lora_dim + - (uint64_t)head * kv_lora_dim); - - float4 low0 = low4[lane]; - float4 low1 = low4[lane + 32u]; - float4 low2 = low4[lane + 64u]; - float4 low3 = low4[lane + 96u]; - float4 qrope = make_float4(0.f, 0.f, 0.f, 0.f); - const uint32_t rope_vecs = qk_rope >> 2; /* 16 */ - if (lane < rope_vecs) { - qrope = *((const float4 *)(qh + qk_nope + lane * 4u)); - } - - float corr_dims[2] = {0.0f, 0.0f}; - if (ext_factor != 0.0f) { - corr_dims[0] = fmaxf(0.0f, - floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig, - beta_fast, freq_base))); - corr_dims[1] = fminf((float)qk_rope - 1.0f, - ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig, - beta_slow, freq_base))); - } - - float M = -FLT_MAX / 2.0f; - float S = 0.0f; - float4 o0 = make_float4(0.f,0.f,0.f,0.f); - float4 o1 = o0, o2 = o0, o3 = o0; - - for (uint32_t ri = 0u; ri < visible; ri++) { - const uint32_t row = selected_rows - ? selected[(uint64_t)token * n_selected + ri] : ri; - if (row >= cache_cap) continue; - const CT *kvrow = kv_lora_cache + (uint64_t)row * kv_lora_dim; - float partial = 0.0f; - { - const float4 k0 = make_float4( - (float)(kvrow[lane*4u+0u]), (float)(kvrow[lane*4u+1u]), - (float)(kvrow[lane*4u+2u]), (float)(kvrow[lane*4u+3u])); - const float4 k1 = make_float4( - (float)(kvrow[(lane+32u)*4u+0u]), (float)(kvrow[(lane+32u)*4u+1u]), - (float)(kvrow[(lane+32u)*4u+2u]), (float)(kvrow[(lane+32u)*4u+3u])); - const float4 k2 = make_float4( - (float)(kvrow[(lane+64u)*4u+0u]), (float)(kvrow[(lane+64u)*4u+1u]), - (float)(kvrow[(lane+64u)*4u+2u]), (float)(kvrow[(lane+64u)*4u+3u])); - const float4 k3 = make_float4( - (float)(kvrow[(lane+96u)*4u+0u]), (float)(kvrow[(lane+96u)*4u+1u]), - (float)(kvrow[(lane+96u)*4u+2u]), (float)(kvrow[(lane+96u)*4u+3u])); - partial += low0.x*k0.x + low0.y*k0.y + low0.z*k0.z + low0.w*k0.w; - partial += low1.x*k1.x + low1.y*k1.y + low1.z*k1.z + low1.w*k1.w; - partial += low2.x*k2.x + low2.y*k2.y + low2.z*k2.z + low2.w*k2.w; - partial += low3.x*k3.x + low3.y*k3.y + low3.z*k3.z + low3.w*k3.w; - if (lane < rope_vecs) { - const uint64_t rope_base = (uint64_t)row * qk_rope; - const uint32_t r = lane * 4u; - const float2 y0 = glm_cache_rope_pair_f16_dev( - k_rope_cache, rope_base, r, row, qk_rope, freq_base, - freq_scale, ext_factor, attn_factor, corr_dims); - const float2 y1 = glm_cache_rope_pair_f16_dev( - k_rope_cache, rope_base, r + 2u, row, qk_rope, - freq_base, freq_scale, ext_factor, attn_factor, - corr_dims); - partial += qrope.x*y0.x + qrope.y*y0.y + - qrope.z*y1.x + qrope.w*y1.y; - } - for (uint32_t off = 16u; off > 0u; off >>= 1u) { - partial += __shfl_xor_sync(0xffffffffu, partial, off); - } - const float score = partial * scale; - const float new_m = fmaxf(M, score); - const float old_scale = expf(M - new_m); - const float row_scale = expf(score - new_m); - o0.x = o0.x*old_scale + k0.x*row_scale; o0.y = o0.y*old_scale + k0.y*row_scale; - o0.z = o0.z*old_scale + k0.z*row_scale; o0.w = o0.w*old_scale + k0.w*row_scale; - o1.x = o1.x*old_scale + k1.x*row_scale; o1.y = o1.y*old_scale + k1.y*row_scale; - o1.z = o1.z*old_scale + k1.z*row_scale; o1.w = o1.w*old_scale + k1.w*row_scale; - o2.x = o2.x*old_scale + k2.x*row_scale; o2.y = o2.y*old_scale + k2.y*row_scale; - o2.z = o2.z*old_scale + k2.z*row_scale; o2.w = o2.w*old_scale + k2.w*row_scale; - o3.x = o3.x*old_scale + k3.x*row_scale; o3.y = o3.y*old_scale + k3.y*row_scale; - o3.z = o3.z*old_scale + k3.z*row_scale; o3.w = o3.w*old_scale + k3.w*row_scale; - S = S*old_scale + row_scale; - M = new_m; - } - } - - const float inv_s = S > 0.0f ? 1.0f / S : 0.0f; - float4 *out4 = (float4 *)(lora_out + - ((uint64_t)token * n_head + head) * kv_lora_dim); - o0.x*=inv_s; o0.y*=inv_s; o0.z*=inv_s; o0.w*=inv_s; - o1.x*=inv_s; o1.y*=inv_s; o1.z*=inv_s; o1.w*=inv_s; - o2.x*=inv_s; o2.y*=inv_s; o2.z*=inv_s; o2.w*=inv_s; - o3.x*=inv_s; o3.y*=inv_s; o3.z*=inv_s; o3.w*=inv_s; - out4[lane] = o0; - out4[lane + 32u] = o1; - out4[lane + 64u] = o2; - out4[lane + 96u] = o3; -} - -extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( - ds4_gpu_tensor *lora_out, - 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, - uint32_t n_tokens, - uint32_t pos0, - 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - n_tokens == 0 || n_head == 0 || kv_lora_dim != 512u || - qk_rope != 64u) { - fprintf(stderr, "ds4: glm attn lora causal: unsupported config " - "(n_tok=%u head=%u lora=%u rope=%u f16=%d)\n", - n_tokens, n_head, kv_lora_dim, qk_rope, (int)cache_f16); - return 0; - } - const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); - dim3 grid((n_head + 7u) / 8u, n_tokens, 1); - if (cache_f16) { - glm_attention_lora_causal_kernel<__half, false><<>>( - (float *)lora_out->ptr, - (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - NULL, cache_cap, - n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope, - qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow, scale); - } else { - glm_attention_lora_causal_kernel<<>>( - (float *)lora_out->ptr, - (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - NULL, cache_cap, - n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope, - qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow, scale); - } - return cuda_ok(cudaGetLastError(), "glm attn lora causal launch"); -} - -extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( - ds4_gpu_tensor *lora_out, - 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 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float); - const uint64_t qk_dim = (uint64_t)qk_nope + qk_rope; - if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - !selected || n_tokens == 0u || n_selected == 0u || n_head == 0u || - kv_lora_dim != 512u || qk_rope != 64u || cache_cap == 0u || - selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t) || - q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || - qk_low->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || - kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem || - k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem || - lora_out->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) { - fprintf(stderr, "ds4: glm attn lora selected: unsupported config " - "(n_tok=%u selected=%u head=%u lora=%u rope=%u f16=%d)\n", - n_tokens, n_selected, n_head, kv_lora_dim, qk_rope, - (int)cache_f16); - return 0; - } - const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); - dim3 grid((n_head + 7u) / 8u, n_tokens, 1); - if (cache_f16) { - glm_attention_lora_causal_kernel<__half, true><<>>( - (float *)lora_out->ptr, - (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, cache_cap, - n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope, - qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow, scale); - } else { - glm_attention_lora_causal_kernel<<>>( - (float *)lora_out->ptr, - (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, cache_cap, - n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope, - qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow, scale); - } - return cuda_ok(cudaGetLastError(), "glm attn lora selected launch"); -} - -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) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_batch_typed_tensor\n"); - return 0; -} - -extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( - ds4_gpu_tensor *heads, - ds4_gpu_tensor *partial_lora, - ds4_gpu_tensor *partial_ms, - 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, - bool selected_rows_valid, - 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, - uint32_t block_rows, - uint32_t n_blocks, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor\n"); - return 0; -} - -__device__ __forceinline__ static float glm_q8_0_dot_row_dev( - const char *row, const float *x, uint32_t n_cols) { - float acc = 0.0f; - const uint32_t nb = n_cols >> 5; - for (uint32_t b = 0; b < nb; b++) { - const char *blk = row + (uint64_t)b * 34u; - const float d = __half2float(*(const __half *)blk); - const int8_t *q = (const int8_t *)(blk + 2); - float s = 0.0f; - #pragma unroll 8 - for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * x[b * 32u + k]; - acc += d * s; - } - return acc; -} - -template -__device__ __forceinline__ static float2 glm_cache_value_pair_dev( - const CT *p) { - return make_float2((float)p[0], (float)p[1]); -} - -template <> -__device__ __forceinline__ float2 glm_cache_value_pair_dev<__half>( - const __half *p) { - return __half22float2(*(const __half2 *)p); -} - -template <> -__device__ __forceinline__ float2 glm_cache_value_pair_dev( - const float *p) { - return *(const float2 *)p; -} - -/* Exact staged decode attention. The original fused kernel owns one block per - * head, which leaves more than half of an L40S idle. These stages preserve the - * fused kernel's arithmetic order for every score, softmax lane, lora output, - * and value-projection row while exposing independent rows/dimensions as - * separate blocks. */ -template -__global__ static void glm_attention_decode_weights_staged_kernel( - float *weights, - float *denom, - const float *q, - const float *qk_low, - const CT *kv_lora_cache, - const CT *k_rope_cache, - const uint32_t *selected, - uint32_t n_selected, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_rope, - float scale, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool score_vec2) { - const uint32_t head = blockIdx.x; - const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u; - const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); - const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u); - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - if (head >= n_head || row_count == 0u) return; - const uint32_t qk_dim = qk_nope + qk_rope; - extern __shared__ float glm_dec_stage_sh[]; - float *red = glm_dec_stage_sh; - float *scores = glm_dec_stage_sh + 256u; - const float *qh = q + - ((uint64_t)token * n_head + head) * qk_dim; - const float *low = qk_low + - ((uint64_t)token * n_head + head) * kv_lora_dim; - - float corr_dims[2] = {0.0f, 0.0f}; - if (ext_factor != 0.0f) { - corr_dims[0] = fmaxf(0.0f, - floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, - (int)n_ctx_orig, beta_fast, freq_base))); - corr_dims[1] = fminf((float)qk_rope - 1.0f, - ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, - (int)n_ctx_orig, beta_slow, freq_base))); - } - - float local_max = -FLT_MAX; - for (uint32_t s = tid; s < row_count; s += nth) { - const uint32_t row = RANGE_TOK2 ? s : selected[s]; - float score = -FLT_MAX; - if (row < cache_cap) { - float dotv = 0.0f; - const uint64_t lora_base = (uint64_t)row * kv_lora_dim; - if (score_vec2) { - for (uint32_t j = 0; j < kv_lora_dim; j += 2u) { - const float2 x = *(const float2 *)(low + j); - const float2 y = glm_cache_value_pair_dev( - kv_lora_cache + lora_base + j); - dotv += x.x * y.x; - dotv += x.y * y.y; - } - } else { - for (uint32_t j = 0; j < kv_lora_dim; j++) { - dotv += low[j] * (float)kv_lora_cache[lora_base + j]; - } - } - const uint64_t rope_base = (uint64_t)row * qk_rope; - for (uint32_t r = 0; r < qk_rope; r += 2u) { - const float2 y = glm_cache_rope_pair_f16_dev( - k_rope_cache, rope_base, r, row, qk_rope, - freq_base, freq_scale, ext_factor, attn_factor, - corr_dims); - dotv += qh[qk_nope + r] * y.x + - qh[qk_nope + r + 1u] * y.y; - } - score = dotv * scale; - } - scores[s] = score; - local_max = fmaxf(local_max, score); - } - red[tid] = local_max; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]); - __syncthreads(); - } - const float max_score = red[0]; - __syncthreads(); - - float local_sum = 0.0f; - for (uint32_t s = tid; s < row_count; s += nth) { - const float w = expf(scores[s] - max_score); - scores[s] = w; - local_sum += w; - } - red[tid] = local_sum; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) red[tid] += red[tid + step]; - __syncthreads(); - } - const uint64_t head_index = (uint64_t)token * n_head + head; - if (tid == 0u) denom[head_index] = fmaxf(red[0], 1.0e-20f); - float *head_weights = weights + head_index * score_stride; - for (uint32_t s = tid; s < row_count; s += nth) { - head_weights[s] = scores[s]; - } -} - -template -__global__ static void glm_attention_decode_lora_staged_kernel( - float *lora_sum, - const float *scores, - const float *denom, - const CT *kv_lora_cache, - const uint32_t *selected, - uint32_t n_selected, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_lora_dim) { - const uint32_t head = blockIdx.y; - const uint32_t token = RANGE_TOK2 ? blockIdx.z : 0u; - const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); - const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u); - const uint32_t pair = blockIdx.x * blockDim.x + threadIdx.x; - const uint32_t j = pair * 2u; - if (head >= n_head || j >= kv_lora_dim) return; - const uint64_t head_index = (uint64_t)token * n_head + head; - const float *head_scores = scores + head_index * score_stride; - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t s = 0; s < row_count; s++) { - const uint32_t row = RANGE_TOK2 ? s : selected[s]; - if (row < cache_cap) { - const float2 v = glm_cache_value_pair_dev( - kv_lora_cache + (uint64_t)row * kv_lora_dim + j); - const float w = head_scores[s]; - acc0 += w * v.x; - acc1 += w * v.y; - } - } - float *out = lora_sum + head_index * kv_lora_dim + j; - out[0] = acc0 / denom[head_index]; - out[1] = acc1 / denom[head_index]; -} - -template -__global__ static void glm_attention_decode_value_staged_kernel( - float *heads, - const float *lora_sum, - const char *value_weight, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t value_dim, - uint32_t value_row_bytes) { - const uint32_t head = blockIdx.y; - const uint32_t token = TOK2 ? blockIdx.z : 0u; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t nwarps = blockDim.x >> 5; - const uint32_t out_warp = blockIdx.x * nwarps + warp; - const uint32_t total_warps = gridDim.x * nwarps; - if (head >= n_head) return; - const float *low = lora_sum + - ((uint64_t)token * n_head + head) * kv_lora_dim; - float *out = heads + - ((uint64_t)token * n_head + head) * value_dim; - const uint32_t nblk = kv_lora_dim >> 5; - for (uint32_t d = out_warp; d < value_dim; d += total_warps) { - const char *row = value_weight + - ((uint64_t)head * value_dim + d) * value_row_bytes; - float acc = 0.0f; - for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) { - const char *b = row + (uint64_t)blk * 34u; - const float dscale = __half2float(*(const __half *)b); - const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u; - const float *xs = low + blk * 32u + (lane & 1u) * 16u; - float s = 0.0f; - #pragma unroll - for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k]; - acc += dscale * s; - } - for (int off = 16; off > 0; off >>= 1) { - acc += __shfl_down_sync(0xffffffffu, acc, off); - } - if (lane == 0u) out[d] = acc; - } -} - -/* Single-token indexed MLA decode attention, one block per head. - * Fuses score (qk_low . kv_lora + q_rope . rope(k_rope@row)), softmax over - * the indexer-selected rows, the weighted kv_lora sum, and the per-head - * value projection (q8_0). Dynamic shared: red[256] + scores[n_selected] + - * lora_sum[kv_lora_dim]. */ -template -__global__ static void glm_attention_indexed_decode_kernel( - float *heads, - const float *q, - const float *qk_low, - const CT *kv_lora_cache, - const CT *k_rope_cache, - const char *value_weight, - const uint32_t *selected, - uint32_t n_selected, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_rope, - uint32_t value_dim, - uint32_t value_row_bytes, - bool lora_vec2, - bool score_vec2, - float scale, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - const uint32_t head = blockIdx.x; - const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u; - const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - if (head >= n_head || row_count == 0u) return; - const uint32_t qk_dim = qk_nope + qk_rope; - - extern __shared__ float glm_dec_sh[]; - float *red = glm_dec_sh; - float *scores = glm_dec_sh + 256u; - float *lora_sum = scores + row_count; - - const float *qh = q + - ((uint64_t)token * n_head + head) * qk_dim; - const float *low = qk_low + - ((uint64_t)token * n_head + head) * kv_lora_dim; - - float corr_dims[2] = {0.0f, 0.0f}; - if (ext_factor != 0.0f) { - corr_dims[0] = fmaxf(0.0f, - floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, - (int)n_ctx_orig, beta_fast, freq_base))); - corr_dims[1] = fminf((float)qk_rope - 1.0f, - ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, - (int)n_ctx_orig, beta_slow, freq_base))); - } - - float local_max = -FLT_MAX; - for (uint32_t s = tid; s < row_count; s += nth) { - const uint32_t row = RANGE_TOK2 ? s : selected[s]; - float score = -FLT_MAX; - if (row < cache_cap) { - float dotv = 0.0f; - const uint64_t lora_base = (uint64_t)row * kv_lora_dim; - if (score_vec2) { - for (uint32_t j = 0; j < kv_lora_dim; j += 2u) { - const float2 x = *(const float2 *)(low + j); - const float2 y = glm_cache_value_pair_dev( - kv_lora_cache + lora_base + j); - dotv += x.x * y.x; - dotv += x.y * y.y; - } - } else { - for (uint32_t j = 0; j < kv_lora_dim; j++) { - dotv += low[j] * (float)kv_lora_cache[lora_base + j]; - } - } - const uint64_t rope_base = (uint64_t)row * qk_rope; - for (uint32_t r = 0; r < qk_rope; r += 2u) { - const float2 y = glm_cache_rope_pair_f16_dev( - k_rope_cache, rope_base, r, row, qk_rope, - freq_base, freq_scale, ext_factor, attn_factor, - corr_dims); - dotv += qh[qk_nope + r] * y.x + qh[qk_nope + r + 1u] * y.y; - } - score = dotv * scale; - } - scores[s] = score; - local_max = fmaxf(local_max, score); - } - red[tid] = local_max; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]); - __syncthreads(); - } - const float max_score = red[0]; - __syncthreads(); - - float local_sum = 0.0f; - for (uint32_t s = tid; s < row_count; s += nth) { - const float w = expf(scores[s] - max_score); - scores[s] = w; - local_sum += w; - } - red[tid] = local_sum; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) red[tid] += red[tid + step]; - __syncthreads(); - } - const float denom = fmaxf(red[0], 1.0e-20f); - __syncthreads(); - - if (lora_vec2) { - for (uint32_t j = tid * 2u; j < kv_lora_dim; j += nth * 2u) { - float acc0 = 0.0f; - float acc1 = 0.0f; - for (uint32_t s = 0; s < row_count; s++) { - const uint32_t row = RANGE_TOK2 ? s : selected[s]; - if (row < cache_cap) { - const float2 v = glm_cache_value_pair_dev( - kv_lora_cache + (uint64_t)row * kv_lora_dim + j); - const float w = scores[s]; - acc0 += w * v.x; - acc1 += w * v.y; - } - } - lora_sum[j] = acc0 / denom; - lora_sum[j + 1u] = acc1 / denom; - } - } else { - for (uint32_t j = tid; j < kv_lora_dim; j += nth) { - float acc = 0.0f; - for (uint32_t s = 0; s < row_count; s++) { - const uint32_t row = RANGE_TOK2 ? s : selected[s]; - if (row < cache_cap) { - acc += scores[s] * - (float)kv_lora_cache[(uint64_t)row * kv_lora_dim + j]; - } - } - lora_sum[j] = acc / denom; - } - } - __syncthreads(); - - float *out = heads + - ((uint64_t)token * n_head + head) * value_dim; - /* Warp-cooperative value projection: one warp per output dim, two - * lanes per q8_0 block (16 cols each). */ - const uint32_t nwarps = nth >> 5; - const uint32_t warp = tid >> 5; - const uint32_t lane = tid & 31u; - const uint32_t nblk = kv_lora_dim >> 5; - for (uint32_t d = warp; d < value_dim; d += nwarps) { - const char *row = value_weight + - ((uint64_t)head * value_dim + d) * value_row_bytes; - float acc = 0.0f; - for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) { - const char *b = row + (uint64_t)blk * 34u; - const float dscale = __half2float(*(const __half *)b); - const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u; - const float *xs = lora_sum + blk * 32u + (lane & 1u) * 16u; - float s = 0.0f; - #pragma unroll - for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k]; - acc += dscale * s; - } - for (int off = 16; off > 0; off >>= 1) { - acc += __shfl_down_sync(0xffffffffu, acc, off); - } - if (lane == 0u) out[d] = acc; - } -} - -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) { - const uint32_t qk_dim = qk_nope + qk_rope; - if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - !model_map || !selected || - n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - n_head == 0 || kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0u || - value_dim == 0) { - return 0; - } - if (value_weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ - fprintf(stderr, - "ds4: glm indexed decode attention: unsupported value type %u\n", - value_weight_type); - return 0; - } - const uint64_t value_row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u; - const uint64_t value_weight_bytes = - (uint64_t)n_head * value_dim * value_row_bytes; - if (value_weight_offset > model_size || - value_weight_bytes > model_size - value_weight_offset) { - return 0; - } - const uint64_t cache_elem = cache_f16 ? 2u : 4u; - if (heads->bytes < (uint64_t)n_head * value_dim * sizeof(float) || - q->bytes < (uint64_t)n_head * qk_dim * sizeof(float) || - qk_low->bytes < (uint64_t)n_head * kv_lora_dim * sizeof(float) || - kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem || - k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem || - selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const char *vw = cuda_resolve_weight_ptr(model_map, value_weight_offset, - value_weight_bytes, logical_tier, "glm_v_b_decode"); - if (!vw) return 0; - const float scale = 1.0f / sqrtf((float)qk_dim); - const bool score_vec2 = getenv("DS4_GLM_ATTN_NO_SCORE_VEC2") == NULL; - const bool range_tok2 = - g_glm_mtp_verify_mode && - getenv("DS4_GLM_MTP_NO_ATTN_TOK2") == NULL && - n_selected < cache_cap && - heads->bytes >= 2u * (uint64_t)n_head * value_dim * sizeof(float) && - q->bytes >= 2u * (uint64_t)n_head * qk_dim * sizeof(float) && - qk_low->bytes >= - 2u * (uint64_t)n_head * kv_lora_dim * sizeof(float); - if (range_tok2 && n_selected < 512u) { - const bool lora_vec2 = - getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL; - const uint32_t shmem = - (256u + n_selected + 1u + kv_lora_dim) * - (uint32_t)sizeof(float); - const dim3 grid(n_head, 2u, 1u); - if (cache_f16) { - glm_attention_indexed_decode_kernel<__half, true> - <<>>( - (float *)heads->ptr, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - vw, (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, value_dim, - (uint32_t)value_row_bytes, - lora_vec2, score_vec2, scale, n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow); - } else { - glm_attention_indexed_decode_kernel - <<>>( - (float *)heads->ptr, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - vw, (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, value_dim, - (uint32_t)value_row_bytes, - lora_vec2, score_vec2, scale, n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow); - } - return cuda_ok(cudaGetLastError(), - "glm indexed decode attention tok2 range"); - } - if (n_selected >= 512u && - getenv("DS4_GLM_ATTN_NO_STAGED_DECODE") == NULL) { - const uint32_t token_count = range_tok2 ? 2u : 1u; - const uint32_t score_stride = n_selected + (range_tok2 ? 1u : 0u); - const uint64_t head_count = (uint64_t)token_count * n_head; - if (head_count > UINT64_MAX / score_stride || - head_count > UINT64_MAX / kv_lora_dim) { - return 0; - } - const uint64_t score_count = head_count * score_stride; - const uint64_t lora_count = head_count * kv_lora_dim; - if (score_count > UINT64_MAX - head_count - lora_count || - score_count + head_count + lora_count > - UINT64_MAX / sizeof(float)) { - return 0; - } - const uint64_t scratch_bytes = - (score_count + head_count + lora_count) * sizeof(float); - float *scratch = (float *)cuda_tmp_alloc_on( - ds4_tensor_device_idx(heads), scratch_bytes, - "glm staged decode attention"); - if (!scratch) return 0; - float *softmax_denom = scratch + score_count; - float *lora_sum = softmax_denom + head_count; - const uint32_t weight_shmem = - (256u + score_stride) * (uint32_t)sizeof(float); - const dim3 weight_grid(n_head, token_count, 1u); - if (cache_f16 && range_tok2) { - glm_attention_decode_weights_staged_kernel<__half, true> - <<>>( - scratch, softmax_denom, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, scale, n_ctx_orig, freq_base, - freq_scale, ext_factor, attn_factor, beta_fast, - beta_slow, score_vec2); - } else if (cache_f16) { - glm_attention_decode_weights_staged_kernel<__half> - <<>>( - scratch, softmax_denom, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, scale, n_ctx_orig, freq_base, - freq_scale, ext_factor, attn_factor, beta_fast, - beta_slow, score_vec2); - } else if (range_tok2) { - glm_attention_decode_weights_staged_kernel - <<>>( - scratch, softmax_denom, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, scale, n_ctx_orig, freq_base, - freq_scale, ext_factor, attn_factor, beta_fast, - beta_slow, score_vec2); - } else { - glm_attention_decode_weights_staged_kernel - <<>>( - scratch, softmax_denom, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, scale, n_ctx_orig, freq_base, - freq_scale, ext_factor, attn_factor, beta_fast, - beta_slow, score_vec2); - } - if (!cuda_ok(cudaGetLastError(), - "glm staged decode weights launch")) { - return 0; - } - dim3 lora_grid((kv_lora_dim / 2u + 63u) / 64u, - n_head, token_count); - if (cache_f16 && range_tok2) { - glm_attention_decode_lora_staged_kernel<__half, true> - <<>>( - lora_sum, scratch, softmax_denom, - (const __half *)kv_lora_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim); - } else if (cache_f16) { - glm_attention_decode_lora_staged_kernel<__half> - <<>>( - lora_sum, scratch, softmax_denom, - (const __half *)kv_lora_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim); - } else if (range_tok2) { - glm_attention_decode_lora_staged_kernel - <<>>( - lora_sum, scratch, softmax_denom, - (const float *)kv_lora_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim); - } else { - glm_attention_decode_lora_staged_kernel - <<>>( - lora_sum, scratch, softmax_denom, - (const float *)kv_lora_cache->ptr, - (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim); - } - if (!cuda_ok(cudaGetLastError(), - "glm staged decode lora launch")) { - return 0; - } - dim3 value_grid((value_dim + 127u) / 128u, - n_head, token_count); - if (range_tok2) { - glm_attention_decode_value_staged_kernel - <<>>( - (float *)heads->ptr, lora_sum, vw, - n_head, kv_lora_dim, value_dim, - (uint32_t)value_row_bytes); - } else { - glm_attention_decode_value_staged_kernel - <<>>( - (float *)heads->ptr, lora_sum, vw, - n_head, kv_lora_dim, value_dim, - (uint32_t)value_row_bytes); - } - return cuda_ok(cudaGetLastError(), - "glm staged decode value launch"); - } - const bool lora_vec2 = getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL; - const uint32_t shmem = - (256u + n_selected + kv_lora_dim) * (uint32_t)sizeof(float); - if (cache_f16) { - glm_attention_indexed_decode_kernel<__half><<>>( - (float *)heads->ptr, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const __half *)kv_lora_cache->ptr, - (const __half *)k_rope_cache->ptr, - vw, (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes, - lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - } else { - glm_attention_indexed_decode_kernel<<>>( - (float *)heads->ptr, (const float *)q->ptr, - (const float *)qk_low->ptr, - (const float *)kv_lora_cache->ptr, - (const float *)k_rope_cache->ptr, - vw, (const uint32_t *)selected->ptr, - n_selected, cache_cap, n_head, kv_lora_dim, - qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes, - lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - } - return cuda_ok(cudaGetLastError(), "glm indexed decode attention"); -} - -extern "C" int ds4_gpu_glm_build_kv_cache_flash_tensor( - ds4_gpu_tensor *key_cache, - ds4_gpu_tensor *value_cache, - const ds4_gpu_tensor *kv_raw, - const ds4_gpu_tensor *k_nope, - const ds4_gpu_tensor *value, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_raw_dim, - 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, - bool cache_f16) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_build_kv_cache_flash_tensor\n"); - return 0; -} - -extern "C" int ds4_gpu_glm_build_kv_cache_tensor( - ds4_gpu_tensor *key_cache, - ds4_gpu_tensor *value_cache, - const ds4_gpu_tensor *kv_raw, - const ds4_gpu_tensor *k_nope, - const ds4_gpu_tensor *value, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_raw_dim, - 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, - bool cache_f16) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_build_kv_cache_tensor\n"); - return 0; -} - -__global__ static void glm_fill_selected_range_batch_kernel( - uint32_t *selected, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_selected, - uint32_t pad_row) { - uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; - const uint32_t total = n_tokens * n_selected; - if (gid >= total || n_selected == 0u) return; - const uint32_t token = gid / n_selected; - const uint32_t slot = gid - token * n_selected; - const uint32_t visible = pos0 + token + 1u; - selected[gid] = slot < visible ? slot : pad_row; -} - -extern "C" int ds4_gpu_glm_fill_selected_range_batch_tensor( - ds4_gpu_tensor *selected, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_selected, - uint32_t pad_row) { - if (!selected || n_tokens == 0 || n_selected == 0 || - selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t)) { - return 0; - } - const uint64_t total = (uint64_t)n_tokens * n_selected; - glm_fill_selected_range_batch_kernel<<<(unsigned)((total + 255) / 256), 256>>>( - (uint32_t *)selected->ptr, n_tokens, pos0, n_selected, pad_row); - return cuda_ok(cudaGetLastError(), "glm fill selected batch launch"); -} - -__global__ static void glm_fill_selected_range_kernel( - uint32_t *selected, uint32_t n_selected) { - uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; - if (gid < n_selected) selected[gid] = gid; -} - -extern "C" int ds4_gpu_glm_fill_selected_range_tensor( - ds4_gpu_tensor *selected, - uint32_t n_selected) { - if (!selected || n_selected == 0 || - selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { - return 0; - } - glm_fill_selected_range_kernel<<<(n_selected + 255) / 256, 256>>>( - (uint32_t *)selected->ptr, n_selected); - return cuda_ok(cudaGetLastError(), "glm fill selected launch"); -} - -static int glm_rope_tail_offset_launch( - ds4_gpu_tensor *x, - uint32_t n_tokens, uint32_t n_head, uint32_t head_dim, - uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0, - uint32_t n_ctx_orig, float freq_base, float freq_scale, - float ext_factor, float attn_factor, - float beta_fast, float beta_slow, const char *what); - -extern "C" int ds4_gpu_glm_indexer_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim, - rot_dim, 0, pos0, n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow, "glm indexer rope tail"); -} - -template -__global__ static void glm_indexer_scores_f32_kernel( - float *scores, - const float *q, - const float *weights, - const CT *indexer_key_cache, - uint32_t n_rows, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool causal) { - const uint32_t row = blockIdx.x; - const uint32_t token = blockIdx.y; - const uint32_t tid = threadIdx.x; - if (row >= n_rows || token >= n_tokens || tid >= 128u) return; - if (causal && row >= min(n_rows, pos0 + token + 1u)) { - if (tid == 0u) scores[(uint64_t)token * n_rows + row] = -INFINITY; - return; - } - - __shared__ float partial[128]; - float total = 0.0f; - const CT *krow = indexer_key_cache + (uint64_t)row * head_dim; - for (uint32_t h = 0; h < n_head; h++) { - const float *qh = q + - ((uint64_t)token * n_head + h) * head_dim; - float dot = tid < head_dim ? qh[tid] * (float)krow[tid] : 0.0f; - partial[tid] = dot; - __syncthreads(); - for (uint32_t stride = 64u; stride > 0u; stride >>= 1u) { - if (tid < stride) partial[tid] += partial[tid + stride]; - __syncthreads(); - } - if (tid == 0u) { - total += fmaxf(partial[0], 0.0f) * - weights[(uint64_t)token * n_head + h]; - } - __syncthreads(); - } - if (tid == 0u) { - scores[(uint64_t)token * n_rows + row] = total * scale; - } -} - -/* 16-token x 128-row indexer tile. Q and cached K are staged as fp16, - * matching the model's compact-cache precision; each head's MMA result and - * the weighted head reduction remain fp32. */ -template -__global__ static void glm_indexer_scores_wmma128_kernel( - float *scores, - const float *q, - const float *weights, - const CT *indexer_key_cache, - uint32_t n_rows, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool causal) { -#if __CUDA_ARCH__ >= 700 - namespace wmma = nvcuda::wmma; - const uint32_t row0 = blockIdx.x * 128u; - const uint32_t token0 = blockIdx.y * 16u; - const uint32_t tid = threadIdx.x; - const uint32_t warp = tid >> 5u; - if (tid >= 256u || head_dim != 128u) return; - - if (causal) { - const uint32_t last_token = min(token0 + 16u, n_tokens); - const uint32_t max_visible = last_token > token0 - ? min(pos0 + last_token, n_rows) : 0u; - if (row0 >= max_visible) { - for (uint32_t i = tid; i < 16u * 128u; i += 256u) { - const uint32_t token = token0 + (i >> 7u); - const uint32_t row = row0 + (i & 127u); - if (token < n_tokens && row < n_rows) { - scores[(uint64_t)token * n_rows + row] = -INFINITY; - } - } - return; - } - } - - __shared__ __half q_sh[16 * 128]; - __shared__ __half k_sh[128 * 128]; - __shared__ float dot_sh[8 * 16 * 16]; - float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, - 0.0f, 0.0f, 0.0f, 0.0f}; - - for (uint32_t i = tid; i < 128u * 128u; i += 256u) { - const uint32_t r = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t row = row0 + r; - const float v = row < n_rows - ? (float)indexer_key_cache[(uint64_t)row * head_dim + d] - : 0.0f; - k_sh[d + r * 128u] = __float2half(v); - } - __syncthreads(); - - for (uint32_t h = 0; h < n_head; h++) { - for (uint32_t i = tid; i < 16u * 128u; i += 256u) { - const uint32_t tr = i >> 7u; - const uint32_t d = i & 127u; - const uint32_t token = token0 + tr; - const float v = token < n_tokens - ? q[((uint64_t)token * n_head + h) * head_dim + d] - : 0.0f; - q_sh[i] = __float2half(v); - } - __syncthreads(); - - wmma::fragment q_frag; - wmma::fragment k_frag; - wmma::fragment dot_frag; - wmma::fill_fragment(dot_frag, 0.0f); - const uint32_t col0 = warp * 16u; - for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { - wmma::load_matrix_sync(q_frag, q_sh + k0, 128); - wmma::load_matrix_sync(k_frag, - k_sh + col0 * 128u + k0, 128); - wmma::mma_sync(dot_frag, q_frag, k_frag, dot_frag); - } - wmma::store_matrix_sync(dot_sh + warp * 16u * 16u, - dot_frag, 16, wmma::mem_row_major); - __syncthreads(); - - const uint32_t local0 = tid & 255u; - const uint32_t token = token0 + (local0 >> 4u); - const float w = token < n_tokens - ? weights[(uint64_t)token * n_head + h] : 0.0f; - uint32_t slot = 0; - for (uint32_t i = tid; i < 8u * 16u * 16u; - i += 256u, slot++) { - const uint32_t row = row0 + (i >> 8u) * 16u + (i & 15u); - if (token < n_tokens && row < n_rows) { - acc[slot] += fmaxf(dot_sh[i], 0.0f) * w; - } - } - __syncthreads(); - } - - uint32_t slot = 0; - for (uint32_t i = tid; i < 8u * 16u * 16u; - i += 256u, slot++) { - const uint32_t local = i & 255u; - const uint32_t token = token0 + (local >> 4u); - const uint32_t row = row0 + (i >> 8u) * 16u + (local & 15u); - if (token < n_tokens && row < n_rows) { - float out = acc[slot] * scale; - if (causal && row >= pos0 + token + 1u) out = -INFINITY; - scores[(uint64_t)token * n_rows + row] = out; - } - } -#endif -} - -static int glm_indexer_scores_launch( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *indexer_key_cache, - uint32_t n_rows, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool cache_f16, - bool causal) { - const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float); - if (!scores || !q || !weights || !indexer_key_cache || n_rows == 0u || - n_tokens == 0u || n_head == 0u || head_dim != 128u || - q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || - weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) || - indexer_key_cache->bytes < (uint64_t)n_rows * head_dim * cache_elem || - scores->bytes < (uint64_t)n_tokens * n_rows * sizeof(float)) { - return 0; - } - if (!g_quality_mode) { - dim3 grid((n_rows + 127u) / 128u, - (n_tokens + 15u) / 16u, 1); - if (cache_f16) { - glm_indexer_scores_wmma128_kernel<__half><<>>( - (float *)scores->ptr, (const float *)q->ptr, - (const float *)weights->ptr, - (const __half *)indexer_key_cache->ptr, - n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); - } else { - glm_indexer_scores_wmma128_kernel<<>>( - (float *)scores->ptr, (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)indexer_key_cache->ptr, - n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); - } - return cuda_ok(cudaGetLastError(), "glm indexer scores wmma launch"); - } - - dim3 grid(n_rows, n_tokens, 1); - if (cache_f16) { - glm_indexer_scores_f32_kernel<__half><<>>( - (float *)scores->ptr, (const float *)q->ptr, - (const float *)weights->ptr, - (const __half *)indexer_key_cache->ptr, - n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); - } else { - glm_indexer_scores_f32_kernel<<>>( - (float *)scores->ptr, (const float *)q->ptr, - (const float *)weights->ptr, - (const float *)indexer_key_cache->ptr, - n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); - } - return cuda_ok(cudaGetLastError(), "glm indexer scores f32 launch"); -} - -extern "C" int ds4_gpu_glm_indexer_score_one_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *indexer_key_cache, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool cache_f16) { - return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache, - n_rows, 1u, 0u, n_head, head_dim, - scale, cache_f16, false); -} - -extern "C" int ds4_gpu_glm_indexer_scores_batch_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *indexer_key_cache, - uint32_t n_rows, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool cache_f16) { - return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache, - n_rows, n_tokens, pos0, n_head, head_dim, - scale, cache_f16, true); -} - -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) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_k_b_project_typed_tensor\n"); - return 0; -} - -__global__ static void glm_kv_lora_rms_norm_kernel( - float *dst, - const float *src, - const float *w, - uint32_t n_tokens, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - float eps) { - const uint32_t row = blockIdx.x; - if (row >= n_tokens) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - const float *x = src + (uint64_t)row * kv_raw_dim; - float *out = dst + (uint64_t)row * kv_lora_dim; - __shared__ float scratch[256]; - float ss = 0.0f; - for (uint32_t i = tid; i < kv_lora_dim; i += nth) { - const float v = x[i]; - ss += v * v; - } - scratch[tid] = ss; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) scratch[tid] += scratch[tid + step]; - __syncthreads(); - } - const float inv = rsqrtf(scratch[0] / (float)kv_lora_dim + eps); - for (uint32_t i = tid; i < kv_lora_dim; i += nth) { - out[i] = x[i] * inv * w[i]; - } -} - -extern "C" int ds4_gpu_glm_kv_lora_rms_norm_tensor( - ds4_gpu_tensor *dst, - const ds4_gpu_tensor *src, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_tokens, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - float eps) { - if (!dst || !src || !model_map || n_tokens == 0 || - kv_lora_dim == 0 || kv_lora_dim > kv_raw_dim) { - return 0; - } - const uint64_t wb = (uint64_t)kv_lora_dim * sizeof(float); - if (weight_offset > model_size || wb > model_size - weight_offset || - src->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) || - dst->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const float *w = (const float *)cuda_resolve_weight_ptr( - model_map, weight_offset, wb, logical_tier, "glm_kv_lora_norm"); - if (!w) return 0; - glm_kv_lora_rms_norm_kernel<<>>( - (float *)dst->ptr, (const float *)src->ptr, w, - n_tokens, kv_raw_dim, kv_lora_dim, eps); - return cuda_ok(cudaGetLastError(), "glm kv lora rms norm launch"); -} - - - -__global__ static void glm_qk_lowrank_q8_0_batch_kernel( - float *qk_low, - const char *weight, - const float *q, - uint32_t n_tokens, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_dim, - uint64_t row_bytes) { - const uint32_t head = blockIdx.x; - const uint32_t token = blockIdx.y; - if (head >= n_head || token >= n_tokens) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - const float *qh = q + (uint64_t)token * n_head * qk_dim + - (uint64_t)head * qk_dim; - float *out = qk_low + (uint64_t)token * n_head * kv_lora_dim + - (uint64_t)head * kv_lora_dim; - for (uint32_t j = tid; j < kv_lora_dim; j += nth) { - const char *row = weight + - ((uint64_t)head * kv_lora_dim + j) * row_bytes; - out[j] = glm_q8_0_dot_row_dev(row, qh, qk_nope); - } -} - -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 (!qk_low || !q || !model_map || n_tokens == 0 || n_head == 0 || - kv_lora_dim == 0 || qk_nope == 0 || (qk_nope & 31u) != 0u) { - return 0; - } - if (weight_type != 8u) { - fprintf(stderr, "ds4: glm qk_lowrank: unsupported type %u\n", - weight_type); - return 0; - } - const uint64_t row_bytes = ((uint64_t)qk_nope / 32u) * 34u; - const uint64_t wbytes = (uint64_t)n_head * kv_lora_dim * row_bytes; - if (weight_offset > model_size || wbytes > model_size - weight_offset || - q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || - qk_low->bytes < - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const char *w = (const char *)cuda_resolve_weight_ptr( - model_map, weight_offset, wbytes, logical_tier, "glm_k_b_qk"); - if (!w) return 0; - if (g_q8_dequant_gemm_enabled && g_cublas_ready && n_tokens >= 128u) { - /* Per-head strided-batched GEMM over a dequantized k_b: the - * per-(token,head) warp kernel was ~65ms/layer at 820 tokens. - * Scratch (executing device): [w_f16][q_f16][out_f32]. */ - const uint64_t wh_bytes = - (uint64_t)n_head * kv_lora_dim * qk_nope * sizeof(__half); - const uint64_t xh_off = (wh_bytes + 255u) & ~255ull; - const uint64_t xh_bytes = - (uint64_t)n_tokens * n_head * qk_dim * sizeof(__half); - const uint64_t oo_off = (xh_off + xh_bytes + 255u) & ~255ull; - const uint64_t oo_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - void *tmp = cuda_tmp_alloc_on(logical_tier, oo_off + oo_bytes, - "glm qk_low gemm"); - if (tmp) { - __half *wh = (__half *)tmp; - __half *xh = (__half *)((char *)tmp + xh_off); - float *oo = (float *)((char *)tmp + oo_off); - const uint64_t total_blocks = - (uint64_t)n_head * kv_lora_dim * (qk_nope / 32u); - q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256>>>( - wh, (const unsigned char *)w, total_blocks, - qk_nope / 32u, qk_nope); - const uint64_t xn = (uint64_t)n_tokens * n_head * qk_dim; - f32_to_f16_kernel<<<(xn + 255u) / 256u, 256>>>( - xh, (const float *)q->ptr, xn); - if (cuda_ok(cudaGetLastError(), "glm qk_low gemm staging")) { - const float alpha = 1.0f; - const float beta = 0.0f; - cublasStatus_t st = cublasGemmStridedBatchedEx( - cuda_cublas_for_tier(logical_tier), - CUBLAS_OP_T, CUBLAS_OP_N, - (int)kv_lora_dim, (int)n_tokens, (int)qk_nope, - &alpha, - wh, CUDA_R_16F, (int)qk_nope, - (long long)((uint64_t)kv_lora_dim * qk_nope), - xh, CUDA_R_16F, (int)(n_head * qk_dim), - (long long)qk_dim, - &beta, - oo, CUDA_R_32F, (int)(n_head * kv_lora_dim), - (long long)kv_lora_dim, - (int)n_head, - CUDA_R_32F, CUBLAS_GEMM_DEFAULT); - if (st == CUBLAS_STATUS_SUCCESS && - cuda_ok(cudaMemcpyAsync(qk_low->ptr, oo, oo_bytes, - cudaMemcpyDeviceToDevice, 0), - "glm qk_low gemm out copy")) { - return 1; - } - fprintf(stderr, - "ds4: glm qk_low gemm failed (status %d); native path\n", - (int)st); - } - } - } - dim3 grid(n_head, n_tokens, 1); - glm_qk_lowrank_q8_0_batch_kernel<<>>( - (float *)qk_low->ptr, w, (const float *)q->ptr, - n_tokens, n_head, kv_lora_dim, qk_nope, qk_dim, row_bytes); - return cuda_ok(cudaGetLastError(), "glm qk lowrank batch launch"); -} - -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) { - return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, q, model_map, - model_size, - weight_offset, - weight_type, 1u, n_head, - kv_lora_dim, qk_nope, - qk_dim); -} - -/* Fused decode-path QKV norm + compact-KV store, one block per - * (token, part): part 0 rms-norms q into q_out, part 1 rms-norms - * kv_raw[:kv_lora_dim] into the kv_lora ring, part 2 copies the - * UNROTATED rope tail into the k_rope ring (roped at attention read). */ -__global__ static void glm_qkv_norm_store_compact_kv_kernel( - float *q_dst, - const float *q_src, - const float *q_w, - uint32_t q_n, - const float *kv_raw, - const float *kv_w, - char *kv_lora_cache, - char *k_rope_cache, - uint32_t pos0, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - int cache_f16, - float eps) { - const uint32_t token = blockIdx.x; - const uint32_t part = blockIdx.y; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - const uint32_t pos = pos0 + token; - - if (part == 2u) { - if (pos >= cache_cap) return; - const float *src = kv_raw + (uint64_t)token * kv_raw_dim + kv_lora_dim; - if (cache_f16) { - __half *dst = (__half *)k_rope_cache + (uint64_t)pos * qk_rope; - for (uint32_t i = tid; i < qk_rope; i += nth) { - dst[i] = __float2half(src[i]); - } - } else { - float *dst = (float *)k_rope_cache + (uint64_t)pos * qk_rope; - for (uint32_t i = tid; i < qk_rope; i += nth) { - dst[i] = src[i]; - } - } - return; - } - - const bool kv_task = part != 0u; - const uint32_t n = kv_task ? kv_lora_dim : q_n; - const float *x = kv_task ? kv_raw + (uint64_t)token * kv_raw_dim - : q_src + (uint64_t)token * q_n; - const float *w = kv_task ? kv_w : q_w; - - __shared__ float sh[32]; - float sumf = 0.0f; - for (uint32_t i = tid; i < n; i += nth) { - const float v = x[i]; - sumf += v * v; - } - for (int off = 16; off > 0; off >>= 1) { - sumf += __shfl_xor_sync(0xffffffffu, sumf, off); - } - if ((tid & 31u) == 0u) sh[tid >> 5] = sumf; - __syncthreads(); - if (tid < 32u) { - sumf = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) { - sumf += __shfl_xor_sync(0xffffffffu, sumf, off); - } - if (tid == 0u) sh[0] = sumf; - } - __syncthreads(); - const float scale = rsqrtf(sh[0] / (float)n + eps); - - if (!kv_task) { - float *y = q_dst + (uint64_t)token * q_n; - for (uint32_t i = tid; i < n; i += nth) { - y[i] = (x[i] * scale) * w[i]; - } - return; - } - - if (pos >= cache_cap) return; - if (cache_f16) { - __half *dst = (__half *)kv_lora_cache + (uint64_t)pos * kv_lora_dim; - for (uint32_t i = tid; i < kv_lora_dim; i += nth) { - dst[i] = __float2half((x[i] * scale) * w[i]); - } - } else { - float *dst = (float *)kv_lora_cache + (uint64_t)pos * kv_lora_dim; - for (uint32_t i = tid; i < kv_lora_dim; i += nth) { - dst[i] = (x[i] * scale) * w[i]; - } - } -} - -extern "C" int ds4_gpu_glm_qkv_norm_store_compact_kv_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_lora_cache, - ds4_gpu_tensor *k_rope_cache, - const ds4_gpu_tensor *kv_raw, - uint64_t kv_weight_offset, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - bool cache_f16, - float eps) { - if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || - !model_map || n_tokens == 0 || q_n == 0 || kv_lora_dim == 0 || - qk_rope == 0 || kv_raw_dim < kv_lora_dim + qk_rope || - q->bytes < (uint64_t)n_tokens * q_n * sizeof(float) || - q_out->bytes < (uint64_t)n_tokens * q_n * sizeof(float) || - kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float)) { - return 0; - } - if (q_weight_offset > model_size || - (uint64_t)q_n * sizeof(float) > model_size - q_weight_offset || - kv_weight_offset > model_size || - (uint64_t)kv_lora_dim * sizeof(float) > model_size - kv_weight_offset) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const float *q_w = (const float *)cuda_resolve_weight_ptr( - model_map, q_weight_offset, (uint64_t)q_n * sizeof(float), - logical_tier, "glm_q_norm"); - const float *kv_w = (const float *)cuda_resolve_weight_ptr( - model_map, kv_weight_offset, - (uint64_t)kv_lora_dim * sizeof(float), - logical_tier, "glm_kv_norm"); - if (!q_w || !kv_w) return 0; - dim3 grid(n_tokens, 3, 1); - glm_qkv_norm_store_compact_kv_kernel<<>>( - (float *)q_out->ptr, - (const float *)q->ptr, - q_w, - q_n, - (const float *)kv_raw->ptr, - kv_w, - (char *)kv_lora_cache->ptr, - (char *)k_rope_cache->ptr, - pos0, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope, - cache_f16 ? 1 : 0, eps); - return cuda_ok(cudaGetLastError(), "glm qkv norm store compact kv"); -} - -/* In-place interleaved-pair yarn rope on a [n_tokens][n_head][head_dim] - * f32 tensor, rotating rot_dim dims starting at rot_offset. Shared by the - * attention q tail (offset = head_dim - rot_dim) and the DSA indexer - * (offset = 0). Grid (n_head, n_tokens). */ -__global__ static void glm_rope_tail_offset_kernel( - float *x, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t rot_offset, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - const uint32_t head = blockIdx.x; - const uint32_t token = blockIdx.y; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - if (rot_dim == 0u || (rot_dim & 1u) != 0u || - rot_offset > head_dim || rot_dim > head_dim - rot_offset) return; - - const uint32_t pos = pos0 + token; - float *row = x + ((uint64_t)token * n_head + head) * head_dim + rot_offset; - - float corr_dims[2] = {0.0f, 0.0f}; - if (ext_factor != 0.0f) { - corr_dims[0] = fmaxf(0.0f, - floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim, - (int)n_ctx_orig, beta_fast, freq_base))); - corr_dims[1] = fminf((float)rot_dim - 1.0f, - ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim, - (int)n_ctx_orig, beta_slow, freq_base))); - } - const float theta_base = (float)pos; - const float inv_ndims = -1.0f / (float)rot_dim; - for (uint32_t i = tid * 2u; i < rot_dim; i += nth * 2u) { - const float theta = - theta_base * powf(freq_base, inv_ndims * (float)i); - float ct, st; - glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i, - ext_factor, attn_factor, &ct, &st); - const float x0 = row[i]; - const float x1 = row[i + 1u]; - row[i] = x0 * ct - x1 * st; - row[i + 1u] = x0 * st + x1 * ct; - } -} - -static int glm_rope_tail_offset_launch( - ds4_gpu_tensor *x, - uint32_t n_tokens, uint32_t n_head, uint32_t head_dim, - uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0, - uint32_t n_ctx_orig, float freq_base, float freq_scale, - float ext_factor, float attn_factor, - float beta_fast, float beta_slow, const char *what) { - if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || - rot_dim == 0 || (rot_dim & 1u) != 0u || rot_offset > head_dim || - rot_dim > head_dim - rot_offset || - x->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float)) { - return 0; - } - dim3 grid(n_head, n_tokens, 1); - glm_rope_tail_offset_kernel<<>>( - (float *)x->ptr, n_head, head_dim, rot_dim, rot_offset, - pos0, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - return cuda_ok(cudaGetLastError(), what); -} - -extern "C" int ds4_gpu_glm_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (rot_dim > head_dim) return 0; - return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim, - rot_dim, head_dim - rot_dim, pos0, n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow, "glm rope tail"); -} - -extern "C" int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor\n"); - return 0; -} - -/* Scalar-correct GLM routed MoE (q2_K experts): per (token, slot) block - * quantizes nothing - dots q2_K rows against a q8_K-quantized activation - * staged in shared memory. Grid: (n_tokens, n_expert). Mid buffer holds - * silu(gate)*up per slot; out accumulates expert_weight-scaled down rows. - */ -__global__ static void glm_routed_moe_batch_q2K_gateup_kernel( - float *mid, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t n_tokens, - uint32_t mid_token_stride) { - const uint32_t tok = blockIdx.x; - const uint32_t slot = blockIdx.y; - if (tok >= n_tokens || slot >= n_expert) return; - const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; - if (expert < 0) return; - const cuda_block_q8_K *xrow = xq + (uint64_t)tok * xq_blocks; - float *mrow = mid + (uint64_t)tok * mid_token_stride + - (uint64_t)slot * expert_mid_dim; - for (uint32_t r = threadIdx.x; r < expert_mid_dim; r += blockDim.x) { - const cuda_block_q2_K *gr = (const cuda_block_q2_K *)(gate_base + - (uint64_t)expert * gate_expert_bytes + (uint64_t)r * gate_row_bytes); - const cuda_block_q2_K *ur = (const cuda_block_q2_K *)(up_base + - (uint64_t)expert * up_expert_bytes + (uint64_t)r * up_row_bytes); - float g = 0.0f, u = 0.0f; - for (uint32_t b = 0; b < xq_blocks; b++) { - g += dev_dot_q2_K_q8_K_block(gr + b, xrow + b); - u += dev_dot_q2_K_q8_K_block(ur + b, xrow + b); - } - mrow[r] = (g / (1.0f + expf(-g))) * u; /* silu(g)*u */ - } -} - -__global__ static void glm_routed_moe_batch_q2K_down_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - const float *weights, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.y; - if (tok >= n_tokens) return; - const uint32_t r = blockIdx.x * blockDim.x + threadIdx.x; - if (r >= out_dim) return; - float acc = 0.0f; - for (uint32_t slot = 0; slot < n_expert; slot++) { - const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; - if (expert < 0) continue; - const float w = weights[(uint64_t)tok * n_expert + slot]; - const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base + - (uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes); - const cuda_block_q8_K *mrow = midq + - ((uint64_t)tok * n_expert + slot) * midq_blocks; - float s = 0.0f; - for (uint32_t b = 0; b < midq_blocks; b++) { - s += dev_dot_q2_K_q8_K_block(dr + b, mrow + b); - } - acc += w * s; - } - out[(uint64_t)tok * out_dim + r] = acc; -} - -/* Warp-per-row routed MoE (q2_K x q8_K). Each block stages the token's - * q8_K activation row in shared memory; one warp produces one mid row - * (gate dot + up dot + silu*mul fused). Grid: - * (expert_mid_dim/warps, n_expert, n_tokens). */ -__global__ static void glm_routed_moe_gateup_warp_kernel( - float *mid, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.z; - const uint32_t slot = blockIdx.y; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - if (tok >= n_tokens || slot >= n_expert) return; - const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; - - extern __shared__ unsigned int glm_moe_sh_u32[]; - { - const unsigned int *src = - (const unsigned int *)(xq + (uint64_t)tok * xq_blocks); - const uint32_t words = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u; - for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) { - glm_moe_sh_u32[i] = src[i]; - } - } - __syncthreads(); - if (expert < 0) return; - const cuda_block_q8_K *xrow = (const cuda_block_q8_K *)glm_moe_sh_u32; - - const uint32_t r = blockIdx.x * warps + warp; - if (r >= expert_mid_dim) return; - const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + - (uint64_t)r * gate_row_bytes; - const char *ur = up_base + (uint64_t)expert * up_expert_bytes + - (uint64_t)r * up_row_bytes; - float g = 0.0f, u = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - g += dev_dot_q2_K_q8_K_block( - (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b); - u += dev_dot_q2_K_q8_K_block( - (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b); - } - for (int off = 16; off > 0; off >>= 1) { - g += __shfl_down_sync(0xffffffffu, g, off); - u += __shfl_down_sync(0xffffffffu, u, off); - } - if (lane == 0u) { - mid[((uint64_t)tok * n_expert + slot) * expert_mid_dim + r] = - (g / (1.0f + expf(-g))) * u; - } -} - -/* Exact two-token gate/up with adjacent-token expert reuse. Token 0 owns an - * expert present in both rows; token 1 only launches work for unmatched - * experts. Each token keeps the native lane assignment and warp reduction. */ -__global__ static void glm_routed_moe_gateup_tok2_reuse_kernel( - float *mid, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *selected, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert) { - const uint32_t owner = blockIdx.y; - const uint32_t tok = owner / n_expert; - const uint32_t slot = owner - tok * n_expert; - const uint32_t warps = blockDim.x >> 5u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t lane = threadIdx.x & 31u; - if (tok >= 2u || slot >= n_expert) return; - - const int32_t expert = selected[owner]; - if (expert < 0) return; - int32_t mate_slot = -1; - for (uint32_t s = 0; s < n_expert; s++) { - if (selected[(uint64_t)(1u - tok) * n_expert + s] == expert) { - mate_slot = (int32_t)s; - break; - } - } - if (tok == 1u && mate_slot >= 0) return; - - const uint32_t np = tok == 0u && mate_slot >= 0 ? 2u : 1u; - const uint32_t pair0 = owner; - const uint32_t pair1 = n_expert + (uint32_t)mate_slot; - extern __shared__ unsigned int glm_moe_tok2_sh_u32[]; - const uint32_t words_per_row = - xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u; - const unsigned int *src0 = (const unsigned int *)( - xq + (uint64_t)tok * xq_blocks); - for (uint32_t i = threadIdx.x; i < words_per_row; - i += blockDim.x) { - glm_moe_tok2_sh_u32[i] = src0[i]; - } - if (np == 2u) { - const unsigned int *src1 = - (const unsigned int *)(xq + xq_blocks); - for (uint32_t i = threadIdx.x; i < words_per_row; - i += blockDim.x) { - glm_moe_tok2_sh_u32[words_per_row + i] = src1[i]; - } - } - __syncthreads(); - - const uint32_t r = blockIdx.x * warps + warp; - if (r >= expert_mid_dim) return; - const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + - (uint64_t)r * gate_row_bytes; - const char *ur = up_base + (uint64_t)expert * up_expert_bytes + - (uint64_t)r * up_row_bytes; - const cuda_block_q8_K *x0 = - (const cuda_block_q8_K *)glm_moe_tok2_sh_u32; - const cuda_block_q8_K *x1 = np == 2u - ? x0 + xq_blocks - : NULL; - float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - dev_dot_q2_K_q8_K_block8( - (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), - x0 + b, np == 2u ? x1 + b : NULL, - NULL, NULL, NULL, NULL, NULL, NULL, np, g); - dev_dot_q2_K_q8_K_block8( - (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), - x0 + b, np == 2u ? x1 + b : NULL, - NULL, NULL, NULL, NULL, NULL, NULL, np, u); - } - for (uint32_t p = 0; p < np; p++) { - for (int off = 16; off > 0; off >>= 1) { - g[p] += __shfl_down_sync(0xffffffffu, g[p], off); - u[p] += __shfl_down_sync(0xffffffffu, u[p], off); - } - if (lane == 0u) { - const uint32_t pair = p == 0u ? pair0 : pair1; - mid[(uint64_t)pair * expert_mid_dim + r] = - (g[p] / (1.0f + expf(-g[p]))) * u[p]; - } - } -} - -/* Warp-per-output-row down projection: stages all n_expert quantized mid - * rows for the token in shared memory, each warp accumulates one out row - * across every selected expert. Grid: (out_dim/warps, n_tokens). */ -__global__ static void glm_routed_moe_down_warp_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *selected, - const float *weights, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.y; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - if (tok >= n_tokens) return; - - extern __shared__ unsigned int glm_moe_sh_u32[]; - { - const unsigned int *src = (const unsigned int *) - (midq + (uint64_t)tok * n_expert * midq_blocks); - const uint32_t words = n_expert * midq_blocks * - (uint32_t)sizeof(cuda_block_q8_K) / 4u; - for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) { - glm_moe_sh_u32[i] = src[i]; - } - } - __syncthreads(); - const cuda_block_q8_K *msh = (const cuda_block_q8_K *)glm_moe_sh_u32; - - const uint32_t r = blockIdx.x * warps + warp; - if (r >= out_dim) return; - const uint32_t units = n_expert * midq_blocks; - float acc = 0.0f; - for (uint32_t idx = lane; idx < units; idx += 32u) { - const uint32_t slot = idx / midq_blocks; - const uint32_t b = idx - slot * midq_blocks; - const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; - if (expert < 0) continue; - const float w = weights[(uint64_t)tok * n_expert + slot]; - const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base + - (uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes); - acc += w * dev_dot_q2_K_q8_K_block(dr + b, msh + slot * midq_blocks + b); - } - for (int off = 16; off > 0; off >>= 1) { - acc += __shfl_down_sync(0xffffffffu, acc, off); - } - if (lane == 0u) out[(uint64_t)tok * out_dim + r] = acc; -} - -/* Expert-major routed MoE for prefill: build per-expert token lists, - * then walk rows expert-by-expert so weights stream once per layer and - * activations hit L2. pair = tok * n_expert + slot indexes selected/ - * weights/mid rows directly. */ -__global__ static void glm_moe_expert_map_kernel( - int32_t *counts, - int32_t *lists, - const int32_t *selected, - uint32_t n_pairs, - uint32_t n_total_expert, - uint32_t cap, - uint32_t pair_base) { - const uint32_t p = blockIdx.x * blockDim.x + threadIdx.x; - if (p >= n_pairs) return; - const uint32_t pair = pair_base + p; - const int32_t e = selected[pair]; - if (e < 0 || (uint32_t)e >= n_total_expert) return; - const int32_t idx = atomicAdd(&counts[e], 1); - lists[(uint64_t)e * cap + idx] = (int32_t)pair; -} - -__global__ static void glm_moe_build_expert_tiles8_kernel( - uint32_t *tile_total, - uint32_t *tile_experts, - uint32_t *tile_starts, - const int32_t *counts, - uint32_t n_total_expert) { - if (blockIdx.x != 0u || threadIdx.x != 0u) return; - uint32_t total = 0; - for (uint32_t e = 0; e < n_total_expert; e++) { - const uint32_t count = counts[e] > 0 ? (uint32_t)counts[e] : 0u; - const uint32_t ntiles = (count + 7u) / 8u; - for (uint32_t t = 0; t < ntiles; t++) { - tile_experts[total] = e; - tile_starts[total] = t * 8u; - total++; - } - } - *tile_total = total; -} - -/* Expert-tiled Q2_K gate/up for GLM prefill. One warp keeps the same - * block-to-lane assignment and reduction tree as the token-major W32 - * kernel, but evaluates eight pairs against each loaded expert row. */ -__global__ static void glm_routed_moe_gateup_expert_tile8_kernel( - float *mid, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *counts, - const int32_t *lists, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t cap) { - const uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t warps = blockDim.x >> 5u; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t expert = tile_experts[tile]; - const uint32_t local_start = tile_starts[tile]; - const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u; - - __shared__ uint32_t pair[8]; - __shared__ uint32_t tok[8]; - __shared__ uint32_t np; - if (threadIdx.x == 0u) { - uint32_t n = count - local_start; - if (n > 8u) n = 8u; - np = n; - for (uint32_t p = 0; p < n; p++) { - const uint32_t pr = - (uint32_t)lists[(uint64_t)expert * cap + local_start + p]; - pair[p] = pr; - tok[p] = pr / n_expert; - } - } - __syncthreads(); - - const uint32_t r = blockIdx.x * warps + warp; - if (r >= expert_mid_dim) return; - const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + - (uint64_t)r * gate_row_bytes; - const char *ur = up_base + (uint64_t)expert * up_expert_bytes + - (uint64_t)r * up_row_bytes; - float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - const cuda_block_q8_K *x0 = np > 0u ? xq + (uint64_t)tok[0] * xq_blocks + b : NULL; - const cuda_block_q8_K *x1 = np > 1u ? xq + (uint64_t)tok[1] * xq_blocks + b : NULL; - const cuda_block_q8_K *x2 = np > 2u ? xq + (uint64_t)tok[2] * xq_blocks + b : NULL; - const cuda_block_q8_K *x3 = np > 3u ? xq + (uint64_t)tok[3] * xq_blocks + b : NULL; - const cuda_block_q8_K *x4 = np > 4u ? xq + (uint64_t)tok[4] * xq_blocks + b : NULL; - const cuda_block_q8_K *x5 = np > 5u ? xq + (uint64_t)tok[5] * xq_blocks + b : NULL; - const cuda_block_q8_K *x6 = np > 6u ? xq + (uint64_t)tok[6] * xq_blocks + b : NULL; - const cuda_block_q8_K *x7 = np > 7u ? xq + (uint64_t)tok[7] * xq_blocks + b : NULL; - dev_dot_q2_K_q8_K_block8( - (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), - x0, x1, x2, x3, x4, x5, x6, x7, np, g); - dev_dot_q2_K_q8_K_block8( - (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), - x0, x1, x2, x3, x4, x5, x6, x7, np, u); - } - for (uint32_t p = 0; p < np; p++) { - for (int off = 16; off > 0; off >>= 1) { - g[p] += __shfl_down_sync(0xffffffffu, g[p], off); - u[p] += __shfl_down_sync(0xffffffffu, u[p], off); - } - if (lane == 0u) { - mid[(uint64_t)pair[p] * expert_mid_dim + r] = - (g[p] / (1.0f + expf(-g[p]))) * u[p]; - } - } -} - -__global__ static void glm_routed_moe_gateup_expert_kernel( - float *mid, - const char *gate_base, - const char *up_base, - const cuda_block_q8_K *xq, - const int32_t *counts, - const int32_t *lists, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint32_t xq_blocks, - uint32_t expert_mid_dim, - uint32_t n_expert, - uint32_t cap) { - const uint32_t e = blockIdx.y; - const int32_t nt = counts[e]; - if (nt == 0) return; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t r = blockIdx.x * warps + warp; - if (r >= expert_mid_dim) return; - const char *gr = gate_base + (uint64_t)e * gate_expert_bytes + - (uint64_t)r * gate_row_bytes; - const char *ur = up_base + (uint64_t)e * up_expert_bytes + - (uint64_t)r * up_row_bytes; - const int32_t *lst = lists + (uint64_t)e * cap; - for (int32_t i = 0; i < nt; i++) { - const uint32_t pair = (uint32_t)lst[i]; - const cuda_block_q8_K *xrow = xq + (uint64_t)(pair / n_expert) * xq_blocks; - float g = 0.0f, u = 0.0f; - for (uint32_t b = lane; b < xq_blocks; b += 32u) { - g += dev_dot_q2_K_q8_K_block( - (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b); - u += dev_dot_q2_K_q8_K_block( - (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b); - } - for (int off = 16; off > 0; off >>= 1) { - g += __shfl_down_sync(0xffffffffu, g, off); - u += __shfl_down_sync(0xffffffffu, u, off); - } - if (lane == 0u) { - mid[(uint64_t)pair * expert_mid_dim + r] = - (g / (1.0f + expf(-g))) * u; - } - } -} - -__global__ static void glm_routed_moe_down_expert_kernel( - float *out, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *counts, - const int32_t *lists, - const float *weights, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t cap) { - const uint32_t e = blockIdx.y; - const int32_t nt = counts[e]; - if (nt == 0) return; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t r = blockIdx.x * warps + warp; - if (r >= out_dim) return; - const char *dr = down_base + (uint64_t)e * down_expert_bytes + - (uint64_t)r * down_row_bytes; - const int32_t *lst = lists + (uint64_t)e * cap; - for (int32_t i = 0; i < nt; i++) { - const uint32_t pair = (uint32_t)lst[i]; - const cuda_block_q8_K *mrow = midq + (uint64_t)pair * midq_blocks; - float s = 0.0f; - for (uint32_t b = lane; b < midq_blocks; b += 32u) { - s += dev_dot_q2_K_q8_K_block( - (const cuda_block_q2_K *)(dr + (uint64_t)b * 84u), mrow + b); - } - for (int off = 16; off > 0; off >>= 1) { - s += __shfl_down_sync(0xffffffffu, s, off); - } - if (lane == 0u) { - atomicAdd(&out[(uint64_t)(pair / n_expert) * out_dim + r], - weights[pair] * s); - } - } -} - -/* Expert-tiled down projection with an exact token-major reduction. The - * first kernel reuses each Q2_K row across eight routed pairs, but materializes - * each block dot. The second kernel applies the router weight and consumes the - * dots with the same lane assignment and warp tree as the native kernel. */ -__global__ static void glm_routed_moe_down_expert_tile8_terms_kernel( - float *terms, - const char *down_base, - const cuda_block_q8_K *midq, - const int32_t *counts, - const int32_t *lists, - const uint32_t *tile_total, - const uint32_t *tile_experts, - const uint32_t *tile_starts, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t cap, - uint32_t pair_base) { - const uint32_t tile = blockIdx.y; - if (tile >= *tile_total) return; - const uint32_t lane = threadIdx.x & 7u; - const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); - const uint32_t expert = tile_experts[tile]; - const uint32_t local_start = tile_starts[tile]; - const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u; - - __shared__ uint32_t pair[8]; - __shared__ uint32_t np; - __shared__ cuda_block_q8_K mq[8][8]; - if (threadIdx.x == 0u) { - uint32_t n = count - local_start; - if (n > 8u) n = 8u; - np = n; - for (uint32_t p = 0; p < n; p++) { - pair[p] = (uint32_t)lists[ - (uint64_t)expert * cap + local_start + p]; - } - } - __syncthreads(); - for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { - const uint32_t p = i / midq_blocks; - const uint32_t b = i - p * midq_blocks; - mq[p][b] = midq[(uint64_t)pair[p] * midq_blocks + b]; - } - __syncthreads(); - if (row >= out_dim || lane >= midq_blocks) return; - - const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + - (uint64_t)expert * down_expert_bytes + - (uint64_t)row * down_row_bytes); - for (uint32_t p = 0; p < np; p++) { - const uint32_t pr = pair[p]; - const float dot = dev_dot_q2_K_q8_K_block(wr + lane, &mq[p][lane]); - terms[((uint64_t)(pr - pair_base) * out_dim + row) * - midq_blocks + lane] = - dot; - } -} - -__global__ static void glm_routed_moe_down_terms_reduce_kernel( - float *out, - const float *terms, - const int32_t *selected, - const float *weights, - uint32_t midq_blocks, - uint32_t out_dim, - uint32_t n_expert, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.y; - const uint32_t warp = threadIdx.x >> 5u; - const uint32_t lane = threadIdx.x & 31u; - const uint32_t row = blockIdx.x * 8u + warp; - if (tok >= n_tokens || row >= out_dim) return; - const uint32_t units = n_expert * midq_blocks; - float acc = 0.0f; - for (uint32_t idx = lane; idx < units; idx += 32u) { - const uint32_t slot = idx / midq_blocks; - const uint32_t b = idx - slot * midq_blocks; - const uint32_t pr = tok * n_expert + slot; - if (selected[pr] >= 0) { - acc += weights[pr] * - terms[((uint64_t)pr * out_dim + row) * midq_blocks + b]; - } - } - for (int off = 16; off > 0; off >>= 1) { - acc += __shfl_down_sync(0xffffffffu, acc, off); - } - if (lane == 0u) out[(uint64_t)tok * out_dim + row] = acc; -} - -static int glm_routed_moe_finish_batch( - ds4_gpu_tensor *out, - float *out_work, - uint64_t out_bytes, - const char *what) { - if (!cuda_ok(cudaGetLastError(), what)) return 0; - if (out_work == (float *)out->ptr) return 1; - return cuda_ok(cudaMemcpyAsync(out->ptr, out_work, out_bytes, - cudaMemcpyDeviceToDevice, 0), - "glm routed moe local output copy"); -} - -extern "C" int ds4_gpu_glm_routed_moe_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride) { - (void)layer_index; (void)n_total_expert; - if (!out || !mid || !x || !selected || !weights || !model_map || - n_tokens == 0 || n_expert == 0 || - (expert_in_dim & 255u) != 0u || (expert_mid_dim & 255u) != 0u) { - return 0; - } - if (gate_type != 10u || up_type != 10u || down_type != 10u) { - fprintf(stderr, "ds4: glm routed moe: unsupported types %u/%u/%u\n", - gate_type, up_type, down_type); - return 0; - } - if (mid_token_stride != n_expert * expert_mid_dim) { - fprintf(stderr, - "ds4: glm routed moe: mid stride %u != %u (packed rows expected)\n", - mid_token_stride, n_expert * expert_mid_dim); - return 0; - } - const int logical_tier = cuda_current_tier(); - const char *gw = (const char *)cuda_resolve_weight_ptr(model_map, - gate_offset, (uint64_t)256 * gate_expert_bytes, logical_tier, - "glm_gate_exps"); - const char *uw = (const char *)cuda_resolve_weight_ptr(model_map, - up_offset, (uint64_t)256 * up_expert_bytes, logical_tier, - "glm_up_exps"); - const char *dw = (const char *)cuda_resolve_weight_ptr(model_map, - down_offset, (uint64_t)256 * down_expert_bytes, logical_tier, - "glm_down_exps"); - if (!gw || !uw || !dw) return 0; - - /* Stage 1: quantize x rows to q8_K (existing kernel). */ - const uint32_t xq_blocks = expert_in_dim / 256u; - const uint32_t midq_blocks = expert_mid_dim / 256u; - static ds4_gpu_tensor *xq_scratch[DS4_MAX_GPUS] = {0}; - static ds4_gpu_tensor *midq_scratch[DS4_MAX_GPUS] = {0}; - int dev = logical_tier; - const int scratch_tier = getenv("DS4_GLM_MOE_SCRATCH_TIER0") ? 0 : dev; - const uint64_t xq_bytes = (uint64_t)n_tokens * xq_blocks * - sizeof(cuda_block_q8_K); - const uint64_t midq_bytes = (uint64_t)n_tokens * n_expert * midq_blocks * - sizeof(cuda_block_q8_K); - if (!xq_scratch[dev] || xq_scratch[dev]->bytes < xq_bytes) { - if (xq_scratch[dev]) ds4_gpu_tensor_free(xq_scratch[dev]); - xq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, xq_bytes); - } - if (!midq_scratch[dev] || midq_scratch[dev]->bytes < midq_bytes) { - if (midq_scratch[dev]) ds4_gpu_tensor_free(midq_scratch[dev]); - midq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, midq_bytes); - } - if (!xq_scratch[dev] || !midq_scratch[dev]) return 0; - - static ds4_gpu_tensor *mid_local[DS4_MAX_GPUS] = {0}; - static ds4_gpu_tensor *out_local[DS4_MAX_GPUS] = {0}; - const uint64_t mid_work_bytes = - (uint64_t)n_tokens * mid_token_stride * sizeof(float); - const uint64_t out_work_bytes = - (uint64_t)n_tokens * out_dim * sizeof(float); - float *mid_work = (float *)mid->ptr; - float *out_work = (float *)out->ptr; - const bool use_local_batch_io = - n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_LOCAL_BATCH_IO"); - if (use_local_batch_io && ds4_tensor_device_idx(mid) != dev) { - if (!mid_local[dev] || mid_local[dev]->bytes < mid_work_bytes) { - if (mid_local[dev]) ds4_gpu_tensor_free(mid_local[dev]); - mid_local[dev] = - ds4_gpu_tensor_alloc_ptr_on(dev, mid_work_bytes); - } - if (!mid_local[dev]) return 0; - mid_work = (float *)mid_local[dev]->ptr; - } - if (use_local_batch_io && ds4_tensor_device_idx(out) != dev) { - if (!out_local[dev] || out_local[dev]->bytes < out_work_bytes) { - if (out_local[dev]) ds4_gpu_tensor_free(out_local[dev]); - out_local[dev] = - ds4_gpu_tensor_alloc_ptr_on(dev, out_work_bytes); - } - if (!out_local[dev]) return 0; - out_work = (float *)out_local[dev]->ptr; - } - { - dim3 gq(xq_blocks, n_tokens, 1); - q8_K_quantize_kernel<<>>( - (cuda_block_q8_K *)xq_scratch[dev]->ptr, - (const float *)x->ptr, expert_in_dim, n_tokens); - } - - static ds4_gpu_tensor *map_scratch[DS4_MAX_GPUS] = {0}; - static ds4_gpu_tensor *down_terms_scratch[DS4_MAX_GPUS] = {0}; - const bool use_expert_tile8 = - n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_EXPERT_TILE8"); - const bool use_expert_major = - n_tokens >= 16u && getenv("DS4_GLM_MOE_EXPERT_MAJOR"); - if (use_expert_tile8 || use_expert_major) { - const uint32_t cap = n_tokens; - const uint32_t n_pairs = n_tokens * n_expert; - const uint64_t counts_bytes = 256u * sizeof(int32_t); - const uint64_t lists_off = (counts_bytes + 255u) & ~255ull; - const uint64_t lists_bytes = - (uint64_t)256u * cap * sizeof(int32_t); - const uint32_t tile_capacity = - (n_pairs + 7u) / 8u + 256u; - const uint64_t tile_total_off = - (lists_off + lists_bytes + 255u) & ~255ull; - const uint64_t tile_experts_off = - (tile_total_off + sizeof(uint32_t) + 255u) & ~255ull; - const uint64_t tile_starts_off = - tile_experts_off + (uint64_t)tile_capacity * sizeof(uint32_t); - const uint64_t map_bytes = use_expert_tile8 - ? tile_starts_off + (uint64_t)tile_capacity * sizeof(uint32_t) - : lists_off + lists_bytes; - if (!map_scratch[dev] || map_scratch[dev]->bytes < map_bytes) { - if (map_scratch[dev]) ds4_gpu_tensor_free(map_scratch[dev]); - map_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(dev, map_bytes); - } - if (map_scratch[dev]) { - int32_t *counts = (int32_t *)map_scratch[dev]->ptr; - int32_t *lists = (int32_t *)((char *)map_scratch[dev]->ptr + lists_off); - cudaMemsetAsync(counts, 0, counts_bytes); - glm_moe_expert_map_kernel<<<(n_pairs + 255u) / 256u, 256>>>( - counts, lists, (const int32_t *)selected->ptr, - n_pairs, 256u, cap, 0u); - if (use_expert_tile8) { - uint32_t *tile_total = (uint32_t *)( - (char *)map_scratch[dev]->ptr + tile_total_off); - uint32_t *tile_experts = (uint32_t *)( - (char *)map_scratch[dev]->ptr + tile_experts_off); - uint32_t *tile_starts = (uint32_t *)( - (char *)map_scratch[dev]->ptr + tile_starts_off); - glm_moe_build_expert_tiles8_kernel<<<1, 1>>>( - tile_total, tile_experts, tile_starts, - counts, 256u); - dim3 ge1((expert_mid_dim + 7u) / 8u, - tile_capacity, 1); - glm_routed_moe_gateup_expert_tile8_kernel<<>>( - mid_work, gw, uw, - (const cuda_block_q8_K *)xq_scratch[dev]->ptr, - counts, lists, - tile_total, tile_experts, tile_starts, - gate_expert_bytes, gate_row_bytes, - up_expert_bytes, up_row_bytes, - xq_blocks, expert_mid_dim, n_expert, cap); - q8_K_quantize_kernel<<< - dim3(midq_blocks, n_tokens * n_expert, 1), 256>>>( - (cuda_block_q8_K *)midq_scratch[dev]->ptr, - mid_work, - expert_mid_dim, n_tokens * n_expert); - if (getenv("DS4_GLM_MOE_NO_DOWN_TILE8_EXACT") == NULL) { - const uint32_t max_chunk_tokens = 512u; - const uint32_t scratch_tokens = - n_tokens < max_chunk_tokens ? n_tokens : max_chunk_tokens; - const uint64_t term_count = - (uint64_t)scratch_tokens * n_expert * - out_dim * midq_blocks; - const uint64_t term_bytes = term_count * sizeof(float); - if (!down_terms_scratch[dev] || - down_terms_scratch[dev]->bytes < term_bytes) { - if (down_terms_scratch[dev]) { - ds4_gpu_tensor_free(down_terms_scratch[dev]); - } - down_terms_scratch[dev] = - ds4_gpu_tensor_alloc_ptr_on(dev, term_bytes); - } - if (down_terms_scratch[dev]) { - for (uint32_t token0 = 0; token0 < n_tokens; - token0 += max_chunk_tokens) { - uint32_t chunk_tokens = n_tokens - token0; - if (chunk_tokens > max_chunk_tokens) { - chunk_tokens = max_chunk_tokens; - } - const uint32_t pair_base = token0 * n_expert; - const uint32_t chunk_pairs = - chunk_tokens * n_expert; - uint32_t chunk_tile_capacity = tile_capacity; - if (n_tokens > max_chunk_tokens) { - chunk_tile_capacity = - (chunk_pairs + 7u) / 8u + 256u; - cudaMemsetAsync(counts, 0, counts_bytes); - glm_moe_expert_map_kernel<<< - (chunk_pairs + 255u) / 256u, 256>>>( - counts, lists, - (const int32_t *)selected->ptr, - chunk_pairs, 256u, cap, pair_base); - glm_moe_build_expert_tiles8_kernel<<<1, 1>>>( - tile_total, tile_experts, tile_starts, - counts, 256u); - } - dim3 gd1((out_dim + 31u) / 32u, - chunk_tile_capacity, 1); - glm_routed_moe_down_expert_tile8_terms_kernel<<< - gd1, 256>>>( - (float *)down_terms_scratch[dev]->ptr, - dw, - (const cuda_block_q8_K *)midq_scratch[dev]->ptr, - counts, lists, - tile_total, tile_experts, tile_starts, - down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, cap, - pair_base); - dim3 gd2((out_dim + 7u) / 8u, - chunk_tokens, 1); - glm_routed_moe_down_terms_reduce_kernel<<< - gd2, 256>>>( - out_work + (uint64_t)token0 * out_dim, - (const float *)down_terms_scratch[dev]->ptr, - (const int32_t *)selected->ptr + pair_base, - (const float *)weights->ptr + pair_base, - midq_blocks, out_dim, n_expert, - chunk_tokens); - } - return glm_routed_moe_finish_batch( - out, out_work, out_work_bytes, - "glm routed moe exact down tile8"); - } - } - const uint32_t warps = 8u; - dim3 ge2((out_dim + warps - 1u) / warps, - n_tokens, 1); - const uint32_t sh2 = n_expert * midq_blocks * - (uint32_t)sizeof(cuda_block_q8_K); - glm_routed_moe_down_warp_kernel<<< - ge2, warps * 32u, sh2>>>( - out_work, dw, - (const cuda_block_q8_K *)midq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, n_tokens); - return glm_routed_moe_finish_batch( - out, out_work, out_work_bytes, - "glm routed moe expert tile8"); - } - dim3 ge1((expert_mid_dim + 7u) / 8u, 256u, 1); - glm_routed_moe_gateup_expert_kernel<<>>( - mid_work, gw, uw, - (const cuda_block_q8_K *)xq_scratch[dev]->ptr, - counts, lists, - gate_expert_bytes, gate_row_bytes, - up_expert_bytes, up_row_bytes, - xq_blocks, expert_mid_dim, n_expert, cap); - q8_K_quantize_kernel<<>>( - (cuda_block_q8_K *)midq_scratch[dev]->ptr, - mid_work, expert_mid_dim, n_tokens * n_expert); - cudaMemsetAsync(out_work, 0, - (uint64_t)n_tokens * out_dim * sizeof(float)); - dim3 ge2((out_dim + 7u) / 8u, 256u, 1); - glm_routed_moe_down_expert_kernel<<>>( - out_work, dw, - (const cuda_block_q8_K *)midq_scratch[dev]->ptr, - counts, lists, (const float *)weights->ptr, - down_expert_bytes, down_row_bytes, - midq_blocks, out_dim, n_expert, cap); - return glm_routed_moe_finish_batch( - out, out_work, out_work_bytes, - "glm routed moe expert-major"); - } - } - if (n_tokens == 2u && - g_glm_mtp_verify_mode && - getenv("DS4_GLM_MTP_NO_MOE_TOK2") == NULL) { - const uint32_t warps = 8u; - dim3 g1((expert_mid_dim + warps - 1u) / warps, - 2u * n_expert, 1u); - const uint32_t sh1 = 2u * xq_blocks * - (uint32_t)sizeof(cuda_block_q8_K); - glm_routed_moe_gateup_tok2_reuse_kernel<<< - g1, warps * 32u, sh1>>>( - mid_work, gw, uw, - (const cuda_block_q8_K *)xq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - gate_expert_bytes, gate_row_bytes, - up_expert_bytes, up_row_bytes, - xq_blocks, expert_mid_dim, n_expert); - } else if (getenv("DS4_GLM_MOE_SCALAR")) { - dim3 g1(n_tokens, n_expert, 1); - glm_routed_moe_batch_q2K_gateup_kernel<<>>( - mid_work, gw, uw, - (const cuda_block_q8_K *)xq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes, - xq_blocks, expert_mid_dim, n_expert, n_tokens, mid_token_stride); - } else { - const uint32_t warps = 8u; - dim3 g1((expert_mid_dim + warps - 1u) / warps, n_expert, n_tokens); - const uint32_t sh1 = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K); - glm_routed_moe_gateup_warp_kernel<<>>( - mid_work, gw, uw, - (const cuda_block_q8_K *)xq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes, - xq_blocks, expert_mid_dim, n_expert, n_tokens); - } - - { - dim3 gq(midq_blocks, n_tokens * n_expert, 1); - q8_K_quantize_kernel<<>>( - (cuda_block_q8_K *)midq_scratch[dev]->ptr, - mid_work, expert_mid_dim, n_tokens * n_expert); - } - - if (getenv("DS4_GLM_MOE_SCALAR")) { - dim3 g2((out_dim + 127u) / 128u, n_tokens, 1); - glm_routed_moe_batch_q2K_down_kernel<<>>( - out_work, dw, - (const cuda_block_q8_K *)midq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - down_expert_bytes, down_row_bytes, midq_blocks, out_dim, - n_expert, n_tokens); - } else { - const uint32_t warps = 8u; - dim3 g2((out_dim + warps - 1u) / warps, n_tokens, 1); - const uint32_t sh2 = n_expert * midq_blocks * - (uint32_t)sizeof(cuda_block_q8_K); - glm_routed_moe_down_warp_kernel<<>>( - out_work, dw, - (const cuda_block_q8_K *)midq_scratch[dev]->ptr, - (const int32_t *)selected->ptr, - (const float *)weights->ptr, - down_expert_bytes, down_row_bytes, midq_blocks, out_dim, - n_expert, n_tokens); - } - return glm_routed_moe_finish_batch( - out, out_work, out_work_bytes, - "glm routed moe batch launch"); -} - -extern "C" int ds4_gpu_glm_routed_moe_one_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - bool force_resident) { - (void)force_resident; - return ds4_gpu_glm_routed_moe_batch_tensor(out, mid, - model_map, model_size, - gate_offset, up_offset, down_offset, - gate_type, up_type, down_type, - gate_expert_bytes, gate_row_bytes, - up_expert_bytes, up_row_bytes, - down_expert_bytes, down_row_bytes, - expert_in_dim, expert_mid_dim, out_dim, - selected, weights, n_total_expert, n_expert, layer_index, - x, 1, n_expert * expert_mid_dim); -} - -/* Parallel router select: 256 threads compute sigmoid probs, then top-k - * via k rounds of shared-memory argmax over probs+bias (value desc, index - * asc tie-break — matches the CPU topk_desc). One block per token. */ -__global__ static void glm_router_select_parallel_kernel( - int32_t *selected, - float *weights_out, - float *probs_out, - const float *bias, - const float *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.x; - const uint32_t tid = threadIdx.x; - if (tok >= n_tokens) return; - const float *lg = logits + (uint64_t)tok * n_expert; - float *probs = probs_out + (uint64_t)tok * n_expert; - int32_t *sel = selected + (uint64_t)tok * n_expert_used; - float *w = weights_out + (uint64_t)tok * n_expert_used; - - __shared__ float sh_v[256]; - __shared__ int sh_i[256]; - __shared__ float sh_sel_v[256]; - __shared__ float sh_sum; - - float my_v = -1e30f; - if (tid < n_expert) { - const float p = 1.0f / (1.0f + expf(-lg[tid])); - probs[tid] = p; - my_v = p + bias[tid]; - } - if (tid == 0u) sh_sum = 0.0f; - sh_sel_v[tid] = my_v; - __syncthreads(); - - for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { - sh_v[tid] = sh_sel_v[tid]; - sh_i[tid] = (int)tid; - __syncthreads(); - for (uint32_t step = 128u; step > 0u; step >>= 1u) { - if (tid < step) { - const float ov = sh_v[tid + step]; - const int oi = sh_i[tid + step]; - if (ov > sh_v[tid] || (ov == sh_v[tid] && oi < sh_i[tid])) { - sh_v[tid] = ov; - sh_i[tid] = oi; - } - } - __syncthreads(); - } - if (tid == 0u) { - const int best = sh_i[0]; - sel[k2] = best; - const float p = probs[best]; - w[k2] = p; - sh_sum += p; - sh_sel_v[best] = -1e30f; - } - __syncthreads(); - } - if (tid == 0u) { - float sum = sh_sum; - if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; - for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { - w[k2] = w[k2] / sum * expert_weight_scale; - } - } -} - -__global__ static void glm_router_select_batch_kernel( - int32_t *selected, - float *weights_out, - float *probs_out, - const float *bias, - const float *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_tokens) { - const uint32_t tok = blockIdx.x; - if (tok >= n_tokens || threadIdx.x != 0u) return; - const float *lg = logits + (uint64_t)tok * n_expert; - float *probs = probs_out + (uint64_t)tok * n_expert; - int32_t *sel = selected + (uint64_t)tok * n_expert_used; - float *w = weights_out + (uint64_t)tok * n_expert_used; - - for (uint32_t i = 0; i < n_expert; i++) { - const float p = 1.0f / (1.0f + expf(-lg[i])); - probs[i] = p; - } - /* top-k over probs+bias, ties by smaller index (matches CPU topk_desc) */ - bool taken[384]; - for (uint32_t i = 0; i < n_expert; i++) taken[i] = false; - float sum = 0.0f; - for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { - int best = -1; float bv = -1e30f; - for (uint32_t i = 0; i < n_expert; i++) { - if (taken[i]) continue; - const float v = probs[i] + bias[i]; - if (v > bv) { bv = v; best = (int)i; } - } - taken[best] = true; - sel[k2] = best; - w[k2] = probs[best]; - sum += probs[best]; - } - if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; - for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { - w[k2] = w[k2] / sum * expert_weight_scale; - } -} - -extern "C" int ds4_gpu_glm_router_select_batch_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - const ds4_gpu_tensor *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_tokens) { - if (!selected || !weights || !probs || !logits || !model_map || - n_expert == 0 || n_expert > 384u || n_expert_used == 0 || - n_tokens == 0) { - return 0; - } - const uint64_t bb = (uint64_t)n_expert * sizeof(float); - if (bias_offset > model_size || bb > model_size - bias_offset || - logits->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) || - selected->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(int32_t) || - weights->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(float) || - probs->bytes < (uint64_t)n_tokens * n_expert * sizeof(float)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const float *bias = (const float *)cuda_resolve_weight_ptr( - model_map, bias_offset, bb, logical_tier, "glm_exp_probs_b"); - if (!bias) return 0; - if (n_expert <= 256u && !getenv("DS4_GLM_ROUTER_SCALAR")) { - glm_router_select_parallel_kernel<<>>( - (int32_t *)selected->ptr, - (float *)weights->ptr, - (float *)probs->ptr, - bias, - (const float *)logits->ptr, - n_expert, n_expert_used, expert_weight_scale, n_tokens); - } else glm_router_select_batch_kernel<<>>( - (int32_t *)selected->ptr, - (float *)weights->ptr, - (float *)probs->ptr, - bias, - (const float *)logits->ptr, - n_expert, n_expert_used, expert_weight_scale, n_tokens); - return cuda_ok(cudaGetLastError(), "glm router select batch launch"); -} - -extern "C" int ds4_gpu_glm_router_select_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - const ds4_gpu_tensor *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale) { - return ds4_gpu_glm_router_select_batch_tensor(selected, weights, probs, - model_map, model_size, - bias_offset, logits, - n_expert, n_expert_used, - expert_weight_scale, 1u); -} - -__global__ static void glm_store_compact_kv_kernel( - char *kv_lora_cache, - char *k_rope_cache, - const float *kv_norm, - const float *kv_raw, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - uint32_t cache_f16) { - const uint32_t token = blockIdx.x; - const uint32_t part = blockIdx.y; - if (token >= n_tokens || part > 1u) return; - const uint32_t pos = pos0 + token; - if (pos >= cache_cap) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - if (part == 0u) { - const float *src = kv_norm + (uint64_t)token * kv_lora_dim; - if (cache_f16) { - __half *dst = (__half *)(kv_lora_cache + - (uint64_t)pos * kv_lora_dim * sizeof(__half)); - for (uint32_t i = tid; i < kv_lora_dim; i += nth) - dst[i] = __float2half(src[i]); - } else { - float *dst = (float *)(kv_lora_cache + - (uint64_t)pos * kv_lora_dim * sizeof(float)); - for (uint32_t i = tid; i < kv_lora_dim; i += nth) - dst[i] = src[i]; - } - } else { - const float *src = kv_raw + - (uint64_t)token * kv_raw_dim + kv_lora_dim; - if (cache_f16) { - __half *dst = (__half *)(k_rope_cache + - (uint64_t)pos * qk_rope * sizeof(__half)); - for (uint32_t i = tid; i < qk_rope; i += nth) - dst[i] = __float2half(src[i]); - } else { - float *dst = (float *)(k_rope_cache + - (uint64_t)pos * qk_rope * sizeof(float)); - for (uint32_t i = tid; i < qk_rope; i += nth) - dst[i] = src[i]; - } - } -} - -extern "C" int ds4_gpu_glm_store_compact_kv_tensor( - ds4_gpu_tensor *kv_lora_cache, - ds4_gpu_tensor *k_rope_cache, - const ds4_gpu_tensor *kv_norm, - const ds4_gpu_tensor *kv_raw, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - bool cache_f16) { - if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || - n_tokens == 0 || kv_lora_dim == 0 || qk_rope == 0 || - kv_lora_dim + qk_rope > kv_raw_dim + qk_rope) { - return 0; - } - const uint64_t es = cache_f16 ? sizeof(__half) : sizeof(float); - if (kv_norm->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float) || - kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) || - kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * es || - k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * es) { - return 0; - } - dim3 grid(n_tokens, 2, 1); - glm_store_compact_kv_kernel<<>>( - (char *)kv_lora_cache->ptr, - (char *)k_rope_cache->ptr, - (const float *)kv_norm->ptr, - (const float *)kv_raw->ptr, - pos0, n_tokens, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope, - cache_f16 ? 1u : 0u); - return cuda_ok(cudaGetLastError(), "glm store compact kv launch"); -} - - - - - - - -__global__ static void glm_store_indexer_k_kernel( - char *cache, - const float *raw_k, - const float *w, - const float *b, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t n_ctx_orig, - float eps, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - uint32_t cache_f16) { - const uint32_t token = blockIdx.x; - if (token >= n_tokens) return; - const uint32_t pos = pos0 + token; - if (pos >= cache_cap) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - const float *src = raw_k + (uint64_t)token * head_dim; - - __shared__ float scratch[256]; - float sum = 0.0f; - for (uint32_t i = tid; i < head_dim; i += nth) sum += src[i]; - scratch[tid] = sum; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) scratch[tid] += scratch[tid + step]; - __syncthreads(); - } - const float mean = scratch[0] / (float)head_dim; - __syncthreads(); - float ss = 0.0f; - for (uint32_t i = tid; i < head_dim; i += nth) { - const float d = src[i] - mean; - ss += d * d; - } - scratch[tid] = ss; - __syncthreads(); - for (uint32_t step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) scratch[tid] += scratch[tid + step]; - __syncthreads(); - } - const float inv = rsqrtf(scratch[0] / (float)head_dim + eps); - - float corr_dims[2] = {0.0f, 0.0f}; - if (ext_factor != 0.0f) { - corr_dims[0] = fmaxf(0.0f, - floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig, - beta_fast, freq_base))); - corr_dims[1] = fminf((float)rot_dim - 1.0f, - ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig, - beta_slow, freq_base))); - } - const float theta_base = (float)pos; - const float inv_ndims = -1.0f / (float)rot_dim; - - for (uint32_t i = tid; i < head_dim; i += nth) { - float v0, v1; bool pair = false; - if (i < rot_dim) { - if ((i & 1u) != 0u) continue; - const float theta = theta_base * powf(freq_base, inv_ndims * (float)i); - float ct, st; - glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i, - ext_factor, attn_factor, &ct, &st); - const float x0 = (src[i] - mean) * inv * w[i] + b[i]; - const float x1 = (src[i + 1u] - mean) * inv * w[i + 1u] + b[i + 1u]; - v0 = x0 * ct - x1 * st; - v1 = x0 * st + x1 * ct; - pair = true; - } else { - v0 = (src[i] - mean) * inv * w[i] + b[i]; - } - if (cache_f16) { - __half *dst = (__half *)(cache + (uint64_t)pos * head_dim * sizeof(__half)); - dst[i] = __float2half(v0); - if (pair) dst[i + 1u] = __float2half(v1); - } else { - float *dst = (float *)(cache + (uint64_t)pos * head_dim * sizeof(float)); - dst[i] = v0; - if (pair) dst[i + 1u] = v1; - } - } -} - -extern "C" int ds4_gpu_glm_store_indexer_k_tensor( - ds4_gpu_tensor *indexer_key_cache, - const ds4_gpu_tensor *raw_k, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint64_t bias_offset, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t n_ctx_orig, - float eps, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool cache_f16) { - if (!indexer_key_cache || !raw_k || !model_map || n_tokens == 0 || - head_dim == 0 || head_dim > 256u || (rot_dim & 1u) != 0u) { - return 0; - } - const uint64_t wb = (uint64_t)head_dim * sizeof(float); - if (weight_offset > model_size || wb > model_size - weight_offset || - bias_offset > model_size || wb > model_size - bias_offset || - raw_k->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || - indexer_key_cache->bytes < - (uint64_t)cache_cap * head_dim * - (cache_f16 ? sizeof(__half) : sizeof(float))) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const float *w = (const float *)cuda_resolve_weight_ptr( - model_map, weight_offset, wb, logical_tier, "glm_indexer_k_norm"); - const float *b = (const float *)cuda_resolve_weight_ptr( - model_map, bias_offset, wb, logical_tier, "glm_indexer_k_norm_b"); - if (!w || !b) return 0; - glm_store_indexer_k_kernel<<>>( - (char *)indexer_key_cache->ptr, - (const float *)raw_k->ptr, - w, b, pos0, n_tokens, cache_cap, head_dim, rot_dim, n_ctx_orig, - eps, freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow, cache_f16 ? 1u : 0u); - return cuda_ok(cudaGetLastError(), "glm store indexer k launch"); -} - -extern "C" int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( - const ds4_gpu_stream_expert_table *table, - const ds4_gpu_tensor *selected, - uint32_t n_selected) { - if (!g_ssd_streaming_mode) return 1; - if (!table || !selected || n_selected == 0 || - selected->bytes < (uint64_t)n_selected * sizeof(int32_t)) { - return 0; - } - std::vector ids; - try { - ids.resize(n_selected); - } catch (...) { - return 0; - } - if (!cuda_ok(cudaMemcpy(ids.data(), selected->ptr, - (size_t)n_selected * sizeof(int32_t), - cudaMemcpyDeviceToHost), - "GLM streaming selected-id read")) { - return 0; - } - return cuda_stream_selected_cache_begin_load(table, ids.data(), n_selected); -} - -__global__ static void glm_value_project_q8_0_batch_heads_kernel( - float *heads, - const char *weight, - const float *lora, - uint32_t n_tokens, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t value_dim, - uint64_t row_bytes) { - const uint32_t head = blockIdx.x; - const uint32_t token = blockIdx.y; - if (head >= n_head || token >= n_tokens) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - extern __shared__ float xsh[]; - const float *src = lora + (uint64_t)token * n_head * kv_lora_dim + - (uint64_t)head * kv_lora_dim; - float *out = heads + (uint64_t)token * n_head * value_dim + - (uint64_t)head * value_dim; - for (uint32_t j = tid; j < kv_lora_dim; j += nth) xsh[j] = src[j]; - __syncthreads(); - for (uint32_t d = tid; d < value_dim; d += nth) { - const char *row = weight + ((uint64_t)head * value_dim + d) * row_bytes; - out[d] = glm_q8_0_dot_row_dev(row, xsh, kv_lora_dim); - } -} - -/* Reuse each Q8 row across a token tile while retaining the scalar kernel's - * block/k accumulation order independently for every output token. */ -template -__global__ static void glm_value_project_q8_0_batch_heads_tiled_kernel( - float *heads, - const char *weight, - const float *lora, - uint32_t n_tokens, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t value_dim, - uint64_t row_bytes) { - const uint32_t head = blockIdx.x; - const uint32_t token0 = blockIdx.y * token_tile; - if (head >= n_head || token0 >= n_tokens) return; - const uint32_t tid = threadIdx.x; - const uint32_t nth = blockDim.x; - extern __shared__ float xsh[]; - -#pragma unroll - for (uint32_t t = 0; t < token_tile; t++) { - const uint32_t token = token0 + t; - if (token >= n_tokens) break; - const float *src = lora + (uint64_t)token * n_head * kv_lora_dim + - (uint64_t)head * kv_lora_dim; - for (uint32_t j = tid; j < kv_lora_dim; j += nth) { - xsh[(uint64_t)t * kv_lora_dim + j] = src[j]; - } - } - __syncthreads(); - - for (uint32_t od = tid; od < value_dim; od += nth) { - const char *row = weight + - ((uint64_t)head * value_dim + od) * row_bytes; - float acc[token_tile] = { 0.0f }; - const uint32_t nb = kv_lora_dim >> 5; - for (uint32_t b = 0; b < nb; b++) { - const char *blk = row + (uint64_t)b * 34u; - const float d = __half2float(*(const __half *)blk); - const int8_t *q = (const int8_t *)(blk + 2); - float s[token_tile] = { 0.0f }; -#pragma unroll 8 - for (uint32_t k = 0; k < 32u; k++) { - const float w = (float)q[k]; -#pragma unroll - for (uint32_t t = 0; t < token_tile; t++) { - if (token0 + t < n_tokens) { - s[t] += w * xsh[(uint64_t)t * kv_lora_dim + - b * 32u + k]; - } - } - } -#pragma unroll - for (uint32_t t = 0; t < token_tile; t++) { - if (token0 + t < n_tokens) acc[t] += d * s[t]; - } - } -#pragma unroll - for (uint32_t t = 0; t < token_tile; t++) { - const uint32_t token = token0 + t; - if (token < n_tokens) { - heads[(uint64_t)token * n_head * value_dim + - (uint64_t)head * value_dim + od] = acc[t]; - } - } - } -} - -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 (!heads || !lora || !model_map || n_tokens == 0 || n_head == 0 || - kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u || value_dim == 0) { - return 0; - } - if (weight_type != 8u) { - fprintf(stderr, "ds4: glm value project: unsupported type %u\n", - weight_type); - return 0; - } - const uint64_t row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u; - const uint64_t wbytes = (uint64_t)n_head * value_dim * row_bytes; - if (weight_offset > model_size || wbytes > model_size - weight_offset || - lora->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || - heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float)) { - return 0; - } - const int logical_tier = cuda_current_tier(); - const char *w = (const char *)cuda_resolve_weight_ptr( - model_map, weight_offset, wbytes, logical_tier, "glm_v_b"); - if (!w) return 0; - if (n_tokens >= 16u && getenv("DS4_GLM_VALUE_NO_TILE16") == NULL) { - dim3 grid(n_head, (n_tokens + 15u) / 16u, 1); - const size_t shmem = 16ull * kv_lora_dim * sizeof(float); - glm_value_project_q8_0_batch_heads_tiled_kernel<16><<>>( - (float *)heads->ptr, w, (const float *)lora->ptr, - n_tokens, n_head, kv_lora_dim, value_dim, row_bytes); - return cuda_ok(cudaGetLastError(), "glm value project tile16 launch"); - } - dim3 grid(n_head, n_tokens, 1); - const size_t shmem = (size_t)kv_lora_dim * sizeof(float); - glm_value_project_q8_0_batch_heads_kernel<<>>( - (float *)heads->ptr, w, (const float *)lora->ptr, - n_tokens, n_head, kv_lora_dim, value_dim, row_bytes); - return cuda_ok(cudaGetLastError(), "glm value project launch"); -} - -/* Decode-time (n_tok small) quant matvec. The Metal "mpp/model-view" - * variant is a bandwidth-tuned matvec; on CUDA the generic quant matmul - * already dispatches per type, so delegate. Revisit in the perf pass. */ -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); - -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) { - 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_matmul_quant_rows_scalar_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) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_matmul_quant_rows_scalar_tensor\n"); - return 0; -} - -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) { - switch (weight_type) { - case 8u: /* Q8_0 */ - return ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, - weight_offset, in_dim, out_dim, - x, n_tok); - case 1u: /* F16 */ - return ds4_gpu_matmul_f16_tensor(out, model_map, model_size, - weight_offset, in_dim, out_dim, - x, n_tok); - default: - fprintf(stderr, "ds4: matmul_quant: unsupported type %u\n", - weight_type); - return 0; - } -} - -extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) { - /* GLM graph memory guard: on this backend the model weights are - * distributed across all devices by the multi-tier placement, so the - * relevant budget is the aggregate VRAM. */ - int n = 0; - if (cudaGetDeviceCount(&n) != cudaSuccess || n <= 0) return 0; - size_t free_b = 0, total_b = 0; - if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) return 0; - return (uint64_t)total_b * (uint64_t)n; -} - -extern "C" int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) { - (void)selected; - (void)n_selected; - return 1; -} - -extern "C" void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { - (void)enabled; /* SSD streaming is not used on the CUDA backend */ -} - -extern "C" void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { - g_glm_mtp_verify_mode = enabled; -} - -extern "C" int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes) { - (void)max_tensor_bytes; - if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) { - return 0; - } - for (uint32_t i = 0; i < count; i++) { - if (offsets[i] > model_size || sizes[i] == 0 || - sizes[i] > model_size - offsets[i]) { - return 0; - } - } - if (!ds4_gpu_set_model_map(model_map, model_size)) return 0; - if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL) { - for (uint32_t i = 0; i < count; i++) { - (void)cuda_model_prefetch_range(model_map, model_size, - offsets[i], sizes[i]); - } - } - return 1; -} - -extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( - ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, - const void *model_map, uint64_t model_size, - uint64_t gate_offset, uint64_t up_offset, - uint64_t in_dim, uint64_t out_dim, - const ds4_gpu_tensor *x, uint64_t n_tok, float clamp); - -extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate, up, mid, - model_map, model_size, gate_offset, up_offset, - in_dim, out_dim, x, 1, clamp); -} - -extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok, - float clamp) { - (void)gate; (void)up; (void)mid; (void)model_map; (void)model_size; - (void)gate_offset; (void)up_offset; (void)in_dim; (void)out_dim; - (void)x; (void)n_tok; (void)clamp; - return 0; -} - -/* Fused single-token shared-expert gate+up+swiglu: one warp per output - * row computes both q8_0 dots against a shared-staged f32 x and writes - * silu(gate)*up directly. Falls back to the split path for n_tok > 1. */ -__global__ static void glm_shared_gate_up_swiglu_one_kernel( - float *mid, - const char *gw, - const char *uw, - const float *x, - uint32_t in_dim, - uint32_t out_dim, - float clamp) { - extern __shared__ float glm_sgu_sh[]; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) { - glm_sgu_sh[i] = x[i]; - } - __syncthreads(); - const uint32_t r = blockIdx.x * warps + warp; - if (r >= out_dim) return; - const uint32_t nblk = in_dim >> 5; - const uint64_t row_bytes = (uint64_t)nblk * 34u; - const char *grow = gw + (uint64_t)r * row_bytes; - const char *urow = uw + (uint64_t)r * row_bytes; - float g = 0.0f, u = 0.0f; - for (uint32_t blk = lane; blk < nblk; blk += 32u) { - const char *gb = grow + (uint64_t)blk * 34u; - const char *ub = urow + (uint64_t)blk * 34u; - const float gd = __half2float(*(const __half *)gb); - const float ud = __half2float(*(const __half *)ub); - const int8_t *gq = (const int8_t *)(gb + 2); - const int8_t *uq = (const int8_t *)(ub + 2); - const float *xs = glm_sgu_sh + blk * 32u; - float gs = 0.0f, us = 0.0f; - #pragma unroll 8 - for (int k = 0; k < 32; k++) { - gs += (float)gq[k] * xs[k]; - us += (float)uq[k] * xs[k]; - } - g += gd * gs; - u += ud * us; - } - for (int off = 16; off > 0; off >>= 1) { - g += __shfl_down_sync(0xffffffffu, g, off); - u += __shfl_down_sync(0xffffffffu, u, off); - } - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (g > clamp) g = clamp; - if (u > clamp) u = clamp; - if (u < -clamp) u = -clamp; - } - mid[r] = (g / (1.0f + expf(-g))) * u; - } -} - -/* Two-token verifier variant of the decode kernel above. Each warp loads a - * gate/up weight row once, while each token keeps the decode kernel's block - * order and warp reduction tree independently. */ -__global__ static void glm_shared_gate_up_swiglu_tok2_exact_kernel( - float *mid, - const char *gw, - const char *uw, - const float *x, - uint32_t in_dim, - uint32_t out_dim, - float clamp) { - extern __shared__ float glm_sgu2_sh[]; - float *x0 = glm_sgu2_sh; - float *x1 = glm_sgu2_sh + in_dim; - const uint32_t warps = blockDim.x >> 5; - const uint32_t warp = threadIdx.x >> 5; - const uint32_t lane = threadIdx.x & 31u; - for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) { - x0[i] = x[i]; - x1[i] = x[in_dim + i]; - } - __syncthreads(); - const uint32_t r = blockIdx.x * warps + warp; - if (r >= out_dim) return; - const uint32_t nblk = in_dim >> 5; - const uint64_t row_bytes = (uint64_t)nblk * 34u; - const char *grow = gw + (uint64_t)r * row_bytes; - const char *urow = uw + (uint64_t)r * row_bytes; - float g0 = 0.0f, u0 = 0.0f; - float g1 = 0.0f, u1 = 0.0f; - for (uint32_t blk = lane; blk < nblk; blk += 32u) { - const char *gb = grow + (uint64_t)blk * 34u; - const char *ub = urow + (uint64_t)blk * 34u; - const float gd = __half2float(*(const __half *)gb); - const float ud = __half2float(*(const __half *)ub); - const int8_t *gq = (const int8_t *)(gb + 2); - const int8_t *uq = (const int8_t *)(ub + 2); - const float *xs0 = x0 + blk * 32u; - const float *xs1 = x1 + blk * 32u; - float gs0 = 0.0f, us0 = 0.0f; - float gs1 = 0.0f, us1 = 0.0f; - #pragma unroll 8 - for (int k = 0; k < 32; k++) { - const float gk = (float)gq[k]; - const float uk = (float)uq[k]; - gs0 += gk * xs0[k]; - us0 += uk * xs0[k]; - gs1 += gk * xs1[k]; - us1 += uk * xs1[k]; - } - g0 += gd * gs0; - u0 += ud * us0; - g1 += gd * gs1; - u1 += ud * us1; - } - for (int off = 16; off > 0; off >>= 1) { - g0 += __shfl_down_sync(0xffffffffu, g0, off); - u0 += __shfl_down_sync(0xffffffffu, u0, off); - g1 += __shfl_down_sync(0xffffffffu, g1, off); - u1 += __shfl_down_sync(0xffffffffu, u1, off); - } - if (lane == 0u) { - if (clamp > 1.0e-6f) { - if (g0 > clamp) g0 = clamp; - if (u0 > clamp) u0 = clamp; - if (u0 < -clamp) u0 = -clamp; - if (g1 > clamp) g1 = clamp; - if (u1 > clamp) u1 = clamp; - if (u1 < -clamp) u1 = -clamp; - } - mid[r] = (g0 / (1.0f + expf(-g0))) * u0; - mid[out_dim + r] = (g1 / (1.0f + expf(-g1))) * u1; - } -} - -extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok, - float clamp) { - if (!gate || !up || !mid || !x || n_tok == 0) return 0; - if (n_tok == 2 && (in_dim & 31u) == 0u && - g_glm_mtp_verify_mode && - getenv("DS4_GLM_MTP_NO_SHARED_TOK2") == NULL && - mid->bytes >= 2u * out_dim * sizeof(float) && - x->bytes >= 2u * in_dim * sizeof(float)) { - const uint64_t row_bytes = (in_dim / 32u) * 34u; - const uint64_t wb = out_dim * row_bytes; - if (gate_offset <= model_size && wb <= model_size - gate_offset && - up_offset <= model_size && wb <= model_size - up_offset) { - const int logical_tier = cuda_current_tier(); - const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset, - wb, logical_tier, "glm_shared_gate"); - const char *uw = cuda_resolve_weight_ptr(model_map, up_offset, - wb, logical_tier, "glm_shared_up"); - if (gw && uw) { - const uint32_t warps = 8u; - const uint32_t sh = 2u * (uint32_t)in_dim * sizeof(float); - glm_shared_gate_up_swiglu_tok2_exact_kernel - <<<(unsigned)((out_dim + warps - 1u) / warps), - warps * 32u, sh>>>( - (float *)mid->ptr, gw, uw, (const float *)x->ptr, - (uint32_t)in_dim, (uint32_t)out_dim, clamp); - return cuda_ok(cudaGetLastError(), - "glm shared swiglu tok2 exact"); - } - } - } - if (n_tok == 1 && (in_dim & 31u) == 0u && - !getenv("DS4_GLM_SHARED_SPLIT") && - mid->bytes >= out_dim * sizeof(float) && - x->bytes >= in_dim * sizeof(float)) { - const uint64_t row_bytes = (in_dim / 32u) * 34u; - const uint64_t wb = out_dim * row_bytes; - if (gate_offset <= model_size && wb <= model_size - gate_offset && - up_offset <= model_size && wb <= model_size - up_offset) { - const int logical_tier = cuda_current_tier(); - const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset, - wb, logical_tier, "glm_shared_gate"); - const char *uw = cuda_resolve_weight_ptr(model_map, up_offset, - wb, logical_tier, "glm_shared_up"); - if (gw && uw) { - const uint32_t warps = 8u; - const uint32_t sh = (uint32_t)in_dim * sizeof(float); - glm_shared_gate_up_swiglu_one_kernel - <<<(unsigned)((out_dim + warps - 1u) / warps), - warps * 32u, sh>>>( - (float *)mid->ptr, gw, uw, (const float *)x->ptr, - (uint32_t)in_dim, (uint32_t)out_dim, clamp); - return cuda_ok(cudaGetLastError(), "glm shared swiglu one"); - } - } - } - if (!ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size, gate_offset, - in_dim, out_dim, x, n_tok) || - !ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size, up_offset, - in_dim, out_dim, x, n_tok)) { - return 0; - } - return ds4_gpu_swiglu_tensor(mid, gate, up, - (uint32_t)(out_dim * n_tok), clamp, 1.0f); -} - -extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_tensor( - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - static ds4_gpu_tensor *gu_scratch[DS4_MAX_GPUS][2] = {{0}}; - const int dev = cuda_current_tier(); - const uint64_t need = out_dim * sizeof(float); - for (int i = 0; i < 2; i++) { - if (!gu_scratch[dev][i] || gu_scratch[dev][i]->bytes < need) { - if (gu_scratch[dev][i]) ds4_gpu_tensor_free(gu_scratch[dev][i]); - gu_scratch[dev][i] = ds4_gpu_tensor_alloc(need); - } - if (!gu_scratch[dev][i]) return 0; - } - return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( - gu_scratch[dev][0], gu_scratch[dev][1], mid, - model_map, model_size, gate_offset, up_offset, - in_dim, out_dim, x, 1, clamp); -} - -extern "C" int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { - if (event_value) *event_value = 1; - return cuda_ok(cudaDeviceSynchronize(), "selected readback signal"); -} - -extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_selected) { - return cuda_stream_selected_cache_begin_load(table, selected_ids, - n_selected); -} - -extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - (void)gate_expert_bytes; - (void)down_expert_bytes; - return 0; -} - -extern "C" int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, - const ds4_gpu_tensor *src, uint64_t src_offset, - uint64_t count) { - if (!dst || !src) return 0; - if (count == 0) return 1; - if (count > UINT64_MAX / sizeof(float) || - count > UINT64_MAX / sizeof(__half)) { - return 0; - } - const uint64_t src_bytes = count * sizeof(float); - const uint64_t dst_bytes = count * sizeof(__half); - if (src_offset > src->bytes || src_bytes > src->bytes - src_offset || - dst_offset > dst->bytes || dst_bytes > dst->bytes - dst_offset || - ds4_tensor_device_idx(dst) != ds4_tensor_device_idx(src)) { - return 0; - } - const int tier = ds4_tensor_device_idx(dst); - if (ds4_gpu_set_current_device(tier) != 0) return 0; - const uint64_t blocks = (count + 255u) / 256u; - if (blocks > UINT32_MAX) return 0; - f32_to_f16_kernel<<<(unsigned)blocks, 256>>>( - (__half *)((char *)dst->ptr + dst_offset), - (const float *)((const char *)src->ptr + src_offset), - count); - return cuda_ok(cudaGetLastError(), "tensor f32-to-f16 copy launch"); -} - -extern "C" int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor, - uint64_t offset, - void *data, - uint64_t bytes, - uint64_t event_value, - const char *label) { - (void)event_value; - if (!tensor || !data || offset > tensor->bytes || - bytes > tensor->bytes - offset) { - return 0; - } - if (!cuda_ok(cudaDeviceSynchronize(), - label ? label : "selected readback wait")) { - return 0; - } - return cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset, - (size_t)bytes, cudaMemcpyDeviceToHost), - "selected tensor read"); -} - -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) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_tp_big_gate_encode\n"); - return 0; -} - -extern "C" int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) { - fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_tp_gate_encode\n"); - return 0; -} - -extern "C" void ds4_gpu_tp_set_attn_head_split(int enabled) { - (void)enabled; /* Mac network-TP head split: no-op on CUDA */ -} - -extern "C" int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) { - (void)event_value; - return cuda_ok(cudaDeviceSynchronize(), - label ? label : "selected readback wait"); -} - -/* Compatibility surface shared with the canonical Metal/ROCm graph. CUDA - * either delegates to its equivalent primitive or reports an unavailable - * optional fast path so the graph can use its established fallback. */ -extern "C" int ds4_gpu_commit_and_wait_selected_readback( - uint64_t event_value, const char *label) { - (void)event_value; - return cuda_ok(cudaDeviceSynchronize(), - label ? label : "selected readback wait"); -} - -extern "C" int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) { - const int ok = ds4_gpu_set_model_fd(fd); - if (ok) g_model_fd_host_base = model_map; - return ok; -} - -extern "C" int ds4_gpu_pro_q4_expert_table_auto_available(void) { - return 0; -} - -extern "C" int ds4_gpu_preload_q4_expert_tables( - const void *model_map, uint64_t model_size, - uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, - uint64_t gate_expert_bytes, uint64_t down_expert_bytes, - uint32_t n_total_expert) { - (void)model_map; (void)model_size; - (void)gate_offset; (void)up_offset; (void)down_offset; - (void)gate_expert_bytes; (void)down_expert_bytes; - (void)n_total_expert; - return 1; -} - -extern "C" void ds4_gpu_set_glm_model(bool enabled) { - (void)enabled; -} - -extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { - g_ssd_streaming_mode = enabled ? 1 : 0; - cuda_stream_selected_cache_invalidate(); - if (!g_ssd_streaming_mode) cuda_stream_selected_cache_release(); -} - -extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { - (void)experts; -} - -extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { - (void)bytes; -} - -extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { - return 0; -} - -extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) { - return g_stream_selected_cache.valid ? - g_stream_selected_cache.compact_count : 0; -} - -extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { -} - -extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) { - cuda_stream_selected_cache_release(); -} - -extern "C" int ds4_gpu_stream_expert_cache_seed_selected( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_selected) { - (void)table; (void)selected_ids; (void)n_selected; - return 1; -} - -extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_tokens, - uint32_t n_selected) { - if (n_tokens == 0 || n_selected == 0 || - (uint64_t)n_tokens * n_selected > UINT32_MAX) { - return 0; - } - return cuda_stream_selected_cache_begin_load( - table, selected_ids, n_tokens * n_selected); -} - -extern "C" int ds4_gpu_stream_expert_cache_seed_experts( - const ds4_gpu_stream_expert_table *table, - const int32_t *expert_ids, - const uint32_t *expert_priorities, - uint32_t n_experts) { - (void)table; (void)expert_ids; (void)expert_priorities; (void)n_experts; - return 1; -} - -extern "C" int ds4_gpu_argmax_tensor( - ds4_gpu_tensor *out_idx, - const ds4_gpu_tensor *logits, - uint32_t n_vocab) { - return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1u, 1u); -} - -extern "C" int ds4_gpu_embed_token_q8_0_tensor( - ds4_gpu_tensor *out, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t token, - uint32_t n_embd) { - return ds4_gpu_embed_token_quant_tensor(out, model_map, model_size, - weight_offset, 8u, n_vocab, - token, n_embd); -} - -extern "C" int ds4_gpu_embed_tokens_q8_0_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *tokens, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd) { - return ds4_gpu_embed_tokens_quant_tensor(out, tokens, model_map, - model_size, weight_offset, 8u, - n_vocab, n_tokens, n_embd); -} - -extern "C" int ds4_gpu_glm_k_b_project_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 n_tokens, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t n_head) { - return ds4_gpu_glm_k_b_project_typed_tensor( - out, kv_norm, model_map, model_size, weight_offset, 8u, - n_tokens, kv_lora_dim, qk_nope, n_head); -} - -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) { - if (!x || x_elem_off > x->bytes / sizeof(float) || - k_cnt > x->bytes / sizeof(float) - x_elem_off) { - return 0; - } - ds4_gpu_tensor x_slice = *x; - x_slice.ptr = (char *)x->ptr + x_elem_off * sizeof(float); - x_slice.bytes = k_cnt * sizeof(float); - x_slice.owner = 0; - return ds4_gpu_matmul_q8_0_kslice_rows_tensor( - out, model_map, model_size, weight_offset, - full_in_dim, out_dim, k_off, k_cnt, &x_slice, 1u); -} - -extern "C" int ds4_gpu_matmul_quant_kslice_tensor( - ds4_gpu_tensor *out, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t weight_type, - 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) { - if (weight_type != 8u) return 0; - return ds4_gpu_matmul_q8_0_kslice_tensor( - out, model_map, model_size, weight_offset, - full_in_dim, k_off, k_cnt, out_dim, x, x_elem_off); -} - -extern "C" int ds4_gpu_matmul_q8_0_f16_out_tensor( - ds4_gpu_tensor *out_h, - 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_tok) { - (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; - (void)in_dim; (void)out_dim; (void)x; (void)n_tok; - return 0; -} - -extern "C" int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( - ds4_gpu_tensor *out, ds4_gpu_tensor *q_half, - 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_tok, uint32_t n_head, uint32_t 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) { - (void)out; (void)q_half; (void)model_map; (void)model_size; - (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; - (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0; - (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale; - (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; - (void)eps; - return 0; -} - -extern "C" int ds4_gpu_attention_prefill_raw_heads_range_tensor( - ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, - uint64_t sinks_offset, const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, uint32_t q_row0, uint32_t n_q, - uint32_t n_kv, uint32_t window, uint32_t n_head, - uint32_t head_dim) { - (void)heads; (void)model_map; (void)model_size; (void)sinks_offset; - (void)q; (void)raw_kv; (void)q_row0; (void)n_q; (void)n_kv; - (void)window; (void)n_head; (void)head_dim; - return 0; -} - -extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( - ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, - uint64_t sinks_offset, const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, uint32_t q_row0, uint32_t n_q, - uint32_t n_tokens, uint32_t n_comp, uint32_t window, - uint32_t ratio, uint32_t n_head, uint32_t head_dim) { - (void)heads; (void)model_map; (void)model_size; (void)sinks_offset; - (void)q; (void)raw_kv; (void)comp_kv; (void)comp_kv_f16; - (void)q_row0; (void)n_q; (void)n_tokens; (void)n_comp; (void)window; - (void)ratio; (void)n_head; (void)head_dim; - return 0; -} - -extern "C" int ds4_gpu_attention_output_q8_batch_f16_tensor( - ds4_gpu_tensor *out_h, 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, - uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) { - (void)out_h; (void)low; (void)model_map; (void)model_size; - (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank; - (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; - return 0; -} - -extern "C" int ds4_gpu_attention_output_q4_K_batch_tensor( - ds4_gpu_tensor *out, ds4_gpu_tensor *low, - ds4_gpu_tensor *group_tmp, ds4_gpu_tensor *low_tmp, - const void *model_map, uint64_t model_size, - uint64_t out_a_offset, uint64_t out_b_offset, uint32_t out_b_type, - uint64_t group_dim, uint64_t rank, uint32_t n_groups, - uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) { - (void)out; (void)low; (void)group_tmp; (void)low_tmp; - (void)model_map; (void)model_size; (void)out_a_offset; - (void)out_b_offset; (void)out_b_type; (void)group_dim; (void)rank; - (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; - return 0; -} - -extern "C" int ds4_gpu_attention_output_low_q4_K_slice_tensor( - ds4_gpu_tensor *low, const void *model_map, uint64_t model_size, - uint64_t out_a_offset, uint64_t group_dim, uint64_t rank, - uint32_t group0, uint32_t group_cnt, - const ds4_gpu_tensor *heads) { - (void)low; (void)model_map; (void)model_size; (void)out_a_offset; - (void)group_dim; (void)rank; (void)group0; (void)group_cnt; - (void)heads; - return 0; -} - -extern "C" int ds4_gpu_hc_expand_split_half_tensor( - ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out_h, - const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, - uint32_t n_embd, uint32_t n_hc) { - (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split; - (void)n_embd; (void)n_hc; - return 0; -} - -extern "C" int ds4_gpu_hc_expand_add_split_half_add_tensor( - ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, - const ds4_gpu_tensor *block_add_h, - const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, - uint32_t n_embd, uint32_t n_hc) { - (void)out_hc; (void)block_out; (void)block_add_h; - (void)residual_hc; (void)split; (void)n_embd; (void)n_hc; - return 0; -} - -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_model_residency_skip(int skip) { - (void)skip; -} - -extern "C" uint64_t ds4_gpu_tp_big_gate_kick( - 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_big_gate_wait(uint64_t seq) { - (void)seq; - return 0; -} - -extern "C" int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) { - (void)layer; (void)rows; - return 0; -} -#pragma GCC diagnostic pop +#include "cuda/runtime.inc" +#include "models/deepseek/cuda/dense_attention.inc" +#include "models/deepseek/cuda/control.inc" +#include "cuda/common_dispatch.inc" +#include "models/deepseek/cuda/moe.inc" +#include "models/deepseek/cuda/hc.inc" +#include "cuda/runtime_services.inc" +#include "models/glm/cuda/kernels.inc" diff --git a/ds4_metal.m b/ds4_metal.m index 39046ba7db..4fc16de54a 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -33,39583 +33,16 @@ * requires it. */ -enum { - DS4_METAL_TENSOR_Q4_0 = 2, - DS4_METAL_TENSOR_Q8_0 = 8, - DS4_METAL_TENSOR_Q2_K = 10, - DS4_METAL_TENSOR_Q4_K = 12, - DS4_METAL_TENSOR_Q5_K = 13, - DS4_METAL_TENSOR_Q6_K = 14, - DS4_METAL_TENSOR_Q8_K = 15, - DS4_METAL_TENSOR_IQ2_XXS = 16, -}; - -@class DS4MetalQ4ExpertTable; - -static id g_device; -static id g_queue; -static id g_library; -static id g_batch_cb; -static id g_batch_enc; -static BOOL g_batch_has_work; -static NSMutableArray> *g_pending_cbs; -static id g_selected_readback_event; -static uint64_t g_selected_readback_event_value; -static id g_set_rows_f32_i32_pipeline; -static id g_get_rows_f32_pipeline; -static id g_get_rows_f16_pipeline; -static id g_get_rows_i32_pipeline; -static id g_get_rows_q8_0_pipeline; -static id g_get_rows_q4_0_pipeline; -static id g_get_rows_q4_K_pipeline; -static id g_repeat_f32_pipeline; -static id g_concat_pipeline; -static id g_cpy_f32_f32_pipeline; -static id g_cpy_f32_f16_pipeline; -static id g_cpy_contig_f32_f16_pipeline; -static id g_cpy_f16_f32_pipeline; -static id g_cpy_f16_f16_pipeline; -static id g_cpy_contig_f16_f32_pipeline; -static id g_cpy_contig_f16_f16_pipeline; -static id g_flash_kv_stage_f16_pipeline; -static id g_swiglu_pipeline; -static id g_swiglu_flat_pipeline; -static id g_add_pipeline; -static id g_add2_pipeline; -static id g_add3_pipeline; -static id g_moe_sum6_pipeline; -static id g_moe_sum8_pipeline; -static id g_mul_pipeline; -static id g_rms_norm_pipeline; -static id g_rms_norm_plain_pipeline; -static id g_add_rms_norm_pipeline; -static id g_rms_norm_scale_pipeline; -static id g_dsv4_qkv_rms_norm_pipeline; -static id g_hc_split_sinkhorn_pipeline; -static id g_hc_split_weighted_sum_pipeline; -static id g_hc_split_weighted_sum_norm_pipeline; -static id g_hc_weighted_sum_pipeline; -static id g_hc_weighted_sum_norm_pipeline; -static id g_output_hc_weights4_pipeline; -static id g_hc_expand_pipeline; -static id g_unary_sigmoid_pipeline; -static id g_unary_silu_pipeline; -static id g_unary_softplus_pipeline; -static id g_unary_sqrt_pipeline; -static id g_unary_clamp_pipeline; -static id g_unary_scale_pipeline; -static id g_unary_fill_pipeline; -static id g_unary_fill_f16_pipeline; -static id g_bin_mul_scalar_pipeline; -static id g_bin_div_row_pipeline; -static id g_moe_mul_mv_id_iq2_xxs_pipeline; -static id g_moe_mul_mv_id_iq2_xxs_pair_pipeline; -static id g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; -static id g_moe_mul_mv_id_q2_k_pipeline; -static id g_moe_mul_mv_id_q2_k_sum6_pipeline; -static id g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; -static id g_moe_mul_mv_id_q4_k_pipeline; -static id g_moe_mul_mv_id_q4_k_pair_pipeline; -static id g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_id_q4_k_sum6_pipeline; -static id g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_group_q4_k_sum6_pipeline; -static id g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_group6_q4_k_sum6_pipeline; -static id g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_group8_q4_k_sum6_pipeline; -static id g_moe_mul_mv_group24_q4_k_id_pipeline; -static id g_moe_mul_mv_group24_q4_k_sum6_pipeline; -static id g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline; -static id g_moe_mul_mv_slots6_q2_k_sum6_pipeline; -static id g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_slots6_q4_k_sum6_pipeline; -static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline; -static id g_moe_mul_mv_addr_iq2_xxs_pipeline; -static id g_moe_mul_mv_addr_q2_k_sum6_pipeline; -static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline; -static id g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline; -static id g_moe_stream_expert_cache_validate_pipeline; -static id g_moe_q4_gather_slots6_pipeline; -static id g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_table_q4_k_sum6_pipeline; -static id g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline; -static id g_moe_mul_mv_addr_q4_k_sum6_pipeline; -static id g_moe_table_q4_pair_gate_encoder; -static id g_moe_table_q4_pair_up_encoder; -static id g_moe_table_q4_sum_down_encoder; -static id g_rope_tail_batch_pipeline; -static id g_rope_tail_inplace_pair_pipeline; -static id g_rope_tail_inplace_pair_shared4_pipeline; -static id g_rope_tail_inplace_pair_affine_pipeline; -static id g_dsv4_fp8_kv_quantize_pipeline; -static id g_dsv4_indexer_qat_pipeline; -static id g_dsv4_kv_fp8_store_pipeline; -static id g_dsv4_ratio4_shift_pipeline; -static id g_dsv4_compressor_pack_ratio4_pipeline; -static id g_dsv4_softmax_pool_ratio4_direct_pipeline; -static id g_dsv4_softmax_pool_pipeline; -static id g_soft_max_f32_pipeline; -static id g_soft_max_f32_4_pipeline; -static id g_argsort_f32_i32_desc_pipeline; -static id g_argsort_merge_f32_i32_desc_pipeline; -static id g_sum_rows_f32_f32_pipeline; -static id g_dsv4_topk_mask_pipeline; -static id g_dsv4_topk_mask_scatter_pipeline; -static id g_dsv4_indexer_weighted_sum_pipeline; -static id g_dsv4_indexer_score_one_direct_pipeline; -static id g_dsv4_compressor_store_one_pipeline; -static id g_dsv4_sort_i32_rows_asc_pipeline; -static id g_dsv4_indexed_attention_heads8_pipeline; -static id g_dsv4_indexed_attention_heads8_rb16_pipeline; -static id g_dsv4_softplus_sqrt_pipeline; -static id g_dsv4_router_finalize_one_pipeline; -static id g_dsv4_router_finalize_one_simd_pipeline; -static id g_dsv4_router_finalize_weights_one_simd_pipeline; -static id g_dsv4_router_weights_one_pipeline; -static id g_glm_router_select_one_pipeline; -static id g_glm_kv_lora_rms_norm_pipeline; -static id g_glm_k_b_project_pipeline; -static id g_glm_store_compact_kv_pipeline; -static id g_glm_qkv_norm_store_compact_kv_pipeline; -static id g_glm_store_indexer_k_pipeline; -static id g_glm_build_kv_cache_pipeline; -static id g_glm_build_kv_cache_decode_group4_pipeline; -static id g_glm_build_kv_cache_flash_pipeline; -static id g_glm_attention_full_pipeline; -static id g_glm_fill_selected_range_pipeline; -static id g_glm_fill_selected_range_batch_pipeline; -static id g_glm_indexer_rope_tail_pipeline; -static id g_glm_indexer_score_one_pipeline; -static id g_glm_indexer_score_one_direct_pipeline; -static id g_glm_indexer_scores_batch_pipeline; -static id g_glm_indexer_scores_tiled_pipeline; -static id g_glm_indexer_scores_tiled_f32_pipeline; -static id g_glm_qk_lowrank_pipeline; -static id g_glm_qk_lowrank_glm52_pipeline; -static id g_glm_qk_lowrank_glm52_sg_pipeline; -static id g_glm_qk_lowrank_batch_pipeline; -static id g_glm_qk_lowrank_batch_glm52_t4_pipeline; -static id g_glm_value_project_q8_0_pipeline; -static id g_glm_value_project_q8_0_batch_heads_pipeline; -static id g_glm_value_project_q8_0_batch_heads_mma_pipeline; -static id g_glm_attention_indexed_decode_pipeline; -static id g_glm_attention_indexed_decode_split_group8_partial_pipeline; -static id g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline; -static id g_glm_attention_indexed_decode_split_group8_reduce_pipeline; -static id g_glm_attention_indexed_decode_split_group8_reduce16_pipeline; -static id g_glm_attention_indexed_batch_pipeline; -static id g_glm_attention_indexed_batch_group2_pipeline; -static id g_glm_attention_indexed_batch_q2_group4_pipeline; -static id g_glm_attention_indexed_batch_group8_pipeline; -static id g_glm_attention_indexed_batch_lora_group8_vec_pipeline; -static id g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline; -static id g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline; -static id g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline; -static id g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline; -static id g_glm_q4_k_pair_swiglu_f32_pipeline; -static id g_glm_q4_k_pair_swiglu2_f32_pipeline; -static id g_glm_q4_k_pair_swiglu4_f32_pipeline; -static id g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline; -static id g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline; -static id g_glm_q2_k_pair_swiglu_f32_pipeline; -static id g_glm_q2_k_addr_pair_swiglu2_f32_pipeline; -static id g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline; -static id g_glm_q4_k_addr_pair_swiglu_f32_pipeline; -static id g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline; -static id g_glm_q2_k_down_f32_pipeline; -static id g_glm_q4_k_down_f32_pipeline; -static id g_glm_q2_k_addr_down_f32_pipeline; -static id g_glm_q4_k_addr_down_f32_pipeline; -static id g_glm_q5_k_pair_swiglu_f32_pipeline; -static id g_glm_q5_k_pair_swiglu_mapped_f32_pipeline; -static id g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline; -static id g_glm_q5_k_down_f32_pipeline; -static id g_glm_q6_k_down_f32_pipeline; -static id g_dsv4_router_weights_batch_pipeline; -static id g_dsv4_hc_expand4_pipeline; -static NSMutableDictionary> *g_pipeline_cache; -static NSMutableDictionary> *g_model_buffer_cache; -static NSMutableDictionary *g_q4_expert_table_cache; -static NSMutableDictionary *g_q4_expert_layer_residency_cache; -static NSMutableArray> *g_transient_buffers; -static id g_model_residency_set; - -typedef struct { - id __strong mask; - id __strong blk; - NSUInteger mask_bytes; - NSUInteger blk_bytes; - uint32_t kind; - uint32_t n_tokens; - uint32_t n_comp; - uint32_t n_keys; - uint32_t window; - uint32_t ratio; - uint32_t nqptg; - uint32_t ncpsg; - bool has_kvpad; - bool bc_mask; - bool valid; - bool blk_ready; -} ds4_gpu_zero_prefix_prefill_mask_cache_entry; - -enum { - DS4_GPU_PREFILL_MASK_CACHE_RAW = 1, - DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2, - DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3, - DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3, -}; - -static ds4_gpu_zero_prefix_prefill_mask_cache_entry - g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS]; -static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void); -static id g_flash_attn_mask_buffer; -static id g_flash_attn_zero_mask_buffer; -static id g_flash_attn_pad_buffer; -static id g_flash_attn_tmp_buffer; -static id g_flash_attn_blk_buffer; -static id g_flash_attn_ring_buffer; -static id g_flash_attn_kv_buffer; -static id g_glm_flash_attn_mask_buffer; -static id g_compressor_pool_kv_buffer; -static id g_compressor_pool_score_buffer; -static id g_compressor_pool_score_cont_buffer; -static id g_compressor_pool_softmax_buffer; -static id g_compressor_pool_product_buffer; -static id g_compressor_store_ape_buffer; -static id g_compressor_store_score_buffer; -static id g_embed_rows_buffer; -static id g_router_selection_buffer; -static id g_router_weight_sum_buffer; -static id g_indexer_head_scores_buffer; -static id g_indexer_topk_buffer; -static id g_indexed_topk_buffer; -static id g_f16_round_scratch_buffer; -static id g_raw_store_round_buffer; -static id g_moe_gate_scratch_buffer; -static id g_moe_down_scratch_buffer; -static id g_moe_id_map_buffer; -static id g_moe_q4_gate_slots_buffer; -static id g_moe_q4_up_slots_buffer; -static id g_moe_q4_down_slots_buffer; -static id g_attn_out_group_ids_buffer; -static int g_model_fd = -1; -static const void *g_model_map_ptr; -static uint64_t g_model_map_size; -static uint64_t g_model_mapped_offset; -static uint64_t g_model_mapped_size; -static uint64_t g_model_mapped_max_tensor_bytes; -static uint64_t g_tensor_alloc_live_bytes; -static uint64_t g_tensor_alloc_peak_bytes; -static pthread_mutex_t g_tensor_mu = PTHREAD_MUTEX_INITIALIZER; -static uintptr_t *g_tensor_live_slots; -static size_t g_tensor_live_cap; -static size_t g_tensor_live_count; -static size_t g_tensor_live_tombs; -static uint64_t g_model_wrap_count; -static uint64_t g_model_wrap_bytes; -static uint64_t g_model_wrap_max_bytes; -static uint64_t g_model_buffer_cache_bytes; -static uint64_t g_model_buffer_cache_evictions; -static int g_model_buffer_cache_over_limit; -static uint64_t g_stream_expert_cache_bytes; -static uint64_t g_stream_expert_cache_expert_bytes; -static uint32_t g_stream_expert_cache_entry_count; -static uint32_t g_stream_expert_cache_budget_override; -static uint64_t g_stream_expert_cache_hits; -static uint64_t g_stream_expert_cache_misses; -static uint64_t g_stream_expert_cache_evictions; -static uint64_t g_stream_expert_cache_wraps; -static uint64_t g_stream_expert_cache_clock; -static uint64_t g_stream_expert_cache_evict_advise_bytes; -static uint64_t g_stream_expert_cache_willneed_advise_bytes; -static uint64_t g_stream_expert_cache_pread_bytes; -static double g_stream_expert_cache_pread_ms; -static uint64_t g_stream_expert_cache_buffer_allocs; -static uint64_t g_stream_expert_cache_buffer_reuses; -static uint64_t g_stream_expert_cache_decode_tokens; -static uint64_t g_stream_expert_cache_hotness_decay_token; -static uint64_t g_stream_expert_timing_selected_calls; -static double g_stream_expert_timing_selected_read_ms; -static double g_stream_expert_timing_selected_sync_ms; -static double g_stream_expert_timing_selected_copy_ms; -static double g_stream_expert_timing_selected_bind_ms; -static uint64_t g_stream_expert_timing_split_layers; -static uint64_t g_stream_expert_timing_split_resident_experts; -static uint64_t g_stream_expert_timing_split_missing_experts; -static double g_stream_expert_timing_split_resident_ms; -static double g_stream_expert_timing_split_missing_ms; -static double g_stream_expert_timing_split_missing_load_ms; -static double g_stream_expert_timing_split_missing_slot_ms; -static double g_stream_expert_timing_split_missing_prune_ms; -static double g_stream_expert_timing_split_missing_addr_ms; -static double g_stream_expert_timing_split_missing_wait_ms; -static uint64_t g_stream_expert_timing_load_calls; -static double g_stream_expert_timing_load_prepare_ms; -static double g_stream_expert_timing_load_pread_ms; -static double g_stream_expert_timing_load_modify_ms; -static double g_stream_expert_timing_load_install_ms; -static uint64_t g_stream_expert_timing_prepare_batch_reuse_calls; -static double g_stream_expert_timing_prepare_batch_reuse_ms; -static uint64_t g_stream_expert_timing_prepare_buffer_calls; -static double g_stream_expert_timing_prepare_buffer_ms; -static uint64_t g_stream_expert_timing_prepare_task_experts; -static double g_stream_expert_timing_prepare_task_ms; -static uint64_t g_stream_expert_timing_reuse_scan_calls; -static uint64_t g_stream_expert_timing_reuse_scan_entries; -static double g_stream_expert_timing_reuse_scan_ms; -static double g_stream_expert_timing_reuse_clear_ms; -static uint64_t g_stream_expert_timing_readahead_calls; -static uint64_t g_stream_expert_timing_readahead_bytes; -static double g_stream_expert_timing_readahead_ms; -static uint64_t g_stream_expert_timing_cache_all_resident_layers; -static uint64_t g_stream_expert_timing_cache_all_missing_layers; -static uint64_t g_stream_expert_timing_cache_mixed_layers; -static uint64_t g_stream_expert_timing_cache_resident_experts; -static uint64_t g_stream_expert_timing_cache_missing_experts; -typedef struct { - uint64_t selected_calls; - double selected_read_ms; - double selected_sync_ms; - double selected_copy_ms; - double selected_bind_ms; - uint64_t split_layers; - uint64_t split_resident_experts; - uint64_t split_missing_experts; - double split_resident_ms; - double split_missing_ms; - double split_missing_load_ms; - double split_missing_slot_ms; - double split_missing_prune_ms; - double split_missing_addr_ms; - double split_missing_wait_ms; - uint64_t load_calls; - double load_prepare_ms; - double load_pread_ms; - double load_modify_ms; - double load_install_ms; - uint64_t prepare_batch_reuse_calls; - double prepare_batch_reuse_ms; - uint64_t prepare_buffer_calls; - double prepare_buffer_ms; - uint64_t prepare_task_experts; - double prepare_task_ms; - uint64_t reuse_scan_calls; - uint64_t reuse_scan_entries; - double reuse_scan_ms; - double reuse_clear_ms; - uint64_t readahead_calls; - uint64_t readahead_bytes; - double readahead_ms; - uint64_t cache_all_resident_layers; - uint64_t cache_all_missing_layers; - uint64_t cache_mixed_layers; - uint64_t cache_resident_experts; - uint64_t cache_missing_experts; -} ds4_gpu_stream_expert_timing_snapshot; -static ds4_gpu_stream_expert_timing_snapshot g_stream_expert_timing_last_report; -static int g_stream_prefill_batch_selected_addr_building; -static int g_glm_stream_expert_addr_table_building; -static uint64_t g_model_residency_count; -static int g_model_residency_added_to_queue; -static int g_glm_model_mode; -static int g_ssd_streaming_mode; -static int g_glm_streaming_prefill_full_layer_runtime; -static int g_metal4_runtime_available; -static int g_metal4_family_supported; -static int g_metal4_queue_supported; -static int g_metal4_m5_neural_accelerators_hint; -static int g_metal4_tensor_api_enabled; -static int g_metal4_tensor_api_compile_supported; -static char g_metal_device_name[128]; -static int ds4_gpu_model_map_log_enabled(void); -static int ds4_gpu_stream_expert_cache_note_expert_size( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes); -static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void); -static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats); -static void ds4_gpu_stream_expert_pending_load_clear(void); -static void ds4_gpu_stream_expert_pread_pool_shutdown(void); -static int ds4_gpu_stream_expert_timing_summary_enabled(void); -static int ds4_gpu_stream_expert_cache_entry_protected( - uint32_t layer, - uint32_t expert, - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect); - -/* The async selected-load worker registers itself so cache paths that would - * flush/wait on command buffers (a race against the encoding thread) fail - * the load instead; the caller then retries on the main thread. */ -static pthread_t g_stream_expert_service_thread; -static int g_stream_expert_service_thread_set; - -void ds4_gpu_stream_expert_cache_note_service_thread(void) { - g_stream_expert_service_thread = pthread_self(); - g_stream_expert_service_thread_set = 1; -} - -static int ds4_gpu_stream_expert_cache_on_service_thread(void) { - return g_stream_expert_service_thread_set && - pthread_equal(pthread_self(), g_stream_expert_service_thread); -} -static NSUInteger g_flash_attn_mask_bytes; -static NSUInteger g_flash_attn_zero_mask_bytes; -static NSUInteger g_flash_attn_pad_bytes; -static NSUInteger g_flash_attn_tmp_bytes; -static NSUInteger g_flash_attn_blk_bytes; -static NSUInteger g_flash_attn_ring_bytes; -static NSUInteger g_flash_attn_kv_bytes; -static NSUInteger g_glm_flash_attn_mask_bytes; -static uint32_t g_glm_flash_attn_mask_pos0; -static uint32_t g_glm_flash_attn_mask_tokens; -static uint32_t g_glm_flash_attn_mask_cache_len; -static int g_glm_flash_attn_mask_valid; -static NSUInteger g_compressor_pool_kv_bytes; -static NSUInteger g_compressor_pool_score_bytes; -static NSUInteger g_compressor_pool_score_cont_bytes; -static NSUInteger g_compressor_pool_softmax_bytes; -static NSUInteger g_compressor_pool_product_bytes; -static NSUInteger g_compressor_store_ape_bytes; -static NSUInteger g_compressor_store_score_bytes; -static NSUInteger g_embed_rows_bytes; -static NSUInteger g_router_selection_bytes; -static NSUInteger g_router_weight_sum_bytes; -static NSUInteger g_indexer_head_scores_bytes; -static NSUInteger g_indexer_topk_bytes; -static NSUInteger g_indexed_topk_bytes; -static NSUInteger g_f16_round_scratch_bytes; -static NSUInteger g_raw_store_round_bytes; -static NSUInteger g_moe_gate_scratch_bytes; -static NSUInteger g_moe_down_scratch_bytes; -static NSUInteger g_moe_id_map_bytes; -static NSUInteger g_moe_q4_gate_slots_bytes; -static NSUInteger g_moe_q4_up_slots_bytes; -static NSUInteger g_moe_q4_down_slots_bytes; -static NSUInteger g_attn_out_group_ids_bytes; -static int g_initialized; -static int g_quality_mode; -static int g_mpp_invalid_env_reported; -#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8 -static int32_t g_routed_moe_selected_override[DS4_METAL_MAX_ROUTED_EXPERT_USED]; -static uint32_t g_routed_moe_selected_override_n; -static int g_moe_selected_trace_record_initialized; -static FILE *g_moe_selected_trace_record_fp; -static uint64_t g_moe_selected_trace_record_count; -static int g_moe_selected_trace_replay_initialized; -static int32_t *g_moe_selected_trace_replay_ids; -static uint64_t g_moe_selected_trace_replay_count; -static uint64_t g_moe_selected_trace_replay_pos; - -static double ds4_gpu_gib(uint64_t bytes); - -static uint64_t ds4_gpu_system_memory_bytes(void) { - uint64_t bytes = 0; - size_t len = sizeof(bytes); - if (sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) != 0) return 0; - return len == sizeof(bytes) ? bytes : 0; -} - -static void ds4_gpu_print_device_summary(void) { - const char *name = g_device.name ? [g_device.name UTF8String] : "unknown Metal device"; - uint64_t mem = ds4_gpu_system_memory_bytes(); - if (mem) { - double gib = (double)mem / 1024.0 / 1024.0 / 1024.0; - fprintf(stderr, "ds4: Metal device %s, %.2f GiB RAM\n", name, gib); - } else { - fprintf(stderr, "ds4: Metal device %s\n", name); - } -} - -#define DS4_METAL_MAX_MODEL_VIEWS 4096 -/* Compatibility fallback for callers that cannot provide a parsed GGUF tensor - * span. The normal DS4 engine passes the exact maximum tensor byte size. */ -#define DS4_METAL_FALLBACK_MAX_TENSOR_BYTES (4ull * 1024ull * 1024ull * 1024ull) - -typedef struct { - __strong id buffer; - const void *model_map; - uint64_t model_size; - uint64_t model_offset; - uint64_t bytes; -} ds4_gpu_model_view; - -static ds4_gpu_model_view g_model_views[DS4_METAL_MAX_MODEL_VIEWS]; -static uint32_t g_model_view_count; - -enum { - DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80, - DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384, - DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED, - DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES = - DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * - DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT, - DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS = 256, - DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16, - DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS = 16, -}; - -typedef struct { - uint32_t layer; - uint32_t expert; - uint64_t hits; -} ds4_gpu_moe_selected_hotlist_entry; - -static int g_moe_selected_hotlist_initialized; -static uint64_t - g_moe_selected_hotlist_counts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; -static uint64_t g_moe_selected_hotlist_records; -static uint64_t g_moe_selected_hotlist_selections; - -typedef struct { - __strong id gate_buffer; - __strong id up_buffer; - __strong id down_buffer; - const void *model_map; - uint64_t model_size; - uint64_t gate_abs_offset; - uint64_t up_abs_offset; - uint64_t down_abs_offset; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; - uint64_t logical_bytes; - uint64_t last_used; - uint64_t use_count; - NSUInteger gate_inner; - NSUInteger up_inner; - NSUInteger down_inner; - uint64_t inflight_seq; - uint32_t slab_slot; - uint8_t valid; - uint8_t slab_backed; -} ds4_gpu_stream_expert_cache_entry; - -typedef struct { - __strong id gate_buffer; - __strong id up_buffer; - __strong id down_buffer; - NSUInteger gate_inner; - NSUInteger up_inner; - NSUInteger down_inner; -} ds4_gpu_stream_expert_reusable_buffers; - -static ds4_gpu_stream_expert_cache_entry - g_stream_expert_cache[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; -static ds4_gpu_stream_expert_cache_entry - g_stream_full_expert_addr_entry[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint32_t g_stream_expert_cache_layer_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static double g_stream_expert_cache_layer_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_last_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_last_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_last_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint64_t g_stream_expert_cache_layer_last_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static double g_stream_expert_cache_layer_last_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static uint32_t - g_stream_expert_cache_route_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; -static id g_stream_expert_cache_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_expert_cache_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_expert_cache_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_expert_cache_slabs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; -static uint32_t g_stream_expert_cache_slab_start_slot[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; -static uint32_t g_stream_expert_cache_slab_slot_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; -static uint32_t g_stream_expert_cache_slab_slots_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; -static uint32_t g_stream_expert_cache_slab_count; -static uint32_t g_stream_expert_cache_slab_total_slots; -static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES]; -static uint32_t g_stream_expert_cache_free_slot_count; -static uint64_t g_stream_expert_cache_slab_slot_bytes; -static uint64_t g_stream_expert_cache_cb_seq; -static uint64_t g_stream_expert_cache_done_seq; -static uint64_t g_stream_expert_cache_batch_seq; -static uint64_t g_stream_expert_cache_owned_seq; -static uint64_t g_stream_expert_cache_pending_max_seq; -static id g_stream_compact_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_compact_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_compact_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_compact_selected_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_selected_id_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; -static id g_stream_expert_validate_status_buffer; - -@interface DS4MetalTensor : NSObject -@property(nonatomic, strong) id buffer; -@property(nonatomic, assign) uint64_t offset; -@property(nonatomic, assign) uint64_t bytes; -@property(nonatomic, assign) uint8_t owner; -@end - -@implementation DS4MetalTensor -@end - -@interface DS4MetalQ4ExpertTable : NSObject -@property(nonatomic, strong) id argumentBuffer; -@property(nonatomic, strong) id addressBuffer; -@property(nonatomic, strong) NSMutableArray> *expertBuffers; -@property(nonatomic, strong) id residencySet; -@property(nonatomic, assign) BOOL residencySetAddedToQueue; -@property(nonatomic, assign) uint32_t nExpert; -@property(nonatomic, assign) uint64_t expertBytes; -@end - -@implementation DS4MetalQ4ExpertTable -- (void)dealloc { -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - if (_residencySet) { - if (_residencySetAddedToQueue && - g_queue && - [g_queue respondsToSelector:@selector(removeResidencySet:)]) { - [g_queue removeResidencySet:_residencySet]; - } - [_residencySet endResidency]; - } - } -#endif -} -@end - -@interface DS4MetalQ4LayerResidency : NSObject -@property(nonatomic, strong) id residencySet; -@property(nonatomic, assign) BOOL addedToQueue; -@end - -@implementation DS4MetalQ4LayerResidency -- (void)dealloc { -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - if (_residencySet) { - if (_addedToQueue && - g_queue && - [g_queue respondsToSelector:@selector(removeResidencySet:)]) { - [g_queue removeResidencySet:_residencySet]; - } - [_residencySet endResidency]; - } - } -#endif -} -@end - -static DS4MetalTensor *ds4_gpu_tensor_obj(ds4_gpu_tensor *tensor) { - return (__bridge DS4MetalTensor *)tensor; -} - -static const DS4MetalTensor *ds4_gpu_tensor_const_obj(const ds4_gpu_tensor *tensor) { - return (__bridge const DS4MetalTensor *)tensor; -} - -/* C code owns ds4_gpu_tensor handles as retained Objective-C objects. Freeing - * the same opaque handle twice would make the second __bridge_transfer release - * an already-deallocated object, which macOS reports as malloc corruption. The - * live table lets free validate a handle before touching Objective-C state; the - * same mutex also serializes the diagnostic allocation counters. */ -static uint64_t ds4_gpu_tensor_ptr_hash(uintptr_t ptr) { - uint64_t x = (uint64_t)(ptr >> 4); - x ^= x >> 33; - x *= UINT64_C(0xff51afd7ed558ccd); - x ^= x >> 33; - x *= UINT64_C(0xc4ceb9fe1a85ec53); - x ^= x >> 33; - return x; -} - -static int ds4_gpu_tensor_live_resize_locked(size_t min_cap) { - size_t new_cap = 1024; - while (new_cap < min_cap) new_cap <<= 1; - - uintptr_t *new_slots = calloc(new_cap, sizeof(new_slots[0])); - if (!new_slots) return 0; - - for (size_t i = 0; i < g_tensor_live_cap; i++) { - const uintptr_t key = g_tensor_live_slots[i]; - if (key == 0 || key == UINTPTR_MAX) continue; - - size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (new_cap - 1); - while (new_slots[idx] != 0) idx = (idx + 1) & (new_cap - 1); - new_slots[idx] = key; - } - - free(g_tensor_live_slots); - g_tensor_live_slots = new_slots; - g_tensor_live_cap = new_cap; - g_tensor_live_tombs = 0; - return 1; -} - -static int ds4_gpu_tensor_live_insert_locked(const void *ptr) { - if (!ptr || (uintptr_t)ptr == UINTPTR_MAX) return 0; - if ((g_tensor_live_count + g_tensor_live_tombs + 1) * 10 >= - g_tensor_live_cap * 7) - { - const size_t min_cap = g_tensor_live_cap ? g_tensor_live_cap * 2 : 1024; - if (!ds4_gpu_tensor_live_resize_locked(min_cap)) return 0; - } - - const uintptr_t key = (uintptr_t)ptr; - size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); - size_t tomb = (size_t)-1; - for (;;) { - const uintptr_t cur = g_tensor_live_slots[idx]; - if (cur == key) return 0; - if (cur == UINTPTR_MAX) { - if (tomb == (size_t)-1) tomb = idx; - } else if (cur == 0) { - if (tomb != (size_t)-1) { - idx = tomb; - g_tensor_live_tombs--; - } - g_tensor_live_slots[idx] = key; - g_tensor_live_count++; - return 1; - } - idx = (idx + 1) & (g_tensor_live_cap - 1); - } -} - -static int ds4_gpu_tensor_live_remove_locked(const void *ptr) { - if (!ptr || g_tensor_live_cap == 0) return 0; - - const uintptr_t key = (uintptr_t)ptr; - size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); - for (;;) { - const uintptr_t cur = g_tensor_live_slots[idx]; - if (cur == 0) return 0; - if (cur == key) { - g_tensor_live_slots[idx] = UINTPTR_MAX; - g_tensor_live_count--; - g_tensor_live_tombs++; - return 1; - } - idx = (idx + 1) & (g_tensor_live_cap - 1); - } -} - -static int ds4_gpu_tensor_track_alloc_locked( - const void *ptr, - uint64_t bytes, - uint64_t *live_snap, - uint64_t *peak_snap) -{ - if (!ds4_gpu_tensor_live_insert_locked(ptr)) return 0; - - g_tensor_alloc_live_bytes += bytes; - if (g_tensor_alloc_live_bytes > g_tensor_alloc_peak_bytes) { - g_tensor_alloc_peak_bytes = g_tensor_alloc_live_bytes; - } - if (live_snap) *live_snap = g_tensor_alloc_live_bytes; - if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; - return 1; -} - -static int ds4_gpu_tensor_track_view_locked(const void *ptr) { - return ds4_gpu_tensor_live_insert_locked(ptr); -} - -static int ds4_gpu_tensor_prepare_free( - ds4_gpu_tensor *tensor, - uint8_t *owner, - uint64_t *bytes, - uint64_t *live_snap, - uint64_t *peak_snap) -{ - pthread_mutex_lock(&g_tensor_mu); - if (!ds4_gpu_tensor_live_remove_locked(tensor)) { - pthread_mutex_unlock(&g_tensor_mu); - fprintf(stderr, - "ds4: Metal tensor free ignored for unknown handle %p\n", - (void *)tensor); - return 0; - } - - DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); - const uint8_t obj_owner = obj.owner; - const uint64_t obj_bytes = obj.bytes; - if (obj_owner) { - if (obj_bytes <= g_tensor_alloc_live_bytes) { - g_tensor_alloc_live_bytes -= obj_bytes; - } else { - g_tensor_alloc_live_bytes = 0; - } - } - if (owner) *owner = obj_owner; - if (bytes) *bytes = obj_bytes; - if (live_snap) *live_snap = g_tensor_alloc_live_bytes; - if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; - pthread_mutex_unlock(&g_tensor_mu); - return 1; -} - -static void ds4_gpu_tensor_tracking_reset(void) { - pthread_mutex_lock(&g_tensor_mu); - if (g_tensor_live_count != 0) { - fprintf(stderr, - "ds4: Metal cleanup discarded %zu live tensor handles\n", - g_tensor_live_count); - } - free(g_tensor_live_slots); - g_tensor_live_slots = NULL; - g_tensor_live_cap = 0; - g_tensor_live_count = 0; - g_tensor_live_tombs = 0; - g_tensor_alloc_live_bytes = 0; - g_tensor_alloc_peak_bytes = 0; - pthread_mutex_unlock(&g_tensor_mu); -} - -static id ds4_gpu_tensor_buffer(const ds4_gpu_tensor *tensor) { - if (!tensor) return nil; - const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); - return obj.buffer; -} - -static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) { - if (!tensor) return 0; - const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); - return (NSUInteger)obj.offset; -} - -static id ds4_gpu_new_command_buffer(void); -static void ds4_gpu_stream_expert_cache_note_owned_created(void); - -static id ds4_gpu_command_buffer(int *owned) { - if (g_batch_cb) { - *owned = 0; - return g_batch_cb; - } - *owned = 1; - id cb = ds4_gpu_new_command_buffer(); - if (cb) ds4_gpu_stream_expert_cache_note_owned_created(); - return cb; -} - -static id ds4_gpu_compute_encoder(id cb) { - if (g_batch_cb && cb == g_batch_cb) { - g_batch_has_work = YES; - if (!g_batch_enc) g_batch_enc = [cb computeCommandEncoder]; - return g_batch_enc; - } - return [cb computeCommandEncoder]; -} - -static void ds4_gpu_end_compute_encoder(id cb, id enc) { - if (!enc) return; - if (g_batch_cb && cb == g_batch_cb && enc == g_batch_enc) return; - [enc endEncoding]; -} - -static void ds4_gpu_close_batch_encoder(void) { - if (!g_batch_enc) return; - [g_batch_enc endEncoding]; - g_batch_enc = nil; -} - -static double g_gpu_busy_accum; -static uint64_t g_gpu_busy_cbs; - -static int ds4_gpu_wait_command_buffer(id cb, const char *label) { - [cb waitUntilCompleted]; - if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) { - const double busy = cb.GPUEndTime - cb.GPUStartTime; - if (busy > 0) g_gpu_busy_accum += busy; - if ((++g_gpu_busy_cbs % 64u) == 0u) { - fprintf(stderr, "ds4: gpu busy accum %.1f ms over %llu cbs\n", - g_gpu_busy_accum * 1000.0, - (unsigned long long)g_gpu_busy_cbs); - } - } - if (cb.status == MTLCommandBufferStatusError) { - fprintf(stderr, "ds4: Metal %s failed: %s\n", - label, [[cb.error localizedDescription] UTF8String]); - return 0; - } - return 1; -} - -static id ds4_gpu_new_command_buffer(void) { - static int initialized; - static int use_unretained; - if (!initialized) { - use_unretained = getenv("DS4_METAL_UNRETAINED_COMMAND_BUFFERS") != NULL; - initialized = 1; - } - if (use_unretained) { - return [g_queue commandBufferWithUnretainedReferences]; - } - return [g_queue commandBuffer]; -} - -static uint64_t ds4_gpu_exact_view_cache_limit_bytes(void) { - static int initialized; - static uint64_t limit_bytes; - if (initialized) return limit_bytes; - - const uint64_t mib = 1024ull * 1024ull; - const uint64_t gib = 1024ull * mib; - limit_bytes = 64ull * gib; - - const char *gib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_GIB"); - if (gib_env && gib_env[0]) { - char *end = NULL; - unsigned long long v = strtoull(gib_env, &end, 10); - if (end != gib_env && *end == '\0') { - limit_bytes = v > UINT64_MAX / gib ? UINT64_MAX : (uint64_t)v * gib; - } - } - - const char *mib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_MIB"); - if (mib_env && mib_env[0]) { - char *end = NULL; - unsigned long long v = strtoull(mib_env, &end, 10); - if (end != mib_env && *end == '\0') { - limit_bytes = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; - } - } - - initialized = 1; - return limit_bytes; -} - -static void ds4_gpu_model_buffer_cache_note_insert(uint64_t bytes) { - if (g_model_buffer_cache_bytes > UINT64_MAX - bytes) { - g_model_buffer_cache_bytes = UINT64_MAX; - } else { - g_model_buffer_cache_bytes += bytes; - } - - const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); - if (limit != 0 && g_model_buffer_cache_bytes > limit) { - g_model_buffer_cache_over_limit = 1; - } -} - -static void ds4_gpu_model_buffer_cache_clear(const char *reason) { - if (!g_model_buffer_cache) { - g_model_buffer_cache_bytes = 0; - g_model_buffer_cache_over_limit = 0; - return; - } - - const NSUInteger entries = [g_model_buffer_cache count]; - if (entries != 0) { - if (getenv("DS4_METAL_EXACT_VIEW_CACHE_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal exact model view cache evict reason=%s entries=%lu bytes=%.2f GiB limit=%.2f GiB\n", - reason ? reason : "unknown", - (unsigned long)entries, - ds4_gpu_gib(g_model_buffer_cache_bytes), - ds4_gpu_gib(ds4_gpu_exact_view_cache_limit_bytes())); - } - [g_model_buffer_cache removeAllObjects]; - g_model_buffer_cache_evictions++; - } - g_model_buffer_cache_bytes = 0; - g_model_buffer_cache_over_limit = 0; -} - -static void ds4_gpu_model_buffer_cache_maybe_evict(const char *reason) { - if (g_model_buffer_cache_over_limit) { - ds4_gpu_model_buffer_cache_clear(reason); - } -} - -static uint64_t ds4_gpu_stream_expert_cache_next_cb_seq(void) { - if (g_stream_expert_cache_cb_seq == UINT64_MAX) { - /* - * A real wrap would require an astronomical number of command buffers. - * Resetting the epoch space is still safer than letting zero become a - * valid in-flight marker. - */ - g_stream_expert_cache_cb_seq = 0; - g_stream_expert_cache_done_seq = 0; - g_stream_expert_cache_batch_seq = 0; - g_stream_expert_cache_owned_seq = 0; - g_stream_expert_cache_pending_max_seq = 0; - for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { - for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { - g_stream_expert_cache[layer][expert].inflight_seq = 0; - } - } - } - return ++g_stream_expert_cache_cb_seq; -} - -static void ds4_gpu_stream_expert_cache_note_batch_created(void) { - g_stream_expert_cache_batch_seq = - ds4_gpu_stream_expert_cache_next_cb_seq(); -} - -static void ds4_gpu_stream_expert_cache_note_batch_committed(void) { - if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_pending_max_seq) { - g_stream_expert_cache_pending_max_seq = g_stream_expert_cache_batch_seq; - } - g_stream_expert_cache_batch_seq = 0; -} - -static void ds4_gpu_stream_expert_cache_note_owned_created(void) { - g_stream_expert_cache_owned_seq = - ds4_gpu_stream_expert_cache_next_cb_seq(); -} - -static void ds4_gpu_stream_expert_cache_note_pending_completed(void) { - if (g_stream_expert_cache_pending_max_seq > g_stream_expert_cache_done_seq) { - g_stream_expert_cache_done_seq = g_stream_expert_cache_pending_max_seq; - } - g_stream_expert_cache_pending_max_seq = 0; -} - -static void ds4_gpu_stream_expert_cache_note_owned_completed(void) { - if (g_stream_expert_cache_owned_seq > g_stream_expert_cache_done_seq) { - g_stream_expert_cache_done_seq = g_stream_expert_cache_owned_seq; - } - g_stream_expert_cache_owned_seq = 0; -} - -static int ds4_gpu_stream_expert_cache_entry_inflight( - const ds4_gpu_stream_expert_cache_entry *e) { - return e && e->valid && e->inflight_seq > g_stream_expert_cache_done_seq; -} - -static int ds4_gpu_stream_expert_cache_mark_inflight( - ds4_gpu_stream_expert_cache_entry *e) { - if (!e || !e->valid) return 0; - const uint64_t seq = g_stream_expert_cache_batch_seq ? - g_stream_expert_cache_batch_seq : - g_stream_expert_cache_owned_seq; - if (seq == 0) return 0; - e->inflight_seq = seq; - return 1; -} - -static int ds4_gpu_stream_expert_cache_mark_entries_inflight( - ds4_gpu_stream_expert_cache_entry * const *entries, - uint32_t n_entries, - uint32_t active_mask) { - if (!entries || n_entries == 0) return 0; - for (uint32_t i = 0; i < n_entries; i++) { - if (active_mask != 0 && (active_mask & (1u << i)) == 0) continue; - if (!ds4_gpu_stream_expert_cache_mark_inflight(entries[i])) return 0; - } - return 1; -} - -static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label); - -static int ds4_gpu_wait_pending_command_buffers(const char *label) { - int ok = 1; - for (id pending in g_pending_cbs) { - if (!ds4_gpu_wait_command_buffer(pending, label)) ok = 0; - } - [g_pending_cbs removeAllObjects]; - ds4_gpu_stream_expert_cache_note_pending_completed(); - if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); - return ok; -} - -static int ds4_gpu_finish_command_buffer(id cb, int owned, const char *label) { - if (!owned) return 1; - - [cb commit]; - int ok = ds4_gpu_wait_pending_command_buffers(label); - if (!ds4_gpu_wait_command_buffer(cb, label)) { - ok = 0; - ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); - } - ds4_gpu_stream_expert_cache_note_owned_completed(); - [g_transient_buffers removeAllObjects]; - ds4_gpu_model_buffer_cache_maybe_evict(label); - return ok; -} - -static int ds4_gpu_device_name_contains(const char *needle); - -static int ds4_gpu_use_m5_private_scratch(void) { - static int initialized; - static int enabled; - if (!initialized) { - enabled = ds4_gpu_device_name_contains("M5"); - initialized = 1; - } - return enabled; -} - -static int ds4_gpu_scratch_needs_cpu_access(const char *label) { - if (!label) return 0; - return strstr(label, "mask") != NULL || - strcmp(label, "ds4_attention_output_group_ids") == 0; -} - -static MTLResourceOptions ds4_gpu_model_resource_options(void) { - MTLResourceOptions options = MTLResourceStorageModeShared; - if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) { - options |= MTLResourceHazardTrackingModeUntracked; - } - return options; -} - -static int ds4_gpu_ensure_scratch_buffer( - id __strong *buffer, - NSUInteger *capacity, - NSUInteger bytes, - const char *label) { - if (*buffer && *capacity >= bytes) return 1; - if (bytes == 0) bytes = 1; - if (bytes > NSUIntegerMax) return 0; - - MTLResourceOptions options = MTLResourceStorageModeShared; - if (ds4_gpu_use_m5_private_scratch() && - !ds4_gpu_scratch_needs_cpu_access(label)) { - /* - * M5 scratch buffers that only flow between Metal kernels do not need - * CPU-visible shared storage. This reduces shared-memory traffic and - * residency pressure for the long prefill scratch pools without - * changing the public buffer lifetime model. Keep default hazard - * tracking because the graph reuses these buffers across dependent - * compute encoders. - */ - options = MTLResourceStorageModePrivate; - } - - *buffer = [g_device newBufferWithLength:bytes options:options]; - if (!*buffer && options != MTLResourceStorageModeShared) { - *buffer = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared]; - } - if (!*buffer) { - fprintf(stderr, "ds4: failed to allocate Metal scratch buffer %s (%llu bytes)\n", - label, (unsigned long long)bytes); - *capacity = 0; - return 0; - } - (*buffer).label = [NSString stringWithUTF8String:label]; - *capacity = bytes; - return 1; -} - -static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) { - const NSUInteger capacity = 8192u * sizeof(uint16_t); - if (bytes > capacity) return 0; - if (g_flash_attn_zero_mask_buffer && - g_flash_attn_zero_mask_bytes >= capacity) { - return 1; - } - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer, - &g_flash_attn_zero_mask_bytes, - capacity, - "ds4_flash_attn_zero_mask")) { - return 0; - } - void *contents = [g_flash_attn_zero_mask_buffer contents]; - if (!contents) return 0; - memset(contents, 0, g_flash_attn_zero_mask_bytes); - return 1; -} - -static uint64_t round_up_u64(uint64_t v, uint64_t align) { - return (v + align - 1) & ~(align - 1); -} - -static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint64_t max_tensor_bytes) { - if (max_tensor_bytes != 0) return max_tensor_bytes; - return map_size < DS4_METAL_FALLBACK_MAX_TENSOR_BYTES ? - map_size : DS4_METAL_FALLBACK_MAX_TENSOR_BYTES; -} - -static id ds4_gpu_get_pipeline(const char *function_name); -static int ds4_gpu_warm_model_views(void); -static double ds4_gpu_gib(uint64_t bytes); - -static double ds4_gpu_now_ms(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0; -} - -static int ds4_gpu_moe_selected_hotlist_cmp(const void *a, const void *b) { - const ds4_gpu_moe_selected_hotlist_entry *ea = a; - const ds4_gpu_moe_selected_hotlist_entry *eb = b; - if (ea->hits < eb->hits) return 1; - if (ea->hits > eb->hits) return -1; - if (ea->layer != eb->layer) return ea->layer < eb->layer ? -1 : 1; - if (ea->expert != eb->expert) return ea->expert < eb->expert ? -1 : 1; - return 0; -} - -static int ds4_gpu_moe_selected_hotlist_merge_requested(void) { - return getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_MERGE") != NULL && - getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_FRESH") == NULL; -} - -static int ds4_gpu_moe_selected_hotlist_load_existing(const char *path) { - if (!path || !path[0] || !ds4_gpu_moe_selected_hotlist_merge_requested()) { - return 1; - } - - FILE *fp = fopen(path, "rb"); - if (!fp) { - if (errno == ENOENT) return 1; - fprintf(stderr, "ds4: failed to open selected hotlist merge file %s\n", path); - return 0; - } - - char line[256]; - uint64_t lineno = 0; - uint64_t loaded_entries = 0; - uint64_t loaded_hits = 0; - uint64_t header_records = UINT64_MAX; - uint64_t header_selections = UINT64_MAX; - while (fgets(line, sizeof(line), fp)) { - lineno++; - char *p = line; - while (*p && isspace((unsigned char)*p)) p++; - if (*p == '\0') continue; - if (*p == '#') { - unsigned long long value = 0; - if (sscanf(p, "# layer_records %llu", &value) == 1) { - header_records = (uint64_t)value; - } else if (sscanf(p, "# selections %llu", &value) == 1) { - header_selections = (uint64_t)value; - } - continue; - } - - errno = 0; - char *end = NULL; - unsigned long layer = strtoul(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - p = end; - while (*p && isspace((unsigned char)*p)) p++; - - errno = 0; - unsigned long expert = strtoul(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - p = end; - while (*p && isspace((unsigned char)*p)) p++; - - errno = 0; - unsigned long long hits = strtoull(p, &end, 10); - if (end == p || errno != 0) goto bad_line; - if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && - hits != 0) { - uint64_t *dst = &g_moe_selected_hotlist_counts[layer][expert]; - if (*dst > UINT64_MAX - (uint64_t)hits) { - *dst = UINT64_MAX; - } else { - *dst += (uint64_t)hits; - } - loaded_entries++; - if (loaded_hits > UINT64_MAX - (uint64_t)hits) { - loaded_hits = UINT64_MAX; - } else { - loaded_hits += (uint64_t)hits; - } - } - continue; - -bad_line: - fprintf(stderr, - "ds4: invalid selected hotlist merge line %" PRIu64 " in %s\n", - lineno, - path); - fclose(fp); - return 0; - } - if (ferror(fp)) { - fprintf(stderr, "ds4: failed to read selected hotlist merge file %s\n", path); - fclose(fp); - return 0; - } - fclose(fp); - - g_moe_selected_hotlist_records = - header_records != UINT64_MAX ? header_records : loaded_hits / 6u; - g_moe_selected_hotlist_selections = - header_selections != UINT64_MAX ? header_selections : loaded_hits; - fprintf(stderr, - "ds4: merged selected-id hotlist %s " - "(%" PRIu64 " entries, %" PRIu64 " hits)\n", - path, - loaded_entries, - loaded_hits); - return 1; -} - -static void ds4_gpu_moe_selected_hotlist_close(void) { - const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); - if (!g_moe_selected_hotlist_initialized || !path || !path[0]) return; - - const size_t cap = - (size_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * - DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - ds4_gpu_moe_selected_hotlist_entry *entries = - malloc(cap * sizeof(entries[0])); - if (!entries) { - fprintf(stderr, "ds4: failed to allocate selected hotlist entries\n"); - return; - } - - size_t n = 0; - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - const uint64_t hits = - g_moe_selected_hotlist_counts[layer][expert]; - if (hits == 0) continue; - entries[n++] = (ds4_gpu_moe_selected_hotlist_entry) { - .layer = layer, - .expert = expert, - .hits = hits, - }; - } - } - qsort(entries, n, sizeof(entries[0]), ds4_gpu_moe_selected_hotlist_cmp); - - FILE *fp = fopen(path, "wb"); - if (!fp) { - fprintf(stderr, "ds4: failed to open selected hotlist file %s\n", path); - free(entries); - return; - } - fprintf(fp, - "# ds4 selected-id hotlist v1\n" - "# layer_records %" PRIu64 "\n" - "# selections %" PRIu64 "\n" - "# columns: layer expert hits weight\n", - g_moe_selected_hotlist_records, - g_moe_selected_hotlist_selections); - for (size_t i = 0; i < n; i++) { - fprintf(fp, - "%u %u %" PRIu64 " 0\n", - entries[i].layer, - entries[i].expert, - entries[i].hits); - } - free(entries); - - if (fclose(fp) != 0) { - fprintf(stderr, "ds4: failed to close selected hotlist file %s\n", path); - } else { - fprintf(stderr, - "ds4: wrote selected-id hotlist to %s " - "(%" PRIu64 " layer records, %" PRIu64 " selections)\n", - path, - g_moe_selected_hotlist_records, - g_moe_selected_hotlist_selections); - } -} - -static int ds4_gpu_moe_selected_hotlist_record( - uint32_t layer, - const int32_t *selected_ids, - uint32_t n_selected, - uint32_t n_total_expert) { - const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); - if (!path || !path[0]) return 1; - if (!selected_ids || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { - return 0; - } - if (!g_moe_selected_hotlist_initialized) { - g_moe_selected_hotlist_initialized = 1; - if (!ds4_gpu_moe_selected_hotlist_load_existing(path)) return 0; - atexit(ds4_gpu_moe_selected_hotlist_close); - } - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 1; - - g_moe_selected_hotlist_records++; - for (uint32_t i = 0; i < n_selected; i++) { - if (selected_ids[i] < 0) continue; - const uint32_t expert = (uint32_t)selected_ids[i]; - if (expert >= n_total_expert || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - continue; - } - g_moe_selected_hotlist_counts[layer][expert]++; - g_moe_selected_hotlist_selections++; - } - return 1; -} - -static void ds4_gpu_moe_selected_trace_record_close(void) { - if (g_moe_selected_trace_record_fp) { - const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); - fclose(g_moe_selected_trace_record_fp); - g_moe_selected_trace_record_fp = NULL; - fprintf(stderr, - "ds4: recorded %" PRIu64 " routed-MoE selected-id entries to %s\n", - g_moe_selected_trace_record_count, - path && path[0] ? path : "(unknown)"); - } -} - -static int ds4_gpu_moe_selected_trace_record( - const int32_t selected_ids[6], - uint32_t n_selected) { - const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); - if (!path || !path[0]) return 1; - if (n_selected != 6) { - fprintf(stderr, "ds4: selected-id recording expects exactly 6 selected experts\n"); - return 0; - } - - if (!g_moe_selected_trace_record_initialized) { - g_moe_selected_trace_record_initialized = 1; - g_moe_selected_trace_record_fp = fopen(path, "wb"); - if (!g_moe_selected_trace_record_fp) { - fprintf(stderr, "ds4: failed to open selected-id record file %s\n", path); - return 0; - } - setvbuf(g_moe_selected_trace_record_fp, NULL, _IOFBF, 1u << 20); - atexit(ds4_gpu_moe_selected_trace_record_close); - } - - if (fwrite(selected_ids, sizeof(selected_ids[0]), n_selected, g_moe_selected_trace_record_fp) != n_selected) { - fprintf(stderr, "ds4: failed to write selected-id record file %s\n", path); - return 0; - } - if (fflush(g_moe_selected_trace_record_fp) != 0) { - fprintf(stderr, "ds4: failed to flush selected-id record file %s\n", path); - return 0; - } - g_moe_selected_trace_record_count++; - return 1; -} - -static int ds4_gpu_moe_selected_trace_replay( - int32_t selected_ids[6], - uint32_t n_selected) { - const char *path = getenv("DS4_MOE_REPLAY_SELECTED_IDS"); - if (!path || !path[0]) return 0; - if (n_selected != 6) { - fprintf(stderr, "ds4: selected-id replay expects exactly 6 selected experts\n"); - return -1; - } - - if (!g_moe_selected_trace_replay_initialized) { - g_moe_selected_trace_replay_initialized = 1; - FILE *fp = fopen(path, "rb"); - if (!fp) { - fprintf(stderr, "ds4: failed to open selected-id replay file %s\n", path); - return -1; - } - if (fseeko(fp, 0, SEEK_END) != 0) { - fprintf(stderr, "ds4: failed to seek selected-id replay file %s\n", path); - fclose(fp); - return -1; - } - const off_t end = ftello(fp); - if (end < 0) { - fprintf(stderr, "ds4: failed to size selected-id replay file %s\n", path); - fclose(fp); - return -1; - } - if (fseeko(fp, 0, SEEK_SET) != 0) { - fprintf(stderr, "ds4: failed to rewind selected-id replay file %s\n", path); - fclose(fp); - return -1; - } - - const uint64_t bytes = (uint64_t)end; - const uint64_t entry_bytes = (uint64_t)n_selected * sizeof(selected_ids[0]); - if (bytes == 0 || (bytes % entry_bytes) != 0) { - fprintf(stderr, - "ds4: selected-id replay file %s has invalid size %" PRIu64 "\n", - path, - bytes); - fclose(fp); - return -1; - } - if (bytes > SIZE_MAX) { - fprintf(stderr, "ds4: selected-id replay file %s is too large\n", path); - fclose(fp); - return -1; - } - g_moe_selected_trace_replay_count = bytes / entry_bytes; - g_moe_selected_trace_replay_ids = malloc((size_t)bytes); - if (!g_moe_selected_trace_replay_ids) { - fprintf(stderr, "ds4: failed to allocate selected-id replay buffer\n"); - fclose(fp); - return -1; - } - if (fread(g_moe_selected_trace_replay_ids, 1, (size_t)bytes, fp) != (size_t)bytes) { - fprintf(stderr, "ds4: failed to read selected-id replay file %s\n", path); - fclose(fp); - free(g_moe_selected_trace_replay_ids); - g_moe_selected_trace_replay_ids = NULL; - return -1; - } - fclose(fp); - fprintf(stderr, - "ds4: loaded %" PRIu64 " routed-MoE selected-id entries from %s\n", - g_moe_selected_trace_replay_count, - path); - } - - if (g_moe_selected_trace_replay_pos >= g_moe_selected_trace_replay_count) { - fprintf(stderr, - "ds4: selected-id replay exhausted after %" PRIu64 " entries\n", - g_moe_selected_trace_replay_pos); - return -1; - } - memcpy(selected_ids, - g_moe_selected_trace_replay_ids + g_moe_selected_trace_replay_pos * n_selected, - (size_t)n_selected * sizeof(selected_ids[0])); - g_moe_selected_trace_replay_pos++; - return 1; -} - -static int ds4_gpu_progress_enabled(void) { - return ds4_log_is_tty(stderr); -} - -static void ds4_gpu_progress_begin(const char *what) { - if (!ds4_gpu_progress_enabled()) return; - fprintf(stderr, "ds4: %s...", what); - fflush(stderr); -} - -static void ds4_gpu_progress_done(void) { - if (!ds4_gpu_progress_enabled()) return; - fputs(" done\n", stderr); - fflush(stderr); -} - -static void ds4_gpu_progress_failed(void) { - if (!ds4_gpu_progress_enabled()) return; - fputs(" failed\n", stderr); - fflush(stderr); -} - -static void ds4_gpu_model_views_clear(void) { - for (uint32_t i = 0; i < g_model_view_count; i++) { - g_model_views[i].buffer = nil; - g_model_views[i].model_map = NULL; - g_model_views[i].model_size = 0; - g_model_views[i].model_offset = 0; - g_model_views[i].bytes = 0; - } - g_model_view_count = 0; -} - -static void ds4_gpu_model_residency_clear(void) { -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - if (g_model_residency_set) { - if (g_model_residency_added_to_queue && - g_queue && - [g_queue respondsToSelector:@selector(removeResidencySet:)]) { - [g_queue removeResidencySet:g_model_residency_set]; - } - [g_model_residency_set endResidency]; - [g_model_residency_set removeAllAllocations]; - g_model_residency_set = nil; - } - } -#endif - g_model_residency_count = 0; - g_model_residency_added_to_queue = 0; -} - -/* TP sharding keeps only this rank's expert ranges warm, - * so whole-view residency requests (which would page in the full file) - * must be skipped; pages fault in lazily through the same view buffers, - * exactly like ssd-streaming mode. */ -static int g_model_residency_skipped; - -void ds4_gpu_model_residency_skip(int skip) { - g_model_residency_skipped = skip; -} - -static int ds4_gpu_model_residency_request_views(void) { - if (g_model_view_count == 0 || - g_ssd_streaming_mode || - g_model_residency_skipped || - getenv("DS4_METAL_NO_RESIDENCY") != NULL) { - return 1; - } - -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - /* - * Register all model views as one residency set before inference. This - * is a GPU residency/budgeting hint, not a request to fault the whole - * 80+ GB file into memory. Its purpose is to make the driver see the - * complete set of large shared allocations during setup instead of - * discovering them lazily from the first measured graph command, where - * VM validation and residency accounting would look like model compute. - */ - MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; - desc.label = @"ds4_model"; - desc.initialCapacity = g_model_view_count; - - NSError *error = nil; - g_model_residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; - if (!g_model_residency_set) { - fprintf(stderr, "ds4: Metal model residency set creation failed: %s\n", - [[error localizedDescription] UTF8String]); - return 0; - } - - for (uint32_t i = 0; i < g_model_view_count; i++) { - [g_model_residency_set addAllocation:g_model_views[i].buffer]; - } - [g_model_residency_set commit]; - [g_model_residency_set requestResidency]; - if (getenv("DS4_METAL_DISABLE_QUEUE_RESIDENCY_SET") == NULL && - g_queue && - [g_queue respondsToSelector:@selector(addResidencySet:)]) { - [g_queue addResidencySet:g_model_residency_set]; - g_model_residency_added_to_queue = 1; - } - g_model_residency_count = g_model_view_count; - } -#endif - - return 1; -} - -static int ds4_gpu_add_model_view_range( - const void *model_map, - uint64_t model_size, - uint64_t map_offset, - uint64_t map_size, - uint64_t max_tensor_bytes, - bool use_default_view_cap, - uint64_t *mapped_model_size_out) { - const uint64_t page = (uint64_t)getpagesize(); - const uintptr_t model_addr = (uintptr_t)model_map; - - if ((model_addr & (uintptr_t)(page - 1)) != 0) { - fprintf(stderr, "ds4: Metal model mmap base is not page aligned\n"); - return 0; - } - if (map_offset > model_size || map_size > model_size - map_offset) { - fprintf(stderr, "ds4: Metal model mapped range is outside the GGUF mapping\n"); - return 0; - } - const uint64_t page_model_offset = map_offset & ~(page - 1); - const uint64_t leading = map_offset - page_model_offset; - if (map_size > UINT64_MAX - leading || - leading + map_size > UINT64_MAX - (page - 1)) - { - fprintf(stderr, "ds4: Metal model mapped range overflows page alignment\n"); - return 0; - } - const uint64_t mapped_model_size = round_up_u64(leading + map_size, page); - uint64_t max_buffer = (uint64_t)[g_device maxBufferLength]; - max_buffer &= ~(page - 1); - - /* - * Wrap only the tensor-data part of the GGUF file. Metadata is parsed by the - * CPU and is never dereferenced by kernels, so exposing it to Metal only - * grows the residency set and the VM range the driver must validate. - * - * Metal buffers have a device-specific maximum length, and this model is - * larger than that maximum on the target machines. Creating one no-copy - * buffer per tensor would avoid the length limit, but it would also move a - * lot of VM-object creation and residency bookkeeping into graph setup. The - * stable shape here is a tiny number of page-aligned views created once. - * - * Adjacent views intentionally overlap by more than the largest tensor, plus - * one page for alignment. That invariant guarantees every tensor lies wholly - * inside at least one view, so hot paths pass one buffer and one inner byte - * offset. We never split a weight tensor across command encoders. - */ - if (max_tensor_bytes > map_size) { - fprintf(stderr, "ds4: Metal model max tensor span is larger than a mapped tensor span\n"); - return 0; - } - if (max_tensor_bytes > UINT64_MAX - (page - 1)) { - fprintf(stderr, "ds4: Metal model max tensor span overflows page alignment\n"); - return 0; - } - const uint64_t max_tensor_rounded = round_up_u64(max_tensor_bytes, page); - if (max_tensor_rounded > UINT64_MAX - page) { - fprintf(stderr, "ds4: Metal model view overlap overflows page slack\n"); - return 0; - } - const uint64_t overlap = max_tensor_rounded + page; - if (max_buffer == 0 || max_buffer <= overlap) { - fprintf(stderr, - "ds4: Metal maxBufferLength is too small for DS4 model views " - "(max tensor %.2f GiB, max buffer %.2f GiB)\n", - ds4_gpu_gib(max_tensor_bytes), - ds4_gpu_gib(max_buffer)); - return 0; - } - - uint64_t view_limit = max_buffer; - const char *view_limit_env = getenv("DS4_METAL_MODEL_VIEW_MAX_GIB"); - if (view_limit_env && view_limit_env[0]) { - char *end = NULL; - unsigned long long gib = strtoull(view_limit_env, &end, 10); - if (end != view_limit_env && gib > 0) { - uint64_t env_limit = gib * 1024ull * 1024ull * 1024ull; - env_limit &= ~(page - 1); - if (env_limit > 0) view_limit = env_limit; - } - } else if (use_default_view_cap && mapped_model_size > max_buffer) { - /* - * Very large no-copy buffers can make Metal's VM validation dominate - * startup or the first graph command on multi-hundred-GiB slices. Keep - * ordinary contiguous model mappings unchanged, but let distributed - * span maps use smaller overlapping views when a range already has to - * be split. - */ - const uint64_t default_limit = 128ull * 1024ull * 1024ull * 1024ull; - if (view_limit > default_limit) view_limit = default_limit; - } - if (view_limit > max_buffer) view_limit = max_buffer; - view_limit &= ~(page - 1); - if (view_limit == 0 || view_limit <= overlap) { - fprintf(stderr, - "ds4: Metal model view cap is too small for DS4 model views " - "(cap %.2f GiB, max tensor %.2f GiB)\n", - ds4_gpu_gib(view_limit), - ds4_gpu_gib(max_tensor_bytes)); - return 0; - } - - const uint64_t step = view_limit - overlap; - uint64_t off = 0; - while (off < mapped_model_size) { - if (g_model_view_count == DS4_METAL_MAX_MODEL_VIEWS) { - fprintf(stderr, "ds4: Metal model needs more mapped views than expected\n"); - return 0; - } - - uint64_t view_bytes = mapped_model_size - off; - if (view_bytes > view_limit) view_bytes = view_limit; - - id buffer = [g_device newBufferWithBytesNoCopy:(void *)(model_addr + page_model_offset + off) - length:(NSUInteger)view_bytes - options:ds4_gpu_model_resource_options() - deallocator:nil]; - if (!buffer) { - fprintf(stderr, - "ds4: Metal could not wrap mmaped model view at %.2f GiB, size %.2f GiB\n", - (double)(page_model_offset + off) / (1024.0 * 1024.0 * 1024.0), - (double)view_bytes / (1024.0 * 1024.0 * 1024.0)); - return 0; - } - buffer.label = [NSString stringWithFormat:@"ds4_model_view_%u", g_model_view_count]; - - g_model_views[g_model_view_count].buffer = buffer; - g_model_views[g_model_view_count].model_map = model_map; - g_model_views[g_model_view_count].model_size = model_size; - g_model_views[g_model_view_count].model_offset = page_model_offset + off; - g_model_views[g_model_view_count].bytes = view_bytes; - g_model_view_count++; - - g_model_wrap_count++; - g_model_wrap_bytes += view_bytes; - if (view_bytes > g_model_wrap_max_bytes) g_model_wrap_max_bytes = view_bytes; - - if (off + view_bytes >= mapped_model_size) break; - off += step; - } - - if (mapped_model_size_out) *mapped_model_size_out += mapped_model_size; - return 1; -} - -static int ds4_gpu_finish_model_views( - double t0, - uint64_t mapped_model_size, - uint64_t display_offset) { - const double t_mapped = ds4_gpu_now_ms(); - const int request_residency = - !g_ssd_streaming_mode && - getenv("DS4_METAL_NO_RESIDENCY") == NULL; - if (request_residency) ds4_gpu_progress_begin("requesting Metal residency (may take tens of seconds)"); - if (!ds4_gpu_model_residency_request_views()) { - if (request_residency) ds4_gpu_progress_failed(); - return 0; - } - if (request_residency) ds4_gpu_progress_done(); - const double t_resident = ds4_gpu_now_ms(); - int warmed = 1; - const double t_warm0 = ds4_gpu_now_ms(); - const int warm_model_views = !g_ssd_streaming_mode && - getenv("DS4_METAL_NO_RESIDENCY") == NULL && - getenv("DS4_METAL_NO_MODEL_WARMUP") == NULL; - if (warm_model_views) { - /* - * The first GPU command touching no-copy mmap storage can pay command - * queue setup, page-table validation, and shared-allocation residency - * costs. Sample each model view here so timed graph execution starts - * after that one-time work. The stride is intentionally coarse: this is - * a validation touch over the VM ranges, not a full model prefetch. A - * dense prefetch would create exactly the kind of memory pressure and - * startup stalls this path is designed to avoid. - */ - if (g_model_residency_skipped) { - /* TP sharding: a single command buffer binding every - * view demands residency of them all and OOMs; the engine's - * CPU-side sharded warm pre-faults the owned bytes instead. */ - warmed = 1; - } else { - ds4_gpu_progress_begin("warming Metal model views"); - warmed = ds4_gpu_warm_model_views(); - if (warmed) ds4_gpu_progress_done(); - else ds4_gpu_progress_failed(); - } - } - const double t_warm = ds4_gpu_now_ms(); - if (ds4_gpu_model_map_log_enabled()) { - fprintf(stderr, - "ds4: Metal model views created in %.3f ms, residency requested in %.3f ms, warmup %.3f ms (mapped %.2f MiB from offset %.2f MiB)\n", - t_mapped - t0, - t_resident - t_mapped, - t_warm - t_warm0, - mapped_model_size / 1024.0 / 1024.0, - display_offset / 1024.0 / 1024.0); - } - if (!warmed) return 0; - return 1; -} - -static int ds4_gpu_map_model_views( - const void *model_map, - uint64_t model_size, - uint64_t map_offset, - uint64_t map_size, - uint64_t max_tensor_bytes) { - const double t0 = ds4_gpu_now_ms(); - uint64_t mapped_model_size = 0; - if (!ds4_gpu_add_model_view_range(model_map, - model_size, - map_offset, - map_size, - max_tensor_bytes, - false, - &mapped_model_size)) { - return 0; - } - return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset); -} - -static id ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) { - if (bytes == 0) bytes = 1; - - id buffer = [g_device newBufferWithLength:bytes - options:MTLResourceStorageModeShared]; - if (!buffer) { - fprintf(stderr, "ds4: failed to allocate Metal transient buffer %s (%llu bytes)\n", - label ? label : "(unnamed)", (unsigned long long)bytes); - return nil; - } - if (label) buffer.label = [NSString stringWithUTF8String:label]; - - /* - * CPU-filled buffers must survive until their command buffer completes. - * A local ObjC strong variable is not enough when the encoder function - * returns before the caller commits the command buffer. - */ - [g_transient_buffers addObject:buffer]; - return buffer; -} - -static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) { - if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || - getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) { - return 0; - } - return ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL; -} - -static ds4_gpu_zero_prefix_prefill_mask_cache_entry * -ds4_gpu_get_zero_prefix_prefill_mask_cache( - uint32_t kind, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t n_keys, - uint32_t window, - uint32_t ratio, - uint32_t nqptg, - uint32_t ncpsg, - bool has_kvpad, - bool bc_mask, - NSUInteger mask_bytes, - NSUInteger blk_bytes, - bool *created) { - if (created) *created = false; - if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() || - mask_bytes == 0 || blk_bytes == 0) { - return NULL; - } - - uint32_t slot = UINT32_MAX; - switch (kind) { - case DS4_GPU_PREFILL_MASK_CACHE_RAW: slot = 0; break; - case DS4_GPU_PREFILL_MASK_CACHE_RATIO4: slot = 1; break; - case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break; - default: return NULL; - } - - ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = - &g_zero_prefix_prefill_mask_cache[slot]; - if (entry->valid && entry->mask && entry->blk && - entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes && - entry->kind == kind && entry->n_tokens == n_tokens && - entry->n_comp == n_comp && entry->n_keys == n_keys && - entry->window == window && entry->ratio == ratio && - entry->nqptg == nqptg && entry->ncpsg == ncpsg && - entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) { - return entry; - } - - id mask = [g_device newBufferWithLength:mask_bytes - options:MTLResourceStorageModeShared]; - id blk = [g_device newBufferWithLength:blk_bytes - options:MTLResourceStorageModePrivate]; - if (!blk) { - blk = [g_device newBufferWithLength:blk_bytes - options:MTLResourceStorageModeShared]; - } - if (!mask || !blk) return NULL; - - mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind]; - blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind]; - - /* A replaced entry may still be referenced by an uncommitted or in-flight - * command buffer. Keep its resources alive with the other batch-scoped - * buffers until command completion instead of mutating or releasing them. */ - if (entry->mask) [g_transient_buffers addObject:entry->mask]; - if (entry->blk) [g_transient_buffers addObject:entry->blk]; - - entry->mask = mask; - entry->blk = blk; - entry->mask_bytes = mask_bytes; - entry->blk_bytes = blk_bytes; - entry->kind = kind; - entry->n_tokens = n_tokens; - entry->n_comp = n_comp; - entry->n_keys = n_keys; - entry->window = window; - entry->ratio = ratio; - entry->nqptg = nqptg; - entry->ncpsg = ncpsg; - entry->has_kvpad = has_kvpad; - entry->bc_mask = bc_mask; - /* The caller publishes the entry only after synchronously filling mask. */ - entry->valid = false; - entry->blk_ready = false; - *created = true; - return entry; -} - -static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) { - for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { - g_zero_prefix_prefill_mask_cache[i].blk_ready = false; - } -} - -static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) { - for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { - ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = - &g_zero_prefix_prefill_mask_cache[i]; - entry->mask = nil; - entry->blk = nil; - entry->mask_bytes = 0; - entry->blk_bytes = 0; - entry->kind = 0; - entry->n_tokens = 0; - entry->n_comp = 0; - entry->n_keys = 0; - entry->window = 0; - entry->ratio = 0; - entry->nqptg = 0; - entry->ncpsg = 0; - entry->has_kvpad = false; - entry->bc_mask = false; - entry->valid = false; - entry->blk_ready = false; - } -} - -void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) { - /* Layer-major prefill waits each layer before advancing, so its final - * release point has no outstanding cache users. Keep the guard here for - * diagnostic callers that may have an open or asynchronously flushed CB. */ - if (!g_initialized || g_batch_cb || - (g_pending_cbs && [g_pending_cbs count] != 0)) { - return; - } - ds4_gpu_clear_zero_prefix_prefill_mask_cache(); -} - -static id ds4_gpu_get_mul_mm_pipeline( - const char *function_name, - bool bc_inp, - bool bc_out) { - NSString *key = [NSString stringWithFormat:@"%s_bci=%d_bco=%d", - function_name, bc_inp ? 1 : 0, bc_out ? 1 : 0]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; - [constants setConstantValue:&bc_out type:MTLDataTypeBool atIndex:701]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_mul_mm_id_pipeline( - const char *function_name, - bool bc_inp) { - NSString *key = [NSString stringWithFormat:@"%s_bci=%d", - function_name, bc_inp ? 1 : 0]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_pipeline( - const char *function_name) { - NSString *key = [NSString stringWithFormat:@"%s", function_name]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found\n", function_name); - return nil; - } - - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static int ds4_gpu_disable_hot_pipeline_statics(void) { - static int initialized; - static int disabled; - if (!initialized) { - disabled = getenv("DS4_METAL_DISABLE_HOT_PIPELINE_STATICS") != NULL; - initialized = 1; - } - return disabled; -} - -static id ds4_gpu_hot_pipeline( - id pipeline, - const char *fallback_name) { - if (!ds4_gpu_disable_hot_pipeline_statics()) return pipeline; - return ds4_gpu_get_pipeline(fallback_name); -} - -static int ds4_gpu_use_compressor_pair_nr4(void) { - static int initialized; - static int enabled; - if (!initialized) { - enabled = getenv("DS4_METAL_COMPRESSOR_PAIR_NR4") != NULL; - initialized = 1; - } - return enabled; -} - -static int ds4_gpu_device_name_contains(const char *needle); - -static int ds4_gpu_env_value_eq(const char *v, size_t n, const char *literal) { - size_t m = strlen(literal); - if (n != m) return 0; - for (size_t i = 0; i < n; i++) { - if (tolower((unsigned char)v[i]) != tolower((unsigned char)literal[i])) return 0; - } - return 1; -} - -static int ds4_gpu_env_bool(const char *name) { - const char *v = getenv(name); - if (!v) return -1; - - while (isspace((unsigned char)*v)) v++; - size_t n = strlen(v); - while (n > 0 && isspace((unsigned char)v[n - 1])) n--; - if (n == 0) return 1; - - if (ds4_gpu_env_value_eq(v, n, "1") || - ds4_gpu_env_value_eq(v, n, "true") || - ds4_gpu_env_value_eq(v, n, "yes") || - ds4_gpu_env_value_eq(v, n, "on")) { - return 1; - } - if (ds4_gpu_env_value_eq(v, n, "0") || - ds4_gpu_env_value_eq(v, n, "false") || - ds4_gpu_env_value_eq(v, n, "no") || - ds4_gpu_env_value_eq(v, n, "off")) { - return 0; - } - - if (!g_mpp_invalid_env_reported) { - fprintf(stderr, - "ds4: invalid Metal boolean environment value %s=%.*s; treating presence as enabled\n", - name, (int)n, v); - g_mpp_invalid_env_reported = 1; - } - return 1; -} - -static uint64_t ds4_gpu_env_u64(const char *name, - uint64_t fallback, - uint64_t min_value, - uint64_t max_value) { - const char *v = getenv(name); - if (!v) return fallback; - while (isspace((unsigned char)*v)) v++; - if (!*v) return fallback; - - errno = 0; - char *end = NULL; - unsigned long long parsed = strtoull(v, &end, 10); - if (end == v || errno == ERANGE) return fallback; - while (isspace((unsigned char)*end)) end++; - if (*end) return fallback; - - if (parsed < min_value) return fallback; - uint64_t value = (uint64_t)parsed; - if (value > max_value) value = max_value; - return value; -} - -static uint32_t ds4_gpu_glm_full_attention_max_cache_len(void) { - /* - * kernel_glm_attention_full stores one score per visible token in - * threadgroup memory, plus 256 reduction slots. Keep the default under - * the 32 KiB envelope used by current Apple GPUs while allowing short - * dense decode to step past the 4096-token prefill boundary. - */ - return 7680u; -} - -static uint32_t ds4_gpu_glm_flash_attention_max_cache_len(void) { - /* - * Staged FlashAttention uses fixed-size threadgroup scratch and separate - * KV staging buffers, so it is not bound by the legacy full-attention - * kernel's one-score-per-token threadgroup-memory envelope. - */ - return 8192u; -} - -static int ds4_gpu_mpp_available(void) { - return g_metal4_tensor_api_enabled && !g_quality_mode; -} - -/* - * Retained Metal4 defaults live here instead of behind user-visible options. - * The public runtime has one automatic accelerated path plus the global - * DS4_METAL_DISABLE_METAL4 comparison switch. Benchmark-only alternatives that - * lost during M5 work are removed or kept out of the dispatch path so future - * changes do not accidentally turn old experiments into new modes. - */ -static int ds4_gpu_use_mpp_attn_out_low_matmul(void) { - return ds4_gpu_mpp_available(); -} - -enum { - DS4_METAL_ATTN_OUT_MPP_TILE_N = 64, -}; - -static void ds4_gpu_warn_mpp_fallback(void) { - static int warned; - if (!warned) { - fprintf(stderr, "ds4: accelerated Metal prefill matmul unavailable; falling back to legacy kernel\n"); - warned = 1; - } -} - -static int ds4_gpu_device_name_contains(const char *needle) { - return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL; -} - -static int ds4_gpu_compile_tensor_probe(void) { -#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 - if (!g_device) return 0; - if (@available(macOS 26.0, *)) { - const char *src = - "#include \n" - "#include \n" - "#include \n" - "using namespace metal;\n" - "using namespace mpp::tensor_ops;\n" - "kernel void ds4_tensor_probe(\n" - " tensor> A [[buffer(0)]],\n" - " tensor> B [[buffer(1)]],\n" - " device float *C [[buffer(2)]],\n" - " uint2 tgid [[threadgroup_position_in_grid]]) {\n" - " auto tA = A.slice(0, (int)tgid.y);\n" - " auto tB = B.slice((int)tgid.x, 0);\n" - " matmul2d> mm;\n" - " auto cT = mm.get_destination_cooperative_tensor();\n" - " auto sA = tA.slice(0, 0);\n" - " auto sB = tB.slice(0, 0);\n" - " mm.run(sB, sA, cT);\n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16));\n" - " cT.store(tC);\n" - "}\n"; - - NSError *error = nil; - NSString *source = [NSString stringWithUTF8String:src]; - id probe_library = [g_device newLibraryWithSource:source options:[MTLCompileOptions new] error:&error]; - if (!probe_library) { - fprintf(stderr, "ds4: Metal 4 tensor API probe compile failed: %s\n", - error ? [[error localizedDescription] UTF8String] : "(unknown)"); - return 0; - } - id fn = [probe_library newFunctionWithName:@"ds4_tensor_probe"]; - if (!fn) { - fprintf(stderr, "ds4: Metal 4 tensor API probe function missing\n"); - return 0; - } - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal 4 tensor API probe pipeline failed: %s\n", - error ? [[error localizedDescription] UTF8String] : "(unknown)"); - return 0; - } - return 1; - } -#endif - return 0; -} - -static void ds4_gpu_detect_metal4_features(void) { - g_metal4_runtime_available = 0; - g_metal4_family_supported = 0; - g_metal4_queue_supported = 0; - g_metal4_m5_neural_accelerators_hint = 0; - g_metal4_tensor_api_enabled = 0; - g_metal4_tensor_api_compile_supported = 0; - g_metal_device_name[0] = '\0'; - - if (!g_device) return; - - const char *name = [[g_device name] UTF8String]; - if (name) { - snprintf(g_metal_device_name, sizeof(g_metal_device_name), "%s", name); - } - - const int metal4_disabled = ds4_gpu_env_bool("DS4_METAL_DISABLE_METAL4") > 0; - -#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 - if (@available(macOS 26.0, *)) { - g_metal4_runtime_available = 1; - g_metal4_family_supported = - !metal4_disabled && [g_device supportsFamily:MTLGPUFamilyMetal4] ? 1 : 0; - g_metal4_queue_supported = [g_device respondsToSelector:@selector(newMTL4CommandQueue)] ? 1 : 0; - - /* - * Apple does not currently expose a separate "Neural Accelerator" bit - * through Metal. On public M5 systems the hardware signal is the device - * generation plus Metal 4 support, so keep this as a conservative hint. - */ - if (g_metal4_family_supported && ds4_gpu_device_name_contains("M5")) { - g_metal4_m5_neural_accelerators_hint = 1; - } - - if (g_metal4_family_supported) { - const int default_enable = - ds4_gpu_device_name_contains("M5") || - ds4_gpu_device_name_contains("M6") || - ds4_gpu_device_name_contains("A19") || - ds4_gpu_device_name_contains("A20"); - - /* - * Metal 4 TensorOps are portable in source, but on pre-M5 hardware - * they can map to ordinary shader fallbacks. Keep the automatic - * fast path restricted to hardware generations where the Neural - * Accelerator/TensorOps path is expected to pay off; older Metal - * machines continue to use the established kernels unless a future - * device is explicitly added here. - */ - if (default_enable) { - g_metal4_tensor_api_compile_supported = ds4_gpu_compile_tensor_probe(); - g_metal4_tensor_api_enabled = g_metal4_tensor_api_compile_supported; - if (!g_metal4_tensor_api_enabled) { - fprintf(stderr, "ds4: Metal 4 tensor API probe failed; using legacy Metal kernels\n"); - } - } else { - fprintf(stderr, "ds4: Metal 4 tensor API disabled for pre-M5/pre-A19 devices\n"); - } - } - } -#endif -} - -static int ds4_gpu_warm_model_views(void) { - if (g_model_view_count == 0) return 1; - - id pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride"); - if (!pipeline) return 0; - - uint64_t stride = 1024ull * 1024ull; - const char *stride_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_MB"); - if (stride_env && stride_env[0]) { - char *end = NULL; - unsigned long long mb = strtoull(stride_env, &end, 10); - if (end != stride_env && mb > 0 && mb <= 1024) { - stride = mb * 1024ull * 1024ull; - } - } - const char *stride_kb_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_KB"); - if (stride_kb_env && stride_kb_env[0]) { - char *end = NULL; - unsigned long long kb = strtoull(stride_kb_env, &end, 10); - if (end != stride_kb_env && kb > 0 && kb <= 1024ull * 1024ull) { - stride = kb * 1024ull; - const uint64_t page = (uint64_t)getpagesize(); - if (stride < page) stride = page; - } - } - - uint64_t total_touches = 0; - for (uint32_t i = 0; i < g_model_view_count; i++) { - total_touches += (g_model_views[i].bytes + stride - 1) / stride; - } - if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0; - - const NSUInteger out_bytes = (NSUInteger)total_touches; - id out = [g_device newBufferWithLength:out_bytes - options:MTLResourceStorageModeShared]; - if (!out) { - fprintf(stderr, "ds4: Metal model warmup scratch allocation failed\n"); - return 0; - } - out.label = @"ds4_model_warmup"; - - id cb = ds4_gpu_new_command_buffer(); - if (!cb) { - fprintf(stderr, "ds4: Metal model warmup command buffer allocation failed\n"); - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - uint64_t dst_offset = 0; - for (uint32_t i = 0; i < g_model_view_count; i++) { - const uint64_t bytes = g_model_views[i].bytes; - const uint64_t n = (bytes + stride - 1) / stride; - [enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0]; - [enc setBuffer:out offset:0 atIndex:1]; - [enc setBytes:&stride length:sizeof(stride) atIndex:2]; - [enc setBytes:&bytes length:sizeof(bytes) atIndex:3]; - [enc setBytes:&dst_offset length:sizeof(dst_offset) atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - dst_offset += n; - } - ds4_gpu_end_compute_encoder(cb, enc); - - [cb commit]; - [cb waitUntilCompleted]; - - if (cb.status == MTLCommandBufferStatusError) { - fprintf(stderr, "ds4: Metal model warmup failed: %s\n", - [[cb.error localizedDescription] UTF8String]); - return 0; - } - - return 1; -} - -static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) { - switch (ne20) { - case 1: return "kernel_mul_mm_id_map0_ne20_1"; - case 2: return "kernel_mul_mm_id_map0_ne20_2"; - case 4: return "kernel_mul_mm_id_map0_ne20_4"; - case 5: return "kernel_mul_mm_id_map0_ne20_5"; - case 6: return "kernel_mul_mm_id_map0_ne20_6"; - case 8: return "kernel_mul_mm_id_map0_ne20_8"; - case 10: return "kernel_mul_mm_id_map0_ne20_10"; - case 16: return "kernel_mul_mm_id_map0_ne20_16"; - case 22: return "kernel_mul_mm_id_map0_ne20_22"; - default: return NULL; - } -} - -static id ds4_gpu_get_mul_mv_pipeline( - const char *function_name, - int16_t nsg) { - NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_mul_mv_ext_pipeline( - const char *function_name, - int16_t nsg, - int16_t nxpsg) { - NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d", - function_name, (int)nsg, (int)nxpsg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; - [constants setConstantValue:&nxpsg type:MTLDataTypeShort atIndex:601]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_flash_attn_pad_pipeline( - bool has_mask, - int32_t ncpsg) { - NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d", - has_mask ? 1 : 0, (int)ncpsg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100]; - [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:125]; - - NSError *error = nil; - id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_pad" - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad function not found: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_flash_attn_blk_pipeline( - int32_t nqptg, - int32_t ncpsg) { - NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d", - (int)nqptg, (int)ncpsg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224]; - [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:225]; - - NSError *error = nil; - id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_blk" - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk function not found: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_flash_attn_pipeline( - const char *function_name, - bool has_mask, - bool has_sinks, - bool has_bias, - bool has_scap, - bool has_kvpad, - bool bc_mask, - int32_t ns10, - int32_t ns20, - int32_t nsg) { - NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d", - function_name, - has_mask ? 1 : 0, - has_sinks ? 1 : 0, - has_bias ? 1 : 0, - has_scap ? 1 : 0, - has_kvpad ? 1 : 0, - bc_mask ? 1 : 0, - (int)ns10, - (int)ns20, - (int)nsg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300]; - [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:301]; - [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:302]; - [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:303]; - [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:304]; - [constants setConstantValue:&bc_mask type:MTLDataTypeBool atIndex:310]; - [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:320]; - [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:321]; - [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:322]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - return pipeline; -} - -static id ds4_gpu_get_flash_attn_vec_pipeline( - const char *function_name, - bool has_mask, - bool has_sinks, - bool has_bias, - bool has_scap, - bool has_kvpad, - bool shared_kvpad, - int32_t ns10, - int32_t ns20, - int32_t nsg, - int32_t nwg) { - /* - * Decode calls this once per layer with identical arguments, so memoize - * the last hit and skip the NSString key + dictionary lookup on the hot - * path. The generic cache below remains the fallback for new variants. - */ - static struct { - const char *fn; - bool m, s, b, c, k, sp; - int32_t n10, n20, sg, wg; - id pipeline; - } memo; - if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 && - memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias && - memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad && - memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) { - return memo.pipeline; - } - - NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", - function_name, - has_mask ? 1 : 0, - has_sinks ? 1 : 0, - has_bias ? 1 : 0, - has_scap ? 1 : 0, - has_kvpad ? 1 : 0, - shared_kvpad ? 1 : 0, - (int)ns10, - (int)ns20, - (int)nsg, - (int)nwg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) { - memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, - has_scap, has_kvpad, shared_kvpad, ns10, ns20, - nsg, nwg, cached }; - return cached; - } - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:400]; - [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:401]; - [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:402]; - [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:403]; - [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404]; - [constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405]; - [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420]; - [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421]; - [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:422]; - [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:423]; - - NSError *error = nil; - NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal %s function not found: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", - function_name, [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, - has_scap, has_kvpad, shared_kvpad, ns10, ns20, - nsg, nwg, pipeline }; - return pipeline; -} - -static id ds4_gpu_get_flash_attn_reduce_pipeline( - int32_t dv, - int32_t nwg) { - /* Same per-layer memo pattern as the vec getter above. */ - static int32_t memo_dv, memo_nwg; - static id memo_pipeline; - if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) { - return memo_pipeline; - } - - NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d", - (int)dv, (int)nwg]; - id cached = [g_pipeline_cache objectForKey:key]; - if (cached) { - memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached; - return cached; - } - - MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; - [constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500]; - [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501]; - - NSError *error = nil; - id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce" - constantValues:constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce function not found: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - error = nil; - id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!pipeline) { - fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - [g_pipeline_cache setObject:pipeline forKey:key]; - memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline; - return pipeline; -} - -static uint32_t ds4_gpu_flash_attn_vec_nsg(uint32_t n_keys, uint32_t nwg, uint32_t ncpsg) { - uint32_t nsg = 1; - while (2u * nwg * nsg * ncpsg < n_keys && nsg < 4u) { - nsg *= 2u; - } - return nsg; -} - -static int ds4_gpu_trace_allocs(void) { - static int initialized; - static int enabled; - if (!initialized) { - enabled = getenv("DS4_METAL_TRACE_ALLOCS") != NULL; - initialized = 1; - } - return enabled; -} - -static double ds4_gpu_mib(uint64_t bytes) { - return (double)bytes / (1024.0 * 1024.0); -} - -static double ds4_gpu_gib(uint64_t bytes) { - return (double)bytes / (1024.0 * 1024.0 * 1024.0); -} - -static ds4_gpu_stream_expert_timing_snapshot -ds4_gpu_stream_expert_timing_current(void) { - return (ds4_gpu_stream_expert_timing_snapshot) { - .selected_calls = g_stream_expert_timing_selected_calls, - .selected_read_ms = g_stream_expert_timing_selected_read_ms, - .selected_sync_ms = g_stream_expert_timing_selected_sync_ms, - .selected_copy_ms = g_stream_expert_timing_selected_copy_ms, - .selected_bind_ms = g_stream_expert_timing_selected_bind_ms, - .split_layers = g_stream_expert_timing_split_layers, - .split_resident_experts = g_stream_expert_timing_split_resident_experts, - .split_missing_experts = g_stream_expert_timing_split_missing_experts, - .split_resident_ms = g_stream_expert_timing_split_resident_ms, - .split_missing_ms = g_stream_expert_timing_split_missing_ms, - .split_missing_load_ms = - g_stream_expert_timing_split_missing_load_ms, - .split_missing_slot_ms = - g_stream_expert_timing_split_missing_slot_ms, - .split_missing_prune_ms = - g_stream_expert_timing_split_missing_prune_ms, - .split_missing_addr_ms = - g_stream_expert_timing_split_missing_addr_ms, - .split_missing_wait_ms = - g_stream_expert_timing_split_missing_wait_ms, - .load_calls = g_stream_expert_timing_load_calls, - .load_prepare_ms = g_stream_expert_timing_load_prepare_ms, - .load_pread_ms = g_stream_expert_timing_load_pread_ms, - .load_modify_ms = g_stream_expert_timing_load_modify_ms, - .load_install_ms = g_stream_expert_timing_load_install_ms, - .prepare_batch_reuse_calls = - g_stream_expert_timing_prepare_batch_reuse_calls, - .prepare_batch_reuse_ms = - g_stream_expert_timing_prepare_batch_reuse_ms, - .prepare_buffer_calls = - g_stream_expert_timing_prepare_buffer_calls, - .prepare_buffer_ms = - g_stream_expert_timing_prepare_buffer_ms, - .prepare_task_experts = - g_stream_expert_timing_prepare_task_experts, - .prepare_task_ms = - g_stream_expert_timing_prepare_task_ms, - .reuse_scan_calls = - g_stream_expert_timing_reuse_scan_calls, - .reuse_scan_entries = - g_stream_expert_timing_reuse_scan_entries, - .reuse_scan_ms = - g_stream_expert_timing_reuse_scan_ms, - .reuse_clear_ms = - g_stream_expert_timing_reuse_clear_ms, - .readahead_calls = - g_stream_expert_timing_readahead_calls, - .readahead_bytes = - g_stream_expert_timing_readahead_bytes, - .readahead_ms = - g_stream_expert_timing_readahead_ms, - .cache_all_resident_layers = - g_stream_expert_timing_cache_all_resident_layers, - .cache_all_missing_layers = - g_stream_expert_timing_cache_all_missing_layers, - .cache_mixed_layers = g_stream_expert_timing_cache_mixed_layers, - .cache_resident_experts = - g_stream_expert_timing_cache_resident_experts, - .cache_missing_experts = - g_stream_expert_timing_cache_missing_experts, - }; -} - -static uint64_t ds4_gpu_stream_expert_timing_delta_u64( - uint64_t current, - uint64_t previous) { - return current >= previous ? current - previous : current; -} - -static double ds4_gpu_stream_expert_timing_delta_f64( - double current, - double previous) { - return current >= previous ? current - previous : current; -} - -static ds4_gpu_stream_expert_timing_snapshot -ds4_gpu_stream_expert_timing_delta( - ds4_gpu_stream_expert_timing_snapshot current, - ds4_gpu_stream_expert_timing_snapshot previous) { - return (ds4_gpu_stream_expert_timing_snapshot) { - .selected_calls = - ds4_gpu_stream_expert_timing_delta_u64(current.selected_calls, - previous.selected_calls), - .selected_read_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.selected_read_ms, - previous.selected_read_ms), - .selected_sync_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.selected_sync_ms, - previous.selected_sync_ms), - .selected_copy_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.selected_copy_ms, - previous.selected_copy_ms), - .selected_bind_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.selected_bind_ms, - previous.selected_bind_ms), - .split_layers = - ds4_gpu_stream_expert_timing_delta_u64(current.split_layers, - previous.split_layers), - .split_resident_experts = - ds4_gpu_stream_expert_timing_delta_u64(current.split_resident_experts, - previous.split_resident_experts), - .split_missing_experts = - ds4_gpu_stream_expert_timing_delta_u64(current.split_missing_experts, - previous.split_missing_experts), - .split_resident_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.split_resident_ms, - previous.split_resident_ms), - .split_missing_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.split_missing_ms, - previous.split_missing_ms), - .split_missing_load_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.split_missing_load_ms, - previous.split_missing_load_ms), - .split_missing_slot_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.split_missing_slot_ms, - previous.split_missing_slot_ms), - .split_missing_prune_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.split_missing_prune_ms, - previous.split_missing_prune_ms), - .split_missing_addr_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.split_missing_addr_ms, - previous.split_missing_addr_ms), - .split_missing_wait_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.split_missing_wait_ms, - previous.split_missing_wait_ms), - .load_calls = - ds4_gpu_stream_expert_timing_delta_u64(current.load_calls, - previous.load_calls), - .load_prepare_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.load_prepare_ms, - previous.load_prepare_ms), - .load_pread_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.load_pread_ms, - previous.load_pread_ms), - .load_modify_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.load_modify_ms, - previous.load_modify_ms), - .load_install_ms = - ds4_gpu_stream_expert_timing_delta_f64(current.load_install_ms, - previous.load_install_ms), - .prepare_batch_reuse_calls = - ds4_gpu_stream_expert_timing_delta_u64( - current.prepare_batch_reuse_calls, - previous.prepare_batch_reuse_calls), - .prepare_batch_reuse_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.prepare_batch_reuse_ms, - previous.prepare_batch_reuse_ms), - .prepare_buffer_calls = - ds4_gpu_stream_expert_timing_delta_u64( - current.prepare_buffer_calls, - previous.prepare_buffer_calls), - .prepare_buffer_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.prepare_buffer_ms, - previous.prepare_buffer_ms), - .prepare_task_experts = - ds4_gpu_stream_expert_timing_delta_u64( - current.prepare_task_experts, - previous.prepare_task_experts), - .prepare_task_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.prepare_task_ms, - previous.prepare_task_ms), - .reuse_scan_calls = - ds4_gpu_stream_expert_timing_delta_u64( - current.reuse_scan_calls, - previous.reuse_scan_calls), - .reuse_scan_entries = - ds4_gpu_stream_expert_timing_delta_u64( - current.reuse_scan_entries, - previous.reuse_scan_entries), - .reuse_scan_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.reuse_scan_ms, - previous.reuse_scan_ms), - .reuse_clear_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.reuse_clear_ms, - previous.reuse_clear_ms), - .readahead_calls = - ds4_gpu_stream_expert_timing_delta_u64( - current.readahead_calls, - previous.readahead_calls), - .readahead_bytes = - ds4_gpu_stream_expert_timing_delta_u64( - current.readahead_bytes, - previous.readahead_bytes), - .readahead_ms = - ds4_gpu_stream_expert_timing_delta_f64( - current.readahead_ms, - previous.readahead_ms), - .cache_all_resident_layers = - ds4_gpu_stream_expert_timing_delta_u64( - current.cache_all_resident_layers, - previous.cache_all_resident_layers), - .cache_all_missing_layers = - ds4_gpu_stream_expert_timing_delta_u64( - current.cache_all_missing_layers, - previous.cache_all_missing_layers), - .cache_mixed_layers = - ds4_gpu_stream_expert_timing_delta_u64(current.cache_mixed_layers, - previous.cache_mixed_layers), - .cache_resident_experts = - ds4_gpu_stream_expert_timing_delta_u64( - current.cache_resident_experts, - previous.cache_resident_experts), - .cache_missing_experts = - ds4_gpu_stream_expert_timing_delta_u64( - current.cache_missing_experts, - previous.cache_missing_experts), - }; -} - -static int ds4_gpu_stream_expert_timing_has_data( - ds4_gpu_stream_expert_timing_snapshot s) { - return s.selected_calls != 0 || - s.split_layers != 0 || - s.load_calls != 0 || - s.cache_all_resident_layers != 0 || - s.cache_all_missing_layers != 0 || - s.cache_mixed_layers != 0; -} - -static void ds4_gpu_stream_expert_timing_print( - const char *scope, - ds4_gpu_stream_expert_timing_snapshot s) { - if (!ds4_gpu_stream_expert_timing_has_data(s)) return; - const double selected_calls = (double)s.selected_calls; - const double split_layers = (double)s.split_layers; - const double selected_read_avg = - selected_calls != 0.0 ? s.selected_read_ms / selected_calls : 0.0; - const double selected_sync_avg = - selected_calls != 0.0 ? s.selected_sync_ms / selected_calls : 0.0; - const double selected_copy_avg = - selected_calls != 0.0 ? s.selected_copy_ms / selected_calls : 0.0; - const double selected_bind_avg = - selected_calls != 0.0 ? s.selected_bind_ms / selected_calls : 0.0; - const double split_resident_avg = - split_layers != 0.0 ? s.split_resident_ms / split_layers : 0.0; - const double split_missing_avg = - split_layers != 0.0 ? s.split_missing_ms / split_layers : 0.0; - const double split_missing_load_avg = - split_layers != 0.0 ? - s.split_missing_load_ms / split_layers : 0.0; - const double split_missing_slot_avg = - split_layers != 0.0 ? - s.split_missing_slot_ms / split_layers : 0.0; - const double split_missing_prune_avg = - split_layers != 0.0 ? - s.split_missing_prune_ms / split_layers : 0.0; - const double split_missing_addr_avg = - split_layers != 0.0 ? - s.split_missing_addr_ms / split_layers : 0.0; - const double split_missing_wait_avg = - split_layers != 0.0 ? - s.split_missing_wait_ms / split_layers : 0.0; - const double load_calls = (double)s.load_calls; - const double load_prepare_avg = - load_calls != 0.0 ? s.load_prepare_ms / load_calls : 0.0; - const double load_pread_avg = - load_calls != 0.0 ? s.load_pread_ms / load_calls : 0.0; - const double load_modify_avg = - load_calls != 0.0 ? s.load_modify_ms / load_calls : 0.0; - const double load_install_avg = - load_calls != 0.0 ? s.load_install_ms / load_calls : 0.0; - const double prepare_batch_reuse_avg = - s.prepare_batch_reuse_calls != 0 ? - s.prepare_batch_reuse_ms / - (double)s.prepare_batch_reuse_calls : 0.0; - const double prepare_buffer_avg = - s.prepare_buffer_calls != 0 ? - s.prepare_buffer_ms / (double)s.prepare_buffer_calls : 0.0; - const double prepare_task_avg = - s.prepare_task_experts != 0 ? - s.prepare_task_ms / (double)s.prepare_task_experts : 0.0; - const double reuse_scan_avg = - s.reuse_scan_calls != 0 ? - s.reuse_scan_ms / (double)s.reuse_scan_calls : 0.0; - const double reuse_scan_entries_avg = - s.reuse_scan_calls != 0 ? - (double)s.reuse_scan_entries / (double)s.reuse_scan_calls : 0.0; - const double readahead_avg = - s.readahead_calls != 0 ? - s.readahead_ms / (double)s.readahead_calls : 0.0; - const double split_resident_experts_avg = - split_layers != 0.0 ? - (double)s.split_resident_experts / split_layers : 0.0; - const double split_missing_experts_avg = - split_layers != 0.0 ? - (double)s.split_missing_experts / split_layers : 0.0; - const uint64_t cache_layers = - s.cache_all_resident_layers + - s.cache_all_missing_layers + - s.cache_mixed_layers; - const double cache_layer_count = (double)cache_layers; - const double cache_resident_experts_avg = - cache_layer_count != 0.0 ? - (double)s.cache_resident_experts / cache_layer_count : 0.0; - const double cache_missing_experts_avg = - cache_layer_count != 0.0 ? - (double)s.cache_missing_experts / cache_layer_count : 0.0; - fprintf(stderr, - "ds4: streaming expert timing %s selected_calls=%llu read_avg=%.3f ms sync_avg=%.3f ms copy_avg=%.3f ms bind_avg=%.3f ms read_total=%.3f ms sync_total=%.3f ms copy_total=%.3f ms bind_total=%.3f ms split_layers=%llu resident_experts_avg=%.2f missing_experts_avg=%.2f resident_submit_avg=%.3f ms missing_bind_avg=%.3f ms resident_submit_total=%.3f ms missing_bind_total=%.3f ms missing_load_avg=%.3f ms missing_slot_avg=%.3f ms missing_prune_avg=%.3f ms missing_addr_avg=%.3f ms missing_wait_avg=%.3f ms missing_wait_total=%.3f ms load_calls=%llu load_prepare_avg=%.3f ms load_pread_avg=%.3f ms load_modify_avg=%.3f ms load_install_avg=%.3f ms prepare_batch_reuse_calls=%llu prepare_batch_reuse_avg=%.3f ms prepare_batch_reuse_total=%.3f ms prepare_buffer_calls=%llu prepare_buffer_avg=%.3f ms prepare_buffer_total=%.3f ms prepare_task_experts=%llu prepare_task_avg=%.3f ms prepare_task_total=%.3f ms reuse_scan_calls=%llu reuse_scan_entries_avg=%.1f reuse_scan_avg=%.3f ms reuse_scan_total=%.3f ms reuse_clear_total=%.3f ms readahead_calls=%llu readahead_avg=%.3f ms readahead_total=%.3f ms readahead_gib=%.2f cache_all_resident=%llu cache_all_missing=%llu cache_mixed=%llu cache_resident_avg=%.2f cache_missing_avg=%.2f\n", - scope ? scope : "total", - (unsigned long long)s.selected_calls, - selected_read_avg, - selected_sync_avg, - selected_copy_avg, - selected_bind_avg, - s.selected_read_ms, - s.selected_sync_ms, - s.selected_copy_ms, - s.selected_bind_ms, - (unsigned long long)s.split_layers, - split_resident_experts_avg, - split_missing_experts_avg, - split_resident_avg, - split_missing_avg, - s.split_resident_ms, - s.split_missing_ms, - split_missing_load_avg, - split_missing_slot_avg, - split_missing_prune_avg, - split_missing_addr_avg, - split_missing_wait_avg, - s.split_missing_wait_ms, - (unsigned long long)s.load_calls, - load_prepare_avg, - load_pread_avg, - load_modify_avg, - load_install_avg, - (unsigned long long)s.prepare_batch_reuse_calls, - prepare_batch_reuse_avg, - s.prepare_batch_reuse_ms, - (unsigned long long)s.prepare_buffer_calls, - prepare_buffer_avg, - s.prepare_buffer_ms, - (unsigned long long)s.prepare_task_experts, - prepare_task_avg, - s.prepare_task_ms, - (unsigned long long)s.reuse_scan_calls, - reuse_scan_entries_avg, - reuse_scan_avg, - s.reuse_scan_ms, - s.reuse_clear_ms, - (unsigned long long)s.readahead_calls, - readahead_avg, - s.readahead_ms, - ds4_gpu_gib(s.readahead_bytes), - (unsigned long long)s.cache_all_resident_layers, - (unsigned long long)s.cache_all_missing_layers, - (unsigned long long)s.cache_mixed_layers, - cache_resident_experts_avg, - cache_missing_experts_avg); -} - -static void ds4_gpu_print_task_memory_report(void) { - task_vm_info_data_t info; - mach_msg_type_number_t count = TASK_VM_INFO_COUNT; - const kern_return_t kr = task_info(mach_task_self(), - TASK_VM_INFO, - (task_info_t)&info, - &count); - if (kr != KERN_SUCCESS) return; - - fprintf(stderr, - "ds4: macOS task memory footprint %.2f GiB, resident %.2f GiB, virtual %.2f GiB\n", - ds4_gpu_gib((uint64_t)info.phys_footprint), - ds4_gpu_gib((uint64_t)info.resident_size), - ds4_gpu_gib((uint64_t)info.virtual_size)); -} - -void ds4_gpu_print_memory_report(const char *label) { - uint64_t cached_prefill_mask_bytes = 0; - uint64_t cached_prefill_blk_bytes = 0; - for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { - const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = - &g_zero_prefix_prefill_mask_cache[i]; - if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes; - if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes; - } - const uint64_t scratch = - (uint64_t)g_flash_attn_mask_bytes + - (uint64_t)g_flash_attn_zero_mask_bytes + - cached_prefill_mask_bytes + - (uint64_t)g_flash_attn_pad_bytes + - (uint64_t)g_flash_attn_tmp_bytes + - (uint64_t)g_flash_attn_blk_bytes + - cached_prefill_blk_bytes + - (uint64_t)g_flash_attn_ring_bytes + - (uint64_t)g_flash_attn_kv_bytes + - (uint64_t)g_glm_flash_attn_mask_bytes + - (uint64_t)g_compressor_pool_kv_bytes + - (uint64_t)g_compressor_pool_score_bytes + - (uint64_t)g_compressor_pool_score_cont_bytes + - (uint64_t)g_compressor_pool_softmax_bytes + - (uint64_t)g_compressor_pool_product_bytes + - (uint64_t)g_compressor_store_ape_bytes + - (uint64_t)g_compressor_store_score_bytes + - (uint64_t)g_embed_rows_bytes + - (uint64_t)g_router_selection_bytes + - (uint64_t)g_router_weight_sum_bytes + - (uint64_t)g_indexer_head_scores_bytes + - (uint64_t)g_indexer_topk_bytes + - (uint64_t)g_indexed_topk_bytes + - (uint64_t)g_f16_round_scratch_bytes + - (uint64_t)g_raw_store_round_bytes + - (uint64_t)g_moe_gate_scratch_bytes + - (uint64_t)g_moe_down_scratch_bytes + - (uint64_t)g_moe_id_map_bytes + - (uint64_t)g_moe_q4_gate_slots_bytes + - (uint64_t)g_moe_q4_up_slots_bytes + - (uint64_t)g_moe_q4_down_slots_bytes; - - pthread_mutex_lock(&g_tensor_mu); - const uint64_t tensor_live_snap = g_tensor_alloc_live_bytes; - const uint64_t tensor_peak_snap = g_tensor_alloc_peak_bytes; - pthread_mutex_unlock(&g_tensor_mu); - - uint64_t tracked_live = tensor_live_snap; - if (tracked_live > UINT64_MAX - g_stream_expert_cache_bytes) { - tracked_live = UINT64_MAX; - } else { - tracked_live += g_stream_expert_cache_bytes; - } - - const bool color = ds4_log_is_tty(stderr); - const char *green = color ? "\x1b[32m" : ""; - const char *bright_green = color ? "\x1b[1;32m" : ""; - const char *reset = color ? "\x1b[0m" : ""; - fprintf(stderr, - "%sds4: Metal memory%s%s: runtime %.2f GiB + streaming experts %.2f GiB = %s%.2f GiB tracked live%s\n", - green, - label && label[0] ? " " : "", - label && label[0] ? label : "", - ds4_gpu_gib(tensor_live_snap), - ds4_gpu_gib(g_stream_expert_cache_bytes), - bright_green, - ds4_gpu_gib(tracked_live), - reset); - if (color) fputs(green, stderr); - fprintf(stderr, - "ds4: runtime tensors live %.2f MiB peak %.2f MiB\n", - ds4_gpu_mib(tensor_live_snap), - ds4_gpu_mib(tensor_peak_snap)); - ds4_gpu_print_task_memory_report(); - fprintf(stderr, - "ds4: mmap model wrapper spans %llu buffers %.2f GiB total, %.2f GiB max (not copied)\n", - (unsigned long long)g_model_wrap_count, - ds4_gpu_gib(g_model_wrap_bytes), - ds4_gpu_gib(g_model_wrap_max_bytes)); - if (g_model_buffer_cache && [g_model_buffer_cache count] != 0) { - const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); - if (limit == 0) { - fprintf(stderr, - "ds4: exact model view cache %lu buffers %.2f GiB unlimited, %llu evictions (not copied)\n", - (unsigned long)[g_model_buffer_cache count], - ds4_gpu_gib(g_model_buffer_cache_bytes), - (unsigned long long)g_model_buffer_cache_evictions); - } else { - fprintf(stderr, - "ds4: exact model view cache %lu buffers %.2f GiB / %.2f GiB, %llu evictions (not copied)\n", - (unsigned long)[g_model_buffer_cache count], - ds4_gpu_gib(g_model_buffer_cache_bytes), - ds4_gpu_gib(limit), - (unsigned long long)g_model_buffer_cache_evictions); - } - } - if (g_stream_expert_cache_hits != 0 || - g_stream_expert_cache_misses != 0 || - g_stream_expert_cache_bytes != 0) { - const uint64_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - uint64_t target_bytes = 0; - if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) { - target_bytes = - budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ? - UINT64_MAX : - budget * g_stream_expert_cache_expert_bytes; - } - const uint64_t lookups = g_stream_expert_cache_hits + g_stream_expert_cache_misses; - const double hit_rate = lookups ? - (double)g_stream_expert_cache_hits / (double)lookups : 0.0; - if (g_stream_expert_cache_evict_advise_bytes != 0 || - g_stream_expert_cache_willneed_advise_bytes != 0 || - g_stream_expert_cache_pread_bytes != 0) { - fprintf(stderr, - "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu evict_dontneed=%.2f GiB miss_willneed=%.2f GiB miss_pread=%.2f GiB pread_ms=%.3f\n", - (unsigned long long)budget, - g_stream_expert_cache_entry_count, - ds4_gpu_mib(g_stream_expert_cache_expert_bytes), - ds4_gpu_gib(target_bytes), - ds4_gpu_gib(g_stream_expert_cache_bytes), - (unsigned long long)g_stream_expert_cache_hits, - (unsigned long long)g_stream_expert_cache_misses, - hit_rate, - (unsigned long long)g_stream_expert_cache_wraps, - (unsigned long long)g_stream_expert_cache_evictions, - (unsigned long long)g_stream_expert_cache_buffer_allocs, - (unsigned long long)g_stream_expert_cache_buffer_reuses, - ds4_gpu_gib(g_stream_expert_cache_evict_advise_bytes), - ds4_gpu_gib(g_stream_expert_cache_willneed_advise_bytes), - ds4_gpu_gib(g_stream_expert_cache_pread_bytes), - g_stream_expert_cache_pread_ms); - } else { - fprintf(stderr, - "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu\n", - (unsigned long long)budget, - g_stream_expert_cache_entry_count, - ds4_gpu_mib(g_stream_expert_cache_expert_bytes), - ds4_gpu_gib(target_bytes), - ds4_gpu_gib(g_stream_expert_cache_bytes), - (unsigned long long)g_stream_expert_cache_hits, - (unsigned long long)g_stream_expert_cache_misses, - hit_rate, - (unsigned long long)g_stream_expert_cache_wraps, - (unsigned long long)g_stream_expert_cache_evictions, - (unsigned long long)g_stream_expert_cache_buffer_allocs, - (unsigned long long)g_stream_expert_cache_buffer_reuses); - } - if (ds4_gpu_stream_expert_timing_summary_enabled()) { - const ds4_gpu_stream_expert_timing_snapshot total = - ds4_gpu_stream_expert_timing_current(); - if (ds4_gpu_stream_expert_timing_has_data(total)) { - const ds4_gpu_stream_expert_timing_snapshot delta = - ds4_gpu_stream_expert_timing_delta( - total, - g_stream_expert_timing_last_report); - ds4_gpu_stream_expert_timing_print("total", total); - ds4_gpu_stream_expert_timing_print("delta", delta); - g_stream_expert_timing_last_report = total; - } - } - if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS") != NULL) { - fprintf(stderr, "ds4: streaming expert cache per-layer stats:\n"); - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - const uint64_t hits = g_stream_expert_cache_layer_hits[layer]; - const uint64_t misses = g_stream_expert_cache_layer_misses[layer]; - const uint64_t lookups = hits + misses; - const uint64_t evictions = - g_stream_expert_cache_layer_evictions[layer]; - const uint64_t pread_bytes = - g_stream_expert_cache_layer_pread_bytes[layer]; - const double pread_ms = - g_stream_expert_cache_layer_pread_ms[layer]; - const uint32_t cached = - g_stream_expert_cache_layer_count[layer]; - const uint32_t layer_slots = - ds4_gpu_stream_expert_cache_configured_count() != 0 ? - DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT : 0; - if (lookups == 0 && evictions == 0 && pread_bytes == 0 && - cached == 0) { - continue; - } - const double layer_hit_rate = lookups ? - (double)hits / (double)lookups : 0.0; - fprintf(stderr, - "ds4: layer=%u layer_slots=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", - layer, - layer_slots, - cached, - (unsigned long long)hits, - (unsigned long long)misses, - layer_hit_rate, - (unsigned long long)evictions, - ds4_gpu_gib(pread_bytes), - pread_ms); - } - if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS_DELTA") != NULL) { - fprintf(stderr, "ds4: streaming expert cache per-layer delta:\n"); - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - const uint64_t hits = - g_stream_expert_cache_layer_hits[layer]; - const uint64_t misses = - g_stream_expert_cache_layer_misses[layer]; - const uint64_t evictions = - g_stream_expert_cache_layer_evictions[layer]; - const uint64_t pread_bytes = - g_stream_expert_cache_layer_pread_bytes[layer]; - const double pread_ms = - g_stream_expert_cache_layer_pread_ms[layer]; - const uint64_t delta_hits = - hits - g_stream_expert_cache_layer_last_hits[layer]; - const uint64_t delta_misses = - misses - g_stream_expert_cache_layer_last_misses[layer]; - const uint64_t delta_evictions = - evictions - - g_stream_expert_cache_layer_last_evictions[layer]; - const uint64_t delta_pread_bytes = - pread_bytes - - g_stream_expert_cache_layer_last_pread_bytes[layer]; - const double delta_pread_ms = - pread_ms - - g_stream_expert_cache_layer_last_pread_ms[layer]; - const uint64_t lookups = delta_hits + delta_misses; - if (lookups == 0 && delta_evictions == 0 && - delta_pread_bytes == 0) { - continue; - } - const double hit_rate = lookups ? - (double)delta_hits / (double)lookups : 0.0; - fprintf(stderr, - "ds4: layer=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", - layer, - g_stream_expert_cache_layer_count[layer], - (unsigned long long)delta_hits, - (unsigned long long)delta_misses, - hit_rate, - (unsigned long long)delta_evictions, - ds4_gpu_gib(delta_pread_bytes), - delta_pread_ms); - g_stream_expert_cache_layer_last_hits[layer] = hits; - g_stream_expert_cache_layer_last_misses[layer] = misses; - g_stream_expert_cache_layer_last_evictions[layer] = - evictions; - g_stream_expert_cache_layer_last_pread_bytes[layer] = - pread_bytes; - g_stream_expert_cache_layer_last_pread_ms[layer] = - pread_ms; - } - } - } - } - fprintf(stderr, - "ds4: model residency requests %llu%s\n", - (unsigned long long)g_model_residency_count, - g_ssd_streaming_mode ? " (ssd-streaming)" : - (getenv("DS4_METAL_NO_RESIDENCY") != NULL ? " (disabled)" : "")); - fprintf(stderr, - "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", - g_metal_device_name[0] ? g_metal_device_name : "(unknown)", - g_metal4_runtime_available ? "yes" : "no", - g_metal4_family_supported ? "yes" : "no", - g_metal4_queue_supported ? "yes" : "no", - g_metal4_tensor_api_enabled ? "enabled" : - (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), - g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); - fprintf(stderr, - "ds4: accelerated Metal path %s%s\n", - ds4_gpu_mpp_available() ? "enabled" : "disabled", - g_quality_mode ? " by --quality" : - (!g_metal4_tensor_api_enabled ? " (tensor API unavailable)" : "")); - fprintf(stderr, - "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", - g_metal_device_name[0] ? g_metal_device_name : "(unknown)", - g_metal4_runtime_available ? "yes" : "no", - g_metal4_family_supported ? "yes" : "no", - g_metal4_queue_supported ? "yes" : "no", - g_metal4_tensor_api_enabled ? "enabled" : - (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), - g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); - fprintf(stderr, - "ds4: scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n", - ds4_gpu_mib(scratch), - ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes + - (uint64_t)g_glm_flash_attn_mask_bytes + - (uint64_t)g_flash_attn_zero_mask_bytes + - cached_prefill_mask_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes + - cached_prefill_blk_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes), - ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes + - (uint64_t)g_compressor_pool_score_bytes + - (uint64_t)g_compressor_pool_score_cont_bytes + - (uint64_t)g_compressor_pool_softmax_bytes + - (uint64_t)g_compressor_pool_product_bytes + - (uint64_t)g_compressor_store_ape_bytes + - (uint64_t)g_compressor_store_score_bytes + - (uint64_t)g_embed_rows_bytes), - ds4_gpu_mib((uint64_t)g_router_selection_bytes + - (uint64_t)g_router_weight_sum_bytes), - ds4_gpu_mib((uint64_t)g_indexer_head_scores_bytes + - (uint64_t)g_indexer_topk_bytes + - (uint64_t)g_indexed_topk_bytes), - ds4_gpu_mib((uint64_t)g_moe_gate_scratch_bytes + - (uint64_t)g_moe_down_scratch_bytes + - (uint64_t)g_moe_id_map_bytes + - (uint64_t)g_moe_q4_gate_slots_bytes + - (uint64_t)g_moe_q4_up_slots_bytes + - (uint64_t)g_moe_q4_down_slots_bytes), - ds4_gpu_mib((uint64_t)g_f16_round_scratch_bytes), - ds4_gpu_mib((uint64_t)g_raw_store_round_bytes)); - if (color) fputs(reset, stderr); -} - -void ds4_gpu_set_quality(bool quality) { - g_quality_mode = quality ? 1 : 0; -} - -void ds4_gpu_set_glm_model(bool enabled) { - g_glm_model_mode = enabled ? 1 : 0; -} - -void ds4_gpu_set_ssd_streaming(bool enabled) { - g_ssd_streaming_mode = enabled ? 1 : 0; - ds4_gpu_stream_expert_cache_clear_all(1); - if (g_ssd_streaming_mode) { - fprintf(stderr, - "ds4: Metal SSD streaming mode enabled; full model residency and warmup are skipped\n"); - } -} - -void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { - g_glm_streaming_prefill_full_layer_runtime = enabled ? 1 : 0; -} - -void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { - if (experts > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { - experts = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; - } - g_stream_expert_cache_budget_override = experts; - ds4_gpu_stream_expert_cache_clear_all(1); -} - -void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { - /* - * Pre-seed the cache's single slab size class with the model's uniform - * per-expert bytes (first routed layer). With a mixed-precision GGUF this - * pins the class to the majority layers so the boosted ones are rejected - * deterministically from startup, instead of depending on which layer - * happens to touch the cache first. - */ - g_stream_expert_cache_expert_bytes = bytes; -} - -uint64_t ds4_gpu_recommended_working_set_size(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_device) return 0; - return (uint64_t)[g_device recommendedMaxWorkingSetSize]; -} - -static int ds4_gpu_model_map_log_enabled(void) { - if (!g_ssd_streaming_mode) return 1; - const char *trace = getenv("DS4_METAL_STREAMING_MAP_TRACE"); - return trace && trace[0] && strcmp(trace, "0") != 0; -} - -static id ds4_gpu_wrap_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset); - -static id ds4_gpu_wrap_model_exact_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset); - -static const char *ds4_gpu_source = -"#include \n" -"#ifdef DS4_METAL_HAS_TENSOR\n" -"#include \n" -"#include \n" -"#endif\n" -"using namespace metal;\n" -"#ifdef DS4_METAL_HAS_TENSOR\n" -"using namespace mpp::tensor_ops;\n" -"#endif\n" -"\n" -"#define MAX(x, y) ((x) > (y) ? (x) : (y))\n" -"#define MIN(x, y) ((x) < (y) ? (x) : (y))\n" -"#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }\n" -"#define QK8_0 32\n" -"#ifndef QK_K\n" -"#define QK_K 256\n" -"#endif\n" -"#define N_SIMDWIDTH 32\n" -"#define N_R0_Q8_0 2\n" -"#define N_SG_Q8_0 4\n" -"#define FC_MUL_MV 600\n" -"#define FC_MUL_MM 700\n" -"#define FC_BIN 1300\n" -"#define FOR_UNROLL(x) _Pragma(\"clang loop unroll(full)\") for (x)\n" -"#define M_PI_F 3.14159265358979323846f\n" -"\n" -"// Reads one byte per stride to warm model-backed pages without copying the\n" -"// model. This is outside inference and exists only to reduce first-use stalls.\n" -"kernel void kernel_touch_u8_stride(\n" -" device const uchar *src [[buffer(0)]],\n" -" device uchar *dst [[buffer(1)]],\n" -" constant ulong &stride [[buffer(2)]],\n" -" constant ulong &bytes [[buffer(3)]],\n" -" constant ulong &dst_offset [[buffer(4)]],\n" -" uint gid [[thread_position_in_grid]]) {\n" -" ulong off = (ulong)gid * stride;\n" -" if (off >= bytes) return;\n" -" dst[dst_offset + (ulong)gid] = src[off];\n" -"}\n" -"\n" -"enum ds4_sort_order {\n" -" DS4_SORT_ORDER_ASC,\n" -" DS4_SORT_ORDER_DESC,\n" -"};\n" -"\n" -"struct block_q8_0 {\n" -" half d;\n" -" int8_t qs[QK8_0];\n" -"};\n" -"\n" -"struct block_q8_K {\n" -" float d;\n" -" int8_t qs[QK_K];\n" -" int16_t bsums[QK_K / 16];\n" -"};\n" -"\n" -"\n"; - -static NSString *ds4_gpu_full_source(void) { - NSString *base = [NSString stringWithUTF8String:ds4_gpu_source]; - NSFileManager *fm = [NSFileManager defaultManager]; - /* - * Kernels are kept as separate files for review, then concatenated into one - * Metal library. Environment overrides are still honored so a diagnostic - * run can swap one source file without changing the executable. - */ - NSArray *> *required_sources = @[ - @[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"], - @[@"DS4_METAL_DENSE_SOURCE", @"metal/dense.metal"], - @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], - @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], - @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], - @[@"DS4_METAL_DSV4_KV_SOURCE", @"metal/dsv4_kv.metal"], - @[@"DS4_METAL_DSV4_ROPE_SOURCE", @"metal/dsv4_rope.metal"], - @[@"DS4_METAL_DSV4_MISC_SOURCE", @"metal/dsv4_misc.metal"], - @[@"DS4_METAL_ARGSORT_SOURCE", @"metal/argsort.metal"], - @[@"DS4_METAL_CPY_SOURCE", @"metal/cpy.metal"], - @[@"DS4_METAL_CONCAT_SOURCE", @"metal/concat.metal"], - @[@"DS4_METAL_GET_ROWS_SOURCE", @"metal/get_rows.metal"], - @[@"DS4_METAL_SUM_ROWS_SOURCE", @"metal/sum_rows.metal"], - @[@"DS4_METAL_SOFTMAX_SOURCE", @"metal/softmax.metal"], - @[@"DS4_METAL_REPEAT_SOURCE", @"metal/repeat.metal"], - @[@"DS4_METAL_GLU_SOURCE", @"metal/glu.metal"], - @[@"DS4_METAL_NORM_SOURCE", @"metal/norm.metal"], - @[@"DS4_METAL_BIN_SOURCE", @"metal/bin.metal"], - @[@"DS4_METAL_SET_ROWS_SOURCE", @"metal/set_rows.metal"], - ]; - - NSMutableString *source = [NSMutableString stringWithString:base]; - for (NSArray *spec in required_sources) { - const char *override_path = getenv([spec[0] UTF8String]); - NSMutableArray *paths = [NSMutableArray array]; - if (override_path && override_path[0]) { - [paths addObject:[NSString stringWithUTF8String:override_path]]; - } - [paths addObject:spec[1]]; - [paths addObject:[@"./" stringByAppendingString:spec[1]]]; - - NSString *loaded = nil; - NSString *loaded_path = nil; - for (NSString *path in paths) { - if (![fm fileExistsAtPath:path]) continue; - - NSError *error = nil; - loaded = [NSString stringWithContentsOfFile:path - encoding:NSUTF8StringEncoding - error:&error]; - if (!loaded) { - fprintf(stderr, "ds4: failed to read Metal source %s: %s\n", - [path UTF8String], [[error localizedDescription] UTF8String]); - return nil; - } - loaded_path = path; - break; - } - - if (!loaded) { - fprintf(stderr, - "ds4: Metal source %s not found (set %s to override)\n", - [spec[1] UTF8String], [spec[0] UTF8String]); - return nil; - } - [source appendFormat:@"\n// appended %@\n%@\n", loaded_path, loaded]; - } - return source; -} - -typedef struct { - int32_t ne00t; - int32_t ne00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne10; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; -} ds4_gpu_get_rows_args; - -typedef struct { - int32_t n_embd; - int32_t n_vocab; - int32_t n_tokens; - uint64_t src_row_bytes; - uint64_t dst_row_bytes; - uint64_t token_stride; -} ds4_gpu_get_rows_q8_0_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; -} ds4_gpu_repeat_args; - -typedef struct { - int32_t nk0; - int32_t ne01; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne11; - int32_t ne12; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; -} ds4_gpu_set_rows_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne10; - int32_t ne11; - int32_t ne12; - int32_t ne13; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - int32_t dim; -} ds4_gpu_concat_args; - -typedef struct { - int64_t nk0; - int64_t ne00; - int64_t ne01; - int64_t ne02; - int64_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int64_t ne0; - int64_t ne1; - int64_t ne2; - int64_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; -} ds4_gpu_cpy_args; - -static ds4_gpu_cpy_args ds4_gpu_make_cpy_1d_args( - uint32_t n, - uint64_t src_elem, - uint64_t dst_elem) { - return (ds4_gpu_cpy_args) { - .nk0 = (int64_t)n, - .ne00 = (int64_t)n, - .ne01 = 1, - .ne02 = 1, - .ne03 = 1, - .nb00 = src_elem, - .nb01 = (uint64_t)n * src_elem, - .nb02 = (uint64_t)n * src_elem, - .nb03 = (uint64_t)n * src_elem, - .ne0 = (int64_t)n, - .ne1 = 1, - .ne2 = 1, - .ne3 = 1, - .nb0 = dst_elem, - .nb1 = (uint64_t)n * dst_elem, - .nb2 = (uint64_t)n * dst_elem, - .nb3 = (uint64_t)n * dst_elem, - }; -} - -static NSUInteger ds4_gpu_cpy_threads(uint32_t n, id pipeline) { - NSUInteger nth = 32u; - const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; - while (nth < (NSUInteger)n && nth < max_threads) nth *= 2u; - if (nth > max_threads) nth = max_threads; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - return nth ? nth : 1u; -} - -static float ds4_gpu_negative_infinity(void) { - union { uint32_t u; float f; } v = { 0xff800000u }; - return v.f; -} - -static float ds4_gpu_positive_infinity(void) { - union { uint32_t u; float f; } v = { 0x7f800000u }; - return v.f; -} - -static int ds4_gpu_encode_cpy_f32_f32_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n); - -static int ds4_gpu_encode_cpy_f32_f32_3d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride); - -static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_col_stride, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride); - -static int ds4_gpu_encode_cpy_f32_f16_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n); - -static int ds4_gpu_encode_cpy_f32_f16_2d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint64_t src_row_stride, - uint64_t dst_row_stride); - -static int ds4_gpu_encode_cpy_f16_f32_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n); - -static int ds4_gpu_encode_fill_f32_rows( - id cb, - id buf, - NSUInteger offset, - uint32_t width, - uint32_t rows, - float value); - -static int ds4_gpu_encode_add_f32_1d( - id cb, - id a, - NSUInteger a_off, - id b, - NSUInteger b_off, - id out, - NSUInteger out_off, - uint32_t n); - -typedef struct { - int32_t ne00; - uint64_t nb01; - int32_t ne10; - uint64_t nb11; - int32_t ne0; - uint64_t nb1; - int32_t i00; - int32_t i10; - float alpha; - float limit; -} ds4_gpu_glu_args; - -typedef struct { - uint32_t n; -} ds4_gpu_add_flat_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne10; - int32_t ne11; - int32_t ne12; - int32_t ne13; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - uint64_t offs; - uint64_t o1[8]; -} ds4_gpu_bin_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - float slope; - float scale; - float bias; - float val; - float min; - float max; -} ds4_gpu_unary_args; - -static ds4_gpu_bin_args ds4_gpu_make_bin_rows_args(uint32_t n, uint32_t rows, uint32_t rhs_n) { - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - const uint64_t rhs_row_bytes = (uint64_t)rhs_n * sizeof(float); - return (ds4_gpu_bin_args) { - .ne00 = (int32_t)n, - .ne01 = (int32_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = row_bytes, - .nb03 = row_bytes, - .ne10 = (int32_t)rhs_n, - .ne11 = 1, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = rhs_row_bytes, - .nb12 = rhs_row_bytes, - .nb13 = rhs_row_bytes, - .ne0 = (int32_t)n, - .ne1 = (int32_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = row_bytes, - .nb2 = row_bytes, - .nb3 = row_bytes, - .offs = 0, - .o1 = { 0 }, - }; -} - -static ds4_gpu_unary_args ds4_gpu_make_unary_rows_args( - uint32_t n, - uint32_t rows, - int c4, - float scale, - float bias) { - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - const uint32_t n_kernel = c4 ? n / 4u : n; - return (ds4_gpu_unary_args) { - .ne00 = (int32_t)n_kernel, - .ne01 = (int32_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = row_bytes, - .nb03 = row_bytes, - .ne0 = (int32_t)n_kernel, - .ne1 = (int32_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = row_bytes, - .nb2 = row_bytes, - .nb3 = row_bytes, - .slope = 0.0f, - .scale = scale, - .bias = bias, - .val = 0.0f, - .min = 0.0f, - .max = 0.0f, - }; -} - -static ds4_gpu_bin_args ds4_gpu_make_bin_same_rows_args(uint32_t n, uint32_t rows) { - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - return (ds4_gpu_bin_args) { - .ne00 = (int32_t)n, - .ne01 = (int32_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = (uint64_t)rows * row_bytes, - .nb03 = (uint64_t)rows * row_bytes, - .ne10 = (int32_t)n, - .ne11 = (int32_t)rows, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = row_bytes, - .nb12 = (uint64_t)rows * row_bytes, - .nb13 = (uint64_t)rows * row_bytes, - .ne0 = (int32_t)n, - .ne1 = (int32_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = row_bytes, - .nb2 = (uint64_t)rows * row_bytes, - .nb3 = (uint64_t)rows * row_bytes, - .offs = 0, - .o1 = { 0 }, - }; -} - -static int ds4_gpu_encode_bin_f32_rows( - id cb, - id pipeline, - const ds4_gpu_bin_args *args, - id a, - NSUInteger a_off, - id b, - NSUInteger b_off, - id out, - NSUInteger out_off); - -static int ds4_gpu_encode_sum_rows_f32( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t width, - uint32_t rows); - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne10; - int32_t ne11; - int32_t ne12; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne0; - int32_t ne1; - int32_t nr0; - int16_t r2; - int16_t r3; -} ds4_gpu_q8_0_matvec_args; - -typedef struct { - int32_t ne00; - int32_t ne02; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne12; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne0; - int32_t ne1; - int16_t r2; - int16_t r3; -} ds4_gpu_mul_mm_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne10; - int32_t ne11; - int32_t ne12; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne0; - int32_t ne1; - int16_t r2; - int16_t r3; -} ds4_gpu_mul_mv_ext_args; - -typedef ds4_gpu_q8_0_matvec_args ds4_gpu_f16_matvec_args; - -static ds4_gpu_q8_0_matvec_args ds4_gpu_make_q8_0_mv_args(uint64_t in_dim, uint64_t out_dim) { - const uint64_t row_bytes = (in_dim / 32u) * 34u; - return (ds4_gpu_q8_0_matvec_args) { - .ne00 = (int32_t)in_dim, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = 34, - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)in_dim, - .ne11 = 1, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * sizeof(float), - .nb13 = in_dim * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = 1, - .nr0 = 2, - .r2 = 1, - .r3 = 1, - }; -} - -static ds4_gpu_f16_matvec_args ds4_gpu_make_f16_mv_args(uint64_t in_dim, uint64_t out_dim) { - const uint64_t row_bytes = in_dim * sizeof(uint16_t); - return (ds4_gpu_f16_matvec_args) { - .ne00 = (int32_t)in_dim, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = sizeof(uint16_t), - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)in_dim, - .ne11 = 1, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * sizeof(float), - .nb13 = in_dim * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = 1, - .nr0 = 2, - .r2 = 1, - .r3 = 1, - }; -} - -static ds4_gpu_q8_0_matvec_args ds4_gpu_make_f32_mv_args( - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_vec) { - const uint64_t row_bytes = in_dim * sizeof(float); - return (ds4_gpu_q8_0_matvec_args) { - .ne00 = (int32_t)in_dim, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)in_dim, - .ne11 = (int32_t)n_vec, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * n_vec * sizeof(float), - .nb13 = in_dim * n_vec * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = (int32_t)n_vec, - .nr0 = 2, - .r2 = 1, - .r3 = 1, - }; -} - -typedef struct { - const char *function_name; - int16_t nsg; - int32_t nr0; - NSUInteger smem; -} ds4_gpu_mv_dispatch; - -static int ds4_gpu_tp_world_is_two(void); - -static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) { - const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u; - const int16_t nsg = - (int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u); - const uint64_t rows = ds4_gpu_env_u64("DS4_METAL_Q8_MV_ROWS", 2u, 2u, 4u); - if (rows >= 4u) { - return (ds4_gpu_mv_dispatch) { - .function_name = "kernel_mul_mv_q8_0_f32_r4", - .nsg = nsg, - .nr0 = 4, - .smem = 32u * 4u * sizeof(float), - }; - } - return (ds4_gpu_mv_dispatch) { - .function_name = "kernel_mul_mv_q8_0_f32", - .nsg = nsg, - .nr0 = 2, - .smem = 32u * 2u * sizeof(float), - }; -} - -static ds4_gpu_mv_dispatch ds4_gpu_make_plain_mv_dispatch( - uint64_t in_dim, - int f32_weights) { - if (in_dim < 32) { - return (ds4_gpu_mv_dispatch) { - .function_name = f32_weights ? "kernel_mul_mv_f32_f32_short" : "kernel_mul_mv_f16_f32_short", - .nsg = 1, - .nr0 = 32, - .smem = 0, - }; - } - - const int16_t nsg = (int16_t)((in_dim + 127u) / 128u > 8u ? 8u : (in_dim + 127u) / 128u); - const int use_4 = (in_dim % 4u) == 0; - return (ds4_gpu_mv_dispatch) { - .function_name = f32_weights - ? (use_4 ? "kernel_mul_mv_f32_f32_4" : "kernel_mul_mv_f32_f32") - : (use_4 ? "kernel_mul_mv_f16_f32_4" : "kernel_mul_mv_f16_f32"), - .nsg = nsg, - .nr0 = 2, - .smem = 32u * 2u * sizeof(float), - }; -} - -static ds4_gpu_mul_mm_args ds4_gpu_make_mm_args( - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t row_bytes) { - return (ds4_gpu_mul_mm_args) { - .ne00 = (int32_t)in_dim, - .ne02 = 1, - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * n_tok * sizeof(float), - .nb13 = in_dim * n_tok * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = (int32_t)n_tok, - .r2 = 1, - .r3 = 1, - }; -} - -static ds4_gpu_mul_mv_ext_args ds4_gpu_make_mv_ext_args( - uint64_t in_dim, - uint64_t out_dim, - uint64_t n_tok, - uint64_t elem_bytes, - uint64_t row_bytes) { - return (ds4_gpu_mul_mv_ext_args) { - .ne00 = (int32_t)in_dim, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = elem_bytes, - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)in_dim, - .ne11 = (int32_t)n_tok, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * n_tok * sizeof(float), - .nb13 = in_dim * n_tok * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = (int32_t)n_tok, - .r2 = 1, - .r3 = 1, - }; -} - -static int16_t ds4_gpu_mv_ext_nxpsg(uint64_t in_dim, uint64_t n_tok) { - if ((in_dim % 256u) == 0 && n_tok < 3) return 16; - if ((in_dim % 128u) == 0) return 8; - return 4; -} - -static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) { - switch (n_tok) { - case 2: return 2; - case 3: - case 6: return 3; - case 4: - case 7: - case 8: return 4; - case 5: return 5; - default: return n_tok > 8 ? 4 : 0; - } -} - -static const char *ds4_gpu_mv_ext_name(int q8, int16_t r1ptg) { - if (q8) { - switch (r1ptg) { - case 2: return "kernel_mul_mv_ext_q8_0_f32_r1_2"; - case 3: return "kernel_mul_mv_ext_q8_0_f32_r1_3"; - case 4: return "kernel_mul_mv_ext_q8_0_f32_r1_4"; - case 5: return "kernel_mul_mv_ext_q8_0_f32_r1_5"; - default: return NULL; - } - } - - switch (r1ptg) { - case 2: return "kernel_mul_mv_ext_f16_f32_r1_2"; - case 3: return "kernel_mul_mv_ext_f16_f32_r1_3"; - case 4: return "kernel_mul_mv_ext_f16_f32_r1_4"; - case 5: return "kernel_mul_mv_ext_f16_f32_r1_5"; - default: return NULL; - } -} - -static const char *ds4_gpu_mv_ext_f32_name(int16_t r1ptg) { - switch (r1ptg) { - case 2: return "kernel_mul_mv_ext_f32_f32_r1_2"; - case 3: return "kernel_mul_mv_ext_f32_f32_r1_3"; - case 4: return "kernel_mul_mv_ext_f32_f32_r1_4"; - case 5: return "kernel_mul_mv_ext_f32_f32_r1_5"; - default: return NULL; - } -} - -static const char *ds4_gpu_mv_ext_q8_pair_swiglu_name(int16_t r1ptg) { - switch (r1ptg) { - case 2: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2"; - case 3: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3"; - case 4: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4"; - case 5: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5"; - default: return NULL; - } -} - -typedef struct { - int32_t ne00; - int32_t ne00_t; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - float eps; - int32_t nef1[3]; - int32_t nef2[3]; - int32_t nef3[3]; - uint64_t nbf1[3]; - uint64_t nbf2[3]; - uint64_t nbf3[3]; -} ds4_gpu_rms_norm_args; - -typedef struct { - int32_t q_n; - int32_t q_n4; - int32_t kv_n; - int32_t kv_n4; - uint64_t q_row_stride; - uint64_t kv_row_stride; - float eps; -} ds4_gpu_qkv_rms_norm_args; - -static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) { - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - return (ds4_gpu_rms_norm_args) { - .ne00 = (int32_t)n, - .ne00_t = (int32_t)(n / 4u), - .nb1 = row_bytes, - .nb2 = row_bytes * rows, - .nb3 = row_bytes * rows, - .eps = eps, - .nef1 = { (int32_t)rows, 1, 1 }, - .nef2 = { 1, 1, 1 }, - .nef3 = { 1, 1, 1 }, - .nbf1 = { row_bytes, row_bytes, row_bytes }, - .nbf2 = { row_bytes * rows, row_bytes, row_bytes }, - .nbf3 = { row_bytes * rows, row_bytes, row_bytes }, - }; -} - -static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_3d_args( - uint32_t n0, - uint32_t n1, - uint32_t n2, - float eps) { - const uint64_t row_bytes = (uint64_t)n0 * sizeof(float); - const uint64_t plane_bytes = row_bytes * n1; - return (ds4_gpu_rms_norm_args) { - .ne00 = (int32_t)n0, - .ne00_t = (int32_t)(n0 / 4u), - .nb1 = row_bytes, - .nb2 = plane_bytes, - .nb3 = plane_bytes * n2, - .eps = eps, - .nef1 = { (int32_t)n1, 1, 1 }, - .nef2 = { (int32_t)n2, 1, 1 }, - .nef3 = { 1, 1, 1 }, - .nbf1 = { row_bytes, row_bytes, row_bytes }, - .nbf2 = { plane_bytes, row_bytes, row_bytes }, - .nbf3 = { plane_bytes * n2, row_bytes, row_bytes }, - }; -} - -static NSUInteger ds4_gpu_rms_norm_threads(uint32_t n) { - NSUInteger ne00_t = n / 4u; - NSUInteger nth = 32u; - while (nth < ne00_t && nth < 1024u) nth *= 2u; - if (nth > ne00_t) nth = ne00_t; - return nth ? nth : 1u; -} - -static NSUInteger ds4_gpu_rms_norm_pipeline_threads( - uint32_t n, - id pipeline) { - NSUInteger ne00_t = n / 4u; - NSUInteger max_threads = pipeline ? [pipeline maxTotalThreadsPerThreadgroup] : 1024u; - NSUInteger nth = 32u; - while (nth < ne00_t && nth < max_threads) nth *= 2u; - if (nth > max_threads) nth = max_threads; - if (nth > ne00_t) nth = ne00_t; - return nth ? nth : 1u; -} - -typedef struct { - int32_t n_hc; - int32_t sinkhorn_iters; - int64_t n_rows; - int64_t mix_hc; - uint64_t nb01; - uint64_t nb1; - float eps; -} ds4_gpu_hc_split_args; - -typedef struct { - int64_t n_embd; - int64_t n_hc; - int64_t n_tokens; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb_w0; - uint64_t nb_w1; - uint64_t nb0; - uint64_t nb1; -} ds4_gpu_hc_weighted_sum_args; - -typedef struct { - int64_t n_embd; - int64_t n_hc; - int64_t n_tokens; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb_w0; - uint64_t nb_w1; - uint64_t nb0; - uint64_t nb1; - uint64_t nb_norm1; - float norm_eps; -} ds4_gpu_hc_weighted_sum_norm_args; - -typedef struct { - float post_scale; - float eps; -} ds4_gpu_output_hc_weights4_args; - -typedef struct { - int64_t n_embd; - int32_t n_hc; - int32_t sinkhorn_iters; - int64_t n_rows; - int64_t mix_hc; - uint64_t nb_mix1; - uint64_t nb_split1; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb0; - uint64_t nb1; - float eps; -} ds4_gpu_hc_split_weighted_sum_args; - -typedef struct { - int64_t n_embd; - int32_t n_hc; - int32_t sinkhorn_iters; - int64_t n_rows; - int64_t mix_hc; - uint64_t nb_mix1; - uint64_t nb_split1; - uint64_t nb_x0; - uint64_t nb_x1; - uint64_t nb_x2; - uint64_t nb0; - uint64_t nb1; - uint64_t nb_norm1; - float eps; - float norm_eps; -} ds4_gpu_hc_split_weighted_sum_norm_args; - -typedef struct { - int64_t n_embd; - int64_t n_hc; - int64_t n_tokens; - uint64_t nb_block0; - uint64_t nb_block1; - uint64_t nb_add0; - uint64_t nb_add1; - uint64_t nb_res0; - uint64_t nb_res1; - uint64_t nb_res2; - uint64_t nb_post0; - uint64_t nb_post1; - uint64_t nb_comb0; - uint64_t nb_comb1; - uint64_t nb_comb2; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - int32_t has_add; -} ds4_gpu_hc_expand_args; - -typedef struct { - int32_t nei0; - int32_t nei1; - uint64_t nbi1; - int32_t ne00; - int32_t ne01; - int32_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - int32_t ne10; - int32_t ne11; - int32_t ne12; - int32_t ne13; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - int32_t ne0; - int32_t ne1; - uint64_t nb1; - int32_t nr0; - /* Tensor-parallel expert ownership; see ds4_metal_args_mul_mv_id in - * metal/moe.metal. Zero (from struct literals) means no split. */ - int32_t tp_rank; - int32_t tp_world; - int32_t tp_addend; - int32_t tp_expert_base; -} ds4_gpu_mul_mv_id_args; - -typedef struct { - uint32_t n_total_expert; - uint32_t n_expert; -} ds4_gpu_stream_expert_validate_args; - -typedef struct { - uint32_t active_mask; - uint32_t accumulate; -} ds4_gpu_stream_expert_split_args; - -typedef struct { - int32_t ne02; - int32_t ne10; - int32_t ne11; - uint64_t nb11; - uint64_t nb12; - int32_t ne21; - int32_t ne20; - uint64_t nb21; -} ds4_gpu_mul_mm_id_map_args; - -typedef struct { - int32_t ne00; - int32_t ne02; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne11; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ne20; - int32_t ne21; - int32_t ne0; - int32_t ne1; - int16_t r2; - int16_t r3; - int32_t tp_rank; - int32_t tp_world; - int32_t tp_expert_base; -} ds4_gpu_mul_mm_id_args; - -static int ds4_gpu_encode_mul_mv_id( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0); - -static int ds4_gpu_encode_attn_out_low_q8_direct( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0); - -static int ds4_gpu_encode_attn_out_low_q8_mpp( - id cb, - id pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off); - -static int ds4_gpu_encode_attn_out_low_q8_mpp( - id cb, - id pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off); - -static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( - uint32_t src0_cols, - uint32_t src0_experts, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens); - -static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( - uint32_t src0_cols, - uint32_t src0_rows, - uint32_t src0_experts, - uint64_t src0_row_bytes, - uint64_t src0_expert_bytes, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens); -static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( - uint32_t src0_cols, - uint32_t src0_rows, - uint32_t src0_experts, - uint64_t src0_row_bytes, - uint64_t src0_expert_bytes, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens, - uint32_t src1_elem_size); - -static int ds4_gpu_encode_mul_mm_id( - id cb, - id map_pipeline, - id mm_pipeline, - const ds4_gpu_mul_mm_id_map_args *map_args, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off); - -static int ds4_gpu_encode_mul_mm_id_map( - id cb, - id map_pipeline, - const ds4_gpu_mul_mm_id_map_args *map_args, - const ds4_gpu_mul_mm_id_args *mm_args, - id ids, - NSUInteger ids_off); - -static int ds4_gpu_encode_mul_mm_id_mapped( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off); -static int ds4_gpu_encode_mul_mm_id_mapped_tile( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes); -static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0_addrs, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes, - ds4_gpu_stream_expert_cache_entry * const *resources, - uint32_t resource_count, - uint32_t resource_kind, - id overflow_resource); - -typedef struct { - int32_t ne11; - int32_t ne_12_2; - int32_t ne_12_3; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - uint64_t nb21; - uint64_t nb22; - uint64_t nb23; - int32_t ne31; - int32_t ne32; - int32_t ne33; - uint64_t nb31; - uint64_t nb32; - uint64_t nb33; -} ds4_gpu_flash_attn_pad_args; - -typedef struct { - uint32_t raw_cap; - uint32_t raw_start; - uint32_t n_raw; - uint32_t n_comp; - uint32_t pad_rows; - uint32_t shared_pad; -} ds4_gpu_flash_kv_stage_f16_args; - -typedef struct { - int32_t ne01; - int32_t ne30; - int32_t ne31; - int32_t ne32; - int32_t ne33; - uint64_t nb31; - uint64_t nb32; - uint64_t nb33; -} ds4_gpu_flash_attn_blk_args; - -typedef struct { - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne11; - int32_t ne_12_2; - int32_t ne_12_3; - int32_t ns10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - int32_t ns20; - uint64_t nb21; - uint64_t nb22; - uint64_t nb23; - int32_t ne31; - int32_t ne32; - int32_t ne33; - uint64_t nb31; - uint64_t nb32; - uint64_t nb33; - int32_t ne1; - int32_t ne2; - int32_t ne3; - float scale; - float max_bias; - float m0; - float m1; - int32_t n_head_log2; - float logit_softcap; -} ds4_gpu_flash_attn_vec_args; - -typedef struct { - int32_t nrows; -} ds4_gpu_flash_attn_reduce_args; - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - int64_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - int32_t n_dims; - int32_t mode; - int32_t n_ctx_orig; - int32_t inverse; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - bool src2; -} ds4_gpu_rope_tail_batch_args; - -typedef struct { - uint64_t row_bytes; - uint64_t token_bytes; - int32_t head_dim; - int32_t n_dims; - int32_t n_ctx_orig; - int32_t inverse; - uint32_t pos0; - uint32_t pos_step; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -} ds4_gpu_rope_affine_pair_args; - -_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64, - "Metal affine RoPE argument ABI changed"); - -static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args( - uint32_t n_tok, - uint32_t n_head, - uint32_t head_dim, - uint32_t n_rot, - 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) { - const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); - const uint64_t tok_bytes = (uint64_t)n_head * row_bytes; - return (ds4_gpu_rope_tail_batch_args) { - .ne00 = head_dim, - .ne01 = n_head, - .ne02 = n_tok, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = tok_bytes, - .nb03 = (uint64_t)n_tok * tok_bytes, - .nb0 = sizeof(float), - .nb1 = row_bytes, - .nb2 = tok_bytes, - .nb3 = (uint64_t)n_tok * tok_bytes, - .n_dims = (int32_t)n_rot, - .mode = 0, - .n_ctx_orig = (int32_t)n_ctx_orig, - .inverse = inverse ? 1 : 0, - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .src2 = false, - }; -} - -static int ds4_gpu_encode_rope_tail_inplace( - id cb, - id xbuf, - NSUInteger xoff, - const ds4_gpu_rope_tail_batch_args *args, - uint32_t n_tok, - uint32_t n_head, - uint32_t head_dim, - uint32_t pos0, - uint32_t pos_step) { - const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u; - const bool lane_compatible = - tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u; - const bool force_affine_position = - getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL; - const bool use_inplace_pair = - g_rope_tail_inplace_pair_pipeline != nil && - args->mode == 0 && !args->src2 && - lane_compatible && - (ds4_gpu_device_name_contains("M3") || - (ds4_gpu_device_name_contains("M5") && n_tok == 1u) || - getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL || - force_affine_position) && - getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL; - const bool use_shared_coeff = - use_inplace_pair && !force_affine_position && - g_rope_tail_inplace_pair_shared4_pipeline != nil && - /* The 256-thread grouped schedule helps long prefill, but reduces the - * per-head parallelism that short batches and decode rely on. */ - tail_threads == 64u && n_head >= 4u && n_tok >= 32u && - getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL; - if (force_affine_position && - g_rope_tail_inplace_pair_affine_pipeline == nil && - getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL && - getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) { - fprintf(stderr, - "ds4: forced affine-position RoPE pipeline is unavailable\n"); - return 0; - } - /* Keep long prefill on the proven shared4 kernel. Reconstructing affine - * positions inside its coefficient cohort perturbs YaRN fast-math codegen; - * the compact affine specialization is exact for the decode pair schedule. */ - const bool use_affine_position = - use_inplace_pair && !use_shared_coeff && - g_rope_tail_inplace_pair_affine_pipeline != nil && - (n_tok == 1u || force_affine_position) && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - force_affine_position) && - getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL; - - int32_t pos_stack[256]; - int32_t *pos = NULL; - id posbuf = nil; - const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t); - if (!use_affine_position) { - pos = pos_stack; - if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { - pos = malloc((size_t)n_tok * sizeof(*pos)); - if (!pos) { - fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n"); - return 0; - } - } - for (uint32_t t = 0; t < n_tok; t++) { - pos[t] = (int32_t)(pos0 + t * pos_step); - } - - if (pos_bytes > 4096u) { - /* - * Metal inline setBytes data is meant for small constants. Long - * prefill RoPE calls need thousands of positions; passing that much - * inline can make the Apple driver abort the process. - */ - posbuf = ds4_gpu_new_transient_buffer( - pos_bytes, "ds4_rope_positions"); - if (!posbuf) { - if (pos != pos_stack) free(pos); - return 0; - } - memcpy([posbuf contents], pos, pos_bytes); - } - } - - ds4_gpu_rope_affine_pair_args affine_args; - if (use_affine_position) { - const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); - affine_args = (ds4_gpu_rope_affine_pair_args) { - .row_bytes = row_bytes, - .token_bytes = (uint64_t)n_head * row_bytes, - .head_dim = (int32_t)head_dim, - .n_dims = args->n_dims, - .n_ctx_orig = args->n_ctx_orig, - .inverse = args->inverse, - .pos0 = pos0, - .pos_step = pos_step, - .freq_base = args->freq_base, - .freq_scale = args->freq_scale, - .ext_factor = args->ext_factor, - .attn_factor = args->attn_factor, - .beta_fast = args->beta_fast, - .beta_slow = args->beta_slow, - }; - } - const NSUInteger reference_nth = - (NSUInteger)(head_dim < 256u ? head_dim : 256u); - const NSUInteger pair_nth = - (NSUInteger)(tail_threads < 256u ? tail_threads : 256u); - const NSUInteger nth = use_shared_coeff ? 256u : - (use_inplace_pair ? pair_nth : reference_nth); - const NSUInteger head_groups = use_shared_coeff ? - (NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head; - id enc = ds4_gpu_compute_encoder(cb); - id pipeline = use_affine_position ? - g_rope_tail_inplace_pair_affine_pipeline : - (use_shared_coeff ? - g_rope_tail_inplace_pair_shared4_pipeline : - (use_inplace_pair ? - g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline)); - [enc setComputePipelineState:pipeline]; - if (use_affine_position) { - [enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0]; - } else { - [enc setBytes:args length:sizeof(*args) atIndex:0]; - } - [enc setBuffer:xbuf offset:xoff atIndex:1]; - if (!use_affine_position) { - if (posbuf) { - [enc setBuffer:posbuf offset:0 atIndex:2]; - } else { - [enc setBytes:pos length:pos_bytes atIndex:2]; - } - [enc setBuffer:xbuf offset:xoff atIndex:3]; - } - [enc setBuffer:xbuf offset:xoff atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1) - threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (pos && pos != pos_stack) free(pos); - return 1; -} - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - int64_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - int32_t n_rot; -} ds4_gpu_dsv4_fp8_kv_quantize_args; - -typedef struct { - int32_t head_dim; - int32_t n_rot; - int32_t raw_row; -} ds4_gpu_dsv4_kv_fp8_store_args; - -typedef struct { - uint32_t n_rows; - uint32_t head_dim; - uint64_t row_stride; -} ds4_gpu_dsv4_indexer_qat_args; - -typedef struct { - uint32_t width; -} ds4_gpu_dsv4_ratio4_shift_args; - -typedef struct { - uint32_t head_dim; - uint32_t n_comp; - uint32_t replay; - uint32_t n_threads; -} ds4_gpu_dsv4_compressor_pack_ratio4_args; - -typedef struct { - int64_t n_rows; - uint32_t head_dim; - uint32_t n_comp; - uint32_t replay; - uint32_t pad; -} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args; - -typedef struct { - uint32_t width; - uint32_t ratio; - uint32_t pos; - uint32_t ape_type; -} ds4_gpu_dsv4_compressor_store_one_args; - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; -} ds4_gpu_dsv4_softmax_pool_args; - -typedef struct { - uint32_t width; - uint32_t ratio; - uint32_t pos0; - uint32_t n_tokens; -} ds4_gpu_dsv4_compressor_score_ape_args; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - int32_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - int32_t top_k; -} ds4_gpu_kargs_argsort; - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - int64_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne0; - int32_t ne1; - int32_t ne2; - int32_t ne3; - int32_t top_k; - int32_t len; -} ds4_gpu_kargs_argsort_merge; - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - int64_t ne03; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int64_t ne0; - int64_t ne1; - int64_t ne2; - int64_t ne3; - uint64_t nb0; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; -} ds4_gpu_kargs_sum_rows; - -typedef struct { - int32_t ne00; - int32_t ne01; - int32_t ne02; - uint64_t nb01; - uint64_t nb02; - uint64_t nb03; - int32_t ne11; - int32_t ne12; - int32_t ne13; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - uint64_t nb1; - uint64_t nb2; - uint64_t nb3; - float scale; - float max_bias; - float m0; - float m1; - int32_t n_head_log2; -} ds4_gpu_softmax_args; - -typedef struct { - int64_t ne00; - int64_t ne01; - uint64_t nb00; - uint64_t nb01; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; -} ds4_gpu_dsv4_topk_mask_args; - -typedef struct { - int64_t ne00; - int64_t ne01; - int64_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - int64_t ne10; - int64_t ne11; - uint64_t nb10; - uint64_t nb11; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; - float scale; -} ds4_gpu_dsv4_indexer_weighted_sum_args; - -typedef struct { - uint32_t has_bias; - uint32_t hash_mode; - uint32_t use_token_buffer; - uint32_t token; - uint32_t hash_rows; -} ds4_gpu_dsv4_router_select_one_args; - -typedef struct { - uint32_t n_expert; - uint32_t n_expert_used; - float expert_weight_scale; - uint32_t pad0; -} ds4_gpu_glm_router_select_one_args; - -typedef struct { - uint32_t n_tokens; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - float eps; -} ds4_gpu_glm_kv_lora_rms_norm_args; - -typedef struct { - uint32_t n_tokens; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t n_head; - uint32_t row_bytes; - uint32_t weight_type; - uint32_t pad1; - uint32_t pad2; -} ds4_gpu_glm_k_b_project_args; - -typedef struct { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t qk_rope; - uint32_t cache_f16; - uint32_t pad1; -} ds4_gpu_glm_store_compact_kv_args; - -typedef struct { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t q_n; - uint32_t q_n4; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t kv_lora_n4; - uint32_t qk_rope; - uint32_t cache_f16; - float eps; - uint32_t pad0; -} ds4_gpu_glm_qkv_norm_store_compact_kv_args; - -typedef struct { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t head_dim; - uint32_t rot_dim; - uint32_t n_ctx_orig; - uint32_t cache_f16; - uint32_t pad0; - float eps; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - float pad1; -} ds4_gpu_glm_store_indexer_k_args; - -typedef struct { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t n_head; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_rope; - uint32_t value_dim; - uint32_t n_ctx_orig; - uint32_t cache_f16; - uint32_t pad0; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -} ds4_gpu_glm_build_kv_cache_args; - -typedef struct { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_len; - uint32_t cache_cap; - uint32_t n_head; - uint32_t qk_dim; - uint32_t value_dim; - uint32_t pad0; - uint32_t cache_f16; - uint32_t pad1; - uint32_t pad2; - float scale; -} ds4_gpu_glm_attention_full_args; - -typedef struct { - uint32_t n_selected; -} ds4_gpu_glm_fill_selected_range_args; - -typedef struct { - uint32_t n_tokens; - uint32_t pos0; - uint32_t n_selected; - uint32_t pad_row; -} ds4_gpu_glm_fill_selected_range_batch_args; - -typedef struct { - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t rot_dim; - uint32_t rot_offset; - uint32_t pos0; - uint32_t n_ctx_orig; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -} ds4_gpu_glm_rope_tail_args; - -typedef struct { - uint32_t n_rows; - uint32_t n_head; - uint32_t head_dim; - uint32_t cache_f16; - float scale; -} ds4_gpu_glm_indexer_score_one_args; - -typedef struct { - uint32_t n_rows; - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t pos0; - uint32_t cache_f16; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t weights_token_stride; - uint64_t score_token_stride; - float scale; -} ds4_gpu_glm_indexer_scores_batch_args; - -typedef struct { - uint32_t n_head; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_dim; - uint32_t row_bytes; - uint32_t weight_type; - uint32_t pad1; - uint32_t pad2; -} ds4_gpu_glm_qk_lowrank_args; - -typedef struct { - uint32_t n_tokens; - uint32_t n_head; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_dim; - uint32_t row_bytes; - uint32_t weight_type; - uint32_t head_base; -} ds4_gpu_glm_qk_lowrank_batch_args; - -typedef struct { - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t value_type; -} ds4_gpu_glm_attention_indexed_decode_args; - -typedef struct { - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - uint32_t block_rows; - uint32_t n_blocks; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t value_type; -} ds4_gpu_glm_attention_indexed_decode_split_args; - -typedef struct { - uint32_t n_tokens; - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - uint32_t value_type; - uint32_t pos0; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t head_base; -} ds4_gpu_glm_attention_indexed_batch_args; - -typedef struct { - uint32_t in_dim; - uint32_t mid_dim; - uint32_t out_dim; - uint32_t n_total_expert; - uint32_t n_expert_used; - uint32_t n_tokens; - uint32_t mid_token_stride; - uint32_t down_type; - int32_t tp_rank; - int32_t tp_world; - int32_t tp_expert_base; - uint64_t gate_expert_bytes; - uint64_t gate_row_bytes; - uint64_t up_expert_bytes; - uint64_t up_row_bytes; - uint64_t down_expert_bytes; - uint64_t down_row_bytes; -} ds4_gpu_glm_routed_moe_args; - -typedef struct { - uint32_t n_tokens; - uint32_t n_head; - uint32_t n_raw; - uint32_t raw_cap; - uint32_t raw_start; - uint32_t n_comp; - uint32_t top_k; - uint32_t pos0; - uint32_t window; - uint32_t ratio; - uint32_t comp_kv_f16; - uint32_t pad0; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t raw_row_stride; - uint64_t comp_row_stride; - uint64_t topk_token_stride; - uint64_t dst_token_stride; - uint64_t dst_head_stride; - float scale; -} ds4_gpu_dsv4_indexed_attention_args; - -typedef struct { - uint32_t n_comp; - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t pos0; - uint32_t ratio; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t weights_token_stride; - uint64_t index_row_stride; - uint64_t score_token_stride; - float scale; -} ds4_gpu_dsv4_indexer_scores_fused_args; - -typedef struct { - uint32_t width; - uint32_t rows; - uint64_t gate_row_stride; - uint64_t up_row_stride; - uint64_t mid_row_stride; - uint64_t weight_stride; - uint32_t write_clamped; - float clamp_value; -} ds4_gpu_dsv4_moe_swiglu_weight_args; - -typedef struct { - uint32_t expert_base; - uint32_t expert_count; - uint32_t accumulate; - uint32_t pad0; -} ds4_gpu_moe_expert_group_args; - -typedef struct { - uint64_t expert_bytes; - uint32_t group_size; - uint32_t n_slots; -} ds4_gpu_q4_gather_slots6_args; - -typedef struct { - uint32_t width; - uint32_t tokens; - uint64_t src_token_stride; - uint64_t dst_token_stride; -} ds4_gpu_dsv4_moe_sum6_args; - -/* Compile the single in-repo Metal source and create the pipelines that every - * session uses. Shape-dependent kernels with function constants are built - * lazily by the small ds4_gpu_get_* caches, so startup stays predictable - * while long-context prefill and decode can still pick specialized variants. */ -int ds4_gpu_init(void) { - if (g_initialized) return 1; - - @autoreleasepool { - g_device = MTLCreateSystemDefaultDevice(); - if (!g_device) { - fprintf(stderr, "ds4: Metal device not available\n"); - return 0; - } - ds4_gpu_print_device_summary(); - ds4_gpu_detect_metal4_features(); - - g_queue = [g_device newCommandQueue]; - if (!g_queue) { - fprintf(stderr, "ds4: failed to create Metal command queue\n"); - g_device = nil; - return 0; - } - g_model_buffer_cache = [NSMutableDictionary dictionary]; - g_model_buffer_cache_bytes = 0; - g_model_buffer_cache_evictions = 0; - g_model_buffer_cache_over_limit = 0; - g_q4_expert_table_cache = [NSMutableDictionary dictionary]; - g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary]; - g_pipeline_cache = [NSMutableDictionary dictionary]; - g_transient_buffers = [NSMutableArray array]; - g_pending_cbs = [NSMutableArray array]; - if (!g_model_buffer_cache || !g_q4_expert_table_cache || - !g_q4_expert_layer_residency_cache || - !g_pipeline_cache || !g_transient_buffers || !g_pending_cbs) { - fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n"); - g_pending_cbs = nil; - g_transient_buffers = nil; - g_pipeline_cache = nil; - g_q4_expert_layer_residency_cache = nil; - g_q4_expert_table_cache = nil; - g_model_buffer_cache = nil; - g_queue = nil; - g_device = nil; - return 0; - } - - NSError *error = nil; - NSString *source = ds4_gpu_full_source(); - if (!source) { - g_queue = nil; - g_device = nil; - return 0; - } - MTLCompileOptions *options = [MTLCompileOptions new]; - NSMutableDictionary *macros = [NSMutableDictionary new]; - if (g_metal4_tensor_api_enabled) { - macros[@"DS4_METAL_HAS_TENSOR"] = @"1"; - fprintf(stderr, "ds4: Metal 4 tensor API enabled for Tensor kernels\n"); - } - - const int drift_hc_stable = ds4_gpu_env_bool("DS4_METAL_HC_STABLE") != 0; // default ON - const int drift_norm_unify = ds4_gpu_env_bool("DS4_METAL_NORM_RSQRT_DISABLE") != 0; // default ON - const int drift_kv_raw_f32 = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32") > 0; // default OFF - const int drift_rope_exp2_log2 = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2") > 0; // default OFF - const int drift_math_safe = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE") > 0; // default OFF - - if (drift_math_safe) { - // MTLCompileOptions.fastMathEnabled defaults to YES and Apple's - // headers explicitly say this "may violate the IEEE 754 standard". - // Different fast-math optimizations get applied across the - // matmul2d cooperative-tensor path and the legacy - // simdgroup_multiply_accumulate path on M5, amplifying the - // mismatch. MTLMathModeSafe pins the entire library to strict - // IEEE-754 semantics. Diagnostic-only: useful to localize drift - // sources but not to ship as a default. - if (@available(macOS 15.0, *)) { - options.mathMode = MTLMathModeSafe; - fprintf(stderr, "ds4: Metal shader library math mode = safe (strict IEEE-754) by DS4_METAL_MATH_SAFE\n"); - } else { -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wdeprecated-declarations" - options.fastMathEnabled = NO; -#pragma clang diagnostic pop - fprintf(stderr, "ds4: Metal shader library fast-math disabled by DS4_METAL_MATH_SAFE (pre-macOS 15)\n"); - } - } - - if (drift_hc_stable) macros[@"DS4_METAL_HC_STABLE"] = @"1"; - if (drift_norm_unify) macros[@"DS4_METAL_NORM_RSQRT_DISABLE"] = @"1"; - if (drift_kv_raw_f32) macros[@"DS4_METAL_KV_RAW_F32"] = @"1"; - if (drift_rope_exp2_log2) macros[@"DS4_METAL_ROPE_EXP2_LOG2"] = @"1"; - fprintf(stderr, - "ds4: drift-patch flags hc_stable=%s norm_unify=%s kv_raw_f32=%s rope_exp2_log2=%s math_safe=%s tensor_matmul=%s\n", - drift_hc_stable ? "on" : "off", - drift_norm_unify ? "on" : "off", - drift_kv_raw_f32 ? "on" : "off", - drift_rope_exp2_log2 ? "on" : "off", - drift_math_safe ? "on" : "off", - g_metal4_tensor_api_enabled ? "on" : "off"); - options.preprocessorMacros = macros; - id library = [g_device newLibraryWithSource:source options:options error:&error]; - if (!library) { - fprintf(stderr, "ds4: Metal shader compilation failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_library = library; - - id fn = [library newFunctionWithName:@"kernel_get_rows_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_get_rows_f16"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_f16 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_f16 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_get_rows_i32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_i32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_i32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_i32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_get_rows_q8_0_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_q8_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_q8_0_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_get_rows_q4_0_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_q4_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_q4_0_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_get_rows_q4_K_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_get_rows_q4_K_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_get_rows_q4_K_pipeline) { - fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_repeat_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_repeat_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_repeat_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_repeat_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_repeat_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_set_rows_f32_i32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_set_rows_f32_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_set_rows_f32_i32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_concat"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_concat function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_concat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_concat_pipeline) { - fprintf(stderr, "ds4: Metal kernel_concat pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_f32_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_f32_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_f32_f16"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_f32_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_contig_f32_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_f16_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_f16_f16"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_f16_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_contig_f16_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_cpy_contig_f16_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"]; - if (fn) { - g_flash_kv_stage_f16_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_flash_kv_stage_f16_pipeline) { - fprintf(stderr, - "ds4: optional Metal gathered KV staging pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, - "ds4: optional Metal gathered KV staging kernel unavailable\n"); - } - - fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_fp8_kv_quantize_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_fp8_kv_quantize_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_indexer_hadamard_fp4_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_indexer_qat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_indexer_qat_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_kv_fp8_store_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_kv_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_kv_fp8_store_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_ratio4_shift_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_ratio4_shift_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_swiglu_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_swiglu_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_swiglu_flat_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - g_swiglu_flat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_swiglu_flat_pipeline) { - fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum6_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_moe_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum8_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - g_moe_sum8_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_sum8_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *bin_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t bin_op = 0; - int16_t bin_f = 1; - bool bin_rb = false; - bool bin_cb = false; - [bin_constants setConstantValue:&bin_op type:MTLDataTypeShort atIndex:1300]; - [bin_constants setConstantValue:&bin_f type:MTLDataTypeShort atIndex:1301]; - [bin_constants setConstantValue:&bin_rb type:MTLDataTypeBool atIndex:1302]; - [bin_constants setConstantValue:&bin_cb type:MTLDataTypeBool atIndex:1303]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" - constantValues:bin_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - g_add_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_add_pipeline) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_add2_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_add2_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - g_add2_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_add2_pipeline) { - fprintf(stderr, "ds4: Metal kernel_add2_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_add3_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_add3_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - g_add3_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_add3_pipeline) { - fprintf(stderr, "ds4: Metal kernel_add3_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *bin_mul_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t bin_mul_plain_op = 2; - int16_t bin_mul_plain_f = 1; - bool bin_mul_plain_rb = false; - bool bin_mul_plain_cb = false; - [bin_mul_constants setConstantValue:&bin_mul_plain_op type:MTLDataTypeShort atIndex:1300]; - [bin_mul_constants setConstantValue:&bin_mul_plain_f type:MTLDataTypeShort atIndex:1301]; - [bin_mul_constants setConstantValue:&bin_mul_plain_rb type:MTLDataTypeBool atIndex:1302]; - [bin_mul_constants setConstantValue:&bin_mul_plain_cb type:MTLDataTypeBool atIndex:1303]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" - constantValues:bin_mul_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - g_mul_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_mul_pipeline) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *bin_mul_scalar_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t bin_mul_op = 2; - int16_t bin_mul_f = 1; - bool bin_mul_rb = false; - bool bin_mul_cb = true; - [bin_mul_scalar_constants setConstantValue:&bin_mul_op type:MTLDataTypeShort atIndex:1300]; - [bin_mul_scalar_constants setConstantValue:&bin_mul_f type:MTLDataTypeShort atIndex:1301]; - [bin_mul_scalar_constants setConstantValue:&bin_mul_rb type:MTLDataTypeBool atIndex:1302]; - [bin_mul_scalar_constants setConstantValue:&bin_mul_cb type:MTLDataTypeBool atIndex:1303]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" - constantValues:bin_mul_scalar_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - g_bin_mul_scalar_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_bin_mul_scalar_pipeline) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *bin_div_row_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t bin_div_op = 3; - int16_t bin_div_f = 1; - bool bin_div_rb = false; - bool bin_div_cb = true; - [bin_div_row_constants setConstantValue:&bin_div_op type:MTLDataTypeShort atIndex:1300]; - [bin_div_row_constants setConstantValue:&bin_div_f type:MTLDataTypeShort atIndex:1301]; - [bin_div_row_constants setConstantValue:&bin_div_rb type:MTLDataTypeBool atIndex:1302]; - [bin_div_row_constants setConstantValue:&bin_div_cb type:MTLDataTypeBool atIndex:1303]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" - constantValues:bin_div_row_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - g_bin_div_row_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_bin_div_row_pipeline) { - fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_rms_norm_mul_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rms_norm_pipeline) { - fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_rms_norm_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_rms_norm_plain_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rms_norm_plain_pipeline) { - fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_add_rms_norm_mul_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_add_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_add_rms_norm_pipeline) { - fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_qkv_rms_norm_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_qkv_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_qkv_rms_norm_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t moe_mv_id_nsg = 2; - [moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_iq2_xxs_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_iq2_xxs_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q2_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q2_k_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q2_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_iq2_xxs_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q2_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_slots6_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_slots6_q2_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_addr_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_iq2_xxs_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_addr_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_q2_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_masked_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_stream_expert_cache_validate"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_stream_expert_cache_validate_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_stream_expert_cache_validate_pipeline) { - fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q4_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q4_k_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q4_k_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q4_k_pair_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_id_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_id_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group6_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group8_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group8_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_id_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group24_q4_k_id_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group24_q4_k_id_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_group24_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_group24_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_slots6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_slots6_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_q4_gather_slots6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_q4_gather_slots6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_table_q4_pair_gate_encoder = [fn newArgumentEncoderWithBufferIndex:2]; - g_moe_table_q4_pair_up_encoder = [fn newArgumentEncoderWithBufferIndex:3]; - if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder) { - fprintf(stderr, "ds4: Metal Q4 expert-table pair argument encoder creation failed\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_table_q4_sum_down_encoder = [fn newArgumentEncoderWithBufferIndex:1]; - if (!g_moe_table_q4_sum_down_encoder) { - fprintf(stderr, "ds4: Metal Q4 expert-table down argument encoder creation failed\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_moe_mul_mv_table_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_table_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_pair_swiglu_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (fn) { - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 function unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_sum6_f32" - constantValues:moe_mv_id_constants - error:&error]; - if (fn) { - g_moe_mul_mv_addr_q4_k_sum6_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_moe_mul_mv_addr_q4_k_sum6_pipeline) { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 function unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - - fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_rope_tail_batch_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rope_tail_batch_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_rope_tail_inplace_pair_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rope_tail_inplace_pair_pipeline) { - fprintf(stderr, - "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"]; - if (!fn) { - fprintf(stderr, - "ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n"); - } else { - g_rope_tail_inplace_pair_shared4_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rope_tail_inplace_pair_shared4_pipeline) { - fprintf(stderr, - "ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n", - [[error localizedDescription] UTF8String]); - } - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"]; - if (fn) { - g_rope_tail_inplace_pair_affine_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_rope_tail_inplace_pair_affine_pipeline) { - fprintf(stderr, - "ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, - "ds4: optional Metal affine-position RoPE pair kernel unavailable\n"); - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_softmax_pool_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_softmax_pool_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_soft_max_f32"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_soft_max_f32 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_soft_max_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_soft_max_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_soft_max_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_soft_max_f32_4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_soft_max_f32_4_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_soft_max_f32_4_pipeline) { - fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_argsort_f32_i32_desc"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_argsort_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_argsort_f32_i32_desc_pipeline) { - fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_argsort_merge_f32_i32_desc"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_argsort_merge_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_argsort_merge_f32_i32_desc_pipeline) { - fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *sum_rows_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t sum_rows_op = 10; - [sum_rows_constants setConstantValue:&sum_rows_op type:MTLDataTypeShort atIndex:1400]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_sum_rows_f32_f32" - constantValues:sum_rows_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_sum_rows_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_sum_rows_f32_f32_pipeline) { - fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_topk_mask_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_topk_mask_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask_scatter"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_topk_mask_scatter_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_topk_mask_scatter_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_indexer_weighted_sum"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_dsv4_indexer_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_dsv4_indexer_weighted_sum_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_sinkhorn"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_hc_split_sinkhorn_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_split_sinkhorn_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_hc_split_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_split_weighted_sum_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum_norm4"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_hc_split_weighted_sum_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_split_weighted_sum_norm_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_hc_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_weighted_sum_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"]; - if (fn) { - g_hc_weighted_sum_norm_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_weighted_sum_norm_pipeline) { - fprintf(stderr, - "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, - "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n"); - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"]; - if (fn) { - g_output_hc_weights4_pipeline = - [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_output_hc_weights4_pipeline) { - fprintf(stderr, - "ds4: optional Metal output HC weights4 pipeline unavailable: %s\n", - [[error localizedDescription] UTF8String]); - } - } else { - fprintf(stderr, - "ds4: optional Metal output HC weights4 kernel unavailable\n"); - } - - MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_sigmoid_op = 102; - bool unary_cnt = false; - [unary_sigmoid_constants setConstantValue:&unary_sigmoid_op type:MTLDataTypeShort atIndex:1200]; - [unary_sigmoid_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_sigmoid_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_sigmoid_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_sigmoid_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_silu_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_silu_op = 106; - [unary_silu_constants setConstantValue:&unary_silu_op type:MTLDataTypeShort atIndex:1200]; - [unary_silu_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_silu_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_silu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_silu_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_softplus_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_softplus_op = 115; - [unary_softplus_constants setConstantValue:&unary_softplus_op type:MTLDataTypeShort atIndex:1200]; - [unary_softplus_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_softplus_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_softplus_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_softplus_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_sqrt_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_sqrt_op = 14; - [unary_sqrt_constants setConstantValue:&unary_sqrt_op type:MTLDataTypeShort atIndex:1200]; - [unary_sqrt_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_sqrt_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_sqrt_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_sqrt_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_clamp_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_clamp_op = 12; - [unary_clamp_constants setConstantValue:&unary_clamp_op type:MTLDataTypeShort atIndex:1200]; - [unary_clamp_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32" - constantValues:unary_clamp_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_clamp_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_clamp_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_scale_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_scale_op = 10; - [unary_scale_constants setConstantValue:&unary_scale_op type:MTLDataTypeShort atIndex:1200]; - [unary_scale_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_scale_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_scale_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_scale_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - MTLFunctionConstantValues *unary_fill_constants = [[MTLFunctionConstantValues alloc] init]; - int16_t unary_fill_op = 11; - [unary_fill_constants setConstantValue:&unary_fill_op type:MTLDataTypeShort atIndex:1200]; - [unary_fill_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" - constantValues:unary_fill_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_fill_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_fill_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - error = nil; - fn = [library newFunctionWithName:@"kernel_unary_f16_f16" - constantValues:unary_fill_constants - error:&error]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill function not found: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - g_unary_fill_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_unary_fill_f16_pipeline) { - fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - fn = [library newFunctionWithName:@"kernel_dsv4_hc_expand"]; - if (!fn) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand function not found\n"); - g_queue = nil; - g_device = nil; - return 0; - } - g_hc_expand_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_hc_expand_pipeline) { - fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand pipeline failed: %s\n", - [[error localizedDescription] UTF8String]); - g_queue = nil; - g_device = nil; - return 0; - } - - g_dsv4_indexer_score_one_direct_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct"); - g_dsv4_compressor_store_one_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one"); - g_dsv4_compressor_pack_ratio4_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4"); - g_dsv4_softmax_pool_ratio4_direct_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct"); - g_rms_norm_scale_pipeline = - ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4"); - g_dsv4_sort_i32_rows_asc_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc"); - g_dsv4_indexed_attention_heads8_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8"); - g_dsv4_indexed_attention_heads8_rb16_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16"); - g_dsv4_softplus_sqrt_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4"); - g_dsv4_router_finalize_one_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one"); - g_dsv4_router_finalize_one_simd_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd"); - g_dsv4_router_finalize_weights_one_simd_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd"); - g_dsv4_router_weights_one_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one"); - g_glm_router_select_one_pipeline = - ds4_gpu_get_pipeline("kernel_glm_router_select_one"); - g_glm_kv_lora_rms_norm_pipeline = - ds4_gpu_get_pipeline("kernel_glm_kv_lora_rms_norm"); - g_glm_k_b_project_pipeline = - ds4_gpu_get_pipeline("kernel_glm_k_b_project_q8_0"); - g_glm_store_compact_kv_pipeline = - ds4_gpu_get_pipeline("kernel_glm_store_compact_kv"); - g_glm_qkv_norm_store_compact_kv_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv"); - g_glm_store_indexer_k_pipeline = - ds4_gpu_get_pipeline("kernel_glm_store_indexer_k"); - g_glm_build_kv_cache_pipeline = - ds4_gpu_get_pipeline("kernel_glm_build_kv_cache"); - g_glm_build_kv_cache_decode_group4_pipeline = - ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_decode_group4"); - g_glm_build_kv_cache_flash_pipeline = - ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_flash"); - g_glm_attention_full_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_full"); - g_glm_fill_selected_range_pipeline = - ds4_gpu_get_pipeline("kernel_glm_fill_selected_range"); - g_glm_fill_selected_range_batch_pipeline = - ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch"); - g_glm_indexer_rope_tail_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32"); - g_glm_indexer_score_one_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_score_one"); - g_glm_indexer_score_one_direct_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct"); - g_glm_indexer_scores_batch_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch"); - g_glm_indexer_scores_tiled_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled"); - g_glm_indexer_scores_tiled_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled_f32"); - g_glm_qk_lowrank_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0"); - g_glm_qk_lowrank_glm52_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52"); - g_glm_qk_lowrank_glm52_sg_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52_sg"); - g_glm_qk_lowrank_batch_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch"); - g_glm_qk_lowrank_batch_glm52_t4_pipeline = - ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4"); - g_glm_value_project_q8_0_pipeline = - ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0"); - g_glm_value_project_q8_0_batch_heads_pipeline = - ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads"); - g_glm_value_project_q8_0_batch_heads_mma_pipeline = - ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads_mma"); - g_glm_attention_indexed_decode_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode"); - g_glm_attention_indexed_decode_split_group8_partial_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial"); - g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads"); - g_glm_attention_indexed_decode_split_group8_reduce_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce"); - g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16"); - g_glm_attention_indexed_batch_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch"); - g_glm_attention_indexed_batch_group2_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group2"); - g_glm_attention_indexed_batch_q2_group4_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_q2_group4"); - g_glm_attention_indexed_batch_group8_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group8"); - g_glm_attention_indexed_batch_lora_group8_vec_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec"); - g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); - g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); - g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); - g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = - ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads"); - g_glm_q4_k_pair_swiglu_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu_f32"); - g_glm_q4_k_pair_swiglu2_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_f32"); - g_glm_q4_k_pair_swiglu4_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu4_f32"); - g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_f32"); - g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_row_f32"); - g_glm_q2_k_pair_swiglu_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q2_K_pair_swiglu_f32"); - g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32"); - g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32_masked"); - g_glm_q4_k_addr_pair_swiglu_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32"); - g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32_masked"); - g_glm_q2_k_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q2_K_down_f32"); - g_glm_q4_k_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_down_simd_f32"); - g_glm_q2_k_addr_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_down_f32"); - g_glm_q4_k_addr_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_down_simd_f32"); - g_glm_q5_k_pair_swiglu_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_f32"); - g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_f32"); - g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_row_f32"); - g_glm_q5_k_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q5_K_down_f32"); - g_glm_q6_k_down_f32_pipeline = - ds4_gpu_get_pipeline("kernel_glm_q6_K_down_f32"); - g_dsv4_router_weights_batch_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch"); - g_dsv4_hc_expand4_pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4"); - if (!g_dsv4_indexer_score_one_direct_pipeline || - !g_dsv4_compressor_store_one_pipeline || - !g_dsv4_sort_i32_rows_asc_pipeline || - !g_dsv4_indexed_attention_heads8_pipeline || - !g_dsv4_indexed_attention_heads8_rb16_pipeline || - !g_dsv4_softplus_sqrt_pipeline || - !g_dsv4_router_finalize_one_pipeline || - !g_dsv4_router_weights_one_pipeline || - !g_glm_router_select_one_pipeline || - !g_glm_kv_lora_rms_norm_pipeline || - !g_glm_k_b_project_pipeline || - !g_glm_store_compact_kv_pipeline || - !g_glm_qkv_norm_store_compact_kv_pipeline || - !g_glm_store_indexer_k_pipeline || - !g_glm_build_kv_cache_pipeline || - !g_glm_build_kv_cache_decode_group4_pipeline || - !g_glm_build_kv_cache_flash_pipeline || - !g_glm_attention_full_pipeline || - !g_glm_fill_selected_range_pipeline || - !g_glm_fill_selected_range_batch_pipeline || - !g_glm_indexer_rope_tail_pipeline || - !g_glm_indexer_score_one_pipeline || - !g_glm_indexer_score_one_direct_pipeline || - !g_glm_indexer_scores_batch_pipeline || - !g_glm_indexer_scores_tiled_pipeline || - !g_glm_indexer_scores_tiled_f32_pipeline || - !g_glm_qk_lowrank_pipeline || - !g_glm_qk_lowrank_glm52_pipeline || - !g_glm_qk_lowrank_glm52_sg_pipeline || - !g_glm_qk_lowrank_batch_pipeline || - !g_glm_qk_lowrank_batch_glm52_t4_pipeline || - !g_glm_value_project_q8_0_pipeline || - !g_glm_value_project_q8_0_batch_heads_pipeline || - !g_glm_value_project_q8_0_batch_heads_mma_pipeline || - !g_glm_attention_indexed_decode_pipeline || - !g_glm_attention_indexed_decode_split_group8_partial_pipeline || - !g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline || - !g_glm_attention_indexed_decode_split_group8_reduce_pipeline || - !g_glm_attention_indexed_decode_split_group8_reduce16_pipeline || - !g_glm_attention_indexed_batch_pipeline || - !g_glm_attention_indexed_batch_group2_pipeline || - !g_glm_attention_indexed_batch_q2_group4_pipeline || - !g_glm_attention_indexed_batch_group8_pipeline || - !g_glm_attention_indexed_batch_lora_group8_vec_pipeline || - !g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline || - !g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline || - !g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline || - !g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline || - !g_glm_q4_k_pair_swiglu_f32_pipeline || - !g_glm_q4_k_pair_swiglu2_f32_pipeline || - !g_glm_q4_k_pair_swiglu4_f32_pipeline || - !g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline || - !g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline || - !g_glm_q2_k_pair_swiglu_f32_pipeline || - !g_glm_q2_k_addr_pair_swiglu2_f32_pipeline || - !g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline || - !g_glm_q4_k_addr_pair_swiglu_f32_pipeline || - !g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline || - !g_glm_q2_k_down_f32_pipeline || - !g_glm_q4_k_down_f32_pipeline || - !g_glm_q2_k_addr_down_f32_pipeline || - !g_glm_q4_k_addr_down_f32_pipeline || - !g_glm_q5_k_pair_swiglu_f32_pipeline || - !g_glm_q5_k_pair_swiglu_mapped_f32_pipeline || - !g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline || - !g_glm_q5_k_down_f32_pipeline || - !g_glm_q6_k_down_f32_pipeline || - !g_dsv4_hc_expand4_pipeline) { - g_queue = nil; - g_device = nil; - return 0; - } - - g_initialized = 1; - } - - return 1; -} - -ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { - if (!g_initialized && !ds4_gpu_init()) return NULL; - if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; - - @autoreleasepool { - DS4MetalTensor *tensor = [DS4MetalTensor new]; - tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes - options:MTLResourceStorageModeShared]; - if (!tensor.buffer) { - return NULL; - } - tensor.offset = 0; - tensor.bytes = bytes; - tensor.owner = 1; - uint64_t live_snap = 0; - uint64_t peak_snap = 0; - pthread_mutex_lock(&g_tensor_mu); - const int tracked = ds4_gpu_tensor_track_alloc_locked( - (__bridge const void *)tensor, - bytes, - &live_snap, - &peak_snap); - pthread_mutex_unlock(&g_tensor_mu); - if (!tracked) { - fprintf(stderr, "ds4: failed to track Metal tensor allocation\n"); - tensor.buffer = nil; - return NULL; - } - if (ds4_gpu_trace_allocs()) { - fprintf(stderr, - "ds4: Metal tensor alloc %.3f MiB live %.3f MiB peak %.3f MiB\n", - (double)bytes / (1024.0 * 1024.0), - (double)live_snap / (1024.0 * 1024.0), - (double)peak_snap / (1024.0 * 1024.0)); - } - return (__bridge_retained ds4_gpu_tensor *)tensor; - } -} - -ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { - return ds4_gpu_tensor_alloc(bytes); -} - -int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) { - (void)kv_cache_bytes; - (void)context_bytes; - return 0; -} - -ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) { - if (!base) return NULL; - const DS4MetalTensor *base_obj = ds4_gpu_tensor_const_obj(base); - if (offset > base_obj.bytes || bytes > base_obj.bytes - offset) return NULL; - if (base_obj.offset > UINT64_MAX - offset) return NULL; - const uint64_t absolute_offset = base_obj.offset + offset; - if (absolute_offset > (uint64_t)NSUIntegerMax) return NULL; - - @autoreleasepool { - DS4MetalTensor *view = [DS4MetalTensor new]; - view.buffer = base_obj.buffer; - view.offset = absolute_offset; - view.bytes = bytes; - view.owner = 0; - pthread_mutex_lock(&g_tensor_mu); - const int tracked = ds4_gpu_tensor_track_view_locked((__bridge const void *)view); - pthread_mutex_unlock(&g_tensor_mu); - if (!tracked) { - fprintf(stderr, "ds4: failed to track Metal tensor view\n"); - view.buffer = nil; - return NULL; - } - return (__bridge_retained ds4_gpu_tensor *)view; - } -} - -void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) { - if (!tensor) return; - @autoreleasepool { - uint8_t owner = 0; - uint64_t bytes = 0; - uint64_t live_snap = 0; - uint64_t peak_snap = 0; - if (!ds4_gpu_tensor_prepare_free(tensor, - &owner, - &bytes, - &live_snap, - &peak_snap)) { - return; - } - DS4MetalTensor *obj = (__bridge_transfer DS4MetalTensor *)tensor; - if (owner) { - if (ds4_gpu_trace_allocs()) { - fprintf(stderr, - "ds4: Metal tensor free %.3f MiB live %.3f MiB peak %.3f MiB\n", - (double)bytes / (1024.0 * 1024.0), - (double)live_snap / (1024.0 * 1024.0), - (double)peak_snap / (1024.0 * 1024.0)); - } - } - obj.buffer = nil; - obj.offset = 0; - obj.bytes = 0; - obj.owner = 0; - } -} - -uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) { - if (!tensor) return 0; - const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); - return obj.bytes; -} - -void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) { - if (!tensor) return NULL; - DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); - return (uint8_t *)[obj.buffer contents] + obj.offset; -} - -int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) { - if (!tensor || count > ds4_gpu_tensor_bytes(tensor) / sizeof(float)) return 0; - float *p = ds4_gpu_tensor_contents(tensor); - if (!p && count != 0) return 0; - for (uint64_t i = 0; i < count; i++) p[i] = value; - return 1; -} - -int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) { - if (!tensor || (!data && bytes != 0)) return 0; - DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); - if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; - if (bytes != 0) { - memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset, data, (size_t)bytes); - } - return 1; -} - -int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) { - if (!tensor || (!data && bytes != 0)) return 0; - const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); - if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; - if (bytes != 0) { - memcpy(data, (const uint8_t *)[obj.buffer contents] + obj.offset + offset, (size_t)bytes); - } - return 1; -} - -int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, - const ds4_gpu_tensor *src, uint64_t src_offset, - uint64_t bytes) { - if (!dst || !src) return 0; - if (!g_initialized && !ds4_gpu_init()) return 0; - DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); - const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); - if (dst_offset > d.bytes || bytes > d.bytes - dst_offset) return 0; - if (src_offset > s.bytes || bytes > s.bytes - src_offset) return 0; - if (bytes == 0) return 1; - if (!g_batch_cb) return 0; - - ds4_gpu_close_batch_encoder(); - g_batch_has_work = YES; - id blit = [g_batch_cb blitCommandEncoder]; - if (!blit) return 0; - [blit copyFromBuffer:s.buffer - sourceOffset:(NSUInteger)(s.offset + src_offset) - toBuffer:d.buffer - destinationOffset:(NSUInteger)(d.offset + dst_offset) - size:(NSUInteger)bytes]; - [blit endEncoding]; - return 1; -} - -int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, - const ds4_gpu_tensor *src, uint64_t src_offset, - uint64_t count) { - if (!dst || !src) return 0; - if (!g_initialized && !ds4_gpu_init()) return 0; - DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); - const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); - if (count == 0) return 1; - if (count > UINT64_MAX / sizeof(float) || - count > UINT64_MAX / sizeof(uint16_t)) { - return 0; - } - const uint64_t src_bytes = count * sizeof(float); - const uint64_t dst_bytes = count * sizeof(uint16_t); - if (src_offset > s.bytes || src_bytes > s.bytes - src_offset || - dst_offset > d.bytes || dst_bytes > d.bytes - dst_offset) { - return 0; - } - - @autoreleasepool { - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - uint64_t done = 0; - int ok = 1; - while (done < count && ok) { - uint64_t chunk64 = count - done; - if (chunk64 > UINT32_MAX) chunk64 = UINT32_MAX; - const uint32_t chunk = (uint32_t)chunk64; - ok = ds4_gpu_encode_cpy_f32_f16_1d( - cb, - s.buffer, - (NSUInteger)(s.offset + src_offset + done * sizeof(float)), - d.buffer, - (NSUInteger)(d.offset + dst_offset + done * sizeof(uint16_t)), - chunk); - done += chunk; - } - if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "tensor f32 to f16 copy"); - return ok; - } -} - -int ds4_gpu_pack_slot_rows_f32_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *slots, - uint32_t n_rows, - uint32_t width, - uint32_t n_slots, - uint32_t slot_cap) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 || - slot_cap == 0 || n_rows > slot_cap) { - return 0; - } - - @autoreleasepool { - id slotsbuf = ds4_gpu_tensor_buffer(slots); - id outbuf = ds4_gpu_tensor_buffer(out); - uint64_t row_bytes = 0; - uint64_t slot_plane_bytes = 0; - uint64_t slots_bytes = 0; - uint64_t out_rows = 0; - uint64_t out_bytes = 0; - if ((uint64_t)width > UINT64_MAX / sizeof(float)) return 0; - row_bytes = (uint64_t)width * sizeof(float); - if ((uint64_t)slot_cap > UINT64_MAX / row_bytes) return 0; - slot_plane_bytes = (uint64_t)slot_cap * row_bytes; - if ((uint64_t)n_slots > UINT64_MAX / slot_plane_bytes) return 0; - slots_bytes = (uint64_t)n_slots * slot_plane_bytes; - if ((uint64_t)n_rows > UINT64_MAX / n_slots) return 0; - out_rows = (uint64_t)n_rows * n_slots; - if (out_rows > UINT64_MAX / row_bytes) return 0; - out_bytes = out_rows * row_bytes; - if (!slotsbuf || !outbuf || - ds4_gpu_tensor_bytes(slots) < slots_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal slot-row pack received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (!ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, - slotsbuf, - ds4_gpu_tensor_offset(slots), - outbuf, - ds4_gpu_tensor_offset(out), - width, - n_slots, - n_rows, - sizeof(float), - slot_plane_bytes, - row_bytes, - row_bytes, - (uint64_t)n_slots * row_bytes)) { - return 0; - } - if (!ds4_gpu_finish_command_buffer(cb, owned, "slot-row pack")) return 0; - } - - return 1; -} - -int ds4_gpu_begin_commands(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (g_batch_cb) return 0; - g_batch_cb = ds4_gpu_new_command_buffer(); - g_batch_has_work = NO; - if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); - return g_batch_cb != nil; -} - -int ds4_gpu_flush_encoder(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_batch_cb) return 0; - ds4_gpu_close_batch_encoder(); - return 1; -} - -int ds4_gpu_flush_commands(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_batch_cb) return 0; - - ds4_gpu_close_batch_encoder(); - id cb = g_batch_cb; - g_batch_cb = nil; - g_batch_has_work = NO; - [cb commit]; - [g_pending_cbs addObject:cb]; - ds4_gpu_stream_expert_cache_note_batch_committed(); - - g_batch_cb = ds4_gpu_new_command_buffer(); - g_batch_has_work = NO; - if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); - if (!g_batch_cb) { - (void)ds4_gpu_wait_pending_command_buffers("command batch"); - [g_transient_buffers removeAllObjects]; - return 0; - } - return 1; -} - -int ds4_gpu_commands_active(void) { - return g_batch_cb != nil; -} - -static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) { - const char *what = label ? label : "streaming expert cache in-flight"; - if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0; - if ([g_pending_cbs count] != 0 && - ds4_gpu_wait_pending_command_buffers(what) == 0) { - return 0; - } - return 1; -} - -int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { - if (!event_value) return 0; - *event_value = 0; - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_batch_cb) return 0; - - if (@available(macOS 12.0, *)) { - if (!g_selected_readback_event) { - g_selected_readback_event = [g_device newSharedEvent]; - if (!g_selected_readback_event) { - fprintf(stderr, "ds4: failed to create Metal shared event for selected-id overlap\n"); - return 0; - } - } - - ds4_gpu_close_batch_encoder(); - const uint64_t value = ++g_selected_readback_event_value; - [g_batch_cb encodeSignalEvent:g_selected_readback_event value:value]; - g_batch_has_work = YES; - *event_value = value; - return 1; - } - - fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); - return 0; -} - -int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_batch_cb || event_value == 0) return 0; - - if (@available(macOS 12.0, *)) { - if (!g_selected_readback_event) return 0; - - ds4_gpu_close_batch_encoder(); - id cb = g_batch_cb; - g_batch_cb = nil; - g_batch_has_work = NO; - [cb commit]; - ds4_gpu_stream_expert_cache_note_batch_committed(); - - const char *what = label ? label : "selected-id overlap"; - const BOOL signaled = - [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; - [g_pending_cbs addObject:cb]; - if (!signaled) { - fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); - (void)ds4_gpu_wait_pending_command_buffers(what); - [g_transient_buffers removeAllObjects]; - return 0; - } - if (cb.status == MTLCommandBufferStatusError) { - fprintf(stderr, "ds4: Metal %s failed: %s\n", - what, - [[cb.error localizedDescription] UTF8String]); - (void)ds4_gpu_wait_pending_command_buffers(what); - [g_transient_buffers removeAllObjects]; - return 0; - } - - g_batch_cb = ds4_gpu_new_command_buffer(); - g_batch_has_work = NO; - if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); - if (!g_batch_cb) { - (void)ds4_gpu_wait_pending_command_buffers(what); - [g_transient_buffers removeAllObjects]; - return 0; - } - return 1; - } - - fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); - return 0; -} - -/* - * Tensor-parallel gates. - * - * A TP gate is a mid-command-stream rendezvous with the peer machine: the - * kernels ahead of the gate leave a partial block output in a slab slot, - * the GPU signals g_tp_gpu_event, and the pre-encoded combine kernel waits - * on g_tp_cpu_event. A dedicated service thread bridges the two: it spins - * until the GPU reaches the gate, runs the transport exchange (RDMA WRITE - * plus flag poll, or a TCP write/read pair — behind the callback), and - * CPU-signals the release. On exchange failure the release is signaled - * anyway so the GPU never deadlocks; the failure latches in g_tp_failed - * and the eval aborts at the next command-buffer boundary. - * - * Gate sequence values increase monotonically per encoded gate. Both ranks - * encode the identical graph, so the values agree by construction and slots - * never need resetting between tokens. - */ -typedef struct { - uint32_t layer; - uint32_t gate; - uint32_t rows; /* 0 = row gate; >0 = verify-block batch gate */ - uint32_t event_arrival; - uint64_t seq; - /* Big batch gates (prefill): exchange big_bytes from big_out into - * big_in directly (CPU-visible bounce buffers), bypassing the slab. */ - const void *big_out; - void *big_in; - uint64_t big_bytes; -} ds4_gpu_tp_request; - -enum { DS4_GPU_TP_QUEUE = 1024 }; - -static id g_tp_gpu_event; /* GPU -> service thread */ -static id g_tp_cpu_event; /* service thread -> GPU */ -/* Batch (verify-block) gates run on their own sequence space and release - * event: the row-gate seq feeds the RDMA pre-posted recv accounting, which - * requires consecutive values, and a shared release event would make a - * small batch value satisfy waits armed against the larger row seq. */ -static id g_tp_batch_gpu_event; -static id g_tp_batch_cpu_event; -static uint64_t g_tp_batch_seq; -/* Batch flag values are tagged so a stale row-gate seq in the reused FFN - * flag word can never satisfy a batch arrival spin (and vice versa). */ -#define DS4_TP_BATCH_FLAG_TAG 0x80000000u -/* Expert-ownership split parameters for routed kernels. World 1 means TP is - * not bound; world 2 assigns each rank one contiguous expert range. */ -static int32_t g_tp_split_rank; -static int32_t g_tp_split_world = 1; -static int32_t g_tp_session_batch_mode; - -static int ds4_gpu_tp_world_is_two(void) { - return g_tp_split_world == 2; -} - -/* Return the contiguous routed-expert range backed by this process. Rank 1 - * owns the high range and receives any odd-count remainder. */ -static void ds4_gpu_tp_expert_range(uint32_t n_total_expert, - uint32_t *first_expert, - uint32_t *n_expert) { - *first_expert = 0; - *n_expert = n_total_expert; - if (g_tp_split_world != 2) return; - - const uint32_t low_experts = n_total_expert / 2u; - if (g_tp_split_rank == 1) { - *first_expert = low_experts; - *n_expert = n_total_expert - low_experts; - } else { - *n_expert = low_experts; - } -} - -/* Attention head split for GLM batch prefill: each rank computes a - * contiguous half of the heads in the qk-low / attention-lora / - * value-project batch kernels; the caller zeroes the unowned head range - * of the heads buffer and combines the attn-output partials over the - * TP big-gate exchange. */ -static int32_t g_tp_attn_head_split; - -void ds4_gpu_tp_set_attn_head_split(int enabled) { - g_tp_attn_head_split = enabled ? 1 : 0; -} - -static void ds4_gpu_tp_attn_head_range(uint32_t n_head, - uint32_t group, - uint32_t *head_base, - uint32_t *head_count) { - *head_base = 0; - *head_count = n_head; - if (!g_tp_attn_head_split || g_tp_split_world != 2) return; - const uint32_t half = n_head / 2u; - if (half == 0u || (half % group) != 0u || (n_head % 2u) != 0u) return; - *head_count = half; - *head_base = g_tp_split_rank == 1 ? half : 0u; -} -/* Flag gates (DS4_TP_FLAG_GATES): the GPU publishes gate arrival by storing - * the sequence number into a slab word instead of signaling the shared - * event; the service thread spin-reads it from shared memory, which wakes - * hundreds of microseconds earlier than signaledValue polling. The - * CPU->GPU release direction stays on the shared event. */ -static bool g_tp_flag_gates; -static id g_tp_slab_buffer; -static NSUInteger g_tp_slab_buffer_off; -static volatile uint32_t *g_tp_gpu_flags; /* CPU view of the flag words */ -static uint64_t g_tp_gpu_flags_off; -static uint64_t g_tp_seq; -static ds4_gpu_tp_exchange_fn g_tp_exchange_fn; -static ds4_gpu_tp_batch_exchange_fn g_tp_batch_exchange_fn; -static ds4_gpu_tp_big_exchange_fn g_tp_big_exchange_fn; - -void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) { - g_tp_big_exchange_fn = fn; -} - -static void *g_tp_exchange_ud; -static pthread_t g_tp_thread; -static int g_tp_thread_running; -static int g_tp_shutdown; -static int g_tp_failed_flag; -static pthread_mutex_t g_tp_mutex = PTHREAD_MUTEX_INITIALIZER; -static pthread_cond_t g_tp_cond = PTHREAD_COND_INITIALIZER; -static ds4_gpu_tp_request g_tp_queue[DS4_GPU_TP_QUEUE]; -static uint32_t g_tp_queue_head; -static uint32_t g_tp_queue_count; - -static uint64_t g_tp_stat_gates; -static double g_tp_stat_gpu_wait_ms; -static double g_tp_stat_exchange_ms; - -/* GPU keep-alive (see kernel_dsv4_tp_keepalive): its own queue and thread, - * alive exactly as long as the TP gate machinery. */ -static id g_tp_keepalive_queue; -static id g_tp_keepalive_buffer; -static pthread_t g_tp_keepalive_thread; -static int g_tp_keepalive_running; - -/* Nonzero while a verify block runs: the GPU is genuinely busy - * there, so the keep-alive is a pure parasite (~2.3ms per 5-row block - * measured against the single-machine verify). */ -static volatile int g_tp_keepalive_paused; - -void ds4_gpu_tp_keepalive_pause(int paused) { - g_tp_keepalive_paused = paused; -} - -void ds4_gpu_tp_set_session_batch_mode(int enabled) { - g_tp_session_batch_mode = enabled ? 1 : 0; -} - -static uint32_t ds4_gpu_tp_keepalive_tgs_from_env(void) { - uint32_t ka_tgs = 1; - const char *tgs_env = getenv("DS4_TP_KEEPALIVE_TGS"); - if (tgs_env) { - int v = atoi(tgs_env); - if (v > 0 && v <= 2048) ka_tgs = (uint32_t)v; - } - return ka_tgs; -} - -static void *ds4_gpu_tp_keepalive_thread(void *arg) { - (void)arg; - /* Swept on the M5 Max pair: too few iterations lets clocks sag. Current - * TP split-resident Flash runs show 1.2M is a small Q4/Q2 decode win over - * 800k, while two threadgroups waste work. */ - uint32_t iters = 1200000; - const char *env = getenv("DS4_TP_KEEPALIVE_ITERS"); - if (env) iters = (uint32_t)atoi(env); - /* One ALU-only threadgroup keeps the GPU from power-gating but does - * not push the frequency governor; solo-vs-engine kernel gaps - * (~1.7x) suggest decode runs well below max clocks. More TGs raise - * apparent utilization without eating memory bandwidth. */ - uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); - id pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive"); - if (!pipeline) { - fprintf(stderr, "ds4: TP keep-alive pipeline missing\n"); - return NULL; - } - while (!g_tp_shutdown) { - if (g_tp_keepalive_paused) { - usleep(200); - continue; - } - @autoreleasepool { - id cb = [g_tp_keepalive_queue commandBuffer]; - id enc = [cb computeCommandEncoder]; - [enc setComputePipelineState:pipeline]; - [enc setBuffer:g_tp_keepalive_buffer offset:0 atIndex:0]; - [enc setBytes:&iters length:sizeof(iters) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)ka_tgs, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - [enc endEncoding]; - [cb commit]; - [cb waitUntilCompleted]; - } - } - return NULL; -} - -static void *ds4_gpu_tp_service_thread(void *arg) { - (void)arg; - const bool profile = getenv("DS4_TP_GATE_PROFILE") != NULL; - while (1) { - pthread_mutex_lock(&g_tp_mutex); - while (g_tp_queue_count == 0 && !g_tp_shutdown) - pthread_cond_wait(&g_tp_cond, &g_tp_mutex); - if (g_tp_shutdown && g_tp_queue_count == 0) { - pthread_mutex_unlock(&g_tp_mutex); - break; - } - ds4_gpu_tp_request req = g_tp_queue[g_tp_queue_head]; - g_tp_queue_head = (g_tp_queue_head + 1) % DS4_GPU_TP_QUEUE; - g_tp_queue_count--; - pthread_mutex_unlock(&g_tp_mutex); - - /* Wait for the GPU to reach this gate. Tight spin: gate arrival is - * on the decode critical path and normally tens of microseconds - * out; yielding here measurably delays the release wake-up. */ - const double t0 = profile ? ds4_gpu_now_ms() : 0.0; - uint32_t spins = 0; - if (req.big_bytes > 0) { - /* Big gates always signal arrival through the batch shared - * event (see ds4_gpu_tp_big_gate_kick): the event completion - * semantics are what guarantee the bounce payload is visible - * before the exchange reads it. */ - while (g_tp_batch_gpu_event.signaledValue < req.seq) { - if (g_tp_shutdown) break; - if (++spins > (1u << 16)) sched_yield(); - } - } else if (!req.event_arrival) { - const uint32_t slot = req.layer * 2u + req.gate; - uint32_t want = (uint32_t)req.seq; - if (req.rows > 0) - want = DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq; - while (__atomic_load_n(&g_tp_gpu_flags[slot], __ATOMIC_ACQUIRE) != want) { - if (g_tp_shutdown) break; - if (++spins > (1u << 20)) { - sched_yield(); - spins = 0; - } - } - } else if (req.rows > 0) { - while (g_tp_batch_gpu_event.signaledValue < req.seq) { - if (g_tp_shutdown) break; - if (++spins > (1u << 16)) sched_yield(); - } - } else { - while (g_tp_gpu_event.signaledValue < req.seq) { - if (g_tp_shutdown) break; - if (++spins > (1u << 16)) sched_yield(); - } - } - const double t1 = profile ? ds4_gpu_now_ms() : 0.0; - int ok = 0; - if (!g_tp_shutdown && !g_tp_failed_flag) { - if (req.big_bytes > 0) { - if (g_tp_big_exchange_fn) - ok = g_tp_big_exchange_fn(g_tp_exchange_ud, req.layer, - req.seq, req.big_out, - req.big_in, req.big_bytes); - } else if (req.rows > 0) { - if (g_tp_batch_exchange_fn) - ok = g_tp_batch_exchange_fn(g_tp_exchange_ud, req.layer, - req.rows, req.seq); - } else if (g_tp_exchange_fn) { - ok = g_tp_exchange_fn(g_tp_exchange_ud, req.layer, req.gate, - req.seq); - } - } - if (!ok && !g_tp_shutdown) { - if (!g_tp_failed_flag) - fprintf(stderr, "ds4: TP gate exchange failed (layer %u gate %u seq %llu)\n", - req.layer, req.gate, (unsigned long long)req.seq); - g_tp_failed_flag = 1; - } - /* Release the GPU even on failure so end_commands can drain. */ - if (req.rows > 0) g_tp_batch_cpu_event.signaledValue = req.seq; - else g_tp_cpu_event.signaledValue = req.seq; - if (profile) { - g_tp_stat_gpu_wait_ms += t1 - t0; - g_tp_stat_exchange_ms += ds4_gpu_now_ms() - t1; - if (++g_tp_stat_gates % 860 == 0) { - fprintf(stderr, - "ds4: TP gates %llu: avg gpu-wait %.1f us, avg exchange %.1f us\n", - (unsigned long long)g_tp_stat_gates, - g_tp_stat_gpu_wait_ms / (double)g_tp_stat_gates * 1000.0, - g_tp_stat_exchange_ms / (double)g_tp_stat_gates * 1000.0); - } - } - } - return NULL; -} - -int ds4_gpu_tp_init(uint32_t rank, - ds4_gpu_tensor *slab, uint64_t gpu_flags_off, - ds4_gpu_tp_exchange_fn fn, void *ud) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (g_tp_thread_running || rank > 1) return 0; - g_tp_split_rank = (int32_t)rank; - g_tp_split_world = 2; - g_tp_slab_buffer = slab ? ds4_gpu_tensor_buffer(slab) : nil; - g_tp_slab_buffer_off = slab ? ds4_gpu_tensor_offset(slab) : 0; - g_tp_gpu_flags_off = gpu_flags_off; - g_tp_gpu_flags = slab ? - (volatile uint32_t *)((uint8_t *)ds4_gpu_tensor_contents(slab) + gpu_flags_off) : NULL; - /* Flag arrival is the default: the slab-word publish detects in ~1-3us - * where signaledValue polling costs 10-20, worth +2.3 t/s on the pair - * (A/B 2026-07-06, byte-identical output). DS4_TP_EVENT_GATES falls - * back to the shared-event arrival path. */ - g_tp_flag_gates = g_tp_gpu_flags != NULL && getenv("DS4_TP_EVENT_GATES") == NULL; - g_tp_gpu_event = [g_device newSharedEvent]; - g_tp_cpu_event = [g_device newSharedEvent]; - g_tp_batch_gpu_event = [g_device newSharedEvent]; - g_tp_batch_cpu_event = [g_device newSharedEvent]; - if (!g_tp_gpu_event || !g_tp_cpu_event || - !g_tp_batch_gpu_event || !g_tp_batch_cpu_event) { - fprintf(stderr, "ds4: failed to create TP shared events\n"); - return 0; - } - g_tp_exchange_fn = fn; - g_tp_exchange_ud = ud; - g_tp_seq = 0; - g_tp_batch_seq = 0; - g_tp_shutdown = 0; - g_tp_failed_flag = 0; - g_tp_queue_head = 0; - g_tp_queue_count = 0; - pthread_attr_t attr; - pthread_attr_init(&attr); - pthread_attr_set_qos_class_np(&attr, QOS_CLASS_USER_INTERACTIVE, 0); - if (pthread_create(&g_tp_thread, &attr, ds4_gpu_tp_service_thread, NULL) != 0) { - pthread_attr_destroy(&attr); - fprintf(stderr, "ds4: failed to start TP gate service thread\n"); - return 0; - } - pthread_attr_destroy(&attr); - g_tp_thread_running = 1; - if (getenv("DS4_TP_NO_KEEPALIVE") == NULL) { - uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); - g_tp_keepalive_queue = [g_device newCommandQueue]; - g_tp_keepalive_buffer = [g_device newBufferWithLength:(NSUInteger)ka_tgs * 256u * sizeof(float) - options:MTLResourceStorageModeShared]; - if (g_tp_keepalive_queue && g_tp_keepalive_buffer && - pthread_create(&g_tp_keepalive_thread, NULL, - ds4_gpu_tp_keepalive_thread, NULL) == 0) { - g_tp_keepalive_running = 1; - } else { - fprintf(stderr, "ds4: TP keep-alive setup failed (continuing without)\n"); - } - } - return 1; -} - -void ds4_gpu_tp_shutdown(void) { - if (!g_tp_thread_running) return; - pthread_mutex_lock(&g_tp_mutex); - g_tp_shutdown = 1; - pthread_cond_broadcast(&g_tp_cond); - pthread_mutex_unlock(&g_tp_mutex); - pthread_join(g_tp_thread, NULL); - g_tp_thread_running = 0; - if (g_tp_keepalive_running) { - pthread_join(g_tp_keepalive_thread, NULL); - g_tp_keepalive_running = 0; - g_tp_keepalive_queue = nil; - g_tp_keepalive_buffer = nil; - } - g_tp_exchange_fn = NULL; - g_tp_batch_exchange_fn = NULL; - g_tp_exchange_ud = NULL; - g_tp_split_rank = 0; - g_tp_split_world = 1; - g_tp_session_batch_mode = 0; -} - -void ds4_gpu_tp_suspend_expert_sharding(int suspend) { - if (!g_tp_thread_running) return; - g_tp_split_world = suspend ? 1 : 2; -} - -int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) { - if (!g_batch_cb) { - fprintf(stderr, "ds4: TP gate encode without an open command batch (layer %u gate %u)\n", - layer, gate); - return 0; - } - if (!g_tp_thread_running) { - fprintf(stderr, "ds4: TP gate encode without the gate service (layer %u)\n", layer); - return 0; - } - const uint64_t seq = ++g_tp_seq; - const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; - if (!event_arrival) { - /* Publish arrival through the slab word; the buffer hazard against - * the partial-output kernels orders the store after the payload. */ - const uint32_t slot = layer * 2u + gate; - const uint32_t value = (uint32_t)seq; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) return 0; - id pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); - if (!pipeline) return 0; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBuffer:g_tp_slab_buffer - offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) - atIndex:0]; - [enc setBytes:&value length:sizeof(value) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - ds4_gpu_close_batch_encoder(); - } else { - ds4_gpu_close_batch_encoder(); - [g_batch_cb encodeSignalEvent:g_tp_gpu_event value:seq]; - } - [g_batch_cb encodeWaitForEvent:g_tp_cpu_event value:seq]; - pthread_mutex_lock(&g_tp_mutex); - if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { - pthread_mutex_unlock(&g_tp_mutex); - fprintf(stderr, "ds4: TP gate queue overflow\n"); - return 0; - } - uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; - g_tp_queue[tail].layer = layer; - g_tp_queue[tail].gate = gate; - g_tp_queue[tail].rows = 0; - g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; - g_tp_queue[tail].seq = seq; - g_tp_queue[tail].big_out = NULL; - g_tp_queue[tail].big_in = NULL; - g_tp_queue[tail].big_bytes = 0; - g_tp_queue_count++; - pthread_cond_signal(&g_tp_cond); - pthread_mutex_unlock(&g_tp_mutex); - return 1; -} - -void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) { - g_tp_batch_exchange_fn = fn; -} - -/* Verify-block batch gate: same arrival/release machinery as the row gate - * (the FFN flag word and event pair are reused — a decode gate and a batch - * gate are never in flight together, and seq values stay globally unique), - * but the service thread runs the multi-row exchange callback. */ -int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) { - if (!g_batch_cb) return 0; - if (!g_tp_thread_running || rows == 0) return 0; - const uint64_t seq = ++g_tp_batch_seq; - const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; - if (!event_arrival) { - const uint32_t slot = layer * 2u + 1u; /* FFN gate slot */ - const uint32_t value = DS4_TP_BATCH_FLAG_TAG | (uint32_t)seq; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) return 0; - id pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); - if (!pipeline) return 0; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBuffer:g_tp_slab_buffer - offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) - atIndex:0]; - [enc setBytes:&value length:sizeof(value) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - ds4_gpu_close_batch_encoder(); - } else { - ds4_gpu_close_batch_encoder(); - [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; - } - [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; - pthread_mutex_lock(&g_tp_mutex); - if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { - pthread_mutex_unlock(&g_tp_mutex); - fprintf(stderr, "ds4: TP gate queue overflow\n"); - return 0; - } - uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; - g_tp_queue[tail].layer = layer; - g_tp_queue[tail].gate = 1u; /* FFN */ - g_tp_queue[tail].rows = rows; - g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; - g_tp_queue[tail].seq = seq; - g_tp_queue[tail].big_out = NULL; - g_tp_queue[tail].big_in = NULL; - g_tp_queue[tail].big_bytes = 0; - g_tp_queue_count++; - pthread_cond_signal(&g_tp_cond); - pthread_mutex_unlock(&g_tp_mutex); - return 1; -} - -/* Prefill batch gate kick: same seq space and release event as the verify - * batch gate, but the service thread exchanges big_bytes directly between - * the two shared bounce buffers instead of slab slots. The kick only - * publishes the GPU arrival marker and queues the exchange; the caller - * encodes the release wait later through ds4_gpu_tp_big_gate_wait, which - * lets it interleave more GPU work with the wire exchange. Arrival always - * uses the batch shared event, NOT the flag word: a flag write carries no - * memory-visibility guarantee for the payload buffer, and once the GPU - * keeps running past the kick (no event wait right behind it) the service - * thread can observe the flag before the producing kernels' stores reach - * CPU-visible memory (measured: stale rows in the first sub-kick). The - * shared-event signal only fires after every preceding command completes, - * which is exactly the payload ordering the exchange needs; the ~10 us - * slower arrival detection is noise against a multi-ms exchange. */ -uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows, - const ds4_gpu_tensor *out_t, - ds4_gpu_tensor *in_t, - uint64_t bytes) { - if (!g_batch_cb) return 0; - if (!g_tp_thread_running || rows == 0 || bytes == 0) return 0; - const void *out_ptr = ds4_gpu_tensor_contents((ds4_gpu_tensor *)out_t); - void *in_ptr = ds4_gpu_tensor_contents(in_t); - if (!out_ptr || !in_ptr) { - fprintf(stderr, "ds4: TP big gate needs CPU-visible bounce buffers\n"); - return 0; - } - const uint64_t seq = ++g_tp_batch_seq; - ds4_gpu_close_batch_encoder(); - [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; - pthread_mutex_lock(&g_tp_mutex); - if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { - pthread_mutex_unlock(&g_tp_mutex); - fprintf(stderr, "ds4: TP gate queue overflow\n"); - return 0; - } - uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; - g_tp_queue[tail].layer = layer; - g_tp_queue[tail].gate = 1u; - g_tp_queue[tail].rows = rows; - g_tp_queue[tail].event_arrival = 1u; - g_tp_queue[tail].seq = seq; - g_tp_queue[tail].big_out = out_ptr; - g_tp_queue[tail].big_in = in_ptr; - g_tp_queue[tail].big_bytes = bytes; - g_tp_queue_count++; - pthread_cond_signal(&g_tp_cond); - pthread_mutex_unlock(&g_tp_mutex); - return seq; -} - -/* Encode the GPU-side release wait for a previously kicked big gate. The - * batch release event is monotonic and the service thread completes queued - * exchanges in kick order, so waiting on the LAST kicked seq of a stage - * also covers every earlier kick. */ -int ds4_gpu_tp_big_gate_wait(uint64_t seq) { - if (!g_batch_cb || seq == 0) return 0; - ds4_gpu_close_batch_encoder(); - [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; - return 1; -} - -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) { - const uint64_t seq = ds4_gpu_tp_big_gate_kick(layer, rows, out_t, in_t, bytes); - if (seq == 0) return 0; - return ds4_gpu_tp_big_gate_wait(seq); -} - -int ds4_gpu_tp_failed(void) { - return g_tp_failed_flag; -} - -int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (event_value == 0) return 0; - - if (@available(macOS 12.0, *)) { - if (!g_selected_readback_event) return 0; - - const char *what = label ? label : "selected-id readback"; - const BOOL signaled = - [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; - if (!signaled) { - fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); - return 0; - } - return 1; - } - - fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); - return 0; -} - -static int ds4_gpu_signal_batch_and_wait_event(const char *label) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!g_batch_cb) return 0; - - if (@available(macOS 12.0, *)) { - if (!g_selected_readback_event) { - g_selected_readback_event = [g_device newSharedEvent]; - if (!g_selected_readback_event) { - fprintf(stderr, "ds4: failed to create Metal shared event for %s\n", - label ? label : "selected-id readback"); - return 0; - } - } - - ds4_gpu_close_batch_encoder(); - id cb = g_batch_cb; - g_batch_cb = nil; - const uint64_t value = ++g_selected_readback_event_value; - [cb encodeSignalEvent:g_selected_readback_event value:value]; - g_batch_has_work = YES; - [cb commit]; - ds4_gpu_stream_expert_cache_note_batch_committed(); - - const BOOL signaled = [g_selected_readback_event waitUntilSignaledValue:value timeoutMS:60000]; - [g_pending_cbs addObject:cb]; - if (!signaled) { - fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", - label ? label : "selected-id readback"); - (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); - [g_transient_buffers removeAllObjects]; - return 0; - } - if (cb.status == MTLCommandBufferStatusError) { - fprintf(stderr, "ds4: Metal %s failed: %s\n", - label ? label : "selected-id readback", - [[cb.error localizedDescription] UTF8String]); - (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); - [g_transient_buffers removeAllObjects]; - return 0; - } - - g_batch_cb = ds4_gpu_new_command_buffer(); - g_batch_has_work = NO; - if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); - if (!g_batch_cb) { - (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); - [g_transient_buffers removeAllObjects]; - return 0; - } - return 1; - } else { - ds4_gpu_close_batch_encoder(); - id cb = g_batch_cb; - g_batch_cb = nil; - g_batch_has_work = NO; - if (ds4_gpu_finish_command_buffer(cb, 1, label ? label : "selected-id readback") == 0) { - return 0; - } - return ds4_gpu_begin_commands(); - } -} - -int ds4_gpu_end_commands(void) { - if (!g_batch_cb) return 0; - ds4_gpu_close_batch_encoder(); - id cb = g_batch_cb; - g_batch_cb = nil; - g_batch_has_work = NO; - g_stream_expert_cache_owned_seq = g_stream_expert_cache_batch_seq; - g_stream_expert_cache_batch_seq = 0; - return ds4_gpu_finish_command_buffer(cb, 1, "command batch"); -} - -static int ds4_gpu_flash_attn_stage_profile_boundary( - id __strong *cbp, - const char *mode, - const char *stage, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t n_keys, - uint32_t n_head, - uint32_t head_dim, - uint32_t window, - uint32_t ratio, - double *stage_t0) { - if (!cbp || !*cbp || !stage_t0 || !stage) return 0; - if (ds4_gpu_end_commands() == 0) return 0; - - const double now_ms = ds4_gpu_now_ms(); - const char *filter = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE_FILTER"); - const int print_stage = - !filter || !filter[0] || - strstr(stage, filter) != NULL || - (mode && strstr(mode, filter) != NULL); - if (print_stage) { - fprintf(stderr, - "ds4: Metal FlashAttention prefill stage mode=%s tokens=%u comp=%u " - "keys=%u heads=%u dim=%u window=%u ratio=%u %s=%.3f ms\n", - mode ? mode : "unknown", - n_tokens, - n_comp, - n_keys, - n_head, - head_dim, - window, - ratio, - stage, - now_ms - *stage_t0); - } - *stage_t0 = now_ms; - - if (ds4_gpu_begin_commands() == 0) return 0; - int owned = 0; - *cbp = ds4_gpu_command_buffer(&owned); - return *cbp != nil && owned == 0; -} - -int ds4_gpu_synchronize(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (g_batch_cb) return ds4_gpu_end_commands(); - if ([g_pending_cbs count] != 0) { - int ok = ds4_gpu_wait_pending_command_buffers("synchronize"); - [g_transient_buffers removeAllObjects]; - ds4_gpu_model_buffer_cache_maybe_evict("synchronize"); - return ok; - } - - id cb = ds4_gpu_new_command_buffer(); - if (!cb) return 0; - return ds4_gpu_finish_command_buffer(cb, 1, "synchronize"); -} - -void ds4_gpu_cleanup(void) { - if (!g_initialized) return; - - @autoreleasepool { - if (g_batch_cb) { - ds4_gpu_close_batch_encoder(); - [g_batch_cb commit]; - [g_batch_cb waitUntilCompleted]; - g_batch_cb = nil; - if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_done_seq) { - g_stream_expert_cache_done_seq = g_stream_expert_cache_batch_seq; - } - g_stream_expert_cache_batch_seq = 0; - } - (void)ds4_gpu_wait_pending_command_buffers("cleanup"); - if (ds4_gpu_stream_expert_timing_summary_enabled() && - getenv("DS4_METAL_MEMORY_REPORT") == NULL) { - ds4_gpu_print_memory_report("at cleanup"); - } - g_selected_readback_event = nil; - g_selected_readback_event_value = 0; - [g_transient_buffers removeAllObjects]; - ds4_gpu_stream_expert_pread_pool_shutdown(); - ds4_gpu_stream_expert_cache_clear_all(1); - for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { - g_stream_expert_cache_gate_addr_buffers[layer] = nil; - g_stream_expert_cache_up_addr_buffers[layer] = nil; - g_stream_expert_cache_down_addr_buffers[layer] = nil; - g_stream_compact_gate_addr_buffers[layer] = nil; - g_stream_compact_up_addr_buffers[layer] = nil; - g_stream_compact_down_addr_buffers[layer] = nil; - g_stream_compact_selected_buffers[layer] = nil; - g_stream_selected_id_buffers[layer] = nil; - } - g_set_rows_f32_i32_pipeline = nil; - g_get_rows_f32_pipeline = nil; - g_get_rows_f16_pipeline = nil; - g_get_rows_i32_pipeline = nil; - g_get_rows_q8_0_pipeline = nil; - g_get_rows_q4_0_pipeline = nil; - g_get_rows_q4_K_pipeline = nil; - g_repeat_f32_pipeline = nil; - g_concat_pipeline = nil; - g_cpy_f32_f32_pipeline = nil; - g_cpy_f32_f16_pipeline = nil; - g_cpy_contig_f32_f16_pipeline = nil; - g_cpy_f16_f32_pipeline = nil; - g_cpy_f16_f16_pipeline = nil; - g_cpy_contig_f16_f32_pipeline = nil; - g_cpy_contig_f16_f16_pipeline = nil; - g_flash_kv_stage_f16_pipeline = nil; - g_swiglu_pipeline = nil; - g_swiglu_flat_pipeline = nil; - g_add_pipeline = nil; - g_add2_pipeline = nil; - g_add3_pipeline = nil; - g_moe_sum6_pipeline = nil; - g_moe_sum8_pipeline = nil; - g_mul_pipeline = nil; - g_bin_mul_scalar_pipeline = nil; - g_bin_div_row_pipeline = nil; - g_unary_sigmoid_pipeline = nil; - g_unary_silu_pipeline = nil; - g_unary_softplus_pipeline = nil; - g_unary_sqrt_pipeline = nil; - g_unary_clamp_pipeline = nil; - g_unary_scale_pipeline = nil; - g_unary_fill_pipeline = nil; - g_unary_fill_f16_pipeline = nil; - g_rms_norm_pipeline = nil; - g_rms_norm_plain_pipeline = nil; - g_add_rms_norm_pipeline = nil; - g_rms_norm_scale_pipeline = nil; - g_dsv4_qkv_rms_norm_pipeline = nil; - g_hc_split_sinkhorn_pipeline = nil; - g_hc_split_weighted_sum_pipeline = nil; - g_hc_split_weighted_sum_norm_pipeline = nil; - g_hc_weighted_sum_pipeline = nil; - g_hc_weighted_sum_norm_pipeline = nil; - g_output_hc_weights4_pipeline = nil; - g_hc_expand_pipeline = nil; - g_moe_mul_mv_id_iq2_xxs_pipeline = nil; - g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil; - g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil; - g_moe_mul_mv_id_q2_k_pipeline = nil; - g_moe_mul_mv_id_q2_k_sum6_pipeline = nil; - g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil; - g_moe_mul_mv_id_q4_k_pipeline = nil; - g_moe_mul_mv_id_q4_k_pair_pipeline = nil; - g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_id_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_group_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_group6_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_group8_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_group24_q4_k_id_pipeline = nil; - g_moe_mul_mv_group24_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = nil; - g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil; - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil; - g_moe_mul_mv_addr_iq2_xxs_pipeline = nil; - g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil; - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = nil; - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = nil; - g_moe_stream_expert_cache_validate_pipeline = nil; - g_moe_q4_gather_slots6_pipeline = nil; - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_table_q4_k_sum6_pipeline = nil; - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = nil; - g_moe_mul_mv_addr_q4_k_sum6_pipeline = nil; - g_moe_table_q4_pair_gate_encoder = nil; - g_moe_table_q4_pair_up_encoder = nil; - g_moe_table_q4_sum_down_encoder = nil; - g_rope_tail_batch_pipeline = nil; - g_rope_tail_inplace_pair_pipeline = nil; - g_rope_tail_inplace_pair_shared4_pipeline = nil; - g_rope_tail_inplace_pair_affine_pipeline = nil; - g_dsv4_fp8_kv_quantize_pipeline = nil; - g_dsv4_indexer_qat_pipeline = nil; - g_dsv4_kv_fp8_store_pipeline = nil; - g_dsv4_ratio4_shift_pipeline = nil; - g_dsv4_compressor_pack_ratio4_pipeline = nil; - g_dsv4_softmax_pool_ratio4_direct_pipeline = nil; - g_dsv4_softmax_pool_pipeline = nil; - g_soft_max_f32_pipeline = nil; - g_soft_max_f32_4_pipeline = nil; - g_argsort_f32_i32_desc_pipeline = nil; - g_argsort_merge_f32_i32_desc_pipeline = nil; - g_sum_rows_f32_f32_pipeline = nil; - g_dsv4_topk_mask_pipeline = nil; - g_dsv4_topk_mask_scatter_pipeline = nil; - g_dsv4_indexer_weighted_sum_pipeline = nil; - g_dsv4_indexer_score_one_direct_pipeline = nil; - g_dsv4_compressor_store_one_pipeline = nil; - g_dsv4_sort_i32_rows_asc_pipeline = nil; - g_dsv4_indexed_attention_heads8_pipeline = nil; - g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; - g_dsv4_softplus_sqrt_pipeline = nil; - g_dsv4_router_finalize_one_pipeline = nil; - g_dsv4_router_finalize_one_simd_pipeline = nil; - g_dsv4_router_finalize_weights_one_simd_pipeline = nil; - g_dsv4_router_weights_one_pipeline = nil; - g_glm_router_select_one_pipeline = nil; - g_glm_kv_lora_rms_norm_pipeline = nil; - g_glm_k_b_project_pipeline = nil; - g_glm_store_compact_kv_pipeline = nil; - g_glm_qkv_norm_store_compact_kv_pipeline = nil; - g_glm_store_indexer_k_pipeline = nil; - g_glm_build_kv_cache_pipeline = nil; - g_glm_build_kv_cache_decode_group4_pipeline = nil; - g_glm_build_kv_cache_flash_pipeline = nil; - g_glm_attention_full_pipeline = nil; - g_glm_fill_selected_range_pipeline = nil; - g_glm_fill_selected_range_batch_pipeline = nil; - g_glm_indexer_rope_tail_pipeline = nil; - g_glm_indexer_score_one_pipeline = nil; - g_glm_indexer_score_one_direct_pipeline = nil; - g_glm_indexer_scores_batch_pipeline = nil; - g_glm_indexer_scores_tiled_pipeline = nil; - g_glm_indexer_scores_tiled_f32_pipeline = nil; - g_glm_qk_lowrank_pipeline = nil; - g_glm_qk_lowrank_glm52_pipeline = nil; - g_glm_qk_lowrank_glm52_sg_pipeline = nil; - g_glm_qk_lowrank_batch_pipeline = nil; - g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil; - g_glm_value_project_q8_0_pipeline = nil; - g_glm_value_project_q8_0_batch_heads_pipeline = nil; - g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil; - g_glm_attention_indexed_decode_pipeline = nil; - g_glm_attention_indexed_decode_split_group8_partial_pipeline = nil; - g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil; - g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil; - g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil; - g_glm_attention_indexed_batch_pipeline = nil; - g_glm_attention_indexed_batch_group2_pipeline = nil; - g_glm_attention_indexed_batch_q2_group4_pipeline = nil; - g_glm_attention_indexed_batch_group8_pipeline = nil; - g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil; - g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil; - g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil; - g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil; - g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil; - g_glm_q4_k_pair_swiglu_f32_pipeline = nil; - g_glm_q4_k_pair_swiglu2_f32_pipeline = nil; - g_glm_q4_k_pair_swiglu4_f32_pipeline = nil; - g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = nil; - g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = nil; - g_glm_q2_k_pair_swiglu_f32_pipeline = nil; - g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = nil; - g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = nil; - g_glm_q4_k_addr_pair_swiglu_f32_pipeline = nil; - g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = nil; - g_glm_q2_k_down_f32_pipeline = nil; - g_glm_q4_k_down_f32_pipeline = nil; - g_glm_q2_k_addr_down_f32_pipeline = nil; - g_glm_q4_k_addr_down_f32_pipeline = nil; - g_glm_q5_k_pair_swiglu_f32_pipeline = nil; - g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = nil; - g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = nil; - g_glm_q5_k_down_f32_pipeline = nil; - g_glm_q6_k_down_f32_pipeline = nil; - g_dsv4_router_weights_batch_pipeline = nil; - g_dsv4_hc_expand4_pipeline = nil; - g_flash_attn_mask_buffer = nil; - g_flash_attn_zero_mask_buffer = nil; - g_flash_attn_pad_buffer = nil; - g_flash_attn_tmp_buffer = nil; - g_flash_attn_blk_buffer = nil; - ds4_gpu_clear_zero_prefix_prefill_mask_cache(); - g_flash_attn_ring_buffer = nil; - g_flash_attn_kv_buffer = nil; - g_glm_flash_attn_mask_buffer = nil; - g_compressor_pool_kv_buffer = nil; - g_compressor_pool_score_buffer = nil; - g_compressor_pool_score_cont_buffer = nil; - g_compressor_pool_softmax_buffer = nil; - g_compressor_pool_product_buffer = nil; - g_compressor_store_ape_buffer = nil; - g_compressor_store_score_buffer = nil; - g_embed_rows_buffer = nil; - g_router_selection_buffer = nil; - g_router_weight_sum_buffer = nil; - g_indexer_head_scores_buffer = nil; - g_indexer_topk_buffer = nil; - g_indexed_topk_buffer = nil; - g_stream_expert_validate_status_buffer = nil; - g_f16_round_scratch_buffer = nil; - g_raw_store_round_buffer = nil; - g_moe_gate_scratch_buffer = nil; - g_moe_down_scratch_buffer = nil; - g_moe_id_map_buffer = nil; - g_moe_q4_gate_slots_buffer = nil; - g_moe_q4_up_slots_buffer = nil; - g_moe_q4_down_slots_buffer = nil; - g_attn_out_group_ids_buffer = nil; - g_model_fd = -1; - g_model_map_ptr = NULL; - g_model_map_size = 0; - g_model_mapped_offset = 0; - g_model_mapped_size = 0; - g_model_mapped_max_tensor_bytes = 0; - ds4_gpu_tensor_tracking_reset(); - g_flash_attn_mask_bytes = 0; - g_flash_attn_zero_mask_bytes = 0; - g_flash_attn_pad_bytes = 0; - g_flash_attn_tmp_bytes = 0; - g_flash_attn_blk_bytes = 0; - g_flash_attn_ring_bytes = 0; - g_flash_attn_kv_bytes = 0; - g_glm_flash_attn_mask_bytes = 0; - g_glm_flash_attn_mask_valid = 0; - g_glm_flash_attn_mask_pos0 = 0; - g_glm_flash_attn_mask_tokens = 0; - g_glm_flash_attn_mask_cache_len = 0; - g_compressor_pool_kv_bytes = 0; - g_compressor_pool_score_bytes = 0; - g_compressor_pool_score_cont_bytes = 0; - g_compressor_pool_softmax_bytes = 0; - g_compressor_pool_product_bytes = 0; - g_compressor_store_ape_bytes = 0; - g_compressor_store_score_bytes = 0; - g_embed_rows_bytes = 0; - g_router_selection_bytes = 0; - g_router_weight_sum_bytes = 0; - g_indexer_head_scores_bytes = 0; - g_indexer_topk_bytes = 0; - g_indexed_topk_bytes = 0; - g_f16_round_scratch_bytes = 0; - g_raw_store_round_bytes = 0; - g_moe_gate_scratch_bytes = 0; - g_moe_down_scratch_bytes = 0; - g_moe_id_map_bytes = 0; - g_moe_q4_gate_slots_bytes = 0; - g_moe_q4_up_slots_bytes = 0; - g_moe_q4_down_slots_bytes = 0; - g_attn_out_group_ids_bytes = 0; - g_model_wrap_count = 0; - g_model_wrap_bytes = 0; - g_model_wrap_max_bytes = 0; - g_model_buffer_cache_bytes = 0; - g_model_buffer_cache_evictions = 0; - g_model_buffer_cache_over_limit = 0; - ds4_gpu_model_residency_clear(); - ds4_gpu_model_views_clear(); - [g_pipeline_cache removeAllObjects]; - g_pipeline_cache = nil; - [g_q4_expert_layer_residency_cache removeAllObjects]; - g_q4_expert_layer_residency_cache = nil; - [g_q4_expert_table_cache removeAllObjects]; - g_q4_expert_table_cache = nil; - [g_model_buffer_cache removeAllObjects]; - g_model_buffer_cache = nil; - g_transient_buffers = nil; - g_pending_cbs = nil; - g_library = nil; - g_queue = nil; - g_device = nil; - g_initialized = 0; - } -} - -static int ds4_gpu_encode_get_rows_f16( - id cb, - id weight, - NSUInteger weight_offset, - id tokens, - NSUInteger tokens_offset, - id out, - NSUInteger out_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd) { - if (!cb || !weight || !tokens || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { - return 0; - } - - const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); - const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); - ds4_gpu_get_rows_args args = { - .ne00t = (int32_t)n_embd, - .ne00 = (int32_t)n_embd, - .nb01 = src_row_bytes, - .nb02 = (uint64_t)n_vocab * src_row_bytes, - .nb03 = (uint64_t)n_vocab * src_row_bytes, - .ne10 = (int32_t)n_tokens, - .nb10 = sizeof(int32_t), - .nb11 = token_bytes, - .nb12 = token_bytes, - .nb1 = dst_row_bytes, - .nb2 = (uint64_t)n_tokens * dst_row_bytes, - .nb3 = (uint64_t)n_tokens * dst_row_bytes, - }; - - NSUInteger nth = (NSUInteger)n_embd; - const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_get_rows_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weight offset:weight_offset atIndex:1]; - [enc setBuffer:tokens offset:tokens_offset atIndex:2]; - [enc setBuffer:out offset:out_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static uint64_t ds4_gpu_q8_0_row_bytes(uint32_t n_embd) { - return (((uint64_t)n_embd + 31u) / 32u) * 34u; -} - -static int ds4_gpu_quant_row_bytes( - uint32_t type, - uint32_t n_embd, - uint64_t *row_bytes_out) { - if (!row_bytes_out || n_embd == 0) return 0; - switch (type) { - case DS4_METAL_TENSOR_Q8_0: - *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 34u; - return 1; - case DS4_METAL_TENSOR_Q4_0: - *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 18u; - return 1; - case DS4_METAL_TENSOR_Q4_K: - if ((n_embd % 256u) != 0) return 0; - *row_bytes_out = ((uint64_t)n_embd / 256u) * 144u; - return 1; - default: - return 0; - } -} - -static int ds4_gpu_q8_0_table_bytes( - uint32_t n_vocab, - uint32_t n_embd, - uint64_t *bytes_out) { - if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; - const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); - if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; - *bytes_out = (uint64_t)n_vocab * row_bytes; - return 1; -} - -static int ds4_gpu_quant_table_bytes( - uint32_t type, - uint32_t n_vocab, - uint32_t n_embd, - uint64_t *bytes_out) { - if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(type, n_embd, &row_bytes)) return 0; - if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; - *bytes_out = (uint64_t)n_vocab * row_bytes; - return 1; -} - -static int ds4_gpu_encode_get_rows_q8_0( - id cb, - id weight, - NSUInteger weight_offset, - id tokens, - NSUInteger tokens_offset, - const int32_t *single_token, - id out, - NSUInteger out_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd) { - if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { - return 0; - } - if (!tokens && (!single_token || n_tokens != 1)) { - return 0; - } - - ds4_gpu_get_rows_q8_0_args args = { - .n_embd = (int32_t)n_embd, - .n_vocab = (int32_t)n_vocab, - .n_tokens = (int32_t)n_tokens, - .src_row_bytes = ds4_gpu_q8_0_row_bytes(n_embd), - .dst_row_bytes = (uint64_t)n_embd * sizeof(float), - .token_stride = sizeof(int32_t), - }; - - NSUInteger nth = 32u; - const NSUInteger max_threads = g_get_rows_q8_0_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - const NSUInteger nblocks = ((NSUInteger)n_embd + 31u) / 32u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_get_rows_q8_0_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weight offset:weight_offset atIndex:1]; - if (tokens) { - [enc setBuffer:tokens offset:tokens_offset atIndex:2]; - } else { - [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; - } - [enc setBuffer:out offset:out_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_get_rows_quant( - id cb, - id weight, - NSUInteger weight_offset, - uint32_t weight_type, - id tokens, - NSUInteger tokens_offset, - const int32_t *single_token, - id out, - NSUInteger out_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd) { - if (weight_type == DS4_METAL_TENSOR_Q8_0) { - return ds4_gpu_encode_get_rows_q8_0(cb, - weight, - weight_offset, - tokens, - tokens_offset, - single_token, - out, - out_offset, - n_vocab, - n_tokens, - n_embd); - } - if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { - return 0; - } - if (!tokens && (!single_token || n_tokens != 1)) { - return 0; - } - - uint64_t src_row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, n_embd, &src_row_bytes)) return 0; - ds4_gpu_get_rows_q8_0_args args = { - .n_embd = (int32_t)n_embd, - .n_vocab = (int32_t)n_vocab, - .n_tokens = (int32_t)n_tokens, - .src_row_bytes = src_row_bytes, - .dst_row_bytes = (uint64_t)n_embd * sizeof(float), - .token_stride = sizeof(int32_t), - }; - - id pipeline = nil; - NSUInteger block_width = 0; - if (weight_type == DS4_METAL_TENSOR_Q4_0) { - pipeline = g_get_rows_q4_0_pipeline; - block_width = 32u; - } else if (weight_type == DS4_METAL_TENSOR_Q4_K) { - pipeline = g_get_rows_q4_K_pipeline; - block_width = 256u; - } - if (!pipeline || block_width == 0) return 0; - - NSUInteger nth = block_width; - const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - const NSUInteger nblocks = ((NSUInteger)n_embd + block_width - 1u) / block_width; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weight offset:weight_offset atIndex:1]; - if (tokens) { - [enc setBuffer:tokens offset:tokens_offset atIndex:2]; - } else { - [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; - } - [enc setBuffer:out offset:out_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_repeat_hc_embedding( - id cb, - id rows, - NSUInteger rows_offset, - id out, - NSUInteger out_offset, - uint32_t n_tokens, - uint32_t n_embd, - uint32_t n_hc) { - if (!cb || !rows || !out || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; - - const uint64_t embd_bytes = (uint64_t)n_embd * sizeof(float); - ds4_gpu_repeat_args args = { - .ne00 = (int32_t)n_embd, - .ne01 = 1, - .ne02 = (int32_t)n_tokens, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = embd_bytes, - .nb02 = embd_bytes, - .nb03 = (uint64_t)n_tokens * embd_bytes, - .ne0 = (int32_t)n_embd, - .ne1 = (int32_t)n_hc, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = embd_bytes, - .nb2 = (uint64_t)n_hc * embd_bytes, - .nb3 = (uint64_t)n_tokens * n_hc * embd_bytes, - }; - - NSUInteger nth = (NSUInteger)n_embd; - const NSUInteger max_threads = g_repeat_f32_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_repeat_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:rows offset:rows_offset atIndex:1]; - [enc setBuffer:out offset:out_offset atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(n_hc, n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -int ds4_gpu_embed_token_q8_0_tensor( - ds4_gpu_tensor *out, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t token, - uint32_t n_embd) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); - if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal Q8_0 embedding received undersized output buffer\n"); - return 0; - } - - uint64_t weight_bytes = 0; - if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || - weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8_0 embedding range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - uint32_t token_for_kernel = token; - id wbuf = nil; - const bool exact_token_row = - getenv("DS4_METAL_DISABLE_TOKEN_EMBED_EXACT_VIEW") == NULL; - if (exact_token_row) { - const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); - const uint64_t token_rel = (uint64_t)token * row_bytes; - if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { - fprintf(stderr, "ds4: Metal Q8_0 embedding token row is outside the mapped table\n"); - return 0; - } - wbuf = ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset + token_rel, - row_bytes, - &inner_offset); - token_for_kernel = 0; - } else { - wbuf = ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - } - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const int32_t token_i32 = (int32_t)token_for_kernel; - if (!ds4_gpu_encode_get_rows_q8_0(cb, - wbuf, - (NSUInteger)inner_offset, - nil, - 0, - &token_i32, - outbuf, - ds4_gpu_tensor_offset(out), - n_vocab, - 1, - n_embd)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed token")) return 0; - } - - return 1; -} - -int ds4_gpu_embed_tokens_q8_0_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *tokens, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out); - id tokbuf = ds4_gpu_tensor_buffer(tokens); - const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); - const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); - if (!outbuf || !tokbuf || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(tokens) < token_bytes) { - fprintf(stderr, "ds4: Metal Q8_0 batched embedding received undersized buffers\n"); - return 0; - } - - uint64_t weight_bytes = 0; - if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || - weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8_0 batched embedding range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_get_rows_q8_0(cb, - wbuf, - (NSUInteger)inner_offset, - tokbuf, - ds4_gpu_tensor_offset(tokens), - NULL, - outbuf, - ds4_gpu_tensor_offset(out), - n_vocab, - n_tokens, - n_embd)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed tokens")) return 0; - } - - return 1; -} - -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 == DS4_METAL_TENSOR_Q8_0) { - return ds4_gpu_embed_token_q8_0_tensor(out, - model_map, - model_size, - weight_offset, - n_vocab, - token, - n_embd); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); - if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal quant embedding received undersized output buffer\n"); - return 0; - } - - uint64_t weight_bytes = 0; - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || - !ds4_gpu_quant_row_bytes(weight_type, n_embd, &row_bytes) || - weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal quant embedding range is outside the mapped model\n"); - return 0; - } - - const uint64_t token_rel = (uint64_t)token * row_bytes; - if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { - fprintf(stderr, "ds4: Metal quant embedding token row is outside the mapped table\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset + token_rel, - row_bytes, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const int32_t token_i32 = 0; - if (!ds4_gpu_encode_get_rows_quant(cb, - wbuf, - (NSUInteger)inner_offset, - weight_type, - nil, - 0, - &token_i32, - outbuf, - ds4_gpu_tensor_offset(out), - n_vocab, - 1, - n_embd)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed token")) return 0; - } - - return 1; -} - -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 == DS4_METAL_TENSOR_Q8_0) { - return ds4_gpu_embed_tokens_q8_0_tensor(out, - tokens, - model_map, - model_size, - weight_offset, - n_vocab, - n_tokens, - n_embd); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out); - id tokbuf = ds4_gpu_tensor_buffer(tokens); - const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); - const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); - if (!outbuf || !tokbuf || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(tokens) < token_bytes) { - fprintf(stderr, "ds4: Metal quant batched embedding received undersized buffers\n"); - return 0; - } - - uint64_t weight_bytes = 0; - if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || - weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal quant batched embedding range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_get_rows_quant(cb, - wbuf, - (NSUInteger)inner_offset, - weight_type, - tokbuf, - ds4_gpu_tensor_offset(tokens), - NULL, - outbuf, - ds4_gpu_tensor_offset(out), - n_vocab, - n_tokens, - n_embd)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed tokens")) return 0; - } - - return 1; -} - -int ds4_gpu_embed_token_hc_tensor( - ds4_gpu_tensor *out_hc, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t token, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0 || n_hc == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out_hc); - const uint64_t out_bytes = (uint64_t)n_embd * n_hc * sizeof(float); - if (!outbuf || ds4_gpu_tensor_bytes(out_hc) < out_bytes) { - fprintf(stderr, "ds4: Metal graph embedding received undersized HC output buffer\n"); - return 0; - } - - const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal graph embedding range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - const NSUInteger row_bytes = (NSUInteger)n_embd * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, - &g_embed_rows_bytes, - row_bytes, - "ds4_embed_rows")) { - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const int32_t token_i32 = (int32_t)token; - const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); - const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); - ds4_gpu_get_rows_args args = { - .ne00t = (int32_t)n_embd, - .ne00 = (int32_t)n_embd, - .nb01 = src_row_bytes, - .nb02 = (uint64_t)n_vocab * src_row_bytes, - .nb03 = (uint64_t)n_vocab * src_row_bytes, - .ne10 = 1, - .nb10 = sizeof(int32_t), - .nb11 = sizeof(int32_t), - .nb12 = sizeof(int32_t), - .nb1 = dst_row_bytes, - .nb2 = dst_row_bytes, - .nb3 = dst_row_bytes, - }; - NSUInteger nth = (NSUInteger)n_embd; - const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_get_rows_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBytes:&token_i32 length:sizeof(token_i32) atIndex:2]; - [enc setBuffer:g_embed_rows_buffer offset:0 atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(nw0, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_encode_repeat_hc_embedding(cb, - g_embed_rows_buffer, - 0, - outbuf, - ds4_gpu_tensor_offset(out_hc), - 1, - n_embd, - n_hc)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed token")) return 0; - } - - return 1; -} - -int ds4_gpu_embed_tokens_hc_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *tokens, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_vocab, - uint32_t n_tokens, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0 || n_hc == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out_hc); - id tokbuf = ds4_gpu_tensor_buffer(tokens); - const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * n_hc * sizeof(float); - const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); - if (!outbuf || !tokbuf || - ds4_gpu_tensor_bytes(out_hc) < out_bytes || - ds4_gpu_tensor_bytes(tokens) < token_bytes) { - fprintf(stderr, "ds4: Metal graph batched embedding received undersized buffers\n"); - return 0; - } - - const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal graph batched embedding range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - const NSUInteger rows_bytes = (NSUInteger)n_tokens * n_embd * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, - &g_embed_rows_bytes, - rows_bytes, - "ds4_embed_rows")) { - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_get_rows_f16(cb, - wbuf, - (NSUInteger)inner_offset, - tokbuf, - ds4_gpu_tensor_offset(tokens), - g_embed_rows_buffer, - 0, - n_vocab, - n_tokens, - n_embd) || - !ds4_gpu_encode_repeat_hc_embedding(cb, - g_embed_rows_buffer, - 0, - outbuf, - ds4_gpu_tensor_offset(out_hc), - n_tokens, - n_embd, - n_hc)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed tokens")) return 0; - } - - return 1; -} - -int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!model_map || model_size == 0) return 0; - if (map_offset > model_size || map_size == 0 || map_size > model_size - map_offset) return 0; - max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, max_tensor_bytes); - - @autoreleasepool { - if (g_model_map_ptr == model_map && - g_model_map_size == model_size && - g_model_mapped_offset == map_offset && - g_model_mapped_size == map_size && - g_model_mapped_max_tensor_bytes == max_tensor_bytes) { - return 1; - } - - for (uint32_t i = 0; i < g_model_view_count; i++) { - if (g_model_views[i].model_map == model_map && - g_model_views[i].model_size == model_size && - map_offset >= g_model_views[i].model_offset && - map_offset + map_size <= g_model_views[i].model_offset + g_model_views[i].bytes) { - return 1; - } - } - - ds4_gpu_model_residency_clear(); - if (!ds4_gpu_map_model_views(model_map, model_size, map_offset, map_size, max_tensor_bytes)) { - ds4_gpu_model_residency_clear(); - return 0; - } - g_model_map_ptr = model_map; - g_model_map_size = model_size; - g_model_mapped_offset = map_offset; - g_model_mapped_size = map_size; - g_model_mapped_max_tensor_bytes = max_tensor_bytes; - if (ds4_gpu_model_map_log_enabled()) { - fprintf(stderr, - "ds4: Metal mapped mmaped model as %u overlapping shared buffers\n", - g_model_view_count); - } - return 1; - } -} - -static int ds4_gpu_model_views_cover_spans( - const void *model_map, - uint64_t model_size, - const uint64_t *offsets, - const uint64_t *sizes, - uint32_t count); - -int ds4_gpu_set_model_map_spans( - const void *model_map, - uint64_t model_size, - const uint64_t *offsets, - const uint64_t *sizes, - uint32_t count, - uint64_t max_tensor_bytes) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0; - if (count == 1) { - return ds4_gpu_set_model_map_range(model_map, - model_size, - offsets[0], - sizes[0], - max_tensor_bytes); - } - if (ds4_gpu_model_views_cover_spans(model_map, model_size, offsets, sizes, count)) { - return 1; - } - - @autoreleasepool { - const double t0 = ds4_gpu_now_ms(); - max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(model_size, max_tensor_bytes); - - ds4_gpu_model_residency_clear(); - ds4_gpu_model_views_clear(); - - uint64_t mapped_total = 0; - uint64_t first_offset = UINT64_MAX; - for (uint32_t i = 0; i < count; i++) { - if (offsets[i] > model_size || sizes[i] == 0 || sizes[i] > model_size - offsets[i]) { - fprintf(stderr, "ds4: Metal model span %u is outside the GGUF mapping\n", i); - ds4_gpu_model_residency_clear(); - ds4_gpu_model_views_clear(); - return 0; - } - if (offsets[i] < first_offset) first_offset = offsets[i]; - uint64_t effective_max = max_tensor_bytes; - if (effective_max > sizes[i]) effective_max = sizes[i]; - if (!ds4_gpu_add_model_view_range(model_map, - model_size, - offsets[i], - sizes[i], - effective_max, - true, - &mapped_total)) { - ds4_gpu_model_residency_clear(); - ds4_gpu_model_views_clear(); - return 0; - } - } - if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) { - ds4_gpu_model_residency_clear(); - ds4_gpu_model_views_clear(); - return 0; - } - g_model_map_ptr = model_map; - g_model_map_size = model_size; - g_model_mapped_offset = first_offset == UINT64_MAX ? 0 : first_offset; - g_model_mapped_size = mapped_total; - g_model_mapped_max_tensor_bytes = max_tensor_bytes; - if (ds4_gpu_model_map_log_enabled()) { - fprintf(stderr, - "ds4: Metal mapped mmaped model as %u disjoint shared buffers across %u tensor spans\n", - g_model_view_count, - count); - } - return 1; - } -} - -int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) { - return ds4_gpu_set_model_map_range(model_map, model_size, 0, model_size, 0); -} - -int ds4_gpu_set_model_fd(int fd) { - g_model_fd = fd; - return 1; -} - -static int ds4_gpu_model_views_cover_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t size) { - if (!model_map || model_size == 0 || size == 0 || - offset > model_size || size > model_size - offset) { - return 0; - } - const uint64_t end = offset + size; - for (uint32_t i = 0; i < g_model_view_count; i++) { - if (g_model_views[i].model_map != model_map || - g_model_views[i].model_size != model_size) { - continue; - } - const uint64_t view_start = g_model_views[i].model_offset; - const uint64_t view_end = view_start + g_model_views[i].bytes; - if (offset >= view_start && end <= view_end) return 1; - } - return 0; -} - -static int ds4_gpu_model_views_cover_spans( - const void *model_map, - uint64_t model_size, - const uint64_t *offsets, - const uint64_t *sizes, - uint32_t count) { - if (!offsets || !sizes || count == 0) return 0; - for (uint32_t i = 0; i < count; i++) { - if (!ds4_gpu_model_views_cover_range(model_map, - model_size, - offsets[i], - sizes[i])) { - return 0; - } - } - return 1; -} - -int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) { - (void)fd; - (void)model_map; - return 1; -} - -static id ds4_gpu_wrap_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset) { - (void)model_map; - if (model_size == 0 || offset > model_size || len > model_size - offset) { - fprintf(stderr, "ds4: Metal model range is outside the mapped model\n"); - return nil; - } - - const uint64_t end = offset + len; - for (uint32_t i = 0; i < g_model_view_count; i++) { - if (g_model_views[i].model_map != model_map || - g_model_views[i].model_size != model_size) { - continue; - } - const uint64_t view_start = g_model_views[i].model_offset; - const uint64_t view_end = view_start + g_model_views[i].bytes; - if (offset >= view_start && end <= view_end) { - *inner_offset = offset - view_start; - return g_model_views[i].buffer; - } - } - - fprintf(stderr, - "ds4: Metal model range %.2f..%.2f GiB is not covered by mapped model views\n", - ds4_gpu_gib(offset), - ds4_gpu_gib(end)); - return nil; -} - -typedef enum { - DS4_GPU_EXACT_VIEW_CACHED, - DS4_GPU_EXACT_VIEW_TRANSIENT, - DS4_GPU_EXACT_VIEW_OWNED, -} ds4_gpu_exact_view_lifetime; - -static id ds4_gpu_wrap_model_exact_range_impl( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset, - ds4_gpu_exact_view_lifetime lifetime) { - const bool cache_view = lifetime == DS4_GPU_EXACT_VIEW_CACHED; - const bool transient_view = lifetime == DS4_GPU_EXACT_VIEW_TRANSIENT; - if (!model_map || !g_device || - (cache_view && !g_model_buffer_cache) || - (transient_view && !g_transient_buffers) || - model_size == 0 || offset > model_size || len > model_size - offset) { - fprintf(stderr, "ds4: Metal exact model range is outside the mapped model\n"); - return nil; - } - - const uint64_t page = (uint64_t)getpagesize(); - const uint64_t page_offset = offset & ~(page - 1); - const uint64_t leading = offset - page_offset; - if (len > UINT64_MAX - leading || - leading + len > UINT64_MAX - (page - 1)) { - fprintf(stderr, "ds4: Metal exact model range overflows page alignment\n"); - return nil; - } - uint64_t view_bytes = round_up_u64(leading + len, page); - if (view_bytes > model_size - page_offset) view_bytes = model_size - page_offset; - if (leading + len > view_bytes) { - fprintf(stderr, "ds4: Metal exact model range alignment exceeds mapped model\n"); - return nil; - } - if (view_bytes > (uint64_t)[g_device maxBufferLength]) { - fprintf(stderr, - "ds4: Metal exact model range %.2f GiB exceeds maxBufferLength %.2f GiB\n", - ds4_gpu_gib(view_bytes), - ds4_gpu_gib((uint64_t)[g_device maxBufferLength])); - return nil; - } - - NSString *key = nil; - id buffer = nil; - if (cache_view) { - key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu", - model_map, - (unsigned long long)model_size, - (unsigned long long)page_offset, - (unsigned long long)view_bytes]; - buffer = [g_model_buffer_cache objectForKey:key]; - } - if (!buffer) { - const uintptr_t base = (uintptr_t)model_map; - buffer = [g_device newBufferWithBytesNoCopy:(void *)(base + page_offset) - length:(NSUInteger)view_bytes - options:ds4_gpu_model_resource_options() - deallocator:nil]; - if (!buffer) { - fprintf(stderr, - "ds4: Metal could not wrap exact mmaped model range at %.2f GiB, size %.2f MiB\n", - ds4_gpu_gib(page_offset), - ds4_gpu_mib(view_bytes)); - return nil; - } - if (cache_view) { - buffer.label = @"ds4_model_exact_view"; - } else if (transient_view) { - buffer.label = @"ds4_model_exact_transient_view"; - } else { - buffer.label = @"ds4_model_exact_owned_view"; - } - if (cache_view) { - [g_model_buffer_cache setObject:buffer forKey:key]; - ds4_gpu_model_buffer_cache_note_insert(view_bytes); - } else if (transient_view) { - [g_transient_buffers addObject:buffer]; - } - } - - if (inner_offset) *inner_offset = leading; - return buffer; -} - -static id ds4_gpu_wrap_model_exact_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset) { - return ds4_gpu_wrap_model_exact_range_impl(model_map, - model_size, - offset, - len, - inner_offset, - DS4_GPU_EXACT_VIEW_CACHED); -} - -static id ds4_gpu_wrap_model_exact_range_transient( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset) { - return ds4_gpu_wrap_model_exact_range_impl(model_map, - model_size, - offset, - len, - inner_offset, - DS4_GPU_EXACT_VIEW_TRANSIENT); -} - -static id ds4_gpu_wrap_model_exact_range_owned( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t *inner_offset) { - return ds4_gpu_wrap_model_exact_range_impl(model_map, - model_size, - offset, - len, - inner_offset, - DS4_GPU_EXACT_VIEW_OWNED); -} - -static id ds4_gpu_wrap_q8_decode_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t n_tokens, - uint64_t *inner_offset) { - const uint64_t exact_decode_max_mib = - ds4_gpu_env_u64("DS4_METAL_Q8_DECODE_EXACT_VIEW_MAX_MIB", - 1024u, - 1u, - 4096u); - const uint64_t exact_decode_max_bytes = - exact_decode_max_mib * 1024ull * 1024ull; - const bool exact_decode_weight_view = - n_tokens == 1u && - len <= exact_decode_max_bytes && - getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; - return exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - offset, - len, - inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - offset, - len, - inner_offset); -} - -static id ds4_gpu_wrap_f32_decode_model_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len, - uint64_t n_tokens, - uint64_t *inner_offset) { - const uint64_t exact_decode_max_mib = - ds4_gpu_env_u64("DS4_METAL_F32_DECODE_EXACT_VIEW_MAX_MIB", - 64u, - 1u, - 4096u); - const uint64_t exact_decode_max_bytes = - exact_decode_max_mib * 1024ull * 1024ull; - const bool exact_decode_weight_view = - n_tokens == 1u && - len <= exact_decode_max_bytes && - getenv("DS4_METAL_ENABLE_F32_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_F32_DECODE_EXACT_VIEWS") == NULL; - return exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - offset, - len, - inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - offset, - len, - inner_offset); -} - -uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { - uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { - budget = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; - } - return budget; -} - -uint32_t ds4_gpu_stream_expert_cache_current_count(void) { - return g_stream_expert_cache_entry_count; -} - -uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes)) { - return 0; - } - return ds4_gpu_stream_expert_cache_configured_budget(); -} - -static int ds4_gpu_stream_expert_cache_note_expert_size( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (gate_expert_bytes == 0 || down_expert_bytes == 0) return 0; - if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { - fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); - return 0; - } - /* - * The cache is a single-size-class slab allocator: the expert byte size is - * frozen on first sight (or pre-seeded at startup from the model's slab - * class) and off-size layers are REJECTED rather than adopted. A rejected - * layer (mixed-precision boost: Q4_K experts among IQ2 layers) falls back - * to the mapped-model per-expert path; last-writer-wins here would instead - * poison the slab size class and deadlock slab reuse. - */ - const uint64_t bytes = gate_expert_bytes * 2ull + down_expert_bytes; - if (g_stream_expert_cache_expert_bytes == 0) { - g_stream_expert_cache_expert_bytes = bytes; - return 1; - } - return bytes == g_stream_expert_cache_expert_bytes; -} - -static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) { - if (!g_ssd_streaming_mode) return 0; - if (g_stream_expert_cache_budget_override != 0) { - return g_stream_expert_cache_budget_override; - } - return 0; -} - -static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) { - return ds4_gpu_stream_expert_cache_requested_budget(); -} - -static uint32_t ds4_gpu_stream_expert_cache_effective_cap( - uint32_t layer, - uint32_t n_total_expert, - uint32_t n_selected) { - (void)layer; - if (ds4_gpu_stream_expert_cache_configured_budget() == 0) return 0; - - /* - * The residency policy is global: every layer can use any expert slot it - * routes to, and global pruning decides which existing entry is least - * valuable. A per-layer cap made cache size depend on model depth rather - * than the actual byte budget. - */ - uint32_t cap = n_total_expert; - if (cap < n_selected) cap = n_selected; - if (cap > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - cap = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - } - return cap; -} - -static int ds4_gpu_stream_expert_timing_summary_enabled(void) { - static int checked = 0; - static int enabled = 0; - if (!checked) { - enabled = - (getenv("DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY") != NULL || - getenv("DS4_METAL_STREAMING_EXPERT_PROFILE_SUMMARY") != NULL) && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_TIMING_SUMMARY") == NULL; - checked = 1; - } - return enabled; -} - -static uint32_t ds4_gpu_stream_expert_popcount(uint32_t mask) { - return (uint32_t)__builtin_popcount(mask); -} - -static int ds4_gpu_stream_expert_split_worthwhile( - uint32_t resident_mask, - uint32_t missing_mask) { - if (resident_mask == 0 || missing_mask == 0) return 0; - /* - * The split path pays an extra command stage and a second routed-expert - * bind. It is worthwhile when several experts are missing and their SSD - * reads can be hidden by resident expert work. With one or two misses, - * especially in large caches, a single unsplit routed pass is faster. - */ - return ds4_gpu_stream_expert_popcount(missing_mask) >= 3u; -} - -static void ds4_gpu_stream_expert_timing_note_selected( - double sync_ms, - double copy_ms, - double bind_ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_selected_calls++; - g_stream_expert_timing_selected_sync_ms += sync_ms; - g_stream_expert_timing_selected_copy_ms += copy_ms; - g_stream_expert_timing_selected_read_ms += sync_ms + copy_ms; - g_stream_expert_timing_selected_bind_ms += bind_ms; -} - -static void ds4_gpu_stream_expert_timing_note_split( - uint32_t resident_mask, - uint32_t missing_mask, - double resident_ms, - double missing_ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_split_layers++; - g_stream_expert_timing_split_resident_experts += - ds4_gpu_stream_expert_popcount(resident_mask); - g_stream_expert_timing_split_missing_experts += - ds4_gpu_stream_expert_popcount(missing_mask); - g_stream_expert_timing_split_resident_ms += resident_ms; - g_stream_expert_timing_split_missing_ms += missing_ms; -} - -static void ds4_gpu_stream_expert_timing_note_split_missing_detail( - double load_ms, - double slot_ms, - double prune_ms, - double addr_ms, - double wait_ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_split_missing_load_ms += load_ms; - g_stream_expert_timing_split_missing_slot_ms += slot_ms; - g_stream_expert_timing_split_missing_prune_ms += prune_ms; - g_stream_expert_timing_split_missing_addr_ms += addr_ms; - g_stream_expert_timing_split_missing_wait_ms += wait_ms; -} - -static void ds4_gpu_stream_expert_timing_note_load_detail( - double prepare_ms, - double pread_ms, - double modify_ms, - double install_ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_load_calls++; - g_stream_expert_timing_load_prepare_ms += prepare_ms; - g_stream_expert_timing_load_pread_ms += pread_ms; - g_stream_expert_timing_load_modify_ms += modify_ms; - g_stream_expert_timing_load_install_ms += install_ms; -} - -static void ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(double ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_prepare_batch_reuse_calls++; - g_stream_expert_timing_prepare_batch_reuse_ms += ms; -} - -static void ds4_gpu_stream_expert_timing_note_prepare_buffer(double ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_prepare_buffer_calls++; - g_stream_expert_timing_prepare_buffer_ms += ms; -} - -static void ds4_gpu_stream_expert_timing_note_prepare_task( - uint32_t experts, - double ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_prepare_task_experts += experts; - g_stream_expert_timing_prepare_task_ms += ms; -} - -static void ds4_gpu_stream_expert_timing_note_reuse_scan( - uint64_t entries, - double ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_reuse_scan_calls++; - g_stream_expert_timing_reuse_scan_entries += entries; - g_stream_expert_timing_reuse_scan_ms += ms; -} - -static void ds4_gpu_stream_expert_timing_note_reuse_clear(double ms) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - g_stream_expert_timing_reuse_clear_ms += ms; -} - -static void ds4_gpu_stream_expert_timing_note_cache_class( - uint32_t resident_mask, - uint32_t missing_mask) { - if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; - const uint32_t resident = ds4_gpu_stream_expert_popcount(resident_mask); - const uint32_t missing = ds4_gpu_stream_expert_popcount(missing_mask); - if (missing == 0) { - g_stream_expert_timing_cache_all_resident_layers++; - } else if (resident == 0) { - g_stream_expert_timing_cache_all_missing_layers++; - } else { - g_stream_expert_timing_cache_mixed_layers++; - } - g_stream_expert_timing_cache_resident_experts += resident; - g_stream_expert_timing_cache_missing_experts += missing; -} - -static int ds4_gpu_stream_expert_readahead_enabled(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_READAHEAD") == NULL; -} - -static void ds4_gpu_stream_expert_readahead_range(uint64_t offset, uint64_t len) { - if (!ds4_gpu_stream_expert_readahead_enabled() || g_model_fd < 0 || len == 0) { - return; - } - -#if defined(F_RDADVISE) - const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); - const double t0 = timing ? ds4_gpu_now_ms() : 0.0; - uint64_t pos = offset; - uint64_t rem = len; - while (rem > 0) { - const uint64_t chunk64 = - rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem; - if (pos > (uint64_t)LLONG_MAX) break; - - struct radvisory ra; - ra.ra_offset = (off_t)pos; - ra.ra_count = (int)chunk64; - (void)fcntl(g_model_fd, F_RDADVISE, &ra); - - pos += chunk64; - rem -= chunk64; - } - if (timing) { - g_stream_expert_timing_readahead_calls++; - g_stream_expert_timing_readahead_bytes += len; - g_stream_expert_timing_readahead_ms += ds4_gpu_now_ms() - t0; - } -#else - (void)offset; - (void)len; -#endif -} - -typedef struct { - uint64_t offset; - uint64_t len; - uint8_t *dst; - uint64_t read_bytes; - double ms; - int ok; -} ds4_gpu_stream_expert_pread_task; - -typedef struct { - int active; - const void *model_map; - uint64_t model_size; - uint32_t layer; - uint32_t n_total_expert; - uint32_t n_selected; - uint32_t missing_mask; - uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t n_loads; - uint32_t n_tasks; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; - int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint64_t gate_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint64_t up_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint64_t down_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; - double start_ms; - double prepare_ms; -} ds4_gpu_stream_expert_pending_load; - -static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load; - -typedef struct { - int active; - const void *model_map; - uint64_t model_size; - uint32_t layer; - uint32_t n_total_expert; - uint32_t n_selected; - uint64_t gate_offset; - uint64_t up_offset; - uint64_t down_offset; - uint64_t gate_expert_bytes; - uint64_t down_expert_bytes; - int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; -} ds4_gpu_glm_stream_selected_prefetch; - -static ds4_gpu_glm_stream_selected_prefetch g_glm_stream_selected_prefetch; - -typedef struct { - ds4_gpu_stream_expert_pread_task *tasks; - uint32_t n_tasks; - uint32_t worker_index; - uint32_t n_workers; -} ds4_gpu_stream_expert_pread_worker_args; - -static void ds4_gpu_stream_expert_cache_note_pread( - uint32_t layer, - uint64_t bytes, - double ms) { - if (g_stream_expert_cache_pread_bytes > UINT64_MAX - bytes) { - g_stream_expert_cache_pread_bytes = UINT64_MAX; - } else { - g_stream_expert_cache_pread_bytes += bytes; - } - g_stream_expert_cache_pread_ms += ms; - if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { - if (g_stream_expert_cache_layer_pread_bytes[layer] > UINT64_MAX - bytes) { - g_stream_expert_cache_layer_pread_bytes[layer] = UINT64_MAX; - } else { - g_stream_expert_cache_layer_pread_bytes[layer] += bytes; - } - g_stream_expert_cache_layer_pread_ms[layer] += ms; - } -} - -static uint32_t ds4_gpu_stream_expert_pread_thread_limit(void) { - uint32_t threads = 9; - const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_THREADS"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0') { - threads = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v; - } - } - if (threads == 0) threads = 1; - if (threads > 18) threads = 18; - return threads; -} - -static uint32_t ds4_gpu_stream_expert_pread_thread_count(uint32_t n_tasks) { - if (n_tasks <= 1) return n_tasks; - uint32_t threads = ds4_gpu_stream_expert_pread_thread_limit(); - if (threads > n_tasks) threads = n_tasks; - return threads; -} - -static int ds4_gpu_stream_expert_pread_into( - uint64_t offset, - uint64_t len, - uint8_t *dst, - uint64_t *read_bytes, - double *ms_out) { - if (read_bytes) *read_bytes = 0; - if (ms_out) *ms_out = 0.0; - if (g_model_fd < 0 || - !dst || - len == 0 || - offset > (uint64_t)LLONG_MAX || - len > (uint64_t)LLONG_MAX - offset) { - return 0; - } - - const double t0 = ds4_gpu_now_ms(); - uint64_t pos = 0; - int ok = 1; - while (pos < len) { - const uint64_t rem = len - pos; - const size_t want = rem > (uint64_t)SSIZE_MAX ? (size_t)SSIZE_MAX : (size_t)rem; - ssize_t nread; - do { - nread = pread(g_model_fd, dst + pos, want, (off_t)(offset + pos)); - } while (nread < 0 && errno == EINTR); - if (nread <= 0) { - ok = 0; - break; - } - pos += (uint64_t)nread; - } - const double dt = ds4_gpu_now_ms() - t0; - if (read_bytes) *read_bytes = pos; - if (ms_out) *ms_out = dt; - if (!ok || pos != len) { - fprintf(stderr, - "ds4: Metal streaming expert explicit pread failed offset=%.2f GiB len=%.2f MiB read=%.2f MiB\n", - ds4_gpu_gib(offset), - ds4_gpu_mib(len), - ds4_gpu_mib(pos)); - return 0; - } - return 1; -} - -static void *ds4_gpu_stream_expert_pread_worker(void *arg) { - ds4_gpu_stream_expert_pread_worker_args *wa = - (ds4_gpu_stream_expert_pread_worker_args *)arg; - for (uint32_t i = wa->worker_index; i < wa->n_tasks; i += wa->n_workers) { - ds4_gpu_stream_expert_pread_task *task = &wa->tasks[i]; - task->ok = ds4_gpu_stream_expert_pread_into(task->offset, - task->len, - task->dst, - &task->read_bytes, - &task->ms); - } - return NULL; -} - -static pthread_mutex_t g_stream_expert_pread_pool_mutex = PTHREAD_MUTEX_INITIALIZER; -static pthread_cond_t g_stream_expert_pread_pool_start_cond = PTHREAD_COND_INITIALIZER; -static pthread_cond_t g_stream_expert_pread_pool_done_cond = PTHREAD_COND_INITIALIZER; -static pthread_t g_stream_expert_pread_pool_threads[18]; -static uint32_t g_stream_expert_pread_pool_thread_count; -static uint32_t g_stream_expert_pread_pool_active_workers; -static uint32_t g_stream_expert_pread_pool_remaining_workers; -static uint32_t g_stream_expert_pread_pool_n_tasks; -static uint32_t g_stream_expert_pread_pool_next_task; -static uint64_t g_stream_expert_pread_pool_generation; -static ds4_gpu_stream_expert_pread_task *g_stream_expert_pread_pool_tasks; -static int g_stream_expert_pread_pool_initialized; -static int g_stream_expert_pread_pool_stopping; - -static int ds4_gpu_stream_expert_pread_pool_enabled(void) { - const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_POOL"); - return !(env && strcmp(env, "0") == 0); -} - -static void *ds4_gpu_stream_expert_pread_pool_worker(void *arg) { - const uint32_t worker_index = (uint32_t)(uintptr_t)arg; - uint64_t seen_generation = 0; - - for (;;) { - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - while (!g_stream_expert_pread_pool_stopping && - g_stream_expert_pread_pool_generation == seen_generation) { - pthread_cond_wait(&g_stream_expert_pread_pool_start_cond, - &g_stream_expert_pread_pool_mutex); - } - if (g_stream_expert_pread_pool_stopping) { - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - break; - } - - seen_generation = g_stream_expert_pread_pool_generation; - if (worker_index >= g_stream_expert_pread_pool_active_workers) { - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - continue; - } - - for (;;) { - const uint32_t task_index = - g_stream_expert_pread_pool_next_task++; - if (task_index >= g_stream_expert_pread_pool_n_tasks) break; - - ds4_gpu_stream_expert_pread_task *task = - &g_stream_expert_pread_pool_tasks[task_index]; - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - - task->ok = ds4_gpu_stream_expert_pread_into(task->offset, - task->len, - task->dst, - &task->read_bytes, - &task->ms); - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - } - - if (g_stream_expert_pread_pool_remaining_workers > 0 && - --g_stream_expert_pread_pool_remaining_workers == 0) { - g_stream_expert_pread_pool_tasks = NULL; - g_stream_expert_pread_pool_n_tasks = 0; - g_stream_expert_pread_pool_active_workers = 0; - pthread_cond_signal(&g_stream_expert_pread_pool_done_cond); - } - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - } - - return NULL; -} - -static int ds4_gpu_stream_expert_pread_pool_init(uint32_t n_threads) { - if (g_stream_expert_pread_pool_initialized) return 1; - if (!ds4_gpu_stream_expert_pread_pool_enabled() || n_threads <= 1) return 0; - if (n_threads > 18) n_threads = 18; - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - g_stream_expert_pread_pool_thread_count = n_threads; - g_stream_expert_pread_pool_stopping = 0; - g_stream_expert_pread_pool_generation = 0; - g_stream_expert_pread_pool_tasks = NULL; - g_stream_expert_pread_pool_n_tasks = 0; - g_stream_expert_pread_pool_next_task = 0; - g_stream_expert_pread_pool_active_workers = 0; - g_stream_expert_pread_pool_remaining_workers = 0; - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - - uint32_t started = 0; - for (uint32_t i = 0; i < n_threads; i++) { - const int rc = pthread_create(&g_stream_expert_pread_pool_threads[i], - NULL, - ds4_gpu_stream_expert_pread_pool_worker, - (void *)(uintptr_t)i); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming expert pread pool thread creation failed: %s\n", - strerror(rc)); - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - g_stream_expert_pread_pool_stopping = 1; - g_stream_expert_pread_pool_generation++; - pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - for (uint32_t j = 0; j < started; j++) { - (void)pthread_join(g_stream_expert_pread_pool_threads[j], NULL); - } - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - g_stream_expert_pread_pool_thread_count = 0; - g_stream_expert_pread_pool_stopping = 0; - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 0; - } - started++; - } - - g_stream_expert_pread_pool_initialized = 1; - return 1; -} - -static int ds4_gpu_stream_expert_pread_pool_begin( - ds4_gpu_stream_expert_pread_task *tasks, - uint32_t n_tasks, - uint32_t n_workers) { - if (n_workers <= 1) return 0; - const uint32_t limit = ds4_gpu_stream_expert_pread_thread_limit(); - if (!ds4_gpu_stream_expert_pread_pool_init(limit)) return 0; - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - if (!g_stream_expert_pread_pool_initialized || - g_stream_expert_pread_pool_stopping || - g_stream_expert_pread_pool_thread_count == 0) { - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 0; - } - if (n_workers > g_stream_expert_pread_pool_thread_count) { - n_workers = g_stream_expert_pread_pool_thread_count; - } - if (n_workers == 0) { - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 0; - } - - if (g_stream_expert_pread_pool_remaining_workers != 0 || - g_stream_expert_pread_pool_tasks != NULL) { - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 0; - } - - g_stream_expert_pread_pool_tasks = tasks; - g_stream_expert_pread_pool_n_tasks = n_tasks; - g_stream_expert_pread_pool_next_task = 0; - g_stream_expert_pread_pool_active_workers = n_workers; - g_stream_expert_pread_pool_remaining_workers = n_workers; - g_stream_expert_pread_pool_generation++; - pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 1; -} - -static int ds4_gpu_stream_expert_pread_pool_wait(void) { - if (!g_stream_expert_pread_pool_initialized) return 0; - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - while (g_stream_expert_pread_pool_remaining_workers != 0) { - pthread_cond_wait(&g_stream_expert_pread_pool_done_cond, - &g_stream_expert_pread_pool_mutex); - } - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - return 1; -} - -static int ds4_gpu_stream_expert_pread_pool_dispatch( - ds4_gpu_stream_expert_pread_task *tasks, - uint32_t n_tasks, - uint32_t n_workers) { - if (!ds4_gpu_stream_expert_pread_pool_begin(tasks, n_tasks, n_workers)) { - return 0; - } - return ds4_gpu_stream_expert_pread_pool_wait(); -} - -static void ds4_gpu_stream_expert_pread_pool_shutdown(void) { - if (!g_stream_expert_pread_pool_initialized) return; - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - g_stream_expert_pread_pool_stopping = 1; - g_stream_expert_pread_pool_generation++; - pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); - - const uint32_t n_threads = g_stream_expert_pread_pool_thread_count; - for (uint32_t i = 0; i < n_threads; i++) { - (void)pthread_join(g_stream_expert_pread_pool_threads[i], NULL); - } - - pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); - g_stream_expert_pread_pool_thread_count = 0; - g_stream_expert_pread_pool_active_workers = 0; - g_stream_expert_pread_pool_remaining_workers = 0; - g_stream_expert_pread_pool_n_tasks = 0; - g_stream_expert_pread_pool_next_task = 0; - g_stream_expert_pread_pool_tasks = NULL; - g_stream_expert_pread_pool_initialized = 0; - g_stream_expert_pread_pool_stopping = 0; - pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); -} - -static int ds4_gpu_stream_expert_pread_tasks( - ds4_gpu_stream_expert_pread_task *tasks, - uint32_t n_tasks, - uint64_t *total_bytes, - double *wall_ms) { - if (total_bytes) *total_bytes = 0; - if (wall_ms) *wall_ms = 0.0; - if (!tasks || n_tasks == 0) return 1; - - const uint32_t n_workers = - ds4_gpu_stream_expert_pread_thread_count(n_tasks); - const double t0 = ds4_gpu_now_ms(); - int ok = 1; - if (n_workers <= 1) { - ds4_gpu_stream_expert_pread_worker_args wa = { - .tasks = tasks, - .n_tasks = n_tasks, - .worker_index = 0, - .n_workers = 1, - }; - (void)ds4_gpu_stream_expert_pread_worker(&wa); - } else if (!ds4_gpu_stream_expert_pread_pool_dispatch(tasks, - n_tasks, - n_workers)) { - pthread_t threads[18]; - ds4_gpu_stream_expert_pread_worker_args args[18]; - uint32_t started = 0; - for (uint32_t i = 0; i < n_workers; i++) { - args[i].tasks = tasks; - args[i].n_tasks = n_tasks; - args[i].worker_index = i; - args[i].n_workers = n_workers; - const int rc = pthread_create(&threads[i], - NULL, - ds4_gpu_stream_expert_pread_worker, - &args[i]); - if (rc != 0) { - fprintf(stderr, - "ds4: Metal streaming expert pread thread creation failed: %s\n", - strerror(rc)); - ok = 0; - break; - } - started++; - } - for (uint32_t i = 0; i < started; i++) { - if (pthread_join(threads[i], NULL) != 0) ok = 0; - } - } - const double dt = ds4_gpu_now_ms() - t0; - - uint64_t bytes = 0; - for (uint32_t i = 0; i < n_tasks; i++) { - if (!tasks[i].ok) ok = 0; - if (bytes > UINT64_MAX - tasks[i].read_bytes) { - bytes = UINT64_MAX; - } else { - bytes += tasks[i].read_bytes; - } - } - if (total_bytes) *total_bytes = bytes; - if (wall_ms) *wall_ms = dt; - return ok; -} - -static id ds4_gpu_stream_expert_alloc_buffer( - uint64_t len, - NSString *label) { - if (!g_device || - len == 0 || - len > (uint64_t)NSUIntegerMax) { - return nil; - } - - id buffer = [g_device newBufferWithLength:(NSUInteger)len - options:MTLResourceStorageModeShared]; - if (!buffer) { - fprintf(stderr, - "ds4: Metal streaming expert explicit buffer allocation failed (%.2f MiB)\n", - ds4_gpu_mib(len)); - return nil; - } - buffer.label = label; - g_stream_expert_cache_buffer_allocs++; - return buffer; -} - -static int ds4_gpu_stream_expert_combined_buffer_enabled(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL; -} - -static int ds4_gpu_stream_expert_slab_enabled(void) { - return ds4_gpu_stream_expert_combined_buffer_enabled() && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_SLABS") == NULL; -} - -/* - * Large PRO caches otherwise create thousands of small shared Metal buffers. - * Slabs keep the buffer object set small while locking pages only for slots - * that actually hold a streamed expert. - */ -static uint64_t ds4_gpu_stream_expert_slab_target_bytes(void) { - const uint64_t mib = 1024ull * 1024ull; - uint64_t target = 4096ull * mib; - const char *env = getenv("DS4_METAL_STREAMING_EXPERT_SLAB_MB"); - if (env && env[0]) { - char *end = NULL; - unsigned long long v = strtoull(env, &end, 10); - if (end != env && *end == '\0' && v != 0) { - target = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; - } - } - return target; -} - -static id ds4_gpu_stream_expert_alloc_slab_buffer( - uint64_t len, - NSString *label) { - if (!g_device || - len == 0 || - len > (uint64_t)NSUIntegerMax) { - return nil; - } - - id buffer = [g_device newBufferWithLength:(NSUInteger)len - options:MTLResourceStorageModeShared]; - if (!buffer) { - fprintf(stderr, - "ds4: Metal streaming expert slab allocation failed (%.2f MiB)\n", - ds4_gpu_mib(len)); - return nil; - } - buffer.label = label; - g_stream_expert_cache_buffer_allocs++; - return buffer; -} - -static int ds4_gpu_stream_expert_slab_slot_range( - uint32_t slot, - uint32_t *slab_index, - uint64_t *slot_base) { - for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { - const uint32_t start = g_stream_expert_cache_slab_start_slot[i]; - const uint32_t count = g_stream_expert_cache_slab_slot_count[i]; - if (slot < start || slot >= start + count) continue; - if (slab_index) *slab_index = i; - if (slot_base) { - *slot_base = - (uint64_t)(slot - start) * g_stream_expert_cache_slab_slot_bytes; - } - return 1; - } - return 0; -} - -static int ds4_gpu_stream_expert_slab_slot_for_buffer( - id buffer, - NSUInteger gate_inner, - uint32_t *slot_out) { - if (!buffer || g_stream_expert_cache_slab_slot_bytes == 0 || !slot_out) { - return 0; - } - for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { - if (g_stream_expert_cache_slabs[i] != buffer) continue; - const uint64_t inner = (uint64_t)gate_inner; - const uint64_t slot_bytes = g_stream_expert_cache_slab_slot_bytes; - if (slot_bytes == 0 || inner % slot_bytes != 0) return 0; - const uint64_t local_slot = inner / slot_bytes; - if (local_slot >= g_stream_expert_cache_slab_slot_count[i]) return 0; - const uint64_t slot = - (uint64_t)g_stream_expert_cache_slab_start_slot[i] + local_slot; - if (slot > UINT32_MAX) return 0; - *slot_out = (uint32_t)slot; - return 1; - } - return 0; -} - -static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) { - if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || - g_stream_expert_cache_free_slot_count >= - DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { - return; - } - g_stream_expert_cache_free_slots[g_stream_expert_cache_free_slot_count++] = - slot; -} - -static int ds4_gpu_stream_expert_slab_slot_buffers( - uint32_t slot, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - __strong id *gate_buf, - __strong id *up_buf, - __strong id *down_buf, - NSUInteger *gate_inner, - NSUInteger *up_inner, - NSUInteger *down_inner) { - uint32_t slab = UINT32_MAX; - uint64_t base = 0; - if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) || - slab >= g_stream_expert_cache_slab_count || - !g_stream_expert_cache_slabs[slab]) { - return 0; - } - if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull || - base > UINT64_MAX - (gate_expert_bytes * 2ull + down_expert_bytes) || - base + gate_expert_bytes * 2ull + down_expert_bytes > - (uint64_t)NSUIntegerMax) { - return 0; - } - id b = g_stream_expert_cache_slabs[slab]; - *gate_buf = b; - *up_buf = b; - *down_buf = b; - *gate_inner = (NSUInteger)base; - *up_inner = (NSUInteger)(base + gate_expert_bytes); - *down_inner = (NSUInteger)(base + gate_expert_bytes * 2ull); - return 1; -} - -static int ds4_gpu_stream_expert_alloc_slab_slot( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - __strong id *gate_buf, - __strong id *up_buf, - __strong id *down_buf, - NSUInteger *gate_inner, - NSUInteger *up_inner, - NSUInteger *down_inner) { - if (!ds4_gpu_stream_expert_slab_enabled() || - !gate_buf || !up_buf || !down_buf || - !gate_inner || !up_inner || !down_inner || - gate_expert_bytes == 0 || down_expert_bytes == 0 || - gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { - return 0; - } - - uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; - if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; - const uint64_t page = (uint64_t)getpagesize(); - if (page != 0) { - slot_bytes = round_up_u64(slot_bytes, page); - if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; - } - if (g_stream_expert_cache_slab_slot_bytes != 0 && - g_stream_expert_cache_slab_slot_bytes != slot_bytes) { - return 0; - } - g_stream_expert_cache_slab_slot_bytes = slot_bytes; - - if (g_stream_expert_cache_free_slot_count != 0) { - const uint32_t slot = - g_stream_expert_cache_free_slots[--g_stream_expert_cache_free_slot_count]; - return ds4_gpu_stream_expert_slab_slot_buffers(slot, - gate_expert_bytes, - down_expert_bytes, - gate_buf, - up_buf, - down_buf, - gate_inner, - up_inner, - down_inner); - } - - uint32_t slab = g_stream_expert_cache_slab_count; - if (slab != 0 && - g_stream_expert_cache_slab_slots_used[slab - 1] < - g_stream_expert_cache_slab_slot_count[slab - 1]) { - slab--; - } else { - if (g_stream_expert_cache_slab_count >= - DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS) { - return 0; - } - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget != 0 && g_stream_expert_cache_slab_total_slots >= budget) { - return 0; - } - uint64_t target = ds4_gpu_stream_expert_slab_target_bytes(); - uint64_t slots64 = target / slot_bytes; - if (slots64 == 0) slots64 = 1; - if (slots64 > UINT32_MAX) slots64 = UINT32_MAX; - uint32_t slots = (uint32_t)slots64; - if (budget != 0) { - const uint32_t remaining = - budget - g_stream_expert_cache_slab_total_slots; - if (slots > remaining) slots = remaining; - } - if (slots == 0) return 0; - id slab_buffer = nil; - while (slots != 0) { - if ((uint64_t)slots <= UINT64_MAX / slot_bytes && - (uint64_t)slots * slot_bytes <= (uint64_t)NSUIntegerMax) { - slab_buffer = - ds4_gpu_stream_expert_alloc_slab_buffer( - (uint64_t)slots * slot_bytes, - @"ds4_stream_expert_slab"); - if (slab_buffer) break; - } - slots /= 2u; - } - if (!slab_buffer || slots == 0) return 0; - - slab = g_stream_expert_cache_slab_count++; - g_stream_expert_cache_slabs[slab] = slab_buffer; - g_stream_expert_cache_slab_start_slot[slab] = - g_stream_expert_cache_slab_total_slots; - g_stream_expert_cache_slab_slot_count[slab] = slots; - g_stream_expert_cache_slab_slots_used[slab] = 0; - g_stream_expert_cache_slab_total_slots += slots; - } - - const uint32_t local_slot = g_stream_expert_cache_slab_slots_used[slab]++; - const uint32_t slot = - g_stream_expert_cache_slab_start_slot[slab] + local_slot; - return ds4_gpu_stream_expert_slab_slot_buffers(slot, - gate_expert_bytes, - down_expert_bytes, - gate_buf, - up_buf, - down_buf, - gate_inner, - up_inner, - down_inner); -} - -static uint64_t ds4_gpu_stream_expert_buffer_object_count( - id gate, - id up, - id down) { - if (!gate || !up || !down) return 0; - if (gate == up && gate == down) return 1; - if (gate == up || gate == down || up == down) return 2; - return 3; -} - -static int ds4_gpu_stream_expert_evict_dontneed_enabled(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_EVICT_DONTNEED") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EVICT_DONTNEED") == NULL; -} - -static void ds4_gpu_stream_expert_evict_dontneed_range( - const void *model_map, - uint64_t model_size, - uint64_t offset, - uint64_t len) { - if (!ds4_gpu_stream_expert_evict_dontneed_enabled() || - !model_map || - model_size == 0 || - offset > model_size || - len == 0 || - len > model_size - offset) { - return; - } - -#if defined(POSIX_MADV_DONTNEED) - const uint64_t page = (uint64_t)getpagesize(); - const uint64_t page_offset = offset & ~(page - 1); - const uint64_t leading = offset - page_offset; - if (len > UINT64_MAX - leading || - leading + len > UINT64_MAX - (page - 1)) { - return; - } - uint64_t advise_bytes = round_up_u64(leading + len, page); - if (advise_bytes > model_size - page_offset) { - advise_bytes = model_size - page_offset; - } - if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) return; - - const uintptr_t base = (uintptr_t)model_map; - if (page_offset > (uint64_t)(UINTPTR_MAX - base)) return; - void *addr = (void *)(base + (uintptr_t)page_offset); - const int rc = posix_madvise(addr, (size_t)advise_bytes, POSIX_MADV_DONTNEED); - if (rc == 0) { - if (g_stream_expert_cache_evict_advise_bytes > UINT64_MAX - advise_bytes) { - g_stream_expert_cache_evict_advise_bytes = UINT64_MAX; - } else { - g_stream_expert_cache_evict_advise_bytes += advise_bytes; - } - } else if (getenv("DS4_METAL_STREAMING_EXPERT_EVICT_DONTNEED_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal streaming expert evict DONTNEED failed offset=%.2f GiB len=%.2f MiB: %s\n", - ds4_gpu_gib(offset), - ds4_gpu_mib(len), - strerror(rc)); - } -#else - (void)model_map; - (void)model_size; - (void)offset; - (void)len; -#endif -} - -static int ds4_gpu_stream_expert_split_requested(void) { - return g_ssd_streaming_mode; -} - -static uint32_t ds4_gpu_stream_expert_split_min_decode_tokens(void) { - return 4; -} - -static uint32_t ds4_gpu_stream_expert_split_min_cached(void) { - uint32_t min_cached = 1024; - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget != 0 && budget < min_cached * 2u) { - min_cached = budget / 2u; - } - return min_cached; -} - -static int ds4_gpu_stream_expert_split_ready(void) { - if (!ds4_gpu_stream_expert_split_requested()) return 0; - if (g_stream_expert_cache_decode_tokens < - ds4_gpu_stream_expert_split_min_decode_tokens()) { - return 0; - } - return g_stream_expert_cache_entry_count >= - ds4_gpu_stream_expert_split_min_cached(); -} - -static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) { - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - g_stream_expert_cache_route_hotness[layer][expert] >>= 1; - } - } -} - -void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { - memset(g_stream_expert_cache_route_hotness, - 0, - sizeof(g_stream_expert_cache_route_hotness)); - g_stream_expert_cache_hotness_decay_token = - g_stream_expert_cache_decode_tokens; -} - -static void ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(void) { - if (g_stream_expert_cache_decode_tokens == 0) return; - if (g_stream_expert_cache_hotness_decay_token == 0) { - g_stream_expert_cache_hotness_decay_token = - g_stream_expert_cache_decode_tokens; - return; - } - while (g_stream_expert_cache_decode_tokens - - g_stream_expert_cache_hotness_decay_token >= - DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS) { - ds4_gpu_stream_expert_cache_decay_route_hotness(); - g_stream_expert_cache_hotness_decay_token += - DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS; - } -} - -static void ds4_gpu_stream_expert_cache_note_route_hotness( - uint32_t layer, - uint32_t expert, - uint32_t amount) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - amount == 0) { - return; - } - uint32_t *hotness = &g_stream_expert_cache_route_hotness[layer][expert]; - if (*hotness > UINT32_MAX - amount) { - *hotness = UINT32_MAX; - } else { - *hotness += amount; - } -} - -static void ds4_gpu_stream_expert_cache_note_selected_hotness( - uint32_t layer, - const int32_t *selected_ids, - uint32_t n_selected) { - if (!selected_ids || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_selected == 0) { - return; - } - ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); - for (uint32_t i = 0; i < n_selected; i++) { - if (selected_ids[i] < 0 || - selected_ids[i] >= - (int32_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - continue; - } - ds4_gpu_stream_expert_cache_note_route_hotness( - layer, - (uint32_t)selected_ids[i], - 1); - } -} - -static void ds4_gpu_stream_expert_cache_note_frequency_hotness( - uint32_t layer, - const uint32_t *frequency, - uint32_t n_total_expert) { - if (!frequency || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { - return; - } - if (n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - n_total_expert = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - } - ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); - for (uint32_t expert = 0; expert < n_total_expert; expert++) { - ds4_gpu_stream_expert_cache_note_route_hotness(layer, - expert, - frequency[expert]); - } -} - -static void ds4_gpu_stream_expert_cache_note_token(uint32_t layer_index) { - if (!g_ssd_streaming_mode || layer_index != 0 || - g_stream_expert_cache_decode_tokens == UINT64_MAX) { - return; - } - g_stream_expert_cache_decode_tokens++; - ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); -} - -static void ds4_gpu_stream_expert_cache_note_decode_token(void) { - if (!g_ssd_streaming_mode || - g_stream_expert_cache_decode_tokens == UINT64_MAX) { - return; - } - g_stream_expert_cache_decode_tokens++; - ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); -} - -static int ds4_gpu_stream_compact_addr_requested(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_ENABLE_STREAMING_COMPACT_ADDR") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_COMPACT_ADDR") == NULL && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; -} - -static int ds4_gpu_stream_expert_addr_table_requested(void) { - return g_ssd_streaming_mode && - (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || - g_stream_prefill_batch_selected_addr_building || - g_glm_stream_expert_addr_table_building || - (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL) || - ds4_gpu_stream_expert_split_requested()) && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; -} - -static int ds4_gpu_stream_expert_addr_table_kernel_requested(void) { - return g_ssd_streaming_mode && - (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || - ds4_gpu_stream_expert_split_ready()) && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; -} - -static int ds4_gpu_stream_expert_masked_addr_requested(void) { - return g_ssd_streaming_mode && - (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || - ds4_gpu_stream_expert_split_ready()) && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_MASKED_ADDR") == NULL; -} - -static int ds4_gpu_stream_expert_hit_validator_requested(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_HIT_VALIDATOR") == NULL; -} - -static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_max( - uint32_t n_total_expert) { - const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX"); - if (env && env[0]) { - char *end = NULL; - const long v = strtol(env, &end, 10); - if (end != env) { - if (v <= 0) return 0; - if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; - return (uint32_t)v; - } - } - if (n_total_expert == 384) return 800u; - if (n_total_expert == 256) return 760u; - return 0; -} - -static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min( - uint32_t n_total_expert) { - const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN"); - if (env && env[0]) { - char *end = NULL; - const long v = strtol(env, &end, 10); - if (end != env) { - if (v <= 0) return 0; - if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; - return (uint32_t)v; - } - } - if (n_total_expert == 384 || n_total_expert == 256) return 2u; - return 0; -} - -static int ds4_gpu_stream_prefill_batch_selected_addr_enabled( - uint32_t n_tokens, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t gate_type, - uint32_t down_type) { - if (!g_ssd_streaming_mode || - n_tokens <= 1 || - n_total_expert == 0 || - n_expert != 6 || - gate_type != DS4_METAL_TENSOR_IQ2_XXS || - down_type != DS4_METAL_TENSOR_Q2_K || - g_quality_mode || - getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) { - return 0; - } - /* All unique experts for one layer must fit simultaneously because the - * address-table kernels consume them in one dispatch. Once the global - * cache fills, preparation reuses entries owned by other layers. */ - if (ds4_gpu_stream_expert_cache_configured_count() < n_total_expert) { - return 0; - } - if (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL) { - return 1; - } - const uint32_t max_tokens = - ds4_gpu_stream_prefill_batch_selected_addr_auto_max(n_total_expert); - const uint32_t min_tokens = - ds4_gpu_stream_prefill_batch_selected_addr_auto_min(n_total_expert); - return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens; -} - -static int ds4_gpu_glm_streaming_prefill_full_layer_active(void) { - return g_glm_streaming_prefill_full_layer_runtime || - getenv("DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER") != NULL; -} - -static int ds4_gpu_stream_full_expert_addr_table_requested(void) { - return g_ssd_streaming_mode && - getenv("DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") == NULL; -} - -static uint64_t ds4_gpu_buffer_address(id buffer, NSUInteger inner) { - if (!buffer) return 0; -#if TARGET_OS_OSX - if (@available(macOS 13.0, *)) { - return (uint64_t)[buffer gpuAddress] + (uint64_t)inner; - } -#endif - return 0; -} - -static int ds4_gpu_stream_compact_addr_ensure_buffers(uint32_t layer) { - if (!g_device || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 0; - - const NSUInteger addr_bytes = 6u * sizeof(uint64_t); - const NSUInteger ids_bytes = 6u * sizeof(int32_t); - for (uint32_t i = 0; i < 4; i++) { - id current = nil; - NSUInteger bytes = addr_bytes; - NSString *label = @"ds4_stream_compact_gate_addresses"; - switch (i) { - case 0: - current = g_stream_compact_gate_addr_buffers[layer]; - label = @"ds4_stream_compact_gate_addresses"; - break; - case 1: - current = g_stream_compact_up_addr_buffers[layer]; - label = @"ds4_stream_compact_up_addresses"; - break; - case 2: - current = g_stream_compact_down_addr_buffers[layer]; - label = @"ds4_stream_compact_down_addresses"; - break; - default: - current = g_stream_compact_selected_buffers[layer]; - bytes = ids_bytes; - label = @"ds4_stream_compact_selected_ids"; - break; - } - if (current) continue; - id b = [g_device newBufferWithLength:bytes - options:MTLResourceStorageModeShared]; - if (!b) { - fprintf(stderr, "ds4: Metal streaming compact address buffer allocation failed\n"); - return 0; - } - b.label = label; - memset([b contents], 0, bytes); - [b didModifyRange:NSMakeRange(0, bytes)]; - switch (i) { - case 0: - g_stream_compact_gate_addr_buffers[layer] = b; - break; - case 1: - g_stream_compact_up_addr_buffers[layer] = b; - break; - case 2: - g_stream_compact_down_addr_buffers[layer] = b; - break; - default: - g_stream_compact_selected_buffers[layer] = b; - break; - } - } - return 1; -} - -static int ds4_gpu_stream_compact_addr_prepare( - uint32_t layer, - ds4_gpu_stream_expert_cache_entry * const entries[6], - uint32_t n_entries, - id *gate_addrs, - id *up_addrs, - id *down_addrs, - id *selected_ids) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - !entries || !gate_addrs || !up_addrs || !down_addrs || !selected_ids || - n_entries == 0 || n_entries > 6) { - return 0; - } - if (!ds4_gpu_stream_compact_addr_ensure_buffers(layer)) return 0; - - uint64_t gate_values[6] = {0, 0, 0, 0, 0, 0}; - uint64_t up_values[6] = {0, 0, 0, 0, 0, 0}; - uint64_t down_values[6] = {0, 0, 0, 0, 0, 0}; - int32_t slot_ids[6] = {0, 1, 2, 3, 4, 5}; - - for (uint32_t i = 0; i < n_entries; i++) { - ds4_gpu_stream_expert_cache_entry *e = entries[i]; - if (!e || !e->gate_buffer || !e->up_buffer || !e->down_buffer) { - return 0; - } - gate_values[i] = ds4_gpu_buffer_address(e->gate_buffer, e->gate_inner); - up_values[i] = ds4_gpu_buffer_address(e->up_buffer, e->up_inner); - down_values[i] = ds4_gpu_buffer_address(e->down_buffer, e->down_inner); - if (gate_values[i] == 0 || up_values[i] == 0 || down_values[i] == 0) { - fprintf(stderr, "ds4: Metal streaming compact address path requires GPU addresses\n"); - return 0; - } - } - - const NSUInteger addr_bytes = 6u * sizeof(uint64_t); - const NSUInteger ids_bytes = 6u * sizeof(int32_t); - id gb = g_stream_compact_gate_addr_buffers[layer]; - id ub = g_stream_compact_up_addr_buffers[layer]; - id db = g_stream_compact_down_addr_buffers[layer]; - id ib = g_stream_compact_selected_buffers[layer]; - memcpy([gb contents], gate_values, addr_bytes); - memcpy([ub contents], up_values, addr_bytes); - memcpy([db contents], down_values, addr_bytes); - memcpy([ib contents], slot_ids, ids_bytes); - [gb didModifyRange:NSMakeRange(0, addr_bytes)]; - [ub didModifyRange:NSMakeRange(0, addr_bytes)]; - [db didModifyRange:NSMakeRange(0, addr_bytes)]; - [ib didModifyRange:NSMakeRange(0, ids_bytes)]; - - *gate_addrs = gb; - *up_addrs = ub; - *down_addrs = db; - *selected_ids = ib; - return 1; -} - -static int ds4_gpu_stream_selected_ids_prepare( - uint32_t layer, - const int32_t *selected_ids, - uint32_t n_selected, - id *selected_buf, - NSUInteger *selected_off) { - if (!g_device || - !selected_ids || - !selected_buf || - !selected_off || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_selected == 0 || - n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED) { - return 0; - } - - const NSUInteger bytes = - (NSUInteger)DS4_METAL_MAX_ROUTED_EXPERT_USED * sizeof(int32_t); - id b = g_stream_selected_id_buffers[layer]; - if (!b) { - b = [g_device newBufferWithLength:bytes - options:MTLResourceStorageModeShared]; - if (!b) { - fprintf(stderr, "ds4: Metal streaming selected-id buffer allocation failed\n"); - return 0; - } - b.label = @"ds4_stream_selected_ids"; - g_stream_selected_id_buffers[layer] = b; - } - - int32_t ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = {0}; - memcpy(ids, selected_ids, (size_t)n_selected * sizeof(ids[0])); - memcpy([b contents], ids, bytes); - [b didModifyRange:NSMakeRange(0, bytes)]; - - *selected_buf = b; - *selected_off = 0; - return 1; -} - -static int ds4_gpu_stream_expert_cache_ensure_addr_buffers(uint32_t layer) { - if (!g_device || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { - return 0; - } - - const NSUInteger bytes = - (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); - id buffers[3] = { - g_stream_expert_cache_gate_addr_buffers[layer], - g_stream_expert_cache_up_addr_buffers[layer], - g_stream_expert_cache_down_addr_buffers[layer], - }; - for (uint32_t i = 0; i < 3; i++) { - if (buffers[i]) continue; - id b = [g_device newBufferWithLength:bytes - options:MTLResourceStorageModeShared]; - if (!b) { - fprintf(stderr, "ds4: Metal streaming expert address table allocation failed\n"); - return 0; - } - b.label = - i == 0 ? @"ds4_stream_expert_gate_addresses" : - (i == 1 ? @"ds4_stream_expert_up_addresses" : - @"ds4_stream_expert_down_addresses"); - memset([b contents], 0, bytes); - [b didModifyRange:NSMakeRange(0, bytes)]; - if (i == 0) { - g_stream_expert_cache_gate_addr_buffers[layer] = b; - } else if (i == 1) { - g_stream_expert_cache_up_addr_buffers[layer] = b; - } else { - g_stream_expert_cache_down_addr_buffers[layer] = b; - } - } - return 1; -} - -static void ds4_gpu_stream_expert_cache_zero_addr_slot(uint32_t layer, uint32_t expert) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - return; - } - - id buffers[3] = { - g_stream_expert_cache_gate_addr_buffers[layer], - g_stream_expert_cache_up_addr_buffers[layer], - g_stream_expert_cache_down_addr_buffers[layer], - }; - const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); - for (uint32_t i = 0; i < 3; i++) { - if (!buffers[i]) continue; - uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); - *addr = 0; - [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; - } -} - -static int ds4_gpu_stream_expert_cache_set_addr_slot_raw( - uint32_t layer, - uint32_t expert, - id gate_buffer, - NSUInteger gate_inner, - id up_buffer, - NSUInteger up_inner, - id down_buffer, - NSUInteger down_inner) { - if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; - - const uint64_t values[3] = { - ds4_gpu_buffer_address(gate_buffer, gate_inner), - ds4_gpu_buffer_address(up_buffer, up_inner), - ds4_gpu_buffer_address(down_buffer, down_inner), - }; - if (values[0] == 0 || values[1] == 0 || values[2] == 0) { - fprintf(stderr, "ds4: Metal streaming expert address table requires GPU addresses\n"); - return 0; - } - - id buffers[3] = { - g_stream_expert_cache_gate_addr_buffers[layer], - g_stream_expert_cache_up_addr_buffers[layer], - g_stream_expert_cache_down_addr_buffers[layer], - }; - const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); - for (uint32_t i = 0; i < 3; i++) { - uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); - *addr = values[i]; - [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; - } - return 1; -} - -static int ds4_gpu_stream_expert_cache_set_addr_slot( - uint32_t layer, - uint32_t expert, - id gate_buffer, - NSUInteger gate_inner, - id up_buffer, - NSUInteger up_inner, - id down_buffer, - NSUInteger down_inner) { - if (!ds4_gpu_stream_expert_addr_table_requested()) return 1; - return ds4_gpu_stream_expert_cache_set_addr_slot_raw(layer, - expert, - gate_buffer, - gate_inner, - up_buffer, - up_inner, - down_buffer, - down_inner); -} - -static int ds4_gpu_stream_expert_cache_addr_buffers( - uint32_t layer, - id *gate, - id *up, - id *down) { - if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; - if (gate) *gate = g_stream_expert_cache_gate_addr_buffers[layer]; - if (up) *up = g_stream_expert_cache_up_addr_buffers[layer]; - if (down) *down = g_stream_expert_cache_down_addr_buffers[layer]; - return g_stream_expert_cache_gate_addr_buffers[layer] && - g_stream_expert_cache_up_addr_buffers[layer] && - g_stream_expert_cache_down_addr_buffers[layer]; -} - -static void ds4_gpu_stream_full_expert_addr_clear_layer(uint32_t layer) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; - ds4_gpu_stream_expert_cache_entry *e = &g_stream_full_expert_addr_entry[layer]; - e->gate_buffer = nil; - e->up_buffer = nil; - e->down_buffer = nil; - e->model_map = NULL; - e->model_size = 0; - e->gate_abs_offset = 0; - e->up_abs_offset = 0; - e->down_abs_offset = 0; - e->gate_expert_bytes = 0; - e->down_expert_bytes = 0; - e->logical_bytes = 0; - e->last_used = 0; - e->use_count = 0; - e->gate_inner = 0; - e->up_inner = 0; - e->down_inner = 0; - e->slab_slot = 0; - e->valid = 0; - e->slab_backed = 0; -} - -static int ds4_gpu_stream_full_expert_addr_table_prepare( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t n_total_expert, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - id *gate_addrs, - id *up_addrs, - id *down_addrs, - ds4_gpu_stream_expert_cache_entry **entry_out) { - if (!ds4_gpu_stream_full_expert_addr_table_requested()) return 0; - if (!model_map || model_size == 0 || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - gate_expert_bytes == 0 || - down_expert_bytes == 0 || - n_total_expert > UINT64_MAX / gate_expert_bytes || - n_total_expert > UINT64_MAX / down_expert_bytes) { - return 0; - } - - const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; - if (gate_abs_offset > model_size || - up_abs_offset > model_size || - down_abs_offset > model_size || - gate_tensor_bytes > model_size - gate_abs_offset || - gate_tensor_bytes > model_size - up_abs_offset || - down_tensor_bytes > model_size - down_abs_offset) { - return 0; - } - - ds4_gpu_stream_expert_cache_entry *entry = - &g_stream_full_expert_addr_entry[layer]; - if (!entry->valid || - entry->model_map != model_map || - entry->model_size != model_size || - entry->gate_abs_offset != gate_abs_offset || - entry->up_abs_offset != up_abs_offset || - entry->down_abs_offset != down_abs_offset || - entry->gate_expert_bytes != gate_expert_bytes || - entry->down_expert_bytes != down_expert_bytes || - !entry->gate_buffer || - !entry->up_buffer || - !entry->down_buffer) { - ds4_gpu_stream_full_expert_addr_clear_layer(layer); - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gate_buf = - ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - gate_abs_offset, - gate_tensor_bytes, - &gate_inner); - id up_buf = - ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - up_abs_offset, - gate_tensor_bytes, - &up_inner); - id down_buf = - ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - down_abs_offset, - down_tensor_bytes, - &down_inner); - if (!gate_buf || !up_buf || !down_buf) return 0; - - entry->gate_buffer = gate_buf; - entry->up_buffer = up_buf; - entry->down_buffer = down_buf; - entry->model_map = model_map; - entry->model_size = model_size; - entry->gate_abs_offset = gate_abs_offset; - entry->up_abs_offset = up_abs_offset; - entry->down_abs_offset = down_abs_offset; - entry->gate_expert_bytes = gate_expert_bytes; - entry->down_expert_bytes = down_expert_bytes; - entry->logical_bytes = gate_tensor_bytes * 2ull + down_tensor_bytes; - entry->gate_inner = (NSUInteger)gate_inner; - entry->up_inner = (NSUInteger)up_inner; - entry->down_inner = (NSUInteger)down_inner; - entry->valid = 1; - - if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { - ds4_gpu_stream_full_expert_addr_clear_layer(layer); - return 0; - } - - const uint64_t gate_base = ds4_gpu_buffer_address(entry->gate_buffer, - entry->gate_inner); - const uint64_t up_base = ds4_gpu_buffer_address(entry->up_buffer, - entry->up_inner); - const uint64_t down_base = ds4_gpu_buffer_address(entry->down_buffer, - entry->down_inner); - if (gate_base == 0 || up_base == 0 || down_base == 0) { - fprintf(stderr, "ds4: Metal full streaming expert address table requires GPU addresses\n"); - ds4_gpu_stream_full_expert_addr_clear_layer(layer); - return 0; - } - - id buffers[3] = { - g_stream_expert_cache_gate_addr_buffers[layer], - g_stream_expert_cache_up_addr_buffers[layer], - g_stream_expert_cache_down_addr_buffers[layer], - }; - uint64_t *gate_addr = (uint64_t *)[buffers[0] contents]; - uint64_t *up_addr = (uint64_t *)[buffers[1] contents]; - uint64_t *down_addr = (uint64_t *)[buffers[2] contents]; - for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { - if (expert < n_total_expert) { - gate_addr[expert] = gate_base + (uint64_t)expert * gate_expert_bytes; - up_addr[expert] = up_base + (uint64_t)expert * gate_expert_bytes; - down_addr[expert] = down_base + (uint64_t)expert * down_expert_bytes; - } else { - gate_addr[expert] = 0; - up_addr[expert] = 0; - down_addr[expert] = 0; - } - } - const NSUInteger bytes = - (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); - for (uint32_t i = 0; i < 3; i++) { - [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; - } - } - - if (gate_addrs) *gate_addrs = g_stream_expert_cache_gate_addr_buffers[layer]; - if (up_addrs) *up_addrs = g_stream_expert_cache_up_addr_buffers[layer]; - if (down_addrs) *down_addrs = g_stream_expert_cache_down_addr_buffers[layer]; - if (entry_out) *entry_out = entry; - return entry->valid && - g_stream_expert_cache_gate_addr_buffers[layer] && - g_stream_expert_cache_up_addr_buffers[layer] && - g_stream_expert_cache_down_addr_buffers[layer]; -} - -static id ds4_gpu_stream_expert_validate_status_buffer(void) { - if (!g_device) return nil; - if (g_stream_expert_validate_status_buffer) { - return g_stream_expert_validate_status_buffer; - } - - const NSUInteger bytes = - (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); - id b = [g_device newBufferWithLength:bytes - options:MTLResourceStorageModeShared]; - if (!b) { - fprintf(stderr, "ds4: Metal streaming expert validator allocation failed\n"); - return nil; - } - b.label = @"ds4_stream_expert_validate_status"; - memset([b contents], 0, bytes); - [b didModifyRange:NSMakeRange(0, bytes)]; - g_stream_expert_validate_status_buffer = b; - return b; -} - -static int ds4_gpu_encode_stream_expert_cache_validate( - id cb, - const ds4_gpu_stream_expert_validate_args *args, - id selected, - NSUInteger selected_off, - id gate_addrs, - id up_addrs, - id down_addrs, - id status) { - if (!cb || !args || !selected || !gate_addrs || !up_addrs || !down_addrs || - !status || !g_moe_stream_expert_cache_validate_pipeline || - args->n_total_expert == 0 || args->n_total_expert > 384 || - args->n_expert == 0 || args->n_expert > 6) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_moe_stream_expert_cache_validate_pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:selected offset:selected_off atIndex:1]; - [enc setBuffer:gate_addrs offset:0 atIndex:2]; - [enc setBuffer:up_addrs offset:0 atIndex:3]; - [enc setBuffer:down_addrs offset:0 atIndex:4]; - [enc setBuffer:status offset:0 atIndex:5]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_stream_expert_cache_validate_selected( - const ds4_gpu_tensor *selected, - id gate_addrs, - id up_addrs, - id down_addrs, - uint32_t n_total_expert, - uint32_t n_expert, - int32_t selected_ids[6], - uint32_t *all_cached, - uint32_t *miss_mask, - uint32_t *invalid_mask) { - if (!selected || !selected_ids || !all_cached || !miss_mask || !invalid_mask) { - return 0; - } - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id status = ds4_gpu_stream_expert_validate_status_buffer(); - if (!selectedbuf || !status) return 0; - - const NSUInteger status_bytes = - (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); - memset([status contents], 0, status_bytes); - [status didModifyRange:NSMakeRange(0, status_bytes)]; - - ds4_gpu_stream_expert_validate_args args = { - .n_total_expert = n_total_expert, - .n_expert = n_expert, - }; - - const int had_batch = g_batch_cb != nil; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (!ds4_gpu_encode_stream_expert_cache_validate(cb, - &args, - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_addrs, - up_addrs, - down_addrs, - status)) { - return 0; - } - - if (had_batch) { - if (ds4_gpu_end_commands() == 0) return 0; - } else if (!ds4_gpu_finish_command_buffer(cb, owned, - "streaming expert cache validator")) { - return 0; - } - - const uint32_t *words = (const uint32_t *)[status contents]; - *all_cached = words[0]; - *miss_mask = words[1]; - *invalid_mask = words[2]; - for (uint32_t i = 0; i < 6; i++) { - selected_ids[i] = (int32_t)words[4 + i]; - } - - if (had_batch && ds4_gpu_begin_commands() == 0) return 0; - return 1; -} - -static void ds4_gpu_stream_expert_cache_clear_entry_internal( - uint32_t layer, - uint32_t expert, - int count_eviction, - int recycle_slab_slot, - ds4_gpu_stream_expert_reusable_buffers *reuse) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - return; - } - if (reuse) { - reuse->gate_buffer = nil; - reuse->up_buffer = nil; - reuse->down_buffer = nil; - reuse->gate_inner = 0; - reuse->up_inner = 0; - reuse->down_inner = 0; - } - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!e->valid) return; - if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { - return; - } - - const uint64_t bytes = e->logical_bytes; - ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, - e->model_size, - e->gate_abs_offset, - e->gate_expert_bytes); - ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, - e->model_size, - e->up_abs_offset, - e->gate_expert_bytes); - ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, - e->model_size, - e->down_abs_offset, - e->down_expert_bytes); - ds4_gpu_stream_expert_cache_zero_addr_slot(layer, expert); - if (reuse) { - reuse->gate_buffer = e->gate_buffer; - reuse->up_buffer = e->up_buffer; - reuse->down_buffer = e->down_buffer; - reuse->gate_inner = e->gate_inner; - reuse->up_inner = e->up_inner; - reuse->down_inner = e->down_inner; - } else if (e->slab_backed && recycle_slab_slot) { - ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot); - } - e->gate_buffer = nil; - e->up_buffer = nil; - e->down_buffer = nil; - e->model_map = NULL; - e->model_size = 0; - e->gate_abs_offset = 0; - e->up_abs_offset = 0; - e->down_abs_offset = 0; - e->gate_expert_bytes = 0; - e->down_expert_bytes = 0; - e->logical_bytes = 0; - e->last_used = 0; - e->use_count = 0; - e->gate_inner = 0; - e->up_inner = 0; - e->down_inner = 0; - e->inflight_seq = 0; - e->slab_slot = 0; - e->valid = 0; - e->slab_backed = 0; - - if (g_stream_expert_cache_layer_count[layer] > 0) { - g_stream_expert_cache_layer_count[layer]--; - } - if (g_stream_expert_cache_entry_count > 0) { - g_stream_expert_cache_entry_count--; - } - if (g_stream_expert_cache_bytes >= bytes) { - g_stream_expert_cache_bytes -= bytes; - } else { - g_stream_expert_cache_bytes = 0; - } - if (count_eviction) { - g_stream_expert_cache_evictions++; - g_stream_expert_cache_layer_evictions[layer]++; - } -} - -static void ds4_gpu_stream_expert_cache_clear_entry( - uint32_t layer, - uint32_t expert, - int count_eviction) { - ds4_gpu_stream_expert_cache_clear_entry_internal(layer, - expert, - count_eviction, - 1, - NULL); -} - -static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) { - ds4_gpu_stream_expert_pending_load_clear(); - g_stream_expert_cache_done_seq = g_stream_expert_cache_cb_seq; - g_stream_expert_cache_batch_seq = 0; - g_stream_expert_cache_owned_seq = 0; - g_stream_expert_cache_pending_max_seq = 0; - for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { - for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { - g_stream_expert_cache[layer][expert].inflight_seq = 0; - ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); - } - ds4_gpu_stream_full_expert_addr_clear_layer(layer); - g_stream_expert_cache_layer_count[layer] = 0; - id buffers[3] = { - g_stream_expert_cache_gate_addr_buffers[layer], - g_stream_expert_cache_up_addr_buffers[layer], - g_stream_expert_cache_down_addr_buffers[layer], - }; - for (uint32_t i = 0; i < 3; i++) { - if (!buffers[i]) continue; - const NSUInteger bytes = - (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); - memset([buffers[i] contents], 0, bytes); - [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; - } - } - g_stream_expert_cache_bytes = 0; - g_stream_expert_cache_entry_count = 0; - for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { - g_stream_expert_cache_slabs[i] = nil; - g_stream_expert_cache_slab_start_slot[i] = 0; - g_stream_expert_cache_slab_slot_count[i] = 0; - g_stream_expert_cache_slab_slots_used[i] = 0; - } - g_stream_expert_cache_slab_count = 0; - g_stream_expert_cache_slab_total_slots = 0; - g_stream_expert_cache_free_slot_count = 0; - g_stream_expert_cache_slab_slot_bytes = 0; - if (reset_stats) { - g_stream_expert_cache_hits = 0; - g_stream_expert_cache_misses = 0; - g_stream_expert_cache_evictions = 0; - g_stream_expert_cache_wraps = 0; - g_stream_expert_cache_clock = 0; - g_stream_expert_cache_evict_advise_bytes = 0; - g_stream_expert_cache_willneed_advise_bytes = 0; - g_stream_expert_cache_pread_bytes = 0; - g_stream_expert_cache_pread_ms = 0.0; - g_stream_expert_cache_buffer_allocs = 0; - g_stream_expert_cache_buffer_reuses = 0; - g_stream_expert_cache_decode_tokens = 0; - g_stream_expert_cache_hotness_decay_token = 0; - memset(g_stream_expert_cache_route_hotness, - 0, - sizeof(g_stream_expert_cache_route_hotness)); - g_stream_expert_timing_selected_calls = 0; - g_stream_expert_timing_selected_read_ms = 0.0; - g_stream_expert_timing_selected_sync_ms = 0.0; - g_stream_expert_timing_selected_copy_ms = 0.0; - g_stream_expert_timing_selected_bind_ms = 0.0; - g_stream_expert_timing_split_layers = 0; - g_stream_expert_timing_split_resident_experts = 0; - g_stream_expert_timing_split_missing_experts = 0; - g_stream_expert_timing_split_resident_ms = 0.0; - g_stream_expert_timing_split_missing_ms = 0.0; - g_stream_expert_timing_split_missing_load_ms = 0.0; - g_stream_expert_timing_split_missing_slot_ms = 0.0; - g_stream_expert_timing_split_missing_prune_ms = 0.0; - g_stream_expert_timing_split_missing_addr_ms = 0.0; - g_stream_expert_timing_split_missing_wait_ms = 0.0; - g_stream_expert_timing_load_calls = 0; - g_stream_expert_timing_load_prepare_ms = 0.0; - g_stream_expert_timing_load_pread_ms = 0.0; - g_stream_expert_timing_load_modify_ms = 0.0; - g_stream_expert_timing_load_install_ms = 0.0; - g_stream_expert_timing_prepare_batch_reuse_calls = 0; - g_stream_expert_timing_prepare_batch_reuse_ms = 0.0; - g_stream_expert_timing_prepare_buffer_calls = 0; - g_stream_expert_timing_prepare_buffer_ms = 0.0; - g_stream_expert_timing_prepare_task_experts = 0; - g_stream_expert_timing_prepare_task_ms = 0.0; - g_stream_expert_timing_reuse_scan_calls = 0; - g_stream_expert_timing_reuse_scan_entries = 0; - g_stream_expert_timing_reuse_scan_ms = 0.0; - g_stream_expert_timing_reuse_clear_ms = 0.0; - g_stream_expert_timing_readahead_calls = 0; - g_stream_expert_timing_readahead_bytes = 0; - g_stream_expert_timing_readahead_ms = 0.0; - g_stream_expert_timing_cache_all_resident_layers = 0; - g_stream_expert_timing_cache_all_missing_layers = 0; - g_stream_expert_timing_cache_mixed_layers = 0; - g_stream_expert_timing_cache_resident_experts = 0; - g_stream_expert_timing_cache_missing_experts = 0; - g_stream_expert_timing_last_report = - (ds4_gpu_stream_expert_timing_snapshot){0}; - memset(g_stream_expert_cache_layer_hits, - 0, - sizeof(g_stream_expert_cache_layer_hits)); - memset(g_stream_expert_cache_layer_misses, - 0, - sizeof(g_stream_expert_cache_layer_misses)); - memset(g_stream_expert_cache_layer_evictions, - 0, - sizeof(g_stream_expert_cache_layer_evictions)); - memset(g_stream_expert_cache_layer_pread_bytes, - 0, - sizeof(g_stream_expert_cache_layer_pread_bytes)); - memset(g_stream_expert_cache_layer_pread_ms, - 0, - sizeof(g_stream_expert_cache_layer_pread_ms)); - } -} - -static int ds4_gpu_stream_expert_cache_is_protected( - uint32_t expert, - const int32_t *protect_ids, - uint32_t n_protect) { - if (!protect_ids) return 0; - for (uint32_t i = 0; i < n_protect; i++) { - if (protect_ids[i] >= 0 && (uint32_t)protect_ids[i] == expert) { - return 1; - } - } - return 0; -} - -static void ds4_gpu_stream_expert_cache_prune_layer( - uint32_t layer, - uint32_t n_total_expert, - uint32_t n_selected, - const int32_t *protect_ids, - uint32_t n_protect) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; - uint32_t cap = ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected); - if (cap == 0) return; - - /* - * Route hotness counts selected experts even when they miss. Hit-count - * LFU penalizes experts that are repeatedly selected but evicted before a - * second hit, which keeps too many decode layers in the mixed-cache path. - */ - while (g_stream_expert_cache_layer_count[layer] > cap) { - uint32_t victim = UINT32_MAX; - uint32_t lowest_hotness = UINT32_MAX; - uint64_t oldest = UINT64_MAX; - for (uint32_t expert = 0; expert < n_total_expert; expert++) { - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!e->valid || - ds4_gpu_stream_expert_cache_entry_inflight(e) || - ds4_gpu_stream_expert_cache_is_protected(expert, protect_ids, n_protect)) { - continue; - } - const uint32_t hotness = - g_stream_expert_cache_route_hotness[layer][expert]; - if (hotness < lowest_hotness || - (hotness == lowest_hotness && e->last_used < oldest)) { - lowest_hotness = hotness; - oldest = e->last_used; - victim = expert; - } - } - if (victim == UINT32_MAX) break; - ds4_gpu_stream_expert_cache_clear_entry(layer, victim, 1); - } -} - -static int ds4_gpu_stream_expert_cache_entry_protected( - uint32_t layer, - uint32_t expert, - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect) { - if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && - ds4_gpu_stream_expert_cache_entry_inflight( - &g_stream_expert_cache[layer][expert])) { - return 1; - } - return layer == protect_layer && - ds4_gpu_stream_expert_cache_is_protected(expert, - protect_ids, - n_protect); -} - -static int ds4_gpu_stream_expert_cache_entry_reusable( - const ds4_gpu_stream_expert_cache_entry *e, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - return e && - e->valid && - e->gate_buffer && - e->up_buffer && - e->down_buffer && - e->gate_expert_bytes == gate_expert_bytes && - e->down_expert_bytes == down_expert_bytes; -} - -static int ds4_gpu_stream_expert_cache_take_reusable( - int force_reuse, - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - ds4_gpu_stream_expert_reusable_buffers *reuse) { - if (!reuse) return 0; - reuse->gate_buffer = nil; - reuse->up_buffer = nil; - reuse->down_buffer = nil; - reuse->gate_inner = 0; - reuse->up_inner = 0; - reuse->down_inner = 0; - - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget == 0 || - (!force_reuse && g_stream_expert_cache_entry_count < budget)) { - return 0; - } - - int waited_inflight = 0; -retry: - ; - uint32_t victim_layer = UINT32_MAX; - uint32_t victim_expert = UINT32_MAX; - uint32_t lowest_hotness = UINT32_MAX; - uint64_t oldest = UINT64_MAX; - int skipped_inflight = 0; - const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); - const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; - uint64_t scan_entries = 0; - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - scan_entries++; - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!ds4_gpu_stream_expert_cache_entry_reusable(e, - gate_expert_bytes, - down_expert_bytes)) { - continue; - } - if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { - skipped_inflight = 1; - continue; - } - if (ds4_gpu_stream_expert_cache_entry_protected(layer, - expert, - protect_layer, - protect_ids, - n_protect)) { - continue; - } - const uint32_t hotness = - g_stream_expert_cache_route_hotness[layer][expert]; - if (hotness < lowest_hotness || - (hotness == lowest_hotness && e->last_used < oldest)) { - lowest_hotness = hotness; - oldest = e->last_used; - victim_layer = layer; - victim_expert = expert; - } - } - } - if (timing) { - ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, - ds4_gpu_now_ms() - scan_t0); - } - - if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) { - if (skipped_inflight && !waited_inflight && - !ds4_gpu_stream_expert_cache_on_service_thread()) { - waited_inflight = 1; - if (!ds4_gpu_stream_expert_cache_wait_inflight( - "streaming expert cache reuse")) { - return 0; - } - goto retry; - } - return 0; - } - const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; - ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer, - victim_expert, - 1, - 1, - reuse); - if (timing) { - ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - - clear_t0); - } - if (!reuse->gate_buffer || !reuse->up_buffer || !reuse->down_buffer) { - reuse->gate_buffer = nil; - reuse->up_buffer = nil; - reuse->down_buffer = nil; - reuse->gate_inner = 0; - reuse->up_inner = 0; - reuse->down_inner = 0; - return 0; - } - g_stream_expert_cache_buffer_reuses += - ds4_gpu_stream_expert_buffer_object_count(reuse->gate_buffer, - reuse->up_buffer, - reuse->down_buffer); - return 1; -} - -static uint32_t ds4_gpu_stream_expert_cache_take_reusable_batch( - uint32_t n_needed, - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - ds4_gpu_stream_expert_reusable_buffers *reuses) { - if (!reuses || n_needed == 0) return 0; - if (n_needed > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { - n_needed = DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; - } - for (uint32_t i = 0; i < n_needed; i++) { - reuses[i].gate_buffer = nil; - reuses[i].up_buffer = nil; - reuses[i].down_buffer = nil; - reuses[i].gate_inner = 0; - reuses[i].up_inner = 0; - reuses[i].down_inner = 0; - } - - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget == 0 || g_stream_expert_cache_entry_count < budget) { - return 0; - } - - int waited_inflight = 0; -retry: - ; - uint32_t victim_layers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t victim_experts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t victim_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint64_t victim_last_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t victim_count = 0; - int skipped_inflight = 0; - const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); - const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; - uint64_t scan_entries = 0; - for (uint32_t i = 0; i < n_needed; i++) { - victim_layers[i] = UINT32_MAX; - victim_experts[i] = UINT32_MAX; - victim_hotness[i] = UINT32_MAX; - victim_last_used[i] = UINT64_MAX; - } - - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - scan_entries++; - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!ds4_gpu_stream_expert_cache_entry_reusable(e, - gate_expert_bytes, - down_expert_bytes)) { - continue; - } - if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { - skipped_inflight = 1; - continue; - } - if (ds4_gpu_stream_expert_cache_entry_protected(layer, - expert, - protect_layer, - protect_ids, - n_protect)) { - continue; - } - - const uint32_t hotness = - g_stream_expert_cache_route_hotness[layer][expert]; - const uint64_t last_used = e->last_used; - if (victim_count < n_needed) { - victim_layers[victim_count] = layer; - victim_experts[victim_count] = expert; - victim_hotness[victim_count] = hotness; - victim_last_used[victim_count] = last_used; - victim_count++; - continue; - } - - uint32_t worst = 0; - for (uint32_t i = 1; i < victim_count; i++) { - if (victim_hotness[i] > victim_hotness[worst] || - (victim_hotness[i] == victim_hotness[worst] && - victim_last_used[i] > victim_last_used[worst])) { - worst = i; - } - } - if (hotness < victim_hotness[worst] || - (hotness == victim_hotness[worst] && - last_used < victim_last_used[worst])) { - victim_layers[worst] = layer; - victim_experts[worst] = expert; - victim_hotness[worst] = hotness; - victim_last_used[worst] = last_used; - } - } - } - if (timing) { - ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, - ds4_gpu_now_ms() - scan_t0); - } - - if (victim_count == 0) { - if (skipped_inflight && !waited_inflight) { - waited_inflight = 1; - if (!ds4_gpu_stream_expert_cache_wait_inflight( - "streaming expert cache batch reuse")) { - return 0; - } - goto retry; - } - return 0; - } - - uint32_t reuse_count = 0; - const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; - for (uint32_t i = 0; i < victim_count; i++) { - if (victim_layers[i] == UINT32_MAX || - victim_experts[i] == UINT32_MAX) { - continue; - } - ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layers[i], - victim_experts[i], - 1, - 1, - &reuses[reuse_count]); - if (!reuses[reuse_count].gate_buffer || - !reuses[reuse_count].up_buffer || - !reuses[reuse_count].down_buffer) { - reuses[reuse_count].gate_buffer = nil; - reuses[reuse_count].up_buffer = nil; - reuses[reuse_count].down_buffer = nil; - reuses[reuse_count].gate_inner = 0; - reuses[reuse_count].up_inner = 0; - reuses[reuse_count].down_inner = 0; - continue; - } - g_stream_expert_cache_buffer_reuses += - ds4_gpu_stream_expert_buffer_object_count( - reuses[reuse_count].gate_buffer, - reuses[reuse_count].up_buffer, - reuses[reuse_count].down_buffer); - reuse_count++; - } - if (timing) { - ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - - clear_t0); - } - return reuse_count; -} - -static int ds4_gpu_stream_expert_batch_reuse_enabled( - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { - return 0; - } - const uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; - /* - * One global victim scan per selected miss is measurable for GLM Q2-size - * slots, but batching larger Q4-size slots regressed short decode on M5. - * Keep larger slots on the older single-victim path until profiling says - * otherwise. - */ - return slot_bytes <= 16ull * 1024ull * 1024ull; -} - -static int ds4_gpu_stream_expert_cache_prepare_load_buffers( - uint32_t layer, - uint32_t expert, - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - int force_reuse, - __strong id *gate_buf, - __strong id *up_buf, - __strong id *down_buf, - NSUInteger *gate_inner, - NSUInteger *up_inner, - NSUInteger *down_inner) { - if (!gate_buf || !up_buf || !down_buf || - !gate_inner || !up_inner || !down_inner || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - return 0; - } - - *gate_buf = nil; - *up_buf = nil; - *down_buf = nil; - *gate_inner = 0; - *up_inner = 0; - *down_inner = 0; - - ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 }; - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (e->valid && ds4_gpu_stream_expert_cache_entry_inflight(e)) { - if (ds4_gpu_stream_expert_cache_on_service_thread()) return 0; - if (!ds4_gpu_stream_expert_cache_wait_inflight( - "streaming expert cache replacement")) { - return 0; - } - e = &g_stream_expert_cache[layer][expert]; - } - if (ds4_gpu_stream_expert_cache_entry_reusable(e, - gate_expert_bytes, - down_expert_bytes)) { - ds4_gpu_stream_expert_cache_clear_entry_internal(layer, - expert, - 0, - 1, - &reuse); - if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { - g_stream_expert_cache_buffer_reuses += - ds4_gpu_stream_expert_buffer_object_count(reuse.gate_buffer, - reuse.up_buffer, - reuse.down_buffer); - } - } else if (e->valid) { - ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); - } - - if (!reuse.gate_buffer || !reuse.up_buffer || !reuse.down_buffer) { - if (!ds4_gpu_stream_expert_cache_take_reusable(force_reuse, - protect_layer, - protect_ids, - n_protect, - gate_expert_bytes, - down_expert_bytes, - &reuse)) { - reuse.gate_buffer = nil; - reuse.up_buffer = nil; - reuse.down_buffer = nil; - } - } - - if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { - *gate_buf = reuse.gate_buffer; - *up_buf = reuse.up_buffer; - *down_buf = reuse.down_buffer; - *gate_inner = reuse.gate_inner; - *up_inner = reuse.up_inner; - *down_inner = reuse.down_inner; - return 1; - } - - if (ds4_gpu_stream_expert_combined_buffer_enabled()) { - if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes || - gate_expert_bytes * 2ull > UINT64_MAX - down_expert_bytes || - gate_expert_bytes > (uint64_t)NSUIntegerMax || - gate_expert_bytes * 2ull > (uint64_t)NSUIntegerMax || - gate_expert_bytes * 2ull + down_expert_bytes > - (uint64_t)NSUIntegerMax) { - return 0; - } - if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes, - down_expert_bytes, - gate_buf, - up_buf, - down_buf, - gate_inner, - up_inner, - down_inner)) { - return 1; - } - const uint64_t up_off = gate_expert_bytes; - const uint64_t down_off = gate_expert_bytes * 2ull; - const uint64_t combined_bytes = down_off + down_expert_bytes; - id combined = - ds4_gpu_stream_expert_alloc_buffer(combined_bytes, - @"ds4_stream_expert_combined"); - if (!combined) return 0; - *gate_buf = combined; - *up_buf = combined; - *down_buf = combined; - *gate_inner = 0; - *up_inner = (NSUInteger)up_off; - *down_inner = (NSUInteger)down_off; - return 1; - } - - *gate_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, - @"ds4_stream_expert_gate"); - *up_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, - @"ds4_stream_expert_up"); - *down_buf = ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes, - @"ds4_stream_expert_down"); - return *gate_buf && *up_buf && *down_buf; -} - -static void ds4_gpu_stream_expert_cache_prune_global( - uint32_t protect_layer, - const int32_t *protect_ids, - uint32_t n_protect) { - const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); - if (budget == 0 || g_stream_expert_cache_entry_count <= budget) return; - - while (g_stream_expert_cache_entry_count > budget) { - uint32_t victim_layer = UINT32_MAX; - uint32_t victim_expert = UINT32_MAX; - uint32_t lowest_hotness = UINT32_MAX; - uint64_t oldest = UINT64_MAX; - for (uint32_t layer = 0; - layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; - layer++) { - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!e->valid || - ds4_gpu_stream_expert_cache_entry_protected(layer, - expert, - protect_layer, - protect_ids, - n_protect)) { - continue; - } - const uint32_t hotness = - g_stream_expert_cache_route_hotness[layer][expert]; - if (hotness < lowest_hotness || - (hotness == lowest_hotness && e->last_used < oldest)) { - lowest_hotness = hotness; - oldest = e->last_used; - victim_layer = layer; - victim_expert = expert; - } - } - } - if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) break; - ds4_gpu_stream_expert_cache_clear_entry(victim_layer, victim_expert, 1); - } -} - -static int ds4_gpu_stream_expert_cache_entry_matches( - const ds4_gpu_stream_expert_cache_entry *e, - const void *model_map, - uint64_t model_size, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - return e && - e->valid && - e->model_map == model_map && - e->model_size == model_size && - e->gate_abs_offset == gate_abs_offset && - e->up_abs_offset == up_abs_offset && - e->down_abs_offset == down_abs_offset && - e->gate_expert_bytes == gate_expert_bytes && - e->down_expert_bytes == down_expert_bytes && - e->gate_buffer && e->up_buffer && e->down_buffer; -} - -static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_peek( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t expert, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - if (!g_ssd_streaming_mode || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - expert >= n_total_expert || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected) == 0) { - return NULL; - } - - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (!ds4_gpu_stream_expert_cache_entry_matches(e, - model_map, - model_size, - gate_abs_offset, - up_abs_offset, - down_abs_offset, - gate_expert_bytes, - down_expert_bytes)) { - return NULL; - } - - e->last_used = ++g_stream_expert_cache_clock; - e->use_count++; - g_stream_expert_cache_hits++; - g_stream_expert_cache_layer_hits[layer]++; - return e; -} - -static ds4_gpu_stream_expert_cache_entry * -ds4_gpu_stream_expert_cache_install_loaded( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t expert, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - id gate_buf, - id up_buf, - id down_buf, - NSUInteger gate_inner, - NSUInteger up_inner, - NSUInteger down_inner) { - if (!gate_buf || !up_buf || !down_buf || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { - return NULL; - } - if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { - fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); - return NULL; - } - const uint64_t logical_bytes = gate_expert_bytes * 2ull + down_expert_bytes; - - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (e->valid) { - if (ds4_gpu_stream_expert_cache_entry_inflight(e) && - ds4_gpu_stream_expert_cache_on_service_thread()) { - return NULL; - } - if (ds4_gpu_stream_expert_cache_entry_inflight(e) && - !ds4_gpu_stream_expert_cache_wait_inflight( - "streaming expert cache install")) { - return NULL; - } - if (ds4_gpu_stream_expert_cache_entry_inflight(e)) return NULL; - ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); - if (e->valid) return NULL; - } - - if (!ds4_gpu_stream_expert_cache_set_addr_slot(layer, - expert, - gate_buf, - gate_inner, - up_buf, - up_inner, - down_buf, - down_inner)) { - return NULL; - } - - e->gate_buffer = gate_buf; - e->up_buffer = up_buf; - e->down_buffer = down_buf; - e->model_map = model_map; - e->model_size = model_size; - e->gate_abs_offset = gate_abs_offset; - e->up_abs_offset = up_abs_offset; - e->down_abs_offset = down_abs_offset; - e->gate_expert_bytes = gate_expert_bytes; - e->down_expert_bytes = down_expert_bytes; - e->logical_bytes = logical_bytes; - e->last_used = ++g_stream_expert_cache_clock; - e->use_count = 1; - e->gate_inner = gate_inner; - e->up_inner = up_inner; - e->down_inner = down_inner; - e->inflight_seq = 0; - uint32_t slab_slot = 0; - if (gate_buf == up_buf && - gate_buf == down_buf && - ds4_gpu_stream_expert_slab_slot_for_buffer(gate_buf, - gate_inner, - &slab_slot)) { - e->slab_backed = 1; - e->slab_slot = slab_slot; - } else { - e->slab_backed = 0; - e->slab_slot = 0; - } - e->valid = 1; - g_stream_expert_cache_layer_count[layer]++; - if (g_stream_expert_cache_entry_count < UINT32_MAX) { - g_stream_expert_cache_entry_count++; - } - if (g_stream_expert_cache_bytes > UINT64_MAX - logical_bytes) { - g_stream_expert_cache_bytes = UINT64_MAX; - } else { - g_stream_expert_cache_bytes += logical_bytes; - } - g_stream_expert_cache_misses++; - g_stream_expert_cache_layer_misses[layer]++; - g_stream_expert_cache_wraps += 3; - return e; -} - -static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get_protected( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t expert, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - const int32_t *protect_ids, - uint32_t n_protect) { - if (!g_ssd_streaming_mode || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - expert >= n_total_expert || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected) == 0) { - return NULL; - } - - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][expert]; - if (ds4_gpu_stream_expert_cache_entry_matches(e, - model_map, - model_size, - gate_abs_offset, - up_abs_offset, - down_abs_offset, - gate_expert_bytes, - down_expert_bytes)) { - e->last_used = ++g_stream_expert_cache_clock; - e->use_count++; - g_stream_expert_cache_hits++; - g_stream_expert_cache_layer_hits[layer]++; - return e; - } - - ds4_gpu_stream_expert_readahead_range(gate_abs_offset, gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(up_abs_offset, gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(down_abs_offset, down_expert_bytes); - - id gate_buf = nil; - id up_buf = nil; - id down_buf = nil; - NSUInteger gate_inner = 0; - NSUInteger up_inner = 0; - NSUInteger down_inner = 0; - const int32_t protect_one = (int32_t)expert; - if (!protect_ids || n_protect == 0) { - protect_ids = &protect_one; - n_protect = 1; - } - const uint32_t cache_budget = - ds4_gpu_stream_expert_cache_configured_budget(); - const int force_reuse = - cache_budget != 0 && g_stream_expert_cache_entry_count >= cache_budget; - if (!ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, - expert, - layer, - protect_ids, - n_protect, - gate_expert_bytes, - down_expert_bytes, - force_reuse, - &gate_buf, - &up_buf, - &down_buf, - &gate_inner, - &up_inner, - &down_inner)) { - return NULL; - } - if (!gate_buf || !up_buf || !down_buf) return NULL; - - uint8_t *gate_dst = (uint8_t *)[gate_buf contents] + gate_inner; - uint8_t *up_dst = (uint8_t *)[up_buf contents] + up_inner; - uint8_t *down_dst = (uint8_t *)[down_buf contents] + down_inner; - if (!gate_dst || !up_dst || !down_dst) return NULL; - - ds4_gpu_stream_expert_pread_task tasks[3] = { - { - .offset = gate_abs_offset, - .len = gate_expert_bytes, - .dst = gate_dst, - }, - { - .offset = up_abs_offset, - .len = gate_expert_bytes, - .dst = up_dst, - }, - { - .offset = down_abs_offset, - .len = down_expert_bytes, - .dst = down_dst, - }, - }; - uint64_t read_bytes = 0; - double read_ms = 0.0; - if (!ds4_gpu_stream_expert_pread_tasks(tasks, 3, &read_bytes, &read_ms)) { - return NULL; - } - ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); - - [gate_buf didModifyRange:NSMakeRange(gate_inner, (NSUInteger)gate_expert_bytes)]; - [up_buf didModifyRange:NSMakeRange(up_inner, (NSUInteger)gate_expert_bytes)]; - [down_buf didModifyRange:NSMakeRange(down_inner, (NSUInteger)down_expert_bytes)]; - if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal streaming expert parallel pread layer=%u experts=1 tensors=3 " - "threads=%u bytes=%.2f GiB wall=%.3f ms\n", - layer, - ds4_gpu_stream_expert_pread_thread_count(3), - ds4_gpu_gib(read_bytes), - read_ms); - } - return ds4_gpu_stream_expert_cache_install_loaded(model_map, - model_size, - layer, - expert, - gate_abs_offset, - up_abs_offset, - down_abs_offset, - gate_expert_bytes, - down_expert_bytes, - gate_buf, - up_buf, - down_buf, - gate_inner, - up_inner, - down_inner); -} - -static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t expert, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_abs_offset, - uint64_t up_abs_offset, - uint64_t down_abs_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - return ds4_gpu_stream_expert_cache_get_protected(model_map, - model_size, - layer, - expert, - n_total_expert, - n_selected, - gate_abs_offset, - up_abs_offset, - down_abs_offset, - gate_expert_bytes, - down_expert_bytes, - NULL, - 0); -} - -static int ds4_gpu_stream_expert_pending_load_profile_enabled(void) { - return getenv("DS4_METAL_STREAMING_EXPERT_EARLY_LOAD_PROFILE") != NULL; -} - -static void ds4_gpu_stream_expert_pending_load_release_buffers( - ds4_gpu_stream_expert_pending_load *p) { - if (!p) return; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - p->gate_bufs[i] = nil; - p->up_bufs[i] = nil; - p->down_bufs[i] = nil; - p->gate_inners[i] = 0; - p->up_inners[i] = 0; - p->down_inners[i] = 0; - } -} - -static int ds4_gpu_stream_expert_pending_load_install( - ds4_gpu_stream_expert_pending_load *p, - ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED], - double elapsed_ms) { - if (!p || p->n_loads == 0) return 1; - - uint64_t read_bytes = 0; - int ok = 1; - for (uint32_t i = 0; i < p->n_tasks; i++) { - if (!p->tasks[i].ok) ok = 0; - if (read_bytes > UINT64_MAX - p->tasks[i].read_bytes) { - read_bytes = UINT64_MAX; - } else { - read_bytes += p->tasks[i].read_bytes; - } - } - if (!ok) return 0; - - ds4_gpu_stream_expert_cache_note_pread(p->layer, read_bytes, elapsed_ms); - const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); - double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - double load_modify_ms = 0.0; - double load_install_ms = 0.0; - for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { - [p->gate_bufs[load_i] didModifyRange:NSMakeRange(p->gate_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; - [p->up_bufs[load_i] didModifyRange:NSMakeRange(p->up_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; - [p->down_bufs[load_i] didModifyRange:NSMakeRange(p->down_inners[load_i], (NSUInteger)p->down_expert_bytes)]; - } - if (load_timing) { - const double now_ms = ds4_gpu_now_ms(); - load_modify_ms = now_ms - load_t0; - load_t0 = now_ms; - } - - ds4_gpu_stream_expert_cache_entry - *loaded_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - loaded_entries[i] = NULL; - } - for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { - const uint32_t slot = p->load_slots[load_i]; - const uint32_t expert = (uint32_t)p->selected_ids[slot]; - ds4_gpu_stream_expert_cache_entry *entry = - ds4_gpu_stream_expert_cache_install_loaded(p->model_map, - p->model_size, - p->layer, - expert, - p->gate_abs_offsets[slot], - p->up_abs_offsets[slot], - p->down_abs_offsets[slot], - p->gate_expert_bytes, - p->down_expert_bytes, - p->gate_bufs[load_i], - p->up_bufs[load_i], - p->down_bufs[load_i], - p->gate_inners[load_i], - p->up_inners[load_i], - p->down_inners[load_i]); - if (!entry) return 0; - loaded_entries[slot] = entry; - if (entries) entries[slot] = entry; - } - for (uint32_t i = 0; i < p->n_selected; i++) { - if ((p->missing_mask & (1u << i)) == 0) continue; - if (entries && entries[i]) continue; - const uint32_t source = p->source_slots[i]; - if (source >= p->n_selected) return 0; - ds4_gpu_stream_expert_cache_entry *entry = entries && entries[source] ? - entries[source] : loaded_entries[source]; - if (!entry) return 0; - entry->use_count++; - if (entries) entries[i] = entry; - } - if (load_timing) { - load_install_ms = ds4_gpu_now_ms() - load_t0; - ds4_gpu_stream_expert_timing_note_load_detail(p->prepare_ms, - elapsed_ms, - load_modify_ms, - load_install_ms); - } - if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { - fprintf(stderr, - "ds4: Metal streaming expert early-load finish layer=%u experts=%u tensors=%u bytes=%.2f GiB wall=%.3f ms\n", - p->layer, - p->n_loads, - p->n_tasks, - ds4_gpu_gib(read_bytes), - elapsed_ms); - } - return 1; -} - -static int ds4_gpu_stream_expert_pending_load_finish( - ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]) { - ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; - if (!p->active) return 1; - - const double start_ms = p->start_ms; - if (!ds4_gpu_stream_expert_pread_pool_wait()) { - ds4_gpu_stream_expert_pending_load_release_buffers(p); - p->active = 0; - return 0; - } - const double elapsed_ms = ds4_gpu_now_ms() - start_ms; - p->active = 0; - const int ok = ds4_gpu_stream_expert_pending_load_install(p, - entries, - elapsed_ms); - ds4_gpu_stream_expert_pending_load_release_buffers(p); - p->n_tasks = 0; - p->n_loads = 0; - p->prepare_ms = 0.0; - return ok; -} - -static void ds4_gpu_stream_expert_pending_load_clear(void) { - if (!g_stream_expert_pending_load.active) { - ds4_gpu_stream_expert_pending_load_release_buffers( - &g_stream_expert_pending_load); - return; - } - (void)ds4_gpu_stream_expert_pending_load_finish(NULL); -} - -static int ds4_gpu_stream_expert_pending_load_matches( - const void *model_map, - uint64_t model_size, - uint32_t layer, - const int32_t *selected_ids, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { - ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; - if (!p->active || - p->model_map != model_map || - p->model_size != model_size || - p->layer != layer || - p->n_total_expert != n_total_expert || - p->n_selected != n_selected || - p->gate_expert_bytes != gate_expert_bytes || - p->down_expert_bytes != down_expert_bytes) { - return 0; - } - for (uint32_t i = 0; i < n_selected; i++) { - if (p->selected_ids[i] != selected_ids[i]) return 0; - } - return 1; -} - -int ds4_gpu_stream_expert_cache_begin_selected_load( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_selected) { - if (!g_ssd_streaming_mode || - getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EARLY_LOAD") != NULL) { - return 1; - } - if (!table) return 0; - const void *model_map = table->model_map; - const uint64_t model_size = table->model_size; - const uint32_t layer = table->layer; - const uint32_t n_total_expert = table->n_total_expert; - const uint64_t gate_offset = table->gate_offset; - const uint64_t up_offset = table->up_offset; - const uint64_t down_offset = table->down_offset; - const uint64_t gate_expert_bytes = table->gate_expert_bytes; - const uint64_t down_expert_bytes = table->down_expert_bytes; - if (!model_map || !selected_ids || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected) == 0) { - return 1; - } - if (!g_initialized && !ds4_gpu_init()) return 0; - - if (ds4_gpu_stream_expert_pending_load_matches(model_map, - model_size, - layer, - selected_ids, - n_total_expert, - n_selected, - gate_expert_bytes, - down_expert_bytes)) { - return 1; - } - - ds4_gpu_stream_expert_pending_load_clear(); - ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; - p->active = 0; - p->model_map = model_map; - p->model_size = model_size; - p->layer = layer; - p->n_total_expert = n_total_expert; - p->n_selected = n_selected; - p->missing_mask = 0; - p->n_loads = 0; - p->n_tasks = 0; - p->gate_expert_bytes = gate_expert_bytes; - p->down_expert_bytes = down_expert_bytes; - p->prepare_ms = 0.0; - const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); - const double load_prepare_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - p->selected_ids[i] = -1; - p->load_slots[i] = 0; - p->source_slots[i] = UINT32_MAX; - p->gate_abs_offsets[i] = 0; - p->up_abs_offsets[i] = 0; - p->down_abs_offsets[i] = 0; - p->gate_bufs[i] = nil; - p->up_bufs[i] = nil; - p->down_bufs[i] = nil; - p->gate_inners[i] = 0; - p->up_inners[i] = 0; - p->down_inners[i] = 0; - } - memset(p->tasks, 0, sizeof(p->tasks)); - - for (uint32_t i = 0; i < n_selected; i++) { - if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { - fprintf(stderr, - "ds4: Metal streaming early-load expert id %d is outside 0..%u\n", - selected_ids[i], - n_total_expert); - return 0; - } - p->selected_ids[i] = selected_ids[i]; - const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); - return 0; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); - return 0; - } - p->gate_abs_offsets[i] = gate_offset + gate_rel; - p->up_abs_offsets[i] = up_offset + gate_rel; - p->down_abs_offsets[i] = down_offset + down_rel; - - ds4_gpu_stream_expert_cache_entry *e = - &g_stream_expert_cache[layer][(uint32_t)selected_ids[i]]; - if (ds4_gpu_stream_expert_cache_entry_matches(e, - model_map, - model_size, - p->gate_abs_offsets[i], - p->up_abs_offsets[i], - p->down_abs_offsets[i], - gate_expert_bytes, - down_expert_bytes)) { - continue; - } - - uint32_t source = UINT32_MAX; - for (uint32_t prev = 0; prev < i; prev++) { - if (selected_ids[prev] == selected_ids[i] && - p->gate_abs_offsets[prev] == p->gate_abs_offsets[i] && - p->up_abs_offsets[prev] == p->up_abs_offsets[i] && - p->down_abs_offsets[prev] == p->down_abs_offsets[i] && - (p->missing_mask & (1u << prev)) != 0) { - source = p->source_slots[prev] != UINT32_MAX ? - p->source_slots[prev] : prev; - break; - } - } - if (source != UINT32_MAX) { - p->source_slots[i] = source; - p->missing_mask |= 1u << i; - continue; - } - p->source_slots[i] = i; - p->missing_mask |= 1u << i; - p->load_slots[p->n_loads++] = i; - } - if (p->n_loads == 0) return 1; - - const uint32_t cache_budget = - ds4_gpu_stream_expert_cache_configured_budget(); - uint32_t reserved_entries = g_stream_expert_cache_entry_count; - ds4_gpu_stream_expert_reusable_buffers - batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - batch_reuse[i] = - (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; - } - uint32_t batch_reuse_count = 0; - if (cache_budget != 0 && - reserved_entries >= cache_budget && - p->n_loads > 1 && - ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, - down_expert_bytes)) { - const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - batch_reuse_count = - ds4_gpu_stream_expert_cache_take_reusable_batch( - p->n_loads, - layer, - selected_ids, - n_selected, - gate_expert_bytes, - down_expert_bytes, - batch_reuse); - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( - ds4_gpu_now_ms() - reuse_t0); - } - } - for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { - const uint32_t slot = p->load_slots[load_i]; - const uint32_t expert = (uint32_t)p->selected_ids[slot]; - const int force_reuse = - cache_budget != 0 && reserved_entries >= cache_budget; - - ds4_gpu_stream_expert_readahead_range(p->gate_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(p->up_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(p->down_abs_offsets[slot], - down_expert_bytes); - - if (load_i < batch_reuse_count && - batch_reuse[load_i].gate_buffer && - batch_reuse[load_i].up_buffer && - batch_reuse[load_i].down_buffer) { - p->gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; - p->up_bufs[load_i] = batch_reuse[load_i].up_buffer; - p->down_bufs[load_i] = batch_reuse[load_i].down_buffer; - p->gate_inners[load_i] = batch_reuse[load_i].gate_inner; - p->up_inners[load_i] = batch_reuse[load_i].up_inner; - p->down_inners[load_i] = batch_reuse[load_i].down_inner; - } else { - const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - const int prepared = - ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, - expert, - layer, - selected_ids, - n_selected, - gate_expert_bytes, - down_expert_bytes, - force_reuse, - &p->gate_bufs[load_i], - &p->up_bufs[load_i], - &p->down_bufs[load_i], - &p->gate_inners[load_i], - &p->up_inners[load_i], - &p->down_inners[load_i]); - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_buffer( - ds4_gpu_now_ms() - buffer_t0); - } - if (!prepared) { - ds4_gpu_stream_expert_pending_load_release_buffers(p); - return 0; - } - } - if (!force_reuse && reserved_entries < UINT32_MAX) { - reserved_entries++; - } - uint8_t *gate_dst = (uint8_t *)[p->gate_bufs[load_i] contents] + - p->gate_inners[load_i]; - uint8_t *up_dst = (uint8_t *)[p->up_bufs[load_i] contents] + - p->up_inners[load_i]; - uint8_t *down_dst = (uint8_t *)[p->down_bufs[load_i] contents] + - p->down_inners[load_i]; - if (!gate_dst || !up_dst || !down_dst) { - ds4_gpu_stream_expert_pending_load_release_buffers(p); - return 0; - } - const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = p->gate_abs_offsets[slot], - .len = gate_expert_bytes, - .dst = gate_dst, - }; - p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = p->up_abs_offsets[slot], - .len = gate_expert_bytes, - .dst = up_dst, - }; - p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = p->down_abs_offsets[slot], - .len = down_expert_bytes, - .dst = down_dst, - }; - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_task( - 1, - ds4_gpu_now_ms() - task_t0); - } - } - - const uint32_t n_workers = - ds4_gpu_stream_expert_pread_thread_count(p->n_tasks); - p->start_ms = ds4_gpu_now_ms(); - if (load_timing) { - p->prepare_ms = p->start_ms - load_prepare_t0; - } - if (ds4_gpu_stream_expert_pread_pool_begin(p->tasks, - p->n_tasks, - n_workers)) { - p->active = 1; - if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { - fprintf(stderr, - "ds4: Metal streaming expert early-load begin layer=%u experts=%u tensors=%u threads=%u\n", - layer, - p->n_loads, - p->n_tasks, - n_workers); - } - return 1; - } - - uint64_t read_bytes = 0; - double read_ms = 0.0; - if (!ds4_gpu_stream_expert_pread_tasks(p->tasks, - p->n_tasks, - &read_bytes, - &read_ms)) { - ds4_gpu_stream_expert_pending_load_release_buffers(p); - return 0; - } - (void)read_bytes; - if (!ds4_gpu_stream_expert_pending_load_install(p, NULL, read_ms)) { - ds4_gpu_stream_expert_pending_load_release_buffers(p); - return 0; - } - ds4_gpu_stream_expert_pending_load_release_buffers(p); - p->n_tasks = 0; - p->n_loads = 0; - p->prepare_ms = 0.0; - return 1; -} - -static int ds4_gpu_stream_expert_cache_load_selected_missing( - const void *model_map, - uint64_t model_size, - uint32_t layer, - const int32_t *selected_ids, - uint32_t n_total_expert, - uint32_t n_selected, - const uint64_t *gate_abs_offsets, - const uint64_t *up_abs_offsets, - const uint64_t *down_abs_offsets, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - uint32_t missing_mask, - ds4_gpu_stream_expert_cache_entry **entries) { - if (!g_ssd_streaming_mode || - !model_map || - !selected_ids || - !gate_abs_offsets || - !up_abs_offsets || - !down_abs_offsets || - !entries || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected) == 0) { - return 0; - } - missing_mask &= (1u << n_selected) - 1u; - if (missing_mask == 0) return 1; - if (g_stream_expert_pending_load.active && - n_selected <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { - if (ds4_gpu_stream_expert_pending_load_matches(model_map, - model_size, - layer, - selected_ids, - n_total_expert, - n_selected, - gate_expert_bytes, - down_expert_bytes)) { - if (!ds4_gpu_stream_expert_pending_load_finish(entries)) return 0; - for (uint32_t i = 0; i < n_selected; i++) { - if ((missing_mask & (1u << i)) != 0 && entries[i]) { - missing_mask &= ~(1u << i); - } - } - if (missing_mask == 0) return 1; - } else { - ds4_gpu_stream_expert_pending_load_clear(); - } - } else if (g_stream_expert_pending_load.active) { - ds4_gpu_stream_expert_pending_load_clear(); - } - - const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); - double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - double load_prepare_ms = 0.0; - double load_modify_ms = 0.0; - double load_install_ms = 0.0; - - uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - load_slots[i] = 0; - source_slots[i] = UINT32_MAX; - } - uint32_t n_loads = 0; - for (uint32_t i = 0; i < n_selected; i++) { - if ((missing_mask & (1u << i)) == 0) continue; - if (entries[i]) { - source_slots[i] = i; - continue; - } - if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { - fprintf(stderr, - "ds4: Metal streaming selected missing expert id %d is outside 0..%u\n", - selected_ids[i], - n_total_expert); - return 0; - } - uint32_t source = UINT32_MAX; - for (uint32_t prev = 0; prev < i; prev++) { - if (selected_ids[prev] == selected_ids[i] && - gate_abs_offsets[prev] == gate_abs_offsets[i] && - up_abs_offsets[prev] == up_abs_offsets[i] && - down_abs_offsets[prev] == down_abs_offsets[i] && - (entries[prev] || (missing_mask & (1u << prev)) != 0)) { - source = source_slots[prev] != UINT32_MAX ? - source_slots[prev] : prev; - break; - } - } - if (source != UINT32_MAX) { - source_slots[i] = source; - continue; - } - source_slots[i] = i; - load_slots[n_loads++] = i; - } - if (n_loads == 0) { - for (uint32_t i = 0; i < n_selected; i++) { - if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; - const uint32_t source = source_slots[i]; - if (source >= n_selected || !entries[source]) return 0; - entries[i] = entries[source]; - entries[i]->use_count++; - } - return 1; - } - - __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - gate_bufs[i] = nil; - up_bufs[i] = nil; - down_bufs[i] = nil; - gate_inners[i] = 0; - up_inners[i] = 0; - down_inners[i] = 0; - } - ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; - memset(tasks, 0, sizeof(tasks)); - uint32_t n_tasks = 0; - const uint32_t cache_budget = - ds4_gpu_stream_expert_cache_configured_budget(); - uint32_t reserved_entries = g_stream_expert_cache_entry_count; - ds4_gpu_stream_expert_reusable_buffers - batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - batch_reuse[i] = - (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; - } - uint32_t batch_reuse_count = 0; - if (cache_budget != 0 && - reserved_entries >= cache_budget && - n_loads > 1 && - ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, - down_expert_bytes)) { - const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - batch_reuse_count = - ds4_gpu_stream_expert_cache_take_reusable_batch( - n_loads, - layer, - selected_ids, - n_selected, - gate_expert_bytes, - down_expert_bytes, - batch_reuse); - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( - ds4_gpu_now_ms() - reuse_t0); - } - } - - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - const uint32_t slot = load_slots[load_i]; - const uint32_t expert = (uint32_t)selected_ids[slot]; - const int force_reuse = - cache_budget != 0 && reserved_entries >= cache_budget; - - ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot], - down_expert_bytes); - - if (load_i < batch_reuse_count && - batch_reuse[load_i].gate_buffer && - batch_reuse[load_i].up_buffer && - batch_reuse[load_i].down_buffer) { - gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; - up_bufs[load_i] = batch_reuse[load_i].up_buffer; - down_bufs[load_i] = batch_reuse[load_i].down_buffer; - gate_inners[load_i] = batch_reuse[load_i].gate_inner; - up_inners[load_i] = batch_reuse[load_i].up_inner; - down_inners[load_i] = batch_reuse[load_i].down_inner; - } else { - const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - const int prepared = - ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, - expert, - layer, - selected_ids, - n_selected, - gate_expert_bytes, - down_expert_bytes, - force_reuse, - &gate_bufs[load_i], - &up_bufs[load_i], - &down_bufs[load_i], - &gate_inners[load_i], - &up_inners[load_i], - &down_inners[load_i]); - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_buffer( - ds4_gpu_now_ms() - buffer_t0); - } - if (!prepared) { - return 0; - } - } - if (!force_reuse && reserved_entries < UINT32_MAX) { - reserved_entries++; - } - if (!gate_bufs[load_i] || !up_bufs[load_i] || !down_bufs[load_i]) { - return 0; - } - - uint8_t *gate_dst = (uint8_t *)[gate_bufs[load_i] contents] + - gate_inners[load_i]; - uint8_t *up_dst = (uint8_t *)[up_bufs[load_i] contents] + - up_inners[load_i]; - uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] + - down_inners[load_i]; - if (!gate_dst || !up_dst || !down_dst) return 0; - - const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = gate_abs_offsets[slot], - .len = gate_expert_bytes, - .dst = gate_dst, - }; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = up_abs_offsets[slot], - .len = gate_expert_bytes, - .dst = up_dst, - }; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = down_abs_offsets[slot], - .len = down_expert_bytes, - .dst = down_dst, - }; - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_task( - 1, - ds4_gpu_now_ms() - task_t0); - } - } - - if (load_timing) { - const double now_ms = ds4_gpu_now_ms(); - load_prepare_ms = now_ms - load_t0; - load_t0 = now_ms; - } - - uint64_t read_bytes = 0; - double read_ms = 0.0; - const int ok = ds4_gpu_stream_expert_pread_tasks(tasks, - n_tasks, - &read_bytes, - &read_ms); - if (!ok) return 0; - if (load_timing) { - load_t0 = ds4_gpu_now_ms(); - } - ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); - - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; - } - if (load_timing) { - const double now_ms = ds4_gpu_now_ms(); - load_modify_ms = now_ms - load_t0; - load_t0 = now_ms; - } - if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u " - "threads=%u bytes=%.2f GiB wall=%.3f ms\n", - layer, - n_loads, - n_tasks, - ds4_gpu_stream_expert_pread_thread_count(n_tasks), - ds4_gpu_gib(read_bytes), - read_ms); - } - - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - const uint32_t slot = load_slots[load_i]; - const uint32_t expert = (uint32_t)selected_ids[slot]; - ds4_gpu_stream_expert_cache_entry *entry = - ds4_gpu_stream_expert_cache_install_loaded(model_map, - model_size, - layer, - expert, - gate_abs_offsets[slot], - up_abs_offsets[slot], - down_abs_offsets[slot], - gate_expert_bytes, - down_expert_bytes, - gate_bufs[load_i], - up_bufs[load_i], - down_bufs[load_i], - gate_inners[load_i], - up_inners[load_i], - down_inners[load_i]); - if (!entry) return 0; - entries[slot] = entry; - } - - for (uint32_t i = 0; i < n_selected; i++) { - if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; - const uint32_t source = source_slots[i]; - if (source >= n_selected || !entries[source]) return 0; - entries[i] = entries[source]; - entries[i]->use_count++; - } - for (uint32_t i = 0; i < n_selected; i++) { - if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0; - } - if (load_timing) { - load_install_ms = ds4_gpu_now_ms() - load_t0; - ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, - read_ms, - load_modify_ms, - load_install_ms); - } - return 1; -} - -static void ds4_gpu_glm_stream_selected_prefetch_set( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_selected) { - g_glm_stream_selected_prefetch.active = 0; - if (!table || !selected_ids || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { - return; - } - g_glm_stream_selected_prefetch.model_map = table->model_map; - g_glm_stream_selected_prefetch.model_size = table->model_size; - g_glm_stream_selected_prefetch.layer = table->layer; - g_glm_stream_selected_prefetch.n_total_expert = table->n_total_expert; - g_glm_stream_selected_prefetch.n_selected = n_selected; - g_glm_stream_selected_prefetch.gate_offset = table->gate_offset; - g_glm_stream_selected_prefetch.up_offset = table->up_offset; - g_glm_stream_selected_prefetch.down_offset = table->down_offset; - g_glm_stream_selected_prefetch.gate_expert_bytes = table->gate_expert_bytes; - g_glm_stream_selected_prefetch.down_expert_bytes = table->down_expert_bytes; - for (uint32_t i = 0; i < n_selected; i++) { - g_glm_stream_selected_prefetch.selected_ids[i] = selected_ids[i]; - } - g_glm_stream_selected_prefetch.active = 1; -} - -static int ds4_gpu_glm_stream_selected_prefetch_take( - const void *model_map, - uint64_t model_size, - uint32_t layer, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - int32_t *selected_ids_out) { - if (!selected_ids_out || !g_glm_stream_selected_prefetch.active) return 0; - ds4_gpu_glm_stream_selected_prefetch *p = &g_glm_stream_selected_prefetch; - if (p->model_map != model_map || - p->model_size != model_size || - p->layer != layer || - p->n_total_expert != n_total_expert || - p->n_selected != n_selected || - p->gate_offset != gate_offset || - p->up_offset != up_offset || - p->down_offset != down_offset || - p->gate_expert_bytes != gate_expert_bytes || - p->down_expert_bytes != down_expert_bytes) { - return 0; - } - for (uint32_t i = 0; i < n_selected; i++) { - selected_ids_out[i] = p->selected_ids[i]; - } - p->active = 0; - return 1; -} - -int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( - const ds4_gpu_stream_expert_table *table, - const ds4_gpu_tensor *selected, - uint32_t n_selected) { - g_glm_stream_selected_prefetch.active = 0; - if (!g_ssd_streaming_mode || - getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) { - return 1; - } - if (!table || !selected || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { - return 1; - } - - const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); - double t0 = timing ? ds4_gpu_now_ms() : 0.0; - double selected_sync_ms = 0.0; - double selected_copy_ms = 0.0; - double selected_bind_ms = 0.0; - const int had_batch = g_batch_cb != nil; - if (had_batch) { - if (ds4_gpu_end_commands() == 0) return 0; - if (timing) { - selected_sync_ms = ds4_gpu_now_ms() - t0; - t0 = ds4_gpu_now_ms(); - } - } - - int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { - selected_ids[i] = -1; - } - int ok = ds4_gpu_tensor_read(selected, - 0, - selected_ids, - (uint64_t)n_selected * sizeof(selected_ids[0])) != 0; - if (timing) { - selected_copy_ms = ds4_gpu_now_ms() - t0; - t0 = ds4_gpu_now_ms(); - } - if (ok) { - ok = ds4_gpu_stream_expert_cache_begin_selected_load(table, - selected_ids, - n_selected) != 0; - } - if (timing) { - selected_bind_ms = ds4_gpu_now_ms() - t0; - ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, - selected_copy_ms, - selected_bind_ms); - } - if (ok) { - ds4_gpu_glm_stream_selected_prefetch_set(table, - selected_ids, - n_selected); - } - if (had_batch && ds4_gpu_begin_commands() == 0) ok = 0; - return ok; -} - -static void ds4_gpu_stream_expert_cache_clear_layer(uint32_t layer) { - if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; - for (uint32_t expert = 0; - expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; - expert++) { - ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); - } - g_stream_expert_cache_layer_count[layer] = 0; -} - -static int ds4_gpu_stream_expert_cache_prepare_selected_batch( - const void *model_map, - uint64_t model_size, - uint32_t layer, - const ds4_gpu_tensor *selected, - uint32_t n_tokens, - uint32_t n_total_expert, - uint32_t n_selected, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint64_t gate_expert_bytes, - uint64_t down_expert_bytes, - id *gate_addrs, - id *up_addrs, - id *down_addrs, - ds4_gpu_stream_expert_cache_entry **resources, - uint32_t *n_resources, - uint32_t *unique_out, - id *overflow_gate, - id *overflow_up, - id *overflow_down) { - if (overflow_gate) *overflow_gate = nil; - if (overflow_up) *overflow_up = nil; - if (overflow_down) *overflow_down = nil; - if (!g_ssd_streaming_mode || - !model_map || - !selected || - !gate_addrs || - !up_addrs || - !down_addrs || - !resources || - !n_resources || - !unique_out || - !overflow_gate || - !overflow_up || - !overflow_down || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_tokens == 0 || - n_selected == 0 || - n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - gate_expert_bytes == 0 || - down_expert_bytes == 0) { - return 0; - } - if (n_tokens > UINT32_MAX / n_selected) return 0; - - const uint64_t n_ids = (uint64_t)n_tokens * n_selected; - if (n_ids > SIZE_MAX / sizeof(int32_t)) return 0; - int32_t *ids = malloc((size_t)n_ids * sizeof(ids[0])); - if (!ids) return 0; - - const bool profile = - getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_PROFILE") != NULL; - const double t0 = profile ? ds4_gpu_now_ms() : 0.0; - int ok = ds4_gpu_tensor_read(selected, - 0, - ids, - n_ids * sizeof(ids[0])); - const double t_read = profile ? ds4_gpu_now_ms() : 0.0; - - bool seen[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { false }; - uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { 0 }; - int32_t unique_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - uint32_t unique_count = 0; - *n_resources = 0; - *unique_out = 0; - if (ok) { - if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { - ok = 0; - } - } - if (ok) { - for (uint64_t i = 0; i < n_ids; i++) { - const int32_t selected_id = ids[i]; - if (selected_id < 0 || (uint32_t)selected_id >= n_total_expert) { - fprintf(stderr, - "ds4: Metal streaming batch selected expert id %d is outside 0..%u at layer %u\n", - selected_id, - n_total_expert, - layer); - ok = 0; - break; - } - frequency[(uint32_t)selected_id]++; - if (!seen[(uint32_t)selected_id]) { - seen[(uint32_t)selected_id] = true; - unique_ids[unique_count++] = selected_id; - } - } - } - if (ok) { - ds4_gpu_stream_expert_cache_note_frequency_hotness(layer, - frequency, - n_total_expert); - } - /* - * When the layer's unique selected set does not fit the cache budget, the - * extra experts are addressed straight into whole-tensor mapped model views - * instead of falling back to a different MoE kernel path. The address-table - * kernels read identical expert bytes either way, so the cache/view split - * does not change the computed logits. - */ - uint32_t view_served = 0; - uint64_t overflow_gate_inner = 0; - uint64_t overflow_up_inner = 0; - uint64_t overflow_down_inner = 0; - - uint64_t unique_gate_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - uint64_t unique_up_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - uint64_t unique_down_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - ds4_gpu_stream_expert_cache_entry - *unique_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - uint32_t load_unique[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - __strong id - gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - __strong id - up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - __strong id - down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; - for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; i++) { - unique_gate_offsets[i] = 0; - unique_up_offsets[i] = 0; - unique_down_offsets[i] = 0; - unique_entries[i] = NULL; - load_unique[i] = 0; - gate_bufs[i] = nil; - up_bufs[i] = nil; - down_bufs[i] = nil; - gate_inners[i] = 0; - up_inners[i] = 0; - down_inners[i] = 0; - } - - ds4_gpu_stream_expert_pread_task *tasks = NULL; - uint32_t n_loads = 0; - uint32_t n_tasks = 0; - double load_prepare_ms = 0.0; - double load_modify_ms = 0.0; - double load_install_ms = 0.0; - double load_timing_t0 = ds4_gpu_stream_expert_timing_summary_enabled() ? - ds4_gpu_now_ms() : 0.0; - const int load_timing = load_timing_t0 != 0.0; - if (ok && unique_count != 0) { - tasks = calloc((size_t)unique_count * 3u, sizeof(tasks[0])); - if (!tasks) ok = 0; - } - if (ok) { - const uint32_t cache_budget = - ds4_gpu_stream_expert_cache_configured_budget(); - uint32_t reserved_entries = g_stream_expert_cache_entry_count; - - for (uint32_t u = 0; u < unique_count; u++) { - const uint32_t expert = (uint32_t)unique_ids[u]; - - if ((uint64_t)expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)expert > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); - ok = 0; - break; - } - const uint64_t gate_rel = (uint64_t)expert * gate_expert_bytes; - const uint64_t down_rel = (uint64_t)expert * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); - ok = 0; - break; - } - unique_gate_offsets[u] = gate_offset + gate_rel; - unique_up_offsets[u] = up_offset + gate_rel; - unique_down_offsets[u] = down_offset + down_rel; - - ds4_gpu_stream_expert_cache_entry *entry = - ds4_gpu_stream_expert_cache_peek(model_map, - model_size, - layer, - expert, - n_total_expert, - n_selected, - unique_gate_offsets[u], - unique_up_offsets[u], - unique_down_offsets[u], - gate_expert_bytes, - down_expert_bytes); - if (entry) { - unique_entries[u] = entry; - continue; - } - - const int force_reuse = - cache_budget != 0 && reserved_entries >= cache_budget; - ds4_gpu_stream_expert_readahead_range(unique_gate_offsets[u], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(unique_up_offsets[u], - gate_expert_bytes); - ds4_gpu_stream_expert_readahead_range(unique_down_offsets[u], - down_expert_bytes); - const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - const int prepared = - ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, - expert, - layer, - unique_ids, - unique_count, - gate_expert_bytes, - down_expert_bytes, - force_reuse, - &gate_bufs[n_loads], - &up_bufs[n_loads], - &down_bufs[n_loads], - &gate_inners[n_loads], - &up_inners[n_loads], - &down_inners[n_loads]); - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_buffer( - ds4_gpu_now_ms() - buffer_t0); - } - if (!prepared) { - ok = 0; - break; - } - if (!force_reuse && reserved_entries < UINT32_MAX) { - reserved_entries++; - } - if (!gate_bufs[n_loads] || - !up_bufs[n_loads] || - !down_bufs[n_loads]) { - ok = 0; - break; - } - - uint8_t *gate_dst = (uint8_t *)[gate_bufs[n_loads] contents] + - gate_inners[n_loads]; - uint8_t *up_dst = (uint8_t *)[up_bufs[n_loads] contents] + - up_inners[n_loads]; - uint8_t *down_dst = (uint8_t *)[down_bufs[n_loads] contents] + - down_inners[n_loads]; - if (!gate_dst || !up_dst || !down_dst) { - ok = 0; - break; - } - - load_unique[n_loads] = u; - const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = unique_gate_offsets[u], - .len = gate_expert_bytes, - .dst = gate_dst, - }; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = unique_up_offsets[u], - .len = gate_expert_bytes, - .dst = up_dst, - }; - tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { - .offset = unique_down_offsets[u], - .len = down_expert_bytes, - .dst = down_dst, - }; - if (load_timing) { - ds4_gpu_stream_expert_timing_note_prepare_task( - 1, - ds4_gpu_now_ms() - task_t0); - } - n_loads++; - } - } - if (ok && n_loads != 0) { - if (load_timing_t0 != 0.0) { - const double now_ms = ds4_gpu_now_ms(); - load_prepare_ms = now_ms - load_timing_t0; - load_timing_t0 = now_ms; - } - uint64_t read_bytes = 0; - double read_ms = 0.0; - ok = ds4_gpu_stream_expert_pread_tasks(tasks, - n_tasks, - &read_bytes, - &read_ms); - if (ok) { - ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); - } - if (load_timing_t0 != 0.0) { - load_timing_t0 = ds4_gpu_now_ms(); - } - if (ok) { - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; - [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; - } - } - if (load_timing_t0 != 0.0) { - const double now_ms = ds4_gpu_now_ms(); - load_modify_ms = now_ms - load_timing_t0; - load_timing_t0 = now_ms; - } - if (ok && getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal streaming batch expert parallel pread layer=%u experts=%u tensors=%u " - "threads=%u bytes=%.2f GiB wall=%.3f ms\n", - layer, - n_loads, - n_tasks, - ds4_gpu_stream_expert_pread_thread_count(n_tasks), - ds4_gpu_gib(read_bytes), - read_ms); - } - if (ok) { - for (uint32_t load_i = 0; load_i < n_loads; load_i++) { - const uint32_t u = load_unique[load_i]; - const uint32_t expert = (uint32_t)unique_ids[u]; - ds4_gpu_stream_expert_cache_entry *entry = - ds4_gpu_stream_expert_cache_install_loaded(model_map, - model_size, - layer, - expert, - unique_gate_offsets[u], - unique_up_offsets[u], - unique_down_offsets[u], - gate_expert_bytes, - down_expert_bytes, - gate_bufs[load_i], - up_bufs[load_i], - down_bufs[load_i], - gate_inners[load_i], - up_inners[load_i], - down_inners[load_i]); - if (!entry) { - ok = 0; - break; - } - unique_entries[u] = entry; - } - } - if (load_timing_t0 != 0.0) { - load_install_ms = ds4_gpu_now_ms() - load_timing_t0; - ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, - read_ms, - load_modify_ms, - load_install_ms); - } - } - if (tasks) free(tasks); - if (ok) { - for (uint32_t u = 0; u < unique_count; u++) { - ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u]; - const uint32_t expert = (uint32_t)unique_ids[u]; - if (!entry) { - const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset; - const uint64_t down_rel = unique_down_offsets[u] - down_offset; - if (!*overflow_gate) { - const uint64_t gate_tensor_bytes = - (uint64_t)n_total_expert * gate_expert_bytes; - const uint64_t down_tensor_bytes = - (uint64_t)n_total_expert * down_expert_bytes; - uint64_t gate_view_inner = 0; - uint64_t up_view_inner = 0; - uint64_t down_view_inner = 0; - id gv = ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset, - gate_tensor_bytes, - &gate_view_inner); - id uv = ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset, - gate_tensor_bytes, - &up_view_inner); - id dv = ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset, - down_tensor_bytes, - &down_view_inner); - if (!gv || !uv || !dv) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr " - "failed to map overflow expert views at layer %u\n", - layer); - ok = 0; - break; - } - *overflow_gate = gv; - *overflow_up = uv; - *overflow_down = dv; - overflow_gate_inner = gate_view_inner; - overflow_up_inner = up_view_inner; - overflow_down_inner = down_view_inner; - } - if (!ds4_gpu_stream_expert_cache_set_addr_slot( - layer, - expert, - *overflow_gate, - (NSUInteger)(overflow_gate_inner + gate_rel), - *overflow_up, - (NSUInteger)(overflow_up_inner + gate_rel), - *overflow_down, - (NSUInteger)(overflow_down_inner + down_rel))) { - ok = 0; - break; - } - view_served++; - continue; - } - const uint32_t extra_uses = - frequency[expert] > 0 ? frequency[expert] - 1u : 0; - if (extra_uses != 0) { - if (entry->use_count > UINT64_MAX - extra_uses) { - entry->use_count = UINT64_MAX; - } else { - entry->use_count += extra_uses; - } - if (g_stream_expert_cache_hits > UINT64_MAX - extra_uses) { - g_stream_expert_cache_hits = UINT64_MAX; - } else { - g_stream_expert_cache_hits += extra_uses; - } - if (g_stream_expert_cache_layer_hits[layer] > - UINT64_MAX - extra_uses) { - g_stream_expert_cache_layer_hits[layer] = UINT64_MAX; - } else { - g_stream_expert_cache_layer_hits[layer] += extra_uses; - } - } - resources[*n_resources] = entry; - (*n_resources)++; - } - } - free(ids); - - if (!ok || (*n_resources == 0 && view_served == 0)) { - ds4_gpu_stream_expert_cache_clear_layer(layer); - return 0; - } - if (!ds4_gpu_stream_expert_cache_addr_buffers(layer, - gate_addrs, - up_addrs, - down_addrs)) { - ds4_gpu_stream_expert_cache_clear_layer(layer); - return 0; - } - *unique_out = *n_resources + view_served; - if (view_served != 0 && - ds4_gpu_stream_expert_timing_summary_enabled()) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "served %u/%u unique experts via mapped views (cache budget %u)\n", - layer, - view_served, - unique_count, - ds4_gpu_stream_expert_cache_configured_budget()); - } - if (profile) { - const double t_done = ds4_gpu_now_ms(); - uint64_t per_expert_bytes = UINT64_MAX; - uint64_t logical_bytes = UINT64_MAX; - if (gate_expert_bytes <= (UINT64_MAX - down_expert_bytes) / 2ull) { - per_expert_bytes = gate_expert_bytes * 2ull + down_expert_bytes; - if (per_expert_bytes == 0 || - *n_resources <= UINT64_MAX / per_expert_bytes) { - logical_bytes = (uint64_t)(*n_resources) * per_expert_bytes; - } - } - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "tokens=%u unique=%u read=%.3f ms wrap=%.3f ms bytes=%.2f GiB\n", - layer, - n_tokens, - *n_resources, - t_read - t0, - t_done - t_read, - ds4_gpu_gib(logical_bytes)); - } - return 1; -} - -int ds4_gpu_stream_expert_cache_seed_selected( - const ds4_gpu_stream_expert_table *table, - const int32_t *selected_ids, - uint32_t n_selected) { - if (!g_ssd_streaming_mode) return 1; - if (!table) return 0; - const void *model_map = table->model_map; - const uint64_t model_size = table->model_size; - const uint32_t layer = table->layer; - const uint32_t n_total_expert = table->n_total_expert; - const uint64_t gate_offset = table->gate_offset; - const uint64_t up_offset = table->up_offset; - const uint64_t down_offset = table->down_offset; - const uint64_t gate_expert_bytes = table->gate_expert_bytes; - const uint64_t down_expert_bytes = table->down_expert_bytes; - if (!model_map || !selected_ids || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_selected == 0 || - n_selected > 6 || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - n_selected) == 0) { - return 1; - } - if (!g_initialized && !ds4_gpu_init()) return 0; - - ds4_gpu_stream_expert_cache_note_selected_hotness(layer, - selected_ids, - n_selected); - for (uint32_t i = 0; i < n_selected; i++) { - if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { - fprintf(stderr, - "ds4: Metal prefill expert-cache seed selected expert id %d is outside 0..%u\n", - selected_ids[i], - n_total_expert); - return 0; - } - const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); - return 0; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); - return 0; - } - - if (!ds4_gpu_stream_expert_cache_get(model_map, - model_size, - layer, - (uint32_t)selected_ids[i], - n_total_expert, - n_selected, - gate_offset + gate_rel, - up_offset + gate_rel, - down_offset + down_rel, - gate_expert_bytes, - down_expert_bytes)) { - return 0; - } - } - ds4_gpu_stream_expert_cache_prune_layer(layer, - n_total_expert, - n_selected, - selected_ids, - n_selected); - ds4_gpu_stream_expert_cache_prune_global(layer, - selected_ids, - n_selected); - return 1; -} - -int ds4_gpu_stream_expert_cache_seed_experts( - const ds4_gpu_stream_expert_table *table, - const int32_t *expert_ids, - const uint32_t *expert_priorities, - uint32_t n_experts) { - if (!g_ssd_streaming_mode) return 1; - if (!table) return 0; - const void *model_map = table->model_map; - const uint64_t model_size = table->model_size; - const uint32_t layer = table->layer; - const uint32_t n_total_expert = table->n_total_expert; - const uint64_t gate_offset = table->gate_offset; - const uint64_t up_offset = table->up_offset; - const uint64_t down_offset = table->down_offset; - const uint64_t gate_expert_bytes = table->gate_expert_bytes; - const uint64_t down_expert_bytes = table->down_expert_bytes; - if (!model_map || !expert_ids || - layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || - n_experts == 0 || - n_total_expert == 0 || - n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || - !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) || - ds4_gpu_stream_expert_cache_effective_cap(layer, - n_total_expert, - 1) == 0) { - return 1; - } - if (!g_initialized && !ds4_gpu_init()) return 0; - - ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); - uint32_t remaining = n_experts; - while (remaining != 0) { - const uint32_t batch = - remaining > 6u ? 6u : remaining; - remaining -= batch; - - int32_t selected_ids[6] = { -1, -1, -1, -1, -1, -1 }; - uint64_t gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - uint64_t up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - uint64_t down_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - ds4_gpu_stream_expert_cache_entry *entries[6] = { - NULL, NULL, NULL, NULL, NULL, NULL - }; - uint32_t missing_mask = 0; - - for (uint32_t i = 0; i < batch; i++) { - const int32_t expert = expert_ids[remaining + i]; - const uint32_t priority = - expert_priorities ? expert_priorities[remaining + i] : 0; - if (expert < 0 || (uint32_t)expert >= n_total_expert) { - fprintf(stderr, - "ds4: Metal streaming hotlist seed expert id %d is outside 0..%u\n", - expert, - n_total_expert); - return 0; - } - const uint64_t expert_id = (uint64_t)(uint32_t)expert; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); - return 0; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); - return 0; - } - - ds4_gpu_stream_expert_cache_note_route_hotness( - layer, - (uint32_t)expert, - priority != 0 ? priority : 1u); - selected_ids[i] = expert; - gate_abs_offsets[i] = gate_offset + gate_rel; - up_abs_offsets[i] = up_offset + gate_rel; - down_abs_offsets[i] = down_offset + down_rel; - entries[i] = ds4_gpu_stream_expert_cache_peek(model_map, - model_size, - layer, - (uint32_t)expert, - n_total_expert, - batch, - gate_abs_offsets[i], - up_abs_offsets[i], - down_abs_offsets[i], - gate_expert_bytes, - down_expert_bytes); - if (!entries[i]) missing_mask |= 1u << i; - if (entries[i] && priority != 0 && - entries[i]->use_count < (uint64_t)priority) { - entries[i]->use_count = (uint64_t)priority; - } - } - - if (missing_mask != 0 && - !ds4_gpu_stream_expert_cache_load_selected_missing(model_map, - model_size, - layer, - selected_ids, - n_total_expert, - batch, - gate_abs_offsets, - up_abs_offsets, - down_abs_offsets, - gate_expert_bytes, - down_expert_bytes, - missing_mask, - entries)) { - return 0; - } - if (expert_priorities) { - for (uint32_t i = 0; i < batch; i++) { - const uint32_t priority = expert_priorities[remaining + i]; - if (entries[i] && priority != 0 && - entries[i]->use_count < (uint64_t)priority) { - entries[i]->use_count = (uint64_t)priority; - } - } - } - } - - const uint32_t protect_n = n_experts < 6u ? n_experts : 6u; - ds4_gpu_stream_expert_cache_prune_layer(layer, - n_total_expert, - 1, - expert_ids, - protect_n); - ds4_gpu_stream_expert_cache_prune_global(layer, - expert_ids, - protect_n); - return 1; -} - -static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) { - const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE"); - if (!env || !env[0]) return 1; - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end == env || *end != '\0' || v < 2 || v > n_total_expert) { - return 1; - } - return (uint32_t)v; -} - -static bool ds4_gpu_q4_table_queue_residency_requested(void) { - return getenv("DS4_METAL_Q4_TABLE_QUEUE_RESIDENCY_SET") != NULL; -} - -static bool ds4_gpu_q4_table_queue_residency_available(void) { -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - return g_queue && [g_queue respondsToSelector:@selector(addResidencySet:)]; - } -#endif - return false; -} - -static bool ds4_gpu_q4_non_streaming_opt_in_enabled(void) { - return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL; -} - -static bool ds4_gpu_q4_selected_paths_allowed(void) { - if (g_ssd_streaming_mode) return true; - if (g_glm_model_mode) return false; - return ds4_gpu_q4_non_streaming_opt_in_enabled(); -} - -int ds4_gpu_pro_q4_expert_table_auto_available(void) { - if (!g_initialized && !ds4_gpu_init()) return 0; - return (g_ssd_streaming_mode || - getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL) && - getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && - ds4_gpu_q4_table_queue_residency_available(); -} - -static bool ds4_gpu_q4_table_queue_residency_enabled(bool auto_queue_residency) { - return (auto_queue_residency || ds4_gpu_q4_table_queue_residency_requested()) && - ds4_gpu_q4_table_queue_residency_available(); -} - -static bool ds4_gpu_q4_table_model_residency_enabled(void) { - return getenv("DS4_METAL_Q4_TABLE_MODEL_RESIDENCY_SET") != NULL; -} - -static bool ds4_gpu_pro_q4_expert_indirect_shape_supported( - uint32_t n_total_expert, - uint32_t n_expert, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes) { - const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; - return n_total_expert == 384 && - n_expert == 6 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes; -} - -static bool ds4_gpu_pro_q4_expert_table_auto_enabled( - uint32_t n_total_expert, - uint32_t n_expert, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes) { - /* - * This path lets the shader choose among exact per-expert resources. - * It is only automatic when a Metal queue residency set can make every - * indirect expert resource visible up front; otherwise the selected - * exact-slice path remains the fallback. - */ - return ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes) && - ds4_gpu_pro_q4_expert_table_auto_available(); -} - -static bool ds4_gpu_pro_q4_expert_address_auto_enabled( - uint32_t n_total_expert, - uint32_t n_expert, - uint64_t gate_tensor_bytes, - uint64_t down_tensor_bytes) { - /* - * GPU-address expert tables are useful for experiments, but they are not - * safe as an automatic path unless every indirect expert resource is made - * visible to Metal. Keep this behind an explicit opt-in while the selected - * active-slice path remains the correctness/performance baseline. - */ - return getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL && - ds4_gpu_q4_selected_paths_allowed() && - ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes) && - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && - getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && - getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && - ds4_gpu_q4_table_queue_residency_available(); -} - -static bool ds4_gpu_q4_table_bind_anchors_enabled(void) { - return getenv("DS4_METAL_Q4_TABLE_BIND_ANCHORS") != NULL; -} - -static int ds4_gpu_use_model_residency_set(id cb) { - if (!ds4_gpu_q4_table_model_residency_enabled()) return 1; -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - if (cb && g_model_residency_set && [cb respondsToSelector:@selector(useResidencySet:)]) { - [cb useResidencySet:g_model_residency_set]; - return 1; - } - } -#endif - fprintf(stderr, "ds4: Metal Q4 table model residency set is not available\n"); - return 0; -} - -static int ds4_gpu_bind_q4_expert_table_anchors( - id enc, - DS4MetalQ4ExpertTable *table, - NSUInteger first_index, - NSUInteger max_count) { - if (!ds4_gpu_q4_table_bind_anchors_enabled()) return 1; - if (!enc || !table || !table.expertBuffers) return 0; - - const NSUInteger count = [table.expertBuffers count]; - if (count > max_count) { - fprintf(stderr, - "ds4: Metal Q4 table anchor count %lu exceeds available slots %lu\n", - (unsigned long)count, - (unsigned long)max_count); - return 0; - } - for (NSUInteger i = 0; i < count; i++) { - [enc setBuffer:[table.expertBuffers objectAtIndex:i] - offset:0 - atIndex:first_index + i]; - } - return 1; -} - -static id ds4_gpu_q4_expert_table_residency_set(NSMutableArray> *buffers) { - if (!buffers || [buffers count] == 0 || - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL || - getenv("DS4_METAL_Q4_TABLE_PER_TENSOR_RESIDENCY_SET") == NULL) { - return nil; - } -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; - desc.label = @"ds4_q4_expert_table"; - desc.initialCapacity = [buffers count]; - NSError *error = nil; - id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; - if (!residency_set) { - fprintf(stderr, "ds4: Metal Q4 expert table residency set creation failed: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - for (id buffer in buffers) { - [residency_set addAllocation:buffer]; - } - [residency_set commit]; - [residency_set requestResidency]; - return residency_set; - } -#endif - return nil; -} - -static void ds4_gpu_q4_residency_add_table(id residency_set, - DS4MetalQ4ExpertTable *table) { -#if TARGET_OS_OSX - if (!residency_set || !table) return; - if (@available(macOS 15.0, *)) { - if (table.argumentBuffer) { - [residency_set addAllocation:table.argumentBuffer]; - } - if (table.addressBuffer) { - [residency_set addAllocation:table.addressBuffer]; - } - for (id buffer in table.expertBuffers) { - [residency_set addAllocation:buffer]; - } - } -#else - (void)residency_set; - (void)table; -#endif -} - -static id ds4_gpu_q4_expert_layer_residency_set(DS4MetalQ4ExpertTable *gate_table, - DS4MetalQ4ExpertTable *up_table, - DS4MetalQ4ExpertTable *down_table, - bool auto_queue_residency) { - const bool queue_residency = - ds4_gpu_q4_table_queue_residency_enabled(auto_queue_residency); - if (!gate_table || !up_table || !down_table || - !g_device || !g_q4_expert_layer_residency_cache || - (getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL && - !queue_residency)) { - return nil; - } -#if TARGET_OS_OSX - if (@available(macOS 15.0, *)) { - NSString *key = [NSString stringWithFormat:@"%p:%p:%p", - gate_table, up_table, down_table]; - DS4MetalQ4LayerResidency *cached = - [g_q4_expert_layer_residency_cache objectForKey:key]; - if (cached) return cached.residencySet; - - const NSUInteger capacity = - [gate_table.expertBuffers count] + - [up_table.expertBuffers count] + - [down_table.expertBuffers count] + 3u; - MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; - desc.label = @"ds4_q4_expert_layer"; - desc.initialCapacity = capacity; - NSError *error = nil; - id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; - if (!residency_set) { - fprintf(stderr, - "ds4: Metal Q4 expert layer residency set creation failed: %s\n", - [[error localizedDescription] UTF8String]); - return nil; - } - - ds4_gpu_q4_residency_add_table(residency_set, gate_table); - ds4_gpu_q4_residency_add_table(residency_set, up_table); - ds4_gpu_q4_residency_add_table(residency_set, down_table); - [residency_set commit]; - [residency_set requestResidency]; - - DS4MetalQ4LayerResidency *entry = [DS4MetalQ4LayerResidency new]; - entry.residencySet = residency_set; - if (queue_residency) { - [g_queue addResidencySet:residency_set]; - entry.addedToQueue = YES; - } - [g_q4_expert_layer_residency_cache setObject:entry forKey:key]; - return residency_set; - } -#endif - return nil; -} - -static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( - const void *model_map, - uint64_t model_size, - uint64_t tensor_offset, - uint64_t expert_bytes, - uint32_t n_total_expert, - id encoder); - -static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( - const void *model_map, - uint64_t model_size, - uint64_t tensor_offset, - uint64_t expert_bytes, - uint32_t n_total_expert); - -int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size, - uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, - uint64_t gate_expert_bytes, uint64_t down_expert_bytes, - uint32_t n_total_expert) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!model_map || model_size == 0 || gate_expert_bytes == 0 || - down_expert_bytes == 0 || n_total_expert == 0) { - return 0; - } - if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert table preload byte size overflow\n"); - return 0; - } - - const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; - const bool address_auto = - ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, - 6, - gate_tensor_bytes, - down_tensor_bytes); - const bool table_auto = - ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, - 6, - gate_tensor_bytes, - down_tensor_bytes); - if (!address_auto && !table_auto) { - return 1; - } - - if (address_auto) { - @autoreleasepool { - DS4MetalQ4ExpertTable *gate_table = - ds4_gpu_q4_expert_address_table(model_map, - model_size, - gate_offset, - gate_expert_bytes, - n_total_expert); - DS4MetalQ4ExpertTable *up_table = - ds4_gpu_q4_expert_address_table(model_map, - model_size, - up_offset, - gate_expert_bytes, - n_total_expert); - DS4MetalQ4ExpertTable *down_table = - ds4_gpu_q4_expert_address_table(model_map, - model_size, - down_offset, - down_expert_bytes, - n_total_expert); - if (!gate_table || !up_table || !down_table) { - return 0; - } - id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, - up_table, - down_table, - true); - if (!residency) { - fprintf(stderr, "ds4: Metal Q4 expert address table preload failed to create queue residency set\n"); - return 0; - } - } - return 1; - } - - if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder || - !g_moe_table_q4_sum_down_encoder) { - fprintf(stderr, "ds4: Metal Q4 expert table preload missing argument encoders\n"); - return 0; - } - - @autoreleasepool { - DS4MetalQ4ExpertTable *gate_table = - ds4_gpu_q4_expert_table(model_map, - model_size, - gate_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_gate_encoder); - DS4MetalQ4ExpertTable *up_table = - ds4_gpu_q4_expert_table(model_map, - model_size, - up_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_up_encoder); - DS4MetalQ4ExpertTable *down_table = - ds4_gpu_q4_expert_table(model_map, - model_size, - down_offset, - down_expert_bytes, - n_total_expert, - g_moe_table_q4_sum_down_encoder); - if (!gate_table || !up_table || !down_table) { - return 0; - } - id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, - up_table, - down_table, - true); - if (!residency) { - fprintf(stderr, "ds4: Metal Q4 expert table preload failed to create queue residency set\n"); - return 0; - } - } - return 1; -} - -static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( - const void *model_map, - uint64_t model_size, - uint64_t tensor_offset, - uint64_t expert_bytes, - uint32_t n_total_expert, - id encoder) { - if (!model_map || !g_device || !g_q4_expert_table_cache || !encoder || - model_size == 0 || expert_bytes == 0 || - n_total_expert == 0 || n_total_expert > 384) { - return nil; - } - if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert table byte size overflow\n"); - return nil; - } - const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; - if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert table is outside the mapped model\n"); - return nil; - } - - const uint32_t table_group_size = - ds4_gpu_q4_expert_table_group_size(n_total_expert); - NSString *key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu:%u:%llu:%u", - model_map, - (unsigned long long)model_size, - (unsigned long long)tensor_offset, - (unsigned long long)expert_bytes, - n_total_expert, - (unsigned long long)[encoder encodedLength], - table_group_size]; - DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; - if (cached) return cached; - - id arg_buffer = - [g_device newBufferWithLength:[encoder encodedLength] - options:MTLResourceStorageModeShared]; - if (!arg_buffer) { - fprintf(stderr, "ds4: Metal Q4 expert table argument buffer allocation failed\n"); - return nil; - } - arg_buffer.label = @"ds4_q4_expert_table"; - [encoder setArgumentBuffer:arg_buffer offset:0]; - - NSMutableArray> *expert_buffers = - [NSMutableArray arrayWithCapacity:table_group_size > 1 ? - (n_total_expert + table_group_size - 1u) / table_group_size : - n_total_expert]; - if (!expert_buffers) return nil; - - if (table_group_size > 1) { - for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { - const uint32_t remaining = n_total_expert - first; - const uint32_t group_n = - remaining < table_group_size ? remaining : table_group_size; - if ((uint64_t)first > UINT64_MAX / expert_bytes || - (uint64_t)group_n > UINT64_MAX / expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert table group byte overflow\n"); - return nil; - } - const uint64_t rel = (uint64_t)first * expert_bytes; - const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; - if (rel > UINT64_MAX - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert table group offset overflow\n"); - return nil; - } - uint64_t inner = 0; - id group_buf = - ds4_gpu_wrap_model_range(model_map, - model_size, - tensor_offset + rel, - group_bytes, - &inner); - if (!group_buf) return nil; - for (uint32_t j = 0; j < group_n; j++) { - const uint64_t expert_inner = inner + (uint64_t)j * expert_bytes; - [encoder setBuffer:group_buf offset:(NSUInteger)expert_inner atIndex:first + j]; - } - [expert_buffers addObject:group_buf]; - } - } else { - for (uint32_t i = 0; i < n_total_expert; i++) { - const uint64_t rel = (uint64_t)i * expert_bytes; - if (rel > UINT64_MAX - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert table offset overflow\n"); - return nil; - } - uint64_t inner = 0; - id expert_buf = - ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - tensor_offset + rel, - expert_bytes, - &inner); - if (!expert_buf) return nil; - [encoder setBuffer:expert_buf offset:(NSUInteger)inner atIndex:i]; - [expert_buffers addObject:expert_buf]; - } - } - [arg_buffer didModifyRange:NSMakeRange(0, [encoder encodedLength])]; - - if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal Q4 expert table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", - n_total_expert, - table_group_size, - (unsigned long)[expert_buffers count], - ds4_gpu_mib(expert_bytes)); - } - - DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; - table.argumentBuffer = arg_buffer; - table.expertBuffers = expert_buffers; - table.residencySet = ds4_gpu_q4_expert_table_residency_set(expert_buffers); - table.nExpert = n_total_expert; - table.expertBytes = expert_bytes; - [g_q4_expert_table_cache setObject:table forKey:key]; - return table; -} - -static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( - const void *model_map, - uint64_t model_size, - uint64_t tensor_offset, - uint64_t expert_bytes, - uint32_t n_total_expert) { - if (!model_map || !g_device || !g_q4_expert_table_cache || - model_size == 0 || expert_bytes == 0 || - n_total_expert == 0 || n_total_expert > 384) { - return nil; - } - if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert address table byte size overflow\n"); - return nil; - } - const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; - if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert address table is outside the mapped model\n"); - return nil; - } - - const uint32_t table_group_size = - ds4_gpu_q4_expert_table_group_size(n_total_expert); - NSString *key = [NSString stringWithFormat:@"addr:%p:%llu:%llu:%llu:%u:%u", - model_map, - (unsigned long long)model_size, - (unsigned long long)tensor_offset, - (unsigned long long)expert_bytes, - n_total_expert, - table_group_size]; - DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; - if (cached) return cached; - - id address_buffer = - [g_device newBufferWithLength:(NSUInteger)n_total_expert * sizeof(uint64_t) - options:MTLResourceStorageModeShared]; - if (!address_buffer) { - fprintf(stderr, "ds4: Metal Q4 expert address table allocation failed\n"); - return nil; - } - address_buffer.label = @"ds4_q4_expert_address_table"; - uint64_t *addresses = (uint64_t *)[address_buffer contents]; - if (!addresses) return nil; - - NSMutableArray> *expert_buffers = - [NSMutableArray arrayWithCapacity:table_group_size > 1 ? - (n_total_expert + table_group_size - 1u) / table_group_size : - n_total_expert]; - if (!expert_buffers) return nil; - -#if TARGET_OS_OSX - if (@available(macOS 13.0, *)) { - for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { - const uint32_t remaining = n_total_expert - first; - const uint32_t group_n = - remaining < table_group_size ? remaining : table_group_size; - if ((uint64_t)first > UINT64_MAX / expert_bytes || - (uint64_t)group_n > UINT64_MAX / expert_bytes) { - fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n"); - return nil; - } - const uint64_t rel = (uint64_t)first * expert_bytes; - const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; - if (rel > UINT64_MAX - tensor_offset) { - fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n"); - return nil; - } - uint64_t inner = 0; - id group_buf = nil; - if (table_group_size > 1) { - group_buf = ds4_gpu_wrap_model_range(model_map, - model_size, - tensor_offset + rel, - group_bytes, - &inner); - } else { - group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map, - model_size, - tensor_offset + rel, - expert_bytes, - &inner); - } - if (!group_buf) return nil; - const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner; - for (uint32_t j = 0; j < group_n; j++) { - addresses[first + j] = base_address + (uint64_t)j * expert_bytes; - } - [expert_buffers addObject:group_buf]; - } - } else -#endif - { - fprintf(stderr, "ds4: Metal GPU addresses require macOS 13 or newer\n"); - return nil; - } - - [address_buffer didModifyRange:NSMakeRange(0, - (NSUInteger)n_total_expert * sizeof(uint64_t))]; - - if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { - fprintf(stderr, - "ds4: Metal Q4 expert address table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", - n_total_expert, - table_group_size, - (unsigned long)[expert_buffers count], - ds4_gpu_mib(expert_bytes)); - } - - DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; - table.addressBuffer = address_buffer; - table.expertBuffers = expert_buffers; - table.nExpert = n_total_expert; - table.expertBytes = expert_bytes; - [g_q4_expert_table_cache setObject:table forKey:key]; - return table; -} - -static void ds4_gpu_use_q4_expert_table_resources( - id cb, - id enc, - DS4MetalQ4ExpertTable *table, - bool queue_residency) { - if (!enc || !table) return; - if (table.argumentBuffer) { - [enc useResource:table.argumentBuffer usage:MTLResourceUsageRead]; - } - if (table.addressBuffer) { - [enc useResource:table.addressBuffer usage:MTLResourceUsageRead]; - } - if (!queue_residency && - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && - table.residencySet && - cb && - [cb respondsToSelector:@selector(useResidencySet:)]) { - [cb useResidencySet:table.residencySet]; - } - if (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") == NULL && - getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") == NULL) { - return; - } - const NSUInteger count = [table.expertBuffers count]; - for (NSUInteger base = 0; base < count; base += 64u) { - const NSUInteger n = count - base < 64u ? count - base : 64u; - __unsafe_unretained id resources[64]; - for (NSUInteger i = 0; i < n; i++) { - resources[i] = [table.expertBuffers objectAtIndex:base + i]; - } - [enc useResources:resources count:n usage:MTLResourceUsageRead]; - } -} - -int ds4_gpu_indexer_score_one_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_head, - uint32_t head_dim, - float scale) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!scores || !q || !weights || !index_comp || - n_comp == 0 || n_head == 0 || head_dim == 0) { - return 0; - } - - @autoreleasepool { - const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); - const uint64_t weight_bytes = (uint64_t)n_head * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t score_bytes = (uint64_t)n_comp * sizeof(float); - id qbuf = ds4_gpu_tensor_buffer(q); - id wbuf = ds4_gpu_tensor_buffer(weights); - id compbuf = ds4_gpu_tensor_buffer(index_comp); - id scorebuf = ds4_gpu_tensor_buffer(scores); - if (!qbuf || !wbuf || !compbuf || !scorebuf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(weights) < weight_bytes || - ds4_gpu_tensor_bytes(index_comp) < comp_bytes || - ds4_gpu_tensor_bytes(scores) < score_bytes) { - fprintf(stderr, "ds4: Metal graph indexer score received undersized buffers\n"); - return 0; - } - - if (n_head == 64 && head_dim == 128) { - id direct_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline, - "kernel_dsv4_indexer_score_one_direct"); - if (!direct_pipeline) return 0; - - ds4_gpu_dsv4_indexer_scores_fused_args args = { - .n_comp = n_comp, - .n_tokens = 1, - .n_head = n_head, - .head_dim = head_dim, - .pos0 = 0, - .ratio = 4, - .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), - .q_head_stride = (uint64_t)head_dim * sizeof(float), - .weights_token_stride = (uint64_t)n_head * sizeof(float), - .index_row_stride = (uint64_t)head_dim * sizeof(float), - .score_token_stride = (uint64_t)n_comp * sizeof(float), - .scale = scale, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:direct_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; - [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; - [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer direct score")) return 0; - return 1; - } - - const uint64_t head_score_bytes = (uint64_t)n_comp * n_head * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_head_scores_buffer, - &g_indexer_head_scores_bytes, - (NSUInteger)head_score_bytes, - "ds4_indexer_head_scores")) { - return 0; - } - - ds4_gpu_q8_0_matvec_args dot_args = - ds4_gpu_make_f32_mv_args(head_dim, n_comp, n_head); - ds4_gpu_mv_dispatch dot_dispatch = - ds4_gpu_make_plain_mv_dispatch(head_dim, 1); - dot_args.nr0 = dot_dispatch.nr0; - id dot_pipeline = - ds4_gpu_get_mul_mv_pipeline(dot_dispatch.function_name, dot_dispatch.nsg); - if (!dot_pipeline) return 0; - ds4_gpu_dsv4_indexer_weighted_sum_args sum_args = { - .ne00 = (int64_t)n_comp, - .ne01 = 1, - .ne02 = (int64_t)n_head, - .nb00 = sizeof(float), - .nb01 = (uint64_t)n_comp * sizeof(float), - .nb02 = (uint64_t)n_comp * sizeof(float), - .ne10 = (int64_t)n_head, - .ne11 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)n_head * sizeof(float), - .ne0 = (int64_t)n_comp, - .ne1 = 1, - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_comp * sizeof(float), - .scale = scale, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:dot_pipeline]; - [enc setBytes:&dot_args length:sizeof(dot_args) atIndex:0]; - [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:1]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; - [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:3]; - if (dot_dispatch.smem) { - [enc setThreadgroupMemoryLength:dot_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + (NSUInteger)dot_dispatch.nr0 - 1u) / (NSUInteger)dot_dispatch.nr0, - n_head, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dot_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_indexer_weighted_sum_pipeline]; - [enc setBytes:&sum_args length:sizeof(sum_args) atIndex:0]; - [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 255u) / 256u, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer score")) return 0; - } - - return 1; -} - -static int ds4_gpu_indexer_scores_batch_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!scores || !q || !weights || !index_comp || - n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || ratio == 0) { - return 0; - } - - @autoreleasepool { - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t weight_bytes = (uint64_t)n_tokens * n_head * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); - id qbuf = ds4_gpu_tensor_buffer(q); - id wbuf = ds4_gpu_tensor_buffer(weights); - id compbuf = ds4_gpu_tensor_buffer(index_comp); - id scorebuf = ds4_gpu_tensor_buffer(scores); - if (!qbuf || !wbuf || !compbuf || !scorebuf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(weights) < weight_bytes || - ds4_gpu_tensor_bytes(index_comp) < comp_bytes || - ds4_gpu_tensor_bytes(scores) < score_bytes) { - fprintf(stderr, "ds4: Metal graph indexer prefill scores received undersized buffers\n"); - return 0; - } - if (head_dim != 128) { - fprintf(stderr, "ds4: Metal fused DS4 indexer scores expect 128-wide rows\n"); - return 0; - } - /* - * The NAX/TensorOps score builder is a prefill-only win. At small - * batches and in one-token decode the setup cost is not amortized, so - * those paths keep the older direct/tiled score kernels. - */ - const bool use_nax = ds4_gpu_mpp_available() && n_tokens >= 16u; - id pipeline = ds4_gpu_get_pipeline( - use_nax ? "kernel_dsv4_indexer_scores_nax" : - (g_quality_mode ? "kernel_dsv4_indexer_scores_tiled_f32" - : "kernel_dsv4_indexer_scores_tiled")); - if (!pipeline) return 0; - - ds4_gpu_dsv4_indexer_scores_fused_args args = { - .n_comp = n_comp, - .n_tokens = n_tokens, - .n_head = n_head, - .head_dim = head_dim, - .pos0 = pos0, - .ratio = ratio, - .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), - .q_head_stride = (uint64_t)head_dim * sizeof(float), - .weights_token_stride = (uint64_t)n_head * sizeof(float), - .index_row_stride = (uint64_t)head_dim * sizeof(float), - .score_token_stride = (uint64_t)n_comp * sizeof(float), - .scale = scale, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; - [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; - if (use_nax) { - const NSUInteger q_shared = 2u * 32u * 32u; - const NSUInteger k_shared = 32u * 128u; - const NSUInteger dot_shared = 32u * 32u; - [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + - dot_shared * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, - ((NSUInteger)n_tokens + 15u) / 16u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - } else if (g_quality_mode) { - const NSUInteger q_shared = 8u * 128u; - const NSUInteger k_shared = 32u * 128u; - const NSUInteger dot_shared = 8u * 32u; - [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, - ((NSUInteger)n_tokens + 7u) / 8u, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - } else { - const NSUInteger q_shared = 8u * 128u; - const NSUInteger k_shared = 32u * 128u; - const NSUInteger dot_shared = 8u * 32u; - [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + - dot_shared * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, - ((NSUInteger)n_tokens + 7u) / 8u, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer prefill scores")) return 0; - } - - return 1; -} - -int ds4_gpu_indexer_scores_prefill_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale) { - return ds4_gpu_indexer_scores_batch_tensor(scores, - q, - weights, - index_comp, - n_comp, - n_tokens, - 0, - n_head, - head_dim, - ratio, - scale); -} - -int ds4_gpu_indexer_scores_decode_batch_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *index_comp, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - uint32_t ratio, - float scale) { - return ds4_gpu_indexer_scores_batch_tensor(scores, - q, - weights, - index_comp, - n_comp, - n_tokens, - pos0, - n_head, - head_dim, - ratio, - scale); -} - -int ds4_gpu_indexer_topk_tensor( - ds4_gpu_tensor *selected, - const ds4_gpu_tensor *scores, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || top_k > n_comp) return 0; - - @autoreleasepool { - const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); - const uint64_t selected_bytes = (uint64_t)top_k * n_tokens * sizeof(uint32_t); - id scorebuf = ds4_gpu_tensor_buffer(scores); - id selbuf = ds4_gpu_tensor_buffer(selected); - if (!scorebuf || !selbuf || - ds4_gpu_tensor_bytes(scores) < score_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n"); - return 0; - } - NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; - if (max_threads == 0) max_threads = 256; - int32_t nth = 1; - while ((uint32_t)nth < n_comp && (uint64_t)2u * (uint64_t)nth <= (uint64_t)max_threads) { - nth *= 2; - } - const int32_t npr = (int32_t)((n_comp + (uint32_t)nth - 1u) / (uint32_t)nth); - const int32_t block_top_k = (int32_t)(top_k < (uint32_t)nth ? top_k : (uint32_t)nth); - int32_t work_width = (int32_t)top_k; - if (npr > 1) { - const int32_t last_block = (int32_t)n_comp - (npr - 1) * nth; - work_width = (npr - 1) * block_top_k + (last_block < block_top_k ? last_block : block_top_k); - } - const uint64_t scratch_row_bytes = (uint64_t)work_width * sizeof(uint32_t); - const bool one_pass = npr <= 1; - const uint64_t scratch_bytes = one_pass ? scratch_row_bytes * n_tokens : - 2u * scratch_row_bytes * n_tokens; - if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_topk_buffer, - &g_indexer_topk_bytes, - (NSUInteger)scratch_bytes, - "ds4_indexer_topk")) { - return 0; - } - - ds4_gpu_kargs_argsort args = { - .ne00 = (int32_t)n_comp, - .ne01 = (int32_t)n_tokens, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = (uint64_t)n_comp * sizeof(float), - .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), - .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), - .ne0 = work_width, - .ne1 = (int32_t)n_tokens, - .ne2 = 1, - .ne3 = 1, - .top_k = block_top_k, - }; - // kernel_argsort_f32_i32_desc stages the block's scores behind the - // index array: nth int32 indices + nth float scores. - const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u; - - NSUInteger cur_off = 0; - NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_argsort_f32_i32_desc_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; - [enc setBuffer:one_pass ? selbuf : g_indexer_topk_buffer - offset:one_pass ? ds4_gpu_tensor_offset(selected) : cur_off - atIndex:2]; - [enc setThreadgroupMemoryLength:smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)npr * n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - int32_t len = block_top_k; - while (len < work_width) { - const int32_t nm = (work_width + 2 * len - 1) / (2 * len); - const bool final_merge = nm == 1; - NSUInteger merge_threads = g_argsort_merge_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; - if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u; - if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len; - if (merge_threads == 0) merge_threads = 1; - - ds4_gpu_kargs_argsort_merge merge_args = { - .ne00 = (int64_t)n_comp, - .ne01 = (int64_t)n_tokens, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = (uint64_t)n_comp * sizeof(float), - .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), - .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), - .ne0 = work_width, - .ne1 = (int32_t)n_tokens, - .ne2 = 1, - .ne3 = 1, - .top_k = nm == 1 ? (int32_t)top_k : work_width, - .len = len, - }; - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_argsort_merge_f32_i32_desc_pipeline]; - [enc setBytes:&merge_args length:sizeof(merge_args) atIndex:0]; - [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; - [enc setBuffer:g_indexer_topk_buffer offset:cur_off atIndex:2]; - [enc setBuffer:final_merge ? selbuf : g_indexer_topk_buffer - offset:final_merge ? ds4_gpu_tensor_offset(selected) : next_off - atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nm * n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(merge_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - const NSUInteger tmp = cur_off; - cur_off = next_off; - next_off = tmp; - len <<= 1; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0; - } - - return 1; -} - -int ds4_gpu_argmax_tensor( - ds4_gpu_tensor *out_idx, - const ds4_gpu_tensor *logits, - uint32_t n_vocab) { - if (!out_idx || !logits || n_vocab == 0) return 0; - if (ds4_gpu_tensor_bytes(out_idx) < sizeof(int32_t) || - ds4_gpu_tensor_bytes(logits) < (uint64_t)n_vocab * sizeof(float)) { - fprintf(stderr, "ds4: Metal graph argmax received undersized buffers\n"); - return 0; - } - - return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1, 1); -} - -int ds4_gpu_dsv4_topk_mask_tensor( - ds4_gpu_tensor *mask, - const ds4_gpu_tensor *topk, - uint32_t n_comp, - uint32_t n_tokens, - uint32_t top_k) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0) return 0; - - @autoreleasepool { - const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); - const uint64_t mask_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); - id topkbuf = ds4_gpu_tensor_buffer(topk); - id maskbuf = ds4_gpu_tensor_buffer(mask); - if (!topkbuf || !maskbuf || - ds4_gpu_tensor_bytes(topk) < topk_bytes || - ds4_gpu_tensor_bytes(mask) < mask_bytes) { - fprintf(stderr, "ds4: Metal dsv4 top-k mask received undersized buffers\n"); - return 0; - } - - ds4_gpu_dsv4_topk_mask_args args = { - .ne00 = (int64_t)top_k, - .ne01 = (int64_t)n_tokens, - .nb00 = sizeof(int32_t), - .nb01 = (uint64_t)top_k * sizeof(int32_t), - .ne0 = (int64_t)n_comp, - .ne1 = (int64_t)n_tokens, - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_comp * sizeof(float), - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_topk_mask_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; - [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)n_comp * n_tokens) + 255u) / 256u, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_topk_mask_scatter_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; - [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)top_k * n_tokens) + 255u) / 256u, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "dsv4 top-k mask")) return 0; - } - - return 1; -} - -static int ds4_gpu_matmul_q8_0_legacy_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_tok, - bool prefer_decode_mpp, - bool force_model_view) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if ((in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal Q8_0 tensor matmul received undersized activation buffers\n"); - return 0; - } - - const uint64_t blocks = in_dim / 32; - const uint64_t row_bytes = blocks * 34; - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8_0 tensor matmul range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = force_model_view ? - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset) : - ds4_gpu_wrap_q8_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); - if (!wbuf) { - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (n_tok == 1) { - if (ds4_gpu_mpp_available() && - (prefer_decode_mpp || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) && - getenv("DS4_METAL_DISABLE_Q8_DECODE_MPP") == NULL && - (in_dim % 64u) == 0) { - const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs"; - id mpp_pipeline = - ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); - if (mpp_pipeline) { - ds4_gpu_mul_mm_args args = - ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:mpp_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1u, - ((NSUInteger)out_dim + 63u) / 64u, - 1u) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode MPP matmul")) { - return 0; - } - return 1; - } - ds4_gpu_warn_mpp_fallback(); - } - - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - if (out_dim > 65536u) mv_dispatch.nsg = 8; - const bool force_output_nr4 = - getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL; - const bool output_shape = - in_dim == 4096u && out_dim == 129280u; - const bool use_output_nr4 = - !g_quality_mode && (out_dim & 3u) == 0u && - (force_output_nr4 || - (output_shape && ds4_gpu_device_name_contains("M3"))) && - getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL; - if (use_output_nr4) { - mv_dispatch.function_name = - "kernel_mul_mv_q8_0_f32_nr4"; - mv_dispatch.nr0 = 4; - mv_dispatch.smem = 32u * 4u * sizeof(float); - } - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matvec")) { - return 0; - } - return 1; - } - - const uint64_t mv_ext_max_tokens = - ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); - if (n_tok <= mv_ext_max_tokens && (in_dim % 128u) == 0) { - const int16_t nsg = 2; - const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); - const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); - const char *fn_name = ds4_gpu_mv_ext_name(1, r1ptg); - id pipeline = - fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; - if (!pipeline) return 0; - - const int16_t nypsg = 32 / nxpsg; - const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; - ds4_gpu_mul_mv_ext_args args = - ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, - ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor mul_mv_ext")) { - return 0; - } - return 1; - } - - /* - * Dense Q8_0 prefill is the cleanest DS4 TensorOps shape: M/N/K are - * aligned and the RHS activation matrix is already dense. The retained - * kernel dequantizes each 64x32 weight tile to half in threadgroup - * memory, then uses direct-RHS MPP for the activation tile. This avoids - * staging RHS into threadgroup memory and was the direct replacement for - * the slower generic MPP prototype. - */ - if (ds4_gpu_mpp_available() && - n_tok >= 32u && - (in_dim % 64u) == 0 && - (out_dim % 64u) == 0 && - (n_tok % 32u) == 0) { - uint64_t nax_tile_n = 32u; - if ((n_tok % 128u) == 0) { - nax_tile_n = 128u; - } else if ((n_tok % 64u) == 0) { - nax_tile_n = 64u; - } - const char *nax_fn = nax_tile_n == 128u - ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128" - : (nax_tile_n == 64u - ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64" - : "kernel_mul_mm_q8_0_f32_nax_direct_rhs"); - id pipeline = - ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); - if (pipeline) { - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), - (NSUInteger)out_dim / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) { - return 0; - } - return 1; - } - ds4_gpu_warn_mpp_fallback(); - } - - const bool bc_inp = (in_dim % 32u) != 0; - const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; - id pipeline = - ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_q8_0_f32", bc_inp, bc_out); - if (!pipeline) return 0; - - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, - ((NSUInteger)out_dim + 63u) / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matmul")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_matmul_q8_0_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_tok) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if ((in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { - return 0; - } - - const int profile_requested = - n_tok > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0; - int profile_prefill = 0; - int split_batch_for_profile = 0; - const char *profile_label = NULL; - char profile_label_buf[128]; - char profile_fallback[128]; - if (profile_requested) { - snprintf(profile_fallback, sizeof(profile_fallback), - "q8 weight_off=%llu in=%llu out=%llu tok=%llu", - (unsigned long long)weight_offset, - (unsigned long long)in_dim, - (unsigned long long)out_dim, - (unsigned long long)n_tok); - snprintf(profile_label_buf, sizeof(profile_label_buf), "%s", profile_fallback); - profile_label = profile_label_buf; - const char *profile_filter = getenv("DS4_METAL_Q8_PREFILL_PROFILE_FILTER"); - profile_prefill = - profile_requested && - (!profile_filter || !profile_filter[0] || - strstr(profile_label, profile_filter) != NULL); - } - if (profile_prefill) { - if (g_batch_cb) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - split_batch_for_profile = 1; - } - } - - const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0; - int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, - weight_offset, in_dim, out_dim, - x, n_tok, false, false); - if (profile_prefill) { - if (split_batch_for_profile && ds4_gpu_end_commands() == 0) { - ok = 0; - } - const double elapsed_ms = ds4_gpu_now_ms() - profile_t0; - fprintf(stderr, - "ds4: Metal Q8_0 prefill profile %s in=%llu out=%llu tok=%llu %.3f ms\n", - profile_label ? profile_label : profile_fallback, - (unsigned long long)in_dim, - (unsigned long long)out_dim, - (unsigned long long)n_tok, - elapsed_ms); - if (split_batch_for_profile && ds4_gpu_begin_commands() == 0) { - ok = 0; - } - } - return ok; -} - -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) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || n_rows == 0 || - n_rows > INT32_MAX || in_dim == 0 || out_dim == 0 || - (in_dim & 31u) != 0 || in_dim > UINT32_MAX || - out_dim > UINT32_MAX || - in_dim > UINT64_MAX / n_rows / sizeof(float) || - out_dim > UINT64_MAX / n_rows / sizeof(float) || - ds4_gpu_tensor_bytes(x) < - (uint64_t)n_rows * in_dim * sizeof(float) || - ds4_gpu_tensor_bytes(out) < - (uint64_t)n_rows * out_dim * sizeof(float)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t blocks = in_dim / 32u; - const uint64_t row_bytes = blocks * 34u; - if (!xbuf || !outbuf || - out_dim > UINT64_MAX / row_bytes) { - return 0; - } - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = ds4_gpu_wrap_q8_decode_model_range( - model_map, model_size, weight_offset, weight_bytes, 1u, - &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - if (out_dim > 65536u) dispatch.nsg = 8; - ds4_gpu_q8_0_matvec_args args = - ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - args.ne11 = (int32_t)n_rows; - args.nb12 = (uint64_t)n_rows * in_dim * sizeof(float); - args.nb13 = args.nb12; - args.ne1 = (int32_t)n_rows; - args.nr0 = dispatch.nr0; - - id pipeline = - ds4_gpu_get_mul_mv_pipeline(dispatch.function_name, dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:dispatch.smem atIndex:0]; - [enc dispatchThreadgroups: - MTLSizeMake(((NSUInteger)out_dim + - (NSUInteger)dispatch.nr0 - 1u) / - (NSUInteger)dispatch.nr0, - (NSUInteger)n_rows, - 1) - threadsPerThreadgroup: - MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return ds4_gpu_finish_command_buffer( - cb, owned, "Q8_0 exact decode-row matvec"); - } -} - -int ds4_gpu_matmul_q8_0_decode_mpp_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_tok) { - return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, - weight_offset, in_dim, out_dim, - x, n_tok, true, false); -} - -int ds4_gpu_matmul_q8_0_decode_mpp_model_view_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_tok) { - return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, - weight_offset, in_dim, out_dim, - x, n_tok, true, true); -} - -static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) { - const char *prefix = NULL; - if (weight_type == DS4_METAL_TENSOR_Q4_K) { - prefix = "kernel_mul_mv_ext_q4_K_f32_r1_"; - } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { - prefix = "kernel_mul_mv_ext_q4_0_f32_r1_"; - } else { - return NULL; - } - switch (r1ptg) { - case 1: return weight_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mv_ext_q4_K_f32_r1_1" : "kernel_mul_mv_ext_q4_0_f32_r1_1"; - case 2: return weight_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mv_ext_q4_K_f32_r1_2" : "kernel_mul_mv_ext_q4_0_f32_r1_2"; - case 3: return weight_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mv_ext_q4_K_f32_r1_3" : "kernel_mul_mv_ext_q4_0_f32_r1_3"; - case 4: return weight_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mv_ext_q4_K_f32_r1_4" : "kernel_mul_mv_ext_q4_0_f32_r1_4"; - case 5: return weight_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mv_ext_q4_K_f32_r1_5" : "kernel_mul_mv_ext_q4_0_f32_r1_5"; - default: - (void)prefix; - return NULL; - } -} - -static const char *ds4_gpu_q4_mm_name(uint32_t weight_type) { - switch (weight_type) { - case DS4_METAL_TENSOR_Q4_0: return "kernel_mul_mm_q4_0_f32"; - case DS4_METAL_TENSOR_Q4_K: return "kernel_mul_mm_q4_K_f32"; - default: return NULL; - } -} - -static const char *ds4_gpu_q4_nax_name(uint32_t weight_type, uint64_t tile_n) { - if (weight_type == DS4_METAL_TENSOR_Q4_0) { - return tile_n == 128u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128" : - tile_n == 64u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64" : - "kernel_mul_mm_q4_0_f32_nax_direct_rhs"; - } - if (weight_type == DS4_METAL_TENSOR_Q4_K) { - return tile_n == 128u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128" : - tile_n == 64u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64" : - "kernel_mul_mm_q4_K_f32_nax_direct_rhs"; - } - return NULL; -} - -static int ds4_gpu_matmul_quant_impl_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, - bool prefer_decode_mpp, - bool force_model_view) { - if (weight_type == DS4_METAL_TENSOR_Q8_0) { - return ds4_gpu_matmul_q8_0_legacy_tensor(out, - model_map, - model_size, - weight_offset, - in_dim, - out_dim, - x, - n_tok, - prefer_decode_mpp, - force_model_view); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || - in_dim == 0 || out_dim == 0 || n_tok == 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { - return 0; - } - - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, (uint32_t)in_dim, &row_bytes)) { - fprintf(stderr, "ds4: Metal quant matmul received unsupported type/dim (%u, in=%llu)\n", - weight_type, - (unsigned long long)in_dim); - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal quant tensor matmul received undersized activation buffers\n"); - return 0; - } - - if (out_dim > UINT64_MAX / row_bytes) return 0; - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal quant tensor matmul range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = force_model_view ? - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - /* - * Small-batch Q4_K goes to the classic (llama.cpp-style) matvec: - * the mul_mv_ext family tops out around 220 GB/s on M5 for the GLM - * DenseQ4 decode shapes while this impl streams 530-650 GB/s - * (misc/q4mv_bench.m). Falls through to ext when unavailable. - */ - if (weight_type == DS4_METAL_TENSOR_Q4_K && - n_tok <= 8u && - (in_dim % 256u) == 0 && - getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { - const int16_t nsg = 2; - id pipeline = - ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); - if (pipeline) { - ds4_gpu_q8_0_matvec_args args = { - .ne00 = (int32_t)in_dim, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = 1, - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)in_dim, - .ne11 = (int32_t)n_tok, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = in_dim * sizeof(float), - .nb12 = in_dim * n_tok * sizeof(float), - .nb13 = in_dim * n_tok * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = (int32_t)n_tok, - .nr0 = 2, - .r2 = 1, - .r3 = 1, - }; - const uint64_t rows_ptg = (uint64_t)nsg * 2u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:32 atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, - (NSUInteger)n_tok, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K classic mul_mv")) return 0; - return 1; - } - } - - if (n_tok <= 8u && (in_dim % 128u) == 0) { - const int16_t nsg = 2; - const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); - const int16_t r1ptg = (n_tok == 1u) ? 1 : ds4_gpu_mv_ext_r1ptg(n_tok); - const char *fn_name = ds4_gpu_q4_mv_ext_name(weight_type, r1ptg); - id pipeline = - fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; - if (pipeline) { - const int16_t nypsg = 32 / nxpsg; - const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; - ds4_gpu_mul_mv_ext_args args = - ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, row_bytes, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, - ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor mul_mv_ext")) return 0; - return 1; - } - } - - if (ds4_gpu_mpp_available() && - n_tok >= 32u && - (in_dim % 64u) == 0 && - (out_dim % 64u) == 0 && - (n_tok % 32u) == 0) { - uint64_t nax_tile_n = 32u; - if ((n_tok % 128u) == 0) { - nax_tile_n = 128u; - } else if ((n_tok % 64u) == 0) { - nax_tile_n = 64u; - } - const char *nax_fn = ds4_gpu_q4_nax_name(weight_type, nax_tile_n); - id pipeline = - nax_fn ? ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false) : nil; - if (pipeline) { - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), - (NSUInteger)out_dim / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 NAX tensor matmul")) return 0; - return 1; - } - ds4_gpu_warn_mpp_fallback(); - } - - const char *mm_fn = ds4_gpu_q4_mm_name(weight_type); - const bool bc_inp = (in_dim % 32u) != 0; - const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; - id pipeline = - mm_fn ? ds4_gpu_get_mul_mm_pipeline(mm_fn, bc_inp, bc_out) : nil; - if (!pipeline) return 0; - - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, - ((NSUInteger)out_dim + 63u) / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor matmul")) return 0; - } - - return 1; -} - -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) { - return ds4_gpu_matmul_quant_impl_tensor(out, - model_map, - model_size, - weight_offset, - weight_type, - in_dim, - out_dim, - x, - n_tok, - false, - false); -} - -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) { - /* GLM decode dense matvecs: the classic Q8_0 kernels stream 570-613 - * GB/s on the GLM shapes while the MPP/nax matrix path measures ~150 - * GB/s at n_tok=1 (TP decode 8.75 -> 16.64 t/s on the IQ2+Q8 gguf). - * MPP stays available as an opt-in via DS4_METAL_Q8_DECODE_MPP. */ - const bool prefer_mpp = getenv("DS4_METAL_Q8_DECODE_MPP") != NULL; - return ds4_gpu_matmul_quant_impl_tensor(out, - model_map, - model_size, - weight_offset, - weight_type, - in_dim, - out_dim, - x, - n_tok, - prefer_mpp, - true); -} - -int ds4_gpu_matmul_quant_rows_scalar_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) { - return ds4_gpu_matmul_quant_tensor(out, - model_map, - model_size, - weight_offset, - weight_type, - in_dim, - out_dim, - x, - n_tok); -} - -int ds4_gpu_matmul_q8_0_rows_scalar_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_tok) { - if (n_tok == 1) { - return ds4_gpu_matmul_q8_0_tensor(out, - model_map, - model_size, - weight_offset, - in_dim, - out_dim, - x, - 1); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || - n_tok == 0 || n_tok > INT32_MAX || - (in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || - in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float)) { - return 0; - } - - const uint64_t x_row_bytes = in_dim * sizeof(float); - const uint64_t out_row_bytes = out_dim * sizeof(float); - if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || - (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t x_bytes = n_tok * x_row_bytes; - const uint64_t out_bytes = n_tok * out_row_bytes; - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul received undersized activation buffers\n"); - return 0; - } - - const uint64_t blocks = in_dim / 32; - if (blocks > UINT64_MAX / 34u) return 0; - const uint64_t row_bytes = blocks * 34u; - if (out_dim > UINT64_MAX / row_bytes) return 0; - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - if (out_dim > 65536u) mv_dispatch.nsg = 8; - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - - const NSUInteger x_base = ds4_gpu_tensor_offset(x); - const NSUInteger out_base = ds4_gpu_tensor_offset(out); - for (uint64_t t = 0; t < n_tok; t++) { - [enc setBuffer:xbuf offset:x_base + (NSUInteger)(t * x_row_bytes) atIndex:2]; - [enc setBuffer:outbuf offset:out_base + (NSUInteger)(t * out_row_bytes) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 scalar-row matvecs")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_matmul_q8_0_pair_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, - uint64_t n_tok) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out0 || !out1 || !model_map || !x || n_tok != 1 || - out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id out0buf = ds4_gpu_tensor_buffer(out0); - id out1buf = ds4_gpu_tensor_buffer(out1); - const uint64_t x_bytes = in_dim * sizeof(float); - const uint64_t out0_bytes = out0_dim * sizeof(float); - const uint64_t out1_bytes = out1_dim * sizeof(float); - if (!xbuf || !out0buf || !out1buf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out0) < out0_bytes || - ds4_gpu_tensor_bytes(out1) < out1_bytes) { - fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n"); - return 0; - } - - const uint64_t row_bytes = (in_dim / 32u) * 34u; - const uint64_t weight0_bytes = out0_dim * row_bytes; - const uint64_t weight1_bytes = out1_dim * row_bytes; - if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset || - weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) { - fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n"); - return 0; - } - - uint64_t inner0 = 0; - uint64_t inner1 = 0; - id weight0buf = - ds4_gpu_wrap_model_range(model_map, model_size, - weight0_offset, weight0_bytes, &inner0); - id weight1buf = - ds4_gpu_wrap_model_range(model_map, model_size, - weight1_offset, weight1_bytes, &inner1); - if (!weight0buf || !weight1buf) return 0; - - ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch(); - ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch(); - if (out0_dim > 65536u) dispatch0.nsg = 8; - if (out1_dim > 65536u) dispatch1.nsg = 8; - /* A common threadgroup shape is required to retain each standalone - * reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */ - if (dispatch0.nsg != dispatch1.nsg) return 0; - - ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim); - ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim); - args0.nr0 = dispatch0.nr0; - args1.nr0 = dispatch1.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args0 length:sizeof(args0) atIndex:0]; - [enc setBytes:&args1 length:sizeof(args1) atIndex:1]; - [enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2]; - [enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; - [enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5]; - [enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6]; - [enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0]; - const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim + - (NSUInteger)dispatch0.nr0 - 1u) / - (NSUInteger)dispatch0.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0; - } - - return 1; -} - -int ds4_gpu_matmul_q8_0_f16_out_tensor( - ds4_gpu_tensor *out_h, - 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_tok) { - (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; - (void)in_dim; (void)out_dim; (void)x; (void)n_tok; - return 0; -} - -static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp, - int store_gate_up, - bool force_model_view) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!mid || !x || !model_map || - (store_gate_up && (!gate || !up)) || - (in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || - !isfinite(clamp) || clamp < 0.0f) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id midbuf = ds4_gpu_tensor_buffer(mid); - id gatebuf = store_gate_up ? - ds4_gpu_tensor_buffer(gate) : midbuf; - id upbuf = store_gate_up ? - ds4_gpu_tensor_buffer(up) : midbuf; - const uint64_t x_bytes = in_dim * sizeof(float); - const uint64_t out_bytes = out_dim * sizeof(float); - if (!xbuf || !gatebuf || !upbuf || !midbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - (store_gate_up && ds4_gpu_tensor_bytes(gate) < out_bytes) || - (store_gate_up && ds4_gpu_tensor_bytes(up) < out_bytes) || - ds4_gpu_tensor_bytes(mid) < out_bytes) { - fprintf(stderr, "ds4: Metal shared expert fused gate/up received undersized activation buffers\n"); - return 0; - } - - const uint64_t blocks = in_dim / 32; - const uint64_t row_bytes = blocks * 34; - const uint64_t weight_bytes = out_dim * row_bytes; - if (gate_offset > model_size || weight_bytes > model_size - gate_offset || - up_offset > model_size || weight_bytes > model_size - up_offset) { - fprintf(stderr, "ds4: Metal shared expert fused gate/up range is outside the mapped model\n"); - return 0; - } - - const bool exact_decode_views = - !force_model_view && - getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - id gate_wbuf = exact_decode_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset, - weight_bytes, - &gate_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset, - weight_bytes, - &gate_inner); - id up_wbuf = exact_decode_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset, - weight_bytes, - &up_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset, - weight_bytes, - &up_inner); - if (!gate_wbuf || !up_wbuf) return 0; - - ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - args.nr0 = mv_dispatch.nr0; - const char *fn_name = store_gate_up ? - (mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : - "kernel_dsv4_shared_gate_up_swiglu_q8_0") : - (mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_mid_swiglu_q8_0_r4" : - "kernel_dsv4_shared_mid_swiglu_q8_0"); - id pipeline = - ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; - [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:gatebuf offset:(store_gate_up ? - ds4_gpu_tensor_offset(gate) : - ds4_gpu_tensor_offset(mid)) atIndex:4]; - [enc setBuffer:upbuf offset:(store_gate_up ? - ds4_gpu_tensor_offset(up) : - ds4_gpu_tensor_offset(mid)) atIndex:5]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; - [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; - [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, - owned, - store_gate_up ? - "shared expert fused gate/up" : - "shared expert fused mid")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, - up, - mid, - model_map, - model_size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp, - 1, - false); -} - -int ds4_gpu_shared_mid_swiglu_q8_0_tensor( - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL, - NULL, - mid, - model_map, - model_size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp, - 0, - true); -} - -int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - float clamp) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, - up, - mid, - model_map, - model_size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp, - 1, - true); -} - -int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok, - float clamp) { - if (n_tok == 1) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, - up, - mid, - model_map, - model_size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!gate || !up || !mid || !x || !model_map || - n_tok == 0 || - (in_dim & 31u) != 0 || (in_dim % 128u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || - !isfinite(clamp) || clamp < 0.0f) { - return 0; - } - - const uint64_t mv_ext_max_tokens = - ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); - if (n_tok > mv_ext_max_tokens) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id gatebuf = ds4_gpu_tensor_buffer(gate); - id upbuf = ds4_gpu_tensor_buffer(up); - id midbuf = ds4_gpu_tensor_buffer(mid); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (!xbuf || !gatebuf || !upbuf || !midbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(gate) < out_bytes || - ds4_gpu_tensor_bytes(up) < out_bytes || - ds4_gpu_tensor_bytes(mid) < out_bytes) { - fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows received undersized activation buffers\n"); - return 0; - } - - const uint64_t blocks = in_dim / 32; - const uint64_t row_bytes = blocks * 34; - const uint64_t weight_bytes = out_dim * row_bytes; - if (gate_offset > model_size || weight_bytes > model_size - gate_offset || - up_offset > model_size || weight_bytes > model_size - up_offset) { - fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows range is outside the mapped model\n"); - return 0; - } - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - id gate_wbuf = - ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); - id up_wbuf = - ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); - if (!gate_wbuf || !up_wbuf) return 0; - - const int16_t nsg = 2; - const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); - const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); - const char *fn_name = ds4_gpu_mv_ext_q8_pair_swiglu_name(r1ptg); - id pipeline = - fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; - if (!pipeline) return 0; - - const int16_t nypsg = 32 / nxpsg; - const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; - ds4_gpu_mul_mv_ext_args args = - ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; - [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; - [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; - [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / - (NSUInteger)r0ptg, - ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / - (NSUInteger)r1ptg, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "fused Q8_0 gate/up rows")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok, - float clamp) { - if (n_tok == 1) { - return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, - up, - mid, - model_map, - model_size, - gate_offset, - up_offset, - in_dim, - out_dim, - x, - clamp); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!gate || !up || !mid || !x || !model_map || - n_tok == 0 || - (in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || - in_dim > UINT64_MAX / sizeof(float) || - out_dim > UINT64_MAX / sizeof(float) || - !isfinite(clamp) || clamp < 0.0f) { - return 0; - } - - const uint64_t x_row_bytes = in_dim * sizeof(float); - const uint64_t out_row_bytes = out_dim * sizeof(float); - if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || - (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id gatebuf = ds4_gpu_tensor_buffer(gate); - id upbuf = ds4_gpu_tensor_buffer(up); - id midbuf = ds4_gpu_tensor_buffer(mid); - const uint64_t x_bytes = n_tok * x_row_bytes; - const uint64_t out_bytes = n_tok * out_row_bytes; - if (!xbuf || !gatebuf || !upbuf || !midbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(gate) < out_bytes || - ds4_gpu_tensor_bytes(up) < out_bytes || - ds4_gpu_tensor_bytes(mid) < out_bytes) { - fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up received undersized activation buffers\n"); - return 0; - } - - const uint64_t blocks = in_dim / 32; - if (blocks > UINT64_MAX / 34u) return 0; - const uint64_t row_bytes = blocks * 34u; - if (out_dim > UINT64_MAX / row_bytes) return 0; - const uint64_t weight_bytes = out_dim * row_bytes; - if (gate_offset > model_size || weight_bytes > model_size - gate_offset || - up_offset > model_size || weight_bytes > model_size - up_offset) { - fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up range is outside the mapped model\n"); - return 0; - } - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - id gate_wbuf = - ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); - id up_wbuf = - ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); - if (!gate_wbuf || !up_wbuf) return 0; - - ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - args.nr0 = mv_dispatch.nr0; - args.ne11 = (int32_t)n_tok; - args.nb12 = n_tok * x_row_bytes; - args.nb13 = args.nb12; - args.ne1 = (int32_t)n_tok; - const char *fn_name = mv_dispatch.nr0 >= 4 ? - "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : - "kernel_dsv4_shared_gate_up_swiglu_q8_0"; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; - [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; - [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; - [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; - [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake( - ((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - (NSUInteger)n_tok, - 1) - threadsPerThreadgroup: - MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "shared expert scalar-row fused gate/up")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_matmul_f16_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_tok) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal F16 tensor matmul received undersized activation buffers\n"); - return 0; - } - - const uint64_t row_bytes = in_dim * sizeof(uint16_t); - const uint64_t weight_bytes = row_bytes * out_dim; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal F16 tensor matmul range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_f32_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (n_tok == 1) { - ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = - ds4_gpu_make_plain_mv_dispatch(in_dim, 0); - if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { - mv_dispatch.nr0 = 4; - mv_dispatch.smem = 32u * 4u * sizeof(float); - } - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - if (mv_dispatch.smem) { - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matvec")) return 0; - return 1; - } - - if (n_tok <= 8 && (in_dim % 128u) == 0) { - const int16_t nsg = 2; - const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); - const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); - const char *fn_name = ds4_gpu_mv_ext_name(0, r1ptg); - id pipeline = - fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; - if (!pipeline) return 0; - - const int16_t nypsg = 32 / nxpsg; - const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; - ds4_gpu_mul_mv_ext_args args = - ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(uint16_t), row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, - ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor mul_mv_ext")) return 0; - return 1; - } - - /* - * Same direct-RHS TensorOps structure as Q8_0, but for F16 model - * matrices. The 128-token RHS tile is kept when the batch alignment - * allows it because the later tile_n=64 retest was neutral/slower. - */ - if (ds4_gpu_mpp_available() && - n_tok >= 32u && - (in_dim % 32u) == 0 && - (out_dim % 64u) == 0 && - (n_tok % 32u) == 0) { - uint64_t nax_tile_n = 32u; - if ((n_tok % 128u) == 0) { - nax_tile_n = 128u; - } else if ((n_tok % 64u) == 0) { - nax_tile_n = 64u; - } - const char *nax_fn = nax_tile_n == 128u - ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n128" - : (nax_tile_n == 64u - ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n64" - : "kernel_mul_mm_f16_f32_mpp_direct_rhs"); - id pipeline = - ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); - if (pipeline) { - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), - (NSUInteger)out_dim / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 NAX tensor matmul")) { - return 0; - } - return 1; - } - ds4_gpu_warn_mpp_fallback(); - } - - const bool bc_inp = (in_dim % 32u) != 0; - const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; - id pipeline = - ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f16_f32", bc_inp, bc_out); - if (!pipeline) return 0; - - ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, - ((NSUInteger)out_dim + 63u) / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matmul")) return 0; - } - - return 1; -} - -int ds4_gpu_matmul_f16_pair_tensor( - ds4_gpu_tensor *out_a, - ds4_gpu_tensor *out_b, - const void *model_map, - uint64_t model_size, - uint64_t weight_a_offset, - uint64_t weight_b_offset, - uint64_t in_dim, - uint64_t out_dim, - const ds4_gpu_tensor *x, - uint64_t n_tok) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok != 1 || (in_dim & 3u) != 0) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outabuf = ds4_gpu_tensor_buffer(out_a); - id outbbuf = ds4_gpu_tensor_buffer(out_b); - const uint64_t x_bytes = in_dim * sizeof(float); - const uint64_t out_bytes = out_dim * sizeof(float); - if (!xbuf || !outabuf || !outbbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out_a) < out_bytes || - ds4_gpu_tensor_bytes(out_b) < out_bytes) { - fprintf(stderr, "ds4: Metal F16 paired matvec received undersized activation buffers\n"); - return 0; - } - - const uint64_t row_bytes = in_dim * sizeof(uint16_t); - const uint64_t weight_bytes = row_bytes * out_dim; - if (weight_a_offset > model_size || weight_bytes > model_size - weight_a_offset || - weight_b_offset > model_size || weight_bytes > model_size - weight_b_offset) { - fprintf(stderr, "ds4: Metal F16 paired matvec range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_a = 0; - uint64_t inner_b = 0; - id wabuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_a_offset, weight_bytes, - &inner_a); - id wbbuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_b_offset, weight_bytes, - &inner_b); - if (!wabuf || !wbbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0); - if (ds4_gpu_use_compressor_pair_nr4() && - (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { - mv_dispatch.nr0 = 4; - mv_dispatch.smem = 32u * 4u * sizeof(float); - } - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_f16_f32_pair_4", mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wabuf offset:(NSUInteger)inner_a atIndex:1]; - [enc setBuffer:wbbuf offset:(NSUInteger)inner_b atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:outabuf offset:ds4_gpu_tensor_offset(out_a) atIndex:4]; - [enc setBuffer:outbbuf offset:ds4_gpu_tensor_offset(out_b) atIndex:5]; - if (mv_dispatch.smem) { - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 paired matvec")) return 0; - } - - return 1; -} - -int ds4_gpu_matmul_f16_pair_compressor_store_tensor( - ds4_gpu_tensor *out_kv, - ds4_gpu_tensor *out_score, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const void *model_map, - uint64_t model_size, - uint64_t weight_kv_offset, - uint64_t weight_score_offset, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t in_dim, - uint32_t width, - const ds4_gpu_tensor *x, - uint32_t ratio, - uint32_t pos) { - if (!g_initialized && !ds4_gpu_init()) return -1; - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL; - if ((g_quality_mode || - (!ds4_gpu_device_name_contains("M3") && - !ds4_gpu_device_name_contains("M5") && !force)) || - getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL || - getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL || - getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) { - return 0; - } - if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x || - in_dim == 0 || width == 0 || ratio == 0 || - (ape_type != 0u && ape_type != 1u)) { - return -1; - } - if (in_dim != 4096u || - (width != 256u && width != 512u && width != 1024u) || - (ratio != 4u && ratio != 128u)) { - return 0; - } - - @autoreleasepool { - const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; - const uint64_t row_bytes = in_dim * sizeof(uint16_t); - const uint64_t weight_bytes = row_bytes * width; - const uint64_t out_bytes = (uint64_t)width * sizeof(float); - const uint64_t state_bytes = - (uint64_t)state_rows * width * sizeof(float); - const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float); - const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem; - if (weight_kv_offset > model_size || - weight_bytes > model_size - weight_kv_offset || - weight_score_offset > model_size || - weight_bytes > model_size - weight_score_offset || - ape_offset > model_size || ape_bytes > model_size - ape_offset) { - return -1; - } - - id xbuf = ds4_gpu_tensor_buffer(x); - id outkvbuf = ds4_gpu_tensor_buffer(out_kv); - id outscorebuf = ds4_gpu_tensor_buffer(out_score); - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf || - ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || - ds4_gpu_tensor_bytes(out_kv) < out_bytes || - ds4_gpu_tensor_bytes(out_score) < out_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes) { - return -1; - } - - uint64_t weight_kv_inner = 0; - uint64_t weight_score_inner = 0; - uint64_t ape_inner = 0; - id weightkvbuf = ds4_gpu_wrap_model_range( - model_map, model_size, weight_kv_offset, weight_bytes, - &weight_kv_inner); - id weightscorebuf = ds4_gpu_wrap_model_range( - model_map, model_size, weight_score_offset, weight_bytes, - &weight_score_inner); - id apebuf = ds4_gpu_wrap_model_range( - model_map, model_size, ape_offset, ape_bytes, &ape_inner); - if (!weightkvbuf || !weightscorebuf || !apebuf) return -1; - - ds4_gpu_f16_matvec_args mv_args = - ds4_gpu_make_f16_mv_args(in_dim, width); - ds4_gpu_mv_dispatch mv_dispatch = - ds4_gpu_make_plain_mv_dispatch(in_dim, 0); - if (ds4_gpu_use_compressor_pair_nr4() && - (width == 512u || width == 1024u) && in_dim >= 4096u) { - mv_dispatch.nr0 = 4; - mv_dispatch.smem = 32u * 4u * sizeof(float); - } - mv_args.nr0 = mv_dispatch.nr0; - ds4_gpu_dsv4_compressor_store_one_args store_args = { - .width = width, - .ratio = ratio, - .pos = pos, - .ape_type = ape_type, - }; - id pipeline = ds4_gpu_get_mul_mv_pipeline( - "kernel_mul_mv_f16_f32_pair_compressor_store_4", - mv_dispatch.nsg); - if (!pipeline) return -1; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return -1; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBytes:&store_args length:sizeof(store_args) atIndex:1]; - [enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2]; - [enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; - [enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5]; - [enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6]; - [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7]; - [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8]; - [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9]; - if (mv_dispatch.smem) { - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake( - ((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake( - 32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer( - cb, owned, "F16 paired matvec compressor state store")) { - return -1; - } - } - - return 1; -} - -int ds4_gpu_matmul_f32_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_tok) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok == 0 || n_tok > UINT32_MAX) return 0; - if (in_dim > UINT64_MAX / n_tok / sizeof(float) || - out_dim > UINT64_MAX / n_tok / sizeof(float)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t x_bytes = n_tok * in_dim * sizeof(float); - const uint64_t out_bytes = n_tok * out_dim * sizeof(float); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal F32 tensor matmul received undersized activation buffers\n"); - return 0; - } - - const uint64_t row_bytes = in_dim * sizeof(float); - const uint64_t weight_bytes = row_bytes * out_dim; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal F32 tensor matmul range is outside the mapped model\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = - ds4_gpu_wrap_f32_decode_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - n_tok, - &inner_offset); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (n_tok == 1) { - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, 1); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - if (mv_dispatch.smem) { - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matvec")) return 0; - return 1; - } - - if (n_tok <= 8 && (in_dim % 128u) == 0) { - const int16_t nsg = 2; - const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); - const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); - const char *fn_name = ds4_gpu_mv_ext_f32_name(r1ptg); - id pipeline = - fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; - if (!pipeline) return 0; - - const int16_t nypsg = 32 / nxpsg; - const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; - ds4_gpu_mul_mv_ext_args args = - ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(float), row_bytes); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, - ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor mul_mv_ext")) return 0; - return 1; - } - - /* Generic multi-row path (GLM prefill shapes: one grid row per - * token through the plain matvec pipeline). */ - { - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, n_tok); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - if (mv_dispatch.smem) { - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - (NSUInteger)n_tok, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matmul")) return 0; - } - } - - return 1; -} - -int ds4_gpu_repeat_hc_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *row, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !row || n_embd == 0 || n_hc == 0) return 0; - - @autoreleasepool { - id rowbuf = ds4_gpu_tensor_buffer(row); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t out_bytes = row_bytes * n_hc; - if (!rowbuf || !outbuf || - ds4_gpu_tensor_bytes(row) < row_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal HC repeat received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (!ds4_gpu_encode_repeat_hc_embedding(cb, - rowbuf, - ds4_gpu_tensor_offset(row), - outbuf, - ds4_gpu_tensor_offset(out), - 1, - n_embd, - n_hc)) { - return 0; - } - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC repeat")) return 0; - } - - return 1; -} - -int ds4_gpu_repeat_hc_rows_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *rows, - uint32_t n_tokens, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; - - @autoreleasepool { - id rowsbuf = ds4_gpu_tensor_buffer(rows); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t rows_bytes = (uint64_t)n_tokens * row_bytes; - const uint64_t out_bytes = rows_bytes * n_hc; - if (!rowsbuf || !outbuf || - ds4_gpu_tensor_bytes(rows) < rows_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal HC row repeat received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (!ds4_gpu_encode_repeat_hc_embedding(cb, - rowsbuf, - ds4_gpu_tensor_offset(rows), - outbuf, - ds4_gpu_tensor_offset(out), - n_tokens, - n_embd, - n_hc)) { - return 0; - } - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC row repeat")) return 0; - } - - return 1; -} - -int ds4_gpu_rms_norm_plain_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *x, - uint32_t n, - float eps) { - return ds4_gpu_rms_norm_plain_rows_tensor(out, x, n, 1, eps); -} - -int ds4_gpu_rms_norm_plain_rows_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *x, - uint32_t n, - uint32_t rows, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t bytes = (uint64_t)n * rows * sizeof(float); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal plain RMS norm received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_rms_norm_plain_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "plain RMS norm")) return 0; - } - - return 1; -} - -static int ds4_gpu_hc_rms_scale_project_mode( - uint32_t in_dim, - uint32_t out_dim, - uint32_t n_rows) { - const bool hard_shape = n_rows > 8u && - (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u; - if (!hard_shape) return 0; - - const bool force = - getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL; - if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL || - (!ds4_gpu_device_name_contains("M3") && !force)) { - return 0; - } - if (g_rms_norm_scale_pipeline == nil) { - return force ? -1 : 0; - } - return 1; -} - -int ds4_gpu_hc_rms_scale_project_f16_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *scale_scratch, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t in_dim, - uint32_t out_dim, - const ds4_gpu_tensor *x, - uint32_t n_rows, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !scale_scratch || !model_map || !x || - in_dim == 0u || out_dim == 0u || n_rows == 0u || - n_rows > INT32_MAX) { - return 0; - } - - const int mode = - ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows); - if (mode < 0) return 0; - if (mode == 0) { - return ds4_gpu_rms_norm_plain_rows_tensor( - scale_scratch, x, in_dim, n_rows, eps) != 0 && - ds4_gpu_matmul_f16_tensor( - out, model_map, model_size, weight_offset, - in_dim, out_dim, scale_scratch, n_rows) != 0; - } - - @autoreleasepool { - const uint64_t x_elems = (uint64_t)in_dim * n_rows; - const uint64_t out_elems = (uint64_t)out_dim * n_rows; - if (x_elems > UINT64_MAX / sizeof(float) || - out_elems > UINT64_MAX / sizeof(float)) { - return 0; - } - const uint64_t x_bytes = x_elems * sizeof(float); - const uint64_t out_bytes = out_elems * sizeof(float); - const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float); - - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - id scalebuf = ds4_gpu_tensor_buffer(scale_scratch); - if (!xbuf || !outbuf || !scalebuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) { - fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n"); - return 0; - } - - const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t); - if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0; - const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes; - if (weight_offset > model_size || - weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n"); - return 0; - } - - const bool bc_inp = (in_dim % 32u) != 0u; - const bool bc_out = (out_dim % 64u) != 0u || - (n_rows % 32u) != 0u; - id mm_pipeline = - ds4_gpu_get_mul_mm_pipeline( - "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out); - if (!mm_pipeline) { - if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) { - return 0; - } - return ds4_gpu_rms_norm_plain_rows_tensor( - scale_scratch, x, in_dim, n_rows, eps) != 0 && - ds4_gpu_matmul_f16_tensor( - out, model_map, model_size, weight_offset, - in_dim, out_dim, scale_scratch, n_rows) != 0; - } - - id scale_pipeline = ds4_gpu_hot_pipeline( - g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4"); - if (!scale_pipeline) return 0; - - uint64_t weight_inner = 0; - id weightbuf = ds4_gpu_wrap_model_range( - model_map, model_size, weight_offset, weight_bytes, &weight_inner); - if (!weightbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_rms_norm_args norm_args = - ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps); - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:scale_pipeline]; - [enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:scalebuf - offset:ds4_gpu_tensor_offset(scale_scratch) - atIndex:2]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake( - ds4_gpu_rms_norm_threads(in_dim), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args( - in_dim, out_dim, n_rows, weight_row_bytes); - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:mm_pipeline]; - [enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setBuffer:scalebuf - offset:ds4_gpu_tensor_offset(scale_scratch) - atIndex:4]; - [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake( - ((NSUInteger)n_rows + 31u) / 32u, - ((NSUInteger)out_dim + 63u) / 64u, - 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer( - cb, owned, "HC RMS scale F16 projection")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_rms_norm_weight_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *x, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n, - float eps) { - return ds4_gpu_rms_norm_weight_rows_tensor(out, x, model_map, model_size, weight_offset, n, 1, eps); -} - -int ds4_gpu_rms_norm_weight_rows_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *x, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n, - uint32_t rows, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - const uint64_t bytes = row_bytes * rows; - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal weighted RMS norm received undersized activation buffers\n"); - return 0; - } - if (weight_offset > model_size || row_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal weighted RMS norm range is outside the mapped model\n"); - return 0; - } - - const bool exact_decode_weight_view = - rows == 1u && - row_bytes <= (1ull << 20) && - getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; - uint64_t inner_offset = 0; - id wbuf = exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_rms_norm_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "weighted RMS norm")) return 0; - } - - return 1; -} - -int ds4_gpu_add_rms_norm_weight_tensor( - ds4_gpu_tensor *norm_out, - ds4_gpu_tensor *sum_out, - const ds4_gpu_tensor *a, - const ds4_gpu_tensor *b, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!norm_out || !sum_out || !a || !b || n == 0 || (n & 3u) != 0) return 0; - - @autoreleasepool { - id abuf = ds4_gpu_tensor_buffer(a); - id bbuf = ds4_gpu_tensor_buffer(b); - id sumbuf = ds4_gpu_tensor_buffer(sum_out); - id normbuf = ds4_gpu_tensor_buffer(norm_out); - const uint64_t row_bytes = (uint64_t)n * sizeof(float); - if (!abuf || !bbuf || !sumbuf || !normbuf || - ds4_gpu_tensor_bytes(a) < row_bytes || - ds4_gpu_tensor_bytes(b) < row_bytes || - ds4_gpu_tensor_bytes(sum_out) < row_bytes || - ds4_gpu_tensor_bytes(norm_out) < row_bytes) { - fprintf(stderr, "ds4: Metal add+RMS norm received undersized activation buffers\n"); - return 0; - } - if (weight_offset > model_size || row_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal add+RMS norm range is outside the mapped model\n"); - return 0; - } - - const bool exact_decode_weight_view = - row_bytes <= (1ull << 20) && - getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; - uint64_t inner_offset = 0; - id wbuf = exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - row_bytes, - &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps); - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_add_rms_norm_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; - [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:3]; - [enc setBuffer:sumbuf offset:ds4_gpu_tensor_offset(sum_out) atIndex:4]; - [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "add+RMS norm")) return 0; - } - - return 1; -} - -int ds4_gpu_dsv4_qkv_rms_norm_rows_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, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!q_out || !q || !kv_out || !kv || q_n == 0 || kv_n == 0 || rows == 0 || - (q_n & 3u) != 0 || (kv_n & 3u) != 0) { - return 0; - } - - @autoreleasepool { - id qbuf = ds4_gpu_tensor_buffer(q); - id qoutbuf = ds4_gpu_tensor_buffer(q_out); - id kvbuf = ds4_gpu_tensor_buffer(kv); - id kvoutbuf = ds4_gpu_tensor_buffer(kv_out); - - const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); - const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float); - if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf || - ds4_gpu_tensor_bytes(q) < q_row_bytes * rows || - ds4_gpu_tensor_bytes(q_out) < q_row_bytes * rows || - ds4_gpu_tensor_bytes(kv) < kv_row_bytes * rows || - ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes * rows) { - fprintf(stderr, "ds4: Metal fused q/kv RMS norm received undersized activation buffers\n"); - return 0; - } - if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || - kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) { - fprintf(stderr, "ds4: Metal fused q/kv RMS norm weight range is outside the mapped model\n"); - return 0; - } - - uint64_t q_inner_offset = 0; - uint64_t kv_inner_offset = 0; - id q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - q_weight_offset, q_row_bytes, - &q_inner_offset); - if (!q_wbuf) return 0; - id kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - kv_weight_offset, kv_row_bytes, - &kv_inner_offset); - if (!kv_wbuf) return 0; - - ds4_gpu_qkv_rms_norm_args args = { - .q_n = (int32_t)q_n, - .q_n4 = (int32_t)(q_n / 4u), - .kv_n = (int32_t)kv_n, - .kv_n4 = (int32_t)(kv_n / 4u), - .q_row_stride = q_row_bytes, - .kv_row_stride = kv_row_bytes, - .eps = eps, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_qkv_rms_norm_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:2]; - [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; - [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:4]; - [enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:5]; - [enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:6]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 2, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv RMS norm")) return 0; - } - - return 1; -} - -int ds4_gpu_head_rms_norm_tensor( - ds4_gpu_tensor *x, - uint32_t n_tok, - uint32_t n_head, - uint32_t head_dim, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || (head_dim & 3u) != 0) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); - if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { - fprintf(stderr, "ds4: Metal head RMS norm received undersized activation buffer\n"); - return 0; - } - - ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_3d_args(head_dim, n_head, n_tok, eps); - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_rms_norm_plain_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_pipeline_threads(head_dim, g_rms_norm_plain_pipeline), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "head RMS norm")) return 0; - } - - return 1; -} - -int ds4_gpu_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tok, - uint32_t n_head, - uint32_t 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) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || n_rot > head_dim || (n_rot & 1u) != 0) { - return 0; - } - if (n_rot == 0) return 1; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); - if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { - fprintf(stderr, "ds4: Metal RoPE received undersized activation buffer\n"); - return 0; - } - - ds4_gpu_rope_tail_batch_args args = ds4_gpu_make_rope_tail_args( - n_tok, n_head, head_dim, n_rot, n_ctx_orig, inverse, - freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_rope_tail_inplace(cb, - xbuf, - ds4_gpu_tensor_offset(x), - &args, - n_tok, - n_head, - head_dim, - pos0, - 1)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "RoPE tail")) return 0; - } - - return 1; -} - -int ds4_gpu_head_rms_norm_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tok, - uint32_t n_head, - uint32_t 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) { - (void)x; (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; - (void)pos0; (void)n_ctx_orig; (void)inverse; (void)freq_base; - (void)freq_scale; (void)ext_factor; (void)attn_factor; - (void)beta_fast; (void)beta_slow; (void)eps; - return 0; -} - -int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *q_half, - 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_tok, - uint32_t n_head, - uint32_t 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) { - (void)out; (void)q_half; (void)model_map; (void)model_size; - (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; - (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0; - (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale; - (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; - (void)eps; - return 0; -} - -int ds4_gpu_dsv4_fp8_kv_quantize_tensor( - ds4_gpu_tensor *x, - uint32_t n_tok, - uint32_t head_dim, - uint32_t n_rot) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || n_tok == 0 || head_dim == 0 || n_rot > head_dim) return 0; - if (n_rot == head_dim) return 1; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - const uint64_t bytes = (uint64_t)n_tok * head_dim * sizeof(float); - if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { - fprintf(stderr, "ds4: Metal DSV4 FP8 KV quantize received undersized activation buffer\n"); - return 0; - } - - ds4_gpu_dsv4_fp8_kv_quantize_args args = { - .ne00 = head_dim, - .ne01 = n_tok, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = (uint64_t)head_dim * sizeof(float), - .nb02 = (uint64_t)n_tok * head_dim * sizeof(float), - .nb03 = (uint64_t)n_tok * head_dim * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)head_dim * sizeof(float), - .nb2 = (uint64_t)n_tok * head_dim * sizeof(float), - .nb3 = (uint64_t)n_tok * head_dim * sizeof(float), - .n_rot = (int32_t)n_rot, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_fp8_kv_quantize_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; - [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_tok, 1, 1) - threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 FP8 KV quantize")) return 0; - } - - return 1; -} - -int ds4_gpu_dsv4_indexer_qat_tensor( - ds4_gpu_tensor *x, - uint32_t n_rows, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || n_rows == 0 || head_dim != 128u) return 0; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - const uint64_t bytes = (uint64_t)n_rows * head_dim * sizeof(float); - if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { - fprintf(stderr, "ds4: Metal DSV4 indexer QAT received undersized activation buffer\n"); - return 0; - } - - ds4_gpu_dsv4_indexer_qat_args args = { - .n_rows = n_rows, - .head_dim = head_dim, - .row_stride = (uint64_t)head_dim * sizeof(float), - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_indexer_qat_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 indexer Hadamard+FP4")) return 0; - } - - return 1; -} - -static void ds4_gpu_set_rows_thread_shape( - uint32_t width, - NSUInteger *nth_out, - NSUInteger *nrptg_out) { - const NSUInteger nk0 = width ? (NSUInteger)width : 1u; - const NSUInteger max_threads = g_set_rows_f32_i32_pipeline - ? (NSUInteger)g_set_rows_f32_i32_pipeline.maxTotalThreadsPerThreadgroup - : 1024u; - - NSUInteger nth = 32u; - while (nth < nk0 && nth < max_threads) { - nth *= 2u; - } - - NSUInteger nrptg = 1u; - if (nth > nk0) { - nrptg = (nth + nk0 - 1u) / nk0; - nth = nk0; - if (nrptg * nth > max_threads) { - nrptg--; - } - } - - if (nth > nk0) nth = nk0; - if (nth == 0u) nth = 1u; - if (nrptg == 0u) nrptg = 1u; - - *nth_out = nth; - *nrptg_out = nrptg; -} - -static int ds4_gpu_encode_f16_round_copy_for_raw_store( - id cb, - const ds4_gpu_tensor *src, - uint32_t n) { - id srcbuf = ds4_gpu_tensor_buffer(src); - const uint64_t src_bytes = (uint64_t)n * sizeof(float); - if (!srcbuf || ds4_gpu_tensor_bytes(src) < src_bytes) { - fprintf(stderr, "ds4: Metal raw KV store received undersized source buffer\n"); - return 0; - } - if (!ds4_gpu_ensure_scratch_buffer(&g_f16_round_scratch_buffer, - &g_f16_round_scratch_bytes, - (NSUInteger)n * sizeof(uint16_t), - "ds4_f16_round_scratch") || - !ds4_gpu_ensure_scratch_buffer(&g_raw_store_round_buffer, - &g_raw_store_round_bytes, - (NSUInteger)n * sizeof(float), - "ds4_raw_store_round")) { - return 0; - } - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - srcbuf, - ds4_gpu_tensor_offset(src), - g_f16_round_scratch_buffer, - 0, - n)) { - return 0; - } - return ds4_gpu_encode_cpy_f16_f32_1d(cb, - g_f16_round_scratch_buffer, - 0, - g_raw_store_round_buffer, - 0, - n); -} - -static int ds4_gpu_encode_set_rows_f32_i32( - id cb, - ds4_gpu_tensor *dst, - id srcbuf, - NSUInteger src_off, - const int32_t *rows, - uint32_t n_rows, - uint32_t dst_rows, - uint32_t width) { - id dstbuf = ds4_gpu_tensor_buffer(dst); - const uint64_t dst_bytes = (uint64_t)dst_rows * width * sizeof(float); - const uint64_t src_bytes = (uint64_t)n_rows * width * sizeof(float); - if (!dstbuf || !srcbuf || !rows || n_rows == 0 || width == 0 || - ds4_gpu_tensor_bytes(dst) < dst_bytes || - src_bytes > NSUIntegerMax - src_off) { - fprintf(stderr, "ds4: Metal DS4 set_rows received invalid buffers\n"); - return 0; - } - - const uint64_t row_bytes = (uint64_t)width * sizeof(float); - const uint64_t rows_bytes = (uint64_t)n_rows * sizeof(int32_t); - ds4_gpu_set_rows_args args = { - .nk0 = (int32_t)width, - .ne01 = (int32_t)n_rows, - .nb01 = row_bytes, - .nb02 = (uint64_t)n_rows * row_bytes, - .nb03 = (uint64_t)n_rows * row_bytes, - .ne11 = 1, - .ne12 = 1, - .nb10 = sizeof(int32_t), - .nb11 = rows_bytes, - .nb12 = rows_bytes, - .nb1 = row_bytes, - .nb2 = (uint64_t)dst_rows * row_bytes, - .nb3 = (uint64_t)dst_rows * row_bytes, - }; - - NSUInteger nth; - NSUInteger nrptg; - ds4_gpu_set_rows_thread_shape(width, &nth, &nrptg); - - id rowsbuf = nil; - if (rows_bytes > 4096u) { - rowsbuf = ds4_gpu_new_transient_buffer((NSUInteger)rows_bytes, "ds4_set_rows_indices"); - if (!rowsbuf) return 0; - memcpy([rowsbuf contents], rows, (NSUInteger)rows_bytes); - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_set_rows_f32_i32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:srcbuf offset:src_off atIndex:1]; - if (rowsbuf) { - [enc setBuffer:rowsbuf offset:0 atIndex:2]; - } else { - [enc setBytes:rows length:(NSUInteger)rows_bytes atIndex:2]; - } - [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + nrptg - 1u) / nrptg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, nrptg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_add_f32_1d( - id cb, - id a, - NSUInteger a_off, - id b, - NSUInteger b_off, - id out, - NSUInteger out_off, - uint32_t n) { - if (!cb || !a || !b || !out || n == 0) return 0; - - ds4_gpu_add_flat_args args = { .n = n }; - NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - if (nth == 0u) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_add2_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:a offset:a_off atIndex:1]; - [enc setBuffer:b offset:b_off atIndex:2]; - [enc setBuffer:out offset:out_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -int ds4_gpu_store_raw_kv_tensor( - ds4_gpu_tensor *raw_cache, - const ds4_gpu_tensor *kv, - uint32_t raw_cap, - uint32_t row, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!raw_cache || !kv || raw_cap == 0 || row >= raw_cap || head_dim == 0 || raw_cap > INT32_MAX) return 0; - - @autoreleasepool { - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { - fprintf(stderr, "ds4: Metal raw KV store received undersized destination buffer\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const int32_t row_i32 = (int32_t)row; - if (!ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, head_dim) || - !ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, - g_raw_store_round_buffer, - 0, - &row_i32, - 1, - raw_cap, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV DS4 set_rows store")) return 0; - } - - return 1; -} - -/* Release decode fused KV finalizer. Reference paths are selected by the C - * graph driver; this Objective-C entry point always means "use the fused - * Metal kernel." */ -int ds4_gpu_kv_fp8_store_raw_tensor( - ds4_gpu_tensor *kv, - ds4_gpu_tensor *raw_cache, - uint32_t raw_cap, - uint32_t row, - uint32_t head_dim, - uint32_t n_rot) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 || - n_rot > head_dim || raw_cap > INT32_MAX) { - return 0; - } - - @autoreleasepool { - id kvbuf = ds4_gpu_tensor_buffer(kv); - id rawbuf = ds4_gpu_tensor_buffer(raw_cache); - const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - if (!kvbuf || !rawbuf || - ds4_gpu_tensor_bytes(kv) < kv_bytes || - ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { - fprintf(stderr, "ds4: Metal fused KV FP8/raw-store received undersized buffers\n"); - return 0; - } - - ds4_gpu_dsv4_kv_fp8_store_args args = { - .head_dim = (int32_t)head_dim, - .n_rot = (int32_t)n_rot, - .raw_row = (int32_t)row, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_kv_fp8_store_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:2]; - [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "KV FP8/raw-store fused")) return 0; - } - - return 1; -} - -int ds4_gpu_store_raw_kv_batch_tensor( - ds4_gpu_tensor *raw_cache, - const ds4_gpu_tensor *kv, - uint32_t raw_cap, - uint32_t pos0, - uint32_t n_tokens, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!raw_cache || !kv || raw_cap == 0 || n_tokens == 0 || head_dim == 0 || raw_cap > INT32_MAX) return 0; - - @autoreleasepool { - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { - fprintf(stderr, "ds4: Metal raw KV batch store received undersized destination buffer\n"); - return 0; - } - - int32_t rows_stack[512]; - int32_t *rows = rows_stack; - if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { - rows = malloc((size_t)n_tokens * sizeof(*rows)); - if (!rows) { - fprintf(stderr, "ds4: failed to allocate raw KV set_rows index list\n"); - return 0; - } - } - for (uint32_t t = 0; t < n_tokens; t++) { - rows[t] = (int32_t)((pos0 + t) % raw_cap); - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) { - if (rows != rows_stack) free(rows); - return 0; - } - - const uint64_t n = (uint64_t)n_tokens * head_dim; - const int ok = n <= UINT32_MAX && - ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, (uint32_t)n) && - ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, - g_raw_store_round_buffer, - 0, - rows, - n_tokens, - raw_cap, - head_dim); - if (rows != rows_stack) free(rows); - if (!ok) return 0; - - if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV batch DS4 set_rows store")) return 0; - } - - return 1; -} - -static int ds4_gpu_encode_compressor_score_with_ape( - id cb, - id score_src, - NSUInteger score_src_offset, - id score_dst, - NSUInteger score_dst_offset, - id apebuf, - NSUInteger ape_offset, - uint32_t ape_type, - uint32_t width, - uint32_t ratio, - uint32_t pos0, - uint32_t n_tokens) { - if (!cb || !score_src || !score_dst || !apebuf || - width == 0 || ratio == 0 || n_tokens == 0 || - (ape_type != 0u && ape_type != 1u)) { - return 0; - } - - const uint64_t total_elems64 = (uint64_t)n_tokens * width; - if (total_elems64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal compressor APE add received too many elements\n"); - return 0; - } - const uint32_t total_elems = (uint32_t)total_elems64; - - const bool force_fused = - getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL; - const bool use_fused = - (ds4_gpu_device_name_contains("M3") || force_fused) && - getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL; - if (use_fused) { - id pipeline = ds4_gpu_get_pipeline( - ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16" - : "kernel_dsv4_compressor_score_ape_f32"); - if (pipeline) { - ds4_gpu_dsv4_compressor_score_ape_args args = { - .width = width, - .ratio = ratio, - .pos0 = pos0, - .n_tokens = n_tokens, - }; - NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth == 0) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:score_src offset:score_src_offset atIndex:1]; - [enc setBuffer:apebuf offset:ape_offset atIndex:2]; - [enc setBuffer:score_dst offset:score_dst_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake( - ((NSUInteger)total_elems + nth - 1u) / nth, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; - } - if (force_fused) return 0; - } - - const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, - &g_compressor_store_ape_bytes, - scratch_bytes, - "ds4_compressor_store_ape")) { - return 0; - } - - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - uint32_t copied_rows = 0; - uint32_t pos_mod = pos0 % ratio; - while (copied_rows < n_tokens) { - uint32_t seg_rows = ratio - pos_mod; - if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; - const uint32_t seg_elems = seg_rows * width; - const NSUInteger src_off = ape_offset + (NSUInteger)pos_mod * width * elem_ape; - const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); - int ok; - if (ape_type == 1u) { - ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, - apebuf, - src_off, - g_compressor_store_ape_buffer, - dst_off, - seg_elems); - } else { - ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, - apebuf, - src_off, - g_compressor_store_ape_buffer, - dst_off, - seg_elems); - } - if (!ok) return 0; - copied_rows += seg_rows; - pos_mod = 0; - } - - return ds4_gpu_encode_add_f32_1d(cb, - score_src, - score_src_offset, - g_compressor_store_ape_buffer, - 0, - score_dst, - score_dst_offset, - total_elems); -} - -static int ds4_gpu_encode_compressor_set_rows_projected( - id cb, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - id kvbuf, - NSUInteger kv_offset, - id scorebuf, - NSUInteger score_offset, - id apebuf, - NSUInteger ape_offset, - uint32_t ape_type, - uint32_t width, - uint32_t ratio, - uint32_t pos0, - const int32_t *rows, - uint32_t n_rows, - uint32_t state_rows) { - if (!cb || !state_kv || !state_score || !kvbuf || !scorebuf || - !apebuf || !rows || width == 0 || n_rows == 0 || state_rows == 0) { - return 0; - } - - const NSUInteger score_scratch_bytes = (NSUInteger)n_rows * width * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, - &g_compressor_store_score_bytes, - score_scratch_bytes, - "ds4_compressor_store_score")) { - return 0; - } - - return ds4_gpu_encode_compressor_score_with_ape(cb, - scorebuf, - score_offset, - g_compressor_store_score_buffer, - 0, - apebuf, - ape_offset, - ape_type, - width, - ratio, - pos0, - n_rows) && - ds4_gpu_encode_set_rows_f32_i32(cb, - state_kv, - kvbuf, - kv_offset, - rows, - n_rows, - state_rows, - width) && - ds4_gpu_encode_set_rows_f32_i32(cb, - state_score, - g_compressor_store_score_buffer, - 0, - rows, - n_rows, - state_rows, - width); -} - -static int ds4_gpu_compressor_store_one_tensor( - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t width, - uint32_t ratio, - uint32_t pos) { - if (!kv || !sc || !state_kv || !state_score || !model_map || - width == 0 || ratio == 0 || (ape_type != 0u && ape_type != 1u)) { - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_dsv4_compressor_store_one_pipeline, - "kernel_dsv4_compressor_store_one"); - if (!pipeline) return 0; - - const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t row_bytes = (uint64_t)width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * row_bytes; - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - ds4_gpu_tensor_bytes(kv) < row_bytes || - ds4_gpu_tensor_bytes(sc) < row_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes) { - return 0; - } - - uint64_t ape_inner = 0; - id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, - ape_offset, ape_bytes, - &ape_inner); - id kvbuf = ds4_gpu_tensor_buffer(kv); - id scbuf = ds4_gpu_tensor_buffer(sc); - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!apebuf || !kvbuf || !scbuf || !statekvbuf || !statescbuf) return 0; - - ds4_gpu_dsv4_compressor_store_one_args args = { - .width = width, - .ratio = ratio, - .pos = pos, - .ape_type = ape_type, - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const NSUInteger nth = 256u; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; - [enc setBuffer:scbuf offset:ds4_gpu_tensor_offset(sc) atIndex:2]; - [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:3]; - [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:4]; - [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:5]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)width + nth - 1u) / nth, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return ds4_gpu_finish_command_buffer(cb, owned, "compressor one-row store"); -} - -int ds4_gpu_compressor_store_batch_tensor( - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos0, - uint32_t n_tokens) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!kv || !sc || !state_kv || !state_score || !model_map || - head_dim == 0 || ratio == 0 || n_tokens == 0 || - (ape_type != 0u && ape_type != 1u)) { - return 0; - } - - @autoreleasepool { - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { - fprintf(stderr, "ds4: Metal compressor batch APE range is outside the mapped model\n"); - return 0; - } - - id kvbuf = ds4_gpu_tensor_buffer(kv); - id scbuf = ds4_gpu_tensor_buffer(sc); - if (!kvbuf || !scbuf || - ds4_gpu_tensor_bytes(kv) < kv_bytes || - ds4_gpu_tensor_bytes(sc) < kv_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes) { - fprintf(stderr, "ds4: Metal compressor batch store received undersized buffers\n"); - return 0; - } - - uint64_t ape_inner = 0; - id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); - if (!apebuf) return 0; - - const uint64_t total_elems64 = (uint64_t)n_tokens * width; - if (total_elems64 > UINT32_MAX || state_rows > INT32_MAX) { - fprintf(stderr, "ds4: Metal compressor batch store received too many elements\n"); - return 0; - } - const uint32_t total_elems = (uint32_t)total_elems64; - const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, - &g_compressor_store_ape_bytes, - scratch_bytes, - "ds4_compressor_store_ape") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, - &g_compressor_store_score_bytes, - scratch_bytes, - "ds4_compressor_store_score")) { - return 0; - } - - int32_t rows_stack[16]; - int32_t *rows = rows_stack; - if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { - rows = malloc((size_t)n_tokens * sizeof(*rows)); - if (!rows) { - fprintf(stderr, "ds4: failed to allocate compressor set_rows index list\n"); - return 0; - } - } - for (uint32_t t = 0; t < n_tokens; t++) { - const uint32_t pos_mod = (pos0 + t) % ratio; - rows[t] = (int32_t)(ratio == 4u ? ratio + pos_mod : pos_mod); - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) { - if (rows != rows_stack) free(rows); - return 0; - } - - int ok = 1; - uint32_t copied_rows = 0; - uint32_t pos_mod = pos0 % ratio; - while (ok && copied_rows < n_tokens) { - uint32_t seg_rows = ratio - pos_mod; - if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; - const uint32_t seg_elems = seg_rows * width; - const NSUInteger src_off = (NSUInteger)ape_inner + - (NSUInteger)pos_mod * width * elem_ape; - const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); - if (ape_type == 1u) { - ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, - apebuf, - src_off, - g_compressor_store_ape_buffer, - dst_off, - seg_elems); - } else { - ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, - apebuf, - src_off, - g_compressor_store_ape_buffer, - dst_off, - seg_elems); - } - copied_rows += seg_rows; - pos_mod = 0; - } - - if (ok) { - ok = ds4_gpu_encode_add_f32_1d(cb, - scbuf, - ds4_gpu_tensor_offset(sc), - g_compressor_store_ape_buffer, - 0, - g_compressor_store_score_buffer, - 0, - total_elems); - } - if (ok) { - ok = ds4_gpu_encode_set_rows_f32_i32(cb, - state_kv, - kvbuf, - ds4_gpu_tensor_offset(kv), - rows, - n_tokens, - state_rows, - width); - } - if (ok) { - ok = ds4_gpu_encode_set_rows_f32_i32(cb, - state_score, - g_compressor_store_score_buffer, - 0, - rows, - n_tokens, - state_rows, - width); - } - if (rows != rows_stack) free(rows); - if (!ok) return 0; - - if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor batch DS4 store")) return 0; - } - - return 1; -} - -static ds4_gpu_bin_args ds4_gpu_make_bin_contiguous_3d_args( - uint32_t cols, - uint32_t rows, - uint32_t planes) { - const uint64_t row_bytes = (uint64_t)cols * sizeof(float); - const uint64_t plane_bytes = (uint64_t)rows * row_bytes; - return (ds4_gpu_bin_args) { - .ne00 = (int32_t)cols, - .ne01 = (int32_t)rows, - .ne02 = (int32_t)planes, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = row_bytes, - .nb02 = plane_bytes, - .nb03 = (uint64_t)planes * plane_bytes, - .ne10 = (int32_t)cols, - .ne11 = (int32_t)rows, - .ne12 = (int32_t)planes, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = row_bytes, - .nb12 = plane_bytes, - .nb13 = (uint64_t)planes * plane_bytes, - .ne0 = (int32_t)cols, - .ne1 = (int32_t)rows, - .ne2 = (int32_t)planes, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = row_bytes, - .nb2 = plane_bytes, - .nb3 = (uint64_t)planes * plane_bytes, - .offs = 0, - .o1 = { 0 }, - }; -} - -static int ds4_gpu_encode_softmax_f32_contiguous( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t width, - uint32_t rows, - uint32_t planes) { - if (!cb || !src || !dst || width == 0 || rows == 0 || planes == 0) return 0; - - const uint64_t row_bytes = (uint64_t)width * sizeof(float); - const uint64_t plane_bytes = (uint64_t)rows * row_bytes; - ds4_gpu_softmax_args args = { - .ne00 = (int32_t)width, - .ne01 = (int32_t)rows, - .ne02 = (int32_t)planes, - .nb01 = row_bytes, - .nb02 = plane_bytes, - .nb03 = (uint64_t)planes * plane_bytes, - .ne11 = (int32_t)width, - .ne12 = (int32_t)rows, - .ne13 = (int32_t)planes, - .nb11 = row_bytes, - .nb12 = plane_bytes, - .nb13 = (uint64_t)planes * plane_bytes, - .nb1 = row_bytes, - .nb2 = plane_bytes, - .nb3 = (uint64_t)planes * plane_bytes, - .scale = 1.0f, - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 1, - }; - - id pipeline = - (width % 4u) == 0 ? g_soft_max_f32_4_pipeline : g_soft_max_f32_pipeline; - if (!pipeline) return 0; - - NSUInteger nth = 32u; - if ((width % 4u) == 0) { - while (nth < (NSUInteger)(width / 4u) && - nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { - nth *= 2u; - } - } else { - while (nth < (NSUInteger)width && - nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { - nth *= 2u; - } - } - const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:src offset:src_off atIndex:2]; - [enc setBuffer:src offset:src_off atIndex:3]; - [enc setBuffer:dst offset:dst_off atIndex:4]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(rows, planes, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( - id cb, - ds4_gpu_tensor *out, - id kvbuf, - NSUInteger kv_offset, - uint64_t kv_nb0, - uint64_t kv_nb1, - uint64_t kv_nb2, - id scorebuf, - NSUInteger score_offset, - uint64_t score_nb0, - uint64_t score_nb1, - uint64_t score_nb2, - uint32_t n_rows, - uint32_t head_dim) { - id outbuf = ds4_gpu_tensor_buffer(out); - if (!cb || !outbuf || !kvbuf || !scorebuf || n_rows == 0 || head_dim == 0 || - ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float)) { - return 0; - } - - const NSUInteger pack_bytes = (NSUInteger)n_rows * head_dim * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer, - &g_compressor_pool_product_bytes, - pack_bytes, - "ds4_compressor_pool_product") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer, - &g_compressor_pool_score_cont_bytes, - pack_bytes, - "ds4_compressor_pool_score_cont") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer, - &g_compressor_pool_softmax_bytes, - pack_bytes, - "ds4_compressor_pool_softmax")) { - return 0; - } - - const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float); - const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride; - - /* - * Keep the n_comp == 1 compressor path as the unfused graph sequence: - * - * score = soft_max(contiguous(score)) - * pooled = sum_rows(contiguous(kv) * score) - * - * The fused DS4 pool kernel is mathematically equivalent, but it reduces in - * a different order. That is enough to create ~1e-6 compressor differences - * and later FP8/routing flips, so this path intentionally keeps the same - * operation boundary and memory layout as the graph. - */ - ds4_gpu_bin_args mul_args = - ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1); - - return - ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, - kvbuf, - kv_offset, - g_compressor_pool_product_buffer, - 0, - n_rows, - head_dim, - 1, - kv_nb0, - kv_nb1, - kv_nb2, - cont_row_stride, - cont_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, - scorebuf, - score_offset, - g_compressor_pool_score_cont_buffer, - 0, - n_rows, - head_dim, - 1, - score_nb0, - score_nb1, - score_nb2, - cont_row_stride, - cont_plane_stride) && - ds4_gpu_encode_softmax_f32_contiguous(cb, - g_compressor_pool_score_cont_buffer, - 0, - g_compressor_pool_softmax_buffer, - 0, - n_rows, - head_dim, - 1) && - ds4_gpu_encode_bin_f32_rows(cb, - g_mul_pipeline, - &mul_args, - g_compressor_pool_product_buffer, - 0, - g_compressor_pool_softmax_buffer, - 0, - g_compressor_pool_product_buffer, - 0) && - ds4_gpu_encode_sum_rows_f32(cb, - g_compressor_pool_product_buffer, - 0, - outbuf, - ds4_gpu_tensor_offset(out), - n_rows, - head_dim); -} - -static int ds4_gpu_encode_dsv4_softmax_pool( - id cb, - ds4_gpu_tensor *out, - id kvbuf, - NSUInteger kv_offset, - uint64_t kv_nb0, - uint64_t kv_nb1, - uint64_t kv_nb2, - id scorebuf, - NSUInteger score_offset, - uint64_t score_nb0, - uint64_t score_nb1, - uint64_t score_nb2, - uint32_t n_rows, - uint32_t head_dim, - uint32_t n_comp) { - id outbuf = ds4_gpu_tensor_buffer(out); - if (!cb || !outbuf || !kvbuf || !scorebuf || - n_rows == 0 || head_dim == 0 || n_comp == 0 || - ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * n_comp * sizeof(float)) { - return 0; - } - - if (n_comp == 1) { - return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(cb, - out, - kvbuf, - kv_offset, - kv_nb0, - kv_nb1, - kv_nb2, - scorebuf, - score_offset, - score_nb0, - score_nb1, - score_nb2, - n_rows, - head_dim); - } - - ds4_gpu_dsv4_softmax_pool_args args = { - .ne00 = (int64_t)n_rows, - .ne01 = (int64_t)head_dim, - .ne02 = (int64_t)n_comp, - .nb00 = kv_nb0, - .nb01 = kv_nb1, - .nb02 = kv_nb2, - .nb10 = score_nb0, - .nb11 = score_nb1, - .nb12 = score_nb2, - .ne0 = (int64_t)head_dim, - .ne1 = (int64_t)n_comp, - .nb0 = sizeof(float), - .nb1 = (uint64_t)head_dim * sizeof(float), - }; - const uint64_t n = (uint64_t)head_dim * n_comp; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_softmax_pool_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; - [enc setBuffer:scorebuf offset:score_offset atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_concat_f32_dim1( - id cb, - id src0, - NSUInteger src0_offset, - uint32_t src0_rows, - uint64_t src0_row_stride, - id src1, - NSUInteger src1_offset, - uint32_t src1_rows, - uint64_t src1_row_stride, - id dst, - NSUInteger dst_offset, - uint32_t cols, - uint64_t dst_row_stride) { - if (!cb || !src0 || !src1 || !dst || cols == 0 || src0_rows == 0 || src1_rows == 0) { - return 0; - } - - const uint32_t rows = src0_rows + src1_rows; - const uint64_t src0_plane = (uint64_t)src0_rows * src0_row_stride; - const uint64_t src1_plane = (uint64_t)src1_rows * src1_row_stride; - const uint64_t dst_plane = (uint64_t)rows * dst_row_stride; - ds4_gpu_concat_args args = { - .ne00 = (int32_t)cols, - .ne01 = (int32_t)src0_rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src0_row_stride, - .nb02 = src0_plane, - .nb03 = src0_plane, - .ne10 = (int32_t)cols, - .ne11 = (int32_t)src1_rows, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = src1_row_stride, - .nb12 = src1_plane, - .nb13 = src1_plane, - .ne0 = (int32_t)cols, - .ne1 = (int32_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = dst_row_stride, - .nb2 = dst_plane, - .nb3 = dst_plane, - .dim = 1, - }; - - NSUInteger nth = cols < 1024u ? (NSUInteger)cols : 1024u; - const NSUInteger max_threads = g_concat_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_concat_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src0 offset:src0_offset atIndex:1]; - [enc setBuffer:src1 offset:src1_offset atIndex:2]; - [enc setBuffer:dst offset:dst_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) { - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL; - const bool default_shape = head_dim == 128u || head_dim == 512u; - if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || - (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { - return 0; - } - if (g_dsv4_compressor_pack_ratio4_pipeline == nil) { - return force ? -1 : 0; - } - return 1; -} - -static bool ds4_gpu_buffer_ranges_overlap( - id a, - NSUInteger a_offset, - uint64_t a_bytes, - id b, - NSUInteger b_offset, - uint64_t b_bytes) { - if (a != b || a_bytes == 0u || b_bytes == 0u) return false; - if (a_offset <= b_offset) { - return a_bytes > (uint64_t)(b_offset - a_offset); - } - return b_bytes > (uint64_t)(a_offset - b_offset); -} - -static int ds4_gpu_compressor_ratio4_direct_pool_mode( - uint32_t head_dim, - uint32_t n_comp) { - // One compressed row intentionally uses the legacy GGML reduction graph. - if (n_comp <= 1u) return 0; - - const bool force = - getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL; - const bool default_shape = head_dim == 128u || head_dim == 512u; - if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || - (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { - return 0; - } - if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) { - return force ? -1 : 0; - } - return 1; -} - -static int ds4_gpu_encode_compressor_ratio4_direct_pool( - id cb, - ds4_gpu_tensor *out, - id kvbuf, - NSUInteger kv_offset, - id scorebuf, - NSUInteger score_offset, - id statekvbuf, - NSUInteger state_kv_offset, - id statescbuf, - NSUInteger state_score_offset, - uint32_t head_dim, - uint32_t n_comp, - bool replay) { - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t n = (uint64_t)head_dim * n_comp; - if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || - head_dim == 0u || n_comp <= 1u || n > UINT32_MAX || - ds4_gpu_tensor_bytes(out) < n * sizeof(float)) { - return 0; - } - - id pipeline = ds4_gpu_hot_pipeline( - g_dsv4_softmax_pool_ratio4_direct_pipeline, - "kernel_dsv4_softmax_pool_ratio4_direct"); - if (!pipeline) return 0; - - ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = { - .n_rows = 8, - .head_dim = head_dim, - .n_comp = n_comp, - .replay = replay ? 1u : 0u, - .pad = 0u, - }; - - NSUInteger nth = 256u; - if (nth > pipeline.maxTotalThreadsPerThreadgroup) { - nth = pipeline.maxTotalThreadsPerThreadgroup; - } - if (nth == 0u) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; - [enc setBuffer:scorebuf offset:score_offset atIndex:2]; - [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; - [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_compressor_pack_ratio4( - id cb, - id kvbuf, - NSUInteger kv_offset, - id scorebuf, - NSUInteger score_offset, - id statekvbuf, - NSUInteger state_kv_offset, - id statescbuf, - NSUInteger state_score_offset, - uint32_t head_dim, - uint32_t n_comp, - bool replay) { - if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || - !g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer || - head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) { - return 0; - } - - const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim; - if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0; - id pipeline = ds4_gpu_hot_pipeline( - g_dsv4_compressor_pack_ratio4_pipeline, - "kernel_dsv4_compressor_pack_ratio4"); - if (!pipeline) return 0; - - NSUInteger nth = head_dim; - if (nth > 256u) nth = 256u; - if (nth > pipeline.maxTotalThreadsPerThreadgroup) { - nth = pipeline.maxTotalThreadsPerThreadgroup; - } - if (nth == 0) return 0; - - ds4_gpu_dsv4_compressor_pack_ratio4_args args = { - .head_dim = head_dim, - .n_comp = n_comp, - .replay = replay ? 1u : 0u, - .n_threads = (uint32_t)nth, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; - [enc setBuffer:scorebuf offset:score_offset atIndex:2]; - [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; - [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; - [enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5]; - [enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6]; - [enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_compressor_pool( - id cb, - ds4_gpu_tensor *out, - const ds4_gpu_tensor *state_kv, - const ds4_gpu_tensor *state_score, - uint32_t head_dim, - uint32_t ratio) { - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0; - - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t rows = coff * ratio; - const uint64_t state_bytes = (uint64_t)width * rows * sizeof(float); - if (ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes) { - return 0; - } - - if (ratio != 4u) { - const uint64_t row_stride = (uint64_t)width * sizeof(float); - return ds4_gpu_encode_dsv4_softmax_pool(cb, - out, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - row_stride, - sizeof(float), - (uint64_t)rows * row_stride, - statescbuf, - ds4_gpu_tensor_offset(state_score), - row_stride, - sizeof(float), - (uint64_t)rows * row_stride, - ratio, - head_dim, - 1); - } - - const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, - &g_compressor_pool_kv_bytes, - packed_bytes, - "ds4_compressor_pool_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, - &g_compressor_pool_score_bytes, - packed_bytes, - "ds4_compressor_pool_score")) { - return 0; - } - - const uint64_t state_row_stride = (uint64_t)width * sizeof(float); - const uint64_t pool_row_stride = (uint64_t)head_dim * sizeof(float); - const NSUInteger curr_offset = (NSUInteger)4u * state_row_stride + - (NSUInteger)head_dim * sizeof(float); - if (!ds4_gpu_encode_concat_f32_dim1(cb, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - 4, - state_row_stride, - statekvbuf, - ds4_gpu_tensor_offset(state_kv) + curr_offset, - 4, - state_row_stride, - g_compressor_pool_kv_buffer, - 0, - head_dim, - pool_row_stride) || - !ds4_gpu_encode_concat_f32_dim1(cb, - statescbuf, - ds4_gpu_tensor_offset(state_score), - 4, - state_row_stride, - statescbuf, - ds4_gpu_tensor_offset(state_score) + curr_offset, - 4, - state_row_stride, - g_compressor_pool_score_buffer, - 0, - head_dim, - pool_row_stride)) { - return 0; - } - - return ds4_gpu_encode_dsv4_softmax_pool(cb, - out, - g_compressor_pool_kv_buffer, - 0, - pool_row_stride, - sizeof(float), - packed_bytes, - g_compressor_pool_score_buffer, - 0, - pool_row_stride, - sizeof(float), - packed_bytes, - 8, - head_dim, - 1); -} - -static int ds4_gpu_encode_compressor_shift_ratio4( - id cb, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - uint32_t width) { - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!cb || !statekvbuf || !statescbuf || !g_dsv4_ratio4_shift_pipeline || width == 0) return 0; - - ds4_gpu_dsv4_ratio4_shift_args args = { .width = width }; - const uint32_t n = 4u * width; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_ratio4_shift_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:1]; - [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -int ds4_gpu_compressor_prefill_tensor( - ds4_gpu_tensor *comp_cache, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_rot, - uint32_t n_ctx_orig, - bool quantize_fp8, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || - head_dim == 0 || ratio == 0 || n_tokens == 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || - norm_type != 0u) { - return 0; - } - - @autoreleasepool { - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint32_t n_comp = n_tokens / ratio; - const uint32_t cutoff = n_comp * ratio; - const uint32_t rem = n_tokens - cutoff; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset) { - fprintf(stderr, "ds4: Metal compressor prefill tensor range is outside the mapped model\n"); - return 0; - } - - id kvbuf = ds4_gpu_tensor_buffer(kv); - id scbuf = ds4_gpu_tensor_buffer(sc); - id compbuf = ds4_gpu_tensor_buffer(comp_cache); - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || - ds4_gpu_tensor_bytes(kv) < kv_bytes || - ds4_gpu_tensor_bytes(sc) < kv_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes || - (n_comp && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { - fprintf(stderr, "ds4: Metal compressor prefill received undersized buffers\n"); - return 0; - } - - uint64_t ape_inner = 0; - id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); - if (!apebuf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - int ok = 1; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) ok = 0; - - if (ok) { - ok = ds4_gpu_encode_fill_f32_rows(cb, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - width, - state_rows, - 0.0f) && - ds4_gpu_encode_fill_f32_rows(cb, - statescbuf, - ds4_gpu_tensor_offset(state_score), - width, - state_rows, - ds4_gpu_negative_infinity()); - } - - if (ok && ratio == 4u) { - int32_t rows_prev[4] = { 0, 1, 2, 3 }; - const int have_prev = cutoff >= ratio ? 1 : 0; - const uint32_t prev_start = rem == 0 ? cutoff - ratio : cutoff - ratio; - if (have_prev) { - ok = ds4_gpu_encode_compressor_set_rows_projected(cb, - state_kv, - state_score, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)prev_start * width * sizeof(float), - scbuf, - ds4_gpu_tensor_offset(sc) + - (NSUInteger)prev_start * width * sizeof(float), - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0 + prev_start, - rows_prev, - 4, - state_rows); - } - if (ok && rem != 0) { - int32_t rows_cur[4]; - for (uint32_t i = 0; i < rem; i++) rows_cur[i] = (int32_t)(ratio + i); - ok = ds4_gpu_encode_compressor_set_rows_projected(cb, - state_kv, - state_score, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)cutoff * width * sizeof(float), - scbuf, - ds4_gpu_tensor_offset(sc) + - (NSUInteger)cutoff * width * sizeof(float), - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0 + cutoff, - rows_cur, - rem, - state_rows); - } - } else if (ok && rem != 0) { - int32_t rows[128]; - if (rem > (uint32_t)(sizeof(rows) / sizeof(rows[0]))) { - fprintf(stderr, "ds4: Metal compressor prefill remainder exceeds local row list\n"); - ok = 0; - } else { - for (uint32_t i = 0; i < rem; i++) rows[i] = (int32_t)i; - ok = ds4_gpu_encode_compressor_set_rows_projected(cb, - state_kv, - state_score, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)cutoff * width * sizeof(float), - scbuf, - ds4_gpu_tensor_offset(sc) + - (NSUInteger)cutoff * width * sizeof(float), - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0 + cutoff, - rows, - rem, - state_rows); - } - } - - if (ok && n_comp != 0) { - const NSUInteger score_bytes = (NSUInteger)cutoff * width * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, - &g_compressor_store_score_bytes, - score_bytes, - "ds4_compressor_store_score")) { - ok = 0; - } - if (ok) { - ok = ds4_gpu_encode_compressor_score_with_ape(cb, - scbuf, - ds4_gpu_tensor_offset(sc), - g_compressor_store_score_buffer, - 0, - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0, - cutoff); - } - - if (ok && ratio == 4u) { - const int direct_pool_mode = - ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); - if (ok && direct_pool_mode < 0) ok = 0; - const uint64_t direct_output_bytes = - (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t direct_input_bytes = - (uint64_t)n_comp * 4u * width * sizeof(float); - const bool direct_pool_overlap = direct_pool_mode > 0 && - ds4_gpu_buffer_ranges_overlap( - compbuf, - ds4_gpu_tensor_offset(comp_cache), - direct_output_bytes, - kvbuf, - ds4_gpu_tensor_offset(kv), - direct_input_bytes); - const bool use_direct_pool = - direct_pool_mode > 0 && !direct_pool_overlap; - const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); - if (ok && !use_direct_pool && - (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, - &g_compressor_pool_kv_bytes, - pack_bytes, - "ds4_compressor_pool_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, - &g_compressor_pool_score_bytes, - pack_bytes, - "ds4_compressor_pool_score"))) { - ok = 0; - } - const int pack_fusion_mode = use_direct_pool ? 0 : - ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); - if (ok && pack_fusion_mode < 0) ok = 0; - const bool use_pack_fusion = pack_fusion_mode > 0; - if (ok && use_direct_pool) { - ok = ds4_gpu_encode_compressor_ratio4_direct_pool( - cb, - comp_cache, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_store_score_buffer, - 0, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - statescbuf, - ds4_gpu_tensor_offset(state_score), - head_dim, - n_comp, - false); - } - if (ok && use_pack_fusion) { - ok = ds4_gpu_encode_compressor_pack_ratio4( - cb, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_store_score_buffer, - 0, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - statescbuf, - ds4_gpu_tensor_offset(state_score), - head_dim, - n_comp, - false); - } - if (ok && !use_direct_pool && !use_pack_fusion) { - ok = ds4_gpu_encode_fill_f32_rows(cb, - g_compressor_pool_kv_buffer, - 0, - head_dim, - 8u * n_comp, - 0.0f) && - ds4_gpu_encode_fill_f32_rows(cb, - g_compressor_pool_score_buffer, - 0, - head_dim, - 8u * n_comp, - ds4_gpu_negative_infinity()); - } - if (ok && !use_direct_pool && !use_pack_fusion) { - const uint64_t src_row_stride = (uint64_t)width * sizeof(float); - const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; - const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); - const uint64_t dst_plane_stride = 8ull * dst_row_stride; - ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)head_dim * sizeof(float), - g_compressor_pool_kv_buffer, - (NSUInteger)4u * head_dim * sizeof(float), - head_dim, - ratio, - n_comp, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d(cb, - g_compressor_store_score_buffer, - (NSUInteger)head_dim * sizeof(float), - g_compressor_pool_score_buffer, - (NSUInteger)4u * head_dim * sizeof(float), - head_dim, - ratio, - n_comp, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride); - } - if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { - const uint64_t src_row_stride = (uint64_t)width * sizeof(float); - const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; - const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); - const uint64_t dst_plane_stride = 8ull * dst_row_stride; - ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_pool_kv_buffer, - dst_plane_stride, - head_dim, - ratio, - n_comp - 1u, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d(cb, - g_compressor_store_score_buffer, - 0, - g_compressor_pool_score_buffer, - dst_plane_stride, - head_dim, - ratio, - n_comp - 1u, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride); - } - if (ok && !use_direct_pool) { - ok = ds4_gpu_encode_dsv4_softmax_pool(cb, - comp_cache, - g_compressor_pool_kv_buffer, - 0, - (uint64_t)head_dim * sizeof(float), - sizeof(float), - 8ull * head_dim * sizeof(float), - g_compressor_pool_score_buffer, - 0, - (uint64_t)head_dim * sizeof(float), - sizeof(float), - 8ull * head_dim * sizeof(float), - 8, - head_dim, - n_comp); - } - } else if (ok) { - const uint64_t row_stride = (uint64_t)width * sizeof(float); - ok = ds4_gpu_encode_dsv4_softmax_pool(cb, - comp_cache, - kvbuf, - ds4_gpu_tensor_offset(kv), - row_stride, - sizeof(float), - (uint64_t)ratio * row_stride, - g_compressor_store_score_buffer, - 0, - row_stride, - sizeof(float), - (uint64_t)ratio * row_stride, - ratio, - head_dim, - n_comp); - } - } - - if (ok && n_comp != 0) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, - comp_cache, - model_map, - model_size, - norm_offset, - head_dim, - n_comp, - rms_eps) != 0; - } - if (ok && n_comp != 0 && n_rot != 0) { - ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( - n_comp, 1, head_dim, n_rot, n_ctx_orig, false, - freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb = ds4_gpu_command_buffer(&owned); - ok = cb && !owned && - ds4_gpu_encode_rope_tail_inplace(cb, - compbuf, - ds4_gpu_tensor_offset(comp_cache), - &rope_args, - n_comp, - 1, - head_dim, - pos0, - ratio); - } - if (ok && n_comp != 0 && quantize_fp8) { - ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; - } - - if (!had_batch) { - const int end_ok = ds4_gpu_end_commands(); - ok = end_ok && ok; - } - return ok ? 1 : 0; - } -} - -int ds4_gpu_compressor_prefill_ratio4_replay_tensor( - ds4_gpu_tensor *comp_cache, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv, - const ds4_gpu_tensor *sc, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t pos0, - uint32_t n_tokens, - uint32_t n_rot, - uint32_t n_ctx_orig, - bool quantize_fp8, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || - head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || - norm_type != 0u) { - return 0; - } - - @autoreleasepool { - const uint32_t ratio = 4u; - const uint32_t width = 2u * head_dim; - const uint32_t state_rows = 8u; - const uint32_t n_comp = n_tokens / ratio; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset) { - fprintf(stderr, "ds4: Metal compressor replay tensor range is outside the mapped model\n"); - return 0; - } - - id kvbuf = ds4_gpu_tensor_buffer(kv); - id scbuf = ds4_gpu_tensor_buffer(sc); - id compbuf = ds4_gpu_tensor_buffer(comp_cache); - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || - ds4_gpu_tensor_bytes(kv) < kv_bytes || - ds4_gpu_tensor_bytes(sc) < kv_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes || - ds4_gpu_tensor_bytes(comp_cache) < comp_bytes) { - fprintf(stderr, "ds4: Metal compressor replay received undersized buffers\n"); - return 0; - } - - uint64_t ape_inner = 0; - id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); - if (!apebuf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - int ok = 1; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) ok = 0; - - const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float); - const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); - const int direct_pool_mode = - ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); - if (ok && direct_pool_mode < 0) ok = 0; - const uint64_t direct_output_bytes = - (uint64_t)n_comp * head_dim * sizeof(float); - const uint64_t direct_input_bytes = - (uint64_t)n_tokens * width * sizeof(float); - const uint64_t direct_state_bytes = - (uint64_t)4u * width * sizeof(float); - const bool direct_pool_overlap = direct_pool_mode > 0 && - (ds4_gpu_buffer_ranges_overlap( - compbuf, - ds4_gpu_tensor_offset(comp_cache), - direct_output_bytes, - kvbuf, - ds4_gpu_tensor_offset(kv), - direct_input_bytes) || - ds4_gpu_buffer_ranges_overlap( - compbuf, - ds4_gpu_tensor_offset(comp_cache), - direct_output_bytes, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - direct_state_bytes) || - ds4_gpu_buffer_ranges_overlap( - compbuf, - ds4_gpu_tensor_offset(comp_cache), - direct_output_bytes, - statescbuf, - ds4_gpu_tensor_offset(state_score), - direct_state_bytes)); - const bool use_direct_pool = - direct_pool_mode > 0 && !direct_pool_overlap; - if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, - &g_compressor_store_score_bytes, - score_bytes, - "ds4_compressor_store_score") || - (!use_direct_pool && - (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, - &g_compressor_pool_kv_bytes, - pack_bytes, - "ds4_compressor_pool_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, - &g_compressor_pool_score_bytes, - pack_bytes, - "ds4_compressor_pool_score"))))) { - ok = 0; - } - - if (ok) { - ok = ds4_gpu_encode_compressor_score_with_ape(cb, - scbuf, - ds4_gpu_tensor_offset(sc), - g_compressor_store_score_buffer, - 0, - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0, - n_tokens); - } - - const uint64_t src_row_stride = (uint64_t)width * sizeof(float); - const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; - const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); - const uint64_t dst_plane_stride = 8ull * dst_row_stride; - const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv); - const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score); - const int pack_fusion_mode = use_direct_pool ? 0 : - ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); - if (ok && pack_fusion_mode < 0) ok = 0; - const bool use_pack_fusion = pack_fusion_mode > 0; - - if (ok && use_direct_pool) { - ok = ds4_gpu_encode_compressor_ratio4_direct_pool( - cb, - comp_cache, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_store_score_buffer, - 0, - statekvbuf, - state_off, - statescbuf, - state_score_off, - head_dim, - n_comp, - true); - } - - if (ok && use_pack_fusion) { - ok = ds4_gpu_encode_compressor_pack_ratio4( - cb, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_store_score_buffer, - 0, - statekvbuf, - state_off, - statescbuf, - state_score_off, - head_dim, - n_comp, - true); - } - - if (ok && !use_direct_pool && !use_pack_fusion) { - ok = ds4_gpu_encode_fill_f32_rows(cb, - g_compressor_pool_kv_buffer, - 0, - head_dim, - 8u * n_comp, - 0.0f) && - ds4_gpu_encode_fill_f32_rows(cb, - g_compressor_pool_score_buffer, - 0, - head_dim, - 8u * n_comp, - ds4_gpu_negative_infinity()); - } - - if (ok && !use_direct_pool && !use_pack_fusion) { - /* - * The aligned nonzero ratio-4 path replays the current ubatch - * compressor, but seeds the first compressed row with the previous - * compressor state. Rows 0..3 are the previous half, rows 4..7 are - * the current half. - */ - ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, - statekvbuf, - state_off, - g_compressor_pool_kv_buffer, - 0, - head_dim, - ratio, - 1, - src_row_stride, - (uint64_t)ratio * src_row_stride, - dst_row_stride, - dst_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d(cb, - statescbuf, - state_score_off, - g_compressor_pool_score_buffer, - 0, - head_dim, - ratio, - 1, - src_row_stride, - (uint64_t)ratio * src_row_stride, - dst_row_stride, - dst_plane_stride); - } - if (ok && !use_direct_pool && !use_pack_fusion) { - ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)head_dim * sizeof(float), - g_compressor_pool_kv_buffer, - (NSUInteger)4u * head_dim * sizeof(float), - head_dim, - ratio, - n_comp, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d(cb, - g_compressor_store_score_buffer, - (NSUInteger)head_dim * sizeof(float), - g_compressor_pool_score_buffer, - (NSUInteger)4u * head_dim * sizeof(float), - head_dim, - ratio, - n_comp, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride); - } - if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { - ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, - kvbuf, - ds4_gpu_tensor_offset(kv), - g_compressor_pool_kv_buffer, - dst_plane_stride, - head_dim, - ratio, - n_comp - 1u, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride) && - ds4_gpu_encode_cpy_f32_f32_3d(cb, - g_compressor_store_score_buffer, - 0, - g_compressor_pool_score_buffer, - dst_plane_stride, - head_dim, - ratio, - n_comp - 1u, - src_row_stride, - src_plane_stride, - dst_row_stride, - dst_plane_stride); - } - if (ok && !use_direct_pool) { - ok = ds4_gpu_encode_dsv4_softmax_pool(cb, - comp_cache, - g_compressor_pool_kv_buffer, - 0, - dst_row_stride, - sizeof(float), - dst_plane_stride, - g_compressor_pool_score_buffer, - 0, - dst_row_stride, - sizeof(float), - dst_plane_stride, - 8, - head_dim, - n_comp); - } - if (ok) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, - comp_cache, - model_map, - model_size, - norm_offset, - head_dim, - n_comp, - rms_eps) != 0; - } - if (ok && n_rot != 0) { - ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( - n_comp, 1, head_dim, n_rot, n_ctx_orig, false, - freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb = ds4_gpu_command_buffer(&owned); - ok = cb && !owned && - ds4_gpu_encode_rope_tail_inplace(cb, - compbuf, - ds4_gpu_tensor_offset(comp_cache), - &rope_args, - n_comp, - 1, - head_dim, - pos0, - ratio); - } - if (ok && quantize_fp8) { - ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; - } - - if (ok) { - ok = ds4_gpu_encode_fill_f32_rows(cb, - statekvbuf, - state_off, - width, - state_rows, - 0.0f) && - ds4_gpu_encode_fill_f32_rows(cb, - statescbuf, - state_score_off, - width, - state_rows, - ds4_gpu_negative_infinity()); - } - if (ok) { - int32_t rows_prev[4] = { 0, 1, 2, 3 }; - const uint32_t prev_start = n_tokens - ratio; - ok = ds4_gpu_encode_compressor_set_rows_projected(cb, - state_kv, - state_score, - kvbuf, - ds4_gpu_tensor_offset(kv) + - (NSUInteger)prev_start * width * sizeof(float), - scbuf, - ds4_gpu_tensor_offset(sc) + - (NSUInteger)prev_start * width * sizeof(float), - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0 + prev_start, - rows_prev, - ratio, - state_rows); - } - - if (!had_batch) { - const int end_ok = ds4_gpu_end_commands(); - ok = end_ok && ok; - } - return ok ? 1 : 0; - } -} - -int ds4_gpu_compressor_prefill_state_ratio4_tensor( - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - const ds4_gpu_tensor *kv_tail, - const ds4_gpu_tensor *sc_tail, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint32_t head_dim, - uint32_t pos0) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map || - head_dim == 0 || (ape_type != 0u && ape_type != 1u)) { - return 0; - } - - @autoreleasepool { - const uint32_t ratio = 4u; - const uint32_t width = 2u * head_dim; - const uint32_t state_rows = 8u; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape; - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { - fprintf(stderr, "ds4: Metal compressor prefill-state APE range is outside the mapped model\n"); - return 0; - } - - id kvbuf = ds4_gpu_tensor_buffer(kv_tail); - id scbuf = ds4_gpu_tensor_buffer(sc_tail); - id statekvbuf = ds4_gpu_tensor_buffer(state_kv); - id statescbuf = ds4_gpu_tensor_buffer(state_score); - if (!kvbuf || !scbuf || !statekvbuf || !statescbuf || - ds4_gpu_tensor_bytes(kv_tail) < tail_bytes || - ds4_gpu_tensor_bytes(sc_tail) < tail_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes) { - fprintf(stderr, "ds4: Metal compressor prefill-state received undersized buffers\n"); - return 0; - } - - uint64_t ape_inner = 0; - id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); - if (!apebuf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - int ok = 1; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) ok = 0; - - if (ok) { - ok = ds4_gpu_encode_fill_f32_rows(cb, - statekvbuf, - ds4_gpu_tensor_offset(state_kv), - width, - state_rows, - 0.0f) && - ds4_gpu_encode_fill_f32_rows(cb, - statescbuf, - ds4_gpu_tensor_offset(state_score), - width, - state_rows, - ds4_gpu_negative_infinity()); - } - if (ok) { - int32_t rows[4] = { 0, 1, 2, 3 }; - ok = ds4_gpu_encode_compressor_set_rows_projected(cb, - state_kv, - state_score, - kvbuf, - ds4_gpu_tensor_offset(kv_tail), - scbuf, - ds4_gpu_tensor_offset(sc_tail), - apebuf, - (NSUInteger)ape_inner, - ape_type, - width, - ratio, - pos0, - rows, - ratio, - state_rows); - } - - if (!had_batch) { - const int end_ok = ds4_gpu_end_commands(); - ok = end_ok && ok; - } - return ok ? 1 : 0; - } -} - -int ds4_gpu_compressor_update_tensor( - const ds4_gpu_tensor *kv_cur, - const ds4_gpu_tensor *sc_cur, - ds4_gpu_tensor *state_kv, - ds4_gpu_tensor *state_score, - ds4_gpu_tensor *comp_cache, - const void *model_map, - uint64_t model_size, - uint64_t ape_offset, - uint32_t ape_type, - uint64_t norm_offset, - uint32_t norm_type, - uint32_t head_dim, - uint32_t ratio, - uint32_t pos, - uint32_t comp_row, - uint32_t n_rot, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - float rms_eps, - bool state_already_stored) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || - !model_map || head_dim == 0 || ratio == 0 || - n_rot > head_dim || (n_rot & 1u) != 0 || - (ape_type != 0u && ape_type != 1u) || - norm_type != 0u) { - return 0; - } - - @autoreleasepool { - const uint32_t coff = ratio == 4u ? 2u : 1u; - const uint32_t width = coff * head_dim; - const uint32_t state_rows = coff * ratio; - const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; - const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; - const uint64_t kv_bytes = (uint64_t)width * sizeof(float); - const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); - const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); - const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - - if (ape_offset > model_size || ape_bytes > model_size - ape_offset || - norm_offset > model_size || norm_bytes > model_size - norm_offset) { - fprintf(stderr, "ds4: Metal compressor tensor range is outside the mapped model\n"); - return 0; - } - - id kvbuf = ds4_gpu_tensor_buffer(kv_cur); - id scbuf = ds4_gpu_tensor_buffer(sc_cur); - id compbuf = ds4_gpu_tensor_buffer(comp_cache); - if (!kvbuf || !scbuf || !compbuf || - ds4_gpu_tensor_bytes(kv_cur) < kv_bytes || - ds4_gpu_tensor_bytes(sc_cur) < kv_bytes || - ds4_gpu_tensor_bytes(state_kv) < state_bytes || - ds4_gpu_tensor_bytes(state_score) < state_bytes || - (emit && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { - fprintf(stderr, "ds4: Metal compressor update received undersized buffers\n"); - return 0; - } - - if (!state_already_stored) { - const bool use_store_one = - getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL; - const int store_ok = use_store_one - ? ds4_gpu_compressor_store_one_tensor(kv_cur, - sc_cur, - state_kv, - state_score, - model_map, - model_size, - ape_offset, - ape_type, - width, - ratio, - pos) - : ds4_gpu_compressor_store_batch_tensor(kv_cur, - sc_cur, - state_kv, - state_score, - model_map, - model_size, - ape_offset, - ape_type, - head_dim, - ratio, - pos, - 1); - if (!store_ok) { - return 0; - } - } - if (!emit) return 1; - - ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( - comp_cache, - (uint64_t)comp_row * head_dim * sizeof(float), - (uint64_t)head_dim * sizeof(float)); - if (!comp_row_view) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - int ok = cb && - ds4_gpu_encode_compressor_pool(cb, - comp_row_view, - state_kv, - state_score, - head_dim, - ratio); - if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool"); - if (ok) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, - comp_row_view, - model_map, - model_size, - norm_offset, - head_dim, - 1, - rms_eps) != 0; - } - if (ok) { - const uint32_t comp_pos = pos + 1u - ratio; - ok = ds4_gpu_rope_tail_tensor(comp_row_view, - 1, - 1, - head_dim, - n_rot, - comp_pos, - n_ctx_orig, - false, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow) != 0; - } - if (ok && ratio == 4u) { - cb = ds4_gpu_command_buffer(&owned); - ok = cb && - ds4_gpu_encode_compressor_shift_ratio4(cb, - state_kv, - state_score, - width); - if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor ratio4 state shift"); - } - ds4_gpu_tensor_free(comp_row_view); - if (!ok) return 0; - } - - return 1; -} - -static int ds4_gpu_encode_fill_f32_rows( - id cb, - id buf, - NSUInteger offset, - uint32_t width, - uint32_t rows, - float value) { - if (!cb || !buf || width == 0 || rows == 0 || (width & 3u) != 0) return 0; - - ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, 1, 0.0f, 0.0f); - args.val = value; - - NSUInteger nth_max = g_unary_fill_pipeline.maxTotalThreadsPerThreadgroup; - if (nth_max > 256u) nth_max = 256u; - NSUInteger nth = (NSUInteger)args.ne00; - if (nth > nth_max) nth = nth_max; - if (nth == 0) nth = 1u; - const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_unary_fill_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:buf offset:offset atIndex:1]; - [enc setBuffer:buf offset:offset atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, - (NSUInteger)args.ne02, - (NSUInteger)args.ne03) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -int ds4_gpu_attention_output_q8_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_tensor *group_tmp, - ds4_gpu_tensor *low_tmp, - 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, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map || - group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || - group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - const uint64_t low_dim = (uint64_t)n_groups * rank; - if ((group_dim % 32u) != 0 || (low_dim % 32u) != 0 || low_dim > UINT32_MAX) { - fprintf(stderr, "ds4: Metal attention output batch received invalid q8 dimensions\n"); - return 0; - } - const uint64_t row_a_bytes = (group_dim / 32u) * 34u; - const uint64_t row_b_bytes = (low_dim / 32u) * 34u; - const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; - const uint64_t out_b_bytes = out_dim * row_b_bytes; - if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || - out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { - fprintf(stderr, "ds4: Metal attention output batch weights are outside the mapped model\n"); - return 0; - } - - const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); - const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(low) < low_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal attention output batch received undersized buffers\n"); - return 0; - } - (void)group_tmp; - (void)low_tmp; - - const bool use_direct_low = - n_tokens < 32u && getenv("DS4_METAL_DISABLE_ATTN_OUT_LOW_DIRECT") == NULL; - /* The exported TensorOps attention-output kernel is a 64-token tile. - * Keep this on full tiles only; smaller multiples of 32 use the legacy - * path instead of relying on cooperative tensor partial RHS bounds. */ - const bool use_mpp_low = - n_tokens >= 32u && - (n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 && - ds4_gpu_use_mpp_attn_out_low_matmul(); - const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); - id group_ids_buffer = nil; - if (!use_direct_low && !use_mpp_low) { - if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { - group_ids_buffer = - ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); - if (!group_ids_buffer) { - return 0; - } - } else { - if (!ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, - &g_attn_out_group_ids_bytes, - ids_bytes, - "ds4_attention_output_group_ids")) { - return 0; - } - group_ids_buffer = g_attn_out_group_ids_buffer; - } - int32_t *ids = (int32_t *)[group_ids_buffer contents]; - for (uint32_t t = 0; t < n_tokens; t++) { - for (uint32_t group = 0; group < n_groups; group++) { - ids[(uint64_t)t * n_groups + group] = (int32_t)group; - } - } - } - - uint64_t out_a_inner = 0; - id out_a_buf = - ds4_gpu_wrap_model_range(model_map, model_size, - out_a_offset, out_a_bytes, - &out_a_inner); - if (!out_a_buf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - bool ok = true; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) { - ok = false; - } - const bool attn_out_profile = - getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") != NULL && g_batch_cb != nil; - if (ok && attn_out_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - ok = false; - } else { - cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) ok = false; - } - } - double attn_out_t0 = attn_out_profile ? ds4_gpu_now_ms() : 0.0; -#define DS4_METAL_PROFILE_ATTN_OUT_STAGE(name) do { \ - if (ok && attn_out_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = false; \ - } else { \ - const double now_ms = ds4_gpu_now_ms(); \ - fprintf(stderr, \ - "ds4: Metal attention output stage tokens=%u %s=%.3f ms\n", \ - n_tokens, (name), now_ms - attn_out_t0); \ - attn_out_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = false; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb || owned) ok = false; \ - } \ - } \ - } \ - } while (0) - - if (ok) { - /* - * Batched attention-output projections switch from the vector - * kernel to the SIMD matrix kernel once the batch has at least 32 - * tokens. This preserves the single-token generation path while - * keeping prefill accumulation stable. - */ - if (use_mpp_low) { - ds4_gpu_mul_mm_id_args mm_args = - ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, - (uint32_t)rank, - n_groups, - row_a_bytes, - (uint64_t)rank * row_a_bytes, - n_groups, - n_groups, - n_tokens); - /* - * Direct RHS lets MPP read the dense low-rank activation tile - * directly from device memory instead of staging a second - * threadgroup tile. The retained attention-output path is the - * 64-token direct-RHS kernel; the older staged-RHS and 32-token - * variants were not kept as alternate runtime modes. - */ - const char *attn_out_pipeline_name = - "kernel_attn_out_low_q8_0_mpp_direct_rhs_n64"; - id mm_pipeline = - ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false); - ok = ds4_gpu_encode_attn_out_low_q8_mpp(cb, - mm_pipeline, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low)) != 0; - if (!ok) { - ds4_gpu_warn_mpp_fallback(); - if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { - if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { - group_ids_buffer = - ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); - } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, - &g_attn_out_group_ids_bytes, - ids_bytes, - "ds4_attention_output_group_ids")) { - group_ids_buffer = g_attn_out_group_ids_buffer; - } - if (group_ids_buffer) { - int32_t *ids = (int32_t *)[group_ids_buffer contents]; - for (uint32_t t = 0; t < n_tokens; t++) { - for (uint32_t group = 0; group < n_groups; group++) { - ids[(uint64_t)t * n_groups + group] = (int32_t)group; - } - } - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, - n_groups, - n_groups, - n_groups, - n_tokens); - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); - id fallback_pipeline = - ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); - ok = ds4_gpu_encode_mul_mm_id(cb, - map_pipeline, - fallback_pipeline, - &map_args, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - group_ids_buffer, - 0) != 0; - } - } - } - } else if (n_tokens >= 32u && ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, - n_groups, - n_groups, - n_groups, - n_tokens); - ds4_gpu_mul_mm_id_args mm_args = - ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, - (uint32_t)rank, - n_groups, - row_a_bytes, - (uint64_t)rank * row_a_bytes, - n_groups, - n_groups, - n_tokens); - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); - id mm_pipeline = - ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); - ok = ds4_gpu_encode_mul_mm_id(cb, - map_pipeline, - mm_pipeline, - &map_args, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - group_ids_buffer, - 0) != 0; - } else if (use_direct_low) { - ds4_gpu_mul_mv_id_args args = { - .nei0 = (int32_t)n_groups, - .nei1 = (int32_t)n_tokens, - .nbi1 = 0, - .ne00 = (int32_t)group_dim, - .ne01 = (int32_t)rank, - .ne02 = (int32_t)n_groups, - .nb00 = 34, - .nb01 = row_a_bytes, - .nb02 = (uint64_t)rank * row_a_bytes, - .ne10 = (int32_t)group_dim, - .ne11 = (int32_t)n_groups, - .ne12 = (int32_t)n_tokens, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)group_dim * sizeof(float), - .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), - .ne0 = (int32_t)rank, - .ne1 = (int32_t)n_groups, - .nb1 = (uint64_t)rank * sizeof(float), - .nr0 = 2, - }; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); - ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, - pipeline, - &args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - 32u * 2u * sizeof(float), - 4, - true) != 0; - } else { - ds4_gpu_mul_mv_id_args args = { - .nei0 = (int32_t)n_groups, - .nei1 = (int32_t)n_tokens, - .nbi1 = (uint64_t)n_groups * sizeof(int32_t), - .ne00 = (int32_t)group_dim, - .ne01 = (int32_t)rank, - .ne02 = (int32_t)n_groups, - .nb00 = 34, - .nb01 = row_a_bytes, - .nb02 = (uint64_t)rank * row_a_bytes, - .ne10 = (int32_t)group_dim, - .ne11 = (int32_t)n_groups, - .ne12 = (int32_t)n_tokens, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)group_dim * sizeof(float), - .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), - .ne0 = (int32_t)rank, - .ne1 = (int32_t)n_groups, - .nb1 = (uint64_t)rank * sizeof(float), - .nr0 = 2, - }; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); - ok = ds4_gpu_encode_mul_mv_id(cb, - pipeline, - &args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - group_ids_buffer, - 0, - 32u * 2u * sizeof(float), - 4, - true) != 0; - } - } - DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj"); - - if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, - out_b_offset, - low_dim, out_dim, low, n_tokens) != 0; - } - DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj"); - - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } -#undef DS4_METAL_PROFILE_ATTN_OUT_STAGE - return ok ? 1 : 0; - } -} - -int ds4_gpu_attention_output_q4_K_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *low, - ds4_gpu_tensor *group_tmp, - ds4_gpu_tensor *low_tmp, - const void *model_map, - uint64_t model_size, - uint64_t out_a_offset, - uint64_t out_b_offset, - uint32_t out_b_type, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !low || !heads || !model_map || - group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || - group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - if (n_tokens < 32u) return 0; - - @autoreleasepool { - const uint64_t low_dim = (uint64_t)n_groups * rank; - if ((group_dim % 256u) != 0 || (low_dim % 256u) != 0 || low_dim > UINT32_MAX) { - return 0; - } - - uint64_t row_a_bytes = 0; - uint64_t row_b_bytes = 0; - if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, (uint32_t)group_dim, &row_a_bytes) || - !ds4_gpu_quant_row_bytes(out_b_type, (uint32_t)low_dim, &row_b_bytes)) { - return 0; - } - - const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; - const uint64_t out_b_bytes = out_dim * row_b_bytes; - if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || - out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { - fprintf(stderr, "ds4: Metal Q4 attention output batch weights are outside the mapped model\n"); - return 0; - } - - const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); - const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(low) < low_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal Q4 attention output batch received undersized buffers\n"); - return 0; - } - (void)group_tmp; - (void)low_tmp; - - const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); - id group_ids_buffer = nil; - if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { - group_ids_buffer = - ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids"); - } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, - &g_attn_out_group_ids_bytes, - ids_bytes, - "ds4_attention_output_group_ids")) { - group_ids_buffer = g_attn_out_group_ids_buffer; - } - if (!group_ids_buffer) return 0; - - int32_t *ids = (int32_t *)[group_ids_buffer contents]; - for (uint32_t t = 0; t < n_tokens; t++) { - for (uint32_t group = 0; group < n_groups; group++) { - ids[(uint64_t)t * n_groups + group] = (int32_t)group; - } - } - - uint64_t out_a_inner = 0; - id out_a_buf = - ds4_gpu_wrap_model_range(model_map, model_size, - out_a_offset, out_a_bytes, - &out_a_inner); - if (!out_a_buf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - bool ok = true; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) ok = false; - - if (ok) { - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, - n_groups, - n_groups, - n_groups, - n_tokens); - ds4_gpu_mul_mm_id_args mm_args = - ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, - (uint32_t)rank, - n_groups, - row_a_bytes, - (uint64_t)rank * row_a_bytes, - n_groups, - n_groups, - n_tokens); - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); - id mm_pipeline = - ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); - ok = ds4_gpu_encode_mul_mm_id(cb, - map_pipeline, - mm_pipeline, - &map_args, - &mm_args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - group_ids_buffer, - 0) != 0; - } - - if (ok) { - ok = ds4_gpu_matmul_quant_tensor(out, - model_map, - model_size, - out_b_offset, - out_b_type, - low_dim, - out_dim, - low, - n_tokens) != 0; - } - - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } - return ok ? 1 : 0; - } -} - -int ds4_gpu_attention_output_q8_batch_f16_tensor( - ds4_gpu_tensor *out_h, - 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, - uint64_t out_dim, - const ds4_gpu_tensor *heads, - uint32_t n_tokens) { - (void)out_h; (void)low; (void)model_map; (void)model_size; - (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank; - (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; - return 0; -} - -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) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if ((full_in_dim & 31u) != 0 || (k_off & 31u) != 0 || (k_cnt & 31u) != 0 || - k_cnt == 0 || k_off + k_cnt > full_in_dim || - full_in_dim > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || - ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { - fprintf(stderr, "ds4: Metal Q8_0 kslice matmul received undersized buffers\n"); - return 0; - } - const uint64_t row_bytes = (full_in_dim / 32u) * 34u; - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8_0 kslice weights are outside the mapped model\n"); - return 0; - } - uint64_t inner = 0; - id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_offset, weight_bytes, &inner); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - /* Same matvec kernel as the full projection: ne00 bounds the k loop - * while nb01/nb02 keep the full-row stride, so each row reads only - * the owned k window. */ - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(full_in_dim, out_dim); - mv_args.ne00 = (int32_t)k_cnt; - mv_args.ne10 = (int32_t)k_cnt; - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - if (out_dim > 65536u) mv_dispatch.nsg = 8; - mv_args.nr0 = mv_dispatch.nr0; - id pipeline = - ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); - if (!pipeline) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)(inner + (k_off / 32u) * 34u) atIndex:1]; - [enc setBuffer:xbuf - offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) - atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 kslice matvec")) { - return 0; - } - return 1; - } -} - -int ds4_gpu_matmul_quant_kslice_tensor( - ds4_gpu_tensor *out, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t weight_type, - 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) { - if (weight_type == DS4_METAL_TENSOR_Q8_0) { - return ds4_gpu_matmul_q8_0_kslice_tensor(out, - model_map, - model_size, - weight_offset, - full_in_dim, - k_off, - k_cnt, - out_dim, - x, - x_elem_off); - } - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || - full_in_dim == 0 || k_cnt == 0 || out_dim == 0 || - k_off + k_cnt > full_in_dim || - full_in_dim > UINT32_MAX || k_cnt > UINT32_MAX || - out_dim > UINT32_MAX) { - return 0; - } - - uint64_t block_elems = 0; - uint64_t block_bytes = 0; - if (weight_type == DS4_METAL_TENSOR_Q4_K) { - block_elems = 256u; - block_bytes = 144u; - } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { - block_elems = 32u; - block_bytes = 18u; - } else { - fprintf(stderr, "ds4: Metal quant kslice received unsupported type %u\n", weight_type); - return 0; - } - if ((full_in_dim % block_elems) != 0 || - (k_off % block_elems) != 0 || - (k_cnt % block_elems) != 0) { - fprintf(stderr, "ds4: Metal quant kslice dimensions are not block aligned\n"); - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id outbuf = ds4_gpu_tensor_buffer(out); - if (!xbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || - ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { - fprintf(stderr, "ds4: Metal quant kslice matmul received undersized buffers\n"); - return 0; - } - - const uint64_t row_bytes = (full_in_dim / block_elems) * block_bytes; - const uint64_t weight_bytes = out_dim * row_bytes; - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal quant kslice weights are outside the mapped model\n"); - return 0; - } - uint64_t inner = 0; - id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_offset, weight_bytes, &inner); - if (!wbuf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (weight_type == DS4_METAL_TENSOR_Q4_K && - (k_cnt % 256u) == 0 && - getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { - const int16_t nsg = 2; - id pipeline = - ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); - if (pipeline) { - ds4_gpu_q8_0_matvec_args args = { - .ne00 = (int32_t)k_cnt, - .ne01 = (int32_t)out_dim, - .ne02 = 1, - .nb00 = 1, - .nb01 = row_bytes, - .nb02 = row_bytes * out_dim, - .nb03 = row_bytes * out_dim, - .ne10 = (int32_t)k_cnt, - .ne11 = 1, - .ne12 = 1, - .nb10 = sizeof(float), - .nb11 = k_cnt * sizeof(float), - .nb12 = k_cnt * sizeof(float), - .nb13 = k_cnt * sizeof(float), - .ne0 = (int32_t)out_dim, - .ne1 = 1, - .nr0 = 2, - .r2 = 1, - .r3 = 1, - }; - const uint64_t rows_ptg = (uint64_t)nsg * 2u; - const uint64_t w_skip = (k_off / block_elems) * block_bytes; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:wbuf offset:(NSUInteger)(inner + w_skip) atIndex:1]; - [enc setBuffer:xbuf - offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) - atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:32 atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K kslice matvec")) return 0; - return 1; - } - } - - fprintf(stderr, "ds4: Metal quant kslice has no kernel for type %u\n", - weight_type); - return 0; - } -} - -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) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !low || !heads || !model_map || - group_dim == 0 || rank == 0 || group_cnt == 0 || - group0 + group_cnt > n_groups_total || - (group_dim % 32u) != 0 || ((rank * group_cnt) % 32u) != 0 || - group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - const uint64_t low_dim_total = (uint64_t)n_groups_total * rank; - const uint64_t row_a_bytes = (group_dim / 32u) * 34u; - const uint64_t a_group_bytes = rank * row_a_bytes; - const uint64_t out_a_bytes = (uint64_t)n_groups_total * a_group_bytes; - if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { - fprintf(stderr, "ds4: Metal TP attention output weights are outside the mapped model\n"); - return 0; - } - /* The heads buffer holds only the owned groups, compact at its - * base (the head slice keeps q/attention output halves packed). */ - if (ds4_gpu_tensor_bytes(heads) < (uint64_t)group_cnt * group_dim * sizeof(float) || - ds4_gpu_tensor_bytes(low) < (uint64_t)group_cnt * rank * sizeof(float) || - ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { - fprintf(stderr, "ds4: Metal TP attention output received undersized buffers\n"); - return 0; - } - - uint64_t out_a_inner = 0; - id out_a_buf = - ds4_gpu_wrap_model_range(model_map, model_size, - out_a_offset, out_a_bytes, &out_a_inner); - if (!out_a_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) return 0; - - /* Low projection for the owned groups only: identical dispatch to - * the single-node direct path with the weight base, heads input and - * group count shifted to the slice. The owned low half lands - * compactly at low[0 .. group_cnt*rank). */ - ds4_gpu_mul_mv_id_args args = { - .nei0 = (int32_t)group_cnt, - .nei1 = 1, - .nbi1 = 0, - .ne00 = (int32_t)group_dim, - .ne01 = (int32_t)rank, - .ne02 = (int32_t)group_cnt, - .nb00 = 34, - .nb01 = row_a_bytes, - .nb02 = a_group_bytes, - .ne10 = (int32_t)group_dim, - .ne11 = (int32_t)group_cnt, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)group_dim * sizeof(float), - .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), - .ne0 = (int32_t)rank, - .ne1 = (int32_t)group_cnt, - .nb1 = (uint64_t)rank * sizeof(float), - .nr0 = 2, - }; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); - int ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, - pipeline, - &args, - out_a_buf, - (NSUInteger)(out_a_inner + (uint64_t)group0 * a_group_bytes), - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - 32u * 2u * sizeof(float), - 4, - true); - if (!ok) return 0; - - /* Expand projection over the owned k window only; the result is this - * rank's partial attention block output. */ - return ds4_gpu_matmul_q8_0_kslice_tensor(out, model_map, model_size, - out_b_offset, - low_dim_total, - (uint64_t)group0 * rank, - (uint64_t)group_cnt * rank, - out_dim, low, 0); - } -} - -int ds4_gpu_attention_output_low_q8_tensor( - ds4_gpu_tensor *low, - const void *model_map, - uint64_t model_size, - uint64_t out_a_offset, - uint64_t group_dim, - uint64_t rank, - uint32_t n_groups, - const ds4_gpu_tensor *heads) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || - n_groups == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - const uint64_t low_dim = (uint64_t)n_groups * rank; - if ((group_dim % 32u) != 0 || low_dim > UINT32_MAX) { - fprintf(stderr, "ds4: Metal attention output low received invalid q8 dimensions\n"); - return 0; - } - - const uint64_t row_a_bytes = (group_dim / 32u) * 34u; - const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; - if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { - fprintf(stderr, "ds4: Metal attention output low weights are outside the mapped model\n"); - return 0; - } - - const uint64_t heads_bytes = (uint64_t)n_groups * group_dim * sizeof(float); - const uint64_t low_bytes = low_dim * sizeof(float); - if (ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(low) < low_bytes) { - fprintf(stderr, "ds4: Metal attention output low received undersized buffers\n"); - return 0; - } - - uint64_t out_a_inner = 0; - id out_a_buf = - ds4_gpu_wrap_model_range(model_map, model_size, - out_a_offset, out_a_bytes, - &out_a_inner); - if (!out_a_buf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - bool ok = true; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) { - ok = false; - } - - if (ok) { - ds4_gpu_mul_mv_id_args args = { - .nei0 = (int32_t)n_groups, - .nei1 = 1, - .nbi1 = 0, - .ne00 = (int32_t)group_dim, - .ne01 = (int32_t)rank, - .ne02 = (int32_t)n_groups, - .nb00 = 34, - .nb01 = row_a_bytes, - .nb02 = (uint64_t)rank * row_a_bytes, - .ne10 = (int32_t)group_dim, - .ne11 = (int32_t)n_groups, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)group_dim * sizeof(float), - .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), - .ne0 = (int32_t)rank, - .ne1 = (int32_t)n_groups, - .nb1 = (uint64_t)rank * sizeof(float), - .nr0 = 2, - }; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); - ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, - pipeline, - &args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - 32u * 2u * sizeof(float), - 4, - true) != 0; - } - - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } - return ok ? 1 : 0; - } -} - -int ds4_gpu_attention_output_low_q4_K_slice_tensor( - ds4_gpu_tensor *low, - const void *model_map, - uint64_t model_size, - uint64_t out_a_offset, - uint64_t group_dim, - uint64_t rank, - uint32_t group0, - uint32_t group_cnt, - const ds4_gpu_tensor *heads) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || - group_cnt == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - if ((group_dim % 256u) != 0) { - fprintf(stderr, "ds4: Metal attention output low received invalid Q4_K dimensions\n"); - return 0; - } - - uint64_t row_a_bytes = 0; - if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, - (uint32_t)group_dim, - &row_a_bytes)) { - return 0; - } - if (rank > UINT64_MAX / row_a_bytes) return 0; - const uint64_t group_weight_bytes = rank * row_a_bytes; - if (group0 > UINT64_MAX / group_weight_bytes || - group_cnt > UINT64_MAX / group_weight_bytes) { - return 0; - } - const uint64_t group_skip = (uint64_t)group0 * group_weight_bytes; - const uint64_t out_a_bytes = (uint64_t)group_cnt * group_weight_bytes; - if (out_a_offset > UINT64_MAX - group_skip || - out_a_offset + group_skip > model_size || - out_a_bytes > model_size - (out_a_offset + group_skip)) { - fprintf(stderr, "ds4: Metal Q4 attention output low weights are outside the mapped model\n"); - return 0; - } - - const uint64_t heads_bytes = (uint64_t)group_cnt * group_dim * sizeof(float); - const uint64_t low_bytes = (uint64_t)group_cnt * rank * sizeof(float); - if (ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(low) < low_bytes) { - fprintf(stderr, "ds4: Metal Q4 attention output low received undersized buffers\n"); - return 0; - } - - uint64_t out_a_inner = 0; - id out_a_buf = - ds4_gpu_wrap_model_range(model_map, model_size, - out_a_offset + group_skip, - out_a_bytes, - &out_a_inner); - if (!out_a_buf) return 0; - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - - bool ok = true; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb || owned) { - ok = false; - } - - if (ok) { - ds4_gpu_mul_mv_id_args args = { - .nei0 = (int32_t)group_cnt, - .nei1 = 1, - .nbi1 = 0, - .ne00 = (int32_t)group_dim, - .ne01 = (int32_t)rank, - .ne02 = (int32_t)group_cnt, - .nb00 = 1, - .nb01 = row_a_bytes, - .nb02 = group_weight_bytes, - .ne10 = (int32_t)group_dim, - .ne11 = (int32_t)group_cnt, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = (uint64_t)group_dim * sizeof(float), - .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), - .ne0 = (int32_t)rank, - .ne1 = (int32_t)group_cnt, - .nb1 = (uint64_t)rank * sizeof(float), - .nr0 = 2, - }; - const NSUInteger nsg = 2; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q4_K_f32", (int16_t)nsg); - ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, - pipeline, - &args, - out_a_buf, - (NSUInteger)out_a_inner, - ds4_gpu_tensor_buffer(heads), - ds4_gpu_tensor_offset(heads), - ds4_gpu_tensor_buffer(low), - ds4_gpu_tensor_offset(low), - 32u, - nsg, - false) != 0; - } - - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } - return ok ? 1 : 0; - } -} - -static NSUInteger ds4_gpu_align_up_ns(NSUInteger value, NSUInteger align) { - return (value + align - 1u) & ~(align - 1u); -} - -static int ds4_gpu_encode_cpy_f32_f32_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n) { - if (!cb || !src || !dst || n == 0) return 0; - - ds4_gpu_cpy_args args = - ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(float)); - const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f32_pipeline); - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f32_f32_3d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride) { - if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; - - ds4_gpu_cpy_args args = { - .nk0 = (int64_t)cols, - .ne00 = (int64_t)cols, - .ne01 = (int64_t)rows, - .ne02 = (int64_t)planes, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src_row_stride, - .nb02 = src_plane_stride, - .nb03 = (uint64_t)planes * src_plane_stride, - .ne0 = (int64_t)cols, - .ne1 = (int64_t)rows, - .ne2 = (int64_t)planes, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = dst_row_stride, - .nb2 = dst_plane_stride, - .nb3 = (uint64_t)planes * dst_plane_stride, - }; - const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); - const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_col_stride, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride) { - if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; - - ds4_gpu_cpy_args args = { - .nk0 = (int64_t)cols, - .ne00 = (int64_t)cols, - .ne01 = (int64_t)rows, - .ne02 = (int64_t)planes, - .ne03 = 1, - .nb00 = src_col_stride, - .nb01 = src_row_stride, - .nb02 = src_plane_stride, - .nb03 = (uint64_t)planes * src_plane_stride, - .ne0 = (int64_t)cols, - .ne1 = (int64_t)rows, - .ne2 = (int64_t)planes, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = dst_row_stride, - .nb2 = dst_plane_stride, - .nb3 = (uint64_t)planes * dst_plane_stride, - }; - const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); - const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f32_f16_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n) { - if (!cb || !src || !dst || n == 0) return 0; - - const int use_contiguous = - ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; - id pipeline = use_contiguous - ? g_cpy_contig_f32_f16_pipeline - : g_cpy_f32_f16_pipeline; - const NSUInteger work_items = use_contiguous - ? ((NSUInteger)n + 3u) / 4u - : (NSUInteger)n; - const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); - const NSUInteger groups = (work_items + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - if (use_contiguous) { - [enc setBytes:&n length:sizeof(n) atIndex:0]; - } else { - ds4_gpu_cpy_args args = - ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); - [enc setBytes:&args length:sizeof(args) atIndex:0]; - } - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f32_f16_2d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint64_t src_row_stride, - uint64_t dst_row_stride) { - if (!cb || !src || !dst || cols == 0 || rows == 0) return 0; - - ds4_gpu_cpy_args args = { - .nk0 = (int64_t)cols, - .ne00 = (int64_t)cols, - .ne01 = (int64_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src_row_stride, - .nb02 = (uint64_t)rows * src_row_stride, - .nb03 = (uint64_t)rows * src_row_stride, - .ne0 = (int64_t)cols, - .ne1 = (int64_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(uint16_t), - .nb1 = dst_row_stride, - .nb2 = (uint64_t)rows * dst_row_stride, - .nb3 = (uint64_t)rows * dst_row_stride, - }; - const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); - const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f32_f16_3d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride) { - if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; - - ds4_gpu_cpy_args args = { - .nk0 = (int64_t)cols, - .ne00 = (int64_t)cols, - .ne01 = (int64_t)rows, - .ne02 = (int64_t)planes, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src_row_stride, - .nb02 = src_plane_stride, - .nb03 = (uint64_t)planes * src_plane_stride, - .ne0 = (int64_t)cols, - .ne1 = (int64_t)rows, - .ne2 = (int64_t)planes, - .ne3 = 1, - .nb0 = sizeof(uint16_t), - .nb1 = dst_row_stride, - .nb2 = dst_plane_stride, - .nb3 = (uint64_t)planes * dst_plane_stride, - }; - const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); - const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f16_f16_3d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t cols, - uint32_t rows, - uint32_t planes, - uint64_t src_row_stride, - uint64_t src_plane_stride, - uint64_t dst_row_stride, - uint64_t dst_plane_stride) { - if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; - - ds4_gpu_cpy_args args = { - .nk0 = (int64_t)cols, - .ne00 = (int64_t)cols, - .ne01 = (int64_t)rows, - .ne02 = (int64_t)planes, - .ne03 = 1, - .nb00 = sizeof(uint16_t), - .nb01 = src_row_stride, - .nb02 = src_plane_stride, - .nb03 = (uint64_t)planes * src_plane_stride, - .ne0 = (int64_t)cols, - .ne1 = (int64_t)rows, - .ne2 = (int64_t)planes, - .ne3 = 1, - .nb0 = sizeof(uint16_t), - .nb1 = dst_row_stride, - .nb2 = dst_plane_stride, - .nb3 = (uint64_t)planes * dst_plane_stride, - }; - const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f16_f16_pipeline); - const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f16_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_cpy_f16_f32_1d( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t n) { - if (!cb || !src || !dst || n == 0) return 0; - - const int use_contiguous = - ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; - id pipeline = use_contiguous - ? g_cpy_contig_f16_f32_pipeline - : g_cpy_f16_f32_pipeline; - const NSUInteger work_items = use_contiguous - ? ((NSUInteger)n + 3u) / 4u - : (NSUInteger)n; - const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); - const NSUInteger groups = (work_items + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - if (use_contiguous) { - [enc setBytes:&n length:sizeof(n) atIndex:0]; - } else { - ds4_gpu_cpy_args args = - ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); - [enc setBytes:&args length:sizeof(args) atIndex:0]; - } - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_copy_to_f16_1d( - id cb, - id src, - NSUInteger src_off, - bool src_is_f16, - id dst, - NSUInteger dst_off, - uint32_t n) { - if (!cb || !src || !dst) return 0; - if (n == 0) return 1; - if (!src_is_f16) { - return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n); - } - - if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) { - const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u; - const NSUInteger nth = ds4_gpu_cpy_threads( - (uint32_t)work_items, - g_cpy_contig_f16_f16_pipeline); - const NSUInteger groups = (work_items + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline]; - [enc setBytes:&n length:sizeof(n) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; - } - - if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder(); - id blit = [cb blitCommandEncoder]; - if (!blit) return 0; - [blit copyFromBuffer:src - sourceOffset:src_off - toBuffer:dst - destinationOffset:dst_off - size:(NSUInteger)n * sizeof(uint16_t)]; - [blit endEncoding]; - return 1; -} - -static int ds4_gpu_encode_copy_raw_ring_to_f16( - id cb, - id raw, - NSUInteger raw_offset, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_raw, - uint32_t head_dim, - id dst, - NSUInteger dst_offset) { - if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap || - n_raw == 0 || n_raw > raw_cap || head_dim == 0) { - return 0; - } - - const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t); - const uint32_t tail_rows = raw_cap - raw_start < n_raw - ? raw_cap - raw_start - : n_raw; - const uint32_t head_rows = n_raw - tail_rows; - const uint64_t tail_count = (uint64_t)tail_rows * head_dim; - const uint64_t head_count = (uint64_t)head_rows * head_dim; - const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes; - const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes; - if (tail_count > UINT32_MAX || head_count > UINT32_MAX || - raw_inner > NSUIntegerMax - raw_offset || - dst_inner > NSUIntegerMax - dst_offset) { - return 0; - } - - if (tail_rows && - !ds4_gpu_encode_cpy_f32_f16_1d( - cb, - raw, - raw_offset + (NSUInteger)raw_inner, - dst, - dst_offset, - (uint32_t)tail_count)) { - return 0; - } - if (head_rows && - !ds4_gpu_encode_cpy_f32_f16_1d( - cb, - raw, - raw_offset, - dst, - dst_offset + (NSUInteger)dst_inner, - (uint32_t)head_count)) { - return 0; - } - return 1; -} - -static int ds4_gpu_encode_flash_kv_stage_f16( - id cb, - id raw, - NSUInteger raw_offset, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_raw, - id comp, - NSUInteger comp_offset, - bool comp_is_f16, - uint32_t n_comp, - uint32_t head_dim, - id dst, - NSUInteger dst_offset, - id mask, - NSUInteger mask_offset, - id pad, - NSUInteger pad_offset, - bool fuse_pad, - bool shared_pad, - bool *did_fuse_pad) { - if (did_fuse_pad) *did_fuse_pad = false; - if (!cb || !raw || !comp || !dst || raw_cap == 0 || - raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap || - n_comp == 0 || head_dim == 0) { - return 0; - } - - const bool force = - getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL; - const bool disabled = - getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL; - const bool require = - getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL; - const bool supported_shape = - comp_is_f16 && head_dim == 512u && raw_cap <= UINT32_MAX / 128u; - const uint64_t row_vecs64 = 128u; - const uint64_t comp_count64 = (uint64_t)n_comp * head_dim; - const uint64_t dst_comp_inner64 = - (uint64_t)n_raw * head_dim * sizeof(uint16_t); - if (comp_count64 > UINT32_MAX || - dst_comp_inner64 > NSUIntegerMax - dst_offset) { - return 0; - } - const uint64_t total_vecs64 = - ((uint64_t)n_raw + n_comp) * row_vecs64; - const bool valid_grid = - total_vecs64 != 0 && total_vecs64 <= UINT32_MAX; - const bool eligible = - supported_shape && valid_grid && !g_quality_mode && !disabled && - g_flash_kv_stage_f16_pipeline != nil && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || force); - const bool component_disabled = eligible && - (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 || - ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0); - const bool use_fusion = eligible && !component_disabled; - const bool use_pad_fusion = - use_fusion && fuse_pad && mask != nil && pad != nil && - getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL; - if (require && supported_shape && !use_fusion) { - fprintf(stderr, - "ds4: required Metal gathered KV staging kernel was not selected\n"); - return 0; - } - - if (use_fusion) { - ds4_gpu_flash_kv_stage_f16_args args = { - .raw_cap = raw_cap, - .raw_start = raw_start, - .n_raw = n_raw, - .n_comp = n_comp, - .pad_rows = use_pad_fusion ? 32u : 0u, - .shared_pad = use_pad_fusion && shared_pad ? 1u : 0u, - }; - const NSUInteger total_vecs = (NSUInteger)total_vecs64 + - (use_pad_fusion ? 32u * 128u + 32u : 0u); - NSUInteger nth = - g_flash_kv_stage_f16_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > total_vecs) nth = total_vecs; - if (nth == 0) return 0; - const NSUInteger groups = (total_vecs + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_flash_kv_stage_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:raw offset:raw_offset atIndex:1]; - [enc setBuffer:comp offset:comp_offset atIndex:2]; - [enc setBuffer:dst offset:dst_offset atIndex:3]; - [enc setBuffer:(use_pad_fusion ? mask : dst) - offset:(use_pad_fusion ? mask_offset : dst_offset) - atIndex:4]; - [enc setBuffer:(use_pad_fusion ? pad : dst) - offset:(use_pad_fusion ? pad_offset : dst_offset) - atIndex:5]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - if (did_fuse_pad) *did_fuse_pad = use_pad_fusion; - return 1; - } - - if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, - raw, - raw_offset, - raw_cap, - raw_start, - n_raw, - head_dim, - dst, - dst_offset)) { - return 0; - } - return ds4_gpu_encode_copy_to_f16_1d( - cb, - comp, - comp_offset, - comp_is_f16, - dst, - dst_offset + (NSUInteger)dst_comp_inner64, - (uint32_t)comp_count64); -} - -int ds4_gpu_flash_kv_stage_f16_tensor( - ds4_gpu_tensor *dst, - const ds4_gpu_tensor *raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_raw, - const ds4_gpu_tensor *comp, - uint32_t comp_is_f16, - uint32_t n_comp, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap || - n_raw == 0 || n_raw > raw_cap || n_comp == 0 || - comp_is_f16 == 0 || head_dim != 512u) { - return 0; - } - - @autoreleasepool { - const uint64_t raw_bytes = - (uint64_t)raw_cap * head_dim * sizeof(float); - const uint64_t comp_bytes = - (uint64_t)n_comp * head_dim * - (comp_is_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t dst_bytes = - ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t); - id rawbuf = ds4_gpu_tensor_buffer(raw); - id compbuf = ds4_gpu_tensor_buffer(comp); - id dstbuf = ds4_gpu_tensor_buffer(dst); - if (!rawbuf || !compbuf || !dstbuf || - ds4_gpu_tensor_bytes(raw) < raw_bytes || - ds4_gpu_tensor_bytes(comp) < comp_bytes || - ds4_gpu_tensor_bytes(dst) < dst_bytes) { - fprintf(stderr, - "ds4: Metal gathered KV staging received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (!ds4_gpu_encode_flash_kv_stage_f16( - cb, - rawbuf, - ds4_gpu_tensor_offset(raw), - raw_cap, - raw_start, - n_raw, - compbuf, - ds4_gpu_tensor_offset(comp), - comp_is_f16 != 0, - n_comp, - head_dim, - dstbuf, - ds4_gpu_tensor_offset(dst), - nil, - 0, - nil, - 0, - false, - false, - NULL)) { - return 0; - } - if (!ds4_gpu_finish_command_buffer( - cb, owned, "gathered KV staging")) { - return 0; - } - } - return 1; -} - -static int ds4_gpu_encode_fill_f16_1d( - id cb, - id buf, - NSUInteger offset, - uint32_t n, - float value) { - if (!cb || !buf || n == 0) return 0; - - ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(n, 1, 0, 0.0f, 0.0f); - args.val = value; - - NSUInteger nth = (NSUInteger)n; - const NSUInteger max_threads = g_unary_fill_f16_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth > 256u) nth = 256u; - if (nth == 0) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_unary_fill_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:buf offset:offset atIndex:1]; - [enc setBuffer:buf offset:offset atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_flash_attention_raw_heads( - id cb, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_head, - uint32_t head_dim) { - if (head_dim != 512 || n_head == 0 || n_raw == 0 || raw_cap < n_raw) { - return 0; - } - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - const uint64_t heads_bytes = q_bytes; - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(heads) < heads_bytes) { - fprintf(stderr, "ds4: Metal DS4 FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t ncpsg = 32; - const uint32_t nwg = 32; - const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_raw, nwg, ncpsg); - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_raw * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; - const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + - (NSUInteger)ncpsg * sizeof(uint16_t); - const NSUInteger nrows = (NSUInteger)n_head; - const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + - nrows * (2u * (NSUInteger)nwg) * sizeof(float); - - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, - &g_flash_attn_tmp_bytes, - tmp_bytes, - "ds4_flash_attn_tmp")) { - return 0; - } - memset([mask_buffer contents], 0, mask_bytes); - - id pad_pipeline = nil; - if ((n_raw % ncpsg) != 0) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id vec_pipeline = - ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, (n_raw % ncpsg) != 0, - false, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg, - (int32_t)nwg); - id reduce_pipeline = - ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); - if (!vec_pipeline || !reduce_pipeline) return 0; - - if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - raw_cap, - raw_start, - n_raw, - head_dim, - g_flash_attn_kv_buffer, - 0)) { - return 0; - } - - if ((n_raw % ncpsg) != 0) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_raw, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_raw * row_bytes_f16, - .nb13 = (uint64_t)n_raw * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_raw * row_bytes_f16, - .nb23 = (uint64_t)n_raw * row_bytes_f16, - .ne31 = 1, - .ne32 = 1, - .ne33 = 1, - .nb31 = mask_bytes, - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - ds4_gpu_flash_attn_vec_args vec_args = { - .ne01 = 1, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_head * row_bytes, - .ne11 = (int32_t)n_raw, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_raw * row_bytes_f16, - .nb13 = (uint64_t)n_raw * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_raw * row_bytes_f16, - .nb23 = (uint64_t)n_raw * row_bytes_f16, - .ne31 = 1, - .ne32 = 1, - .ne33 = 1, - .nb31 = mask_bytes, - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = 1, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + - 4u * ncpsg + - 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:vec_pipeline]; - [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_flash_attn_reduce_args reduce_args = { - .nrows = (int32_t)nrows, - }; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; - [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -/* Rectangular causal/window mask for raw prefill: q covers rows - * [q_row0, q_row0 + n_q) of an n_kv-token block whose keys all live in - * rows [0, n_kv). The square prefill case is q_row0 == 0, n_q == n_kv. */ -static void ds4_gpu_fill_raw_prefill_mask( - uint16_t *mask, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_kv, - uint32_t window) { - const uint16_t neg_inf_half = 0xfc00u; - for (uint32_t q = 0; q < n_q; q++) { - const uint32_t qpos = q_row0 + q; - uint16_t *row = mask + (uint64_t)q * n_kv; - for (uint32_t k = 0; k < n_kv; k++) { - const bool causal = k <= qpos; - const bool in_window = window == 0 || qpos - k < window; - row[k] = causal && in_window ? 0u : neg_inf_half; - } - } -} - -static void ds4_gpu_fill_glm_prefill_mask( - uint16_t *mask, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len) { - const uint16_t neg_inf_half = 0xfc00u; - for (uint32_t q = 0; q < n_tokens; q++) { - const uint32_t qpos = pos0 + q; - uint16_t *row = mask + (uint64_t)q * cache_len; - for (uint32_t k = 0; k < cache_len; k++) { - row[k] = k <= qpos ? 0u : neg_inf_half; - } - } -} - -static id ds4_gpu_glm_prefill_mask_buffer( - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - NSUInteger mask_bytes) { - const int same_shape = - g_glm_flash_attn_mask_valid && - g_glm_flash_attn_mask_buffer && - g_glm_flash_attn_mask_bytes >= mask_bytes && - g_glm_flash_attn_mask_pos0 == pos0 && - g_glm_flash_attn_mask_tokens == n_tokens && - g_glm_flash_attn_mask_cache_len == cache_len; - if (same_shape) return g_glm_flash_attn_mask_buffer; - - if (g_glm_flash_attn_mask_buffer) { - [g_transient_buffers addObject:g_glm_flash_attn_mask_buffer]; - g_glm_flash_attn_mask_buffer = nil; - } - g_glm_flash_attn_mask_bytes = 0; - g_glm_flash_attn_mask_valid = 0; - if (!ds4_gpu_ensure_scratch_buffer(&g_glm_flash_attn_mask_buffer, - &g_glm_flash_attn_mask_bytes, - mask_bytes, - "ds4_glm_flash_attn_mask")) { - return nil; - } - - ds4_gpu_fill_glm_prefill_mask((uint16_t *)[g_glm_flash_attn_mask_buffer contents], - pos0, - n_tokens, - cache_len); - g_glm_flash_attn_mask_pos0 = pos0; - g_glm_flash_attn_mask_tokens = n_tokens; - g_glm_flash_attn_mask_cache_len = cache_len; - g_glm_flash_attn_mask_valid = 1; - return g_glm_flash_attn_mask_buffer; -} - -static void ds4_gpu_fill_raw_decode_batch_mask( - uint16_t *mask, - uint32_t n_tokens, - uint32_t n_raw, - uint32_t pos0, - uint32_t window) { - const uint16_t neg_inf_half = 0xfc00u; - const uint32_t last_pos = pos0 + n_tokens - 1u; - /* The caller has already copied the SWA ring into logical order when it - * wraps, so key row k represents first_raw_pos + k. */ - const uint32_t first_raw_pos = last_pos + 1u - n_raw; - for (uint32_t q = 0; q < n_tokens; q++) { - const uint32_t qpos = pos0 + q; - uint16_t *row = mask + (uint64_t)q * n_raw; - for (uint32_t k = 0; k < n_raw; k++) { - const uint32_t kpos = first_raw_pos + k; - const bool causal = kpos <= qpos; - const bool in_window = causal && (window == 0 || qpos - kpos < window); - row[k] = causal && in_window ? 0u : neg_inf_half; - } - } -} - -static void ds4_gpu_fill_mixed_decode_batch_mask( - uint16_t *mask, - uint32_t n_tokens, - uint32_t n_raw, - uint32_t n_comp, - uint32_t pos0, - uint32_t window, - uint32_t ratio) { - const uint16_t neg_inf_half = 0xfc00u; - const uint32_t n_keys = n_raw + n_comp; - const uint32_t last_pos = pos0 + n_tokens - 1u; - /* Raw keys are laid out by logical position; compressed keys follow them. */ - const uint32_t first_raw_pos = last_pos + 1u - n_raw; - for (uint32_t q = 0; q < n_tokens; q++) { - const uint32_t qpos = pos0 + q; - uint16_t *row = mask + (uint64_t)q * n_keys; - for (uint32_t k = 0; k < n_raw; k++) { - const uint32_t kpos = first_raw_pos + k; - const bool causal = kpos <= qpos; - const bool in_window = causal && (window == 0 || qpos - kpos < window); - row[k] = causal && in_window ? 0u : neg_inf_half; - } - const uint32_t n_visible = (qpos + 1u) / ratio; - for (uint32_t c = 0; c < n_comp; c++) { - row[n_raw + c] = c < n_visible ? 0u : neg_inf_half; - } - } -} - -/* Rectangular causal/window + compressed-key visibility mask: q covers rows - * [q_row0, q_row0 + n_q) of an n_tokens-token chunk whose raw keys all stay - * resident, followed by n_comp compressed keys. The square prefill case is - * q_row0 == 0, n_q == n_tokens. */ -static void ds4_gpu_fill_static_mixed_prefill_mask( - uint16_t *mask, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio) { - const uint16_t neg_inf_half = 0xfc00u; - const uint32_t n_keys = n_tokens + n_comp; - for (uint32_t q = 0; q < n_q; q++) { - const uint32_t qpos = q_row0 + q; - uint16_t *row = mask + (uint64_t)q * n_keys; - for (uint32_t k = 0; k < n_tokens; k++) { - const bool causal = k <= qpos; - const bool in_window = window == 0 || qpos - k < window; - row[k] = causal && in_window ? 0u : neg_inf_half; - } - - const uint32_t n_visible = (qpos + 1u) / ratio; - for (uint32_t c = 0; c < n_comp; c++) { - row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half; - } - } -} - -/* Static-mixed prefill FlashAttention over a rectangular problem: q holds - * n_q query rows for token positions [q_row0, q_row0 + n_q) of the chunk, - * while the keys stay full (all n_tokens raw rows plus n_comp compressed - * rows). The classic square prefill is q_row0 == 0, n_q == n_tokens. */ -static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long( - id __strong *cbp, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!cbp || !*cbp) return 0; - id cb = *cbp; - if (head_dim != 512 || n_head == 0 || n_q == 0 || n_tokens == 0 || ratio == 0) { - return 0; - } - - const uint32_t n_keys = n_tokens + n_comp; - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; - id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * - (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t comp_mask_bytes = use_comp_mask - ? (uint64_t)n_comp * (q_row0 + n_q) * sizeof(float) : 0u; - if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || - (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal prefill static mixed DS4 non-vector FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t nqptg = 8; - const uint32_t ncpsg = 64; - const uint32_t nsg = head_dim >= 512 ? 8u : 4u; - const bool has_kvpad = (n_keys % ncpsg) != 0; - const bool bc_mask = (n_q % nqptg) != 0; - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_q * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) - : 1u; - const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; - const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; - const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - - const uint32_t mask_cache_kind = ratio == 4u - ? DS4_GPU_PREFILL_MASK_CACHE_RATIO4 - : (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u); - bool mask_cache_created = false; - ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = - use_comp_mask == 0u && mask_cache_kind != 0u && - q_row0 == 0u && n_q == n_tokens - ? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind, - n_tokens, - n_comp, - n_keys, - window, - ratio, - nqptg, - ncpsg, - has_kvpad, - bc_mask, - mask_bytes, - blk_bytes, - &mask_cache_created) - : NULL; - id mask_buffer = mask_cache - ? mask_cache->mask - : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (mask_cache && mask_cache_created) { - ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], - 0u, - n_tokens, - n_tokens, - n_comp, - window, - ratio); - mask_cache->valid = true; - } - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - (!mask_cache && - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk"))) { - return 0; - } - id blk_buffer = mask_cache - ? mask_cache->blk - : g_flash_attn_blk_buffer; - - const bool flash_stage_profile = - getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; - double flash_stage_t0 = 0.0; - if (flash_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - int profile_owned = 0; - cb = ds4_gpu_command_buffer(&profile_owned); - if (!cb || profile_owned) return 0; - *cbp = cb; - flash_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ - if (flash_stage_profile) { \ - if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ - "static_mixed_nonvec", (name), n_q, n_comp, n_keys, \ - n_head, head_dim, window, ratio, &flash_stage_t0)) { \ - return 0; \ - } \ - cb = *cbp; \ - } \ - } while (0) - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_kv_buffer, - 0, - n_tokens * head_dim)) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); - if (n_comp && - !ds4_gpu_encode_copy_to_f16_1d(cb, - compbuf, - ds4_gpu_tensor_offset(comp_kv), - comp_kv_f16 != 0, - g_flash_attn_kv_buffer, - (NSUInteger)n_tokens * row_bytes_f16, - n_comp * head_dim)) { - return 0; - } - if (n_comp) { - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); - } - - if (!mask_cache) { - ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], - q_row0, - n_q, - n_tokens, - n_comp, - window, - ratio); - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); - if (use_comp_mask && n_comp != 0) { - if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, - maskbuf, - ds4_gpu_tensor_offset(comp_mask) + - (NSUInteger)((uint64_t)q_row0 * n_comp * sizeof(float)), - mask_buffer, - (NSUInteger)n_tokens * sizeof(uint16_t), - n_comp, - n_q, - (uint64_t)n_comp * sizeof(float), - (uint64_t)n_keys * sizeof(uint16_t))) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); - } - - id pad_pipeline = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id blk_pipeline = - ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); - id attn_pipeline = - ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", - true, true, false, false, has_kvpad, bc_mask, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg); - if (!blk_pipeline || !attn_pipeline) return 0; - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); - } - - ds4_gpu_flash_attn_blk_args blk_args = { - .ne01 = (int32_t)n_q, - .ne30 = (int32_t)n_keys, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = nil; - if (!mask_cache || !mask_cache->blk_ready) { - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - if (mask_cache) mask_cache->blk_ready = true; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); - - ds4_gpu_flash_attn_vec_args args = { - .ne01 = (int32_t)n_q, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_q * n_head * row_bytes, - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_q, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); - const NSUInteger shared_elems = (NSUInteger)nqptg * - ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:blk_buffer offset:0 atIndex:7]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); - -#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE - return 1; -} - -static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec( - id __strong *cbp, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!cbp || !*cbp) return 0; - id cb = *cbp; - if (head_dim != 512 || n_head == 0 || n_tokens == 0 || ratio == 0) { - return 0; - } - - const uint32_t n_keys = n_tokens + n_comp; - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; - id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * - (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; - if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || - (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal prefill static mixed DS4 FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t ncpsg = 32; - const uint32_t nwg = 32; - const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; - const bool has_kvpad = (n_keys % ncpsg) != 0; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) - : 1u; - const NSUInteger nrows = (NSUInteger)n_tokens * n_head; - const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + - nrows * (2u * (NSUInteger)nwg) * sizeof(float); - - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, - &g_flash_attn_tmp_bytes, - tmp_bytes, - "ds4_flash_attn_tmp")) { - return 0; - } - - const bool flash_stage_profile = - getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; - double flash_stage_t0 = 0.0; - if (flash_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - int profile_owned = 0; - cb = ds4_gpu_command_buffer(&profile_owned); - if (!cb || profile_owned) return 0; - *cbp = cb; - flash_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ - if (flash_stage_profile) { \ - if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ - "static_mixed_vec", (name), n_tokens, n_comp, n_keys, \ - n_head, head_dim, window, ratio, &flash_stage_t0)) { \ - return 0; \ - } \ - cb = *cbp; \ - } \ - } while (0) - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_kv_buffer, - 0, - n_tokens * head_dim)) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); - if (n_comp) { - if (!ds4_gpu_encode_copy_to_f16_1d(cb, - compbuf, - ds4_gpu_tensor_offset(comp_kv), - comp_kv_f16 != 0, - g_flash_attn_kv_buffer, - (NSUInteger)n_tokens * row_bytes_f16, - n_comp * head_dim)) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); - } - - ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], - 0, - n_tokens, - n_tokens, - n_comp, - window, - ratio); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); - if (use_comp_mask && n_comp != 0) { - if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, - maskbuf, - ds4_gpu_tensor_offset(comp_mask), - mask_buffer, - (NSUInteger)n_tokens * sizeof(uint16_t), - n_comp, - n_tokens, - (uint64_t)n_comp * sizeof(float), - (uint64_t)n_keys * sizeof(uint16_t))) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); - } - - id pad_pipeline = nil; - id enc = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id vec_pipeline = - ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, has_kvpad, - false, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg, - (int32_t)nwg); - id reduce_pipeline = - ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); - if (!vec_pipeline || !reduce_pipeline) return 0; - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); - } - - ds4_gpu_flash_attn_vec_args vec_args = { - .ne01 = (int32_t)n_tokens, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_tokens * n_head * row_bytes, - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + - 4u * ncpsg + - 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:vec_pipeline]; - [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); - - ds4_gpu_flash_attn_reduce_args reduce_args = { - .nrows = (int32_t)nrows, - }; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; - [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); - -#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE - return 1; -} - -static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec( - id __strong *cbp, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - /* The vector sibling below only handles the small square case; every - * rectangular (TP row-split) problem goes through the long path. */ - if (n_tokens >= 20 || q_row0 != 0 || n_q != n_tokens) { - return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(cbp, - heads, - sinks_buf, - sinks_offset, - q, - raw_kv, - comp_kv, - comp_kv_f16, - comp_mask, - use_comp_mask, - q_row0, - n_q, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim); - } - return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(cbp, - heads, - sinks_buf, - sinks_offset, - q, - raw_kv, - comp_kv, - comp_kv_f16, - comp_mask, - use_comp_mask, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim); -} - -/* Raw prefill FlashAttention over a rectangular problem: q holds n_q query - * rows that correspond to token positions [q_row0, q_row0 + n_q) of the - * chunk, raw_kv holds all n_kv key rows, and heads receives one output row - * per query row. The classic square prefill is q_row0 == 0, n_q == n_kv. */ -static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( - id __strong *cbp, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_kv, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - if (!cbp || !*cbp) return 0; - id cb = *cbp; - if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { - return 0; - } - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * sizeof(float); - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal prefill raw DS4 non-vector FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t nqptg = 8; - const uint32_t ncpsg = 64; - const uint32_t nsg = head_dim >= 512 ? 8u : 4u; - const bool has_kvpad = (n_kv % ncpsg) != 0; - const bool bc_mask = (n_q % nqptg) != 0; - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_q * (NSUInteger)n_kv * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_kv * row_bytes_f16; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) - : 1u; - const NSUInteger nblk0 = ((NSUInteger)n_kv + ncpsg - 1u) / ncpsg; - const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; - const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - - bool mask_cache_created = false; - ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = - q_row0 == 0u && n_q == n_kv - ? ds4_gpu_get_zero_prefix_prefill_mask_cache( - DS4_GPU_PREFILL_MASK_CACHE_RAW, - n_kv, - 0u, - n_kv, - window, - 0u, - nqptg, - ncpsg, - has_kvpad, - bc_mask, - mask_bytes, - blk_bytes, - &mask_cache_created) - : NULL; - id mask_buffer = mask_cache - ? mask_cache->mask - : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (mask_cache && mask_cache_created) { - ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], - 0u, n_kv, n_kv, window); - mask_cache->valid = true; - } - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - (!mask_cache && - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk"))) { - return 0; - } - id blk_buffer = mask_cache - ? mask_cache->blk - : g_flash_attn_blk_buffer; - - const bool flash_stage_profile = - getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; - double flash_stage_t0 = 0.0; - if (flash_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - int profile_owned = 0; - cb = ds4_gpu_command_buffer(&profile_owned); - if (!cb || profile_owned) return 0; - *cbp = cb; - flash_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ - if (flash_stage_profile) { \ - if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ - "raw_nonvec", (name), n_q, 0, n_kv, \ - n_head, head_dim, window, 0, &flash_stage_t0)) { \ - return 0; \ - } \ - cb = *cbp; \ - } \ - } while (0) - - if (!mask_cache) { - ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], - q_row0, n_q, n_kv, window); - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); - - id pad_pipeline = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id blk_pipeline = - ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); - id attn_pipeline = - ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", - true, true, false, false, has_kvpad, bc_mask, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg); - if (!blk_pipeline || !attn_pipeline) return 0; - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_kv_buffer, - 0, - n_kv * head_dim)) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_kv, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_kv * row_bytes_f16, - .nb13 = (uint64_t)n_kv * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_kv * row_bytes_f16, - .nb23 = (uint64_t)n_kv * row_bytes_f16, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_kv * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); - } - - ds4_gpu_flash_attn_blk_args blk_args = { - .ne01 = (int32_t)n_q, - .ne30 = (int32_t)n_kv, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_kv * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = nil; - if (!mask_cache || !mask_cache->blk_ready) { - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - if (mask_cache) mask_cache->blk_ready = true; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); - - ds4_gpu_flash_attn_vec_args args = { - .ne01 = (int32_t)n_q, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_q * n_head * row_bytes, - .ne11 = (int32_t)n_kv, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_kv * row_bytes_f16, - .nb13 = (uint64_t)n_kv * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_kv * row_bytes_f16, - .nb23 = (uint64_t)n_kv * row_bytes_f16, - .ne31 = (int32_t)n_q, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_kv * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_q, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); - const NSUInteger shared_elems = (NSUInteger)nqptg * - ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:blk_buffer offset:0 atIndex:7]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); - -#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE - return 1; -} - -static int ds4_gpu_encode_flash_attention_prefill_raw_heads( - id __strong *cbp, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_kv, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - if (!cbp || !*cbp) return 0; - id cb = *cbp; - if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { - return 0; - } - /* The vector sibling below only handles the small square case; every - * rectangular (TP row-split) problem goes through the non-vector path. */ - if (n_kv >= 20 || q_row0 != 0 || n_q != n_kv) { - return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp, - heads, - sinks_buf, - sinks_offset, - q, - raw_kv, - q_row0, - n_q, - n_kv, - window, - n_head, - head_dim); - } - const uint32_t n_tokens = n_q; - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal prefill raw DS4 FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t ncpsg = 32; - const uint32_t nwg = 32; - const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_tokens, nwg, ncpsg); - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)n_tokens * sizeof(uint16_t); - const NSUInteger kv_f16_offset = 0; - const NSUInteger kv_f16_bytes = (NSUInteger)n_tokens * row_bytes_f16; - const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + - (NSUInteger)ncpsg * (NSUInteger)n_tokens * sizeof(uint16_t); - const NSUInteger nrows = (NSUInteger)n_tokens * n_head; - const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + - nrows * (2u * (NSUInteger)nwg) * sizeof(float); - - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_f16_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, - &g_flash_attn_tmp_bytes, - tmp_bytes, - "ds4_flash_attn_tmp")) { - return 0; - } - - const bool flash_stage_profile = - getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; - double flash_stage_t0 = 0.0; - if (flash_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - int profile_owned = 0; - cb = ds4_gpu_command_buffer(&profile_owned); - if (!cb || profile_owned) return 0; - *cbp = cb; - flash_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ - if (flash_stage_profile) { \ - if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ - "raw_vec", (name), n_tokens, 0, n_tokens, \ - n_head, head_dim, window, 0, &flash_stage_t0)) { \ - return 0; \ - } \ - cb = *cbp; \ - } \ - } while (0) - - ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], 0, n_tokens, n_tokens, window); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); - - id pad_pipeline = nil; - if ((n_tokens % ncpsg) != 0) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id vec_pipeline = - ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, true, - false, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg, - (int32_t)nwg); - id reduce_pipeline = - ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); - if (!vec_pipeline || !reduce_pipeline) return 0; - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_kv_buffer, - kv_f16_offset, - n_tokens * head_dim)) { - return 0; - } - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); - - if ((n_tokens % ncpsg) != 0) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_tokens, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_tokens * row_bytes_f16, - .nb13 = (uint64_t)n_tokens * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_tokens * row_bytes_f16, - .nb23 = (uint64_t)n_tokens * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); - } - - ds4_gpu_flash_attn_vec_args vec_args = { - .ne01 = (int32_t)n_tokens, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_tokens * n_head * row_bytes, - .ne11 = (int32_t)n_tokens, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_tokens * row_bytes_f16, - .nb13 = (uint64_t)n_tokens * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_tokens * row_bytes_f16, - .nb23 = (uint64_t)n_tokens * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + - 4u * ncpsg + - 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:vec_pipeline]; - [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); - - ds4_gpu_flash_attn_reduce_args reduce_args = { - .nrows = (int32_t)nrows, - }; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; - [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); - -#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE - return 1; -} - -static int ds4_gpu_encode_flash_attention_gathered_heads( - id cb, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_comp, - const ds4_gpu_tensor *comp_mask, - uint32_t use_mask, - uint32_t n_head, - uint32_t head_dim) { - const uint32_t n_keys = n_raw + n_comp; - if (head_dim != 512 || n_head == 0 || n_raw == 0 || n_keys == 0 || - raw_cap < n_raw || n_keys < n_raw) { - return 0; - } - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : nil; - id headsbuf = ds4_gpu_tensor_buffer(heads); - id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : nil; - const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * - (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || - (n_comp && !compbuf) || - (use_mask && !maskbuf) || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || - ds4_gpu_tensor_bytes(heads) < q_bytes || - (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes)) { - fprintf(stderr, "ds4: Metal gathered DS4 FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t ncpsg = 32; - const uint32_t nwg = 32; - const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_keys * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; - const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + - (NSUInteger)ncpsg * sizeof(uint16_t); - const NSUInteger nrows = (NSUInteger)n_head; - const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + - nrows * (2u * (NSUInteger)nwg) * sizeof(float); - - const bool use_persistent_zero_mask = - use_mask == 0u && - (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) && - getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL; - - if (!(use_persistent_zero_mask - ? ds4_gpu_ensure_zero_attention_mask(mask_bytes) - : ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer, - &g_flash_attn_mask_bytes, - mask_bytes, - "ds4_flash_attn_mask")) || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, - &g_flash_attn_tmp_bytes, - tmp_bytes, - "ds4_flash_attn_tmp")) { - return 0; - } - id flash_mask_buffer = use_persistent_zero_mask - ? g_flash_attn_zero_mask_buffer - : g_flash_attn_mask_buffer; - - const bool has_kvpad = (n_keys % ncpsg) != 0; - const bool use_shared_kvpad = - has_kvpad && - (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) && - getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL; - id vec_pipeline = - ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, has_kvpad, - use_shared_kvpad, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg, - (int32_t)nwg); - id reduce_pipeline = - ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); - if (!vec_pipeline || !reduce_pipeline) return 0; - - if (!use_persistent_zero_mask && - !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) { - return 0; - } - if (use_mask && n_comp && - !ds4_gpu_encode_cpy_f32_f16_1d(cb, - maskbuf, - ds4_gpu_tensor_offset(comp_mask), - flash_mask_buffer, - (NSUInteger)n_raw * sizeof(uint16_t), - n_comp)) { - return 0; - } - - bool pad_fused = false; - if (!ds4_gpu_encode_flash_kv_stage_f16( - cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - raw_cap, - raw_start, - n_raw, - compbuf, - ds4_gpu_tensor_offset(comp_kv), - comp_kv_f16 != 0, - n_comp, - head_dim, - g_flash_attn_kv_buffer, - 0, - flash_mask_buffer, - 0, - g_flash_attn_pad_buffer, - 0, - has_kvpad, - use_shared_kvpad, - &pad_fused)) { - return 0; - } - - id pad_pipeline = nil; - if (has_kvpad && !pad_fused) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - - if (has_kvpad && !pad_fused) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = 1, - .ne32 = 1, - .ne33 = 1, - .nb31 = mask_bytes, - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:flash_mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - ds4_gpu_flash_attn_vec_args vec_args = { - .ne01 = 1, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_head * row_bytes, - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = 1, - .ne32 = 1, - .ne33 = 1, - .nb31 = mask_bytes, - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = 1, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + - 4u * ncpsg + - 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:vec_pipeline]; - [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_flash_attn_reduce_args reduce_args = { - .nrows = (int32_t)nrows, - }; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; - [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; - [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads( - id cb, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t window, - uint32_t n_head, - uint32_t head_dim, - bool noncausal) { - if (head_dim != 512 || n_head == 0 || n_tokens == 0 || - n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { - return 0; - } - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal decode raw batch FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t nqptg = 8; - const uint32_t ncpsg = 64; - const uint32_t nsg = head_dim >= 512 ? 8u : 4u; - const bool has_kvpad = (n_raw % ncpsg) != 0; - const bool bc_mask = (n_tokens % nqptg) != 0; - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_raw * (NSUInteger)n_tokens * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) - : 1u; - const NSUInteger nblk0 = ((NSUInteger)n_raw + ncpsg - 1u) / ncpsg; - const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; - const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk")) { - return 0; - } - - if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - raw_cap, - raw_start, - n_raw, - head_dim, - g_flash_attn_kv_buffer, - 0)) { - return 0; - } - - if (noncausal) { - memset([mask_buffer contents], 0, mask_bytes); - } else { - ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents], - n_tokens, - n_raw, - pos0, - window); - } - - id pad_pipeline = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id blk_pipeline = - ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); - id attn_pipeline = - ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", - true, true, false, false, has_kvpad, bc_mask, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg); - if (!blk_pipeline || !attn_pipeline) return 0; - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_raw, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_raw * row_bytes_f16, - .nb13 = (uint64_t)n_raw * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_raw * row_bytes_f16, - .nb23 = (uint64_t)n_raw * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_raw * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - ds4_gpu_flash_attn_blk_args blk_args = { - .ne01 = (int32_t)n_tokens, - .ne30 = (int32_t)n_raw, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_raw * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_flash_attn_vec_args args = { - .ne01 = (int32_t)n_tokens, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_tokens * n_head * row_bytes, - .ne11 = (int32_t)n_raw, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_raw * row_bytes_f16, - .nb13 = (uint64_t)n_raw * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_raw * row_bytes_f16, - .nb23 = (uint64_t)n_raw * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_raw * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); - const NSUInteger shared_elems = (NSUInteger)nqptg * - ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads( - id cb, - ds4_gpu_tensor *heads, - id sinks_buf, - NSUInteger sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (n_comp == 0) { - return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, - heads, - sinks_buf, - sinks_offset, - q, - raw_kv, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - window, - n_head, - head_dim, - false); - } - if (head_dim != 512 || n_head == 0 || n_tokens == 0 || - n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || - ratio == 0 || !comp_kv || (use_comp_mask && !comp_mask)) { - return 0; - } - - const uint32_t n_keys = n_raw + n_comp; - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = ds4_gpu_tensor_buffer(comp_kv); - id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * - (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; - if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || - (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal decode mixed batch FlashAttention received undersized buffers\n"); - return 0; - } - - const uint32_t nqptg = 8; - const uint32_t ncpsg = 64; - const uint32_t nsg = head_dim >= 512 ? 8u : 4u; - const bool has_kvpad = (n_keys % ncpsg) != 0; - const bool bc_mask = (n_tokens % nqptg) != 0; - const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); - const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); - const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) - : 1u; - const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; - const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; - const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); - if (!mask_buffer || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_bytes, - "ds4_flash_attn_kv_f16") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk")) { - return 0; - } - - if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - raw_cap, - raw_start, - n_raw, - head_dim, - g_flash_attn_kv_buffer, - 0) || - !ds4_gpu_encode_copy_to_f16_1d(cb, - compbuf, - ds4_gpu_tensor_offset(comp_kv), - comp_kv_f16 != 0, - g_flash_attn_kv_buffer, - (NSUInteger)n_raw * row_bytes_f16, - n_comp * head_dim)) { - return 0; - } - - ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents], - n_tokens, - n_raw, - n_comp, - pos0, - window, - ratio); - if (use_comp_mask) { - if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, - maskbuf, - ds4_gpu_tensor_offset(comp_mask), - mask_buffer, - (NSUInteger)n_raw * sizeof(uint16_t), - n_comp, - n_tokens, - (uint64_t)n_comp * sizeof(float), - (uint64_t)n_keys * sizeof(uint16_t))) { - return 0; - } - } - - id pad_pipeline = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id blk_pipeline = - ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); - id attn_pipeline = - ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", - true, true, false, false, has_kvpad, bc_mask, - (int32_t)head_dim, - (int32_t)head_dim, - (int32_t)nsg); - if (!blk_pipeline || !attn_pipeline) return 0; - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - ds4_gpu_flash_attn_blk_args blk_args = { - .ne01 = (int32_t)n_tokens, - .ne30 = (int32_t)n_keys, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_flash_attn_vec_args args = { - .ne01 = (int32_t)n_tokens, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * row_bytes, - .nb02 = row_bytes, - .nb03 = (uint64_t)n_tokens * n_head * row_bytes, - .ne11 = (int32_t)n_keys, - .ne_12_2 = 1, - .ne_12_3 = 1, - .ns10 = (int32_t)head_dim, - .nb11 = row_bytes_f16, - .nb12 = (uint64_t)n_keys * row_bytes_f16, - .nb13 = (uint64_t)n_keys * row_bytes_f16, - .ns20 = (int32_t)head_dim, - .nb21 = row_bytes_f16, - .nb22 = (uint64_t)n_keys * row_bytes_f16, - .nb23 = (uint64_t)n_keys * row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)n_keys * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)head_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); - const NSUInteger shared_elems = (NSUInteger)nqptg * - ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -int ds4_gpu_attention_prefill_raw_heads_range_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_kv, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_kv == 0) return 0; - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_prefill_raw_heads(&cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - q_row0, - n_q, - n_kv, - window, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill raw attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_prefill_raw_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - return ds4_gpu_attention_prefill_raw_heads_range_tensor(heads, - model_map, - model_size, - sinks_offset, - q, - raw_kv, - 0, - n_tokens, - n_tokens, - window, - n_head, - head_dim); -} - -int ds4_gpu_attention_decode_raw_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t window, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || - n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { - return 0; - } - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - window, - n_head, - head_dim, - false)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode raw batch attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_tokens, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !model_map || - n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || - raw_start >= raw_cap || n_head == 0 || head_dim == 0) { - return 0; - } - - @autoreleasepool { - const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); - if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal noncausal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - sink_bytes, - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - n_tokens, - 0, - n_raw, - raw_cap, - raw_start, - 0, - n_head, - head_dim, - true)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph noncausal raw batch attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_decode_mixed_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t use_comp_mask, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || - n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || - ratio == 0 || (n_comp != 0 && !comp_kv) || - (use_comp_mask != 0 && !comp_mask)) { - return 0; - } - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - comp_kv, - comp_kv_f16, - comp_mask, - use_comp_mask, - n_tokens, - pos0, - n_raw, - raw_cap, - raw_start, - n_comp, - window, - ratio, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode mixed batch attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *topk, - uint32_t n_tokens, - uint32_t pos0, - 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 n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !model_map || !q || !raw_kv || !comp_kv || !topk || - n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || - n_comp == 0 || top_k == 0 || top_k > n_comp || (top_k & (top_k - 1u)) != 0 || - ratio == 0 || n_head == 0 || head_dim != 512) { - return 0; - } - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal indexed attention sinks range is outside the mapped model\n"); - return 0; - } - - const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); - const uint64_t row_bytes_f16 = (uint64_t)head_dim * sizeof(uint16_t); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * row_bytes; - const uint64_t raw_bytes = (uint64_t)raw_cap * row_bytes; - const uint64_t comp_bytes = (uint64_t)n_comp * (comp_kv_f16 ? row_bytes_f16 : row_bytes); - const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = ds4_gpu_tensor_buffer(comp_kv); - id topkbuf = ds4_gpu_tensor_buffer(topk); - id headsbuf = ds4_gpu_tensor_buffer(heads); - if (!qbuf || !rawbuf || !compbuf || !topkbuf || !headsbuf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || - ds4_gpu_tensor_bytes(topk) < topk_bytes || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal indexed mixed attention received undersized buffers\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - id sort_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, - "kernel_dsv4_sort_i32_rows_asc"); - const bool decode_one_token = n_tokens == 1u; - id attn_pipeline = - decode_one_token ? - ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline, - "kernel_dsv4_indexed_mixed_attention_heads8_rb16") : - ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline, - "kernel_dsv4_indexed_mixed_attention_heads8"); - if (!sort_pipeline || !attn_pipeline) return 0; - if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) { - fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n"); - return 0; - } - /* - * Fast decode attends to the same full top-k compressed rows but keeps - * them in score order, avoiding a chronological sort dispatch. - * --quality restores the sorted order for stricter reproducibility. - */ - const bool skip_decode_sort = !g_quality_mode && decode_one_token; - if (!skip_decode_sort && - !ds4_gpu_ensure_scratch_buffer(&g_indexed_topk_buffer, - &g_indexed_topk_bytes, - (NSUInteger)topk_bytes, - "ds4_indexed_topk_sorted")) { - return 0; - } - - ds4_gpu_dsv4_topk_mask_args sort_args = { - .ne00 = (int64_t)top_k, - .ne01 = (int64_t)n_tokens, - .nb00 = sizeof(int32_t), - .nb01 = (uint64_t)top_k * sizeof(int32_t), - .ne0 = (int64_t)top_k, - .ne1 = (int64_t)n_tokens, - .nb0 = sizeof(int32_t), - .nb1 = (uint64_t)top_k * sizeof(int32_t), - }; - ds4_gpu_dsv4_indexed_attention_args attn_args = { - .n_tokens = n_tokens, - .n_head = n_head, - .n_raw = n_raw, - .raw_cap = raw_cap, - .raw_start = raw_start, - .n_comp = n_comp, - .top_k = top_k, - .pos0 = pos0, - .window = window, - .ratio = ratio, - .comp_kv_f16 = comp_kv_f16 ? 1u : 0u, - .pad0 = 0, - .q_token_stride = (uint64_t)n_head * row_bytes, - .q_head_stride = row_bytes, - .raw_row_stride = row_bytes, - .comp_row_stride = comp_kv_f16 ? row_bytes_f16 : row_bytes, - .topk_token_stride = (uint64_t)top_k * sizeof(int32_t), - .dst_token_stride = (uint64_t)n_head * row_bytes, - .dst_head_stride = row_bytes, - .scale = 1.0f / sqrtf((float)head_dim), - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = nil; - if (!skip_decode_sort) { - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:sort_pipeline]; - [enc setBytes:&sort_args length:sizeof(sort_args) atIndex:0]; - [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; - [enc setBuffer:g_indexed_topk_buffer offset:0 atIndex:2]; - [enc setThreadgroupMemoryLength:(NSUInteger)top_k * sizeof(int32_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(top_k, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; - [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; - [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer - offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 - atIndex:4]; - [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; - [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * - 128u * 4u * sizeof(uint16_t) - atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, ((NSUInteger)n_head + 7u) / 8u, 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t q_row0, - uint32_t n_q, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_tokens == 0 || - ratio == 0 || (n_comp != 0 && !comp_kv)) { - return 0; - } - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - comp_kv, - comp_kv_f16, - NULL, - 0, - q_row0, - n_q, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill static mixed attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_prefill_static_mixed_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - return ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(heads, - model_map, - model_size, - sinks_offset, - q, - raw_kv, - comp_kv, - comp_kv_f16, - 0, - n_tokens, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim); -} - -int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - const ds4_gpu_tensor *comp_mask, - uint32_t n_tokens, - uint32_t n_comp, - uint32_t window, - uint32_t ratio, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !raw_kv || !comp_kv || !comp_mask || !model_map || - n_tokens == 0 || n_comp == 0 || ratio == 0) { - return 0; - } - - @autoreleasepool { - if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, - sinks_offset, - (uint64_t)n_head * sizeof(float), - &sinks_inner); - if (!sinks_buf) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - comp_kv, - comp_kv_f16, - comp_mask, - 1, - 0, - n_tokens, - n_tokens, - n_comp, - window, - ratio, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill masked mixed attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_attention_decode_heads_tensor( - ds4_gpu_tensor *heads, - const void *model_map, - uint64_t model_size, - uint64_t sinks_offset, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *raw_kv, - uint32_t n_raw, - uint32_t raw_cap, - uint32_t raw_start, - const ds4_gpu_tensor *comp_kv, - uint32_t comp_kv_f16, - uint32_t n_comp, - const ds4_gpu_tensor *comp_mask, - uint32_t use_mask, - uint32_t n_head, - uint32_t head_dim) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !model_map || !q || !raw_kv || - n_raw == 0 || n_head == 0 || head_dim == 0 || - raw_cap < n_raw || raw_start >= raw_cap || - n_raw > UINT32_MAX - n_comp || n_raw + n_comp > 8192u || - (n_comp != 0 && !comp_kv) || - (use_mask != 0 && !comp_mask)) { - return 0; - } - - @autoreleasepool { - const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); - const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); - const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * - (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); - if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { - fprintf(stderr, "ds4: Metal graph attention heads sink range is outside the mapped model\n"); - return 0; - } - - id qbuf = ds4_gpu_tensor_buffer(q); - id rawbuf = ds4_gpu_tensor_buffer(raw_kv); - id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; - id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; - id headsbuf = ds4_gpu_tensor_buffer(heads); - const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; - if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || - (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || - (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || - ds4_gpu_tensor_bytes(heads) < q_bytes) { - fprintf(stderr, "ds4: Metal graph attention heads received undersized buffers\n"); - return 0; - } - - uint64_t sinks_inner = 0; - id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner); - if (!sinks_buf) return 0; - - if (n_comp == 0) { - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_raw_heads(cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - n_raw, - raw_cap, - raw_start, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph raw attention heads")) return 0; - return 1; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!ds4_gpu_encode_flash_attention_gathered_heads(cb, - heads, - sinks_buf, - (NSUInteger)sinks_inner, - q, - raw_kv, - n_raw, - raw_cap, - raw_start, - comp_kv, - comp_kv_f16, - n_comp, - comp_mask, - use_mask, - n_head, - head_dim)) { - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "graph attention heads")) return 0; - } - - return 1; -} - -int ds4_gpu_swiglu_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *gate, - const ds4_gpu_tensor *up, - uint32_t n, - float clamp, - float weight) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !gate || !up || n == 0) return 0; - if (!isfinite(clamp) || clamp < 0.0f || !isfinite(weight)) return 0; - - @autoreleasepool { - id gatebuf = ds4_gpu_tensor_buffer(gate); - id upbuf = ds4_gpu_tensor_buffer(up); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t bytes = (uint64_t)n * sizeof(float); - if (!gatebuf || !upbuf || !outbuf || - ds4_gpu_tensor_bytes(gate) < bytes || - ds4_gpu_tensor_bytes(up) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal SwiGLU received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glu_args args = { - .ne00 = (int32_t)n, - .nb01 = (uint64_t)n * sizeof(float), - .ne10 = (int32_t)n, - .nb11 = (uint64_t)n * sizeof(float), - .ne0 = (int32_t)n, - .nb1 = (uint64_t)n * sizeof(float), - .i00 = 0, - .i10 = 0, - .alpha = weight, - .limit = clamp, - }; - NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - if (nth == 0u) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_swiglu_flat_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:1]; - [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "SwiGLU")) return 0; - } - - return 1; -} - -int ds4_gpu_add_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *a, - const ds4_gpu_tensor *b, - uint32_t n) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !a || !b || n == 0) return 0; - - @autoreleasepool { - id abuf = ds4_gpu_tensor_buffer(a); - id bbuf = ds4_gpu_tensor_buffer(b); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t bytes = (uint64_t)n * sizeof(float); - if (!abuf || !bbuf || !outbuf || - ds4_gpu_tensor_bytes(a) < bytes || - ds4_gpu_tensor_bytes(b) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal tensor add received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_add_flat_args args = { .n = n }; - NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - if (nth == 0u) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_add2_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; - [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add")) return 0; - } - - return 1; -} - -int ds4_gpu_add3_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *a, - const ds4_gpu_tensor *b, - const ds4_gpu_tensor *c, - uint32_t n) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !a || !b || !c || n == 0) return 0; - - @autoreleasepool { - id abuf = ds4_gpu_tensor_buffer(a); - id bbuf = ds4_gpu_tensor_buffer(b); - id cbuf = ds4_gpu_tensor_buffer(c); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t bytes = (uint64_t)n * sizeof(float); - if (!abuf || !bbuf || !cbuf || !outbuf || - ds4_gpu_tensor_bytes(a) < bytes || - ds4_gpu_tensor_bytes(b) < bytes || - ds4_gpu_tensor_bytes(c) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal tensor add3 received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_add_flat_args args = { .n = n }; - NSUInteger nth = g_add3_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - if (nth == 0u) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_add3_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; - [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; - [enc setBuffer:cbuf offset:ds4_gpu_tensor_offset(c) atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add3")) return 0; - } - - return 1; -} - -typedef struct { - uint32_t width; - uint32_t rows; - uint32_t layer; - uint32_t n_threads; - float scale; -} ds4_gpu_directional_steering_project_args; - -int ds4_gpu_directional_steering_project_tensor( - ds4_gpu_tensor *x, - const ds4_gpu_tensor *directions, - uint32_t layer, - uint32_t width, - uint32_t rows, - float scale) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0; - - @autoreleasepool { - id pipeline = - ds4_gpu_get_pipeline("kernel_dsv4_directional_steering_project_f32"); - if (!pipeline) return 0; - - id xbuf = ds4_gpu_tensor_buffer(x); - id dbuf = ds4_gpu_tensor_buffer(directions); - const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float); - const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float); - if (!xbuf || !dbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(directions) < dir_bytes) { - fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n"); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - while (nth > width && nth > 1u) nth >>= 1; - if (nth == 0) nth = 1; - - ds4_gpu_directional_steering_project_args args = { - .width = width, - .rows = rows, - .layer = layer, - .n_threads = (uint32_t)nth, - .scale = scale, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc setBuffer:dbuf offset:ds4_gpu_tensor_offset(directions) atIndex:2]; - [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "directional steering")) return 0; - } - - return 1; -} - -static NSUInteger ds4_gpu_bin_threads(uint32_t width, id pipeline) { - NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; - if (nth_max > 256u) nth_max = 256u; - NSUInteger nth = 1u; - while (2u * nth < (NSUInteger)width && nth < nth_max) nth *= 2u; - return nth ? nth : 1u; -} - -static int ds4_gpu_encode_unary_f32_rows( - id cb, - id pipeline, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t width, - uint32_t rows, - int c4, - float min, - float max) { - if (!cb || !pipeline || !src || !dst || width == 0 || rows == 0) return 0; - if (c4 && (width & 3u) != 0) return 0; - - ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, c4, 0.0f, 0.0f); - args.min = min; - args.max = max; - - NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; - if (nth_max > 256u) nth_max = 256u; - NSUInteger nth = (NSUInteger)args.ne00; - if (nth > nth_max) nth = nth_max; - if (nth == 0) nth = 1u; - const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, - (NSUInteger)args.ne02, - (NSUInteger)args.ne03) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_bin_f32_rows( - id cb, - id pipeline, - const ds4_gpu_bin_args *args, - id a, - NSUInteger a_off, - id b, - NSUInteger b_off, - id out, - NSUInteger out_off) { - if (!cb || !pipeline || !args || !a || !b || !out || args->ne0 <= 0 || args->ne1 <= 0) { - return 0; - } - - const NSUInteger nth = ds4_gpu_bin_threads((uint32_t)args->ne0, pipeline); - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:a offset:a_off atIndex:1]; - [enc setBuffer:b offset:b_off atIndex:2]; - [enc setBuffer:out offset:out_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)args->ne1, - (NSUInteger)args->ne2, - (NSUInteger)args->ne3) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static ds4_gpu_bin_args ds4_gpu_make_bin_rowwise_scalar_args(uint32_t width, uint32_t rows) { - const uint64_t lhs_row_bytes = (uint64_t)width * sizeof(float); - const uint64_t rhs_row_bytes = sizeof(float); - return (ds4_gpu_bin_args) { - .ne00 = (int32_t)width, - .ne01 = (int32_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = lhs_row_bytes, - .nb02 = (uint64_t)rows * lhs_row_bytes, - .nb03 = (uint64_t)rows * lhs_row_bytes, - .ne10 = 1, - .ne11 = (int32_t)rows, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = rhs_row_bytes, - .nb12 = (uint64_t)rows * rhs_row_bytes, - .nb13 = (uint64_t)rows * rhs_row_bytes, - .ne0 = (int32_t)width, - .ne1 = (int32_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = lhs_row_bytes, - .nb2 = (uint64_t)rows * lhs_row_bytes, - .nb3 = (uint64_t)rows * lhs_row_bytes, - .offs = 0, - .o1 = { 0 }, - }; -} - -static ds4_gpu_mul_mv_id_args ds4_gpu_make_mul_mv_id_args( - uint32_t src0_cols, - uint32_t src0_rows, - uint32_t src0_experts, - uint64_t src0_row_bytes, - uint64_t src0_expert_bytes, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens, - uint32_t nr0) { - const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); - const uint64_t src0_blocks = src0_cols / 256u; - const uint64_t src0_block_bytes = src0_blocks ? src0_row_bytes / src0_blocks : 1u; - return (ds4_gpu_mul_mv_id_args) { - .nei0 = (int32_t)selected_experts, - .nei1 = (int32_t)n_tokens, - .nbi1 = (uint64_t)selected_experts * sizeof(int32_t), - .ne00 = (int32_t)src0_cols, - .ne01 = (int32_t)src0_rows, - .ne02 = (int32_t)src0_experts, - .nb00 = src0_block_bytes, - .nb01 = src0_row_bytes, - .nb02 = src0_expert_bytes, - .ne10 = (int32_t)src0_cols, - .ne11 = (int32_t)src1_expert_rows, - .ne12 = (int32_t)n_tokens, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = src1_row_bytes, - .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, - .ne0 = (int32_t)src0_rows, - .ne1 = (int32_t)selected_experts, - .nb1 = (uint64_t)src0_rows * sizeof(float), - .nr0 = (int32_t)nr0, - }; -} - -static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( - uint32_t src0_cols, - uint32_t src0_experts, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens) { - const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); - return (ds4_gpu_mul_mm_id_map_args) { - .ne02 = (int32_t)src0_experts, - .ne10 = (int32_t)src0_cols, - .ne11 = (int32_t)src1_expert_rows, - .nb11 = src1_row_bytes, - .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, - .ne21 = (int32_t)n_tokens, - .ne20 = (int32_t)selected_experts, - .nb21 = (uint64_t)selected_experts * sizeof(int32_t), - }; -} - -static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( - uint32_t src0_cols, - uint32_t src0_rows, - uint32_t src0_experts, - uint64_t src0_row_bytes, - uint64_t src0_expert_bytes, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens) { - return ds4_gpu_make_mul_mm_id_args_src1_size(src0_cols, - src0_rows, - src0_experts, - src0_row_bytes, - src0_expert_bytes, - src1_expert_rows, - selected_experts, - n_tokens, - sizeof(float)); -} - -static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( - uint32_t src0_cols, - uint32_t src0_rows, - uint32_t src0_experts, - uint64_t src0_row_bytes, - uint64_t src0_expert_bytes, - uint32_t src1_expert_rows, - uint32_t selected_experts, - uint32_t n_tokens, - uint32_t src1_elem_size) { - const uint64_t src1_row_bytes = (uint64_t)src0_cols * src1_elem_size; - return (ds4_gpu_mul_mm_id_args) { - .ne00 = (int32_t)src0_cols, - .ne02 = (int32_t)src0_experts, - .nb01 = src0_row_bytes, - .nb02 = src0_expert_bytes, - .nb03 = (uint64_t)src0_experts * src0_expert_bytes, - .ne11 = (int32_t)src1_expert_rows, - .nb10 = src1_elem_size, - .nb11 = src1_row_bytes, - .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, - .nb13 = (uint64_t)n_tokens * (uint64_t)src1_expert_rows * src1_row_bytes, - .ne20 = (int32_t)selected_experts, - .ne21 = (int32_t)n_tokens, - .ne0 = (int32_t)src0_rows, - .ne1 = (int32_t)selected_experts, - .r2 = 1, - .r3 = 1, - }; -} - -static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q8_0: return 2; - case DS4_METAL_TENSOR_Q8_K: return 2; - case DS4_METAL_TENSOR_Q4_K: return 2; - case DS4_METAL_TENSOR_Q2_K: - case DS4_METAL_TENSOR_IQ2_XXS: return 4; - default: return 0; - } -} - -static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_IQ2_XXS: return "iq2_xxs"; - case DS4_METAL_TENSOR_Q2_K: return "q2_k"; - case DS4_METAL_TENSOR_Q4_K: return "q4_k"; - case DS4_METAL_TENSOR_Q5_K: return "q5_k"; - case DS4_METAL_TENSOR_Q6_K: return "q6_k"; - default: return "unknown"; - } -} - -static const char *ds4_gpu_trim_env_value(const char *env, size_t *len_out) { - if (len_out) *len_out = 0; - if (!env) return NULL; - - while (isspace((unsigned char)*env)) env++; - size_t n = strlen(env); - while (n > 0 && isspace((unsigned char)env[n - 1])) n--; - if (len_out) *len_out = n; - return env; -} - -static bool ds4_gpu_profile_layer_value_match(const char *env, uint32_t layer_index) { - size_t env_len = 0; - env = ds4_gpu_trim_env_value(env, &env_len); - if (!env || env_len == 0) return true; - if (ds4_gpu_env_value_eq(env, env_len, "all")) return true; - - const char *p = env; - const char *end_env = env + env_len; - while (p < end_env) { - while (p < end_env && (*p == ' ' || *p == '\t' || *p == ',')) p++; - if (p >= end_env) break; - - char *end = NULL; - const unsigned long first = strtoul(p, &end, 10); - if (end == p || end > end_env || first > UINT32_MAX) return false; - - unsigned long last = first; - p = end; - if (p < end_env && *p == '-') { - p++; - last = strtoul(p, &end, 10); - if (end == p || end > end_env || last > UINT32_MAX) return false; - p = end; - } - - if (first <= layer_index && layer_index <= last) return true; - while (p < end_env && (*p == ' ' || *p == '\t')) p++; - if (p < end_env && *p != ',') return false; - } - return false; -} - -static bool ds4_gpu_stage_profile_enabled_for_layer(const char *flag_env_name, - const char *layer_env_name, - uint32_t layer_index) { - size_t flag_len = 0; - const char *flag = ds4_gpu_trim_env_value(getenv(flag_env_name), &flag_len); - if (!flag) return false; - - size_t layer_len = 0; - const char *layer = ds4_gpu_trim_env_value(getenv(layer_env_name), &layer_len); - const bool has_layer_filter = layer && layer_len != 0; - - if (flag_len != 0) { - if (ds4_gpu_env_value_eq(flag, flag_len, "0") || - ds4_gpu_env_value_eq(flag, flag_len, "false") || - ds4_gpu_env_value_eq(flag, flag_len, "no") || - ds4_gpu_env_value_eq(flag, flag_len, "off")) { - return false; - } - if (!has_layer_filter && - !ds4_gpu_env_value_eq(flag, flag_len, "1") && - !ds4_gpu_env_value_eq(flag, flag_len, "true") && - !ds4_gpu_env_value_eq(flag, flag_len, "yes") && - !ds4_gpu_env_value_eq(flag, flag_len, "on") && - !ds4_gpu_env_value_eq(flag, flag_len, "all")) { - return ds4_gpu_profile_layer_value_match(flag, layer_index); - } - } - - return ds4_gpu_profile_layer_value_match(layer, layer_index); -} - -static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) { - if (type == DS4_METAL_TENSOR_Q8_0) { - return 32u * 2u * sizeof(float); - } - if (type == DS4_METAL_TENSOR_IQ2_XXS) { - return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t); - } - return 0; -} - -static NSUInteger ds4_gpu_routed_mv_nsg(uint32_t type) { - return type == DS4_METAL_TENSOR_Q8_0 ? 4u : 2u; -} - -static bool ds4_gpu_routed_mv_rows_per_group_is_nr0(uint32_t type) { - return type == DS4_METAL_TENSOR_Q8_0; -} - -static id ds4_gpu_routed_mv_pipeline(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q8_0: - return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); - case DS4_METAL_TENSOR_Q8_K: - return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_K_f32", 2); - case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline; - case DS4_METAL_TENSOR_Q2_K: return g_moe_mul_mv_id_q2_k_pipeline; - case DS4_METAL_TENSOR_Q4_K: return g_moe_mul_mv_id_q4_k_pipeline; - default: return nil; - } -} - -static id ds4_gpu_routed_mm_pipeline(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q8_0: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); - case DS4_METAL_TENSOR_Q8_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f32", false); - case DS4_METAL_TENSOR_IQ2_XXS: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32", false); - case DS4_METAL_TENSOR_Q2_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f32", false); - case DS4_METAL_TENSOR_Q4_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); - case DS4_METAL_TENSOR_Q5_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false); - case DS4_METAL_TENSOR_Q6_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false); - default: - return nil; - } -} - -static id ds4_gpu_routed_mm_addr_pipeline(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q2_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false); - case DS4_METAL_TENSOR_Q4_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false); - default: - return nil; - } -} - -static id ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q8_0: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f16", false); - case DS4_METAL_TENSOR_Q8_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f16", false); - case DS4_METAL_TENSOR_IQ2_XXS: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f16", false); - case DS4_METAL_TENSOR_Q2_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f16", false); - case DS4_METAL_TENSOR_Q4_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16", false); - case DS4_METAL_TENSOR_Q5_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false); - case DS4_METAL_TENSOR_Q6_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false); - default: - return nil; - } -} - -static id ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint32_t type) { - switch (type) { - case DS4_METAL_TENSOR_Q2_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false); - case DS4_METAL_TENSOR_Q4_K: - return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false); - default: - return nil; - } -} - -static int ds4_gpu_encode_mul_mv_id( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:src0 offset:src0_off atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_attn_out_low_q8_direct( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !src0 || !src1 || !dst || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { - return 0; - } - - /* Two row conventions in the classic matvec family: Q8_0 k-splits one - * nr0-row group across all nsg simdgroups (cross-simdgroup reduce), so a - * threadgroup covers nr0 rows; Q4_K gives each simdgroup its own nr0 - * rows, covering nr0*nsg. Dispatching Q8 with the Q4 stride leaves - * (nsg-1)/nsg of the output rows unwritten. */ - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? - (NSUInteger)args->nr0 : (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:src0 offset:src0_off atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_id_pair( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0_a, - NSUInteger src0_a_off, - id src0_b, - NSUInteger src0_b_off, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:src0_a offset:src0_a_off atIndex:1]; - [enc setBuffer:src0_b offset:src0_b_off atIndex:2]; - [enc setBuffer:src1 offset:src1_off atIndex:3]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:4]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:5]; - [enc setBuffer:ids offset:ids_off atIndex:6]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_id_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - id src0_a, - NSUInteger src0_a_off, - id src0_b, - NSUInteger src0_b_off, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || - !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBuffer:src0_a offset:src0_a_off atIndex:2]; - [enc setBuffer:src0_b offset:src0_b_off atIndex:3]; - [enc setBuffer:src1 offset:src1_off atIndex:4]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; - [enc setBuffer:ids offset:ids_off atIndex:8]; - [enc setBuffer:weights offset:weights_off atIndex:9]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_table_q4_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - DS4MetalQ4ExpertTable *gate_table, - DS4MetalQ4ExpertTable *up_table, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0, - bool queue_residency) { - if (!cb || !pipeline || !args || !act || !gate_table || !up_table || - !gate_table.argumentBuffer || !up_table.argumentBuffer || - !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBuffer:gate_table.argumentBuffer offset:0 atIndex:2]; - [enc setBuffer:up_table.argumentBuffer offset:0 atIndex:3]; - [enc setBuffer:src1 offset:src1_off atIndex:4]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; - [enc setBuffer:ids offset:ids_off atIndex:8]; - [enc setBuffer:weights offset:weights_off atIndex:9]; - if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || - !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { - ds4_gpu_end_compute_encoder(cb, enc); - return 0; - } - ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, queue_residency); - ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, queue_residency); - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - DS4MetalQ4ExpertTable *gate_table, - DS4MetalQ4ExpertTable *up_table, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !gate_table || !up_table || - !gate_table.addressBuffer || !up_table.addressBuffer || - !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBuffer:gate_table.addressBuffer offset:0 atIndex:2]; - [enc setBuffer:up_table.addressBuffer offset:0 atIndex:3]; - [enc setBuffer:src1 offset:src1_off atIndex:4]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; - [enc setBuffer:ids offset:ids_off atIndex:8]; - [enc setBuffer:weights offset:weights_off atIndex:9]; - if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || - !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { - ds4_gpu_end_compute_encoder(cb, enc); - return 0; - } - if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { - ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, false); - ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, false); - } - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_table_q4_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - DS4MetalQ4ExpertTable *table, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool queue_residency) { - if (!cb || !pipeline || !args || !table || !table.argumentBuffer || - !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:table.argumentBuffer offset:0 atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { - ds4_gpu_end_compute_encoder(cb, enc); - return 0; - } - ds4_gpu_use_q4_expert_table_resources(cb, enc, table, queue_residency); - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_q4_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - DS4MetalQ4ExpertTable *table, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !table || !table.addressBuffer || - !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:table.addressBuffer offset:0 atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { - ds4_gpu_end_compute_encoder(cb, enc); - return 0; - } - if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { - ds4_gpu_use_q4_expert_table_resources(cb, enc, table, false); - } - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static uint32_t ds4_gpu_q4_expert_group_size(uint32_t n_total_expert) { - uint32_t group_size = 32; - const char *env = getenv("DS4_METAL_Q4_EXPERT_GROUP_SIZE"); - if (env && env[0]) { - char *end = NULL; - unsigned long v = strtoul(env, &end, 10); - if (end != env && *end == '\0' && v > 0 && v <= UINT32_MAX) { - group_size = (uint32_t)v; - } - } - if (group_size == 0) group_size = 1; - if (group_size > n_total_expert) group_size = n_total_expert; - return group_size; -} - -static int ds4_gpu_encode_mul_mv_group_q4_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - const ds4_gpu_moe_expert_group_args *group, - id src0_a, - NSUInteger src0_a_off, - id src0_b, - NSUInteger src0_b_off, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !group || - !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0 || - group->expert_count == 0) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBytes:group length:sizeof(*group) atIndex:2]; - [enc setBuffer:src0_a offset:src0_a_off atIndex:3]; - [enc setBuffer:src0_b offset:src0_b_off atIndex:4]; - [enc setBuffer:src1 offset:src1_off atIndex:5]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; - [enc setBuffer:ids offset:ids_off atIndex:9]; - [enc setBuffer:weights offset:weights_off atIndex:10]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group_q4_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_moe_expert_group_args *group, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !group || !src0 || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - group->expert_count == 0) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:group length:sizeof(*group) atIndex:1]; - [enc setBuffer:src0 offset:src0_off atIndex:2]; - [enc setBuffer:src1 offset:src1_off atIndex:3]; - [enc setBuffer:dst offset:dst_off atIndex:4]; - [enc setBuffer:ids offset:ids_off atIndex:5]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_id_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - id add_in, - NSUInteger add_in_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || - args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:src0 offset:src0_off atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - [enc setBuffer:(add_in ? add_in : dst) offset:(add_in ? add_in_off : dst_off) atIndex:5]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_q4_gather_slots6( - id cb, - id pipeline, - const ds4_gpu_q4_gather_slots6_args *args, - __unsafe_unretained id src_groups[6], - const NSUInteger src_group_offsets[6], - id ids, - NSUInteger ids_off, - id dst, - NSUInteger dst_off) { - if (!cb || !pipeline || !args || !src_groups || !src_group_offsets || !ids || !dst || - args->expert_bytes == 0 || (args->expert_bytes & 15u) != 0 || - args->group_size == 0 || args->n_slots == 0 || args->n_slots > 6) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if (!src_groups[i]) return 0; - } - - const uint64_t chunks = args->expert_bytes >> 4; - if (chunks == 0 || chunks > NSUIntegerMax) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src_groups[i] offset:src_group_offsets[i] atIndex:1 + i]; - } - [enc setBuffer:ids offset:ids_off atIndex:7]; - [enc setBuffer:dst offset:dst_off atIndex:8]; - const NSUInteger threads = 256u; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)chunks + threads - 1u) / threads, - (NSUInteger)args->n_slots, - 1) - threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_slots6_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - __unsafe_unretained id src0_a[6], - const NSUInteger src0_a_off[6], - __unsafe_unretained id src0_b[6], - const NSUInteger src0_b_off[6], - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || - !src1 || !dst_a || !dst_b || !dst_mid || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if (!src0_a[i] || !src0_b[i]) return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; - } - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:14]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; - [enc setBuffer:weights offset:weights_off atIndex:18]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_slots6_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - __unsafe_unretained id src0[6], - const NSUInteger src0_off[6], - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if (!src0[i]) return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:7]; - [enc setBuffer:dst offset:dst_off atIndex:8]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group6_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - __unsafe_unretained id src0_a[6], - const NSUInteger src0_a_off[6], - __unsafe_unretained id src0_b[6], - const NSUInteger src0_b_off[6], - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || - !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if (!src0_a[i] || !src0_b[i]) return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; - } - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:14]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; - [enc setBuffer:ids offset:ids_off atIndex:18]; - [enc setBuffer:weights offset:weights_off atIndex:19]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group6_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - __unsafe_unretained id src0[6], - const NSUInteger src0_off[6], - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if (!src0[i]) return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 6; i++) { - [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:7]; - [enc setBuffer:dst offset:dst_off atIndex:8]; - [enc setBuffer:ids offset:ids_off atIndex:9]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - ds4_gpu_stream_expert_cache_entry * const *entries, - uint32_t n_entries, - id gate_addrs, - id up_addrs, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0, - id overflow_gate, - id overflow_up) { - if (!cb || !pipeline || !args || !act || !entries || - (n_entries == 0 && !overflow_gate) || - !gate_addrs || !up_addrs || - !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || - args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || - args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - for (uint32_t i = 0; i < n_entries; i++) { - if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; - } - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, - n_entries, - 0)) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBuffer:gate_addrs offset:0 atIndex:2]; - [enc setBuffer:up_addrs offset:0 atIndex:3]; - [enc setBuffer:src1 offset:src1_off atIndex:4]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; - [enc setBuffer:ids offset:ids_off atIndex:8]; - [enc setBuffer:weights offset:weights_off atIndex:9]; - for (uint32_t i = 0; i < n_entries; i++) { - [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; - [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; - } - /* Overflow experts are addressed straight into the mapped model views - * when a layer's unique selected set exceeds the cache budget. */ - if (overflow_gate) [enc useResource:overflow_gate usage:MTLResourceUsageRead]; - if (overflow_up) [enc useResource:overflow_up usage:MTLResourceUsageRead]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_iq2( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - ds4_gpu_stream_expert_cache_entry * const *entries, - uint32_t n_entries, - id addrs, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !entries || n_entries == 0 || - n_entries > DS4_METAL_MAX_ROUTED_EXPERT_USED || - !addrs || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || - args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || - args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - for (uint32_t i = 0; i < n_entries; i++) { - if (!entries[i] || !entries[i]->down_buffer) return 0; - } - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, - n_entries, - 0)) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:addrs offset:0 atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - for (uint32_t i = 0; i < n_entries; i++) { - [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; - } - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_q2_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - ds4_gpu_stream_expert_cache_entry * const *entries, - uint32_t n_entries, - id addrs, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - id overflow_down) { - if (!cb || !pipeline || !args || !entries || - (n_entries == 0 && !overflow_down) || - !addrs || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - for (uint32_t i = 0; i < n_entries; i++) { - if (!entries[i] || !entries[i]->down_buffer) return 0; - } - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, - n_entries, - 0)) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:addrs offset:0 atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setBuffer:ids offset:ids_off atIndex:4]; - for (uint32_t i = 0; i < n_entries; i++) { - [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; - } - if (overflow_down) [enc useResource:overflow_down usage:MTLResourceUsageRead]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - const ds4_gpu_stream_expert_split_args *split, - ds4_gpu_stream_expert_cache_entry * const entries[6], - id gate_addrs, - id up_addrs, - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !split || !entries || - !gate_addrs || !up_addrs || !src1 || !dst_a || !dst_b || !dst_mid || - !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if ((split->active_mask & (1u << i)) == 0) continue; - if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; - } - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, - 6, - split->active_mask)) { - return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - [enc setBytes:split length:sizeof(*split) atIndex:2]; - [enc setBuffer:gate_addrs offset:0 atIndex:3]; - [enc setBuffer:up_addrs offset:0 atIndex:4]; - [enc setBuffer:src1 offset:src1_off atIndex:5]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; - [enc setBuffer:ids offset:ids_off atIndex:9]; - [enc setBuffer:weights offset:weights_off atIndex:10]; - for (uint32_t i = 0; i < 6; i++) { - if ((split->active_mask & (1u << i)) == 0) continue; - [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; - [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; - } - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_addr_q2_sum6_masked( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_stream_expert_split_args *split, - ds4_gpu_stream_expert_cache_entry * const entries[6], - id addrs, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !split || !entries || !addrs || - !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 <= 0 || args->ne02 > 384) { - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - if ((split->active_mask & (1u << i)) == 0) continue; - if (!entries[i] || !entries[i]->down_buffer) return 0; - } - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, - 6, - split->active_mask)) { - return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:split length:sizeof(*split) atIndex:1]; - [enc setBuffer:addrs offset:0 atIndex:2]; - [enc setBuffer:src1 offset:src1_off atIndex:3]; - [enc setBuffer:dst offset:dst_off atIndex:4]; - [enc setBuffer:ids offset:ids_off atIndex:5]; - for (uint32_t i = 0; i < 6; i++) { - if ((split->active_mask & (1u << i)) == 0) continue; - [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; - } - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group8_pair_swiglu( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act, - __unsafe_unretained id src0_a[8], - const NSUInteger src0_a_off[8], - __unsafe_unretained id src0_b[8], - const NSUInteger src0_b_off[8], - id src1, - NSUInteger src1_off, - id dst_a, - NSUInteger dst_a_off, - id dst_b, - NSUInteger dst_b_off, - id dst_mid, - NSUInteger dst_mid_off, - id ids, - NSUInteger ids_off, - id weights, - NSUInteger weights_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || - !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 8; i++) { - if (!src0_a[i] || !src0_b[i]) return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBytes:act length:sizeof(*act) atIndex:1]; - for (uint32_t i = 0; i < 8; i++) { - [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; - } - for (uint32_t i = 0; i < 8; i++) { - [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:10 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:18]; - [enc setBuffer:dst_a offset:dst_a_off atIndex:19]; - [enc setBuffer:dst_b offset:dst_b_off atIndex:20]; - [enc setBuffer:dst_mid offset:dst_mid_off atIndex:21]; - [enc setBuffer:ids offset:ids_off atIndex:22]; - [enc setBuffer:weights offset:weights_off atIndex:23]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group8_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - __unsafe_unretained id src0[8], - const NSUInteger src0_off[8], - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 8; i++) { - if (!src0[i]) return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 8; i++) { - [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:9]; - [enc setBuffer:dst offset:dst_off atIndex:10]; - [enc setBuffer:ids offset:ids_off atIndex:11]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group24_id( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - __unsafe_unretained id src0[24], - const NSUInteger src0_off[24], - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg, - bool rows_per_group_is_nr0) { - if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 24; i++) { - if (!src0[i]) return 0; - } - - const NSUInteger nr0 = (NSUInteger)args->nr0; - const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 24; i++) { - [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:25]; - [enc setBuffer:dst offset:dst_off atIndex:26]; - [enc setBuffer:ids offset:ids_off atIndex:27]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mv_group24_sum6( - id cb, - id pipeline, - const ds4_gpu_mul_mv_id_args *args, - __unsafe_unretained id src0[24], - const NSUInteger src0_off[24], - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off, - NSUInteger threadgroup_bytes, - NSUInteger nsg) { - if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || - args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || - args->ne02 != 384) { - return 0; - } - for (uint32_t i = 0; i < 24; i++) { - if (!src0[i]) return 0; - } - - const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; - const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - for (uint32_t i = 0; i < 24; i++) { - [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; - } - [enc setBuffer:src1 offset:src1_off atIndex:25]; - [enc setBuffer:dst offset:dst_off atIndex:26]; - [enc setBuffer:ids offset:ids_off atIndex:27]; - if (threadgroup_bytes != 0) { - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mm_id( - id cb, - id map_pipeline, - id mm_pipeline, - const ds4_gpu_mul_mm_id_map_args *map_args, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - id ids, - NSUInteger ids_off) { - if (!cb || !map_pipeline || !mm_pipeline || !map_args || !mm_args || - !src0 || !src1 || !dst || !ids || - mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || - mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { - return 0; - } - - return ds4_gpu_encode_mul_mm_id_map(cb, - map_pipeline, - map_args, - mm_args, - ids, - ids_off) && - ds4_gpu_encode_mul_mm_id_mapped(cb, - mm_pipeline, - mm_args, - src0, - src0_off, - src1, - src1_off, - dst, - dst_off); -} - -static int ds4_gpu_encode_mul_mm_id_map( - id cb, - id map_pipeline, - const ds4_gpu_mul_mm_id_map_args *map_args, - const ds4_gpu_mul_mm_id_args *mm_args, - id ids, - NSUInteger ids_off) { - if (!cb || !map_pipeline || !map_args || !mm_args || !ids || - mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { - return 0; - } - - const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); - const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0; - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer, - &g_moe_id_map_bytes, - tpe_bytes + hids_bytes, - "ds4_moe_id_map")) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:map_pipeline]; - [enc setBytes:map_args length:sizeof(*map_args) atIndex:0]; - [enc setBuffer:ids offset:ids_off atIndex:1]; - [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2]; - [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3]; - [enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mm_id_mapped_tile( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes) { - if (!cb || !mm_pipeline || !mm_args || !src0 || !src1 || !dst || - !g_moe_id_map_buffer || - mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || - mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { - return 0; - } - /* - * The routed MoE grouped matmul uses the legacy 32-token expert-major tile. - * The removed TensorOps variant was not semantically stable on evals, so keep - * this encoder tied to the tested simdgroup kernel shape. - */ - const NSUInteger tile_n = 32u; - const bool use_resource_hints = - getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL && - getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL; - - const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); - const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes || - g_moe_id_map_bytes < tpe_bytes + hids_bytes) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:mm_pipeline]; - [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; - [enc setBuffer:src0 offset:src0_off atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; - [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; - [enc setBuffer:dst offset:dst_off atIndex:5]; - if (use_resource_hints) { - [enc useResource:src0 usage:MTLResourceUsageRead]; - } - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, - ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0_addrs, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off, - NSUInteger threadgroup_bytes, - ds4_gpu_stream_expert_cache_entry * const *resources, - uint32_t resource_count, - uint32_t resource_kind, - id overflow_resource) { - if (!cb || !mm_pipeline || !mm_args || !src0_addrs || !src1 || !dst || - !g_moe_id_map_buffer || - mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || - mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { - return 0; - } - - const NSUInteger tile_n = 32u; - const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); - const NSUInteger hids_bytes = - (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes || - g_moe_id_map_bytes < tpe_bytes + hids_bytes) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:mm_pipeline]; - [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; - [enc setBuffer:src0_addrs offset:0 atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; - [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; - [enc setBuffer:dst offset:dst_off atIndex:5]; - [enc useResource:src0_addrs usage:MTLResourceUsageRead]; - for (uint32_t i = 0; resources && i < resource_count; i++) { - ds4_gpu_stream_expert_cache_entry *entry = resources[i]; - if (!entry) continue; - id b = - resource_kind == 0 ? entry->gate_buffer : - resource_kind == 1 ? entry->up_buffer : - entry->down_buffer; - if (b) [enc useResource:b usage:MTLResourceUsageRead]; - } - if (overflow_resource) { - [enc useResource:overflow_resource usage:MTLResourceUsageRead]; - } - [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, - ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( - id cb, - id pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args, - id gate_src0, - NSUInteger gate_src0_off, - id up_src0, - NSUInteger up_src0_off, - id src1, - NSUInteger src1_off, - id mid, - NSUInteger mid_off, - id weights, - NSUInteger weights_off) { - if (!cb || !pipeline || !mm_args || !act_args || - !gate_src0 || !up_src0 || !src1 || !mid || !weights || - !g_moe_id_map_buffer || - mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || - mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { - return 0; - } - - const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); - const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); - if (tpe_bytes > NSUIntegerMax - hids_bytes || - g_moe_id_map_bytes < tpe_bytes + hids_bytes) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; - [enc setBytes:act_args length:sizeof(*act_args) atIndex:1]; - [enc setBuffer:gate_src0 offset:gate_src0_off atIndex:2]; - [enc setBuffer:up_src0 offset:up_src0_off atIndex:3]; - [enc setBuffer:src1 offset:src1_off atIndex:4]; - [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:5]; - [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6]; - [enc setBuffer:mid offset:mid_off atIndex:7]; - [enc setBuffer:weights offset:weights_off atIndex:8]; - [enc setThreadgroupMemoryLength:16384u atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + 31u) / 32u, - ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_mul_mm_id_mapped( - id cb, - id mm_pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off) { - return ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - mm_pipeline, - mm_args, - src0, - src0_off, - src1, - src1_off, - dst, - dst_off, - 8192u); -} - -static int ds4_gpu_encode_attn_out_low_q8_mpp( - id cb, - id pipeline, - const ds4_gpu_mul_mm_id_args *mm_args, - id src0, - NSUInteger src0_off, - id src1, - NSUInteger src1_off, - id dst, - NSUInteger dst_off) { - if (!cb || !pipeline || !mm_args || !src0 || !src1 || !dst || - mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || - mm_args->ne02 <= 0 || mm_args->ne1 <= 0 || mm_args->ne21 <= 0) { - return 0; - } - - const uint32_t tile_n = DS4_METAL_ATTN_OUT_MPP_TILE_N; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; - [enc setBuffer:src0 offset:src0_off atIndex:1]; - [enc setBuffer:src1 offset:src1_off atIndex:2]; - [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setThreadgroupMemoryLength:8192u atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n, - ((NSUInteger)mm_args->ne0 + 63u) / 64u, - (NSUInteger)mm_args->ne02) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_swiglu_flat( - id cb, - id gate, - NSUInteger gate_off, - id up, - NSUInteger up_off, - id out, - NSUInteger out_off, - uint32_t n) { - if (!cb || !gate || !up || !out || n == 0) return 0; - - ds4_gpu_glu_args args = { - .ne00 = (int32_t)n, - .nb01 = (uint64_t)n * sizeof(float), - .ne10 = (int32_t)n, - .nb11 = (uint64_t)n * sizeof(float), - .ne0 = (int32_t)n, - .nb1 = (uint64_t)n * sizeof(float), - .i00 = 0, - .i10 = 0, - .alpha = 1.0f, - .limit = 0.0f, - }; - NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n) nth = (NSUInteger)n; - if (nth == 0u) nth = 1u; - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_swiglu_flat_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gate offset:gate_off atIndex:1]; - [enc setBuffer:up offset:up_off atIndex:2]; - [enc setBuffer:out offset:out_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_moe_swiglu_weight( - id cb, - id gate, - NSUInteger gate_off, - id up, - NSUInteger up_off, - id mid, - NSUInteger mid_off, - id weights, - NSUInteger weights_off, - uint32_t width, - uint32_t rows, - float clamp_value, - bool mid_f16) { - if (!cb || !gate || !up || !mid || !weights || width == 0 || rows == 0) return 0; - - id pipeline = - ds4_gpu_get_pipeline(mid_f16 ? "kernel_dsv4_moe_swiglu_weight_f16" : - "kernel_dsv4_moe_swiglu_weight"); - if (!pipeline) return 0; - - ds4_gpu_dsv4_moe_swiglu_weight_args args = { - .width = width, - .rows = rows, - .gate_row_stride = (uint64_t)width * sizeof(float), - .up_row_stride = (uint64_t)width * sizeof(float), - .mid_row_stride = (uint64_t)width * (mid_f16 ? sizeof(uint16_t) : sizeof(float)), - .weight_stride = sizeof(float), - .write_clamped = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ? 1u : 0u, - .clamp_value = clamp_value, - }; - - NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > width) nth = width; - if (nth == 0) nth = 1u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:gate offset:gate_off atIndex:1]; - [enc setBuffer:up offset:up_off atIndex:2]; - [enc setBuffer:mid offset:mid_off atIndex:3]; - [enc setBuffer:weights offset:weights_off atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_moe_sum6( - id cb, - id experts, - NSUInteger experts_off, - id out, - NSUInteger out_off, - uint32_t out_dim, - uint32_t n_tokens) { - if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; - - if (!g_moe_sum6_pipeline) return 0; - - const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); - ds4_gpu_dsv4_moe_sum6_args args = { - .width = out_dim, - .tokens = n_tokens, - .src_token_stride = 6u * out_row_bytes, - .dst_token_stride = out_row_bytes, - }; - - NSUInteger nth = g_moe_sum6_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > out_dim) nth = out_dim; - if (nth == 0) nth = 1u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_moe_sum6_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:experts offset:experts_off atIndex:1]; - [enc setBuffer:out offset:out_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_moe_sum8( - id cb, - id experts, - NSUInteger experts_off, - id out, - NSUInteger out_off, - uint32_t out_dim, - uint32_t n_tokens) { - if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; - - if (!g_moe_sum8_pipeline) return 0; - - const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); - ds4_gpu_dsv4_moe_sum6_args args = { - .width = out_dim, - .tokens = n_tokens, - .src_token_stride = 8u * out_row_bytes, - .dst_token_stride = out_row_bytes, - }; - - NSUInteger nth = g_moe_sum8_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > out_dim) nth = out_dim; - if (nth == 0) nth = 1u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_moe_sum8_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:experts offset:experts_off atIndex:1]; - [enc setBuffer:out offset:out_off atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static ds4_gpu_bin_args ds4_gpu_make_moe_add_args( - uint32_t out_dim, - uint32_t n_tokens, - uint64_t src0_token_stride, - uint64_t src1_token_stride, - uint64_t dst_token_stride) { - return (ds4_gpu_bin_args) { - .ne00 = (int32_t)out_dim, - .ne01 = (int32_t)n_tokens, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src0_token_stride, - .nb02 = (uint64_t)n_tokens * src0_token_stride, - .nb03 = (uint64_t)n_tokens * src0_token_stride, - .ne10 = (int32_t)out_dim, - .ne11 = (int32_t)n_tokens, - .ne12 = 1, - .ne13 = 1, - .nb10 = sizeof(float), - .nb11 = src1_token_stride, - .nb12 = (uint64_t)n_tokens * src1_token_stride, - .nb13 = (uint64_t)n_tokens * src1_token_stride, - .ne0 = (int32_t)out_dim, - .ne1 = (int32_t)n_tokens, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = dst_token_stride, - .nb2 = (uint64_t)n_tokens * dst_token_stride, - .nb3 = (uint64_t)n_tokens * dst_token_stride, - .offs = 0, - .o1 = { 0 }, - }; -} - -static int ds4_gpu_encode_moe_sum_experts( - id cb, - id experts, - NSUInteger experts_off, - id out, - NSUInteger out_off, - uint32_t out_dim, - uint32_t n_expert, - uint32_t n_tokens) { - if (!cb || !experts || !out || out_dim == 0 || n_expert < 2 || n_tokens == 0) return 0; - - const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); - const uint64_t expert_token_stride = (uint64_t)n_expert * out_row_bytes; - - if (n_expert == 6 && - ds4_gpu_encode_moe_sum6(cb, - experts, - experts_off, - out, - out_off, - out_dim, - n_tokens)) { - return 1; - } - - if (n_expert == 8 && - ds4_gpu_encode_moe_sum8(cb, - experts, - experts_off, - out, - out_off, - out_dim, - n_tokens)) { - return 1; - } - - ds4_gpu_bin_args first = - ds4_gpu_make_moe_add_args(out_dim, n_tokens, expert_token_stride, expert_token_stride, out_row_bytes); - if (!ds4_gpu_encode_bin_f32_rows(cb, - g_add_pipeline, - &first, - experts, - experts_off, - experts, - experts_off + (NSUInteger)out_row_bytes, - out, - out_off)) { - return 0; - } - - ds4_gpu_bin_args accum = - ds4_gpu_make_moe_add_args(out_dim, n_tokens, out_row_bytes, expert_token_stride, out_row_bytes); - for (uint32_t slot = 2; slot < n_expert; slot++) { - if (!ds4_gpu_encode_bin_f32_rows(cb, - g_add_pipeline, - &accum, - out, - out_off, - experts, - experts_off + (NSUInteger)((uint64_t)slot * out_row_bytes), - out, - out_off)) { - return 0; - } - } - return 1; -} - -static int ds4_gpu_encode_get_rows_i32_token_rows( - id cb, - id table, - NSUInteger table_off, - id tokens, - NSUInteger tokens_off, - const int32_t *token_inline, - id selected, - NSUInteger selected_off, - uint32_t hash_rows, - uint32_t n_cols, - uint32_t n_tokens) { - if (!cb || !table || !selected || hash_rows == 0 || n_cols == 0 || n_tokens == 0) return 0; - if (!tokens && !token_inline) return 0; - - const uint64_t table_row_bytes = (uint64_t)n_cols * sizeof(int32_t); - const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); - ds4_gpu_get_rows_args args = { - .ne00t = (int64_t)n_cols, - .ne00 = (int64_t)n_cols, - .nb01 = table_row_bytes, - .nb02 = (uint64_t)hash_rows * table_row_bytes, - .nb03 = (uint64_t)hash_rows * table_row_bytes, - .ne10 = (int32_t)n_tokens, - .nb10 = sizeof(int32_t), - .nb11 = token_bytes, - .nb12 = token_bytes, - .nb1 = table_row_bytes, - .nb2 = (uint64_t)n_tokens * table_row_bytes, - .nb3 = (uint64_t)n_tokens * table_row_bytes, - }; - - NSUInteger nth = (NSUInteger)n_cols; - const NSUInteger max_threads = g_get_rows_i32_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > max_threads) nth = max_threads; - if (nth == 0) nth = 1u; - const NSUInteger nw0 = ((NSUInteger)n_cols + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_get_rows_i32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:table offset:table_off atIndex:1]; - if (tokens) { - [enc setBuffer:tokens offset:tokens_off atIndex:2]; - } else { - [enc setBytes:token_inline length:sizeof(*token_inline) atIndex:2]; - } - [enc setBuffer:selected offset:selected_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_get_rows_f32_router_weights( - id cb, - id probs, - NSUInteger probs_off, - id selected, - NSUInteger selected_off, - id weights, - NSUInteger weights_off, - uint32_t n_expert, - uint32_t n_expert_used, - uint32_t n_tokens) { - if (!cb || !probs || !selected || !weights || n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; - - const uint64_t probs_token_bytes = (uint64_t)n_expert * sizeof(float); - const uint64_t selected_row_bytes = (uint64_t)n_expert_used * sizeof(int32_t); - const uint64_t weights_row_bytes = (uint64_t)n_expert_used * sizeof(float); - ds4_gpu_get_rows_args args = { - .ne00t = 1, - .ne00 = 1, - .nb01 = sizeof(float), - .nb02 = probs_token_bytes, - .nb03 = (uint64_t)n_tokens * probs_token_bytes, - .ne10 = (int64_t)n_expert_used, - .nb10 = sizeof(int32_t), - .nb11 = selected_row_bytes, - .nb12 = (uint64_t)n_tokens * selected_row_bytes, - .nb1 = sizeof(float), - .nb2 = weights_row_bytes, - .nb3 = (uint64_t)n_tokens * weights_row_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_get_rows_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:probs offset:probs_off atIndex:1]; - [enc setBuffer:selected offset:selected_off atIndex:2]; - [enc setBuffer:weights offset:weights_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_expert_used, n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_sum_rows_f32( - id cb, - id src, - NSUInteger src_off, - id dst, - NSUInteger dst_off, - uint32_t width, - uint32_t rows) { - if (!cb || !src || !dst || width == 0 || rows == 0) return 0; - - const uint64_t src_row_bytes = (uint64_t)width * sizeof(float); - ds4_gpu_kargs_sum_rows args = { - .ne00 = (int64_t)width, - .ne01 = (int64_t)rows, - .ne02 = 1, - .ne03 = 1, - .nb00 = sizeof(float), - .nb01 = src_row_bytes, - .nb02 = (uint64_t)rows * src_row_bytes, - .nb03 = (uint64_t)rows * src_row_bytes, - .ne0 = 1, - .ne1 = (int64_t)rows, - .ne2 = 1, - .ne3 = 1, - .nb0 = sizeof(float), - .nb1 = sizeof(float), - .nb2 = (uint64_t)rows * sizeof(float), - .nb3 = (uint64_t)rows * sizeof(float), - }; - - NSUInteger nth = 32u; - const NSUInteger max_threads = g_sum_rows_f32_f32_pipeline.maxTotalThreadsPerThreadgroup; - while (nth < (NSUInteger)args.ne00 && nth < max_threads) nth *= 2u; - if (nth > max_threads) nth = max_threads; - if (nth > (NSUInteger)args.ne00) nth = (NSUInteger)args.ne00; - if (nth == 0) nth = 1u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_sum_rows_f32_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:src offset:src_off atIndex:1]; - [enc setBuffer:dst offset:dst_off atIndex:2]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; -} - -static int ds4_gpu_encode_router_select( - id cb, - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - id logitsbuf, - NSUInteger logits_off, - id biasbuf, - NSUInteger bias_off, - id hashbuf, - NSUInteger hash_off, - id tokensbuf, - NSUInteger tokens_off, - const int32_t *single_token, - uint32_t hash_rows, - uint32_t n_tokens, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - bool has_bias, - bool hash_mode) { - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id probsbuf = ds4_gpu_tensor_buffer(probs); - const NSUInteger selected_off = ds4_gpu_tensor_offset(selected); - const NSUInteger weights_off = ds4_gpu_tensor_offset(weights); - const NSUInteger probs_off = ds4_gpu_tensor_offset(probs); - - if (!cb || !selectedbuf || !weightsbuf || !probsbuf || !logitsbuf || - n_tokens == 0 || n_expert == 0 || n_expert_used == 0) return 0; - - const NSUInteger probs_bytes = (NSUInteger)n_tokens * (NSUInteger)n_expert * sizeof(float); - const bool flash_router_fast_path = - n_expert == 256u && - n_expert_used == 6u && - fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f; - - int ok = 0; - if (flash_router_fast_path && - !g_quality_mode && n_tokens == 1 && - getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) { - const bool force_simd_weights_fusion = - getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL; - const bool use_simd_finalize = - !hash_mode && - g_dsv4_router_finalize_one_simd_pipeline != nil && - g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u && - g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL || - force_simd_weights_fusion) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL; - const bool use_simd_weights_fusion = - use_simd_finalize && - g_dsv4_router_finalize_weights_one_simd_pipeline != nil && - g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && - g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && - (ds4_gpu_device_name_contains("M3") || - ds4_gpu_device_name_contains("M5") || - force_simd_weights_fusion) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL; - id softplus_sqrt_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, - "kernel_dsv4_softplus_sqrt_f32_4"); - id router_finalize_pipeline = - ds4_gpu_hot_pipeline( - use_simd_weights_fusion - ? g_dsv4_router_finalize_weights_one_simd_pipeline - : use_simd_finalize - ? g_dsv4_router_finalize_one_simd_pipeline - : g_dsv4_router_finalize_one_pipeline, - use_simd_weights_fusion - ? "kernel_dsv4_router_finalize_weights_one_simd" - : use_simd_finalize - ? "kernel_dsv4_router_finalize_one_simd" - : "kernel_dsv4_router_finalize_one"); - id router_weights_pipeline = use_simd_weights_fusion - ? nil - : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, - "kernel_dsv4_router_weights_one"); - if (!softplus_sqrt_pipeline || !router_finalize_pipeline || - (!use_simd_weights_fusion && !router_weights_pipeline)) return 0; - - ok = ds4_gpu_encode_unary_f32_rows(cb, - softplus_sqrt_pipeline, - logitsbuf, - logits_off, - probsbuf, - probs_off, - n_expert, - 1, - 1, - 0.0f, - 0.0f); - if (!ok) return 0; - - const bool use_token_buffer = single_token == NULL; - ds4_gpu_dsv4_router_select_one_args args = { - .has_bias = has_bias ? 1u : 0u, - .hash_mode = hash_mode ? 1u : 0u, - .use_token_buffer = use_token_buffer ? 1u : 0u, - .token = single_token ? (uint32_t)*single_token : 0u, - .hash_rows = hash_rows, - }; - - const float zero_f32 = 0.0f; - const int32_t zero_i32 = 0; - if ((has_bias && !biasbuf) || - (hash_mode && !hashbuf) || - (use_token_buffer && !tokensbuf)) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:router_finalize_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:probsbuf offset:probs_off atIndex:1]; - if (has_bias) { - [enc setBuffer:biasbuf offset:bias_off atIndex:2]; - } else { - [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2]; - } - if (hash_mode) { - [enc setBuffer:hashbuf offset:hash_off atIndex:3]; - } else { - [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3]; - } - if (use_token_buffer) { - [enc setBuffer:tokensbuf offset:tokens_off atIndex:4]; - } else { - [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; - } - [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; - if (use_simd_weights_fusion) { - [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; - } - const NSUInteger router_finalize_scratch_bytes = use_simd_finalize - ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t)) - : 256u * sizeof(float) + 256u * sizeof(int32_t); - [enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (use_simd_weights_fusion) return 1; - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:router_weights_pipeline]; - [enc setBuffer:probsbuf offset:probs_off atIndex:0]; - [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; - [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; - [enc dispatchThreads:MTLSizeMake(6, 1, 1) - threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; - } - - if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { - id softplus_sqrt_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, - "kernel_dsv4_softplus_sqrt_f32_4"); - ok = softplus_sqrt_pipeline && - ds4_gpu_encode_unary_f32_rows(cb, - softplus_sqrt_pipeline, - logitsbuf, - logits_off, - probsbuf, - probs_off, - n_expert, - 1, - 1, - 0.0f, - 0.0f); - } else { - ok = ds4_gpu_encode_unary_f32_rows(cb, - g_unary_softplus_pipeline, - logitsbuf, - logits_off, - probsbuf, - probs_off, - n_expert, - n_tokens, - 1, - 0.0f, - 0.0f) && - ds4_gpu_encode_unary_f32_rows(cb, - g_unary_sqrt_pipeline, - probsbuf, - probs_off, - probsbuf, - probs_off, - n_expert, - n_tokens, - 1, - 0.0f, - 0.0f); - } - if (!ok) return 0; - - if (hash_mode) { - ok = ds4_gpu_encode_get_rows_i32_token_rows(cb, - hashbuf, - hash_off, - tokensbuf, - tokens_off, - single_token, - selectedbuf, - selected_off, - hash_rows, - n_expert_used, - n_tokens); - } else { - ds4_gpu_tensor *score_tensor = probs; - DS4MetalTensor *selection_view = nil; - - if (has_bias) { - if (!biasbuf || - !ds4_gpu_ensure_scratch_buffer(&g_router_selection_buffer, - &g_router_selection_bytes, - probs_bytes, - "ds4_router_selection")) { - return 0; - } - - ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_expert, n_tokens, n_expert); - ok = ds4_gpu_encode_bin_f32_rows(cb, - g_add_pipeline, - &add_args, - probsbuf, - probs_off, - biasbuf, - bias_off, - g_router_selection_buffer, - 0); - if (!ok) return 0; - - selection_view = [DS4MetalTensor new]; - selection_view.buffer = g_router_selection_buffer; - selection_view.offset = 0; - selection_view.bytes = probs_bytes; - selection_view.owner = 0; - score_tensor = (__bridge ds4_gpu_tensor *)selection_view; - } - - ok = ds4_gpu_indexer_topk_tensor(selected, score_tensor, n_expert, n_tokens, n_expert_used) != 0; - } - if (!ok) return 0; - - const bool use_batch_weights_fusion = - flash_router_fast_path && !g_quality_mode && n_tokens > 1u && - g_dsv4_router_weights_batch_pipeline != nil && - (ds4_gpu_device_name_contains("M3") || - getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) && - getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && - getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL; - if (use_batch_weights_fusion) { - const float scale = expert_weight_scale; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline]; - [enc setBytes:&scale length:sizeof(scale) atIndex:0]; - [enc setBuffer:probsbuf offset:probs_off atIndex:1]; - [enc setBuffer:selectedbuf offset:selected_off atIndex:2]; - [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; - [enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; - } - - if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { - id router_weights_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, - "kernel_dsv4_router_weights_one"); - if (!router_weights_pipeline) return 0; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:router_weights_pipeline]; - [enc setBuffer:probsbuf offset:probs_off atIndex:0]; - [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; - [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; - [enc dispatchThreads:MTLSizeMake(6, 1, 1) - threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - return 1; - } - - const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer, - &g_router_weight_sum_bytes, - sum_bytes, - "ds4_router_weight_sum")) { - return 0; - } - - ok = ds4_gpu_encode_get_rows_f32_router_weights(cb, - probsbuf, - probs_off, - selectedbuf, - selected_off, - weightsbuf, - weights_off, - n_expert, - n_expert_used, - n_tokens) && - ds4_gpu_encode_sum_rows_f32(cb, - weightsbuf, - weights_off, - g_router_weight_sum_buffer, - 0, - n_expert_used, - n_tokens) && - ds4_gpu_encode_unary_f32_rows(cb, - g_unary_clamp_pipeline, - g_router_weight_sum_buffer, - 0, - g_router_weight_sum_buffer, - 0, - 1, - n_tokens, - 0, - 6.103515625e-5f, - ds4_gpu_positive_infinity()); - if (!ok) return 0; - - ds4_gpu_bin_args div_args = ds4_gpu_make_bin_rowwise_scalar_args(n_expert_used, n_tokens); - const float scale = expert_weight_scale; - ds4_gpu_bin_args scale_args = ds4_gpu_make_bin_rows_args(n_expert_used, n_tokens, 1); - - ok = ds4_gpu_encode_bin_f32_rows(cb, - g_bin_div_row_pipeline, - &div_args, - weightsbuf, - weights_off, - g_router_weight_sum_buffer, - 0, - weightsbuf, - weights_off); - if (!ok) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; - [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; - [enc setBuffer:weightsbuf offset:weights_off atIndex:1]; - [enc setBytes:&scale length:sizeof(scale) atIndex:2]; - [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)scale_args.ne1, - (NSUInteger)scale_args.ne2, - (NSUInteger)scale_args.ne3) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_bin_threads(n_expert_used, g_bin_mul_scalar_pipeline), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; -} - -int ds4_gpu_glm_kv_lora_rms_norm_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *kv_raw, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint32_t n_tokens, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !kv_raw || !model_map || - n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || - kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 || - !isfinite(eps) || eps < 0.0f) { - return 0; - } - - @autoreleasepool { - id rawbuf = ds4_gpu_tensor_buffer(kv_raw); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); - const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); - if (!rawbuf || !outbuf || - ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal GLM KV RMS norm received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal GLM KV RMS norm weight range is outside the mapped model\n"); - return 0; - } - - const bool exact_decode_weight_view = - n_tokens == 1u && - weight_bytes <= (1ull << 20) && - getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; - uint64_t weight_inner = 0; - id weightbuf = exact_decode_weight_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner); - if (!weightbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_kv_lora_rms_norm_pipeline, - "kernel_glm_kv_lora_rms_norm"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_kv_lora_rms_norm_args args = { - .n_tokens = n_tokens, - .kv_raw_dim = kv_raw_dim, - .kv_lora_dim = kv_lora_dim, - .eps = eps, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(kv_lora_dim); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV RMS norm")) return 0; - } - - return 1; -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !kv_norm || !model_map || - n_tokens == 0 || kv_lora_dim == 0 || - qk_nope == 0 || n_head == 0) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(out); - id kvbuf = ds4_gpu_tensor_buffer(kv_norm); - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { - fprintf(stderr, "ds4: Metal GLM k_b projection received unsupported weight type\n"); - return 0; - } - const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; - const uint64_t weight_bytes = weight_rows * row_bytes; - const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); - if (!outbuf || !kvbuf || - ds4_gpu_tensor_bytes(kv_norm) < kv_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes) { - fprintf(stderr, "ds4: Metal GLM k_b projection received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal GLM k_b projection range is outside the mapped model\n"); - return 0; - } - const NSUInteger q_blocks = ((NSUInteger)qk_nope + 31u) / 32u; - if (q_blocks > 8u) { - fprintf(stderr, "ds4: Metal GLM k_b projection q width is too large for the tiled kernel\n"); - return 0; - } - - uint64_t weight_inner = 0; - id weightbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner); - if (!weightbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_k_b_project_pipeline, - "kernel_glm_k_b_project_q8_0"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_k_b_project_args args = { - .n_tokens = n_tokens, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .n_head = n_head, - .row_bytes = (uint32_t)row_bytes, - .weight_type = weight_type, - .pad1 = 0, - .pad2 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; - [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, q_blocks, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM k_b projection")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_k_b_project_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 n_tokens, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t n_head) { - return ds4_gpu_glm_k_b_project_typed_tensor(out, - kv_norm, - model_map, - model_size, - weight_offset, - DS4_METAL_TENSOR_Q8_0, - n_tokens, - kv_lora_dim, - qk_nope, - n_head); -} - -int ds4_gpu_glm_store_compact_kv_tensor( - ds4_gpu_tensor *kv_lora_cache, - ds4_gpu_tensor *k_rope_cache, - const ds4_gpu_tensor *kv_norm, - const ds4_gpu_tensor *kv_raw, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || - n_tokens == 0 || cache_cap == 0 || - kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || - kv_lora_dim > kv_raw_dim || - qk_rope > kv_raw_dim - kv_lora_dim || - pos0 > cache_cap || n_tokens > cache_cap - pos0) { - return 0; - } - - @autoreleasepool { - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id kvnormbuf = ds4_gpu_tensor_buffer(kv_norm); - id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t kv_cache_bytes = - (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = - (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t kv_norm_bytes = - (uint64_t)n_tokens * kv_lora_dim * sizeof(float); - const uint64_t kv_raw_bytes = - (uint64_t)n_tokens * kv_raw_dim * sizeof(float); - if (!kvcachebuf || !ropecachebuf || !kvnormbuf || !kvrawbuf || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(kv_norm) < kv_norm_bytes || - ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { - fprintf(stderr, "ds4: Metal GLM compact KV store received undersized buffers\n"); - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_store_compact_kv_pipeline, - "kernel_glm_store_compact_kv"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_store_compact_kv_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_cap = cache_cap, - .kv_raw_dim = kv_raw_dim, - .kv_lora_dim = kv_lora_dim, - .qk_rope = qk_rope, - .cache_f16 = cache_f16 ? 1u : 0u, - .pad1 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:kvnormbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:1]; - [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 2, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM compact KV store")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_qkv_norm_store_compact_kv_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_lora_cache, - ds4_gpu_tensor *k_rope_cache, - const ds4_gpu_tensor *kv_raw, - uint64_t kv_weight_offset, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t kv_raw_dim, - uint32_t kv_lora_dim, - uint32_t qk_rope, - bool cache_f16, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || - !model_map || n_tokens == 0 || cache_cap == 0 || - q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || - (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 || - kv_lora_dim > kv_raw_dim || - qk_rope > kv_raw_dim - kv_lora_dim || - pos0 > cache_cap || n_tokens > cache_cap - pos0) { - return 0; - } - - @autoreleasepool { - id qbuf = ds4_gpu_tensor_buffer(q); - id qoutbuf = ds4_gpu_tensor_buffer(q_out); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); - const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = - (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = - (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t kv_raw_bytes = - (uint64_t)n_tokens * kv_raw_dim * sizeof(float); - if (!qbuf || !qoutbuf || !kvcachebuf || !ropecachebuf || !kvrawbuf || - ds4_gpu_tensor_bytes(q) < q_row_bytes * n_tokens || - ds4_gpu_tensor_bytes(q_out) < q_row_bytes * n_tokens || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { - fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store received undersized buffers\n"); - return 0; - } - if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || - kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) { - fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store weight range is outside the mapped model\n"); - return 0; - } - - uint64_t q_weight_inner = 0; - id q_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, - q_weight_offset, q_row_bytes, - &q_weight_inner); - if (!q_weightbuf) return 0; - uint64_t kv_weight_inner = 0; - id kv_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, - kv_weight_offset, kv_weight_bytes, - &kv_weight_inner); - if (!kv_weightbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_qkv_norm_store_compact_kv_pipeline, - "kernel_glm_qkv_norm_store_compact_kv"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_qkv_norm_store_compact_kv_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_cap = cache_cap, - .q_n = q_n, - .q_n4 = q_n / 4u, - .kv_raw_dim = kv_raw_dim, - .kv_lora_dim = kv_lora_dim, - .kv_lora_n4 = kv_lora_dim / 4u, - .qk_rope = qk_rope, - .cache_f16 = cache_f16 ? 1u : 0u, - .eps = eps, - .pad0 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:q_weightbuf offset:(NSUInteger)q_weight_inner atIndex:2]; - [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; - [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:4]; - [enc setBuffer:kv_weightbuf offset:(NSUInteger)kv_weight_inner atIndex:5]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:6]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:7]; - [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 3, 1) - threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM fused q/kv norm compact store")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_store_indexer_k_tensor( - ds4_gpu_tensor *indexer_key_cache, - const ds4_gpu_tensor *raw_k, - const void *model_map, - uint64_t model_size, - uint64_t weight_offset, - uint64_t bias_offset, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t n_ctx_orig, - float eps, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!indexer_key_cache || !raw_k || !model_map || - n_tokens == 0 || cache_cap == 0 || - head_dim == 0 || rot_dim == 0 || - rot_dim > head_dim || (rot_dim & 1u) != 0 || - pos0 > cache_cap || n_tokens > cache_cap - pos0) { - return 0; - } - - @autoreleasepool { - id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); - id rawbuf = ds4_gpu_tensor_buffer(raw_k); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t cache_bytes = - (uint64_t)cache_cap * head_dim * cache_elem_bytes; - const uint64_t raw_bytes = - (uint64_t)n_tokens * head_dim * sizeof(float); - const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); - if (!cachebuf || !rawbuf || - ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes || - ds4_gpu_tensor_bytes(raw_k) < raw_bytes) { - fprintf(stderr, "ds4: Metal GLM indexer K store received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || norm_bytes > model_size - weight_offset || - bias_offset > model_size || norm_bytes > model_size - bias_offset) { - fprintf(stderr, "ds4: Metal GLM indexer K norm range is outside the mapped model\n"); - return 0; - } - - uint64_t weight_inner = 0; - uint64_t bias_inner = 0; - id weightbuf = - ds4_gpu_wrap_model_range(model_map, model_size, - weight_offset, norm_bytes, - &weight_inner); - id biasbuf = - ds4_gpu_wrap_model_range(model_map, model_size, - bias_offset, norm_bytes, - &bias_inner); - if (!weightbuf || !biasbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_store_indexer_k_pipeline, - "kernel_glm_store_indexer_k"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_store_indexer_k_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_cap = cache_cap, - .head_dim = head_dim, - .rot_dim = rot_dim, - .n_ctx_orig = n_ctx_orig, - .cache_f16 = cache_f16 ? 1u : 0u, - .pad0 = 0, - .eps = eps, - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .pad1 = 0.0f, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; - [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:3]; - [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:4]; - [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer K store")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_build_kv_cache_tensor( - ds4_gpu_tensor *key_cache, - ds4_gpu_tensor *value_cache, - const ds4_gpu_tensor *kv_raw, - const ds4_gpu_tensor *k_nope, - const ds4_gpu_tensor *value, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_raw_dim, - 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, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || - n_tokens == 0 || cache_cap == 0 || n_head == 0 || - kv_raw_dim == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || value_dim == 0 || - kv_lora_dim + qk_rope > kv_raw_dim || - qk_dim < qk_nope || (qk_rope & 1u) != 0 || - pos0 > cache_cap || n_tokens > cache_cap - pos0 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id keybuf = ds4_gpu_tensor_buffer(key_cache); - id valbuf = ds4_gpu_tensor_buffer(value_cache); - id rawbuf = ds4_gpu_tensor_buffer(kv_raw); - id knbuf = ds4_gpu_tensor_buffer(k_nope); - id vbuf = ds4_gpu_tensor_buffer(value); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; - const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; - const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); - const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); - const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || - ds4_gpu_tensor_bytes(key_cache) < key_bytes || - ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || - ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || - ds4_gpu_tensor_bytes(k_nope) < kn_bytes || - ds4_gpu_tensor_bytes(value) < value_bytes) { - fprintf(stderr, "ds4: Metal GLM KV cache builder received undersized buffers\n"); - return 0; - } - - const bool decode_group4 = - n_tokens == 1u && - n_head >= 4u && - ds4_gpu_env_bool("DS4_METAL_DISABLE_GLM_DECODE_KV_GROUP4") <= 0; - id pipeline = - decode_group4 ? - ds4_gpu_hot_pipeline(g_glm_build_kv_cache_decode_group4_pipeline, - "kernel_glm_build_kv_cache_decode_group4") : - ds4_gpu_hot_pipeline(g_glm_build_kv_cache_pipeline, - "kernel_glm_build_kv_cache"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_build_kv_cache_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_cap = cache_cap, - .n_head = n_head, - .kv_raw_dim = kv_raw_dim, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = value_dim, - .n_ctx_orig = n_ctx_orig, - .cache_f16 = cache_f16 ? 1u : 0u, - .pad0 = 0u, - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; - [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; - [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; - [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; - [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; - const NSUInteger group_y = decode_group4 ? - ((NSUInteger)n_head + 3u) / 4u : - (NSUInteger)n_head; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, group_y, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV cache build")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_build_kv_cache_flash_tensor( - ds4_gpu_tensor *key_cache, - ds4_gpu_tensor *value_cache, - const ds4_gpu_tensor *kv_raw, - const ds4_gpu_tensor *k_nope, - const ds4_gpu_tensor *value, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_cap, - uint32_t n_head, - uint32_t kv_raw_dim, - 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, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || - pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 || - kv_raw_dim == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || value_dim == 0 || - kv_lora_dim + qk_rope > kv_raw_dim || - qk_dim < qk_nope || (qk_rope & 1u) != 0 || - n_tokens > cache_cap || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id keybuf = ds4_gpu_tensor_buffer(key_cache); - id valbuf = ds4_gpu_tensor_buffer(value_cache); - id rawbuf = ds4_gpu_tensor_buffer(kv_raw); - id knbuf = ds4_gpu_tensor_buffer(k_nope); - id vbuf = ds4_gpu_tensor_buffer(value); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; - const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; - const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); - const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); - const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || - ds4_gpu_tensor_bytes(key_cache) < key_bytes || - ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || - ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || - ds4_gpu_tensor_bytes(k_nope) < kn_bytes || - ds4_gpu_tensor_bytes(value) < value_bytes) { - fprintf(stderr, "ds4: Metal GLM staged KV cache builder received undersized buffers\n"); - return 0; - } - - const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); - const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); - const NSUInteger key_f16_offset = 0; - const NSUInteger key_f16_bytes = - (NSUInteger)n_tokens * (NSUInteger)n_head * q_row_bytes_f16; - const NSUInteger value_f16_offset = key_f16_bytes; - const NSUInteger value_f16_bytes = - (NSUInteger)n_tokens * (NSUInteger)n_head * v_row_bytes_f16; - const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_f16_bytes, - "ds4_glm_flash_attn_kv_f16")) { - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_build_kv_cache_flash_pipeline, - "kernel_glm_build_kv_cache_flash"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_build_kv_cache_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_cap = cache_cap, - .n_head = n_head, - .kv_raw_dim = kv_raw_dim, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = value_dim, - .n_ctx_orig = n_ctx_orig, - .cache_f16 = cache_f16 ? 1u : 0u, - .pad0 = 0u, - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; - [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; - [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; - [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; - [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; - [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:6]; - [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:7]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM staged KV cache build")) return 0; - } - - return 1; -} - -static int ds4_gpu_glm_attention_flash_tensor_impl( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16, - int kv_pre_staged) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !key_cache || !value_cache || - n_tokens == 0 || cache_len == 0 || cache_cap == 0 || - n_head == 0 || qk_dim != 256u || value_dim != 256u || - cache_len > cache_cap || - pos0 > cache_len || n_tokens > cache_len - pos0 || - cache_len > ds4_gpu_glm_flash_attention_max_cache_len()) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id qbuf = ds4_gpu_tensor_buffer(q); - id keybuf = ds4_gpu_tensor_buffer(key_cache); - id valbuf = ds4_gpu_tensor_buffer(value_cache); - const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; - const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; - if (!headsbuf || !qbuf || !keybuf || !valbuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(key_cache) < key_bytes || - ds4_gpu_tensor_bytes(value_cache) < value_bytes) { - fprintf(stderr, "ds4: Metal GLM FlashAttention received undersized buffers\n"); - return 0; - } - const uint64_t key_elems = (uint64_t)cache_len * n_head * qk_dim; - const uint64_t value_elems = (uint64_t)cache_len * n_head * value_dim; - if (key_elems > UINT32_MAX || value_elems > UINT32_MAX) { - return 0; - } - - const uint32_t nqptg = 8; - const uint32_t ncpsg = 64; - const uint32_t nsg = 4; - const bool has_kvpad = (cache_len % ncpsg) != 0; - const bool bc_mask = (n_tokens % nqptg) != 0; - const NSUInteger q_row_bytes = (NSUInteger)qk_dim * sizeof(float); - const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); - const NSUInteger v_row_bytes = (NSUInteger)value_dim * sizeof(float); - const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); - const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)cache_len * sizeof(uint16_t); - const NSUInteger key_f16_offset = 0; - const NSUInteger key_f16_bytes = - (NSUInteger)cache_len * (NSUInteger)n_head * q_row_bytes_f16; - const NSUInteger value_f16_offset = key_f16_bytes; - const NSUInteger value_f16_bytes = - (NSUInteger)cache_len * (NSUInteger)n_head * v_row_bytes_f16; - const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; - const NSUInteger pad_bytes = has_kvpad - ? (NSUInteger)ncpsg * ((NSUInteger)n_head * (q_row_bytes_f16 + v_row_bytes_f16) + - (NSUInteger)n_tokens * sizeof(uint16_t)) - : 1u; - const NSUInteger nblk0 = ((NSUInteger)cache_len + ncpsg - 1u) / ncpsg; - const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; - const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - - id mask_buffer = - ds4_gpu_glm_prefill_mask_buffer(pos0, n_tokens, cache_len, mask_bytes); - if (!mask_buffer) return 0; - if (kv_pre_staged) { - if (!g_flash_attn_kv_buffer || g_flash_attn_kv_bytes < kv_f16_bytes) { - fprintf(stderr, "ds4: GLM staged FlashAttention KV scratch is missing\n"); - return 0; - } - } else if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, - &g_flash_attn_kv_bytes, - kv_f16_bytes, - "ds4_glm_flash_attn_kv_f16")) { - return 0; - } - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, - &g_flash_attn_pad_bytes, - pad_bytes, - "ds4_glm_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_glm_flash_attn_blk")) { - return 0; - } - - id pad_pipeline = nil; - if (has_kvpad) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } - id blk_pipeline = - ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); - id attn_pipeline = - ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk256_dv256", - true, false, false, false, has_kvpad, bc_mask, - (int32_t)qk_dim, - (int32_t)value_dim, - (int32_t)nsg); - if (!blk_pipeline || !attn_pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - if (!kv_pre_staged) { - const bool copied = cache_f16 ? - (ds4_gpu_encode_cpy_f16_f16_3d(cb, - keybuf, - ds4_gpu_tensor_offset(key_cache), - g_flash_attn_kv_buffer, - key_f16_offset, - qk_dim, - cache_len, - n_head, - (uint64_t)n_head * q_row_bytes_f16, - q_row_bytes_f16, - q_row_bytes_f16, - (uint64_t)cache_len * q_row_bytes_f16) && - ds4_gpu_encode_cpy_f16_f16_3d(cb, - valbuf, - ds4_gpu_tensor_offset(value_cache), - g_flash_attn_kv_buffer, - value_f16_offset, - value_dim, - cache_len, - n_head, - (uint64_t)n_head * v_row_bytes_f16, - v_row_bytes_f16, - v_row_bytes_f16, - (uint64_t)cache_len * v_row_bytes_f16)) : - (ds4_gpu_encode_cpy_f32_f16_3d(cb, - keybuf, - ds4_gpu_tensor_offset(key_cache), - g_flash_attn_kv_buffer, - key_f16_offset, - qk_dim, - cache_len, - n_head, - (uint64_t)n_head * q_row_bytes, - q_row_bytes, - q_row_bytes_f16, - (uint64_t)cache_len * q_row_bytes_f16) && - ds4_gpu_encode_cpy_f32_f16_3d(cb, - valbuf, - ds4_gpu_tensor_offset(value_cache), - g_flash_attn_kv_buffer, - value_f16_offset, - value_dim, - cache_len, - n_head, - (uint64_t)n_head * v_row_bytes, - v_row_bytes, - v_row_bytes_f16, - (uint64_t)cache_len * v_row_bytes_f16)); - if (!copied) { - return 0; - } - } - - if (has_kvpad) { - ds4_gpu_flash_attn_pad_args pad_args = { - .ne11 = (int32_t)cache_len, - .ne_12_2 = (int32_t)n_head, - .ne_12_3 = 1, - .nb11 = q_row_bytes_f16, - .nb12 = (uint64_t)cache_len * q_row_bytes_f16, - .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, - .nb21 = v_row_bytes_f16, - .nb22 = (uint64_t)cache_len * v_row_bytes_f16, - .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)cache_len * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pad_pipeline]; - [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; - [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:2]; - [enc setBuffer:mask_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(ncpsg, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - - ds4_gpu_flash_attn_blk_args blk_args = { - .ne01 = (int32_t)n_tokens, - .ne30 = (int32_t)cache_len, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)cache_len * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - ds4_gpu_flash_attn_vec_args args = { - .ne01 = (int32_t)n_tokens, - .ne02 = (int32_t)n_head, - .ne03 = 1, - .nb01 = (uint64_t)n_head * q_row_bytes, - .nb02 = q_row_bytes, - .nb03 = (uint64_t)n_tokens * n_head * q_row_bytes, - .ne11 = (int32_t)cache_len, - .ne_12_2 = (int32_t)n_head, - .ne_12_3 = 1, - .ns10 = (int32_t)qk_dim, - .nb11 = q_row_bytes_f16, - .nb12 = (uint64_t)cache_len * q_row_bytes_f16, - .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, - .ns20 = (int32_t)value_dim, - .nb21 = v_row_bytes_f16, - .nb22 = (uint64_t)cache_len * v_row_bytes_f16, - .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, - .ne31 = (int32_t)n_tokens, - .ne32 = 1, - .ne33 = 1, - .nb31 = (uint64_t)cache_len * sizeof(uint16_t), - .nb32 = mask_bytes, - .nb33 = mask_bytes, - .ne1 = (int32_t)n_head, - .ne2 = (int32_t)n_tokens, - .ne3 = 1, - .scale = 1.0f / sqrtf((float)qk_dim), - .max_bias = 0.0f, - .m0 = 0.0f, - .m1 = 0.0f, - .n_head_log2 = 0, - .logit_softcap = 0.0f, - }; - - const NSUInteger padded_v = ds4_gpu_align_up_ns(value_dim, 64u); - const NSUInteger shared_elems = (NSUInteger)nqptg * - ((NSUInteger)qk_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); - const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:attn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:2]; - [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:3]; - [enc setBuffer:mask_buffer offset:0 atIndex:4]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:5]; - [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) - threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM FlashAttention")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_flash_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - return ds4_gpu_glm_attention_flash_tensor_impl(heads, - q, - key_cache, - value_cache, - pos0, - n_tokens, - cache_len, - cache_cap, - n_head, - qk_dim, - value_dim, - cache_f16, - 0); -} - -int ds4_gpu_glm_attention_flash_staged_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - if (pos0 != 0 || n_tokens != cache_len) return 0; - return ds4_gpu_glm_attention_flash_tensor_impl(heads, - q, - key_cache, - value_cache, - pos0, - n_tokens, - cache_len, - cache_cap, - n_head, - qk_dim, - value_dim, - cache_f16, - 1); -} - -int ds4_gpu_glm_attention_full_tensor( - ds4_gpu_tensor *heads, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *key_cache, - const ds4_gpu_tensor *value_cache, - uint32_t pos0, - uint32_t n_tokens, - uint32_t cache_len, - uint32_t cache_cap, - uint32_t n_head, - uint32_t qk_dim, - uint32_t value_dim, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !q || !key_cache || !value_cache || - n_tokens == 0 || cache_len == 0 || cache_cap == 0 || - n_head == 0 || qk_dim == 0 || value_dim == 0 || - (qk_dim & 3u) != 0 || - cache_len > cache_cap || - pos0 > cache_len || n_tokens > cache_len - pos0 || - cache_len > ds4_gpu_glm_full_attention_max_cache_len()) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id qbuf = ds4_gpu_tensor_buffer(q); - id keybuf = ds4_gpu_tensor_buffer(key_cache); - id valbuf = ds4_gpu_tensor_buffer(value_cache); - const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; - const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; - if (!headsbuf || !qbuf || !keybuf || !valbuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(key_cache) < key_bytes || - ds4_gpu_tensor_bytes(value_cache) < value_bytes) { - fprintf(stderr, "ds4: Metal GLM attention received undersized buffers\n"); - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_attention_full_pipeline, - "kernel_glm_attention_full"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const uint32_t full_attn_mode = 2u; - ds4_gpu_glm_attention_full_args args = { - .pos0 = pos0, - .n_tokens = n_tokens, - .cache_len = cache_len, - .cache_cap = cache_cap, - .n_head = n_head, - .qk_dim = qk_dim, - .value_dim = value_dim, - .pad0 = full_attn_mode, - .cache_f16 = cache_f16 ? 1u : 0u, - .pad1 = 0u, - .pad2 = 0u, - .scale = 1.0f / sqrtf((float)qk_dim), - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:2]; - [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:3]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; - [enc setThreadgroupMemoryLength:(256u + (NSUInteger)cache_len) * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM full attention")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_fill_selected_range_tensor( - ds4_gpu_tensor *selected, - uint32_t n_selected) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || n_selected == 0) return 0; - - @autoreleasepool { - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); - if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM selected range received undersized buffer\n"); - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_fill_selected_range_pipeline, - "kernel_glm_fill_selected_range"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_fill_selected_range_args args = { - .n_selected = n_selected, - }; - const NSUInteger nth = 256u; - const NSUInteger n_groups = ((NSUInteger)n_selected + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_fill_selected_range_batch_tensor( - ds4_gpu_tensor *selected, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_selected, - uint32_t pad_row) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || n_tokens == 0 || n_selected == 0) return 0; - - @autoreleasepool { - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t total = (uint64_t)n_tokens * n_selected; - if (n_tokens != 0 && total / n_tokens != n_selected) return 0; - if (total > UINT64_MAX / sizeof(uint32_t)) return 0; - const uint64_t selected_bytes = total * sizeof(uint32_t); - if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM selected range batch received undersized buffer\n"); - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_fill_selected_range_batch_pipeline, - "kernel_glm_fill_selected_range_batch"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_fill_selected_range_batch_args args = { - .n_tokens = n_tokens, - .pos0 = pos0, - .n_selected = n_selected, - .pad_row = pad_row, - }; - const NSUInteger nth = 256u; - const NSUInteger n_groups = ((NSUInteger)total + nth - 1u) / nth; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range batch")) return 0; - } - - return 1; -} - -static int ds4_gpu_glm_rope_tail_offset_tensor( - ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t rot_offset, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - const char *label) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || - rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset || - (rot_dim & 1u) != 0 || - pos0 > UINT32_MAX - n_tokens || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { - fprintf(stderr, "ds4: Metal %s received undersized buffer\n", - label ? label : "GLM RoPE"); - return 0; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_indexer_rope_tail_pipeline, - "kernel_glm_indexer_rope_tail_f32"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_rope_tail_args args = { - .n_tokens = n_tokens, - .n_head = n_head, - .head_dim = head_dim, - .rot_dim = rot_dim, - .rot_offset = rot_offset, - .pos0 = pos0, - .n_ctx_orig = n_ctx_orig, - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(rot_dim); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, label ? label : "GLM RoPE")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (rot_dim > head_dim) return 0; - return ds4_gpu_glm_rope_tail_offset_tensor(x, - n_tokens, - n_head, - head_dim, - rot_dim, - head_dim - rot_dim, - pos0, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow, - "GLM RoPE"); -} - -int ds4_gpu_glm_indexer_rope_tail_tensor( - ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t n_head, - uint32_t head_dim, - uint32_t rot_dim, - uint32_t pos0, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ds4_gpu_glm_rope_tail_offset_tensor(x, - n_tokens, - n_head, - head_dim, - rot_dim, - 0, - pos0, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow, - "GLM indexer RoPE"); -} - -int ds4_gpu_glm_indexer_score_one_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *indexer_key_cache, - uint32_t n_rows, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!scores || !q || !weights || !indexer_key_cache || - n_rows == 0 || n_head == 0 || head_dim == 0 || - !isfinite(scale) || scale <= 0.0f) { - return 0; - } - - @autoreleasepool { - id scoresbuf = ds4_gpu_tensor_buffer(scores); - id qbuf = ds4_gpu_tensor_buffer(q); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t score_bytes = (uint64_t)n_rows * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); - const uint64_t weights_bytes = (uint64_t)n_head * sizeof(float); - const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; - if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || - ds4_gpu_tensor_bytes(scores) < score_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes || - ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { - fprintf(stderr, "ds4: Metal GLM indexer score received undersized buffers\n"); - return 0; - } - - ds4_gpu_glm_indexer_score_one_args args = { - .n_rows = n_rows, - .n_head = n_head, - .head_dim = head_dim, - .cache_f16 = cache_f16 ? 1u : 0u, - .scale = scale, - }; - - if (n_head == 32u && head_dim == 128u) { - id direct_pipeline = - ds4_gpu_hot_pipeline(g_glm_indexer_score_one_direct_pipeline, - "kernel_glm_indexer_score_one_direct"); - if (!direct_pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:direct_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; - [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; - [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer direct score")) return 0; - return 1; - } - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_indexer_score_one_pipeline, - "kernel_glm_indexer_score_one"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; - [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; - [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer score")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_indexer_scores_batch_tensor( - ds4_gpu_tensor *scores, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *weights, - const ds4_gpu_tensor *indexer_key_cache, - uint32_t n_rows, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_head, - uint32_t head_dim, - float scale, - bool cache_f16) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!scores || !q || !weights || !indexer_key_cache || - n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128 || - pos0 >= n_rows || n_tokens > n_rows - pos0 || - !isfinite(scale) || scale <= 0.0f) { - return 0; - } - - @autoreleasepool { - id scoresbuf = ds4_gpu_tensor_buffer(scores); - id qbuf = ds4_gpu_tensor_buffer(q); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t score_bytes = (uint64_t)n_rows * n_tokens * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); - const uint64_t weights_bytes = (uint64_t)n_tokens * n_head * sizeof(float); - const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; - if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || - ds4_gpu_tensor_bytes(scores) < score_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes || - ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { - fprintf(stderr, "ds4: Metal GLM indexer batch scores received undersized buffers\n"); - return 0; - } - - const bool force_scalar = g_quality_mode; - const bool use_tiled_f32 = false; - const bool use_tiled = !force_scalar && n_tokens >= 8u && - n_head == 32u && head_dim == 128u; - id pipeline = - use_tiled - ? ds4_gpu_hot_pipeline(use_tiled_f32 ? g_glm_indexer_scores_tiled_f32_pipeline - : g_glm_indexer_scores_tiled_pipeline, - use_tiled_f32 ? "kernel_glm_indexer_scores_tiled_f32" - : "kernel_glm_indexer_scores_tiled") - : ds4_gpu_hot_pipeline(g_glm_indexer_scores_batch_pipeline, - "kernel_glm_indexer_scores_batch"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_indexer_scores_batch_args args = { - .n_rows = n_rows, - .n_tokens = n_tokens, - .n_head = n_head, - .head_dim = head_dim, - .pos0 = pos0, - .cache_f16 = cache_f16 ? 1u : 0u, - .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), - .q_head_stride = (uint64_t)head_dim * sizeof(float), - .weights_token_stride = (uint64_t)n_head * sizeof(float), - .score_token_stride = (uint64_t)n_rows * sizeof(float), - .scale = scale, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; - [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; - if (use_tiled) { - const NSUInteger q_shared = 8u * 128u; - const NSUInteger k_shared = 32u * 128u; - const NSUInteger dot_shared = 8u * 32u; - if (use_tiled_f32) { - [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * - sizeof(float) atIndex:0]; - } else { - [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + - dot_shared * sizeof(float) atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + 31u) / 32u, - ((NSUInteger)n_tokens + 7u) / 8u, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - } else { - [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer batch scores")) return 0; - } - - return 1; -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - if (!qk_low || !q || !model_map || - n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_nope > qk_dim) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(qk_low); - id qbuf = ds4_gpu_tensor_buffer(q); - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { - fprintf(stderr, "ds4: Metal GLM qk lowrank received unsupported weight type\n"); - return 0; - } - const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; - const uint64_t weight_bytes = weight_rows * row_bytes; - const uint64_t out_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); - if (!outbuf || !qbuf || - ds4_gpu_tensor_bytes(qk_low) < out_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes) { - fprintf(stderr, "ds4: Metal GLM qk lowrank received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal GLM qk lowrank range is outside the mapped model\n"); - return 0; - } - - uint64_t weight_inner = 0; - id weightbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner); - if (!weightbuf) return 0; - - const int use_glm52 = - n_head == 64u && - kv_lora_dim == 512u && - qk_nope == 192u && - qk_dim == 256u && - row_bytes == 204u && - weight_type == DS4_METAL_TENSOR_Q8_0; - /* Coalesced simdgroup variant: lanes split the 192-dot so weight - * reads coalesce, and 2048 threadgroups replace 64. The thread- - * per-row kernels measured ~7.5x off the weight-bandwidth floor - * (7.2ms of the decode token by skip-ablation). Covers the GLM 5.2 - * shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8 - * fast path above never engaged there). */ - const int use_glm52_sg = - n_head == 64u && - kv_lora_dim == 512u && - qk_nope == 192u && - qk_dim == 256u && - ((weight_type == DS4_METAL_TENSOR_Q8_0 && row_bytes == 204u) || - (weight_type == DS4_METAL_TENSOR_Q4_0 && row_bytes == 108u)) && - getenv("DS4_METAL_DISABLE_GLM_QKLOW_SG") == NULL && - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, - "kernel_glm_qk_lowrank_q8_0_glm52_sg") != nil; - if (getenv("DS4_METAL_GLM_QKLOW_DEBUG")) { - static int printed = 0; - if (!printed) { - printed = 1; - fprintf(stderr, "ds4: qk_lowrank decode path: use_glm52=%d sg=%d n_head=%u kv=%u nope=%u dim=%u rb=%llu type=%u\n", - use_glm52, use_glm52_sg, n_head, kv_lora_dim, qk_nope, qk_dim, - (unsigned long long)row_bytes, weight_type); - } - } - id pipeline = - use_glm52_sg ? - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, - "kernel_glm_qk_lowrank_q8_0_glm52_sg") : - use_glm52 ? - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_pipeline, - "kernel_glm_qk_lowrank_q8_0_glm52") : - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_pipeline, - "kernel_glm_qk_lowrank_q8_0"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_qk_lowrank_args args = { - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_dim = qk_dim, - .row_bytes = (uint32_t)row_bytes, - .weight_type = weight_type, - .pad1 = 0, - .pad2 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; - if (use_glm52_sg) { - /* 8 simdgroups x 2 rows per threadgroup: (64, 512/16) grid. */ - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, - (NSUInteger)(kv_lora_dim / 16u), - 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - } else { - if (use_glm52) { - [enc setThreadgroupMemoryLength:192u * sizeof(float) atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM qk lowrank")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_qk_lowrank_q8_0_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 n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_dim) { - return ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low, - q, - model_map, - model_size, - weight_offset, - DS4_METAL_TENSOR_Q8_0, - n_head, - kv_lora_dim, - qk_nope, - qk_dim); -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - if (!qk_low || !q || !model_map || - n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_nope > qk_dim) { - return 0; - } - - @autoreleasepool { - id outbuf = ds4_gpu_tensor_buffer(qk_low); - id qbuf = ds4_gpu_tensor_buffer(q); - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { - fprintf(stderr, "ds4: Metal GLM batch qk lowrank received unsupported weight type\n"); - return 0; - } - const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; - const uint64_t weight_bytes = weight_rows * row_bytes; - const uint64_t out_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t q_bytes = - (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - if (!outbuf || !qbuf || - ds4_gpu_tensor_bytes(qk_low) < out_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes) { - fprintf(stderr, "ds4: Metal GLM batch qk lowrank received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal GLM batch qk lowrank range is outside the mapped model\n"); - return 0; - } - - uint64_t weight_inner = 0; - id weightbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner); - if (!weightbuf) return 0; - - const int use_glm52_t4 = - n_tokens >= 4u && - n_head == 64u && - kv_lora_dim == 512u && - qk_nope == 192u && - qk_dim == 256u && - row_bytes == 204u && - weight_type == DS4_METAL_TENSOR_Q8_0; - id pipeline = - use_glm52_t4 ? - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline, - "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") : - ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline, - "kernel_glm_qk_lowrank_q8_0_batch"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - uint32_t head_base = 0; - uint32_t head_count = n_head; - ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); - ds4_gpu_glm_qk_lowrank_batch_args args = { - .n_tokens = n_tokens, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_dim = qk_dim, - .row_bytes = (uint32_t)row_bytes, - .weight_type = weight_type, - .head_base = head_base, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; - if (use_glm52_t4) { - [enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, - ((NSUInteger)n_tokens + 3u) / 4u, - 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } else { - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch qk lowrank")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_qk_lowrank_q8_0_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 n_tokens, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_dim) { - return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, - q, - model_map, - model_size, - weight_offset, - DS4_METAL_TENSOR_Q8_0, - n_tokens, - n_head, - kv_lora_dim, - qk_nope, - qk_dim); -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - if (!heads || !lora || !model_map || - n_tokens == 0 || n_head == 0 || - kv_lora_dim == 0 || value_dim == 0) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id lorabuf = ds4_gpu_tensor_buffer(lora); - uint64_t row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(weight_type, kv_lora_dim, &row_bytes)) { - fprintf(stderr, "ds4: Metal GLM batch value project received unsupported weight type\n"); - return 0; - } - const uint64_t weight_rows = (uint64_t)n_head * value_dim; - const uint64_t weight_bytes = weight_rows * row_bytes; - const uint64_t heads_bytes = - (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - const uint64_t lora_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - if (!headsbuf || !lorabuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(lora) < lora_bytes) { - fprintf(stderr, "ds4: Metal GLM batch value project received undersized buffers\n"); - return 0; - } - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal GLM batch value project range is outside the mapped model\n"); - return 0; - } - - uint64_t weight_inner = 0; - id weightbuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - weight_offset, - weight_bytes, - &weight_inner); - if (!weightbuf) return 0; - - const int use_mma = - n_head == 64u && - kv_lora_dim == 512u && - value_dim == 256u && - row_bytes == 544u && - weight_type == DS4_METAL_TENSOR_Q8_0 && - n_tokens >= 32u; - id pipeline = - use_mma ? - ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_mma_pipeline, - "kernel_glm_value_project_q8_0_batch_heads_mma") : - ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_pipeline, - "kernel_glm_value_project_q8_0_batch_heads"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - uint32_t head_base = 0; - uint32_t head_count = n_head; - ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); - ds4_gpu_glm_qk_lowrank_batch_args args = { - .n_tokens = n_tokens, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = 0, - .qk_dim = value_dim, - .row_bytes = (uint32_t)row_bytes, - .weight_type = weight_type, - .head_base = head_base, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; - [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:2]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3]; - if (use_mma) { - [enc setThreadgroupMemoryLength:16u * 1024u atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tokens + 31u) / 32u, - ((NSUInteger)value_dim + 63u) / 64u, - (NSUInteger)head_count) - threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; - } else { - [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) - atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch value project")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_value_project_q8_0_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 n_tokens, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t value_dim) { - return ds4_gpu_glm_value_project_typed_batch_heads_tensor(heads, - lora, - model_map, - model_size, - weight_offset, - DS4_METAL_TENSOR_Q8_0, - n_tokens, - n_head, - kv_lora_dim, - value_dim); -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - !model_map || !selected || - n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || - value_dim == 0 || qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id qbuf = ds4_gpu_tensor_buffer(q); - id lowbuf = ds4_gpu_tensor_buffer(qk_low); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - uint64_t value_row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { - fprintf(stderr, "ds4: Metal GLM indexed attention received unsupported value type\n"); - return 0; - } - const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; - const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; - const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); - const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); - if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(qk_low) < low_bytes || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM indexed attention received undersized buffers\n"); - return 0; - } - if (value_weight_offset > model_size || - value_weight_bytes > model_size - value_weight_offset) { - fprintf(stderr, "ds4: Metal GLM indexed attention value range is outside the mapped model\n"); - return 0; - } - - uint64_t value_inner = 0; - id valuebuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - value_weight_offset, - value_weight_bytes, - &value_inner); - if (!valuebuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_pipeline, - "kernel_glm_attention_indexed_decode"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - ds4_gpu_glm_attention_indexed_decode_args args = { - .n_selected = n_selected, - .cache_cap = cache_cap, - .cache_f16 = cache_f16 ? 1u : 0u, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = value_dim, - .n_ctx_orig = n_ctx_orig, - .value_row_bytes = (uint32_t)value_row_bytes, - .scale = 1.0f / sqrtf((float)qk_dim), - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .value_type = value_weight_type, - }; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; - const NSUInteger scratch_floats = - 256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim; - [enc setThreadgroupMemoryLength:scratch_floats * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed attention decode")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_indexed_decode_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, - 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) { - return ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - model_map, - model_size, - value_weight_offset, - DS4_METAL_TENSOR_Q8_0, - 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); -} - -int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( - ds4_gpu_tensor *heads, - ds4_gpu_tensor *partial_lora, - ds4_gpu_tensor *partial_ms, - 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, - bool selected_rows_valid, - 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, - uint32_t block_rows, - uint32_t n_blocks, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - const uint32_t needed_blocks = - block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; - if (!heads || !partial_lora || !partial_ms || !q || !qk_low || - !kv_lora_cache || !k_rope_cache || !model_map || !selected || - n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - n_head == 0 || (n_head % 8u) != 0 || - kv_lora_dim != 512u || - qk_nope == 0 || qk_rope != 64u || - value_dim == 0 || qk_dim < qk_nope || - block_rows == 0u || needed_blocks == 0u || - n_blocks < needed_blocks || n_blocks > 64u || - !cache_f16 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id partial_lorabuf = ds4_gpu_tensor_buffer(partial_lora); - id partial_msbuf = ds4_gpu_tensor_buffer(partial_ms); - id qbuf = ds4_gpu_tensor_buffer(q); - id lowbuf = ds4_gpu_tensor_buffer(qk_low); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t cache_elem_bytes = sizeof(uint16_t); - uint64_t value_row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { - fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received unsupported value type\n"); - return 0; - } - const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; - const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; - const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); - const uint64_t partial_lora_bytes = - (uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float); - const uint64_t partial_ms_bytes = - (uint64_t)n_blocks * n_head * 2u * sizeof(float); - const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); - const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); - if (!headsbuf || !partial_lorabuf || !partial_msbuf || - !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(partial_lora) < partial_lora_bytes || - ds4_gpu_tensor_bytes(partial_ms) < partial_ms_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(qk_low) < low_bytes || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received undersized buffers\n"); - return 0; - } - if (value_weight_offset > model_size || - value_weight_bytes > model_size - value_weight_offset) { - fprintf(stderr, "ds4: Metal GLM split grouped indexed attention value range is outside the mapped model\n"); - return 0; - } - - uint64_t value_inner = 0; - id valuebuf = - ds4_gpu_wrap_model_range(model_map, - model_size, - value_weight_offset, - value_weight_bytes, - &value_inner); - if (!valuebuf) return 0; - - const bool use_valid_fullheads = - selected_rows_valid && (n_head % 8u) == 0u; - id partial_pipeline = - ds4_gpu_hot_pipeline(use_valid_fullheads ? - g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline : - g_glm_attention_indexed_decode_split_group8_partial_pipeline, - use_valid_fullheads ? - "kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads" : - "kernel_glm_attention_indexed_decode_split_group8_partial"); - const bool use_reduce16 = - n_blocks == 16u && block_rows == 128u && - n_selected == 2048u && value_dim == 256u; - id reduce_pipeline = - ds4_gpu_hot_pipeline(use_reduce16 ? - g_glm_attention_indexed_decode_split_group8_reduce16_pipeline : - g_glm_attention_indexed_decode_split_group8_reduce_pipeline, - use_reduce16 ? - "kernel_glm_attention_indexed_decode_split_group8_reduce16" : - "kernel_glm_attention_indexed_decode_split_group8_reduce"); - if (!partial_pipeline || !reduce_pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_attention_indexed_decode_split_args args = { - .n_selected = n_selected, - .cache_cap = cache_cap, - .cache_f16 = 1u, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = value_dim, - .n_ctx_orig = n_ctx_orig, - .value_row_bytes = (uint32_t)value_row_bytes, - .block_rows = block_rows, - .n_blocks = n_blocks, - .scale = 1.0f / sqrtf((float)qk_dim), - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .value_type = value_weight_type, - }; - const NSUInteger stage_rows = 16u; - const NSUInteger kv_vecs = (NSUInteger)kv_lora_dim / 4u; - const NSUInteger rope_vecs = (NSUInteger)qk_rope / 4u; - const NSUInteger partial_scratch_bytes = - stage_rows * kv_vecs * sizeof(uint16_t) * 4u + - stage_rows * rope_vecs * sizeof(float) * 4u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:partial_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; - [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:6]; - [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:7]; - [enc setThreadgroupMemoryLength:partial_scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head / 8u, (NSUInteger)n_blocks, 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - const NSUInteger reduce_scratch_floats = 256u + 64u + (NSUInteger)kv_lora_dim; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:reduce_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:1]; - [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:2]; - [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:3]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; - [enc setThreadgroupMemoryLength:reduce_scratch_floats * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM split grouped indexed attention decode")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( - ds4_gpu_tensor *heads, - ds4_gpu_tensor *partial_lora, - ds4_gpu_tensor *partial_ms, - 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, - const ds4_gpu_tensor *selected, - uint32_t n_selected, - bool selected_rows_valid, - 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, - uint32_t block_rows, - uint32_t n_blocks, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(heads, - partial_lora, - partial_ms, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - model_map, - model_size, - value_weight_offset, - DS4_METAL_TENSOR_Q8_0, - selected, - n_selected, - selected_rows_valid, - cache_cap, - cache_f16, - n_head, - kv_lora_dim, - qk_nope, - qk_rope, - value_dim, - n_ctx_orig, - block_rows, - n_blocks, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow); -} - -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 (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - !model_map || !selected || - n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || - value_dim == 0 || qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id headsbuf = ds4_gpu_tensor_buffer(heads); - id qbuf = ds4_gpu_tensor_buffer(q); - id lowbuf = ds4_gpu_tensor_buffer(qk_low); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - uint64_t value_row_bytes = 0; - if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { - fprintf(stderr, "ds4: Metal GLM indexed batch attention received unsupported value type\n"); - return 0; - } - const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; - const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; - const uint64_t heads_bytes = - (uint64_t)n_tokens * n_head * value_dim * sizeof(float); - const uint64_t q_bytes = - (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - const uint64_t low_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t selected_bytes = - (uint64_t)n_tokens * n_selected * sizeof(uint32_t); - if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || - ds4_gpu_tensor_bytes(heads) < heads_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(qk_low) < low_bytes || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM indexed batch attention received undersized buffers\n"); - return 0; - } - if (value_weight_offset > model_size || - value_weight_bytes > model_size - value_weight_offset) { - fprintf(stderr, "ds4: Metal GLM indexed batch attention value range is outside the mapped model\n"); - return 0; - } - - uint64_t value_inner = 0; - id valuebuf = - ds4_gpu_wrap_model_range(model_map, model_size, - value_weight_offset, - value_weight_bytes, - &value_inner); - if (!valuebuf) return 0; - - const uint64_t grouped_attn = (n_tokens >= 128u) ? 8u : 2u; - const NSUInteger q2_bit_words = ((NSUInteger)cache_cap + 31u) / 32u; - const NSUInteger q2_scratch_bytes = - q2_bit_words * sizeof(uint32_t) + - 4u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + - 8u * (NSUInteger)kv_lora_dim * sizeof(float); - const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; - const bool q2_fits = max_tg_mem == 0 || q2_scratch_bytes <= max_tg_mem; - const bool use_q2_group4 = grouped_attn == 4u && n_tokens >= 2u && q2_fits; - const bool use_group8 = - grouped_attn == 8u || - (grouped_attn == 4u && !use_q2_group4); - const bool use_group2 = grouped_attn == 1u || grouped_attn == 2u; - id pipeline = use_q2_group4 ? - ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_q2_group4_pipeline, - "kernel_glm_attention_indexed_batch_q2_group4") : - (use_group8 ? - ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, - "kernel_glm_attention_indexed_batch_group8") : - (use_group2 ? - ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group2_pipeline, - "kernel_glm_attention_indexed_batch_group2") : - ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_pipeline, - "kernel_glm_attention_indexed_batch"))); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_attention_indexed_batch_args args = { - .n_tokens = n_tokens, - .n_selected = n_selected, - .cache_cap = cache_cap, - .cache_f16 = cache_f16 ? 1u : 0u, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = value_dim, - .n_ctx_orig = n_ctx_orig, - .value_row_bytes = (uint32_t)value_row_bytes, - .value_type = value_weight_type, - .scale = 1.0f / sqrtf((float)qk_dim), - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .head_base = 0, - }; - const NSUInteger scratch_bytes = use_q2_group4 ? - q2_scratch_bytes : - (use_group8 ? - (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + - 8u * (NSUInteger)kv_lora_dim * sizeof(float)) : - ((use_group2 ? - (512u + 2u * (NSUInteger)n_selected + 2u * (NSUInteger)kv_lora_dim) : - (256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim)) * sizeof(float))); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; - [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; - [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; - if (use_q2_group4) { - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 3u) / 4u, - ((NSUInteger)n_tokens + 1u) / 2u, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; - } else if (use_group8) { - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 7u) / 8u, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - } else if (use_group2) { - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 1u) / 2u, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } else { - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - } - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_indexed_batch_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, - 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) { - return ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - model_map, - model_size, - value_weight_offset, - DS4_METAL_TENSOR_Q8_0, - 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); -} - -int ds4_gpu_sort_i32_rows_asc_tensor( - ds4_gpu_tensor *dst, - const ds4_gpu_tensor *src, - uint32_t row_width, - uint32_t n_rows) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!dst || !src || row_width == 0 || n_rows == 0 || - (row_width & (row_width - 1u)) != 0) { - return 0; - } - - @autoreleasepool { - id pipeline = - ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, - "kernel_dsv4_sort_i32_rows_asc"); - if (!pipeline) return 0; - - const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t); - id srcbuf = ds4_gpu_tensor_buffer(src); - id dstbuf = ds4_gpu_tensor_buffer(dst); - if (!srcbuf || !dstbuf || - ds4_gpu_tensor_bytes(src) < bytes || - ds4_gpu_tensor_bytes(dst) < bytes) { - fprintf(stderr, "ds4: Metal row sort received undersized buffers\n"); - return 0; - } - - NSUInteger threads = (NSUInteger)row_width; - const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; - if (max_threads != 0 && threads > max_threads) threads = max_threads; - if (threads == 0) return 0; - - const NSUInteger scratch_bytes = (NSUInteger)row_width * sizeof(int32_t); - const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; - if (max_tg_mem != 0 && scratch_bytes > max_tg_mem) { - fprintf(stderr, "ds4: Metal row sort scratch exceeds threadgroup memory limit\n"); - return 0; - } - - ds4_gpu_dsv4_topk_mask_args args = { - .ne00 = (int64_t)row_width, - .ne01 = (int64_t)n_rows, - .nb00 = sizeof(int32_t), - .nb01 = (uint64_t)row_width * sizeof(int32_t), - .ne0 = (int64_t)row_width, - .ne1 = (int64_t)n_rows, - .nb0 = sizeof(int32_t), - .nb1 = (uint64_t)row_width * sizeof(int32_t), - }; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1]; - [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:2]; - [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) - threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "sort i32 rows asc")) return 0; - } - - return 1; -} - -static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( - ds4_gpu_tensor *lora_out, - 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 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool selected_rows_valid) { - if (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected || - n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - n_head == 0 || kv_lora_dim == 0 || - qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || - qk_dim < qk_nope || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id lorabuf = ds4_gpu_tensor_buffer(lora_out); - id qbuf = ds4_gpu_tensor_buffer(q); - id lowbuf = ds4_gpu_tensor_buffer(qk_low); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); - const uint64_t lora_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t q_bytes = - (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - const uint64_t low_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = - (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = - (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - const uint64_t selected_bytes = - (uint64_t)n_tokens * n_selected * sizeof(uint32_t); - if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || - ds4_gpu_tensor_bytes(lora_out) < lora_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(qk_low) < low_bytes || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes) { - fprintf(stderr, "ds4: Metal GLM indexed batch attention-lora received undersized buffers\n"); - return 0; - } - - const bool use_vec_lora = - cache_f16 && kv_lora_dim == 512u && qk_rope == 64u; - const bool full_head_groups = (n_head % 8u) == 0u; - id pipeline = nil; - if (use_vec_lora && selected_rows_valid && full_head_groups) { - pipeline = ds4_gpu_hot_pipeline( - g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline, - "kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); - } else if (use_vec_lora && selected_rows_valid) { - pipeline = ds4_gpu_hot_pipeline( - g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline, - "kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); - } else if (use_vec_lora) { - pipeline = ds4_gpu_hot_pipeline( - g_glm_attention_indexed_batch_lora_group8_vec_pipeline, - "kernel_glm_attention_indexed_batch_lora_group8_vec"); - } else { - pipeline = ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, - "kernel_glm_attention_indexed_batch_group8"); - } - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_attention_indexed_batch_args args = { - .n_tokens = n_tokens, - .n_selected = n_selected, - .cache_cap = cache_cap, - .cache_f16 = cache_f16 ? 1u : 0u, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = kv_lora_dim, - .n_ctx_orig = n_ctx_orig, - .value_row_bytes = 0, - .value_type = 1u, - .scale = 1.0f / sqrtf((float)qk_dim), - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .head_base = 0, - }; - uint32_t head_count = n_head; - ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); - const NSUInteger scratch_bytes = use_vec_lora ? - (16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + - 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) : - (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + - 8u * (NSUInteger)kv_lora_dim * sizeof(float)); - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - if (use_vec_lora) { - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; - [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:6]; - } else { - [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; - [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7]; - } - [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention-lora")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( - ds4_gpu_tensor *lora_out, - 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, - uint32_t n_tokens, - uint32_t pos0, - 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (!g_initialized && !ds4_gpu_init()) return 0; - const uint32_t qk_dim = qk_nope + qk_rope; - if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || - n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || - pos0 > n_selected || n_tokens > n_selected - pos0 || - n_head == 0 || kv_lora_dim != 512u || - qk_nope == 0 || qk_rope != 64u || - qk_dim < qk_nope || !cache_f16 || - !isfinite(freq_base) || freq_base <= 0.0f || - !isfinite(freq_scale) || freq_scale <= 0.0f || - !isfinite(ext_factor) || !isfinite(attn_factor) || - !isfinite(beta_fast) || !isfinite(beta_slow)) { - return 0; - } - - @autoreleasepool { - id lorabuf = ds4_gpu_tensor_buffer(lora_out); - id qbuf = ds4_gpu_tensor_buffer(q); - id lowbuf = ds4_gpu_tensor_buffer(qk_low); - id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); - id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); - const uint64_t cache_elem_bytes = sizeof(uint16_t); - const uint64_t lora_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t q_bytes = - (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); - const uint64_t low_bytes = - (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); - const uint64_t kv_cache_bytes = - (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; - const uint64_t rope_cache_bytes = - (uint64_t)cache_cap * qk_rope * cache_elem_bytes; - if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || - ds4_gpu_tensor_bytes(lora_out) < lora_bytes || - ds4_gpu_tensor_bytes(q) < q_bytes || - ds4_gpu_tensor_bytes(qk_low) < low_bytes || - ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || - ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes) { - fprintf(stderr, "ds4: Metal GLM causal batch attention-lora received undersized buffers\n"); - return 0; - } - - const bool full_head_groups = (n_head % 8u) == 0u; - id pipeline = full_head_groups ? - ds4_gpu_hot_pipeline( - g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline, - "kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads") : - ds4_gpu_hot_pipeline( - g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline, - "kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_attention_indexed_batch_args args = { - .n_tokens = n_tokens, - .n_selected = n_selected, - .cache_cap = cache_cap, - .cache_f16 = 1u, - .n_head = n_head, - .kv_lora_dim = kv_lora_dim, - .qk_nope = qk_nope, - .qk_rope = qk_rope, - .value_dim = kv_lora_dim, - .n_ctx_orig = n_ctx_orig, - .value_row_bytes = 0, - .value_type = 1u, - .pos0 = pos0, - .scale = 1.0f / sqrtf((float)qk_dim), - .freq_base = freq_base, - .freq_scale = freq_scale, - .ext_factor = ext_factor, - .attn_factor = attn_factor, - .beta_fast = beta_fast, - .beta_slow = beta_slow, - .head_base = 0, - }; - uint32_t head_count = n_head; - ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); - const NSUInteger scratch_bytes = - 16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + - 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; - [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; - [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; - [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; - [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; - [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM causal indexed batch attention-lora")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_attention_indexed_batch_lora_tensor( - ds4_gpu_tensor *lora_out, - 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 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - selected, - n_tokens, - n_selected, - cache_cap, - cache_f16, - n_head, - kv_lora_dim, - qk_nope, - qk_rope, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow, - false); -} - -int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( - ds4_gpu_tensor *lora_out, - 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 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 n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - selected, - n_tokens, - n_selected, - cache_cap, - cache_f16, - n_head, - kv_lora_dim, - qk_nope, - qk_rope, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow, - true); -} - -int ds4_gpu_glm_router_select_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - const ds4_gpu_tensor *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || !weights || !probs || !logits || !model_map || - n_expert == 0 || n_expert > 256u || - n_expert_used == 0 || n_expert_used > n_expert) { - return 0; - } - - @autoreleasepool { - id logitsbuf = ds4_gpu_tensor_buffer(logits); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id probsbuf = ds4_gpu_tensor_buffer(probs); - if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || - ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || - ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t) || - ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || - ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { - fprintf(stderr, "ds4: Metal GLM router received undersized buffers\n"); - return 0; - } - - const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); - if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { - fprintf(stderr, "ds4: Metal GLM router bias range is outside the mapped model\n"); - return 0; - } - const bool exact_bias_view = - bias_bytes <= (1ull << 20) && - getenv("DS4_METAL_DISABLE_DECODE_ROUTER_BIAS_EXACT_VIEWS") == NULL; - uint64_t bias_inner = 0; - id biasbuf = exact_bias_view ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - bias_offset, - bias_bytes, - &bias_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - bias_offset, - bias_bytes, - &bias_inner); - if (!biasbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, - "kernel_glm_router_select_one"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_router_select_one_args args = { - .n_expert = n_expert, - .n_expert_used = n_expert_used, - .expert_weight_scale = expert_weight_scale, - .pad0 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; - [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; - [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0; - } - - return 1; -} - -int ds4_gpu_glm_router_select_batch_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - const ds4_gpu_tensor *logits, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_tokens) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || !weights || !probs || !logits || !model_map || - n_tokens == 0 || - n_expert == 0 || n_expert > 256u || - n_expert_used == 0 || n_expert_used > n_expert) { - return 0; - } - - @autoreleasepool { - id logitsbuf = ds4_gpu_tensor_buffer(logits); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id probsbuf = ds4_gpu_tensor_buffer(probs); - const uint64_t logits_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); - const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(int32_t); - const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(float); - const uint64_t probs_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); - if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || - ds4_gpu_tensor_bytes(logits) < logits_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes || - ds4_gpu_tensor_bytes(probs) < probs_bytes) { - fprintf(stderr, "ds4: Metal GLM batch router received undersized buffers\n"); - return 0; - } - - const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); - if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { - fprintf(stderr, "ds4: Metal GLM batch router bias range is outside the mapped model\n"); - return 0; - } - uint64_t bias_inner = 0; - id biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, - bias_offset, bias_bytes, - &bias_inner); - if (!biasbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, - "kernel_glm_router_select_one"); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - ds4_gpu_glm_router_select_one_args args = { - .n_expert = n_expert, - .n_expert_used = n_expert_used, - .expert_weight_scale = expert_weight_scale, - .pad0 = 0, - }; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; - [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; - [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0; - } - - return 1; -} - -static bool ds4_gpu_glm_gate_pair_type_supported( - uint32_t gate_type, - uint32_t up_type) { - return gate_type == up_type && - (gate_type == DS4_METAL_TENSOR_Q2_K || - gate_type == DS4_METAL_TENSOR_Q4_K || - gate_type == DS4_METAL_TENSOR_Q5_K); -} - -static bool ds4_gpu_glm_down_type_supported(uint32_t down_type) { - return down_type == DS4_METAL_TENSOR_Q2_K || - down_type == DS4_METAL_TENSOR_Q4_K || - down_type == DS4_METAL_TENSOR_Q5_K || - down_type == DS4_METAL_TENSOR_Q6_K; -} - -int ds4_gpu_glm_routed_moe_one_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - bool force_resident) { - if (!g_initialized && !ds4_gpu_init()) return 0; - /* TP sharding: only the owned contiguous expert range is mapped, - * so bind from the owned base, validate only its bytes, and tell the - * kernels the first expert id present at that base. */ - uint32_t first_expert = 0; - uint32_t n_bind_expert = 0; - ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); - const int32_t tp_expert_base_host = (int32_t)first_expert; - gate_offset += (uint64_t)first_expert * gate_expert_bytes; - up_offset += (uint64_t)first_expert * up_expert_bytes; - down_offset += (uint64_t)first_expert * down_expert_bytes; - - if (!out || !mid || !model_map || !selected || !weights || !x || - n_total_expert == 0 || n_expert == 0 || n_expert > 256u || - n_expert > n_total_expert || - expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || - gate_expert_bytes == 0 || gate_row_bytes == 0 || - up_expert_bytes == 0 || up_row_bytes == 0 || - down_expert_bytes == 0 || down_row_bytes == 0 || - (expert_in_dim % 256u) != 0 || - (expert_mid_dim % 256u) != 0 || - !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || - !ds4_gpu_glm_down_type_supported(down_type)) { - return 0; - } - - if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || - (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || - (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || - (uint64_t)out_dim > UINT64_MAX / down_row_bytes) { - fprintf(stderr, "ds4: Metal GLM routed MoE tensor byte size overflow\n"); - return 0; - } - - const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; - const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; - if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || - up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || - down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { - fprintf(stderr, "ds4: Metal GLM routed MoE received inconsistent expert strides\n"); - return 0; - } - if (layer_index == 3u) { - ds4_gpu_stream_expert_cache_note_decode_token(); - } - if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset || - up_offset > model_size || up_tensor_bytes > model_size - up_offset || - down_offset > model_size || down_tensor_bytes > model_size - down_offset) { - fprintf(stderr, "ds4: Metal GLM routed MoE tensor range is outside the mapped model\n"); - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); - const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); - if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int32_t) || - ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { - fprintf(stderr, "ds4: Metal GLM routed MoE received undersized activation buffers\n"); - return 0; - } - - const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; - const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; - const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; - const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; - const BOOL down_simd_q4 = down_scalar_q4; - const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; - const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; - const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; - const BOOL stream_addr_q2 = - gate_pair_q2 && down_scalar_q2 && - g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && - g_glm_q2_k_addr_down_f32_pipeline != nil; - const BOOL stream_addr_q4 = - !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && - g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && - g_glm_q4_k_addr_down_f32_pipeline != nil; - BOOL use_stream_expert_addr_table = - g_ssd_streaming_mode && - !force_resident && - getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL && - (stream_addr_q2 || stream_addr_q4) && - layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && - n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && - n_expert <= 8u && - ds4_gpu_stream_expert_cache_configured_budget() >= n_expert && - ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) && - ds4_gpu_stream_expert_cache_effective_cap(layer_index, - n_total_expert, - n_expert) != 0; - int32_t stream_selected_ids[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - ds4_gpu_stream_expert_cache_entry *stream_entries[8] = { - NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL - }; - uint64_t stream_gate_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint64_t stream_up_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint64_t stream_down_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; - uint32_t stream_missing_mask = 0; - uint32_t stream_entry_count = 0; - uint32_t stream_resident_mask = 0; - BOOL use_stream_split_deferred = false; - id stream_gate_addr_buf = nil; - id stream_up_addr_buf = nil; - id stream_down_addr_buf = nil; - const int glm_stream_timing = - ds4_gpu_stream_expert_timing_summary_enabled(); - - if (use_stream_expert_addr_table) { - const int had_batch = g_batch_cb != nil; - int stream_ok = 1; - const int have_prefetched_selected = - ds4_gpu_glm_stream_selected_prefetch_take(model_map, - model_size, - layer_index, - n_total_expert, - n_expert, - gate_offset, - up_offset, - down_offset, - gate_expert_bytes, - down_expert_bytes, - stream_selected_ids); - if (have_prefetched_selected) { - if (had_batch && g_batch_has_work && - g_stream_expert_pending_load.active && - ds4_gpu_flush_commands() == 0) { - stream_ok = 0; - } - } else { - if (had_batch && ds4_gpu_end_commands() == 0) { - stream_ok = 0; - } - if (stream_ok && - ds4_gpu_tensor_read(selected, - 0, - stream_selected_ids, - (uint64_t)n_expert * sizeof(stream_selected_ids[0])) == 0) { - stream_ok = 0; - } - } - for (uint32_t i = 0; stream_ok && i < n_expert; i++) { - if (stream_selected_ids[i] < 0 || - (uint32_t)stream_selected_ids[i] >= n_total_expert) { - fprintf(stderr, - "ds4: Metal GLM routed MoE selected expert id %d is outside 0..%u\n", - stream_selected_ids[i], - n_total_expert); - stream_ok = 0; - } - } - if (stream_ok) { - ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, - stream_selected_ids, - n_expert); - if (!ds4_gpu_moe_selected_hotlist_record(layer_index, - stream_selected_ids, - n_expert, - n_total_expert)) { - stream_ok = 0; - } - } - if (stream_ok) { - g_glm_stream_expert_addr_table_building++; - for (uint32_t i = 0; stream_ok && i < n_expert; i++) { - const uint64_t expert_id = (uint64_t)(uint32_t)stream_selected_ids[i]; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); - stream_ok = 0; - break; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); - stream_ok = 0; - break; - } - stream_gate_abs_offsets[i] = gate_offset + gate_rel; - stream_up_abs_offsets[i] = up_offset + gate_rel; - stream_down_abs_offsets[i] = down_offset + down_rel; - stream_entries[i] = - ds4_gpu_stream_expert_cache_peek(model_map, - model_size, - layer_index, - (uint32_t)stream_selected_ids[i], - n_total_expert, - n_expert, - stream_gate_abs_offsets[i], - stream_up_abs_offsets[i], - stream_down_abs_offsets[i], - gate_expert_bytes, - down_expert_bytes); - if (!stream_entries[i]) { - stream_missing_mask |= 1u << i; - } else { - stream_resident_mask |= 1u << i; - } - } - if (stream_ok && glm_stream_timing) { - ds4_gpu_stream_expert_timing_note_cache_class( - stream_resident_mask, - stream_missing_mask); - } - use_stream_split_deferred = - stream_ok && - getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_SPLIT") == NULL && - stream_missing_mask != 0 && - stream_resident_mask != 0 && - ds4_gpu_stream_expert_split_worthwhile(stream_resident_mask, - stream_missing_mask) && - ds4_gpu_stream_expert_split_ready() && - ((gate_pair_q2 && - g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline != nil) || - (!gate_pair_q2 && !gate_pair_q5 && - g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline != nil)); - if (use_stream_split_deferred) { - const ds4_gpu_stream_expert_table table = { - .model_map = model_map, - .model_size = model_size, - .layer = layer_index, - .n_total_expert = n_total_expert, - .gate_offset = gate_offset, - .up_offset = up_offset, - .down_offset = down_offset, - .gate_expert_bytes = gate_expert_bytes, - .down_expert_bytes = down_expert_bytes, - }; - if (!ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - stream_selected_ids, - n_expert)) { - stream_ok = 0; - } - } - if (stream_ok && stream_missing_mask != 0 && - !use_stream_split_deferred && - !ds4_gpu_stream_expert_cache_load_selected_missing( - model_map, - model_size, - layer_index, - stream_selected_ids, - n_total_expert, - n_expert, - stream_gate_abs_offsets, - stream_up_abs_offsets, - stream_down_abs_offsets, - gate_expert_bytes, - down_expert_bytes, - stream_missing_mask, - stream_entries)) { - fprintf(stderr, - "ds4: Metal GLM streaming expert cache failed to load " - "layer=%u missing=0x%x budget=%u\n", - layer_index, - stream_missing_mask, - ds4_gpu_stream_expert_cache_configured_budget()); - stream_ok = 0; - } - for (uint32_t i = 0; stream_ok && i < n_expert; i++) { - ds4_gpu_stream_expert_cache_entry *entry = stream_entries[i]; - if (!entry) { - if (use_stream_split_deferred && - (stream_missing_mask & (1u << i)) != 0) { - continue; - } - stream_ok = 0; - break; - } - if (!ds4_gpu_stream_expert_cache_set_addr_slot( - layer_index, - (uint32_t)stream_selected_ids[i], - entry->gate_buffer, - entry->gate_inner, - entry->up_buffer, - entry->up_inner, - entry->down_buffer, - entry->down_inner)) { - stream_ok = 0; - break; - } - stream_entry_count++; - } - if (stream_ok && - !ds4_gpu_stream_expert_cache_addr_buffers(layer_index, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf)) { - stream_ok = 0; - } - g_glm_stream_expert_addr_table_building--; - } - if (!have_prefetched_selected && had_batch && - ds4_gpu_begin_commands() == 0) { - stream_ok = 0; - } - if (!stream_ok) return 0; - ds4_gpu_stream_expert_cache_prune_layer(layer_index, - n_total_expert, - n_expert, - stream_selected_ids, - n_expert); - ds4_gpu_stream_expert_cache_prune_global(layer_index, - stream_selected_ids, - n_expert); - } - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gatebuf = nil; - id upbuf = nil; - id downbuf = nil; - if (!use_stream_expert_addr_table) { - gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, - gate_offset, gate_tensor_bytes, - &gate_inner); - upbuf = ds4_gpu_wrap_model_range(model_map, model_size, - up_offset, up_tensor_bytes, - &up_inner); - downbuf = ds4_gpu_wrap_model_range(model_map, model_size, - down_offset, down_tensor_bytes, - &down_inner); - if (!gatebuf || !upbuf || !downbuf) return 0; - } - - id pair_pipeline = - (use_stream_split_deferred ? - (gate_pair_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline, - "kernel_glm_q2_K_addr_pair_swiglu2_f32_masked") : - ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline, - "kernel_glm_q4_K_addr_pair_swiglu_f32_masked")) : - use_stream_expert_addr_table ? - (gate_pair_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, - "kernel_glm_q2_K_addr_pair_swiglu2_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, - "kernel_glm_q4_K_addr_pair_swiglu_f32")) : - gate_pair_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, - "kernel_glm_q2_K_pair_swiglu_f32") : - gate_pair_q5 ? - ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, - "kernel_glm_q5_K_pair_swiglu_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, - "kernel_glm_q4_K_pair_swiglu2_f32")); - id down_pipeline = - (use_stream_expert_addr_table ? - (down_scalar_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, - "kernel_glm_q2_K_addr_down_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, - "kernel_glm_q4_K_addr_down_f32")) : - down_scalar_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, - "kernel_glm_q2_K_down_f32") : - down_scalar_q4 ? - ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, - "kernel_glm_q4_K_down_f32") : - down_simd_q5 ? - ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, - "kernel_glm_q5_K_down_f32") : - ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, - "kernel_glm_q6_K_down_f32")); - if (!pair_pipeline || !down_pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const bool glm_moe_stage_profile = - g_batch_cb != nil && - ds4_gpu_stage_profile_enabled_for_layer( - "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE", - "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE_LAYER", - layer_index); - const char *glm_moe_stage_filter = - getenv("DS4_METAL_GLM_MOE_STAGE_PROFILE_FILTER"); - const char *glm_pair_path = - use_stream_expert_addr_table ? - (gate_pair_q2 ? "q2_stream_addr_swiglu" : - "q4_stream_addr_swiglu") : - gate_pair_q2 ? "q2_scalar_swiglu" : - gate_pair_q5 ? "q5_pair_simd_swiglu" : "q4_pair2_simd_swiglu"; - const char *glm_down_path = - use_stream_expert_addr_table ? - (down_scalar_q2 ? "q2_stream_addr_down" : - "q4_stream_addr_down_simd") : - down_scalar_q2 ? "q2_down_simd" : - down_scalar_q4 ? "q4_down_simd" : - down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; - double glm_moe_stage_t0 = 0.0; - if (glm_moe_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - glm_moe_stage_t0 = ds4_gpu_now_ms(); - } - int ok = 1; -#define DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE(name) do { \ - if (ok && glm_moe_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ - strstr(stage_name, glm_moe_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal GLM routed MoE one stage layer=%u tokens=1 experts=%u " \ - "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ - layer_index, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - glm_pair_path, glm_down_path, \ - stage_name, now_ms - glm_moe_stage_t0); \ - } \ - glm_moe_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - - ds4_gpu_glm_routed_moe_args args = { - .tp_rank = g_tp_split_rank, - .tp_world = g_tp_split_world, - .tp_expert_base = tp_expert_base_host, - .in_dim = expert_in_dim, - .mid_dim = expert_mid_dim, - .out_dim = out_dim, - .n_total_expert = n_total_expert, - .n_expert_used = n_expert, - .n_tokens = 1, - .mid_token_stride = n_expert * expert_mid_dim, - .down_type = down_type, - .gate_expert_bytes = gate_expert_bytes, - .gate_row_bytes = gate_row_bytes, - .up_expert_bytes = up_expert_bytes, - .up_row_bytes = up_row_bytes, - .down_expert_bytes = down_expert_bytes, - .down_row_bytes = down_row_bytes, - }; - const NSUInteger pair_x_groups = - gate_pair_q2 ? (use_stream_expert_addr_table ? - (NSUInteger)((expert_mid_dim + 1u) / 2u) : - (NSUInteger)((expert_mid_dim + 7u) / 8u)) : - gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : - use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : - (NSUInteger)((expert_mid_dim + 1u) / 2u); - const NSUInteger pair_threadgroup_bytes = 0u; - const NSUInteger pair_threads = 64u; - const NSUInteger down_x_groups = - down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : - down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : - down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : - down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : - (NSUInteger)out_dim; - const NSUInteger down_threadgroup_bytes = - (down_scalar_q2 || down_simd) ? 0u : 256u * sizeof(float); - const NSUInteger down_threads = - (down_scalar_q2 || down_simd) ? 64u : 256u; - if (use_stream_expert_addr_table && - !ds4_gpu_stream_expert_cache_mark_entries_inflight( - stream_entries, - use_stream_split_deferred ? n_expert : stream_entry_count, - use_stream_split_deferred ? stream_resident_mask : 0)) { - return 0; - } - - const int glm_stream_split_timing = - use_stream_split_deferred && glm_stream_timing; - double glm_stream_split_t0 = - glm_stream_split_timing ? ds4_gpu_now_ms() : 0.0; - double glm_stream_split_resident_ms = 0.0; - double glm_stream_split_missing_load_ms = 0.0; - double glm_stream_split_missing_wait_ms = 0.0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pair_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - if (use_stream_split_deferred) { - [enc setBytes:&stream_resident_mask length:sizeof(stream_resident_mask) atIndex:1]; - [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; - [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; - } else { - [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner - atIndex:1]; - [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner - atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; - } - if (use_stream_expert_addr_table) { - const uint32_t use_count = - use_stream_split_deferred ? n_expert : stream_entry_count; - const uint32_t use_mask = - use_stream_split_deferred ? stream_resident_mask : 0; - for (uint32_t i = 0; i < use_count; i++) { - if (use_mask != 0 && (use_mask & (1u << i)) == 0) continue; - [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; - [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; - } - } - if (pair_threadgroup_bytes != 0u) { - [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, - (NSUInteger)n_expert, - 1) - threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("pair"); - - if (ok && use_stream_split_deferred) { - id resident_cb = nil; - const int resident_owned = owned; - if (owned) { - resident_cb = cb; - [resident_cb commit]; - cb = nil; - } else { - ok = ds4_gpu_flush_commands(); - } - if (glm_stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - glm_stream_split_resident_ms = now_ms - glm_stream_split_t0; - glm_stream_split_t0 = now_ms; - } - if (ok) { - ok = ds4_gpu_stream_expert_pending_load_finish(stream_entries); - } - if (glm_stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - glm_stream_split_missing_load_ms = now_ms - glm_stream_split_t0; - glm_stream_split_t0 = now_ms; - } - if (ok) { - ds4_gpu_stream_expert_cache_prune_layer(layer_index, - n_total_expert, - n_expert, - stream_selected_ids, - n_expert); - ds4_gpu_stream_expert_cache_prune_global(layer_index, - stream_selected_ids, - n_expert); - ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf); - } - if (ok) { - if (resident_owned) { - ok = ds4_gpu_wait_command_buffer( - resident_cb, - "GLM streaming split resident pair"); - ds4_gpu_stream_expert_cache_note_owned_completed(); - } else { - ok = ds4_gpu_wait_pending_command_buffers( - "GLM streaming split resident pair"); - } - } - if (glm_stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - glm_stream_split_missing_wait_ms = now_ms - glm_stream_split_t0; - glm_stream_split_t0 = now_ms; - ds4_gpu_stream_expert_timing_note_split( - stream_resident_mask, - stream_missing_mask, - glm_stream_split_resident_ms, - glm_stream_split_missing_load_ms + - glm_stream_split_missing_wait_ms); - ds4_gpu_stream_expert_timing_note_split_missing_detail( - glm_stream_split_missing_load_ms, - 0.0, - 0.0, - 0.0, - glm_stream_split_missing_wait_ms); - } - if (ok && - !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_entries, - n_expert, - stream_missing_mask)) { - ok = 0; - } - if (ok) { - cb = ds4_gpu_command_buffer(&owned); - if (!cb) ok = 0; - } - if (ok) { - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pair_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBytes:&stream_missing_mask length:sizeof(stream_missing_mask) atIndex:1]; - [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; - [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; - for (uint32_t i = 0; i < n_expert; i++) { - if ((stream_missing_mask & (1u << i)) == 0) continue; - [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; - [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; - } - if (pair_threadgroup_bytes != 0u) { - [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, - (NSUInteger)n_expert, - 1) - threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - } - } - - if (!ok) return 0; - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:down_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner - atIndex:1]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - if (use_stream_expert_addr_table) { - const uint32_t use_count = - use_stream_split_deferred ? n_expert : stream_entry_count; - for (uint32_t i = 0; i < use_count; i++) { - [enc useResource:stream_entries[i]->down_buffer usage:MTLResourceUsageRead]; - } - } - if (down_threadgroup_bytes != 0u) { - [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, 1, 1) - threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("down"); - - if (!ok) return 0; - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed MoE")) return 0; -#undef DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE - } - - return 1; -} - -static bool ds4_gpu_glm_grouped_moe_fast_default(void) { - return !g_quality_mode; -} - -static bool ds4_gpu_glm_routed_moe_batch_grouped_available( - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint32_t n_expert, - uint32_t n_tokens) { - if (n_tokens < 32u || - !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || - !ds4_gpu_glm_down_type_supported(down_type) || - ds4_gpu_mul_mm_id_map0_name(n_expert) == NULL) { - return false; - } - const bool fast_default = ds4_gpu_glm_grouped_moe_fast_default(); - if (!fast_default) { - return false; - } - if (n_tokens < 96u) return false; - - return ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)) != nil && - ds4_gpu_routed_mm_pipeline(gate_type) != nil && - ds4_gpu_routed_mm_pipeline(up_type) != nil && - ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) != nil; -} - -static bool ds4_gpu_glm_grouped_moe_layer_enabled(uint32_t layer_index) { - (void)layer_index; - return true; -} - -static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens) { - if (!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, - up_type, - down_type, - n_expert, - n_tokens)) { - return 0; - } - if (n_expert > UINT32_MAX / n_tokens || - (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { - return 0; - } - - const uint32_t pair_rows = n_tokens * n_expert; - if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || - (uint64_t)pair_rows > UINT64_MAX / out_dim || - (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || - (uint64_t)n_tokens > UINT64_MAX / out_dim) { - return 0; - } - - const bool mid_f16 = true; - const NSUInteger mm_id_threadgroup_bytes = 8192u; - const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; - const uint64_t down_values = (uint64_t)pair_rows * out_dim; - const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; - const uint64_t out_values = (uint64_t)n_tokens * out_dim; - if (compact_mid_values > UINT64_MAX / sizeof(float) || - compact_mid_values > UINT64_MAX / (mid_f16 ? sizeof(uint16_t) : sizeof(float)) || - down_values > UINT64_MAX / sizeof(float) || - x_values > UINT64_MAX / sizeof(float) || - out_values > UINT64_MAX / sizeof(float)) { - return 0; - } - - const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); - const uint64_t mid_bytes = compact_mid_values * (mid_f16 ? sizeof(uint16_t) : sizeof(float)); - const uint64_t down_scratch_bytes = down_values * sizeof(float); - const uint64_t x_bytes = x_values * sizeof(float); - const uint64_t out_bytes = out_values * sizeof(float); - const uint64_t selected_values = (uint64_t)n_tokens * n_expert; - const uint64_t selected_bytes = selected_values * sizeof(int32_t); - const uint64_t weights_bytes = selected_values * sizeof(float); - if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || - gate_scratch_bytes > NSUIntegerMax || - gate_scratch_bytes * 2ull > NSUIntegerMax || - down_scratch_bytes > NSUIntegerMax) { - return 0; - } - - uint32_t first_expert = 0; - uint32_t n_bind_expert = 0; - ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); - gate_offset += (uint64_t)first_expert * gate_expert_bytes; - up_offset += (uint64_t)first_expert * up_expert_bytes; - down_offset += (uint64_t)first_expert * down_expert_bytes; - const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; - const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes) { - fprintf(stderr, "ds4: Metal GLM grouped routed MoE received undersized activation buffers\n"); - return 0; - } - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, - &g_moe_gate_scratch_bytes, - (NSUInteger)(gate_scratch_bytes * 2ull), - "ds4_glm_moe_gate_up_scratch")) { - return 0; - } - if (n_expert > 1 && - !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, - &g_moe_down_scratch_bytes, - (NSUInteger)down_scratch_bytes, - "ds4_glm_moe_down_scratch")) { - return 0; - } - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, - gate_offset, gate_tensor_bytes, - &gate_inner); - id upbuf = ds4_gpu_wrap_model_range(model_map, model_size, - up_offset, up_tensor_bytes, - &up_inner); - id downbuf = ds4_gpu_wrap_model_range(model_map, model_size, - down_offset, down_tensor_bytes, - &down_inner); - if (!gatebuf || !upbuf || !downbuf) return 0; - - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); - id gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); - id up_pipeline = ds4_gpu_routed_mm_pipeline(up_type); - id down_pipeline = - ds4_gpu_routed_mm_f16_rhs_pipeline(down_type); - if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { - return 0; - } - - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args gate_args = - ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - gate_row_bytes, gate_expert_bytes, - 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args up_args = - ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - up_row_bytes, up_expert_bytes, - 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args down_args = - ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, - down_row_bytes, down_expert_bytes, - n_expert, n_expert, n_tokens, - mid_f16 ? sizeof(uint16_t) : sizeof(float)); - gate_args.tp_rank = g_tp_split_rank; - gate_args.tp_world = g_tp_split_world; - gate_args.tp_expert_base = (int32_t)first_expert; - up_args.tp_rank = g_tp_split_rank; - up_args.tp_world = g_tp_split_world; - up_args.tp_expert_base = (int32_t)first_expert; - down_args.tp_rank = g_tp_split_rank; - down_args.tp_world = g_tp_split_world; - down_args.tp_expert_base = (int32_t)first_expert; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const bool glm_moe_stage_profile = false; - const char *glm_moe_stage_filter = NULL; - double glm_moe_stage_t0 = 0.0; - if (glm_moe_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - glm_moe_stage_t0 = ds4_gpu_now_ms(); - } - - int ok = 1; -#define DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE(name) do { \ - if (ok && glm_moe_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ - strstr(stage_name, glm_moe_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal GLM grouped routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ - "gate=%s down=%s mid=%s %s=%.3f ms\n", \ - layer_index, n_tokens, pair_rows, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - mid_f16 ? "f16" : "f32", \ - stage_name, now_ms - glm_moe_stage_t0); \ - } \ - glm_moe_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - - ok = ds4_gpu_encode_mul_mm_id_map(cb, - map_pipeline, - &map_args, - &gate_args, - selectedbuf, - ds4_gpu_tensor_offset(selected)); - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("map"); - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - gate_pipeline, - &gate_args, - gatebuf, - (NSUInteger)gate_inner, - xbuf, - ds4_gpu_tensor_offset(x), - g_moe_gate_scratch_buffer, - 0, - mm_id_threadgroup_bytes); - } - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("gate"); - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - up_pipeline, - &up_args, - upbuf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - g_moe_gate_scratch_buffer, - (NSUInteger)gate_scratch_bytes, - mm_id_threadgroup_bytes); - } - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("up"); - if (ok) { - ok = ds4_gpu_encode_moe_swiglu_weight(cb, - g_moe_gate_scratch_buffer, - 0, - g_moe_gate_scratch_buffer, - (NSUInteger)gate_scratch_bytes, - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - expert_mid_dim, - pair_rows, - 0.0f, - mid_f16); - } - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight"); - - id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; - NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - down_pipeline, - &down_args, - downbuf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - mm_id_threadgroup_bytes); - } - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("down"); - if (ok && n_expert > 1) { - ok = ds4_gpu_encode_moe_sum_experts(cb, - down_dst, - down_dst_off, - outbuf, - ds4_gpu_tensor_offset(out), - out_dim, - n_expert, - n_tokens); - } - DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("sum"); - if (!ok) return 0; - - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM grouped routed batch MoE")) { - return 0; - } -#undef DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE - } - - return 1; -} - -static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - id gate_addrs, - id up_addrs, - id down_addrs, - id overflow_gate, - id overflow_up, - id overflow_down, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - ds4_gpu_stream_expert_cache_entry * const *resources, - uint32_t resource_count) { - if (n_expert > UINT32_MAX / n_tokens || - !gate_addrs || !up_addrs || !down_addrs || - !ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, - up_type, - down_type, - n_expert, - n_tokens)) { - return 0; - } - - const uint32_t pair_rows = n_tokens * n_expert; - if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || - (uint64_t)pair_rows > UINT64_MAX / out_dim || - (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || - (uint64_t)n_tokens > UINT64_MAX / out_dim) { - return 0; - } - - const bool mid_f16 = true; - const NSUInteger mm_id_threadgroup_bytes = 8192u; - const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; - const uint64_t down_values = (uint64_t)pair_rows * out_dim; - const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; - const uint64_t out_values = (uint64_t)n_tokens * out_dim; - if (compact_mid_values > UINT64_MAX / sizeof(float) || - compact_mid_values > UINT64_MAX / sizeof(uint16_t) || - down_values > UINT64_MAX / sizeof(float) || - x_values > UINT64_MAX / sizeof(float) || - out_values > UINT64_MAX / sizeof(float)) { - return 0; - } - - const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); - const uint64_t mid_bytes = compact_mid_values * sizeof(uint16_t); - const uint64_t down_scratch_bytes = down_values * sizeof(float); - const uint64_t x_bytes = x_values * sizeof(float); - const uint64_t out_bytes = out_values * sizeof(float); - const uint64_t selected_values = (uint64_t)n_tokens * n_expert; - const uint64_t selected_bytes = selected_values * sizeof(int32_t); - const uint64_t weights_bytes = selected_values * sizeof(float); - if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || - gate_scratch_bytes > NSUIntegerMax || - gate_scratch_bytes * 2ull > NSUIntegerMax || - down_scratch_bytes > NSUIntegerMax) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes) { - fprintf(stderr, "ds4: Metal GLM grouped-address routed MoE received undersized activation buffers\n"); - return 0; - } - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, - &g_moe_gate_scratch_bytes, - (NSUInteger)(gate_scratch_bytes * 2ull), - "ds4_glm_moe_gate_up_scratch")) { - return 0; - } - if (n_expert > 1 && - !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, - &g_moe_down_scratch_bytes, - (NSUInteger)down_scratch_bytes, - "ds4_glm_moe_down_scratch")) { - return 0; - } - - id map_pipeline = - ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); - id gate_pipeline = - ds4_gpu_routed_mm_addr_pipeline(gate_type); - id up_pipeline = - ds4_gpu_routed_mm_addr_pipeline(up_type); - id down_pipeline = - ds4_gpu_routed_mm_addr_f16_rhs_pipeline(down_type); - if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { - return 0; - } - - ds4_gpu_mul_mm_id_map_args map_args = - ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args gate_args = - ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - gate_row_bytes, gate_expert_bytes, - 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args up_args = - ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - up_row_bytes, up_expert_bytes, - 1, n_expert, n_tokens); - ds4_gpu_mul_mm_id_args down_args = - ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, - down_row_bytes, down_expert_bytes, - n_expert, n_expert, n_tokens, - sizeof(uint16_t)); - if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(resources, - resource_count, - 0)) { - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const bool glm_moe_stage_profile = false; - const char *glm_moe_stage_filter = NULL; - double glm_moe_stage_t0 = 0.0; - if (glm_moe_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - glm_moe_stage_t0 = ds4_gpu_now_ms(); - } - - int ok = 1; -#define DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE(name) do { \ - if (ok && glm_moe_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ - strstr(stage_name, glm_moe_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal GLM grouped-address routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ - "gate=%s down=%s mid=f16 %s=%.3f ms\n", \ - layer_index, n_tokens, pair_rows, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - stage_name, now_ms - glm_moe_stage_t0); \ - } \ - glm_moe_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - - ok = ds4_gpu_encode_mul_mm_id_map(cb, - map_pipeline, - &map_args, - &gate_args, - selectedbuf, - ds4_gpu_tensor_offset(selected)); - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("map"); - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, - gate_pipeline, - &gate_args, - gate_addrs, - xbuf, - ds4_gpu_tensor_offset(x), - g_moe_gate_scratch_buffer, - 0, - mm_id_threadgroup_bytes, - resources, - resource_count, - 0, - overflow_gate); - } - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("gate"); - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, - up_pipeline, - &up_args, - up_addrs, - xbuf, - ds4_gpu_tensor_offset(x), - g_moe_gate_scratch_buffer, - (NSUInteger)gate_scratch_bytes, - mm_id_threadgroup_bytes, - resources, - resource_count, - 1, - overflow_up); - } - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("up"); - if (ok) { - ok = ds4_gpu_encode_moe_swiglu_weight(cb, - g_moe_gate_scratch_buffer, - 0, - g_moe_gate_scratch_buffer, - (NSUInteger)gate_scratch_bytes, - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - expert_mid_dim, - pair_rows, - 0.0f, - mid_f16); - } - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight"); - - id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; - NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; - if (ok) { - ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, - down_pipeline, - &down_args, - down_addrs, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - mm_id_threadgroup_bytes, - resources, - resource_count, - 2, - overflow_down); - } - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("down"); - if (ok && n_expert > 1) { - ok = ds4_gpu_encode_moe_sum_experts(cb, - down_dst, - down_dst_off, - outbuf, - ds4_gpu_tensor_offset(out), - out_dim, - n_expert, - n_tokens); - } - DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("sum"); - if (!ok) return 0; - - if (!ds4_gpu_finish_command_buffer(cb, owned, - "GLM grouped-address routed batch MoE")) { - return 0; - } -#undef DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE - } - - return 1; -} - -static int ds4_gpu_glm_routed_moe_batch_tensor_impl( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride, - bool allow_grouped, - bool force_scalar_q4_pair) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !mid || !model_map || !selected || !weights || !x || - n_tokens == 0 || - n_total_expert == 0 || n_expert == 0 || n_expert > 256u || - n_expert > n_total_expert || - expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || - gate_expert_bytes == 0 || gate_row_bytes == 0 || - up_expert_bytes == 0 || up_row_bytes == 0 || - down_expert_bytes == 0 || down_row_bytes == 0 || - (expert_in_dim % 256u) != 0 || - (expert_mid_dim % 256u) != 0 || - !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || - !ds4_gpu_glm_down_type_supported(down_type)) { - return 0; - } - - const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim; - if ((uint64_t)mid_token_stride < per_token_mid || - (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || - (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || - (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || - (uint64_t)out_dim > UINT64_MAX / down_row_bytes || - (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || - (uint64_t)n_tokens > UINT64_MAX / out_dim) { - fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor byte size overflow\n"); - return 0; - } - - const uint64_t full_gate_tensor_bytes = - (uint64_t)n_total_expert * gate_expert_bytes; - const uint64_t full_up_tensor_bytes = - (uint64_t)n_total_expert * up_expert_bytes; - const uint64_t full_down_tensor_bytes = - (uint64_t)n_total_expert * down_expert_bytes; - if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || - up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || - down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { - fprintf(stderr, "ds4: Metal GLM routed batch MoE received inconsistent expert strides\n"); - return 0; - } - if (gate_offset > model_size || full_gate_tensor_bytes > model_size - gate_offset || - up_offset > model_size || full_up_tensor_bytes > model_size - up_offset || - down_offset > model_size || full_down_tensor_bytes > model_size - down_offset) { - fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor range is outside the mapped model\n"); - return 0; - } - - if (allow_grouped && - (!g_ssd_streaming_mode || - ds4_gpu_glm_streaming_prefill_full_layer_active()) && - ds4_gpu_glm_grouped_moe_layer_enabled(layer_index) && - ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, - up_type, - down_type, - n_expert, - n_tokens)) { - return ds4_gpu_glm_routed_moe_batch_grouped_tensor(out, - mid, - model_map, - model_size, - gate_offset, - up_offset, - down_offset, - gate_type, - up_type, - down_type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - selected, - weights, - n_total_expert, - n_expert, - layer_index, - x, - n_tokens); - } - - uint32_t first_expert = 0; - uint32_t n_bind_expert = 0; - ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); - gate_offset += (uint64_t)first_expert * gate_expert_bytes; - up_offset += (uint64_t)first_expert * up_expert_bytes; - down_offset += (uint64_t)first_expert * down_expert_bytes; - const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; - const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; - const uint64_t out_values = (uint64_t)n_tokens * out_dim; - const uint64_t mid_values = - (uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid; - const uint64_t selected_values = (uint64_t)n_tokens * n_expert; - if (x_values > UINT64_MAX / sizeof(float) || - out_values > UINT64_MAX / sizeof(float) || - mid_values > UINT64_MAX / sizeof(float) || - selected_values > UINT64_MAX / sizeof(int32_t)) { - fprintf(stderr, "ds4: Metal GLM routed batch MoE activation byte size overflow\n"); - return 0; - } - const uint64_t x_bytes = x_values * sizeof(float); - const uint64_t mid_bytes = mid_values * sizeof(float); - const uint64_t out_bytes = out_values * sizeof(float); - const uint64_t selected_bytes = selected_values * sizeof(int32_t); - const uint64_t weights_bytes = selected_values * sizeof(float); - if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes) { - fprintf(stderr, "ds4: Metal GLM routed batch MoE received undersized activation buffers\n"); - return 0; - } - - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gatebuf = nil; - id upbuf = nil; - id downbuf = nil; - id stream_gate_addr_buf = nil; - id stream_up_addr_buf = nil; - id stream_down_addr_buf = nil; - id stream_overflow_gate = nil; - id stream_overflow_up = nil; - id stream_overflow_down = nil; - ds4_gpu_stream_expert_cache_entry - *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; - uint32_t stream_resource_count = 0; - uint32_t stream_unique = 0; - - const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; - const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; - const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; - const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; - const BOOL down_simd_q4 = down_scalar_q4; - const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; - const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; - const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; - const BOOL stream_addr_q2 = - gate_pair_q2 && down_scalar_q2 && - g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && - g_glm_q2_k_addr_down_f32_pipeline != nil; - const BOOL stream_addr_q4 = - !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && - g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && - g_glm_q4_k_addr_down_f32_pipeline != nil; - BOOL use_stream_expert_addr_table = - g_ssd_streaming_mode && - !ds4_gpu_glm_streaming_prefill_full_layer_active() && - n_tokens > 1 && - (stream_addr_q2 || stream_addr_q4) && - layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && - n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && - n_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED && - ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes) && - ds4_gpu_stream_expert_cache_effective_cap(layer_index, - n_total_expert, - n_expert) != 0; - const BOOL use_stream_grouped_addr_table = - use_stream_expert_addr_table && - allow_grouped && - getenv("DS4_METAL_GLM_DISABLE_STREAMING_GROUPED_ADDR_PREFILL") == NULL && - ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, - up_type, - down_type, - n_expert, - n_tokens); - const BOOL enable_q4_pair4 = true; - const BOOL q4_scalar_pair = false; - const BOOL q4_pair2 = - !use_stream_expert_addr_table && - !gate_pair_q5 && !q4_scalar_pair && - (force_scalar_q4_pair || !enable_q4_pair4); - id pair_pipeline = - use_stream_expert_addr_table ? - (gate_pair_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, - "kernel_glm_q2_K_addr_pair_swiglu2_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, - "kernel_glm_q4_K_addr_pair_swiglu_f32")) : - gate_pair_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, - "kernel_glm_q2_K_pair_swiglu_f32") : - gate_pair_q5 ? - ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, - "kernel_glm_q5_K_pair_swiglu_f32") : - (q4_scalar_pair ? - ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu_f32_pipeline, - "kernel_glm_q4_K_pair_swiglu_f32") : - q4_pair2 ? - ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, - "kernel_glm_q4_K_pair_swiglu2_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu4_f32_pipeline, - "kernel_glm_q4_K_pair_swiglu4_f32")); - id down_pipeline = - use_stream_expert_addr_table ? - (down_scalar_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, - "kernel_glm_q2_K_addr_down_f32") : - ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, - "kernel_glm_q4_K_addr_down_f32")) : - down_scalar_q2 ? - ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, - "kernel_glm_q2_K_down_f32") : - down_scalar_q4 ? - ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, - "kernel_glm_q4_K_down_f32") : - down_simd_q5 ? - ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, - "kernel_glm_q5_K_down_f32") : - ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, - "kernel_glm_q6_K_down_f32"); - if (!pair_pipeline || !down_pipeline) return 0; - - if (use_stream_expert_addr_table) { - const int had_batch = g_batch_cb != nil; - if (had_batch && ds4_gpu_end_commands() == 0) { - return 0; - } - if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( - model_map, - model_size, - layer_index, - selected, - n_tokens, - n_total_expert, - n_expert, - gate_offset, - up_offset, - down_offset, - gate_expert_bytes, - down_expert_bytes, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf, - stream_resources, - &stream_resource_count, - &stream_unique, - &stream_overflow_gate, - &stream_overflow_up, - &stream_overflow_down)) { - return 0; - } - if (stream_unique == 0) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - if (had_batch && ds4_gpu_begin_commands() == 0) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - if (use_stream_grouped_addr_table) { - return ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( - out, - mid, - stream_gate_addr_buf, - stream_up_addr_buf, - stream_down_addr_buf, - stream_overflow_gate, - stream_overflow_up, - stream_overflow_down, - gate_type, - up_type, - down_type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - selected, - weights, - n_total_expert, - n_expert, - layer_index, - x, - n_tokens, - stream_resources, - stream_resource_count); - } - } else { - gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, - gate_offset, gate_tensor_bytes, - &gate_inner); - upbuf = ds4_gpu_wrap_model_range(model_map, model_size, - up_offset, up_tensor_bytes, - &up_inner); - downbuf = ds4_gpu_wrap_model_range(model_map, model_size, - down_offset, down_tensor_bytes, - &down_inner); - if (!gatebuf || !upbuf || !downbuf) return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - const bool glm_moe_stage_profile = false; - const char *glm_moe_stage_filter = NULL; - const char *glm_pair_path = use_stream_expert_addr_table ? - (gate_pair_q2 ? "q2_stream_addr_swiglu" : - "q4_stream_addr_swiglu") : - gate_pair_q2 ? "q2_scalar_swiglu" : - gate_pair_q5 ? "q5_pair_simd_swiglu" : - (q4_scalar_pair ? "q4_scalar_swiglu" : - (q4_pair2 ? "q4_pair2_simd_swiglu" : - "q4_pair4_simd_swiglu")); - const char *glm_down_path = - use_stream_expert_addr_table ? - (down_scalar_q2 ? "q2_stream_addr_down" : "q4_stream_addr_down_simd") : - down_scalar_q2 ? "q2_down_scalar" : - down_scalar_q4 ? "q4_down_simd" : - down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; - double glm_moe_stage_t0 = 0.0; - if (glm_moe_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - glm_moe_stage_t0 = ds4_gpu_now_ms(); - } - int ok = 1; -#define DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE(name) do { \ - if (ok && glm_moe_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ - strstr(stage_name, glm_moe_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal GLM routed MoE batch stage layer=%u tokens=%u experts=%u " \ - "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ - layer_index, n_tokens, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - glm_pair_path, glm_down_path, \ - stage_name, now_ms - glm_moe_stage_t0); \ - } \ - glm_moe_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - - ds4_gpu_glm_routed_moe_args args = { - .tp_rank = g_tp_split_rank, - .tp_world = g_tp_split_world, - .tp_expert_base = (int32_t)first_expert, - .in_dim = expert_in_dim, - .mid_dim = expert_mid_dim, - .out_dim = out_dim, - .n_total_expert = n_total_expert, - .n_expert_used = n_expert, - .n_tokens = n_tokens, - .mid_token_stride = mid_token_stride, - .down_type = down_type, - .gate_expert_bytes = gate_expert_bytes, - .gate_row_bytes = gate_row_bytes, - .up_expert_bytes = up_expert_bytes, - .up_row_bytes = up_row_bytes, - .down_expert_bytes = down_expert_bytes, - .down_row_bytes = down_row_bytes, - }; - const NSUInteger pair_x_groups = - gate_pair_q2 ? (use_stream_expert_addr_table ? - (NSUInteger)((expert_mid_dim + 1u) / 2u) : - (NSUInteger)((expert_mid_dim + 7u) / 8u)) : - gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : - use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : - q4_scalar_pair ? (NSUInteger)expert_mid_dim : - q4_pair2 ? (NSUInteger)((expert_mid_dim + 1u) / 2u) : - (NSUInteger)((expert_mid_dim + 7u) / 8u); - const NSUInteger pair_threadgroup_bytes = - q4_scalar_pair ? 512u * sizeof(float) : 0u; - const NSUInteger pair_threads = - q4_scalar_pair ? 256u : 64u; - const NSUInteger down_x_groups = - down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : - down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : - down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : - down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : - (NSUInteger)out_dim; - const NSUInteger down_threadgroup_bytes = - down_simd ? 0u : 256u * sizeof(float); - const NSUInteger down_threads = - down_simd ? 64u : 256u; - if (use_stream_expert_addr_table && - !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_resources, - stream_resource_count, - 0)) { - return 0; - } - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pair_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner - atIndex:1]; - [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner - atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; - [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; - if (use_stream_expert_addr_table) { - for (uint32_t i = 0; i < stream_resource_count; i++) { - [enc useResource:stream_resources[i]->gate_buffer usage:MTLResourceUsageRead]; - [enc useResource:stream_resources[i]->up_buffer usage:MTLResourceUsageRead]; - } - if (stream_overflow_gate) [enc useResource:stream_overflow_gate usage:MTLResourceUsageRead]; - if (stream_overflow_up) [enc useResource:stream_overflow_up usage:MTLResourceUsageRead]; - } - if (pair_threadgroup_bytes != 0u) { - [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, - (NSUInteger)n_expert, - (NSUInteger)n_tokens) - threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("pair"); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:down_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf - offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner - atIndex:1]; - [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - if (use_stream_expert_addr_table) { - for (uint32_t i = 0; i < stream_resource_count; i++) { - [enc useResource:stream_resources[i]->down_buffer usage:MTLResourceUsageRead]; - } - if (stream_overflow_down) [enc useResource:stream_overflow_down usage:MTLResourceUsageRead]; - } - if (down_threadgroup_bytes != 0u) { - [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; - } - [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, - (NSUInteger)n_tokens, - 1) - threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("down"); - - if (!ok) return 0; - if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed batch MoE")) return 0; -#undef DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE - } - - return 1; -} - -int ds4_gpu_glm_routed_moe_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride, - bool force_resident) { - (void)force_resident; - return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, - mid, - model_map, - model_size, - gate_offset, - up_offset, - down_offset, - gate_type, - up_type, - down_type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - selected, - weights, - n_total_expert, - n_expert, - layer_index, - x, - n_tokens, - mid_token_stride, - true, - false); -} - -int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *mid, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t up_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t up_expert_bytes, - uint64_t up_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - uint32_t layer_index, - const ds4_gpu_tensor *x, - uint32_t n_tokens, - uint32_t mid_token_stride) { - return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, - mid, - model_map, - model_size, - gate_offset, - up_offset, - down_offset, - gate_type, - up_type, - down_type, - gate_expert_bytes, - gate_row_bytes, - up_expert_bytes, - up_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - selected, - weights, - n_total_expert, - n_expert, - layer_index, - x, - n_tokens, - mid_token_stride, - false, - false); -} - -int ds4_gpu_router_select_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - uint64_t hash_offset, - uint32_t hash_rows, - uint32_t token, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_expert_groups, - uint32_t n_group_used, - bool has_bias, - bool hash_mode, - const ds4_gpu_tensor *logits) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || !weights || !probs || !logits || !model_map || - n_expert == 0 || n_expert_used == 0) return 0; - if (hash_mode && token >= hash_rows) return 0; - if (n_expert_groups > 1u || n_group_used > 0u) { - fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); - return 0; - } - - @autoreleasepool { - id logitsbuf = ds4_gpu_tensor_buffer(logits); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id probsbuf = ds4_gpu_tensor_buffer(probs); - if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || - ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || - ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) || - ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || - ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { - fprintf(stderr, "ds4: Metal router select received undersized buffers\n"); - return 0; - } - - uint64_t bias_inner = 0; - uint64_t hash_inner = 0; - id biasbuf = nil; - id hashbuf = nil; - NSUInteger bias_set_offset = 0; - NSUInteger hash_set_offset = 0; - if (has_bias && !hash_mode) { - const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); - biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); - if (!biasbuf) return 0; - bias_set_offset = (NSUInteger)bias_inner; - } - if (hash_mode) { - const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); - hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); - if (!hashbuf) return 0; - hash_set_offset = (NSUInteger)hash_inner; - } - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - const int32_t token_i32 = (int32_t)token; - int ok = cb && - ds4_gpu_encode_router_select(cb, - selected, - weights, - probs, - logitsbuf, - ds4_gpu_tensor_offset(logits), - biasbuf, - bias_set_offset, - hashbuf, - hash_set_offset, - nil, - 0, - &token_i32, - hash_rows, - 1, - n_expert, - n_expert_used, - expert_weight_scale, - has_bias && !hash_mode, - hash_mode); - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } - if (!ok) return 0; - } - - return 1; -} - -int ds4_gpu_router_select_batch_tensor( - ds4_gpu_tensor *selected, - ds4_gpu_tensor *weights, - ds4_gpu_tensor *probs, - const void *model_map, - uint64_t model_size, - uint64_t bias_offset, - uint64_t hash_offset, - uint32_t hash_rows, - uint32_t n_expert_groups, - uint32_t n_group_used, - bool has_bias, - bool hash_mode, - const ds4_gpu_tensor *logits, - const ds4_gpu_tensor *tokens, - uint32_t n_expert, - uint32_t n_expert_used, - float expert_weight_scale, - uint32_t n_tokens) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!selected || !weights || !probs || !logits || !tokens || !model_map || - n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; - if (n_expert_groups > 1u || n_group_used > 0u) { - fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); - return 0; - } - - @autoreleasepool { - id logitsbuf = ds4_gpu_tensor_buffer(logits); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - id probsbuf = ds4_gpu_tensor_buffer(probs); - id tokensbuf = ds4_gpu_tensor_buffer(tokens); - if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf || - ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) || - ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) || - ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) || - ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) || - ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) { - fprintf(stderr, "ds4: Metal router batch select received undersized buffers\n"); - return 0; - } - - uint64_t bias_inner = 0; - uint64_t hash_inner = 0; - id biasbuf = nil; - id hashbuf = nil; - NSUInteger bias_set_offset = 0; - NSUInteger hash_set_offset = 0; - if (has_bias && !hash_mode) { - const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); - biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); - if (!biasbuf) return 0; - bias_set_offset = (NSUInteger)bias_inner; - } - if (hash_mode) { - const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); - hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); - if (!hashbuf) return 0; - hash_set_offset = (NSUInteger)hash_inner; - } - - const bool had_batch = g_batch_cb != nil; - if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - int ok = cb && - ds4_gpu_encode_router_select(cb, - selected, - weights, - probs, - logitsbuf, - ds4_gpu_tensor_offset(logits), - biasbuf, - bias_set_offset, - hashbuf, - hash_set_offset, - tokensbuf, - ds4_gpu_tensor_offset(tokens), - NULL, - hash_rows, - n_tokens, - n_expert, - n_expert_used, - expert_weight_scale, - has_bias && !hash_mode, - hash_mode); - if (!had_batch) { - ok = ds4_gpu_end_commands() != 0 && ok; - } - if (!ok) return 0; - } - - return 1; -} - -int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) { - if (n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED || - (!selected && n_selected != 0)) return 0; - for (uint32_t i = 0; i < n_selected; i++) { - g_routed_moe_selected_override[i] = selected[i]; - } - g_routed_moe_selected_override_n = n_selected; - return 1; -} - -int ds4_gpu_routed_moe_one_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - ds4_gpu_tensor *experts, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - float clamp, - const ds4_gpu_tensor *x, - const ds4_gpu_tensor *add_in, - uint32_t layer_index, - bool force_resident) { - if (!g_initialized && !ds4_gpu_init()) return 0; - /* TP sharding: only the owned contiguous expert range is mapped, - * so bind from the owned base, validate only its bytes, and tell the - * kernels the first expert id present at that base. */ - uint32_t first_expert = 0; - uint32_t n_bind_expert = 0; - ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); - const int32_t tp_expert_base_host = (int32_t)first_expert; - gate_offset += (uint64_t)first_expert * gate_expert_bytes; - up_offset += (uint64_t)first_expert * gate_expert_bytes; - down_offset += (uint64_t)first_expert * down_expert_bytes; - - if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || - n_total_expert == 0 || n_expert == 0 || - n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED || - gate_expert_bytes == 0 || down_expert_bytes == 0 || - gate_row_bytes == 0 || down_row_bytes == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32044); - return 0; - } - if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32047); return 0; } - ds4_gpu_stream_expert_cache_note_token(layer_index); - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id gatebuf = ds4_gpu_tensor_buffer(gate); - id upbuf = ds4_gpu_tensor_buffer(up); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id expertsbuf = ds4_gpu_tensor_buffer(experts); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id selected_exec_buf = selectedbuf; - NSUInteger selected_exec_off = ds4_gpu_tensor_offset(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); - const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); - if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(gate) < mid_bytes || - ds4_gpu_tensor_bytes(up) < mid_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int) || - ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { - fprintf(stderr, "ds4: Metal routed tensor MoE received undersized activation buffers\n"); - return 0; - } - if (n_expert > 1 && - (!expertsbuf || - ds4_gpu_tensor_bytes(experts) < (uint64_t)n_expert * out_dim * sizeof(float))) { - fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n"); - return 0; - } - - if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed MoE tensor byte size overflow\n"); - return 0; - } - const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gate_buf = nil; - id up_buf = nil; - id down_buf = nil; - __unsafe_unretained id gate_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; - __unsafe_unretained id up_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; - __unsafe_unretained id down_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; - ds4_gpu_stream_expert_cache_entry *stream_slot_entries[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { NULL }; - NSUInteger gate_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - NSUInteger up_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - NSUInteger down_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - uint64_t stream_gate_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - uint64_t stream_up_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - uint64_t stream_down_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - id stream_gate_addr_buf = nil; - id stream_up_addr_buf = nil; - id stream_down_addr_buf = nil; - bool use_stream_expert_addr_table = false; - bool use_stream_expert_masked_addr_table = false; - bool use_stream_compact_addr_table = false; - bool use_stream_expert_cache = false; - bool use_stream_expert_split_candidate = false; - bool use_stream_expert_split_deferred = false; - bool stream_expert_split_completed = false; - uint32_t stream_expert_resident_mask = 0; - uint32_t stream_expert_missing_mask = 0; - __unsafe_unretained id gate_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; - __unsafe_unretained id up_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; - __unsafe_unretained id down_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; - NSUInteger gate_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - NSUInteger up_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - NSUInteger down_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; - __unsafe_unretained id gate_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; - __unsafe_unretained id up_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; - __unsafe_unretained id down_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; - NSUInteger gate_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; - NSUInteger up_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; - NSUInteger down_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; - __unsafe_unretained id gate_group24_bufs[24] = { nil }; - __unsafe_unretained id up_group24_bufs[24] = { nil }; - __unsafe_unretained id down_group24_bufs[24] = { nil }; - NSUInteger gate_group24_offsets[24] = { 0 }; - NSUInteger up_group24_offsets[24] = { 0 }; - NSUInteger down_group24_offsets[24] = { 0 }; - DS4MetalQ4ExpertTable *gate_table = nil; - DS4MetalQ4ExpertTable *up_table = nil; - DS4MetalQ4ExpertTable *down_table = nil; - id q4_table_layer_residency = nil; - int32_t selected_ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; - - const uint32_t n_tokens = 1; - const uint32_t pair_rows = n_tokens * n_expert; - const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); - if ((n_expert > 1 && !expertsbuf && - !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, - &g_moe_down_scratch_bytes, - (NSUInteger)down_scratch_bytes, - "ds4_moe_down_scratch"))) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32148); - return 0; - } - - const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); - const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); - id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); - id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); - if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { - fprintf(stderr, "ds4: unsupported Metal routed MoE quant types gate=%u down=%u\n", - gate_type, down_type); - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32158); - return 0; - } - - ds4_gpu_mul_mv_id_args gate_args = - ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - gate_row_bytes, gate_expert_bytes, - 1, n_expert, n_tokens, gate_nr0); - ds4_gpu_mul_mv_id_args down_args = - ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, - down_row_bytes, down_expert_bytes, - n_expert, n_expert, n_tokens, down_nr0); - /* Tensor-parallel expert ownership; non-TP calls keep tp_world at 1. */ - gate_args.tp_rank = g_tp_split_rank; - gate_args.tp_world = g_tp_split_world; - gate_args.tp_expert_base = tp_expert_base_host; - down_args.tp_rank = g_tp_split_rank; - down_args.tp_world = g_tp_split_world; - down_args.tp_addend = add_in != NULL; - down_args.tp_expert_base = tp_expert_base_host; - - const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); - const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); - const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); - const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); - const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); - const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); - int ok = 1; - const bool write_clamped_moe = - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL; - id pair_swiglu_pipeline = nil; - if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { - pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; - } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { - pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; - } - const bool fuse_pair_swiglu = - !g_quality_mode && - !write_clamped_moe && - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && - pair_swiglu_pipeline != nil; - id down_sum6_pipeline = nil; - if (down_type == DS4_METAL_TENSOR_Q2_K) { - down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; - } else if (down_type == DS4_METAL_TENSOR_Q4_K) { - down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; - } else if (down_type == DS4_METAL_TENSOR_IQ2_XXS && - g_tp_split_world == 2) { - /* IQ2 down-sum exists for the GLM TP resident split only; the - * streaming paths must keep their addr/masked chain (expert - * bytes are not at their model-map offsets when streamed). */ - down_sum6_pipeline = g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; - } - const bool direct_down_sum = - !g_quality_mode && - (n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2)) && - n_tokens == 1 && - down_sum6_pipeline != nil; - /* The expert-ownership split lives only in the fused id pair+sum6 - * kernels; every other routed variant would silently compute full - * sums on both ranks and double the combine. Fail fast instead. */ - if ((g_tp_split_world > 1 || add_in) && !(fuse_pair_swiglu && direct_down_sum)) { - fprintf(stderr, "ds4: tensor-parallel routed MoE requires the fused pair+sum6 decode path\n"); - return 0; - } - const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; - /* - * The grouped Q4 experiment keeps selected IDs on GPU, but it also walks - * every expert window in the layer. On PRO Q4 this measured far slower - * than the active selected-slot path, so keep it opt-in for profiling. - */ - const bool enable_q4_grouped_experts = - getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL; - const bool use_q4_grouped_experts = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert >= 128 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group_q4_k_sum6_pipeline != nil && - enable_q4_grouped_experts; - const uint32_t q4_expert_group_size = - use_q4_grouped_experts ? ds4_gpu_q4_expert_group_size(n_total_expert) : 0; - const bool q4_grouped_boundary = - use_q4_grouped_experts && - g_batch_cb != nil && - getenv("DS4_METAL_DISABLE_Q4_GROUPED_BOUNDARY") == NULL; - const bool q4_grouped_cache_views = - getenv("DS4_METAL_Q4_GROUPED_CACHE_VIEWS") != NULL; - const uint32_t q4_group6_expert_group_size = 64; - const bool use_q4_group6_experts = - !use_q4_grouped_experts && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == q4_group6_expert_group_size * 6u && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil && - getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL; - const uint32_t q4_group8_expert_group_size = 48; - const bool use_q4_group8_experts = - !use_q4_grouped_experts && - !use_q4_group6_experts && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == q4_group8_expert_group_size * 8u && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil && - getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL; - const uint32_t q4_group24_expert_group_size = 16; - const bool use_q4_group24_experts = - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == q4_group24_expert_group_size * 24u && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - direct_down_sum && - g_moe_mul_mv_group24_q4_k_id_pipeline != nil && - g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil && - getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL; - const bool q4_group24_exact_views = - use_q4_group24_experts && - getenv("DS4_METAL_Q4_GROUP24_EXACT_VIEWS") != NULL && - getenv("DS4_METAL_Q4_GROUP24_BASE_VIEWS") == NULL; - const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; - const bool can_wrap_q4_exact_tensors = - max_buffer_len != 0 && - gate_tensor_bytes <= max_buffer_len && - down_tensor_bytes <= max_buffer_len; - /* - * The full-tensor Q4 ID path is the closest arithmetic analogue to IQ2, - * but PRO Q4 routed tensors are multi-GiB. Keep it opt-in: even with - * per-layer command boundaries it is much slower than binding only the - * six active experts on current M3 Ultra Metal. - */ - const bool enable_q4_exact_tensor_id = - getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && - getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; - const bool use_q4_exact_tensor_id = - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - !use_q4_group24_experts && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == 384 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - can_wrap_q4_exact_tensors && - fuse_pair_swiglu && - direct_down_sum && - enable_q4_exact_tensor_id; - const bool q4_exact_boundary = - use_q4_exact_tensor_id && - g_batch_cb != nil && - getenv("DS4_METAL_DISABLE_Q4_EXACT_BOUNDARY") == NULL; - const bool q4_expert_table_auto = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes); - const bool q4_expert_address_auto = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes); - const bool q4_table_queue_residency = - ds4_gpu_q4_table_queue_residency_enabled(q4_expert_table_auto || - q4_expert_address_auto); - const bool use_q4_expert_address_table = - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - !use_q4_group24_experts && - !use_q4_exact_tensor_id && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == 384 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && - (getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || - q4_expert_address_auto) && - (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL || - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || - q4_table_queue_residency || - ds4_gpu_q4_table_model_residency_enabled() || - getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL; - const bool enable_q4_expert_table = - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || - q4_expert_table_auto; - const bool use_q4_expert_table = - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - !use_q4_group24_experts && - !use_q4_exact_tensor_id && - !use_q4_expert_address_table && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == 384 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && - g_moe_table_q4_pair_gate_encoder != nil && - g_moe_table_q4_pair_up_encoder != nil && - g_moe_table_q4_sum_down_encoder != nil && - enable_q4_expert_table && - (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || - q4_table_queue_residency || - ds4_gpu_q4_table_model_residency_enabled() || - getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL; - const bool q4_table_boundary = - use_q4_expert_table && - g_batch_cb != nil && - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && - !q4_table_queue_residency && - getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; - const bool enable_q4_gather_slots = - getenv("DS4_METAL_ENABLE_Q4_GATHER_SLOTS") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GATHER_SLOTS") == NULL; - const bool use_q4_gather_slots = - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - !use_q4_group24_experts && - !use_q4_exact_tensor_id && - !use_q4_expert_address_table && - !use_q4_expert_table && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert == q4_group6_expert_group_size * 6u && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - fuse_pair_swiglu && - direct_down_sum && - g_moe_q4_gather_slots6_pipeline != nil && - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && - enable_q4_gather_slots; - const bool use_q4_selected_slots = - !force_resident && - !use_q4_grouped_experts && - !use_q4_group6_experts && - !use_q4_group8_experts && - !use_q4_group24_experts && - !use_q4_exact_tensor_id && - !use_q4_expert_address_table && - !use_q4_expert_table && - !use_q4_gather_slots && - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens == 1 && - n_total_expert >= 128 && - (g_ssd_streaming_mode || - (gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes)) && - fuse_pair_swiglu && - direct_down_sum && - ds4_gpu_q4_selected_paths_allowed() && - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && - getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL; - const bool use_iq2_selected_slots = - !force_resident && - g_ssd_streaming_mode && - gate_type == DS4_METAL_TENSOR_IQ2_XXS && - down_type == DS4_METAL_TENSOR_Q2_K && - n_expert == 6 && - n_tokens == 1 && - fuse_pair_swiglu && - direct_down_sum && - g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil && - getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; - const bool use_iq2_stream_addr_table = - !force_resident && - g_ssd_streaming_mode && - gate_type == DS4_METAL_TENSOR_IQ2_XXS && - down_type == DS4_METAL_TENSOR_IQ2_XXS && - n_expert <= DS4_METAL_MAX_ROUTED_EXPERT_USED && - n_tokens == 1 && - fuse_pair_swiglu && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_iq2_xxs_pipeline != nil && - getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL; - const bool use_selected_slots = - use_q4_selected_slots || use_iq2_selected_slots || use_iq2_stream_addr_table; - id slots_pair_swiglu_pipeline = - use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline : - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; - id slots_sum6_pipeline = - use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline : - g_moe_mul_mv_slots6_q4_k_sum6_pipeline; - const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE"); - if (!selected_profile_env) { - selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE"); - } - const char *selected_profile_layer_env = getenv("DS4_METAL_SELECTED_PROFILE_LAYER"); - if (!selected_profile_layer_env) { - selected_profile_layer_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE_LAYER"); - } - bool selected_profile_layer_match = true; - if (selected_profile_layer_env && selected_profile_layer_env[0]) { - char *end = NULL; - long layer = strtol(selected_profile_layer_env, &end, 10); - selected_profile_layer_match = - end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index; - } - const bool selected_profile = - use_selected_slots && - selected_profile_env != NULL && - selected_profile_layer_match; - const bool q4_selected_shared_event = - use_q4_selected_slots && - getenv("DS4_METAL_Q4_SELECTED_SHARED_EVENT") != NULL; - const bool q4_selected_base_views = - use_q4_selected_slots && - getenv("DS4_METAL_Q4_SELECTED_USE_BASE_VIEWS") != NULL && - getenv("DS4_METAL_Q4_SELECTED_EXACT_VIEWS") == NULL; - const bool q4_selected_transient_views = - use_q4_selected_slots && - !q4_selected_base_views && - getenv("DS4_METAL_Q4_SELECTED_TRANSIENT_VIEWS") != NULL; - const char *q4_selected_view_mode = - q4_selected_base_views ? "base" : - (q4_selected_transient_views ? "transient" : "cached"); - if (!use_selected_slots) { - g_routed_moe_selected_override_n = 0; - } - if (use_q4_expert_address_table) { - gate_table = ds4_gpu_q4_expert_address_table(model_map, - model_size, - gate_offset, - gate_expert_bytes, - n_total_expert); - up_table = ds4_gpu_q4_expert_address_table(model_map, - model_size, - up_offset, - gate_expert_bytes, - n_total_expert); - down_table = ds4_gpu_q4_expert_address_table(model_map, - model_size, - down_offset, - down_expert_bytes, - n_total_expert); - if (!gate_table || !up_table || !down_table) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32546); - return 0; - } - q4_table_layer_residency = - ds4_gpu_q4_expert_layer_residency_set(gate_table, - up_table, - down_table, - q4_expert_address_auto); - } else if (use_q4_expert_table) { - gate_table = ds4_gpu_q4_expert_table(model_map, - model_size, - gate_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_gate_encoder); - up_table = ds4_gpu_q4_expert_table(model_map, - model_size, - up_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_up_encoder); - down_table = ds4_gpu_q4_expert_table(model_map, - model_size, - down_offset, - down_expert_bytes, - n_total_expert, - g_moe_table_q4_sum_down_encoder); - if (!gate_table || !up_table || !down_table) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32573); - return 0; - } - q4_table_layer_residency = - ds4_gpu_q4_expert_layer_residency_set(gate_table, - up_table, - down_table, - q4_expert_table_auto); - } else if (use_q4_exact_tensor_id) { - gate_buf = ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset, - gate_tensor_bytes, - &gate_inner); - up_buf = ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset, - gate_tensor_bytes, - &up_inner); - down_buf = ds4_gpu_wrap_model_exact_range(model_map, - model_size, - down_offset, - down_tensor_bytes, - &down_inner); - if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32601); return 0; } - } else if (use_q4_group6_experts || use_q4_gather_slots) { - if ((uint64_t)q4_group6_expert_group_size > UINT64_MAX / gate_expert_bytes || - (uint64_t)q4_group6_expert_group_size > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group6 byte size overflow\n"); - return 0; - } - const uint64_t gate_group_bytes = - (uint64_t)q4_group6_expert_group_size * gate_expert_bytes; - const uint64_t down_group_bytes = - (uint64_t)q4_group6_expert_group_size * down_expert_bytes; - for (uint32_t i = 0; i < 6; i++) { - const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; - const uint64_t down_rel = (uint64_t)i * down_group_bytes; - if (gate_rel > gate_tensor_bytes || - gate_group_bytes > gate_tensor_bytes - gate_rel || - down_rel > down_tensor_bytes || - down_group_bytes > down_tensor_bytes - down_rel || - gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group6 offset overflow\n"); - return 0; - } - - uint64_t group_inner = 0; - gate_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset + gate_rel, - gate_group_bytes, - &group_inner); - gate_group6_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - up_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset + gate_rel, - gate_group_bytes, - &group_inner); - up_group6_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - down_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset + down_rel, - down_group_bytes, - &group_inner); - down_group6_offsets[i] = (NSUInteger)group_inner; - if (!gate_group6_bufs[i] || !up_group6_bufs[i] || !down_group6_bufs[i]) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32643); - return 0; - } - } - if (use_q4_gather_slots) { - if (gate_expert_bytes > NSUIntegerMax || - down_expert_bytes > NSUIntegerMax || - gate_expert_bytes > UINT64_MAX / 6u || - down_expert_bytes > UINT64_MAX / 6u || - 6ull * gate_expert_bytes > NSUIntegerMax || - 6ull * down_expert_bytes > NSUIntegerMax) { - fprintf(stderr, "ds4: Metal routed MoE Q4 gather scratch byte size overflow\n"); - return 0; - } - const NSUInteger gate_slots_bytes = (NSUInteger)(6ull * gate_expert_bytes); - const NSUInteger down_slots_bytes = (NSUInteger)(6ull * down_expert_bytes); - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_gate_slots_buffer, - &g_moe_q4_gate_slots_bytes, - gate_slots_bytes, - "ds4_moe_q4_gate_slots") || - !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_up_slots_buffer, - &g_moe_q4_up_slots_bytes, - gate_slots_bytes, - "ds4_moe_q4_up_slots") || - !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_down_slots_buffer, - &g_moe_q4_down_slots_bytes, - down_slots_bytes, - "ds4_moe_q4_down_slots")) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32670); - return 0; - } - for (uint32_t i = 0; i < 6; i++) { - gate_slot_bufs[i] = g_moe_q4_gate_slots_buffer; - up_slot_bufs[i] = g_moe_q4_up_slots_buffer; - down_slot_bufs[i] = g_moe_q4_down_slots_buffer; - gate_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); - up_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); - down_slot_offsets[i] = (NSUInteger)((uint64_t)i * down_expert_bytes); - } - } - } else if (use_q4_group8_experts) { - if ((uint64_t)q4_group8_expert_group_size > UINT64_MAX / gate_expert_bytes || - (uint64_t)q4_group8_expert_group_size > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group8 byte size overflow\n"); - return 0; - } - const uint64_t gate_group_bytes = - (uint64_t)q4_group8_expert_group_size * gate_expert_bytes; - const uint64_t down_group_bytes = - (uint64_t)q4_group8_expert_group_size * down_expert_bytes; - for (uint32_t i = 0; i < 8; i++) { - const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; - const uint64_t down_rel = (uint64_t)i * down_group_bytes; - if (gate_rel > gate_tensor_bytes || - gate_group_bytes > gate_tensor_bytes - gate_rel || - down_rel > down_tensor_bytes || - down_group_bytes > down_tensor_bytes - down_rel || - gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group8 offset overflow\n"); - return 0; - } - - uint64_t group_inner = 0; - gate_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset + gate_rel, - gate_group_bytes, - &group_inner); - gate_group8_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - up_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset + gate_rel, - gate_group_bytes, - &group_inner); - up_group8_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - down_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset + down_rel, - down_group_bytes, - &group_inner); - down_group8_offsets[i] = (NSUInteger)group_inner; - if (!gate_group8_bufs[i] || !up_group8_bufs[i] || !down_group8_bufs[i]) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32727); - return 0; - } - } - } else if (use_q4_group24_experts) { - if ((uint64_t)q4_group24_expert_group_size > UINT64_MAX / gate_expert_bytes || - (uint64_t)q4_group24_expert_group_size > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group24 byte size overflow\n"); - return 0; - } - const uint64_t gate_group_bytes = - (uint64_t)q4_group24_expert_group_size * gate_expert_bytes; - const uint64_t down_group_bytes = - (uint64_t)q4_group24_expert_group_size * down_expert_bytes; - for (uint32_t i = 0; i < 24; i++) { - const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; - const uint64_t down_rel = (uint64_t)i * down_group_bytes; - if (gate_rel > gate_tensor_bytes || - gate_group_bytes > gate_tensor_bytes - gate_rel || - down_rel > down_tensor_bytes || - down_group_bytes > down_tensor_bytes - down_rel || - gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal routed MoE Q4 group24 offset overflow\n"); - return 0; - } - - uint64_t group_inner = 0; - gate_group24_bufs[i] = q4_group24_exact_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset + gate_rel, - gate_group_bytes, - &group_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset + gate_rel, - gate_group_bytes, - &group_inner); - gate_group24_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - up_group24_bufs[i] = q4_group24_exact_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset + gate_rel, - gate_group_bytes, - &group_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset + gate_rel, - gate_group_bytes, - &group_inner); - up_group24_offsets[i] = (NSUInteger)group_inner; - group_inner = 0; - down_group24_bufs[i] = q4_group24_exact_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - down_offset + down_rel, - down_group_bytes, - &group_inner) : - ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset + down_rel, - down_group_bytes, - &group_inner); - down_group24_offsets[i] = (NSUInteger)group_inner; - if (!gate_group24_bufs[i] || !up_group24_bufs[i] || !down_group24_bufs[i]) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32794); - return 0; - } - } - } else if (use_selected_slots) { - const bool selected_timing = - selected_profile || - ds4_gpu_stream_expert_timing_summary_enabled(); - double selected_t0 = selected_timing ? ds4_gpu_now_ms() : 0.0; - double selected_read_ms = 0.0; - double selected_sync_ms = 0.0; - double selected_copy_ms = 0.0; - double selected_wrap_ms = 0.0; - uint64_t selected_cache_hits0 = g_stream_expert_cache_hits; - uint64_t selected_cache_misses0 = g_stream_expert_cache_misses; - uint64_t selected_cache_wraps0 = g_stream_expert_cache_wraps; - uint64_t selected_cache_evictions0 = g_stream_expert_cache_evictions; - const char *selected_id_source = "readback"; - bool selected_ids_available = true; - bool selected_exec_ids_from_host = false; - const int stream_expert_cache_size_known = - ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, - down_expert_bytes); - const bool use_iq2_full_expert_addr_table = - use_iq2_selected_slots && - ds4_gpu_stream_full_expert_addr_table_requested() && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; - use_stream_expert_cache = - !use_iq2_full_expert_addr_table && - (use_iq2_selected_slots || use_iq2_stream_addr_table || use_q4_selected_slots) && - stream_expert_cache_size_known && - ds4_gpu_stream_expert_cache_effective_cap(layer_index, - n_total_expert, - n_expert) != 0; - if (use_iq2_stream_addr_table && !use_stream_expert_cache) { - fprintf(stderr, - "ds4: Metal IQ2/IQ2 streaming decode requires a non-empty expert cache\n"); - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32831); - return 0; - } - const bool stream_split_ready = - use_stream_expert_cache && - ds4_gpu_stream_expert_split_ready(); - const bool use_stream_compact_addr = - use_stream_expert_cache && - use_iq2_selected_slots && - ds4_gpu_stream_compact_addr_requested() && - !stream_split_ready && - !ds4_gpu_stream_expert_masked_addr_requested() && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; - use_stream_expert_split_candidate = - use_stream_expert_cache && - use_iq2_selected_slots && - !use_stream_compact_addr && - stream_split_ready && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; - const bool use_stream_hit_validator = - use_stream_expert_cache && - use_iq2_selected_slots && - ds4_gpu_stream_expert_hit_validator_requested() && - g_moe_stream_expert_cache_validate_pipeline != nil && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil && - ds4_gpu_stream_expert_cache_addr_buffers(layer_index, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf); - if (use_iq2_full_expert_addr_table) { - selected_id_source = "gpu-full-addr"; - selected_ids_available = false; - g_routed_moe_selected_override_n = 0; - ds4_gpu_stream_expert_cache_entry *full_entry = NULL; - if (!ds4_gpu_stream_full_expert_addr_table_prepare(model_map, - model_size, - layer_index, - n_total_expert, - gate_offset, - up_offset, - down_offset, - gate_expert_bytes, - down_expert_bytes, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf, - &full_entry)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32880); - return 0; - } - for (uint32_t i = 0; i < n_expert; i++) { - stream_slot_entries[i] = full_entry; - } - use_stream_expert_addr_table = true; - } else { - const int replayed_selected_ids = - ds4_gpu_moe_selected_trace_replay(selected_ids, n_expert); - if (replayed_selected_ids < 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32890); - return 0; - } else if (replayed_selected_ids > 0) { - selected_id_source = "replay"; - selected_exec_ids_from_host = true; - g_routed_moe_selected_override_n = 0; - } else if (g_routed_moe_selected_override_n == n_expert) { - memcpy(selected_ids, - g_routed_moe_selected_override, - (size_t)n_expert * sizeof(selected_ids[0])); - selected_id_source = "override"; - selected_exec_ids_from_host = true; - g_routed_moe_selected_override_n = 0; - } else if (use_stream_hit_validator) { - g_routed_moe_selected_override_n = 0; - uint32_t validator_all_cached = 0; - uint32_t validator_miss_mask = 0; - uint32_t validator_invalid_mask = 0; - if (!ds4_gpu_stream_expert_cache_validate_selected(selected, - stream_gate_addr_buf, - stream_up_addr_buf, - stream_down_addr_buf, - n_total_expert, - n_expert, - selected_ids, - &validator_all_cached, - &validator_miss_mask, - &validator_invalid_mask)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32917); - return 0; - } - selected_id_source = - validator_invalid_mask != 0 ? "validator-invalid" : - (validator_miss_mask != 0 ? "validator-miss" : - (validator_all_cached != 0 ? "validator-hit" : "validator-miss")); - } else { - g_routed_moe_selected_override_n = 0; - if (g_batch_cb != nil) { - double selected_boundary_t0 = - selected_timing ? ds4_gpu_now_ms() : 0.0; - if (q4_selected_shared_event) { - if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32942); return 0; } - } else if (ds4_gpu_end_commands() == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32931); - return 0; - } - if (selected_timing) { - selected_sync_ms += - ds4_gpu_now_ms() - selected_boundary_t0; - } - double selected_copy_t0 = - selected_timing ? ds4_gpu_now_ms() : 0.0; - if (ds4_gpu_tensor_read(selected, - 0, - selected_ids, - (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32943); - return 0; - } - if (selected_timing) { - selected_copy_ms += - ds4_gpu_now_ms() - selected_copy_t0; - } - if (!q4_selected_shared_event) { - selected_boundary_t0 = - selected_timing ? ds4_gpu_now_ms() : 0.0; - if (ds4_gpu_begin_commands() == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32967); return 0; } - if (selected_timing) { - selected_sync_ms += - ds4_gpu_now_ms() - selected_boundary_t0; - } - } - } else { - double selected_copy_t0 = - selected_timing ? ds4_gpu_now_ms() : 0.0; - if (ds4_gpu_tensor_read(selected, - 0, - selected_ids, - (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32965); - return 0; - } - if (selected_timing) { - selected_copy_ms += - ds4_gpu_now_ms() - selected_copy_t0; - } - } - } - } - if (selected_timing) { - selected_read_ms = selected_sync_ms + selected_copy_ms; - selected_t0 = ds4_gpu_now_ms(); - } - - if (selected_ids_available) { - for (uint32_t i = 0; i < n_expert; i++) { - if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { - fprintf(stderr, - "ds4: Metal routed MoE selected expert id %d is outside 0..%u\n", - selected_ids[i], - n_total_expert); - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32986); - return 0; - } - } - ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, - selected_ids, - n_expert); - if (!ds4_gpu_moe_selected_trace_record(selected_ids, n_expert) || - !ds4_gpu_moe_selected_hotlist_record(layer_index, - selected_ids, - n_expert, - n_total_expert)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32997); - return 0; - } - - for (uint32_t i = 0; i < n_expert; i++) { - const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; - if (expert_id > UINT64_MAX / gate_expert_bytes || - expert_id > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); - return 0; - } - const uint64_t gate_rel = expert_id * gate_expert_bytes; - const uint64_t down_rel = expert_id * down_expert_bytes; - if (gate_rel > UINT64_MAX - gate_offset || - gate_rel > UINT64_MAX - up_offset || - down_rel > UINT64_MAX - down_offset) { - fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); - return 0; - } - stream_gate_abs_offsets[i] = gate_offset + gate_rel; - stream_up_abs_offsets[i] = up_offset + gate_rel; - stream_down_abs_offsets[i] = down_offset + down_rel; - - if (use_stream_expert_cache) { - ds4_gpu_stream_expert_cache_entry *entry = NULL; - entry = ds4_gpu_stream_expert_cache_peek(model_map, - model_size, - layer_index, - (uint32_t)selected_ids[i], - n_total_expert, - n_expert, - stream_gate_abs_offsets[i], - stream_up_abs_offsets[i], - stream_down_abs_offsets[i], - gate_expert_bytes, - down_expert_bytes); - if (!entry) { - stream_expert_missing_mask |= 1u << i; - continue; - } - stream_expert_resident_mask |= 1u << i; - stream_slot_entries[i] = entry; - gate_slot_bufs[i] = entry->gate_buffer; - gate_slot_offsets[i] = entry->gate_inner; - up_slot_bufs[i] = entry->up_buffer; - up_slot_offsets[i] = entry->up_inner; - down_slot_bufs[i] = entry->down_buffer; - down_slot_offsets[i] = entry->down_inner; - continue; - } - - uint64_t slot_inner = 0; - gate_slot_bufs[i] = q4_selected_base_views ? - ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset + gate_rel, - gate_expert_bytes, - &slot_inner) : - (q4_selected_transient_views ? - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - gate_offset + gate_rel, - gate_expert_bytes, - &slot_inner) : - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset + gate_rel, - gate_expert_bytes, - &slot_inner)); - gate_slot_offsets[i] = (NSUInteger)slot_inner; - slot_inner = 0; - up_slot_bufs[i] = q4_selected_base_views ? - ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset + gate_rel, - gate_expert_bytes, - &slot_inner) : - (q4_selected_transient_views ? - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - up_offset + gate_rel, - gate_expert_bytes, - &slot_inner) : - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset + gate_rel, - gate_expert_bytes, - &slot_inner)); - up_slot_offsets[i] = (NSUInteger)slot_inner; - slot_inner = 0; - down_slot_bufs[i] = q4_selected_base_views ? - ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset + down_rel, - down_expert_bytes, - &slot_inner) : - (q4_selected_transient_views ? - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - down_offset + down_rel, - down_expert_bytes, - &slot_inner) : - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - down_offset + down_rel, - down_expert_bytes, - &slot_inner)); - down_slot_offsets[i] = (NSUInteger)slot_inner; - if (!gate_slot_bufs[i] || !up_slot_bufs[i] || !down_slot_bufs[i]) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33105); - return 0; - } - } - } - if (use_stream_expert_cache) { - use_stream_expert_addr_table = - ((use_iq2_selected_slots && - ds4_gpu_stream_expert_addr_table_kernel_requested() && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil) || - use_iq2_stream_addr_table) && - ds4_gpu_stream_expert_cache_addr_buffers(layer_index, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf); - if (use_iq2_stream_addr_table && !use_stream_expert_addr_table) { - fprintf(stderr, - "ds4: Metal IQ2/IQ2 streaming decode could not prepare expert address buffers\n"); - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33123); - return 0; - } - use_stream_expert_masked_addr_table = - use_stream_expert_addr_table && - use_iq2_selected_slots && - ds4_gpu_stream_expert_masked_addr_requested() && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; - use_stream_expert_split_deferred = - use_stream_expert_split_candidate && - use_stream_expert_masked_addr_table && - stream_expert_resident_mask != 0 && - stream_expert_missing_mask != 0 && - ds4_gpu_stream_expert_split_worthwhile(stream_expert_resident_mask, - stream_expert_missing_mask) && - g_batch_cb != nil && - getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL; - if (use_stream_expert_split_deferred) { - const ds4_gpu_stream_expert_table table = { - .model_map = model_map, - .model_size = model_size, - .layer = layer_index, - .n_total_expert = n_total_expert, - .gate_offset = gate_offset, - .up_offset = up_offset, - .down_offset = down_offset, - .gate_expert_bytes = gate_expert_bytes, - .down_expert_bytes = down_expert_bytes, - }; - if (!ds4_gpu_stream_expert_cache_begin_selected_load( - &table, - selected_ids, - n_expert)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33156); - return 0; - } - } - if (stream_expert_missing_mask != 0 && - !use_stream_expert_split_deferred) { - if (!ds4_gpu_stream_expert_cache_load_selected_missing( - model_map, - model_size, - layer_index, - selected_ids, - n_total_expert, - n_expert, - stream_gate_abs_offsets, - stream_up_abs_offsets, - stream_down_abs_offsets, - gate_expert_bytes, - down_expert_bytes, - stream_expert_missing_mask, - stream_slot_entries)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33175); - return 0; - } - for (uint32_t i = 0; i < n_expert; i++) { - if ((stream_expert_missing_mask & (1u << i)) == 0) continue; - ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; - if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); return 0; } - gate_slot_bufs[i] = entry->gate_buffer; - gate_slot_offsets[i] = entry->gate_inner; - up_slot_bufs[i] = entry->up_buffer; - up_slot_offsets[i] = entry->up_inner; - down_slot_bufs[i] = entry->down_buffer; - down_slot_offsets[i] = entry->down_inner; - } - } - if (use_iq2_stream_addr_table && use_stream_expert_addr_table) { - for (uint32_t i = 0; i < n_expert; i++) { - ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; - if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33214); return 0; } - if (!ds4_gpu_stream_expert_cache_set_addr_slot_raw( - layer_index, - (uint32_t)selected_ids[i], - entry->gate_buffer, - entry->gate_inner, - entry->up_buffer, - entry->up_inner, - entry->down_buffer, - entry->down_inner)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); - return 0; - } - } - } - ds4_gpu_stream_expert_cache_prune_layer(layer_index, - n_total_expert, - n_expert, - selected_ids, - n_expert); - ds4_gpu_stream_expert_cache_prune_global(layer_index, - selected_ids, - n_expert); - if (use_stream_compact_addr) { - id compact_selected = nil; - if (!ds4_gpu_stream_compact_addr_prepare(layer_index, - stream_slot_entries, - n_expert, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf, - &compact_selected)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33223); - return 0; - } - selected_exec_buf = compact_selected; - selected_exec_off = 0; - use_stream_expert_addr_table = true; - use_stream_expert_masked_addr_table = false; - use_stream_compact_addr_table = true; - } - if (use_stream_expert_addr_table && - !use_stream_compact_addr_table && - selected_exec_ids_from_host) { - if (!ds4_gpu_stream_selected_ids_prepare(layer_index, - selected_ids, - n_expert, - &selected_exec_buf, - &selected_exec_off)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33239); - return 0; - } - } - } - if (selected_timing) { - selected_wrap_ms = ds4_gpu_now_ms() - selected_t0; - if (use_stream_expert_cache) { - ds4_gpu_stream_expert_timing_note_cache_class( - stream_expert_resident_mask, - stream_expert_missing_mask); - } - ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, - selected_copy_ms, - selected_wrap_ms); - } - if (selected_profile) { - const uint64_t selected_cache_hits = - g_stream_expert_cache_hits - selected_cache_hits0; - const uint64_t selected_cache_misses = - g_stream_expert_cache_misses - selected_cache_misses0; - const uint64_t selected_cache_wraps = - g_stream_expert_cache_wraps - selected_cache_wraps0; - const uint64_t selected_cache_evictions = - g_stream_expert_cache_evictions - selected_cache_evictions0; - const char *selected_path = - use_iq2_stream_addr_table ? "iq2/iq2" : - (use_iq2_selected_slots ? "iq2/q2" : "q4/q4"); - const char *selected_view_mode = - use_stream_expert_split_deferred ? "stream-split" : - use_stream_expert_masked_addr_table ? "stream-addr-mask" : - use_stream_compact_addr_table ? "stream-compact-addr" : - use_stream_expert_addr_table ? "stream-addr" : - (use_stream_expert_cache ? "stream-cache" : - (use_iq2_selected_slots ? "exact-cache" : q4_selected_view_mode)); - fprintf(stderr, - "ds4: Metal selected views layer=%u path=%s mode=%s ids=%s " - "experts=%d,%d,%d,%d,%d,%d expert_gate=%.2f MiB " - "expert_down=%.2f MiB read=%.3f ms bind=%.3f ms " - "cache_hits=%llu cache_misses=%llu cache_wraps=%llu cache_evictions=%llu\n", - layer_index, - selected_path, - selected_view_mode, - selected_id_source, - selected_ids_available ? selected_ids[0] : -1, - selected_ids_available ? selected_ids[1] : -1, - selected_ids_available ? selected_ids[2] : -1, - selected_ids_available ? selected_ids[3] : -1, - selected_ids_available ? selected_ids[4] : -1, - selected_ids_available ? selected_ids[5] : -1, - ds4_gpu_mib(gate_expert_bytes), - ds4_gpu_mib(down_expert_bytes), - selected_read_ms, - selected_wrap_ms, - (unsigned long long)selected_cache_hits, - (unsigned long long)selected_cache_misses, - (unsigned long long)selected_cache_wraps, - (unsigned long long)selected_cache_evictions); - } - } else if (!use_q4_grouped_experts) { - gate_buf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, gate_tensor_bytes, &gate_inner); - up_buf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, gate_tensor_bytes, &up_inner); - down_buf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, down_tensor_bytes, &down_inner); - if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33326); return 0; } - } - if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33305); - return 0; - } - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33337); return 0; } - if ((use_q4_expert_address_table || use_q4_expert_table) && - !ds4_gpu_use_model_residency_set(cb)) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33314); - return 0; - } - if (!q4_table_queue_residency && - q4_table_layer_residency && - [cb respondsToSelector:@selector(useResidencySet:)]) { - [cb useResidencySet:q4_table_layer_residency]; - } - - const bool moe_one_stage_profile = - g_batch_cb != nil && - ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_ONE_STAGE_PROFILE", - "DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER", - layer_index); - const char *moe_one_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); - const char *moe_one_path = - use_q4_grouped_experts ? "q4_grouped_pair_swiglu" : - use_q4_group6_experts ? "q4_group6_pair_swiglu" : - use_q4_group8_experts ? "q4_group8_pair_swiglu" : - use_q4_group24_experts ? "q4_group24_split_gate_up" : - use_q4_exact_tensor_id ? "q4_exact_pair_swiglu" : - use_q4_expert_address_table ? "q4_addr_pair_swiglu" : - use_q4_expert_table ? "q4_table_pair_swiglu" : - use_q4_gather_slots ? "q4_gather_slots6_pair_swiglu" : - use_stream_expert_split_deferred ? "iq2_stream_split_pair_swiglu" : - use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" : - use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" : - use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" : - use_q4_selected_slots ? "q4_slots6_pair_swiglu" : - (fuse_pair_swiglu ? "pair_swiglu" : - ((!g_quality_mode && - ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || - (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline))) ? "pair" : "single")); - double moe_one_stage_t0 = moe_one_stage_profile ? ds4_gpu_now_ms() : 0.0; - if (moe_one_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33349); - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33380); return 0; } - moe_one_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_MOE_ONE_STAGE(name) do { \ - if (ok && moe_one_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !moe_one_stage_filter || !moe_one_stage_filter[0] || \ - strstr(stage_name, moe_one_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal routed MoE one stage layer=%u pairs=%u experts=%u " \ - "gate=%s down=%s path=%s %s=%.3f ms\n", \ - layer_index, pair_rows, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - moe_one_path, \ - stage_name, now_ms - moe_one_stage_t0); \ - } \ - moe_one_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - if (use_q4_gather_slots) { - ds4_gpu_q4_gather_slots6_args gate_gather_args = { - .expert_bytes = gate_expert_bytes, - .group_size = q4_group6_expert_group_size, - .n_slots = n_expert, - }; - ds4_gpu_q4_gather_slots6_args down_gather_args = { - .expert_bytes = down_expert_bytes, - .group_size = q4_group6_expert_group_size, - .n_slots = n_expert, - }; - ok = ds4_gpu_encode_q4_gather_slots6(cb, - g_moe_q4_gather_slots6_pipeline, - &gate_gather_args, - gate_group6_bufs, - gate_group6_offsets, - selectedbuf, - ds4_gpu_tensor_offset(selected), - g_moe_q4_gate_slots_buffer, - 0) && - ds4_gpu_encode_q4_gather_slots6(cb, - g_moe_q4_gather_slots6_pipeline, - &gate_gather_args, - up_group6_bufs, - up_group6_offsets, - selectedbuf, - ds4_gpu_tensor_offset(selected), - g_moe_q4_up_slots_buffer, - 0) && - ds4_gpu_encode_q4_gather_slots6(cb, - g_moe_q4_gather_slots6_pipeline, - &down_gather_args, - down_group6_bufs, - down_group6_offsets, - selectedbuf, - ds4_gpu_tensor_offset(selected), - g_moe_q4_down_slots_buffer, - 0); - } - if (use_q4_gather_slots) { - DS4_METAL_PROFILE_MOE_ONE_STAGE("q4_gather"); - } - if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33455); return 0; } - if (use_q4_grouped_experts) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { - const uint32_t expert_count = - q4_expert_group_size < n_total_expert - expert_base ? - q4_expert_group_size : n_total_expert - expert_base; - if ((uint64_t)expert_base > UINT64_MAX / gate_expert_bytes || - (uint64_t)expert_count > UINT64_MAX / gate_expert_bytes) { - ok = 0; - break; - } - const uint64_t group_rel = (uint64_t)expert_base * gate_expert_bytes; - const uint64_t group_bytes = (uint64_t)expert_count * gate_expert_bytes; - if (group_rel > UINT64_MAX - gate_offset || - group_rel > UINT64_MAX - up_offset) { - ok = 0; - break; - } - uint64_t gate_group_inner = 0; - uint64_t up_group_inner = 0; - id gate_group_buf = - q4_grouped_cache_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - gate_offset + group_rel, - group_bytes, - &gate_group_inner) : - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - gate_offset + group_rel, - group_bytes, - &gate_group_inner); - id up_group_buf = - q4_grouped_cache_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - up_offset + group_rel, - group_bytes, - &up_group_inner) : - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - up_offset + group_rel, - group_bytes, - &up_group_inner); - if (!gate_group_buf || !up_group_buf) { - ok = 0; - break; - } - ds4_gpu_moe_expert_group_args group_args = { - .expert_base = expert_base, - .expert_count = expert_count, - .accumulate = 0, - .pad0 = 0, - }; - ok = ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(cb, - g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - &group_args, - gate_group_buf, - (NSUInteger)gate_group_inner, - up_group_buf, - (NSUInteger)up_group_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } - } else if (use_q4_expert_address_table) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(cb, - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_table, - up_table, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } else if (use_q4_expert_table) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_table, - up_table, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false, - q4_table_queue_residency); - } else if (use_q4_group6_experts) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_group6_pair_swiglu(cb, - g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_group6_bufs, - gate_group6_offsets, - up_group6_bufs, - up_group6_offsets, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } else if (use_q4_group8_experts) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_group8_pair_swiglu(cb, - g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_group8_bufs, - gate_group8_offsets, - up_group8_bufs, - up_group8_offsets, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } else if (use_q4_group24_experts) { - ok = ds4_gpu_encode_mul_mv_group24_id(cb, - g_moe_mul_mv_group24_q4_k_id_pipeline, - &gate_args, - gate_group24_bufs, - gate_group24_offsets, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - 2, - false) && - ds4_gpu_encode_mul_mv_group24_id(cb, - g_moe_mul_mv_group24_q4_k_id_pipeline, - &gate_args, - up_group24_bufs, - up_group24_offsets, - xbuf, - ds4_gpu_tensor_offset(x), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - 2, - false); - } else if (use_q4_gather_slots || use_selected_slots) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - if (use_stream_expert_addr_table) { - if (use_stream_expert_masked_addr_table) { - if (use_stream_expert_split_deferred) { - const bool stream_split_profile = - getenv("DS4_METAL_STREAMING_EXPERT_SPLIT_PROFILE") != NULL; - const bool stream_split_timing = - stream_split_profile || - ds4_gpu_stream_expert_timing_summary_enabled(); - double stream_split_t0 = - stream_split_timing ? ds4_gpu_now_ms() : 0.0; - ds4_gpu_stream_expert_split_args resident_pair_args = { - .active_mask = stream_expert_resident_mask, - .accumulate = 0u, - }; - ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, - &gate_args, - &act_args, - &resident_pair_args, - stream_slot_entries, - stream_gate_addr_buf, - stream_up_addr_buf, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selected_exec_buf, - selected_exec_off, - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - if (ok) { - ok = ds4_gpu_flush_commands(); - if (ok) { - cb = ds4_gpu_command_buffer(&owned); - if (!cb) ok = 0; - } - } - const double stream_split_resident_ms = - stream_split_timing ? ds4_gpu_now_ms() - stream_split_t0 : 0.0; - if (stream_split_timing) stream_split_t0 = ds4_gpu_now_ms(); - const double stream_split_missing_start_ms = stream_split_t0; - double stream_split_missing_load_ms = 0.0; - double stream_split_missing_slot_ms = 0.0; - double stream_split_missing_prune_ms = 0.0; - double stream_split_missing_addr_ms = 0.0; - double stream_split_missing_wait_ms = 0.0; - if (ok) { - ok = ds4_gpu_stream_expert_cache_load_selected_missing( - model_map, - model_size, - layer_index, - selected_ids, - n_total_expert, - n_expert, - stream_gate_abs_offsets, - stream_up_abs_offsets, - stream_down_abs_offsets, - gate_expert_bytes, - down_expert_bytes, - stream_expert_missing_mask, - stream_slot_entries); - if (stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - stream_split_missing_load_ms = - now_ms - stream_split_t0; - stream_split_t0 = now_ms; - } - if (ok) { - for (uint32_t i = 0; i < n_expert; i++) { - if ((stream_expert_missing_mask & (1u << i)) == 0) continue; - ds4_gpu_stream_expert_cache_entry *entry = - stream_slot_entries[i]; - if (!entry) { - ok = 0; - break; - } - gate_slot_bufs[i] = entry->gate_buffer; - gate_slot_offsets[i] = entry->gate_inner; - up_slot_bufs[i] = entry->up_buffer; - up_slot_offsets[i] = entry->up_inner; - down_slot_bufs[i] = entry->down_buffer; - down_slot_offsets[i] = entry->down_inner; - } - } - } - if (stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - stream_split_missing_slot_ms = - now_ms - stream_split_t0; - stream_split_t0 = now_ms; - } - if (ok) { - ds4_gpu_stream_expert_cache_prune_layer(layer_index, - n_total_expert, - n_expert, - selected_ids, - n_expert); - ds4_gpu_stream_expert_cache_prune_global(layer_index, - selected_ids, - n_expert); - if (stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - stream_split_missing_prune_ms = - now_ms - stream_split_t0; - stream_split_t0 = now_ms; - } - ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf); - if (stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - stream_split_missing_addr_ms = - now_ms - stream_split_t0; - stream_split_t0 = now_ms; - } - } - if (ok) { - /* - * The resident stage was submitted before the - * CPU read of missing experts so I/O can overlap - * with GPU work. The missing stage reuses the same - * gate/up/mid scratch buffers, so it must not - * execute until the resident command buffer has - * finished. The down/sum pass is issued once after - * all six mid slots exist; this keeps the final - * accumulation order stable regardless of the - * resident/missing split. - */ - ok = ds4_gpu_wait_pending_command_buffers( - "streaming expert split resident"); - if (stream_split_timing) { - const double now_ms = ds4_gpu_now_ms(); - stream_split_missing_wait_ms = - now_ms - stream_split_t0; - stream_split_t0 = now_ms; - } - } - const double stream_split_missing_ms = - stream_split_timing ? - ds4_gpu_now_ms() - stream_split_missing_start_ms : - 0.0; - ds4_gpu_stream_expert_split_args missing_pair_args = { - .active_mask = stream_expert_missing_mask, - .accumulate = 0u, - }; - ds4_gpu_stream_expert_split_args all_down_args = { - .active_mask = 0x3fu, - .accumulate = 0u, - }; - if (ok) { - ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, - &gate_args, - &act_args, - &missing_pair_args, - stream_slot_entries, - stream_gate_addr_buf, - stream_up_addr_buf, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selected_exec_buf, - selected_exec_off, - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } - if (ok) { - ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, - &down_args, - &all_down_args, - stream_slot_entries, - stream_down_addr_buf, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selected_exec_buf, - selected_exec_off, - down_smem, - 2); - } - stream_expert_split_completed = ok; - if (stream_split_timing) { - ds4_gpu_stream_expert_timing_note_split( - stream_expert_resident_mask, - stream_expert_missing_mask, - stream_split_resident_ms, - stream_split_missing_ms); - ds4_gpu_stream_expert_timing_note_split_missing_detail( - stream_split_missing_load_ms, - stream_split_missing_slot_ms, - stream_split_missing_prune_ms, - stream_split_missing_addr_ms, - stream_split_missing_wait_ms); - } - if (stream_split_profile) { - fprintf(stderr, - "ds4: Metal streaming expert split layer=%u " - "resident=0x%02x missing=0x%02x resident_submit=%.3f ms " - "missing_bind=%.3f ms\n", - layer_index, - stream_expert_resident_mask, - stream_expert_missing_mask, - stream_split_resident_ms, - stream_split_missing_ms); - } - } else { - ds4_gpu_stream_expert_split_args split_args = { - .active_mask = 0x3fu, - .accumulate = 0u, - }; - ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, - &gate_args, - &act_args, - &split_args, - stream_slot_entries, - stream_gate_addr_buf, - stream_up_addr_buf, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selected_exec_buf, - selected_exec_off, - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } - } else { - ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb, - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, - &gate_args, - &act_args, - stream_slot_entries, - n_expert, - stream_gate_addr_buf, - stream_up_addr_buf, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selected_exec_buf, - selected_exec_off, - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false, - nil, - nil); - } - } else { - ok = (!use_stream_expert_cache || - ds4_gpu_stream_expert_cache_mark_entries_inflight( - stream_slot_entries, - n_expert, - 0)) && - ds4_gpu_encode_mul_mv_slots6_pair_swiglu(cb, - slots_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_slot_bufs, - gate_slot_offsets, - up_slot_bufs, - up_slot_offsets, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } - } else if (fuse_pair_swiglu) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, - pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } else if (!g_quality_mode && - gate_type == DS4_METAL_TENSOR_IQ2_XXS && - g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { - ok = ds4_gpu_encode_mul_mv_id_pair(cb, - g_moe_mul_mv_id_iq2_xxs_pair_pipeline, - &gate_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - 2, - false); - } else if (!g_quality_mode && - gate_type == DS4_METAL_TENSOR_Q4_K && - g_moe_mul_mv_id_q4_k_pair_pipeline) { - ok = ds4_gpu_encode_mul_mv_id_pair(cb, - g_moe_mul_mv_id_q4_k_pair_pipeline, - &gate_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - 2, - false); - } else { - ok = ds4_gpu_encode_mul_mv_id(cb, - gate_mv_pipeline, - &gate_args, - gate_buf, - (NSUInteger)gate_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - gate_nsg, - gate_rows_per_group_is_nr0) && - ds4_gpu_encode_mul_mv_id(cb, - gate_mv_pipeline, - &gate_args, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - gate_nsg, - gate_rows_per_group_is_nr0); - } - DS4_METAL_PROFILE_MOE_ONE_STAGE("gate_up"); - if (ok && (!fuse_pair_swiglu || use_q4_group24_experts)) { - ok = ds4_gpu_encode_moe_swiglu_weight(cb, - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - expert_mid_dim, - pair_rows, - clamp, - false); - } - DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight"); - - id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); - NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : - (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); - if (ok && stream_expert_split_completed) { - /* The split path already wrote the resident partial output and - * accumulated the missing experts into out. */ - } else if (ok && use_q4_grouped_experts) { - bool first_group = true; - for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { - const uint32_t expert_count = - q4_expert_group_size < n_total_expert - expert_base ? - q4_expert_group_size : n_total_expert - expert_base; - if ((uint64_t)expert_base > UINT64_MAX / down_expert_bytes || - (uint64_t)expert_count > UINT64_MAX / down_expert_bytes) { - ok = 0; - break; - } - const uint64_t group_rel = (uint64_t)expert_base * down_expert_bytes; - const uint64_t group_bytes = (uint64_t)expert_count * down_expert_bytes; - if (group_rel > UINT64_MAX - down_offset) { - ok = 0; - break; - } - uint64_t down_group_inner = 0; - id down_group_buf = - q4_grouped_cache_views ? - ds4_gpu_wrap_model_exact_range(model_map, - model_size, - down_offset + group_rel, - group_bytes, - &down_group_inner) : - ds4_gpu_wrap_model_exact_range_transient(model_map, - model_size, - down_offset + group_rel, - group_bytes, - &down_group_inner); - if (!down_group_buf) { - ok = 0; - break; - } - ds4_gpu_moe_expert_group_args group_args = { - .expert_base = expert_base, - .expert_count = expert_count, - .accumulate = first_group ? 0u : 1u, - .pad0 = 0, - }; - ok = ds4_gpu_encode_mul_mv_group_q4_sum6(cb, - g_moe_mul_mv_group_q4_k_sum6_pipeline, - &down_args, - &group_args, - down_group_buf, - (NSUInteger)down_group_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2); - first_group = false; - } - } else if (ok && use_q4_expert_address_table) { - ok = ds4_gpu_encode_mul_mv_addr_q4_sum6(cb, - g_moe_mul_mv_addr_q4_k_sum6_pipeline, - &down_args, - down_table, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2); - } else if (ok && use_q4_expert_table) { - ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, - g_moe_mul_mv_table_q4_k_sum6_pipeline, - &down_args, - down_table, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2, - q4_table_queue_residency); - } else if (ok && use_q4_group6_experts) { - ok = ds4_gpu_encode_mul_mv_group6_sum6(cb, - g_moe_mul_mv_group6_q4_k_sum6_pipeline, - &down_args, - down_group6_bufs, - down_group6_offsets, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2); - } else if (ok && use_q4_group8_experts) { - ok = ds4_gpu_encode_mul_mv_group8_sum6(cb, - g_moe_mul_mv_group8_q4_k_sum6_pipeline, - &down_args, - down_group8_bufs, - down_group8_offsets, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2); - } else if (ok && use_q4_group24_experts) { - ok = ds4_gpu_encode_mul_mv_group24_sum6(cb, - g_moe_mul_mv_group24_q4_k_sum6_pipeline, - &down_args, - down_group24_bufs, - down_group24_offsets, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2); - } else if (ok && (use_q4_gather_slots || use_selected_slots)) { - if (use_stream_expert_addr_table) { - if (down_type == DS4_METAL_TENSOR_IQ2_XXS) { - ok = ds4_gpu_encode_mul_mv_addr_iq2(cb, - g_moe_mul_mv_addr_iq2_xxs_pipeline, - &down_args, - stream_slot_entries, - n_expert, - stream_down_addr_buf, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - selected_exec_buf, - selected_exec_off, - down_smem, - 2, - false); - } else if (use_stream_expert_masked_addr_table) { - ds4_gpu_stream_expert_split_args split_args = { - .active_mask = 0x3fu, - .accumulate = 0u, - }; - ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, - g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, - &down_args, - &split_args, - stream_slot_entries, - stream_down_addr_buf, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selected_exec_buf, - selected_exec_off, - down_smem, - 2); - } else { - ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(cb, - g_moe_mul_mv_addr_q2_k_sum6_pipeline, - &down_args, - stream_slot_entries, - n_expert, - stream_down_addr_buf, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selected_exec_buf, - selected_exec_off, - down_smem, - 2, - nil); - } - } else { - ok = (!use_stream_expert_cache || - ds4_gpu_stream_expert_cache_mark_entries_inflight( - stream_slot_entries, - n_expert, - 0)) && - ds4_gpu_encode_mul_mv_slots6_sum6(cb, - slots_sum6_pipeline, - &down_args, - down_slot_bufs, - down_slot_offsets, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - down_smem, - 2); - } - } else if (ok && direct_down_sum) { - ok = ds4_gpu_encode_mul_mv_id_sum6(cb, - down_sum6_pipeline, - &down_args, - down_buf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - add_in ? ds4_gpu_tensor_buffer(add_in) : nil, - add_in ? ds4_gpu_tensor_offset(add_in) : 0, - down_smem, - 2); - } else if (ok) { - ok = ds4_gpu_encode_mul_mv_id(cb, - down_mv_pipeline, - &down_args, - down_buf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - down_nsg, - down_rows_per_group_is_nr0); - } - DS4_METAL_PROFILE_MOE_ONE_STAGE("down"); - if (ok && n_expert > 1 && !direct_down_sum && !stream_expert_split_completed) { - ok = ds4_gpu_encode_moe_sum_experts(cb, - down_dst, - down_dst_off, - outbuf, - ds4_gpu_tensor_offset(out), - out_dim, - n_expert, - n_tokens); - } - DS4_METAL_PROFILE_MOE_ONE_STAGE("sum"); - if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; } - if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34372); - return 0; - } - } -#undef DS4_METAL_PROFILE_MOE_ONE_STAGE - } - - return 1; -} - -int ds4_gpu_routed_moe_batch_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *gate, - ds4_gpu_tensor *up, - ds4_gpu_tensor *mid, - ds4_gpu_tensor *experts, - const void *model_map, - uint64_t model_size, - uint64_t gate_offset, - uint64_t up_offset, - uint64_t down_offset, - uint32_t gate_type, - uint32_t down_type, - uint64_t gate_expert_bytes, - uint64_t gate_row_bytes, - uint64_t down_expert_bytes, - uint64_t down_row_bytes, - uint32_t expert_in_dim, - uint32_t expert_mid_dim, - uint32_t out_dim, - const ds4_gpu_tensor *selected, - const ds4_gpu_tensor *weights, - uint32_t n_total_expert, - uint32_t n_expert, - float clamp, - const ds4_gpu_tensor *x, - uint32_t layer_index, - uint32_t n_tokens, - bool *mid_is_f16, - bool force_resident) { - (void)force_resident; - if (!g_initialized && !ds4_gpu_init()) return 0; - /* TP sharding (see ds4_gpu_routed_moe_one_tensor): bind from the owned - * expert range and rebase ids in the kernels. */ - uint32_t first_expert = 0; - uint32_t n_bind_expert = 0; - ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); - const int32_t tp_expert_base_host = (int32_t)first_expert; - gate_offset += (uint64_t)first_expert * gate_expert_bytes; - up_offset += (uint64_t)first_expert * gate_expert_bytes; - down_offset += (uint64_t)first_expert * down_expert_bytes; - if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || - n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || - n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED) { - return 0; - } - if (gate_expert_bytes == 0 || down_expert_bytes == 0 || - gate_row_bytes == 0 || down_row_bytes == 0) { - return 0; - } - if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) return 0; - if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || - (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { - fprintf(stderr, "ds4: Metal routed batch MoE tensor byte size overflow\n"); - return 0; - } - const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; - const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; - - /* - * PRO Q4 routed expert tensors are multi-GiB per layer. A one-token - * layer-slice prefill should use the one-token path so it can either bind - * exact tensor views with GPU-selected IDs or fall back to selected-expert - * views. Keep this guarded by tensor size so Flash and mixed-Flash Q4 keep - * their existing fast path. - */ - const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; - const bool can_single_token_q4_grouped = - getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL && - g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group_q4_k_sum6_pipeline != nil; - const bool can_single_token_q4_selected_slots = - ds4_gpu_q4_selected_paths_allowed() && - getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL && - g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil; - const bool can_single_token_q4_group6 = - n_total_expert == 384 && - getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL && - g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil; - const bool can_single_token_q4_group8 = - n_total_expert == 384 && - getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL && - g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil; - const bool can_single_token_q4_group24 = - n_total_expert == 384 && - getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL && - g_moe_mul_mv_group24_q4_k_id_pipeline != nil && - g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil; - const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; - const bool can_single_token_q4_exact_tensor_id = - n_total_expert == 384 && - max_buffer_len != 0 && - gate_tensor_bytes <= max_buffer_len && - down_tensor_bytes <= max_buffer_len && - getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && - getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; - const bool enable_single_token_q4_expert_table = - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || - ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes); - const bool can_single_token_q4_expert_table = - n_total_expert == 384 && - enable_single_token_q4_expert_table && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && - g_moe_table_q4_pair_gate_encoder != nil && - g_moe_table_q4_pair_up_encoder != nil && - g_moe_table_q4_sum_down_encoder != nil; - const bool can_single_token_q4_expert_address_table = - n_total_expert == 384 && - getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && - g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil; - const bool use_single_token_q4_one_tensor = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_tokens == 1 && - n_expert == 6 && - n_total_expert >= 128 && - (g_ssd_streaming_mode || - (gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes)) && - !g_quality_mode && - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && - (can_single_token_q4_grouped || - can_single_token_q4_group6 || - can_single_token_q4_group8 || - can_single_token_q4_group24 || - can_single_token_q4_exact_tensor_id || - can_single_token_q4_expert_address_table || - can_single_token_q4_expert_table || - can_single_token_q4_selected_slots); - if (use_single_token_q4_one_tensor) { - if (mid_is_f16) *mid_is_f16 = false; - return ds4_gpu_routed_moe_one_tensor(out, - gate, - up, - mid, - experts, - model_map, - model_size, - gate_offset, - up_offset, - down_offset, - gate_type, - down_type, - gate_expert_bytes, - gate_row_bytes, - down_expert_bytes, - down_row_bytes, - expert_in_dim, - expert_mid_dim, - out_dim, - selected, - weights, - n_total_expert, - n_expert, - clamp, - x, - NULL, - layer_index, - false); - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id gatebuf = ds4_gpu_tensor_buffer(gate); - id upbuf = ds4_gpu_tensor_buffer(up); - id midbuf = ds4_gpu_tensor_buffer(mid); - id outbuf = ds4_gpu_tensor_buffer(out); - id expertsbuf = ds4_gpu_tensor_buffer(experts); - id selectedbuf = ds4_gpu_tensor_buffer(selected); - id weightsbuf = ds4_gpu_tensor_buffer(weights); - const uint64_t x_bytes = (uint64_t)n_tokens * expert_in_dim * sizeof(float); - const uint64_t mid_bytes = (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float); - const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); - const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert * sizeof(int); - const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); - if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(gate) < mid_bytes || - ds4_gpu_tensor_bytes(up) < mid_bytes || - ds4_gpu_tensor_bytes(mid) < mid_bytes || - ds4_gpu_tensor_bytes(out) < out_bytes || - ds4_gpu_tensor_bytes(selected) < selected_bytes || - ds4_gpu_tensor_bytes(weights) < weights_bytes) { - fprintf(stderr, "ds4: Metal routed batch MoE received undersized activation buffers\n"); - return 0; - } - uint64_t gate_inner = 0; - uint64_t up_inner = 0; - uint64_t down_inner = 0; - id gate_buf = nil; - id up_buf = nil; - id down_buf = nil; - DS4MetalQ4ExpertTable *gate_table = nil; - DS4MetalQ4ExpertTable *up_table = nil; - DS4MetalQ4ExpertTable *down_table = nil; - id q4_table_layer_residency = nil; - id stream_gate_addr_buf = nil; - id stream_up_addr_buf = nil; - id stream_down_addr_buf = nil; - id stream_overflow_gate = nil; - id stream_overflow_up = nil; - id stream_overflow_down = nil; - ds4_gpu_stream_expert_cache_entry - *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; - uint32_t stream_resource_count = 0; - uint32_t stream_unique = 0; - - const uint32_t pair_rows = n_tokens * n_expert; - const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); - - const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); - const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); - id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); - id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); - id gate_mm_pipeline = nil; - id up_mm_pipeline = nil; - id down_mm_pipeline = nil; - id pair_swiglu_mm_pipeline = nil; - if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { - fprintf(stderr, "ds4: unsupported Metal routed batch MoE quant types gate=%u down=%u\n", - gate_type, down_type); - return 0; - } - const bool use_iq2_batch_selected_addr = - ds4_gpu_stream_prefill_batch_selected_addr_enabled(n_tokens, - n_total_expert, - n_expert, - gate_type, - down_type) && - n_tokens > 1 && - n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && - !g_quality_mode && - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && - g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; - - ds4_gpu_mul_mv_id_args gate_args = - ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - gate_row_bytes, gate_expert_bytes, - 1, n_expert, n_tokens, gate_nr0); - gate_args.tp_rank = g_tp_split_rank; - gate_args.tp_world = g_tp_split_world; - gate_args.tp_expert_base = tp_expert_base_host; - ds4_gpu_mul_mv_id_args down_args = - ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, - down_row_bytes, down_expert_bytes, - n_expert, n_expert, n_tokens, down_nr0); - down_args.tp_rank = g_tp_split_rank; - down_args.tp_world = g_tp_split_world; - down_args.tp_expert_base = tp_expert_base_host; - const bool q4_batch_expert_table_auto = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, - n_expert, - gate_tensor_bytes, - down_tensor_bytes); - const bool q4_batch_table_queue_residency = - ds4_gpu_q4_table_queue_residency_enabled(q4_batch_expert_table_auto); - const bool enable_q4_batch_expert_table = - getenv("DS4_METAL_ENABLE_Q4_BATCH_EXPERT_TABLE") != NULL || - getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || - q4_batch_expert_table_auto; - const bool use_q4_batch_expert_table = - gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K && - n_expert == 6 && - n_tokens > 1 && - n_total_expert == 384 && - gate_tensor_bytes >= q4_selected_min_tensor_bytes && - down_tensor_bytes >= q4_selected_min_tensor_bytes && - !g_quality_mode && - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && - getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && - getenv("DS4_METAL_DISABLE_Q4_BATCH_EXPERT_TABLE") == NULL && - getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && - g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && - g_moe_table_q4_pair_gate_encoder != nil && - g_moe_table_q4_pair_up_encoder != nil && - g_moe_table_q4_sum_down_encoder != nil && - enable_q4_batch_expert_table && - (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || - q4_batch_table_queue_residency || - ds4_gpu_q4_table_model_residency_enabled()); - const bool use_mm_id = - !use_q4_batch_expert_table && - !use_iq2_batch_selected_addr && - n_tokens >= 32u && - ds4_gpu_mul_mm_id_map0_name(n_expert) != NULL; - /* - * MTP verification is neither normal decode nor large prefill: the - * target model must verify a tiny suffix (up to DSpark's 5-token - * block) in one layer-major pass. For that shape the prefill - * expert-major GEMM path - * is too large, but the decode pair kernels are exactly the right - * primitive: they read the same activation once and compute routed - * gate/up together for every selected expert row. Keep this limited to - * tiny batches so ordinary prefill keeps using the higher-throughput - * grouped matmul path. - */ - const bool use_tiny_pair_mv = - !g_quality_mode && - n_tokens <= 5u && - !use_q4_batch_expert_table && - !use_mm_id && - ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || - (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline)); - id tiny_pair_swiglu_pipeline = nil; - if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { - tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; - } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { - tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; - } - const bool use_tiny_pair_swiglu = - use_tiny_pair_mv && - tiny_pair_swiglu_pipeline != nil && - getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL && - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL; - ds4_gpu_mul_mm_id_map_args gate_map_args = { 0 }; - ds4_gpu_mul_mm_id_args gate_mm_args = { 0 }; - ds4_gpu_mul_mm_id_args down_mm_args = { 0 }; - id map_pipeline = nil; - /* - * The grouped routed-MoE matmul loads activation tiles as half before - * using SIMD-group MMA. Store the SwiGLU/route-weight intermediate in - * that same precision so the down projection avoids a large F32 mid - * write/read. --quality keeps the older F32 intermediate. - */ - const bool request_mid_f16 = - !g_quality_mode && - !use_q4_batch_expert_table && - !use_iq2_batch_selected_addr; - /* - * Fused gate+up grouped matmul with the SwiGLU epilogue. The IQ2 - * variant stays opt-in behind its env flag; the Q4_K variant is the - * default path — same MMA accumulation order and epilogue math as the - * separate GEMMs + swiglu pass, so the mid tensor is bit-identical. - */ - const bool use_mm_id_pair_swiglu = - use_mm_id && - g_tp_split_world != 2 && /* pair-swiglu mm kernel lacks expert ownership */ - request_mid_f16 && - n_expert == 6 && - ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && - down_type == DS4_METAL_TENSOR_Q2_K && - getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) || - (gate_type == DS4_METAL_TENSOR_Q4_K && - down_type == DS4_METAL_TENSOR_Q4_K)) && - getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && - getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && - getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL; - if (use_mm_id) { - gate_map_args = - ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); - gate_mm_args = - ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, - gate_row_bytes, gate_expert_bytes, - 1, n_expert, n_tokens); - down_mm_args = - ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, - down_row_bytes, down_expert_bytes, - n_expert, n_expert, n_tokens, - request_mid_f16 ? sizeof(uint16_t) : sizeof(float)); - gate_mm_args.tp_rank = g_tp_split_rank; - gate_mm_args.tp_world = g_tp_split_world; - gate_mm_args.tp_expert_base = tp_expert_base_host; - down_mm_args.tp_rank = g_tp_split_rank; - down_mm_args.tp_world = g_tp_split_world; - down_mm_args.tp_expert_base = tp_expert_base_host; - - map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); - gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); - up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); - down_mm_pipeline = request_mid_f16 ? - ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) : - ds4_gpu_routed_mm_pipeline(down_type); - if (use_mm_id_pair_swiglu) { - pair_swiglu_mm_pipeline = - ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ? - "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : - "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); - } - if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline || - (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) { - return 0; - } - } - - if (use_iq2_batch_selected_addr) { - const int had_batch = g_batch_cb != nil; - if (had_batch && ds4_gpu_end_commands() == 0) { - return 0; - } - g_stream_prefill_batch_selected_addr_building++; - if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( - model_map, - model_size, - layer_index, - selected, - n_tokens, - n_total_expert, - n_expert, - gate_offset, - up_offset, - down_offset, - gate_expert_bytes, - down_expert_bytes, - &stream_gate_addr_buf, - &stream_up_addr_buf, - &stream_down_addr_buf, - stream_resources, - &stream_resource_count, - &stream_unique, - &stream_overflow_gate, - &stream_overflow_up, - &stream_overflow_down)) { - g_stream_prefill_batch_selected_addr_building--; - return 0; - } - g_stream_prefill_batch_selected_addr_building--; - if (stream_unique == 0) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "produced no resident experts\n", - layer_index); - return 0; - } - for (uint32_t i = 0; i < stream_resource_count; i++) { - ds4_gpu_stream_expert_cache_entry *entry = stream_resources[i]; - if (!entry || - !entry->valid || - !entry->gate_buffer || - !entry->up_buffer || - !entry->down_buffer) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "lost resident expert resource %u/%u during preparation\n", - layer_index, - i, - stream_resource_count); - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - } - if (had_batch && ds4_gpu_begin_commands() == 0) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "failed to reopen command batch after preparation\n", - layer_index); - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - } - - if (use_q4_batch_expert_table) { - gate_table = ds4_gpu_q4_expert_table(model_map, - model_size, - gate_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_gate_encoder); - up_table = ds4_gpu_q4_expert_table(model_map, - model_size, - up_offset, - gate_expert_bytes, - n_total_expert, - g_moe_table_q4_pair_up_encoder); - down_table = ds4_gpu_q4_expert_table(model_map, - model_size, - down_offset, - down_expert_bytes, - n_total_expert, - g_moe_table_q4_sum_down_encoder); - if (!gate_table || !up_table || !down_table) { - return 0; - } - q4_table_layer_residency = - ds4_gpu_q4_expert_layer_residency_set(gate_table, - up_table, - down_table, - q4_batch_expert_table_auto); - if ((q4_batch_table_queue_residency || - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL) && - !q4_table_layer_residency && - !ds4_gpu_q4_table_model_residency_enabled()) { - fprintf(stderr, "ds4: Metal Q4 batch expert table residency set is not available\n"); - return 0; - } - } else if (use_iq2_batch_selected_addr) { - if (n_expert > 1 && (!expertsbuf || - ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, - &g_moe_down_scratch_bytes, - (NSUInteger)down_scratch_bytes, - "ds4_moe_down_scratch")) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - } - } else { - if (n_expert > 1 && (!expertsbuf || - ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { - if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, - &g_moe_down_scratch_bytes, - (NSUInteger)down_scratch_bytes, - "ds4_moe_down_scratch")) { - return 0; - } - } - gate_buf = ds4_gpu_wrap_model_range(model_map, - model_size, - gate_offset, - gate_tensor_bytes, - &gate_inner); - up_buf = ds4_gpu_wrap_model_range(model_map, - model_size, - up_offset, - gate_tensor_bytes, - &up_inner); - down_buf = ds4_gpu_wrap_model_range(model_map, - model_size, - down_offset, - down_tensor_bytes, - &down_inner); - if (!gate_buf || !up_buf || !down_buf) return 0; - if (getenv("DS4_GLM_TP_DEBUG") && layer_index == 3) { - fprintf(stderr, - "ds4: batch mv binds l=%u rank=%d base=%d gate=%llu+%llu " - "up=%llu down=%llu inner=%llu/%llu/%llu ne02=%d nei0=%d nr0=%d\n", - layer_index, g_tp_split_rank, tp_expert_base_host, - (unsigned long long)gate_offset, - (unsigned long long)gate_tensor_bytes, - (unsigned long long)up_offset, - (unsigned long long)down_offset, - (unsigned long long)gate_inner, - (unsigned long long)up_inner, - (unsigned long long)down_inner, - gate_args.ne02, gate_args.nei0, gate_args.nr0); - } - } - - const bool q4_batch_table_boundary = - use_q4_batch_expert_table && - g_batch_cb != nil && - getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && - !q4_batch_table_queue_residency && - getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; - if (q4_batch_table_boundary) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - if (use_q4_batch_expert_table && !ds4_gpu_use_model_residency_set(cb)) { - return 0; - } - if (!q4_batch_table_queue_residency && - q4_table_layer_residency && - [cb respondsToSelector:@selector(useResidencySet:)]) { - [cb useResidencySet:q4_table_layer_residency]; - } - const bool moe_stage_profile = - g_batch_cb != nil && - ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_STAGE_PROFILE", - "DS4_METAL_MOE_STAGE_PROFILE_LAYER", - layer_index); - const char *moe_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); - const char *moe_path = - use_q4_batch_expert_table ? "q4_table_pair_swiglu" : - use_iq2_batch_selected_addr ? "iq2_batch_stream_addr" : - use_mm_id_pair_swiglu ? "mm_id_pair_swiglu" : - use_mm_id ? "mm_id" : - use_tiny_pair_swiglu ? "tiny_pair_swiglu" : - (use_tiny_pair_mv ? "tiny_pair_mv" : "mv"); - double moe_stage_t0 = moe_stage_profile ? ds4_gpu_now_ms() : 0.0; - if (moe_stage_profile) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - moe_stage_t0 = ds4_gpu_now_ms(); - } -#define DS4_METAL_PROFILE_MOE_STAGE(name) do { \ - if (ok && moe_stage_profile) { \ - if (ds4_gpu_end_commands() == 0) { \ - ok = 0; \ - } else { \ - const char *stage_name = (name); \ - const double now_ms = ds4_gpu_now_ms(); \ - const int print_stage = \ - !moe_stage_filter || !moe_stage_filter[0] || \ - strstr(stage_name, moe_stage_filter) != NULL; \ - if (print_stage) { \ - fprintf(stderr, \ - "ds4: Metal routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ - "gate=%s down=%s path=%s mid=%s %s=%.3f ms\n", \ - layer_index, n_tokens, pair_rows, n_expert, \ - ds4_gpu_metal_tensor_type_name(gate_type), \ - ds4_gpu_metal_tensor_type_name(down_type), \ - moe_path, \ - request_mid_f16 ? "f16" : "f32", \ - stage_name, now_ms - moe_stage_t0); \ - } \ - moe_stage_t0 = now_ms; \ - if (ds4_gpu_begin_commands() == 0) { \ - ok = 0; \ - } else { \ - cb = ds4_gpu_command_buffer(&owned); \ - if (!cb) ok = 0; \ - } \ - } \ - } \ - } while (0) - - const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); - const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); - const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); - const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); - const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); - const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); - id down_sum6_pipeline = nil; - if (down_type == DS4_METAL_TENSOR_Q2_K) { - down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; - } else if (down_type == DS4_METAL_TENSOR_Q4_K) { - down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; - } - const bool direct_down_sum = - !g_quality_mode && - !use_q4_batch_expert_table && - !use_mm_id && - n_expert == 6 && - n_tokens <= 4u && - down_sum6_pipeline != nil; - int ok = 0; - if (use_iq2_batch_selected_addr) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( - cb, - g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, - &gate_args, - &act_args, - stream_resources, - stream_resource_count, - stream_gate_addr_buf, - stream_up_addr_buf, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false, - stream_overflow_gate, - stream_overflow_up); - if (!ok) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "failed to encode gate/up path tokens=%u unique=%u\n", - layer_index, - n_tokens, - stream_unique); - } - } else if (use_q4_batch_expert_table) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, - g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_table, - up_table, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false, - q4_batch_table_queue_residency); - } else if (use_mm_id) { - /* - * The routed pair ids are the same for gate, up, and down. Build - * the expert-major work map once, then reuse it for all three - * batched expert matmuls. - */ - ok = ds4_gpu_encode_mul_mm_id_map(cb, - map_pipeline, - &gate_map_args, - &gate_mm_args, - selectedbuf, - ds4_gpu_tensor_offset(selected)); - DS4_METAL_PROFILE_MOE_STAGE("map"); - if (ok && use_mm_id_pair_swiglu) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(uint16_t), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb, - pair_swiglu_mm_pipeline, - &gate_mm_args, - &act_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights)); - DS4_METAL_PROFILE_MOE_STAGE("gate_up_fused"); - } else if (ok) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - gate_mm_pipeline, - &gate_mm_args, - gate_buf, - (NSUInteger)gate_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - 8192u); - DS4_METAL_PROFILE_MOE_STAGE("gate"); - } - if (ok && !use_mm_id_pair_swiglu) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - up_mm_pipeline, - &gate_mm_args, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - upbuf, - ds4_gpu_tensor_offset(up), - 8192u); - DS4_METAL_PROFILE_MOE_STAGE("up"); - } - } else if (use_tiny_pair_swiglu) { - ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { - .width = expert_mid_dim, - .rows = pair_rows, - .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), - .weight_stride = sizeof(float), - .write_clamped = 0, - .clamp_value = clamp, - }; - ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, - tiny_pair_swiglu_pipeline, - &gate_args, - &act_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - selectedbuf, - ds4_gpu_tensor_offset(selected), - weightsbuf, - ds4_gpu_tensor_offset(weights), - gate_smem, - 2, - false); - } else if (use_tiny_pair_mv) { - id pair_pipeline = - gate_type == DS4_METAL_TENSOR_IQ2_XXS ? - g_moe_mul_mv_id_iq2_xxs_pair_pipeline : - g_moe_mul_mv_id_q4_k_pair_pipeline; - ok = ds4_gpu_encode_mul_mv_id_pair(cb, - pair_pipeline, - &gate_args, - gate_buf, - (NSUInteger)gate_inner, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - 2, - false); - } else { - ok = ds4_gpu_encode_mul_mv_id(cb, - gate_mv_pipeline, - &gate_args, - gate_buf, - (NSUInteger)gate_inner, - xbuf, - ds4_gpu_tensor_offset(x), - gatebuf, - ds4_gpu_tensor_offset(gate), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - gate_nsg, - gate_rows_per_group_is_nr0) && - ds4_gpu_encode_mul_mv_id(cb, - gate_mv_pipeline, - &gate_args, - up_buf, - (NSUInteger)up_inner, - xbuf, - ds4_gpu_tensor_offset(x), - upbuf, - ds4_gpu_tensor_offset(up), - selectedbuf, - ds4_gpu_tensor_offset(selected), - gate_smem, - gate_nsg, - gate_rows_per_group_is_nr0); - } - DS4_METAL_PROFILE_MOE_STAGE("gate_up"); - const bool use_fused_activation = !g_quality_mode && !use_q4_batch_expert_table; - const bool use_mid_f16 = - use_mm_id && - use_fused_activation && - request_mid_f16; - if (mid_is_f16) *mid_is_f16 = use_mid_f16; - if (ok && use_iq2_batch_selected_addr) { - /* The address-table pair kernel already wrote weighted SwiGLU rows into mid. */ - } else if (ok && use_q4_batch_expert_table) { - /* The table pair kernel already wrote weighted SwiGLU rows into mid. */ - } else if (ok && use_mm_id_pair_swiglu) { - /* The fused batch mm_id pair kernel already wrote weighted f16 SwiGLU rows into mid. */ - } else if (ok && use_tiny_pair_swiglu) { - /* The fused tiny pair kernel already wrote weighted F32 SwiGLU rows into mid. */ - } else if (ok && use_fused_activation) { - ok = ds4_gpu_encode_moe_swiglu_weight(cb, - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - expert_mid_dim, - pair_rows, - clamp, - use_mid_f16); - } else if (ok && clamp > 1.0e-6f) { - ok = ds4_gpu_encode_unary_f32_rows(cb, - g_unary_clamp_pipeline, - gatebuf, - ds4_gpu_tensor_offset(gate), - gatebuf, - ds4_gpu_tensor_offset(gate), - expert_mid_dim, - pair_rows, - 0, - -FLT_MAX, - clamp); - if (ok) { - ok = ds4_gpu_encode_unary_f32_rows(cb, - g_unary_silu_pipeline, - gatebuf, - ds4_gpu_tensor_offset(gate), - midbuf, - ds4_gpu_tensor_offset(mid), - expert_mid_dim, - pair_rows, - 1, - 0.0f, - 0.0f); - } - if (ok) { - ok = ds4_gpu_encode_unary_f32_rows(cb, - g_unary_clamp_pipeline, - upbuf, - ds4_gpu_tensor_offset(up), - upbuf, - ds4_gpu_tensor_offset(up), - expert_mid_dim, - pair_rows, - 0, - -clamp, - clamp); - } - if (ok) { - ds4_gpu_bin_args mul_args = - ds4_gpu_make_bin_same_rows_args(expert_mid_dim, pair_rows); - ok = ds4_gpu_encode_bin_f32_rows(cb, - g_mul_pipeline, - &mul_args, - midbuf, - ds4_gpu_tensor_offset(mid), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid)); - } - } else if (ok) { - ok = ds4_gpu_encode_swiglu_flat(cb, - gatebuf, - ds4_gpu_tensor_offset(gate), - upbuf, - ds4_gpu_tensor_offset(up), - midbuf, - ds4_gpu_tensor_offset(mid), - (uint32_t)((uint64_t)pair_rows * expert_mid_dim)); - } - if (ok && !use_fused_activation && !use_q4_batch_expert_table) { - ds4_gpu_bin_args weight_args = - ds4_gpu_make_bin_rowwise_scalar_args(expert_mid_dim, pair_rows); - ok = ds4_gpu_encode_bin_f32_rows(cb, - g_bin_mul_scalar_pipeline, - &weight_args, - midbuf, - ds4_gpu_tensor_offset(mid), - weightsbuf, - ds4_gpu_tensor_offset(weights), - midbuf, - ds4_gpu_tensor_offset(mid)); - } - DS4_METAL_PROFILE_MOE_STAGE("activation_weight"); - - id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); - NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : - (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); - if (ok) { - if (use_iq2_batch_selected_addr) { - ok = ds4_gpu_encode_mul_mv_addr_q2_sum6( - cb, - g_moe_mul_mv_addr_q2_k_sum6_pipeline, - &down_args, - stream_resources, - stream_resource_count, - stream_down_addr_buf, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2, - stream_overflow_down); - if (!ok) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "failed to encode down path tokens=%u unique=%u\n", - layer_index, - n_tokens, - stream_unique); - } - } else if (use_q4_batch_expert_table) { - ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, - g_moe_mul_mv_table_q4_k_sum6_pipeline, - &down_args, - down_table, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - 2, - q4_batch_table_queue_residency); - } else if (direct_down_sum) { - ok = ds4_gpu_encode_mul_mv_id_sum6(cb, - down_sum6_pipeline, - &down_args, - down_buf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - outbuf, - ds4_gpu_tensor_offset(out), - selectedbuf, - ds4_gpu_tensor_offset(selected), - nil, - 0, - down_smem, - 2); - } else if (use_mm_id) { - ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, - down_mm_pipeline, - &down_mm_args, - down_buf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - 8192u); - } else { - ok = ds4_gpu_encode_mul_mv_id(cb, - down_mv_pipeline, - &down_args, - down_buf, - (NSUInteger)down_inner, - midbuf, - ds4_gpu_tensor_offset(mid), - down_dst, - down_dst_off, - selectedbuf, - ds4_gpu_tensor_offset(selected), - down_smem, - down_nsg, - down_rows_per_group_is_nr0); - } - } - DS4_METAL_PROFILE_MOE_STAGE("down"); - if (ok && - n_expert > 1 && - !direct_down_sum && - !use_q4_batch_expert_table && - !use_iq2_batch_selected_addr) { - ok = ds4_gpu_encode_moe_sum_experts(cb, - down_dst, - down_dst_off, - outbuf, - ds4_gpu_tensor_offset(out), - out_dim, - n_expert, - n_tokens); - } - DS4_METAL_PROFILE_MOE_STAGE("sum"); - if (!ok) { - fprintf(stderr, - "ds4: Metal routed batch MoE failed before submit layer=%u tokens=%u " - "gate=%s down=%s path=%s\n", - layer_index, - n_tokens, - ds4_gpu_metal_tensor_type_name(gate_type), - ds4_gpu_metal_tensor_type_name(down_type), - moe_path); - return 0; - } - - if (!ds4_gpu_finish_command_buffer(cb, owned, "routed batch MoE")) { - if (use_iq2_batch_selected_addr) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - } - return 0; - } - if (use_iq2_batch_selected_addr) { - if (!owned) { - if (ds4_gpu_end_commands() == 0) { - ds4_gpu_stream_expert_cache_clear_layer(layer_index); - return 0; - } - if (ds4_gpu_begin_commands() == 0) { - fprintf(stderr, - "ds4: Metal streaming prefill batch selected addr layer=%u " - "failed to reopen command batch after execution\n", - layer_index); - return 0; - } - } - } - if (q4_batch_table_boundary) { - if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { - return 0; - } - } -#undef DS4_METAL_PROFILE_MOE_STAGE - } - - return 1; -} - -int ds4_gpu_hc_split_sinkhorn_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *mix, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint32_t n_hc, - uint32_t sinkhorn_iters, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (n_hc == 0 || n_hc > 16) return 0; - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t mix_bytes = mix_hc * sizeof(float); - const uint64_t scale_bytes = 3ull * sizeof(float); - - @autoreleasepool { - id mixbuf = ds4_gpu_tensor_buffer(mix); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t mix_tensor_bytes = ds4_gpu_tensor_bytes(mix); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); - if (!mixbuf || !outbuf || - mix_tensor_bytes < mix_bytes || - out_tensor_bytes < mix_bytes) { - fprintf(stderr, "ds4: Metal HC split received undersized activation buffers\n"); - return 0; - } - if (scale_offset > model_size || scale_bytes > model_size - scale_offset || - base_offset > model_size || mix_bytes > model_size - base_offset) { - fprintf(stderr, "ds4: Metal HC split parameter range is outside the mapped model\n"); - return 0; - } - - uint64_t scale_inner = 0; - uint64_t base_inner = 0; - id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); - id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); - if (!scalebuf || !basebuf) return 0; - - uint64_t n_rows64 = mix_tensor_bytes / mix_bytes; - const uint64_t out_rows64 = out_tensor_bytes / mix_bytes; - if (out_rows64 < n_rows64) n_rows64 = out_rows64; - if (n_rows64 == 0 || n_rows64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC split row count is outside supported range\n"); - return 0; - } - - ds4_gpu_hc_split_args args = { - .n_hc = (int32_t)n_hc, - .sinkhorn_iters = (int32_t)sinkhorn_iters, - .n_rows = (int64_t)n_rows64, - .mix_hc = (int64_t)mix_hc, - .nb01 = mix_bytes, - .nb1 = mix_bytes, - .eps = eps, - }; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_rows64)); - const NSUInteger n_tg = ((NSUInteger)n_rows64 + nth - 1u) / nth; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_hc_split_sinkhorn_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; - [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; - [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sinkhorn")) return 0; - } - - return 1; -} - -static int ds4_gpu_hc_weighted_sum_strided( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *weights, - uint64_t weight_offset, - uint64_t weight_row_stride, - uint32_t n_embd, - uint32_t n_hc, - const char *label) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0 || - weight_row_stride < (uint64_t)n_hc * sizeof(float)) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(residual_hc); - id wbuf = ds4_gpu_tensor_buffer(weights); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); - if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || out_tensor_bytes % out_row_bytes != 0) { - fprintf(stderr, "ds4: Metal HC weighted sum output size is not a whole token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / out_row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC weighted sum token count is outside supported range\n"); - return 0; - } - - const uint64_t x_row_values = (uint64_t)n_hc * n_embd; - if (x_row_values == 0 || - x_row_values > UINT64_MAX / sizeof(float) || - n_tokens64 > UINT64_MAX / (x_row_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / ((uint64_t)n_hc * sizeof(float))) { - fprintf(stderr, "ds4: Metal HC weighted sum activation size overflow\n"); - return 0; - } - - const uint64_t x_bytes = n_tokens64 * x_row_values * sizeof(float); - const uint64_t w_last = weight_offset + - (n_tokens64 - 1u) * weight_row_stride + - (uint64_t)n_hc * sizeof(float); - if (!xbuf || !wbuf || !outbuf || - ds4_gpu_tensor_bytes(residual_hc) < x_bytes || - ds4_gpu_tensor_bytes(weights) < w_last) { - fprintf(stderr, "ds4: Metal HC weighted sum received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_hc_weighted_sum_args args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = (int64_t)n_tokens64, - .nb_x0 = sizeof(float), - .nb_x1 = (uint64_t)n_embd * sizeof(float), - .nb_x2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_w0 = sizeof(float), - .nb_w1 = weight_row_stride, - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - }; - const uint64_t n_elem = (uint64_t)n_embd * n_tokens64; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); - const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_hc_weighted_sum_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) + (NSUInteger)weight_offset atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, label)) return 0; - } - - return 1; -} - -int ds4_gpu_hc_weighted_sum_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *weights, - uint32_t n_embd, - uint32_t n_hc) { - return ds4_gpu_hc_weighted_sum_strided(out, - residual_hc, - weights, - 0, - (uint64_t)n_hc * sizeof(float), - n_embd, - n_hc, - "HC weighted sum"); -} - -int ds4_gpu_hc_weighted_sum_norm_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *weights, - const void *model_map, - uint64_t model_size, - uint64_t norm_weight_offset, - uint32_t n_embd, - uint32_t n_hc, - float norm_eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0; - - const bool force = - getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool disabled = - getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool require = - getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; - const bool supported_shape = - n_hc == 4u && (n_embd == 4096u || n_embd == 7168u); - const bool auto_shape = - n_embd == 4096u && ds4_gpu_device_name_contains("M3"); - const bool use_fusion = - supported_shape && !g_quality_mode && !disabled && - g_hc_weighted_sum_norm_pipeline != nil && - (auto_shape || force); - if (require && supported_shape && !use_fusion) { - fprintf(stderr, - "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n"); - return 0; - } - if (!use_fusion) return 0; - - @autoreleasepool { - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes; - const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); - id xbuf = ds4_gpu_tensor_buffer(residual_hc); - id wbuf = ds4_gpu_tensor_buffer(weights); - id outbuf = ds4_gpu_tensor_buffer(out); - id normbuf = ds4_gpu_tensor_buffer(norm_out); - if (!xbuf || !wbuf || !outbuf || !normbuf || - ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || - ds4_gpu_tensor_bytes(weights) < weight_bytes || - ds4_gpu_tensor_bytes(out) != out_row_bytes || - ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) { - fprintf(stderr, - "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n"); - return 0; - } - if (norm_weight_offset > model_size || - out_row_bytes > model_size - norm_weight_offset) { - fprintf(stderr, - "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n"); - return 0; - } - - uint64_t norm_inner = 0; - id normwbuf = ds4_gpu_wrap_model_range( - model_map, model_size, norm_weight_offset, - out_row_bytes, &norm_inner); - if (!normwbuf) return 0; - - ds4_gpu_hc_weighted_sum_norm_args args = { - .n_embd = (int64_t)n_embd, - .n_hc = (int64_t)n_hc, - .n_tokens = 1, - .nb_x0 = sizeof(float), - .nb_x1 = out_row_bytes, - .nb_x2 = residual_bytes, - .nb_w0 = sizeof(float), - .nb_w1 = weight_bytes, - .nb0 = sizeof(float), - .nb1 = out_row_bytes, - .nb_norm1 = out_row_bytes, - .norm_eps = norm_eps, - }; - const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); - const NSUInteger shared_bytes = - ((NSUInteger)n_embd + 32u) * sizeof(float); - if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup || - shared_bytes > [g_device maxThreadgroupMemoryLength]) { - if (require) { - fprintf(stderr, - "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n"); - } - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; - [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4]; - [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer( - cb, owned, "output HC sum/RMSNorm fused")) { - return 0; - } - } - - return 1; -} - -int ds4_gpu_hc_weighted_sum_split_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - return ds4_gpu_hc_weighted_sum_strided(out, - residual_hc, - split, - 0, - mix_hc * sizeof(float), - n_embd, - n_hc, - "HC weighted sum split"); -} - -/* Release decode fused HC pre-sublayer operation. The graph driver owns the - * optional reference fallback so this function stays a direct fused dispatch. */ -int ds4_gpu_hc_split_weighted_sum_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *split, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint32_t n_embd, - uint32_t n_hc, - uint32_t sinkhorn_iters, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !split || !mix || !residual_hc || !model_map || - n_embd == 0 || n_hc == 0) { - return 0; - } - if (n_hc != 4) { - fprintf(stderr, "ds4: Metal fused HC split/sum is specialized for HC=4\n"); - return 0; - } - - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t mix_bytes = mix_hc * sizeof(float); - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t scale_bytes = 3ull * sizeof(float); - - @autoreleasepool { - id mixbuf = ds4_gpu_tensor_buffer(mix); - id splitbuf = ds4_gpu_tensor_buffer(split); - id xbuf = ds4_gpu_tensor_buffer(residual_hc); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); - if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || - out_tensor_bytes % out_row_bytes != 0) { - fprintf(stderr, "ds4: Metal fused HC split/sum output size is not a whole token row\n"); - return 0; - } - - const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; - if (n_rows64 == 0 || n_rows64 > UINT32_MAX || - n_rows64 > UINT64_MAX / mix_bytes || - n_rows64 > UINT64_MAX / residual_row_bytes) { - fprintf(stderr, "ds4: Metal fused HC split/sum row count is outside supported range\n"); - return 0; - } - - const uint64_t mix_total_bytes = n_rows64 * mix_bytes; - const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; - if (!mixbuf || !splitbuf || !xbuf || !outbuf || - ds4_gpu_tensor_bytes(mix) < mix_total_bytes || - ds4_gpu_tensor_bytes(split) < mix_total_bytes || - ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes) { - fprintf(stderr, "ds4: Metal fused HC split/sum received undersized activation buffers\n"); - return 0; - } - - if (scale_offset > model_size || scale_bytes > model_size - scale_offset || - base_offset > model_size || mix_bytes > model_size - base_offset) { - fprintf(stderr, "ds4: Metal fused HC split/sum parameter range is outside the mapped model\n"); - return 0; - } - - uint64_t scale_inner = 0; - uint64_t base_inner = 0; - id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); - id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); - if (!scalebuf || !basebuf) return 0; - - ds4_gpu_hc_split_weighted_sum_args args = { - .n_embd = (int64_t)n_embd, - .n_hc = (int32_t)n_hc, - .sinkhorn_iters = (int32_t)sinkhorn_iters, - .n_rows = (int64_t)n_rows64, - .mix_hc = (int64_t)mix_hc, - .nb_mix1 = mix_bytes, - .nb_split1 = mix_bytes, - .nb_x0 = sizeof(float), - .nb_x1 = (uint64_t)n_embd * sizeof(float), - .nb_x2 = residual_row_bytes, - .nb0 = sizeof(float), - .nb1 = out_row_bytes, - .eps = eps, - }; - - NSUInteger nth = g_hc_split_weighted_sum_pipeline.maxTotalThreadsPerThreadgroup; - if (nth > 256u) nth = 256u; - if (nth > (NSUInteger)n_embd) nth = (NSUInteger)n_embd; - if (nth == 0) nth = 1u; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_hc_split_weighted_sum_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; - [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; - [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; - [enc setThreadgroupMemoryLength:(NSUInteger)n_hc * sizeof(float) atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum fused")) return 0; - } - - return 1; -} - -/* HC-pre plus the immediately following weighted RMSNorm, specialized for - * DS4's HC=4 shape. Both decode and batched prefill use this implementation; - * the kernel preserves their established single-row and batched scale formulas. */ -int ds4_gpu_hc_split_weighted_sum_norm_tensor( - ds4_gpu_tensor *out, - ds4_gpu_tensor *norm_out, - ds4_gpu_tensor *split, - const ds4_gpu_tensor *mix, - const ds4_gpu_tensor *residual_hc, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint64_t norm_weight_offset, - uint32_t n_embd, - uint32_t n_hc, - uint32_t sinkhorn_iters, - float eps, - float norm_eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !norm_out || !split || !mix || !residual_hc || !model_map || - n_hc != 4 || (n_embd & 3u) != 0) { - return 0; - } - - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t mix_bytes = mix_hc * sizeof(float); - const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); - const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t scale_bytes = 3ull * sizeof(float); - - @autoreleasepool { - id mixbuf = ds4_gpu_tensor_buffer(mix); - id splitbuf = ds4_gpu_tensor_buffer(split); - id xbuf = ds4_gpu_tensor_buffer(residual_hc); - id outbuf = ds4_gpu_tensor_buffer(out); - id normbuf = ds4_gpu_tensor_buffer(norm_out); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); - if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || - out_tensor_bytes % out_row_bytes != 0) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm output size is not a whole token row\n"); - return 0; - } - - const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; - if (n_rows64 == 0 || n_rows64 > UINT32_MAX || - n_rows64 > UINT64_MAX / mix_bytes || - n_rows64 > UINT64_MAX / residual_row_bytes) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm row count is outside supported range\n"); - return 0; - } - - const uint64_t mix_total_bytes = n_rows64 * mix_bytes; - const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; - const uint64_t out_total_bytes = n_rows64 * out_row_bytes; - if (!mixbuf || !splitbuf || !xbuf || !outbuf || !normbuf || - ds4_gpu_tensor_bytes(mix) < mix_total_bytes || - ds4_gpu_tensor_bytes(split) < mix_total_bytes || - ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes || - ds4_gpu_tensor_bytes(norm_out) < out_total_bytes) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm received undersized activation buffers\n"); - return 0; - } - - if (scale_offset > model_size || scale_bytes > model_size - scale_offset || - base_offset > model_size || mix_bytes > model_size - base_offset || - norm_weight_offset > model_size || out_row_bytes > model_size - norm_weight_offset) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm parameter range is outside the mapped model\n"); - return 0; - } - - uint64_t scale_inner = 0; - uint64_t base_inner = 0; - uint64_t norm_inner = 0; - id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); - id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); - id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, out_row_bytes, &norm_inner); - if (!scalebuf || !basebuf || !normwbuf) return 0; - - id pipeline = - ds4_gpu_hot_pipeline(g_hc_split_weighted_sum_norm_pipeline, - "kernel_dsv4_hc_split_weighted_sum_norm4"); - if (!pipeline) return 0; - - ds4_gpu_hc_split_weighted_sum_norm_args args = { - .n_embd = (int64_t)n_embd, - .n_hc = (int32_t)n_hc, - .sinkhorn_iters = (int32_t)sinkhorn_iters, - .n_rows = (int64_t)n_rows64, - .mix_hc = (int64_t)mix_hc, - .nb_mix1 = mix_bytes, - .nb_split1 = mix_bytes, - .nb_x0 = sizeof(float), - .nb_x1 = (uint64_t)n_embd * sizeof(float), - .nb_x2 = residual_row_bytes, - .nb0 = sizeof(float), - .nb1 = out_row_bytes, - .nb_norm1 = out_row_bytes, - .eps = eps, - .norm_eps = norm_eps, - }; - - NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); - if (nth > pipeline.maxTotalThreadsPerThreadgroup) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu threads but pipeline supports %lu\n", - (unsigned long)nth, - (unsigned long)pipeline.maxTotalThreadsPerThreadgroup); - return 0; - } - - const NSUInteger shared_bytes = ((NSUInteger)n_embd + 4u + 32u) * sizeof(float); - const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength]; - if (max_shared != 0 && shared_bytes > max_shared) { - fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu bytes of threadgroup memory but device supports %lu\n", - (unsigned long)shared_bytes, - (unsigned long)max_shared); - return 0; - } - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; - [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; - [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; - [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:7]; - [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:8]; - [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum/norm fused")) return 0; - } - - return 1; -} - -int ds4_gpu_output_hc_weights_tensor( - ds4_gpu_tensor *out, - const ds4_gpu_tensor *pre, - const void *model_map, - uint64_t model_size, - uint64_t scale_offset, - uint64_t base_offset, - uint32_t n_hc, - float eps) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !pre || !model_map || n_hc == 0) return 0; - - @autoreleasepool { - if ((n_hc % 4u) != 0) { - fprintf(stderr, "ds4: Metal output HC weights requires a multiple-of-4 HC width\n"); - return 0; - } - - id prebuf = ds4_gpu_tensor_buffer(pre); - id outbuf = ds4_gpu_tensor_buffer(out); - const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); - if (row_bytes == 0 || out_tensor_bytes < row_bytes || out_tensor_bytes % row_bytes != 0) { - fprintf(stderr, "ds4: Metal output HC weights size is not a whole token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX || - n_tokens64 > UINT64_MAX / row_bytes) { - fprintf(stderr, "ds4: Metal output HC weights token count is outside supported range\n"); - return 0; - } - - const uint64_t bytes = n_tokens64 * row_bytes; - if (!prebuf || !outbuf || - ds4_gpu_tensor_bytes(pre) < bytes || - ds4_gpu_tensor_bytes(out) < bytes) { - fprintf(stderr, "ds4: Metal output HC weights received undersized buffers\n"); - return 0; - } - - uint64_t scale_inner = 0; - uint64_t base_inner = 0; - id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, - scale_offset, sizeof(float), - &scale_inner); - id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, - base_offset, row_bytes, - &base_inner); - if (!scalebuf || !basebuf) return 0; - - const bool force_weights4 = - getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL; - const bool disable_weights4 = - getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL; - const bool require_weights4 = - getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL; - const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u; - const bool use_weights4 = - weights4_shape && !g_quality_mode && !disable_weights4 && - g_output_hc_weights4_pipeline != nil && - (ds4_gpu_device_name_contains("M3") || force_weights4) && - g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u; - if (require_weights4 && weights4_shape && !use_weights4) { - fprintf(stderr, - "ds4: required Metal output HC weights4 kernel was not selected\n"); - return 0; - } - - if (use_weights4) { - ds4_gpu_output_hc_weights4_args args = { - .post_scale = 1.0f, - .eps = eps, - }; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_output_hc_weights4_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; - [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; - [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) - threadsPerThreadgroup:MTLSizeMake(2, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer( - cb, owned, "output HC weights4")) { - return 0; - } - return 1; - } - - const uint32_t n_tokens = (uint32_t)n_tokens64; - ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1); - ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc); - ds4_gpu_unary_args sigmoid_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 0.0f, 0.0f); - ds4_gpu_unary_args scale_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 1.0f, eps); - - NSUInteger mul_nth_max = g_bin_mul_scalar_pipeline.maxTotalThreadsPerThreadgroup; - if (mul_nth_max > 256u) mul_nth_max = 256u; - NSUInteger mul_nth = 1u; - while (2u * mul_nth < (NSUInteger)mul_args.ne0 && mul_nth < mul_nth_max) { - mul_nth *= 2u; - } - - NSUInteger add_nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup; - if (add_nth_max > 256u) add_nth_max = 256u; - NSUInteger add_nth = 1u; - while (2u * add_nth < (NSUInteger)add_args.ne0 && add_nth < add_nth_max) { - add_nth *= 2u; - } - - NSUInteger unary_nth_max = g_unary_sigmoid_pipeline.maxTotalThreadsPerThreadgroup; - if (unary_nth_max > 256u) unary_nth_max = 256u; - NSUInteger unary_nth = (NSUInteger)sigmoid_args.ne00; - if (unary_nth > unary_nth_max) unary_nth = unary_nth_max; - if (unary_nth == 0) unary_nth = 1u; - const NSUInteger unary_nk0 = ((NSUInteger)sigmoid_args.ne00 + unary_nth - 1u) / unary_nth; - const NSUInteger out_offset = ds4_gpu_tensor_offset(out); - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - - [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; - [enc setBytes:&mul_args length:sizeof(mul_args) atIndex:0]; - [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; - [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; - [enc setBuffer:outbuf offset:out_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)mul_args.ne01, - (NSUInteger)mul_args.ne02, - (NSUInteger)mul_args.ne03) - threadsPerThreadgroup:MTLSizeMake(mul_nth, 1, 1)]; - - [enc setComputePipelineState:g_add_pipeline]; - [enc setBytes:&add_args length:sizeof(add_args) atIndex:0]; - [enc setBuffer:outbuf offset:out_offset atIndex:1]; - [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:2]; - [enc setBuffer:outbuf offset:out_offset atIndex:3]; - [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)add_args.ne01, - (NSUInteger)add_args.ne02, - (NSUInteger)add_args.ne03) - threadsPerThreadgroup:MTLSizeMake(add_nth, 1, 1)]; - - [enc setComputePipelineState:g_unary_sigmoid_pipeline]; - [enc setBytes:&sigmoid_args length:sizeof(sigmoid_args) atIndex:0]; - [enc setBuffer:outbuf offset:out_offset atIndex:1]; - [enc setBuffer:outbuf offset:out_offset atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)sigmoid_args.ne01, - (NSUInteger)sigmoid_args.ne02, - (NSUInteger)sigmoid_args.ne03) - threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; - - [enc setComputePipelineState:g_unary_scale_pipeline]; - [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; - [enc setBuffer:outbuf offset:out_offset atIndex:1]; - [enc setBuffer:outbuf offset:out_offset atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)scale_args.ne01, - (NSUInteger)scale_args.ne02, - (NSUInteger)scale_args.ne03) - threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "output HC weights")) return 0; - } - - return 1; -} - -int ds4_gpu_hc_expand_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *block_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *post, - const ds4_gpu_tensor *comb, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (n_embd == 0 || n_hc == 0) return 0; - - @autoreleasepool { - id blockbuf = ds4_gpu_tensor_buffer(block_out); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id postbuf = ds4_gpu_tensor_buffer(post); - id combbuf = ds4_gpu_tensor_buffer(comb); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); - if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { - fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); - return 0; - } - - const uint64_t block_values = (uint64_t)n_embd; - const uint64_t hc_values = (uint64_t)n_hc * n_embd; - const uint64_t comb_values = (uint64_t)n_hc * n_hc; - if (hc_values == 0 || - hc_values > UINT64_MAX / sizeof(float) || - comb_values > UINT64_MAX / sizeof(float) || - n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { - fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); - return 0; - } - - const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); - const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); - const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); - const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); - if (!blockbuf || !resbuf || !postbuf || !combbuf || !outbuf || - ds4_gpu_tensor_bytes(block_out) < block_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(post) < post_bytes || - ds4_gpu_tensor_bytes(comb) < comb_bytes) { - fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_hc_expand_args args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = (int64_t)n_tokens64, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = (uint64_t)n_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 0, - }; - id expand_pipeline = g_hc_expand_pipeline; - uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; - if (n_hc == 4) { - expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, - "kernel_dsv4_hc_expand4"); - n_elem = (uint64_t)n_embd * n_tokens64; - } - if (!expand_pipeline) return 0; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); - const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:expand_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; - [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; - [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand")) return 0; - } - - return 1; -} - -/* Expand of the SUM of two block vectors — the TP attention combine folded - * into the HC expand (the kernel's has_add path adds them element-wise - * before the post/comb mixing, canonical rank order preserved by argument - * position). */ -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) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (n_embd == 0 || n_hc != 4) return 0; - - @autoreleasepool { - id blockbuf = ds4_gpu_tensor_buffer(block_out); - id addbuf = ds4_gpu_tensor_buffer(block_add); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id postbuf = ds4_gpu_tensor_buffer(post); - id combbuf = ds4_gpu_tensor_buffer(comb); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); - if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { - fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); - return 0; - } - - const uint64_t block_values = (uint64_t)n_embd; - const uint64_t hc_values = (uint64_t)n_hc * n_embd; - const uint64_t comb_values = (uint64_t)n_hc * n_hc; - if (hc_values == 0 || - hc_values > UINT64_MAX / sizeof(float) || - comb_values > UINT64_MAX / sizeof(float) || - n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { - fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); - return 0; - } - - const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); - const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); - const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); - const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); - if (!blockbuf || !addbuf || !resbuf || !postbuf || !combbuf || !outbuf || - ds4_gpu_tensor_bytes(block_out) < block_bytes || - ds4_gpu_tensor_bytes(block_add) < block_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(post) < post_bytes || - ds4_gpu_tensor_bytes(comb) < comb_bytes) { - fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_hc_expand_args args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = (int64_t)n_tokens64, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = (uint64_t)n_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 1, - }; - id expand_pipeline = g_hc_expand_pipeline; - uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; - if (n_hc == 4) { - expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, - "kernel_dsv4_hc_expand4"); - n_elem = (uint64_t)n_embd * n_tokens64; - } - if (!expand_pipeline) return 0; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); - const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:expand_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; - [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; - [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; - [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add")) return 0; - } - - return 1; -} - -int ds4_gpu_hc_expand_split_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *block_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - - @autoreleasepool { - id blockbuf = ds4_gpu_tensor_buffer(block_out); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id splitbuf = ds4_gpu_tensor_buffer(split); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); - if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { - fprintf(stderr, "ds4: Metal HC expand split output size is not a whole HC token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC expand split token count is outside supported range\n"); - return 0; - } - - const uint64_t block_values = (uint64_t)n_embd; - const uint64_t hc_values = (uint64_t)n_hc * n_embd; - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - if (hc_values == 0 || - hc_values > UINT64_MAX / sizeof(float) || - mix_hc > UINT64_MAX / sizeof(float) || - n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { - fprintf(stderr, "ds4: Metal HC expand split activation size overflow\n"); - return 0; - } - - const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); - const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); - const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); - if (!blockbuf || !resbuf || !splitbuf || !outbuf || - ds4_gpu_tensor_bytes(block_out) < block_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(split) < split_bytes) { - fprintf(stderr, "ds4: Metal HC expand split received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_hc_expand_args args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = (int64_t)n_tokens64, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = mix_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = mix_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 0, - }; - id expand_pipeline = g_hc_expand_pipeline; - uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; - if (n_hc == 4) { - expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, - "kernel_dsv4_hc_expand4"); - n_elem = (uint64_t)n_embd * n_tokens64; - } - if (!expand_pipeline) return 0; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); - const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:expand_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand split")) return 0; - } - - return 1; -} - -int ds4_gpu_hc_expand_split_half_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *block_out_h, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split; - (void)n_embd; (void)n_hc; - return 0; -} - -int ds4_gpu_hc_expand_add_split_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 *split, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; - - @autoreleasepool { - id blockbuf = ds4_gpu_tensor_buffer(block_out); - id addbuf = ds4_gpu_tensor_buffer(block_add); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id splitbuf = ds4_gpu_tensor_buffer(split); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); - if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { - fprintf(stderr, "ds4: Metal HC expand add split output size is not a whole HC token row\n"); - return 0; - } - - const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; - if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { - fprintf(stderr, "ds4: Metal HC expand add split token count is outside supported range\n"); - return 0; - } - - const uint64_t block_values = (uint64_t)n_embd; - const uint64_t hc_values = (uint64_t)n_hc * n_embd; - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - if (hc_values == 0 || - hc_values > UINT64_MAX / sizeof(float) || - mix_hc > UINT64_MAX / sizeof(float) || - n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || - n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { - fprintf(stderr, "ds4: Metal HC expand add split activation size overflow\n"); - return 0; - } - - const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); - const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); - const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); - if (!blockbuf || !addbuf || !resbuf || !splitbuf || !outbuf || - ds4_gpu_tensor_bytes(block_out) < block_bytes || - ds4_gpu_tensor_bytes(block_add) < block_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(split) < split_bytes) { - fprintf(stderr, "ds4: Metal HC expand add split received undersized activation buffers\n"); - return 0; - } - - ds4_gpu_hc_expand_args args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = (int64_t)n_tokens64, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = mix_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = mix_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 1, - }; - id expand_pipeline = g_hc_expand_pipeline; - uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; - if (n_hc == 4) { - expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, - "kernel_dsv4_hc_expand4"); - n_elem = (uint64_t)n_embd * n_tokens64; - } - if (!expand_pipeline) return 0; - const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); - const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:expand_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; - [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; - [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add split")) return 0; - } - - return 1; -} - -int ds4_gpu_hc_expand_add_split_half_add_tensor( - ds4_gpu_tensor *out_hc, - const ds4_gpu_tensor *block_out, - const ds4_gpu_tensor *block_add_h, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - (void)out_hc; (void)block_out; (void)block_add_h; (void)residual_hc; - (void)split; (void)n_embd; (void)n_hc; - return 0; -} - -int ds4_gpu_shared_down_hc_expand_q8_0_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *shared_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 *shared_mid, - const ds4_gpu_tensor *routed_out, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !shared_out || !model_map || !shared_mid || !routed_out || - !residual_hc || !split || n_embd == 0 || n_hc == 0 || - n_hc != 4 || out_dim != n_embd || (in_dim & 31u) != 0 || - in_dim > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - id midbuf = ds4_gpu_tensor_buffer(shared_mid); - id sharedbuf = ds4_gpu_tensor_buffer(shared_out); - id routedbuf = ds4_gpu_tensor_buffer(routed_out); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id splitbuf = ds4_gpu_tensor_buffer(split); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - - const uint64_t row_bytes = (in_dim / 32u) * 34u; - const uint64_t weight_bytes = out_dim * row_bytes; - const uint64_t shared_mid_bytes = in_dim * sizeof(float); - const uint64_t embd_bytes = out_dim * sizeof(float); - const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t split_bytes = mix_hc * sizeof(float); - - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal shared-down HC fusion weight range is outside the mapped model\n"); - return 0; - } - if (!midbuf || !sharedbuf || !routedbuf || !resbuf || !splitbuf || !outbuf || - ds4_gpu_tensor_bytes(shared_mid) < shared_mid_bytes || - ds4_gpu_tensor_bytes(shared_out) < embd_bytes || - ds4_gpu_tensor_bytes(routed_out) < embd_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(split) < split_bytes || - ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { - fprintf(stderr, "ds4: Metal shared-down HC fusion received undersized buffers\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_offset, weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - mv_args.nr0 = mv_dispatch.nr0; - - ds4_gpu_hc_expand_args hc_args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = 1, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = mix_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = mix_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 1, - }; - - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_hc_expand4_q8_0", - mv_dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; - [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(shared_mid) atIndex:3]; - [enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:4]; - [enc setBuffer:routedbuf offset:ds4_gpu_tensor_offset(routed_out) atIndex:5]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:7]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:8]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:9]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "shared-down HC expand fused")) return 0; - } - - return 1; -} - -int ds4_gpu_matmul_q8_0_hc_expand_tensor( - ds4_gpu_tensor *out_hc, - ds4_gpu_tensor *block_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, - const ds4_gpu_tensor *residual_hc, - const ds4_gpu_tensor *split, - uint32_t n_embd, - uint32_t n_hc) { - if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out_hc || !block_out || !model_map || !x || !residual_hc || !split || - n_embd == 0 || n_hc == 0 || n_hc != 4 || out_dim != n_embd || - (in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX) { - return 0; - } - - @autoreleasepool { - id xbuf = ds4_gpu_tensor_buffer(x); - id blockbuf = ds4_gpu_tensor_buffer(block_out); - id resbuf = ds4_gpu_tensor_buffer(residual_hc); - id splitbuf = ds4_gpu_tensor_buffer(split); - id outbuf = ds4_gpu_tensor_buffer(out_hc); - - const uint64_t row_bytes = (in_dim / 32u) * 34u; - const uint64_t weight_bytes = out_dim * row_bytes; - const uint64_t x_bytes = in_dim * sizeof(float); - const uint64_t embd_bytes = out_dim * sizeof(float); - const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); - const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; - const uint64_t split_bytes = mix_hc * sizeof(float); - - if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { - fprintf(stderr, "ds4: Metal Q8 HC fusion weight range is outside the mapped model\n"); - return 0; - } - if (!xbuf || !blockbuf || !resbuf || !splitbuf || !outbuf || - ds4_gpu_tensor_bytes(x) < x_bytes || - ds4_gpu_tensor_bytes(block_out) < embd_bytes || - ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || - ds4_gpu_tensor_bytes(split) < split_bytes || - ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { - fprintf(stderr, "ds4: Metal Q8 HC fusion received undersized buffers\n"); - return 0; - } - - uint64_t inner_offset = 0; - id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, - weight_offset, weight_bytes, - &inner_offset); - if (!wbuf) return 0; - - ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); - ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); - mv_args.nr0 = mv_dispatch.nr0; - - ds4_gpu_hc_expand_args hc_args = { - .n_embd = n_embd, - .n_hc = n_hc, - .n_tokens = 1, - .nb_block0 = sizeof(float), - .nb_block1 = (uint64_t)n_embd * sizeof(float), - .nb_add0 = sizeof(float), - .nb_add1 = (uint64_t)n_embd * sizeof(float), - .nb_res0 = sizeof(float), - .nb_res1 = (uint64_t)n_embd * sizeof(float), - .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), - .nb_post0 = sizeof(float), - .nb_post1 = mix_hc * sizeof(float), - .nb_comb0 = sizeof(float), - .nb_comb1 = (uint64_t)n_hc * sizeof(float), - .nb_comb2 = mix_hc * sizeof(float), - .nb0 = sizeof(float), - .nb1 = (uint64_t)n_embd * sizeof(float), - .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), - .has_add = 0, - }; - - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_q8_hc_expand4_q8_0", - mv_dispatch.nsg); - if (!pipeline) return 0; - - int owned = 0; - id cb = ds4_gpu_command_buffer(&owned); - if (!cb) return 0; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:pipeline]; - [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; - [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; - [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; - [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; - [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4]; - [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:6]; - [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7]; - [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8]; - [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; - [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / - (NSUInteger)mv_dispatch.nr0, - 1, - 1) - threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8 HC expand fused")) return 0; - } - - return 1; -} - -void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { - (void)enabled; -} +#include "metal/runtime.inc" +#include "metal/embedding.inc" +#include "metal/model_io.inc" +#include "metal/expert_streaming.inc" +#include "models/deepseek/metal/host/indexer.inc" +#include "metal/dense_norm.inc" +#include "models/deepseek/metal/host/attention.inc" +#include "metal/elementwise.inc" +#include "metal/moe_dispatch.inc" +#include "models/glm/metal/host/kernels.inc" +#include "models/deepseek/metal/host/moe.inc" +#include "models/deepseek/metal/host/hc.inc" +#include "metal/compat.inc" diff --git a/ds4_model_provider.c b/ds4_model_provider.c new file mode 100644 index 0000000000..cfae54084e --- /dev/null +++ b/ds4_model_provider.c @@ -0,0 +1,26 @@ +#include "ds4_model_provider_builtin.h" + +bool ds4_model_provider_valid(const ds4_model_provider_v1 *provider) { + return provider && + provider->abi_version == DS4_MODEL_PROVIDER_ABI_VERSION && + provider->struct_size >= sizeof(ds4_model_provider_v1) && + provider->id && + provider->session_create && + provider->session_destroy && + provider->session_sync && + provider->session_eval && + provider->sessions_eval_batch && + provider->sessions_eval_batch_with_prefill && + provider->session_eval_speculative && + provider->session_invalidate && + provider->session_rewind && + provider->session_layer_slice_reset && + provider->session_eval_output_head && + provider->session_eval_layer_slice && + provider->session_payload_bytes && + provider->session_save_payload && + provider->session_load_payload && + provider->session_layer_payload_bytes && + provider->session_save_layer_payload && + provider->session_load_layer_payload; +} diff --git a/ds4_model_provider.h b/ds4_model_provider.h new file mode 100644 index 0000000000..a5dfb347a2 --- /dev/null +++ b/ds4_model_provider.h @@ -0,0 +1,133 @@ +#ifndef DS4_MODEL_PROVIDER_H +#define DS4_MODEL_PROVIDER_H + +#include +#include +#include +#include + +#include "ds4.h" + +/* + * Source-level boundary between the engine core and a model integration. + * + * A provider owns whole-model orchestration. It may call its custom kernels + * directly and keep model-specific graph state private; the core only enters + * through these lifecycle operations. This is intentionally not an operator + * or individual-kernel interface. + */ +#define DS4_MODEL_PROVIDER_ABI_VERSION 1u + +typedef struct ds4_model_provider_v1 { + uint32_t abi_version; + uint32_t struct_size; + const char *id; + + int (*session_create)(ds4_session **out, + ds4_engine *engine, + int context_size); + + /* + * Release only provider-owned state embedded in the session. The core + * releases transport, checkpoint, sampling, and session storage. + */ + void (*session_destroy)(ds4_session *session); + + int (*session_sync)(ds4_session *session, + const ds4_tokens *prompt, + char *err, + size_t errlen); + + int (*session_eval)(ds4_session *session, + int token, + bool probe_support_model, + char *err, + size_t errlen); + + int (*sessions_eval_batch)(ds4_decode_item *items, + int count, + char *err, + size_t errlen); + + int (*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 (*session_eval_speculative)(ds4_session *session, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen); + + /* Update provider-private cache frontiers after core checkpoint changes. */ + void (*session_invalidate)(ds4_session *session); + void (*session_rewind)(ds4_session *session, int position); + + /* Pipeline-parallel execution of a contiguous transformer slice. */ + int (*session_layer_slice_reset)(ds4_session *session, + char *err, + size_t errlen); + + int (*session_eval_output_head)(ds4_session *session, + const float *hidden_state, + uint32_t token_count, + float *logits, + char *err, + size_t errlen); + + int (*session_eval_layer_slice)(ds4_session *session, + const int *tokens, + uint32_t token_count, + uint32_t position, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hidden_state, + float *output_hidden_state, + bool output_logits, + float *logits, + char *err, + size_t errlen); + + uint64_t (*session_payload_bytes)(ds4_session *session); + + int (*session_save_payload)(ds4_session *session, + FILE *file, + char *err, + size_t errlen); + + int (*session_load_payload)(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + char *err, + size_t errlen); + + uint64_t (*session_layer_payload_bytes)(ds4_session *session, + uint32_t layer_start, + uint32_t layer_end); + + int (*session_save_layer_payload)(ds4_session *session, + FILE *file, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); + + int (*session_load_layer_payload)(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + const int *tokens, + uint32_t token_count, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); +} ds4_model_provider_v1; + +#endif diff --git a/ds4_model_provider_builtin.h b/ds4_model_provider_builtin.h new file mode 100644 index 0000000000..fe056a4954 --- /dev/null +++ b/ds4_model_provider_builtin.h @@ -0,0 +1,25 @@ +#ifndef DS4_MODEL_PROVIDER_BUILTIN_H +#define DS4_MODEL_PROVIDER_BUILTIN_H + +#include "ds4_model_provider.h" + +/* + * Shared entry points used by the built-in providers. Model-specific + * lifecycle declarations live beside each provider under models//. + */ +bool ds4_model_provider_valid(const ds4_model_provider_v1 *provider); + +int ds4_builtin_sessions_eval_batch(ds4_decode_item *items, + int count, + char *err, + size_t errlen); + +int ds4_builtin_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); + +#endif diff --git a/ds4_rocm.cu b/ds4_rocm.cu index 46ba223df6..742645dea7 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -97,38 +97,38 @@ typedef struct { #include "rocm/ds4_rocm_norm_rope.cuh" -#include "rocm/ds4_rocm_fp8_kv.cuh" +#include "models/deepseek/rocm/fp8_kv.cuh" -#include "rocm/ds4_rocm_attention.cuh" +#include "models/deepseek/rocm/attention.cuh" -#include "rocm/ds4_rocm_hc.cuh" +#include "models/deepseek/rocm/hc.cuh" -#include "rocm/ds4_rocm_output.cuh" +#include "models/deepseek/rocm/output.cuh" -#include "rocm/ds4_rocm_indexer.cuh" +#include "models/deepseek/rocm/indexer.cuh" #include "rocm/ds4_rocm_embedding_launch.cuh" #include "rocm/ds4_rocm_matmul.cuh" -#include "rocm/ds4_rocm_fp8_kv_launch.cuh" +#include "models/deepseek/rocm/fp8_kv_launch.cuh" -#include "rocm/ds4_rocm_compressor.cuh" +#include "models/deepseek/rocm/compressor.cuh" -#include "rocm/ds4_rocm_attention_launch.cuh" +#include "models/deepseek/rocm/attention_launch.cuh" #include "rocm/ds4_rocm_shared_expert.cuh" #include "rocm/ds4_rocm_misc_launch.cuh" -#include "rocm/ds4_rocm_router.cuh" +#include "models/deepseek/rocm/router.cuh" #include "rocm/ds4_rocm_moe.cuh" #include "rocm/ds4_rocm_moe_launch.cuh" -#include "rocm/ds4_rocm_glm.cuh" +#include "models/glm/rocm/kernels.cuh" -#include "rocm/ds4_rocm_hc_output_launch.cuh" +#include "models/deepseek/rocm/hc_output_launch.cuh" #include "rocm/ds4_rocm_current_api_compat.cuh" diff --git a/kernels/README.md b/kernels/README.md new file mode 100644 index 0000000000..d9e3630174 --- /dev/null +++ b/kernels/README.md @@ -0,0 +1,12 @@ +# Kernel implementation units + +The CPU kernels are grouped by concrete responsibility: + +- `cpu_quant.inc`: scalar conversion, quantization formats, and quantized dot + products. +- `cpu_matmul.inc`: embedding lookup, dense matrix-vector operations, and + routed-expert matrix products. + +These are implementation fragments included exactly once by `ds4.c`. They +remain specialized for the supported model tensor layouts; this directory is +an ownership boundary, not a generic kernel API. diff --git a/kernels/cpu_matmul.inc b/kernels/cpu_matmul.inc new file mode 100644 index 0000000000..b54be27f19 --- /dev/null +++ b/kernels/cpu_matmul.inc @@ -0,0 +1,3022 @@ +/* + * CPU embedding, matrix-vector, and routed-expert kernels. + * + * Included exactly once by ds4.c. These remain concrete kernels specialized + * for the supported tensor layouts. + */ + +static void weights_free(ds4_weights *w) { + memset(w, 0, sizeof(*w)); +} + +/* Load one token embedding row and expand it to float activations. */ +static void embed_token_f16(const ds4_model *m, const ds4_weights *w, int token, float *out) { + ds4_tensor *te = w->token_embd; + if (te->type != DS4_TENSOR_F16 || te->ndim != 2) { + ds4_die("expected a 2D F16 token embedding tensor"); + } + if (token < 0 || (uint64_t)token >= te->dim[1]) { + ds4_die("token id is outside the embedding table"); + } + + const uint16_t *base = tensor_data(m, te); + const uint64_t stride = te->dim[0]; + const uint16_t *row = base + (uint64_t)token * stride; + + for (uint64_t i = 0; i < stride; i++) { + out[i] = f16_to_f32(row[i]); + } +} + +static void embed_token_q8_0(const ds4_model *m, const ds4_weights *w, int token, float *out) { + ds4_tensor *te = w->token_embd; + if (te->type != DS4_TENSOR_Q8_0 || te->ndim != 2) { + ds4_die("expected a 2D Q8_0 token embedding tensor"); + } + if (token < 0 || (uint64_t)token >= te->dim[1]) { + ds4_die("token id is outside the embedding table"); + } + + const uint64_t n = te->dim[0]; + const uint64_t blocks = (n + 31) / 32; + const uint8_t *row = (const uint8_t *)tensor_data(m, te) + + (uint64_t)token * blocks * 34; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); + const float scale = f16_to_f32(scale_bits); + const int8_t *qs = (const int8_t *)(row + b * 34 + 2); + const uint64_t i0 = b * 32; + const uint64_t bn = n - i0 < 32 ? n - i0 : 32; + for (uint64_t i = 0; i < bn; i++) { + out[i0 + i] = scale * (float)qs[i]; + } + } +} + +static void embed_token_any(const ds4_model *m, const ds4_weights *w, int token, float *out) { + if (!w->token_embd) ds4_die("token embedding tensor is missing"); + switch (w->token_embd->type) { + case DS4_TENSOR_F16: + embed_token_f16(m, w, token, out); + break; + case DS4_TENSOR_Q8_0: + embed_token_q8_0(m, w, token, out); + break; + default: + ds4_die("unsupported token embedding tensor type"); + } +} + +/* RMSNorm without a learned scale, used by hyper-connection control vectors. */ +static void rms_norm_no_weight(float *out, const float *x, uint64_t n, float eps) { + double ss = 0.0; + for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; + + const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); + for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale; +} + +/* Standard DS4 RMSNorm with learned per-channel scale. */ +static void rms_norm_weight(float *out, const float *x, const float *weight, uint64_t n, float eps) { + double ss = 0.0; + for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; + + const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); + for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale * weight[i]; +} + +/* Normalize each attention head independently after Q projection. */ +static void head_rms_norm_inplace(float *x, uint32_t n_head, uint32_t head_dim, float eps) { + for (uint32_t h = 0; h < n_head; h++) { + float *head = x + (uint64_t)h * head_dim; + double ss = 0.0; + for (uint32_t i = 0; i < head_dim; i++) ss += (double)head[i] * head[i]; + + const float scale = 1.0f / sqrtf((float)(ss / (double)head_dim) + eps); + for (uint32_t i = 0; i < head_dim; i++) head[i] *= scale; + } +} + +typedef struct { + float *out; + const uint16_t *data; + const float *x; + uint64_t in_dim; +} matvec_f16_ctx; + +static inline float dot_f16_row(const uint16_t *row, const float *x, uint64_t n) { +#if defined(__ARM_NEON) + uint64_t i = 0; + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + for (; i + 8 <= n; i += 8) { + const float16x8_t hv = vreinterpretq_f16_u16(vld1q_u16(row + i)); + const float32x4_t h0 = vcvt_f32_f16(vget_low_f16(hv)); + const float32x4_t h1 = vcvt_f32_f16(vget_high_f16(hv)); + acc0 = vfmaq_f32(acc0, h0, vld1q_f32(x + i)); + acc1 = vfmaq_f32(acc1, h1, vld1q_f32(x + i + 4)); + } + + float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); + for (; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; + return acc; +#else + float acc = 0.0f; + for (uint64_t i = 0; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; + return acc; +#endif +} + +static void matvec_f16_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_f16_ctx *ctx = vctx; + + for (uint64_t o = row0; o < row1; o++) { + const uint16_t *row = ctx->data + o * ctx->in_dim; + ctx->out[o] = dot_f16_row(row, ctx->x, ctx->in_dim); + } +} + +/* Dense F16 matvec for small control projections such as HC and router heads. */ +static void matvec_f16(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { + if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t out_dim = w->dim[1]; + matvec_f16_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .x = x, + .in_dim = in_dim, + }; + + const uint64_t ops = in_dim * out_dim; + const uint64_t min_rows = ops >= 262144 ? 1 : 512; + ds4_parallel_for_min_rows(out_dim, matvec_f16_worker, &ctx, min_rows); +} + +static void matvec_f16_serial(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { + if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t out_dim = w->dim[1]; + const uint16_t *data = tensor_data(m, w); + for (uint64_t o = 0; o < out_dim; o++) { + out[o] = dot_f16_row(data + o * in_dim, x, in_dim); + } +} + +typedef struct { + float *out; + const uint8_t *data; + const int8_t *xq; + const float *xscale; + uint64_t in_dim; + uint64_t row0; + uint64_t blocks; +} matvec_q8_0_ctx; + +typedef struct { + float *out0; + float *out1; + const uint8_t *data0; + const uint8_t *data1; + const int8_t *xq; + const float *xscale; + uint64_t in_dim; + uint64_t blocks; +} matvec_q8_0_pair_ctx; + +typedef struct { + float *out; + const uint8_t *data; + const int8_t *xq; + const float *xscale; + uint64_t in_dim; + uint64_t blocks; + uint64_t rank; +} matvec_q8_0_grouped_ctx; + +typedef struct { + float *out; + const uint8_t *data; + const int8_t *xq; + const float *xscale; + uint64_t n_tok; + uint64_t n_groups; + uint64_t group_dim; + uint64_t blocks; + uint64_t rank; +} matmul_q8_0_grouped_batch_ctx; + +typedef struct { + float *out; + const uint8_t *data; + const int8_t *xq; + const float *xscale; + uint64_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t blocks; +} matmul_q8_0_batch_ctx; + +typedef struct { + float *out0; + float *out1; + const uint8_t *data0; + const uint8_t *data1; + const int8_t *xq; + const float *xscale; + uint64_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t blocks; +} matmul_q8_0_pair_batch_ctx; + +typedef struct { + const float *x; + int8_t *xq; + float *xscale; + uint64_t in_dim; + uint64_t blocks; +} quantize_q8_0_batch_ctx; + +static inline int32_t dot_i8_32(const int8_t *a, const int8_t *b, uint64_t n) { +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + if (n == 32) { + int32x4_t acc = vdupq_n_s32(0); + acc = vdotq_s32(acc, vld1q_s8(a), vld1q_s8(b)); + acc = vdotq_s32(acc, vld1q_s8(a + 16), vld1q_s8(b + 16)); + return vaddvq_s32(acc); + } +#endif + int32_t sum = 0; + for (uint64_t i = 0; i < n; i++) sum += (int32_t)a[i] * (int32_t)b[i]; + return sum; +} + +static inline float dot_q8_0_row( + const uint8_t *row, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t blocks) { +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + if ((in_dim & 31u) == 0) { + float32x4_t accv0 = vdupq_n_f32(0.0f); + float32x4_t accv1 = vdupq_n_f32(0.0f); + + uint64_t b = 0; + for (; b + 1 < blocks; b += 2) { + uint16_t scale_bits0; + uint16_t scale_bits1; + memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); + memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); + + const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); + const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); + const int8_t *xq0 = xq + b * 32; + const int8_t *xq1 = xq + (b + 1) * 32; + + int32x4_t dot0 = vdupq_n_s32(0); + dot0 = vdotq_s32(dot0, vld1q_s8(qs0), vld1q_s8(xq0)); + dot0 = vdotq_s32(dot0, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + 16)); + + int32x4_t dot1 = vdupq_n_s32(0); + dot1 = vdotq_s32(dot1, vld1q_s8(qs1), vld1q_s8(xq1)); + dot1 = vdotq_s32(dot1, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + 16)); + + accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot0), f16_to_f32(scale_bits0) * xscale[b]); + accv1 = vfmaq_n_f32(accv1, vcvtq_f32_s32(dot1), f16_to_f32(scale_bits1) * xscale[b + 1]); + } + + if (b < blocks) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); + const int8_t *qs = (const int8_t *)(row + b * 34 + 2); + const int8_t *xqb = xq + b * 32; + int32x4_t dot = vdupq_n_s32(0); + dot = vdotq_s32(dot, vld1q_s8(qs), vld1q_s8(xqb)); + dot = vdotq_s32(dot, vld1q_s8(qs + 16), vld1q_s8(xqb + 16)); + accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot), f16_to_f32(scale_bits) * xscale[b]); + } + + return vaddvq_f32(vaddq_f32(accv0, accv1)); + } +#endif + + float acc = 0.0f; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); + const int8_t *qs = (const int8_t *)(row + b * 34 + 2); + + const uint64_t i0 = b * 32; + const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; + acc += f16_to_f32(scale_bits) * xscale[b] * (float)dot_i8_32(qs, xq + i0, n); + } + return acc; +} + +static inline void dot_q8_0_row_2( + const uint8_t *row, + const int8_t *xq0, + const float *xscale0, + const int8_t *xq1, + const float *xscale1, + uint64_t in_dim, + uint64_t blocks, + float *out0, + float *out1) { +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + if ((in_dim & 31u) == 0) { + float32x4_t acc00 = vdupq_n_f32(0.0f); + float32x4_t acc01 = vdupq_n_f32(0.0f); + float32x4_t acc10 = vdupq_n_f32(0.0f); + float32x4_t acc11 = vdupq_n_f32(0.0f); + + uint64_t b = 0; + for (; b + 1 < blocks; b += 2) { + uint16_t scale_bits0; + uint16_t scale_bits1; + memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); + memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); + + const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); + const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); + + int32x4_t d00 = vdupq_n_s32(0); + d00 = vdotq_s32(d00, vld1q_s8(qs0), vld1q_s8(xq0 + b * 32)); + d00 = vdotq_s32(d00, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + b * 32 + 16)); + int32x4_t d01 = vdupq_n_s32(0); + d01 = vdotq_s32(d01, vld1q_s8(qs1), vld1q_s8(xq0 + (b + 1) * 32)); + d01 = vdotq_s32(d01, vld1q_s8(qs1 + 16), vld1q_s8(xq0 + (b + 1) * 32 + 16)); + + int32x4_t d10 = vdupq_n_s32(0); + d10 = vdotq_s32(d10, vld1q_s8(qs0), vld1q_s8(xq1 + b * 32)); + d10 = vdotq_s32(d10, vld1q_s8(qs0 + 16), vld1q_s8(xq1 + b * 32 + 16)); + int32x4_t d11 = vdupq_n_s32(0); + d11 = vdotq_s32(d11, vld1q_s8(qs1), vld1q_s8(xq1 + (b + 1) * 32)); + d11 = vdotq_s32(d11, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + (b + 1) * 32 + 16)); + + const float s0 = f16_to_f32(scale_bits0); + const float s1 = f16_to_f32(scale_bits1); + acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), s0 * xscale0[b]); + acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), s1 * xscale0[b + 1]); + acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), s0 * xscale1[b]); + acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), s1 * xscale1[b + 1]); + } + + if (b < blocks) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); + const int8_t *qs = (const int8_t *)(row + b * 34 + 2); + + int32x4_t d0 = vdupq_n_s32(0); + d0 = vdotq_s32(d0, vld1q_s8(qs), vld1q_s8(xq0 + b * 32)); + d0 = vdotq_s32(d0, vld1q_s8(qs + 16), vld1q_s8(xq0 + b * 32 + 16)); + int32x4_t d1 = vdupq_n_s32(0); + d1 = vdotq_s32(d1, vld1q_s8(qs), vld1q_s8(xq1 + b * 32)); + d1 = vdotq_s32(d1, vld1q_s8(qs + 16), vld1q_s8(xq1 + b * 32 + 16)); + + const float s0 = f16_to_f32(scale_bits); + acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), s0 * xscale0[b]); + acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), s0 * xscale1[b]); + } + + *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); + *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); + return; + } +#endif + + *out0 = dot_q8_0_row(row, xq0, xscale0, in_dim, blocks); + *out1 = dot_q8_0_row(row, xq1, xscale1, in_dim, blocks); +} + +static inline DS4_MAYBE_UNUSED void dot_q8_0_row_pair( + const uint8_t *row0, + const uint8_t *row1, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t blocks, + float *out0, + float *out1) { +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + if ((in_dim & 31u) == 0) { + float32x4_t acc00 = vdupq_n_f32(0.0f); + float32x4_t acc01 = vdupq_n_f32(0.0f); + float32x4_t acc10 = vdupq_n_f32(0.0f); + float32x4_t acc11 = vdupq_n_f32(0.0f); + + uint64_t b = 0; + for (; b + 1 < blocks; b += 2) { + uint16_t s00, s01, s10, s11; + memcpy(&s00, row0 + b * 34, sizeof(s00)); + memcpy(&s01, row0 + (b + 1) * 34, sizeof(s01)); + memcpy(&s10, row1 + b * 34, sizeof(s10)); + memcpy(&s11, row1 + (b + 1) * 34, sizeof(s11)); + + const int8_t *xq0 = xq + b * 32; + const int8_t *xq1 = xq + (b + 1) * 32; + const int8x16_t xv00 = vld1q_s8(xq0); + const int8x16_t xv01 = vld1q_s8(xq0 + 16); + const int8x16_t xv10 = vld1q_s8(xq1); + const int8x16_t xv11 = vld1q_s8(xq1 + 16); + + const int8_t *q00 = (const int8_t *)(row0 + b * 34 + 2); + const int8_t *q01 = (const int8_t *)(row0 + (b + 1) * 34 + 2); + const int8_t *q10 = (const int8_t *)(row1 + b * 34 + 2); + const int8_t *q11 = (const int8_t *)(row1 + (b + 1) * 34 + 2); + + int32x4_t d00 = vdupq_n_s32(0); + d00 = vdotq_s32(d00, vld1q_s8(q00), xv00); + d00 = vdotq_s32(d00, vld1q_s8(q00 + 16), xv01); + int32x4_t d01 = vdupq_n_s32(0); + d01 = vdotq_s32(d01, vld1q_s8(q01), xv10); + d01 = vdotq_s32(d01, vld1q_s8(q01 + 16), xv11); + int32x4_t d10 = vdupq_n_s32(0); + d10 = vdotq_s32(d10, vld1q_s8(q10), xv00); + d10 = vdotq_s32(d10, vld1q_s8(q10 + 16), xv01); + int32x4_t d11 = vdupq_n_s32(0); + d11 = vdotq_s32(d11, vld1q_s8(q11), xv10); + d11 = vdotq_s32(d11, vld1q_s8(q11 + 16), xv11); + + acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), f16_to_f32(s00) * xscale[b]); + acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), f16_to_f32(s01) * xscale[b + 1]); + acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), f16_to_f32(s10) * xscale[b]); + acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), f16_to_f32(s11) * xscale[b + 1]); + } + + if (b < blocks) { + uint16_t s0, s1; + memcpy(&s0, row0 + b * 34, sizeof(s0)); + memcpy(&s1, row1 + b * 34, sizeof(s1)); + const int8_t *xqb = xq + b * 32; + const int8x16_t xv0 = vld1q_s8(xqb); + const int8x16_t xv1 = vld1q_s8(xqb + 16); + const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); + const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); + int32x4_t d0 = vdupq_n_s32(0); + d0 = vdotq_s32(d0, vld1q_s8(q0), xv0); + d0 = vdotq_s32(d0, vld1q_s8(q0 + 16), xv1); + int32x4_t d1 = vdupq_n_s32(0); + d1 = vdotq_s32(d1, vld1q_s8(q1), xv0); + d1 = vdotq_s32(d1, vld1q_s8(q1 + 16), xv1); + acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), f16_to_f32(s0) * xscale[b]); + acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), f16_to_f32(s1) * xscale[b]); + } + + *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); + *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); + return; + } +#endif + + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t s0_bits; + uint16_t s1_bits; + memcpy(&s0_bits, row0 + b * 34, sizeof(s0_bits)); + memcpy(&s1_bits, row1 + b * 34, sizeof(s1_bits)); + const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); + const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); + const uint64_t i0 = b * 32; + const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; + acc0 += f16_to_f32(s0_bits) * xscale[b] * (float)dot_i8_32(q0, xq + i0, n); + acc1 += f16_to_f32(s1_bits) * xscale[b] * (float)dot_i8_32(q1, xq + i0, n); + } + *out0 = acc0; + *out1 = acc1; +} + +static void quantize_q8_0_activation(const float *x, int8_t *xq, float *scale, uint64_t n) { + const uint64_t blocks = (n + 31) / 32; + for (uint64_t b = 0; b < blocks; b++) { + const uint64_t i0 = b * 32; + const uint64_t bn = n - i0 < 32 ? n - i0 : 32; + float amax = 0.0f; + for (uint64_t i = 0; i < bn; i++) { + const float ax = fabsf(x[i0 + i]); + if (ax > amax) amax = ax; + } + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + scale[b] = d; + for (uint64_t i = 0; i < bn; i++) { + int v = (int)lrintf(x[i0 + i] * id); + if (v > 127) v = 127; + if (v < -128) v = -128; + xq[i0 + i] = (int8_t)v; + } + for (uint64_t i = bn; i < 32 && i0 + i < blocks * 32; i++) { + xq[i0 + i] = 0; + } + } +} + +static void quantize_q8_0_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + quantize_q8_0_batch_ctx *ctx = vctx; + for (uint64_t t = t0; t < t1; t++) { + quantize_q8_0_activation(ctx->x + t * ctx->in_dim, + ctx->xq + t * ctx->blocks * 32, + ctx->xscale + t * ctx->blocks, + ctx->in_dim); + } +} + +static void quantize_q8_0_activation_batch( + const float *x, + int8_t *xq, + float *xscale, + uint64_t n_tok, + uint64_t in_dim) { + quantize_q8_0_batch_ctx ctx = { + .x = x, + .xq = xq, + .xscale = xscale, + .in_dim = in_dim, + .blocks = (in_dim + 31) / 32, + }; + ds4_parallel_for(n_tok, quantize_q8_0_batch_worker, &ctx); +} + +static void matvec_q8_0_worker(void *vctx, uint64_t r0, uint64_t r1) { + matvec_q8_0_ctx *ctx = vctx; + + for (uint64_t r = r0; r < r1; r++) { + const uint64_t o = ctx->row0 + r; + const uint8_t *row = ctx->data + o * ctx->blocks * 34; + ctx->out[r] = dot_q8_0_row(row, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks); + } +} + +static void matvec_q8_0_pair_worker(void *vctx, uint64_t r0, uint64_t r1) { + matvec_q8_0_pair_ctx *ctx = vctx; + + for (uint64_t r = r0; r < r1; r++) { + const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; + const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; + dot_q8_0_row_pair(row0, row1, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks, + ctx->out0 + r, ctx->out1 + r); + } +} + +static void matvec_q8_0_grouped_worker(void *vctx, uint64_t r0, uint64_t r1) { + matvec_q8_0_grouped_ctx *ctx = vctx; + + for (uint64_t idx = r0; idx < r1; idx++) { + const uint64_t group = idx / ctx->rank; + const uint64_t row_in_group = idx - group * ctx->rank; + const uint64_t tensor_row = group * ctx->rank + row_in_group; + const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; + const int8_t *xq = ctx->xq + group * ctx->blocks * 32; + const float *xscale = ctx->xscale + group * ctx->blocks; + ctx->out[idx] = dot_q8_0_row(row, xq, xscale, ctx->in_dim, ctx->blocks); + } +} + +static void matmul_q8_0_grouped_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { + matmul_q8_0_grouped_batch_ctx *ctx = vctx; + + for (uint64_t idx = r0; idx < r1; idx++) { + const uint64_t group = idx / ctx->rank; + const uint64_t row_in_group = idx - group * ctx->rank; + const uint64_t tensor_row = group * ctx->rank + row_in_group; + const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; + + uint64_t t = 0; + for (; t + 1 < ctx->n_tok; t += 2) { + const uint64_t xbase0 = (t * ctx->n_groups + group) * ctx->blocks; + const uint64_t xbase1 = ((t + 1) * ctx->n_groups + group) * ctx->blocks; + dot_q8_0_row_2(row, + ctx->xq + xbase0 * 32, + ctx->xscale + xbase0, + ctx->xq + xbase1 * 32, + ctx->xscale + xbase1, + ctx->group_dim, + ctx->blocks, + ctx->out + t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group, + ctx->out + (t + 1) * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group); + } + for (; t < ctx->n_tok; t++) { + const uint64_t xbase = (t * ctx->n_groups + group) * ctx->blocks; + ctx->out[t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group] = + dot_q8_0_row(row, + ctx->xq + xbase * 32, + ctx->xscale + xbase, + ctx->group_dim, + ctx->blocks); + } + } +} + +static void matmul_q8_0_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { + matmul_q8_0_batch_ctx *ctx = vctx; + + for (uint64_t r = r0; r < r1; r++) { + const uint8_t *row = ctx->data + r * ctx->blocks * 34; + uint64_t t = 0; + for (; t + 1 < ctx->n_tok; t += 2) { + dot_q8_0_row_2(row, + ctx->xq + t * ctx->blocks * 32, + ctx->xscale + t * ctx->blocks, + ctx->xq + (t + 1) * ctx->blocks * 32, + ctx->xscale + (t + 1) * ctx->blocks, + ctx->in_dim, + ctx->blocks, + ctx->out + t * ctx->out_dim + r, + ctx->out + (t + 1) * ctx->out_dim + r); + } + for (; t < ctx->n_tok; t++) { + ctx->out[t * ctx->out_dim + r] = + dot_q8_0_row(row, + ctx->xq + t * ctx->blocks * 32, + ctx->xscale + t * ctx->blocks, + ctx->in_dim, + ctx->blocks); + } + } +} + +static void matmul_q8_0_pair_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { + matmul_q8_0_pair_batch_ctx *ctx = vctx; + + for (uint64_t r = r0; r < r1; r++) { + const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; + const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; + uint64_t t = 0; + for (; t + 1 < ctx->n_tok; t += 2) { + const int8_t *xq0 = ctx->xq + t * ctx->blocks * 32; + const float *xscale0 = ctx->xscale + t * ctx->blocks; + const int8_t *xq1 = ctx->xq + (t + 1) * ctx->blocks * 32; + const float *xscale1 = ctx->xscale + (t + 1) * ctx->blocks; + dot_q8_0_row_2(row0, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, + ctx->out0 + t * ctx->out_dim + r, + ctx->out0 + (t + 1) * ctx->out_dim + r); + dot_q8_0_row_2(row1, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, + ctx->out1 + t * ctx->out_dim + r, + ctx->out1 + (t + 1) * ctx->out_dim + r); + } + for (; t < ctx->n_tok; t++) { + const int8_t *xq = ctx->xq + t * ctx->blocks * 32; + const float *xscale = ctx->xscale + t * ctx->blocks; + dot_q8_0_row_pair(row0, row1, xq, xscale, ctx->in_dim, ctx->blocks, + ctx->out0 + t * ctx->out_dim + r, + ctx->out1 + t * ctx->out_dim + r); + } + } +} + +/* Multiply selected Q8_0 rows by an activation that has already been quantized + * once. This avoids repeated activation quantization for paired projections. */ +static void matvec_q8_0_rows_prequant( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const int8_t * xq, + const float * xscale, + uint64_t row0, + uint64_t n_rows) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t out_dim = w->dim[1]; + if (row0 > out_dim || n_rows > out_dim - row0) ds4_die("Q8_0 row range is outside tensor"); + const uint64_t ctx_blocks = (in_dim + 31) / 32; + + matvec_q8_0_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .xq = xq, + .xscale = xscale, + .in_dim = in_dim, + .row0 = row0, + .blocks = ctx_blocks, + }; + ds4_parallel_for(n_rows, matvec_q8_0_worker, &ctx); +} + +static DS4_MAYBE_UNUSED void matvec_q8_0_prequant( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const int8_t * xq, + const float * xscale) { + matvec_q8_0_rows_prequant(out, m, w, xq, xscale, 0, w->dim[1]); +} + +static void matvec_q8_0_3d_slice_prequant( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const int8_t * xq, + const float * xscale, + uint64_t slice) { + if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); + if (slice >= w->dim[2]) ds4_die("Q8_0 slice is outside tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t out_dim = w->dim[1]; + const uint64_t blocks = (in_dim + 31) / 32; + const uint64_t slice_bytes = out_dim * blocks * 34; + const uint8_t *data = (const uint8_t *)tensor_data(m, w) + slice * slice_bytes; + + matvec_q8_0_ctx ctx = { + .out = out, + .data = data, + .xq = xq, + .xscale = xscale, + .in_dim = in_dim, + .row0 = 0, + .blocks = blocks, + }; + ds4_parallel_for(out_dim, matvec_q8_0_worker, &ctx); +} + +static DS4_MAYBE_UNUSED void matvec_q8_0_3d_slice( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint64_t slice) { + if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)blocks * 32); + float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); + + quantize_q8_0_activation(x, xq, xscale, in_dim); + matvec_q8_0_3d_slice_prequant(out, m, w, xq, xscale, slice); + + free(xscale); + free(xq); +} + +/* Compute two Q8_0 projections from the same input, used by gate/up and + * compressor kv/score pairs. */ +static void matvec_q8_0_pair_prequant( + float * out0, + float * out1, + const ds4_model * m, + const ds4_tensor * w0, + const ds4_tensor * w1, + const int8_t * xq, + const float * xscale) { + if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { + ds4_die("expected two 2D Q8_0 tensors"); + } + if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { + ds4_die("paired Q8_0 tensors do not have the same shape"); + } + + const uint64_t in_dim = w0->dim[0]; + matvec_q8_0_pair_ctx ctx = { + .out0 = out0, + .out1 = out1, + .data0 = tensor_data(m, w0), + .data1 = tensor_data(m, w1), + .xq = xq, + .xscale = xscale, + .in_dim = in_dim, + .blocks = (in_dim + 31) / 32, + }; + ds4_parallel_for(w0->dim[1], matvec_q8_0_pair_worker, &ctx); +} + +static void matmul_q8_0_batch_prequant( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const int8_t * xq, + const float * xscale, + uint64_t n_tok) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + + matmul_q8_0_batch_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .xq = xq, + .xscale = xscale, + .n_tok = n_tok, + .in_dim = w->dim[0], + .out_dim = w->dim[1], + .blocks = (w->dim[0] + 31) / 32, + }; + ds4_parallel_for(ctx.out_dim, matmul_q8_0_batch_worker, &ctx); +} + +static void matmul_q8_0_pair_batch_prequant( + float * out0, + float * out1, + const ds4_model * m, + const ds4_tensor * w0, + const ds4_tensor * w1, + const int8_t * xq, + const float * xscale, + uint64_t n_tok) { + if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { + ds4_die("expected two 2D Q8_0 tensors"); + } + if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { + ds4_die("paired Q8_0 tensors do not have the same shape"); + } + + matmul_q8_0_pair_batch_ctx ctx = { + .out0 = out0, + .out1 = out1, + .data0 = tensor_data(m, w0), + .data1 = tensor_data(m, w1), + .xq = xq, + .xscale = xscale, + .n_tok = n_tok, + .in_dim = w0->dim[0], + .out_dim = w0->dim[1], + .blocks = (w0->dim[0] + 31) / 32, + }; + ds4_parallel_for(ctx.out_dim, matmul_q8_0_pair_batch_worker, &ctx); +} + +/* Batched Q8_0 matmul for prefill: quantize all token activations, then scan + * weight rows once per output channel. */ +static void matmul_q8_0_batch( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint64_t n_tok) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); + float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); + + quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); + matmul_q8_0_batch_prequant(out, m, w, xq, xscale, n_tok); + + free(xscale); + free(xq); +} + +static void matmul_q8_0_pair_batch( + float * out0, + float * out1, + const ds4_model * m, + const ds4_tensor * w0, + const ds4_tensor * w1, + const float * x, + uint64_t n_tok) { + if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { + ds4_die("expected two 2D Q8_0 tensors"); + } + if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { + ds4_die("paired Q8_0 tensors do not have the same shape"); + } + + const uint64_t in_dim = w0->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); + float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); + + quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); + matmul_q8_0_pair_batch_prequant(out0, out1, m, w0, w1, xq, xscale, n_tok); + + free(xscale); + free(xq); +} + +static void matvec_q8_0_rows( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint64_t row0, + uint64_t n_rows) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + + const uint64_t in_dim = w->dim[0]; + const uint64_t ctx_blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)ctx_blocks * 32); + float *xscale = xmalloc((size_t)ctx_blocks * sizeof(xscale[0])); + + quantize_q8_0_activation(x, xq, xscale, in_dim); + matvec_q8_0_rows_prequant(out, m, w, xq, xscale, row0, n_rows); + + free(xscale); + free(xq); +} + +/* Single-token Q8_0 matvec, used heavily in decode. */ +static void matvec_q8_0(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { + matvec_q8_0_rows(out, m, w, x, 0, w->dim[1]); +} + +static inline float dot_q8_0_row_f32_ref( + const uint8_t *row, + const float *x, + uint64_t in_dim, + uint64_t blocks) { + float acc = 0.0f; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); + const int8_t *qs = (const int8_t *)(row + b * 34 + 2); + const float d = f16_to_f32(scale_bits); + const uint64_t i0 = b * 32; + const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; + for (uint64_t i = 0; i < n; i++) { + acc += d * (float)qs[i] * x[i0 + i]; + } + } + return acc; +} + +typedef struct { + float *out; + const uint8_t *data; + const float *x; + uint64_t in_dim; + uint64_t blocks; +} matvec_q8_0_f32_ref_ctx; + +static void matvec_q8_0_f32_ref_worker(void *vctx, uint64_t r0, uint64_t r1) { + matvec_q8_0_f32_ref_ctx *ctx = vctx; + const uint64_t row_bytes = ctx->blocks * 34; + for (uint64_t r = r0; r < r1; r++) { + ctx->out[r] = dot_q8_0_row_f32_ref(ctx->data + r * row_bytes, + ctx->x, + ctx->in_dim, + ctx->blocks); + } +} + +static void matvec_q8_0_f32_ref( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const float *x) { + if (w->type != DS4_TENSOR_Q8_0 || w->ndim < 2 || w->dim[0] == 0) { + ds4_die("expected a Q8_0 tensor with matrix rows"); + } + matvec_q8_0_f32_ref_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .x = x, + .in_dim = w->dim[0], + .blocks = (w->dim[0] + 31) / 32, + }; + ds4_parallel_for(w->elements / w->dim[0], matvec_q8_0_f32_ref_worker, &ctx); +} + +static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x); + +/* Decode scratch owns this temporary activation quantization so generation + * can assert that the hot path performs no malloc. */ +static void cpu_decode_quantize_q8_0( + ds4_cpu_decode_scratch * scratch, + const float * x, + uint64_t in_dim) { + if (in_dim > scratch->q8_cap) ds4_die("CPU decode Q8_0 scratch buffer is too small"); + quantize_q8_0_activation(x, scratch->q8_xq, scratch->q8_xscale, in_dim); +} + +static void matvec_q8_0_decode_scratch( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + ds4_cpu_decode_scratch * scratch) { + cpu_decode_quantize_q8_0(scratch, x, w->dim[0]); + matvec_q8_0_prequant(out, m, w, scratch->q8_xq, scratch->q8_xscale); +} + +static void matvec_q8_0_pair_decode_scratch( + float * out0, + float * out1, + const ds4_model * m, + const ds4_tensor * w0, + const ds4_tensor * w1, + const float * x, + ds4_cpu_decode_scratch * scratch) { + cpu_decode_quantize_q8_0(scratch, x, w0->dim[0]); + matvec_q8_0_pair_prequant(out0, out1, m, w0, w1, scratch->q8_xq, scratch->q8_xscale); +} + +static void matvec_any_decode_scratch( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + ds4_cpu_decode_scratch * scratch) { + if (w->type == 8) { + matvec_q8_0_decode_scratch(out, m, w, x, scratch); + } else { + matvec_any(out, m, w, x); + } +} + +static void matvec_q8_0_grouped_rows( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint32_t n_groups, + uint64_t group_dim, + uint64_t rank) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { + ds4_die("grouped Q8_0 tensor has an unexpected layout"); + } + + const uint64_t blocks = (group_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)n_groups * blocks * 32); + float *xscale = xmalloc((size_t)n_groups * blocks * sizeof(xscale[0])); + + for (uint32_t g = 0; g < n_groups; g++) { + quantize_q8_0_activation(x + (uint64_t)g * group_dim, + xq + (uint64_t)g * blocks * 32, + xscale + (uint64_t)g * blocks, + group_dim); + } + + matvec_q8_0_grouped_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .xq = xq, + .xscale = xscale, + .in_dim = group_dim, + .blocks = blocks, + .rank = rank, + }; + ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); + + free(xscale); + free(xq); +} + +static void matvec_q8_0_grouped_rows_decode_scratch( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint32_t n_groups, + uint64_t group_dim, + uint64_t rank, + ds4_cpu_decode_scratch * scratch) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { + ds4_die("grouped Q8_0 tensor has an unexpected layout"); + } + if ((uint64_t)n_groups * group_dim > scratch->q8_cap) { + ds4_die("CPU decode grouped Q8_0 scratch buffer is too small"); + } + + const uint64_t blocks = (group_dim + 31) / 32; + for (uint32_t g = 0; g < n_groups; g++) { + quantize_q8_0_activation(x + (uint64_t)g * group_dim, + scratch->q8_xq + (uint64_t)g * blocks * 32, + scratch->q8_xscale + (uint64_t)g * blocks, + group_dim); + } + + matvec_q8_0_grouped_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .xq = scratch->q8_xq, + .xscale = scratch->q8_xscale, + .in_dim = group_dim, + .blocks = blocks, + .rank = rank, + }; + ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); +} + +static void matmul_q8_0_grouped_batch( + float * out, + const ds4_model * m, + const ds4_tensor * w, + const float * x, + uint64_t n_tok, + uint32_t n_groups, + uint64_t group_dim, + uint64_t rank) { + if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); + if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { + ds4_die("grouped Q8_0 tensor has an unexpected layout"); + } + + const uint64_t blocks = (group_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)n_tok * n_groups * blocks * 32); + float *xscale = xmalloc((size_t)n_tok * n_groups * blocks * sizeof(xscale[0])); + + for (uint64_t t = 0; t < n_tok; t++) { + for (uint32_t g = 0; g < n_groups; g++) { + const uint64_t xbase = (t * n_groups + g) * blocks; + quantize_q8_0_activation(x + t * n_groups * group_dim + (uint64_t)g * group_dim, + xq + xbase * 32, + xscale + xbase, + group_dim); + } + } + + matmul_q8_0_grouped_batch_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .xq = xq, + .xscale = xscale, + .n_tok = n_tok, + .n_groups = n_groups, + .group_dim = group_dim, + .blocks = blocks, + .rank = rank, + }; + ds4_parallel_for((uint64_t)n_groups * rank, matmul_q8_0_grouped_batch_worker, &ctx); + + free(xscale); + free(xq); +} + +typedef struct { + float *out; + const float *data; + const float *x; + uint64_t in_dim; +} matvec_f32_ctx; + +static void matvec_f32_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_f32_ctx *ctx = vctx; + + for (uint64_t o = row0; o < row1; o++) { + double acc = 0.0; + const float *row = ctx->data + o * ctx->in_dim; + for (uint64_t i = 0; i < ctx->in_dim; i++) { + acc += (double)row[i] * ctx->x[i]; + } + ctx->out[o] = (float)acc; + } +} + +static void matvec_f32(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { + if (w->type != 0 || w->ndim != 2) ds4_die("expected a 2D F32 tensor"); + + matvec_f32_ctx ctx = { + .out = out, + .data = tensor_data(m, w), + .x = x, + .in_dim = w->dim[0], + }; + ds4_parallel_for(w->dim[1], matvec_f32_worker, &ctx); +} + +/* Dispatch for dense F32/F16/Q8_0 tensors used by auxiliary projections. */ +static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { + switch (w->type) { + case 0: matvec_f32(out, m, w, x); break; + case 1: matvec_f16(out, m, w, x); break; + case 8: matvec_q8_0(out, m, w, x); break; + default: + ds4_die("unsupported tensor type for dense matvec"); + } +} + +static float tensor_1d_value(const ds4_model *m, const ds4_tensor *t, uint64_t i) { + if (i >= t->elements) ds4_die("tensor scalar index is out of bounds"); + if (t->type == 0) { + const float *p = tensor_data(m, t); + return p[i]; + } + if (t->type == 1) { + const uint16_t *p = tensor_data(m, t); + return f16_to_f32(p[i]); + } + ds4_die("unsupported tensor scalar type"); + return 0.0f; +} + +static float tensor_2d_value(const ds4_model *m, const ds4_tensor *t, uint64_t x, uint64_t y) { + if (t->ndim != 2 || x >= t->dim[0] || y >= t->dim[1]) { + ds4_die("tensor 2D index is out of bounds"); + } + return tensor_1d_value(m, t, y * t->dim[0] + x); +} + +/* Locate one expert's 2D matrix inside a 3D GGUF expert tensor. */ +static const uint8_t *tensor_expert_bytes( + const ds4_model *m, + const ds4_tensor *w, + uint32_t expert, + uint64_t *in_dim, + uint64_t *out_dim, + uint64_t *row_bytes) { + if (w->ndim != 3) ds4_die("expected a 3D expert tensor"); + if (expert >= w->dim[2]) ds4_die("expert id is outside expert tensor"); + + *in_dim = w->dim[0]; + *out_dim = w->dim[1]; + + const gguf_type_info *info = tensor_type(w->type); + if (!info || info->block_elems == 0) ds4_die("unsupported expert tensor type"); + const uint64_t blocks = (*in_dim + info->block_elems - 1) / info->block_elems; + *row_bytes = blocks * info->block_bytes; + + const uint64_t expert_bytes = *out_dim * *row_bytes; + return (const uint8_t *)tensor_data(m, w) + (uint64_t)expert * expert_bytes; +} + +typedef struct { + float *out0; + float *out1; + const uint8_t *base0; + const uint8_t *base1; + const block_q8_K *xq; + uint64_t in_dim; + uint64_t row_bytes0; + uint64_t row_bytes1; +} matvec_iq2_xxs_pair_ctx; + +static void matvec_iq2_xxs_pair_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_iq2_xxs_pair_ctx *ctx = vctx; + for (uint64_t row = row0; row < row1; row++) { + const block_iq2_xxs *br0 = (const block_iq2_xxs *)(ctx->base0 + row * ctx->row_bytes0); + const block_iq2_xxs *br1 = (const block_iq2_xxs *)(ctx->base1 + row * ctx->row_bytes1); + ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &ctx->out0[row], &ctx->out1[row], br0, br1, ctx->xq); + } +} + +/* Project one routed expert's gate and up matrices. Both are IQ2_XXS and + * share the same Q8_K activation. */ +static void matvec_iq2_xxs_expert_pair_prequant( + float *out0, + float *out1, + const ds4_model *m, + const ds4_tensor *w0, + const ds4_tensor *w1, + const block_q8_K *xq, + uint32_t expert) { + if (w0->type != 16 || w1->type != 16) ds4_die("expected IQ2_XXS expert tensors"); + + uint64_t in_dim0, out_dim0, row_bytes0; + uint64_t in_dim1, out_dim1, row_bytes1; + const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &row_bytes0); + const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &row_bytes1); + if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired IQ2_XXS expert tensors do not match"); + if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); + + matvec_iq2_xxs_pair_ctx ctx = { + .out0 = out0, + .out1 = out1, + .base0 = base0, + .base1 = base1, + .xq = xq, + .in_dim = in_dim0, + .row_bytes0 = row_bytes0, + .row_bytes1 = row_bytes1, + }; + ds4_parallel_for(out_dim0, matvec_iq2_xxs_pair_worker, &ctx); +} + +static float silu(float x); + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_iq2_xxs_mid_ctx; + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const int8_t *xq; + const float *xscale; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t blocks; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q8_0_mid_ctx; + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q8_k_mid_ctx; + +static void matvec_iq2_xxs_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_iq2_xxs_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); + const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); + ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +static void matvec_q8_0_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_0_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + const uint8_t *gate_row = ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]; + const uint8_t *up_row = ctx->up_base[slot] + row * ctx->up_row_bytes[slot]; + dot_q8_0_row_pair(gate_row, up_row, ctx->xq, ctx->xscale, + ctx->in_dim, ctx->blocks, &gate, &up); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +static void matvec_q8_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_k_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); + const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); + ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +/* Build all selected expert hidden vectors: IQ2_XXS gate/up, clamp, SwiGLU, + * and router weight. The down projection runs later on the quantized mids. */ +static void matvec_iq2_xxs_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type != 16 || up_w->type != 16) ds4_die("expected IQ2_XXS expert tensors"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_iq2_xxs_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .clamp = clamp, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired IQ2_XXS expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("IQ2_XXS expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_iq2_xxs_mid_worker, &ctx); +} + +static DS4_MAYBE_UNUSED void matvec_q8_0_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const int8_t *xq, + const float *xscale, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type != DS4_TENSOR_Q8_0 || up_w->type != DS4_TENSOR_Q8_0) { + ds4_die("expected Q8_0 expert tensors"); + } + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q8_0_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .xscale = xscale, + .clamp = clamp, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired Q8_0 expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("Q8_0 expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ctx.blocks = in_dim0 / 32u; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_0_mid_worker, &ctx); +} + +static DS4_MAYBE_UNUSED void matvec_q8_k_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type != DS4_TENSOR_Q8_K || up_w->type != DS4_TENSOR_Q8_K) { + ds4_die("expected Q8_K expert tensors"); + } + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q8_k_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .clamp = clamp, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired Q8_K expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("Q8_K expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_k_mid_worker, &ctx); +} + +typedef struct { + float *out; + const uint8_t *base; + const block_q8_K *xq; + uint64_t in_dim; + uint64_t row_bytes; +} matvec_q2_k_ctx; + +static void matvec_q2_k_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q2_k_ctx *ctx = vctx; + for (uint64_t row = row0; row < row1; row++) { + const block_q2_K *br = (const block_q2_K *)(ctx->base + row * ctx->row_bytes); + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &ctx->out[row], br, ctx->xq); + } +} + +/* Single expert Q2_K down projection, kept mostly for tracing and diagnostics. */ +static void matvec_q2_k_expert( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const float *x, + uint32_t expert) { + if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); + + uint64_t in_dim, out_dim, row_bytes; + const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); + if (in_dim % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); + + block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); + ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); + + matvec_q2_k_ctx ctx = { + .out = out, + .base = base, + .xq = xq, + .in_dim = in_dim, + .row_bytes = row_bytes, + }; + ds4_parallel_for(out_dim, matvec_q2_k_worker, &ctx); + + free(xq); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q2_k_accum_ctx; + +static void matvec_q2_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q2_k_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + float v = 0.0f; + const block_q2_K *br = (const block_q2_K *)(ctx->base[i] + row * ctx->row_bytes[i]); + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); + acc += v; + } + ctx->out[row] = acc; + } +} + +/* Accumulate all selected experts' Q2_K down projections directly into the + * 4096-wide MoE output. */ +static void matvec_q2_k_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("Q2_K expert tensors do not share a layout"); + } + } + if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); + + const uint64_t n_blocks = in_dim0 / QK_K; + matvec_q2_k_accum_ctx ctx = { + .out = out, + .in_dim = in_dim0, + .n_expert = n_expert, + }; + for (int i = 0; i < n_expert; i++) { + ctx.base[i] = base[i]; + ctx.row_bytes[i] = row_bytes[i]; + ctx.xq[i] = xq + (uint64_t)i * n_blocks; + } + + ds4_parallel_for(out_dim0, matvec_q2_k_accum_worker, &ctx); +} + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q2_k_mid_ctx; + +static void matvec_q2_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q2_k_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); + + const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +static void matvec_q2_k_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type != DS4_TENSOR_Q2_K || up_w->type != DS4_TENSOR_Q2_K) { + ds4_die("expected Q2_K expert tensors"); + } + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q2_k_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .clamp = clamp, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired Q2_K expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("Q2_K expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q2_k_mid_worker, &ctx); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const int8_t *xq[DS4_MAX_EXPERT_USED]; + const float *xscale[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t blocks; + int n_expert; +} matvec_q8_0_accum_ctx; + +static void matvec_q8_0_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_0_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + const uint8_t *br = ctx->base[i] + row * ctx->row_bytes[i]; + acc += dot_q8_0_row(br, ctx->xq[i], ctx->xscale[i], + ctx->in_dim, ctx->blocks); + } + ctx->out[row] = acc; + } +} + +static void matvec_q8_0_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const int8_t *xq, + const float *xscale, + const int *selected, + int n_expert) { + if (w->type != DS4_TENSOR_Q8_0) ds4_die("expected a Q8_0 expert tensor"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q8_0_accum_ctx ctx = { + .out = out, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], + &in_dim, &out_dim, &ctx.row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("Q8_0 expert tensors do not share a layout"); + } + } + if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); + + const uint64_t blocks0 = in_dim0 / 32u; + ctx.in_dim = in_dim0; + ctx.blocks = blocks0; + for (int i = 0; i < n_expert; i++) { + ctx.xq[i] = xq + (uint64_t)i * blocks0 * 32u; + ctx.xscale[i] = xscale + (uint64_t)i * blocks0; + } + + ds4_parallel_for(out_dim0, matvec_q8_0_accum_worker, &ctx); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q8_k_accum_ctx; + +static void matvec_q8_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_k_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + float v = 0.0f; + const block_q8_K *br = (const block_q8_K *)(ctx->base[i] + row * ctx->row_bytes[i]); + ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); + acc += v; + } + ctx->out[row] = acc; + } +} + +static void matvec_q8_k_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type != DS4_TENSOR_Q8_K) ds4_die("expected a Q8_K expert tensor"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q8_k_accum_ctx ctx = { + .out = out, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], + &in_dim, &out_dim, &ctx.row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("Q8_K expert tensors do not share a layout"); + } + } + if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); + + const uint64_t n_blocks = in_dim0 / QK_K; + ctx.in_dim = in_dim0; + for (int i = 0; i < n_expert; i++) { + ctx.xq[i] = xq + (uint64_t)i * n_blocks; + } + + ds4_parallel_for(out_dim0, matvec_q8_k_accum_worker, &ctx); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_iq2_xxs_accum_ctx; + +static void matvec_iq2_xxs_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_iq2_xxs_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + float v = 0.0f; + const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[i] + row * ctx->row_bytes[i]); + ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); + acc += v; + } + ctx->out[row] = acc; + } +} + +static void matvec_iq2_xxs_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type != DS4_TENSOR_IQ2_XXS) ds4_die("expected an IQ2_XXS expert tensor"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("IQ2_XXS expert tensors do not share a layout"); + } + } + if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); + + const uint64_t n_blocks = in_dim0 / QK_K; + matvec_iq2_xxs_accum_ctx ctx = { + .out = out, + .in_dim = in_dim0, + .n_expert = n_expert, + }; + for (int i = 0; i < n_expert; i++) { + ctx.base[i] = base[i]; + ctx.row_bytes[i] = row_bytes[i]; + ctx.xq[i] = xq + (uint64_t)i * n_blocks; + } + + ds4_parallel_for(out_dim0, matvec_iq2_xxs_accum_worker, &ctx); +} + +typedef struct { + uint32_t token; + uint32_t slot; +} ds4_expert_pair; + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT]; + const uint8_t *up_base[DS4_MAX_EXPERT]; + const block_q8_K *xq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + const float *pair_weight; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT]; + uint64_t up_row_bytes[DS4_MAX_EXPERT]; + uint64_t xq_blocks; +} matvec_q2_k_batch_mid_ctx; + +static void matvec_q2_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_q2_k_batch_mid_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + + const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); + const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; + float gate = 0.0f; + float up = 0.0f; + + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + + ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; + } + } +} + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT]; + const uint8_t *up_base[DS4_MAX_EXPERT]; + const block_q8_K *xq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + const float *pair_weight; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT]; + uint64_t up_row_bytes[DS4_MAX_EXPERT]; + uint64_t xq_blocks; +} matvec_iq2_xxs_batch_mid_ctx; + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT]; + const uint8_t *up_base[DS4_MAX_EXPERT]; + const block_q8_K *xq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + const float *pair_weight; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t xq_blocks; + uint64_t gate_row_bytes[DS4_MAX_EXPERT]; + uint64_t up_row_bytes[DS4_MAX_EXPERT]; +} matvec_q8_k_batch_mid_ctx; + +static void matvec_iq2_xxs_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_iq2_xxs_batch_mid_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + + const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); + const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; + float gate = 0.0f; + float up = 0.0f; + + ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + + ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; + } + } +} + +static void matvec_q8_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_q8_k_batch_mid_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + + const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); + const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; + float gate = 0.0f; + float up = 0.0f; + + ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + + ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; + } + } +} + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT]; + const uint8_t *up_base[DS4_MAX_EXPERT]; + const int8_t *xq; + const float *xscale; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + const float *pair_weight; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t blocks; + uint64_t gate_row_bytes[DS4_MAX_EXPERT]; + uint64_t up_row_bytes[DS4_MAX_EXPERT]; +} matvec_q8_0_batch_mid_ctx; + +static void matvec_q8_0_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_q8_0_batch_mid_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + + const uint8_t *gate_row = ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]; + const uint8_t *up_row = ctx->up_base[expert] + row * ctx->up_row_bytes[expert]; + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + float gate = 0.0f; + float up = 0.0f; + + dot_q8_0_row_pair(gate_row, up_row, + ctx->xq + (uint64_t)pair.token * ctx->blocks * 32u, + ctx->xscale + (uint64_t)pair.token * ctx->blocks, + ctx->in_dim, ctx->blocks, &gate, &up); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + + ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; + } + } +} + +typedef struct { + const float *mid; + block_q8_K *midq; + uint64_t down_in_dim; + uint64_t down_blocks; +} quantize_mid_pairs_ctx; + +static void quantize_mid_pairs_worker(void *vctx, uint64_t p0, uint64_t p1) { + quantize_mid_pairs_ctx *ctx = vctx; + for (uint64_t p = p0; p < p1; p++) { + ds4_quantize_row_q8_K(ctx->mid + p * ctx->down_in_dim, + ctx->midq + p * ctx->down_blocks, + (int64_t)ctx->down_in_dim); + } +} + +typedef struct { + float *down_pair; + const uint8_t *base[DS4_MAX_EXPERT]; + const block_q8_K *midq; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t midq_blocks; +} matvec_q2_k_batch_down_ctx; + +static DS4_MAYBE_UNUSED void matvec_q2_k_batch_down_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_q2_k_batch_down_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, + ctx->down_pair + (uint64_t)pair_id * ctx->out_dim + row, + br, xq); + } + } +} + +typedef struct { + float *moe; + const uint8_t *base[DS4_MAX_EXPERT]; + const block_q8_K *midq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint32_t n_active; + uint32_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t midq_blocks; +} matvec_q2_k_batch_accum_rows_ctx; + +static void matvec_q2_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q2_k_batch_accum_rows_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + for (uint32_t t = 0; t < ctx->n_tok; t++) { + ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; + } + + for (uint32_t ai = 0; ai < ctx->n_active; ai++) { + const uint32_t expert = ctx->active_expert[ai]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; + float v = 0.0f; + + ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, xq); + ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; + } + } + } +} + +/* ========================================================================= + * Q4_K routed expert matrix-vector products. + * ========================================================================= */ + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q4_k_mid_ctx; + +static void matvec_q4_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q4_k_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); + + const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +static void matvec_q4_k_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type != DS4_TENSOR_Q4_K || up_w->type != DS4_TENSOR_Q4_K) + ds4_die("expected Q4_K expert tensors"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q4_k_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .clamp = clamp, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired Q4_K expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("Q4_K expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q4_k_mid_worker, &ctx); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + int n_expert; +} matvec_q4_k_accum_ctx; + +static void matvec_q4_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q4_k_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + float v = 0.0f; + const block_q4_K *br = (const block_q4_K *)(ctx->base[i] + row * ctx->row_bytes[i]); + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); + acc += v; + } + ctx->out[row] = acc; + } +} + +static void matvec_q4_k_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type != DS4_TENSOR_Q4_K) ds4_die("expected a Q4_K expert tensor"); + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("Q4_K expert tensors do not share a layout"); + } + } + if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); + + const uint64_t n_blocks = in_dim0 / QK_K; + matvec_q4_k_accum_ctx ctx = { + .out = out, + .in_dim = in_dim0, + .n_expert = n_expert, + }; + for (int i = 0; i < n_expert; i++) { + ctx.base[i] = base[i]; + ctx.row_bytes[i] = row_bytes[i]; + ctx.xq[i] = xq + (uint64_t)i * n_blocks; + } + + ds4_parallel_for(out_dim0, matvec_q4_k_accum_worker, &ctx); +} + +static inline void ds4_vec_dot_q5_q6_K_q8_K( + uint32_t type, + int n, + float *s, + const uint8_t *x, + const block_q8_K *y) { + if (type == DS4_TENSOR_Q5_K) { + ds4_vec_dot_q5_K_q8_K(n, s, (const block_q5_K *)x, y); + } else if (type == DS4_TENSOR_Q6_K) { + ds4_vec_dot_q6_K_q8_K(n, s, (const block_q6_K *)x, y); + } else { + ds4_die("expected a Q5_K or Q6_K tensor"); + } +} + +typedef struct { + float *out; + const uint8_t *base; + const block_q8_K *xq; + uint64_t in_dim; + uint64_t row_bytes; + uint32_t type; +} matvec_q5_q6_k_ctx; + +static void matvec_q5_q6_k_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q5_q6_k_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &ctx->out[row], + ctx->base + row * ctx->row_bytes, ctx->xq); + } +} + +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq; + float expert_weight[DS4_MAX_EXPERT_USED]; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + uint32_t gate_type; + uint32_t up_type; + int n_expert; +} matvec_q5_q6_k_mid_ctx; + +static void matvec_q5_q6_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q5_q6_k_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const int slot = (int)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + float gate = 0.0f; + float up = 0.0f; + + ds4_vec_dot_q5_q6_K_q8_K(ctx->gate_type, (int)ctx->in_dim, &gate, + ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot], + ctx->xq); + ds4_vec_dot_q5_q6_K_q8_K(ctx->up_type, (int)ctx->in_dim, &up, + ctx->up_base[slot] + row * ctx->up_row_bytes[slot], + ctx->xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +static void matvec_q5_q6_k_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if ((gate_w->type != DS4_TENSOR_Q5_K && gate_w->type != DS4_TENSOR_Q6_K) || + (up_w->type != DS4_TENSOR_Q5_K && up_w->type != DS4_TENSOR_Q6_K)) { + ds4_die("expected Q5_K/Q6_K expert tensors"); + } + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + matvec_q5_q6_k_mid_ctx ctx = { + .mid = mid, + .xq = xq, + .clamp = clamp, + .gate_type = gate_w->type, + .up_type = up_w->type, + .n_expert = n_expert, + }; + + for (int i = 0; i < n_expert; i++) { + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], + &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); + ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], + &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); + if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { + ds4_die("paired Q5_K/Q6_K expert tensors do not match"); + } + if (i == 0) { + in_dim0 = gate_in_dim; + out_dim0 = gate_out_dim; + } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { + ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); + } + ctx.expert_weight[i] = expert_weight[i]; + } + if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); + + ctx.in_dim = in_dim0; + ctx.out_dim = out_dim0; + ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q5_q6_k_mid_worker, &ctx); +} + +static void matvec_q5_q6_k_expert( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const float *x, + uint32_t expert) { + if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { + ds4_die("expected a Q5_K or Q6_K expert tensor"); + } + + uint64_t in_dim, out_dim, row_bytes; + const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); + if (in_dim % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); + + block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); + ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); + + matvec_q5_q6_k_ctx ctx = { + .out = out, + .base = base, + .xq = xq, + .in_dim = in_dim, + .row_bytes = row_bytes, + .type = w->type, + }; + ds4_parallel_for(out_dim, matvec_q5_q6_k_worker, &ctx); + + free(xq); +} + +typedef struct { + float *out; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + const block_q8_K *xq[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + uint32_t type; + int n_expert; +} matvec_q5_q6_k_accum_ctx; + +static void matvec_q5_q6_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q5_q6_k_accum_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (int i = 0; i < ctx->n_expert; i++) { + float v = 0.0f; + ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &v, + ctx->base[i] + row * ctx->row_bytes[i], + ctx->xq[i]); + acc += v; + } + ctx->out[row] = acc; + } +} + +static void matvec_q5_q6_k_experts_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { + ds4_die("expected a Q5_K or Q6_K expert tensor"); + } + if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); + + uint64_t in_dim0 = 0; + uint64_t out_dim0 = 0; + const uint8_t *base[DS4_MAX_EXPERT_USED]; + uint64_t row_bytes[DS4_MAX_EXPERT_USED]; + + for (int i = 0; i < n_expert; i++) { + uint64_t in_dim, out_dim; + base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); + if (i == 0) { + in_dim0 = in_dim; + out_dim0 = out_dim; + } else if (in_dim != in_dim0 || out_dim != out_dim0) { + ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); + } + } + if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); + + const uint64_t n_blocks = in_dim0 / QK_K; + matvec_q5_q6_k_accum_ctx ctx = { + .out = out, + .in_dim = in_dim0, + .type = w->type, + .n_expert = n_expert, + }; + for (int i = 0; i < n_expert; i++) { + ctx.base[i] = base[i]; + ctx.row_bytes[i] = row_bytes[i]; + ctx.xq[i] = xq + (uint64_t)i * n_blocks; + } + + ds4_parallel_for(out_dim0, matvec_q5_q6_k_accum_worker, &ctx); +} + +/* Q4_K batch mid worker: same structure as IQ2_XXS batch but uses Q4_K dot. */ +typedef struct { + float *mid; + const uint8_t *gate_base[DS4_MAX_EXPERT]; + const uint8_t *up_base[DS4_MAX_EXPERT]; + const block_q8_K *xq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + const float *pair_weight; + float clamp; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT]; + uint64_t up_row_bytes[DS4_MAX_EXPERT]; + uint64_t xq_blocks; +} matvec_q4_k_batch_mid_ctx; + +static void matvec_q4_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { + matvec_q4_k_batch_mid_ctx *ctx = vctx; + + for (uint64_t task = task0; task < task1; task++) { + const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); + const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; + const uint32_t expert = ctx->active_expert[active_idx]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + + const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); + const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; + float gate = 0.0f; + float up = 0.0f; + + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, xq); + + if (ctx->clamp > 1.0e-6f) { + if (gate > ctx->clamp) gate = ctx->clamp; + if (up > ctx->clamp) up = ctx->clamp; + if (up < -ctx->clamp) up = -ctx->clamp; + } + + ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; + } + } +} + +/* Q4_K batch down accum worker: same structure as Q2_K batch but uses Q4_K dot. */ +typedef struct { + float *moe; + const uint8_t *base[DS4_MAX_EXPERT]; + const block_q8_K *midq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint32_t n_active; + uint32_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t midq_blocks; +} matvec_q4_k_batch_accum_rows_ctx; + +static void matvec_q4_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q4_k_batch_accum_rows_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + for (uint32_t t = 0; t < ctx->n_tok; t++) { + ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; + } + + for (uint32_t ai = 0; ai < ctx->n_active; ai++) { + const uint32_t expert = ctx->active_expert[ai]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const block_q4_K *br = (const block_q4_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; + float v = 0.0f; + + ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, xq); + ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; + } + } + } +} + +typedef struct { + float *moe; + const uint8_t *base[DS4_MAX_EXPERT]; + const block_q8_K *midq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint32_t n_active; + uint32_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t midq_blocks; +} matvec_iq2_xxs_batch_accum_rows_ctx; + +static void matvec_iq2_xxs_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_iq2_xxs_batch_accum_rows_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + for (uint32_t t = 0; t < ctx->n_tok; t++) { + ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; + } + + for (uint32_t ai = 0; ai < ctx->n_active; ai++) { + const uint32_t expert = ctx->active_expert[ai]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[expert] + row * ctx->row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; + float v = 0.0f; + + ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, xq); + ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; + } + } + } +} + +/* Dispatch: call the right gate/up mid builder based on tensor type. */ +static void matvec_experts_mid_prequant( + float *mid, + const ds4_model *m, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + const block_q8_K *xq, + const int *selected, + const float *expert_weight, + int n_expert, + float clamp) { + if (gate_w->type == DS4_TENSOR_IQ2_XXS) { + matvec_iq2_xxs_experts_mid_prequant(mid, m, gate_w, up_w, xq, + selected, expert_weight, n_expert, clamp); + } else if (gate_w->type == DS4_TENSOR_Q2_K) { + matvec_q2_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, + selected, expert_weight, n_expert, clamp); + } else if (gate_w->type == DS4_TENSOR_Q4_K) { + matvec_q4_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, + selected, expert_weight, n_expert, clamp); + } else if (gate_w->type == DS4_TENSOR_Q5_K || gate_w->type == DS4_TENSOR_Q6_K) { + matvec_q5_q6_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, + selected, expert_weight, n_expert, clamp); + } else { + ds4_die("unsupported gate/up expert tensor type"); + } +} + +/* Dispatch: call the right down-projection accumulator based on tensor type. */ +static void matvec_experts_down_accum_prequant( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const block_q8_K *xq, + const int *selected, + int n_expert) { + if (w->type == DS4_TENSOR_IQ2_XXS) { + matvec_iq2_xxs_experts_accum_prequant(out, m, w, xq, selected, n_expert); + } else if (w->type == DS4_TENSOR_Q2_K) { + matvec_q2_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); + } else if (w->type == DS4_TENSOR_Q4_K) { + matvec_q4_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); + } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { + matvec_q5_q6_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); + } else { + ds4_die("unsupported down expert tensor type"); + } +} + +/* Dispatch: single-expert gate/up pair for tracing. */ +static void matvec_expert_pair_prequant( + float *out0, + float *out1, + const ds4_model *m, + const ds4_tensor *w0, + const ds4_tensor *w1, + const block_q8_K *xq, + uint32_t expert) { + if (w0->type == DS4_TENSOR_IQ2_XXS) { + matvec_iq2_xxs_expert_pair_prequant(out0, out1, m, w0, w1, xq, expert); + } else if (w0->type == DS4_TENSOR_Q2_K) { + uint64_t in_dim0, out_dim0, rb0; + uint64_t in_dim1, out_dim1, rb1; + const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); + const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); + if (w1->type != DS4_TENSOR_Q2_K || + in_dim0 != in_dim1 || + out_dim0 != out_dim1) { + ds4_die("paired Q2_K expert tensors do not match"); + } + + for (uint64_t row = 0; row < out_dim0; row++) { + const block_q2_K *gr = (const block_q2_K *)(base0 + row * rb0); + ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out0[row], gr, xq); + const block_q2_K *ur = (const block_q2_K *)(base1 + row * rb1); + ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out1[row], ur, xq); + } + } else if (w0->type == DS4_TENSOR_Q4_K) { + uint64_t in_dim0, out_dim0, rb0; + uint64_t in_dim1, out_dim1, rb1; + const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); + const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); + if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q4_K expert tensors do not match"); + + for (uint64_t row = 0; row < out_dim0; row++) { + const block_q4_K *gr = (const block_q4_K *)(base0 + row * rb0); + ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out0[row], gr, xq); + const block_q4_K *ur = (const block_q4_K *)(base1 + row * rb1); + ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out1[row], ur, xq); + } + } else if (w0->type == DS4_TENSOR_Q5_K || w0->type == DS4_TENSOR_Q6_K) { + uint64_t in_dim0, out_dim0, rb0; + uint64_t in_dim1, out_dim1, rb1; + const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); + const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); + if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q5_K/Q6_K expert tensors do not match"); + + for (uint64_t row = 0; row < out_dim0; row++) { + ds4_vec_dot_q5_q6_K_q8_K(w0->type, (int)in_dim0, &out0[row], base0 + row * rb0, xq); + ds4_vec_dot_q5_q6_K_q8_K(w1->type, (int)in_dim0, &out1[row], base1 + row * rb1, xq); + } + } else { + ds4_die("unsupported gate/up expert tensor type"); + } +} + +/* Dispatch: single-expert down projection for tracing. */ +static void matvec_expert_down( + float *out, + const ds4_model *m, + const ds4_tensor *w, + const float *x, + uint32_t expert) { + if (w->type == DS4_TENSOR_IQ2_XXS) { + uint64_t in_dim, out_dim, row_bytes; + const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); + if (in_dim % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); + + block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); + ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); + + for (uint64_t row = 0; row < out_dim; row++) { + const block_iq2_xxs *br = (const block_iq2_xxs *)(base + row * row_bytes); + ds4_vec_dot_iq2_xxs_q8_K((int)in_dim, &out[row], br, xq); + } + free(xq); + } else if (w->type == DS4_TENSOR_Q2_K) { + matvec_q2_k_expert(out, m, w, x, expert); + } else if (w->type == DS4_TENSOR_Q4_K) { + uint64_t in_dim, out_dim, row_bytes; + const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); + if (in_dim % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); + + block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); + ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); + + for (uint64_t row = 0; row < out_dim; row++) { + const block_q4_K *br = (const block_q4_K *)(base + row * row_bytes); + ds4_vec_dot_q4_K_q8_K((int)in_dim, &out[row], br, xq); + } + free(xq); + } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { + matvec_q5_q6_k_expert(out, m, w, x, expert); + } else { + ds4_die("unsupported down expert tensor type"); + } +} + +typedef struct { + float *moe; + const uint8_t *base[DS4_MAX_EXPERT]; + const block_q8_K *midq; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint32_t n_active; + uint32_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t midq_blocks; +} matvec_q8_k_batch_accum_rows_ctx; + +static void matvec_q8_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_k_batch_accum_rows_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + for (uint32_t t = 0; t < ctx->n_tok; t++) { + ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; + } + + for (uint32_t ai = 0; ai < ctx->n_active; ai++) { + const uint32_t expert = ctx->active_expert[ai]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const block_q8_K *br = (const block_q8_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; + float v = 0.0f; + + ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, xq); + ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; + } + } + } +} + +typedef struct { + float *moe; + const uint8_t *base[DS4_MAX_EXPERT]; + const int8_t *midq; + const float *midscale; + const ds4_expert_pair *pairs; + const uint32_t *pair_ids; + const uint32_t *expert_offset; + const uint32_t *active_expert; + uint32_t n_active; + uint32_t n_tok; + uint64_t in_dim; + uint64_t out_dim; + uint64_t row_bytes[DS4_MAX_EXPERT]; + uint64_t blocks; +} matvec_q8_0_batch_accum_rows_ctx; + +static void matvec_q8_0_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { + matvec_q8_0_batch_accum_rows_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + for (uint32_t t = 0; t < ctx->n_tok; t++) { + ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; + } + + for (uint32_t ai = 0; ai < ctx->n_active; ai++) { + const uint32_t expert = ctx->active_expert[ai]; + const uint32_t begin = ctx->expert_offset[expert]; + const uint32_t end = ctx->expert_offset[expert + 1]; + const uint8_t *br = ctx->base[expert] + row * ctx->row_bytes[expert]; + + for (uint32_t i = begin; i < end; i++) { + const uint32_t pair_id = ctx->pair_ids[i]; + const ds4_expert_pair pair = ctx->pairs[pair_id]; + const int8_t *xq = ctx->midq + (uint64_t)pair_id * ctx->blocks * 32u; + const float *xscale = ctx->midscale + (uint64_t)pair_id * ctx->blocks; + const float v = dot_q8_0_row(br, xq, xscale, ctx->in_dim, ctx->blocks); + ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; + } + } + } +} + +typedef struct { + float *moe; + const float *down_pair; + uint32_t n_tok; + uint64_t out_dim; +} sum_down_pairs_ctx; + +static DS4_MAYBE_UNUSED void sum_down_pairs_worker(void *vctx, uint64_t row0, uint64_t row1) { + sum_down_pairs_ctx *ctx = vctx; + for (uint64_t idx = row0; idx < row1; idx++) { + const uint32_t token = (uint32_t)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)token * ctx->out_dim; + float acc = 0.0f; + for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { + const uint64_t pair_id = (uint64_t)token * DS4_N_EXPERT_USED + slot; + acc += ctx->down_pair[pair_id * ctx->out_dim + row]; + } + ctx->moe[idx] = acc; + } +} diff --git a/kernels/cpu_quant.inc b/kernels/cpu_quant.inc new file mode 100644 index 0000000000..3875269646 --- /dev/null +++ b/kernels/cpu_quant.inc @@ -0,0 +1,1075 @@ +/* + * CPU scalar conversion and quantized tensor kernels. + * + * Included exactly once by ds4.c. Keeping this as an implementation fragment + * preserves private static linkage without creating a kernel abstraction. + */ + +/* ========================================================================= + * Scalar Conversion and Quantized Tensor Kernels. + * ========================================================================= + * + * These functions are the CPU reference math used by the C backend and by + * Metal diagnostics. They implement only the tensor formats present in the + * DeepSeek V4 Flash GGUF: F16, F32, Q8_0, Q2_K, IQ2_XXS, and Q8_K activation + * blocks used for expert dot products. + */ + +static inline float f16_to_f32(uint16_t h) { +#if defined(__ARM_NEON) + const float16x4_t hv = vreinterpret_f16_u16(vdup_n_u16(h)); + return vgetq_lane_f32(vcvt_f32_f16(hv), 0); +#else + uint32_t sign = (uint32_t)(h & 0x8000) << 16; + uint32_t exp = (h >> 10) & 0x1f; + uint32_t mant = h & 0x03ff; + uint32_t bits; + + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + exp = 1; + while ((mant & 0x0400) == 0) { + mant <<= 1; + exp--; + } + mant &= 0x03ff; + bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); + } + } else if (exp == 31) { + bits = sign | 0x7f800000u | (mant << 13); + } else { + bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); + } + + float f; + memcpy(&f, &bits, sizeof(f)); + return f; +#endif +} + +static inline uint16_t f32_to_f16(float f) { +#if defined(__ARM_NEON) + const float32x4_t fv = vdupq_n_f32(f); + const float16x4_t hv = vcvt_f16_f32(fv); + return vget_lane_u16(vreinterpret_u16_f16(hv), 0); +#else + uint32_t bits; + memcpy(&bits, &f, sizeof(bits)); + + const uint32_t sign = (bits >> 16) & 0x8000u; + int32_t exp = (int32_t)((bits >> 23) & 0xffu) - 127 + 15; + uint32_t mant = bits & 0x7fffffu; + + if (exp <= 0) { + if (exp < -10) return (uint16_t)sign; + mant |= 0x800000u; + const uint32_t shift = (uint32_t)(14 - exp); + uint32_t half_mant = mant >> shift; + const uint32_t round_bit = (mant >> (shift - 1)) & 1u; + const uint32_t sticky = mant & ((1u << (shift - 1)) - 1u); + if (round_bit && (sticky || (half_mant & 1u))) half_mant++; + return (uint16_t)(sign | half_mant); + } + + if (exp >= 31) { + if (((bits >> 23) & 0xffu) == 0xffu && mant != 0) { + return (uint16_t)(sign | 0x7e00u); + } + return (uint16_t)(sign | 0x7c00u); + } + + uint32_t half = sign | ((uint32_t)exp << 10) | (mant >> 13); + const uint32_t round = mant & 0x1fffu; + if (round > 0x1000u || (round == 0x1000u && (half & 1u))) half++; + return (uint16_t)half; +#endif +} + +static void f16_round_inplace_cpu(float *x, uint32_t n) { + for (uint32_t i = 0; i < n; i++) x[i] = f16_to_f32(f32_to_f16(x[i])); +} + +static float dsv4_e4m3fn_value_cpu(int i) { + static const float exp_scale[16] = { + 0.0f, 0.015625f, 0.03125f, 0.0625f, + 0.125f, 0.25f, 0.5f, 1.0f, + 2.0f, 4.0f, 8.0f, 16.0f, + 32.0f, 64.0f, 128.0f, 256.0f, + }; + + const int exp = (i >> 3) & 0x0f; + const int mant = i & 0x07; + return exp == 0 + ? (float)mant * 0.001953125f + : (1.0f + (float)mant * 0.125f) * exp_scale[exp]; +} + +static float dsv4_e4m3fn_dequant_cpu(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 448.0f); + + int lo = 0; + int hi = 126; + while (lo < hi) { + const int mid = (lo + hi + 1) >> 1; + if (dsv4_e4m3fn_value_cpu(mid) <= ax) { + lo = mid; + } else { + hi = mid - 1; + } + } + + int best = lo; + if (best < 126) { + const float best_diff = fabsf(ax - dsv4_e4m3fn_value_cpu(best)); + const float next_diff = fabsf(ax - dsv4_e4m3fn_value_cpu(best + 1)); + if (next_diff < best_diff || (next_diff == best_diff && ((best + 1) & 1) == 0 && (best & 1) != 0)) { + best++; + } + } + + return sign * dsv4_e4m3fn_value_cpu(best); +} + +/* DeepSeek V4 stores the non-RoPE part of compressed KV through an E4M3-style + * round trip. Keeping this in the CPU reference makes cache values comparable + * to the Metal graph's compressed-cache behavior. */ +static void dsv4_fp8_kv_quantize_row_inplace_cpu(float *x, uint32_t head_dim, uint32_t n_rot) { + const uint32_t n_nope = head_dim - n_rot; + for (uint32_t off = 0; off < n_nope; off += 64) { + float amax = 0.0f; + for (uint32_t i = 0; i < 64; i++) { + const float av = fabsf(x[off + i]); + if (av > amax) amax = av; + } + + if (amax < 1.0e-4f) amax = 1.0e-4f; + const float scale = ldexpf(1.0f, (int)ceilf(log2f(amax / 448.0f))); + for (uint32_t i = 0; i < 64; i++) { + float v = x[off + i] / scale; + if (v > 448.0f) v = 448.0f; + if (v < -448.0f) v = -448.0f; + x[off + i] = dsv4_e4m3fn_dequant_cpu(v) * scale; + } + } +} + +static float dsv4_e2m1fn_value_cpu(int i) { + static const float values[8] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + }; + return values[i & 7]; +} + +static float dsv4_e2m1fn_dequant_cpu(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 6.0f); + int best = 0; + float best_diff = fabsf(ax - dsv4_e2m1fn_value_cpu(0)); + for (int i = 1; i < 8; i++) { + const float diff = fabsf(ax - dsv4_e2m1fn_value_cpu(i)); + if (diff < best_diff || (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * dsv4_e2m1fn_value_cpu(best); +} + +static void dsv4_hadamard128_inplace_cpu(float *x) { + for (uint32_t stride = 1; stride < 128; stride <<= 1) { + for (uint32_t base = 0; base < 128; base += 2u * stride) { + for (uint32_t i = 0; i < stride; i++) { + const float a = x[base + i]; + const float b = x[base + stride + i]; + x[base + i] = a + b; + x[base + stride + i] = a - b; + } + } + } + const float scale = 0.08838834764831845f; + for (uint32_t i = 0; i < 128; i++) x[i] *= scale; +} + +static void dsv4_fp4_act_quantize_row_inplace_cpu(float *x, uint32_t n) { + if ((n % 32u) != 0) ds4_die("DSV4 FP4 activation quantization requires 32-aligned rows"); + for (uint32_t off = 0; off < n; off += 32) { + float amax = 0.0f; + for (uint32_t i = 0; i < 32; i++) { + const float av = fabsf(x[off + i]); + if (av > amax) amax = av; + } + + if (amax < 7.052966104933725e-38f) amax = 7.052966104933725e-38f; + const float scale = ldexpf(1.0f, (int)ceilf(log2f(amax / 6.0f))); + for (uint32_t i = 0; i < 32; i++) { + float v = x[off + i] / scale; + if (v > 6.0f) v = 6.0f; + if (v < -6.0f) v = -6.0f; + x[off + i] = dsv4_e2m1fn_dequant_cpu(v) * scale; + } + } +} + +/* The official DeepSeek V4 graph rotates indexer activations with a 128-wide + * Hadamard transform and immediately runs the FP4 activation-simulation + * round trip. This applies to both indexer Q and the indexer compressor KV; + * without it, the top-k compressed-row selection is not the model's graph. */ +static void dsv4_indexer_qat_row_inplace_cpu(float *x, uint32_t head_dim) { + if (head_dim != 128) ds4_die("DSV4 indexer QAT expects 128-wide indexer rows"); + dsv4_hadamard128_inplace_cpu(x); + dsv4_fp4_act_quantize_row_inplace_cpu(x, head_dim); +} + +static void dsv4_indexer_qat_rows_inplace_cpu(float *x, uint32_t rows, uint32_t head_dim) { + for (uint32_t r = 0; r < rows; r++) { + dsv4_indexer_qat_row_inplace_cpu(x + (uint64_t)r * head_dim, head_dim); + } +} + +/* Quantize a float activation into Q8_K blocks so GGUF Q2_K/IQ2_XXS expert + * kernels can reuse the same activation for many expert rows. */ +static void ds4_quantize_row_q8_K(const float *x, block_q8_K *y, int64_t k) { + if (k % QK_K != 0) ds4_die("Q8_K quantization length is not QK_K aligned"); + const int64_t nb = k / QK_K; + + for (int64_t b = 0; b < nb; b++) { + float max = 0.0f; + float amax = 0.0f; + for (int j = 0; j < QK_K; j++) { + const float ax = fabsf(x[j]); + if (ax > amax) { + amax = ax; + max = x[j]; + } + } + + if (amax == 0.0f) { + y[b].d = 0.0f; + memset(y[b].qs, 0, sizeof(y[b].qs)); + memset(y[b].bsums, 0, sizeof(y[b].bsums)); + x += QK_K; + continue; + } + + const float iscale = -127.0f / max; + for (int j = 0; j < QK_K; j++) { + int v = (int)lrintf(iscale * x[j]); + if (v > 127) v = 127; + if (v < -128) v = -128; + y[b].qs[j] = (int8_t)v; + } + for (int j = 0; j < QK_K / 16; j++) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += y[b].qs[j * 16 + i]; + y[b].bsums[j] = (int16_t)sum; + } + y[b].d = 1.0f / iscale; + x += QK_K; + } +} + +static void ds4_vec_dot_q2_K_q8_K(int n, float *s, const block_q2_K *x, const block_q8_K *y) { + const int nb = n / QK_K; + +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + const uint8x16_t m3 = vdupq_n_u8(0x03); + const uint8x16_t m4 = vdupq_n_u8(0x0f); + const int32x4_t zero = vdupq_n_s32(0); + float sum = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = y[i].d * f16_to_f32(x[i].d); + const float dmin = -y[i].d * f16_to_f32(x[i].dmin); + + const uint8_t *q2 = x[i].qs; + const int8_t *q8 = y[i].qs; + const uint8_t *sc = x[i].scales; + + const uint8x16_t mins_and_scales = vld1q_u8(sc); + const uint8x16_t scales = vandq_u8(mins_and_scales, m4); + uint8_t scale_lanes[16]; + vst1q_u8(scale_lanes, scales); + + const uint8x16_t mins = vshrq_n_u8(mins_and_scales, 4); + const int16x8x2_t q8sums = vld1q_s16_x2(y[i].bsums); + const int16x8x2_t mins16 = {{ + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(mins))), + vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(mins))), + }}; + const int32x4_t s0 = vaddq_s32( + vmull_s16(vget_low_s16(mins16.val[0]), vget_low_s16(q8sums.val[0])), + vmull_s16(vget_high_s16(mins16.val[0]), vget_high_s16(q8sums.val[0]))); + const int32x4_t s1 = vaddq_s32( + vmull_s16(vget_low_s16(mins16.val[1]), vget_low_s16(q8sums.val[1])), + vmull_s16(vget_high_s16(mins16.val[1]), vget_high_s16(q8sums.val[1]))); + sum += dmin * (float)vaddvq_s32(vaddq_s32(s0, s1)); + + int isum = 0; + int is = 0; + for (int j = 0; j < QK_K / 128; j++) { + const uint8x16x2_t q2bits = vld1q_u8_x2(q2); + q2 += 32; + +#define DS4_Q2_DOT_NOSHIFT(scale_index) do { \ + const int8x16x2_t q8bytes = vld1q_s8_x2(q8); \ + q8 += 32; \ + const int8x16_t q2lo = vreinterpretq_s8_u8(vandq_u8(q2bits.val[0], m3));\ + const int8x16_t q2hi = vreinterpretq_s8_u8(vandq_u8(q2bits.val[1], m3));\ + isum += vaddvq_s32(vdotq_s32(zero, q2lo, q8bytes.val[0])) * \ + scale_lanes[is + (scale_index)]; \ + isum += vaddvq_s32(vdotq_s32(zero, q2hi, q8bytes.val[1])) * \ + scale_lanes[is + 1 + (scale_index)]; \ + } while (0) + +#define DS4_Q2_DOT_SHIFT(shift, scale_index) do { \ + const int8x16x2_t q8bytes = vld1q_s8_x2(q8); \ + q8 += 32; \ + const int8x16_t q2lo = vreinterpretq_s8_u8( \ + vandq_u8(vshrq_n_u8(q2bits.val[0], (shift)), m3)); \ + const int8x16_t q2hi = vreinterpretq_s8_u8( \ + vandq_u8(vshrq_n_u8(q2bits.val[1], (shift)), m3)); \ + isum += vaddvq_s32(vdotq_s32(zero, q2lo, q8bytes.val[0])) * \ + scale_lanes[is + (scale_index)]; \ + isum += vaddvq_s32(vdotq_s32(zero, q2hi, q8bytes.val[1])) * \ + scale_lanes[is + 1 + (scale_index)]; \ + } while (0) + + DS4_Q2_DOT_NOSHIFT(0); + DS4_Q2_DOT_SHIFT(2, 2); + DS4_Q2_DOT_SHIFT(4, 4); + DS4_Q2_DOT_SHIFT(6, 6); + is += 8; + +#undef DS4_Q2_DOT_NOSHIFT +#undef DS4_Q2_DOT_SHIFT + } + + sum += d * (float)isum; + } + + *s = sum; +#else + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const uint8_t *q2 = x[i].qs; + const int8_t *q8 = y[i].qs; + const uint8_t *sc = x[i].scales; + + int summs = 0; + for (int j = 0; j < 16; j++) { + summs += y[i].bsums[j] * (sc[j] >> 4); + } + + const float dall = y[i].d * f16_to_f32(x[i].d); + const float dmin = y[i].d * f16_to_f32(x[i].dmin); + + int isum = 0; + int is = 0; + for (int k = 0; k < QK_K / 128; k++) { + int shift = 0; + for (int j = 0; j < 4; j++) { + int d = sc[is++] & 0x0f; + int isuml = dot_q2_16(q2, q8, shift); + isum += d * isuml; + + d = sc[is++] & 0x0f; + isuml = dot_q2_16(q2 + 16, q8 + 16, shift); + isum += d * isuml; + + shift += 2; + q8 += 32; + } + q2 += 32; + } + sumf += dall * (float)isum - dmin * (float)summs; + } + *s = sumf; +#endif +} + +static inline float q2_k_value_f32(const block_q2_K *blocks, uint32_t k) { + const uint32_t block = k / QK_K; + const uint32_t idx = k - block * QK_K; + const block_q2_K *xb = blocks + block; + const uint32_t group = idx / 16u; + const uint32_t l = idx - group * 16u; + const uint32_t q_base = 32u * (group / 8u) + 16u * (group & 1u); + const uint32_t shift = ((group / 2u) & 3u) * 2u; + const uint32_t q = ((uint32_t)xb->qs[q_base + l] >> shift) & 0x03u; + const uint32_t sc = xb->scales[group]; + return f16_to_f32(xb->d) * (float)(sc & 0x0fu) * (float)q - + f16_to_f32(xb->dmin) * (float)(sc >> 4u); +} + +static float ds4_vec_dot_q2_K_f32(int n, const block_q2_K *x, const float *y) { + float sum = 0.0f; + for (int k = 0; k < n; k++) { + sum += q2_k_value_f32(x, (uint32_t)k) * y[k]; + } + return sum; +} + +static inline void q4_k_get_scale_min(int j, const uint8_t *q, uint8_t *sc, uint8_t *m) { + if (j < 4) { + *sc = q[j] & 63; + *m = q[j + 4] & 63; + } else { + *sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + *m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +static void ds4_vec_dot_q4_K_q8_K(int n, float *s, const block_q4_K *x, const block_q8_K *y) { + const int nb = n / QK_K; + +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + const int32x4_t zero = vdupq_n_s32(0); + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = y[i].d * f16_to_f32(x[i].d); + const float dm = -y[i].d * f16_to_f32(x[i].dmin); + + const uint8_t *qs = x[i].qs; + const uint8_t *sc = x[i].scales; + const int8_t *q8 = y[i].qs; + + int32_t summs = 0; + for (int j = 0; j < QK_K / 32; j++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(j, sc, &sc_val, &m_val); + int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1]; + summs += m_val * gsum; + } + + int isum = 0; + for (int j = 0; j < QK_K / 32; j++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(j, sc, &sc_val, &m_val); + + const int byte_off = (j >> 1) * 32; + const int shift = (j & 1) * 4; + + /* Load 32 q8 values for this group */ + const int8x16x2_t q8v = vld1q_s8_x2(q8 + j * 32); + + /* Unpack 32 q4 values from 32 bytes at qs[byte_off] with shift */ + uint8_t q4_u[32]; + if (shift == 0) { + for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] & 0xF; + } else { + for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] >> 4; + } + + const int8x16_t q4a = vreinterpretq_s8_u8(vld1q_u8(q4_u)); + const int8x16_t q4b = vreinterpretq_s8_u8(vld1q_u8(q4_u + 16)); + + isum += vaddvq_s32(vdotq_s32(zero, q4a, q8v.val[0])) * sc_val; + isum += vaddvq_s32(vdotq_s32(zero, q4b, q8v.val[1])) * sc_val; + } + + sumf += d * (float)isum + dm * (float)summs; + } + + *s = sumf; +#else + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = y[i].d * f16_to_f32(x[i].d); + const float dm = -y[i].d * f16_to_f32(x[i].dmin); + + const uint8_t *qs = x[i].qs; + const uint8_t *sc = x[i].scales; + const int8_t *q8 = y[i].qs; + + int summs = 0; + for (int j = 0; j < QK_K / 32; j++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(j, sc, &sc_val, &m_val); + int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1]; + summs += m_val * gsum; + } + + int isum = 0; + for (int j = 0; j < QK_K / 32; j++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(j, sc, &sc_val, &m_val); + + const int byte_off = (j >> 1) * 32; + const int shift = (j & 1) * 4; + + for (int l = 0; l < 32; l++) { + isum += ((qs[byte_off + l] >> shift) & 0xF) * (int)q8[j * 32 + l] * sc_val; + } + } + + sumf += d * (float)isum + dm * (float)summs; + } + + *s = sumf; +#endif +} + +static void ds4_vec_dot_q5_K_q8_K(int n, float *s, const block_q5_K *x, const block_q8_K *y) { + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = y[i].d * f16_to_f32(x[i].d); + const float dmin = y[i].d * f16_to_f32(x[i].dmin); + const uint8_t *ql = x[i].qs; + const uint8_t *qh = x[i].qh; + const int8_t *q8 = y[i].qs; + const uint8_t *scales = x[i].scales; + + int64_t isum = 0; + int64_t summs = 0; + int is = 0; + uint8_t u1 = 1; + uint8_t u2 = 2; + + for (int j = 0; j < QK_K; j += 64) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(is, scales, &sc_val, &m_val); + summs += (int64_t)m_val * ((int32_t)y[i].bsums[2 * is] + (int32_t)y[i].bsums[2 * is + 1]); + for (int l = 0; l < 32; l++) { + const int q = (int)(ql[l] & 0x0F) + ((qh[l] & u1) ? 16 : 0); + isum += (int64_t)sc_val * q * (int)q8[j + l]; + } + + q4_k_get_scale_min(is + 1, scales, &sc_val, &m_val); + summs += (int64_t)m_val * ((int32_t)y[i].bsums[2 * (is + 1)] + (int32_t)y[i].bsums[2 * (is + 1) + 1]); + for (int l = 0; l < 32; l++) { + const int q = (int)(ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0); + isum += (int64_t)sc_val * q * (int)q8[j + 32 + l]; + } + + ql += 32; + is += 2; + u1 = (uint8_t)(u1 << 2); + u2 = (uint8_t)(u2 << 2); + } + + sumf += d * (float)isum - dmin * (float)summs; + } + + *s = sumf; +} + +static void ds4_vec_dot_q6_K_q8_K(int n, float *s, const block_q6_K *x, const block_q8_K *y) { + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = y[i].d * f16_to_f32(x[i].d); + const uint8_t *ql = x[i].ql; + const uint8_t *qh = x[i].qh; + const int8_t *scales = x[i].scales; + const int8_t *q8 = y[i].qs; + int64_t isum = 0; + + for (int n128 = 0; n128 < QK_K; n128 += 128) { + for (int l = 0; l < 32; l++) { + const int is = l / 16; + const int q1 = ((int)(ql[l + 0] & 0x0F) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int q2 = ((int)(ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int q3 = ((int)(ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int q4 = ((int)(ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + + isum += (int64_t)scales[is + 0] * q1 * (int)q8[n128 + l + 0]; + isum += (int64_t)scales[is + 2] * q2 * (int)q8[n128 + l + 32]; + isum += (int64_t)scales[is + 4] * q3 * (int)q8[n128 + l + 64]; + isum += (int64_t)scales[is + 6] * q4 * (int)q8[n128 + l + 96]; + } + + ql += 64; + qh += 32; + scales += 8; + } + + sumf += d * (float)isum; + } + + *s = sumf; +} + +static float ds4_vec_dot_q4_K_f32(int n, const block_q4_K *x, const float *y) { + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d); + const float dmin = f16_to_f32(x[i].dmin); + const uint8_t *qs = x[i].qs; + const uint8_t *scales = x[i].scales; + const float *yb = y + (uint64_t)i * QK_K; + + for (int j = 0; j < QK_K / 32; j++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(j, scales, &sc_val, &m_val); + + const int byte_off = (j >> 1) * 32; + const int shift = (j & 1) * 4; + const float scale = d * (float)sc_val; + const float minv = dmin * (float)m_val; + for (int l = 0; l < 32; l++) { + const int q = (qs[byte_off + l] >> shift) & 0x0F; + sumf += (scale * (float)q - minv) * yb[j * 32 + l]; + } + } + } + + return sumf; +} + +static float ds4_vec_dot_q5_K_f32(int n, const block_q5_K *x, const float *y) { + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d); + const float dmin = f16_to_f32(x[i].dmin); + const uint8_t *ql = x[i].qs; + const uint8_t *qh = x[i].qh; + const uint8_t *scales = x[i].scales; + const float *yb = y + (uint64_t)i * QK_K; + + for (int group = 0; group < QK_K / 32; group++) { + uint8_t sc_val, m_val; + q4_k_get_scale_min(group, scales, &sc_val, &m_val); + + const int ql_base = (group >> 1) * 32; + const int shift = (group & 1) * 4; + const uint8_t hmask = (uint8_t)(1u << group); + const float scale = d * (float)sc_val; + const float minv = dmin * (float)m_val; + for (int l = 0; l < 32; l++) { + const int q = ((ql[ql_base + l] >> shift) & 0x0F) + + ((qh[l] & hmask) ? 16 : 0); + sumf += (scale * (float)q - minv) * yb[group * 32 + l]; + } + } + } + + return sumf; +} + +static float ds4_vec_dot_q6_K_f32(int n, const block_q6_K *x, const float *y) { + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d); + const uint8_t *ql = x[i].ql; + const uint8_t *qh = x[i].qh; + const int8_t *scales = x[i].scales; + const float *yb = y + (uint64_t)i * QK_K; + + for (int n128 = 0; n128 < QK_K; n128 += 128) { + for (int l = 0; l < 32; l++) { + const int is = l / 16; + const int q1 = ((int)(ql[l + 0] & 0x0F) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int q2 = ((int)(ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int q3 = ((int)(ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int q4 = ((int)(ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + + sumf += d * (float)scales[is + 0] * (float)q1 * yb[n128 + l + 0]; + sumf += d * (float)scales[is + 2] * (float)q2 * yb[n128 + l + 32]; + sumf += d * (float)scales[is + 4] * (float)q3 * yb[n128 + l + 64]; + sumf += d * (float)scales[is + 6] * (float)q4 * yb[n128 + l + 96]; + } + + ql += 64; + qh += 32; + scales += 8; + } + } + + return sumf; +} + +static inline float ds4_vec_dot_q5_q6_K_f32(uint32_t type, int n, const uint8_t *x, const float *y) { + if (type == DS4_TENSOR_Q5_K) { + return ds4_vec_dot_q5_K_f32(n, (const block_q5_K *)x, y); + } else if (type == DS4_TENSOR_Q6_K) { + return ds4_vec_dot_q6_K_f32(n, (const block_q6_K *)x, y); + } else { + ds4_die("expected a Q5_K or Q6_K tensor"); + } + return 0.0f; +} + +static float ds4_vec_dot_iq2_xxs_f32(int n, const block_iq2_xxs *x, const float *y) { + pthread_once(&iq2xxs_signed_grid_once, iq2xxs_signed_grid_init); + + const int nb = n / QK_K; + float sumf = 0.0f; + uint32_t aux32[2]; + const uint8_t *aux8 = (const uint8_t *)aux32; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d); + const uint16_t *q2 = x[i].qs; + const float *yb = y + (uint64_t)i * QK_K; + + for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { + memcpy(aux32, q2, 2 * sizeof(uint32_t)); + q2 += 4; + + const float scale = 0.125f * d * (float)(2u * (aux32[1] >> 28) + 1u); + const uint32_t base = (uint32_t)ib32 * 32u; + for (int l = 0; l < 4; l++) { + const uint32_t sign_idx = (aux32[1] >> (7 * l)) & 127u; + const int8_t *grid = iq2xxs_signed_grid[aux8[l]][sign_idx]; + const float *yf = yb + base + (uint32_t)l * 8u; + for (int j = 0; j < 8; j++) { + sumf += scale * (float)grid[j] * yf[j]; + } + } + } + } + + return sumf; +} + +static void ds4_vec_dot_q8_K_q8_K(int n, float *s, + const block_q8_K *x, + const block_q8_K *y) { + const int nb = n / QK_K; + float sum = 0.0f; +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + for (int i = 0; i < nb; i++) { + int32x4_t isum = vdupq_n_s32(0); + for (int j = 0; j < QK_K; j += 16) { + isum = vdotq_s32(isum, vld1q_s8(x[i].qs + j), + vld1q_s8(y[i].qs + j)); + } + sum += x[i].d * y[i].d * (float)vaddvq_s32(isum); + } +#else + for (int i = 0; i < nb; i++) { + int isum = 0; + for (int j = 0; j < QK_K; j++) { + isum += (int)x[i].qs[j] * (int)y[i].qs[j]; + } + sum += x[i].d * y[i].d * (float)isum; + } +#endif + *s = sum; +} + +static void ds4_vec_dot_q8_K_pair_q8_K( + int n, float *s0, float *s1, + const block_q8_K *x0, const block_q8_K *x1, + const block_q8_K *y) { + const int nb = n / QK_K; + float sum0 = 0.0f; + float sum1 = 0.0f; +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + for (int i = 0; i < nb; i++) { + int32x4_t isum0 = vdupq_n_s32(0); + int32x4_t isum1 = vdupq_n_s32(0); + for (int j = 0; j < QK_K; j += 16) { + const int8x16_t yv = vld1q_s8(y[i].qs + j); + isum0 = vdotq_s32(isum0, vld1q_s8(x0[i].qs + j), yv); + isum1 = vdotq_s32(isum1, vld1q_s8(x1[i].qs + j), yv); + } + sum0 += x0[i].d * y[i].d * (float)vaddvq_s32(isum0); + sum1 += x1[i].d * y[i].d * (float)vaddvq_s32(isum1); + } +#else + for (int i = 0; i < nb; i++) { + int isum0 = 0; + int isum1 = 0; + for (int j = 0; j < QK_K; j++) { + const int yv = (int)y[i].qs[j]; + isum0 += (int)x0[i].qs[j] * yv; + isum1 += (int)x1[i].qs[j] * yv; + } + sum0 += x0[i].d * y[i].d * (float)isum0; + sum1 += x1[i].d * y[i].d * (float)isum1; + } +#endif + *s0 = sum0; + *s1 = sum1; +} + +static DS4_MAYBE_UNUSED void ds4_vec_dot_iq2_xxs_q8_K(int n, float *s, const block_iq2_xxs *x, const block_q8_K *y) { + const int nb = n / QK_K; + +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d) * y[i].d; + const uint16_t *q2 = x[i].qs; + const int8_t *q8 = y[i].qs; + float sumf1 = 0.0f; + float sumf2 = 0.0f; + + for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { + int8x16x4_t q8b = vld1q_s8_x4(q8); + q8 += 64; + + uint32_t aux32[4]; + memcpy(aux32, q2, sizeof(aux32)); + q2 += 8; + const uint8_t *aux8 = (const uint8_t *)aux32; + + int8x16_t q2u0 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[0])), + vld1_s8((const int8_t *)(iq2xxs_grid + aux8[1]))); + int8x16_t q2u1 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[2])), + vld1_s8((const int8_t *)(iq2xxs_grid + aux8[3]))); + int8x16_t q2u2 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[8])), + vld1_s8((const int8_t *)(iq2xxs_grid + aux8[9]))); + int8x16_t q2u3 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + aux8[10])), + vld1_s8((const int8_t *)(iq2xxs_grid + aux8[11]))); + + const int8x16_t q2s0 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[1] >> 0) & 127]), + vld1_s8(iq2xxs_signs[(aux32[1] >> 7) & 127])); + const int8x16_t q2s1 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[1] >> 14) & 127]), + vld1_s8(iq2xxs_signs[(aux32[1] >> 21) & 127])); + const int8x16_t q2s2 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[3] >> 0) & 127]), + vld1_s8(iq2xxs_signs[(aux32[3] >> 7) & 127])); + const int8x16_t q2s3 = vcombine_s8(vld1_s8(iq2xxs_signs[(aux32[3] >> 14) & 127]), + vld1_s8(iq2xxs_signs[(aux32[3] >> 21) & 127])); + + q2u0 = vmulq_s8(q2u0, q2s0); + q2u1 = vmulq_s8(q2u1, q2s1); + q2u2 = vmulq_s8(q2u2, q2s2); + q2u3 = vmulq_s8(q2u3, q2s3); + + const int32x4_t p1 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), q2u0, q8b.val[0]), q2u1, q8b.val[1]); + const int32x4_t p2 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), q2u2, q8b.val[2]), q2u3, q8b.val[3]); + + sumf1 += (float)vaddvq_s32(p1) * (0.5f + (float)(aux32[1] >> 28)); + sumf2 += (float)vaddvq_s32(p2) * (0.5f + (float)(aux32[3] >> 28)); + } + + sumf += d * (sumf1 + sumf2); + } + + *s = 0.25f * sumf; +#else + uint32_t aux32[2]; + const uint8_t *aux8 = (const uint8_t *)aux32; + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d = f16_to_f32(x[i].d) * y[i].d; + const uint16_t *q2 = x[i].qs; + const int8_t *q8 = y[i].qs; + int32_t bsum = 0; + + for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { + memcpy(aux32, q2, 2 * sizeof(uint32_t)); + q2 += 4; + + const uint32_t ls = 2 * (aux32[1] >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; l += 2) { + const uint32_t sign_idx0 = (aux32[1] >> (7 * l)) & 127; + const uint32_t sign_idx1 = (aux32[1] >> (7 * (l + 1))) & 127; + sumi += dot_iq2_pair_16(iq2xxs_signed_grid[aux8[l]][sign_idx0], + iq2xxs_signed_grid[aux8[l + 1]][sign_idx1], + q8); + q8 += 16; + } + bsum += sumi * (int32_t)ls; + } + sumf += d * (float)bsum; + } + *s = 0.125f * sumf; +#endif +} + +static void ds4_vec_dot_iq2_xxs_pair_q8_K( + int n, + float *s0, + float *s1, + const block_iq2_xxs *x0, + const block_iq2_xxs *x1, + const block_q8_K *y) { +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + const int nb = n / QK_K; + float total0 = 0.0f; + float total1 = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d0 = f16_to_f32(x0[i].d) * y[i].d; + const float d1 = f16_to_f32(x1[i].d) * y[i].d; + const uint16_t *q20 = x0[i].qs; + const uint16_t *q21 = x1[i].qs; + const int8_t *q8 = y[i].qs; + float sum01 = 0.0f; + float sum02 = 0.0f; + float sum11 = 0.0f; + float sum12 = 0.0f; + + for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { + const int8x16x4_t q8b = vld1q_s8_x4(q8); + q8 += 64; + + uint32_t aux0[4]; + uint32_t aux1[4]; + memcpy(aux0, q20, sizeof(aux0)); + memcpy(aux1, q21, sizeof(aux1)); + q20 += 8; + q21 += 8; + const uint8_t *a0 = (const uint8_t *)aux0; + const uint8_t *a1 = (const uint8_t *)aux1; + +#define DS4_IQ2_PAIR_DOT(aux, aux8, accum_a, accum_b) do { \ + int8x16_t u0 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[0])), \ + vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[1]))); \ + int8x16_t u1 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[2])), \ + vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[3]))); \ + int8x16_t u2 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[8])), \ + vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[9]))); \ + int8x16_t u3 = vcombine_s8(vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[10])), \ + vld1_s8((const int8_t *)(iq2xxs_grid + (aux8)[11]))); \ + const int8x16_t sgn0 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[1] >> 0) & 127]), \ + vld1_s8(iq2xxs_signs[((aux)[1] >> 7) & 127])); \ + const int8x16_t sgn1 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[1] >> 14) & 127]), \ + vld1_s8(iq2xxs_signs[((aux)[1] >> 21) & 127])); \ + const int8x16_t sgn2 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[3] >> 0) & 127]), \ + vld1_s8(iq2xxs_signs[((aux)[3] >> 7) & 127])); \ + const int8x16_t sgn3 = vcombine_s8(vld1_s8(iq2xxs_signs[((aux)[3] >> 14) & 127]), \ + vld1_s8(iq2xxs_signs[((aux)[3] >> 21) & 127])); \ + u0 = vmulq_s8(u0, sgn0); \ + u1 = vmulq_s8(u1, sgn1); \ + u2 = vmulq_s8(u2, sgn2); \ + u3 = vmulq_s8(u3, sgn3); \ + const int32x4_t p1 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), u0, q8b.val[0]), u1, q8b.val[1]); \ + const int32x4_t p2 = vdotq_s32(vdotq_s32(vdupq_n_s32(0), u2, q8b.val[2]), u3, q8b.val[3]); \ + (accum_a) += (float)vaddvq_s32(p1) * (0.5f + (float)((aux)[1] >> 28)); \ + (accum_b) += (float)vaddvq_s32(p2) * (0.5f + (float)((aux)[3] >> 28)); \ + } while (0) + + DS4_IQ2_PAIR_DOT(aux0, a0, sum01, sum02); + DS4_IQ2_PAIR_DOT(aux1, a1, sum11, sum12); + +#undef DS4_IQ2_PAIR_DOT + } + + total0 += d0 * (sum01 + sum02); + total1 += d1 * (sum11 + sum12); + } + + *s0 = 0.25f * total0; + *s1 = 0.25f * total1; +#else + ds4_vec_dot_iq2_xxs_q8_K(n, s0, x0, y); + ds4_vec_dot_iq2_xxs_q8_K(n, s1, x1, y); +#endif +} + +typedef struct { + ds4_tensor *hc_attn_fn; + ds4_tensor *hc_attn_scale; + ds4_tensor *hc_attn_base; + ds4_tensor *attn_norm; + ds4_tensor *attn_q_a; + ds4_tensor *attn_q_a_norm; + ds4_tensor *attn_q_b; + ds4_tensor *attn_kv; + ds4_tensor *attn_kv_a_mqa; + ds4_tensor *attn_kv_a_norm; + ds4_tensor *attn_k_b; + ds4_tensor *attn_v_b; + ds4_tensor *attn_sinks; + ds4_tensor *attn_output; + ds4_tensor *attn_output_a; + ds4_tensor *attn_output_b; + ds4_tensor *attn_compressor_ape; + ds4_tensor *attn_compressor_kv; + ds4_tensor *attn_compressor_gate; + ds4_tensor *attn_compressor_norm; + ds4_tensor *indexer_attn_q_b; + ds4_tensor *indexer_attn_k; + ds4_tensor *indexer_k_norm; + ds4_tensor *indexer_k_norm_b; + ds4_tensor *indexer_proj; + ds4_tensor *indexer_compressor_ape; + ds4_tensor *indexer_compressor_kv; + ds4_tensor *indexer_compressor_gate; + ds4_tensor *indexer_compressor_norm; + ds4_tensor *hc_ffn_fn; + ds4_tensor *hc_ffn_scale; + ds4_tensor *hc_ffn_base; + ds4_tensor *ffn_norm; + ds4_tensor *ffn_gate_tid2eid; + ds4_tensor *ffn_gate; + ds4_tensor *ffn_up; + ds4_tensor *ffn_down; + ds4_tensor *ffn_gate_inp; + ds4_tensor *ffn_exp_probs_b; + ds4_tensor *ffn_gate_exps; + ds4_tensor *ffn_up_exps; + ds4_tensor *ffn_down_exps; + ds4_tensor *ffn_gate_shexp; + ds4_tensor *ffn_up_shexp; + ds4_tensor *ffn_down_shexp; + ds4_tensor *nextn_eh_proj; + ds4_tensor *nextn_enorm; + ds4_tensor *nextn_hnorm; + ds4_tensor *nextn_shared_head_norm; +} ds4_layer_weights; + +typedef struct { + ds4_tensor *token_embd; + ds4_tensor *output_hc_base; + ds4_tensor *output_hc_fn; + ds4_tensor *output_hc_scale; + ds4_tensor *output_norm; + ds4_tensor *output; + ds4_layer_weights layer[DS4_MAX_LAYER]; +} ds4_weights; + +typedef struct { + ds4_tensor *e_proj; + ds4_tensor *h_proj; + ds4_tensor *enorm; + ds4_tensor *hnorm; + ds4_tensor *norm; + ds4_tensor *hc_head_base; + ds4_tensor *hc_head_fn; + ds4_tensor *hc_head_scale; + ds4_layer_weights block; +} ds4_mtp_weights; + +typedef struct { + ds4_tensor *main_proj; + ds4_tensor *main_norm; + ds4_tensor *norm; + ds4_tensor *hc_head_base; + ds4_tensor *hc_head_fn; + ds4_tensor *hc_head_scale; + ds4_tensor *markov_w1; + ds4_tensor *markov_w2; + ds4_tensor *confidence_proj; + ds4_layer_weights block; +} ds4_dspark_stage_weights; + +typedef struct { + uint32_t n_stages; + uint32_t block_size; + uint32_t markov_rank; + uint32_t noise_token_id; + uint32_t target_layer_count; + uint32_t target_layers[DS4_DSPARK_MAX_TARGET_LAYERS]; + uint32_t present_tensors; + uint32_t missing_tensors; + uint32_t invalid_tensors; + uint32_t metadata_errors; + bool has_block_size; + bool has_markov_rank; + bool has_noise_token_id; + bool has_target_layers; + ds4_dspark_stage_weights stage[DS4_DSPARK_MAX_STAGES]; +} ds4_dspark_weights; diff --git a/metal/README.md b/metal/README.md new file mode 100644 index 0000000000..598b4d7a99 --- /dev/null +++ b/metal/README.md @@ -0,0 +1,22 @@ +# Metal implementation units + +This directory owns Metal runtime code, shared host launch paths, and reusable +device primitives: + +- `runtime.inc`, `model_io.inc`, and `expert_streaming.inc`: backend lifetime, + command submission, mapped weights, and streamed expert residency. +- `embedding.inc`, `dense_norm.inc`, and `elementwise.inc`: shared concrete + kernel launch paths. +- `moe_dispatch.inc`: shared low-level MoE encoding helpers. +- `compat.inc`: compatibility entry points required by the common GPU API. +- `*.metal`: shared device primitives. `model_abi.metal` defines the structs + used while concatenating the model shader sources into the library. + +Model-specific host launch paths and shaders live under +`models//metal/host/` and `models//metal/shaders/`. The runtime +concatenates the shared and selected model shader sources into one library; +`ds4_metal.m` likewise includes all host fragments into one Objective-C +translation unit. + +The directory split is ownership only. Model implementations may duplicate +code when that keeps their kernel paths direct and tunable. diff --git a/metal/compat.inc b/metal/compat.inc new file mode 100644 index 0000000000..01d602550f --- /dev/null +++ b/metal/compat.inc @@ -0,0 +1,3 @@ +void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { + (void)enabled; +} diff --git a/metal/dense_norm.inc b/metal/dense_norm.inc new file mode 100644 index 0000000000..5032717eef --- /dev/null +++ b/metal/dense_norm.inc @@ -0,0 +1,2664 @@ +int ds4_gpu_matmul_q8_0_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_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + const int profile_requested = + n_tok > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0; + int profile_prefill = 0; + int split_batch_for_profile = 0; + const char *profile_label = NULL; + char profile_label_buf[128]; + char profile_fallback[128]; + if (profile_requested) { + snprintf(profile_fallback, sizeof(profile_fallback), + "q8 weight_off=%llu in=%llu out=%llu tok=%llu", + (unsigned long long)weight_offset, + (unsigned long long)in_dim, + (unsigned long long)out_dim, + (unsigned long long)n_tok); + snprintf(profile_label_buf, sizeof(profile_label_buf), "%s", profile_fallback); + profile_label = profile_label_buf; + const char *profile_filter = getenv("DS4_METAL_Q8_PREFILL_PROFILE_FILTER"); + profile_prefill = + profile_requested && + (!profile_filter || !profile_filter[0] || + strstr(profile_label, profile_filter) != NULL); + } + if (profile_prefill) { + if (g_batch_cb) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + split_batch_for_profile = 1; + } + } + + const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0; + int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, false, false); + if (profile_prefill) { + if (split_batch_for_profile && ds4_gpu_end_commands() == 0) { + ok = 0; + } + const double elapsed_ms = ds4_gpu_now_ms() - profile_t0; + fprintf(stderr, + "ds4: Metal Q8_0 prefill profile %s in=%llu out=%llu tok=%llu %.3f ms\n", + profile_label ? profile_label : profile_fallback, + (unsigned long long)in_dim, + (unsigned long long)out_dim, + (unsigned long long)n_tok, + elapsed_ms); + if (split_batch_for_profile && ds4_gpu_begin_commands() == 0) { + ok = 0; + } + } + return ok; +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || n_rows == 0 || + n_rows > INT32_MAX || in_dim == 0 || out_dim == 0 || + (in_dim & 31u) != 0 || in_dim > UINT32_MAX || + out_dim > UINT32_MAX || + in_dim > UINT64_MAX / n_rows / sizeof(float) || + out_dim > UINT64_MAX / n_rows / sizeof(float) || + ds4_gpu_tensor_bytes(x) < + (uint64_t)n_rows * in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out) < + (uint64_t)n_rows * out_dim * sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t blocks = in_dim / 32u; + const uint64_t row_bytes = blocks * 34u; + if (!xbuf || !outbuf || + out_dim > UINT64_MAX / row_bytes) { + return 0; + } + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_q8_decode_model_range( + model_map, model_size, weight_offset, weight_bytes, 1u, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) dispatch.nsg = 8; + ds4_gpu_q8_0_matvec_args args = + ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + args.ne11 = (int32_t)n_rows; + args.nb12 = (uint64_t)n_rows * in_dim * sizeof(float); + args.nb13 = args.nb12; + args.ne1 = (int32_t)n_rows; + args.nr0 = dispatch.nr0; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline(dispatch.function_name, dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:dispatch.smem atIndex:0]; + [enc dispatchThreadgroups: + MTLSizeMake(((NSUInteger)out_dim + + (NSUInteger)dispatch.nr0 - 1u) / + (NSUInteger)dispatch.nr0, + (NSUInteger)n_rows, + 1) + threadsPerThreadgroup: + MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return ds4_gpu_finish_command_buffer( + cb, owned, "Q8_0 exact decode-row matvec"); + } +} + +int ds4_gpu_matmul_q8_0_decode_mpp_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_tok) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, true, false); +} + +int ds4_gpu_matmul_q8_0_decode_mpp_model_view_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_tok) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, true, true); +} + +static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) { + const char *prefix = NULL; + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + prefix = "kernel_mul_mv_ext_q4_K_f32_r1_"; + } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { + prefix = "kernel_mul_mv_ext_q4_0_f32_r1_"; + } else { + return NULL; + } + switch (r1ptg) { + case 1: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_1" : "kernel_mul_mv_ext_q4_0_f32_r1_1"; + case 2: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_2" : "kernel_mul_mv_ext_q4_0_f32_r1_2"; + case 3: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_3" : "kernel_mul_mv_ext_q4_0_f32_r1_3"; + case 4: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_4" : "kernel_mul_mv_ext_q4_0_f32_r1_4"; + case 5: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_5" : "kernel_mul_mv_ext_q4_0_f32_r1_5"; + default: + (void)prefix; + return NULL; + } +} + +static const char *ds4_gpu_q4_mm_name(uint32_t weight_type) { + switch (weight_type) { + case DS4_METAL_TENSOR_Q4_0: return "kernel_mul_mm_q4_0_f32"; + case DS4_METAL_TENSOR_Q4_K: return "kernel_mul_mm_q4_K_f32"; + default: return NULL; + } +} + +static const char *ds4_gpu_q4_nax_name(uint32_t weight_type, uint64_t tile_n) { + if (weight_type == DS4_METAL_TENSOR_Q4_0) { + return tile_n == 128u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128" : + tile_n == 64u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64" : + "kernel_mul_mm_q4_0_f32_nax_direct_rhs"; + } + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + return tile_n == 128u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128" : + tile_n == 64u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64" : + "kernel_mul_mm_q4_K_f32_nax_direct_rhs"; + } + return NULL; +} + +static int ds4_gpu_matmul_quant_impl_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, + bool prefer_decode_mpp, + bool force_model_view) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, + model_map, + model_size, + weight_offset, + in_dim, + out_dim, + x, + n_tok, + prefer_decode_mpp, + force_model_view); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + in_dim == 0 || out_dim == 0 || n_tok == 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, (uint32_t)in_dim, &row_bytes)) { + fprintf(stderr, "ds4: Metal quant matmul received unsupported type/dim (%u, in=%llu)\n", + weight_type, + (unsigned long long)in_dim); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal quant tensor matmul received undersized activation buffers\n"); + return 0; + } + + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = force_model_view ? + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + /* + * Small-batch Q4_K goes to the classic (llama.cpp-style) matvec: + * the mul_mv_ext family tops out around 220 GB/s on M5 for the GLM + * DenseQ4 decode shapes while this impl streams 530-650 GB/s + * (misc/q4mv_bench.m). Falls through to ext when unavailable. + */ + if (weight_type == DS4_METAL_TENSOR_Q4_K && + n_tok <= 8u && + (in_dim % 256u) == 0 && + getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { + const int16_t nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); + if (pipeline) { + ds4_gpu_q8_0_matvec_args args = { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_tok, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; + const uint64_t rows_ptg = (uint64_t)nsg * 2u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:32 atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K classic mul_mv")) return 0; + return 1; + } + } + + if (n_tok <= 8u && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = (n_tok == 1u) ? 1 : ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_q4_mv_ext_name(weight_type, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (pipeline) { + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, row_bytes, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor mul_mv_ext")) return 0; + return 1; + } + } + + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 64u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = ds4_gpu_q4_nax_name(weight_type, nax_tile_n); + id pipeline = + nax_fn ? ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false) : nil; + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 NAX tensor matmul")) return 0; + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const char *mm_fn = ds4_gpu_q4_mm_name(weight_type); + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + mm_fn ? ds4_gpu_get_mul_mm_pipeline(mm_fn, bc_inp, bc_out) : nil; + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor matmul")) return 0; + } + + return 1; +} + +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) { + return ds4_gpu_matmul_quant_impl_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok, + false, + false); +} + +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) { + /* GLM decode dense matvecs: the classic Q8_0 kernels stream 570-613 + * GB/s on the GLM shapes while the MPP/nax matrix path measures ~150 + * GB/s at n_tok=1 (TP decode 8.75 -> 16.64 t/s on the IQ2+Q8 gguf). + * MPP stays available as an opt-in via DS4_METAL_Q8_DECODE_MPP. */ + const bool prefer_mpp = getenv("DS4_METAL_Q8_DECODE_MPP") != NULL; + return ds4_gpu_matmul_quant_impl_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok, + prefer_mpp, + true); +} + +int ds4_gpu_matmul_quant_rows_scalar_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) { + return ds4_gpu_matmul_quant_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok); +} + +int ds4_gpu_matmul_q8_0_rows_scalar_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_tok) { + if (n_tok == 1) { + return ds4_gpu_matmul_q8_0_tensor(out, + model_map, + model_size, + weight_offset, + in_dim, + out_dim, + x, + 1); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + n_tok == 0 || n_tok > INT32_MAX || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || + in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t x_row_bytes = in_dim * sizeof(float); + const uint64_t out_row_bytes = out_dim * sizeof(float); + if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || + (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * x_row_bytes; + const uint64_t out_bytes = n_tok * out_row_bytes; + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + if (blocks > UINT64_MAX / 34u) return 0; + const uint64_t row_bytes = blocks * 34u; + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + + const NSUInteger x_base = ds4_gpu_tensor_offset(x); + const NSUInteger out_base = ds4_gpu_tensor_offset(out); + for (uint64_t t = 0; t < n_tok; t++) { + [enc setBuffer:xbuf offset:x_base + (NSUInteger)(t * x_row_bytes) atIndex:2]; + [enc setBuffer:outbuf offset:out_base + (NSUInteger)(t * out_row_bytes) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 scalar-row matvecs")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_pair_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, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out0 || !out1 || !model_map || !x || n_tok != 1 || + out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id out0buf = ds4_gpu_tensor_buffer(out0); + id out1buf = ds4_gpu_tensor_buffer(out1); + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out0_bytes = out0_dim * sizeof(float); + const uint64_t out1_bytes = out1_dim * sizeof(float); + if (!xbuf || !out0buf || !out1buf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out0) < out0_bytes || + ds4_gpu_tensor_bytes(out1) < out1_bytes) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight0_bytes = out0_dim * row_bytes; + const uint64_t weight1_bytes = out1_dim * row_bytes; + if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset || + weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n"); + return 0; + } + + uint64_t inner0 = 0; + uint64_t inner1 = 0; + id weight0buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight0_offset, weight0_bytes, &inner0); + id weight1buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight1_offset, weight1_bytes, &inner1); + if (!weight0buf || !weight1buf) return 0; + + ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch(); + ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch(); + if (out0_dim > 65536u) dispatch0.nsg = 8; + if (out1_dim > 65536u) dispatch1.nsg = 8; + /* A common threadgroup shape is required to retain each standalone + * reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */ + if (dispatch0.nsg != dispatch1.nsg) return 0; + + ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim); + ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim); + args0.nr0 = dispatch0.nr0; + args1.nr0 = dispatch1.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args0 length:sizeof(args0) atIndex:0]; + [enc setBytes:&args1 length:sizeof(args1) atIndex:1]; + [enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2]; + [enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5]; + [enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6]; + [enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0]; + const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim + + (NSUInteger)dispatch0.nr0 - 1u) / + (NSUInteger)dispatch0.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_f16_out_tensor( + ds4_gpu_tensor *out_h, + 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_tok) { + (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; + (void)in_dim; (void)out_dim; (void)x; (void)n_tok; + return 0; +} + +static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + int store_gate_up, + bool force_model_view) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!mid || !x || !model_map || + (store_gate_up && (!gate || !up)) || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id gatebuf = store_gate_up ? + ds4_gpu_tensor_buffer(gate) : midbuf; + id upbuf = store_gate_up ? + ds4_gpu_tensor_buffer(up) : midbuf; + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out_bytes = out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + (store_gate_up && ds4_gpu_tensor_bytes(gate) < out_bytes) || + (store_gate_up && ds4_gpu_tensor_bytes(up) < out_bytes) || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal shared expert fused gate/up received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal shared expert fused gate/up range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_views = + !force_model_view && + getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = exact_decode_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset, + weight_bytes, + &gate_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + weight_bytes, + &gate_inner); + id up_wbuf = exact_decode_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset, + weight_bytes, + &up_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + weight_bytes, + &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + const char *fn_name = store_gate_up ? + (mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : + "kernel_dsv4_shared_gate_up_swiglu_q8_0") : + (mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_mid_swiglu_q8_0_r4" : + "kernel_dsv4_shared_mid_swiglu_q8_0"); + id pipeline = + ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:(store_gate_up ? + ds4_gpu_tensor_offset(gate) : + ds4_gpu_tensor_offset(mid)) atIndex:4]; + [enc setBuffer:upbuf offset:(store_gate_up ? + ds4_gpu_tensor_offset(up) : + ds4_gpu_tensor_offset(mid)) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, + owned, + store_gate_up ? + "shared expert fused gate/up" : + "shared expert fused mid")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 1, + false); +} + +int ds4_gpu_shared_mid_swiglu_q8_0_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL, + NULL, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 0, + true); +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 1, + true); +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + if (n_tok == 1) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!gate || !up || !mid || !x || !model_map || + n_tok == 0 || + (in_dim & 31u) != 0 || (in_dim % 128u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + const uint64_t mv_ext_max_tokens = + ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); + if (n_tok > mv_ext_max_tokens) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < out_bytes || + ds4_gpu_tensor_bytes(up) < out_bytes || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows range is outside the mapped model\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_q8_pair_swiglu_name(r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / + (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / + (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused Q8_0 gate/up rows")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + if (n_tok == 1) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!gate || !up || !mid || !x || !model_map || + n_tok == 0 || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || + in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float) || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + const uint64_t x_row_bytes = in_dim * sizeof(float); + const uint64_t out_row_bytes = out_dim * sizeof(float); + if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || + (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + const uint64_t x_bytes = n_tok * x_row_bytes; + const uint64_t out_bytes = n_tok * out_row_bytes; + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < out_bytes || + ds4_gpu_tensor_bytes(up) < out_bytes || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + if (blocks > UINT64_MAX / 34u) return 0; + const uint64_t row_bytes = blocks * 34u; + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up range is outside the mapped model\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + args.ne11 = (int32_t)n_tok; + args.nb12 = n_tok * x_row_bytes; + args.nb13 = args.nb12; + args.ne1 = (int32_t)n_tok; + const char *fn_name = mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : + "kernel_dsv4_shared_gate_up_swiglu_q8_0"; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup: + MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "shared expert scalar-row fused gate/up")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_matmul_f16_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_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal F16 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal F16 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_f32_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = + ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matvec")) return 0; + return 1; + } + + if (n_tok <= 8 && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_name(0, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(uint16_t), row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor mul_mv_ext")) return 0; + return 1; + } + + /* + * Same direct-RHS TensorOps structure as Q8_0, but for F16 model + * matrices. The 128-token RHS tile is kept when the batch alignment + * allows it because the later tile_n=64 retest was neutral/slower. + */ + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 32u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = nax_tile_n == 128u + ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n128" + : (nax_tile_n == 64u + ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n64" + : "kernel_mul_mm_f16_f32_mpp_direct_rhs"); + id pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 NAX tensor matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f16_f32", bc_inp, bc_out); + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matmul")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_f16_pair_tensor( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok != 1 || (in_dim & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outabuf = ds4_gpu_tensor_buffer(out_a); + id outbbuf = ds4_gpu_tensor_buffer(out_b); + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out_bytes = out_dim * sizeof(float); + if (!xbuf || !outabuf || !outbbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out_a) < out_bytes || + ds4_gpu_tensor_bytes(out_b) < out_bytes) { + fprintf(stderr, "ds4: Metal F16 paired matvec received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_a_offset > model_size || weight_bytes > model_size - weight_a_offset || + weight_b_offset > model_size || weight_bytes > model_size - weight_b_offset) { + fprintf(stderr, "ds4: Metal F16 paired matvec range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_a = 0; + uint64_t inner_b = 0; + id wabuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_a_offset, weight_bytes, + &inner_a); + id wbbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_b_offset, weight_bytes, + &inner_b); + if (!wabuf || !wbbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (ds4_gpu_use_compressor_pair_nr4() && + (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_f16_f32_pair_4", mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wabuf offset:(NSUInteger)inner_a atIndex:1]; + [enc setBuffer:wbbuf offset:(NSUInteger)inner_b atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outabuf offset:ds4_gpu_tensor_offset(out_a) atIndex:4]; + [enc setBuffer:outbbuf offset:ds4_gpu_tensor_offset(out_b) atIndex:5]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 paired matvec")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_f16_pair_compressor_store_tensor( + ds4_gpu_tensor *out_kv, + ds4_gpu_tensor *out_score, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t weight_kv_offset, + uint64_t weight_score_offset, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t in_dim, + uint32_t width, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos) { + if (!g_initialized && !ds4_gpu_init()) return -1; + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL; + if ((g_quality_mode || + (!ds4_gpu_device_name_contains("M3") && + !ds4_gpu_device_name_contains("M5") && !force)) || + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) { + return 0; + } + if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x || + in_dim == 0 || width == 0 || ratio == 0 || + (ape_type != 0u && ape_type != 1u)) { + return -1; + } + if (in_dim != 4096u || + (width != 256u && width != 512u && width != 1024u) || + (ratio != 4u && ratio != 128u)) { + return 0; + } + + @autoreleasepool { + const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * width; + const uint64_t out_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = + (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem; + if (weight_kv_offset > model_size || + weight_bytes > model_size - weight_kv_offset || + weight_score_offset > model_size || + weight_bytes > model_size - weight_score_offset || + ape_offset > model_size || ape_bytes > model_size - ape_offset) { + return -1; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id outkvbuf = ds4_gpu_tensor_buffer(out_kv); + id outscorebuf = ds4_gpu_tensor_buffer(out_score); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out_kv) < out_bytes || + ds4_gpu_tensor_bytes(out_score) < out_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return -1; + } + + uint64_t weight_kv_inner = 0; + uint64_t weight_score_inner = 0; + uint64_t ape_inner = 0; + id weightkvbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_kv_offset, weight_bytes, + &weight_kv_inner); + id weightscorebuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_score_offset, weight_bytes, + &weight_score_inner); + id apebuf = ds4_gpu_wrap_model_range( + model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!weightkvbuf || !weightscorebuf || !apebuf) return -1; + + ds4_gpu_f16_matvec_args mv_args = + ds4_gpu_make_f16_mv_args(in_dim, width); + ds4_gpu_mv_dispatch mv_dispatch = + ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (ds4_gpu_use_compressor_pair_nr4() && + (width == 512u || width == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + ds4_gpu_dsv4_compressor_store_one_args store_args = { + .width = width, + .ratio = ratio, + .pos = pos, + .ape_type = ape_type, + }; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_mul_mv_f16_f32_pair_compressor_store_4", + mv_dispatch.nsg); + if (!pipeline) return -1; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return -1; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&store_args length:sizeof(store_args) atIndex:1]; + [enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2]; + [enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5]; + [enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake( + 32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "F16 paired matvec compressor state store")) { + return -1; + } + } + + return 1; +} + +int ds4_gpu_matmul_f32_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_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok == 0 || n_tok > UINT32_MAX) return 0; + if (in_dim > UINT64_MAX / n_tok / sizeof(float) || + out_dim > UINT64_MAX / n_tok / sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal F32 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(float); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal F32 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_f32_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, 1); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matvec")) return 0; + return 1; + } + + if (n_tok <= 8 && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_f32_name(r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(float), row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor mul_mv_ext")) return 0; + return 1; + } + + /* Generic multi-row path (GLM prefill shapes: one grid row per + * token through the plain matvec pipeline). */ + { + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, n_tok); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matmul")) return 0; + } + } + + return 1; +} + +int ds4_gpu_repeat_hc_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *row, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !row || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id rowbuf = ds4_gpu_tensor_buffer(row); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t out_bytes = row_bytes * n_hc; + if (!rowbuf || !outbuf || + ds4_gpu_tensor_bytes(row) < row_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal HC repeat received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + rowbuf, + ds4_gpu_tensor_offset(row), + outbuf, + ds4_gpu_tensor_offset(out), + 1, + n_embd, + n_hc)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC repeat")) return 0; + } + + return 1; +} + +int ds4_gpu_repeat_hc_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *rows, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id rowsbuf = ds4_gpu_tensor_buffer(rows); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t rows_bytes = (uint64_t)n_tokens * row_bytes; + const uint64_t out_bytes = rows_bytes * n_hc; + if (!rowsbuf || !outbuf || + ds4_gpu_tensor_bytes(rows) < rows_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal HC row repeat received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + rowsbuf, + ds4_gpu_tensor_offset(rows), + outbuf, + ds4_gpu_tensor_offset(out), + n_tokens, + n_embd, + n_hc)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC row repeat")) return 0; + } + + return 1; +} + +int ds4_gpu_rms_norm_plain_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + float eps) { + return ds4_gpu_rms_norm_plain_rows_tensor(out, x, n, 1, eps); +} + +int ds4_gpu_rms_norm_plain_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + uint32_t rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * rows * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal plain RMS norm received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_plain_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "plain RMS norm")) return 0; + } + + return 1; +} + +static int ds4_gpu_hc_rms_scale_project_mode( + uint32_t in_dim, + uint32_t out_dim, + uint32_t n_rows) { + const bool hard_shape = n_rows > 8u && + (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u; + if (!hard_shape) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL; + if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL || + (!ds4_gpu_device_name_contains("M3") && !force)) { + return 0; + } + if (g_rms_norm_scale_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +int ds4_gpu_hc_rms_scale_project_f16_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *scale_scratch, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !scale_scratch || !model_map || !x || + in_dim == 0u || out_dim == 0u || n_rows == 0u || + n_rows > INT32_MAX) { + return 0; + } + + const int mode = + ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows); + if (mode < 0) return 0; + if (mode == 0) { + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + @autoreleasepool { + const uint64_t x_elems = (uint64_t)in_dim * n_rows; + const uint64_t out_elems = (uint64_t)out_dim * n_rows; + if (x_elems > UINT64_MAX / sizeof(float) || + out_elems > UINT64_MAX / sizeof(float)) { + return 0; + } + const uint64_t x_bytes = x_elems * sizeof(float); + const uint64_t out_bytes = out_elems * sizeof(float); + const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float); + + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + id scalebuf = ds4_gpu_tensor_buffer(scale_scratch); + if (!xbuf || !outbuf || !scalebuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) { + fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n"); + return 0; + } + + const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t); + if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0; + const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes; + if (weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n"); + return 0; + } + + const bool bc_inp = (in_dim % 32u) != 0u; + const bool bc_out = (out_dim % 64u) != 0u || + (n_rows % 32u) != 0u; + id mm_pipeline = + ds4_gpu_get_mul_mm_pipeline( + "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out); + if (!mm_pipeline) { + if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) { + return 0; + } + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + id scale_pipeline = ds4_gpu_hot_pipeline( + g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4"); + if (!scale_pipeline) return 0; + + uint64_t weight_inner = 0; + id weightbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_offset, weight_bytes, &weight_inner); + if (!weightbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_rms_norm_args norm_args = + ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps); + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:scale_pipeline]; + [enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:2]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake( + ds4_gpu_rms_norm_threads(in_dim), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args( + in_dim, out_dim, n_rows, weight_row_bytes); + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:4]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)n_rows + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "HC RMS scale F16 projection")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_rms_norm_weight_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps) { + return ds4_gpu_rms_norm_weight_rows_tensor(out, x, model_map, model_size, weight_offset, n, 1, eps); +} + +int ds4_gpu_rms_norm_weight_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + uint32_t rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint64_t bytes = row_bytes * rows; + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal weighted RMS norm received undersized activation buffers\n"); + return 0; + } + if (weight_offset > model_size || row_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal weighted RMS norm range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + rows == 1u && + row_bytes <= (1ull << 20) && + getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t inner_offset = 0; + id wbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "weighted RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_add_rms_norm_weight_tensor( + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *sum_out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!norm_out || !sum_out || !a || !b || n == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id sumbuf = ds4_gpu_tensor_buffer(sum_out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !sumbuf || !normbuf || + ds4_gpu_tensor_bytes(a) < row_bytes || + ds4_gpu_tensor_bytes(b) < row_bytes || + ds4_gpu_tensor_bytes(sum_out) < row_bytes || + ds4_gpu_tensor_bytes(norm_out) < row_bytes) { + fprintf(stderr, "ds4: Metal add+RMS norm received undersized activation buffers\n"); + return 0; + } + if (weight_offset > model_size || row_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal add+RMS norm range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + row_bytes <= (1ull << 20) && + getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t inner_offset = 0; + id wbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:3]; + [enc setBuffer:sumbuf offset:ds4_gpu_tensor_offset(sum_out) atIndex:4]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "add+RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_dsv4_qkv_rms_norm_rows_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, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q_out || !q || !kv_out || !kv || q_n == 0 || kv_n == 0 || rows == 0 || + (q_n & 3u) != 0 || (kv_n & 3u) != 0) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id qoutbuf = ds4_gpu_tensor_buffer(q_out); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id kvoutbuf = ds4_gpu_tensor_buffer(kv_out); + + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float); + if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf || + ds4_gpu_tensor_bytes(q) < q_row_bytes * rows || + ds4_gpu_tensor_bytes(q_out) < q_row_bytes * rows || + ds4_gpu_tensor_bytes(kv) < kv_row_bytes * rows || + ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes * rows) { + fprintf(stderr, "ds4: Metal fused q/kv RMS norm received undersized activation buffers\n"); + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) { + fprintf(stderr, "ds4: Metal fused q/kv RMS norm weight range is outside the mapped model\n"); + return 0; + } + + uint64_t q_inner_offset = 0; + uint64_t kv_inner_offset = 0; + id q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + q_weight_offset, q_row_bytes, + &q_inner_offset); + if (!q_wbuf) return 0; + id kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + kv_weight_offset, kv_row_bytes, + &kv_inner_offset); + if (!kv_wbuf) return 0; + + ds4_gpu_qkv_rms_norm_args args = { + .q_n = (int32_t)q_n, + .q_n4 = (int32_t)(q_n / 4u), + .kv_n = (int32_t)kv_n, + .kv_n4 = (int32_t)(kv_n / 4u), + .q_row_stride = q_row_bytes, + .kv_row_stride = kv_row_bytes, + .eps = eps, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_qkv_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:2]; + [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:4]; + [enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:5]; + [enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:6]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 2, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_head_rms_norm_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || (head_dim & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal head RMS norm received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_3d_args(head_dim, n_head, n_tok, eps); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_plain_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_pipeline_threads(head_dim, g_rms_norm_plain_pipeline), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "head RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t 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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || n_rot > head_dim || (n_rot & 1u) != 0) { + return 0; + } + if (n_rot == 0) return 1; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal RoPE received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_rope_tail_batch_args args = ds4_gpu_make_rope_tail_args( + n_tok, n_head, head_dim, n_rot, n_ctx_orig, inverse, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_rope_tail_inplace(cb, + xbuf, + ds4_gpu_tensor_offset(x), + &args, + n_tok, + n_head, + head_dim, + pos0, + 1)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "RoPE tail")) return 0; + } + + return 1; +} + +int ds4_gpu_head_rms_norm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t 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) { + (void)x; (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; + (void)pos0; (void)n_ctx_orig; (void)inverse; (void)freq_base; + (void)freq_scale; (void)ext_factor; (void)attn_factor; + (void)beta_fast; (void)beta_slow; (void)eps; + return 0; +} + +int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *q_half, + 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_tok, + uint32_t n_head, + uint32_t 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) { + (void)out; (void)q_half; (void)model_map; (void)model_size; + (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; + (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0; + (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale; + (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; + (void)eps; + return 0; +} + +int ds4_gpu_dsv4_fp8_kv_quantize_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t head_dim, + uint32_t n_rot) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || head_dim == 0 || n_rot > head_dim) return 0; + if (n_rot == head_dim) return 1; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal DSV4 FP8 KV quantize received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_dsv4_fp8_kv_quantize_args args = { + .ne00 = head_dim, + .ne01 = n_tok, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)head_dim * sizeof(float), + .nb02 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb03 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)head_dim * sizeof(float), + .nb2 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb3 = (uint64_t)n_tok * head_dim * sizeof(float), + .n_rot = (int32_t)n_rot, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_fp8_kv_quantize_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tok, 1, 1) + threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 FP8 KV quantize")) return 0; + } + + return 1; +} + +int ds4_gpu_dsv4_indexer_qat_tensor( + ds4_gpu_tensor *x, + uint32_t n_rows, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_rows == 0 || head_dim != 128u) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_rows * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal DSV4 indexer QAT received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_dsv4_indexer_qat_args args = { + .n_rows = n_rows, + .head_dim = head_dim, + .row_stride = (uint64_t)head_dim * sizeof(float), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_indexer_qat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 indexer Hadamard+FP4")) return 0; + } + + return 1; +} + +static void ds4_gpu_set_rows_thread_shape( + uint32_t width, + NSUInteger *nth_out, + NSUInteger *nrptg_out) { + const NSUInteger nk0 = width ? (NSUInteger)width : 1u; + const NSUInteger max_threads = g_set_rows_f32_i32_pipeline + ? (NSUInteger)g_set_rows_f32_i32_pipeline.maxTotalThreadsPerThreadgroup + : 1024u; + + NSUInteger nth = 32u; + while (nth < nk0 && nth < max_threads) { + nth *= 2u; + } + + NSUInteger nrptg = 1u; + if (nth > nk0) { + nrptg = (nth + nk0 - 1u) / nk0; + nth = nk0; + if (nrptg * nth > max_threads) { + nrptg--; + } + } + + if (nth > nk0) nth = nk0; + if (nth == 0u) nth = 1u; + if (nrptg == 0u) nrptg = 1u; + + *nth_out = nth; + *nrptg_out = nrptg; +} + +static int ds4_gpu_encode_f16_round_copy_for_raw_store( + id cb, + const ds4_gpu_tensor *src, + uint32_t n) { + id srcbuf = ds4_gpu_tensor_buffer(src); + const uint64_t src_bytes = (uint64_t)n * sizeof(float); + if (!srcbuf || ds4_gpu_tensor_bytes(src) < src_bytes) { + fprintf(stderr, "ds4: Metal raw KV store received undersized source buffer\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_f16_round_scratch_buffer, + &g_f16_round_scratch_bytes, + (NSUInteger)n * sizeof(uint16_t), + "ds4_f16_round_scratch") || + !ds4_gpu_ensure_scratch_buffer(&g_raw_store_round_buffer, + &g_raw_store_round_bytes, + (NSUInteger)n * sizeof(float), + "ds4_raw_store_round")) { + return 0; + } + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + srcbuf, + ds4_gpu_tensor_offset(src), + g_f16_round_scratch_buffer, + 0, + n)) { + return 0; + } + return ds4_gpu_encode_cpy_f16_f32_1d(cb, + g_f16_round_scratch_buffer, + 0, + g_raw_store_round_buffer, + 0, + n); +} + +static int ds4_gpu_encode_set_rows_f32_i32( + id cb, + ds4_gpu_tensor *dst, + id srcbuf, + NSUInteger src_off, + const int32_t *rows, + uint32_t n_rows, + uint32_t dst_rows, + uint32_t width) { + id dstbuf = ds4_gpu_tensor_buffer(dst); + const uint64_t dst_bytes = (uint64_t)dst_rows * width * sizeof(float); + const uint64_t src_bytes = (uint64_t)n_rows * width * sizeof(float); + if (!dstbuf || !srcbuf || !rows || n_rows == 0 || width == 0 || + ds4_gpu_tensor_bytes(dst) < dst_bytes || + src_bytes > NSUIntegerMax - src_off) { + fprintf(stderr, "ds4: Metal DS4 set_rows received invalid buffers\n"); + return 0; + } + + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t rows_bytes = (uint64_t)n_rows * sizeof(int32_t); + ds4_gpu_set_rows_args args = { + .nk0 = (int32_t)width, + .ne01 = (int32_t)n_rows, + .nb01 = row_bytes, + .nb02 = (uint64_t)n_rows * row_bytes, + .nb03 = (uint64_t)n_rows * row_bytes, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(int32_t), + .nb11 = rows_bytes, + .nb12 = rows_bytes, + .nb1 = row_bytes, + .nb2 = (uint64_t)dst_rows * row_bytes, + .nb3 = (uint64_t)dst_rows * row_bytes, + }; + + NSUInteger nth; + NSUInteger nrptg; + ds4_gpu_set_rows_thread_shape(width, &nth, &nrptg); + + id rowsbuf = nil; + if (rows_bytes > 4096u) { + rowsbuf = ds4_gpu_new_transient_buffer((NSUInteger)rows_bytes, "ds4_set_rows_indices"); + if (!rowsbuf) return 0; + memcpy([rowsbuf contents], rows, (NSUInteger)rows_bytes); + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_set_rows_f32_i32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:srcbuf offset:src_off atIndex:1]; + if (rowsbuf) { + [enc setBuffer:rowsbuf offset:0 atIndex:2]; + } else { + [enc setBytes:rows length:(NSUInteger)rows_bytes atIndex:2]; + } + [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + nrptg - 1u) / nrptg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, nrptg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_add_f32_1d( + id cb, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off, + uint32_t n) { + if (!cb || !a || !b || !out || n == 0) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add2_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:a offset:a_off atIndex:1]; + [enc setBuffer:b offset:b_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} diff --git a/metal/elementwise.inc b/metal/elementwise.inc new file mode 100644 index 0000000000..6f2a479714 --- /dev/null +++ b/metal/elementwise.inc @@ -0,0 +1,228 @@ +int ds4_gpu_swiglu_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *gate, + const ds4_gpu_tensor *up, + uint32_t n, + float clamp, + float weight) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !gate || !up || n == 0) return 0; + if (!isfinite(clamp) || clamp < 0.0f || !isfinite(weight)) return 0; + + @autoreleasepool { + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!gatebuf || !upbuf || !outbuf || + ds4_gpu_tensor_bytes(gate) < bytes || + ds4_gpu_tensor_bytes(up) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal SwiGLU received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glu_args args = { + .ne00 = (int32_t)n, + .nb01 = (uint64_t)n * sizeof(float), + .ne10 = (int32_t)n, + .nb11 = (uint64_t)n * sizeof(float), + .ne0 = (int32_t)n, + .nb1 = (uint64_t)n * sizeof(float), + .i00 = 0, + .i10 = 0, + .alpha = weight, + .limit = clamp, + }; + NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_swiglu_flat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:1]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "SwiGLU")) return 0; + } + + return 1; +} + +int ds4_gpu_add_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + uint32_t n) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !a || !b || n == 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !outbuf || + ds4_gpu_tensor_bytes(a) < bytes || + ds4_gpu_tensor_bytes(b) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal tensor add received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add2_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add")) return 0; + } + + return 1; +} + +int ds4_gpu_add3_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + uint32_t n) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !a || !b || !c || n == 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id cbuf = ds4_gpu_tensor_buffer(c); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !cbuf || !outbuf || + ds4_gpu_tensor_bytes(a) < bytes || + ds4_gpu_tensor_bytes(b) < bytes || + ds4_gpu_tensor_bytes(c) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal tensor add3 received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add3_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add3_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:cbuf offset:ds4_gpu_tensor_offset(c) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add3")) return 0; + } + + return 1; +} + +typedef struct { + uint32_t width; + uint32_t rows; + uint32_t layer; + uint32_t n_threads; + float scale; +} ds4_gpu_directional_steering_project_args; + +int ds4_gpu_directional_steering_project_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_tensor *directions, + uint32_t layer, + uint32_t width, + uint32_t rows, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0; + + @autoreleasepool { + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_directional_steering_project_f32"); + if (!pipeline) return 0; + + id xbuf = ds4_gpu_tensor_buffer(x); + id dbuf = ds4_gpu_tensor_buffer(directions); + const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float); + const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float); + if (!xbuf || !dbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(directions) < dir_bytes) { + fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + while (nth > width && nth > 1u) nth >>= 1; + if (nth == 0) nth = 1; + + ds4_gpu_directional_steering_project_args args = { + .width = width, + .rows = rows, + .layer = layer, + .n_threads = (uint32_t)nth, + .scale = scale, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:dbuf offset:ds4_gpu_tensor_offset(directions) atIndex:2]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "directional steering")) return 0; + } + + return 1; +} diff --git a/metal/embedding.inc b/metal/embedding.inc new file mode 100644 index 0000000000..717d14711c --- /dev/null +++ b/metal/embedding.inc @@ -0,0 +1,761 @@ +static int ds4_gpu_encode_get_rows_f16( + id cb, + id weight, + NSUInteger weight_offset, + id tokens, + NSUInteger tokens_offset, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!cb || !weight || !tokens || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); + const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + ds4_gpu_get_rows_args args = { + .ne00t = (int32_t)n_embd, + .ne00 = (int32_t)n_embd, + .nb01 = src_row_bytes, + .nb02 = (uint64_t)n_vocab * src_row_bytes, + .nb03 = (uint64_t)n_vocab * src_row_bytes, + .ne10 = (int32_t)n_tokens, + .nb10 = sizeof(int32_t), + .nb11 = token_bytes, + .nb12 = token_bytes, + .nb1 = dst_row_bytes, + .nb2 = (uint64_t)n_tokens * dst_row_bytes, + .nb3 = (uint64_t)n_tokens * dst_row_bytes, + }; + + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static uint64_t ds4_gpu_q8_0_row_bytes(uint32_t n_embd) { + return (((uint64_t)n_embd + 31u) / 32u) * 34u; +} + +static int ds4_gpu_quant_row_bytes( + uint32_t type, + uint32_t n_embd, + uint64_t *row_bytes_out) { + if (!row_bytes_out || n_embd == 0) return 0; + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 34u; + return 1; + case DS4_METAL_TENSOR_Q4_0: + *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 18u; + return 1; + case DS4_METAL_TENSOR_Q4_K: + if ((n_embd % 256u) != 0) return 0; + *row_bytes_out = ((uint64_t)n_embd / 256u) * 144u; + return 1; + default: + return 0; + } +} + +static int ds4_gpu_q8_0_table_bytes( + uint32_t n_vocab, + uint32_t n_embd, + uint64_t *bytes_out) { + if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; + const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); + if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; + *bytes_out = (uint64_t)n_vocab * row_bytes; + return 1; +} + +static int ds4_gpu_quant_table_bytes( + uint32_t type, + uint32_t n_vocab, + uint32_t n_embd, + uint64_t *bytes_out) { + if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(type, n_embd, &row_bytes)) return 0; + if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; + *bytes_out = (uint64_t)n_vocab * row_bytes; + return 1; +} + +static int ds4_gpu_encode_get_rows_q8_0( + id cb, + id weight, + NSUInteger weight_offset, + id tokens, + NSUInteger tokens_offset, + const int32_t *single_token, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + if (!tokens && (!single_token || n_tokens != 1)) { + return 0; + } + + ds4_gpu_get_rows_q8_0_args args = { + .n_embd = (int32_t)n_embd, + .n_vocab = (int32_t)n_vocab, + .n_tokens = (int32_t)n_tokens, + .src_row_bytes = ds4_gpu_q8_0_row_bytes(n_embd), + .dst_row_bytes = (uint64_t)n_embd * sizeof(float), + .token_stride = sizeof(int32_t), + }; + + NSUInteger nth = 32u; + const NSUInteger max_threads = g_get_rows_q8_0_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nblocks = ((NSUInteger)n_embd + 31u) / 32u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_q8_0_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + } else { + [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; + } + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_get_rows_quant( + id cb, + id weight, + NSUInteger weight_offset, + uint32_t weight_type, + id tokens, + NSUInteger tokens_offset, + const int32_t *single_token, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_encode_get_rows_q8_0(cb, + weight, + weight_offset, + tokens, + tokens_offset, + single_token, + out, + out_offset, + n_vocab, + n_tokens, + n_embd); + } + if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + if (!tokens && (!single_token || n_tokens != 1)) { + return 0; + } + + uint64_t src_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, n_embd, &src_row_bytes)) return 0; + ds4_gpu_get_rows_q8_0_args args = { + .n_embd = (int32_t)n_embd, + .n_vocab = (int32_t)n_vocab, + .n_tokens = (int32_t)n_tokens, + .src_row_bytes = src_row_bytes, + .dst_row_bytes = (uint64_t)n_embd * sizeof(float), + .token_stride = sizeof(int32_t), + }; + + id pipeline = nil; + NSUInteger block_width = 0; + if (weight_type == DS4_METAL_TENSOR_Q4_0) { + pipeline = g_get_rows_q4_0_pipeline; + block_width = 32u; + } else if (weight_type == DS4_METAL_TENSOR_Q4_K) { + pipeline = g_get_rows_q4_K_pipeline; + block_width = 256u; + } + if (!pipeline || block_width == 0) return 0; + + NSUInteger nth = block_width; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nblocks = ((NSUInteger)n_embd + block_width - 1u) / block_width; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + } else { + [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; + } + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_repeat_hc_embedding( + id cb, + id rows, + NSUInteger rows_offset, + id out, + NSUInteger out_offset, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!cb || !rows || !out || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; + + const uint64_t embd_bytes = (uint64_t)n_embd * sizeof(float); + ds4_gpu_repeat_args args = { + .ne00 = (int32_t)n_embd, + .ne01 = 1, + .ne02 = (int32_t)n_tokens, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = embd_bytes, + .nb02 = embd_bytes, + .nb03 = (uint64_t)n_tokens * embd_bytes, + .ne0 = (int32_t)n_embd, + .ne1 = (int32_t)n_hc, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = embd_bytes, + .nb2 = (uint64_t)n_hc * embd_bytes, + .nb3 = (uint64_t)n_tokens * n_hc * embd_bytes, + }; + + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_repeat_f32_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_repeat_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rows offset:rows_offset atIndex:1]; + [enc setBuffer:out offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(n_hc, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_embed_token_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 embedding received undersized output buffer\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + uint32_t token_for_kernel = token; + id wbuf = nil; + const bool exact_token_row = + getenv("DS4_METAL_DISABLE_TOKEN_EMBED_EXACT_VIEW") == NULL; + if (exact_token_row) { + const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); + const uint64_t token_rel = (uint64_t)token * row_bytes; + if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { + fprintf(stderr, "ds4: Metal Q8_0 embedding token row is outside the mapped table\n"); + return 0; + } + wbuf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset + token_rel, + row_bytes, + &inner_offset); + token_for_kernel = 0; + } else { + wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + } + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = (int32_t)token_for_kernel; + if (!ds4_gpu_encode_get_rows_q8_0(cb, + wbuf, + (NSUInteger)inner_offset, + nil, + 0, + &token_i32, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + 1, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed token")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_tokens_q8_0_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 batched embedding received undersized buffers\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_q8_0(cb, + wbuf, + (NSUInteger)inner_offset, + tokbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + n_tokens, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed tokens")) return 0; + } + + return 1; +} + +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 == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_embed_token_q8_0_tensor(out, + model_map, + model_size, + weight_offset, + n_vocab, + token, + n_embd); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal quant embedding received undersized output buffer\n"); + return 0; + } + + uint64_t weight_bytes = 0; + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || + !ds4_gpu_quant_row_bytes(weight_type, n_embd, &row_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant embedding range is outside the mapped model\n"); + return 0; + } + + const uint64_t token_rel = (uint64_t)token * row_bytes; + if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { + fprintf(stderr, "ds4: Metal quant embedding token row is outside the mapped table\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset + token_rel, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = 0; + if (!ds4_gpu_encode_get_rows_quant(cb, + wbuf, + (NSUInteger)inner_offset, + weight_type, + nil, + 0, + &token_i32, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + 1, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed token")) return 0; + } + + return 1; +} + +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 == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_embed_tokens_q8_0_tensor(out, + tokens, + model_map, + model_size, + weight_offset, + n_vocab, + n_tokens, + n_embd); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal quant batched embedding received undersized buffers\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_quant(cb, + wbuf, + (NSUInteger)inner_offset, + weight_type, + tokbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + n_tokens, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed tokens")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_token_hc_tensor( + ds4_gpu_tensor *out_hc, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0 || n_hc == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t out_bytes = (uint64_t)n_embd * n_hc * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out_hc) < out_bytes) { + fprintf(stderr, "ds4: Metal graph embedding received undersized HC output buffer\n"); + return 0; + } + + const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal graph embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + const NSUInteger row_bytes = (NSUInteger)n_embd * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, + &g_embed_rows_bytes, + row_bytes, + "ds4_embed_rows")) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = (int32_t)token; + const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); + const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); + ds4_gpu_get_rows_args args = { + .ne00t = (int32_t)n_embd, + .ne00 = (int32_t)n_embd, + .nb01 = src_row_bytes, + .nb02 = (uint64_t)n_vocab * src_row_bytes, + .nb03 = (uint64_t)n_vocab * src_row_bytes, + .ne10 = 1, + .nb10 = sizeof(int32_t), + .nb11 = sizeof(int32_t), + .nb12 = sizeof(int32_t), + .nb1 = dst_row_bytes, + .nb2 = dst_row_bytes, + .nb3 = dst_row_bytes, + }; + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBytes:&token_i32 length:sizeof(token_i32) atIndex:2]; + [enc setBuffer:g_embed_rows_buffer offset:0 atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + g_embed_rows_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out_hc), + 1, + n_embd, + n_hc)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed token")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_tokens_hc_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0 || n_hc == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out_hc); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * n_hc * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out_hc) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal graph batched embedding received undersized buffers\n"); + return 0; + } + + const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal graph batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + const NSUInteger rows_bytes = (NSUInteger)n_tokens * n_embd * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, + &g_embed_rows_bytes, + rows_bytes, + "ds4_embed_rows")) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_f16(cb, + wbuf, + (NSUInteger)inner_offset, + tokbuf, + ds4_gpu_tensor_offset(tokens), + g_embed_rows_buffer, + 0, + n_vocab, + n_tokens, + n_embd) || + !ds4_gpu_encode_repeat_hc_embedding(cb, + g_embed_rows_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out_hc), + n_tokens, + n_embd, + n_hc)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed tokens")) return 0; + } + + return 1; +} diff --git a/metal/expert_streaming.inc b/metal/expert_streaming.inc new file mode 100644 index 0000000000..b1ecf543c0 --- /dev/null +++ b/metal/expert_streaming.inc @@ -0,0 +1,5215 @@ +uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { + uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + budget = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; + } + return budget; +} + +uint32_t ds4_gpu_stream_expert_cache_current_count(void) { + return g_stream_expert_cache_entry_count; +} + +uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes)) { + return 0; + } + return ds4_gpu_stream_expert_cache_configured_budget(); +} + +static int ds4_gpu_stream_expert_cache_note_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (gate_expert_bytes == 0 || down_expert_bytes == 0) return 0; + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); + return 0; + } + /* + * The cache is a single-size-class slab allocator: the expert byte size is + * frozen on first sight (or pre-seeded at startup from the model's slab + * class) and off-size layers are REJECTED rather than adopted. A rejected + * layer (mixed-precision boost: Q4_K experts among IQ2 layers) falls back + * to the mapped-model per-expert path; last-writer-wins here would instead + * poison the slab size class and deadlock slab reuse. + */ + const uint64_t bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (g_stream_expert_cache_expert_bytes == 0) { + g_stream_expert_cache_expert_bytes = bytes; + return 1; + } + return bytes == g_stream_expert_cache_expert_bytes; +} + +static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) { + if (!g_ssd_streaming_mode) return 0; + if (g_stream_expert_cache_budget_override != 0) { + return g_stream_expert_cache_budget_override; + } + return 0; +} + +static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) { + return ds4_gpu_stream_expert_cache_requested_budget(); +} + +static uint32_t ds4_gpu_stream_expert_cache_effective_cap( + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected) { + (void)layer; + if (ds4_gpu_stream_expert_cache_configured_budget() == 0) return 0; + + /* + * The residency policy is global: every layer can use any expert slot it + * routes to, and global pruning decides which existing entry is least + * valuable. A per-layer cap made cache size depend on model depth rather + * than the actual byte budget. + */ + uint32_t cap = n_total_expert; + if (cap < n_selected) cap = n_selected; + if (cap > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + cap = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + } + return cap; +} + +static int ds4_gpu_stream_expert_timing_summary_enabled(void) { + static int checked = 0; + static int enabled = 0; + if (!checked) { + enabled = + (getenv("DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY") != NULL || + getenv("DS4_METAL_STREAMING_EXPERT_PROFILE_SUMMARY") != NULL) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_TIMING_SUMMARY") == NULL; + checked = 1; + } + return enabled; +} + +static uint32_t ds4_gpu_stream_expert_popcount(uint32_t mask) { + return (uint32_t)__builtin_popcount(mask); +} + +static int ds4_gpu_stream_expert_split_worthwhile( + uint32_t resident_mask, + uint32_t missing_mask) { + if (resident_mask == 0 || missing_mask == 0) return 0; + /* + * The split path pays an extra command stage and a second routed-expert + * bind. It is worthwhile when several experts are missing and their SSD + * reads can be hidden by resident expert work. With one or two misses, + * especially in large caches, a single unsplit routed pass is faster. + */ + return ds4_gpu_stream_expert_popcount(missing_mask) >= 3u; +} + +static void ds4_gpu_stream_expert_timing_note_selected( + double sync_ms, + double copy_ms, + double bind_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_selected_calls++; + g_stream_expert_timing_selected_sync_ms += sync_ms; + g_stream_expert_timing_selected_copy_ms += copy_ms; + g_stream_expert_timing_selected_read_ms += sync_ms + copy_ms; + g_stream_expert_timing_selected_bind_ms += bind_ms; +} + +static void ds4_gpu_stream_expert_timing_note_split( + uint32_t resident_mask, + uint32_t missing_mask, + double resident_ms, + double missing_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_split_layers++; + g_stream_expert_timing_split_resident_experts += + ds4_gpu_stream_expert_popcount(resident_mask); + g_stream_expert_timing_split_missing_experts += + ds4_gpu_stream_expert_popcount(missing_mask); + g_stream_expert_timing_split_resident_ms += resident_ms; + g_stream_expert_timing_split_missing_ms += missing_ms; +} + +static void ds4_gpu_stream_expert_timing_note_split_missing_detail( + double load_ms, + double slot_ms, + double prune_ms, + double addr_ms, + double wait_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_split_missing_load_ms += load_ms; + g_stream_expert_timing_split_missing_slot_ms += slot_ms; + g_stream_expert_timing_split_missing_prune_ms += prune_ms; + g_stream_expert_timing_split_missing_addr_ms += addr_ms; + g_stream_expert_timing_split_missing_wait_ms += wait_ms; +} + +static void ds4_gpu_stream_expert_timing_note_load_detail( + double prepare_ms, + double pread_ms, + double modify_ms, + double install_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_load_calls++; + g_stream_expert_timing_load_prepare_ms += prepare_ms; + g_stream_expert_timing_load_pread_ms += pread_ms; + g_stream_expert_timing_load_modify_ms += modify_ms; + g_stream_expert_timing_load_install_ms += install_ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_batch_reuse_calls++; + g_stream_expert_timing_prepare_batch_reuse_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_buffer(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_buffer_calls++; + g_stream_expert_timing_prepare_buffer_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_task( + uint32_t experts, + double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_task_experts += experts; + g_stream_expert_timing_prepare_task_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_reuse_scan( + uint64_t entries, + double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_reuse_scan_calls++; + g_stream_expert_timing_reuse_scan_entries += entries; + g_stream_expert_timing_reuse_scan_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_reuse_clear(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_reuse_clear_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_cache_class( + uint32_t resident_mask, + uint32_t missing_mask) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + const uint32_t resident = ds4_gpu_stream_expert_popcount(resident_mask); + const uint32_t missing = ds4_gpu_stream_expert_popcount(missing_mask); + if (missing == 0) { + g_stream_expert_timing_cache_all_resident_layers++; + } else if (resident == 0) { + g_stream_expert_timing_cache_all_missing_layers++; + } else { + g_stream_expert_timing_cache_mixed_layers++; + } + g_stream_expert_timing_cache_resident_experts += resident; + g_stream_expert_timing_cache_missing_experts += missing; +} + +static int ds4_gpu_stream_expert_readahead_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_READAHEAD") == NULL; +} + +static void ds4_gpu_stream_expert_readahead_range(uint64_t offset, uint64_t len) { + if (!ds4_gpu_stream_expert_readahead_enabled() || g_model_fd < 0 || len == 0) { + return; + } + +#if defined(F_RDADVISE) + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t pos = offset; + uint64_t rem = len; + while (rem > 0) { + const uint64_t chunk64 = + rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem; + if (pos > (uint64_t)LLONG_MAX) break; + + struct radvisory ra; + ra.ra_offset = (off_t)pos; + ra.ra_count = (int)chunk64; + (void)fcntl(g_model_fd, F_RDADVISE, &ra); + + pos += chunk64; + rem -= chunk64; + } + if (timing) { + g_stream_expert_timing_readahead_calls++; + g_stream_expert_timing_readahead_bytes += len; + g_stream_expert_timing_readahead_ms += ds4_gpu_now_ms() - t0; + } +#else + (void)offset; + (void)len; +#endif +} + +typedef struct { + uint64_t offset; + uint64_t len; + uint8_t *dst; + uint64_t read_bytes; + double ms; + int ok; +} ds4_gpu_stream_expert_pread_task; + +typedef struct { + int active; + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint32_t n_selected; + uint32_t missing_mask; + uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t n_loads; + uint32_t n_tasks; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t gate_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t up_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t down_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; + double start_ms; + double prepare_ms; +} ds4_gpu_stream_expert_pending_load; + +static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load; + +typedef struct { + int active; + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint32_t n_selected; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; +} ds4_gpu_glm_stream_selected_prefetch; + +static ds4_gpu_glm_stream_selected_prefetch g_glm_stream_selected_prefetch; + +typedef struct { + ds4_gpu_stream_expert_pread_task *tasks; + uint32_t n_tasks; + uint32_t worker_index; + uint32_t n_workers; +} ds4_gpu_stream_expert_pread_worker_args; + +static void ds4_gpu_stream_expert_cache_note_pread( + uint32_t layer, + uint64_t bytes, + double ms) { + if (g_stream_expert_cache_pread_bytes > UINT64_MAX - bytes) { + g_stream_expert_cache_pread_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_pread_bytes += bytes; + } + g_stream_expert_cache_pread_ms += ms; + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + if (g_stream_expert_cache_layer_pread_bytes[layer] > UINT64_MAX - bytes) { + g_stream_expert_cache_layer_pread_bytes[layer] = UINT64_MAX; + } else { + g_stream_expert_cache_layer_pread_bytes[layer] += bytes; + } + g_stream_expert_cache_layer_pread_ms[layer] += ms; + } +} + +static uint32_t ds4_gpu_stream_expert_pread_thread_limit(void) { + uint32_t threads = 9; + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_THREADS"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0') { + threads = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v; + } + } + if (threads == 0) threads = 1; + if (threads > 18) threads = 18; + return threads; +} + +static uint32_t ds4_gpu_stream_expert_pread_thread_count(uint32_t n_tasks) { + if (n_tasks <= 1) return n_tasks; + uint32_t threads = ds4_gpu_stream_expert_pread_thread_limit(); + if (threads > n_tasks) threads = n_tasks; + return threads; +} + +static int ds4_gpu_stream_expert_pread_into( + uint64_t offset, + uint64_t len, + uint8_t *dst, + uint64_t *read_bytes, + double *ms_out) { + if (read_bytes) *read_bytes = 0; + if (ms_out) *ms_out = 0.0; + if (g_model_fd < 0 || + !dst || + len == 0 || + offset > (uint64_t)LLONG_MAX || + len > (uint64_t)LLONG_MAX - offset) { + return 0; + } + + const double t0 = ds4_gpu_now_ms(); + uint64_t pos = 0; + int ok = 1; + while (pos < len) { + const uint64_t rem = len - pos; + const size_t want = rem > (uint64_t)SSIZE_MAX ? (size_t)SSIZE_MAX : (size_t)rem; + ssize_t nread; + do { + nread = pread(g_model_fd, dst + pos, want, (off_t)(offset + pos)); + } while (nread < 0 && errno == EINTR); + if (nread <= 0) { + ok = 0; + break; + } + pos += (uint64_t)nread; + } + const double dt = ds4_gpu_now_ms() - t0; + if (read_bytes) *read_bytes = pos; + if (ms_out) *ms_out = dt; + if (!ok || pos != len) { + fprintf(stderr, + "ds4: Metal streaming expert explicit pread failed offset=%.2f GiB len=%.2f MiB read=%.2f MiB\n", + ds4_gpu_gib(offset), + ds4_gpu_mib(len), + ds4_gpu_mib(pos)); + return 0; + } + return 1; +} + +static void *ds4_gpu_stream_expert_pread_worker(void *arg) { + ds4_gpu_stream_expert_pread_worker_args *wa = + (ds4_gpu_stream_expert_pread_worker_args *)arg; + for (uint32_t i = wa->worker_index; i < wa->n_tasks; i += wa->n_workers) { + ds4_gpu_stream_expert_pread_task *task = &wa->tasks[i]; + task->ok = ds4_gpu_stream_expert_pread_into(task->offset, + task->len, + task->dst, + &task->read_bytes, + &task->ms); + } + return NULL; +} + +static pthread_mutex_t g_stream_expert_pread_pool_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_stream_expert_pread_pool_start_cond = PTHREAD_COND_INITIALIZER; +static pthread_cond_t g_stream_expert_pread_pool_done_cond = PTHREAD_COND_INITIALIZER; +static pthread_t g_stream_expert_pread_pool_threads[18]; +static uint32_t g_stream_expert_pread_pool_thread_count; +static uint32_t g_stream_expert_pread_pool_active_workers; +static uint32_t g_stream_expert_pread_pool_remaining_workers; +static uint32_t g_stream_expert_pread_pool_n_tasks; +static uint32_t g_stream_expert_pread_pool_next_task; +static uint64_t g_stream_expert_pread_pool_generation; +static ds4_gpu_stream_expert_pread_task *g_stream_expert_pread_pool_tasks; +static int g_stream_expert_pread_pool_initialized; +static int g_stream_expert_pread_pool_stopping; + +static int ds4_gpu_stream_expert_pread_pool_enabled(void) { + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_POOL"); + return !(env && strcmp(env, "0") == 0); +} + +static void *ds4_gpu_stream_expert_pread_pool_worker(void *arg) { + const uint32_t worker_index = (uint32_t)(uintptr_t)arg; + uint64_t seen_generation = 0; + + for (;;) { + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + while (!g_stream_expert_pread_pool_stopping && + g_stream_expert_pread_pool_generation == seen_generation) { + pthread_cond_wait(&g_stream_expert_pread_pool_start_cond, + &g_stream_expert_pread_pool_mutex); + } + if (g_stream_expert_pread_pool_stopping) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + break; + } + + seen_generation = g_stream_expert_pread_pool_generation; + if (worker_index >= g_stream_expert_pread_pool_active_workers) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + continue; + } + + for (;;) { + const uint32_t task_index = + g_stream_expert_pread_pool_next_task++; + if (task_index >= g_stream_expert_pread_pool_n_tasks) break; + + ds4_gpu_stream_expert_pread_task *task = + &g_stream_expert_pread_pool_tasks[task_index]; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + task->ok = ds4_gpu_stream_expert_pread_into(task->offset, + task->len, + task->dst, + &task->read_bytes, + &task->ms); + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + } + + if (g_stream_expert_pread_pool_remaining_workers > 0 && + --g_stream_expert_pread_pool_remaining_workers == 0) { + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_active_workers = 0; + pthread_cond_signal(&g_stream_expert_pread_pool_done_cond); + } + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + } + + return NULL; +} + +static int ds4_gpu_stream_expert_pread_pool_init(uint32_t n_threads) { + if (g_stream_expert_pread_pool_initialized) return 1; + if (!ds4_gpu_stream_expert_pread_pool_enabled() || n_threads <= 1) return 0; + if (n_threads > 18) n_threads = 18; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = n_threads; + g_stream_expert_pread_pool_stopping = 0; + g_stream_expert_pread_pool_generation = 0; + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_active_workers = 0; + g_stream_expert_pread_pool_remaining_workers = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + uint32_t started = 0; + for (uint32_t i = 0; i < n_threads; i++) { + const int rc = pthread_create(&g_stream_expert_pread_pool_threads[i], + NULL, + ds4_gpu_stream_expert_pread_pool_worker, + (void *)(uintptr_t)i); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming expert pread pool thread creation failed: %s\n", + strerror(rc)); + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_stopping = 1; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + for (uint32_t j = 0; j < started; j++) { + (void)pthread_join(g_stream_expert_pread_pool_threads[j], NULL); + } + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = 0; + g_stream_expert_pread_pool_stopping = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + started++; + } + + g_stream_expert_pread_pool_initialized = 1; + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_begin( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint32_t n_workers) { + if (n_workers <= 1) return 0; + const uint32_t limit = ds4_gpu_stream_expert_pread_thread_limit(); + if (!ds4_gpu_stream_expert_pread_pool_init(limit)) return 0; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + if (!g_stream_expert_pread_pool_initialized || + g_stream_expert_pread_pool_stopping || + g_stream_expert_pread_pool_thread_count == 0) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + if (n_workers > g_stream_expert_pread_pool_thread_count) { + n_workers = g_stream_expert_pread_pool_thread_count; + } + if (n_workers == 0) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + + if (g_stream_expert_pread_pool_remaining_workers != 0 || + g_stream_expert_pread_pool_tasks != NULL) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + + g_stream_expert_pread_pool_tasks = tasks; + g_stream_expert_pread_pool_n_tasks = n_tasks; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_active_workers = n_workers; + g_stream_expert_pread_pool_remaining_workers = n_workers; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_wait(void) { + if (!g_stream_expert_pread_pool_initialized) return 0; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + while (g_stream_expert_pread_pool_remaining_workers != 0) { + pthread_cond_wait(&g_stream_expert_pread_pool_done_cond, + &g_stream_expert_pread_pool_mutex); + } + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_dispatch( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint32_t n_workers) { + if (!ds4_gpu_stream_expert_pread_pool_begin(tasks, n_tasks, n_workers)) { + return 0; + } + return ds4_gpu_stream_expert_pread_pool_wait(); +} + +static void ds4_gpu_stream_expert_pread_pool_shutdown(void) { + if (!g_stream_expert_pread_pool_initialized) return; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_stopping = 1; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + const uint32_t n_threads = g_stream_expert_pread_pool_thread_count; + for (uint32_t i = 0; i < n_threads; i++) { + (void)pthread_join(g_stream_expert_pread_pool_threads[i], NULL); + } + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = 0; + g_stream_expert_pread_pool_active_workers = 0; + g_stream_expert_pread_pool_remaining_workers = 0; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_initialized = 0; + g_stream_expert_pread_pool_stopping = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); +} + +static int ds4_gpu_stream_expert_pread_tasks( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint64_t *total_bytes, + double *wall_ms) { + if (total_bytes) *total_bytes = 0; + if (wall_ms) *wall_ms = 0.0; + if (!tasks || n_tasks == 0) return 1; + + const uint32_t n_workers = + ds4_gpu_stream_expert_pread_thread_count(n_tasks); + const double t0 = ds4_gpu_now_ms(); + int ok = 1; + if (n_workers <= 1) { + ds4_gpu_stream_expert_pread_worker_args wa = { + .tasks = tasks, + .n_tasks = n_tasks, + .worker_index = 0, + .n_workers = 1, + }; + (void)ds4_gpu_stream_expert_pread_worker(&wa); + } else if (!ds4_gpu_stream_expert_pread_pool_dispatch(tasks, + n_tasks, + n_workers)) { + pthread_t threads[18]; + ds4_gpu_stream_expert_pread_worker_args args[18]; + uint32_t started = 0; + for (uint32_t i = 0; i < n_workers; i++) { + args[i].tasks = tasks; + args[i].n_tasks = n_tasks; + args[i].worker_index = i; + args[i].n_workers = n_workers; + const int rc = pthread_create(&threads[i], + NULL, + ds4_gpu_stream_expert_pread_worker, + &args[i]); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming expert pread thread creation failed: %s\n", + strerror(rc)); + ok = 0; + break; + } + started++; + } + for (uint32_t i = 0; i < started; i++) { + if (pthread_join(threads[i], NULL) != 0) ok = 0; + } + } + const double dt = ds4_gpu_now_ms() - t0; + + uint64_t bytes = 0; + for (uint32_t i = 0; i < n_tasks; i++) { + if (!tasks[i].ok) ok = 0; + if (bytes > UINT64_MAX - tasks[i].read_bytes) { + bytes = UINT64_MAX; + } else { + bytes += tasks[i].read_bytes; + } + } + if (total_bytes) *total_bytes = bytes; + if (wall_ms) *wall_ms = dt; + return ok; +} + +static id ds4_gpu_stream_expert_alloc_buffer( + uint64_t len, + NSString *label) { + if (!g_device || + len == 0 || + len > (uint64_t)NSUIntegerMax) { + return nil; + } + + id buffer = [g_device newBufferWithLength:(NSUInteger)len + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal streaming expert explicit buffer allocation failed (%.2f MiB)\n", + ds4_gpu_mib(len)); + return nil; + } + buffer.label = label; + g_stream_expert_cache_buffer_allocs++; + return buffer; +} + +static int ds4_gpu_stream_expert_combined_buffer_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL; +} + +static int ds4_gpu_stream_expert_slab_enabled(void) { + return ds4_gpu_stream_expert_combined_buffer_enabled() && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_SLABS") == NULL; +} + +/* + * Large PRO caches otherwise create thousands of small shared Metal buffers. + * Slabs keep the buffer object set small while locking pages only for slots + * that actually hold a streamed expert. + */ +static uint64_t ds4_gpu_stream_expert_slab_target_bytes(void) { + const uint64_t mib = 1024ull * 1024ull; + uint64_t target = 4096ull * mib; + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_SLAB_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env && *end == '\0' && v != 0) { + target = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; + } + } + return target; +} + +static id ds4_gpu_stream_expert_alloc_slab_buffer( + uint64_t len, + NSString *label) { + if (!g_device || + len == 0 || + len > (uint64_t)NSUIntegerMax) { + return nil; + } + + id buffer = [g_device newBufferWithLength:(NSUInteger)len + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal streaming expert slab allocation failed (%.2f MiB)\n", + ds4_gpu_mib(len)); + return nil; + } + buffer.label = label; + g_stream_expert_cache_buffer_allocs++; + return buffer; +} + +static int ds4_gpu_stream_expert_slab_slot_range( + uint32_t slot, + uint32_t *slab_index, + uint64_t *slot_base) { + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + const uint32_t start = g_stream_expert_cache_slab_start_slot[i]; + const uint32_t count = g_stream_expert_cache_slab_slot_count[i]; + if (slot < start || slot >= start + count) continue; + if (slab_index) *slab_index = i; + if (slot_base) { + *slot_base = + (uint64_t)(slot - start) * g_stream_expert_cache_slab_slot_bytes; + } + return 1; + } + return 0; +} + +static int ds4_gpu_stream_expert_slab_slot_for_buffer( + id buffer, + NSUInteger gate_inner, + uint32_t *slot_out) { + if (!buffer || g_stream_expert_cache_slab_slot_bytes == 0 || !slot_out) { + return 0; + } + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + if (g_stream_expert_cache_slabs[i] != buffer) continue; + const uint64_t inner = (uint64_t)gate_inner; + const uint64_t slot_bytes = g_stream_expert_cache_slab_slot_bytes; + if (slot_bytes == 0 || inner % slot_bytes != 0) return 0; + const uint64_t local_slot = inner / slot_bytes; + if (local_slot >= g_stream_expert_cache_slab_slot_count[i]) return 0; + const uint64_t slot = + (uint64_t)g_stream_expert_cache_slab_start_slot[i] + local_slot; + if (slot > UINT32_MAX) return 0; + *slot_out = (uint32_t)slot; + return 1; + } + return 0; +} + +static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) { + if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || + g_stream_expert_cache_free_slot_count >= + DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + return; + } + g_stream_expert_cache_free_slots[g_stream_expert_cache_free_slot_count++] = + slot; +} + +static int ds4_gpu_stream_expert_slab_slot_buffers( + uint32_t slot, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + uint32_t slab = UINT32_MAX; + uint64_t base = 0; + if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) || + slab >= g_stream_expert_cache_slab_count || + !g_stream_expert_cache_slabs[slab]) { + return 0; + } + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull || + base > UINT64_MAX - (gate_expert_bytes * 2ull + down_expert_bytes) || + base + gate_expert_bytes * 2ull + down_expert_bytes > + (uint64_t)NSUIntegerMax) { + return 0; + } + id b = g_stream_expert_cache_slabs[slab]; + *gate_buf = b; + *up_buf = b; + *down_buf = b; + *gate_inner = (NSUInteger)base; + *up_inner = (NSUInteger)(base + gate_expert_bytes); + *down_inner = (NSUInteger)(base + gate_expert_bytes * 2ull); + return 1; +} + +static int ds4_gpu_stream_expert_alloc_slab_slot( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + if (!ds4_gpu_stream_expert_slab_enabled() || + !gate_buf || !up_buf || !down_buf || + !gate_inner || !up_inner || !down_inner || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + + uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; + const uint64_t page = (uint64_t)getpagesize(); + if (page != 0) { + slot_bytes = round_up_u64(slot_bytes, page); + if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; + } + if (g_stream_expert_cache_slab_slot_bytes != 0 && + g_stream_expert_cache_slab_slot_bytes != slot_bytes) { + return 0; + } + g_stream_expert_cache_slab_slot_bytes = slot_bytes; + + if (g_stream_expert_cache_free_slot_count != 0) { + const uint32_t slot = + g_stream_expert_cache_free_slots[--g_stream_expert_cache_free_slot_count]; + return ds4_gpu_stream_expert_slab_slot_buffers(slot, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); + } + + uint32_t slab = g_stream_expert_cache_slab_count; + if (slab != 0 && + g_stream_expert_cache_slab_slots_used[slab - 1] < + g_stream_expert_cache_slab_slot_count[slab - 1]) { + slab--; + } else { + if (g_stream_expert_cache_slab_count >= + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS) { + return 0; + } + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget != 0 && g_stream_expert_cache_slab_total_slots >= budget) { + return 0; + } + uint64_t target = ds4_gpu_stream_expert_slab_target_bytes(); + uint64_t slots64 = target / slot_bytes; + if (slots64 == 0) slots64 = 1; + if (slots64 > UINT32_MAX) slots64 = UINT32_MAX; + uint32_t slots = (uint32_t)slots64; + if (budget != 0) { + const uint32_t remaining = + budget - g_stream_expert_cache_slab_total_slots; + if (slots > remaining) slots = remaining; + } + if (slots == 0) return 0; + id slab_buffer = nil; + while (slots != 0) { + if ((uint64_t)slots <= UINT64_MAX / slot_bytes && + (uint64_t)slots * slot_bytes <= (uint64_t)NSUIntegerMax) { + slab_buffer = + ds4_gpu_stream_expert_alloc_slab_buffer( + (uint64_t)slots * slot_bytes, + @"ds4_stream_expert_slab"); + if (slab_buffer) break; + } + slots /= 2u; + } + if (!slab_buffer || slots == 0) return 0; + + slab = g_stream_expert_cache_slab_count++; + g_stream_expert_cache_slabs[slab] = slab_buffer; + g_stream_expert_cache_slab_start_slot[slab] = + g_stream_expert_cache_slab_total_slots; + g_stream_expert_cache_slab_slot_count[slab] = slots; + g_stream_expert_cache_slab_slots_used[slab] = 0; + g_stream_expert_cache_slab_total_slots += slots; + } + + const uint32_t local_slot = g_stream_expert_cache_slab_slots_used[slab]++; + const uint32_t slot = + g_stream_expert_cache_slab_start_slot[slab] + local_slot; + return ds4_gpu_stream_expert_slab_slot_buffers(slot, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); +} + +static uint64_t ds4_gpu_stream_expert_buffer_object_count( + id gate, + id up, + id down) { + if (!gate || !up || !down) return 0; + if (gate == up && gate == down) return 1; + if (gate == up || gate == down || up == down) return 2; + return 3; +} + +static int ds4_gpu_stream_expert_evict_dontneed_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_EVICT_DONTNEED") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EVICT_DONTNEED") == NULL; +} + +static void ds4_gpu_stream_expert_evict_dontneed_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len) { + if (!ds4_gpu_stream_expert_evict_dontneed_enabled() || + !model_map || + model_size == 0 || + offset > model_size || + len == 0 || + len > model_size - offset) { + return; + } + +#if defined(POSIX_MADV_DONTNEED) + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t page_offset = offset & ~(page - 1); + const uint64_t leading = offset - page_offset; + if (len > UINT64_MAX - leading || + leading + len > UINT64_MAX - (page - 1)) { + return; + } + uint64_t advise_bytes = round_up_u64(leading + len, page); + if (advise_bytes > model_size - page_offset) { + advise_bytes = model_size - page_offset; + } + if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) return; + + const uintptr_t base = (uintptr_t)model_map; + if (page_offset > (uint64_t)(UINTPTR_MAX - base)) return; + void *addr = (void *)(base + (uintptr_t)page_offset); + const int rc = posix_madvise(addr, (size_t)advise_bytes, POSIX_MADV_DONTNEED); + if (rc == 0) { + if (g_stream_expert_cache_evict_advise_bytes > UINT64_MAX - advise_bytes) { + g_stream_expert_cache_evict_advise_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_evict_advise_bytes += advise_bytes; + } + } else if (getenv("DS4_METAL_STREAMING_EXPERT_EVICT_DONTNEED_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert evict DONTNEED failed offset=%.2f GiB len=%.2f MiB: %s\n", + ds4_gpu_gib(offset), + ds4_gpu_mib(len), + strerror(rc)); + } +#else + (void)model_map; + (void)model_size; + (void)offset; + (void)len; +#endif +} + +static int ds4_gpu_stream_expert_split_requested(void) { + return g_ssd_streaming_mode; +} + +static uint32_t ds4_gpu_stream_expert_split_min_decode_tokens(void) { + return 4; +} + +static uint32_t ds4_gpu_stream_expert_split_min_cached(void) { + uint32_t min_cached = 1024; + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget != 0 && budget < min_cached * 2u) { + min_cached = budget / 2u; + } + return min_cached; +} + +static int ds4_gpu_stream_expert_split_ready(void) { + if (!ds4_gpu_stream_expert_split_requested()) return 0; + if (g_stream_expert_cache_decode_tokens < + ds4_gpu_stream_expert_split_min_decode_tokens()) { + return 0; + } + return g_stream_expert_cache_entry_count >= + ds4_gpu_stream_expert_split_min_cached(); +} + +static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) { + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + g_stream_expert_cache_route_hotness[layer][expert] >>= 1; + } + } +} + +void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { + memset(g_stream_expert_cache_route_hotness, + 0, + sizeof(g_stream_expert_cache_route_hotness)); + g_stream_expert_cache_hotness_decay_token = + g_stream_expert_cache_decode_tokens; +} + +static void ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(void) { + if (g_stream_expert_cache_decode_tokens == 0) return; + if (g_stream_expert_cache_hotness_decay_token == 0) { + g_stream_expert_cache_hotness_decay_token = + g_stream_expert_cache_decode_tokens; + return; + } + while (g_stream_expert_cache_decode_tokens - + g_stream_expert_cache_hotness_decay_token >= + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS) { + ds4_gpu_stream_expert_cache_decay_route_hotness(); + g_stream_expert_cache_hotness_decay_token += + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS; + } +} + +static void ds4_gpu_stream_expert_cache_note_route_hotness( + uint32_t layer, + uint32_t expert, + uint32_t amount) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + amount == 0) { + return; + } + uint32_t *hotness = &g_stream_expert_cache_route_hotness[layer][expert]; + if (*hotness > UINT32_MAX - amount) { + *hotness = UINT32_MAX; + } else { + *hotness += amount; + } +} + +static void ds4_gpu_stream_expert_cache_note_selected_hotness( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0) { + return; + } + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || + selected_ids[i] >= + (int32_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + continue; + } + ds4_gpu_stream_expert_cache_note_route_hotness( + layer, + (uint32_t)selected_ids[i], + 1); + } +} + +static void ds4_gpu_stream_expert_cache_note_frequency_hotness( + uint32_t layer, + const uint32_t *frequency, + uint32_t n_total_expert) { + if (!frequency || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + return; + } + if (n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + n_total_expert = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + } + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + for (uint32_t expert = 0; expert < n_total_expert; expert++) { + ds4_gpu_stream_expert_cache_note_route_hotness(layer, + expert, + frequency[expert]); + } +} + +static void ds4_gpu_stream_expert_cache_note_token(uint32_t layer_index) { + if (!g_ssd_streaming_mode || layer_index != 0 || + g_stream_expert_cache_decode_tokens == UINT64_MAX) { + return; + } + g_stream_expert_cache_decode_tokens++; + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); +} + +static void ds4_gpu_stream_expert_cache_note_decode_token(void) { + if (!g_ssd_streaming_mode || + g_stream_expert_cache_decode_tokens == UINT64_MAX) { + return; + } + g_stream_expert_cache_decode_tokens++; + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); +} + +static int ds4_gpu_stream_compact_addr_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_COMPACT_ADDR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_COMPACT_ADDR") == NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_addr_table_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + g_stream_prefill_batch_selected_addr_building || + g_glm_stream_expert_addr_table_building || + (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL) || + ds4_gpu_stream_expert_split_requested()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_addr_table_kernel_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + ds4_gpu_stream_expert_split_ready()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_masked_addr_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + ds4_gpu_stream_expert_split_ready()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_MASKED_ADDR") == NULL; +} + +static int ds4_gpu_stream_expert_hit_validator_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_HIT_VALIDATOR") == NULL; +} + +static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_max( + uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } + if (n_total_expert == 384) return 800u; + if (n_total_expert == 256) return 760u; + return 0; +} + +static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min( + uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } + if (n_total_expert == 384 || n_total_expert == 256) return 2u; + return 0; +} + +static int ds4_gpu_stream_prefill_batch_selected_addr_enabled( + uint32_t n_tokens, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t gate_type, + uint32_t down_type) { + if (!g_ssd_streaming_mode || + n_tokens <= 1 || + n_total_expert == 0 || + n_expert != 6 || + gate_type != DS4_METAL_TENSOR_IQ2_XXS || + down_type != DS4_METAL_TENSOR_Q2_K || + g_quality_mode || + getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) { + return 0; + } + /* All unique experts for one layer must fit simultaneously because the + * address-table kernels consume them in one dispatch. Once the global + * cache fills, preparation reuses entries owned by other layers. */ + if (ds4_gpu_stream_expert_cache_configured_count() < n_total_expert) { + return 0; + } + if (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL) { + return 1; + } + const uint32_t max_tokens = + ds4_gpu_stream_prefill_batch_selected_addr_auto_max(n_total_expert); + const uint32_t min_tokens = + ds4_gpu_stream_prefill_batch_selected_addr_auto_min(n_total_expert); + return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens; +} + +static int ds4_gpu_glm_streaming_prefill_full_layer_active(void) { + return g_glm_streaming_prefill_full_layer_runtime || + getenv("DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER") != NULL; +} + +static int ds4_gpu_stream_full_expert_addr_table_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") == NULL; +} + +static uint64_t ds4_gpu_buffer_address(id buffer, NSUInteger inner) { + if (!buffer) return 0; +#if TARGET_OS_OSX + if (@available(macOS 13.0, *)) { + return (uint64_t)[buffer gpuAddress] + (uint64_t)inner; + } +#endif + return 0; +} + +static int ds4_gpu_stream_compact_addr_ensure_buffers(uint32_t layer) { + if (!g_device || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 0; + + const NSUInteger addr_bytes = 6u * sizeof(uint64_t); + const NSUInteger ids_bytes = 6u * sizeof(int32_t); + for (uint32_t i = 0; i < 4; i++) { + id current = nil; + NSUInteger bytes = addr_bytes; + NSString *label = @"ds4_stream_compact_gate_addresses"; + switch (i) { + case 0: + current = g_stream_compact_gate_addr_buffers[layer]; + label = @"ds4_stream_compact_gate_addresses"; + break; + case 1: + current = g_stream_compact_up_addr_buffers[layer]; + label = @"ds4_stream_compact_up_addresses"; + break; + case 2: + current = g_stream_compact_down_addr_buffers[layer]; + label = @"ds4_stream_compact_down_addresses"; + break; + default: + current = g_stream_compact_selected_buffers[layer]; + bytes = ids_bytes; + label = @"ds4_stream_compact_selected_ids"; + break; + } + if (current) continue; + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming compact address buffer allocation failed\n"); + return 0; + } + b.label = label; + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + switch (i) { + case 0: + g_stream_compact_gate_addr_buffers[layer] = b; + break; + case 1: + g_stream_compact_up_addr_buffers[layer] = b; + break; + case 2: + g_stream_compact_down_addr_buffers[layer] = b; + break; + default: + g_stream_compact_selected_buffers[layer] = b; + break; + } + } + return 1; +} + +static int ds4_gpu_stream_compact_addr_prepare( + uint32_t layer, + ds4_gpu_stream_expert_cache_entry * const entries[6], + uint32_t n_entries, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + id *selected_ids) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + !entries || !gate_addrs || !up_addrs || !down_addrs || !selected_ids || + n_entries == 0 || n_entries > 6) { + return 0; + } + if (!ds4_gpu_stream_compact_addr_ensure_buffers(layer)) return 0; + + uint64_t gate_values[6] = {0, 0, 0, 0, 0, 0}; + uint64_t up_values[6] = {0, 0, 0, 0, 0, 0}; + uint64_t down_values[6] = {0, 0, 0, 0, 0, 0}; + int32_t slot_ids[6] = {0, 1, 2, 3, 4, 5}; + + for (uint32_t i = 0; i < n_entries; i++) { + ds4_gpu_stream_expert_cache_entry *e = entries[i]; + if (!e || !e->gate_buffer || !e->up_buffer || !e->down_buffer) { + return 0; + } + gate_values[i] = ds4_gpu_buffer_address(e->gate_buffer, e->gate_inner); + up_values[i] = ds4_gpu_buffer_address(e->up_buffer, e->up_inner); + down_values[i] = ds4_gpu_buffer_address(e->down_buffer, e->down_inner); + if (gate_values[i] == 0 || up_values[i] == 0 || down_values[i] == 0) { + fprintf(stderr, "ds4: Metal streaming compact address path requires GPU addresses\n"); + return 0; + } + } + + const NSUInteger addr_bytes = 6u * sizeof(uint64_t); + const NSUInteger ids_bytes = 6u * sizeof(int32_t); + id gb = g_stream_compact_gate_addr_buffers[layer]; + id ub = g_stream_compact_up_addr_buffers[layer]; + id db = g_stream_compact_down_addr_buffers[layer]; + id ib = g_stream_compact_selected_buffers[layer]; + memcpy([gb contents], gate_values, addr_bytes); + memcpy([ub contents], up_values, addr_bytes); + memcpy([db contents], down_values, addr_bytes); + memcpy([ib contents], slot_ids, ids_bytes); + [gb didModifyRange:NSMakeRange(0, addr_bytes)]; + [ub didModifyRange:NSMakeRange(0, addr_bytes)]; + [db didModifyRange:NSMakeRange(0, addr_bytes)]; + [ib didModifyRange:NSMakeRange(0, ids_bytes)]; + + *gate_addrs = gb; + *up_addrs = ub; + *down_addrs = db; + *selected_ids = ib; + return 1; +} + +static int ds4_gpu_stream_selected_ids_prepare( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected, + id *selected_buf, + NSUInteger *selected_off) { + if (!g_device || + !selected_ids || + !selected_buf || + !selected_off || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED) { + return 0; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_MAX_ROUTED_EXPERT_USED * sizeof(int32_t); + id b = g_stream_selected_id_buffers[layer]; + if (!b) { + b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming selected-id buffer allocation failed\n"); + return 0; + } + b.label = @"ds4_stream_selected_ids"; + g_stream_selected_id_buffers[layer] = b; + } + + int32_t ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = {0}; + memcpy(ids, selected_ids, (size_t)n_selected * sizeof(ids[0])); + memcpy([b contents], ids, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + + *selected_buf = b; + *selected_off = 0; + return 1; +} + +static int ds4_gpu_stream_expert_cache_ensure_addr_buffers(uint32_t layer) { + if (!g_device || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + return 0; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + for (uint32_t i = 0; i < 3; i++) { + if (buffers[i]) continue; + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming expert address table allocation failed\n"); + return 0; + } + b.label = + i == 0 ? @"ds4_stream_expert_gate_addresses" : + (i == 1 ? @"ds4_stream_expert_up_addresses" : + @"ds4_stream_expert_down_addresses"); + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + if (i == 0) { + g_stream_expert_cache_gate_addr_buffers[layer] = b; + } else if (i == 1) { + g_stream_expert_cache_up_addr_buffers[layer] = b; + } else { + g_stream_expert_cache_down_addr_buffers[layer] = b; + } + } + return 1; +} + +static void ds4_gpu_stream_expert_cache_zero_addr_slot(uint32_t layer, uint32_t expert) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + if (!buffers[i]) continue; + uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); + *addr = 0; + [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; + } +} + +static int ds4_gpu_stream_expert_cache_set_addr_slot_raw( + uint32_t layer, + uint32_t expert, + id gate_buffer, + NSUInteger gate_inner, + id up_buffer, + NSUInteger up_inner, + id down_buffer, + NSUInteger down_inner) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; + + const uint64_t values[3] = { + ds4_gpu_buffer_address(gate_buffer, gate_inner), + ds4_gpu_buffer_address(up_buffer, up_inner), + ds4_gpu_buffer_address(down_buffer, down_inner), + }; + if (values[0] == 0 || values[1] == 0 || values[2] == 0) { + fprintf(stderr, "ds4: Metal streaming expert address table requires GPU addresses\n"); + return 0; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); + *addr = values[i]; + [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; + } + return 1; +} + +static int ds4_gpu_stream_expert_cache_set_addr_slot( + uint32_t layer, + uint32_t expert, + id gate_buffer, + NSUInteger gate_inner, + id up_buffer, + NSUInteger up_inner, + id down_buffer, + NSUInteger down_inner) { + if (!ds4_gpu_stream_expert_addr_table_requested()) return 1; + return ds4_gpu_stream_expert_cache_set_addr_slot_raw(layer, + expert, + gate_buffer, + gate_inner, + up_buffer, + up_inner, + down_buffer, + down_inner); +} + +static int ds4_gpu_stream_expert_cache_addr_buffers( + uint32_t layer, + id *gate, + id *up, + id *down) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; + if (gate) *gate = g_stream_expert_cache_gate_addr_buffers[layer]; + if (up) *up = g_stream_expert_cache_up_addr_buffers[layer]; + if (down) *down = g_stream_expert_cache_down_addr_buffers[layer]; + return g_stream_expert_cache_gate_addr_buffers[layer] && + g_stream_expert_cache_up_addr_buffers[layer] && + g_stream_expert_cache_down_addr_buffers[layer]; +} + +static void ds4_gpu_stream_full_expert_addr_clear_layer(uint32_t layer) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + ds4_gpu_stream_expert_cache_entry *e = &g_stream_full_expert_addr_entry[layer]; + e->gate_buffer = nil; + e->up_buffer = nil; + e->down_buffer = nil; + e->model_map = NULL; + e->model_size = 0; + e->gate_abs_offset = 0; + e->up_abs_offset = 0; + e->down_abs_offset = 0; + e->gate_expert_bytes = 0; + e->down_expert_bytes = 0; + e->logical_bytes = 0; + e->last_used = 0; + e->use_count = 0; + e->gate_inner = 0; + e->up_inner = 0; + e->down_inner = 0; + e->slab_slot = 0; + e->valid = 0; + e->slab_backed = 0; +} + +static int ds4_gpu_stream_full_expert_addr_table_prepare( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + ds4_gpu_stream_expert_cache_entry **entry_out) { + if (!ds4_gpu_stream_full_expert_addr_table_requested()) return 0; + if (!model_map || model_size == 0 || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + gate_expert_bytes == 0 || + down_expert_bytes == 0 || + n_total_expert > UINT64_MAX / gate_expert_bytes || + n_total_expert > UINT64_MAX / down_expert_bytes) { + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; + if (gate_abs_offset > model_size || + up_abs_offset > model_size || + down_abs_offset > model_size || + gate_tensor_bytes > model_size - gate_abs_offset || + gate_tensor_bytes > model_size - up_abs_offset || + down_tensor_bytes > model_size - down_abs_offset) { + return 0; + } + + ds4_gpu_stream_expert_cache_entry *entry = + &g_stream_full_expert_addr_entry[layer]; + if (!entry->valid || + entry->model_map != model_map || + entry->model_size != model_size || + entry->gate_abs_offset != gate_abs_offset || + entry->up_abs_offset != up_abs_offset || + entry->down_abs_offset != down_abs_offset || + entry->gate_expert_bytes != gate_expert_bytes || + entry->down_expert_bytes != down_expert_bytes || + !entry->gate_buffer || + !entry->up_buffer || + !entry->down_buffer) { + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + gate_abs_offset, + gate_tensor_bytes, + &gate_inner); + id up_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + up_abs_offset, + gate_tensor_bytes, + &up_inner); + id down_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + down_abs_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) return 0; + + entry->gate_buffer = gate_buf; + entry->up_buffer = up_buf; + entry->down_buffer = down_buf; + entry->model_map = model_map; + entry->model_size = model_size; + entry->gate_abs_offset = gate_abs_offset; + entry->up_abs_offset = up_abs_offset; + entry->down_abs_offset = down_abs_offset; + entry->gate_expert_bytes = gate_expert_bytes; + entry->down_expert_bytes = down_expert_bytes; + entry->logical_bytes = gate_tensor_bytes * 2ull + down_tensor_bytes; + entry->gate_inner = (NSUInteger)gate_inner; + entry->up_inner = (NSUInteger)up_inner; + entry->down_inner = (NSUInteger)down_inner; + entry->valid = 1; + + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + return 0; + } + + const uint64_t gate_base = ds4_gpu_buffer_address(entry->gate_buffer, + entry->gate_inner); + const uint64_t up_base = ds4_gpu_buffer_address(entry->up_buffer, + entry->up_inner); + const uint64_t down_base = ds4_gpu_buffer_address(entry->down_buffer, + entry->down_inner); + if (gate_base == 0 || up_base == 0 || down_base == 0) { + fprintf(stderr, "ds4: Metal full streaming expert address table requires GPU addresses\n"); + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + return 0; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + uint64_t *gate_addr = (uint64_t *)[buffers[0] contents]; + uint64_t *up_addr = (uint64_t *)[buffers[1] contents]; + uint64_t *down_addr = (uint64_t *)[buffers[2] contents]; + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + if (expert < n_total_expert) { + gate_addr[expert] = gate_base + (uint64_t)expert * gate_expert_bytes; + up_addr[expert] = up_base + (uint64_t)expert * gate_expert_bytes; + down_addr[expert] = down_base + (uint64_t)expert * down_expert_bytes; + } else { + gate_addr[expert] = 0; + up_addr[expert] = 0; + down_addr[expert] = 0; + } + } + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; + } + } + + if (gate_addrs) *gate_addrs = g_stream_expert_cache_gate_addr_buffers[layer]; + if (up_addrs) *up_addrs = g_stream_expert_cache_up_addr_buffers[layer]; + if (down_addrs) *down_addrs = g_stream_expert_cache_down_addr_buffers[layer]; + if (entry_out) *entry_out = entry; + return entry->valid && + g_stream_expert_cache_gate_addr_buffers[layer] && + g_stream_expert_cache_up_addr_buffers[layer] && + g_stream_expert_cache_down_addr_buffers[layer]; +} + +static id ds4_gpu_stream_expert_validate_status_buffer(void) { + if (!g_device) return nil; + if (g_stream_expert_validate_status_buffer) { + return g_stream_expert_validate_status_buffer; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming expert validator allocation failed\n"); + return nil; + } + b.label = @"ds4_stream_expert_validate_status"; + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + g_stream_expert_validate_status_buffer = b; + return b; +} + +static int ds4_gpu_encode_stream_expert_cache_validate( + id cb, + const ds4_gpu_stream_expert_validate_args *args, + id selected, + NSUInteger selected_off, + id gate_addrs, + id up_addrs, + id down_addrs, + id status) { + if (!cb || !args || !selected || !gate_addrs || !up_addrs || !down_addrs || + !status || !g_moe_stream_expert_cache_validate_pipeline || + args->n_total_expert == 0 || args->n_total_expert > 384 || + args->n_expert == 0 || args->n_expert > 6) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_stream_expert_cache_validate_pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:selected offset:selected_off atIndex:1]; + [enc setBuffer:gate_addrs offset:0 atIndex:2]; + [enc setBuffer:up_addrs offset:0 atIndex:3]; + [enc setBuffer:down_addrs offset:0 atIndex:4]; + [enc setBuffer:status offset:0 atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_stream_expert_cache_validate_selected( + const ds4_gpu_tensor *selected, + id gate_addrs, + id up_addrs, + id down_addrs, + uint32_t n_total_expert, + uint32_t n_expert, + int32_t selected_ids[6], + uint32_t *all_cached, + uint32_t *miss_mask, + uint32_t *invalid_mask) { + if (!selected || !selected_ids || !all_cached || !miss_mask || !invalid_mask) { + return 0; + } + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id status = ds4_gpu_stream_expert_validate_status_buffer(); + if (!selectedbuf || !status) return 0; + + const NSUInteger status_bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); + memset([status contents], 0, status_bytes); + [status didModifyRange:NSMakeRange(0, status_bytes)]; + + ds4_gpu_stream_expert_validate_args args = { + .n_total_expert = n_total_expert, + .n_expert = n_expert, + }; + + const int had_batch = g_batch_cb != nil; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_stream_expert_cache_validate(cb, + &args, + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_addrs, + up_addrs, + down_addrs, + status)) { + return 0; + } + + if (had_batch) { + if (ds4_gpu_end_commands() == 0) return 0; + } else if (!ds4_gpu_finish_command_buffer(cb, owned, + "streaming expert cache validator")) { + return 0; + } + + const uint32_t *words = (const uint32_t *)[status contents]; + *all_cached = words[0]; + *miss_mask = words[1]; + *invalid_mask = words[2]; + for (uint32_t i = 0; i < 6; i++) { + selected_ids[i] = (int32_t)words[4 + i]; + } + + if (had_batch && ds4_gpu_begin_commands() == 0) return 0; + return 1; +} + +static void ds4_gpu_stream_expert_cache_clear_entry_internal( + uint32_t layer, + uint32_t expert, + int count_eviction, + int recycle_slab_slot, + ds4_gpu_stream_expert_reusable_buffers *reuse) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return; + } + if (reuse) { + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + } + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid) return; + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + return; + } + + const uint64_t bytes = e->logical_bytes; + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->gate_abs_offset, + e->gate_expert_bytes); + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->up_abs_offset, + e->gate_expert_bytes); + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->down_abs_offset, + e->down_expert_bytes); + ds4_gpu_stream_expert_cache_zero_addr_slot(layer, expert); + if (reuse) { + reuse->gate_buffer = e->gate_buffer; + reuse->up_buffer = e->up_buffer; + reuse->down_buffer = e->down_buffer; + reuse->gate_inner = e->gate_inner; + reuse->up_inner = e->up_inner; + reuse->down_inner = e->down_inner; + } else if (e->slab_backed && recycle_slab_slot) { + ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot); + } + e->gate_buffer = nil; + e->up_buffer = nil; + e->down_buffer = nil; + e->model_map = NULL; + e->model_size = 0; + e->gate_abs_offset = 0; + e->up_abs_offset = 0; + e->down_abs_offset = 0; + e->gate_expert_bytes = 0; + e->down_expert_bytes = 0; + e->logical_bytes = 0; + e->last_used = 0; + e->use_count = 0; + e->gate_inner = 0; + e->up_inner = 0; + e->down_inner = 0; + e->inflight_seq = 0; + e->slab_slot = 0; + e->valid = 0; + e->slab_backed = 0; + + if (g_stream_expert_cache_layer_count[layer] > 0) { + g_stream_expert_cache_layer_count[layer]--; + } + if (g_stream_expert_cache_entry_count > 0) { + g_stream_expert_cache_entry_count--; + } + if (g_stream_expert_cache_bytes >= bytes) { + g_stream_expert_cache_bytes -= bytes; + } else { + g_stream_expert_cache_bytes = 0; + } + if (count_eviction) { + g_stream_expert_cache_evictions++; + g_stream_expert_cache_layer_evictions[layer]++; + } +} + +static void ds4_gpu_stream_expert_cache_clear_entry( + uint32_t layer, + uint32_t expert, + int count_eviction) { + ds4_gpu_stream_expert_cache_clear_entry_internal(layer, + expert, + count_eviction, + 1, + NULL); +} + +static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) { + ds4_gpu_stream_expert_pending_load_clear(); + g_stream_expert_cache_done_seq = g_stream_expert_cache_cb_seq; + g_stream_expert_cache_batch_seq = 0; + g_stream_expert_cache_owned_seq = 0; + g_stream_expert_cache_pending_max_seq = 0; + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + g_stream_expert_cache[layer][expert].inflight_seq = 0; + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + g_stream_expert_cache_layer_count[layer] = 0; + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + for (uint32_t i = 0; i < 3; i++) { + if (!buffers[i]) continue; + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + memset([buffers[i] contents], 0, bytes); + [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; + } + } + g_stream_expert_cache_bytes = 0; + g_stream_expert_cache_entry_count = 0; + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + g_stream_expert_cache_slabs[i] = nil; + g_stream_expert_cache_slab_start_slot[i] = 0; + g_stream_expert_cache_slab_slot_count[i] = 0; + g_stream_expert_cache_slab_slots_used[i] = 0; + } + g_stream_expert_cache_slab_count = 0; + g_stream_expert_cache_slab_total_slots = 0; + g_stream_expert_cache_free_slot_count = 0; + g_stream_expert_cache_slab_slot_bytes = 0; + if (reset_stats) { + g_stream_expert_cache_hits = 0; + g_stream_expert_cache_misses = 0; + g_stream_expert_cache_evictions = 0; + g_stream_expert_cache_wraps = 0; + g_stream_expert_cache_clock = 0; + g_stream_expert_cache_evict_advise_bytes = 0; + g_stream_expert_cache_willneed_advise_bytes = 0; + g_stream_expert_cache_pread_bytes = 0; + g_stream_expert_cache_pread_ms = 0.0; + g_stream_expert_cache_buffer_allocs = 0; + g_stream_expert_cache_buffer_reuses = 0; + g_stream_expert_cache_decode_tokens = 0; + g_stream_expert_cache_hotness_decay_token = 0; + memset(g_stream_expert_cache_route_hotness, + 0, + sizeof(g_stream_expert_cache_route_hotness)); + g_stream_expert_timing_selected_calls = 0; + g_stream_expert_timing_selected_read_ms = 0.0; + g_stream_expert_timing_selected_sync_ms = 0.0; + g_stream_expert_timing_selected_copy_ms = 0.0; + g_stream_expert_timing_selected_bind_ms = 0.0; + g_stream_expert_timing_split_layers = 0; + g_stream_expert_timing_split_resident_experts = 0; + g_stream_expert_timing_split_missing_experts = 0; + g_stream_expert_timing_split_resident_ms = 0.0; + g_stream_expert_timing_split_missing_ms = 0.0; + g_stream_expert_timing_split_missing_load_ms = 0.0; + g_stream_expert_timing_split_missing_slot_ms = 0.0; + g_stream_expert_timing_split_missing_prune_ms = 0.0; + g_stream_expert_timing_split_missing_addr_ms = 0.0; + g_stream_expert_timing_split_missing_wait_ms = 0.0; + g_stream_expert_timing_load_calls = 0; + g_stream_expert_timing_load_prepare_ms = 0.0; + g_stream_expert_timing_load_pread_ms = 0.0; + g_stream_expert_timing_load_modify_ms = 0.0; + g_stream_expert_timing_load_install_ms = 0.0; + g_stream_expert_timing_prepare_batch_reuse_calls = 0; + g_stream_expert_timing_prepare_batch_reuse_ms = 0.0; + g_stream_expert_timing_prepare_buffer_calls = 0; + g_stream_expert_timing_prepare_buffer_ms = 0.0; + g_stream_expert_timing_prepare_task_experts = 0; + g_stream_expert_timing_prepare_task_ms = 0.0; + g_stream_expert_timing_reuse_scan_calls = 0; + g_stream_expert_timing_reuse_scan_entries = 0; + g_stream_expert_timing_reuse_scan_ms = 0.0; + g_stream_expert_timing_reuse_clear_ms = 0.0; + g_stream_expert_timing_readahead_calls = 0; + g_stream_expert_timing_readahead_bytes = 0; + g_stream_expert_timing_readahead_ms = 0.0; + g_stream_expert_timing_cache_all_resident_layers = 0; + g_stream_expert_timing_cache_all_missing_layers = 0; + g_stream_expert_timing_cache_mixed_layers = 0; + g_stream_expert_timing_cache_resident_experts = 0; + g_stream_expert_timing_cache_missing_experts = 0; + g_stream_expert_timing_last_report = + (ds4_gpu_stream_expert_timing_snapshot){0}; + memset(g_stream_expert_cache_layer_hits, + 0, + sizeof(g_stream_expert_cache_layer_hits)); + memset(g_stream_expert_cache_layer_misses, + 0, + sizeof(g_stream_expert_cache_layer_misses)); + memset(g_stream_expert_cache_layer_evictions, + 0, + sizeof(g_stream_expert_cache_layer_evictions)); + memset(g_stream_expert_cache_layer_pread_bytes, + 0, + sizeof(g_stream_expert_cache_layer_pread_bytes)); + memset(g_stream_expert_cache_layer_pread_ms, + 0, + sizeof(g_stream_expert_cache_layer_pread_ms)); + } +} + +static int ds4_gpu_stream_expert_cache_is_protected( + uint32_t expert, + const int32_t *protect_ids, + uint32_t n_protect) { + if (!protect_ids) return 0; + for (uint32_t i = 0; i < n_protect; i++) { + if (protect_ids[i] >= 0 && (uint32_t)protect_ids[i] == expert) { + return 1; + } + } + return 0; +} + +static void ds4_gpu_stream_expert_cache_prune_layer( + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected, + const int32_t *protect_ids, + uint32_t n_protect) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + uint32_t cap = ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected); + if (cap == 0) return; + + /* + * Route hotness counts selected experts even when they miss. Hit-count + * LFU penalizes experts that are repeatedly selected but evicted before a + * second hit, which keeps too many decode layers in the mixed-cache path. + */ + while (g_stream_expert_cache_layer_count[layer] > cap) { + uint32_t victim = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + for (uint32_t expert = 0; expert < n_total_expert; expert++) { + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid || + ds4_gpu_stream_expert_cache_entry_inflight(e) || + ds4_gpu_stream_expert_cache_is_protected(expert, protect_ids, n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim = expert; + } + } + if (victim == UINT32_MAX) break; + ds4_gpu_stream_expert_cache_clear_entry(layer, victim, 1); + } +} + +static int ds4_gpu_stream_expert_cache_entry_protected( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect) { + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + ds4_gpu_stream_expert_cache_entry_inflight( + &g_stream_expert_cache[layer][expert])) { + return 1; + } + return layer == protect_layer && + ds4_gpu_stream_expert_cache_is_protected(expert, + protect_ids, + n_protect); +} + +static int ds4_gpu_stream_expert_cache_entry_reusable( + const ds4_gpu_stream_expert_cache_entry *e, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return e && + e->valid && + e->gate_buffer && + e->up_buffer && + e->down_buffer && + e->gate_expert_bytes == gate_expert_bytes && + e->down_expert_bytes == down_expert_bytes; +} + +static int ds4_gpu_stream_expert_cache_take_reusable( + int force_reuse, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + ds4_gpu_stream_expert_reusable_buffers *reuse) { + if (!reuse) return 0; + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || + (!force_reuse && g_stream_expert_cache_entry_count < budget)) { + return 0; + } + + int waited_inflight = 0; +retry: + ; + uint32_t victim_layer = UINT32_MAX; + uint32_t victim_expert = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + int skipped_inflight = 0; + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t scan_entries = 0; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + scan_entries++; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + skipped_inflight = 1; + continue; + } + if (ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim_layer = layer; + victim_expert = expert; + } + } + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, + ds4_gpu_now_ms() - scan_t0); + } + + if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) { + if (skipped_inflight && !waited_inflight && + !ds4_gpu_stream_expert_cache_on_service_thread()) { + waited_inflight = 1; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache reuse")) { + return 0; + } + goto retry; + } + return 0; + } + const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; + ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer, + victim_expert, + 1, + 1, + reuse); + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - + clear_t0); + } + if (!reuse->gate_buffer || !reuse->up_buffer || !reuse->down_buffer) { + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + return 0; + } + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count(reuse->gate_buffer, + reuse->up_buffer, + reuse->down_buffer); + return 1; +} + +static uint32_t ds4_gpu_stream_expert_cache_take_reusable_batch( + uint32_t n_needed, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + ds4_gpu_stream_expert_reusable_buffers *reuses) { + if (!reuses || n_needed == 0) return 0; + if (n_needed > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + n_needed = DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; + } + for (uint32_t i = 0; i < n_needed; i++) { + reuses[i].gate_buffer = nil; + reuses[i].up_buffer = nil; + reuses[i].down_buffer = nil; + reuses[i].gate_inner = 0; + reuses[i].up_inner = 0; + reuses[i].down_inner = 0; + } + + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || g_stream_expert_cache_entry_count < budget) { + return 0; + } + + int waited_inflight = 0; +retry: + ; + uint32_t victim_layers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_experts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t victim_last_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_count = 0; + int skipped_inflight = 0; + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t scan_entries = 0; + for (uint32_t i = 0; i < n_needed; i++) { + victim_layers[i] = UINT32_MAX; + victim_experts[i] = UINT32_MAX; + victim_hotness[i] = UINT32_MAX; + victim_last_used[i] = UINT64_MAX; + } + + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + scan_entries++; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + skipped_inflight = 1; + continue; + } + if (ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + const uint64_t last_used = e->last_used; + if (victim_count < n_needed) { + victim_layers[victim_count] = layer; + victim_experts[victim_count] = expert; + victim_hotness[victim_count] = hotness; + victim_last_used[victim_count] = last_used; + victim_count++; + continue; + } + + uint32_t worst = 0; + for (uint32_t i = 1; i < victim_count; i++) { + if (victim_hotness[i] > victim_hotness[worst] || + (victim_hotness[i] == victim_hotness[worst] && + victim_last_used[i] > victim_last_used[worst])) { + worst = i; + } + } + if (hotness < victim_hotness[worst] || + (hotness == victim_hotness[worst] && + last_used < victim_last_used[worst])) { + victim_layers[worst] = layer; + victim_experts[worst] = expert; + victim_hotness[worst] = hotness; + victim_last_used[worst] = last_used; + } + } + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, + ds4_gpu_now_ms() - scan_t0); + } + + if (victim_count == 0) { + if (skipped_inflight && !waited_inflight) { + waited_inflight = 1; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache batch reuse")) { + return 0; + } + goto retry; + } + return 0; + } + + uint32_t reuse_count = 0; + const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; + for (uint32_t i = 0; i < victim_count; i++) { + if (victim_layers[i] == UINT32_MAX || + victim_experts[i] == UINT32_MAX) { + continue; + } + ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layers[i], + victim_experts[i], + 1, + 1, + &reuses[reuse_count]); + if (!reuses[reuse_count].gate_buffer || + !reuses[reuse_count].up_buffer || + !reuses[reuse_count].down_buffer) { + reuses[reuse_count].gate_buffer = nil; + reuses[reuse_count].up_buffer = nil; + reuses[reuse_count].down_buffer = nil; + reuses[reuse_count].gate_inner = 0; + reuses[reuse_count].up_inner = 0; + reuses[reuse_count].down_inner = 0; + continue; + } + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count( + reuses[reuse_count].gate_buffer, + reuses[reuse_count].up_buffer, + reuses[reuse_count].down_buffer); + reuse_count++; + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - + clear_t0); + } + return reuse_count; +} + +static int ds4_gpu_stream_expert_batch_reuse_enabled( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + const uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + /* + * One global victim scan per selected miss is measurable for GLM Q2-size + * slots, but batching larger Q4-size slots regressed short decode on M5. + * Keep larger slots on the older single-victim path until profiling says + * otherwise. + */ + return slot_bytes <= 16ull * 1024ull * 1024ull; +} + +static int ds4_gpu_stream_expert_cache_prepare_load_buffers( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + int force_reuse, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + if (!gate_buf || !up_buf || !down_buf || + !gate_inner || !up_inner || !down_inner || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return 0; + } + + *gate_buf = nil; + *up_buf = nil; + *down_buf = nil; + *gate_inner = 0; + *up_inner = 0; + *down_inner = 0; + + ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 }; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (e->valid && ds4_gpu_stream_expert_cache_entry_inflight(e)) { + if (ds4_gpu_stream_expert_cache_on_service_thread()) return 0; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache replacement")) { + return 0; + } + e = &g_stream_expert_cache[layer][expert]; + } + if (ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + ds4_gpu_stream_expert_cache_clear_entry_internal(layer, + expert, + 0, + 1, + &reuse); + if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count(reuse.gate_buffer, + reuse.up_buffer, + reuse.down_buffer); + } + } else if (e->valid) { + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + + if (!reuse.gate_buffer || !reuse.up_buffer || !reuse.down_buffer) { + if (!ds4_gpu_stream_expert_cache_take_reusable(force_reuse, + protect_layer, + protect_ids, + n_protect, + gate_expert_bytes, + down_expert_bytes, + &reuse)) { + reuse.gate_buffer = nil; + reuse.up_buffer = nil; + reuse.down_buffer = nil; + } + } + + if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { + *gate_buf = reuse.gate_buffer; + *up_buf = reuse.up_buffer; + *down_buf = reuse.down_buffer; + *gate_inner = reuse.gate_inner; + *up_inner = reuse.up_inner; + *down_inner = reuse.down_inner; + return 1; + } + + if (ds4_gpu_stream_expert_combined_buffer_enabled()) { + if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes || + gate_expert_bytes * 2ull > UINT64_MAX - down_expert_bytes || + gate_expert_bytes > (uint64_t)NSUIntegerMax || + gate_expert_bytes * 2ull > (uint64_t)NSUIntegerMax || + gate_expert_bytes * 2ull + down_expert_bytes > + (uint64_t)NSUIntegerMax) { + return 0; + } + if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner)) { + return 1; + } + const uint64_t up_off = gate_expert_bytes; + const uint64_t down_off = gate_expert_bytes * 2ull; + const uint64_t combined_bytes = down_off + down_expert_bytes; + id combined = + ds4_gpu_stream_expert_alloc_buffer(combined_bytes, + @"ds4_stream_expert_combined"); + if (!combined) return 0; + *gate_buf = combined; + *up_buf = combined; + *down_buf = combined; + *gate_inner = 0; + *up_inner = (NSUInteger)up_off; + *down_inner = (NSUInteger)down_off; + return 1; + } + + *gate_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_gate"); + *up_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_up"); + *down_buf = ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes, + @"ds4_stream_expert_down"); + return *gate_buf && *up_buf && *down_buf; +} + +static void ds4_gpu_stream_expert_cache_prune_global( + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect) { + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || g_stream_expert_cache_entry_count <= budget) return; + + while (g_stream_expert_cache_entry_count > budget) { + uint32_t victim_layer = UINT32_MAX; + uint32_t victim_expert = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid || + ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim_layer = layer; + victim_expert = expert; + } + } + } + if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) break; + ds4_gpu_stream_expert_cache_clear_entry(victim_layer, victim_expert, 1); + } +} + +static int ds4_gpu_stream_expert_cache_entry_matches( + const ds4_gpu_stream_expert_cache_entry *e, + const void *model_map, + uint64_t model_size, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return e && + e->valid && + e->model_map == model_map && + e->model_size == model_size && + e->gate_abs_offset == gate_abs_offset && + e->up_abs_offset == up_abs_offset && + e->down_abs_offset == down_abs_offset && + e->gate_expert_bytes == gate_expert_bytes && + e->down_expert_bytes == down_expert_bytes && + e->gate_buffer && e->up_buffer && e->down_buffer; +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_peek( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!g_ssd_streaming_mode || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + expert >= n_total_expert || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return NULL; + } + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes)) { + return NULL; + } + + e->last_used = ++g_stream_expert_cache_clock; + e->use_count++; + g_stream_expert_cache_hits++; + g_stream_expert_cache_layer_hits[layer]++; + return e; +} + +static ds4_gpu_stream_expert_cache_entry * +ds4_gpu_stream_expert_cache_install_loaded( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id gate_buf, + id up_buf, + id down_buf, + NSUInteger gate_inner, + NSUInteger up_inner, + NSUInteger down_inner) { + if (!gate_buf || !up_buf || !down_buf || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return NULL; + } + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); + return NULL; + } + const uint64_t logical_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (e->valid) { + if (ds4_gpu_stream_expert_cache_entry_inflight(e) && + ds4_gpu_stream_expert_cache_on_service_thread()) { + return NULL; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e) && + !ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache install")) { + return NULL; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) return NULL; + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + if (e->valid) return NULL; + } + + if (!ds4_gpu_stream_expert_cache_set_addr_slot(layer, + expert, + gate_buf, + gate_inner, + up_buf, + up_inner, + down_buf, + down_inner)) { + return NULL; + } + + e->gate_buffer = gate_buf; + e->up_buffer = up_buf; + e->down_buffer = down_buf; + e->model_map = model_map; + e->model_size = model_size; + e->gate_abs_offset = gate_abs_offset; + e->up_abs_offset = up_abs_offset; + e->down_abs_offset = down_abs_offset; + e->gate_expert_bytes = gate_expert_bytes; + e->down_expert_bytes = down_expert_bytes; + e->logical_bytes = logical_bytes; + e->last_used = ++g_stream_expert_cache_clock; + e->use_count = 1; + e->gate_inner = gate_inner; + e->up_inner = up_inner; + e->down_inner = down_inner; + e->inflight_seq = 0; + uint32_t slab_slot = 0; + if (gate_buf == up_buf && + gate_buf == down_buf && + ds4_gpu_stream_expert_slab_slot_for_buffer(gate_buf, + gate_inner, + &slab_slot)) { + e->slab_backed = 1; + e->slab_slot = slab_slot; + } else { + e->slab_backed = 0; + e->slab_slot = 0; + } + e->valid = 1; + g_stream_expert_cache_layer_count[layer]++; + if (g_stream_expert_cache_entry_count < UINT32_MAX) { + g_stream_expert_cache_entry_count++; + } + if (g_stream_expert_cache_bytes > UINT64_MAX - logical_bytes) { + g_stream_expert_cache_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_bytes += logical_bytes; + } + g_stream_expert_cache_misses++; + g_stream_expert_cache_layer_misses[layer]++; + g_stream_expert_cache_wraps += 3; + return e; +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get_protected( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + const int32_t *protect_ids, + uint32_t n_protect) { + if (!g_ssd_streaming_mode || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + expert >= n_total_expert || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return NULL; + } + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes)) { + e->last_used = ++g_stream_expert_cache_clock; + e->use_count++; + g_stream_expert_cache_hits++; + g_stream_expert_cache_layer_hits[layer]++; + return e; + } + + ds4_gpu_stream_expert_readahead_range(gate_abs_offset, gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(up_abs_offset, gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(down_abs_offset, down_expert_bytes); + + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + NSUInteger gate_inner = 0; + NSUInteger up_inner = 0; + NSUInteger down_inner = 0; + const int32_t protect_one = (int32_t)expert; + if (!protect_ids || n_protect == 0) { + protect_ids = &protect_one; + n_protect = 1; + } + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + const int force_reuse = + cache_budget != 0 && g_stream_expert_cache_entry_count >= cache_budget; + if (!ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + protect_ids, + n_protect, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_buf, + &up_buf, + &down_buf, + &gate_inner, + &up_inner, + &down_inner)) { + return NULL; + } + if (!gate_buf || !up_buf || !down_buf) return NULL; + + uint8_t *gate_dst = (uint8_t *)[gate_buf contents] + gate_inner; + uint8_t *up_dst = (uint8_t *)[up_buf contents] + up_inner; + uint8_t *down_dst = (uint8_t *)[down_buf contents] + down_inner; + if (!gate_dst || !up_dst || !down_dst) return NULL; + + ds4_gpu_stream_expert_pread_task tasks[3] = { + { + .offset = gate_abs_offset, + .len = gate_expert_bytes, + .dst = gate_dst, + }, + { + .offset = up_abs_offset, + .len = gate_expert_bytes, + .dst = up_dst, + }, + { + .offset = down_abs_offset, + .len = down_expert_bytes, + .dst = down_dst, + }, + }; + uint64_t read_bytes = 0; + double read_ms = 0.0; + if (!ds4_gpu_stream_expert_pread_tasks(tasks, 3, &read_bytes, &read_ms)) { + return NULL; + } + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + + [gate_buf didModifyRange:NSMakeRange(gate_inner, (NSUInteger)gate_expert_bytes)]; + [up_buf didModifyRange:NSMakeRange(up_inner, (NSUInteger)gate_expert_bytes)]; + [down_buf didModifyRange:NSMakeRange(down_inner, (NSUInteger)down_expert_bytes)]; + if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert parallel pread layer=%u experts=1 tensors=3 " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + ds4_gpu_stream_expert_pread_thread_count(3), + ds4_gpu_gib(read_bytes), + read_ms); + } + return ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return ds4_gpu_stream_expert_cache_get_protected(model_map, + model_size, + layer, + expert, + n_total_expert, + n_selected, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes, + NULL, + 0); +} + +static int ds4_gpu_stream_expert_pending_load_profile_enabled(void) { + return getenv("DS4_METAL_STREAMING_EXPERT_EARLY_LOAD_PROFILE") != NULL; +} + +static void ds4_gpu_stream_expert_pending_load_release_buffers( + ds4_gpu_stream_expert_pending_load *p) { + if (!p) return; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + p->gate_bufs[i] = nil; + p->up_bufs[i] = nil; + p->down_bufs[i] = nil; + p->gate_inners[i] = 0; + p->up_inners[i] = 0; + p->down_inners[i] = 0; + } +} + +static int ds4_gpu_stream_expert_pending_load_install( + ds4_gpu_stream_expert_pending_load *p, + ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED], + double elapsed_ms) { + if (!p || p->n_loads == 0) return 1; + + uint64_t read_bytes = 0; + int ok = 1; + for (uint32_t i = 0; i < p->n_tasks; i++) { + if (!p->tasks[i].ok) ok = 0; + if (read_bytes > UINT64_MAX - p->tasks[i].read_bytes) { + read_bytes = UINT64_MAX; + } else { + read_bytes += p->tasks[i].read_bytes; + } + } + if (!ok) return 0; + + ds4_gpu_stream_expert_cache_note_pread(p->layer, read_bytes, elapsed_ms); + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + [p->gate_bufs[load_i] didModifyRange:NSMakeRange(p->gate_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; + [p->up_bufs[load_i] didModifyRange:NSMakeRange(p->up_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; + [p->down_bufs[load_i] didModifyRange:NSMakeRange(p->down_inners[load_i], (NSUInteger)p->down_expert_bytes)]; + } + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_t0; + load_t0 = now_ms; + } + + ds4_gpu_stream_expert_cache_entry + *loaded_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + loaded_entries[i] = NULL; + } + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + const uint32_t slot = p->load_slots[load_i]; + const uint32_t expert = (uint32_t)p->selected_ids[slot]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(p->model_map, + p->model_size, + p->layer, + expert, + p->gate_abs_offsets[slot], + p->up_abs_offsets[slot], + p->down_abs_offsets[slot], + p->gate_expert_bytes, + p->down_expert_bytes, + p->gate_bufs[load_i], + p->up_bufs[load_i], + p->down_bufs[load_i], + p->gate_inners[load_i], + p->up_inners[load_i], + p->down_inners[load_i]); + if (!entry) return 0; + loaded_entries[slot] = entry; + if (entries) entries[slot] = entry; + } + for (uint32_t i = 0; i < p->n_selected; i++) { + if ((p->missing_mask & (1u << i)) == 0) continue; + if (entries && entries[i]) continue; + const uint32_t source = p->source_slots[i]; + if (source >= p->n_selected) return 0; + ds4_gpu_stream_expert_cache_entry *entry = entries && entries[source] ? + entries[source] : loaded_entries[source]; + if (!entry) return 0; + entry->use_count++; + if (entries) entries[i] = entry; + } + if (load_timing) { + load_install_ms = ds4_gpu_now_ms() - load_t0; + ds4_gpu_stream_expert_timing_note_load_detail(p->prepare_ms, + elapsed_ms, + load_modify_ms, + load_install_ms); + } + if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { + fprintf(stderr, + "ds4: Metal streaming expert early-load finish layer=%u experts=%u tensors=%u bytes=%.2f GiB wall=%.3f ms\n", + p->layer, + p->n_loads, + p->n_tasks, + ds4_gpu_gib(read_bytes), + elapsed_ms); + } + return 1; +} + +static int ds4_gpu_stream_expert_pending_load_finish( + ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]) { + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + if (!p->active) return 1; + + const double start_ms = p->start_ms; + if (!ds4_gpu_stream_expert_pread_pool_wait()) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->active = 0; + return 0; + } + const double elapsed_ms = ds4_gpu_now_ms() - start_ms; + p->active = 0; + const int ok = ds4_gpu_stream_expert_pending_load_install(p, + entries, + elapsed_ms); + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->n_tasks = 0; + p->n_loads = 0; + p->prepare_ms = 0.0; + return ok; +} + +static void ds4_gpu_stream_expert_pending_load_clear(void) { + if (!g_stream_expert_pending_load.active) { + ds4_gpu_stream_expert_pending_load_release_buffers( + &g_stream_expert_pending_load); + return; + } + (void)ds4_gpu_stream_expert_pending_load_finish(NULL); +} + +static int ds4_gpu_stream_expert_pending_load_matches( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + if (!p->active || + p->model_map != model_map || + p->model_size != model_size || + p->layer != layer || + p->n_total_expert != n_total_expert || + p->n_selected != n_selected || + p->gate_expert_bytes != gate_expert_bytes || + p->down_expert_bytes != down_expert_bytes) { + return 0; + } + for (uint32_t i = 0; i < n_selected; i++) { + if (p->selected_ids[i] != selected_ids[i]) return 0; + } + return 1; +} + +int ds4_gpu_stream_expert_cache_begin_selected_load( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!g_ssd_streaming_mode || + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EARLY_LOAD") != NULL) { + return 1; + } + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + if (ds4_gpu_stream_expert_pending_load_matches(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + n_selected, + gate_expert_bytes, + down_expert_bytes)) { + return 1; + } + + ds4_gpu_stream_expert_pending_load_clear(); + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + p->active = 0; + p->model_map = model_map; + p->model_size = model_size; + p->layer = layer; + p->n_total_expert = n_total_expert; + p->n_selected = n_selected; + p->missing_mask = 0; + p->n_loads = 0; + p->n_tasks = 0; + p->gate_expert_bytes = gate_expert_bytes; + p->down_expert_bytes = down_expert_bytes; + p->prepare_ms = 0.0; + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double load_prepare_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + p->selected_ids[i] = -1; + p->load_slots[i] = 0; + p->source_slots[i] = UINT32_MAX; + p->gate_abs_offsets[i] = 0; + p->up_abs_offsets[i] = 0; + p->down_abs_offsets[i] = 0; + p->gate_bufs[i] = nil; + p->up_bufs[i] = nil; + p->down_bufs[i] = nil; + p->gate_inners[i] = 0; + p->up_inners[i] = 0; + p->down_inners[i] = 0; + } + memset(p->tasks, 0, sizeof(p->tasks)); + + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming early-load expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + p->selected_ids[i] = selected_ids[i]; + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); + return 0; + } + p->gate_abs_offsets[i] = gate_offset + gate_rel; + p->up_abs_offsets[i] = up_offset + gate_rel; + p->down_abs_offsets[i] = down_offset + down_rel; + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][(uint32_t)selected_ids[i]]; + if (ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + p->gate_abs_offsets[i], + p->up_abs_offsets[i], + p->down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + + uint32_t source = UINT32_MAX; + for (uint32_t prev = 0; prev < i; prev++) { + if (selected_ids[prev] == selected_ids[i] && + p->gate_abs_offsets[prev] == p->gate_abs_offsets[i] && + p->up_abs_offsets[prev] == p->up_abs_offsets[i] && + p->down_abs_offsets[prev] == p->down_abs_offsets[i] && + (p->missing_mask & (1u << prev)) != 0) { + source = p->source_slots[prev] != UINT32_MAX ? + p->source_slots[prev] : prev; + break; + } + } + if (source != UINT32_MAX) { + p->source_slots[i] = source; + p->missing_mask |= 1u << i; + continue; + } + p->source_slots[i] = i; + p->missing_mask |= 1u << i; + p->load_slots[p->n_loads++] = i; + } + if (p->n_loads == 0) return 1; + + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + ds4_gpu_stream_expert_reusable_buffers + batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + batch_reuse[i] = + (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; + } + uint32_t batch_reuse_count = 0; + if (cache_budget != 0 && + reserved_entries >= cache_budget && + p->n_loads > 1 && + ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, + down_expert_bytes)) { + const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + batch_reuse_count = + ds4_gpu_stream_expert_cache_take_reusable_batch( + p->n_loads, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + batch_reuse); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( + ds4_gpu_now_ms() - reuse_t0); + } + } + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + const uint32_t slot = p->load_slots[load_i]; + const uint32_t expert = (uint32_t)p->selected_ids[slot]; + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + + ds4_gpu_stream_expert_readahead_range(p->gate_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(p->up_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(p->down_abs_offsets[slot], + down_expert_bytes); + + if (load_i < batch_reuse_count && + batch_reuse[load_i].gate_buffer && + batch_reuse[load_i].up_buffer && + batch_reuse[load_i].down_buffer) { + p->gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; + p->up_bufs[load_i] = batch_reuse[load_i].up_buffer; + p->down_bufs[load_i] = batch_reuse[load_i].down_buffer; + p->gate_inners[load_i] = batch_reuse[load_i].gate_inner; + p->up_inners[load_i] = batch_reuse[load_i].up_inner; + p->down_inners[load_i] = batch_reuse[load_i].down_inner; + } else { + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &p->gate_bufs[load_i], + &p->up_bufs[load_i], + &p->down_bufs[load_i], + &p->gate_inners[load_i], + &p->up_inners[load_i], + &p->down_inners[load_i]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + uint8_t *gate_dst = (uint8_t *)[p->gate_bufs[load_i] contents] + + p->gate_inners[load_i]; + uint8_t *up_dst = (uint8_t *)[p->up_bufs[load_i] contents] + + p->up_inners[load_i]; + uint8_t *down_dst = (uint8_t *)[p->down_bufs[load_i] contents] + + p->down_inners[load_i]; + if (!gate_dst || !up_dst || !down_dst) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->gate_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->up_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = up_dst, + }; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->down_abs_offsets[slot], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + } + + const uint32_t n_workers = + ds4_gpu_stream_expert_pread_thread_count(p->n_tasks); + p->start_ms = ds4_gpu_now_ms(); + if (load_timing) { + p->prepare_ms = p->start_ms - load_prepare_t0; + } + if (ds4_gpu_stream_expert_pread_pool_begin(p->tasks, + p->n_tasks, + n_workers)) { + p->active = 1; + if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { + fprintf(stderr, + "ds4: Metal streaming expert early-load begin layer=%u experts=%u tensors=%u threads=%u\n", + layer, + p->n_loads, + p->n_tasks, + n_workers); + } + return 1; + } + + uint64_t read_bytes = 0; + double read_ms = 0.0; + if (!ds4_gpu_stream_expert_pread_tasks(p->tasks, + p->n_tasks, + &read_bytes, + &read_ms)) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + (void)read_bytes; + if (!ds4_gpu_stream_expert_pending_load_install(p, NULL, read_ms)) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->n_tasks = 0; + p->n_loads = 0; + p->prepare_ms = 0.0; + return 1; +} + +static int ds4_gpu_stream_expert_cache_load_selected_missing( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_total_expert, + uint32_t n_selected, + const uint64_t *gate_abs_offsets, + const uint64_t *up_abs_offsets, + const uint64_t *down_abs_offsets, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint32_t missing_mask, + ds4_gpu_stream_expert_cache_entry **entries) { + if (!g_ssd_streaming_mode || + !model_map || + !selected_ids || + !gate_abs_offsets || + !up_abs_offsets || + !down_abs_offsets || + !entries || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 0; + } + missing_mask &= (1u << n_selected) - 1u; + if (missing_mask == 0) return 1; + if (g_stream_expert_pending_load.active && + n_selected <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + if (ds4_gpu_stream_expert_pending_load_matches(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + n_selected, + gate_expert_bytes, + down_expert_bytes)) { + if (!ds4_gpu_stream_expert_pending_load_finish(entries)) return 0; + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) != 0 && entries[i]) { + missing_mask &= ~(1u << i); + } + } + if (missing_mask == 0) return 1; + } else { + ds4_gpu_stream_expert_pending_load_clear(); + } + } else if (g_stream_expert_pending_load.active) { + ds4_gpu_stream_expert_pending_load_clear(); + } + + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + double load_prepare_ms = 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + + uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + load_slots[i] = 0; + source_slots[i] = UINT32_MAX; + } + uint32_t n_loads = 0; + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0) continue; + if (entries[i]) { + source_slots[i] = i; + continue; + } + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming selected missing expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + uint32_t source = UINT32_MAX; + for (uint32_t prev = 0; prev < i; prev++) { + if (selected_ids[prev] == selected_ids[i] && + gate_abs_offsets[prev] == gate_abs_offsets[i] && + up_abs_offsets[prev] == up_abs_offsets[i] && + down_abs_offsets[prev] == down_abs_offsets[i] && + (entries[prev] || (missing_mask & (1u << prev)) != 0)) { + source = source_slots[prev] != UINT32_MAX ? + source_slots[prev] : prev; + break; + } + } + if (source != UINT32_MAX) { + source_slots[i] = source; + continue; + } + source_slots[i] = i; + load_slots[n_loads++] = i; + } + if (n_loads == 0) { + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; + const uint32_t source = source_slots[i]; + if (source >= n_selected || !entries[source]) return 0; + entries[i] = entries[source]; + entries[i]->use_count++; + } + return 1; + } + + __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + gate_bufs[i] = nil; + up_bufs[i] = nil; + down_bufs[i] = nil; + gate_inners[i] = 0; + up_inners[i] = 0; + down_inners[i] = 0; + } + ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; + memset(tasks, 0, sizeof(tasks)); + uint32_t n_tasks = 0; + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + ds4_gpu_stream_expert_reusable_buffers + batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + batch_reuse[i] = + (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; + } + uint32_t batch_reuse_count = 0; + if (cache_budget != 0 && + reserved_entries >= cache_budget && + n_loads > 1 && + ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, + down_expert_bytes)) { + const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + batch_reuse_count = + ds4_gpu_stream_expert_cache_take_reusable_batch( + n_loads, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + batch_reuse); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( + ds4_gpu_now_ms() - reuse_t0); + } + } + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t slot = load_slots[load_i]; + const uint32_t expert = (uint32_t)selected_ids[slot]; + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + + ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot], + down_expert_bytes); + + if (load_i < batch_reuse_count && + batch_reuse[load_i].gate_buffer && + batch_reuse[load_i].up_buffer && + batch_reuse[load_i].down_buffer) { + gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; + up_bufs[load_i] = batch_reuse[load_i].up_buffer; + down_bufs[load_i] = batch_reuse[load_i].down_buffer; + gate_inners[load_i] = batch_reuse[load_i].gate_inner; + up_inners[load_i] = batch_reuse[load_i].up_inner; + down_inners[load_i] = batch_reuse[load_i].down_inner; + } else { + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_bufs[load_i], + &up_bufs[load_i], + &down_bufs[load_i], + &gate_inners[load_i], + &up_inners[load_i], + &down_inners[load_i]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + return 0; + } + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + if (!gate_bufs[load_i] || !up_bufs[load_i] || !down_bufs[load_i]) { + return 0; + } + + uint8_t *gate_dst = (uint8_t *)[gate_bufs[load_i] contents] + + gate_inners[load_i]; + uint8_t *up_dst = (uint8_t *)[up_bufs[load_i] contents] + + up_inners[load_i]; + uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] + + down_inners[load_i]; + if (!gate_dst || !up_dst || !down_dst) return 0; + + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = gate_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = up_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = up_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = down_abs_offsets[slot], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + } + + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_prepare_ms = now_ms - load_t0; + load_t0 = now_ms; + } + + uint64_t read_bytes = 0; + double read_ms = 0.0; + const int ok = ds4_gpu_stream_expert_pread_tasks(tasks, + n_tasks, + &read_bytes, + &read_ms); + if (!ok) return 0; + if (load_timing) { + load_t0 = ds4_gpu_now_ms(); + } + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + } + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_t0; + load_t0 = now_ms; + } + if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + n_loads, + n_tasks, + ds4_gpu_stream_expert_pread_thread_count(n_tasks), + ds4_gpu_gib(read_bytes), + read_ms); + } + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t slot = load_slots[load_i]; + const uint32_t expert = (uint32_t)selected_ids[slot]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + gate_abs_offsets[slot], + up_abs_offsets[slot], + down_abs_offsets[slot], + gate_expert_bytes, + down_expert_bytes, + gate_bufs[load_i], + up_bufs[load_i], + down_bufs[load_i], + gate_inners[load_i], + up_inners[load_i], + down_inners[load_i]); + if (!entry) return 0; + entries[slot] = entry; + } + + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; + const uint32_t source = source_slots[i]; + if (source >= n_selected || !entries[source]) return 0; + entries[i] = entries[source]; + entries[i]->use_count++; + } + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0; + } + if (load_timing) { + load_install_ms = ds4_gpu_now_ms() - load_t0; + ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, + read_ms, + load_modify_ms, + load_install_ms); + } + return 1; +} + +static void ds4_gpu_glm_stream_selected_prefetch_set( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + g_glm_stream_selected_prefetch.active = 0; + if (!table || !selected_ids || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return; + } + g_glm_stream_selected_prefetch.model_map = table->model_map; + g_glm_stream_selected_prefetch.model_size = table->model_size; + g_glm_stream_selected_prefetch.layer = table->layer; + g_glm_stream_selected_prefetch.n_total_expert = table->n_total_expert; + g_glm_stream_selected_prefetch.n_selected = n_selected; + g_glm_stream_selected_prefetch.gate_offset = table->gate_offset; + g_glm_stream_selected_prefetch.up_offset = table->up_offset; + g_glm_stream_selected_prefetch.down_offset = table->down_offset; + g_glm_stream_selected_prefetch.gate_expert_bytes = table->gate_expert_bytes; + g_glm_stream_selected_prefetch.down_expert_bytes = table->down_expert_bytes; + for (uint32_t i = 0; i < n_selected; i++) { + g_glm_stream_selected_prefetch.selected_ids[i] = selected_ids[i]; + } + g_glm_stream_selected_prefetch.active = 1; +} + +static int ds4_gpu_glm_stream_selected_prefetch_take( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + int32_t *selected_ids_out) { + if (!selected_ids_out || !g_glm_stream_selected_prefetch.active) return 0; + ds4_gpu_glm_stream_selected_prefetch *p = &g_glm_stream_selected_prefetch; + if (p->model_map != model_map || + p->model_size != model_size || + p->layer != layer || + p->n_total_expert != n_total_expert || + p->n_selected != n_selected || + p->gate_offset != gate_offset || + p->up_offset != up_offset || + p->down_offset != down_offset || + p->gate_expert_bytes != gate_expert_bytes || + p->down_expert_bytes != down_expert_bytes) { + return 0; + } + for (uint32_t i = 0; i < n_selected; i++) { + selected_ids_out[i] = p->selected_ids[i]; + } + p->active = 0; + return 1; +} + +int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_selected) { + g_glm_stream_selected_prefetch.active = 0; + if (!g_ssd_streaming_mode || + getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) { + return 1; + } + if (!table || !selected || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return 1; + } + + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double t0 = timing ? ds4_gpu_now_ms() : 0.0; + double selected_sync_ms = 0.0; + double selected_copy_ms = 0.0; + double selected_bind_ms = 0.0; + const int had_batch = g_batch_cb != nil; + if (had_batch) { + if (ds4_gpu_end_commands() == 0) return 0; + if (timing) { + selected_sync_ms = ds4_gpu_now_ms() - t0; + t0 = ds4_gpu_now_ms(); + } + } + + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + selected_ids[i] = -1; + } + int ok = ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_selected * sizeof(selected_ids[0])) != 0; + if (timing) { + selected_copy_ms = ds4_gpu_now_ms() - t0; + t0 = ds4_gpu_now_ms(); + } + if (ok) { + ok = ds4_gpu_stream_expert_cache_begin_selected_load(table, + selected_ids, + n_selected) != 0; + } + if (timing) { + selected_bind_ms = ds4_gpu_now_ms() - t0; + ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, + selected_copy_ms, + selected_bind_ms); + } + if (ok) { + ds4_gpu_glm_stream_selected_prefetch_set(table, + selected_ids, + n_selected); + } + if (had_batch && ds4_gpu_begin_commands() == 0) ok = 0; + return ok; +} + +static void ds4_gpu_stream_expert_cache_clear_layer(uint32_t layer) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + g_stream_expert_cache_layer_count[layer] = 0; +} + +static int ds4_gpu_stream_expert_cache_prepare_selected_batch( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + ds4_gpu_stream_expert_cache_entry **resources, + uint32_t *n_resources, + uint32_t *unique_out, + id *overflow_gate, + id *overflow_up, + id *overflow_down) { + if (overflow_gate) *overflow_gate = nil; + if (overflow_up) *overflow_up = nil; + if (overflow_down) *overflow_down = nil; + if (!g_ssd_streaming_mode || + !model_map || + !selected || + !gate_addrs || + !up_addrs || + !down_addrs || + !resources || + !n_resources || + !unique_out || + !overflow_gate || + !overflow_up || + !overflow_down || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_tokens == 0 || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + gate_expert_bytes == 0 || + down_expert_bytes == 0) { + return 0; + } + if (n_tokens > UINT32_MAX / n_selected) return 0; + + const uint64_t n_ids = (uint64_t)n_tokens * n_selected; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return 0; + int32_t *ids = malloc((size_t)n_ids * sizeof(ids[0])); + if (!ids) return 0; + + const bool profile = + getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_PROFILE") != NULL; + const double t0 = profile ? ds4_gpu_now_ms() : 0.0; + int ok = ds4_gpu_tensor_read(selected, + 0, + ids, + n_ids * sizeof(ids[0])); + const double t_read = profile ? ds4_gpu_now_ms() : 0.0; + + bool seen[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { false }; + uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { 0 }; + int32_t unique_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint32_t unique_count = 0; + *n_resources = 0; + *unique_out = 0; + if (ok) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { + ok = 0; + } + } + if (ok) { + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t selected_id = ids[i]; + if (selected_id < 0 || (uint32_t)selected_id >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming batch selected expert id %d is outside 0..%u at layer %u\n", + selected_id, + n_total_expert, + layer); + ok = 0; + break; + } + frequency[(uint32_t)selected_id]++; + if (!seen[(uint32_t)selected_id]) { + seen[(uint32_t)selected_id] = true; + unique_ids[unique_count++] = selected_id; + } + } + } + if (ok) { + ds4_gpu_stream_expert_cache_note_frequency_hotness(layer, + frequency, + n_total_expert); + } + /* + * When the layer's unique selected set does not fit the cache budget, the + * extra experts are addressed straight into whole-tensor mapped model views + * instead of falling back to a different MoE kernel path. The address-table + * kernels read identical expert bytes either way, so the cache/view split + * does not change the computed logits. + */ + uint32_t view_served = 0; + uint64_t overflow_gate_inner = 0; + uint64_t overflow_up_inner = 0; + uint64_t overflow_down_inner = 0; + + uint64_t unique_gate_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint64_t unique_up_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint64_t unique_down_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + ds4_gpu_stream_expert_cache_entry + *unique_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint32_t load_unique[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; i++) { + unique_gate_offsets[i] = 0; + unique_up_offsets[i] = 0; + unique_down_offsets[i] = 0; + unique_entries[i] = NULL; + load_unique[i] = 0; + gate_bufs[i] = nil; + up_bufs[i] = nil; + down_bufs[i] = nil; + gate_inners[i] = 0; + up_inners[i] = 0; + down_inners[i] = 0; + } + + ds4_gpu_stream_expert_pread_task *tasks = NULL; + uint32_t n_loads = 0; + uint32_t n_tasks = 0; + double load_prepare_ms = 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + double load_timing_t0 = ds4_gpu_stream_expert_timing_summary_enabled() ? + ds4_gpu_now_ms() : 0.0; + const int load_timing = load_timing_t0 != 0.0; + if (ok && unique_count != 0) { + tasks = calloc((size_t)unique_count * 3u, sizeof(tasks[0])); + if (!tasks) ok = 0; + } + if (ok) { + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + + for (uint32_t u = 0; u < unique_count; u++) { + const uint32_t expert = (uint32_t)unique_ids[u]; + + if ((uint64_t)expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); + ok = 0; + break; + } + const uint64_t gate_rel = (uint64_t)expert * gate_expert_bytes; + const uint64_t down_rel = (uint64_t)expert * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); + ok = 0; + break; + } + unique_gate_offsets[u] = gate_offset + gate_rel; + unique_up_offsets[u] = up_offset + gate_rel; + unique_down_offsets[u] = down_offset + down_rel; + + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer, + expert, + n_total_expert, + n_selected, + unique_gate_offsets[u], + unique_up_offsets[u], + unique_down_offsets[u], + gate_expert_bytes, + down_expert_bytes); + if (entry) { + unique_entries[u] = entry; + continue; + } + + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + ds4_gpu_stream_expert_readahead_range(unique_gate_offsets[u], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(unique_up_offsets[u], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(unique_down_offsets[u], + down_expert_bytes); + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + unique_ids, + unique_count, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_bufs[n_loads], + &up_bufs[n_loads], + &down_bufs[n_loads], + &gate_inners[n_loads], + &up_inners[n_loads], + &down_inners[n_loads]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + ok = 0; + break; + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + if (!gate_bufs[n_loads] || + !up_bufs[n_loads] || + !down_bufs[n_loads]) { + ok = 0; + break; + } + + uint8_t *gate_dst = (uint8_t *)[gate_bufs[n_loads] contents] + + gate_inners[n_loads]; + uint8_t *up_dst = (uint8_t *)[up_bufs[n_loads] contents] + + up_inners[n_loads]; + uint8_t *down_dst = (uint8_t *)[down_bufs[n_loads] contents] + + down_inners[n_loads]; + if (!gate_dst || !up_dst || !down_dst) { + ok = 0; + break; + } + + load_unique[n_loads] = u; + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_gate_offsets[u], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_up_offsets[u], + .len = gate_expert_bytes, + .dst = up_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_down_offsets[u], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + n_loads++; + } + } + if (ok && n_loads != 0) { + if (load_timing_t0 != 0.0) { + const double now_ms = ds4_gpu_now_ms(); + load_prepare_ms = now_ms - load_timing_t0; + load_timing_t0 = now_ms; + } + uint64_t read_bytes = 0; + double read_ms = 0.0; + ok = ds4_gpu_stream_expert_pread_tasks(tasks, + n_tasks, + &read_bytes, + &read_ms); + if (ok) { + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + } + if (load_timing_t0 != 0.0) { + load_timing_t0 = ds4_gpu_now_ms(); + } + if (ok) { + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + } + } + if (load_timing_t0 != 0.0) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_timing_t0; + load_timing_t0 = now_ms; + } + if (ok && getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming batch expert parallel pread layer=%u experts=%u tensors=%u " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + n_loads, + n_tasks, + ds4_gpu_stream_expert_pread_thread_count(n_tasks), + ds4_gpu_gib(read_bytes), + read_ms); + } + if (ok) { + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t u = load_unique[load_i]; + const uint32_t expert = (uint32_t)unique_ids[u]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + unique_gate_offsets[u], + unique_up_offsets[u], + unique_down_offsets[u], + gate_expert_bytes, + down_expert_bytes, + gate_bufs[load_i], + up_bufs[load_i], + down_bufs[load_i], + gate_inners[load_i], + up_inners[load_i], + down_inners[load_i]); + if (!entry) { + ok = 0; + break; + } + unique_entries[u] = entry; + } + } + if (load_timing_t0 != 0.0) { + load_install_ms = ds4_gpu_now_ms() - load_timing_t0; + ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, + read_ms, + load_modify_ms, + load_install_ms); + } + } + if (tasks) free(tasks); + if (ok) { + for (uint32_t u = 0; u < unique_count; u++) { + ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u]; + const uint32_t expert = (uint32_t)unique_ids[u]; + if (!entry) { + const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset; + const uint64_t down_rel = unique_down_offsets[u] - down_offset; + if (!*overflow_gate) { + const uint64_t gate_tensor_bytes = + (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = + (uint64_t)n_total_expert * down_expert_bytes; + uint64_t gate_view_inner = 0; + uint64_t up_view_inner = 0; + uint64_t down_view_inner = 0; + id gv = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_view_inner); + id uv = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_view_inner); + id dv = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_view_inner); + if (!gv || !uv || !dv) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr " + "failed to map overflow expert views at layer %u\n", + layer); + ok = 0; + break; + } + *overflow_gate = gv; + *overflow_up = uv; + *overflow_down = dv; + overflow_gate_inner = gate_view_inner; + overflow_up_inner = up_view_inner; + overflow_down_inner = down_view_inner; + } + if (!ds4_gpu_stream_expert_cache_set_addr_slot( + layer, + expert, + *overflow_gate, + (NSUInteger)(overflow_gate_inner + gate_rel), + *overflow_up, + (NSUInteger)(overflow_up_inner + gate_rel), + *overflow_down, + (NSUInteger)(overflow_down_inner + down_rel))) { + ok = 0; + break; + } + view_served++; + continue; + } + const uint32_t extra_uses = + frequency[expert] > 0 ? frequency[expert] - 1u : 0; + if (extra_uses != 0) { + if (entry->use_count > UINT64_MAX - extra_uses) { + entry->use_count = UINT64_MAX; + } else { + entry->use_count += extra_uses; + } + if (g_stream_expert_cache_hits > UINT64_MAX - extra_uses) { + g_stream_expert_cache_hits = UINT64_MAX; + } else { + g_stream_expert_cache_hits += extra_uses; + } + if (g_stream_expert_cache_layer_hits[layer] > + UINT64_MAX - extra_uses) { + g_stream_expert_cache_layer_hits[layer] = UINT64_MAX; + } else { + g_stream_expert_cache_layer_hits[layer] += extra_uses; + } + } + resources[*n_resources] = entry; + (*n_resources)++; + } + } + free(ids); + + if (!ok || (*n_resources == 0 && view_served == 0)) { + ds4_gpu_stream_expert_cache_clear_layer(layer); + return 0; + } + if (!ds4_gpu_stream_expert_cache_addr_buffers(layer, + gate_addrs, + up_addrs, + down_addrs)) { + ds4_gpu_stream_expert_cache_clear_layer(layer); + return 0; + } + *unique_out = *n_resources + view_served; + if (view_served != 0 && + ds4_gpu_stream_expert_timing_summary_enabled()) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "served %u/%u unique experts via mapped views (cache budget %u)\n", + layer, + view_served, + unique_count, + ds4_gpu_stream_expert_cache_configured_budget()); + } + if (profile) { + const double t_done = ds4_gpu_now_ms(); + uint64_t per_expert_bytes = UINT64_MAX; + uint64_t logical_bytes = UINT64_MAX; + if (gate_expert_bytes <= (UINT64_MAX - down_expert_bytes) / 2ull) { + per_expert_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (per_expert_bytes == 0 || + *n_resources <= UINT64_MAX / per_expert_bytes) { + logical_bytes = (uint64_t)(*n_resources) * per_expert_bytes; + } + } + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "tokens=%u unique=%u read=%.3f ms wrap=%.3f ms bytes=%.2f GiB\n", + layer, + n_tokens, + *n_resources, + t_read - t0, + t_done - t_read, + ds4_gpu_gib(logical_bytes)); + } + return 1; +} + +int ds4_gpu_stream_expert_cache_seed_selected( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!g_ssd_streaming_mode) return 1; + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > 6 || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + ds4_gpu_stream_expert_cache_note_selected_hotness(layer, + selected_ids, + n_selected); + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal prefill expert-cache seed selected expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); + return 0; + } + + if (!ds4_gpu_stream_expert_cache_get(model_map, + model_size, + layer, + (uint32_t)selected_ids[i], + n_total_expert, + n_selected, + gate_offset + gate_rel, + up_offset + gate_rel, + down_offset + down_rel, + gate_expert_bytes, + down_expert_bytes)) { + return 0; + } + } + ds4_gpu_stream_expert_cache_prune_layer(layer, + n_total_expert, + n_selected, + selected_ids, + n_selected); + ds4_gpu_stream_expert_cache_prune_global(layer, + selected_ids, + n_selected); + return 1; +} + +int ds4_gpu_stream_expert_cache_seed_experts( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts) { + if (!g_ssd_streaming_mode) return 1; + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !expert_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_experts == 0 || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + 1) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + uint32_t remaining = n_experts; + while (remaining != 0) { + const uint32_t batch = + remaining > 6u ? 6u : remaining; + remaining -= batch; + + int32_t selected_ids[6] = { -1, -1, -1, -1, -1, -1 }; + uint64_t gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + uint64_t up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + uint64_t down_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + ds4_gpu_stream_expert_cache_entry *entries[6] = { + NULL, NULL, NULL, NULL, NULL, NULL + }; + uint32_t missing_mask = 0; + + for (uint32_t i = 0; i < batch; i++) { + const int32_t expert = expert_ids[remaining + i]; + const uint32_t priority = + expert_priorities ? expert_priorities[remaining + i] : 0; + if (expert < 0 || (uint32_t)expert >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming hotlist seed expert id %d is outside 0..%u\n", + expert, + n_total_expert); + return 0; + } + const uint64_t expert_id = (uint64_t)(uint32_t)expert; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); + return 0; + } + + ds4_gpu_stream_expert_cache_note_route_hotness( + layer, + (uint32_t)expert, + priority != 0 ? priority : 1u); + selected_ids[i] = expert; + gate_abs_offsets[i] = gate_offset + gate_rel; + up_abs_offsets[i] = up_offset + gate_rel; + down_abs_offsets[i] = down_offset + down_rel; + entries[i] = ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer, + (uint32_t)expert, + n_total_expert, + batch, + gate_abs_offsets[i], + up_abs_offsets[i], + down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!entries[i]) missing_mask |= 1u << i; + if (entries[i] && priority != 0 && + entries[i]->use_count < (uint64_t)priority) { + entries[i]->use_count = (uint64_t)priority; + } + } + + if (missing_mask != 0 && + !ds4_gpu_stream_expert_cache_load_selected_missing(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + batch, + gate_abs_offsets, + up_abs_offsets, + down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + missing_mask, + entries)) { + return 0; + } + if (expert_priorities) { + for (uint32_t i = 0; i < batch; i++) { + const uint32_t priority = expert_priorities[remaining + i]; + if (entries[i] && priority != 0 && + entries[i]->use_count < (uint64_t)priority) { + entries[i]->use_count = (uint64_t)priority; + } + } + } + } + + const uint32_t protect_n = n_experts < 6u ? n_experts : 6u; + ds4_gpu_stream_expert_cache_prune_layer(layer, + n_total_expert, + 1, + expert_ids, + protect_n); + ds4_gpu_stream_expert_cache_prune_global(layer, + expert_ids, + protect_n); + return 1; +} + +static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE"); + if (!env || !env[0]) return 1; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0' || v < 2 || v > n_total_expert) { + return 1; + } + return (uint32_t)v; +} + +static bool ds4_gpu_q4_table_queue_residency_requested(void) { + return getenv("DS4_METAL_Q4_TABLE_QUEUE_RESIDENCY_SET") != NULL; +} + +static bool ds4_gpu_q4_table_queue_residency_available(void) { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + return g_queue && [g_queue respondsToSelector:@selector(addResidencySet:)]; + } +#endif + return false; +} + +static bool ds4_gpu_q4_non_streaming_opt_in_enabled(void) { + return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL; +} + +static bool ds4_gpu_q4_selected_paths_allowed(void) { + if (g_ssd_streaming_mode) return true; + if (g_glm_model_mode) return false; + return ds4_gpu_q4_non_streaming_opt_in_enabled(); +} + +int ds4_gpu_pro_q4_expert_table_auto_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + return (g_ssd_streaming_mode || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL) && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_queue_residency_enabled(bool auto_queue_residency) { + return (auto_queue_residency || ds4_gpu_q4_table_queue_residency_requested()) && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_model_residency_enabled(void) { + return getenv("DS4_METAL_Q4_TABLE_MODEL_RESIDENCY_SET") != NULL; +} + +static bool ds4_gpu_pro_q4_expert_indirect_shape_supported( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + return n_total_expert == 384 && + n_expert == 6 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes; +} + +static bool ds4_gpu_pro_q4_expert_table_auto_enabled( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + /* + * This path lets the shader choose among exact per-expert resources. + * It is only automatic when a Metal queue residency set can make every + * indirect expert resource visible up front; otherwise the selected + * exact-slice path remains the fallback. + */ + return ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes) && + ds4_gpu_pro_q4_expert_table_auto_available(); +} + +static bool ds4_gpu_pro_q4_expert_address_auto_enabled( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + /* + * GPU-address expert tables are useful for experiments, but they are not + * safe as an automatic path unless every indirect expert resource is made + * visible to Metal. Keep this behind an explicit opt-in while the selected + * active-slice path remains the correctness/performance baseline. + */ + return getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL && + ds4_gpu_q4_selected_paths_allowed() && + ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes) && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_bind_anchors_enabled(void) { + return getenv("DS4_METAL_Q4_TABLE_BIND_ANCHORS") != NULL; +} + +static int ds4_gpu_use_model_residency_set(id cb) { + if (!ds4_gpu_q4_table_model_residency_enabled()) return 1; +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (cb && g_model_residency_set && [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:g_model_residency_set]; + return 1; + } + } +#endif + fprintf(stderr, "ds4: Metal Q4 table model residency set is not available\n"); + return 0; +} + +static int ds4_gpu_bind_q4_expert_table_anchors( + id enc, + DS4MetalQ4ExpertTable *table, + NSUInteger first_index, + NSUInteger max_count) { + if (!ds4_gpu_q4_table_bind_anchors_enabled()) return 1; + if (!enc || !table || !table.expertBuffers) return 0; + + const NSUInteger count = [table.expertBuffers count]; + if (count > max_count) { + fprintf(stderr, + "ds4: Metal Q4 table anchor count %lu exceeds available slots %lu\n", + (unsigned long)count, + (unsigned long)max_count); + return 0; + } + for (NSUInteger i = 0; i < count; i++) { + [enc setBuffer:[table.expertBuffers objectAtIndex:i] + offset:0 + atIndex:first_index + i]; + } + return 1; +} + +static id ds4_gpu_q4_expert_table_residency_set(NSMutableArray> *buffers) { + if (!buffers || [buffers count] == 0 || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL || + getenv("DS4_METAL_Q4_TABLE_PER_TENSOR_RESIDENCY_SET") == NULL) { + return nil; + } +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_q4_expert_table"; + desc.initialCapacity = [buffers count]; + NSError *error = nil; + id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!residency_set) { + fprintf(stderr, "ds4: Metal Q4 expert table residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + for (id buffer in buffers) { + [residency_set addAllocation:buffer]; + } + [residency_set commit]; + [residency_set requestResidency]; + return residency_set; + } +#endif + return nil; +} + +static void ds4_gpu_q4_residency_add_table(id residency_set, + DS4MetalQ4ExpertTable *table) { +#if TARGET_OS_OSX + if (!residency_set || !table) return; + if (@available(macOS 15.0, *)) { + if (table.argumentBuffer) { + [residency_set addAllocation:table.argumentBuffer]; + } + if (table.addressBuffer) { + [residency_set addAllocation:table.addressBuffer]; + } + for (id buffer in table.expertBuffers) { + [residency_set addAllocation:buffer]; + } + } +#else + (void)residency_set; + (void)table; +#endif +} + +static id ds4_gpu_q4_expert_layer_residency_set(DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + DS4MetalQ4ExpertTable *down_table, + bool auto_queue_residency) { + const bool queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(auto_queue_residency); + if (!gate_table || !up_table || !down_table || + !g_device || !g_q4_expert_layer_residency_cache || + (getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL && + !queue_residency)) { + return nil; + } +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + NSString *key = [NSString stringWithFormat:@"%p:%p:%p", + gate_table, up_table, down_table]; + DS4MetalQ4LayerResidency *cached = + [g_q4_expert_layer_residency_cache objectForKey:key]; + if (cached) return cached.residencySet; + + const NSUInteger capacity = + [gate_table.expertBuffers count] + + [up_table.expertBuffers count] + + [down_table.expertBuffers count] + 3u; + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_q4_expert_layer"; + desc.initialCapacity = capacity; + NSError *error = nil; + id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!residency_set) { + fprintf(stderr, + "ds4: Metal Q4 expert layer residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + ds4_gpu_q4_residency_add_table(residency_set, gate_table); + ds4_gpu_q4_residency_add_table(residency_set, up_table); + ds4_gpu_q4_residency_add_table(residency_set, down_table); + [residency_set commit]; + [residency_set requestResidency]; + + DS4MetalQ4LayerResidency *entry = [DS4MetalQ4LayerResidency new]; + entry.residencySet = residency_set; + if (queue_residency) { + [g_queue addResidencySet:residency_set]; + entry.addedToQueue = YES; + } + [g_q4_expert_layer_residency_cache setObject:entry forKey:key]; + return residency_set; + } +#endif + return nil; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert, + id encoder); + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert); + +int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, + uint64_t gate_expert_bytes, uint64_t down_expert_bytes, + uint32_t n_total_expert) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0 || gate_expert_bytes == 0 || + down_expert_bytes == 0 || n_total_expert == 0) { + return 0; + } + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table preload byte size overflow\n"); + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; + const bool address_auto = + ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, + 6, + gate_tensor_bytes, + down_tensor_bytes); + const bool table_auto = + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + 6, + gate_tensor_bytes, + down_tensor_bytes); + if (!address_auto && !table_auto) { + return 1; + } + + if (address_auto) { + @autoreleasepool { + DS4MetalQ4ExpertTable *gate_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert); + DS4MetalQ4ExpertTable *up_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert); + DS4MetalQ4ExpertTable *down_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert); + if (!gate_table || !up_table || !down_table) { + return 0; + } + id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + true); + if (!residency) { + fprintf(stderr, "ds4: Metal Q4 expert address table preload failed to create queue residency set\n"); + return 0; + } + } + return 1; + } + + if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder || + !g_moe_table_q4_sum_down_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert table preload missing argument encoders\n"); + return 0; + } + + @autoreleasepool { + DS4MetalQ4ExpertTable *gate_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + DS4MetalQ4ExpertTable *up_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + DS4MetalQ4ExpertTable *down_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + return 0; + } + id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + true); + if (!residency) { + fprintf(stderr, "ds4: Metal Q4 expert table preload failed to create queue residency set\n"); + return 0; + } + } + return 1; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert, + id encoder) { + if (!model_map || !g_device || !g_q4_expert_table_cache || !encoder || + model_size == 0 || expert_bytes == 0 || + n_total_expert == 0 || n_total_expert > 384) { + return nil; + } + if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table byte size overflow\n"); + return nil; + } + const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; + if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table is outside the mapped model\n"); + return nil; + } + + const uint32_t table_group_size = + ds4_gpu_q4_expert_table_group_size(n_total_expert); + NSString *key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu:%u:%llu:%u", + model_map, + (unsigned long long)model_size, + (unsigned long long)tensor_offset, + (unsigned long long)expert_bytes, + n_total_expert, + (unsigned long long)[encoder encodedLength], + table_group_size]; + DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; + if (cached) return cached; + + id arg_buffer = + [g_device newBufferWithLength:[encoder encodedLength] + options:MTLResourceStorageModeShared]; + if (!arg_buffer) { + fprintf(stderr, "ds4: Metal Q4 expert table argument buffer allocation failed\n"); + return nil; + } + arg_buffer.label = @"ds4_q4_expert_table"; + [encoder setArgumentBuffer:arg_buffer offset:0]; + + NSMutableArray> *expert_buffers = + [NSMutableArray arrayWithCapacity:table_group_size > 1 ? + (n_total_expert + table_group_size - 1u) / table_group_size : + n_total_expert]; + if (!expert_buffers) return nil; + + if (table_group_size > 1) { + for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { + const uint32_t remaining = n_total_expert - first; + const uint32_t group_n = + remaining < table_group_size ? remaining : table_group_size; + if ((uint64_t)first > UINT64_MAX / expert_bytes || + (uint64_t)group_n > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table group byte overflow\n"); + return nil; + } + const uint64_t rel = (uint64_t)first * expert_bytes; + const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table group offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id group_buf = + ds4_gpu_wrap_model_range(model_map, + model_size, + tensor_offset + rel, + group_bytes, + &inner); + if (!group_buf) return nil; + for (uint32_t j = 0; j < group_n; j++) { + const uint64_t expert_inner = inner + (uint64_t)j * expert_bytes; + [encoder setBuffer:group_buf offset:(NSUInteger)expert_inner atIndex:first + j]; + } + [expert_buffers addObject:group_buf]; + } + } else { + for (uint32_t i = 0; i < n_total_expert; i++) { + const uint64_t rel = (uint64_t)i * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id expert_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + tensor_offset + rel, + expert_bytes, + &inner); + if (!expert_buf) return nil; + [encoder setBuffer:expert_buf offset:(NSUInteger)inner atIndex:i]; + [expert_buffers addObject:expert_buf]; + } + } + [arg_buffer didModifyRange:NSMakeRange(0, [encoder encodedLength])]; + + if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal Q4 expert table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", + n_total_expert, + table_group_size, + (unsigned long)[expert_buffers count], + ds4_gpu_mib(expert_bytes)); + } + + DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; + table.argumentBuffer = arg_buffer; + table.expertBuffers = expert_buffers; + table.residencySet = ds4_gpu_q4_expert_table_residency_set(expert_buffers); + table.nExpert = n_total_expert; + table.expertBytes = expert_bytes; + [g_q4_expert_table_cache setObject:table forKey:key]; + return table; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert) { + if (!model_map || !g_device || !g_q4_expert_table_cache || + model_size == 0 || expert_bytes == 0 || + n_total_expert == 0 || n_total_expert > 384) { + return nil; + } + if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert address table byte size overflow\n"); + return nil; + } + const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; + if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert address table is outside the mapped model\n"); + return nil; + } + + const uint32_t table_group_size = + ds4_gpu_q4_expert_table_group_size(n_total_expert); + NSString *key = [NSString stringWithFormat:@"addr:%p:%llu:%llu:%llu:%u:%u", + model_map, + (unsigned long long)model_size, + (unsigned long long)tensor_offset, + (unsigned long long)expert_bytes, + n_total_expert, + table_group_size]; + DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; + if (cached) return cached; + + id address_buffer = + [g_device newBufferWithLength:(NSUInteger)n_total_expert * sizeof(uint64_t) + options:MTLResourceStorageModeShared]; + if (!address_buffer) { + fprintf(stderr, "ds4: Metal Q4 expert address table allocation failed\n"); + return nil; + } + address_buffer.label = @"ds4_q4_expert_address_table"; + uint64_t *addresses = (uint64_t *)[address_buffer contents]; + if (!addresses) return nil; + + NSMutableArray> *expert_buffers = + [NSMutableArray arrayWithCapacity:table_group_size > 1 ? + (n_total_expert + table_group_size - 1u) / table_group_size : + n_total_expert]; + if (!expert_buffers) return nil; + +#if TARGET_OS_OSX + if (@available(macOS 13.0, *)) { + for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { + const uint32_t remaining = n_total_expert - first; + const uint32_t group_n = + remaining < table_group_size ? remaining : table_group_size; + if ((uint64_t)first > UINT64_MAX / expert_bytes || + (uint64_t)group_n > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n"); + return nil; + } + const uint64_t rel = (uint64_t)first * expert_bytes; + const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id group_buf = nil; + if (table_group_size > 1) { + group_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + tensor_offset + rel, + group_bytes, + &inner); + } else { + group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + tensor_offset + rel, + expert_bytes, + &inner); + } + if (!group_buf) return nil; + const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner; + for (uint32_t j = 0; j < group_n; j++) { + addresses[first + j] = base_address + (uint64_t)j * expert_bytes; + } + [expert_buffers addObject:group_buf]; + } + } else +#endif + { + fprintf(stderr, "ds4: Metal GPU addresses require macOS 13 or newer\n"); + return nil; + } + + [address_buffer didModifyRange:NSMakeRange(0, + (NSUInteger)n_total_expert * sizeof(uint64_t))]; + + if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal Q4 expert address table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", + n_total_expert, + table_group_size, + (unsigned long)[expert_buffers count], + ds4_gpu_mib(expert_bytes)); + } + + DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; + table.addressBuffer = address_buffer; + table.expertBuffers = expert_buffers; + table.nExpert = n_total_expert; + table.expertBytes = expert_bytes; + [g_q4_expert_table_cache setObject:table forKey:key]; + return table; +} + +static void ds4_gpu_use_q4_expert_table_resources( + id cb, + id enc, + DS4MetalQ4ExpertTable *table, + bool queue_residency) { + if (!enc || !table) return; + if (table.argumentBuffer) { + [enc useResource:table.argumentBuffer usage:MTLResourceUsageRead]; + } + if (table.addressBuffer) { + [enc useResource:table.addressBuffer usage:MTLResourceUsageRead]; + } + if (!queue_residency && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + table.residencySet && + cb && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:table.residencySet]; + } + if (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") == NULL && + getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") == NULL) { + return; + } + const NSUInteger count = [table.expertBuffers count]; + for (NSUInteger base = 0; base < count; base += 64u) { + const NSUInteger n = count - base < 64u ? count - base : 64u; + __unsafe_unretained id resources[64]; + for (NSUInteger i = 0; i < n; i++) { + resources[i] = [table.expertBuffers objectAtIndex:base + i]; + } + [enc useResources:resources count:n usage:MTLResourceUsageRead]; + } +} diff --git a/metal/model_abi.metal b/metal/model_abi.metal new file mode 100644 index 0000000000..47b102c599 --- /dev/null +++ b/metal/model_abi.metal @@ -0,0 +1,368 @@ +struct ds4_metal_args_dsv4_topk_mask { + int64_t ne00; + int64_t ne01; + uint64_t nb00; + uint64_t nb01; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +}; + +struct ds4_metal_args_dsv4_indexer_weighted_sum { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int64_t ne10; + int64_t ne11; + uint64_t nb10; + uint64_t nb11; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; + float scale; +}; + +struct ds4_metal_args_dsv4_softmax_pool { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +}; + +struct ds4_metal_args_dsv4_softmax_pool_ratio4_direct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +}; + +struct ds4_metal_args_dsv4_compressor_score_ape { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +}; + +struct ds4_metal_args_dsv4_indexed_attention { + uint32_t n_tokens; + uint32_t n_head; + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; + uint32_t top_k; + uint32_t pos0; + uint32_t window; + uint32_t ratio; + uint32_t comp_kv_f16; + uint32_t pad0; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t raw_row_stride; + uint64_t comp_row_stride; + uint64_t topk_token_stride; + uint64_t dst_token_stride; + uint64_t dst_head_stride; + float scale; +}; + +struct ds4_metal_args_dsv4_indexer_scores_fused { + uint32_t n_comp; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t ratio; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t index_row_stride; + uint64_t score_token_stride; + float scale; +}; + +struct ds4_metal_args_dsv4_router_select_one { + uint32_t has_bias; + uint32_t hash_mode; + uint32_t use_token_buffer; + uint32_t token; + uint32_t hash_rows; +}; + +struct ds4_metal_args_glm_router_select_one { + uint32_t n_expert; + uint32_t n_expert_used; + float expert_weight_scale; + uint32_t pad0; +}; + +struct ds4_metal_args_glm_kv_lora_rms_norm { + uint32_t n_tokens; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + float eps; +}; + +struct ds4_metal_args_glm_k_b_project { + uint32_t n_tokens; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t n_head; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +}; + +struct ds4_metal_args_glm_build_kv_cache { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + +struct ds4_metal_args_glm_store_compact_kv { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_rope; + uint32_t cache_f16; + uint32_t pad1; +}; + +struct ds4_metal_args_glm_qkv_norm_store_compact_kv { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t q_n; + uint32_t q_n4; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t kv_lora_n4; + uint32_t qk_rope; + uint32_t cache_f16; + float eps; + uint32_t pad0; +}; + +struct ds4_metal_args_glm_store_indexer_k { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + float pad1; +}; + +struct ds4_metal_args_glm_attention_full { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_len; + uint32_t cache_cap; + uint32_t n_head; + uint32_t qk_dim; + uint32_t value_dim; + uint32_t pad0; + uint32_t cache_f16; + uint32_t pad1; + uint32_t pad2; + float scale; +}; + +struct ds4_metal_args_glm_fill_selected_range { + uint32_t n_selected; +}; + +struct ds4_metal_args_glm_fill_selected_range_batch { + uint32_t n_tokens; + uint32_t pos0; + uint32_t n_selected; + uint32_t pad_row; +}; + +struct ds4_metal_args_glm_indexer_rope_tail { + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t rot_offset; + uint32_t pos0; + uint32_t n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + +struct ds4_metal_args_glm_indexer_score_one { + uint32_t n_rows; + uint32_t n_head; + uint32_t head_dim; + uint32_t cache_f16; + float scale; +}; + +struct ds4_metal_args_glm_indexer_scores_batch { + uint32_t n_rows; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t cache_f16; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t score_token_stride; + float scale; +}; + +struct ds4_metal_args_glm_qk_lowrank { + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +}; + +struct ds4_metal_args_glm_qk_lowrank_batch { + uint32_t n_tokens; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + /* First head this dispatch computes: under tensor-parallel head split + * each rank covers a contiguous half of the heads; buffers and weights + * keep full-model layout and are indexed by absolute head. */ + uint32_t head_base; +}; + +struct ds4_metal_args_glm_attention_indexed_decode { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +}; + +struct ds4_metal_args_glm_attention_indexed_decode_split { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + uint32_t block_rows; + uint32_t n_blocks; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +}; + +struct ds4_metal_args_glm_attention_indexed_batch { + uint32_t n_tokens; + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t pos0; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t head_base; +}; + +struct ds4_metal_args_dsv4_directional_steering_project { + uint32_t width; + uint32_t rows; + uint32_t layer; + uint32_t n_threads; + float scale; +}; + +// Optional directional steering projection. +// +// Each threadgroup owns one 4096-wide token row, computes +// dot(row, direction[layer]), then subtracts scale * direction * dot in-place. +// Positive scales remove a concept direction; negative scales amplify it. The +// kernel is not used unless a steering file and nonzero scale are provided. diff --git a/metal/model_io.inc b/metal/model_io.inc new file mode 100644 index 0000000000..b099e7fb3d --- /dev/null +++ b/metal/model_io.inc @@ -0,0 +1,401 @@ +int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0) return 0; + if (map_offset > model_size || map_size == 0 || map_size > model_size - map_offset) return 0; + max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, max_tensor_bytes); + + @autoreleasepool { + if (g_model_map_ptr == model_map && + g_model_map_size == model_size && + g_model_mapped_offset == map_offset && + g_model_mapped_size == map_size && + g_model_mapped_max_tensor_bytes == max_tensor_bytes) { + return 1; + } + + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size && + map_offset >= g_model_views[i].model_offset && + map_offset + map_size <= g_model_views[i].model_offset + g_model_views[i].bytes) { + return 1; + } + } + + ds4_gpu_model_residency_clear(); + if (!ds4_gpu_map_model_views(model_map, model_size, map_offset, map_size, max_tensor_bytes)) { + ds4_gpu_model_residency_clear(); + return 0; + } + g_model_map_ptr = model_map; + g_model_map_size = model_size; + g_model_mapped_offset = map_offset; + g_model_mapped_size = map_size; + g_model_mapped_max_tensor_bytes = max_tensor_bytes; + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal mapped mmaped model as %u overlapping shared buffers\n", + g_model_view_count); + } + return 1; + } +} + +static int ds4_gpu_model_views_cover_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count); + +int ds4_gpu_set_model_map_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count, + uint64_t max_tensor_bytes) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0; + if (count == 1) { + return ds4_gpu_set_model_map_range(model_map, + model_size, + offsets[0], + sizes[0], + max_tensor_bytes); + } + if (ds4_gpu_model_views_cover_spans(model_map, model_size, offsets, sizes, count)) { + return 1; + } + + @autoreleasepool { + const double t0 = ds4_gpu_now_ms(); + max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(model_size, max_tensor_bytes); + + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + + uint64_t mapped_total = 0; + uint64_t first_offset = UINT64_MAX; + for (uint32_t i = 0; i < count; i++) { + if (offsets[i] > model_size || sizes[i] == 0 || sizes[i] > model_size - offsets[i]) { + fprintf(stderr, "ds4: Metal model span %u is outside the GGUF mapping\n", i); + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + if (offsets[i] < first_offset) first_offset = offsets[i]; + uint64_t effective_max = max_tensor_bytes; + if (effective_max > sizes[i]) effective_max = sizes[i]; + if (!ds4_gpu_add_model_view_range(model_map, + model_size, + offsets[i], + sizes[i], + effective_max, + true, + &mapped_total)) { + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + } + if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) { + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + g_model_map_ptr = model_map; + g_model_map_size = model_size; + g_model_mapped_offset = first_offset == UINT64_MAX ? 0 : first_offset; + g_model_mapped_size = mapped_total; + g_model_mapped_max_tensor_bytes = max_tensor_bytes; + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal mapped mmaped model as %u disjoint shared buffers across %u tensor spans\n", + g_model_view_count, + count); + } + return 1; + } +} + +int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) { + return ds4_gpu_set_model_map_range(model_map, model_size, 0, model_size, 0); +} + +int ds4_gpu_set_model_fd(int fd) { + g_model_fd = fd; + return 1; +} + +static int ds4_gpu_model_views_cover_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t size) { + if (!model_map || model_size == 0 || size == 0 || + offset > model_size || size > model_size - offset) { + return 0; + } + const uint64_t end = offset + size; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map != model_map || + g_model_views[i].model_size != model_size) { + continue; + } + const uint64_t view_start = g_model_views[i].model_offset; + const uint64_t view_end = view_start + g_model_views[i].bytes; + if (offset >= view_start && end <= view_end) return 1; + } + return 0; +} + +static int ds4_gpu_model_views_cover_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count) { + if (!offsets || !sizes || count == 0) return 0; + for (uint32_t i = 0; i < count; i++) { + if (!ds4_gpu_model_views_cover_range(model_map, + model_size, + offsets[i], + sizes[i])) { + return 0; + } + } + return 1; +} + +int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) { + (void)fd; + (void)model_map; + return 1; +} + +static id ds4_gpu_wrap_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + (void)model_map; + if (model_size == 0 || offset > model_size || len > model_size - offset) { + fprintf(stderr, "ds4: Metal model range is outside the mapped model\n"); + return nil; + } + + const uint64_t end = offset + len; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map != model_map || + g_model_views[i].model_size != model_size) { + continue; + } + const uint64_t view_start = g_model_views[i].model_offset; + const uint64_t view_end = view_start + g_model_views[i].bytes; + if (offset >= view_start && end <= view_end) { + *inner_offset = offset - view_start; + return g_model_views[i].buffer; + } + } + + fprintf(stderr, + "ds4: Metal model range %.2f..%.2f GiB is not covered by mapped model views\n", + ds4_gpu_gib(offset), + ds4_gpu_gib(end)); + return nil; +} + +typedef enum { + DS4_GPU_EXACT_VIEW_CACHED, + DS4_GPU_EXACT_VIEW_TRANSIENT, + DS4_GPU_EXACT_VIEW_OWNED, +} ds4_gpu_exact_view_lifetime; + +static id ds4_gpu_wrap_model_exact_range_impl( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset, + ds4_gpu_exact_view_lifetime lifetime) { + const bool cache_view = lifetime == DS4_GPU_EXACT_VIEW_CACHED; + const bool transient_view = lifetime == DS4_GPU_EXACT_VIEW_TRANSIENT; + if (!model_map || !g_device || + (cache_view && !g_model_buffer_cache) || + (transient_view && !g_transient_buffers) || + model_size == 0 || offset > model_size || len > model_size - offset) { + fprintf(stderr, "ds4: Metal exact model range is outside the mapped model\n"); + return nil; + } + + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t page_offset = offset & ~(page - 1); + const uint64_t leading = offset - page_offset; + if (len > UINT64_MAX - leading || + leading + len > UINT64_MAX - (page - 1)) { + fprintf(stderr, "ds4: Metal exact model range overflows page alignment\n"); + return nil; + } + uint64_t view_bytes = round_up_u64(leading + len, page); + if (view_bytes > model_size - page_offset) view_bytes = model_size - page_offset; + if (leading + len > view_bytes) { + fprintf(stderr, "ds4: Metal exact model range alignment exceeds mapped model\n"); + return nil; + } + if (view_bytes > (uint64_t)[g_device maxBufferLength]) { + fprintf(stderr, + "ds4: Metal exact model range %.2f GiB exceeds maxBufferLength %.2f GiB\n", + ds4_gpu_gib(view_bytes), + ds4_gpu_gib((uint64_t)[g_device maxBufferLength])); + return nil; + } + + NSString *key = nil; + id buffer = nil; + if (cache_view) { + key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu", + model_map, + (unsigned long long)model_size, + (unsigned long long)page_offset, + (unsigned long long)view_bytes]; + buffer = [g_model_buffer_cache objectForKey:key]; + } + if (!buffer) { + const uintptr_t base = (uintptr_t)model_map; + buffer = [g_device newBufferWithBytesNoCopy:(void *)(base + page_offset) + length:(NSUInteger)view_bytes + options:ds4_gpu_model_resource_options() + deallocator:nil]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal could not wrap exact mmaped model range at %.2f GiB, size %.2f MiB\n", + ds4_gpu_gib(page_offset), + ds4_gpu_mib(view_bytes)); + return nil; + } + if (cache_view) { + buffer.label = @"ds4_model_exact_view"; + } else if (transient_view) { + buffer.label = @"ds4_model_exact_transient_view"; + } else { + buffer.label = @"ds4_model_exact_owned_view"; + } + if (cache_view) { + [g_model_buffer_cache setObject:buffer forKey:key]; + ds4_gpu_model_buffer_cache_note_insert(view_bytes); + } else if (transient_view) { + [g_transient_buffers addObject:buffer]; + } + } + + if (inner_offset) *inner_offset = leading; + return buffer; +} + +static id ds4_gpu_wrap_model_exact_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_CACHED); +} + +static id ds4_gpu_wrap_model_exact_range_transient( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_TRANSIENT); +} + +static id ds4_gpu_wrap_model_exact_range_owned( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_OWNED); +} + +static id ds4_gpu_wrap_q8_decode_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t n_tokens, + uint64_t *inner_offset) { + const uint64_t exact_decode_max_mib = + ds4_gpu_env_u64("DS4_METAL_Q8_DECODE_EXACT_VIEW_MAX_MIB", + 1024u, + 1u, + 4096u); + const uint64_t exact_decode_max_bytes = + exact_decode_max_mib * 1024ull * 1024ull; + const bool exact_decode_weight_view = + n_tokens == 1u && + len <= exact_decode_max_bytes && + getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; + return exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + offset, + len, + inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + offset, + len, + inner_offset); +} + +static id ds4_gpu_wrap_f32_decode_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t n_tokens, + uint64_t *inner_offset) { + const uint64_t exact_decode_max_mib = + ds4_gpu_env_u64("DS4_METAL_F32_DECODE_EXACT_VIEW_MAX_MIB", + 64u, + 1u, + 4096u); + const uint64_t exact_decode_max_bytes = + exact_decode_max_mib * 1024ull * 1024ull; + const bool exact_decode_weight_view = + n_tokens == 1u && + len <= exact_decode_max_bytes && + getenv("DS4_METAL_ENABLE_F32_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_F32_DECODE_EXACT_VIEWS") == NULL; + return exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + offset, + len, + inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + offset, + len, + inner_offset); +} diff --git a/metal/moe_dispatch.inc b/metal/moe_dispatch.inc new file mode 100644 index 0000000000..8685eded0e --- /dev/null +++ b/metal/moe_dispatch.inc @@ -0,0 +1,2796 @@ +static NSUInteger ds4_gpu_bin_threads(uint32_t width, id pipeline) { + NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = 1u; + while (2u * nth < (NSUInteger)width && nth < nth_max) nth *= 2u; + return nth ? nth : 1u; +} + +static int ds4_gpu_encode_unary_f32_rows( + id cb, + id pipeline, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows, + int c4, + float min, + float max) { + if (!cb || !pipeline || !src || !dst || width == 0 || rows == 0) return 0; + if (c4 && (width & 3u) != 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, c4, 0.0f, 0.0f); + args.min = min; + args.max = max; + + NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = (NSUInteger)args.ne00; + if (nth > nth_max) nth = nth_max; + if (nth == 0) nth = 1u; + const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, + (NSUInteger)args.ne02, + (NSUInteger)args.ne03) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_bin_f32_rows( + id cb, + id pipeline, + const ds4_gpu_bin_args *args, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off) { + if (!cb || !pipeline || !args || !a || !b || !out || args->ne0 <= 0 || args->ne1 <= 0) { + return 0; + } + + const NSUInteger nth = ds4_gpu_bin_threads((uint32_t)args->ne0, pipeline); + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:a offset:a_off atIndex:1]; + [enc setBuffer:b offset:b_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)args->ne1, + (NSUInteger)args->ne2, + (NSUInteger)args->ne3) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_rowwise_scalar_args(uint32_t width, uint32_t rows) { + const uint64_t lhs_row_bytes = (uint64_t)width * sizeof(float); + const uint64_t rhs_row_bytes = sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)width, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = lhs_row_bytes, + .nb02 = (uint64_t)rows * lhs_row_bytes, + .nb03 = (uint64_t)rows * lhs_row_bytes, + .ne10 = 1, + .ne11 = (int32_t)rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = rhs_row_bytes, + .nb12 = (uint64_t)rows * rhs_row_bytes, + .nb13 = (uint64_t)rows * rhs_row_bytes, + .ne0 = (int32_t)width, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = lhs_row_bytes, + .nb2 = (uint64_t)rows * lhs_row_bytes, + .nb3 = (uint64_t)rows * lhs_row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static ds4_gpu_mul_mv_id_args ds4_gpu_make_mul_mv_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t nr0) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); + const uint64_t src0_blocks = src0_cols / 256u; + const uint64_t src0_block_bytes = src0_blocks ? src0_row_bytes / src0_blocks : 1u; + return (ds4_gpu_mul_mv_id_args) { + .nei0 = (int32_t)selected_experts, + .nei1 = (int32_t)n_tokens, + .nbi1 = (uint64_t)selected_experts * sizeof(int32_t), + .ne00 = (int32_t)src0_cols, + .ne01 = (int32_t)src0_rows, + .ne02 = (int32_t)src0_experts, + .nb00 = src0_block_bytes, + .nb01 = src0_row_bytes, + .nb02 = src0_expert_bytes, + .ne10 = (int32_t)src0_cols, + .ne11 = (int32_t)src1_expert_rows, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .ne0 = (int32_t)src0_rows, + .ne1 = (int32_t)selected_experts, + .nb1 = (uint64_t)src0_rows * sizeof(float), + .nr0 = (int32_t)nr0, + }; +} + +static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( + uint32_t src0_cols, + uint32_t src0_experts, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); + return (ds4_gpu_mul_mm_id_map_args) { + .ne02 = (int32_t)src0_experts, + .ne10 = (int32_t)src0_cols, + .ne11 = (int32_t)src1_expert_rows, + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .ne21 = (int32_t)n_tokens, + .ne20 = (int32_t)selected_experts, + .nb21 = (uint64_t)selected_experts * sizeof(int32_t), + }; +} + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens) { + return ds4_gpu_make_mul_mm_id_args_src1_size(src0_cols, + src0_rows, + src0_experts, + src0_row_bytes, + src0_expert_bytes, + src1_expert_rows, + selected_experts, + n_tokens, + sizeof(float)); +} + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t src1_elem_size) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * src1_elem_size; + return (ds4_gpu_mul_mm_id_args) { + .ne00 = (int32_t)src0_cols, + .ne02 = (int32_t)src0_experts, + .nb01 = src0_row_bytes, + .nb02 = src0_expert_bytes, + .nb03 = (uint64_t)src0_experts * src0_expert_bytes, + .ne11 = (int32_t)src1_expert_rows, + .nb10 = src1_elem_size, + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .nb13 = (uint64_t)n_tokens * (uint64_t)src1_expert_rows * src1_row_bytes, + .ne20 = (int32_t)selected_experts, + .ne21 = (int32_t)n_tokens, + .ne0 = (int32_t)src0_rows, + .ne1 = (int32_t)selected_experts, + .r2 = 1, + .r3 = 1, + }; +} + +static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: return 2; + case DS4_METAL_TENSOR_Q8_K: return 2; + case DS4_METAL_TENSOR_Q4_K: return 2; + case DS4_METAL_TENSOR_Q2_K: + case DS4_METAL_TENSOR_IQ2_XXS: return 4; + default: return 0; + } +} + +static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_IQ2_XXS: return "iq2_xxs"; + case DS4_METAL_TENSOR_Q2_K: return "q2_k"; + case DS4_METAL_TENSOR_Q4_K: return "q4_k"; + case DS4_METAL_TENSOR_Q5_K: return "q5_k"; + case DS4_METAL_TENSOR_Q6_K: return "q6_k"; + default: return "unknown"; + } +} + +static const char *ds4_gpu_trim_env_value(const char *env, size_t *len_out) { + if (len_out) *len_out = 0; + if (!env) return NULL; + + while (isspace((unsigned char)*env)) env++; + size_t n = strlen(env); + while (n > 0 && isspace((unsigned char)env[n - 1])) n--; + if (len_out) *len_out = n; + return env; +} + +static bool ds4_gpu_profile_layer_value_match(const char *env, uint32_t layer_index) { + size_t env_len = 0; + env = ds4_gpu_trim_env_value(env, &env_len); + if (!env || env_len == 0) return true; + if (ds4_gpu_env_value_eq(env, env_len, "all")) return true; + + const char *p = env; + const char *end_env = env + env_len; + while (p < end_env) { + while (p < end_env && (*p == ' ' || *p == '\t' || *p == ',')) p++; + if (p >= end_env) break; + + char *end = NULL; + const unsigned long first = strtoul(p, &end, 10); + if (end == p || end > end_env || first > UINT32_MAX) return false; + + unsigned long last = first; + p = end; + if (p < end_env && *p == '-') { + p++; + last = strtoul(p, &end, 10); + if (end == p || end > end_env || last > UINT32_MAX) return false; + p = end; + } + + if (first <= layer_index && layer_index <= last) return true; + while (p < end_env && (*p == ' ' || *p == '\t')) p++; + if (p < end_env && *p != ',') return false; + } + return false; +} + +static bool ds4_gpu_stage_profile_enabled_for_layer(const char *flag_env_name, + const char *layer_env_name, + uint32_t layer_index) { + size_t flag_len = 0; + const char *flag = ds4_gpu_trim_env_value(getenv(flag_env_name), &flag_len); + if (!flag) return false; + + size_t layer_len = 0; + const char *layer = ds4_gpu_trim_env_value(getenv(layer_env_name), &layer_len); + const bool has_layer_filter = layer && layer_len != 0; + + if (flag_len != 0) { + if (ds4_gpu_env_value_eq(flag, flag_len, "0") || + ds4_gpu_env_value_eq(flag, flag_len, "false") || + ds4_gpu_env_value_eq(flag, flag_len, "no") || + ds4_gpu_env_value_eq(flag, flag_len, "off")) { + return false; + } + if (!has_layer_filter && + !ds4_gpu_env_value_eq(flag, flag_len, "1") && + !ds4_gpu_env_value_eq(flag, flag_len, "true") && + !ds4_gpu_env_value_eq(flag, flag_len, "yes") && + !ds4_gpu_env_value_eq(flag, flag_len, "on") && + !ds4_gpu_env_value_eq(flag, flag_len, "all")) { + return ds4_gpu_profile_layer_value_match(flag, layer_index); + } + } + + return ds4_gpu_profile_layer_value_match(layer, layer_index); +} + +static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) { + if (type == DS4_METAL_TENSOR_Q8_0) { + return 32u * 2u * sizeof(float); + } + if (type == DS4_METAL_TENSOR_IQ2_XXS) { + return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t); + } + return 0; +} + +static NSUInteger ds4_gpu_routed_mv_nsg(uint32_t type) { + return type == DS4_METAL_TENSOR_Q8_0 ? 4u : 2u; +} + +static bool ds4_gpu_routed_mv_rows_per_group_is_nr0(uint32_t type) { + return type == DS4_METAL_TENSOR_Q8_0; +} + +static id ds4_gpu_routed_mv_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_K_f32", 2); + case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline; + case DS4_METAL_TENSOR_Q2_K: return g_moe_mul_mv_id_q2_k_pipeline; + case DS4_METAL_TENSOR_Q4_K: return g_moe_mul_mv_id_q4_k_pipeline; + default: return nil; + } +} + +static id ds4_gpu_routed_mm_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f32", false); + case DS4_METAL_TENSOR_IQ2_XXS: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32", false); + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f32", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); + case DS4_METAL_TENSOR_Q5_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false); + case DS4_METAL_TENSOR_Q6_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_addr_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f16", false); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f16", false); + case DS4_METAL_TENSOR_IQ2_XXS: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f16", false); + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f16", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16", false); + case DS4_METAL_TENSOR_Q5_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false); + case DS4_METAL_TENSOR_Q6_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false); + default: + return nil; + } +} + +static int ds4_gpu_encode_mul_mv_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_attn_out_low_q8_direct( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + /* Two row conventions in the classic matvec family: Q8_0 k-splits one + * nr0-row group across all nsg simdgroups (cross-simdgroup reduce), so a + * threadgroup covers nr0 rows; Q4_K gives each simdgroup its own nr0 + * rows, covering nr0*nsg. Dispatching Q8 with the Q4 stride leaves + * (nsg-1)/nsg of the output rows unwritten. */ + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? + (NSUInteger)args->nr0 : (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_pair( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:1]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:4]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:5]; + [enc setBuffer:ids offset:ids_off atIndex:6]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || + !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:2]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_table_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0, + bool queue_residency) { + if (!cb || !pipeline || !args || !act || !gate_table || !up_table || + !gate_table.argumentBuffer || !up_table.argumentBuffer || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_table.argumentBuffer offset:0 atIndex:2]; + [enc setBuffer:up_table.argumentBuffer offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || + !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, queue_residency); + ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, queue_residency); + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !gate_table || !up_table || + !gate_table.addressBuffer || !up_table.addressBuffer || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_table.addressBuffer offset:0 atIndex:2]; + [enc setBuffer:up_table.addressBuffer offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || + !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { + ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, false); + ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, false); + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_table_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + DS4MetalQ4ExpertTable *table, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool queue_residency) { + if (!cb || !pipeline || !args || !table || !table.argumentBuffer || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:table.argumentBuffer offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + ds4_gpu_use_q4_expert_table_resources(cb, enc, table, queue_residency); + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + DS4MetalQ4ExpertTable *table, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !table || !table.addressBuffer || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:table.addressBuffer offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { + ds4_gpu_use_q4_expert_table_resources(cb, enc, table, false); + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static uint32_t ds4_gpu_q4_expert_group_size(uint32_t n_total_expert) { + uint32_t group_size = 32; + const char *env = getenv("DS4_METAL_Q4_EXPERT_GROUP_SIZE"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0' && v > 0 && v <= UINT32_MAX) { + group_size = (uint32_t)v; + } + } + if (group_size == 0) group_size = 1; + if (group_size > n_total_expert) group_size = n_total_expert; + return group_size; +} + +static int ds4_gpu_encode_mul_mv_group_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + const ds4_gpu_moe_expert_group_args *group, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !group || + !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0 || + group->expert_count == 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBytes:group length:sizeof(*group) atIndex:2]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:3]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:4]; + [enc setBuffer:src1 offset:src1_off atIndex:5]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + [enc setBuffer:weights offset:weights_off atIndex:10]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_moe_expert_group_args *group, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !group || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + group->expert_count == 0) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:group length:sizeof(*group) atIndex:1]; + [enc setBuffer:src0 offset:src0_off atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setBuffer:ids offset:ids_off atIndex:5]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + id add_in, + NSUInteger add_in_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + [enc setBuffer:(add_in ? add_in : dst) offset:(add_in ? add_in_off : dst_off) atIndex:5]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_q4_gather_slots6( + id cb, + id pipeline, + const ds4_gpu_q4_gather_slots6_args *args, + __unsafe_unretained id src_groups[6], + const NSUInteger src_group_offsets[6], + id ids, + NSUInteger ids_off, + id dst, + NSUInteger dst_off) { + if (!cb || !pipeline || !args || !src_groups || !src_group_offsets || !ids || !dst || + args->expert_bytes == 0 || (args->expert_bytes & 15u) != 0 || + args->group_size == 0 || args->n_slots == 0 || args->n_slots > 6) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src_groups[i]) return 0; + } + + const uint64_t chunks = args->expert_bytes >> 4; + if (chunks == 0 || chunks > NSUIntegerMax) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src_groups[i] offset:src_group_offsets[i] atIndex:1 + i]; + } + [enc setBuffer:ids offset:ids_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + const NSUInteger threads = 256u; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)chunks + threads - 1u) / threads, + (NSUInteger)args->n_slots, + 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_slots6_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[6], + const NSUInteger src0_a_off[6], + __unsafe_unretained id src0_b[6], + const NSUInteger src0_b_off[6], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:14]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; + [enc setBuffer:weights offset:weights_off atIndex:18]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_slots6_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[6], + const NSUInteger src0_off[6], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group6_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[6], + const NSUInteger src0_a_off[6], + __unsafe_unretained id src0_b[6], + const NSUInteger src0_b_off[6], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:14]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; + [enc setBuffer:ids offset:ids_off atIndex:18]; + [enc setBuffer:weights offset:weights_off atIndex:19]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group6_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[6], + const NSUInteger src0_off[6], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id gate_addrs, + id up_addrs, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0, + id overflow_gate, + id overflow_up) { + if (!cb || !pipeline || !args || !act || !entries || + (n_entries == 0 && !overflow_gate) || + !gate_addrs || !up_addrs || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || + args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_addrs offset:0 atIndex:2]; + [enc setBuffer:up_addrs offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + /* Overflow experts are addressed straight into the mapped model views + * when a layer's unique selected set exceeds the cache budget. */ + if (overflow_gate) [enc useResource:overflow_gate usage:MTLResourceUsageRead]; + if (overflow_up) [enc useResource:overflow_up usage:MTLResourceUsageRead]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !entries || n_entries == 0 || + n_entries > DS4_METAL_MAX_ROUTED_EXPERT_USED || + !addrs || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || + args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q2_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + id overflow_down) { + if (!cb || !pipeline || !args || !entries || + (n_entries == 0 && !overflow_down) || + !addrs || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (overflow_down) [enc useResource:overflow_down usage:MTLResourceUsageRead]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + const ds4_gpu_stream_expert_split_args *split, + ds4_gpu_stream_expert_cache_entry * const entries[6], + id gate_addrs, + id up_addrs, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !split || !entries || + !gate_addrs || !up_addrs || !src1 || !dst_a || !dst_b || !dst_mid || + !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + 6, + split->active_mask)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBytes:split length:sizeof(*split) atIndex:2]; + [enc setBuffer:gate_addrs offset:0 atIndex:3]; + [enc setBuffer:up_addrs offset:0 atIndex:4]; + [enc setBuffer:src1 offset:src1_off atIndex:5]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + [enc setBuffer:weights offset:weights_off atIndex:10]; + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q2_sum6_masked( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_stream_expert_split_args *split, + ds4_gpu_stream_expert_cache_entry * const entries[6], + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !split || !entries || !addrs || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + 6, + split->active_mask)) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:split length:sizeof(*split) atIndex:1]; + [enc setBuffer:addrs offset:0 atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setBuffer:ids offset:ids_off atIndex:5]; + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group8_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[8], + const NSUInteger src0_a_off[8], + __unsafe_unretained id src0_b[8], + const NSUInteger src0_b_off[8], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 8; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:10 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:18]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:19]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:20]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:21]; + [enc setBuffer:ids offset:ids_off atIndex:22]; + [enc setBuffer:weights offset:weights_off atIndex:23]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group8_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[8], + const NSUInteger src0_off[8], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 8; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:9]; + [enc setBuffer:dst offset:dst_off atIndex:10]; + [enc setBuffer:ids offset:ids_off atIndex:11]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group24_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[24], + const NSUInteger src0_off[24], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 24; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 24; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:25]; + [enc setBuffer:dst offset:dst_off atIndex:26]; + [enc setBuffer:ids offset:ids_off atIndex:27]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group24_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[24], + const NSUInteger src0_off[24], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 24; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 24; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:25]; + [enc setBuffer:dst offset:dst_off atIndex:26]; + [enc setBuffer:ids offset:ids_off atIndex:27]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id( + id cb, + id map_pipeline, + id mm_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off) { + if (!cb || !map_pipeline || !mm_pipeline || !map_args || !mm_args || + !src0 || !src1 || !dst || !ids || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + return ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + map_args, + mm_args, + ids, + ids_off) && + ds4_gpu_encode_mul_mm_id_mapped(cb, + mm_pipeline, + mm_args, + src0, + src0_off, + src1, + src1_off, + dst, + dst_off); +} + +static int ds4_gpu_encode_mul_mm_id_map( + id cb, + id map_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id ids, + NSUInteger ids_off) { + if (!cb || !map_pipeline || !map_args || !mm_args || !ids || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0; + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer, + &g_moe_id_map_bytes, + tpe_bytes + hids_bytes, + "ds4_moe_id_map")) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:map_pipeline]; + [enc setBytes:map_args length:sizeof(*map_args) atIndex:0]; + [enc setBuffer:ids offset:ids_off atIndex:1]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3]; + [enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes) { + if (!cb || !mm_pipeline || !mm_args || !src0 || !src1 || !dst || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + /* + * The routed MoE grouped matmul uses the legacy 32-token expert-major tile. + * The removed TensorOps variant was not semantically stable on evals, so keep + * this encoder tied to the tested simdgroup kernel shape. + */ + const NSUInteger tile_n = 32u; + const bool use_resource_hints = + getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL && + getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL; + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; + [enc setBuffer:dst offset:dst_off atIndex:5]; + if (use_resource_hints) { + [enc useResource:src0 usage:MTLResourceUsageRead]; + } + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0_addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count, + uint32_t resource_kind, + id overflow_resource) { + if (!cb || !mm_pipeline || !mm_args || !src0_addrs || !src1 || !dst || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tile_n = 32u; + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = + (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0_addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; + [enc setBuffer:dst offset:dst_off atIndex:5]; + [enc useResource:src0_addrs usage:MTLResourceUsageRead]; + for (uint32_t i = 0; resources && i < resource_count; i++) { + ds4_gpu_stream_expert_cache_entry *entry = resources[i]; + if (!entry) continue; + id b = + resource_kind == 0 ? entry->gate_buffer : + resource_kind == 1 ? entry->up_buffer : + entry->down_buffer; + if (b) [enc useResource:b usage:MTLResourceUsageRead]; + } + if (overflow_resource) { + [enc useResource:overflow_resource usage:MTLResourceUsageRead]; + } + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args, + id gate_src0, + NSUInteger gate_src0_off, + id up_src0, + NSUInteger up_src0_off, + id src1, + NSUInteger src1_off, + id mid, + NSUInteger mid_off, + id weights, + NSUInteger weights_off) { + if (!cb || !pipeline || !mm_args || !act_args || + !gate_src0 || !up_src0 || !src1 || !mid || !weights || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBytes:act_args length:sizeof(*act_args) atIndex:1]; + [enc setBuffer:gate_src0 offset:gate_src0_off atIndex:2]; + [enc setBuffer:up_src0 offset:up_src0_off atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:5]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6]; + [enc setBuffer:mid offset:mid_off atIndex:7]; + [enc setBuffer:weights offset:weights_off atIndex:8]; + [enc setThreadgroupMemoryLength:16384u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + 31u) / 32u, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_mapped( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off) { + return ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + mm_pipeline, + mm_args, + src0, + src0_off, + src1, + src1_off, + dst, + dst_off, + 8192u); +} + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off) { + if (!cb || !pipeline || !mm_args || !src0 || !src1 || !dst || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne02 <= 0 || mm_args->ne1 <= 0 || mm_args->ne21 <= 0) { + return 0; + } + + const uint32_t tile_n = DS4_METAL_ATTN_OUT_MPP_TILE_N; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setThreadgroupMemoryLength:8192u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_swiglu_flat( + id cb, + id gate, + NSUInteger gate_off, + id up, + NSUInteger up_off, + id out, + NSUInteger out_off, + uint32_t n) { + if (!cb || !gate || !up || !out || n == 0) return 0; + + ds4_gpu_glu_args args = { + .ne00 = (int32_t)n, + .nb01 = (uint64_t)n * sizeof(float), + .ne10 = (int32_t)n, + .nb11 = (uint64_t)n * sizeof(float), + .ne0 = (int32_t)n, + .nb1 = (uint64_t)n * sizeof(float), + .i00 = 0, + .i10 = 0, + .alpha = 1.0f, + .limit = 0.0f, + }; + NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_swiglu_flat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate offset:gate_off atIndex:1]; + [enc setBuffer:up offset:up_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_swiglu_weight( + id cb, + id gate, + NSUInteger gate_off, + id up, + NSUInteger up_off, + id mid, + NSUInteger mid_off, + id weights, + NSUInteger weights_off, + uint32_t width, + uint32_t rows, + float clamp_value, + bool mid_f16) { + if (!cb || !gate || !up || !mid || !weights || width == 0 || rows == 0) return 0; + + id pipeline = + ds4_gpu_get_pipeline(mid_f16 ? "kernel_dsv4_moe_swiglu_weight_f16" : + "kernel_dsv4_moe_swiglu_weight"); + if (!pipeline) return 0; + + ds4_gpu_dsv4_moe_swiglu_weight_args args = { + .width = width, + .rows = rows, + .gate_row_stride = (uint64_t)width * sizeof(float), + .up_row_stride = (uint64_t)width * sizeof(float), + .mid_row_stride = (uint64_t)width * (mid_f16 ? sizeof(uint16_t) : sizeof(float)), + .weight_stride = sizeof(float), + .write_clamped = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ? 1u : 0u, + .clamp_value = clamp_value, + }; + + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > width) nth = width; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate offset:gate_off atIndex:1]; + [enc setBuffer:up offset:up_off atIndex:2]; + [enc setBuffer:mid offset:mid_off atIndex:3]; + [enc setBuffer:weights offset:weights_off atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_sum6( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; + + if (!g_moe_sum6_pipeline) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + ds4_gpu_dsv4_moe_sum6_args args = { + .width = out_dim, + .tokens = n_tokens, + .src_token_stride = 6u * out_row_bytes, + .dst_token_stride = out_row_bytes, + }; + + NSUInteger nth = g_moe_sum6_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > out_dim) nth = out_dim; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_sum6_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:experts offset:experts_off atIndex:1]; + [enc setBuffer:out offset:out_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_sum8( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; + + if (!g_moe_sum8_pipeline) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + ds4_gpu_dsv4_moe_sum6_args args = { + .width = out_dim, + .tokens = n_tokens, + .src_token_stride = 8u * out_row_bytes, + .dst_token_stride = out_row_bytes, + }; + + NSUInteger nth = g_moe_sum8_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > out_dim) nth = out_dim; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_sum8_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:experts offset:experts_off atIndex:1]; + [enc setBuffer:out offset:out_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_moe_add_args( + uint32_t out_dim, + uint32_t n_tokens, + uint64_t src0_token_stride, + uint64_t src1_token_stride, + uint64_t dst_token_stride) { + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)out_dim, + .ne01 = (int32_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src0_token_stride, + .nb02 = (uint64_t)n_tokens * src0_token_stride, + .nb03 = (uint64_t)n_tokens * src0_token_stride, + .ne10 = (int32_t)out_dim, + .ne11 = (int32_t)n_tokens, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_token_stride, + .nb12 = (uint64_t)n_tokens * src1_token_stride, + .nb13 = (uint64_t)n_tokens * src1_token_stride, + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_token_stride, + .nb2 = (uint64_t)n_tokens * dst_token_stride, + .nb3 = (uint64_t)n_tokens * dst_token_stride, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_moe_sum_experts( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_expert < 2 || n_tokens == 0) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + const uint64_t expert_token_stride = (uint64_t)n_expert * out_row_bytes; + + if (n_expert == 6 && + ds4_gpu_encode_moe_sum6(cb, + experts, + experts_off, + out, + out_off, + out_dim, + n_tokens)) { + return 1; + } + + if (n_expert == 8 && + ds4_gpu_encode_moe_sum8(cb, + experts, + experts_off, + out, + out_off, + out_dim, + n_tokens)) { + return 1; + } + + ds4_gpu_bin_args first = + ds4_gpu_make_moe_add_args(out_dim, n_tokens, expert_token_stride, expert_token_stride, out_row_bytes); + if (!ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &first, + experts, + experts_off, + experts, + experts_off + (NSUInteger)out_row_bytes, + out, + out_off)) { + return 0; + } + + ds4_gpu_bin_args accum = + ds4_gpu_make_moe_add_args(out_dim, n_tokens, out_row_bytes, expert_token_stride, out_row_bytes); + for (uint32_t slot = 2; slot < n_expert; slot++) { + if (!ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &accum, + out, + out_off, + experts, + experts_off + (NSUInteger)((uint64_t)slot * out_row_bytes), + out, + out_off)) { + return 0; + } + } + return 1; +} + +static int ds4_gpu_encode_get_rows_i32_token_rows( + id cb, + id table, + NSUInteger table_off, + id tokens, + NSUInteger tokens_off, + const int32_t *token_inline, + id selected, + NSUInteger selected_off, + uint32_t hash_rows, + uint32_t n_cols, + uint32_t n_tokens) { + if (!cb || !table || !selected || hash_rows == 0 || n_cols == 0 || n_tokens == 0) return 0; + if (!tokens && !token_inline) return 0; + + const uint64_t table_row_bytes = (uint64_t)n_cols * sizeof(int32_t); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + ds4_gpu_get_rows_args args = { + .ne00t = (int64_t)n_cols, + .ne00 = (int64_t)n_cols, + .nb01 = table_row_bytes, + .nb02 = (uint64_t)hash_rows * table_row_bytes, + .nb03 = (uint64_t)hash_rows * table_row_bytes, + .ne10 = (int32_t)n_tokens, + .nb10 = sizeof(int32_t), + .nb11 = token_bytes, + .nb12 = token_bytes, + .nb1 = table_row_bytes, + .nb2 = (uint64_t)n_tokens * table_row_bytes, + .nb3 = (uint64_t)n_tokens * table_row_bytes, + }; + + NSUInteger nth = (NSUInteger)n_cols; + const NSUInteger max_threads = g_get_rows_i32_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1u; + const NSUInteger nw0 = ((NSUInteger)n_cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_i32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:table offset:table_off atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_off atIndex:2]; + } else { + [enc setBytes:token_inline length:sizeof(*token_inline) atIndex:2]; + } + [enc setBuffer:selected offset:selected_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_get_rows_f32_router_weights( + id cb, + id probs, + NSUInteger probs_off, + id selected, + NSUInteger selected_off, + id weights, + NSUInteger weights_off, + uint32_t n_expert, + uint32_t n_expert_used, + uint32_t n_tokens) { + if (!cb || !probs || !selected || !weights || n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; + + const uint64_t probs_token_bytes = (uint64_t)n_expert * sizeof(float); + const uint64_t selected_row_bytes = (uint64_t)n_expert_used * sizeof(int32_t); + const uint64_t weights_row_bytes = (uint64_t)n_expert_used * sizeof(float); + ds4_gpu_get_rows_args args = { + .ne00t = 1, + .ne00 = 1, + .nb01 = sizeof(float), + .nb02 = probs_token_bytes, + .nb03 = (uint64_t)n_tokens * probs_token_bytes, + .ne10 = (int64_t)n_expert_used, + .nb10 = sizeof(int32_t), + .nb11 = selected_row_bytes, + .nb12 = (uint64_t)n_tokens * selected_row_bytes, + .nb1 = sizeof(float), + .nb2 = weights_row_bytes, + .nb3 = (uint64_t)n_tokens * weights_row_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:probs offset:probs_off atIndex:1]; + [enc setBuffer:selected offset:selected_off atIndex:2]; + [enc setBuffer:weights offset:weights_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_expert_used, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_sum_rows_f32( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows) { + if (!cb || !src || !dst || width == 0 || rows == 0) return 0; + + const uint64_t src_row_bytes = (uint64_t)width * sizeof(float); + ds4_gpu_kargs_sum_rows args = { + .ne00 = (int64_t)width, + .ne01 = (int64_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_bytes, + .nb02 = (uint64_t)rows * src_row_bytes, + .nb03 = (uint64_t)rows * src_row_bytes, + .ne0 = 1, + .ne1 = (int64_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = sizeof(float), + .nb2 = (uint64_t)rows * sizeof(float), + .nb3 = (uint64_t)rows * sizeof(float), + }; + + NSUInteger nth = 32u; + const NSUInteger max_threads = g_sum_rows_f32_f32_pipeline.maxTotalThreadsPerThreadgroup; + while (nth < (NSUInteger)args.ne00 && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > (NSUInteger)args.ne00) nth = (NSUInteger)args.ne00; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_sum_rows_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_router_select( + id cb, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + id logitsbuf, + NSUInteger logits_off, + id biasbuf, + NSUInteger bias_off, + id hashbuf, + NSUInteger hash_off, + id tokensbuf, + NSUInteger tokens_off, + const int32_t *single_token, + uint32_t hash_rows, + uint32_t n_tokens, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + bool has_bias, + bool hash_mode) { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + const NSUInteger selected_off = ds4_gpu_tensor_offset(selected); + const NSUInteger weights_off = ds4_gpu_tensor_offset(weights); + const NSUInteger probs_off = ds4_gpu_tensor_offset(probs); + + if (!cb || !selectedbuf || !weightsbuf || !probsbuf || !logitsbuf || + n_tokens == 0 || n_expert == 0 || n_expert_used == 0) return 0; + + const NSUInteger probs_bytes = (NSUInteger)n_tokens * (NSUInteger)n_expert * sizeof(float); + const bool flash_router_fast_path = + n_expert == 256u && + n_expert_used == 6u && + fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f; + + int ok = 0; + if (flash_router_fast_path && + !g_quality_mode && n_tokens == 1 && + getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) { + const bool force_simd_weights_fusion = + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL; + const bool use_simd_finalize = + !hash_mode && + g_dsv4_router_finalize_one_simd_pipeline != nil && + g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL || + force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL; + const bool use_simd_weights_fusion = + use_simd_finalize && + g_dsv4_router_finalize_weights_one_simd_pipeline != nil && + g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL; + id softplus_sqrt_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, + "kernel_dsv4_softplus_sqrt_f32_4"); + id router_finalize_pipeline = + ds4_gpu_hot_pipeline( + use_simd_weights_fusion + ? g_dsv4_router_finalize_weights_one_simd_pipeline + : use_simd_finalize + ? g_dsv4_router_finalize_one_simd_pipeline + : g_dsv4_router_finalize_one_pipeline, + use_simd_weights_fusion + ? "kernel_dsv4_router_finalize_weights_one_simd" + : use_simd_finalize + ? "kernel_dsv4_router_finalize_one_simd" + : "kernel_dsv4_router_finalize_one"); + id router_weights_pipeline = use_simd_weights_fusion + ? nil + : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, + "kernel_dsv4_router_weights_one"); + if (!softplus_sqrt_pipeline || !router_finalize_pipeline || + (!use_simd_weights_fusion && !router_weights_pipeline)) return 0; + + ok = ds4_gpu_encode_unary_f32_rows(cb, + softplus_sqrt_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + 1, + 1, + 0.0f, + 0.0f); + if (!ok) return 0; + + const bool use_token_buffer = single_token == NULL; + ds4_gpu_dsv4_router_select_one_args args = { + .has_bias = has_bias ? 1u : 0u, + .hash_mode = hash_mode ? 1u : 0u, + .use_token_buffer = use_token_buffer ? 1u : 0u, + .token = single_token ? (uint32_t)*single_token : 0u, + .hash_rows = hash_rows, + }; + + const float zero_f32 = 0.0f; + const int32_t zero_i32 = 0; + if ((has_bias && !biasbuf) || + (hash_mode && !hashbuf) || + (use_token_buffer && !tokensbuf)) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_finalize_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + if (has_bias) { + [enc setBuffer:biasbuf offset:bias_off atIndex:2]; + } else { + [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2]; + } + if (hash_mode) { + [enc setBuffer:hashbuf offset:hash_off atIndex:3]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3]; + } + if (use_token_buffer) { + [enc setBuffer:tokensbuf offset:tokens_off atIndex:4]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; + } + [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; + if (use_simd_weights_fusion) { + [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; + } + const NSUInteger router_finalize_scratch_bytes = use_simd_finalize + ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t)) + : 256u * sizeof(float) + 256u * sizeof(int32_t); + [enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (use_simd_weights_fusion) return 1; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_weights_pipeline]; + [enc setBuffer:probsbuf offset:probs_off atIndex:0]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; + [enc dispatchThreads:MTLSizeMake(6, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { + id softplus_sqrt_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, + "kernel_dsv4_softplus_sqrt_f32_4"); + ok = softplus_sqrt_pipeline && + ds4_gpu_encode_unary_f32_rows(cb, + softplus_sqrt_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + 1, + 1, + 0.0f, + 0.0f); + } else { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_softplus_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + n_tokens, + 1, + 0.0f, + 0.0f) && + ds4_gpu_encode_unary_f32_rows(cb, + g_unary_sqrt_pipeline, + probsbuf, + probs_off, + probsbuf, + probs_off, + n_expert, + n_tokens, + 1, + 0.0f, + 0.0f); + } + if (!ok) return 0; + + if (hash_mode) { + ok = ds4_gpu_encode_get_rows_i32_token_rows(cb, + hashbuf, + hash_off, + tokensbuf, + tokens_off, + single_token, + selectedbuf, + selected_off, + hash_rows, + n_expert_used, + n_tokens); + } else { + ds4_gpu_tensor *score_tensor = probs; + DS4MetalTensor *selection_view = nil; + + if (has_bias) { + if (!biasbuf || + !ds4_gpu_ensure_scratch_buffer(&g_router_selection_buffer, + &g_router_selection_bytes, + probs_bytes, + "ds4_router_selection")) { + return 0; + } + + ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_expert, n_tokens, n_expert); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &add_args, + probsbuf, + probs_off, + biasbuf, + bias_off, + g_router_selection_buffer, + 0); + if (!ok) return 0; + + selection_view = [DS4MetalTensor new]; + selection_view.buffer = g_router_selection_buffer; + selection_view.offset = 0; + selection_view.bytes = probs_bytes; + selection_view.owner = 0; + score_tensor = (__bridge ds4_gpu_tensor *)selection_view; + } + + ok = ds4_gpu_indexer_topk_tensor(selected, score_tensor, n_expert, n_tokens, n_expert_used) != 0; + } + if (!ok) return 0; + + const bool use_batch_weights_fusion = + flash_router_fast_path && !g_quality_mode && n_tokens > 1u && + g_dsv4_router_weights_batch_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && + getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL; + if (use_batch_weights_fusion) { + const float scale = expert_weight_scale; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline]; + [enc setBytes:&scale length:sizeof(scale) atIndex:0]; + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:2]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; + [enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { + id router_weights_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, + "kernel_dsv4_router_weights_one"); + if (!router_weights_pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_weights_pipeline]; + [enc setBuffer:probsbuf offset:probs_off atIndex:0]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; + [enc dispatchThreads:MTLSizeMake(6, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer, + &g_router_weight_sum_bytes, + sum_bytes, + "ds4_router_weight_sum")) { + return 0; + } + + ok = ds4_gpu_encode_get_rows_f32_router_weights(cb, + probsbuf, + probs_off, + selectedbuf, + selected_off, + weightsbuf, + weights_off, + n_expert, + n_expert_used, + n_tokens) && + ds4_gpu_encode_sum_rows_f32(cb, + weightsbuf, + weights_off, + g_router_weight_sum_buffer, + 0, + n_expert_used, + n_tokens) && + ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + g_router_weight_sum_buffer, + 0, + g_router_weight_sum_buffer, + 0, + 1, + n_tokens, + 0, + 6.103515625e-5f, + ds4_gpu_positive_infinity()); + if (!ok) return 0; + + ds4_gpu_bin_args div_args = ds4_gpu_make_bin_rowwise_scalar_args(n_expert_used, n_tokens); + const float scale = expert_weight_scale; + ds4_gpu_bin_args scale_args = ds4_gpu_make_bin_rows_args(n_expert_used, n_tokens, 1); + + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_bin_div_row_pipeline, + &div_args, + weightsbuf, + weights_off, + g_router_weight_sum_buffer, + 0, + weightsbuf, + weights_off); + if (!ok) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; + [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:1]; + [enc setBytes:&scale length:sizeof(scale) atIndex:2]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)scale_args.ne1, + (NSUInteger)scale_args.ne2, + (NSUInteger)scale_args.ne3) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_bin_threads(n_expert_used, g_bin_mul_scalar_pipeline), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} diff --git a/metal/runtime.inc b/metal/runtime.inc new file mode 100644 index 0000000000..4f066b7000 --- /dev/null +++ b/metal/runtime.inc @@ -0,0 +1,9269 @@ +enum { + DS4_METAL_TENSOR_Q4_0 = 2, + DS4_METAL_TENSOR_Q8_0 = 8, + DS4_METAL_TENSOR_Q2_K = 10, + DS4_METAL_TENSOR_Q4_K = 12, + DS4_METAL_TENSOR_Q5_K = 13, + DS4_METAL_TENSOR_Q6_K = 14, + DS4_METAL_TENSOR_Q8_K = 15, + DS4_METAL_TENSOR_IQ2_XXS = 16, +}; + +@class DS4MetalQ4ExpertTable; + +static id g_device; +static id g_queue; +static id g_library; +static id g_batch_cb; +static id g_batch_enc; +static BOOL g_batch_has_work; +static NSMutableArray> *g_pending_cbs; +static id g_selected_readback_event; +static uint64_t g_selected_readback_event_value; +static id g_set_rows_f32_i32_pipeline; +static id g_get_rows_f32_pipeline; +static id g_get_rows_f16_pipeline; +static id g_get_rows_i32_pipeline; +static id g_get_rows_q8_0_pipeline; +static id g_get_rows_q4_0_pipeline; +static id g_get_rows_q4_K_pipeline; +static id g_repeat_f32_pipeline; +static id g_concat_pipeline; +static id g_cpy_f32_f32_pipeline; +static id g_cpy_f32_f16_pipeline; +static id g_cpy_contig_f32_f16_pipeline; +static id g_cpy_f16_f32_pipeline; +static id g_cpy_f16_f16_pipeline; +static id g_cpy_contig_f16_f32_pipeline; +static id g_cpy_contig_f16_f16_pipeline; +static id g_flash_kv_stage_f16_pipeline; +static id g_swiglu_pipeline; +static id g_swiglu_flat_pipeline; +static id g_add_pipeline; +static id g_add2_pipeline; +static id g_add3_pipeline; +static id g_moe_sum6_pipeline; +static id g_moe_sum8_pipeline; +static id g_mul_pipeline; +static id g_rms_norm_pipeline; +static id g_rms_norm_plain_pipeline; +static id g_add_rms_norm_pipeline; +static id g_rms_norm_scale_pipeline; +static id g_dsv4_qkv_rms_norm_pipeline; +static id g_hc_split_sinkhorn_pipeline; +static id g_hc_split_weighted_sum_pipeline; +static id g_hc_split_weighted_sum_norm_pipeline; +static id g_hc_weighted_sum_pipeline; +static id g_hc_weighted_sum_norm_pipeline; +static id g_output_hc_weights4_pipeline; +static id g_hc_expand_pipeline; +static id g_unary_sigmoid_pipeline; +static id g_unary_silu_pipeline; +static id g_unary_softplus_pipeline; +static id g_unary_sqrt_pipeline; +static id g_unary_clamp_pipeline; +static id g_unary_scale_pipeline; +static id g_unary_fill_pipeline; +static id g_unary_fill_f16_pipeline; +static id g_bin_mul_scalar_pipeline; +static id g_bin_div_row_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pair_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_q2_k_pipeline; +static id g_moe_mul_mv_id_q2_k_sum6_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; +static id g_moe_mul_mv_id_q4_k_pipeline; +static id g_moe_mul_mv_id_q4_k_pair_pipeline; +static id g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group6_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group8_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group24_q4_k_id_pipeline; +static id g_moe_mul_mv_group24_q4_k_sum6_pipeline; +static id g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_slots6_q2_k_sum6_pipeline; +static id g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_slots6_q4_k_sum6_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pipeline; +static id g_moe_mul_mv_addr_q2_k_sum6_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline; +static id g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline; +static id g_moe_stream_expert_cache_validate_pipeline; +static id g_moe_q4_gather_slots6_pipeline; +static id g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_table_q4_k_sum6_pipeline; +static id g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_addr_q4_k_sum6_pipeline; +static id g_moe_table_q4_pair_gate_encoder; +static id g_moe_table_q4_pair_up_encoder; +static id g_moe_table_q4_sum_down_encoder; +static id g_rope_tail_batch_pipeline; +static id g_rope_tail_inplace_pair_pipeline; +static id g_rope_tail_inplace_pair_shared4_pipeline; +static id g_rope_tail_inplace_pair_affine_pipeline; +static id g_dsv4_fp8_kv_quantize_pipeline; +static id g_dsv4_indexer_qat_pipeline; +static id g_dsv4_kv_fp8_store_pipeline; +static id g_dsv4_ratio4_shift_pipeline; +static id g_dsv4_compressor_pack_ratio4_pipeline; +static id g_dsv4_softmax_pool_ratio4_direct_pipeline; +static id g_dsv4_softmax_pool_pipeline; +static id g_soft_max_f32_pipeline; +static id g_soft_max_f32_4_pipeline; +static id g_argsort_f32_i32_desc_pipeline; +static id g_argsort_merge_f32_i32_desc_pipeline; +static id g_sum_rows_f32_f32_pipeline; +static id g_dsv4_topk_mask_pipeline; +static id g_dsv4_topk_mask_scatter_pipeline; +static id g_dsv4_indexer_weighted_sum_pipeline; +static id g_dsv4_indexer_score_one_direct_pipeline; +static id g_dsv4_compressor_store_one_pipeline; +static id g_dsv4_sort_i32_rows_asc_pipeline; +static id g_dsv4_indexed_attention_heads8_pipeline; +static id g_dsv4_indexed_attention_heads8_rb16_pipeline; +static id g_dsv4_softplus_sqrt_pipeline; +static id g_dsv4_router_finalize_one_pipeline; +static id g_dsv4_router_finalize_one_simd_pipeline; +static id g_dsv4_router_finalize_weights_one_simd_pipeline; +static id g_dsv4_router_weights_one_pipeline; +static id g_glm_router_select_one_pipeline; +static id g_glm_kv_lora_rms_norm_pipeline; +static id g_glm_k_b_project_pipeline; +static id g_glm_store_compact_kv_pipeline; +static id g_glm_qkv_norm_store_compact_kv_pipeline; +static id g_glm_store_indexer_k_pipeline; +static id g_glm_build_kv_cache_pipeline; +static id g_glm_build_kv_cache_decode_group4_pipeline; +static id g_glm_build_kv_cache_flash_pipeline; +static id g_glm_attention_full_pipeline; +static id g_glm_fill_selected_range_pipeline; +static id g_glm_fill_selected_range_batch_pipeline; +static id g_glm_indexer_rope_tail_pipeline; +static id g_glm_indexer_score_one_pipeline; +static id g_glm_indexer_score_one_direct_pipeline; +static id g_glm_indexer_scores_batch_pipeline; +static id g_glm_indexer_scores_tiled_pipeline; +static id g_glm_indexer_scores_tiled_f32_pipeline; +static id g_glm_qk_lowrank_pipeline; +static id g_glm_qk_lowrank_glm52_pipeline; +static id g_glm_qk_lowrank_glm52_sg_pipeline; +static id g_glm_qk_lowrank_batch_pipeline; +static id g_glm_qk_lowrank_batch_glm52_t4_pipeline; +static id g_glm_value_project_q8_0_pipeline; +static id g_glm_value_project_q8_0_batch_heads_pipeline; +static id g_glm_value_project_q8_0_batch_heads_mma_pipeline; +static id g_glm_attention_indexed_decode_pipeline; +static id g_glm_attention_indexed_decode_split_group8_partial_pipeline; +static id g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline; +static id g_glm_attention_indexed_decode_split_group8_reduce_pipeline; +static id g_glm_attention_indexed_decode_split_group8_reduce16_pipeline; +static id g_glm_attention_indexed_batch_pipeline; +static id g_glm_attention_indexed_batch_group2_pipeline; +static id g_glm_attention_indexed_batch_q2_group4_pipeline; +static id g_glm_attention_indexed_batch_group8_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline; +static id g_glm_q4_k_pair_swiglu_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_f32_pipeline; +static id g_glm_q4_k_pair_swiglu4_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline; +static id g_glm_q2_k_pair_swiglu_f32_pipeline; +static id g_glm_q2_k_addr_pair_swiglu2_f32_pipeline; +static id g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline; +static id g_glm_q4_k_addr_pair_swiglu_f32_pipeline; +static id g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline; +static id g_glm_q2_k_down_f32_pipeline; +static id g_glm_q4_k_down_f32_pipeline; +static id g_glm_q2_k_addr_down_f32_pipeline; +static id g_glm_q4_k_addr_down_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_mapped_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline; +static id g_glm_q5_k_down_f32_pipeline; +static id g_glm_q6_k_down_f32_pipeline; +static id g_dsv4_router_weights_batch_pipeline; +static id g_dsv4_hc_expand4_pipeline; +static NSMutableDictionary> *g_pipeline_cache; +static NSMutableDictionary> *g_model_buffer_cache; +static NSMutableDictionary *g_q4_expert_table_cache; +static NSMutableDictionary *g_q4_expert_layer_residency_cache; +static NSMutableArray> *g_transient_buffers; +static id g_model_residency_set; + +typedef struct { + id __strong mask; + id __strong blk; + NSUInteger mask_bytes; + NSUInteger blk_bytes; + uint32_t kind; + uint32_t n_tokens; + uint32_t n_comp; + uint32_t n_keys; + uint32_t window; + uint32_t ratio; + uint32_t nqptg; + uint32_t ncpsg; + bool has_kvpad; + bool bc_mask; + bool valid; + bool blk_ready; +} ds4_gpu_zero_prefix_prefill_mask_cache_entry; + +enum { + DS4_GPU_PREFILL_MASK_CACHE_RAW = 1, + DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2, + DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3, + DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3, +}; + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry + g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS]; +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void); +static id g_flash_attn_mask_buffer; +static id g_flash_attn_zero_mask_buffer; +static id g_flash_attn_pad_buffer; +static id g_flash_attn_tmp_buffer; +static id g_flash_attn_blk_buffer; +static id g_flash_attn_ring_buffer; +static id g_flash_attn_kv_buffer; +static id g_glm_flash_attn_mask_buffer; +static id g_compressor_pool_kv_buffer; +static id g_compressor_pool_score_buffer; +static id g_compressor_pool_score_cont_buffer; +static id g_compressor_pool_softmax_buffer; +static id g_compressor_pool_product_buffer; +static id g_compressor_store_ape_buffer; +static id g_compressor_store_score_buffer; +static id g_embed_rows_buffer; +static id g_router_selection_buffer; +static id g_router_weight_sum_buffer; +static id g_indexer_head_scores_buffer; +static id g_indexer_topk_buffer; +static id g_indexed_topk_buffer; +static id g_f16_round_scratch_buffer; +static id g_raw_store_round_buffer; +static id g_moe_gate_scratch_buffer; +static id g_moe_down_scratch_buffer; +static id g_moe_id_map_buffer; +static id g_moe_q4_gate_slots_buffer; +static id g_moe_q4_up_slots_buffer; +static id g_moe_q4_down_slots_buffer; +static id g_attn_out_group_ids_buffer; +static int g_model_fd = -1; +static const void *g_model_map_ptr; +static uint64_t g_model_map_size; +static uint64_t g_model_mapped_offset; +static uint64_t g_model_mapped_size; +static uint64_t g_model_mapped_max_tensor_bytes; +static uint64_t g_tensor_alloc_live_bytes; +static uint64_t g_tensor_alloc_peak_bytes; +static pthread_mutex_t g_tensor_mu = PTHREAD_MUTEX_INITIALIZER; +static uintptr_t *g_tensor_live_slots; +static size_t g_tensor_live_cap; +static size_t g_tensor_live_count; +static size_t g_tensor_live_tombs; +static uint64_t g_model_wrap_count; +static uint64_t g_model_wrap_bytes; +static uint64_t g_model_wrap_max_bytes; +static uint64_t g_model_buffer_cache_bytes; +static uint64_t g_model_buffer_cache_evictions; +static int g_model_buffer_cache_over_limit; +static uint64_t g_stream_expert_cache_bytes; +static uint64_t g_stream_expert_cache_expert_bytes; +static uint32_t g_stream_expert_cache_entry_count; +static uint32_t g_stream_expert_cache_budget_override; +static uint64_t g_stream_expert_cache_hits; +static uint64_t g_stream_expert_cache_misses; +static uint64_t g_stream_expert_cache_evictions; +static uint64_t g_stream_expert_cache_wraps; +static uint64_t g_stream_expert_cache_clock; +static uint64_t g_stream_expert_cache_evict_advise_bytes; +static uint64_t g_stream_expert_cache_willneed_advise_bytes; +static uint64_t g_stream_expert_cache_pread_bytes; +static double g_stream_expert_cache_pread_ms; +static uint64_t g_stream_expert_cache_buffer_allocs; +static uint64_t g_stream_expert_cache_buffer_reuses; +static uint64_t g_stream_expert_cache_decode_tokens; +static uint64_t g_stream_expert_cache_hotness_decay_token; +static uint64_t g_stream_expert_timing_selected_calls; +static double g_stream_expert_timing_selected_read_ms; +static double g_stream_expert_timing_selected_sync_ms; +static double g_stream_expert_timing_selected_copy_ms; +static double g_stream_expert_timing_selected_bind_ms; +static uint64_t g_stream_expert_timing_split_layers; +static uint64_t g_stream_expert_timing_split_resident_experts; +static uint64_t g_stream_expert_timing_split_missing_experts; +static double g_stream_expert_timing_split_resident_ms; +static double g_stream_expert_timing_split_missing_ms; +static double g_stream_expert_timing_split_missing_load_ms; +static double g_stream_expert_timing_split_missing_slot_ms; +static double g_stream_expert_timing_split_missing_prune_ms; +static double g_stream_expert_timing_split_missing_addr_ms; +static double g_stream_expert_timing_split_missing_wait_ms; +static uint64_t g_stream_expert_timing_load_calls; +static double g_stream_expert_timing_load_prepare_ms; +static double g_stream_expert_timing_load_pread_ms; +static double g_stream_expert_timing_load_modify_ms; +static double g_stream_expert_timing_load_install_ms; +static uint64_t g_stream_expert_timing_prepare_batch_reuse_calls; +static double g_stream_expert_timing_prepare_batch_reuse_ms; +static uint64_t g_stream_expert_timing_prepare_buffer_calls; +static double g_stream_expert_timing_prepare_buffer_ms; +static uint64_t g_stream_expert_timing_prepare_task_experts; +static double g_stream_expert_timing_prepare_task_ms; +static uint64_t g_stream_expert_timing_reuse_scan_calls; +static uint64_t g_stream_expert_timing_reuse_scan_entries; +static double g_stream_expert_timing_reuse_scan_ms; +static double g_stream_expert_timing_reuse_clear_ms; +static uint64_t g_stream_expert_timing_readahead_calls; +static uint64_t g_stream_expert_timing_readahead_bytes; +static double g_stream_expert_timing_readahead_ms; +static uint64_t g_stream_expert_timing_cache_all_resident_layers; +static uint64_t g_stream_expert_timing_cache_all_missing_layers; +static uint64_t g_stream_expert_timing_cache_mixed_layers; +static uint64_t g_stream_expert_timing_cache_resident_experts; +static uint64_t g_stream_expert_timing_cache_missing_experts; +typedef struct { + uint64_t selected_calls; + double selected_read_ms; + double selected_sync_ms; + double selected_copy_ms; + double selected_bind_ms; + uint64_t split_layers; + uint64_t split_resident_experts; + uint64_t split_missing_experts; + double split_resident_ms; + double split_missing_ms; + double split_missing_load_ms; + double split_missing_slot_ms; + double split_missing_prune_ms; + double split_missing_addr_ms; + double split_missing_wait_ms; + uint64_t load_calls; + double load_prepare_ms; + double load_pread_ms; + double load_modify_ms; + double load_install_ms; + uint64_t prepare_batch_reuse_calls; + double prepare_batch_reuse_ms; + uint64_t prepare_buffer_calls; + double prepare_buffer_ms; + uint64_t prepare_task_experts; + double prepare_task_ms; + uint64_t reuse_scan_calls; + uint64_t reuse_scan_entries; + double reuse_scan_ms; + double reuse_clear_ms; + uint64_t readahead_calls; + uint64_t readahead_bytes; + double readahead_ms; + uint64_t cache_all_resident_layers; + uint64_t cache_all_missing_layers; + uint64_t cache_mixed_layers; + uint64_t cache_resident_experts; + uint64_t cache_missing_experts; +} ds4_gpu_stream_expert_timing_snapshot; +static ds4_gpu_stream_expert_timing_snapshot g_stream_expert_timing_last_report; +static int g_stream_prefill_batch_selected_addr_building; +static int g_glm_stream_expert_addr_table_building; +static uint64_t g_model_residency_count; +static int g_model_residency_added_to_queue; +static int g_glm_model_mode; +static int g_ssd_streaming_mode; +static int g_glm_streaming_prefill_full_layer_runtime; +static int g_metal4_runtime_available; +static int g_metal4_family_supported; +static int g_metal4_queue_supported; +static int g_metal4_m5_neural_accelerators_hint; +static int g_metal4_tensor_api_enabled; +static int g_metal4_tensor_api_compile_supported; +static char g_metal_device_name[128]; +static int ds4_gpu_model_map_log_enabled(void); +static int ds4_gpu_stream_expert_cache_note_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes); +static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void); +static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats); +static void ds4_gpu_stream_expert_pending_load_clear(void); +static void ds4_gpu_stream_expert_pread_pool_shutdown(void); +static int ds4_gpu_stream_expert_timing_summary_enabled(void); +static int ds4_gpu_stream_expert_cache_entry_protected( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect); + +/* The async selected-load worker registers itself so cache paths that would + * flush/wait on command buffers (a race against the encoding thread) fail + * the load instead; the caller then retries on the main thread. */ +static pthread_t g_stream_expert_service_thread; +static int g_stream_expert_service_thread_set; + +void ds4_gpu_stream_expert_cache_note_service_thread(void) { + g_stream_expert_service_thread = pthread_self(); + g_stream_expert_service_thread_set = 1; +} + +static int ds4_gpu_stream_expert_cache_on_service_thread(void) { + return g_stream_expert_service_thread_set && + pthread_equal(pthread_self(), g_stream_expert_service_thread); +} +static NSUInteger g_flash_attn_mask_bytes; +static NSUInteger g_flash_attn_zero_mask_bytes; +static NSUInteger g_flash_attn_pad_bytes; +static NSUInteger g_flash_attn_tmp_bytes; +static NSUInteger g_flash_attn_blk_bytes; +static NSUInteger g_flash_attn_ring_bytes; +static NSUInteger g_flash_attn_kv_bytes; +static NSUInteger g_glm_flash_attn_mask_bytes; +static uint32_t g_glm_flash_attn_mask_pos0; +static uint32_t g_glm_flash_attn_mask_tokens; +static uint32_t g_glm_flash_attn_mask_cache_len; +static int g_glm_flash_attn_mask_valid; +static NSUInteger g_compressor_pool_kv_bytes; +static NSUInteger g_compressor_pool_score_bytes; +static NSUInteger g_compressor_pool_score_cont_bytes; +static NSUInteger g_compressor_pool_softmax_bytes; +static NSUInteger g_compressor_pool_product_bytes; +static NSUInteger g_compressor_store_ape_bytes; +static NSUInteger g_compressor_store_score_bytes; +static NSUInteger g_embed_rows_bytes; +static NSUInteger g_router_selection_bytes; +static NSUInteger g_router_weight_sum_bytes; +static NSUInteger g_indexer_head_scores_bytes; +static NSUInteger g_indexer_topk_bytes; +static NSUInteger g_indexed_topk_bytes; +static NSUInteger g_f16_round_scratch_bytes; +static NSUInteger g_raw_store_round_bytes; +static NSUInteger g_moe_gate_scratch_bytes; +static NSUInteger g_moe_down_scratch_bytes; +static NSUInteger g_moe_id_map_bytes; +static NSUInteger g_moe_q4_gate_slots_bytes; +static NSUInteger g_moe_q4_up_slots_bytes; +static NSUInteger g_moe_q4_down_slots_bytes; +static NSUInteger g_attn_out_group_ids_bytes; +static int g_initialized; +static int g_quality_mode; +static int g_mpp_invalid_env_reported; +#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8 +static int32_t g_routed_moe_selected_override[DS4_METAL_MAX_ROUTED_EXPERT_USED]; +static uint32_t g_routed_moe_selected_override_n; +static int g_moe_selected_trace_record_initialized; +static FILE *g_moe_selected_trace_record_fp; +static uint64_t g_moe_selected_trace_record_count; +static int g_moe_selected_trace_replay_initialized; +static int32_t *g_moe_selected_trace_replay_ids; +static uint64_t g_moe_selected_trace_replay_count; +static uint64_t g_moe_selected_trace_replay_pos; + +static double ds4_gpu_gib(uint64_t bytes); + +static uint64_t ds4_gpu_system_memory_bytes(void) { + uint64_t bytes = 0; + size_t len = sizeof(bytes); + if (sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) != 0) return 0; + return len == sizeof(bytes) ? bytes : 0; +} + +static void ds4_gpu_print_device_summary(void) { + const char *name = g_device.name ? [g_device.name UTF8String] : "unknown Metal device"; + uint64_t mem = ds4_gpu_system_memory_bytes(); + if (mem) { + double gib = (double)mem / 1024.0 / 1024.0 / 1024.0; + fprintf(stderr, "ds4: Metal device %s, %.2f GiB RAM\n", name, gib); + } else { + fprintf(stderr, "ds4: Metal device %s\n", name); + } +} + +#define DS4_METAL_MAX_MODEL_VIEWS 4096 +/* Compatibility fallback for callers that cannot provide a parsed GGUF tensor + * span. The normal DS4 engine passes the exact maximum tensor byte size. */ +#define DS4_METAL_FALLBACK_MAX_TENSOR_BYTES (4ull * 1024ull * 1024ull * 1024ull) + +typedef struct { + __strong id buffer; + const void *model_map; + uint64_t model_size; + uint64_t model_offset; + uint64_t bytes; +} ds4_gpu_model_view; + +static ds4_gpu_model_view g_model_views[DS4_METAL_MAX_MODEL_VIEWS]; +static uint32_t g_model_view_count; + +enum { + DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES = + DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS = 256, + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16, + DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS = 16, +}; + +typedef struct { + uint32_t layer; + uint32_t expert; + uint64_t hits; +} ds4_gpu_moe_selected_hotlist_entry; + +static int g_moe_selected_hotlist_initialized; +static uint64_t + g_moe_selected_hotlist_counts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static uint64_t g_moe_selected_hotlist_records; +static uint64_t g_moe_selected_hotlist_selections; + +typedef struct { + __strong id gate_buffer; + __strong id up_buffer; + __strong id down_buffer; + const void *model_map; + uint64_t model_size; + uint64_t gate_abs_offset; + uint64_t up_abs_offset; + uint64_t down_abs_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + uint64_t logical_bytes; + uint64_t last_used; + uint64_t use_count; + NSUInteger gate_inner; + NSUInteger up_inner; + NSUInteger down_inner; + uint64_t inflight_seq; + uint32_t slab_slot; + uint8_t valid; + uint8_t slab_backed; +} ds4_gpu_stream_expert_cache_entry; + +typedef struct { + __strong id gate_buffer; + __strong id up_buffer; + __strong id down_buffer; + NSUInteger gate_inner; + NSUInteger up_inner; + NSUInteger down_inner; +} ds4_gpu_stream_expert_reusable_buffers; + +static ds4_gpu_stream_expert_cache_entry + g_stream_expert_cache[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static ds4_gpu_stream_expert_cache_entry + g_stream_full_expert_addr_entry[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint32_t g_stream_expert_cache_layer_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static double g_stream_expert_cache_layer_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static double g_stream_expert_cache_layer_last_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint32_t + g_stream_expert_cache_route_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static id g_stream_expert_cache_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_slabs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_start_slot[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_slot_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_slots_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_count; +static uint32_t g_stream_expert_cache_slab_total_slots; +static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES]; +static uint32_t g_stream_expert_cache_free_slot_count; +static uint64_t g_stream_expert_cache_slab_slot_bytes; +static uint64_t g_stream_expert_cache_cb_seq; +static uint64_t g_stream_expert_cache_done_seq; +static uint64_t g_stream_expert_cache_batch_seq; +static uint64_t g_stream_expert_cache_owned_seq; +static uint64_t g_stream_expert_cache_pending_max_seq; +static id g_stream_compact_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_selected_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_selected_id_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_validate_status_buffer; + +@interface DS4MetalTensor : NSObject +@property(nonatomic, strong) id buffer; +@property(nonatomic, assign) uint64_t offset; +@property(nonatomic, assign) uint64_t bytes; +@property(nonatomic, assign) uint8_t owner; +@end + +@implementation DS4MetalTensor +@end + +@interface DS4MetalQ4ExpertTable : NSObject +@property(nonatomic, strong) id argumentBuffer; +@property(nonatomic, strong) id addressBuffer; +@property(nonatomic, strong) NSMutableArray> *expertBuffers; +@property(nonatomic, strong) id residencySet; +@property(nonatomic, assign) BOOL residencySetAddedToQueue; +@property(nonatomic, assign) uint32_t nExpert; +@property(nonatomic, assign) uint64_t expertBytes; +@end + +@implementation DS4MetalQ4ExpertTable +- (void)dealloc { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (_residencySet) { + if (_residencySetAddedToQueue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:_residencySet]; + } + [_residencySet endResidency]; + } + } +#endif +} +@end + +@interface DS4MetalQ4LayerResidency : NSObject +@property(nonatomic, strong) id residencySet; +@property(nonatomic, assign) BOOL addedToQueue; +@end + +@implementation DS4MetalQ4LayerResidency +- (void)dealloc { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (_residencySet) { + if (_addedToQueue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:_residencySet]; + } + [_residencySet endResidency]; + } + } +#endif +} +@end + +static DS4MetalTensor *ds4_gpu_tensor_obj(ds4_gpu_tensor *tensor) { + return (__bridge DS4MetalTensor *)tensor; +} + +static const DS4MetalTensor *ds4_gpu_tensor_const_obj(const ds4_gpu_tensor *tensor) { + return (__bridge const DS4MetalTensor *)tensor; +} + +/* C code owns ds4_gpu_tensor handles as retained Objective-C objects. Freeing + * the same opaque handle twice would make the second __bridge_transfer release + * an already-deallocated object, which macOS reports as malloc corruption. The + * live table lets free validate a handle before touching Objective-C state; the + * same mutex also serializes the diagnostic allocation counters. */ +static uint64_t ds4_gpu_tensor_ptr_hash(uintptr_t ptr) { + uint64_t x = (uint64_t)(ptr >> 4); + x ^= x >> 33; + x *= UINT64_C(0xff51afd7ed558ccd); + x ^= x >> 33; + x *= UINT64_C(0xc4ceb9fe1a85ec53); + x ^= x >> 33; + return x; +} + +static int ds4_gpu_tensor_live_resize_locked(size_t min_cap) { + size_t new_cap = 1024; + while (new_cap < min_cap) new_cap <<= 1; + + uintptr_t *new_slots = calloc(new_cap, sizeof(new_slots[0])); + if (!new_slots) return 0; + + for (size_t i = 0; i < g_tensor_live_cap; i++) { + const uintptr_t key = g_tensor_live_slots[i]; + if (key == 0 || key == UINTPTR_MAX) continue; + + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (new_cap - 1); + while (new_slots[idx] != 0) idx = (idx + 1) & (new_cap - 1); + new_slots[idx] = key; + } + + free(g_tensor_live_slots); + g_tensor_live_slots = new_slots; + g_tensor_live_cap = new_cap; + g_tensor_live_tombs = 0; + return 1; +} + +static int ds4_gpu_tensor_live_insert_locked(const void *ptr) { + if (!ptr || (uintptr_t)ptr == UINTPTR_MAX) return 0; + if ((g_tensor_live_count + g_tensor_live_tombs + 1) * 10 >= + g_tensor_live_cap * 7) + { + const size_t min_cap = g_tensor_live_cap ? g_tensor_live_cap * 2 : 1024; + if (!ds4_gpu_tensor_live_resize_locked(min_cap)) return 0; + } + + const uintptr_t key = (uintptr_t)ptr; + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); + size_t tomb = (size_t)-1; + for (;;) { + const uintptr_t cur = g_tensor_live_slots[idx]; + if (cur == key) return 0; + if (cur == UINTPTR_MAX) { + if (tomb == (size_t)-1) tomb = idx; + } else if (cur == 0) { + if (tomb != (size_t)-1) { + idx = tomb; + g_tensor_live_tombs--; + } + g_tensor_live_slots[idx] = key; + g_tensor_live_count++; + return 1; + } + idx = (idx + 1) & (g_tensor_live_cap - 1); + } +} + +static int ds4_gpu_tensor_live_remove_locked(const void *ptr) { + if (!ptr || g_tensor_live_cap == 0) return 0; + + const uintptr_t key = (uintptr_t)ptr; + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); + for (;;) { + const uintptr_t cur = g_tensor_live_slots[idx]; + if (cur == 0) return 0; + if (cur == key) { + g_tensor_live_slots[idx] = UINTPTR_MAX; + g_tensor_live_count--; + g_tensor_live_tombs++; + return 1; + } + idx = (idx + 1) & (g_tensor_live_cap - 1); + } +} + +static int ds4_gpu_tensor_track_alloc_locked( + const void *ptr, + uint64_t bytes, + uint64_t *live_snap, + uint64_t *peak_snap) +{ + if (!ds4_gpu_tensor_live_insert_locked(ptr)) return 0; + + g_tensor_alloc_live_bytes += bytes; + if (g_tensor_alloc_live_bytes > g_tensor_alloc_peak_bytes) { + g_tensor_alloc_peak_bytes = g_tensor_alloc_live_bytes; + } + if (live_snap) *live_snap = g_tensor_alloc_live_bytes; + if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; + return 1; +} + +static int ds4_gpu_tensor_track_view_locked(const void *ptr) { + return ds4_gpu_tensor_live_insert_locked(ptr); +} + +static int ds4_gpu_tensor_prepare_free( + ds4_gpu_tensor *tensor, + uint8_t *owner, + uint64_t *bytes, + uint64_t *live_snap, + uint64_t *peak_snap) +{ + pthread_mutex_lock(&g_tensor_mu); + if (!ds4_gpu_tensor_live_remove_locked(tensor)) { + pthread_mutex_unlock(&g_tensor_mu); + fprintf(stderr, + "ds4: Metal tensor free ignored for unknown handle %p\n", + (void *)tensor); + return 0; + } + + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + const uint8_t obj_owner = obj.owner; + const uint64_t obj_bytes = obj.bytes; + if (obj_owner) { + if (obj_bytes <= g_tensor_alloc_live_bytes) { + g_tensor_alloc_live_bytes -= obj_bytes; + } else { + g_tensor_alloc_live_bytes = 0; + } + } + if (owner) *owner = obj_owner; + if (bytes) *bytes = obj_bytes; + if (live_snap) *live_snap = g_tensor_alloc_live_bytes; + if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; + pthread_mutex_unlock(&g_tensor_mu); + return 1; +} + +static void ds4_gpu_tensor_tracking_reset(void) { + pthread_mutex_lock(&g_tensor_mu); + if (g_tensor_live_count != 0) { + fprintf(stderr, + "ds4: Metal cleanup discarded %zu live tensor handles\n", + g_tensor_live_count); + } + free(g_tensor_live_slots); + g_tensor_live_slots = NULL; + g_tensor_live_cap = 0; + g_tensor_live_count = 0; + g_tensor_live_tombs = 0; + g_tensor_alloc_live_bytes = 0; + g_tensor_alloc_peak_bytes = 0; + pthread_mutex_unlock(&g_tensor_mu); +} + +static id ds4_gpu_tensor_buffer(const ds4_gpu_tensor *tensor) { + if (!tensor) return nil; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return obj.buffer; +} + +static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) { + if (!tensor) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return (NSUInteger)obj.offset; +} + +static id ds4_gpu_new_command_buffer(void); +static void ds4_gpu_stream_expert_cache_note_owned_created(void); + +static id ds4_gpu_command_buffer(int *owned) { + if (g_batch_cb) { + *owned = 0; + return g_batch_cb; + } + *owned = 1; + id cb = ds4_gpu_new_command_buffer(); + if (cb) ds4_gpu_stream_expert_cache_note_owned_created(); + return cb; +} + +static id ds4_gpu_compute_encoder(id cb) { + if (g_batch_cb && cb == g_batch_cb) { + g_batch_has_work = YES; + if (!g_batch_enc) g_batch_enc = [cb computeCommandEncoder]; + return g_batch_enc; + } + return [cb computeCommandEncoder]; +} + +static void ds4_gpu_end_compute_encoder(id cb, id enc) { + if (!enc) return; + if (g_batch_cb && cb == g_batch_cb && enc == g_batch_enc) return; + [enc endEncoding]; +} + +static void ds4_gpu_close_batch_encoder(void) { + if (!g_batch_enc) return; + [g_batch_enc endEncoding]; + g_batch_enc = nil; +} + +static double g_gpu_busy_accum; +static uint64_t g_gpu_busy_cbs; + +static int ds4_gpu_wait_command_buffer(id cb, const char *label) { + [cb waitUntilCompleted]; + if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) { + const double busy = cb.GPUEndTime - cb.GPUStartTime; + if (busy > 0) g_gpu_busy_accum += busy; + if ((++g_gpu_busy_cbs % 64u) == 0u) { + fprintf(stderr, "ds4: gpu busy accum %.1f ms over %llu cbs\n", + g_gpu_busy_accum * 1000.0, + (unsigned long long)g_gpu_busy_cbs); + } + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + label, [[cb.error localizedDescription] UTF8String]); + return 0; + } + return 1; +} + +static id ds4_gpu_new_command_buffer(void) { + static int initialized; + static int use_unretained; + if (!initialized) { + use_unretained = getenv("DS4_METAL_UNRETAINED_COMMAND_BUFFERS") != NULL; + initialized = 1; + } + if (use_unretained) { + return [g_queue commandBufferWithUnretainedReferences]; + } + return [g_queue commandBuffer]; +} + +static uint64_t ds4_gpu_exact_view_cache_limit_bytes(void) { + static int initialized; + static uint64_t limit_bytes; + if (initialized) return limit_bytes; + + const uint64_t mib = 1024ull * 1024ull; + const uint64_t gib = 1024ull * mib; + limit_bytes = 64ull * gib; + + const char *gib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_GIB"); + if (gib_env && gib_env[0]) { + char *end = NULL; + unsigned long long v = strtoull(gib_env, &end, 10); + if (end != gib_env && *end == '\0') { + limit_bytes = v > UINT64_MAX / gib ? UINT64_MAX : (uint64_t)v * gib; + } + } + + const char *mib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_MIB"); + if (mib_env && mib_env[0]) { + char *end = NULL; + unsigned long long v = strtoull(mib_env, &end, 10); + if (end != mib_env && *end == '\0') { + limit_bytes = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; + } + } + + initialized = 1; + return limit_bytes; +} + +static void ds4_gpu_model_buffer_cache_note_insert(uint64_t bytes) { + if (g_model_buffer_cache_bytes > UINT64_MAX - bytes) { + g_model_buffer_cache_bytes = UINT64_MAX; + } else { + g_model_buffer_cache_bytes += bytes; + } + + const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); + if (limit != 0 && g_model_buffer_cache_bytes > limit) { + g_model_buffer_cache_over_limit = 1; + } +} + +static void ds4_gpu_model_buffer_cache_clear(const char *reason) { + if (!g_model_buffer_cache) { + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_over_limit = 0; + return; + } + + const NSUInteger entries = [g_model_buffer_cache count]; + if (entries != 0) { + if (getenv("DS4_METAL_EXACT_VIEW_CACHE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal exact model view cache evict reason=%s entries=%lu bytes=%.2f GiB limit=%.2f GiB\n", + reason ? reason : "unknown", + (unsigned long)entries, + ds4_gpu_gib(g_model_buffer_cache_bytes), + ds4_gpu_gib(ds4_gpu_exact_view_cache_limit_bytes())); + } + [g_model_buffer_cache removeAllObjects]; + g_model_buffer_cache_evictions++; + } + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_over_limit = 0; +} + +static void ds4_gpu_model_buffer_cache_maybe_evict(const char *reason) { + if (g_model_buffer_cache_over_limit) { + ds4_gpu_model_buffer_cache_clear(reason); + } +} + +static uint64_t ds4_gpu_stream_expert_cache_next_cb_seq(void) { + if (g_stream_expert_cache_cb_seq == UINT64_MAX) { + /* + * A real wrap would require an astronomical number of command buffers. + * Resetting the epoch space is still safer than letting zero become a + * valid in-flight marker. + */ + g_stream_expert_cache_cb_seq = 0; + g_stream_expert_cache_done_seq = 0; + g_stream_expert_cache_batch_seq = 0; + g_stream_expert_cache_owned_seq = 0; + g_stream_expert_cache_pending_max_seq = 0; + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + g_stream_expert_cache[layer][expert].inflight_seq = 0; + } + } + } + return ++g_stream_expert_cache_cb_seq; +} + +static void ds4_gpu_stream_expert_cache_note_batch_created(void) { + g_stream_expert_cache_batch_seq = + ds4_gpu_stream_expert_cache_next_cb_seq(); +} + +static void ds4_gpu_stream_expert_cache_note_batch_committed(void) { + if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_pending_max_seq) { + g_stream_expert_cache_pending_max_seq = g_stream_expert_cache_batch_seq; + } + g_stream_expert_cache_batch_seq = 0; +} + +static void ds4_gpu_stream_expert_cache_note_owned_created(void) { + g_stream_expert_cache_owned_seq = + ds4_gpu_stream_expert_cache_next_cb_seq(); +} + +static void ds4_gpu_stream_expert_cache_note_pending_completed(void) { + if (g_stream_expert_cache_pending_max_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_pending_max_seq; + } + g_stream_expert_cache_pending_max_seq = 0; +} + +static void ds4_gpu_stream_expert_cache_note_owned_completed(void) { + if (g_stream_expert_cache_owned_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_owned_seq; + } + g_stream_expert_cache_owned_seq = 0; +} + +static int ds4_gpu_stream_expert_cache_entry_inflight( + const ds4_gpu_stream_expert_cache_entry *e) { + return e && e->valid && e->inflight_seq > g_stream_expert_cache_done_seq; +} + +static int ds4_gpu_stream_expert_cache_mark_inflight( + ds4_gpu_stream_expert_cache_entry *e) { + if (!e || !e->valid) return 0; + const uint64_t seq = g_stream_expert_cache_batch_seq ? + g_stream_expert_cache_batch_seq : + g_stream_expert_cache_owned_seq; + if (seq == 0) return 0; + e->inflight_seq = seq; + return 1; +} + +static int ds4_gpu_stream_expert_cache_mark_entries_inflight( + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + uint32_t active_mask) { + if (!entries || n_entries == 0) return 0; + for (uint32_t i = 0; i < n_entries; i++) { + if (active_mask != 0 && (active_mask & (1u << i)) == 0) continue; + if (!ds4_gpu_stream_expert_cache_mark_inflight(entries[i])) return 0; + } + return 1; +} + +static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label); + +static int ds4_gpu_wait_pending_command_buffers(const char *label) { + int ok = 1; + for (id pending in g_pending_cbs) { + if (!ds4_gpu_wait_command_buffer(pending, label)) ok = 0; + } + [g_pending_cbs removeAllObjects]; + ds4_gpu_stream_expert_cache_note_pending_completed(); + if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); + return ok; +} + +static int ds4_gpu_finish_command_buffer(id cb, int owned, const char *label) { + if (!owned) return 1; + + [cb commit]; + int ok = ds4_gpu_wait_pending_command_buffers(label); + if (!ds4_gpu_wait_command_buffer(cb, label)) { + ok = 0; + ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); + } + ds4_gpu_stream_expert_cache_note_owned_completed(); + [g_transient_buffers removeAllObjects]; + ds4_gpu_model_buffer_cache_maybe_evict(label); + return ok; +} + +static int ds4_gpu_device_name_contains(const char *needle); + +static int ds4_gpu_use_m5_private_scratch(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = ds4_gpu_device_name_contains("M5"); + initialized = 1; + } + return enabled; +} + +static int ds4_gpu_scratch_needs_cpu_access(const char *label) { + if (!label) return 0; + return strstr(label, "mask") != NULL || + strcmp(label, "ds4_attention_output_group_ids") == 0; +} + +static MTLResourceOptions ds4_gpu_model_resource_options(void) { + MTLResourceOptions options = MTLResourceStorageModeShared; + if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) { + options |= MTLResourceHazardTrackingModeUntracked; + } + return options; +} + +static int ds4_gpu_ensure_scratch_buffer( + id __strong *buffer, + NSUInteger *capacity, + NSUInteger bytes, + const char *label) { + if (*buffer && *capacity >= bytes) return 1; + if (bytes == 0) bytes = 1; + if (bytes > NSUIntegerMax) return 0; + + MTLResourceOptions options = MTLResourceStorageModeShared; + if (ds4_gpu_use_m5_private_scratch() && + !ds4_gpu_scratch_needs_cpu_access(label)) { + /* + * M5 scratch buffers that only flow between Metal kernels do not need + * CPU-visible shared storage. This reduces shared-memory traffic and + * residency pressure for the long prefill scratch pools without + * changing the public buffer lifetime model. Keep default hazard + * tracking because the graph reuses these buffers across dependent + * compute encoders. + */ + options = MTLResourceStorageModePrivate; + } + + *buffer = [g_device newBufferWithLength:bytes options:options]; + if (!*buffer && options != MTLResourceStorageModeShared) { + *buffer = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared]; + } + if (!*buffer) { + fprintf(stderr, "ds4: failed to allocate Metal scratch buffer %s (%llu bytes)\n", + label, (unsigned long long)bytes); + *capacity = 0; + return 0; + } + (*buffer).label = [NSString stringWithUTF8String:label]; + *capacity = bytes; + return 1; +} + +static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) { + const NSUInteger capacity = 8192u * sizeof(uint16_t); + if (bytes > capacity) return 0; + if (g_flash_attn_zero_mask_buffer && + g_flash_attn_zero_mask_bytes >= capacity) { + return 1; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer, + &g_flash_attn_zero_mask_bytes, + capacity, + "ds4_flash_attn_zero_mask")) { + return 0; + } + void *contents = [g_flash_attn_zero_mask_buffer contents]; + if (!contents) return 0; + memset(contents, 0, g_flash_attn_zero_mask_bytes); + return 1; +} + +static uint64_t round_up_u64(uint64_t v, uint64_t align) { + return (v + align - 1) & ~(align - 1); +} + +static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint64_t max_tensor_bytes) { + if (max_tensor_bytes != 0) return max_tensor_bytes; + return map_size < DS4_METAL_FALLBACK_MAX_TENSOR_BYTES ? + map_size : DS4_METAL_FALLBACK_MAX_TENSOR_BYTES; +} + +static id ds4_gpu_get_pipeline(const char *function_name); +static int ds4_gpu_warm_model_views(void); +static double ds4_gpu_gib(uint64_t bytes); + +static double ds4_gpu_now_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0; +} + +static int ds4_gpu_moe_selected_hotlist_cmp(const void *a, const void *b) { + const ds4_gpu_moe_selected_hotlist_entry *ea = a; + const ds4_gpu_moe_selected_hotlist_entry *eb = b; + if (ea->hits < eb->hits) return 1; + if (ea->hits > eb->hits) return -1; + if (ea->layer != eb->layer) return ea->layer < eb->layer ? -1 : 1; + if (ea->expert != eb->expert) return ea->expert < eb->expert ? -1 : 1; + return 0; +} + +static int ds4_gpu_moe_selected_hotlist_merge_requested(void) { + return getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_MERGE") != NULL && + getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_FRESH") == NULL; +} + +static int ds4_gpu_moe_selected_hotlist_load_existing(const char *path) { + if (!path || !path[0] || !ds4_gpu_moe_selected_hotlist_merge_requested()) { + return 1; + } + + FILE *fp = fopen(path, "rb"); + if (!fp) { + if (errno == ENOENT) return 1; + fprintf(stderr, "ds4: failed to open selected hotlist merge file %s\n", path); + return 0; + } + + char line[256]; + uint64_t lineno = 0; + uint64_t loaded_entries = 0; + uint64_t loaded_hits = 0; + uint64_t header_records = UINT64_MAX; + uint64_t header_selections = UINT64_MAX; + while (fgets(line, sizeof(line), fp)) { + lineno++; + char *p = line; + while (*p && isspace((unsigned char)*p)) p++; + if (*p == '\0') continue; + if (*p == '#') { + unsigned long long value = 0; + if (sscanf(p, "# layer_records %llu", &value) == 1) { + header_records = (uint64_t)value; + } else if (sscanf(p, "# selections %llu", &value) == 1) { + header_selections = (uint64_t)value; + } + continue; + } + + errno = 0; + char *end = NULL; + unsigned long layer = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long expert = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long long hits = strtoull(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + hits != 0) { + uint64_t *dst = &g_moe_selected_hotlist_counts[layer][expert]; + if (*dst > UINT64_MAX - (uint64_t)hits) { + *dst = UINT64_MAX; + } else { + *dst += (uint64_t)hits; + } + loaded_entries++; + if (loaded_hits > UINT64_MAX - (uint64_t)hits) { + loaded_hits = UINT64_MAX; + } else { + loaded_hits += (uint64_t)hits; + } + } + continue; + +bad_line: + fprintf(stderr, + "ds4: invalid selected hotlist merge line %" PRIu64 " in %s\n", + lineno, + path); + fclose(fp); + return 0; + } + if (ferror(fp)) { + fprintf(stderr, "ds4: failed to read selected hotlist merge file %s\n", path); + fclose(fp); + return 0; + } + fclose(fp); + + g_moe_selected_hotlist_records = + header_records != UINT64_MAX ? header_records : loaded_hits / 6u; + g_moe_selected_hotlist_selections = + header_selections != UINT64_MAX ? header_selections : loaded_hits; + fprintf(stderr, + "ds4: merged selected-id hotlist %s " + "(%" PRIu64 " entries, %" PRIu64 " hits)\n", + path, + loaded_entries, + loaded_hits); + return 1; +} + +static void ds4_gpu_moe_selected_hotlist_close(void) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); + if (!g_moe_selected_hotlist_initialized || !path || !path[0]) return; + + const size_t cap = + (size_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + ds4_gpu_moe_selected_hotlist_entry *entries = + malloc(cap * sizeof(entries[0])); + if (!entries) { + fprintf(stderr, "ds4: failed to allocate selected hotlist entries\n"); + return; + } + + size_t n = 0; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + const uint64_t hits = + g_moe_selected_hotlist_counts[layer][expert]; + if (hits == 0) continue; + entries[n++] = (ds4_gpu_moe_selected_hotlist_entry) { + .layer = layer, + .expert = expert, + .hits = hits, + }; + } + } + qsort(entries, n, sizeof(entries[0]), ds4_gpu_moe_selected_hotlist_cmp); + + FILE *fp = fopen(path, "wb"); + if (!fp) { + fprintf(stderr, "ds4: failed to open selected hotlist file %s\n", path); + free(entries); + return; + } + fprintf(fp, + "# ds4 selected-id hotlist v1\n" + "# layer_records %" PRIu64 "\n" + "# selections %" PRIu64 "\n" + "# columns: layer expert hits weight\n", + g_moe_selected_hotlist_records, + g_moe_selected_hotlist_selections); + for (size_t i = 0; i < n; i++) { + fprintf(fp, + "%u %u %" PRIu64 " 0\n", + entries[i].layer, + entries[i].expert, + entries[i].hits); + } + free(entries); + + if (fclose(fp) != 0) { + fprintf(stderr, "ds4: failed to close selected hotlist file %s\n", path); + } else { + fprintf(stderr, + "ds4: wrote selected-id hotlist to %s " + "(%" PRIu64 " layer records, %" PRIu64 " selections)\n", + path, + g_moe_selected_hotlist_records, + g_moe_selected_hotlist_selections); + } +} + +static int ds4_gpu_moe_selected_hotlist_record( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected, + uint32_t n_total_expert) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); + if (!path || !path[0]) return 1; + if (!selected_ids || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return 0; + } + if (!g_moe_selected_hotlist_initialized) { + g_moe_selected_hotlist_initialized = 1; + if (!ds4_gpu_moe_selected_hotlist_load_existing(path)) return 0; + atexit(ds4_gpu_moe_selected_hotlist_close); + } + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 1; + + g_moe_selected_hotlist_records++; + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0) continue; + const uint32_t expert = (uint32_t)selected_ids[i]; + if (expert >= n_total_expert || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + continue; + } + g_moe_selected_hotlist_counts[layer][expert]++; + g_moe_selected_hotlist_selections++; + } + return 1; +} + +static void ds4_gpu_moe_selected_trace_record_close(void) { + if (g_moe_selected_trace_record_fp) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); + fclose(g_moe_selected_trace_record_fp); + g_moe_selected_trace_record_fp = NULL; + fprintf(stderr, + "ds4: recorded %" PRIu64 " routed-MoE selected-id entries to %s\n", + g_moe_selected_trace_record_count, + path && path[0] ? path : "(unknown)"); + } +} + +static int ds4_gpu_moe_selected_trace_record( + const int32_t selected_ids[6], + uint32_t n_selected) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); + if (!path || !path[0]) return 1; + if (n_selected != 6) { + fprintf(stderr, "ds4: selected-id recording expects exactly 6 selected experts\n"); + return 0; + } + + if (!g_moe_selected_trace_record_initialized) { + g_moe_selected_trace_record_initialized = 1; + g_moe_selected_trace_record_fp = fopen(path, "wb"); + if (!g_moe_selected_trace_record_fp) { + fprintf(stderr, "ds4: failed to open selected-id record file %s\n", path); + return 0; + } + setvbuf(g_moe_selected_trace_record_fp, NULL, _IOFBF, 1u << 20); + atexit(ds4_gpu_moe_selected_trace_record_close); + } + + if (fwrite(selected_ids, sizeof(selected_ids[0]), n_selected, g_moe_selected_trace_record_fp) != n_selected) { + fprintf(stderr, "ds4: failed to write selected-id record file %s\n", path); + return 0; + } + if (fflush(g_moe_selected_trace_record_fp) != 0) { + fprintf(stderr, "ds4: failed to flush selected-id record file %s\n", path); + return 0; + } + g_moe_selected_trace_record_count++; + return 1; +} + +static int ds4_gpu_moe_selected_trace_replay( + int32_t selected_ids[6], + uint32_t n_selected) { + const char *path = getenv("DS4_MOE_REPLAY_SELECTED_IDS"); + if (!path || !path[0]) return 0; + if (n_selected != 6) { + fprintf(stderr, "ds4: selected-id replay expects exactly 6 selected experts\n"); + return -1; + } + + if (!g_moe_selected_trace_replay_initialized) { + g_moe_selected_trace_replay_initialized = 1; + FILE *fp = fopen(path, "rb"); + if (!fp) { + fprintf(stderr, "ds4: failed to open selected-id replay file %s\n", path); + return -1; + } + if (fseeko(fp, 0, SEEK_END) != 0) { + fprintf(stderr, "ds4: failed to seek selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + const off_t end = ftello(fp); + if (end < 0) { + fprintf(stderr, "ds4: failed to size selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + if (fseeko(fp, 0, SEEK_SET) != 0) { + fprintf(stderr, "ds4: failed to rewind selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + + const uint64_t bytes = (uint64_t)end; + const uint64_t entry_bytes = (uint64_t)n_selected * sizeof(selected_ids[0]); + if (bytes == 0 || (bytes % entry_bytes) != 0) { + fprintf(stderr, + "ds4: selected-id replay file %s has invalid size %" PRIu64 "\n", + path, + bytes); + fclose(fp); + return -1; + } + if (bytes > SIZE_MAX) { + fprintf(stderr, "ds4: selected-id replay file %s is too large\n", path); + fclose(fp); + return -1; + } + g_moe_selected_trace_replay_count = bytes / entry_bytes; + g_moe_selected_trace_replay_ids = malloc((size_t)bytes); + if (!g_moe_selected_trace_replay_ids) { + fprintf(stderr, "ds4: failed to allocate selected-id replay buffer\n"); + fclose(fp); + return -1; + } + if (fread(g_moe_selected_trace_replay_ids, 1, (size_t)bytes, fp) != (size_t)bytes) { + fprintf(stderr, "ds4: failed to read selected-id replay file %s\n", path); + fclose(fp); + free(g_moe_selected_trace_replay_ids); + g_moe_selected_trace_replay_ids = NULL; + return -1; + } + fclose(fp); + fprintf(stderr, + "ds4: loaded %" PRIu64 " routed-MoE selected-id entries from %s\n", + g_moe_selected_trace_replay_count, + path); + } + + if (g_moe_selected_trace_replay_pos >= g_moe_selected_trace_replay_count) { + fprintf(stderr, + "ds4: selected-id replay exhausted after %" PRIu64 " entries\n", + g_moe_selected_trace_replay_pos); + return -1; + } + memcpy(selected_ids, + g_moe_selected_trace_replay_ids + g_moe_selected_trace_replay_pos * n_selected, + (size_t)n_selected * sizeof(selected_ids[0])); + g_moe_selected_trace_replay_pos++; + return 1; +} + +static int ds4_gpu_progress_enabled(void) { + return ds4_log_is_tty(stderr); +} + +static void ds4_gpu_progress_begin(const char *what) { + if (!ds4_gpu_progress_enabled()) return; + fprintf(stderr, "ds4: %s...", what); + fflush(stderr); +} + +static void ds4_gpu_progress_done(void) { + if (!ds4_gpu_progress_enabled()) return; + fputs(" done\n", stderr); + fflush(stderr); +} + +static void ds4_gpu_progress_failed(void) { + if (!ds4_gpu_progress_enabled()) return; + fputs(" failed\n", stderr); + fflush(stderr); +} + +static void ds4_gpu_model_views_clear(void) { + for (uint32_t i = 0; i < g_model_view_count; i++) { + g_model_views[i].buffer = nil; + g_model_views[i].model_map = NULL; + g_model_views[i].model_size = 0; + g_model_views[i].model_offset = 0; + g_model_views[i].bytes = 0; + } + g_model_view_count = 0; +} + +static void ds4_gpu_model_residency_clear(void) { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (g_model_residency_set) { + if (g_model_residency_added_to_queue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:g_model_residency_set]; + } + [g_model_residency_set endResidency]; + [g_model_residency_set removeAllAllocations]; + g_model_residency_set = nil; + } + } +#endif + g_model_residency_count = 0; + g_model_residency_added_to_queue = 0; +} + +/* TP sharding keeps only this rank's expert ranges warm, + * so whole-view residency requests (which would page in the full file) + * must be skipped; pages fault in lazily through the same view buffers, + * exactly like ssd-streaming mode. */ +static int g_model_residency_skipped; + +void ds4_gpu_model_residency_skip(int skip) { + g_model_residency_skipped = skip; +} + +static int ds4_gpu_model_residency_request_views(void) { + if (g_model_view_count == 0 || + g_ssd_streaming_mode || + g_model_residency_skipped || + getenv("DS4_METAL_NO_RESIDENCY") != NULL) { + return 1; + } + +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + /* + * Register all model views as one residency set before inference. This + * is a GPU residency/budgeting hint, not a request to fault the whole + * 80+ GB file into memory. Its purpose is to make the driver see the + * complete set of large shared allocations during setup instead of + * discovering them lazily from the first measured graph command, where + * VM validation and residency accounting would look like model compute. + */ + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_model"; + desc.initialCapacity = g_model_view_count; + + NSError *error = nil; + g_model_residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!g_model_residency_set) { + fprintf(stderr, "ds4: Metal model residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return 0; + } + + for (uint32_t i = 0; i < g_model_view_count; i++) { + [g_model_residency_set addAllocation:g_model_views[i].buffer]; + } + [g_model_residency_set commit]; + [g_model_residency_set requestResidency]; + if (getenv("DS4_METAL_DISABLE_QUEUE_RESIDENCY_SET") == NULL && + g_queue && + [g_queue respondsToSelector:@selector(addResidencySet:)]) { + [g_queue addResidencySet:g_model_residency_set]; + g_model_residency_added_to_queue = 1; + } + g_model_residency_count = g_model_view_count; + } +#endif + + return 1; +} + +static int ds4_gpu_add_model_view_range( + const void *model_map, + uint64_t model_size, + uint64_t map_offset, + uint64_t map_size, + uint64_t max_tensor_bytes, + bool use_default_view_cap, + uint64_t *mapped_model_size_out) { + const uint64_t page = (uint64_t)getpagesize(); + const uintptr_t model_addr = (uintptr_t)model_map; + + if ((model_addr & (uintptr_t)(page - 1)) != 0) { + fprintf(stderr, "ds4: Metal model mmap base is not page aligned\n"); + return 0; + } + if (map_offset > model_size || map_size > model_size - map_offset) { + fprintf(stderr, "ds4: Metal model mapped range is outside the GGUF mapping\n"); + return 0; + } + const uint64_t page_model_offset = map_offset & ~(page - 1); + const uint64_t leading = map_offset - page_model_offset; + if (map_size > UINT64_MAX - leading || + leading + map_size > UINT64_MAX - (page - 1)) + { + fprintf(stderr, "ds4: Metal model mapped range overflows page alignment\n"); + return 0; + } + const uint64_t mapped_model_size = round_up_u64(leading + map_size, page); + uint64_t max_buffer = (uint64_t)[g_device maxBufferLength]; + max_buffer &= ~(page - 1); + + /* + * Wrap only the tensor-data part of the GGUF file. Metadata is parsed by the + * CPU and is never dereferenced by kernels, so exposing it to Metal only + * grows the residency set and the VM range the driver must validate. + * + * Metal buffers have a device-specific maximum length, and this model is + * larger than that maximum on the target machines. Creating one no-copy + * buffer per tensor would avoid the length limit, but it would also move a + * lot of VM-object creation and residency bookkeeping into graph setup. The + * stable shape here is a tiny number of page-aligned views created once. + * + * Adjacent views intentionally overlap by more than the largest tensor, plus + * one page for alignment. That invariant guarantees every tensor lies wholly + * inside at least one view, so hot paths pass one buffer and one inner byte + * offset. We never split a weight tensor across command encoders. + */ + if (max_tensor_bytes > map_size) { + fprintf(stderr, "ds4: Metal model max tensor span is larger than a mapped tensor span\n"); + return 0; + } + if (max_tensor_bytes > UINT64_MAX - (page - 1)) { + fprintf(stderr, "ds4: Metal model max tensor span overflows page alignment\n"); + return 0; + } + const uint64_t max_tensor_rounded = round_up_u64(max_tensor_bytes, page); + if (max_tensor_rounded > UINT64_MAX - page) { + fprintf(stderr, "ds4: Metal model view overlap overflows page slack\n"); + return 0; + } + const uint64_t overlap = max_tensor_rounded + page; + if (max_buffer == 0 || max_buffer <= overlap) { + fprintf(stderr, + "ds4: Metal maxBufferLength is too small for DS4 model views " + "(max tensor %.2f GiB, max buffer %.2f GiB)\n", + ds4_gpu_gib(max_tensor_bytes), + ds4_gpu_gib(max_buffer)); + return 0; + } + + uint64_t view_limit = max_buffer; + const char *view_limit_env = getenv("DS4_METAL_MODEL_VIEW_MAX_GIB"); + if (view_limit_env && view_limit_env[0]) { + char *end = NULL; + unsigned long long gib = strtoull(view_limit_env, &end, 10); + if (end != view_limit_env && gib > 0) { + uint64_t env_limit = gib * 1024ull * 1024ull * 1024ull; + env_limit &= ~(page - 1); + if (env_limit > 0) view_limit = env_limit; + } + } else if (use_default_view_cap && mapped_model_size > max_buffer) { + /* + * Very large no-copy buffers can make Metal's VM validation dominate + * startup or the first graph command on multi-hundred-GiB slices. Keep + * ordinary contiguous model mappings unchanged, but let distributed + * span maps use smaller overlapping views when a range already has to + * be split. + */ + const uint64_t default_limit = 128ull * 1024ull * 1024ull * 1024ull; + if (view_limit > default_limit) view_limit = default_limit; + } + if (view_limit > max_buffer) view_limit = max_buffer; + view_limit &= ~(page - 1); + if (view_limit == 0 || view_limit <= overlap) { + fprintf(stderr, + "ds4: Metal model view cap is too small for DS4 model views " + "(cap %.2f GiB, max tensor %.2f GiB)\n", + ds4_gpu_gib(view_limit), + ds4_gpu_gib(max_tensor_bytes)); + return 0; + } + + const uint64_t step = view_limit - overlap; + uint64_t off = 0; + while (off < mapped_model_size) { + if (g_model_view_count == DS4_METAL_MAX_MODEL_VIEWS) { + fprintf(stderr, "ds4: Metal model needs more mapped views than expected\n"); + return 0; + } + + uint64_t view_bytes = mapped_model_size - off; + if (view_bytes > view_limit) view_bytes = view_limit; + + id buffer = [g_device newBufferWithBytesNoCopy:(void *)(model_addr + page_model_offset + off) + length:(NSUInteger)view_bytes + options:ds4_gpu_model_resource_options() + deallocator:nil]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal could not wrap mmaped model view at %.2f GiB, size %.2f GiB\n", + (double)(page_model_offset + off) / (1024.0 * 1024.0 * 1024.0), + (double)view_bytes / (1024.0 * 1024.0 * 1024.0)); + return 0; + } + buffer.label = [NSString stringWithFormat:@"ds4_model_view_%u", g_model_view_count]; + + g_model_views[g_model_view_count].buffer = buffer; + g_model_views[g_model_view_count].model_map = model_map; + g_model_views[g_model_view_count].model_size = model_size; + g_model_views[g_model_view_count].model_offset = page_model_offset + off; + g_model_views[g_model_view_count].bytes = view_bytes; + g_model_view_count++; + + g_model_wrap_count++; + g_model_wrap_bytes += view_bytes; + if (view_bytes > g_model_wrap_max_bytes) g_model_wrap_max_bytes = view_bytes; + + if (off + view_bytes >= mapped_model_size) break; + off += step; + } + + if (mapped_model_size_out) *mapped_model_size_out += mapped_model_size; + return 1; +} + +static int ds4_gpu_finish_model_views( + double t0, + uint64_t mapped_model_size, + uint64_t display_offset) { + const double t_mapped = ds4_gpu_now_ms(); + const int request_residency = + !g_ssd_streaming_mode && + getenv("DS4_METAL_NO_RESIDENCY") == NULL; + if (request_residency) ds4_gpu_progress_begin("requesting Metal residency (may take tens of seconds)"); + if (!ds4_gpu_model_residency_request_views()) { + if (request_residency) ds4_gpu_progress_failed(); + return 0; + } + if (request_residency) ds4_gpu_progress_done(); + const double t_resident = ds4_gpu_now_ms(); + int warmed = 1; + const double t_warm0 = ds4_gpu_now_ms(); + const int warm_model_views = !g_ssd_streaming_mode && + getenv("DS4_METAL_NO_RESIDENCY") == NULL && + getenv("DS4_METAL_NO_MODEL_WARMUP") == NULL; + if (warm_model_views) { + /* + * The first GPU command touching no-copy mmap storage can pay command + * queue setup, page-table validation, and shared-allocation residency + * costs. Sample each model view here so timed graph execution starts + * after that one-time work. The stride is intentionally coarse: this is + * a validation touch over the VM ranges, not a full model prefetch. A + * dense prefetch would create exactly the kind of memory pressure and + * startup stalls this path is designed to avoid. + */ + if (g_model_residency_skipped) { + /* TP sharding: a single command buffer binding every + * view demands residency of them all and OOMs; the engine's + * CPU-side sharded warm pre-faults the owned bytes instead. */ + warmed = 1; + } else { + ds4_gpu_progress_begin("warming Metal model views"); + warmed = ds4_gpu_warm_model_views(); + if (warmed) ds4_gpu_progress_done(); + else ds4_gpu_progress_failed(); + } + } + const double t_warm = ds4_gpu_now_ms(); + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal model views created in %.3f ms, residency requested in %.3f ms, warmup %.3f ms (mapped %.2f MiB from offset %.2f MiB)\n", + t_mapped - t0, + t_resident - t_mapped, + t_warm - t_warm0, + mapped_model_size / 1024.0 / 1024.0, + display_offset / 1024.0 / 1024.0); + } + if (!warmed) return 0; + return 1; +} + +static int ds4_gpu_map_model_views( + const void *model_map, + uint64_t model_size, + uint64_t map_offset, + uint64_t map_size, + uint64_t max_tensor_bytes) { + const double t0 = ds4_gpu_now_ms(); + uint64_t mapped_model_size = 0; + if (!ds4_gpu_add_model_view_range(model_map, + model_size, + map_offset, + map_size, + max_tensor_bytes, + false, + &mapped_model_size)) { + return 0; + } + return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset); +} + +static id ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) { + if (bytes == 0) bytes = 1; + + id buffer = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, "ds4: failed to allocate Metal transient buffer %s (%llu bytes)\n", + label ? label : "(unnamed)", (unsigned long long)bytes); + return nil; + } + if (label) buffer.label = [NSString stringWithUTF8String:label]; + + /* + * CPU-filled buffers must survive until their command buffer completes. + * A local ObjC strong variable is not enough when the encoder function + * returns before the caller commits the command buffer. + */ + [g_transient_buffers addObject:buffer]; + return buffer; +} + +static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) { + if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) { + return 0; + } + return ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL; +} + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry * +ds4_gpu_get_zero_prefix_prefill_mask_cache( + uint32_t kind, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t n_keys, + uint32_t window, + uint32_t ratio, + uint32_t nqptg, + uint32_t ncpsg, + bool has_kvpad, + bool bc_mask, + NSUInteger mask_bytes, + NSUInteger blk_bytes, + bool *created) { + if (created) *created = false; + if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() || + mask_bytes == 0 || blk_bytes == 0) { + return NULL; + } + + uint32_t slot = UINT32_MAX; + switch (kind) { + case DS4_GPU_PREFILL_MASK_CACHE_RAW: slot = 0; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO4: slot = 1; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break; + default: return NULL; + } + + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[slot]; + if (entry->valid && entry->mask && entry->blk && + entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes && + entry->kind == kind && entry->n_tokens == n_tokens && + entry->n_comp == n_comp && entry->n_keys == n_keys && + entry->window == window && entry->ratio == ratio && + entry->nqptg == nqptg && entry->ncpsg == ncpsg && + entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) { + return entry; + } + + id mask = [g_device newBufferWithLength:mask_bytes + options:MTLResourceStorageModeShared]; + id blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModePrivate]; + if (!blk) { + blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModeShared]; + } + if (!mask || !blk) return NULL; + + mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind]; + blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind]; + + /* A replaced entry may still be referenced by an uncommitted or in-flight + * command buffer. Keep its resources alive with the other batch-scoped + * buffers until command completion instead of mutating or releasing them. */ + if (entry->mask) [g_transient_buffers addObject:entry->mask]; + if (entry->blk) [g_transient_buffers addObject:entry->blk]; + + entry->mask = mask; + entry->blk = blk; + entry->mask_bytes = mask_bytes; + entry->blk_bytes = blk_bytes; + entry->kind = kind; + entry->n_tokens = n_tokens; + entry->n_comp = n_comp; + entry->n_keys = n_keys; + entry->window = window; + entry->ratio = ratio; + entry->nqptg = nqptg; + entry->ncpsg = ncpsg; + entry->has_kvpad = has_kvpad; + entry->bc_mask = bc_mask; + /* The caller publishes the entry only after synchronously filling mask. */ + entry->valid = false; + entry->blk_ready = false; + *created = true; + return entry; +} + +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + g_zero_prefix_prefill_mask_cache[i].blk_ready = false; + } +} + +static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + entry->mask = nil; + entry->blk = nil; + entry->mask_bytes = 0; + entry->blk_bytes = 0; + entry->kind = 0; + entry->n_tokens = 0; + entry->n_comp = 0; + entry->n_keys = 0; + entry->window = 0; + entry->ratio = 0; + entry->nqptg = 0; + entry->ncpsg = 0; + entry->has_kvpad = false; + entry->bc_mask = false; + entry->valid = false; + entry->blk_ready = false; + } +} + +void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) { + /* Layer-major prefill waits each layer before advancing, so its final + * release point has no outstanding cache users. Keep the guard here for + * diagnostic callers that may have an open or asynchronously flushed CB. */ + if (!g_initialized || g_batch_cb || + (g_pending_cbs && [g_pending_cbs count] != 0)) { + return; + } + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); +} + +static id ds4_gpu_get_mul_mm_pipeline( + const char *function_name, + bool bc_inp, + bool bc_out) { + NSString *key = [NSString stringWithFormat:@"%s_bci=%d_bco=%d", + function_name, bc_inp ? 1 : 0, bc_out ? 1 : 0]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; + [constants setConstantValue:&bc_out type:MTLDataTypeBool atIndex:701]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_mul_mm_id_pipeline( + const char *function_name, + bool bc_inp) { + NSString *key = [NSString stringWithFormat:@"%s_bci=%d", + function_name, bc_inp ? 1 : 0]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_pipeline( + const char *function_name) { + NSString *key = [NSString stringWithFormat:@"%s", function_name]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found\n", function_name); + return nil; + } + + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static int ds4_gpu_disable_hot_pipeline_statics(void) { + static int initialized; + static int disabled; + if (!initialized) { + disabled = getenv("DS4_METAL_DISABLE_HOT_PIPELINE_STATICS") != NULL; + initialized = 1; + } + return disabled; +} + +static id ds4_gpu_hot_pipeline( + id pipeline, + const char *fallback_name) { + if (!ds4_gpu_disable_hot_pipeline_statics()) return pipeline; + return ds4_gpu_get_pipeline(fallback_name); +} + +static int ds4_gpu_use_compressor_pair_nr4(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = getenv("DS4_METAL_COMPRESSOR_PAIR_NR4") != NULL; + initialized = 1; + } + return enabled; +} + +static int ds4_gpu_device_name_contains(const char *needle); + +static int ds4_gpu_env_value_eq(const char *v, size_t n, const char *literal) { + size_t m = strlen(literal); + if (n != m) return 0; + for (size_t i = 0; i < n; i++) { + if (tolower((unsigned char)v[i]) != tolower((unsigned char)literal[i])) return 0; + } + return 1; +} + +static int ds4_gpu_env_bool(const char *name) { + const char *v = getenv(name); + if (!v) return -1; + + while (isspace((unsigned char)*v)) v++; + size_t n = strlen(v); + while (n > 0 && isspace((unsigned char)v[n - 1])) n--; + if (n == 0) return 1; + + if (ds4_gpu_env_value_eq(v, n, "1") || + ds4_gpu_env_value_eq(v, n, "true") || + ds4_gpu_env_value_eq(v, n, "yes") || + ds4_gpu_env_value_eq(v, n, "on")) { + return 1; + } + if (ds4_gpu_env_value_eq(v, n, "0") || + ds4_gpu_env_value_eq(v, n, "false") || + ds4_gpu_env_value_eq(v, n, "no") || + ds4_gpu_env_value_eq(v, n, "off")) { + return 0; + } + + if (!g_mpp_invalid_env_reported) { + fprintf(stderr, + "ds4: invalid Metal boolean environment value %s=%.*s; treating presence as enabled\n", + name, (int)n, v); + g_mpp_invalid_env_reported = 1; + } + return 1; +} + +static uint64_t ds4_gpu_env_u64(const char *name, + uint64_t fallback, + uint64_t min_value, + uint64_t max_value) { + const char *v = getenv(name); + if (!v) return fallback; + while (isspace((unsigned char)*v)) v++; + if (!*v) return fallback; + + errno = 0; + char *end = NULL; + unsigned long long parsed = strtoull(v, &end, 10); + if (end == v || errno == ERANGE) return fallback; + while (isspace((unsigned char)*end)) end++; + if (*end) return fallback; + + if (parsed < min_value) return fallback; + uint64_t value = (uint64_t)parsed; + if (value > max_value) value = max_value; + return value; +} + +static uint32_t ds4_gpu_glm_full_attention_max_cache_len(void) { + /* + * kernel_glm_attention_full stores one score per visible token in + * threadgroup memory, plus 256 reduction slots. Keep the default under + * the 32 KiB envelope used by current Apple GPUs while allowing short + * dense decode to step past the 4096-token prefill boundary. + */ + return 7680u; +} + +static uint32_t ds4_gpu_glm_flash_attention_max_cache_len(void) { + /* + * Staged FlashAttention uses fixed-size threadgroup scratch and separate + * KV staging buffers, so it is not bound by the legacy full-attention + * kernel's one-score-per-token threadgroup-memory envelope. + */ + return 8192u; +} + +static int ds4_gpu_mpp_available(void) { + return g_metal4_tensor_api_enabled && !g_quality_mode; +} + +/* + * Retained Metal4 defaults live here instead of behind user-visible options. + * The public runtime has one automatic accelerated path plus the global + * DS4_METAL_DISABLE_METAL4 comparison switch. Benchmark-only alternatives that + * lost during M5 work are removed or kept out of the dispatch path so future + * changes do not accidentally turn old experiments into new modes. + */ +static int ds4_gpu_use_mpp_attn_out_low_matmul(void) { + return ds4_gpu_mpp_available(); +} + +enum { + DS4_METAL_ATTN_OUT_MPP_TILE_N = 64, +}; + +static void ds4_gpu_warn_mpp_fallback(void) { + static int warned; + if (!warned) { + fprintf(stderr, "ds4: accelerated Metal prefill matmul unavailable; falling back to legacy kernel\n"); + warned = 1; + } +} + +static int ds4_gpu_device_name_contains(const char *needle) { + return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL; +} + +static int ds4_gpu_compile_tensor_probe(void) { +#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 + if (!g_device) return 0; + if (@available(macOS 26.0, *)) { + const char *src = + "#include \n" + "#include \n" + "#include \n" + "using namespace metal;\n" + "using namespace mpp::tensor_ops;\n" + "kernel void ds4_tensor_probe(\n" + " tensor> A [[buffer(0)]],\n" + " tensor> B [[buffer(1)]],\n" + " device float *C [[buffer(2)]],\n" + " uint2 tgid [[threadgroup_position_in_grid]]) {\n" + " auto tA = A.slice(0, (int)tgid.y);\n" + " auto tB = B.slice((int)tgid.x, 0);\n" + " matmul2d> mm;\n" + " auto cT = mm.get_destination_cooperative_tensor();\n" + " auto sA = tA.slice(0, 0);\n" + " auto sB = tB.slice(0, 0);\n" + " mm.run(sB, sA, cT);\n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16));\n" + " cT.store(tC);\n" + "}\n"; + + NSError *error = nil; + NSString *source = [NSString stringWithUTF8String:src]; + id probe_library = [g_device newLibraryWithSource:source options:[MTLCompileOptions new] error:&error]; + if (!probe_library) { + fprintf(stderr, "ds4: Metal 4 tensor API probe compile failed: %s\n", + error ? [[error localizedDescription] UTF8String] : "(unknown)"); + return 0; + } + id fn = [probe_library newFunctionWithName:@"ds4_tensor_probe"]; + if (!fn) { + fprintf(stderr, "ds4: Metal 4 tensor API probe function missing\n"); + return 0; + } + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal 4 tensor API probe pipeline failed: %s\n", + error ? [[error localizedDescription] UTF8String] : "(unknown)"); + return 0; + } + return 1; + } +#endif + return 0; +} + +static void ds4_gpu_detect_metal4_features(void) { + g_metal4_runtime_available = 0; + g_metal4_family_supported = 0; + g_metal4_queue_supported = 0; + g_metal4_m5_neural_accelerators_hint = 0; + g_metal4_tensor_api_enabled = 0; + g_metal4_tensor_api_compile_supported = 0; + g_metal_device_name[0] = '\0'; + + if (!g_device) return; + + const char *name = [[g_device name] UTF8String]; + if (name) { + snprintf(g_metal_device_name, sizeof(g_metal_device_name), "%s", name); + } + + const int metal4_disabled = ds4_gpu_env_bool("DS4_METAL_DISABLE_METAL4") > 0; + +#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 + if (@available(macOS 26.0, *)) { + g_metal4_runtime_available = 1; + g_metal4_family_supported = + !metal4_disabled && [g_device supportsFamily:MTLGPUFamilyMetal4] ? 1 : 0; + g_metal4_queue_supported = [g_device respondsToSelector:@selector(newMTL4CommandQueue)] ? 1 : 0; + + /* + * Apple does not currently expose a separate "Neural Accelerator" bit + * through Metal. On public M5 systems the hardware signal is the device + * generation plus Metal 4 support, so keep this as a conservative hint. + */ + if (g_metal4_family_supported && ds4_gpu_device_name_contains("M5")) { + g_metal4_m5_neural_accelerators_hint = 1; + } + + if (g_metal4_family_supported) { + const int default_enable = + ds4_gpu_device_name_contains("M5") || + ds4_gpu_device_name_contains("M6") || + ds4_gpu_device_name_contains("A19") || + ds4_gpu_device_name_contains("A20"); + + /* + * Metal 4 TensorOps are portable in source, but on pre-M5 hardware + * they can map to ordinary shader fallbacks. Keep the automatic + * fast path restricted to hardware generations where the Neural + * Accelerator/TensorOps path is expected to pay off; older Metal + * machines continue to use the established kernels unless a future + * device is explicitly added here. + */ + if (default_enable) { + g_metal4_tensor_api_compile_supported = ds4_gpu_compile_tensor_probe(); + g_metal4_tensor_api_enabled = g_metal4_tensor_api_compile_supported; + if (!g_metal4_tensor_api_enabled) { + fprintf(stderr, "ds4: Metal 4 tensor API probe failed; using legacy Metal kernels\n"); + } + } else { + fprintf(stderr, "ds4: Metal 4 tensor API disabled for pre-M5/pre-A19 devices\n"); + } + } + } +#endif +} + +static int ds4_gpu_warm_model_views(void) { + if (g_model_view_count == 0) return 1; + + id pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride"); + if (!pipeline) return 0; + + uint64_t stride = 1024ull * 1024ull; + const char *stride_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_MB"); + if (stride_env && stride_env[0]) { + char *end = NULL; + unsigned long long mb = strtoull(stride_env, &end, 10); + if (end != stride_env && mb > 0 && mb <= 1024) { + stride = mb * 1024ull * 1024ull; + } + } + const char *stride_kb_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_KB"); + if (stride_kb_env && stride_kb_env[0]) { + char *end = NULL; + unsigned long long kb = strtoull(stride_kb_env, &end, 10); + if (end != stride_kb_env && kb > 0 && kb <= 1024ull * 1024ull) { + stride = kb * 1024ull; + const uint64_t page = (uint64_t)getpagesize(); + if (stride < page) stride = page; + } + } + + uint64_t total_touches = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + total_touches += (g_model_views[i].bytes + stride - 1) / stride; + } + if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0; + + const NSUInteger out_bytes = (NSUInteger)total_touches; + id out = [g_device newBufferWithLength:out_bytes + options:MTLResourceStorageModeShared]; + if (!out) { + fprintf(stderr, "ds4: Metal model warmup scratch allocation failed\n"); + return 0; + } + out.label = @"ds4_model_warmup"; + + id cb = ds4_gpu_new_command_buffer(); + if (!cb) { + fprintf(stderr, "ds4: Metal model warmup command buffer allocation failed\n"); + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + uint64_t dst_offset = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + const uint64_t bytes = g_model_views[i].bytes; + const uint64_t n = (bytes + stride - 1) / stride; + [enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0]; + [enc setBuffer:out offset:0 atIndex:1]; + [enc setBytes:&stride length:sizeof(stride) atIndex:2]; + [enc setBytes:&bytes length:sizeof(bytes) atIndex:3]; + [enc setBytes:&dst_offset length:sizeof(dst_offset) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + dst_offset += n; + } + ds4_gpu_end_compute_encoder(cb, enc); + + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal model warmup failed: %s\n", + [[cb.error localizedDescription] UTF8String]); + return 0; + } + + return 1; +} + +static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) { + switch (ne20) { + case 1: return "kernel_mul_mm_id_map0_ne20_1"; + case 2: return "kernel_mul_mm_id_map0_ne20_2"; + case 4: return "kernel_mul_mm_id_map0_ne20_4"; + case 5: return "kernel_mul_mm_id_map0_ne20_5"; + case 6: return "kernel_mul_mm_id_map0_ne20_6"; + case 8: return "kernel_mul_mm_id_map0_ne20_8"; + case 10: return "kernel_mul_mm_id_map0_ne20_10"; + case 16: return "kernel_mul_mm_id_map0_ne20_16"; + case 22: return "kernel_mul_mm_id_map0_ne20_22"; + default: return NULL; + } +} + +static id ds4_gpu_get_mul_mv_pipeline( + const char *function_name, + int16_t nsg) { + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_mul_mv_ext_pipeline( + const char *function_name, + int16_t nsg, + int16_t nxpsg) { + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d", + function_name, (int)nsg, (int)nxpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; + [constants setConstantValue:&nxpsg type:MTLDataTypeShort atIndex:601]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_pad_pipeline( + bool has_mask, + int32_t ncpsg) { + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d", + has_mask ? 1 : 0, (int)ncpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100]; + [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:125]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_pad" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_blk_pipeline( + int32_t nqptg, + int32_t ncpsg) { + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d", + (int)nqptg, (int)ncpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224]; + [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:225]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_blk" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_pipeline( + const char *function_name, + bool has_mask, + bool has_sinks, + bool has_bias, + bool has_scap, + bool has_kvpad, + bool bc_mask, + int32_t ns10, + int32_t ns20, + int32_t nsg) { + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d", + function_name, + has_mask ? 1 : 0, + has_sinks ? 1 : 0, + has_bias ? 1 : 0, + has_scap ? 1 : 0, + has_kvpad ? 1 : 0, + bc_mask ? 1 : 0, + (int)ns10, + (int)ns20, + (int)nsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300]; + [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:301]; + [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:302]; + [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:303]; + [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:304]; + [constants setConstantValue:&bc_mask type:MTLDataTypeBool atIndex:310]; + [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:320]; + [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:321]; + [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:322]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_vec_pipeline( + const char *function_name, + bool has_mask, + bool has_sinks, + bool has_bias, + bool has_scap, + bool has_kvpad, + bool shared_kvpad, + int32_t ns10, + int32_t ns20, + int32_t nsg, + int32_t nwg) { + /* + * Decode calls this once per layer with identical arguments, so memoize + * the last hit and skip the NSString key + dictionary lookup on the hot + * path. The generic cache below remains the fallback for new variants. + */ + static struct { + const char *fn; + bool m, s, b, c, k, sp; + int32_t n10, n20, sg, wg; + id pipeline; + } memo; + if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 && + memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias && + memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad && + memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) { + return memo.pipeline; + } + + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", + function_name, + has_mask ? 1 : 0, + has_sinks ? 1 : 0, + has_bias ? 1 : 0, + has_scap ? 1 : 0, + has_kvpad ? 1 : 0, + shared_kvpad ? 1 : 0, + (int)ns10, + (int)ns20, + (int)nsg, + (int)nwg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) { + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, cached }; + return cached; + } + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:400]; + [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:401]; + [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:402]; + [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:403]; + [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404]; + [constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405]; + [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420]; + [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421]; + [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:422]; + [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:423]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, pipeline }; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_reduce_pipeline( + int32_t dv, + int32_t nwg) { + /* Same per-layer memo pattern as the vec getter above. */ + static int32_t memo_dv, memo_nwg; + static id memo_pipeline; + if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) { + return memo_pipeline; + } + + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d", + (int)dv, (int)nwg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) { + memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached; + return cached; + } + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500]; + [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline; + return pipeline; +} + +static uint32_t ds4_gpu_flash_attn_vec_nsg(uint32_t n_keys, uint32_t nwg, uint32_t ncpsg) { + uint32_t nsg = 1; + while (2u * nwg * nsg * ncpsg < n_keys && nsg < 4u) { + nsg *= 2u; + } + return nsg; +} + +static int ds4_gpu_trace_allocs(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = getenv("DS4_METAL_TRACE_ALLOCS") != NULL; + initialized = 1; + } + return enabled; +} + +static double ds4_gpu_mib(uint64_t bytes) { + return (double)bytes / (1024.0 * 1024.0); +} + +static double ds4_gpu_gib(uint64_t bytes) { + return (double)bytes / (1024.0 * 1024.0 * 1024.0); +} + +static ds4_gpu_stream_expert_timing_snapshot +ds4_gpu_stream_expert_timing_current(void) { + return (ds4_gpu_stream_expert_timing_snapshot) { + .selected_calls = g_stream_expert_timing_selected_calls, + .selected_read_ms = g_stream_expert_timing_selected_read_ms, + .selected_sync_ms = g_stream_expert_timing_selected_sync_ms, + .selected_copy_ms = g_stream_expert_timing_selected_copy_ms, + .selected_bind_ms = g_stream_expert_timing_selected_bind_ms, + .split_layers = g_stream_expert_timing_split_layers, + .split_resident_experts = g_stream_expert_timing_split_resident_experts, + .split_missing_experts = g_stream_expert_timing_split_missing_experts, + .split_resident_ms = g_stream_expert_timing_split_resident_ms, + .split_missing_ms = g_stream_expert_timing_split_missing_ms, + .split_missing_load_ms = + g_stream_expert_timing_split_missing_load_ms, + .split_missing_slot_ms = + g_stream_expert_timing_split_missing_slot_ms, + .split_missing_prune_ms = + g_stream_expert_timing_split_missing_prune_ms, + .split_missing_addr_ms = + g_stream_expert_timing_split_missing_addr_ms, + .split_missing_wait_ms = + g_stream_expert_timing_split_missing_wait_ms, + .load_calls = g_stream_expert_timing_load_calls, + .load_prepare_ms = g_stream_expert_timing_load_prepare_ms, + .load_pread_ms = g_stream_expert_timing_load_pread_ms, + .load_modify_ms = g_stream_expert_timing_load_modify_ms, + .load_install_ms = g_stream_expert_timing_load_install_ms, + .prepare_batch_reuse_calls = + g_stream_expert_timing_prepare_batch_reuse_calls, + .prepare_batch_reuse_ms = + g_stream_expert_timing_prepare_batch_reuse_ms, + .prepare_buffer_calls = + g_stream_expert_timing_prepare_buffer_calls, + .prepare_buffer_ms = + g_stream_expert_timing_prepare_buffer_ms, + .prepare_task_experts = + g_stream_expert_timing_prepare_task_experts, + .prepare_task_ms = + g_stream_expert_timing_prepare_task_ms, + .reuse_scan_calls = + g_stream_expert_timing_reuse_scan_calls, + .reuse_scan_entries = + g_stream_expert_timing_reuse_scan_entries, + .reuse_scan_ms = + g_stream_expert_timing_reuse_scan_ms, + .reuse_clear_ms = + g_stream_expert_timing_reuse_clear_ms, + .readahead_calls = + g_stream_expert_timing_readahead_calls, + .readahead_bytes = + g_stream_expert_timing_readahead_bytes, + .readahead_ms = + g_stream_expert_timing_readahead_ms, + .cache_all_resident_layers = + g_stream_expert_timing_cache_all_resident_layers, + .cache_all_missing_layers = + g_stream_expert_timing_cache_all_missing_layers, + .cache_mixed_layers = g_stream_expert_timing_cache_mixed_layers, + .cache_resident_experts = + g_stream_expert_timing_cache_resident_experts, + .cache_missing_experts = + g_stream_expert_timing_cache_missing_experts, + }; +} + +static uint64_t ds4_gpu_stream_expert_timing_delta_u64( + uint64_t current, + uint64_t previous) { + return current >= previous ? current - previous : current; +} + +static double ds4_gpu_stream_expert_timing_delta_f64( + double current, + double previous) { + return current >= previous ? current - previous : current; +} + +static ds4_gpu_stream_expert_timing_snapshot +ds4_gpu_stream_expert_timing_delta( + ds4_gpu_stream_expert_timing_snapshot current, + ds4_gpu_stream_expert_timing_snapshot previous) { + return (ds4_gpu_stream_expert_timing_snapshot) { + .selected_calls = + ds4_gpu_stream_expert_timing_delta_u64(current.selected_calls, + previous.selected_calls), + .selected_read_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_read_ms, + previous.selected_read_ms), + .selected_sync_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_sync_ms, + previous.selected_sync_ms), + .selected_copy_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_copy_ms, + previous.selected_copy_ms), + .selected_bind_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_bind_ms, + previous.selected_bind_ms), + .split_layers = + ds4_gpu_stream_expert_timing_delta_u64(current.split_layers, + previous.split_layers), + .split_resident_experts = + ds4_gpu_stream_expert_timing_delta_u64(current.split_resident_experts, + previous.split_resident_experts), + .split_missing_experts = + ds4_gpu_stream_expert_timing_delta_u64(current.split_missing_experts, + previous.split_missing_experts), + .split_resident_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.split_resident_ms, + previous.split_resident_ms), + .split_missing_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.split_missing_ms, + previous.split_missing_ms), + .split_missing_load_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_load_ms, + previous.split_missing_load_ms), + .split_missing_slot_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_slot_ms, + previous.split_missing_slot_ms), + .split_missing_prune_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_prune_ms, + previous.split_missing_prune_ms), + .split_missing_addr_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_addr_ms, + previous.split_missing_addr_ms), + .split_missing_wait_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_wait_ms, + previous.split_missing_wait_ms), + .load_calls = + ds4_gpu_stream_expert_timing_delta_u64(current.load_calls, + previous.load_calls), + .load_prepare_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_prepare_ms, + previous.load_prepare_ms), + .load_pread_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_pread_ms, + previous.load_pread_ms), + .load_modify_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_modify_ms, + previous.load_modify_ms), + .load_install_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_install_ms, + previous.load_install_ms), + .prepare_batch_reuse_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_batch_reuse_calls, + previous.prepare_batch_reuse_calls), + .prepare_batch_reuse_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_batch_reuse_ms, + previous.prepare_batch_reuse_ms), + .prepare_buffer_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_buffer_calls, + previous.prepare_buffer_calls), + .prepare_buffer_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_buffer_ms, + previous.prepare_buffer_ms), + .prepare_task_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_task_experts, + previous.prepare_task_experts), + .prepare_task_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_task_ms, + previous.prepare_task_ms), + .reuse_scan_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.reuse_scan_calls, + previous.reuse_scan_calls), + .reuse_scan_entries = + ds4_gpu_stream_expert_timing_delta_u64( + current.reuse_scan_entries, + previous.reuse_scan_entries), + .reuse_scan_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.reuse_scan_ms, + previous.reuse_scan_ms), + .reuse_clear_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.reuse_clear_ms, + previous.reuse_clear_ms), + .readahead_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.readahead_calls, + previous.readahead_calls), + .readahead_bytes = + ds4_gpu_stream_expert_timing_delta_u64( + current.readahead_bytes, + previous.readahead_bytes), + .readahead_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.readahead_ms, + previous.readahead_ms), + .cache_all_resident_layers = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_all_resident_layers, + previous.cache_all_resident_layers), + .cache_all_missing_layers = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_all_missing_layers, + previous.cache_all_missing_layers), + .cache_mixed_layers = + ds4_gpu_stream_expert_timing_delta_u64(current.cache_mixed_layers, + previous.cache_mixed_layers), + .cache_resident_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_resident_experts, + previous.cache_resident_experts), + .cache_missing_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_missing_experts, + previous.cache_missing_experts), + }; +} + +static int ds4_gpu_stream_expert_timing_has_data( + ds4_gpu_stream_expert_timing_snapshot s) { + return s.selected_calls != 0 || + s.split_layers != 0 || + s.load_calls != 0 || + s.cache_all_resident_layers != 0 || + s.cache_all_missing_layers != 0 || + s.cache_mixed_layers != 0; +} + +static void ds4_gpu_stream_expert_timing_print( + const char *scope, + ds4_gpu_stream_expert_timing_snapshot s) { + if (!ds4_gpu_stream_expert_timing_has_data(s)) return; + const double selected_calls = (double)s.selected_calls; + const double split_layers = (double)s.split_layers; + const double selected_read_avg = + selected_calls != 0.0 ? s.selected_read_ms / selected_calls : 0.0; + const double selected_sync_avg = + selected_calls != 0.0 ? s.selected_sync_ms / selected_calls : 0.0; + const double selected_copy_avg = + selected_calls != 0.0 ? s.selected_copy_ms / selected_calls : 0.0; + const double selected_bind_avg = + selected_calls != 0.0 ? s.selected_bind_ms / selected_calls : 0.0; + const double split_resident_avg = + split_layers != 0.0 ? s.split_resident_ms / split_layers : 0.0; + const double split_missing_avg = + split_layers != 0.0 ? s.split_missing_ms / split_layers : 0.0; + const double split_missing_load_avg = + split_layers != 0.0 ? + s.split_missing_load_ms / split_layers : 0.0; + const double split_missing_slot_avg = + split_layers != 0.0 ? + s.split_missing_slot_ms / split_layers : 0.0; + const double split_missing_prune_avg = + split_layers != 0.0 ? + s.split_missing_prune_ms / split_layers : 0.0; + const double split_missing_addr_avg = + split_layers != 0.0 ? + s.split_missing_addr_ms / split_layers : 0.0; + const double split_missing_wait_avg = + split_layers != 0.0 ? + s.split_missing_wait_ms / split_layers : 0.0; + const double load_calls = (double)s.load_calls; + const double load_prepare_avg = + load_calls != 0.0 ? s.load_prepare_ms / load_calls : 0.0; + const double load_pread_avg = + load_calls != 0.0 ? s.load_pread_ms / load_calls : 0.0; + const double load_modify_avg = + load_calls != 0.0 ? s.load_modify_ms / load_calls : 0.0; + const double load_install_avg = + load_calls != 0.0 ? s.load_install_ms / load_calls : 0.0; + const double prepare_batch_reuse_avg = + s.prepare_batch_reuse_calls != 0 ? + s.prepare_batch_reuse_ms / + (double)s.prepare_batch_reuse_calls : 0.0; + const double prepare_buffer_avg = + s.prepare_buffer_calls != 0 ? + s.prepare_buffer_ms / (double)s.prepare_buffer_calls : 0.0; + const double prepare_task_avg = + s.prepare_task_experts != 0 ? + s.prepare_task_ms / (double)s.prepare_task_experts : 0.0; + const double reuse_scan_avg = + s.reuse_scan_calls != 0 ? + s.reuse_scan_ms / (double)s.reuse_scan_calls : 0.0; + const double reuse_scan_entries_avg = + s.reuse_scan_calls != 0 ? + (double)s.reuse_scan_entries / (double)s.reuse_scan_calls : 0.0; + const double readahead_avg = + s.readahead_calls != 0 ? + s.readahead_ms / (double)s.readahead_calls : 0.0; + const double split_resident_experts_avg = + split_layers != 0.0 ? + (double)s.split_resident_experts / split_layers : 0.0; + const double split_missing_experts_avg = + split_layers != 0.0 ? + (double)s.split_missing_experts / split_layers : 0.0; + const uint64_t cache_layers = + s.cache_all_resident_layers + + s.cache_all_missing_layers + + s.cache_mixed_layers; + const double cache_layer_count = (double)cache_layers; + const double cache_resident_experts_avg = + cache_layer_count != 0.0 ? + (double)s.cache_resident_experts / cache_layer_count : 0.0; + const double cache_missing_experts_avg = + cache_layer_count != 0.0 ? + (double)s.cache_missing_experts / cache_layer_count : 0.0; + fprintf(stderr, + "ds4: streaming expert timing %s selected_calls=%llu read_avg=%.3f ms sync_avg=%.3f ms copy_avg=%.3f ms bind_avg=%.3f ms read_total=%.3f ms sync_total=%.3f ms copy_total=%.3f ms bind_total=%.3f ms split_layers=%llu resident_experts_avg=%.2f missing_experts_avg=%.2f resident_submit_avg=%.3f ms missing_bind_avg=%.3f ms resident_submit_total=%.3f ms missing_bind_total=%.3f ms missing_load_avg=%.3f ms missing_slot_avg=%.3f ms missing_prune_avg=%.3f ms missing_addr_avg=%.3f ms missing_wait_avg=%.3f ms missing_wait_total=%.3f ms load_calls=%llu load_prepare_avg=%.3f ms load_pread_avg=%.3f ms load_modify_avg=%.3f ms load_install_avg=%.3f ms prepare_batch_reuse_calls=%llu prepare_batch_reuse_avg=%.3f ms prepare_batch_reuse_total=%.3f ms prepare_buffer_calls=%llu prepare_buffer_avg=%.3f ms prepare_buffer_total=%.3f ms prepare_task_experts=%llu prepare_task_avg=%.3f ms prepare_task_total=%.3f ms reuse_scan_calls=%llu reuse_scan_entries_avg=%.1f reuse_scan_avg=%.3f ms reuse_scan_total=%.3f ms reuse_clear_total=%.3f ms readahead_calls=%llu readahead_avg=%.3f ms readahead_total=%.3f ms readahead_gib=%.2f cache_all_resident=%llu cache_all_missing=%llu cache_mixed=%llu cache_resident_avg=%.2f cache_missing_avg=%.2f\n", + scope ? scope : "total", + (unsigned long long)s.selected_calls, + selected_read_avg, + selected_sync_avg, + selected_copy_avg, + selected_bind_avg, + s.selected_read_ms, + s.selected_sync_ms, + s.selected_copy_ms, + s.selected_bind_ms, + (unsigned long long)s.split_layers, + split_resident_experts_avg, + split_missing_experts_avg, + split_resident_avg, + split_missing_avg, + s.split_resident_ms, + s.split_missing_ms, + split_missing_load_avg, + split_missing_slot_avg, + split_missing_prune_avg, + split_missing_addr_avg, + split_missing_wait_avg, + s.split_missing_wait_ms, + (unsigned long long)s.load_calls, + load_prepare_avg, + load_pread_avg, + load_modify_avg, + load_install_avg, + (unsigned long long)s.prepare_batch_reuse_calls, + prepare_batch_reuse_avg, + s.prepare_batch_reuse_ms, + (unsigned long long)s.prepare_buffer_calls, + prepare_buffer_avg, + s.prepare_buffer_ms, + (unsigned long long)s.prepare_task_experts, + prepare_task_avg, + s.prepare_task_ms, + (unsigned long long)s.reuse_scan_calls, + reuse_scan_entries_avg, + reuse_scan_avg, + s.reuse_scan_ms, + s.reuse_clear_ms, + (unsigned long long)s.readahead_calls, + readahead_avg, + s.readahead_ms, + ds4_gpu_gib(s.readahead_bytes), + (unsigned long long)s.cache_all_resident_layers, + (unsigned long long)s.cache_all_missing_layers, + (unsigned long long)s.cache_mixed_layers, + cache_resident_experts_avg, + cache_missing_experts_avg); +} + +static void ds4_gpu_print_task_memory_report(void) { + task_vm_info_data_t info; + mach_msg_type_number_t count = TASK_VM_INFO_COUNT; + const kern_return_t kr = task_info(mach_task_self(), + TASK_VM_INFO, + (task_info_t)&info, + &count); + if (kr != KERN_SUCCESS) return; + + fprintf(stderr, + "ds4: macOS task memory footprint %.2f GiB, resident %.2f GiB, virtual %.2f GiB\n", + ds4_gpu_gib((uint64_t)info.phys_footprint), + ds4_gpu_gib((uint64_t)info.resident_size), + ds4_gpu_gib((uint64_t)info.virtual_size)); +} + +void ds4_gpu_print_memory_report(const char *label) { + uint64_t cached_prefill_mask_bytes = 0; + uint64_t cached_prefill_blk_bytes = 0; + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes; + if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes; + } + const uint64_t scratch = + (uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes + + (uint64_t)g_flash_attn_pad_bytes + + (uint64_t)g_flash_attn_tmp_bytes + + (uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes + + (uint64_t)g_flash_attn_ring_bytes + + (uint64_t)g_flash_attn_kv_bytes + + (uint64_t)g_glm_flash_attn_mask_bytes + + (uint64_t)g_compressor_pool_kv_bytes + + (uint64_t)g_compressor_pool_score_bytes + + (uint64_t)g_compressor_pool_score_cont_bytes + + (uint64_t)g_compressor_pool_softmax_bytes + + (uint64_t)g_compressor_pool_product_bytes + + (uint64_t)g_compressor_store_ape_bytes + + (uint64_t)g_compressor_store_score_bytes + + (uint64_t)g_embed_rows_bytes + + (uint64_t)g_router_selection_bytes + + (uint64_t)g_router_weight_sum_bytes + + (uint64_t)g_indexer_head_scores_bytes + + (uint64_t)g_indexer_topk_bytes + + (uint64_t)g_indexed_topk_bytes + + (uint64_t)g_f16_round_scratch_bytes + + (uint64_t)g_raw_store_round_bytes + + (uint64_t)g_moe_gate_scratch_bytes + + (uint64_t)g_moe_down_scratch_bytes + + (uint64_t)g_moe_id_map_bytes + + (uint64_t)g_moe_q4_gate_slots_bytes + + (uint64_t)g_moe_q4_up_slots_bytes + + (uint64_t)g_moe_q4_down_slots_bytes; + + pthread_mutex_lock(&g_tensor_mu); + const uint64_t tensor_live_snap = g_tensor_alloc_live_bytes; + const uint64_t tensor_peak_snap = g_tensor_alloc_peak_bytes; + pthread_mutex_unlock(&g_tensor_mu); + + uint64_t tracked_live = tensor_live_snap; + if (tracked_live > UINT64_MAX - g_stream_expert_cache_bytes) { + tracked_live = UINT64_MAX; + } else { + tracked_live += g_stream_expert_cache_bytes; + } + + const bool color = ds4_log_is_tty(stderr); + const char *green = color ? "\x1b[32m" : ""; + const char *bright_green = color ? "\x1b[1;32m" : ""; + const char *reset = color ? "\x1b[0m" : ""; + fprintf(stderr, + "%sds4: Metal memory%s%s: runtime %.2f GiB + streaming experts %.2f GiB = %s%.2f GiB tracked live%s\n", + green, + label && label[0] ? " " : "", + label && label[0] ? label : "", + ds4_gpu_gib(tensor_live_snap), + ds4_gpu_gib(g_stream_expert_cache_bytes), + bright_green, + ds4_gpu_gib(tracked_live), + reset); + if (color) fputs(green, stderr); + fprintf(stderr, + "ds4: runtime tensors live %.2f MiB peak %.2f MiB\n", + ds4_gpu_mib(tensor_live_snap), + ds4_gpu_mib(tensor_peak_snap)); + ds4_gpu_print_task_memory_report(); + fprintf(stderr, + "ds4: mmap model wrapper spans %llu buffers %.2f GiB total, %.2f GiB max (not copied)\n", + (unsigned long long)g_model_wrap_count, + ds4_gpu_gib(g_model_wrap_bytes), + ds4_gpu_gib(g_model_wrap_max_bytes)); + if (g_model_buffer_cache && [g_model_buffer_cache count] != 0) { + const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); + if (limit == 0) { + fprintf(stderr, + "ds4: exact model view cache %lu buffers %.2f GiB unlimited, %llu evictions (not copied)\n", + (unsigned long)[g_model_buffer_cache count], + ds4_gpu_gib(g_model_buffer_cache_bytes), + (unsigned long long)g_model_buffer_cache_evictions); + } else { + fprintf(stderr, + "ds4: exact model view cache %lu buffers %.2f GiB / %.2f GiB, %llu evictions (not copied)\n", + (unsigned long)[g_model_buffer_cache count], + ds4_gpu_gib(g_model_buffer_cache_bytes), + ds4_gpu_gib(limit), + (unsigned long long)g_model_buffer_cache_evictions); + } + } + if (g_stream_expert_cache_hits != 0 || + g_stream_expert_cache_misses != 0 || + g_stream_expert_cache_bytes != 0) { + const uint64_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + uint64_t target_bytes = 0; + if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) { + target_bytes = + budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ? + UINT64_MAX : + budget * g_stream_expert_cache_expert_bytes; + } + const uint64_t lookups = g_stream_expert_cache_hits + g_stream_expert_cache_misses; + const double hit_rate = lookups ? + (double)g_stream_expert_cache_hits / (double)lookups : 0.0; + if (g_stream_expert_cache_evict_advise_bytes != 0 || + g_stream_expert_cache_willneed_advise_bytes != 0 || + g_stream_expert_cache_pread_bytes != 0) { + fprintf(stderr, + "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu evict_dontneed=%.2f GiB miss_willneed=%.2f GiB miss_pread=%.2f GiB pread_ms=%.3f\n", + (unsigned long long)budget, + g_stream_expert_cache_entry_count, + ds4_gpu_mib(g_stream_expert_cache_expert_bytes), + ds4_gpu_gib(target_bytes), + ds4_gpu_gib(g_stream_expert_cache_bytes), + (unsigned long long)g_stream_expert_cache_hits, + (unsigned long long)g_stream_expert_cache_misses, + hit_rate, + (unsigned long long)g_stream_expert_cache_wraps, + (unsigned long long)g_stream_expert_cache_evictions, + (unsigned long long)g_stream_expert_cache_buffer_allocs, + (unsigned long long)g_stream_expert_cache_buffer_reuses, + ds4_gpu_gib(g_stream_expert_cache_evict_advise_bytes), + ds4_gpu_gib(g_stream_expert_cache_willneed_advise_bytes), + ds4_gpu_gib(g_stream_expert_cache_pread_bytes), + g_stream_expert_cache_pread_ms); + } else { + fprintf(stderr, + "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu\n", + (unsigned long long)budget, + g_stream_expert_cache_entry_count, + ds4_gpu_mib(g_stream_expert_cache_expert_bytes), + ds4_gpu_gib(target_bytes), + ds4_gpu_gib(g_stream_expert_cache_bytes), + (unsigned long long)g_stream_expert_cache_hits, + (unsigned long long)g_stream_expert_cache_misses, + hit_rate, + (unsigned long long)g_stream_expert_cache_wraps, + (unsigned long long)g_stream_expert_cache_evictions, + (unsigned long long)g_stream_expert_cache_buffer_allocs, + (unsigned long long)g_stream_expert_cache_buffer_reuses); + } + if (ds4_gpu_stream_expert_timing_summary_enabled()) { + const ds4_gpu_stream_expert_timing_snapshot total = + ds4_gpu_stream_expert_timing_current(); + if (ds4_gpu_stream_expert_timing_has_data(total)) { + const ds4_gpu_stream_expert_timing_snapshot delta = + ds4_gpu_stream_expert_timing_delta( + total, + g_stream_expert_timing_last_report); + ds4_gpu_stream_expert_timing_print("total", total); + ds4_gpu_stream_expert_timing_print("delta", delta); + g_stream_expert_timing_last_report = total; + } + } + if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS") != NULL) { + fprintf(stderr, "ds4: streaming expert cache per-layer stats:\n"); + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + const uint64_t hits = g_stream_expert_cache_layer_hits[layer]; + const uint64_t misses = g_stream_expert_cache_layer_misses[layer]; + const uint64_t lookups = hits + misses; + const uint64_t evictions = + g_stream_expert_cache_layer_evictions[layer]; + const uint64_t pread_bytes = + g_stream_expert_cache_layer_pread_bytes[layer]; + const double pread_ms = + g_stream_expert_cache_layer_pread_ms[layer]; + const uint32_t cached = + g_stream_expert_cache_layer_count[layer]; + const uint32_t layer_slots = + ds4_gpu_stream_expert_cache_configured_count() != 0 ? + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT : 0; + if (lookups == 0 && evictions == 0 && pread_bytes == 0 && + cached == 0) { + continue; + } + const double layer_hit_rate = lookups ? + (double)hits / (double)lookups : 0.0; + fprintf(stderr, + "ds4: layer=%u layer_slots=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", + layer, + layer_slots, + cached, + (unsigned long long)hits, + (unsigned long long)misses, + layer_hit_rate, + (unsigned long long)evictions, + ds4_gpu_gib(pread_bytes), + pread_ms); + } + if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS_DELTA") != NULL) { + fprintf(stderr, "ds4: streaming expert cache per-layer delta:\n"); + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + const uint64_t hits = + g_stream_expert_cache_layer_hits[layer]; + const uint64_t misses = + g_stream_expert_cache_layer_misses[layer]; + const uint64_t evictions = + g_stream_expert_cache_layer_evictions[layer]; + const uint64_t pread_bytes = + g_stream_expert_cache_layer_pread_bytes[layer]; + const double pread_ms = + g_stream_expert_cache_layer_pread_ms[layer]; + const uint64_t delta_hits = + hits - g_stream_expert_cache_layer_last_hits[layer]; + const uint64_t delta_misses = + misses - g_stream_expert_cache_layer_last_misses[layer]; + const uint64_t delta_evictions = + evictions - + g_stream_expert_cache_layer_last_evictions[layer]; + const uint64_t delta_pread_bytes = + pread_bytes - + g_stream_expert_cache_layer_last_pread_bytes[layer]; + const double delta_pread_ms = + pread_ms - + g_stream_expert_cache_layer_last_pread_ms[layer]; + const uint64_t lookups = delta_hits + delta_misses; + if (lookups == 0 && delta_evictions == 0 && + delta_pread_bytes == 0) { + continue; + } + const double hit_rate = lookups ? + (double)delta_hits / (double)lookups : 0.0; + fprintf(stderr, + "ds4: layer=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", + layer, + g_stream_expert_cache_layer_count[layer], + (unsigned long long)delta_hits, + (unsigned long long)delta_misses, + hit_rate, + (unsigned long long)delta_evictions, + ds4_gpu_gib(delta_pread_bytes), + delta_pread_ms); + g_stream_expert_cache_layer_last_hits[layer] = hits; + g_stream_expert_cache_layer_last_misses[layer] = misses; + g_stream_expert_cache_layer_last_evictions[layer] = + evictions; + g_stream_expert_cache_layer_last_pread_bytes[layer] = + pread_bytes; + g_stream_expert_cache_layer_last_pread_ms[layer] = + pread_ms; + } + } + } + } + fprintf(stderr, + "ds4: model residency requests %llu%s\n", + (unsigned long long)g_model_residency_count, + g_ssd_streaming_mode ? " (ssd-streaming)" : + (getenv("DS4_METAL_NO_RESIDENCY") != NULL ? " (disabled)" : "")); + fprintf(stderr, + "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", + g_metal_device_name[0] ? g_metal_device_name : "(unknown)", + g_metal4_runtime_available ? "yes" : "no", + g_metal4_family_supported ? "yes" : "no", + g_metal4_queue_supported ? "yes" : "no", + g_metal4_tensor_api_enabled ? "enabled" : + (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), + g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); + fprintf(stderr, + "ds4: accelerated Metal path %s%s\n", + ds4_gpu_mpp_available() ? "enabled" : "disabled", + g_quality_mode ? " by --quality" : + (!g_metal4_tensor_api_enabled ? " (tensor API unavailable)" : "")); + fprintf(stderr, + "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", + g_metal_device_name[0] ? g_metal_device_name : "(unknown)", + g_metal4_runtime_available ? "yes" : "no", + g_metal4_family_supported ? "yes" : "no", + g_metal4_queue_supported ? "yes" : "no", + g_metal4_tensor_api_enabled ? "enabled" : + (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), + g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); + fprintf(stderr, + "ds4: scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n", + ds4_gpu_mib(scratch), + ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_glm_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes), + ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes + + (uint64_t)g_compressor_pool_score_bytes + + (uint64_t)g_compressor_pool_score_cont_bytes + + (uint64_t)g_compressor_pool_softmax_bytes + + (uint64_t)g_compressor_pool_product_bytes + + (uint64_t)g_compressor_store_ape_bytes + + (uint64_t)g_compressor_store_score_bytes + + (uint64_t)g_embed_rows_bytes), + ds4_gpu_mib((uint64_t)g_router_selection_bytes + + (uint64_t)g_router_weight_sum_bytes), + ds4_gpu_mib((uint64_t)g_indexer_head_scores_bytes + + (uint64_t)g_indexer_topk_bytes + + (uint64_t)g_indexed_topk_bytes), + ds4_gpu_mib((uint64_t)g_moe_gate_scratch_bytes + + (uint64_t)g_moe_down_scratch_bytes + + (uint64_t)g_moe_id_map_bytes + + (uint64_t)g_moe_q4_gate_slots_bytes + + (uint64_t)g_moe_q4_up_slots_bytes + + (uint64_t)g_moe_q4_down_slots_bytes), + ds4_gpu_mib((uint64_t)g_f16_round_scratch_bytes), + ds4_gpu_mib((uint64_t)g_raw_store_round_bytes)); + if (color) fputs(reset, stderr); +} + +void ds4_gpu_set_quality(bool quality) { + g_quality_mode = quality ? 1 : 0; +} + +void ds4_gpu_set_glm_model(bool enabled) { + g_glm_model_mode = enabled ? 1 : 0; +} + +void ds4_gpu_set_ssd_streaming(bool enabled) { + g_ssd_streaming_mode = enabled ? 1 : 0; + ds4_gpu_stream_expert_cache_clear_all(1); + if (g_ssd_streaming_mode) { + fprintf(stderr, + "ds4: Metal SSD streaming mode enabled; full model residency and warmup are skipped\n"); + } +} + +void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { + g_glm_streaming_prefill_full_layer_runtime = enabled ? 1 : 0; +} + +void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { + if (experts > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + experts = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; + } + g_stream_expert_cache_budget_override = experts; + ds4_gpu_stream_expert_cache_clear_all(1); +} + +void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { + /* + * Pre-seed the cache's single slab size class with the model's uniform + * per-expert bytes (first routed layer). With a mixed-precision GGUF this + * pins the class to the majority layers so the boosted ones are rejected + * deterministically from startup, instead of depending on which layer + * happens to touch the cache first. + */ + g_stream_expert_cache_expert_bytes = bytes; +} + +uint64_t ds4_gpu_recommended_working_set_size(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_device) return 0; + return (uint64_t)[g_device recommendedMaxWorkingSetSize]; +} + +static int ds4_gpu_model_map_log_enabled(void) { + if (!g_ssd_streaming_mode) return 1; + const char *trace = getenv("DS4_METAL_STREAMING_MAP_TRACE"); + return trace && trace[0] && strcmp(trace, "0") != 0; +} + +static id ds4_gpu_wrap_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset); + +static id ds4_gpu_wrap_model_exact_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset); + +static const char *ds4_gpu_source = +"#include \n" +"#ifdef DS4_METAL_HAS_TENSOR\n" +"#include \n" +"#include \n" +"#endif\n" +"using namespace metal;\n" +"#ifdef DS4_METAL_HAS_TENSOR\n" +"using namespace mpp::tensor_ops;\n" +"#endif\n" +"\n" +"#define MAX(x, y) ((x) > (y) ? (x) : (y))\n" +"#define MIN(x, y) ((x) < (y) ? (x) : (y))\n" +"#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }\n" +"#define QK8_0 32\n" +"#ifndef QK_K\n" +"#define QK_K 256\n" +"#endif\n" +"#define N_SIMDWIDTH 32\n" +"#define N_R0_Q8_0 2\n" +"#define N_SG_Q8_0 4\n" +"#define FC_MUL_MV 600\n" +"#define FC_MUL_MM 700\n" +"#define FC_BIN 1300\n" +"#define FOR_UNROLL(x) _Pragma(\"clang loop unroll(full)\") for (x)\n" +"#define M_PI_F 3.14159265358979323846f\n" +"\n" +"// Reads one byte per stride to warm model-backed pages without copying the\n" +"// model. This is outside inference and exists only to reduce first-use stalls.\n" +"kernel void kernel_touch_u8_stride(\n" +" device const uchar *src [[buffer(0)]],\n" +" device uchar *dst [[buffer(1)]],\n" +" constant ulong &stride [[buffer(2)]],\n" +" constant ulong &bytes [[buffer(3)]],\n" +" constant ulong &dst_offset [[buffer(4)]],\n" +" uint gid [[thread_position_in_grid]]) {\n" +" ulong off = (ulong)gid * stride;\n" +" if (off >= bytes) return;\n" +" dst[dst_offset + (ulong)gid] = src[off];\n" +"}\n" +"\n" +"enum ds4_sort_order {\n" +" DS4_SORT_ORDER_ASC,\n" +" DS4_SORT_ORDER_DESC,\n" +"};\n" +"\n" +"struct block_q8_0 {\n" +" half d;\n" +" int8_t qs[QK8_0];\n" +"};\n" +"\n" +"struct block_q8_K {\n" +" float d;\n" +" int8_t qs[QK_K];\n" +" int16_t bsums[QK_K / 16];\n" +"};\n" +"\n" +"\n"; + +static NSString *ds4_gpu_full_source(void) { + NSString *base = [NSString stringWithUTF8String:ds4_gpu_source]; + NSFileManager *fm = [NSFileManager defaultManager]; + /* + * Kernels are kept as separate files for review, then concatenated into one + * Metal library. Environment overrides are still honored so a diagnostic + * run can swap one source file without changing the executable. + */ + NSArray *> *required_sources = @[ + @[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"], + @[@"DS4_METAL_DENSE_SOURCE", @"metal/dense.metal"], + @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], + @[@"DS4_METAL_DSV4_HC_SOURCE", @"models/deepseek/metal/shaders/hc.metal"], + @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], + @[@"DS4_METAL_DSV4_KV_SOURCE", @"models/deepseek/metal/shaders/kv.metal"], + @[@"DS4_METAL_DSV4_ROPE_SOURCE", @"models/deepseek/metal/shaders/rope.metal"], + @[@"DS4_METAL_MODEL_ABI_SOURCE", @"metal/model_abi.metal"], + @[@"DS4_METAL_DSV4_CONTROL_SOURCE", + @"models/deepseek/metal/shaders/control.metal"], + @[@"DS4_METAL_GLM_SOURCE", @"models/glm/metal/shaders/kernels.metal"], + @[@"DS4_METAL_DSV4_ATTN_INDEXER_SOURCE", + @"models/deepseek/metal/shaders/attention_indexer.metal"], + @[@"DS4_METAL_ARGSORT_SOURCE", @"metal/argsort.metal"], + @[@"DS4_METAL_CPY_SOURCE", @"metal/cpy.metal"], + @[@"DS4_METAL_CONCAT_SOURCE", @"metal/concat.metal"], + @[@"DS4_METAL_GET_ROWS_SOURCE", @"metal/get_rows.metal"], + @[@"DS4_METAL_SUM_ROWS_SOURCE", @"metal/sum_rows.metal"], + @[@"DS4_METAL_SOFTMAX_SOURCE", @"metal/softmax.metal"], + @[@"DS4_METAL_REPEAT_SOURCE", @"metal/repeat.metal"], + @[@"DS4_METAL_GLU_SOURCE", @"metal/glu.metal"], + @[@"DS4_METAL_NORM_SOURCE", @"metal/norm.metal"], + @[@"DS4_METAL_BIN_SOURCE", @"metal/bin.metal"], + @[@"DS4_METAL_SET_ROWS_SOURCE", @"metal/set_rows.metal"], + ]; + + NSMutableString *source = [NSMutableString stringWithString:base]; + for (NSArray *spec in required_sources) { + const char *override_path = getenv([spec[0] UTF8String]); + NSMutableArray *paths = [NSMutableArray array]; + if (override_path && override_path[0]) { + [paths addObject:[NSString stringWithUTF8String:override_path]]; + } + [paths addObject:spec[1]]; + [paths addObject:[@"./" stringByAppendingString:spec[1]]]; + + NSString *loaded = nil; + NSString *loaded_path = nil; + for (NSString *path in paths) { + if (![fm fileExistsAtPath:path]) continue; + + NSError *error = nil; + loaded = [NSString stringWithContentsOfFile:path + encoding:NSUTF8StringEncoding + error:&error]; + if (!loaded) { + fprintf(stderr, "ds4: failed to read Metal source %s: %s\n", + [path UTF8String], [[error localizedDescription] UTF8String]); + return nil; + } + loaded_path = path; + break; + } + + if (!loaded) { + fprintf(stderr, + "ds4: Metal source %s not found (set %s to override)\n", + [spec[1] UTF8String], [spec[0] UTF8String]); + return nil; + } + [source appendFormat:@"\n// appended %@\n%@\n", loaded_path, loaded]; + } + return source; +} + +typedef struct { + int32_t ne00t; + int32_t ne00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_get_rows_args; + +typedef struct { + int32_t n_embd; + int32_t n_vocab; + int32_t n_tokens; + uint64_t src_row_bytes; + uint64_t dst_row_bytes; + uint64_t token_stride; +} ds4_gpu_get_rows_q8_0_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_repeat_args; + +typedef struct { + int32_t nk0; + int32_t ne01; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_set_rows_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t dim; +} ds4_gpu_concat_args; + +typedef struct { + int64_t nk0; + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_cpy_args; + +static ds4_gpu_cpy_args ds4_gpu_make_cpy_1d_args( + uint32_t n, + uint64_t src_elem, + uint64_t dst_elem) { + return (ds4_gpu_cpy_args) { + .nk0 = (int64_t)n, + .ne00 = (int64_t)n, + .ne01 = 1, + .ne02 = 1, + .ne03 = 1, + .nb00 = src_elem, + .nb01 = (uint64_t)n * src_elem, + .nb02 = (uint64_t)n * src_elem, + .nb03 = (uint64_t)n * src_elem, + .ne0 = (int64_t)n, + .ne1 = 1, + .ne2 = 1, + .ne3 = 1, + .nb0 = dst_elem, + .nb1 = (uint64_t)n * dst_elem, + .nb2 = (uint64_t)n * dst_elem, + .nb3 = (uint64_t)n * dst_elem, + }; +} + +static NSUInteger ds4_gpu_cpy_threads(uint32_t n, id pipeline) { + NSUInteger nth = 32u; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + while (nth < (NSUInteger)n && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + return nth ? nth : 1u; +} + +static float ds4_gpu_negative_infinity(void) { + union { uint32_t u; float f; } v = { 0xff800000u }; + return v.f; +} + +static float ds4_gpu_positive_infinity(void) { + union { uint32_t u; float f; } v = { 0x7f800000u }; + return v.f; +} + +static int ds4_gpu_encode_cpy_f32_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_cpy_f32_f32_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride); + +static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_col_stride, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride); + +static int ds4_gpu_encode_cpy_f32_f16_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_cpy_f32_f16_2d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint64_t src_row_stride, + uint64_t dst_row_stride); + +static int ds4_gpu_encode_cpy_f16_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_fill_f32_rows( + id cb, + id buf, + NSUInteger offset, + uint32_t width, + uint32_t rows, + float value); + +static int ds4_gpu_encode_add_f32_1d( + id cb, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off, + uint32_t n); + +typedef struct { + int32_t ne00; + uint64_t nb01; + int32_t ne10; + uint64_t nb11; + int32_t ne0; + uint64_t nb1; + int32_t i00; + int32_t i10; + float alpha; + float limit; +} ds4_gpu_glu_args; + +typedef struct { + uint32_t n; +} ds4_gpu_add_flat_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + uint64_t offs; + uint64_t o1[8]; +} ds4_gpu_bin_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float slope; + float scale; + float bias; + float val; + float min; + float max; +} ds4_gpu_unary_args; + +static ds4_gpu_bin_args ds4_gpu_make_bin_rows_args(uint32_t n, uint32_t rows, uint32_t rhs_n) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint64_t rhs_row_bytes = (uint64_t)rhs_n * sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)n, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes, + .nb03 = row_bytes, + .ne10 = (int32_t)rhs_n, + .ne11 = 1, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = rhs_row_bytes, + .nb12 = rhs_row_bytes, + .nb13 = rhs_row_bytes, + .ne0 = (int32_t)n, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = row_bytes, + .nb3 = row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static ds4_gpu_unary_args ds4_gpu_make_unary_rows_args( + uint32_t n, + uint32_t rows, + int c4, + float scale, + float bias) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint32_t n_kernel = c4 ? n / 4u : n; + return (ds4_gpu_unary_args) { + .ne00 = (int32_t)n_kernel, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes, + .nb03 = row_bytes, + .ne0 = (int32_t)n_kernel, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = row_bytes, + .nb3 = row_bytes, + .slope = 0.0f, + .scale = scale, + .bias = bias, + .val = 0.0f, + .min = 0.0f, + .max = 0.0f, + }; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_same_rows_args(uint32_t n, uint32_t rows) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)n, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = (uint64_t)rows * row_bytes, + .nb03 = (uint64_t)rows * row_bytes, + .ne10 = (int32_t)n, + .ne11 = (int32_t)rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = row_bytes, + .nb12 = (uint64_t)rows * row_bytes, + .nb13 = (uint64_t)rows * row_bytes, + .ne0 = (int32_t)n, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = (uint64_t)rows * row_bytes, + .nb3 = (uint64_t)rows * row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_bin_f32_rows( + id cb, + id pipeline, + const ds4_gpu_bin_args *args, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off); + +static int ds4_gpu_encode_sum_rows_f32( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows); + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t nr0; + int16_t r2; + int16_t r3; +} ds4_gpu_q8_0_matvec_args; + +typedef struct { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +} ds4_gpu_mul_mm_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +} ds4_gpu_mul_mv_ext_args; + +typedef ds4_gpu_q8_0_matvec_args ds4_gpu_f16_matvec_args; + +static ds4_gpu_q8_0_matvec_args ds4_gpu_make_q8_0_mv_args(uint64_t in_dim, uint64_t out_dim) { + const uint64_t row_bytes = (in_dim / 32u) * 34u; + return (ds4_gpu_q8_0_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 34, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * sizeof(float), + .nb13 = in_dim * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_f16_matvec_args ds4_gpu_make_f16_mv_args(uint64_t in_dim, uint64_t out_dim) { + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + return (ds4_gpu_f16_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = sizeof(uint16_t), + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * sizeof(float), + .nb13 = in_dim * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_q8_0_matvec_args ds4_gpu_make_f32_mv_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_vec) { + const uint64_t row_bytes = in_dim * sizeof(float); + return (ds4_gpu_q8_0_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_vec, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_vec * sizeof(float), + .nb13 = in_dim * n_vec * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_vec, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +typedef struct { + const char *function_name; + int16_t nsg; + int32_t nr0; + NSUInteger smem; +} ds4_gpu_mv_dispatch; + +static int ds4_gpu_tp_world_is_two(void); + +static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) { + const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u; + const int16_t nsg = + (int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u); + const uint64_t rows = ds4_gpu_env_u64("DS4_METAL_Q8_MV_ROWS", 2u, 2u, 4u); + if (rows >= 4u) { + return (ds4_gpu_mv_dispatch) { + .function_name = "kernel_mul_mv_q8_0_f32_r4", + .nsg = nsg, + .nr0 = 4, + .smem = 32u * 4u * sizeof(float), + }; + } + return (ds4_gpu_mv_dispatch) { + .function_name = "kernel_mul_mv_q8_0_f32", + .nsg = nsg, + .nr0 = 2, + .smem = 32u * 2u * sizeof(float), + }; +} + +static ds4_gpu_mv_dispatch ds4_gpu_make_plain_mv_dispatch( + uint64_t in_dim, + int f32_weights) { + if (in_dim < 32) { + return (ds4_gpu_mv_dispatch) { + .function_name = f32_weights ? "kernel_mul_mv_f32_f32_short" : "kernel_mul_mv_f16_f32_short", + .nsg = 1, + .nr0 = 32, + .smem = 0, + }; + } + + const int16_t nsg = (int16_t)((in_dim + 127u) / 128u > 8u ? 8u : (in_dim + 127u) / 128u); + const int use_4 = (in_dim % 4u) == 0; + return (ds4_gpu_mv_dispatch) { + .function_name = f32_weights + ? (use_4 ? "kernel_mul_mv_f32_f32_4" : "kernel_mul_mv_f32_f32") + : (use_4 ? "kernel_mul_mv_f16_f32_4" : "kernel_mul_mv_f16_f32"), + .nsg = nsg, + .nr0 = 2, + .smem = 32u * 2u * sizeof(float), + }; +} + +static ds4_gpu_mul_mm_args ds4_gpu_make_mm_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t row_bytes) { + return (ds4_gpu_mul_mm_args) { + .ne00 = (int32_t)in_dim, + .ne02 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_mul_mv_ext_args ds4_gpu_make_mv_ext_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t elem_bytes, + uint64_t row_bytes) { + return (ds4_gpu_mul_mv_ext_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = elem_bytes, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_tok, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .r2 = 1, + .r3 = 1, + }; +} + +static int16_t ds4_gpu_mv_ext_nxpsg(uint64_t in_dim, uint64_t n_tok) { + if ((in_dim % 256u) == 0 && n_tok < 3) return 16; + if ((in_dim % 128u) == 0) return 8; + return 4; +} + +static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) { + switch (n_tok) { + case 2: return 2; + case 3: + case 6: return 3; + case 4: + case 7: + case 8: return 4; + case 5: return 5; + default: return n_tok > 8 ? 4 : 0; + } +} + +static const char *ds4_gpu_mv_ext_name(int q8, int16_t r1ptg) { + if (q8) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_q8_0_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_q8_0_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_q8_0_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_q8_0_f32_r1_5"; + default: return NULL; + } + } + + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_f16_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_f16_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_f16_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_f16_f32_r1_5"; + default: return NULL; + } +} + +static const char *ds4_gpu_mv_ext_f32_name(int16_t r1ptg) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_f32_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_f32_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_f32_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_f32_f32_r1_5"; + default: return NULL; + } +} + +static const char *ds4_gpu_mv_ext_q8_pair_swiglu_name(int16_t r1ptg) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5"; + default: return NULL; + } +} + +typedef struct { + int32_t ne00; + int32_t ne00_t; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float eps; + int32_t nef1[3]; + int32_t nef2[3]; + int32_t nef3[3]; + uint64_t nbf1[3]; + uint64_t nbf2[3]; + uint64_t nbf3[3]; +} ds4_gpu_rms_norm_args; + +typedef struct { + int32_t q_n; + int32_t q_n4; + int32_t kv_n; + int32_t kv_n4; + uint64_t q_row_stride; + uint64_t kv_row_stride; + float eps; +} ds4_gpu_qkv_rms_norm_args; + +static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + return (ds4_gpu_rms_norm_args) { + .ne00 = (int32_t)n, + .ne00_t = (int32_t)(n / 4u), + .nb1 = row_bytes, + .nb2 = row_bytes * rows, + .nb3 = row_bytes * rows, + .eps = eps, + .nef1 = { (int32_t)rows, 1, 1 }, + .nef2 = { 1, 1, 1 }, + .nef3 = { 1, 1, 1 }, + .nbf1 = { row_bytes, row_bytes, row_bytes }, + .nbf2 = { row_bytes * rows, row_bytes, row_bytes }, + .nbf3 = { row_bytes * rows, row_bytes, row_bytes }, + }; +} + +static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_3d_args( + uint32_t n0, + uint32_t n1, + uint32_t n2, + float eps) { + const uint64_t row_bytes = (uint64_t)n0 * sizeof(float); + const uint64_t plane_bytes = row_bytes * n1; + return (ds4_gpu_rms_norm_args) { + .ne00 = (int32_t)n0, + .ne00_t = (int32_t)(n0 / 4u), + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = plane_bytes * n2, + .eps = eps, + .nef1 = { (int32_t)n1, 1, 1 }, + .nef2 = { (int32_t)n2, 1, 1 }, + .nef3 = { 1, 1, 1 }, + .nbf1 = { row_bytes, row_bytes, row_bytes }, + .nbf2 = { plane_bytes, row_bytes, row_bytes }, + .nbf3 = { plane_bytes * n2, row_bytes, row_bytes }, + }; +} + +static NSUInteger ds4_gpu_rms_norm_threads(uint32_t n) { + NSUInteger ne00_t = n / 4u; + NSUInteger nth = 32u; + while (nth < ne00_t && nth < 1024u) nth *= 2u; + if (nth > ne00_t) nth = ne00_t; + return nth ? nth : 1u; +} + +static NSUInteger ds4_gpu_rms_norm_pipeline_threads( + uint32_t n, + id pipeline) { + NSUInteger ne00_t = n / 4u; + NSUInteger max_threads = pipeline ? [pipeline maxTotalThreadsPerThreadgroup] : 1024u; + NSUInteger nth = 32u; + while (nth < ne00_t && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > ne00_t) nth = ne00_t; + return nth ? nth : 1u; +} + +typedef struct { + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb01; + uint64_t nb1; + float eps; +} ds4_gpu_hc_split_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_hc_weighted_sum_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float norm_eps; +} ds4_gpu_hc_weighted_sum_norm_args; + +typedef struct { + float post_scale; + float eps; +} ds4_gpu_output_hc_weights4_args; + +typedef struct { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + float eps; +} ds4_gpu_hc_split_weighted_sum_args; + +typedef struct { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float eps; + float norm_eps; +} ds4_gpu_hc_split_weighted_sum_norm_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_block0; + uint64_t nb_block1; + uint64_t nb_add0; + uint64_t nb_add1; + uint64_t nb_res0; + uint64_t nb_res1; + uint64_t nb_res2; + uint64_t nb_post0; + uint64_t nb_post1; + uint64_t nb_comb0; + uint64_t nb_comb1; + uint64_t nb_comb2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + int32_t has_add; +} ds4_gpu_hc_expand_args; + +typedef struct { + int32_t nei0; + int32_t nei1; + uint64_t nbi1; + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int32_t ne0; + int32_t ne1; + uint64_t nb1; + int32_t nr0; + /* Tensor-parallel expert ownership; see ds4_metal_args_mul_mv_id in + * metal/moe.metal. Zero (from struct literals) means no split. */ + int32_t tp_rank; + int32_t tp_world; + int32_t tp_addend; + int32_t tp_expert_base; +} ds4_gpu_mul_mv_id_args; + +typedef struct { + uint32_t n_total_expert; + uint32_t n_expert; +} ds4_gpu_stream_expert_validate_args; + +typedef struct { + uint32_t active_mask; + uint32_t accumulate; +} ds4_gpu_stream_expert_split_args; + +typedef struct { + int32_t ne02; + int32_t ne10; + int32_t ne11; + uint64_t nb11; + uint64_t nb12; + int32_t ne21; + int32_t ne20; + uint64_t nb21; +} ds4_gpu_mul_mm_id_map_args; + +typedef struct { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne20; + int32_t ne21; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; +} ds4_gpu_mul_mm_id_args; + +static int ds4_gpu_encode_mul_mv_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0); + +static int ds4_gpu_encode_attn_out_low_q8_direct( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0); + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); + +static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( + uint32_t src0_cols, + uint32_t src0_experts, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens); + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens); +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t src1_elem_size); + +static int ds4_gpu_encode_mul_mm_id( + id cb, + id map_pipeline, + id mm_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off); + +static int ds4_gpu_encode_mul_mm_id_map( + id cb, + id map_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id ids, + NSUInteger ids_off); + +static int ds4_gpu_encode_mul_mm_id_mapped( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); +static int ds4_gpu_encode_mul_mm_id_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes); +static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0_addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count, + uint32_t resource_kind, + id overflow_resource); + +typedef struct { + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +} ds4_gpu_flash_attn_pad_args; + +typedef struct { + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_raw; + uint32_t n_comp; + uint32_t pad_rows; + uint32_t shared_pad; +} ds4_gpu_flash_kv_stage_f16_args; + +typedef struct { + int32_t ne01; + int32_t ne30; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +} ds4_gpu_flash_attn_blk_args; + +typedef struct { + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + int32_t ns10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ns20; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; + int32_t ne1; + int32_t ne2; + int32_t ne3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; + float logit_softcap; +} ds4_gpu_flash_attn_vec_args; + +typedef struct { + int32_t nrows; +} ds4_gpu_flash_attn_reduce_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t n_dims; + int32_t mode; + int32_t n_ctx_orig; + int32_t inverse; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + bool src2; +} ds4_gpu_rope_tail_batch_args; + +typedef struct { + uint64_t row_bytes; + uint64_t token_bytes; + int32_t head_dim; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t inverse; + uint32_t pos0; + uint32_t pos_step; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_rope_affine_pair_args; + +_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64, + "Metal affine RoPE argument ABI changed"); + +static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args( + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + 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) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t tok_bytes = (uint64_t)n_head * row_bytes; + return (ds4_gpu_rope_tail_batch_args) { + .ne00 = head_dim, + .ne01 = n_head, + .ne02 = n_tok, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = tok_bytes, + .nb03 = (uint64_t)n_tok * tok_bytes, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = tok_bytes, + .nb3 = (uint64_t)n_tok * tok_bytes, + .n_dims = (int32_t)n_rot, + .mode = 0, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = inverse ? 1 : 0, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .src2 = false, + }; +} + +static int ds4_gpu_encode_rope_tail_inplace( + id cb, + id xbuf, + NSUInteger xoff, + const ds4_gpu_rope_tail_batch_args *args, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t pos0, + uint32_t pos_step) { + const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u; + const bool lane_compatible = + tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u; + const bool force_affine_position = + getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL; + const bool use_inplace_pair = + g_rope_tail_inplace_pair_pipeline != nil && + args->mode == 0 && !args->src2 && + lane_compatible && + (ds4_gpu_device_name_contains("M3") || + (ds4_gpu_device_name_contains("M5") && n_tok == 1u) || + getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL || + force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL; + const bool use_shared_coeff = + use_inplace_pair && !force_affine_position && + g_rope_tail_inplace_pair_shared4_pipeline != nil && + /* The 256-thread grouped schedule helps long prefill, but reduces the + * per-head parallelism that short batches and decode rely on. */ + tail_threads == 64u && n_head >= 4u && n_tok >= 32u && + getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL; + if (force_affine_position && + g_rope_tail_inplace_pair_affine_pipeline == nil && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) { + fprintf(stderr, + "ds4: forced affine-position RoPE pipeline is unavailable\n"); + return 0; + } + /* Keep long prefill on the proven shared4 kernel. Reconstructing affine + * positions inside its coefficient cohort perturbs YaRN fast-math codegen; + * the compact affine specialization is exact for the decode pair schedule. */ + const bool use_affine_position = + use_inplace_pair && !use_shared_coeff && + g_rope_tail_inplace_pair_affine_pipeline != nil && + (n_tok == 1u || force_affine_position) && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL; + + int32_t pos_stack[256]; + int32_t *pos = NULL; + id posbuf = nil; + const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t); + if (!use_affine_position) { + pos = pos_stack; + if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { + pos = malloc((size_t)n_tok * sizeof(*pos)); + if (!pos) { + fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tok; t++) { + pos[t] = (int32_t)(pos0 + t * pos_step); + } + + if (pos_bytes > 4096u) { + /* + * Metal inline setBytes data is meant for small constants. Long + * prefill RoPE calls need thousands of positions; passing that much + * inline can make the Apple driver abort the process. + */ + posbuf = ds4_gpu_new_transient_buffer( + pos_bytes, "ds4_rope_positions"); + if (!posbuf) { + if (pos != pos_stack) free(pos); + return 0; + } + memcpy([posbuf contents], pos, pos_bytes); + } + } + + ds4_gpu_rope_affine_pair_args affine_args; + if (use_affine_position) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + affine_args = (ds4_gpu_rope_affine_pair_args) { + .row_bytes = row_bytes, + .token_bytes = (uint64_t)n_head * row_bytes, + .head_dim = (int32_t)head_dim, + .n_dims = args->n_dims, + .n_ctx_orig = args->n_ctx_orig, + .inverse = args->inverse, + .pos0 = pos0, + .pos_step = pos_step, + .freq_base = args->freq_base, + .freq_scale = args->freq_scale, + .ext_factor = args->ext_factor, + .attn_factor = args->attn_factor, + .beta_fast = args->beta_fast, + .beta_slow = args->beta_slow, + }; + } + const NSUInteger reference_nth = + (NSUInteger)(head_dim < 256u ? head_dim : 256u); + const NSUInteger pair_nth = + (NSUInteger)(tail_threads < 256u ? tail_threads : 256u); + const NSUInteger nth = use_shared_coeff ? 256u : + (use_inplace_pair ? pair_nth : reference_nth); + const NSUInteger head_groups = use_shared_coeff ? + (NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head; + id enc = ds4_gpu_compute_encoder(cb); + id pipeline = use_affine_position ? + g_rope_tail_inplace_pair_affine_pipeline : + (use_shared_coeff ? + g_rope_tail_inplace_pair_shared4_pipeline : + (use_inplace_pair ? + g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline)); + [enc setComputePipelineState:pipeline]; + if (use_affine_position) { + [enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0]; + } else { + [enc setBytes:args length:sizeof(*args) atIndex:0]; + } + [enc setBuffer:xbuf offset:xoff atIndex:1]; + if (!use_affine_position) { + if (posbuf) { + [enc setBuffer:posbuf offset:0 atIndex:2]; + } else { + [enc setBytes:pos length:pos_bytes atIndex:2]; + } + [enc setBuffer:xbuf offset:xoff atIndex:3]; + } + [enc setBuffer:xbuf offset:xoff atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (pos && pos != pos_stack) free(pos); + return 1; +} + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t n_rot; +} ds4_gpu_dsv4_fp8_kv_quantize_args; + +typedef struct { + int32_t head_dim; + int32_t n_rot; + int32_t raw_row; +} ds4_gpu_dsv4_kv_fp8_store_args; + +typedef struct { + uint32_t n_rows; + uint32_t head_dim; + uint64_t row_stride; +} ds4_gpu_dsv4_indexer_qat_args; + +typedef struct { + uint32_t width; +} ds4_gpu_dsv4_ratio4_shift_args; + +typedef struct { + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t n_threads; +} ds4_gpu_dsv4_compressor_pack_ratio4_args; + +typedef struct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args; + +typedef struct { + uint32_t width; + uint32_t ratio; + uint32_t pos; + uint32_t ape_type; +} ds4_gpu_dsv4_compressor_store_one_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_dsv4_softmax_pool_args; + +typedef struct { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +} ds4_gpu_dsv4_compressor_score_ape_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; +} ds4_gpu_kargs_argsort; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; + int32_t len; +} ds4_gpu_kargs_argsort_merge; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_kargs_sum_rows; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; +} ds4_gpu_softmax_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + uint64_t nb00; + uint64_t nb01; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_dsv4_topk_mask_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int64_t ne10; + int64_t ne11; + uint64_t nb10; + uint64_t nb11; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; + float scale; +} ds4_gpu_dsv4_indexer_weighted_sum_args; + +typedef struct { + uint32_t has_bias; + uint32_t hash_mode; + uint32_t use_token_buffer; + uint32_t token; + uint32_t hash_rows; +} ds4_gpu_dsv4_router_select_one_args; + +typedef struct { + uint32_t n_expert; + uint32_t n_expert_used; + float expert_weight_scale; + uint32_t pad0; +} ds4_gpu_glm_router_select_one_args; + +typedef struct { + uint32_t n_tokens; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + float eps; +} ds4_gpu_glm_kv_lora_rms_norm_args; + +typedef struct { + uint32_t n_tokens; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t n_head; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +} ds4_gpu_glm_k_b_project_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_rope; + uint32_t cache_f16; + uint32_t pad1; +} ds4_gpu_glm_store_compact_kv_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t q_n; + uint32_t q_n4; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t kv_lora_n4; + uint32_t qk_rope; + uint32_t cache_f16; + float eps; + uint32_t pad0; +} ds4_gpu_glm_qkv_norm_store_compact_kv_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + float pad1; +} ds4_gpu_glm_store_indexer_k_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_glm_build_kv_cache_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_len; + uint32_t cache_cap; + uint32_t n_head; + uint32_t qk_dim; + uint32_t value_dim; + uint32_t pad0; + uint32_t cache_f16; + uint32_t pad1; + uint32_t pad2; + float scale; +} ds4_gpu_glm_attention_full_args; + +typedef struct { + uint32_t n_selected; +} ds4_gpu_glm_fill_selected_range_args; + +typedef struct { + uint32_t n_tokens; + uint32_t pos0; + uint32_t n_selected; + uint32_t pad_row; +} ds4_gpu_glm_fill_selected_range_batch_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t rot_offset; + uint32_t pos0; + uint32_t n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_glm_rope_tail_args; + +typedef struct { + uint32_t n_rows; + uint32_t n_head; + uint32_t head_dim; + uint32_t cache_f16; + float scale; +} ds4_gpu_glm_indexer_score_one_args; + +typedef struct { + uint32_t n_rows; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t cache_f16; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t score_token_stride; + float scale; +} ds4_gpu_glm_indexer_scores_batch_args; + +typedef struct { + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +} ds4_gpu_glm_qk_lowrank_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t head_base; +} ds4_gpu_glm_qk_lowrank_batch_args; + +typedef struct { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +} ds4_gpu_glm_attention_indexed_decode_args; + +typedef struct { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + uint32_t block_rows; + uint32_t n_blocks; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +} ds4_gpu_glm_attention_indexed_decode_split_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_selected; + uint32_t cache_cap; + uint32_t 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; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t pos0; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t head_base; +} ds4_gpu_glm_attention_indexed_batch_args; + +typedef struct { + uint32_t in_dim; + uint32_t mid_dim; + uint32_t out_dim; + uint32_t n_total_expert; + uint32_t n_expert_used; + uint32_t n_tokens; + uint32_t mid_token_stride; + uint32_t down_type; + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; + uint64_t gate_expert_bytes; + uint64_t gate_row_bytes; + uint64_t up_expert_bytes; + uint64_t up_row_bytes; + uint64_t down_expert_bytes; + uint64_t down_row_bytes; +} ds4_gpu_glm_routed_moe_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; + uint32_t top_k; + uint32_t pos0; + uint32_t window; + uint32_t ratio; + uint32_t comp_kv_f16; + uint32_t pad0; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t raw_row_stride; + uint64_t comp_row_stride; + uint64_t topk_token_stride; + uint64_t dst_token_stride; + uint64_t dst_head_stride; + float scale; +} ds4_gpu_dsv4_indexed_attention_args; + +typedef struct { + uint32_t n_comp; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t ratio; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t index_row_stride; + uint64_t score_token_stride; + float scale; +} ds4_gpu_dsv4_indexer_scores_fused_args; + +typedef struct { + uint32_t width; + uint32_t rows; + uint64_t gate_row_stride; + uint64_t up_row_stride; + uint64_t mid_row_stride; + uint64_t weight_stride; + uint32_t write_clamped; + float clamp_value; +} ds4_gpu_dsv4_moe_swiglu_weight_args; + +typedef struct { + uint32_t expert_base; + uint32_t expert_count; + uint32_t accumulate; + uint32_t pad0; +} ds4_gpu_moe_expert_group_args; + +typedef struct { + uint64_t expert_bytes; + uint32_t group_size; + uint32_t n_slots; +} ds4_gpu_q4_gather_slots6_args; + +typedef struct { + uint32_t width; + uint32_t tokens; + uint64_t src_token_stride; + uint64_t dst_token_stride; +} ds4_gpu_dsv4_moe_sum6_args; + +/* Compile the single in-repo Metal source and create the pipelines that every + * session uses. Shape-dependent kernels with function constants are built + * lazily by the small ds4_gpu_get_* caches, so startup stays predictable + * while long-context prefill and decode can still pick specialized variants. */ +int ds4_gpu_init(void) { + if (g_initialized) return 1; + + @autoreleasepool { + g_device = MTLCreateSystemDefaultDevice(); + if (!g_device) { + fprintf(stderr, "ds4: Metal device not available\n"); + return 0; + } + ds4_gpu_print_device_summary(); + ds4_gpu_detect_metal4_features(); + + g_queue = [g_device newCommandQueue]; + if (!g_queue) { + fprintf(stderr, "ds4: failed to create Metal command queue\n"); + g_device = nil; + return 0; + } + g_model_buffer_cache = [NSMutableDictionary dictionary]; + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_evictions = 0; + g_model_buffer_cache_over_limit = 0; + g_q4_expert_table_cache = [NSMutableDictionary dictionary]; + g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary]; + g_pipeline_cache = [NSMutableDictionary dictionary]; + g_transient_buffers = [NSMutableArray array]; + g_pending_cbs = [NSMutableArray array]; + if (!g_model_buffer_cache || !g_q4_expert_table_cache || + !g_q4_expert_layer_residency_cache || + !g_pipeline_cache || !g_transient_buffers || !g_pending_cbs) { + fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n"); + g_pending_cbs = nil; + g_transient_buffers = nil; + g_pipeline_cache = nil; + g_q4_expert_layer_residency_cache = nil; + g_q4_expert_table_cache = nil; + g_model_buffer_cache = nil; + g_queue = nil; + g_device = nil; + return 0; + } + + NSError *error = nil; + NSString *source = ds4_gpu_full_source(); + if (!source) { + g_queue = nil; + g_device = nil; + return 0; + } + MTLCompileOptions *options = [MTLCompileOptions new]; + NSMutableDictionary *macros = [NSMutableDictionary new]; + if (g_metal4_tensor_api_enabled) { + macros[@"DS4_METAL_HAS_TENSOR"] = @"1"; + fprintf(stderr, "ds4: Metal 4 tensor API enabled for Tensor kernels\n"); + } + + const int drift_hc_stable = ds4_gpu_env_bool("DS4_METAL_HC_STABLE") != 0; // default ON + const int drift_norm_unify = ds4_gpu_env_bool("DS4_METAL_NORM_RSQRT_DISABLE") != 0; // default ON + const int drift_kv_raw_f32 = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32") > 0; // default OFF + const int drift_rope_exp2_log2 = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2") > 0; // default OFF + const int drift_math_safe = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE") > 0; // default OFF + + if (drift_math_safe) { + // MTLCompileOptions.fastMathEnabled defaults to YES and Apple's + // headers explicitly say this "may violate the IEEE 754 standard". + // Different fast-math optimizations get applied across the + // matmul2d cooperative-tensor path and the legacy + // simdgroup_multiply_accumulate path on M5, amplifying the + // mismatch. MTLMathModeSafe pins the entire library to strict + // IEEE-754 semantics. Diagnostic-only: useful to localize drift + // sources but not to ship as a default. + if (@available(macOS 15.0, *)) { + options.mathMode = MTLMathModeSafe; + fprintf(stderr, "ds4: Metal shader library math mode = safe (strict IEEE-754) by DS4_METAL_MATH_SAFE\n"); + } else { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" + options.fastMathEnabled = NO; +#pragma clang diagnostic pop + fprintf(stderr, "ds4: Metal shader library fast-math disabled by DS4_METAL_MATH_SAFE (pre-macOS 15)\n"); + } + } + + if (drift_hc_stable) macros[@"DS4_METAL_HC_STABLE"] = @"1"; + if (drift_norm_unify) macros[@"DS4_METAL_NORM_RSQRT_DISABLE"] = @"1"; + if (drift_kv_raw_f32) macros[@"DS4_METAL_KV_RAW_F32"] = @"1"; + if (drift_rope_exp2_log2) macros[@"DS4_METAL_ROPE_EXP2_LOG2"] = @"1"; + fprintf(stderr, + "ds4: drift-patch flags hc_stable=%s norm_unify=%s kv_raw_f32=%s rope_exp2_log2=%s math_safe=%s tensor_matmul=%s\n", + drift_hc_stable ? "on" : "off", + drift_norm_unify ? "on" : "off", + drift_kv_raw_f32 ? "on" : "off", + drift_rope_exp2_log2 ? "on" : "off", + drift_math_safe ? "on" : "off", + g_metal4_tensor_api_enabled ? "on" : "off"); + options.preprocessorMacros = macros; + id library = [g_device newLibraryWithSource:source options:options error:&error]; + if (!library) { + fprintf(stderr, "ds4: Metal shader compilation failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_library = library; + + id fn = [library newFunctionWithName:@"kernel_get_rows_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_i32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_i32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_i32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_i32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q8_0_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q8_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q8_0_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q4_0_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q4_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q4_0_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q4_K_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q4_K_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q4_K_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_repeat_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_repeat_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_repeat_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_repeat_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_repeat_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_set_rows_f32_i32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_set_rows_f32_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_set_rows_f32_i32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_concat"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_concat function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_concat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_concat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_concat pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f32_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f32_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f32_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f32_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f32_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f16_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f16_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f16_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"]; + if (fn) { + g_flash_kv_stage_f16_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_flash_kv_stage_f16_pipeline) { + fprintf(stderr, + "ds4: optional Metal gathered KV staging pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal gathered KV staging kernel unavailable\n"); + } + + fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_fp8_kv_quantize_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_fp8_kv_quantize_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_indexer_hadamard_fp4_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_indexer_qat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_indexer_qat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_kv_fp8_store_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_kv_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_kv_fp8_store_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_ratio4_shift_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_ratio4_shift_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_swiglu_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_swiglu_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_swiglu_flat_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_swiglu_flat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_swiglu_flat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum6_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_moe_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum8_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_moe_sum8_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_sum8_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_op = 0; + int16_t bin_f = 1; + bool bin_rb = false; + bool bin_cb = false; + [bin_constants setConstantValue:&bin_op type:MTLDataTypeShort atIndex:1300]; + [bin_constants setConstantValue:&bin_f type:MTLDataTypeShort atIndex:1301]; + [bin_constants setConstantValue:&bin_rb type:MTLDataTypeBool atIndex:1302]; + [bin_constants setConstantValue:&bin_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_add_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add2_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add2_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_add2_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add2_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add2_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add3_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add3_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_add3_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add3_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add3_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_mul_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_mul_plain_op = 2; + int16_t bin_mul_plain_f = 1; + bool bin_mul_plain_rb = false; + bool bin_mul_plain_cb = false; + [bin_mul_constants setConstantValue:&bin_mul_plain_op type:MTLDataTypeShort atIndex:1300]; + [bin_mul_constants setConstantValue:&bin_mul_plain_f type:MTLDataTypeShort atIndex:1301]; + [bin_mul_constants setConstantValue:&bin_mul_plain_rb type:MTLDataTypeBool atIndex:1302]; + [bin_mul_constants setConstantValue:&bin_mul_plain_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_mul_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_mul_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_mul_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_mul_scalar_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_mul_op = 2; + int16_t bin_mul_f = 1; + bool bin_mul_rb = false; + bool bin_mul_cb = true; + [bin_mul_scalar_constants setConstantValue:&bin_mul_op type:MTLDataTypeShort atIndex:1300]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_f type:MTLDataTypeShort atIndex:1301]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_rb type:MTLDataTypeBool atIndex:1302]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_mul_scalar_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_bin_mul_scalar_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_bin_mul_scalar_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_div_row_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_div_op = 3; + int16_t bin_div_f = 1; + bool bin_div_rb = false; + bool bin_div_cb = true; + [bin_div_row_constants setConstantValue:&bin_div_op type:MTLDataTypeShort atIndex:1300]; + [bin_div_row_constants setConstantValue:&bin_div_f type:MTLDataTypeShort atIndex:1301]; + [bin_div_row_constants setConstantValue:&bin_div_rb type:MTLDataTypeBool atIndex:1302]; + [bin_div_row_constants setConstantValue:&bin_div_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_div_row_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_bin_div_row_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_bin_div_row_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_rms_norm_mul_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_rms_norm_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rms_norm_plain_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rms_norm_plain_pipeline) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add_rms_norm_mul_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_add_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_qkv_rms_norm_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_qkv_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_qkv_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t moe_mv_id_nsg = 2; + [moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q2_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q2_k_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_masked_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_stream_expert_cache_validate"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_stream_expert_cache_validate_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_stream_expert_cache_validate_pipeline) { + fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pair_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group6_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group8_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group8_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_id_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group24_q4_k_id_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group24_q4_k_id_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group24_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group24_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_q4_gather_slots6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_q4_gather_slots6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_table_q4_pair_gate_encoder = [fn newArgumentEncoderWithBufferIndex:2]; + g_moe_table_q4_pair_up_encoder = [fn newArgumentEncoderWithBufferIndex:3]; + if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert-table pair argument encoder creation failed\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_table_q4_sum_down_encoder = [fn newArgumentEncoderWithBufferIndex:1]; + if (!g_moe_table_q4_sum_down_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert-table down argument encoder creation failed\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_table_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_table_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (fn) { + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 function unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (fn) { + g_moe_mul_mv_addr_q4_k_sum6_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 function unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rope_tail_batch_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_batch_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rope_tail_inplace_pair_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_pipeline) { + fprintf(stderr, + "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"]; + if (!fn) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n"); + } else { + g_rope_tail_inplace_pair_shared4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_shared4_pipeline) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n", + [[error localizedDescription] UTF8String]); + } + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"]; + if (fn) { + g_rope_tail_inplace_pair_affine_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_affine_pipeline) { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair kernel unavailable\n"); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_softmax_pool_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_softmax_pool_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_soft_max_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_soft_max_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_soft_max_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_soft_max_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_soft_max_f32_4_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_soft_max_f32_4_pipeline) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_argsort_f32_i32_desc"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_argsort_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_argsort_f32_i32_desc_pipeline) { + fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_argsort_merge_f32_i32_desc"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_argsort_merge_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_argsort_merge_f32_i32_desc_pipeline) { + fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *sum_rows_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t sum_rows_op = 10; + [sum_rows_constants setConstantValue:&sum_rows_op type:MTLDataTypeShort atIndex:1400]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_sum_rows_f32_f32" + constantValues:sum_rows_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_sum_rows_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_sum_rows_f32_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_topk_mask_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_topk_mask_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask_scatter"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_topk_mask_scatter_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_topk_mask_scatter_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_indexer_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_indexer_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_indexer_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_sinkhorn"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_sinkhorn_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_sinkhorn_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum_norm4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_weighted_sum_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_weighted_sum_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"]; + if (fn) { + g_hc_weighted_sum_norm_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_weighted_sum_norm_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n"); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"]; + if (fn) { + g_output_hc_weights4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_output_hc_weights4_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC weights4 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC weights4 kernel unavailable\n"); + } + + MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_sigmoid_op = 102; + bool unary_cnt = false; + [unary_sigmoid_constants setConstantValue:&unary_sigmoid_op type:MTLDataTypeShort atIndex:1200]; + [unary_sigmoid_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_sigmoid_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_sigmoid_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_sigmoid_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_silu_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_silu_op = 106; + [unary_silu_constants setConstantValue:&unary_silu_op type:MTLDataTypeShort atIndex:1200]; + [unary_silu_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_silu_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_silu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_silu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_softplus_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_softplus_op = 115; + [unary_softplus_constants setConstantValue:&unary_softplus_op type:MTLDataTypeShort atIndex:1200]; + [unary_softplus_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_softplus_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_softplus_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_softplus_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_sqrt_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_sqrt_op = 14; + [unary_sqrt_constants setConstantValue:&unary_sqrt_op type:MTLDataTypeShort atIndex:1200]; + [unary_sqrt_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_sqrt_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_sqrt_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_sqrt_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_clamp_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_clamp_op = 12; + [unary_clamp_constants setConstantValue:&unary_clamp_op type:MTLDataTypeShort atIndex:1200]; + [unary_clamp_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32" + constantValues:unary_clamp_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_clamp_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_clamp_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_scale_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_scale_op = 10; + [unary_scale_constants setConstantValue:&unary_scale_op type:MTLDataTypeShort atIndex:1200]; + [unary_scale_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_scale_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_scale_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_scale_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_fill_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_fill_op = 11; + [unary_fill_constants setConstantValue:&unary_fill_op type:MTLDataTypeShort atIndex:1200]; + [unary_fill_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_fill_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_fill_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_fill_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f16_f16" + constantValues:unary_fill_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_fill_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_fill_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_expand"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_expand_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_expand_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_dsv4_indexer_score_one_direct_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct"); + g_dsv4_compressor_store_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one"); + g_dsv4_compressor_pack_ratio4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4"); + g_dsv4_softmax_pool_ratio4_direct_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct"); + g_rms_norm_scale_pipeline = + ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4"); + g_dsv4_sort_i32_rows_asc_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc"); + g_dsv4_indexed_attention_heads8_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8"); + g_dsv4_indexed_attention_heads8_rb16_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16"); + g_dsv4_softplus_sqrt_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4"); + g_dsv4_router_finalize_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one"); + g_dsv4_router_finalize_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd"); + g_dsv4_router_finalize_weights_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd"); + g_dsv4_router_weights_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one"); + g_glm_router_select_one_pipeline = + ds4_gpu_get_pipeline("kernel_glm_router_select_one"); + g_glm_kv_lora_rms_norm_pipeline = + ds4_gpu_get_pipeline("kernel_glm_kv_lora_rms_norm"); + g_glm_k_b_project_pipeline = + ds4_gpu_get_pipeline("kernel_glm_k_b_project_q8_0"); + g_glm_store_compact_kv_pipeline = + ds4_gpu_get_pipeline("kernel_glm_store_compact_kv"); + g_glm_qkv_norm_store_compact_kv_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv"); + g_glm_store_indexer_k_pipeline = + ds4_gpu_get_pipeline("kernel_glm_store_indexer_k"); + g_glm_build_kv_cache_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache"); + g_glm_build_kv_cache_decode_group4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_decode_group4"); + g_glm_build_kv_cache_flash_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_flash"); + g_glm_attention_full_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_full"); + g_glm_fill_selected_range_pipeline = + ds4_gpu_get_pipeline("kernel_glm_fill_selected_range"); + g_glm_fill_selected_range_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch"); + g_glm_indexer_rope_tail_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32"); + g_glm_indexer_score_one_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_score_one"); + g_glm_indexer_score_one_direct_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct"); + g_glm_indexer_scores_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch"); + g_glm_indexer_scores_tiled_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled"); + g_glm_indexer_scores_tiled_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled_f32"); + g_glm_qk_lowrank_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0"); + g_glm_qk_lowrank_glm52_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52"); + g_glm_qk_lowrank_glm52_sg_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52_sg"); + g_glm_qk_lowrank_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch"); + g_glm_qk_lowrank_batch_glm52_t4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4"); + g_glm_value_project_q8_0_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0"); + g_glm_value_project_q8_0_batch_heads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads"); + g_glm_value_project_q8_0_batch_heads_mma_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads_mma"); + g_glm_attention_indexed_decode_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode"); + g_glm_attention_indexed_decode_split_group8_partial_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial"); + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads"); + g_glm_attention_indexed_decode_split_group8_reduce_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce"); + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16"); + g_glm_attention_indexed_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch"); + g_glm_attention_indexed_batch_group2_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group2"); + g_glm_attention_indexed_batch_q2_group4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_q2_group4"); + g_glm_attention_indexed_batch_group8_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group8"); + g_glm_attention_indexed_batch_lora_group8_vec_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec"); + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads"); + g_glm_q4_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu_f32"); + g_glm_q4_k_pair_swiglu2_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_f32"); + g_glm_q4_k_pair_swiglu4_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu4_f32"); + g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_f32"); + g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_row_f32"); + g_glm_q2_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_pair_swiglu_f32"); + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32"); + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32_masked"); + g_glm_q4_k_addr_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32"); + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32_masked"); + g_glm_q2_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_down_f32"); + g_glm_q4_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_down_simd_f32"); + g_glm_q2_k_addr_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_down_f32"); + g_glm_q4_k_addr_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_down_simd_f32"); + g_glm_q5_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_f32"); + g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_f32"); + g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_row_f32"); + g_glm_q5_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_down_f32"); + g_glm_q6_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q6_K_down_f32"); + g_dsv4_router_weights_batch_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch"); + g_dsv4_hc_expand4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4"); + if (!g_dsv4_indexer_score_one_direct_pipeline || + !g_dsv4_compressor_store_one_pipeline || + !g_dsv4_sort_i32_rows_asc_pipeline || + !g_dsv4_indexed_attention_heads8_pipeline || + !g_dsv4_indexed_attention_heads8_rb16_pipeline || + !g_dsv4_softplus_sqrt_pipeline || + !g_dsv4_router_finalize_one_pipeline || + !g_dsv4_router_weights_one_pipeline || + !g_glm_router_select_one_pipeline || + !g_glm_kv_lora_rms_norm_pipeline || + !g_glm_k_b_project_pipeline || + !g_glm_store_compact_kv_pipeline || + !g_glm_qkv_norm_store_compact_kv_pipeline || + !g_glm_store_indexer_k_pipeline || + !g_glm_build_kv_cache_pipeline || + !g_glm_build_kv_cache_decode_group4_pipeline || + !g_glm_build_kv_cache_flash_pipeline || + !g_glm_attention_full_pipeline || + !g_glm_fill_selected_range_pipeline || + !g_glm_fill_selected_range_batch_pipeline || + !g_glm_indexer_rope_tail_pipeline || + !g_glm_indexer_score_one_pipeline || + !g_glm_indexer_score_one_direct_pipeline || + !g_glm_indexer_scores_batch_pipeline || + !g_glm_indexer_scores_tiled_pipeline || + !g_glm_indexer_scores_tiled_f32_pipeline || + !g_glm_qk_lowrank_pipeline || + !g_glm_qk_lowrank_glm52_pipeline || + !g_glm_qk_lowrank_glm52_sg_pipeline || + !g_glm_qk_lowrank_batch_pipeline || + !g_glm_qk_lowrank_batch_glm52_t4_pipeline || + !g_glm_value_project_q8_0_pipeline || + !g_glm_value_project_q8_0_batch_heads_pipeline || + !g_glm_value_project_q8_0_batch_heads_mma_pipeline || + !g_glm_attention_indexed_decode_pipeline || + !g_glm_attention_indexed_decode_split_group8_partial_pipeline || + !g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline || + !g_glm_attention_indexed_decode_split_group8_reduce_pipeline || + !g_glm_attention_indexed_decode_split_group8_reduce16_pipeline || + !g_glm_attention_indexed_batch_pipeline || + !g_glm_attention_indexed_batch_group2_pipeline || + !g_glm_attention_indexed_batch_q2_group4_pipeline || + !g_glm_attention_indexed_batch_group8_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline || + !g_glm_q4_k_pair_swiglu_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_f32_pipeline || + !g_glm_q4_k_pair_swiglu4_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline || + !g_glm_q2_k_pair_swiglu_f32_pipeline || + !g_glm_q2_k_addr_pair_swiglu2_f32_pipeline || + !g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline || + !g_glm_q4_k_addr_pair_swiglu_f32_pipeline || + !g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline || + !g_glm_q2_k_down_f32_pipeline || + !g_glm_q4_k_down_f32_pipeline || + !g_glm_q2_k_addr_down_f32_pipeline || + !g_glm_q4_k_addr_down_f32_pipeline || + !g_glm_q5_k_pair_swiglu_f32_pipeline || + !g_glm_q5_k_pair_swiglu_mapped_f32_pipeline || + !g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline || + !g_glm_q5_k_down_f32_pipeline || + !g_glm_q6_k_down_f32_pipeline || + !g_dsv4_hc_expand4_pipeline) { + g_queue = nil; + g_device = nil; + return 0; + } + + g_initialized = 1; + } + + return 1; +} + +ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { + if (!g_initialized && !ds4_gpu_init()) return NULL; + if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; + + @autoreleasepool { + DS4MetalTensor *tensor = [DS4MetalTensor new]; + tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes + options:MTLResourceStorageModeShared]; + if (!tensor.buffer) { + return NULL; + } + tensor.offset = 0; + tensor.bytes = bytes; + tensor.owner = 1; + uint64_t live_snap = 0; + uint64_t peak_snap = 0; + pthread_mutex_lock(&g_tensor_mu); + const int tracked = ds4_gpu_tensor_track_alloc_locked( + (__bridge const void *)tensor, + bytes, + &live_snap, + &peak_snap); + pthread_mutex_unlock(&g_tensor_mu); + if (!tracked) { + fprintf(stderr, "ds4: failed to track Metal tensor allocation\n"); + tensor.buffer = nil; + return NULL; + } + if (ds4_gpu_trace_allocs()) { + fprintf(stderr, + "ds4: Metal tensor alloc %.3f MiB live %.3f MiB peak %.3f MiB\n", + (double)bytes / (1024.0 * 1024.0), + (double)live_snap / (1024.0 * 1024.0), + (double)peak_snap / (1024.0 * 1024.0)); + } + return (__bridge_retained ds4_gpu_tensor *)tensor; + } +} + +ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { + return ds4_gpu_tensor_alloc(bytes); +} + +int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) { + (void)kv_cache_bytes; + (void)context_bytes; + return 0; +} + +ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) { + if (!base) return NULL; + const DS4MetalTensor *base_obj = ds4_gpu_tensor_const_obj(base); + if (offset > base_obj.bytes || bytes > base_obj.bytes - offset) return NULL; + if (base_obj.offset > UINT64_MAX - offset) return NULL; + const uint64_t absolute_offset = base_obj.offset + offset; + if (absolute_offset > (uint64_t)NSUIntegerMax) return NULL; + + @autoreleasepool { + DS4MetalTensor *view = [DS4MetalTensor new]; + view.buffer = base_obj.buffer; + view.offset = absolute_offset; + view.bytes = bytes; + view.owner = 0; + pthread_mutex_lock(&g_tensor_mu); + const int tracked = ds4_gpu_tensor_track_view_locked((__bridge const void *)view); + pthread_mutex_unlock(&g_tensor_mu); + if (!tracked) { + fprintf(stderr, "ds4: failed to track Metal tensor view\n"); + view.buffer = nil; + return NULL; + } + return (__bridge_retained ds4_gpu_tensor *)view; + } +} + +void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) { + if (!tensor) return; + @autoreleasepool { + uint8_t owner = 0; + uint64_t bytes = 0; + uint64_t live_snap = 0; + uint64_t peak_snap = 0; + if (!ds4_gpu_tensor_prepare_free(tensor, + &owner, + &bytes, + &live_snap, + &peak_snap)) { + return; + } + DS4MetalTensor *obj = (__bridge_transfer DS4MetalTensor *)tensor; + if (owner) { + if (ds4_gpu_trace_allocs()) { + fprintf(stderr, + "ds4: Metal tensor free %.3f MiB live %.3f MiB peak %.3f MiB\n", + (double)bytes / (1024.0 * 1024.0), + (double)live_snap / (1024.0 * 1024.0), + (double)peak_snap / (1024.0 * 1024.0)); + } + } + obj.buffer = nil; + obj.offset = 0; + obj.bytes = 0; + obj.owner = 0; + } +} + +uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) { + if (!tensor) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return obj.bytes; +} + +void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) { + if (!tensor) return NULL; + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + return (uint8_t *)[obj.buffer contents] + obj.offset; +} + +int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) { + if (!tensor || count > ds4_gpu_tensor_bytes(tensor) / sizeof(float)) return 0; + float *p = ds4_gpu_tensor_contents(tensor); + if (!p && count != 0) return 0; + for (uint64_t i = 0; i < count; i++) p[i] = value; + return 1; +} + +int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) { + if (!tensor || (!data && bytes != 0)) return 0; + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; + if (bytes != 0) { + memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset, data, (size_t)bytes); + } + return 1; +} + +int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) { + if (!tensor || (!data && bytes != 0)) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; + if (bytes != 0) { + memcpy(data, (const uint8_t *)[obj.buffer contents] + obj.offset + offset, (size_t)bytes); + } + return 1; +} + +int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t bytes) { + if (!dst || !src) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); + const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); + if (dst_offset > d.bytes || bytes > d.bytes - dst_offset) return 0; + if (src_offset > s.bytes || bytes > s.bytes - src_offset) return 0; + if (bytes == 0) return 1; + if (!g_batch_cb) return 0; + + ds4_gpu_close_batch_encoder(); + g_batch_has_work = YES; + id blit = [g_batch_cb blitCommandEncoder]; + if (!blit) return 0; + [blit copyFromBuffer:s.buffer + sourceOffset:(NSUInteger)(s.offset + src_offset) + toBuffer:d.buffer + destinationOffset:(NSUInteger)(d.offset + dst_offset) + size:(NSUInteger)bytes]; + [blit endEncoding]; + return 1; +} + +int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t count) { + if (!dst || !src) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); + const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); + if (count == 0) return 1; + if (count > UINT64_MAX / sizeof(float) || + count > UINT64_MAX / sizeof(uint16_t)) { + return 0; + } + const uint64_t src_bytes = count * sizeof(float); + const uint64_t dst_bytes = count * sizeof(uint16_t); + if (src_offset > s.bytes || src_bytes > s.bytes - src_offset || + dst_offset > d.bytes || dst_bytes > d.bytes - dst_offset) { + return 0; + } + + @autoreleasepool { + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint64_t done = 0; + int ok = 1; + while (done < count && ok) { + uint64_t chunk64 = count - done; + if (chunk64 > UINT32_MAX) chunk64 = UINT32_MAX; + const uint32_t chunk = (uint32_t)chunk64; + ok = ds4_gpu_encode_cpy_f32_f16_1d( + cb, + s.buffer, + (NSUInteger)(s.offset + src_offset + done * sizeof(float)), + d.buffer, + (NSUInteger)(d.offset + dst_offset + done * sizeof(uint16_t)), + chunk); + done += chunk; + } + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "tensor f32 to f16 copy"); + return ok; + } +} + +int ds4_gpu_pack_slot_rows_f32_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t n_rows, + uint32_t width, + uint32_t n_slots, + uint32_t slot_cap) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 || + slot_cap == 0 || n_rows > slot_cap) { + return 0; + } + + @autoreleasepool { + id slotsbuf = ds4_gpu_tensor_buffer(slots); + id outbuf = ds4_gpu_tensor_buffer(out); + uint64_t row_bytes = 0; + uint64_t slot_plane_bytes = 0; + uint64_t slots_bytes = 0; + uint64_t out_rows = 0; + uint64_t out_bytes = 0; + if ((uint64_t)width > UINT64_MAX / sizeof(float)) return 0; + row_bytes = (uint64_t)width * sizeof(float); + if ((uint64_t)slot_cap > UINT64_MAX / row_bytes) return 0; + slot_plane_bytes = (uint64_t)slot_cap * row_bytes; + if ((uint64_t)n_slots > UINT64_MAX / slot_plane_bytes) return 0; + slots_bytes = (uint64_t)n_slots * slot_plane_bytes; + if ((uint64_t)n_rows > UINT64_MAX / n_slots) return 0; + out_rows = (uint64_t)n_rows * n_slots; + if (out_rows > UINT64_MAX / row_bytes) return 0; + out_bytes = out_rows * row_bytes; + if (!slotsbuf || !outbuf || + ds4_gpu_tensor_bytes(slots) < slots_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal slot-row pack received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + slotsbuf, + ds4_gpu_tensor_offset(slots), + outbuf, + ds4_gpu_tensor_offset(out), + width, + n_slots, + n_rows, + sizeof(float), + slot_plane_bytes, + row_bytes, + row_bytes, + (uint64_t)n_slots * row_bytes)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "slot-row pack")) return 0; + } + + return 1; +} + +int ds4_gpu_begin_commands(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_batch_cb) return 0; + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + return g_batch_cb != nil; +} + +int ds4_gpu_flush_encoder(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + ds4_gpu_close_batch_encoder(); + return 1; +} + +int ds4_gpu_flush_commands(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + [cb commit]; + [g_pending_cbs addObject:cb]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers("command batch"); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; +} + +int ds4_gpu_commands_active(void) { + return g_batch_cb != nil; +} + +static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) { + const char *what = label ? label : "streaming expert cache in-flight"; + if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0; + if ([g_pending_cbs count] != 0 && + ds4_gpu_wait_pending_command_buffers(what) == 0) { + return 0; + } + return 1; +} + +int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { + if (!event_value) return 0; + *event_value = 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) { + g_selected_readback_event = [g_device newSharedEvent]; + if (!g_selected_readback_event) { + fprintf(stderr, "ds4: failed to create Metal shared event for selected-id overlap\n"); + return 0; + } + } + + ds4_gpu_close_batch_encoder(); + const uint64_t value = ++g_selected_readback_event_value; + [g_batch_cb encodeSignalEvent:g_selected_readback_event value:value]; + g_batch_has_work = YES; + *event_value = value; + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb || event_value == 0) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) return 0; + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + [cb commit]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + const char *what = label ? label : "selected-id overlap"; + const BOOL signaled = + [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; + [g_pending_cbs addObject:cb]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + what, + [[cb.error localizedDescription] UTF8String]); + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +/* + * Tensor-parallel gates. + * + * A TP gate is a mid-command-stream rendezvous with the peer machine: the + * kernels ahead of the gate leave a partial block output in a slab slot, + * the GPU signals g_tp_gpu_event, and the pre-encoded combine kernel waits + * on g_tp_cpu_event. A dedicated service thread bridges the two: it spins + * until the GPU reaches the gate, runs the transport exchange (RDMA WRITE + * plus flag poll, or a TCP write/read pair — behind the callback), and + * CPU-signals the release. On exchange failure the release is signaled + * anyway so the GPU never deadlocks; the failure latches in g_tp_failed + * and the eval aborts at the next command-buffer boundary. + * + * Gate sequence values increase monotonically per encoded gate. Both ranks + * encode the identical graph, so the values agree by construction and slots + * never need resetting between tokens. + */ +typedef struct { + uint32_t layer; + uint32_t gate; + uint32_t rows; /* 0 = row gate; >0 = verify-block batch gate */ + uint32_t event_arrival; + uint64_t seq; + /* Big batch gates (prefill): exchange big_bytes from big_out into + * big_in directly (CPU-visible bounce buffers), bypassing the slab. */ + const void *big_out; + void *big_in; + uint64_t big_bytes; +} ds4_gpu_tp_request; + +enum { DS4_GPU_TP_QUEUE = 1024 }; + +static id g_tp_gpu_event; /* GPU -> service thread */ +static id g_tp_cpu_event; /* service thread -> GPU */ +/* Batch (verify-block) gates run on their own sequence space and release + * event: the row-gate seq feeds the RDMA pre-posted recv accounting, which + * requires consecutive values, and a shared release event would make a + * small batch value satisfy waits armed against the larger row seq. */ +static id g_tp_batch_gpu_event; +static id g_tp_batch_cpu_event; +static uint64_t g_tp_batch_seq; +/* Batch flag values are tagged so a stale row-gate seq in the reused FFN + * flag word can never satisfy a batch arrival spin (and vice versa). */ +#define DS4_TP_BATCH_FLAG_TAG 0x80000000u +/* Expert-ownership split parameters for routed kernels. World 1 means TP is + * not bound; world 2 assigns each rank one contiguous expert range. */ +static int32_t g_tp_split_rank; +static int32_t g_tp_split_world = 1; +static int32_t g_tp_session_batch_mode; + +static int ds4_gpu_tp_world_is_two(void) { + return g_tp_split_world == 2; +} + +/* Return the contiguous routed-expert range backed by this process. Rank 1 + * owns the high range and receives any odd-count remainder. */ +static void ds4_gpu_tp_expert_range(uint32_t n_total_expert, + uint32_t *first_expert, + uint32_t *n_expert) { + *first_expert = 0; + *n_expert = n_total_expert; + if (g_tp_split_world != 2) return; + + const uint32_t low_experts = n_total_expert / 2u; + if (g_tp_split_rank == 1) { + *first_expert = low_experts; + *n_expert = n_total_expert - low_experts; + } else { + *n_expert = low_experts; + } +} + +/* Attention head split for GLM batch prefill: each rank computes a + * contiguous half of the heads in the qk-low / attention-lora / + * value-project batch kernels; the caller zeroes the unowned head range + * of the heads buffer and combines the attn-output partials over the + * TP big-gate exchange. */ +static int32_t g_tp_attn_head_split; + +void ds4_gpu_tp_set_attn_head_split(int enabled) { + g_tp_attn_head_split = enabled ? 1 : 0; +} + +static void ds4_gpu_tp_attn_head_range(uint32_t n_head, + uint32_t group, + uint32_t *head_base, + uint32_t *head_count) { + *head_base = 0; + *head_count = n_head; + if (!g_tp_attn_head_split || g_tp_split_world != 2) return; + const uint32_t half = n_head / 2u; + if (half == 0u || (half % group) != 0u || (n_head % 2u) != 0u) return; + *head_count = half; + *head_base = g_tp_split_rank == 1 ? half : 0u; +} +/* Flag gates (DS4_TP_FLAG_GATES): the GPU publishes gate arrival by storing + * the sequence number into a slab word instead of signaling the shared + * event; the service thread spin-reads it from shared memory, which wakes + * hundreds of microseconds earlier than signaledValue polling. The + * CPU->GPU release direction stays on the shared event. */ +static bool g_tp_flag_gates; +static id g_tp_slab_buffer; +static NSUInteger g_tp_slab_buffer_off; +static volatile uint32_t *g_tp_gpu_flags; /* CPU view of the flag words */ +static uint64_t g_tp_gpu_flags_off; +static uint64_t g_tp_seq; +static ds4_gpu_tp_exchange_fn g_tp_exchange_fn; +static ds4_gpu_tp_batch_exchange_fn g_tp_batch_exchange_fn; +static ds4_gpu_tp_big_exchange_fn g_tp_big_exchange_fn; + +void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) { + g_tp_big_exchange_fn = fn; +} + +static void *g_tp_exchange_ud; +static pthread_t g_tp_thread; +static int g_tp_thread_running; +static int g_tp_shutdown; +static int g_tp_failed_flag; +static pthread_mutex_t g_tp_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_tp_cond = PTHREAD_COND_INITIALIZER; +static ds4_gpu_tp_request g_tp_queue[DS4_GPU_TP_QUEUE]; +static uint32_t g_tp_queue_head; +static uint32_t g_tp_queue_count; + +static uint64_t g_tp_stat_gates; +static double g_tp_stat_gpu_wait_ms; +static double g_tp_stat_exchange_ms; + +/* GPU keep-alive (see kernel_dsv4_tp_keepalive): its own queue and thread, + * alive exactly as long as the TP gate machinery. */ +static id g_tp_keepalive_queue; +static id g_tp_keepalive_buffer; +static pthread_t g_tp_keepalive_thread; +static int g_tp_keepalive_running; + +/* Nonzero while a verify block runs: the GPU is genuinely busy + * there, so the keep-alive is a pure parasite (~2.3ms per 5-row block + * measured against the single-machine verify). */ +static volatile int g_tp_keepalive_paused; + +void ds4_gpu_tp_keepalive_pause(int paused) { + g_tp_keepalive_paused = paused; +} + +void ds4_gpu_tp_set_session_batch_mode(int enabled) { + g_tp_session_batch_mode = enabled ? 1 : 0; +} + +static uint32_t ds4_gpu_tp_keepalive_tgs_from_env(void) { + uint32_t ka_tgs = 1; + const char *tgs_env = getenv("DS4_TP_KEEPALIVE_TGS"); + if (tgs_env) { + int v = atoi(tgs_env); + if (v > 0 && v <= 2048) ka_tgs = (uint32_t)v; + } + return ka_tgs; +} + +static void *ds4_gpu_tp_keepalive_thread(void *arg) { + (void)arg; + /* Swept on the M5 Max pair: too few iterations lets clocks sag. Current + * TP split-resident Flash runs show 1.2M is a small Q4/Q2 decode win over + * 800k, while two threadgroups waste work. */ + uint32_t iters = 1200000; + const char *env = getenv("DS4_TP_KEEPALIVE_ITERS"); + if (env) iters = (uint32_t)atoi(env); + /* One ALU-only threadgroup keeps the GPU from power-gating but does + * not push the frequency governor; solo-vs-engine kernel gaps + * (~1.7x) suggest decode runs well below max clocks. More TGs raise + * apparent utilization without eating memory bandwidth. */ + uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive"); + if (!pipeline) { + fprintf(stderr, "ds4: TP keep-alive pipeline missing\n"); + return NULL; + } + while (!g_tp_shutdown) { + if (g_tp_keepalive_paused) { + usleep(200); + continue; + } + @autoreleasepool { + id cb = [g_tp_keepalive_queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_keepalive_buffer offset:0 atIndex:0]; + [enc setBytes:&iters length:sizeof(iters) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)ka_tgs, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + } + } + return NULL; +} + +static void *ds4_gpu_tp_service_thread(void *arg) { + (void)arg; + const bool profile = getenv("DS4_TP_GATE_PROFILE") != NULL; + while (1) { + pthread_mutex_lock(&g_tp_mutex); + while (g_tp_queue_count == 0 && !g_tp_shutdown) + pthread_cond_wait(&g_tp_cond, &g_tp_mutex); + if (g_tp_shutdown && g_tp_queue_count == 0) { + pthread_mutex_unlock(&g_tp_mutex); + break; + } + ds4_gpu_tp_request req = g_tp_queue[g_tp_queue_head]; + g_tp_queue_head = (g_tp_queue_head + 1) % DS4_GPU_TP_QUEUE; + g_tp_queue_count--; + pthread_mutex_unlock(&g_tp_mutex); + + /* Wait for the GPU to reach this gate. Tight spin: gate arrival is + * on the decode critical path and normally tens of microseconds + * out; yielding here measurably delays the release wake-up. */ + const double t0 = profile ? ds4_gpu_now_ms() : 0.0; + uint32_t spins = 0; + if (req.big_bytes > 0) { + /* Big gates always signal arrival through the batch shared + * event (see ds4_gpu_tp_big_gate_kick): the event completion + * semantics are what guarantee the bounce payload is visible + * before the exchange reads it. */ + while (g_tp_batch_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } else if (!req.event_arrival) { + const uint32_t slot = req.layer * 2u + req.gate; + uint32_t want = (uint32_t)req.seq; + if (req.rows > 0) + want = DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq; + while (__atomic_load_n(&g_tp_gpu_flags[slot], __ATOMIC_ACQUIRE) != want) { + if (g_tp_shutdown) break; + if (++spins > (1u << 20)) { + sched_yield(); + spins = 0; + } + } + } else if (req.rows > 0) { + while (g_tp_batch_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } else { + while (g_tp_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } + const double t1 = profile ? ds4_gpu_now_ms() : 0.0; + int ok = 0; + if (!g_tp_shutdown && !g_tp_failed_flag) { + if (req.big_bytes > 0) { + if (g_tp_big_exchange_fn) + ok = g_tp_big_exchange_fn(g_tp_exchange_ud, req.layer, + req.seq, req.big_out, + req.big_in, req.big_bytes); + } else if (req.rows > 0) { + if (g_tp_batch_exchange_fn) + ok = g_tp_batch_exchange_fn(g_tp_exchange_ud, req.layer, + req.rows, req.seq); + } else if (g_tp_exchange_fn) { + ok = g_tp_exchange_fn(g_tp_exchange_ud, req.layer, req.gate, + req.seq); + } + } + if (!ok && !g_tp_shutdown) { + if (!g_tp_failed_flag) + fprintf(stderr, "ds4: TP gate exchange failed (layer %u gate %u seq %llu)\n", + req.layer, req.gate, (unsigned long long)req.seq); + g_tp_failed_flag = 1; + } + /* Release the GPU even on failure so end_commands can drain. */ + if (req.rows > 0) g_tp_batch_cpu_event.signaledValue = req.seq; + else g_tp_cpu_event.signaledValue = req.seq; + if (profile) { + g_tp_stat_gpu_wait_ms += t1 - t0; + g_tp_stat_exchange_ms += ds4_gpu_now_ms() - t1; + if (++g_tp_stat_gates % 860 == 0) { + fprintf(stderr, + "ds4: TP gates %llu: avg gpu-wait %.1f us, avg exchange %.1f us\n", + (unsigned long long)g_tp_stat_gates, + g_tp_stat_gpu_wait_ms / (double)g_tp_stat_gates * 1000.0, + g_tp_stat_exchange_ms / (double)g_tp_stat_gates * 1000.0); + } + } + } + return NULL; +} + +int ds4_gpu_tp_init(uint32_t rank, + ds4_gpu_tensor *slab, uint64_t gpu_flags_off, + ds4_gpu_tp_exchange_fn fn, void *ud) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_tp_thread_running || rank > 1) return 0; + g_tp_split_rank = (int32_t)rank; + g_tp_split_world = 2; + g_tp_slab_buffer = slab ? ds4_gpu_tensor_buffer(slab) : nil; + g_tp_slab_buffer_off = slab ? ds4_gpu_tensor_offset(slab) : 0; + g_tp_gpu_flags_off = gpu_flags_off; + g_tp_gpu_flags = slab ? + (volatile uint32_t *)((uint8_t *)ds4_gpu_tensor_contents(slab) + gpu_flags_off) : NULL; + /* Flag arrival is the default: the slab-word publish detects in ~1-3us + * where signaledValue polling costs 10-20, worth +2.3 t/s on the pair + * (A/B 2026-07-06, byte-identical output). DS4_TP_EVENT_GATES falls + * back to the shared-event arrival path. */ + g_tp_flag_gates = g_tp_gpu_flags != NULL && getenv("DS4_TP_EVENT_GATES") == NULL; + g_tp_gpu_event = [g_device newSharedEvent]; + g_tp_cpu_event = [g_device newSharedEvent]; + g_tp_batch_gpu_event = [g_device newSharedEvent]; + g_tp_batch_cpu_event = [g_device newSharedEvent]; + if (!g_tp_gpu_event || !g_tp_cpu_event || + !g_tp_batch_gpu_event || !g_tp_batch_cpu_event) { + fprintf(stderr, "ds4: failed to create TP shared events\n"); + return 0; + } + g_tp_exchange_fn = fn; + g_tp_exchange_ud = ud; + g_tp_seq = 0; + g_tp_batch_seq = 0; + g_tp_shutdown = 0; + g_tp_failed_flag = 0; + g_tp_queue_head = 0; + g_tp_queue_count = 0; + pthread_attr_t attr; + pthread_attr_init(&attr); + pthread_attr_set_qos_class_np(&attr, QOS_CLASS_USER_INTERACTIVE, 0); + if (pthread_create(&g_tp_thread, &attr, ds4_gpu_tp_service_thread, NULL) != 0) { + pthread_attr_destroy(&attr); + fprintf(stderr, "ds4: failed to start TP gate service thread\n"); + return 0; + } + pthread_attr_destroy(&attr); + g_tp_thread_running = 1; + if (getenv("DS4_TP_NO_KEEPALIVE") == NULL) { + uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); + g_tp_keepalive_queue = [g_device newCommandQueue]; + g_tp_keepalive_buffer = [g_device newBufferWithLength:(NSUInteger)ka_tgs * 256u * sizeof(float) + options:MTLResourceStorageModeShared]; + if (g_tp_keepalive_queue && g_tp_keepalive_buffer && + pthread_create(&g_tp_keepalive_thread, NULL, + ds4_gpu_tp_keepalive_thread, NULL) == 0) { + g_tp_keepalive_running = 1; + } else { + fprintf(stderr, "ds4: TP keep-alive setup failed (continuing without)\n"); + } + } + return 1; +} + +void ds4_gpu_tp_shutdown(void) { + if (!g_tp_thread_running) return; + pthread_mutex_lock(&g_tp_mutex); + g_tp_shutdown = 1; + pthread_cond_broadcast(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + pthread_join(g_tp_thread, NULL); + g_tp_thread_running = 0; + if (g_tp_keepalive_running) { + pthread_join(g_tp_keepalive_thread, NULL); + g_tp_keepalive_running = 0; + g_tp_keepalive_queue = nil; + g_tp_keepalive_buffer = nil; + } + g_tp_exchange_fn = NULL; + g_tp_batch_exchange_fn = NULL; + g_tp_exchange_ud = NULL; + g_tp_split_rank = 0; + g_tp_split_world = 1; + g_tp_session_batch_mode = 0; +} + +void ds4_gpu_tp_suspend_expert_sharding(int suspend) { + if (!g_tp_thread_running) return; + g_tp_split_world = suspend ? 1 : 2; +} + +int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) { + if (!g_batch_cb) { + fprintf(stderr, "ds4: TP gate encode without an open command batch (layer %u gate %u)\n", + layer, gate); + return 0; + } + if (!g_tp_thread_running) { + fprintf(stderr, "ds4: TP gate encode without the gate service (layer %u)\n", layer); + return 0; + } + const uint64_t seq = ++g_tp_seq; + const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; + if (!event_arrival) { + /* Publish arrival through the slab word; the buffer hazard against + * the partial-output kernels orders the store after the payload. */ + const uint32_t slot = layer * 2u + gate; + const uint32_t value = (uint32_t)seq; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); + if (!pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_slab_buffer + offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) + atIndex:0]; + [enc setBytes:&value length:sizeof(value) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + ds4_gpu_close_batch_encoder(); + } else { + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_gpu_event value:seq]; + } + [g_batch_cb encodeWaitForEvent:g_tp_cpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = gate; + g_tp_queue[tail].rows = 0; + g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = NULL; + g_tp_queue[tail].big_in = NULL; + g_tp_queue[tail].big_bytes = 0; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return 1; +} + +void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) { + g_tp_batch_exchange_fn = fn; +} + +/* Verify-block batch gate: same arrival/release machinery as the row gate + * (the FFN flag word and event pair are reused — a decode gate and a batch + * gate are never in flight together, and seq values stay globally unique), + * but the service thread runs the multi-row exchange callback. */ +int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) { + if (!g_batch_cb) return 0; + if (!g_tp_thread_running || rows == 0) return 0; + const uint64_t seq = ++g_tp_batch_seq; + const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; + if (!event_arrival) { + const uint32_t slot = layer * 2u + 1u; /* FFN gate slot */ + const uint32_t value = DS4_TP_BATCH_FLAG_TAG | (uint32_t)seq; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); + if (!pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_slab_buffer + offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) + atIndex:0]; + [enc setBytes:&value length:sizeof(value) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + ds4_gpu_close_batch_encoder(); + } else { + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; + } + [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = 1u; /* FFN */ + g_tp_queue[tail].rows = rows; + g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = NULL; + g_tp_queue[tail].big_in = NULL; + g_tp_queue[tail].big_bytes = 0; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return 1; +} + +/* Prefill batch gate kick: same seq space and release event as the verify + * batch gate, but the service thread exchanges big_bytes directly between + * the two shared bounce buffers instead of slab slots. The kick only + * publishes the GPU arrival marker and queues the exchange; the caller + * encodes the release wait later through ds4_gpu_tp_big_gate_wait, which + * lets it interleave more GPU work with the wire exchange. Arrival always + * uses the batch shared event, NOT the flag word: a flag write carries no + * memory-visibility guarantee for the payload buffer, and once the GPU + * keeps running past the kick (no event wait right behind it) the service + * thread can observe the flag before the producing kernels' stores reach + * CPU-visible memory (measured: stale rows in the first sub-kick). The + * shared-event signal only fires after every preceding command completes, + * which is exactly the payload ordering the exchange needs; the ~10 us + * slower arrival detection is noise against a multi-ms exchange. */ +uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows, + const ds4_gpu_tensor *out_t, + ds4_gpu_tensor *in_t, + uint64_t bytes) { + if (!g_batch_cb) return 0; + if (!g_tp_thread_running || rows == 0 || bytes == 0) return 0; + const void *out_ptr = ds4_gpu_tensor_contents((ds4_gpu_tensor *)out_t); + void *in_ptr = ds4_gpu_tensor_contents(in_t); + if (!out_ptr || !in_ptr) { + fprintf(stderr, "ds4: TP big gate needs CPU-visible bounce buffers\n"); + return 0; + } + const uint64_t seq = ++g_tp_batch_seq; + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = 1u; + g_tp_queue[tail].rows = rows; + g_tp_queue[tail].event_arrival = 1u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = out_ptr; + g_tp_queue[tail].big_in = in_ptr; + g_tp_queue[tail].big_bytes = bytes; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return seq; +} + +/* Encode the GPU-side release wait for a previously kicked big gate. The + * batch release event is monotonic and the service thread completes queued + * exchanges in kick order, so waiting on the LAST kicked seq of a stage + * also covers every earlier kick. */ +int ds4_gpu_tp_big_gate_wait(uint64_t seq) { + if (!g_batch_cb || seq == 0) return 0; + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; + return 1; +} + +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) { + const uint64_t seq = ds4_gpu_tp_big_gate_kick(layer, rows, out_t, in_t, bytes); + if (seq == 0) return 0; + return ds4_gpu_tp_big_gate_wait(seq); +} + +int ds4_gpu_tp_failed(void) { + return g_tp_failed_flag; +} + +int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (event_value == 0) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) return 0; + + const char *what = label ? label : "selected-id readback"; + const BOOL signaled = + [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); + return 0; + } + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +static int ds4_gpu_signal_batch_and_wait_event(const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) { + g_selected_readback_event = [g_device newSharedEvent]; + if (!g_selected_readback_event) { + fprintf(stderr, "ds4: failed to create Metal shared event for %s\n", + label ? label : "selected-id readback"); + return 0; + } + } + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + const uint64_t value = ++g_selected_readback_event_value; + [cb encodeSignalEvent:g_selected_readback_event value:value]; + g_batch_has_work = YES; + [cb commit]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + const BOOL signaled = [g_selected_readback_event waitUntilSignaledValue:value timeoutMS:60000]; + [g_pending_cbs addObject:cb]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", + label ? label : "selected-id readback"); + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + label ? label : "selected-id readback", + [[cb.error localizedDescription] UTF8String]); + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; + } else { + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + if (ds4_gpu_finish_command_buffer(cb, 1, label ? label : "selected-id readback") == 0) { + return 0; + } + return ds4_gpu_begin_commands(); + } +} + +int ds4_gpu_end_commands(void) { + if (!g_batch_cb) return 0; + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + g_stream_expert_cache_owned_seq = g_stream_expert_cache_batch_seq; + g_stream_expert_cache_batch_seq = 0; + return ds4_gpu_finish_command_buffer(cb, 1, "command batch"); +} + +static int ds4_gpu_flash_attn_stage_profile_boundary( + id __strong *cbp, + const char *mode, + const char *stage, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t n_keys, + uint32_t n_head, + uint32_t head_dim, + uint32_t window, + uint32_t ratio, + double *stage_t0) { + if (!cbp || !*cbp || !stage_t0 || !stage) return 0; + if (ds4_gpu_end_commands() == 0) return 0; + + const double now_ms = ds4_gpu_now_ms(); + const char *filter = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE_FILTER"); + const int print_stage = + !filter || !filter[0] || + strstr(stage, filter) != NULL || + (mode && strstr(mode, filter) != NULL); + if (print_stage) { + fprintf(stderr, + "ds4: Metal FlashAttention prefill stage mode=%s tokens=%u comp=%u " + "keys=%u heads=%u dim=%u window=%u ratio=%u %s=%.3f ms\n", + mode ? mode : "unknown", + n_tokens, + n_comp, + n_keys, + n_head, + head_dim, + window, + ratio, + stage, + now_ms - *stage_t0); + } + *stage_t0 = now_ms; + + if (ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + *cbp = ds4_gpu_command_buffer(&owned); + return *cbp != nil && owned == 0; +} + +int ds4_gpu_synchronize(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_batch_cb) return ds4_gpu_end_commands(); + if ([g_pending_cbs count] != 0) { + int ok = ds4_gpu_wait_pending_command_buffers("synchronize"); + [g_transient_buffers removeAllObjects]; + ds4_gpu_model_buffer_cache_maybe_evict("synchronize"); + return ok; + } + + id cb = ds4_gpu_new_command_buffer(); + if (!cb) return 0; + return ds4_gpu_finish_command_buffer(cb, 1, "synchronize"); +} + +void ds4_gpu_cleanup(void) { + if (!g_initialized) return; + + @autoreleasepool { + if (g_batch_cb) { + ds4_gpu_close_batch_encoder(); + [g_batch_cb commit]; + [g_batch_cb waitUntilCompleted]; + g_batch_cb = nil; + if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_batch_seq; + } + g_stream_expert_cache_batch_seq = 0; + } + (void)ds4_gpu_wait_pending_command_buffers("cleanup"); + if (ds4_gpu_stream_expert_timing_summary_enabled() && + getenv("DS4_METAL_MEMORY_REPORT") == NULL) { + ds4_gpu_print_memory_report("at cleanup"); + } + g_selected_readback_event = nil; + g_selected_readback_event_value = 0; + [g_transient_buffers removeAllObjects]; + ds4_gpu_stream_expert_pread_pool_shutdown(); + ds4_gpu_stream_expert_cache_clear_all(1); + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + g_stream_expert_cache_gate_addr_buffers[layer] = nil; + g_stream_expert_cache_up_addr_buffers[layer] = nil; + g_stream_expert_cache_down_addr_buffers[layer] = nil; + g_stream_compact_gate_addr_buffers[layer] = nil; + g_stream_compact_up_addr_buffers[layer] = nil; + g_stream_compact_down_addr_buffers[layer] = nil; + g_stream_compact_selected_buffers[layer] = nil; + g_stream_selected_id_buffers[layer] = nil; + } + g_set_rows_f32_i32_pipeline = nil; + g_get_rows_f32_pipeline = nil; + g_get_rows_f16_pipeline = nil; + g_get_rows_i32_pipeline = nil; + g_get_rows_q8_0_pipeline = nil; + g_get_rows_q4_0_pipeline = nil; + g_get_rows_q4_K_pipeline = nil; + g_repeat_f32_pipeline = nil; + g_concat_pipeline = nil; + g_cpy_f32_f32_pipeline = nil; + g_cpy_f32_f16_pipeline = nil; + g_cpy_contig_f32_f16_pipeline = nil; + g_cpy_f16_f32_pipeline = nil; + g_cpy_f16_f16_pipeline = nil; + g_cpy_contig_f16_f32_pipeline = nil; + g_cpy_contig_f16_f16_pipeline = nil; + g_flash_kv_stage_f16_pipeline = nil; + g_swiglu_pipeline = nil; + g_swiglu_flat_pipeline = nil; + g_add_pipeline = nil; + g_add2_pipeline = nil; + g_add3_pipeline = nil; + g_moe_sum6_pipeline = nil; + g_moe_sum8_pipeline = nil; + g_mul_pipeline = nil; + g_bin_mul_scalar_pipeline = nil; + g_bin_div_row_pipeline = nil; + g_unary_sigmoid_pipeline = nil; + g_unary_silu_pipeline = nil; + g_unary_softplus_pipeline = nil; + g_unary_sqrt_pipeline = nil; + g_unary_clamp_pipeline = nil; + g_unary_scale_pipeline = nil; + g_unary_fill_pipeline = nil; + g_unary_fill_f16_pipeline = nil; + g_rms_norm_pipeline = nil; + g_rms_norm_plain_pipeline = nil; + g_add_rms_norm_pipeline = nil; + g_rms_norm_scale_pipeline = nil; + g_dsv4_qkv_rms_norm_pipeline = nil; + g_hc_split_sinkhorn_pipeline = nil; + g_hc_split_weighted_sum_pipeline = nil; + g_hc_split_weighted_sum_norm_pipeline = nil; + g_hc_weighted_sum_pipeline = nil; + g_hc_weighted_sum_norm_pipeline = nil; + g_output_hc_weights4_pipeline = nil; + g_hc_expand_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_q2_k_pipeline = nil; + g_moe_mul_mv_id_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil; + g_moe_mul_mv_id_q4_k_pipeline = nil; + g_moe_mul_mv_id_q4_k_pair_pipeline = nil; + g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group6_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group8_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group24_q4_k_id_pipeline = nil; + g_moe_mul_mv_group24_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pipeline = nil; + g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = nil; + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = nil; + g_moe_stream_expert_cache_validate_pipeline = nil; + g_moe_q4_gather_slots6_pipeline = nil; + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_table_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_addr_q4_k_sum6_pipeline = nil; + g_moe_table_q4_pair_gate_encoder = nil; + g_moe_table_q4_pair_up_encoder = nil; + g_moe_table_q4_sum_down_encoder = nil; + g_rope_tail_batch_pipeline = nil; + g_rope_tail_inplace_pair_pipeline = nil; + g_rope_tail_inplace_pair_shared4_pipeline = nil; + g_rope_tail_inplace_pair_affine_pipeline = nil; + g_dsv4_fp8_kv_quantize_pipeline = nil; + g_dsv4_indexer_qat_pipeline = nil; + g_dsv4_kv_fp8_store_pipeline = nil; + g_dsv4_ratio4_shift_pipeline = nil; + g_dsv4_compressor_pack_ratio4_pipeline = nil; + g_dsv4_softmax_pool_ratio4_direct_pipeline = nil; + g_dsv4_softmax_pool_pipeline = nil; + g_soft_max_f32_pipeline = nil; + g_soft_max_f32_4_pipeline = nil; + g_argsort_f32_i32_desc_pipeline = nil; + g_argsort_merge_f32_i32_desc_pipeline = nil; + g_sum_rows_f32_f32_pipeline = nil; + g_dsv4_topk_mask_pipeline = nil; + g_dsv4_topk_mask_scatter_pipeline = nil; + g_dsv4_indexer_weighted_sum_pipeline = nil; + g_dsv4_indexer_score_one_direct_pipeline = nil; + g_dsv4_compressor_store_one_pipeline = nil; + g_dsv4_sort_i32_rows_asc_pipeline = nil; + g_dsv4_indexed_attention_heads8_pipeline = nil; + g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; + g_dsv4_softplus_sqrt_pipeline = nil; + g_dsv4_router_finalize_one_pipeline = nil; + g_dsv4_router_finalize_one_simd_pipeline = nil; + g_dsv4_router_finalize_weights_one_simd_pipeline = nil; + g_dsv4_router_weights_one_pipeline = nil; + g_glm_router_select_one_pipeline = nil; + g_glm_kv_lora_rms_norm_pipeline = nil; + g_glm_k_b_project_pipeline = nil; + g_glm_store_compact_kv_pipeline = nil; + g_glm_qkv_norm_store_compact_kv_pipeline = nil; + g_glm_store_indexer_k_pipeline = nil; + g_glm_build_kv_cache_pipeline = nil; + g_glm_build_kv_cache_decode_group4_pipeline = nil; + g_glm_build_kv_cache_flash_pipeline = nil; + g_glm_attention_full_pipeline = nil; + g_glm_fill_selected_range_pipeline = nil; + g_glm_fill_selected_range_batch_pipeline = nil; + g_glm_indexer_rope_tail_pipeline = nil; + g_glm_indexer_score_one_pipeline = nil; + g_glm_indexer_score_one_direct_pipeline = nil; + g_glm_indexer_scores_batch_pipeline = nil; + g_glm_indexer_scores_tiled_pipeline = nil; + g_glm_indexer_scores_tiled_f32_pipeline = nil; + g_glm_qk_lowrank_pipeline = nil; + g_glm_qk_lowrank_glm52_pipeline = nil; + g_glm_qk_lowrank_glm52_sg_pipeline = nil; + g_glm_qk_lowrank_batch_pipeline = nil; + g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil; + g_glm_value_project_q8_0_pipeline = nil; + g_glm_value_project_q8_0_batch_heads_pipeline = nil; + g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil; + g_glm_attention_indexed_decode_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_partial_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil; + g_glm_attention_indexed_batch_pipeline = nil; + g_glm_attention_indexed_batch_group2_pipeline = nil; + g_glm_attention_indexed_batch_q2_group4_pipeline = nil; + g_glm_attention_indexed_batch_group8_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil; + g_glm_q4_k_pair_swiglu_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu4_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = nil; + g_glm_q2_k_pair_swiglu_f32_pipeline = nil; + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = nil; + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = nil; + g_glm_q4_k_addr_pair_swiglu_f32_pipeline = nil; + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = nil; + g_glm_q2_k_down_f32_pipeline = nil; + g_glm_q4_k_down_f32_pipeline = nil; + g_glm_q2_k_addr_down_f32_pipeline = nil; + g_glm_q4_k_addr_down_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = nil; + g_glm_q5_k_down_f32_pipeline = nil; + g_glm_q6_k_down_f32_pipeline = nil; + g_dsv4_router_weights_batch_pipeline = nil; + g_dsv4_hc_expand4_pipeline = nil; + g_flash_attn_mask_buffer = nil; + g_flash_attn_zero_mask_buffer = nil; + g_flash_attn_pad_buffer = nil; + g_flash_attn_tmp_buffer = nil; + g_flash_attn_blk_buffer = nil; + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); + g_flash_attn_ring_buffer = nil; + g_flash_attn_kv_buffer = nil; + g_glm_flash_attn_mask_buffer = nil; + g_compressor_pool_kv_buffer = nil; + g_compressor_pool_score_buffer = nil; + g_compressor_pool_score_cont_buffer = nil; + g_compressor_pool_softmax_buffer = nil; + g_compressor_pool_product_buffer = nil; + g_compressor_store_ape_buffer = nil; + g_compressor_store_score_buffer = nil; + g_embed_rows_buffer = nil; + g_router_selection_buffer = nil; + g_router_weight_sum_buffer = nil; + g_indexer_head_scores_buffer = nil; + g_indexer_topk_buffer = nil; + g_indexed_topk_buffer = nil; + g_stream_expert_validate_status_buffer = nil; + g_f16_round_scratch_buffer = nil; + g_raw_store_round_buffer = nil; + g_moe_gate_scratch_buffer = nil; + g_moe_down_scratch_buffer = nil; + g_moe_id_map_buffer = nil; + g_moe_q4_gate_slots_buffer = nil; + g_moe_q4_up_slots_buffer = nil; + g_moe_q4_down_slots_buffer = nil; + g_attn_out_group_ids_buffer = nil; + g_model_fd = -1; + g_model_map_ptr = NULL; + g_model_map_size = 0; + g_model_mapped_offset = 0; + g_model_mapped_size = 0; + g_model_mapped_max_tensor_bytes = 0; + ds4_gpu_tensor_tracking_reset(); + g_flash_attn_mask_bytes = 0; + g_flash_attn_zero_mask_bytes = 0; + g_flash_attn_pad_bytes = 0; + g_flash_attn_tmp_bytes = 0; + g_flash_attn_blk_bytes = 0; + g_flash_attn_ring_bytes = 0; + g_flash_attn_kv_bytes = 0; + g_glm_flash_attn_mask_bytes = 0; + g_glm_flash_attn_mask_valid = 0; + g_glm_flash_attn_mask_pos0 = 0; + g_glm_flash_attn_mask_tokens = 0; + g_glm_flash_attn_mask_cache_len = 0; + g_compressor_pool_kv_bytes = 0; + g_compressor_pool_score_bytes = 0; + g_compressor_pool_score_cont_bytes = 0; + g_compressor_pool_softmax_bytes = 0; + g_compressor_pool_product_bytes = 0; + g_compressor_store_ape_bytes = 0; + g_compressor_store_score_bytes = 0; + g_embed_rows_bytes = 0; + g_router_selection_bytes = 0; + g_router_weight_sum_bytes = 0; + g_indexer_head_scores_bytes = 0; + g_indexer_topk_bytes = 0; + g_indexed_topk_bytes = 0; + g_f16_round_scratch_bytes = 0; + g_raw_store_round_bytes = 0; + g_moe_gate_scratch_bytes = 0; + g_moe_down_scratch_bytes = 0; + g_moe_id_map_bytes = 0; + g_moe_q4_gate_slots_bytes = 0; + g_moe_q4_up_slots_bytes = 0; + g_moe_q4_down_slots_bytes = 0; + g_attn_out_group_ids_bytes = 0; + g_model_wrap_count = 0; + g_model_wrap_bytes = 0; + g_model_wrap_max_bytes = 0; + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_evictions = 0; + g_model_buffer_cache_over_limit = 0; + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + [g_pipeline_cache removeAllObjects]; + g_pipeline_cache = nil; + [g_q4_expert_layer_residency_cache removeAllObjects]; + g_q4_expert_layer_residency_cache = nil; + [g_q4_expert_table_cache removeAllObjects]; + g_q4_expert_table_cache = nil; + [g_model_buffer_cache removeAllObjects]; + g_model_buffer_cache = nil; + g_transient_buffers = nil; + g_pending_cbs = nil; + g_library = nil; + g_queue = nil; + g_device = nil; + g_initialized = 0; + } +} diff --git a/models/README.md b/models/README.md new file mode 100644 index 0000000000..8261680abd --- /dev/null +++ b/models/README.md @@ -0,0 +1,34 @@ +# Model integrations + +Each directory owns one complete, tailored model integration: + +```text +models// +├── provider.c / provider.h +├── cpu.inc +├── graph.inc +├── cuda/ +├── metal/ +│ ├── host/ +│ └── shaders/ +└── rocm/ +``` + +The engine-facing boundary is the whole-model `ds4_model_provider_v1` +lifecycle. A provider owns its session orchestration and calls its custom +kernels directly; there is intentionally no generic kernel, operator, graph, +or tensor interface between them. + +The `.inc` implementation fragments are still included exactly once by the +engine or backend entry point. This preserves the existing single translation +units, private types, static linkage, and compiler visibility. The directory +split expresses ownership without adding wrappers to hot paths. + +Code belongs under `models//` when its semantics, tensor layout, or +launch sequence are specific to that model. Backend runtime, memory management, +and genuinely reused low-level primitives remain under `cuda/`, `metal/`, +`rocm/`, and `kernels/`. + +To add a model, implement its provider and the backend paths it supports. It is +fine to duplicate kernels when separate implementations are easier to tune or +understand. diff --git a/models/deepseek/README.md b/models/deepseek/README.md new file mode 100644 index 0000000000..c52bb9131a --- /dev/null +++ b/models/deepseek/README.md @@ -0,0 +1,14 @@ +# DeepSeek V4 integration + +This directory owns the DeepSeek V4 model provider and its tailored inference +implementation: + +- `provider.c` exposes the whole-model lifecycle to the engine core. +- `cpu.inc` is the CPU reference and decode path. +- `graph.inc` owns GPU graph state, allocation, prefill, decode, checkpoint, + and layer-slice orchestration. +- `cuda/`, `metal/`, and `rocm/` contain DeepSeek-specific host and device + implementations. + +The provider calls these concrete paths directly. They do not implement a +generic kernel interface. diff --git a/models/deepseek/cpu.inc b/models/deepseek/cpu.inc new file mode 100644 index 0000000000..d87e7774a0 --- /dev/null +++ b/models/deepseek/cpu.inc @@ -0,0 +1,4373 @@ +/* + * DeepSeek V4 CPU inference pipeline. + * + * Included exactly once by ds4.c so model-private helpers stay static and the + * compiler sees the complete specialized inference path. + */ + +/* ========================================================================= + * Hyper-Connection Transforms. + * ========================================================================= + * + * DeepSeek V4 Flash keeps four hyper-connection streams per token. Before + * attention or FFN, a learned small projection chooses how to reduce the HC + * state into the 4096-wide sublayer input. After the sublayer, the post and + * combine weights expand the result back into the four-stream HC state. + */ + +/* Decode the HC control projection. The output contains pre weights, post + * gates, and a small doubly-normalized combine matrix. */ +static void hc_split_sinkhorn_one( + float * out, + const float * mix, + const float * scale, + const float * base, + int n_hc, + int iters, + float eps) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + for (int i = 0; i < n_hc; i++) { + const float z = mix[i] * pre_scale + base[i]; + out[i] = 1.0f / (1.0f + expf(-z)) + eps; + } + + for (int i = 0; i < n_hc; i++) { + const int off = n_hc + i; + const float z = mix[off] * post_scale + base[off]; + out[off] = 2.0f / (1.0f + expf(-z)); + } + + float c[16 * 16]; + + for (int dst = 0; dst < n_hc; dst++) { + float row_max = DS4_NEG_INF; + for (int src = 0; src < n_hc; src++) { + const int idx = src + dst * n_hc; + const int off = 2 * n_hc + idx; + const float v = mix[off] * comb_scale + base[off]; + c[idx] = v; + if (v > row_max) row_max = v; + } + + float row_sum = 0.0f; + for (int src = 0; src < n_hc; src++) { + const int idx = src + dst * n_hc; + const float v = expf(c[idx] - row_max); + c[idx] = v; + row_sum += v; + } + + const float inv = 1.0f / row_sum; + for (int src = 0; src < n_hc; src++) { + const int idx = src + dst * n_hc; + c[idx] = c[idx] * inv + eps; + } + } + + for (int src = 0; src < n_hc; src++) { + float sum = 0.0f; + for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; + + const float inv = 1.0f / (sum + eps); + for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; + } + + for (int iter = 1; iter < iters; iter++) { + for (int dst = 0; dst < n_hc; dst++) { + float sum = 0.0f; + for (int src = 0; src < n_hc; src++) sum += c[src + dst * n_hc]; + + const float inv = 1.0f / (sum + eps); + for (int src = 0; src < n_hc; src++) c[src + dst * n_hc] *= inv; + } + + for (int src = 0; src < n_hc; src++) { + float sum = 0.0f; + for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; + + const float inv = 1.0f / (sum + eps); + for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; + } + } + + for (int i = 0; i < n_hc * n_hc; i++) out[2 * n_hc + i] = c[i]; +} + +/* Reduce the four HC streams into the plain embedding vector consumed by a + * normal attention or FFN sublayer. */ +static void hc_weighted_sum_one( + float * out, + const float * x, + const float * weights, + uint32_t n_embd, + uint32_t n_hc) { + for (uint32_t d = 0; d < n_embd; d++) { + float acc = 0.0f; + for (uint32_t h = 0; h < n_hc; h++) { + acc += x[(uint64_t)h * n_embd + d] * weights[h]; + } + out[d] = acc; + } +} + +/* HC pre step for one token. It normalizes the HC state, projects the control + * vector, runs the Sinkhorn split, and emits the sublayer input plus post data. */ +static void hc_pre_from_state_one_scratch( + const ds4_model * model, + const ds4_tensor * fn, + const ds4_tensor * scale_tensor, + const ds4_tensor * base_tensor, + const float * residual_hc, + float * out, + float * post, + float * comb, + float * flat, + bool serial_fn) { + const uint32_t n_hc = DS4_N_HC; + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; + + float mix[24]; + float split[24]; + + rms_norm_no_weight(flat, residual_hc, hc_dim, DS4_RMS_EPS); + if (serial_fn) { + matvec_f16_serial(mix, model, fn, flat); + } else { + matvec_f16(mix, model, fn, flat); + } + + const float *scale = tensor_data(model, scale_tensor); + const float *base = tensor_data(model, base_tensor); + hc_split_sinkhorn_one(split, mix, scale, base, (int)n_hc, DS4_N_HC_SINKHORN_ITER, 1.0e-6f); + hc_weighted_sum_one(out, residual_hc, split, DS4_N_EMBD, n_hc); + + memcpy(post, split + n_hc, n_hc * sizeof(post[0])); + memcpy(comb, split + 2 * n_hc, n_hc * n_hc * sizeof(comb[0])); +} + +static void hc_pre_from_state_one( + const ds4_model * model, + const ds4_tensor * fn, + const ds4_tensor * scale_tensor, + const ds4_tensor * base_tensor, + const float * residual_hc, + float * out, + float * post, + float * comb) { + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; + float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); + + hc_pre_from_state_one_scratch(model, + fn, scale_tensor, base_tensor, + residual_hc, out, post, comb, + flat, false); + free(flat); +} + +static void layer_attn_pre_one( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * token_embd, + float * out, + float * residual_hc, + float * post, + float * comb) { + const uint32_t n_hc = DS4_N_HC; + + for (uint32_t h = 0; h < n_hc; h++) { + memcpy(residual_hc + (uint64_t)h * DS4_N_EMBD, token_embd, (size_t)DS4_N_EMBD * sizeof(token_embd[0])); + } + + hc_pre_from_state_one(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + residual_hc, out, post, comb); +} + +/* The input embedding starts all HC streams with the same token vector. */ +static void hc_from_plain_embedding(float *out_hc, const float *x, uint32_t n_embd, uint32_t n_hc) { + for (uint32_t h = 0; h < n_hc; h++) { + memcpy(out_hc + (uint64_t)h * n_embd, x, (size_t)n_embd * sizeof(x[0])); + } +} + +/* HC post step for one sublayer output. It injects the new block output and + * mixes the previous HC streams through the learned combine matrix. */ +static void hc_post_one( + float * out_hc, + const float * block_out, + const float * residual_hc, + const float * post, + const float * comb, + uint32_t n_embd, + uint32_t n_hc) { + for (uint32_t dst = 0; dst < n_hc; dst++) { + for (uint32_t d = 0; d < n_embd; d++) { + float acc = block_out[d] * post[dst]; + + for (uint32_t src = 0; src < n_hc; src++) { + /* The HC combine matrix is addressed as [dst_hc, src_hc]. */ + acc += comb[dst + src * n_hc] * residual_hc[(uint64_t)src * n_embd + d]; + } + + out_hc[(uint64_t)dst * n_embd + d] = acc; + } + } +} + +typedef struct { + float *out_hc; + const float *block_out; + const float *residual_hc; + const float *post; + const float *comb; + uint64_t hc_dim; + uint32_t n_embd; + uint32_t n_hc; +} hc_post_batch_ctx; + +static void hc_post_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + hc_post_batch_ctx *ctx = vctx; + for (uint64_t t = t0; t < t1; t++) { + hc_post_one(ctx->out_hc + t * ctx->hc_dim, + ctx->block_out + t * ctx->n_embd, + ctx->residual_hc + t * ctx->hc_dim, + ctx->post + t * ctx->n_hc, + ctx->comb + t * ctx->n_hc * ctx->n_hc, + ctx->n_embd, + ctx->n_hc); + } +} + +static void hc_post_batch( + float * out_hc, + const float * block_out, + const float * residual_hc, + const float * post, + const float * comb, + uint32_t n_tok, + uint32_t n_embd, + uint32_t n_hc) { + hc_post_batch_ctx ctx = { + .out_hc = out_hc, + .block_out = block_out, + .residual_hc = residual_hc, + .post = post, + .comb = comb, + .hc_dim = (uint64_t)n_hc * n_embd, + .n_embd = n_embd, + .n_hc = n_hc, + }; + ds4_parallel_for_min_rows(n_tok, hc_post_batch_worker, &ctx, 1); +} + +typedef struct { + float *out_hc; + const float *moe; + const float *shared; + const float *residual_hc; + const float *post; + const float *comb; + uint64_t hc_dim; + uint32_t n_embd; + uint32_t n_hc; +} hc_post_sum_batch_ctx; + +static void hc_post_sum_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + hc_post_sum_batch_ctx *ctx = vctx; + for (uint64_t t = t0; t < t1; t++) { + const float *moe = ctx->moe + t * ctx->n_embd; + const float *shared = ctx->shared + t * ctx->n_embd; + const float *residual = ctx->residual_hc + t * ctx->hc_dim; + const float *post = ctx->post + t * ctx->n_hc; + const float *comb = ctx->comb + t * ctx->n_hc * ctx->n_hc; + float *out = ctx->out_hc + t * ctx->hc_dim; + + for (uint32_t dst = 0; dst < ctx->n_hc; dst++) { + for (uint32_t d = 0; d < ctx->n_embd; d++) { + float acc = (moe[d] + shared[d]) * post[dst]; + for (uint32_t src = 0; src < ctx->n_hc; src++) { + acc += comb[dst + src * ctx->n_hc] * + residual[(uint64_t)src * ctx->n_embd + d]; + } + out[(uint64_t)dst * ctx->n_embd + d] = acc; + } + } + } +} + +static void hc_post_sum_batch( + float * out_hc, + const float * moe, + const float * shared, + const float * residual_hc, + const float * post, + const float * comb, + uint32_t n_tok, + uint32_t n_embd, + uint32_t n_hc) { + hc_post_sum_batch_ctx ctx = { + .out_hc = out_hc, + .moe = moe, + .shared = shared, + .residual_hc = residual_hc, + .post = post, + .comb = comb, + .hc_dim = (uint64_t)n_hc * n_embd, + .n_embd = n_embd, + .n_hc = n_hc, + }; + ds4_parallel_for_min_rows(n_tok, hc_post_sum_batch_worker, &ctx, 1); +} + +typedef struct { + const ds4_model *model; + const ds4_tensor *fn; + const ds4_tensor *scale; + const ds4_tensor *base; + const ds4_tensor *norm_w; + const float *inp_hc; + float *residual_hc; + float *cur; + float *norm; + float *post; + float *comb; + uint64_t hc_dim; + uint32_t n_hc; +} hc_pre_norm_batch_ctx; + +static void hc_pre_norm_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + hc_pre_norm_batch_ctx *ctx = vctx; + const float *norm_w = tensor_data(ctx->model, ctx->norm_w); + float *flat = xmalloc((size_t)ctx->hc_dim * sizeof(flat[0])); + + for (uint64_t t = t0; t < t1; t++) { + const float *residual = ctx->inp_hc + t * ctx->hc_dim; + if (ctx->residual_hc) { + float *dst = ctx->residual_hc + t * ctx->hc_dim; + memcpy(dst, residual, (size_t)ctx->hc_dim * sizeof(dst[0])); + residual = dst; + } + + hc_pre_from_state_one_scratch(ctx->model, + ctx->fn, + ctx->scale, + ctx->base, + residual, + ctx->cur + t * DS4_N_EMBD, + ctx->post + t * ctx->n_hc, + ctx->comb + t * ctx->n_hc * ctx->n_hc, + flat, + true); + rms_norm_weight(ctx->norm + t * DS4_N_EMBD, + ctx->cur + t * DS4_N_EMBD, + norm_w, + DS4_N_EMBD, + DS4_RMS_EPS); + } + + free(flat); +} + +/* Batched HC pre plus RMSNorm. Prefill uses this to keep the layer-major + * token batch in contiguous arrays. */ +static void hc_pre_norm_batch( + const ds4_model * model, + const ds4_tensor * fn, + const ds4_tensor * scale, + const ds4_tensor * base, + const ds4_tensor * norm_w, + const float * inp_hc, + float * residual_hc, + float * cur, + float * norm, + float * post, + float * comb, + uint32_t n_tok) { + hc_pre_norm_batch_ctx ctx = { + .model = model, + .fn = fn, + .scale = scale, + .base = base, + .norm_w = norm_w, + .inp_hc = inp_hc, + .residual_hc = residual_hc, + .cur = cur, + .norm = norm, + .post = post, + .comb = comb, + .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, + .n_hc = DS4_N_HC, + }; + ds4_parallel_for_min_rows(n_tok, hc_pre_norm_batch_worker, &ctx, 1); +} + +static void layer_attn_norm_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + const float *attn_norm = tensor_data(model, layer->attn_norm); + rms_norm_weight(out, x, attn_norm, DS4_N_EMBD, DS4_RMS_EPS); +} + +/* ========================================================================= + * Attention Projections, RoPE, and Attention Output. + * ========================================================================= + * + * This block performs the attention half of a transformer layer: HC pre, + * attention RMSNorm, Q and KV projections, layer-specific RoPE, sink-aware + * attention over raw and compressed KV rows, and the grouped LoRA output + * projection back to embedding width. + */ + +/* Q projection is low-rank: Q8_0 into the model-specific LoRA-Q rank, + * RMSNorm, then Q8_0 back to all attention heads. */ +static void layer_q_projection_normed_one( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * norm, + float * q) { + const uint32_t q_rank = DS4_N_LORA_Q; + float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); + float *qr_norm = xmalloc((size_t)q_rank * sizeof(qr_norm[0])); + + const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); + + matvec_q8_0(qr, model, layer->attn_q_a, norm); + rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); + matvec_q8_0(q, model, layer->attn_q_b, qr_norm); + head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); + + free(qr_norm); + free(qr); +} + +static void layer_q_projection_with_lora_one( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * norm, + float * q, + float * qr_norm) { + const uint32_t q_rank = DS4_N_LORA_Q; + float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); + const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); + + matvec_q8_0(qr, model, layer->attn_q_a, norm); + rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); + matvec_q8_0(q, model, layer->attn_q_b, qr_norm); + head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); + + free(qr); +} + +/* KV projection has one KV head of width 512, followed by a learned RMSNorm. */ +static void layer_kv_projection_normed_one( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * normed, + float * kv) { + float *raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(raw[0])); + + const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); + + matvec_q8_0(raw, model, layer->attn_kv, normed); + rms_norm_weight(kv, raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); + + free(raw); +} + +static void layer_q_projection_with_lora_one_decode_scratch( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * norm, + float * q, + float * qr_norm, + ds4_cpu_decode_scratch * scratch) { + const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); + + matvec_q8_0_decode_scratch(scratch->qr, model, layer->attn_q_a, norm, scratch); + rms_norm_weight(qr_norm, scratch->qr, q_a_norm, DS4_N_LORA_Q, DS4_RMS_EPS); + matvec_q8_0_decode_scratch(q, model, layer->attn_q_b, qr_norm, scratch); + head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); +} + +static void layer_kv_projection_normed_one_decode_scratch( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * normed, + float * kv, + ds4_cpu_decode_scratch * scratch) { + const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); + + matvec_q8_0_decode_scratch(scratch->kv_raw, model, layer->attn_kv, normed, scratch); + rms_norm_weight(kv, scratch->kv_raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); +} + +static float rope_yarn_ramp(float low, float high, int i0) { + const float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low); + return 1.0f - fminf(1.0f, fmaxf(0.0f, y)); +} + +static float rope_yarn_corr_dim(int n_dims, uint64_t n_ctx_orig, float n_rot, float base) { + return (float)n_dims * logf((float)n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) / (2.0f * logf(base)); +} + +static void rope_yarn_corr_dims(int n_dims, uint64_t n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]) { + const float start = floorf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); + const float end = ceilf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); + dims[0] = fmaxf(0.0f, start); + dims[1] = fminf((float)(n_dims - 1), end); +} + +/* Apply DS4 RoPE only to the tail of each head. Compressed layers use the + * long-context frequency base and scale; inverse mode rotates attention output + * back before the grouped output projection. */ +static void rope_tail_ext_inplace( + float * x, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos, + uint64_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool inverse) { + const uint32_t n_nope = head_dim - n_rot; + const float theta_scale = powf(freq_base, -2.0f / (float)n_rot); + const float sin_sign = inverse ? -1.0f : 1.0f; + float corr_dims[2] = { 0.0f, 0.0f }; + if (ext_factor != 0.0f) { + rope_yarn_corr_dims((int)n_rot, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + } + + for (uint32_t h = 0; h < n_head; h++) { + float *tail = x + (uint64_t)h * head_dim + n_nope; + float theta_extrap = (float)pos; + + for (uint32_t i = 0; i < n_rot; i += 2) { + const float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = attn_factor; + + if (ext_factor != 0.0f) { + const float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], (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); + } + + const float c = cosf(theta) * mscale; + const float s = sin_sign * sinf(theta) * mscale; + const float x0 = tail[i + 0]; + const float x1 = tail[i + 1]; + + tail[i + 0] = x0 * c - x1 * s; + tail[i + 1] = x0 * s + x1 * c; + + theta_extrap *= theta_scale; + } + } +} + +/* Dense layers and compressed layers use different RoPE bases. */ +static float layer_rope_freq_base(uint32_t il) { + return ds4_layer_compress_ratio(il) != 0 && DS4_COMPRESS_ROPE_FREQ_BASE > 0.0f + ? DS4_COMPRESS_ROPE_FREQ_BASE + : DS4_ROPE_FREQ_BASE; +} + +static float layer_rope_freq_scale(uint32_t il) { + if (ds4_layer_compress_ratio(il) == 0 || DS4_ROPE_SCALE_FACTOR <= 0.0f) { + return 1.0f; + } + return 1.0f / DS4_ROPE_SCALE_FACTOR; +} + +static void rope_tail_layer_inplace( + float * x, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos, + uint32_t il, + bool inverse) { + const bool compressed = ds4_layer_compress_ratio(il) != 0; + const float freq_base = layer_rope_freq_base(il); + const float freq_scale = layer_rope_freq_scale(il); + const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; + float attn_factor = 1.0f; + if (ext_factor != 0.0f && freq_scale > 0.0f) { + /* + * This YaRN helper applies magnitude scaling internally. DeepSeek V4 + * reference RoPE uses interpolation without that magnitude change, so + * pass the inverse factor here and let the helper cancel itself out. + */ + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + + rope_tail_ext_inplace(x, n_head, head_dim, n_rot, pos, + compressed ? DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + inverse); +} + +typedef struct { + float *x; + uint64_t stride; + uint32_t n_head; + uint32_t head_dim; + uint32_t n_rot; + uint32_t pos0; + uint32_t il; + bool inverse; +} rope_tail_batch_ctx; + +static void rope_tail_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + rope_tail_batch_ctx *ctx = vctx; + for (uint64_t tt = t0; tt < t1; tt++) { + rope_tail_layer_inplace(ctx->x + tt * ctx->stride, + ctx->n_head, + ctx->head_dim, + ctx->n_rot, + ctx->pos0 + (uint32_t)tt, + ctx->il, + ctx->inverse); + } +} + +static void rope_tail_layer_batch_inplace( + float *x, + uint64_t stride, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t il, + bool inverse, + uint32_t n_tok) { + rope_tail_batch_ctx ctx = { + .x = x, + .stride = stride, + .n_head = n_head, + .head_dim = head_dim, + .n_rot = n_rot, + .pos0 = pos0, + .il = il, + .inverse = inverse, + }; + ds4_parallel_for_min_rows(n_tok, rope_tail_batch_worker, &ctx, 1); +} + +static inline float dot_f32(const float *a, const float *b, uint32_t n) { +#if defined(__ARM_NEON) + uint32_t i = 0; + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + for (; i + 8 <= n; i += 8) { + acc0 = vfmaq_f32(acc0, vld1q_f32(a + i), vld1q_f32(b + i)); + acc1 = vfmaq_f32(acc1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); + } + float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); + for (; i < n; i++) acc += a[i] * b[i]; + return acc; +#else + float acc = 0.0f; + for (uint32_t i = 0; i < n; i++) acc += a[i] * b[i]; + return acc; +#endif +} + +static inline void axpy_f32(float *y, const float *x, float a, uint32_t n) { +#if defined(__ARM_NEON) + uint32_t i = 0; + const float32x4_t av = vdupq_n_f32(a); + for (; i + 8 <= n; i += 8) { + vst1q_f32(y + i, vfmaq_f32(vld1q_f32(y + i), av, vld1q_f32(x + i))); + vst1q_f32(y + i + 4, vfmaq_f32(vld1q_f32(y + i + 4), av, vld1q_f32(x + i + 4))); + } + for (; i < n; i++) y[i] += a * x[i]; +#else + for (uint32_t i = 0; i < n; i++) y[i] += a * x[i]; +#endif +} + +static inline void scale_f32(float *x, float a, uint32_t n) { +#if defined(__ARM_NEON) + uint32_t i = 0; + const float32x4_t av = vdupq_n_f32(a); + for (; i + 8 <= n; i += 8) { + vst1q_f32(x + i, vmulq_f32(vld1q_f32(x + i), av)); + vst1q_f32(x + i + 4, vmulq_f32(vld1q_f32(x + i + 4), av)); + } + for (; i < n; i++) x[i] *= a; +#else + for (uint32_t i = 0; i < n; i++) x[i] *= a; +#endif +} + +static float sigmoid_stable(float x) { + if (x >= 0.0f) { + const float e = expf(-x); + return 1.0f / (1.0f + e); + } else { + const float e = expf(x); + return e / (1.0f + e); + } +} + +/* Sink-aware attention over a set of KV rows. The learned sink logit is part + * of the softmax denominator but contributes no value vector. */ +static void layer_attention_rows_one( + float * out_heads, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * q, + const float * kv_rows, + uint32_t n_kv) { + const float *sinks = tensor_data(model, layer->attn_sinks); + const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); + float score_stack[512]; + float *score = n_kv <= 512 ? score_stack : xmalloc((size_t)n_kv * sizeof(score[0])); + + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; + + float max_score = sinks[h]; + for (uint32_t r = 0; r < n_kv; r++) { + const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; + score[r] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[r] > max_score) max_score = score[r]; + } + + float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; + memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); + + float denom = expf(sinks[h] - max_score); + for (uint32_t r = 0; r < n_kv; r++) { + const float weight = expf(score[r] - max_score); + const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + + const float inv = 1.0f / denom; + scale_f32(oh, inv, DS4_N_HEAD_DIM); + } + + if (score != score_stack) free(score); +} + +static void layer_attention_one( + float * out_heads, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * q, + const float * kv) { + layer_attention_rows_one(out_heads, model, layer, q, kv, 1); +} + +/* Attention output projection is grouped: each group first maps its heads to + * a 1024-rank low vector, then all groups are projected back to 4096. */ +static void layer_grouped_out_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * heads) { + const uint32_t n_groups = 8; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = 1024; + + float *low = xcalloc((size_t)n_groups * rank, sizeof(low[0])); + + matvec_q8_0_grouped_rows(low, model, layer->attn_output_a, heads, n_groups, group_dim, rank); + + matvec_q8_0(out, model, layer->attn_output_b, low); + free(low); +} + +static void layer_grouped_out_one_decode_scratch( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * heads, + ds4_cpu_decode_scratch * scratch) { + const uint32_t n_groups = 8; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = 1024; + + memset(scratch->attn_low, 0, (size_t)n_groups * rank * sizeof(scratch->attn_low[0])); + matvec_q8_0_grouped_rows_decode_scratch(scratch->attn_low, model, layer->attn_output_a, + heads, n_groups, group_dim, rank, scratch); + matvec_q8_0_decode_scratch(out, model, layer->attn_output_b, scratch->attn_low, scratch); +} + +static void layer_grouped_out_batch( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * heads, + uint32_t n_tok) { + const uint32_t n_groups = 8; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = 1024; + + float *low = xcalloc((size_t)n_tok * n_groups * rank, sizeof(low[0])); + + matmul_q8_0_grouped_batch(low, model, layer->attn_output_a, heads, + n_tok, n_groups, group_dim, rank); + matmul_q8_0_batch(out, model, layer->attn_output_b, low, n_tok); + + free(low); +} + +/* ========================================================================= + * Mixture-of-Experts FFN. + * ========================================================================= + * + * This is the FFN half of each layer. It includes the shared expert, routed + * expert selection, IQ2_XXS gate/up projections, SwiGLU, Q2_K down projection, + * and the HC post step that returns the result to four-stream state. + */ + +static float silu(float x) { + return x * sigmoid_stable(x); +} + +static float softplus_stable(float x) { + if (x > 20.0f) return x; + if (x < -20.0f) return expf(x); + return log1pf(expf(x)); +} + +static void swiglu(float *out, const float *gate, const float *up, uint64_t n, float clamp) { + for (uint64_t i = 0; i < n; i++) { + float g = gate[i]; + float u = up[i]; + if (clamp > 1.0e-6f) { + if (g > clamp) g = clamp; + if (u > clamp) u = clamp; + if (u < -clamp) u = -clamp; + } + out[i] = silu(g) * u; + } +} + +/* The shared expert is a normal Q8_0 SwiGLU MLP that runs for every token. */ +static void layer_shared_ffn_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + float *gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])); + float *up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])); + float *mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])); + const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)blocks * 32); + float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); + + if (layer->ffn_up_shexp->type != 8 || + layer->ffn_gate_shexp->type != 8 || + layer->ffn_up_shexp->dim[0] != in_dim) { + ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); + } + + quantize_q8_0_activation(x, xq, xscale, in_dim); + matvec_q8_0_pair_prequant(gate, up, model, + layer->ffn_gate_shexp, + layer->ffn_up_shexp, + xq, xscale); + swiglu(mid, gate, up, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP); + matvec_q8_0(out, model, layer->ffn_down_shexp, mid); + + free(xscale); + free(xq); + free(mid); + free(up); + free(gate); +} + +static void layer_shared_ffn_one_decode_scratch( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + ds4_cpu_decode_scratch * scratch) { + const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; + if (layer->ffn_up_shexp->type != 8 || + layer->ffn_gate_shexp->type != 8 || + layer->ffn_up_shexp->dim[0] != in_dim) { + ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); + } + + matvec_q8_0_pair_decode_scratch(scratch->shared_gate, + scratch->shared_up, + model, + layer->ffn_gate_shexp, + layer->ffn_up_shexp, + x, + scratch); + swiglu(scratch->shared_mid, scratch->shared_gate, scratch->shared_up, DS4_N_FF_EXP, + DS4_SWIGLU_CLAMP_EXP); + matvec_q8_0_decode_scratch(out, model, layer->ffn_down_shexp, scratch->shared_mid, scratch); +} + +typedef struct { + float *mid; + const float *gate; + const float *up; + uint64_t n; + float clamp; +} swiglu_batch_ctx; + +static void swiglu_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { + swiglu_batch_ctx *ctx = vctx; + for (uint64_t t = t0; t < t1; t++) { + swiglu(ctx->mid + t * ctx->n, + ctx->gate + t * ctx->n, + ctx->up + t * ctx->n, + ctx->n, + ctx->clamp); + } +} + +static void layer_shared_ffn_batch( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t n_tok) { + const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; + const uint64_t hidden = layer->ffn_gate_shexp->dim[1]; + + if (layer->ffn_up_shexp->type != 8 || + layer->ffn_gate_shexp->type != 8 || + layer->ffn_down_shexp->type != 8 || + layer->ffn_up_shexp->dim[0] != in_dim || + layer->ffn_up_shexp->dim[1] != hidden || + layer->ffn_down_shexp->dim[0] != hidden) { + ds4_die("shared expert tensors do not share the expected Q8_0 layout"); + } + + float *gate = xmalloc((size_t)n_tok * hidden * sizeof(gate[0])); + float *up = xmalloc((size_t)n_tok * hidden * sizeof(up[0])); + float *mid = xmalloc((size_t)n_tok * hidden * sizeof(mid[0])); + + matmul_q8_0_pair_batch(gate, up, model, + layer->ffn_gate_shexp, + layer->ffn_up_shexp, + x, + n_tok); + + swiglu_batch_ctx swiglu_ctx = { + .mid = mid, + .gate = gate, + .up = up, + .n = hidden, + .clamp = DS4_SWIGLU_CLAMP_EXP, + }; + ds4_parallel_for(n_tok, swiglu_batch_worker, &swiglu_ctx); + + matmul_q8_0_batch(out, model, layer->ffn_down_shexp, mid, n_tok); + + free(mid); + free(up); + free(gate); +} + +/* Early DS4 layers use token-id hash routing instead of top-k routing. */ +static void layer_hash_selected_experts( + int selected[DS4_MAX_EXPERT_USED], + const ds4_model *model, + const ds4_layer_weights *layer, + int token) { + ds4_tensor *t = layer->ffn_gate_tid2eid; + if (!t) ds4_die("hash routing table is missing for this layer"); + if (t->type != 26 || t->ndim != 2 || t->dim[0] != DS4_N_EXPERT_USED) { + ds4_die("ffn_gate_tid2eid.weight has an unexpected layout"); + } + if (token < 0 || (uint64_t)token >= t->dim[1]) { + ds4_die("token id is outside the hash routing table"); + } + + const int32_t *table = tensor_data(model, t); + const int32_t *row = table + (uint64_t)token * DS4_N_EXPERT_USED; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) selected[i] = row[i]; +} + +/* Router scores use sqrt(softplus(logit)); normalization happens only after + * the six selected experts are known. */ +static void layer_router_probs_one( + float probs[DS4_MAX_EXPERT], + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + float logits[DS4_MAX_EXPERT]; + + matvec_any(logits, model, layer->ffn_gate_inp, x); + for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { + probs[i] = sqrtf(softplus_stable(logits[i])); + } +} + +static void layer_hash_router_weights_from_probs( + float weights_out[DS4_MAX_EXPERT_USED], + const float probs[DS4_MAX_EXPERT], + const int selected[DS4_MAX_EXPERT_USED]) { + float sum = 0.0f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) ds4_die("hash-selected expert is outside router range"); + weights_out[i] = probs[selected[i]]; + sum += weights_out[i]; + } + + if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + weights_out[i] = weights_out[i] / sum * DS4_EXPERT_WEIGHT_SCALE; + } +} + +static void layer_hash_router_weights_one( + float weights_out[DS4_MAX_EXPERT_USED], + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + const int selected[DS4_MAX_EXPERT_USED]) { + float probs[DS4_MAX_EXPERT]; + + layer_router_probs_one(probs, model, layer, x); + layer_hash_router_weights_from_probs(weights_out, probs, selected); +} + +static void topk_desc(const float *score, int n, int k, int *idx) { + for (int i = 0; i < k; i++) idx[i] = -1; + + for (int i = 0; i < n; i++) { + for (int j = 0; j < k; j++) { + if (idx[j] < 0 || score[i] > score[idx[j]]) { + for (int m = k - 1; m > j; m--) idx[m] = idx[m - 1]; + idx[j] = i; + break; + } + } + } +} + +/* Later layers choose the six experts by biased top-k, but weight them using + * the unbiased router probabilities. */ +static void layer_topk_selected_experts_from_probs( + int selected[DS4_MAX_EXPERT_USED], + float expert_weight[DS4_MAX_EXPERT_USED], + const ds4_model *model, + const ds4_layer_weights *layer, + const float probs[DS4_MAX_EXPERT]); + +static void layer_topk_selected_experts( + int selected[DS4_MAX_EXPERT_USED], + float expert_weight[DS4_MAX_EXPERT_USED], + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x) { + float probs[DS4_MAX_EXPERT] = {0}; + + layer_router_probs_one(probs, model, layer, x); + layer_topk_selected_experts_from_probs(selected, expert_weight, model, layer, probs); +} + +static void layer_topk_selected_experts_from_probs( + int selected[DS4_MAX_EXPERT_USED], + float expert_weight[DS4_MAX_EXPERT_USED], + const ds4_model *model, + const ds4_layer_weights *layer, + const float probs[DS4_MAX_EXPERT]) { + float selection[DS4_MAX_EXPERT]; + + memcpy(selection, probs, sizeof(selection)); + + if (layer->ffn_exp_probs_b) { + const float *bias = tensor_data(model, layer->ffn_exp_probs_b); + for (uint32_t i = 0; i < DS4_N_EXPERT; i++) selection[i] += bias[i]; + } + + topk_desc(selection, (int)DS4_N_EXPERT, (int)DS4_N_EXPERT_USED, selected); + + float sum = 0.0f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + expert_weight[i] = probs[selected[i]]; + sum += expert_weight[i]; + } + if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + expert_weight[i] = expert_weight[i] / sum * DS4_EXPERT_WEIGHT_SCALE; + } +} + +static void print_vec_stats(const char *name, const float *x, uint64_t n); + +/* Single-token routed MoE. It selects six experts, runs IQ2_XXS gate/up, + * applies SwiGLU and router weights, then accumulates Q2_K down projections. */ +static void layer_routed_moe_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il, + int token, + float clamp, + bool trace) { + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + float *gate = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])) : NULL; + float *up = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])) : NULL; + float *mid = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])) : NULL; + float *mid_all = trace ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid_all[0])); + float *down = trace ? xmalloc((size_t)DS4_N_EMBD * sizeof(down[0])) : NULL; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const bool routed_q8_k = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; + if (routed_q8_0) { + if (trace) ds4_die("Q8_0 routed trace mode is not supported"); + if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); + if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { + ds4_die("Q8_0 expert input has an unexpected layout"); + } + } else { + if (routed_q8_k && trace) ds4_die("Q8_K routed trace mode is not supported"); + if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); + if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { + ds4_die("routed expert down input has an unexpected layout"); + } + } + block_q8_K *xq = routed_q8_0 ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0])); + block_q8_K *midq = (trace || routed_q8_0) ? NULL : + xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0])); + + memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); + if (!routed_q8_0) { + ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); + } + + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(selected, model, layer, token); + layer_hash_router_weights_one(expert_weight, model, layer, x, selected); + } else { + layer_topk_selected_experts(selected, expert_weight, model, layer, x); + } + + if (routed_q8_0) { + const uint64_t x_blocks = expert_in_dim / 32u; + int8_t *xq8 = xmalloc((size_t)x_blocks * 32u); + float *xscale8 = xmalloc((size_t)x_blocks * sizeof(float)); + quantize_q8_0_activation(x, xq8, xscale8, expert_in_dim); + matvec_q8_0_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq8, xscale8, selected, + expert_weight, + DS4_N_EXPERT_USED, clamp); + + const uint64_t mid_blocks = down_in_dim / 32u; + int8_t *midq8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * 32u); + float *midscale8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * sizeof(float)); + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, + midq8 + (uint64_t)i * mid_blocks * 32u, + midscale8 + (uint64_t)i * mid_blocks, + down_in_dim); + } + matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, + midq8, midscale8, selected, + DS4_N_EXPERT_USED); + free(midscale8); + free(midq8); + free(xscale8); + free(xq8); + } else if (routed_q8_k) { + matvec_q8_k_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, selected, expert_weight, + DS4_N_EXPERT_USED, clamp); + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, + midq + (uint64_t)i * (down_in_dim / QK_K), + (int64_t)down_in_dim); + } + matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, + midq, selected, DS4_N_EXPERT_USED); + } else if (!trace) { + matvec_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, + selected, + expert_weight, + DS4_N_EXPERT_USED, + clamp); + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, + midq + (uint64_t)i * (down_in_dim / QK_K), + (int64_t)down_in_dim); + } + matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); + } else { + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + const uint32_t expert = (uint32_t)selected[i]; + + matvec_expert_pair_prequant(gate, up, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, + expert); + char name[64]; + snprintf(name, sizeof(name), "blk.%u expert %u gate", il, expert); + print_vec_stats(name, gate, DS4_N_FF_EXP); + snprintf(name, sizeof(name), "blk.%u expert %u up", il, expert); + print_vec_stats(name, up, DS4_N_FF_EXP); + + /* + * DeepSeek V4 clamps routed expert gate/up values before SwiGLU and + * applies the router weight before the down projection. + */ + const float limit = clamp; + for (uint32_t j = 0; j < DS4_N_FF_EXP; j++) { + if (limit > 1.0e-6f) { + if (gate[j] > limit) gate[j] = limit; + if (up[j] > limit) up[j] = limit; + if (up[j] < -limit) up[j] = -limit; + } + mid[j] = silu(gate[j]) * up[j] * expert_weight[i]; + } + + snprintf(name, sizeof(name), "blk.%u expert %u mid", il, expert); + print_vec_stats(name, mid, DS4_N_FF_EXP); + + matvec_expert_down(down, model, layer->ffn_down_exps, mid, expert); + snprintf(name, sizeof(name), "blk.%u expert %u down", il, expert); + print_vec_stats(name, down, DS4_N_EMBD); + for (uint32_t j = 0; j < DS4_N_EMBD; j++) out[j] += down[j]; + } + } + + free(midq); + free(xq); + free(down); + free(mid_all); + free(mid); + free(up); + free(gate); +} + +/* Decode version of routed MoE: same math as layer_routed_moe_one(), but all + * large temporaries come from the persistent scratch arena. */ +static void layer_routed_moe_one_prealloc( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il, + int token, + float clamp, + float * mid_all, + block_q8_K * xq, + block_q8_K * midq, + int8_t * q8_xq, + float * q8_xscale, + int8_t * q8_midq, + float * q8_midscale) { + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const bool routed_q8_k = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; + + if (routed_q8_0) { + if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); + if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { + ds4_die("Q8_0 expert input has an unexpected layout"); + } + } else { + if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); + if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { + ds4_die("routed expert down input has an unexpected layout"); + } + } + + memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); + + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(selected, model, layer, token); + layer_hash_router_weights_one(expert_weight, model, layer, x, selected); + } else { + layer_topk_selected_experts(selected, expert_weight, model, layer, x); + } + + if (routed_q8_0) { + if (!q8_xq || !q8_xscale || !q8_midq || !q8_midscale) { + ds4_die("missing Q8_0 routed decode scratch"); + } + quantize_q8_0_activation(x, q8_xq, q8_xscale, expert_in_dim); + matvec_q8_0_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + q8_xq, q8_xscale, selected, + expert_weight, + DS4_N_EXPERT_USED, clamp); + const uint64_t mid_blocks = down_in_dim / 32u; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, + q8_midq + (uint64_t)i * mid_blocks * 32u, + q8_midscale + (uint64_t)i * mid_blocks, + down_in_dim); + } + matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, + q8_midq, q8_midscale, selected, + DS4_N_EXPERT_USED); + (void)il; + return; + } + + if (!mid_all || !xq || !midq) ds4_die("missing routed decode scratch"); + ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); + + if (routed_q8_k) { + matvec_q8_k_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, selected, expert_weight, + DS4_N_EXPERT_USED, clamp); + } else { + matvec_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, selected, expert_weight, + DS4_N_EXPERT_USED, clamp); + } + + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, + midq + (uint64_t)i * (down_in_dim / QK_K), + (int64_t)down_in_dim); + } + if (routed_q8_k) { + matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, + midq, selected, DS4_N_EXPERT_USED); + } else { + matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, + midq, selected, DS4_N_EXPERT_USED); + } + + (void)il; +} + +/* Prefill MoE groups token/expert pairs by expert so each active expert's + * rows are scanned once for the whole token batch. */ +static void layer_routed_moe_batch( + float * moe, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * norm, + const int * token_ids, + uint32_t n_tok, + uint32_t il, + float clamp) { + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t expert_out_dim = layer->ffn_gate_exps->dim[1]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const uint64_t down_out_dim = layer->ffn_down_exps->dim[1]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const bool routed_q8_k = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; + if (routed_q8_0) { + if ((expert_in_dim % 32u) != 0 || (down_in_dim % 32u) != 0) { + ds4_die("Q8_0 routed expert input is not QK8_0 aligned"); + } + } else { + if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); + if (down_in_dim % QK_K != 0) ds4_die("routed expert down input is not QK_K aligned"); + } + if (expert_out_dim != down_in_dim || down_out_dim != DS4_N_EMBD) { + ds4_die("routed expert tensor layout is unexpected"); + } + + const uint32_t total_pairs = n_tok * DS4_N_EXPERT_USED; + uint32_t counts[DS4_MAX_EXPERT + 1] = {0}; + uint32_t cursor[DS4_MAX_EXPERT] = {0}; + uint32_t active_expert[DS4_MAX_EXPERT]; + uint32_t n_active = 0; + + int *selected = xmalloc((size_t)total_pairs * sizeof(selected[0])); + float *pair_weight = xmalloc((size_t)total_pairs * sizeof(pair_weight[0])); + ds4_expert_pair *pairs = xmalloc((size_t)total_pairs * sizeof(pairs[0])); + + for (uint32_t t = 0; t < n_tok; t++) { + int sel[DS4_MAX_EXPERT_USED]; + float weights[DS4_MAX_EXPERT_USED]; + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(sel, model, layer, token_ids[t]); + layer_hash_router_weights_one(weights, model, layer, norm + (uint64_t)t * expert_in_dim, sel); + } else { + layer_topk_selected_experts(sel, weights, model, layer, norm + (uint64_t)t * expert_in_dim); + } + + for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { + const uint32_t pair_id = t * DS4_N_EXPERT_USED + slot; + selected[pair_id] = sel[slot]; + pair_weight[pair_id] = weights[slot]; + pairs[pair_id] = (ds4_expert_pair){ .token = t, .slot = slot }; + if (sel[slot] < 0 || (uint32_t)sel[slot] >= DS4_N_EXPERT) ds4_die("selected expert is outside range"); + counts[(uint32_t)sel[slot] + 1]++; + } + } + + for (uint32_t e = 0; e < DS4_N_EXPERT; e++) { + counts[e + 1] += counts[e]; + cursor[e] = counts[e]; + if (counts[e + 1] != counts[e]) active_expert[n_active++] = e; + } + + uint32_t *pair_ids = xmalloc((size_t)total_pairs * sizeof(pair_ids[0])); + for (uint32_t p = 0; p < total_pairs; p++) { + const uint32_t e = (uint32_t)selected[p]; + pair_ids[cursor[e]++] = p; + } + + if (routed_q8_0) { + const uint64_t x_blocks = expert_in_dim / 32u; + int8_t *xq8 = xmalloc((size_t)n_tok * x_blocks * 32u); + float *xscale8 = xmalloc((size_t)n_tok * x_blocks * sizeof(float)); + for (uint32_t t = 0; t < n_tok; t++) { + quantize_q8_0_activation(norm + (uint64_t)t * expert_in_dim, + xq8 + (uint64_t)t * x_blocks * 32u, + xscale8 + (uint64_t)t * x_blocks, + expert_in_dim); + } + + float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); + matvec_q8_0_batch_mid_ctx mid_ctx = { + .mid = mid, + .xq = xq8, + .xscale = xscale8, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .pair_weight = pair_weight, + .clamp = clamp, + .in_dim = expert_in_dim, + .out_dim = expert_out_dim, + .blocks = x_blocks, + }; + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, + &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); + mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, + &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); + if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || + gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { + ds4_die("Q8_0 batch expert tensor layout mismatch"); + } + } + ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_0_batch_mid_worker, &mid_ctx); + + const uint64_t mid_blocks = down_in_dim / 32u; + int8_t *midq8 = xmalloc((size_t)total_pairs * mid_blocks * 32u); + float *midscale8 = xmalloc((size_t)total_pairs * mid_blocks * sizeof(float)); + for (uint32_t p = 0; p < total_pairs; p++) { + quantize_q8_0_activation(mid + (uint64_t)p * down_in_dim, + midq8 + (uint64_t)p * mid_blocks * 32u, + midscale8 + (uint64_t)p * mid_blocks, + down_in_dim); + } + free(mid); + + matvec_q8_0_batch_accum_rows_ctx down_ctx = { + .moe = moe, + .midq = midq8, + .midscale = midscale8, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .n_active = n_active, + .n_tok = n_tok, + .in_dim = down_in_dim, + .out_dim = down_out_dim, + .blocks = mid_blocks, + }; + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t in_dim, out_dim; + down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, + &in_dim, &out_dim, &down_ctx.row_bytes[e]); + if (in_dim != down_in_dim || out_dim != down_out_dim) { + ds4_die("Q8_0 batch down expert tensor layout mismatch"); + } + } + ds4_parallel_for(down_out_dim, matvec_q8_0_batch_accum_rows_worker, &down_ctx); + + free(midscale8); + free(midq8); + free(xscale8); + free(xq8); + free(pair_ids); + free(pairs); + free(pair_weight); + free(selected); + (void)il; + return; + } + + if (routed_q8_k) { + const uint64_t xq_blocks = expert_in_dim / QK_K; + block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); + for (uint32_t t = 0; t < n_tok; t++) { + ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, + xq + (uint64_t)t * xq_blocks, + (int64_t)expert_in_dim); + } + + float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); + + matvec_q8_k_batch_mid_ctx mid_ctx = { + .mid = mid, + .xq = xq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .pair_weight = pair_weight, + .clamp = clamp, + .in_dim = expert_in_dim, + .out_dim = expert_out_dim, + .xq_blocks = xq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, + &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); + mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, + &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); + if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || + gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { + ds4_die("Q8_K batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_k_batch_mid_worker, &mid_ctx); + + const uint64_t midq_blocks = down_in_dim / QK_K; + block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); + quantize_mid_pairs_ctx quant_ctx = { + .mid = mid, + .midq = midq, + .down_in_dim = down_in_dim, + .down_blocks = midq_blocks, + }; + ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); + free(mid); + + matvec_q8_k_batch_accum_rows_ctx down_ctx = { + .moe = moe, + .midq = midq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .n_active = n_active, + .n_tok = n_tok, + .in_dim = down_in_dim, + .out_dim = down_out_dim, + .midq_blocks = midq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t in_dim, out_dim; + down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, + &in_dim, &out_dim, &down_ctx.row_bytes[e]); + if (in_dim != down_in_dim || out_dim != down_out_dim) { + ds4_die("Q8_K batch down expert tensor layout mismatch"); + } + } + + ds4_parallel_for(down_out_dim, matvec_q8_k_batch_accum_rows_worker, &down_ctx); + + free(midq); + free(pair_ids); + free(xq); + free(pairs); + free(pair_weight); + free(selected); + + (void)il; + return; + } + + const uint64_t xq_blocks = expert_in_dim / QK_K; + block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); + for (uint32_t t = 0; t < n_tok; t++) { + ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, + xq + (uint64_t)t * xq_blocks, + (int64_t)expert_in_dim); + } + + float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); + + const uint32_t gate_type = layer->ffn_gate_exps->type; + + /* Build mid vectors: dispatch based on gate/up tensor type. */ + if (gate_type == DS4_TENSOR_IQ2_XXS) { + matvec_iq2_xxs_batch_mid_ctx mid_ctx = { + .mid = mid, + .xq = xq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .pair_weight = pair_weight, + .clamp = clamp, + .in_dim = expert_in_dim, + .out_dim = expert_out_dim, + .xq_blocks = xq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, + &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); + mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, + &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); + if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || + gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_iq2_xxs_batch_mid_worker, &mid_ctx); + } else if (gate_type == DS4_TENSOR_Q2_K) { + matvec_q2_k_batch_mid_ctx mid_ctx = { + .mid = mid, + .xq = xq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .pair_weight = pair_weight, + .clamp = clamp, + .in_dim = expert_in_dim, + .out_dim = expert_out_dim, + .xq_blocks = xq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, + &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); + mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, + &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); + if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || + gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q2_k_batch_mid_worker, &mid_ctx); + } else if (gate_type == DS4_TENSOR_Q4_K) { + matvec_q4_k_batch_mid_ctx mid_ctx = { + .mid = mid, + .xq = xq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .pair_weight = pair_weight, + .clamp = clamp, + .in_dim = expert_in_dim, + .out_dim = expert_out_dim, + .xq_blocks = xq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t gate_in_dim, gate_out_dim; + uint64_t up_in_dim, up_out_dim; + mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, + &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); + mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, + &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); + if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || + gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q4_k_batch_mid_worker, &mid_ctx); + } else { + ds4_die("unsupported gate/up expert tensor type for batch"); + } + + const uint64_t midq_blocks = down_in_dim / QK_K; + block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); + quantize_mid_pairs_ctx quant_ctx = { + .mid = mid, + .midq = midq, + .down_in_dim = down_in_dim, + .down_blocks = midq_blocks, + }; + ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); + free(mid); + + /* Down projection: dispatch based on down tensor type. */ + const uint32_t down_type = layer->ffn_down_exps->type; + + if (down_type == DS4_TENSOR_IQ2_XXS) { + matvec_iq2_xxs_batch_accum_rows_ctx down_ctx = { + .moe = moe, + .midq = midq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .n_active = n_active, + .n_tok = n_tok, + .in_dim = down_in_dim, + .out_dim = down_out_dim, + .midq_blocks = midq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t in_dim, out_dim; + down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, + &in_dim, &out_dim, &down_ctx.row_bytes[e]); + if (in_dim != down_in_dim || out_dim != down_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for(down_out_dim, matvec_iq2_xxs_batch_accum_rows_worker, &down_ctx); + } else if (down_type == DS4_TENSOR_Q2_K) { + matvec_q2_k_batch_accum_rows_ctx down_ctx = { + .moe = moe, + .midq = midq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .n_active = n_active, + .n_tok = n_tok, + .in_dim = down_in_dim, + .out_dim = down_out_dim, + .midq_blocks = midq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t in_dim, out_dim; + down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, + &in_dim, &out_dim, &down_ctx.row_bytes[e]); + if (in_dim != down_in_dim || out_dim != down_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for(down_out_dim, matvec_q2_k_batch_accum_rows_worker, &down_ctx); + } else if (down_type == DS4_TENSOR_Q4_K) { + matvec_q4_k_batch_accum_rows_ctx down_ctx = { + .moe = moe, + .midq = midq, + .pairs = pairs, + .pair_ids = pair_ids, + .expert_offset = counts, + .active_expert = active_expert, + .n_active = n_active, + .n_tok = n_tok, + .in_dim = down_in_dim, + .out_dim = down_out_dim, + .midq_blocks = midq_blocks, + }; + + for (uint32_t ai = 0; ai < n_active; ai++) { + const uint32_t e = active_expert[ai]; + uint64_t in_dim, out_dim; + down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, + &in_dim, &out_dim, &down_ctx.row_bytes[e]); + if (in_dim != down_in_dim || out_dim != down_out_dim) { + ds4_die("batch expert tensor layout mismatch"); + } + } + + ds4_parallel_for(down_out_dim, matvec_q4_k_batch_accum_rows_worker, &down_ctx); + } else { + ds4_die("unsupported down expert tensor type for batch"); + } + + free(midq); + free(pair_ids); + free(xq); + free(pairs); + free(pair_weight); + free(selected); + + (void)il; +} + +static void print_vec_stats(const char *name, const float *x, uint64_t n); + +/* Full FFN sublayer for one token: HC pre, RMSNorm, routed MoE, shared expert, + * sum, and HC post. */ +static void layer_ffn_one( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + uint32_t il, + int token, + const float * steering_dirs, + float steering_scale, + bool trace) { + const uint32_t n_hc = DS4_N_HC; + const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; + const double t_start = profile ? now_sec() : 0.0; + double t_hc = 0.0; + double t_norm = 0.0; + double t_routed = 0.0; + double t_shared = 0.0; + double t_post = 0.0; + float *ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_cur[0])); + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); + float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); + float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); + float post[4]; + float comb[16]; + + double t0 = profile ? now_sec() : 0.0; + hc_pre_from_state_one(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + inp_hc, ffn_cur, post, comb); + if (profile) t_hc = now_sec() - t0; + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u ffn_cur", il); + print_vec_stats(name, ffn_cur, DS4_N_EMBD); + } + + t0 = profile ? now_sec() : 0.0; + const float *ffn_norm = tensor_data(model, layer->ffn_norm); + rms_norm_weight(norm, ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); + if (profile) t_norm = now_sec() - t0; + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u ffn_norm", il); + print_vec_stats(name, norm, DS4_N_EMBD); + } + + t0 = profile ? now_sec() : 0.0; + layer_routed_moe_one(moe, model, layer, norm, il, token, DS4_SWIGLU_CLAMP_EXP, trace); + if (profile) t_routed = now_sec() - t0; + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u routed_moe", il); + print_vec_stats(name, moe, DS4_N_EMBD); + } + t0 = profile ? now_sec() : 0.0; + layer_shared_ffn_one(shared, model, layer, norm); + if (profile) t_shared = now_sec() - t0; + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u shared_ffn", il); + print_vec_stats(name, shared, DS4_N_EMBD); + } + + t0 = profile ? now_sec() : 0.0; + for (uint32_t i = 0; i < DS4_N_EMBD; i++) { + ffn_out[i] = moe[i] + shared[i]; + } + cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, 1, steering_scale); + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u ffn_out", il); + print_vec_stats(name, ffn_out, DS4_N_EMBD); + } + + hc_post_one(out_hc, ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); + if (profile) t_post = now_sec() - t0; + if (trace) { + char name[64]; + snprintf(name, sizeof(name), "blk.%u ffn_post_hc", il); + print_vec_stats(name, out_hc, (uint64_t)n_hc * DS4_N_EMBD); + } + + if (profile) { + fprintf(stderr, + "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", + il, + t_hc * 1000.0, + t_norm * 1000.0, + t_routed * 1000.0, + t_shared * 1000.0, + t_post * 1000.0, + (now_sec() - t_start) * 1000.0); + } + + free(ffn_out); + free(shared); + free(moe); + free(norm); + free(ffn_cur); +} + +/* Allocation-free decode FFN using the persistent CPU scratch buffers. */ +static void layer_ffn_one_decode_scratch( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + uint32_t il, + int token, + const float * steering_dirs, + float steering_scale, + ds4_cpu_decode_scratch * scratch) { + const uint32_t n_hc = DS4_N_HC; + const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; + const double t_start = profile ? now_sec() : 0.0; + double t_hc = 0.0; + double t_norm = 0.0; + double t_routed = 0.0; + double t_shared = 0.0; + double t_post = 0.0; + float post[4]; + float comb[16]; + + double t0 = profile ? now_sec() : 0.0; + hc_pre_from_state_one_scratch(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + inp_hc, scratch->ffn_cur, post, comb, + scratch->hc_flat, + false); + if (profile) t_hc = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + const float *ffn_norm = tensor_data(model, layer->ffn_norm); + rms_norm_weight(scratch->ffn_norm, scratch->ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); + if (profile) t_norm = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_routed_moe_one_prealloc(scratch->ffn_moe, + model, + layer, + scratch->ffn_norm, + il, + token, + DS4_SWIGLU_CLAMP_EXP, + scratch->routed_mid_all, + scratch->routed_xq, + scratch->routed_midq, + scratch->routed_q8_xq, + scratch->routed_q8_xscale, + scratch->routed_q8_midq, + scratch->routed_q8_midscale); + if (profile) t_routed = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_shared_ffn_one_decode_scratch(scratch->ffn_shared, model, layer, scratch->ffn_norm, scratch); + if (profile) t_shared = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + for (uint32_t i = 0; i < DS4_N_EMBD; i++) { + scratch->ffn_out[i] = scratch->ffn_moe[i] + scratch->ffn_shared[i]; + } + cpu_directional_steering_project_rows(scratch->ffn_out, steering_dirs, il, 1, steering_scale); + hc_post_one(out_hc, scratch->ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); + if (profile) t_post = now_sec() - t0; + + if (profile) { + fprintf(stderr, + "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", + il, + t_hc * 1000.0, + t_norm * 1000.0, + t_routed * 1000.0, + t_shared * 1000.0, + t_post * 1000.0, + (now_sec() - t_start) * 1000.0); + } +} + +static void layer_ffn_batch( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + const int * token_ids, + uint32_t n_tok, + uint32_t il, + const float * steering_dirs, + float steering_scale) { + if (n_tok == 0) return; + const uint32_t n_hc = DS4_N_HC; + const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; + float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); + float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); + float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); + float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); + float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); + float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); + const float *ffn_norm = tensor_data(model, layer->ffn_norm); + + for (uint32_t t = 0; t < n_tok; t++) { + hc_pre_from_state_one(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + inp_hc + (uint64_t)t * hc_dim, + ffn_cur + (uint64_t)t * DS4_N_EMBD, + post + (uint64_t)t * n_hc, + comb + (uint64_t)t * n_hc * n_hc); + rms_norm_weight(norm + (uint64_t)t * DS4_N_EMBD, + ffn_cur + (uint64_t)t * DS4_N_EMBD, + ffn_norm, + DS4_N_EMBD, + DS4_RMS_EPS); + } + + layer_routed_moe_batch(moe, model, layer, norm, token_ids, n_tok, il, DS4_SWIGLU_CLAMP_EXP); + layer_shared_ffn_batch(shared, model, layer, norm, n_tok); + + if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { + float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); + for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { + ffn_out[i] = moe[i] + shared[i]; + } + cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); + hc_post_batch(out_hc, + ffn_out, + inp_hc, + post, + comb, + n_tok, + DS4_N_EMBD, + n_hc); + free(ffn_out); + } else { + hc_post_sum_batch(out_hc, + moe, + shared, + inp_hc, + post, + comb, + n_tok, + DS4_N_EMBD, + n_hc); + } + + free(comb); + free(post); + free(shared); + free(moe); + free(norm); + free(ffn_cur); +} + +typedef struct { + float *moe; + const ds4_model *model; + const ds4_layer_weights *layer; + const float *norm; + const int *token_ids; + uint64_t expert_in_dim; + uint64_t down_in_dim; + uint32_t il; + bool routed_q8_0; +} routed_moe_tokens_ctx; + +static void routed_moe_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { + routed_moe_tokens_ctx *ctx = vctx; + const uint64_t q8_x_blocks = ctx->expert_in_dim / 32u; + const uint64_t q8_mid_blocks = ctx->down_in_dim / 32u; + float *routed_mid = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); + block_q8_K *routed_xq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)(ctx->expert_in_dim / QK_K) * sizeof(routed_xq[0])); + block_q8_K *routed_midq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (ctx->down_in_dim / QK_K) * sizeof(routed_midq[0])); + int8_t *routed_q8_xq = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * 32u) : NULL; + float *routed_q8_xscale = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; + int8_t *routed_q8_midq = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * 32u) : NULL; + float *routed_q8_midscale = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; + + for (uint64_t t = t0; t < t1; t++) { + layer_routed_moe_one_prealloc(ctx->moe + t * DS4_N_EMBD, + ctx->model, + ctx->layer, + ctx->norm + t * DS4_N_EMBD, + ctx->il, + ctx->token_ids[t], + DS4_SWIGLU_CLAMP_EXP, + routed_mid, + routed_xq, + routed_midq, + routed_q8_xq, + routed_q8_xscale, + routed_q8_midq, + routed_q8_midscale); + } + + free(routed_q8_midscale); + free(routed_q8_midq); + free(routed_q8_xscale); + free(routed_q8_xq); + free(routed_midq); + free(routed_xq); + free(routed_mid); +} + +static void layer_routed_moe_tokens_parallel( + float * moe, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * norm, + const int * token_ids, + uint32_t n_tok, + uint32_t il) { + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const bool routed_q8_k = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; + if (routed_q8_k) { + if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); + if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { + ds4_die("Q8_K expert input has an unexpected layout"); + } + } + routed_moe_tokens_ctx ctx = { + .moe = moe, + .model = model, + .layer = layer, + .norm = norm, + .token_ids = token_ids, + .expert_in_dim = expert_in_dim, + .down_in_dim = down_in_dim, + .il = il, + .routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0, + }; + ds4_parallel_for_min_rows(n_tok, routed_moe_tokens_worker, &ctx, 1); +} + +/* Default prefill FFN path. HC and shared expert are batched, while routed + * experts can run either token-parallel or expert-grouped depending on size. */ +static void layer_ffn_shared_batch( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + const int * token_ids, + uint32_t n_tok, + uint32_t il, + const float * steering_dirs, + float steering_scale) { + const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; + const double t_start = profile ? now_sec() : 0.0; + double t_hc_norm = 0.0; + double t_routed = 0.0; + double t_shared = 0.0; + double t_post = 0.0; + const uint32_t n_hc = DS4_N_HC; + float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); + float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); + float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); + float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); + float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); + float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const bool routed_q8_k = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; + if (routed_q8_k) { + if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); + if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { + ds4_die("Q8_K expert input has an unexpected layout"); + } + } + const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; + const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; + const bool routed_token_parallel = + getenv("DS4_ROUTED_TOKEN_PARALLEL") != NULL || + (getenv("DS4_NO_ROUTED_TOKEN_PARALLEL") == NULL && n_tok >= 64); + float *routed_mid = routed_token_parallel ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); + block_q8_K *routed_xq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(routed_xq[0])); + block_q8_K *routed_midq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(routed_midq[0])); + int8_t *routed_q8_xq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; + float *routed_q8_xscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; + int8_t *routed_q8_midq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; + float *routed_q8_midscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; + + double t0 = profile ? now_sec() : 0.0; + hc_pre_norm_batch(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + layer->ffn_norm, + inp_hc, + NULL, + ffn_cur, + norm, + post, + comb, + n_tok); + if (profile) t_hc_norm = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + if (routed_token_parallel) { + layer_routed_moe_tokens_parallel(moe, model, layer, norm, token_ids, n_tok, il); + } else { + for (uint32_t t = 0; t < n_tok; t++) { + layer_routed_moe_one_prealloc(moe + (uint64_t)t * DS4_N_EMBD, + model, + layer, + norm + (uint64_t)t * DS4_N_EMBD, + il, + token_ids[t], + DS4_SWIGLU_CLAMP_EXP, + routed_mid, + routed_xq, + routed_midq, + routed_q8_xq, + routed_q8_xscale, + routed_q8_midq, + routed_q8_midscale); + } + } + if (profile) t_routed = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_shared_ffn_batch(shared, model, layer, norm, n_tok); + if (profile) t_shared = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { + float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); + for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { + ffn_out[i] = moe[i] + shared[i]; + } + cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); + hc_post_batch(out_hc, + ffn_out, + inp_hc, + post, + comb, + n_tok, + DS4_N_EMBD, + n_hc); + free(ffn_out); + } else { + hc_post_sum_batch(out_hc, + moe, + shared, + inp_hc, + post, + comb, + n_tok, + DS4_N_EMBD, + n_hc); + } + if (profile) t_post = now_sec() - t0; + + if (profile) { + fprintf(stderr, + "ds4: prefill detail layer %u ffn hc_norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f\n", + il, t_hc_norm, t_routed, t_shared, t_post, now_sec() - t_start); + } + + free(comb); + free(post); + free(routed_q8_midscale); + free(routed_q8_midq); + free(routed_q8_xscale); + free(routed_q8_xq); + free(routed_midq); + free(routed_xq); + free(routed_mid); + free(shared); + free(moe); + free(norm); + free(ffn_cur); +} + +typedef struct { + float *out_hc; + const ds4_model *model; + const ds4_layer_weights *layer; + const float *inp_hc; + const int *token_ids; + const float *steering_dirs; + float steering_scale; + uint64_t hc_dim; + uint32_t il; +} layer_ffn_tokens_ctx; + +static void layer_ffn_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { + layer_ffn_tokens_ctx *ctx = vctx; + for (uint64_t t = t0; t < t1; t++) { + layer_ffn_one(ctx->out_hc + t * ctx->hc_dim, + ctx->model, + ctx->layer, + ctx->inp_hc + t * ctx->hc_dim, + ctx->il, + ctx->token_ids[t], + ctx->steering_dirs, + ctx->steering_scale, + false); + } +} + +static void layer_ffn_tokens_parallel( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + const int * token_ids, + uint32_t n_tok, + uint32_t il, + const float * steering_dirs, + float steering_scale) { + layer_ffn_tokens_ctx ctx = { + .out_hc = out_hc, + .model = model, + .layer = layer, + .inp_hc = inp_hc, + .token_ids = token_ids, + .steering_dirs = steering_dirs, + .steering_scale = steering_scale, + .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, + .il = il, + }; + ds4_parallel_for(n_tok, layer_ffn_tokens_worker, &ctx); +} + +static void output_logits_one( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + const float * inp_hc); + +/* ========================================================================= + * KV Cache, Compressors, and CPU Layer Execution. + * ========================================================================= + * + * The CPU path is the correctness reference. It maintains raw SWA KV rows, + * optional compressed KV rows, the indexer mask for ratio-4 layers, and a + * reusable decode scratch arena so token generation does not allocate in the + * hot loop. + */ + +typedef struct { + float *raw_kv; + uint32_t n_raw; + uint32_t cap_raw; + + uint32_t compress_ratio; + uint32_t comp_cap; + uint32_t n_comp; + float *attn_comp_kv; + float *attn_state_kv; + float *attn_state_score; + + uint32_t n_index_comp; + float *index_comp_kv; + float *index_state_kv; + float *index_state_score; +} ds4_layer_cache; + +typedef struct { + ds4_layer_cache layer[DS4_MAX_LAYER]; + uint32_t head_dim; +} ds4_kv_cache; + +static uint32_t ds4_default_raw_cap(uint32_t ctx_size) { + uint32_t raw_cap = DS4_N_SWA; + if (raw_cap > ctx_size) raw_cap = ctx_size; + if (raw_cap == 0) raw_cap = 1; + return raw_cap; +} + +#define DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK 2048u + +static uint32_t ds4_effective_prefill_chunk(bool cuda_tensor_parallel, + uint32_t requested_chunk) { + if (requested_chunk != 0) return requested_chunk; + return cuda_tensor_parallel ? DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK : 0; +} + +static uint32_t ds4_prefill_cap_for_prompt(int prompt_len, + uint32_t requested_chunk) { + if (prompt_len <= 0) return 1; + uint32_t cap = (uint32_t)prompt_len; + + if (requested_chunk != 0) { + cap = requested_chunk; + } else { + const char *env = getenv("DS4_METAL_PREFILL_CHUNK"); + if (env && env[0]) { + char *endp = NULL; + const long v = strtol(env, &endp, 10); + if (endp != env) { + if (v <= 0) return cap; + cap = (uint32_t)v; + } + } else if (prompt_len > 4096) { + cap = DS4_MODEL_VARIANT == DS4_VARIANT_PRO ? 8192u : 4096u; + } + } + + if (cap == 0) cap = 1; + if (cap > (uint32_t)prompt_len) cap = (uint32_t)prompt_len; + return cap; +} + +/* Allocate all CPU decode temporaries once. This keeps generation deterministic + * from the VM's point of view and makes accidental hot-loop malloc visible. */ +static void cpu_decode_scratch_init(ds4_cpu_decode_scratch *scratch, uint32_t ctx_size) { + memset(scratch, 0, sizeof(*scratch)); + if (ctx_size == 0) ctx_size = 1; + const uint32_t raw_cap = ds4_default_raw_cap(ctx_size); + const uint32_t comp_cap = ctx_size / 4 + 2; + const uint32_t attn_score_cap = raw_cap + comp_cap; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t q8_cap = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t q8_blocks = (q8_cap + 31u) / 32u; + if ((DS4_N_EMBD % 32u) != 0 || (DS4_N_FF_EXP % 32u) != 0) { + ds4_die("Q8_0 routed decode scratch dimensions are not QK8_0 aligned"); + } + const uint64_t routed_q8_x_blocks = DS4_N_EMBD / 32u; + const uint64_t routed_q8_mid_blocks = DS4_N_FF_EXP / 32u; + + /* + * The CPU decode path used to malloc/free dozens of medium-sized buffers + * for every layer of every generated token. On macOS this can drive the VM + * system through repeated map/unmap bookkeeping while the huge model mmap is + * also being streamed, and we have observed kernel panics in VM accounting. + * Keep decode scratch resident for the whole generation instead. + */ + scratch->ctx_size = ctx_size; + scratch->comp_cap = comp_cap; + scratch->attn_score_cap = attn_score_cap; + scratch->q8_cap = (uint32_t)q8_cap; + + scratch->plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->cur = xmalloc((size_t)hc_dim * sizeof(float)); + scratch->next = xmalloc((size_t)hc_dim * sizeof(float)); + + scratch->attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->attn_residual = xmalloc((size_t)hc_dim * sizeof(float)); + scratch->q = xmalloc((size_t)q_dim * sizeof(float)); + scratch->qr = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); + scratch->qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); + scratch->kv_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + scratch->kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + scratch->heads = xmalloc((size_t)q_dim * sizeof(float)); + scratch->attn_low = xmalloc((size_t)DS4_N_OUT_GROUP * DS4_N_LORA_O * sizeof(float)); + scratch->attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + scratch->attn_score = xmalloc((size_t)attn_score_cap * sizeof(float)); + + scratch->comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + scratch->index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + scratch->comp_kv_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); + scratch->comp_sc_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); + scratch->comp_pooled = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + + scratch->index_allowed = xmalloc((size_t)comp_cap * sizeof(bool)); + scratch->index_q = xmalloc((size_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + scratch->index_weights = xmalloc((size_t)DS4_N_INDEXER_HEAD * sizeof(float)); + scratch->index_scores = xmalloc((size_t)comp_cap * sizeof(float)); + + scratch->ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->ffn_moe = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->ffn_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->shared_gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); + scratch->shared_up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); + scratch->shared_mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); + scratch->routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); + scratch->routed_xq = xmalloc((size_t)(DS4_N_EMBD / QK_K) * sizeof(block_q8_K)); + scratch->routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (DS4_N_FF_EXP / QK_K) * sizeof(block_q8_K)); + scratch->routed_q8_xq = xmalloc((size_t)routed_q8_x_blocks * 32u); + scratch->routed_q8_xscale = xmalloc((size_t)routed_q8_x_blocks * sizeof(float)); + scratch->routed_q8_midq = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u); + scratch->routed_q8_midscale = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(float)); + + scratch->q8_xq = xmalloc((size_t)q8_blocks * 32u); + scratch->q8_xscale = xmalloc((size_t)q8_blocks * sizeof(float)); + + scratch->hc_flat = xmalloc((size_t)hc_dim * sizeof(float)); + scratch->output_flat = xmalloc((size_t)hc_dim * sizeof(float)); + scratch->output_pre = xmalloc((size_t)DS4_N_HC * sizeof(float)); + scratch->output_weights = xmalloc((size_t)DS4_N_HC * sizeof(float)); + scratch->output_embd = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + scratch->output_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); +} + +static void cpu_decode_scratch_free(ds4_cpu_decode_scratch *scratch) { + if (!scratch) return; + free(scratch->output_norm); + free(scratch->output_embd); + free(scratch->output_weights); + free(scratch->output_pre); + free(scratch->output_flat); + free(scratch->hc_flat); + free(scratch->q8_xscale); + free(scratch->q8_xq); + free(scratch->routed_q8_midscale); + free(scratch->routed_q8_midq); + free(scratch->routed_q8_xscale); + free(scratch->routed_q8_xq); + free(scratch->routed_midq); + free(scratch->routed_xq); + free(scratch->routed_mid_all); + free(scratch->shared_mid); + free(scratch->shared_up); + free(scratch->shared_gate); + free(scratch->ffn_out); + free(scratch->ffn_shared); + free(scratch->ffn_moe); + free(scratch->ffn_norm); + free(scratch->ffn_cur); + free(scratch->index_scores); + free(scratch->index_weights); + free(scratch->index_q); + free(scratch->index_allowed); + free(scratch->comp_pooled); + free(scratch->comp_sc_cur); + free(scratch->comp_kv_cur); + free(scratch->index_comp); + free(scratch->comp); + free(scratch->attn_score); + free(scratch->after_attn_hc); + free(scratch->attn_out); + free(scratch->attn_low); + free(scratch->heads); + free(scratch->kv); + free(scratch->kv_raw); + free(scratch->qr_norm); + free(scratch->qr); + free(scratch->q); + free(scratch->attn_residual); + free(scratch->attn_norm); + free(scratch->attn_cur); + free(scratch->next); + free(scratch->cur); + free(scratch->plain); + memset(scratch, 0, sizeof(*scratch)); +} + +/* Allocate per-layer KV state: a raw sliding window for all layers, plus + * compressed attention/indexer caches for layers whose ratio is nonzero. */ +static void kv_cache_init(ds4_kv_cache *cache, uint32_t ctx_size, uint32_t raw_cap) { + memset(cache, 0, sizeof(*cache)); + if (raw_cap == 0) raw_cap = ds4_default_raw_cap(ctx_size); + if (raw_cap > ctx_size) raw_cap = ctx_size; + if (raw_cap == 0) raw_cap = 1; + + cache->head_dim = DS4_N_HEAD_DIM; + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + cache->layer[il].cap_raw = raw_cap; + cache->layer[il].raw_kv = xmalloc_zeroed((size_t)raw_cap * DS4_N_HEAD_DIM, sizeof(float)); + cache->layer[il].compress_ratio = ratio; + + if (ratio != 0) { + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_cap = ctx_size / ratio + 2; + const uint32_t attn_width = coff * DS4_N_HEAD_DIM; + const uint32_t attn_rows = coff * ratio; + + cache->layer[il].comp_cap = comp_cap; + cache->layer[il].attn_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_HEAD_DIM, sizeof(float)); + cache->layer[il].attn_state_kv = xmalloc_zeroed((size_t)attn_width * attn_rows, sizeof(float)); + cache->layer[il].attn_state_score = xmalloc((size_t)attn_width * attn_rows * sizeof(float)); + for (uint64_t i = 0; i < (uint64_t)attn_width * attn_rows; i++) { + cache->layer[il].attn_state_score[i] = DS4_NEG_INF; + } + + if (ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + const uint32_t index_rows = coff * ratio; + cache->layer[il].index_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_INDEXER_HEAD_DIM, sizeof(float)); + cache->layer[il].index_state_kv = xmalloc_zeroed((size_t)index_width * index_rows, sizeof(float)); + cache->layer[il].index_state_score = xmalloc((size_t)index_width * index_rows * sizeof(float)); + for (uint64_t i = 0; i < (uint64_t)index_width * index_rows; i++) { + cache->layer[il].index_state_score[i] = DS4_NEG_INF; + } + } + } + } +} + +static void kv_cache_free(ds4_kv_cache *cache) { + if (!cache) return; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + free(cache->layer[il].raw_kv); + free(cache->layer[il].attn_comp_kv); + free(cache->layer[il].attn_state_kv); + free(cache->layer[il].attn_state_score); + free(cache->layer[il].index_comp_kv); + free(cache->layer[il].index_state_kv); + free(cache->layer[il].index_state_score); + } + memset(cache, 0, sizeof(*cache)); +} + +/* Append to the raw SWA cache. Once full, it slides by one row. */ +static void kv_cache_push_raw(ds4_layer_cache *cache, const float *kv) { + if (cache->n_raw < cache->cap_raw) { + float *dst = cache->raw_kv + (uint64_t)cache->n_raw * DS4_N_HEAD_DIM; + for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); + cache->n_raw++; + return; + } + + memmove(cache->raw_kv, + cache->raw_kv + DS4_N_HEAD_DIM, + (size_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM * sizeof(cache->raw_kv[0])); + float *dst = cache->raw_kv + (uint64_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM; + for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); +} + +static void kv_cache_push_comp(float *rows, uint32_t *n_rows, uint32_t cap_rows, uint32_t row_dim, const float *kv) { + if (*n_rows >= cap_rows) ds4_die("compressed KV cache capacity exceeded"); + float *dst = rows + (uint64_t)(*n_rows) * row_dim; + for (uint32_t i = 0; i < row_dim; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); + (*n_rows)++; +} + +/* After prefill, clear unused compressor state rows so decode starts from the + * same partial-window state the streaming path would have produced. */ +static void compressor_finish_prefill_state_cpu( + float * state_kv, + float * state_score, + uint32_t head_dim, + uint32_t compress_ratio, + uint32_t n_tokens) { + if (!state_kv || !state_score || head_dim == 0 || compress_ratio == 0) return; + + const uint32_t coff = compress_ratio == 4 ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t rem = n_tokens % compress_ratio; + const uint32_t clear_start = compress_ratio == 4 ? compress_ratio + rem : rem; + const uint32_t clear_end = compress_ratio == 4 ? 2u * compress_ratio : compress_ratio; + + for (uint32_t row = clear_start; row < clear_end; row++) { + float *kv = state_kv + (uint64_t)row * width; + float *score = state_score + (uint64_t)row * width; + memset(kv, 0, (size_t)width * sizeof(kv[0])); + for (uint32_t i = 0; i < width; i++) score[i] = DS4_NEG_INF; + } +} + +static void kv_cache_finish_prefill_states(ds4_kv_cache *cache, uint32_t n_tokens) { + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_layer_cache *layer = &cache->layer[il]; + const uint32_t ratio = layer->compress_ratio; + if (ratio == 0) continue; + + compressor_finish_prefill_state_cpu(layer->attn_state_kv, + layer->attn_state_score, + DS4_N_HEAD_DIM, + ratio, + n_tokens); + if (ratio == 4) { + compressor_finish_prefill_state_cpu(layer->index_state_kv, + layer->index_state_score, + DS4_N_INDEXER_HEAD_DIM, + ratio, + n_tokens); + } + } +} + +/* Pool the current compression window with a softmax over per-dimension scores. + * Ratio-4 layers keep two lanes: attention compression and indexer compression. */ +static void compressor_pool_decode_state( + float * out, + float * state_kv, + float * state_score, + uint32_t head_dim, + uint32_t compress_ratio) { + const uint32_t coff = compress_ratio == 4 ? 2u : 1u; + const uint32_t width = coff * head_dim; + + for (uint32_t j = 0; j < head_dim; j++) { + float max_score = DS4_NEG_INF; + + if (compress_ratio == 4) { + for (uint32_t r = 0; r < compress_ratio; r++) { + const float sp = state_score[(uint64_t)r * width + j]; + const float sc = state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j]; + if (sp > max_score) max_score = sp; + if (sc > max_score) max_score = sc; + } + } else { + for (uint32_t r = 0; r < compress_ratio; r++) { + const float s = state_score[(uint64_t)r * width + j]; + if (s > max_score) max_score = s; + } + } + + if (max_score <= DS4_NEG_INF * 0.5f) { + out[j] = 0.0f; + continue; + } + + float denom = 0.0f; + float sum = 0.0f; + if (compress_ratio == 4) { + for (uint32_t r = 0; r < compress_ratio; r++) { + const float wp = expf(state_score[(uint64_t)r * width + j] - max_score); + const float wc = expf(state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j] - max_score); + denom += wp + wc; + sum += wp * state_kv[(uint64_t)r * width + j]; + sum += wc * state_kv[(uint64_t)(compress_ratio + r) * width + head_dim + j]; + } + } else { + for (uint32_t r = 0; r < compress_ratio; r++) { + const float w = expf(state_score[(uint64_t)r * width + j] - max_score); + denom += w; + sum += w * state_kv[(uint64_t)r * width + j]; + } + } + + out[j] = denom > 0.0f ? sum / denom : 0.0f; + } +} + +/* Streaming compressor update for one token. It projects kv/score rows, + * updates the rolling state, and emits a compressed KV row on ratio boundaries. */ +static bool compressor_decode_one( + float * out_comp, + const ds4_model * model, + const ds4_tensor * wkv, + const ds4_tensor * wgate, + const ds4_tensor * ape, + const ds4_tensor * norm, + const float * x, + float * state_kv, + float * state_score, + uint32_t head_dim, + uint32_t compress_ratio, + uint32_t il, + uint32_t pos) { + const uint32_t coff = compress_ratio == 4 ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t pos_mod = pos % compress_ratio; + const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; + const bool should_compress = ((pos + 1) % compress_ratio) == 0; + + float *kv_cur = xmalloc((size_t)width * sizeof(kv_cur[0])); + float *sc_cur = xmalloc((size_t)width * sizeof(sc_cur[0])); + if (wkv->type == 8 && + wgate->type == 8 && + wkv->ndim == 2 && + wgate->ndim == 2 && + wkv->dim[0] == wgate->dim[0]) { + const uint64_t in_dim = wkv->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + int8_t *xq = xmalloc((size_t)blocks * 32); + float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); + + quantize_q8_0_activation(x, xq, xscale, in_dim); + matvec_q8_0_pair_prequant(kv_cur, sc_cur, model, wkv, wgate, xq, xscale); + + free(xscale); + free(xq); + } else { + matvec_any(kv_cur, model, wkv, x); + matvec_any(sc_cur, model, wgate, x); + } + + for (uint32_t j = 0; j < width; j++) { + sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); + } + + memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); + memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); + + free(sc_cur); + free(kv_cur); + + if (!should_compress) { + return false; + } + + float *pooled = xmalloc((size_t)head_dim * sizeof(pooled[0])); + compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); + + double ss = 0.0; + for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; + const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); + for (uint32_t i = 0; i < head_dim; i++) { + out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); + } + + const uint32_t comp_pos = pos + 1 - compress_ratio; + rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); + if (head_dim == DS4_N_HEAD_DIM) { + dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); + } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { + dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); + } + + if (compress_ratio == 4) { + for (uint32_t r = 0; r < compress_ratio; r++) { + memcpy(state_kv + (uint64_t)r * width, + state_kv + (uint64_t)(compress_ratio + r) * width, + (size_t)width * sizeof(state_kv[0])); + memcpy(state_score + (uint64_t)r * width, + state_score + (uint64_t)(compress_ratio + r) * width, + (size_t)width * sizeof(state_score[0])); + } + for (uint32_t r = 0; r < compress_ratio; r++) { + memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, + state_kv + (uint64_t)r * width, + (size_t)width * sizeof(state_kv[0])); + memcpy(state_score + (uint64_t)(compress_ratio + r) * width, + state_score + (uint64_t)r * width, + (size_t)width * sizeof(state_score[0])); + } + } + + free(pooled); + return true; +} + +static bool compressor_decode_one_decode_scratch( + float * out_comp, + const ds4_model * model, + const ds4_tensor * wkv, + const ds4_tensor * wgate, + const ds4_tensor * ape, + const ds4_tensor * norm, + const float * x, + float * state_kv, + float * state_score, + uint32_t head_dim, + uint32_t compress_ratio, + uint32_t il, + uint32_t pos, + ds4_cpu_decode_scratch * scratch) { + const uint32_t coff = compress_ratio == 4 ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t pos_mod = pos % compress_ratio; + const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; + const bool should_compress = ((pos + 1) % compress_ratio) == 0; + + if (width > 2u * DS4_N_HEAD_DIM) ds4_die("compressor scratch width is outside the fixed model layout"); + float *kv_cur = scratch->comp_kv_cur; + float *sc_cur = scratch->comp_sc_cur; + + if (wkv->type == 8 && + wgate->type == 8 && + wkv->ndim == 2 && + wgate->ndim == 2 && + wkv->dim[0] == wgate->dim[0]) { + matvec_q8_0_pair_decode_scratch(kv_cur, sc_cur, model, wkv, wgate, x, scratch); + } else { + matvec_any_decode_scratch(kv_cur, model, wkv, x, scratch); + matvec_any_decode_scratch(sc_cur, model, wgate, x, scratch); + } + + for (uint32_t j = 0; j < width; j++) { + sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); + } + + memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); + memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); + + if (!should_compress) { + return false; + } + + float *pooled = scratch->comp_pooled; + compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); + + double ss = 0.0; + for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; + const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); + for (uint32_t i = 0; i < head_dim; i++) { + out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); + } + + const uint32_t comp_pos = pos + 1 - compress_ratio; + rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); + if (head_dim == DS4_N_HEAD_DIM) { + dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); + } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { + dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); + } + + if (compress_ratio == 4) { + for (uint32_t r = 0; r < compress_ratio; r++) { + memcpy(state_kv + (uint64_t)r * width, + state_kv + (uint64_t)(compress_ratio + r) * width, + (size_t)width * sizeof(state_kv[0])); + memcpy(state_score + (uint64_t)r * width, + state_score + (uint64_t)(compress_ratio + r) * width, + (size_t)width * sizeof(state_score[0])); + } + for (uint32_t r = 0; r < compress_ratio; r++) { + memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, + state_kv + (uint64_t)r * width, + (size_t)width * sizeof(state_kv[0])); + memcpy(state_score + (uint64_t)(compress_ratio + r) * width, + state_score + (uint64_t)r * width, + (size_t)width * sizeof(state_score[0])); + } + } + + return true; +} + +/* Attention over raw SWA rows plus optional compressed rows. Ratio-4 layers + * pass an indexer mask to hide compressed rows not selected for this token. */ +static void layer_attention_mixed_one( + float * out_heads, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * q, + const float * raw_kv, + uint32_t n_raw, + const float * comp_kv, + uint32_t n_comp, + const bool * comp_allowed) { + const float *sinks = tensor_data(model, layer->attn_sinks); + const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); + const uint32_t n_total = n_raw + n_comp; + float score_stack[512]; + float *score = n_total <= 512 ? score_stack : xmalloc((size_t)n_total * sizeof(score[0])); + + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; + float max_score = sinks[h]; + uint32_t idx = 0; + + for (uint32_t r = 0; r < n_raw; r++, idx++) { + const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; + score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[idx] > max_score) max_score = score[idx]; + } + for (uint32_t r = 0; r < n_comp; r++, idx++) { + if (comp_allowed && !comp_allowed[r]) { + score[idx] = DS4_NEG_INF; + continue; + } + const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; + score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[idx] > max_score) max_score = score[idx]; + } + + float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; + memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); + + float denom = expf(sinks[h] - max_score); + idx = 0; + for (uint32_t r = 0; r < n_raw; r++, idx++) { + const float weight = expf(score[idx] - max_score); + const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + for (uint32_t r = 0; r < n_comp; r++, idx++) { + if (score[idx] <= DS4_NEG_INF * 0.5f) continue; + const float weight = expf(score[idx] - max_score); + const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + + const float inv = 1.0f / denom; + scale_f32(oh, inv, DS4_N_HEAD_DIM); + } + + if (score != score_stack) free(score); +} + +static void layer_attention_mixed_one_decode_scratch( + float * out_heads, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * q, + const float * raw_kv, + uint32_t n_raw, + const float * comp_kv, + uint32_t n_comp, + const bool * comp_allowed, + ds4_cpu_decode_scratch * scratch) { + const float *sinks = tensor_data(model, layer->attn_sinks); + const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); + const uint32_t n_total = n_raw + n_comp; + if (n_total > scratch->attn_score_cap) ds4_die("CPU decode attention score scratch buffer is too small"); + float *score = scratch->attn_score; + + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; + float max_score = sinks[h]; + uint32_t idx = 0; + + for (uint32_t r = 0; r < n_raw; r++, idx++) { + const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; + score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[idx] > max_score) max_score = score[idx]; + } + for (uint32_t r = 0; r < n_comp; r++, idx++) { + if (comp_allowed && !comp_allowed[r]) { + score[idx] = DS4_NEG_INF; + continue; + } + const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; + score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[idx] > max_score) max_score = score[idx]; + } + + float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; + memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); + + float denom = expf(sinks[h] - max_score); + idx = 0; + for (uint32_t r = 0; r < n_raw; r++, idx++) { + const float weight = expf(score[idx] - max_score); + const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + for (uint32_t r = 0; r < n_comp; r++, idx++) { + if (score[idx] <= DS4_NEG_INF * 0.5f) continue; + const float weight = expf(score[idx] - max_score); + const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + + const float inv = 1.0f / denom; + scale_f32(oh, inv, DS4_N_HEAD_DIM); + } +} + +typedef struct { + float * out_heads; + const ds4_model * model; + const ds4_layer_weights * layer; + const float * q; + const float * raw_kv; + const float * comp_kv; + const uint32_t * comp_counts; + const uint8_t * allowed_mask; + const uint8_t * allowed_bits; + uint64_t allowed_stride; + uint32_t n_tok; + uint32_t raw_cap; +} layer_attention_prefix_batch_ctx; + +static inline bool attention_prefix_comp_allowed( + const layer_attention_prefix_batch_ctx *ctx, + uint32_t t, + uint32_t c) { + if (!ctx->allowed_bits || !ctx->allowed_mask || !ctx->allowed_mask[t]) return true; + const uint8_t *bits = ctx->allowed_bits + (uint64_t)t * ctx->allowed_stride; + return (bits[c >> 3] & (uint8_t)(1u << (c & 7u))) != 0; +} + +static void layer_attention_prefix_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { + layer_attention_prefix_batch_ctx *ctx = vctx; + const float *sinks = tensor_data(ctx->model, ctx->layer->attn_sinks); + const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); + const uint32_t max_comp = ctx->comp_counts ? ctx->comp_counts[ctx->n_tok - 1] : 0; + const uint32_t max_total = ctx->raw_cap + max_comp; + float score_stack[2048]; + float *score = max_total <= 2048 ? score_stack : xmalloc((size_t)max_total * sizeof(score[0])); + + for (uint64_t idx = r0; idx < r1; idx++) { + const uint32_t t = (uint32_t)(idx / DS4_N_HEAD); + const uint32_t h = (uint32_t)(idx - (uint64_t)t * DS4_N_HEAD); + const uint32_t raw_count = t + 1 < ctx->raw_cap ? t + 1 : ctx->raw_cap; + const uint32_t raw_start = t + 1 - raw_count; + const uint32_t comp_count = ctx->comp_counts ? ctx->comp_counts[t] : 0; + const float *qh = ctx->q + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; + + float max_score = sinks[h]; + uint32_t sidx = 0; + for (uint32_t r = 0; r < raw_count; r++, sidx++) { + const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; + score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[sidx] > max_score) max_score = score[sidx]; + } + for (uint32_t c = 0; c < comp_count; c++, sidx++) { + if (!attention_prefix_comp_allowed(ctx, t, c)) { + score[sidx] = DS4_NEG_INF; + continue; + } + const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; + score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; + if (score[sidx] > max_score) max_score = score[sidx]; + } + + float *oh = ctx->out_heads + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; + memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); + + float denom = expf(sinks[h] - max_score); + sidx = 0; + for (uint32_t r = 0; r < raw_count; r++, sidx++) { + const float weight = expf(score[sidx] - max_score); + const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + for (uint32_t c = 0; c < comp_count; c++, sidx++) { + if (score[sidx] <= DS4_NEG_INF * 0.5f) continue; + const float weight = expf(score[sidx] - max_score); + const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; + denom += weight; + axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); + } + + scale_f32(oh, 1.0f / denom, DS4_N_HEAD_DIM); + } + + if (score != score_stack) free(score); +} + +/* Prefix prefill attention for a fresh prompt. It computes each token's view + * of the raw window and compressed rows without running the decode loop. */ +static void layer_attention_prefix_batch( + float * out_heads, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * q, + const float * raw_kv, + const float * comp_kv, + const uint32_t * comp_counts, + const uint8_t * allowed_mask, + const uint8_t * allowed_bits, + uint64_t allowed_stride, + uint32_t n_tok, + uint32_t raw_cap) { + layer_attention_prefix_batch_ctx ctx = { + .out_heads = out_heads, + .model = model, + .layer = layer, + .q = q, + .raw_kv = raw_kv, + .comp_kv = comp_kv, + .comp_counts = comp_counts, + .allowed_mask = allowed_mask, + .allowed_bits = allowed_bits, + .allowed_stride = allowed_stride, + .n_tok = n_tok, + .raw_cap = raw_cap, + }; + ds4_parallel_for_min_rows((uint64_t)n_tok * DS4_N_HEAD, + layer_attention_prefix_batch_worker, + &ctx, + 1); +} + +/* Ratio-4 layers use an auxiliary indexer to select which compressed rows are + * visible to attention. This is the CPU allocation-owning helper. */ +static bool *indexer_allowed_decode_one( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * cur, + const float * qr_norm, + const float * index_comp, + uint32_t n_comp, + uint32_t il, + uint32_t pos) { + if (n_comp == 0) return NULL; + + bool *allowed = xcalloc(n_comp, sizeof(allowed[0])); + const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; + if (top_k == n_comp) { + for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; + return allowed; + } + + const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; + const uint32_t n_head = DS4_N_INDEXER_HEAD; + float *q = xmalloc((size_t)head_dim * n_head * sizeof(q[0])); + float *weights = xmalloc((size_t)n_head * sizeof(weights[0])); + float *scores = xmalloc((size_t)n_comp * sizeof(scores[0])); + + matvec_any(q, model, layer->indexer_attn_q_b, qr_norm); + rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); + dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); + + matvec_any(weights, model, layer->indexer_proj, cur); + const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); + for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; + + for (uint32_t c = 0; c < n_comp; c++) { + const float *kv = index_comp + (uint64_t)c * head_dim; + float s = 0.0f; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + (uint64_t)h * head_dim; + float dot = dot_f32(kv, qh, head_dim); + if (dot < 0.0f) dot = 0.0f; + s += dot * weights[h]; + } + scores[c] = s; + } + + for (uint32_t k = 0; k < top_k; k++) { + uint32_t best = 0; + float best_score = DS4_NEG_INF; + for (uint32_t c = 0; c < n_comp; c++) { + if (!allowed[c] && scores[c] > best_score) { + best = c; + best_score = scores[c]; + } + } + allowed[best] = true; + } + + free(scores); + free(weights); + free(q); + return allowed; +} + +/* Scratch-backed indexer selection for decode. */ +static bool *indexer_allowed_decode_one_decode_scratch( + const ds4_model * model, + const ds4_layer_weights * layer, + const float * cur, + const float * qr_norm, + const float * index_comp, + uint32_t n_comp, + uint32_t il, + uint32_t pos, + ds4_cpu_decode_scratch * scratch) { + if (n_comp == 0) return NULL; + if (n_comp > scratch->comp_cap) ds4_die("CPU decode indexer scratch buffer is too small"); + + bool *allowed = scratch->index_allowed; + memset(allowed, 0, (size_t)n_comp * sizeof(allowed[0])); + const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; + if (top_k == n_comp) { + for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; + return allowed; + } + + const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; + const uint32_t n_head = DS4_N_INDEXER_HEAD; + float *q = scratch->index_q; + float *weights = scratch->index_weights; + float *scores = scratch->index_scores; + + matvec_any_decode_scratch(q, model, layer->indexer_attn_q_b, qr_norm, scratch); + rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); + dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); + + matvec_any_decode_scratch(weights, model, layer->indexer_proj, cur, scratch); + const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); + for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; + + for (uint32_t c = 0; c < n_comp; c++) { + const float *kv = index_comp + (uint64_t)c * head_dim; + float s = 0.0f; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + (uint64_t)h * head_dim; + float dot = dot_f32(kv, qh, head_dim); + if (dot < 0.0f) dot = 0.0f; + s += dot * weights[h]; + } + scores[c] = s; + } + + for (uint32_t k = 0; k < top_k; k++) { + uint32_t best = 0; + float best_score = DS4_NEG_INF; + for (uint32_t c = 0; c < n_comp; c++) { + if (!allowed[c] && scores[c] > best_score) { + best = c; + best_score = scores[c]; + } + } + allowed[best] = true; + } + + return allowed; +} + +/* Single-token attention sublayer with raw SWA cache and DS4 compression. */ +static void layer_attention_raw_swa_one( + float * after_attn_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + ds4_layer_cache * cache, + const float * inp_hc, + uint32_t il, + uint32_t pos, + const float * steering_dirs, + float steering_scale) { + const uint32_t n_hc = DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + + float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); + float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); + float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); + float *q = xmalloc((size_t)q_dim * sizeof(q[0])); + float *qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(qr_norm[0])); + float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); + float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); + float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); + bool *comp_allowed = NULL; + float post[4]; + float comb[16]; + + memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); + hc_pre_from_state_one(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + attn_residual, attn_cur, post, comb); + + layer_attn_norm_one(attn_norm, model, layer, attn_cur); + layer_q_projection_with_lora_one(model, layer, attn_norm, q, qr_norm); + layer_kv_projection_normed_one(model, layer, attn_norm, kv); + + rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); + + kv_cache_push_raw(cache, kv); + + const uint32_t ratio = cache->compress_ratio; + if (ratio != 0) { + float *comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp[0])); + if (compressor_decode_one(comp, model, + layer->attn_compressor_kv, + layer->attn_compressor_gate, + layer->attn_compressor_ape, + layer->attn_compressor_norm, + attn_norm, + cache->attn_state_kv, + cache->attn_state_score, + DS4_N_HEAD_DIM, + ratio, + il, + pos)) { + kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); + } + free(comp); + + if (ratio == 4) { + float *index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp[0])); + if (compressor_decode_one(index_comp, model, + layer->indexer_compressor_kv, + layer->indexer_compressor_gate, + layer->indexer_compressor_ape, + layer->indexer_compressor_norm, + attn_norm, + cache->index_state_kv, + cache->index_state_score, + DS4_N_INDEXER_HEAD_DIM, + ratio, + il, + pos)) { + kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); + } + free(index_comp); + + comp_allowed = indexer_allowed_decode_one(model, layer, + attn_norm, qr_norm, + cache->index_comp_kv, + cache->n_index_comp, + il, pos); + } + + layer_attention_mixed_one(heads, model, layer, q, + cache->raw_kv, cache->n_raw, + cache->attn_comp_kv, cache->n_comp, + comp_allowed); + } else { + layer_attention_rows_one(heads, model, layer, q, cache->raw_kv, cache->n_raw); + } + + rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); + layer_grouped_out_one(attn_out, model, layer, heads); + cpu_directional_steering_project_rows(attn_out, steering_dirs, il, 1, steering_scale); + hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); + + free(comp_allowed); + free(attn_out); + free(heads); + free(kv); + free(qr_norm); + free(q); + free(attn_residual); + free(attn_norm); + free(attn_cur); +} + +/* Batched prefill attention. It projects Q/KV for all tokens, streams them + * through the same raw/compressed cache updates, then runs prefix attention. */ +static void layer_attention_raw_swa_batch( + float * after_attn_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + ds4_layer_cache * cache, + const float * inp_hc, + uint32_t n_tok, + uint32_t il, + uint32_t pos0, + const float * steering_dirs, + float steering_scale) { + const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; + const double t_start = profile ? now_sec() : 0.0; + double t_hc_norm = 0.0; + double t_q = 0.0; + double t_kv = 0.0; + double t_token_loop = 0.0; + double t_tl_rope_cache = 0.0; + double t_tl_compress = 0.0; + double t_tl_indexer = 0.0; + double t_tl_attn_rows = 0.0; + double t_tl_inv_rope = 0.0; + double t_out = 0.0; + const uint32_t n_hc = DS4_N_HC; + const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + + float *attn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_cur[0])); + float *attn_norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_norm[0])); + float *attn_residual = xmalloc((size_t)n_tok * hc_dim * sizeof(attn_residual[0])); + const uint32_t q_rank = DS4_N_LORA_Q; + float *qr = xmalloc((size_t)n_tok * q_rank * sizeof(qr[0])); + float *qr_norm = xmalloc((size_t)n_tok * q_rank * sizeof(qr_norm[0])); + float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); + float *kv_raw = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv_raw[0])); + float *kv = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv[0])); + float *heads = NULL; + float *attn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_out[0])); + float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); + float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); + + const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); + const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); + + double t0 = profile ? now_sec() : 0.0; + hc_pre_norm_batch(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + layer->attn_norm, + inp_hc, + attn_residual, + attn_cur, + attn_norm, + post, + comb, + n_tok); + if (profile) t_hc_norm = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + matmul_q8_0_batch(qr, model, layer->attn_q_a, attn_norm, n_tok); + for (uint32_t t = 0; t < n_tok; t++) { + rms_norm_weight(qr_norm + (uint64_t)t * q_rank, + qr + (uint64_t)t * q_rank, + q_a_norm, + q_rank, + DS4_RMS_EPS); + } + matmul_q8_0_batch(q, model, layer->attn_q_b, qr_norm, n_tok); + for (uint32_t t = 0; t < n_tok; t++) { + head_rms_norm_inplace(q + (uint64_t)t * q_dim, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_RMS_EPS); + } + if (profile) t_q = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + matmul_q8_0_batch(kv_raw, model, layer->attn_kv, attn_norm, n_tok); + for (uint32_t t = 0; t < n_tok; t++) { + rms_norm_weight(kv + (uint64_t)t * DS4_N_HEAD_DIM, + kv_raw + (uint64_t)t * DS4_N_HEAD_DIM, + kv_norm, + DS4_N_HEAD_DIM, + DS4_RMS_EPS); + } + if (profile) t_kv = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + const uint32_t ratio = cache->compress_ratio; + const bool prefer_parallel_attn = getenv("DS4_PARALLEL_ATTN_ROWS") != NULL; + const bool prefix_batch_attn = + prefer_parallel_attn && + getenv("DS4_NO_PARALLEL_ATTN_ROWS") == NULL && + cache->n_raw == 0 && + pos0 == 0; + if (!prefix_batch_attn) { + heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); + } + uint32_t batch_rope_max = 4096; + const char *batch_rope_max_env = getenv("DS4_BATCHED_ROPE_MAX"); + if (batch_rope_max_env && batch_rope_max_env[0]) { + long v = strtol(batch_rope_max_env, NULL, 10); + if (v >= 0 && v <= 65536) batch_rope_max = (uint32_t)v; + } + const bool batch_prefix_rope = + prefix_batch_attn && + getenv("DS4_NO_BATCHED_ROPE") == NULL && + n_tok <= batch_rope_max; + uint32_t *comp_counts = prefix_batch_attn ? + xcalloc((size_t)n_tok, sizeof(comp_counts[0])) : NULL; + uint8_t *allowed_mask = prefix_batch_attn && ratio == 4 ? + xcalloc((size_t)n_tok, sizeof(allowed_mask[0])) : NULL; + uint8_t *allowed_bits = NULL; + const uint64_t allowed_stride = ratio == 4 ? ((uint64_t)cache->comp_cap + 7u) / 8u : 0; + float *comp_scratch = NULL; + float *index_comp_scratch = NULL; + + if (ratio != 0) { + comp_scratch = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp_scratch[0])); + + if (ratio == 4) { + index_comp_scratch = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp_scratch[0])); + } + } + + if (batch_prefix_rope) { + double tx = profile ? now_sec() : 0.0; + rope_tail_layer_batch_inplace(q, + q_dim, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + il, + false, + n_tok); + rope_tail_layer_batch_inplace(kv, + DS4_N_HEAD_DIM, + DS4_N_HEAD_KV, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + il, + false, + n_tok); + if (profile) t_tl_rope_cache += now_sec() - tx; + } + + for (uint32_t t = 0; t < n_tok; t++) { + const uint32_t pos = pos0 + t; + float *q_t = q + (uint64_t)t * q_dim; + float *kv_t = kv + (uint64_t)t * DS4_N_HEAD_DIM; + bool *comp_allowed = NULL; + + double tx = profile ? now_sec() : 0.0; + if (!batch_prefix_rope) { + rope_tail_layer_inplace(q_t, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + rope_tail_layer_inplace(kv_t, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + } + dsv4_fp8_kv_quantize_row_inplace_cpu(kv_t, DS4_N_HEAD_DIM, DS4_N_ROT); + + kv_cache_push_raw(cache, kv_t); + if (profile) t_tl_rope_cache += now_sec() - tx; + + if (ratio != 0) { + tx = profile ? now_sec() : 0.0; + float *comp = comp_scratch; + const bool have_comp = compressor_decode_one(comp, model, + layer->attn_compressor_kv, + layer->attn_compressor_gate, + layer->attn_compressor_ape, + layer->attn_compressor_norm, + attn_norm + (uint64_t)t * DS4_N_EMBD, + cache->attn_state_kv, + cache->attn_state_score, + DS4_N_HEAD_DIM, + ratio, + il, + pos); + if (have_comp) { + kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); + } + + if (ratio == 4) { + float *index_comp = index_comp_scratch; + const bool have_index_comp = compressor_decode_one(index_comp, model, + layer->indexer_compressor_kv, + layer->indexer_compressor_gate, + layer->indexer_compressor_ape, + layer->indexer_compressor_norm, + attn_norm + (uint64_t)t * DS4_N_EMBD, + cache->index_state_kv, + cache->index_state_score, + DS4_N_INDEXER_HEAD_DIM, + ratio, + il, + pos); + if (have_index_comp) { + kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); + } + if (profile) t_tl_compress += now_sec() - tx; + + tx = profile ? now_sec() : 0.0; + comp_allowed = indexer_allowed_decode_one(model, layer, + attn_norm + (uint64_t)t * DS4_N_EMBD, + qr_norm + (uint64_t)t * q_rank, + cache->index_comp_kv, + cache->n_index_comp, + il, pos); + if (profile) t_tl_indexer += now_sec() - tx; + } else { + if (profile) t_tl_compress += now_sec() - tx; + } + + if (comp_counts) comp_counts[t] = cache->n_comp; + if (prefix_batch_attn && comp_allowed) { + if (!allowed_bits) { + allowed_bits = xcalloc((size_t)n_tok * allowed_stride, sizeof(allowed_bits[0])); + } + allowed_mask[t] = 1; + uint8_t *bits = allowed_bits + (uint64_t)t * allowed_stride; + for (uint32_t c = 0; c < cache->n_comp; c++) { + if (comp_allowed[c]) bits[c >> 3] |= (uint8_t)(1u << (c & 7u)); + } + } + + if (!prefix_batch_attn) { + tx = profile ? now_sec() : 0.0; + layer_attention_mixed_one(heads + (uint64_t)t * q_dim, model, layer, q_t, + cache->raw_kv, cache->n_raw, + cache->attn_comp_kv, cache->n_comp, + comp_allowed); + if (profile) t_tl_attn_rows += now_sec() - tx; + } + } else { + if (!prefix_batch_attn) { + tx = profile ? now_sec() : 0.0; + layer_attention_rows_one(heads + (uint64_t)t * q_dim, model, layer, q_t, cache->raw_kv, cache->n_raw); + if (profile) t_tl_attn_rows += now_sec() - tx; + } + } + + if (!prefix_batch_attn) { + tx = profile ? now_sec() : 0.0; + rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + il, + true); + if (profile) t_tl_inv_rope += now_sec() - tx; + } + + free(comp_allowed); + } + + if (prefix_batch_attn) { + double tx = profile ? now_sec() : 0.0; + const float *comp_kv_for_prefix = cache->attn_comp_kv ? cache->attn_comp_kv : kv; + if (!heads) { + heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); + } + layer_attention_prefix_batch(heads, model, layer, + q, + kv, + comp_kv_for_prefix, + comp_counts, + allowed_mask, + allowed_bits, + allowed_stride, + n_tok, + cache->cap_raw); + if (profile) t_tl_attn_rows += now_sec() - tx; + tx = profile ? now_sec() : 0.0; + if (batch_prefix_rope) { + rope_tail_layer_batch_inplace(heads, + q_dim, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + il, + true, + n_tok); + } else { + for (uint32_t t = 0; t < n_tok; t++) { + rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0 + t, + il, + true); + } + } + if (profile) t_tl_inv_rope += now_sec() - tx; + } + if (profile) t_token_loop = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_grouped_out_batch(attn_out, model, layer, heads, n_tok); + cpu_directional_steering_project_rows(attn_out, steering_dirs, il, n_tok, steering_scale); + + hc_post_batch(after_attn_hc, + attn_out, + attn_residual, + post, + comb, + n_tok, + DS4_N_EMBD, + n_hc); + if (profile) t_out = now_sec() - t0; + + if (profile) { + fprintf(stderr, + "ds4: prefill detail layer %u attn hc_norm=%.3f q=%.3f kv=%.3f token_loop=%.3f out=%.3f total=%.3f\n", + il, t_hc_norm, t_q, t_kv, t_token_loop, t_out, now_sec() - t_start); + if (getenv("DS4_PREFILL_PROFILE_TOKEN") != NULL) { + fprintf(stderr, + "ds4: prefill token detail layer %u rope_cache=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f\n", + il, t_tl_rope_cache, t_tl_compress, t_tl_indexer, t_tl_attn_rows, t_tl_inv_rope); + } + } + + free(allowed_bits); + free(allowed_mask); + free(comp_counts); + free(index_comp_scratch); + free(comp_scratch); + free(comb); + free(post); + free(attn_out); + free(heads); + free(kv); + free(kv_raw); + free(q); + free(qr_norm); + free(qr); + free(attn_residual); + free(attn_norm); + free(attn_cur); +} + +/* Full transformer layer for one decode token: attention sublayer followed by + * FFN sublayer, both operating on the HC state. */ +static void layer_forward_raw_swa_one( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + ds4_layer_cache * cache, + const float * inp_hc, + uint32_t il, + uint32_t pos, + int token, + const float * steering_dirs, + float steering_attn_scale, + float steering_ffn_scale, + ds4_cpu_decode_scratch * scratch) { + const uint32_t n_hc = DS4_N_HC; + const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; + const double t_start = profile ? now_sec() : 0.0; + double t_hc = 0.0; + double t_q = 0.0; + double t_kv = 0.0; + double t_rope_cache = 0.0; + double t_compress = 0.0; + double t_indexer = 0.0; + double t_attn_rows = 0.0; + double t_inv_rope = 0.0; + double t_out = 0.0; + double t_post = 0.0; + double t_ffn = 0.0; + + bool *comp_allowed = NULL; + float post[4]; + float comb[16]; + + double t0 = profile ? now_sec() : 0.0; + memcpy(scratch->attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); + hc_pre_from_state_one_scratch(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + scratch->attn_residual, scratch->attn_cur, post, comb, + scratch->hc_flat, + false); + if (profile) t_hc = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_attn_norm_one(scratch->attn_norm, model, layer, scratch->attn_cur); + const uint32_t ratio = cache->compress_ratio; + layer_q_projection_with_lora_one_decode_scratch(model, layer, + scratch->attn_norm, + scratch->q, + scratch->qr_norm, + scratch); + if (profile) t_q = now_sec() - t0; + t0 = profile ? now_sec() : 0.0; + layer_kv_projection_normed_one_decode_scratch(model, layer, + scratch->attn_norm, + scratch->kv, + scratch); + if (profile) t_kv = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + rope_tail_layer_inplace(scratch->q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + rope_tail_layer_inplace(scratch->kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + dsv4_fp8_kv_quantize_row_inplace_cpu(scratch->kv, DS4_N_HEAD_DIM, DS4_N_ROT); + + kv_cache_push_raw(cache, scratch->kv); + if (profile) t_rope_cache = now_sec() - t0; + + if (ratio != 0) { + t0 = profile ? now_sec() : 0.0; + if (compressor_decode_one_decode_scratch(scratch->comp, model, + layer->attn_compressor_kv, + layer->attn_compressor_gate, + layer->attn_compressor_ape, + layer->attn_compressor_norm, + scratch->attn_norm, + cache->attn_state_kv, + cache->attn_state_score, + DS4_N_HEAD_DIM, + ratio, + il, + pos, + scratch)) { + kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, scratch->comp); + } + + if (ratio == 4) { + if (compressor_decode_one_decode_scratch(scratch->index_comp, model, + layer->indexer_compressor_kv, + layer->indexer_compressor_gate, + layer->indexer_compressor_ape, + layer->indexer_compressor_norm, + scratch->attn_norm, + cache->index_state_kv, + cache->index_state_score, + DS4_N_INDEXER_HEAD_DIM, + ratio, + il, + pos, + scratch)) { + kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, + DS4_N_INDEXER_HEAD_DIM, scratch->index_comp); + } + if (profile) t_compress = now_sec() - t0; + } else if (profile) { + t_compress = now_sec() - t0; + } + } + if (ratio == 4) { + t0 = profile ? now_sec() : 0.0; + comp_allowed = indexer_allowed_decode_one_decode_scratch(model, layer, + scratch->attn_norm, + scratch->qr_norm, + cache->index_comp_kv, + cache->n_index_comp, + il, pos, + scratch); + if (profile) t_indexer = now_sec() - t0; + } + + t0 = profile ? now_sec() : 0.0; + if (ratio != 0) { + layer_attention_mixed_one_decode_scratch(scratch->heads, model, layer, scratch->q, + cache->raw_kv, cache->n_raw, + cache->attn_comp_kv, cache->n_comp, + comp_allowed, + scratch); + } else { + layer_attention_rows_one(scratch->heads, model, layer, scratch->q, cache->raw_kv, cache->n_raw); + } + if (profile) t_attn_rows = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + rope_tail_layer_inplace(scratch->heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); + if (profile) t_inv_rope = now_sec() - t0; + t0 = profile ? now_sec() : 0.0; + layer_grouped_out_one_decode_scratch(scratch->attn_out, model, layer, scratch->heads, scratch); + cpu_directional_steering_project_rows(scratch->attn_out, steering_dirs, il, 1, steering_attn_scale); + if (profile) t_out = now_sec() - t0; + t0 = profile ? now_sec() : 0.0; + hc_post_one(scratch->after_attn_hc, scratch->attn_out, scratch->attn_residual, post, comb, DS4_N_EMBD, n_hc); + if (profile) t_post = now_sec() - t0; + + t0 = profile ? now_sec() : 0.0; + layer_ffn_one_decode_scratch(out_hc, model, layer, scratch->after_attn_hc, il, token, + steering_dirs, steering_ffn_scale, scratch); + if (profile) t_ffn = now_sec() - t0; + + if (profile) { + fprintf(stderr, + "ds4: decode detail layer %u attn hc=%.3f q=%.3f kv=%.3f rope=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f out=%.3f post=%.3f ffn=%.3f total=%.3f ms\n", + il, + t_hc * 1000.0, + t_q * 1000.0, + t_kv * 1000.0, + t_rope_cache * 1000.0, + t_compress * 1000.0, + t_indexer * 1000.0, + t_attn_rows * 1000.0, + t_inv_rope * 1000.0, + t_out * 1000.0, + t_post * 1000.0, + t_ffn * 1000.0, + (now_sec() - t_start) * 1000.0); + } + +} + +static void output_logits_one_decode_scratch( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + const float * inp_hc, + ds4_cpu_decode_scratch * scratch); + +/* CPU decode for one token through all 43 layers. The caller owns scratch and + * cache lifetimes so no per-token allocations are needed. */ +static void forward_token_raw_swa_cpu_decode_scratch( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + ds4_kv_cache * cache, + int token, + uint32_t pos, + const float * steering_dirs, + float steering_attn_scale, + float steering_ffn_scale, + ds4_cpu_decode_scratch * scratch) { + float *cur = scratch->cur; + float *next = scratch->next; + + embed_token_f16(model, weights, token, scratch->plain); + hc_from_plain_embedding(cur, scratch->plain, DS4_N_EMBD, DS4_N_HC); + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + layer_forward_raw_swa_one(next, model, &weights->layer[il], &cache->layer[il], + cur, il, pos, token, + steering_dirs, + steering_attn_scale, + steering_ffn_scale, + scratch); + float *tmp = cur; + cur = next; + next = tmp; + } + + if (logits) { + output_logits_one_decode_scratch(logits, model, weights, cur, scratch); + } +} + +#ifndef DS4_NO_GPU +static void forward_token_raw_swa_cpu( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + ds4_kv_cache * cache, + int token, + uint32_t pos) { + ds4_cpu_decode_scratch scratch; + uint32_t ctx_guess = pos + 1; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = cache->layer[il].compress_ratio; + if (ratio != 0 && cache->layer[il].comp_cap > 2) { + const uint32_t ctx_from_comp = (cache->layer[il].comp_cap - 2u) * ratio; + if (ctx_guess < ctx_from_comp) ctx_guess = ctx_from_comp; + } + } + cpu_decode_scratch_init(&scratch, ctx_guess); + forward_token_raw_swa_cpu_decode_scratch(logits, model, weights, cache, token, pos, + NULL, 0.0f, 0.0f, &scratch); + cpu_decode_scratch_free(&scratch); +} +#endif + +/* CPU prefill in layer-major order. All prompt tokens pass through layer 0, + * then layer 1, etc., which exposes batch matmul opportunities. */ +static void prefill_layer_major_cpu( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + ds4_kv_cache * cache, + const token_vec * prompt, + const float * steering_dirs, + float steering_attn_scale, + float steering_ffn_scale) { + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t n_tok = (uint64_t)prompt->len; + float *cur = xmalloc((size_t)n_tok * hc_dim * sizeof(cur[0])); + float *next = xmalloc((size_t)n_tok * hc_dim * sizeof(next[0])); + float *attn = xmalloc((size_t)n_tok * hc_dim * sizeof(attn[0])); + float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); + uint32_t ffn_batch = 128; + const bool batched_attn = getenv("DS4_NO_BATCHED_ATTN") == NULL; + const bool batched_ffn = getenv("DS4_BATCHED_FFN") != NULL; + const bool parallel_ffn = getenv("DS4_PARALLEL_FFN") != NULL; + const bool shared_batch_ffn = getenv("DS4_NO_SHARED_BATCH_FFN") == NULL; + const char *batch_env = getenv("DS4_PREFILL_BATCH"); + ds4_cpu_decode_scratch decode_scratch; + bool decode_scratch_ready = false; + if (batch_env && batch_env[0]) { + long v = strtol(batch_env, NULL, 10); + if (v > 0 && v < 4096) ffn_batch = (uint32_t)v; + } + + for (uint64_t t = 0; t < n_tok; t++) { + embed_token_f16(model, weights, prompt->v[t], plain); + hc_from_plain_embedding(cur + t * hc_dim, plain, DS4_N_EMBD, DS4_N_HC); + } + + free(plain); + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + fprintf(stderr, "ds4: prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); + fflush(stderr); + + if (batched_attn) { + layer_attention_raw_swa_batch(attn, + model, + &weights->layer[il], + &cache->layer[il], + cur, + (uint32_t)n_tok, + il, + 0, + steering_dirs, + steering_attn_scale); + + if (batched_ffn) { + for (uint64_t t = 0; t < n_tok; t += ffn_batch) { + uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); + layer_ffn_batch(next + t * hc_dim, + model, + &weights->layer[il], + attn + t * hc_dim, + prompt->v + t, + nb, + il, + steering_dirs, + steering_ffn_scale); + } + } else if (shared_batch_ffn) { + layer_ffn_shared_batch(next, + model, + &weights->layer[il], + attn, + prompt->v, + (uint32_t)n_tok, + il, + steering_dirs, + steering_ffn_scale); + } else if (parallel_ffn) { + layer_ffn_tokens_parallel(next, + model, + &weights->layer[il], + attn, + prompt->v, + (uint32_t)n_tok, + il, + steering_dirs, + steering_ffn_scale); + } else { + for (uint64_t t = 0; t < n_tok; t++) { + layer_ffn_one(next + t * hc_dim, + model, + &weights->layer[il], + attn + t * hc_dim, + il, + prompt->v[t], + steering_dirs, + steering_ffn_scale, + false); + } + } + } else if (batched_ffn) { + for (uint64_t t = 0; t < n_tok; t++) { + layer_attention_raw_swa_one(attn + t * hc_dim, + model, + &weights->layer[il], + &cache->layer[il], + cur + t * hc_dim, + il, + (uint32_t)t, + steering_dirs, + steering_attn_scale); + } + + for (uint64_t t = 0; t < n_tok; t += ffn_batch) { + uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); + layer_ffn_batch(next + t * hc_dim, + model, + &weights->layer[il], + attn + t * hc_dim, + prompt->v + t, + nb, + il, + steering_dirs, + steering_ffn_scale); + } + } else { + if (!decode_scratch_ready) { + cpu_decode_scratch_init(&decode_scratch, (uint32_t)n_tok); + decode_scratch_ready = true; + } + for (uint64_t t = 0; t < n_tok; t++) { + layer_forward_raw_swa_one(next + t * hc_dim, + model, + &weights->layer[il], + &cache->layer[il], + cur + t * hc_dim, + il, + (uint32_t)t, + prompt->v[t], + steering_dirs, + steering_attn_scale, + steering_ffn_scale, + &decode_scratch); + } + } + + float *tmp = cur; + cur = next; + next = tmp; + } + + kv_cache_finish_prefill_states(cache, (uint32_t)n_tok); + + if (logits) { + output_logits_one(logits, model, weights, cur + (n_tok - 1) * hc_dim); + } + + if (decode_scratch_ready) cpu_decode_scratch_free(&decode_scratch); + free(next); + free(cur); + free(attn); +} + +/* Diagnostic first-token layer without cache history: the token attends only + * to itself, useful for checking a minimal end-to-end slice. */ +static void layer_forward_self_one( + float * out_hc, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * inp_hc, + uint32_t il, + uint32_t pos, + int token) { + const uint32_t n_hc = DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + + float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); + float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); + float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); + float *q = xmalloc((size_t)q_dim * sizeof(q[0])); + float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); + float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); + float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); + float *after_attn_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(after_attn_hc[0])); + float post[4]; + float comb[16]; + + memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); + hc_pre_from_state_one(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + attn_residual, attn_cur, post, comb); + + layer_attn_norm_one(attn_norm, model, layer, attn_cur); + layer_q_projection_normed_one(model, layer, attn_norm, q); + layer_kv_projection_normed_one(model, layer, attn_norm, kv); + rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); + dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); + f16_round_inplace_cpu(kv, DS4_N_HEAD_DIM); + + layer_attention_one(heads, model, layer, q, kv); + rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); + layer_grouped_out_one(attn_out, model, layer, heads); + hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); + + layer_ffn_one(out_hc, model, layer, after_attn_hc, il, token, + NULL, 0.0f, false); + + free(after_attn_hc); + free(attn_out); + free(heads); + free(kv); + free(q); + free(attn_residual); + free(attn_norm); + free(attn_cur); +} + +static void forward_first_token_cpu( + float * out_hc, + const ds4_model * model, + const ds4_weights * weights, + int token) { + float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); + float *cur = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(cur[0])); + float *next = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(next[0])); + + embed_token_f16(model, weights, token, plain); + hc_from_plain_embedding(cur, plain, DS4_N_EMBD, DS4_N_HC); + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + layer_forward_self_one(next, model, &weights->layer[il], cur, il, 0, token); + float *tmp = cur; + cur = next; + next = tmp; + } + + memcpy(out_hc, cur, (size_t)DS4_N_HC * DS4_N_EMBD * sizeof(out_hc[0])); + + free(next); + free(cur); + free(plain); +} + +/* Collapse final HC streams into the ordinary embedding vector before the + * output norm and vocabulary projection. */ +static void output_hc_head_one( + float * out, + const ds4_model * model, + const ds4_weights * weights, + const float * inp_hc) { + const uint32_t n_hc = DS4_N_HC; + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; + float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); + float *pre = xmalloc((size_t)n_hc * sizeof(pre[0])); + float *w = xmalloc((size_t)n_hc * sizeof(w[0])); + + rms_norm_no_weight(flat, inp_hc, hc_dim, DS4_RMS_EPS); + matvec_f16(pre, model, weights->output_hc_fn, flat); + + const float *scale = tensor_data(model, weights->output_hc_scale); + const float *base = tensor_data(model, weights->output_hc_base); + for (uint32_t i = 0; i < n_hc; i++) { + w[i] = sigmoid_stable(pre[i] * scale[0] + base[i]) + DS4_HC_EPS; + } + + hc_weighted_sum_one(out, inp_hc, w, DS4_N_EMBD, n_hc); + + free(w); + free(pre); + free(flat); +} + +/* Final language-model head: HC collapse, RMSNorm, and Q8_0 vocab projection. */ +static void output_logits_one( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + const float * inp_hc) { + float *embd = xmalloc((size_t)DS4_N_EMBD * sizeof(embd[0])); + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + + output_hc_head_one(embd, model, weights, inp_hc); + rms_norm_weight(norm, embd, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); + + matvec_q8_0(logits, model, weights->output, norm); + + free(norm); + free(embd); +} + +/* Allocation-free logits head for CPU decode. */ +static void output_logits_one_decode_scratch( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + const float * inp_hc, + ds4_cpu_decode_scratch * scratch) { + const uint32_t n_hc = DS4_N_HC; + const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; + + rms_norm_no_weight(scratch->output_flat, inp_hc, hc_dim, DS4_RMS_EPS); + matvec_f16(scratch->output_pre, model, weights->output_hc_fn, scratch->output_flat); + + const float *scale = tensor_data(model, weights->output_hc_scale); + const float *base = tensor_data(model, weights->output_hc_base); + for (uint32_t i = 0; i < n_hc; i++) { + scratch->output_weights[i] = sigmoid_stable(scratch->output_pre[i] * scale[0] + base[i]) + DS4_HC_EPS; + } + + hc_weighted_sum_one(scratch->output_embd, inp_hc, scratch->output_weights, DS4_N_EMBD, n_hc); + rms_norm_weight(scratch->output_norm, scratch->output_embd, + tensor_data(model, weights->output_norm), + DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0_decode_scratch(logits, model, weights->output, scratch->output_norm, scratch); +} diff --git a/models/deepseek/cuda/control.inc b/models/deepseek/cuda/control.inc new file mode 100644 index 0000000000..cfa84e31d5 --- /dev/null +++ b/models/deepseek/cuda/control.inc @@ -0,0 +1,2434 @@ +__device__ static void hc4_split_one(float *out, const float *mix, const float *scale, const float *base, uint32_t sinkhorn_iters, float epsv) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + for (int i = 0; i < 4; i++) { + float z = mix[i] * pre_scale + base[i]; + out[i] = 1.0f / (1.0f + expf(-z)) + epsv; + } + for (int i = 0; i < 4; i++) { + float z = mix[4 + i] * post_scale + base[4 + i]; + out[4 + i] = 2.0f / (1.0f + expf(-z)); + } + float c[16]; + for (int r = 0; r < 4; r++) { + float m = -INFINITY; + for (int col = 0; col < 4; col++) { + float v = mix[8 + r * 4 + col] * comb_scale + base[8 + r * 4 + col]; + c[r * 4 + col] = v; + m = fmaxf(m, v); + } + float s = 0.0f; + for (int col = 0; col < 4; col++) { + float v = expf(c[r * 4 + col] - m); + c[r * 4 + col] = v; + s += v; + } + for (int col = 0; col < 4; col++) c[r * 4 + col] = c[r * 4 + col] / s + epsv; + } + for (int col = 0; col < 4; col++) { + float s = epsv; + for (int r = 0; r < 4; r++) s += c[r * 4 + col]; + for (int r = 0; r < 4; r++) c[r * 4 + col] /= s; + } + for (uint32_t iter = 1; iter < sinkhorn_iters; iter++) { + for (int r = 0; r < 4; r++) { + float s = epsv; + for (int col = 0; col < 4; col++) s += c[r * 4 + col]; + for (int col = 0; col < 4; col++) c[r * 4 + col] /= s; + } + for (int col = 0; col < 4; col++) { + float s = epsv; + for (int r = 0; r < 4; r++) s += c[r * 4 + col]; + for (int r = 0; r < 4; r++) c[r * 4 + col] /= s; + } + } + for (int i = 0; i < 16; i++) out[8 + i] = c[i]; +} + +__global__ static void hc_split_sinkhorn_kernel(float *out, const float *mix, const float *scale, const float *base, uint32_t n_rows, uint32_t sinkhorn_iters, float epsv) { + uint32_t row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= n_rows) return; + hc4_split_one(out + (uint64_t)row * 24, mix + (uint64_t)row * 24, scale, base, sinkhorn_iters, epsv); +} + +__global__ static void hc_weighted_sum_kernel(float *out, const float *x, const float *w, uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens, uint32_t weight_stride_f32) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_embd * n_tokens; + if (gid >= n) return; + uint32_t d = gid % n_embd; + uint32_t t = gid / n_embd; + float acc = 0.0f; + for (uint32_t h = 0; h < n_hc; h++) { + acc += x[(uint64_t)t * n_hc * n_embd + (uint64_t)h * n_embd + d] * + w[(uint64_t)t * weight_stride_f32 + h]; + } + out[(uint64_t)t * n_embd + d] = acc; +} + +__global__ static void hc_expand_kernel( + float *out_hc, + const float *block_out, + const float *block_add, + const float *block_add2, + const float *residual_hc, + const float *post, + const float *comb, + uint32_t n_embd, + uint32_t n_hc, + uint32_t n_tokens, + uint32_t post_stride, + uint32_t comb_stride, + int has_add, + int has_add2) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + if (gid >= n_elem) return; + uint32_t d = gid % n_embd; + uint64_t tmp = gid / n_embd; + uint32_t dst_hc = tmp % n_hc; + uint32_t t = tmp / n_hc; + + float block_v = block_out[(uint64_t)t * n_embd + d]; + if (has_add) { + float add_v = block_add[(uint64_t)t * n_embd + d]; + if (has_add2) add_v += block_add2[(uint64_t)t * n_embd + d]; + block_v += add_v; + } + float acc = block_v * post[(uint64_t)t * post_stride + dst_hc]; + for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) { + float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc]; + float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d]; + acc += comb_v * res_v; + } + out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc; +} + +__global__ static void hc_split_weighted_sum_fused_kernel( + float *out, + float *split, + const float *mix, + const float *residual_hc, + const float *scale, + const float *base, + uint32_t n_embd, + uint32_t n_hc, + uint32_t n_rows, + uint32_t sinkhorn_iters, + float epsv) { + uint32_t t = blockIdx.x; + uint32_t d = threadIdx.x; + if (t >= n_rows || n_hc != 4) return; + const uint32_t mix_hc = 24; + float *sp = split + (uint64_t)t * mix_hc; + if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv); + __syncthreads(); + for (uint32_t col = d; col < n_embd; col += blockDim.x) { + float acc = 0.0f; + for (uint32_t h = 0; h < 4; h++) { + acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h]; + } + out[(uint64_t)t * n_embd + col] = acc; + } +} + +__global__ static void hc_split_weighted_sum_norm_fused_kernel( + float *out, + float *norm_out, + float *split, + const float *mix, + const float *residual_hc, + const float *scale, + const float *base, + const float *norm_w, + uint32_t n_embd, + uint32_t n_hc, + uint32_t n_rows, + uint32_t sinkhorn_iters, + float epsv, + float norm_eps) { + const uint32_t t = blockIdx.x; + const uint32_t d = threadIdx.x; + if (t >= n_rows || n_hc != 4) return; + const uint32_t mix_hc = 24; + float *sp = split + (uint64_t)t * mix_hc; + if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv); + __syncthreads(); + + float sum = 0.0f; + for (uint32_t col = d; col < n_embd; col += blockDim.x) { + float acc = 0.0f; + for (uint32_t h = 0; h < 4; h++) { + acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h]; + } + out[(uint64_t)t * n_embd + col] = acc; + sum += acc * acc; + } + + __shared__ float partial[256]; + partial[d] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (d < stride) partial[d] += partial[d + stride]; + __syncthreads(); + } + const float norm_scale = rsqrtf(partial[0] / (float)n_embd + norm_eps); + for (uint32_t col = d; col < n_embd; col += blockDim.x) { + const float v = out[(uint64_t)t * n_embd + col]; + norm_out[(uint64_t)t * n_embd + col] = v * norm_scale * norm_w[col]; + } +} + +__global__ static void output_hc_weights_kernel( + float *out, + const float *pre, + const float *scale, + const float *base, + uint32_t n_hc, + uint32_t n_tokens, + float epsv) { + uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t n = n_tokens * n_hc; + if (gid >= n) return; + uint32_t h = gid % n_hc; + float z = pre[gid] * scale[0] + base[h]; + out[gid] = 1.0f / (1.0f + expf(-z)) + epsv; +} + +__global__ static void fill_f32_kernel(float *x, uint64_t n, float v) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) x[i] = v; +} + +__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) { + 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); +} + +__global__ static void compressor_set_rows_kernel( + float *state_kv, + float *state_score, + const float *kv, + const float *sc, + const void *model_map, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos0, + uint32_t src0, + uint32_t dst0, + uint32_t rows) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)rows * width; + if (gid >= n) return; + uint32_t r = gid / width; + uint32_t j = gid - (uint64_t)r * width; + uint32_t src = src0 + r; + uint32_t dst = dst0 + r; + uint32_t phase = (pos0 + src) % ratio; + state_kv[(uint64_t)dst * width + j] = kv[(uint64_t)src * width + j]; + state_score[(uint64_t)dst * width + j] = + sc[(uint64_t)src * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)phase * width + j); +} + +__global__ static void compressor_prefill_pool_kernel( + float *comp, + const float *kv, + const float *sc, + const float *state_kv, + const 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_comp, + uint32_t replay) { + uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t c = blockIdx.y; + if (d >= head_dim || c >= n_comp) return; + uint32_t coff = ratio == 4u ? 2u : 1u; + uint32_t width = coff * head_dim; + float vals[128]; + float scores[128]; + float max_s = -INFINITY; + uint32_t n_cand = 0; + if (ratio == 4u) { + if (replay && c == 0) { + for (uint32_t r = 0; r < 4; r++) { + vals[n_cand] = state_kv[(uint64_t)r * width + d]; + scores[n_cand] = state_score[(uint64_t)r * width + d]; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } else if (c > 0) { + uint32_t base = (c - 1u) * ratio; + for (uint32_t r = 0; r < 4; r++) { + uint32_t t = base + r; + float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d); + vals[n_cand] = kv[(uint64_t)t * width + d]; + scores[n_cand] = sc[(uint64_t)t * width + d] + ape; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } + uint32_t base = c * ratio; + for (uint32_t r = 0; r < 4; r++) { + uint32_t t = base + r; + float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + head_dim + d); + vals[n_cand] = kv[(uint64_t)t * width + head_dim + d]; + scores[n_cand] = sc[(uint64_t)t * width + head_dim + d] + ape; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } else { + uint32_t base = c * ratio; + for (uint32_t r = 0; r < ratio; r++) { + uint32_t t = base + r; + float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d); + vals[n_cand] = kv[(uint64_t)t * width + d]; + scores[n_cand] = sc[(uint64_t)t * width + d] + ape; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } + float den = 0.0f, acc = 0.0f; + for (uint32_t i = 0; i < n_cand; i++) { + float w = expf(scores[i] - max_s); + den += w; + acc += vals[i] * w; + } + comp[(uint64_t)c * head_dim + d] = den != 0.0f ? acc / den : 0.0f; +} + +__global__ static void compressor_update_pool_kernel( + float *row, + const float *state_kv, + const float *state_score, + uint32_t head_dim, + uint32_t ratio) { + uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= head_dim) return; + uint32_t coff = ratio == 4u ? 2u : 1u; + uint32_t width = coff * head_dim; + float vals[128]; + float scores[128]; + float max_s = -INFINITY; + uint32_t n_cand = 0; + if (ratio == 4u) { + for (uint32_t r = 0; r < 4; r++) { + vals[n_cand] = state_kv[(uint64_t)r * width + d]; + scores[n_cand] = state_score[(uint64_t)r * width + d]; + max_s = fmaxf(max_s, scores[n_cand++]); + } + for (uint32_t r = 0; r < 4; r++) { + vals[n_cand] = state_kv[(uint64_t)(ratio + r) * width + head_dim + d]; + scores[n_cand] = state_score[(uint64_t)(ratio + r) * width + head_dim + d]; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } else { + for (uint32_t r = 0; r < ratio; r++) { + vals[n_cand] = state_kv[(uint64_t)r * width + d]; + scores[n_cand] = state_score[(uint64_t)r * width + d]; + max_s = fmaxf(max_s, scores[n_cand++]); + } + } + float den = 0.0f, acc = 0.0f; + for (uint32_t i = 0; i < n_cand; i++) { + float w = expf(scores[i] - max_s); + den += w; + acc += vals[i] * w; + } + row[d] = den != 0.0f ? acc / den : 0.0f; +} + +__global__ static void compressor_shift_ratio4_kernel(float *state_kv, float *state_score, uint32_t width) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t half = 4ull * width; + if (i >= half) return; + float v = state_kv[half + i]; + float s = state_score[half + i]; + state_kv[i] = v; + state_score[i] = s; + state_kv[half + i] = v; + state_score[half + i] = 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)); +} + +__global__ static void router_select_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) { + uint32_t t = blockIdx.x; + if (t >= n_tokens || threadIdx.x != 0) return; + const float *log = logits + (uint64_t)t * 256; + float *prob = probs + (uint64_t)t * 256; + int32_t *sel = selected + (uint64_t)t * 6; + float *w = weights + (uint64_t)t * 6; + + for (int i = 0; i < 256; i++) prob[i] = sqrtf(softplus_dev(log[i])); + + if (hash_mode) { + int32_t tok = tokens ? tokens[t] : token_scalar; + if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0; + const int32_t *row = hash + (uint64_t)tok * 6; + for (int i = 0; i < 6; i++) sel[i] = row[i]; + } else { + for (int i = 0; i < 6; i++) sel[i] = -1; + for (int i = 0; i < 256; i++) { + float score = prob[i] + (has_bias ? bias[i] : 0.0f); + for (int j = 0; j < 6; j++) { + if (sel[j] < 0 || score > prob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) { + for (int k = 5; k > j; k--) sel[k] = sel[k - 1]; + sel[j] = i; + break; + } + } + } + } + + float sum = 0.0f; + for (int i = 0; i < 6; i++) { + int e = sel[i]; + float v = (e >= 0 && e < 256) ? prob[e] : 0.0f; + w[i] = v; + sum += v; + } + sum = fmaxf(sum, 6.103515625e-5f); + for (int i = 0; i < 6; i++) w[i] = w[i] / sum * 1.5f; +} + +__global__ static void router_select_parallel_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) { + uint32_t t = blockIdx.x; + uint32_t i = threadIdx.x; + if (t >= n_tokens || i >= 256u) return; + const float *log = logits + (uint64_t)t * 256; + float *prob = probs + (uint64_t)t * 256; + int32_t *sel = selected + (uint64_t)t * 6; + float *w = weights + (uint64_t)t * 6; + __shared__ float sprob[256]; + + const float p = sqrtf(softplus_dev(log[i])); + sprob[i] = p; + prob[i] = p; + __syncthreads(); + + if (i != 0) return; + if (hash_mode) { + int32_t tok = tokens ? tokens[t] : token_scalar; + if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0; + const int32_t *row = hash + (uint64_t)tok * 6; + for (int j = 0; j < 6; j++) sel[j] = row[j]; + } else { + for (int j = 0; j < 6; j++) sel[j] = -1; + for (int e = 0; e < 256; e++) { + float score = sprob[e] + (has_bias ? bias[e] : 0.0f); + for (int j = 0; j < 6; j++) { + if (sel[j] < 0 || score > sprob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) { + for (int k = 5; k > j; k--) sel[k] = sel[k - 1]; + sel[j] = e; + break; + } + } + } + } + + float sum = 0.0f; + for (int j = 0; j < 6; j++) { + int e = sel[j]; + float v = (e >= 0 && e < 256) ? sprob[e] : 0.0f; + w[j] = v; + sum += v; + } + sum = fmaxf(sum, 6.103515625e-5f); + for (int j = 0; j < 6; j++) w[j] = w[j] / sum * 1.5f; +} + +__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 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] : 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; + } +} + +__global__ static void swiglu_kernel(float *out, const float *gate, const float *up, uint32_t n, float clamp, float weight) { + uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + float g = gate[i]; + float u = up[i]; + if (clamp > 1.0e-6f) { + g = fminf(g, clamp); + u = fminf(fmaxf(u, -clamp), clamp); + } + float s = g / (1.0f + expf(-g)); + out[i] = s * u * weight; +} + +__global__ static void add_kernel(float *out, const float *a, const float *b, uint32_t n) { + uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + out[i] = a[i] + b[i]; +} + +__global__ static void directional_steering_project_kernel( + float *x, + const float *directions, + uint32_t layer, + uint32_t width, + uint32_t rows, + float scale) { + const uint32_t row = blockIdx.x; + if (row >= rows || width == 0) return; + + float *xr = x + (uint64_t)row * width; + const float *dir = directions + (uint64_t)layer * width; + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) { + sum += xr[i] * dir[i]; + } + + __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 coeff = scale * partial[0]; + for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) { + xr[i] -= coeff * dir[i]; + } +} + +__global__ static void zero_kernel(float *out, uint64_t n) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = 0.0f; +} + +__global__ static void indexer_scores_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + int causal) { + uint32_t c = blockIdx.x; + uint32_t t = blockIdx.y; + if (c >= n_comp || t >= n_tokens) return; + if (causal) { + uint32_t n_visible = (pos0 + t + 1u) / ratio; + if (c >= n_visible) { + if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = -INFINITY; + return; + } + } + float total = 0.0f; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + const float *kh = index_comp + (uint64_t)c * head_dim; + float dot = 0.0f; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) dot += qh[d] * kh[d]; + __shared__ float partial[256]; + partial[threadIdx.x] = dot; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + total += fmaxf(partial[0], 0.0f) * weights[(uint64_t)t * n_head + h]; + __syncthreads(); + } + if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = total * scale; +} + +__global__ static void indexer_score_one_direct_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t pos0, + uint32_t ratio, + float scale, + int causal) { + const uint32_t c = blockIdx.x; + const uint32_t tid = threadIdx.x; + const uint32_t lane = tid & 31u; + const uint32_t warp = tid >> 5u; + if (c >= n_comp || tid >= 128u) return; + if (causal) { + const uint32_t visible = ratio ? (pos0 + 1u) / ratio : n_comp; + if (c >= visible) { + if (tid == 0) scores[c] = -INFINITY; + return; + } + } + + __shared__ float krow[128]; + __shared__ float partial[4]; + if (tid < 128u) krow[tid] = index_comp[(uint64_t)c * 128u + tid]; + __syncthreads(); + + float total = 0.0f; + for (uint32_t h0 = 0; h0 < 64u; h0 += 4u) { + const uint32_t h = h0 + warp; + const float4 qv = ((const float4 *)(q + (uint64_t)h * 128u))[lane]; + const float4 kv = ((const float4 *)krow)[lane]; + float dot = qv.x * kv.x + qv.y * kv.y + qv.z * kv.z + qv.w * kv.w; + dot = warp_sum_f32(dot); + if (lane == 0) partial[warp] = fmaxf(dot, 0.0f) * weights[h] * scale; + __syncthreads(); + if (tid == 0) total += partial[0] + partial[1] + partial[2] + partial[3]; + __syncthreads(); + } + if (tid == 0) scores[c] = total; +} + +__global__ static void indexer_scores_wmma_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + int causal) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + const uint32_t tile_c = blockIdx.x * 16u; + const uint32_t tile_t = blockIdx.y * 16u; + const uint32_t tid = threadIdx.x; + if (tid >= 32u || head_dim != 128u) return; + + if (causal) { + const uint32_t last_token = min(tile_t + 16u, n_tokens); + const uint32_t max_visible = last_token > tile_t + ? min((pos0 + last_token) / ratio, n_comp) + : 0u; + if (tile_c >= max_visible) { + for (uint32_t i = tid; i < 16u * 16u; i += 32u) { + const uint32_t r = i >> 4u; + const uint32_t c = i & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + c; + if (token < n_tokens && comp < n_comp) { + scores[(uint64_t)token * n_comp + comp] = -INFINITY; + } + } + return; + } + } + + __shared__ __half a_sh[16 * 128]; + __shared__ __half b_sh[16 * 128]; + __shared__ float c_sh[16 * 16]; + __shared__ float acc_sh[16 * 16]; + + for (uint32_t i = tid; i < 16u * 16u; i += 32u) acc_sh[i] = 0.0f; + for (uint32_t i = tid; i < 16u * 128u; i += 32u) { + const uint32_t c = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t comp = tile_c + c; + float v = 0.0f; + if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; + b_sh[d + c * 128u] = __float2half(v); + } + __syncthreads(); + + for (uint32_t h = 0; h < n_head; h++) { + for (uint32_t i = tid; i < 16u * 128u; i += 32u) { + const uint32_t r = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t token = tile_t + r; + float v = 0.0f; + if (token < n_tokens) { + v = q[((uint64_t)token * n_head + h) * head_dim + d]; + } + a_sh[i] = __float2half(v); + } + __syncthreads(); + + wmma::fragment a_frag; + wmma::fragment b_frag; + wmma::fragment c_frag; + wmma::fill_fragment(c_frag, 0.0f); + for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { + wmma::load_matrix_sync(a_frag, a_sh + k0, 128); + wmma::load_matrix_sync(b_frag, b_sh + k0, 128); + wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + wmma::store_matrix_sync(c_sh, c_frag, 16, wmma::mem_row_major); + __syncthreads(); + + for (uint32_t i = tid; i < 16u * 16u; i += 32u) { + const uint32_t r = i >> 4u; + const uint32_t token = tile_t + r; + if (token < n_tokens) { + const float w = weights[(uint64_t)token * n_head + h]; + acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; + } + } + __syncthreads(); + } + + for (uint32_t i = tid; i < 16u * 16u; i += 32u) { + const uint32_t r = i >> 4u; + const uint32_t c = i & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + c; + if (token < n_tokens && comp < n_comp) { + float out = acc_sh[i] * scale; + if (causal) { + const uint32_t visible = (pos0 + token + 1u) / ratio; + if (comp >= visible) out = -INFINITY; + } + scores[(uint64_t)token * n_comp + comp] = out; + } + } +#endif +} + +__global__ static void indexer_scores_wmma32_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + int causal) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + const uint32_t tile_c = blockIdx.x * 32u; + const uint32_t tile_t = blockIdx.y * 16u; + const uint32_t tid = threadIdx.x; + const uint32_t warp = tid >> 5u; + if (tid >= 64u || head_dim != 128u) return; + + if (causal) { + const uint32_t last_token = min(tile_t + 16u, n_tokens); + const uint32_t max_visible = last_token > tile_t + ? min((pos0 + last_token) / ratio, n_comp) + : 0u; + if (tile_c >= max_visible) { + for (uint32_t i = tid; i < 16u * 32u; i += 64u) { + const uint32_t r = i >> 5u; + const uint32_t c = i & 31u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + c; + if (token < n_tokens && comp < n_comp) { + scores[(uint64_t)token * n_comp + comp] = -INFINITY; + } + } + return; + } + } + + __shared__ __half a_sh[16 * 128]; + __shared__ __half b_sh[32 * 128]; + __shared__ float c_sh[2 * 16 * 16]; + __shared__ float acc_sh[2 * 16 * 16]; + + for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) acc_sh[i] = 0.0f; + for (uint32_t i = tid; i < 32u * 128u; i += 64u) { + const uint32_t c = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t comp = tile_c + c; + float v = 0.0f; + if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; + b_sh[d + c * 128u] = __float2half(v); + } + __syncthreads(); + + for (uint32_t h = 0; h < n_head; h++) { + for (uint32_t i = tid; i < 16u * 128u; i += 64u) { + const uint32_t r = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t token = tile_t + r; + float v = 0.0f; + if (token < n_tokens) { + v = q[((uint64_t)token * n_head + h) * head_dim + d]; + } + a_sh[i] = __float2half(v); + } + __syncthreads(); + + wmma::fragment a_frag; + wmma::fragment b_frag; + wmma::fragment c_frag; + wmma::fill_fragment(c_frag, 0.0f); + const uint32_t col0 = warp * 16u; + for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { + wmma::load_matrix_sync(a_frag, a_sh + k0, 128); + wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); + wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); + __syncthreads(); + + for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + const float w = weights[(uint64_t)token * n_head + h]; + acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; + } + } + __syncthreads(); + } + + for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + float out = acc_sh[i] * scale; + if (causal) { + const uint32_t visible = (pos0 + token + 1u) / ratio; + if (comp >= visible) out = -INFINITY; + } + scores[(uint64_t)token * n_comp + comp] = out; + } + } +#endif +} + +__global__ static void indexer_scores_wmma64_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + int causal) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + const uint32_t tile_c = blockIdx.x * 64u; + const uint32_t tile_t = blockIdx.y * 16u; + const uint32_t tid = threadIdx.x; + const uint32_t warp = tid >> 5u; + if (tid >= 128u || head_dim != 128u) return; + + if (causal) { + const uint32_t last_token = min(tile_t + 16u, n_tokens); + const uint32_t max_visible = last_token > tile_t + ? min((pos0 + last_token) / ratio, n_comp) + : 0u; + if (tile_c >= max_visible) { + for (uint32_t i = tid; i < 16u * 64u; i += 128u) { + const uint32_t r = i >> 6u; + const uint32_t c = i & 63u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + c; + if (token < n_tokens && comp < n_comp) { + scores[(uint64_t)token * n_comp + comp] = -INFINITY; + } + } + return; + } + } + + __shared__ __half a_sh[16 * 128]; + __shared__ __half b_sh[64 * 128]; + __shared__ float c_sh[4 * 16 * 16]; + __shared__ float acc_sh[4 * 16 * 16]; + + for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) acc_sh[i] = 0.0f; + for (uint32_t i = tid; i < 64u * 128u; i += 128u) { + const uint32_t c = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t comp = tile_c + c; + float v = 0.0f; + if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; + b_sh[d + c * 128u] = __float2half(v); + } + __syncthreads(); + + for (uint32_t h = 0; h < n_head; h++) { + for (uint32_t i = tid; i < 16u * 128u; i += 128u) { + const uint32_t r = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t token = tile_t + r; + float v = 0.0f; + if (token < n_tokens) { + v = q[((uint64_t)token * n_head + h) * head_dim + d]; + } + a_sh[i] = __float2half(v); + } + __syncthreads(); + + wmma::fragment a_frag; + wmma::fragment b_frag; + wmma::fragment c_frag; + wmma::fill_fragment(c_frag, 0.0f); + const uint32_t col0 = warp * 16u; + for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { + wmma::load_matrix_sync(a_frag, a_sh + k0, 128); + wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); + wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); + __syncthreads(); + + for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + const float w = weights[(uint64_t)token * n_head + h]; + acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w; + } + } + __syncthreads(); + } + + for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + float out = acc_sh[i] * scale; + if (causal) { + const uint32_t visible = (pos0 + token + 1u) / ratio; + if (comp >= visible) out = -INFINITY; + } + scores[(uint64_t)token * n_comp + comp] = out; + } + } +#endif +} + +__global__ static void indexer_scores_wmma128_kernel( + float *scores, + const float *q, + const float *weights, + const float *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + int causal) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + const uint32_t tile_c = blockIdx.x * 128u; + const uint32_t tile_t = blockIdx.y * 16u; + const uint32_t tid = threadIdx.x; + const uint32_t warp = tid >> 5u; + if (tid >= 256u || head_dim != 128u) return; + + if (causal) { + const uint32_t last_token = min(tile_t + 16u, n_tokens); + const uint32_t max_visible = last_token > tile_t + ? min((pos0 + last_token) / ratio, n_comp) + : 0u; + if (tile_c >= max_visible) { + for (uint32_t i = tid; i < 16u * 128u; i += 256u) { + const uint32_t r = i >> 7u; + const uint32_t c = i & 127u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + c; + if (token < n_tokens && comp < n_comp) { + scores[(uint64_t)token * n_comp + comp] = -INFINITY; + } + } + return; + } + } + + __shared__ __half a_sh[16 * 128]; + __shared__ __half b_sh[128 * 128]; + __shared__ float c_sh[8 * 16 * 16]; + + float acc[8]; +#pragma unroll + for (uint32_t i = 0; i < 8u; i++) acc[i] = 0.0f; + + for (uint32_t i = tid; i < 128u * 128u; i += 256u) { + const uint32_t c = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t comp = tile_c + c; + float v = 0.0f; + if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d]; + b_sh[d + c * 128u] = __float2half(v); + } + __syncthreads(); + + for (uint32_t h = 0; h < n_head; h++) { + for (uint32_t i = tid; i < 16u * 128u; i += 256u) { + const uint32_t r = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t token = tile_t + r; + float v = 0.0f; + if (token < n_tokens) { + v = q[((uint64_t)token * n_head + h) * head_dim + d]; + } + a_sh[i] = __float2half(v); + } + __syncthreads(); + + wmma::fragment a_frag; + wmma::fragment b_frag; + wmma::fragment c_frag; + wmma::fill_fragment(c_frag, 0.0f); + const uint32_t col0 = warp * 16u; + for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { + wmma::load_matrix_sync(a_frag, a_sh + k0, 128); + wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128); + wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major); + __syncthreads(); + + const uint32_t local0 = tid & 255u; + const uint32_t token0 = tile_t + (local0 >> 4u); + const float w0 = token0 < n_tokens ? weights[(uint64_t)token0 * n_head + h] : 0.0f; + uint32_t slot = 0; + for (uint32_t i = tid; i < 8u * 16u * 16u; i += 256u, slot++) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + acc[slot] += fmaxf(c_sh[i], 0.0f) * w0; + } + } + __syncthreads(); + } + + uint32_t slot = 0; + for (uint32_t i = tid; i < 8u * 16u * 16u; i += 256u, slot++) { + const uint32_t wtile = i >> 8u; + const uint32_t local = i & 255u; + const uint32_t r = local >> 4u; + const uint32_t c = local & 15u; + const uint32_t token = tile_t + r; + const uint32_t comp = tile_c + wtile * 16u + c; + if (token < n_tokens && comp < n_comp) { + float out = acc[slot] * scale; + if (causal) { + const uint32_t visible = (pos0 + token + 1u) / ratio; + if (comp >= visible) out = -INFINITY; + } + scores[(uint64_t)token * n_comp + comp] = out; + } + } +#endif +} + +__global__ static void indexer_topk_kernel(uint32_t *selected, const float *scores, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) { + uint32_t t = blockIdx.x; + if (t >= n_tokens || threadIdx.x != 0) return; + const float *row = scores + (uint64_t)t * n_comp; + uint32_t *sel = selected + (uint64_t)t * top_k; + for (uint32_t k = 0; k < top_k; k++) sel[k] = 0; + for (uint32_t c = 0; c < n_comp; c++) { + float v = row[c]; + for (uint32_t k = 0; k < top_k; k++) { + if ((k >= c) || v > row[sel[k]]) { + for (uint32_t j = top_k - 1; j > k; j--) sel[j] = sel[j - 1]; + sel[k] = c; + break; + } + } + } +} + +__device__ __forceinline__ static bool topk_score_better(float av, uint32_t ai, float bv, uint32_t bi) { + return av > bv || (av == bv && ai < bi); +} + +__device__ __forceinline__ static void top2_insert_candidate( + float v, + uint32_t i, + float *v0, + uint32_t *i0, + float *v1, + uint32_t *i1) { + if (i == *i0 || i == *i1) return; + if (topk_score_better(v, i, *v0, *i0)) { + *v1 = *v0; + *i1 = *i0; + *v0 = v; + *i0 = i; + } else if (topk_score_better(v, i, *v1, *i1)) { + *v1 = v; + *i1 = i; + } +} + +/* DSpark markov chain step: out = argmax_i(logits[i] + dot(w2[i], w1[prev])) + * over the vocab, entirely on-device (logits row stays resident; the chain + * loop only reads back 4 bytes per draft). w1/w2 are q8_0 with 272-byte rows + * (8 blocks of 32). Single block; ~35 MB w2 read per step. */ +__global__ static void dspark_markov_argmax_kernel( + unsigned long long *out_key, + const float *logits, + const unsigned char *w1_row, + const unsigned char *w2, + uint32_t vocab, + uint32_t rank_blocks) { + __shared__ float state[256]; + const uint32_t tid = threadIdx.x; + if (tid < rank_blocks * 32u) { + const uint32_t b = tid >> 5, k = tid & 31u; + const unsigned char *blk = w1_row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + state[tid] = d * (float)((const int8_t *)(blk + 2))[k]; + } + __syncthreads(); + + float best_v = -INFINITY; + uint32_t best_i = 0; + for (uint32_t i = blockIdx.x * blockDim.x + tid; i < vocab; + i += gridDim.x * blockDim.x) { + const unsigned char *row = w2 + (uint64_t)i * rank_blocks * 34u; + float acc = 0.0f; + for (uint32_t b = 0; b < rank_blocks; b++) { + const unsigned char *blk = row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + const int8_t *q = (const int8_t *)(blk + 2); + float s = 0.0f; + #pragma unroll + for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * state[b * 32u + k]; + acc += d * s; + } + const float v = logits[i] + acc; + if (topk_score_better(v, i, best_v, best_i)) { + best_v = v; + best_i = i; + } + } + + __shared__ float vals[256]; + __shared__ uint32_t idxs[256]; + vals[tid] = best_v; + idxs[tid] = best_i; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0u; stride >>= 1u) { + if (tid < stride) { + if (topk_score_better(vals[tid + stride], idxs[tid + stride], + vals[tid], idxs[tid])) { + vals[tid] = vals[tid + stride]; + idxs[tid] = idxs[tid + stride]; + } + } + __syncthreads(); + } + if (tid == 0u) { + /* Monotonic float key; ~idx in the low bits makes ties resolve to + * the smaller index under atomicMax (matches topk_score_better). */ + const unsigned int f = __float_as_uint(vals[0]); + const unsigned int fkey = (f & 0x80000000u) ? ~f : (f | 0x80000000u); + const unsigned long long key = + ((unsigned long long)fkey << 32) | (unsigned int)(~idxs[0]); + atomicMax(out_key, key); + } +} + +__global__ static void indexer_top1_kernel( + uint32_t *selected, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens) { + const uint32_t t = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= 1024u) return; + + const float *row = scores + (uint64_t)t * n_comp; + float best_v = -INFINITY; + uint32_t best_i = 0; + for (uint32_t i = tid; i < n_comp; i += 1024u) { + const float v = row[i]; + if (topk_score_better(v, i, best_v, best_i)) { + best_v = v; + best_i = i; + } + } + + __shared__ float vals[1024]; + __shared__ uint32_t idxs[1024]; + vals[tid] = best_v; + idxs[tid] = best_i; + __syncthreads(); + + for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + const float ov = vals[tid + stride]; + const uint32_t oi = idxs[tid + stride]; + if (topk_score_better(ov, oi, vals[tid], idxs[tid])) { + vals[tid] = ov; + idxs[tid] = oi; + } + } + __syncthreads(); + } + + if (tid == 0u) selected[t] = idxs[0]; +} + +__global__ static void indexer_top1_value_kernel( + uint32_t *selected, + float *values, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t index_offset) { + const uint32_t t = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= 1024u) return; + + const float *row = scores + (uint64_t)t * n_comp; + float best_v = -INFINITY; + uint32_t best_i = 0; + for (uint32_t i = tid; i < n_comp; i += 1024u) { + const float v = row[i]; + const uint32_t gi = index_offset + i; + const uint32_t best_gi = index_offset + best_i; + if (topk_score_better(v, gi, best_v, best_gi)) { + best_v = v; + best_i = i; + } + } + + __shared__ float vals[1024]; + __shared__ uint32_t idxs[1024]; + vals[tid] = best_v; + idxs[tid] = best_i; + __syncthreads(); + + for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + const float ov = vals[tid + stride]; + const uint32_t oi = idxs[tid + stride]; + const uint32_t ogi = index_offset + oi; + const uint32_t gi = index_offset + idxs[tid]; + if (topk_score_better(ov, ogi, vals[tid], gi)) { + vals[tid] = ov; + idxs[tid] = oi; + } + } + __syncthreads(); + } + + if (tid == 0u) { + selected[t] = index_offset + idxs[0]; + values[t] = vals[0]; + } +} + +__global__ static void indexer_top2_value_kernel( + uint32_t *selected, + float *values, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t index_offset) { + const uint32_t t = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= 1024u) return; + + const float *row = scores + (uint64_t)t * n_comp; + float best0_v = -INFINITY; + float best1_v = -INFINITY; + uint32_t best0_i = UINT32_MAX; + uint32_t best1_i = UINT32_MAX; + for (uint32_t i = tid; i < n_comp; i += 1024u) { + const uint32_t gi = index_offset + i; + top2_insert_candidate(row[i], gi, + &best0_v, &best0_i, + &best1_v, &best1_i); + } + + __shared__ float vals0[1024]; + __shared__ float vals1[1024]; + __shared__ uint32_t idxs0[1024]; + __shared__ uint32_t idxs1[1024]; + vals0[tid] = best0_v; + vals1[tid] = best1_v; + idxs0[tid] = best0_i; + idxs1[tid] = best1_i; + __syncthreads(); + + for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + top2_insert_candidate(vals0[tid + stride], idxs0[tid + stride], + &vals0[tid], &idxs0[tid], + &vals1[tid], &idxs1[tid]); + top2_insert_candidate(vals1[tid + stride], idxs1[tid + stride], + &vals0[tid], &idxs0[tid], + &vals1[tid], &idxs1[tid]); + } + __syncthreads(); + } + + if (tid == 0u) { + selected[(uint64_t)t * 2u + 0u] = idxs0[0]; + selected[(uint64_t)t * 2u + 1u] = idxs1[0]; + values[(uint64_t)t * 2u + 0u] = vals0[0]; + values[(uint64_t)t * 2u + 1u] = vals1[0]; + } +} + +__device__ __forceinline__ static uint32_t topk_float_ordered_key(float v) { + const uint32_t u = __float_as_uint(v); + return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); +} + +__device__ __forceinline__ static uint64_t topk_pack_key(float v, uint32_t idx) { + return ((uint64_t)topk_float_ordered_key(v) << 32u) | (uint64_t)(0xffffffffu - idx); +} + +__global__ static void indexer_topk_8192_cub_kernel( + uint32_t *selected, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + constexpr uint32_t BLOCK_THREADS = 512u; + constexpr uint32_t ITEMS_PER_THREAD = 16u; + using BlockSort = cub::BlockRadixSort; + extern __shared__ __align__(16) unsigned char sort_smem[]; + typename BlockSort::TempStorage &sort_storage = + *reinterpret_cast(sort_smem); + + const uint32_t t = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= BLOCK_THREADS) return; + + const float *row = scores + (uint64_t)t * n_comp; + uint64_t keys[ITEMS_PER_THREAD]; +#pragma unroll + for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) { + const uint32_t i = tid * ITEMS_PER_THREAD + item; + if (i < n_comp) { + keys[item] = topk_pack_key(row[i], i); + } else { + keys[item] = topk_pack_key(-INFINITY, UINT32_MAX); + } + } + + BlockSort(sort_storage).SortDescending(keys); + +#pragma unroll + for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) { + const uint32_t i = tid * ITEMS_PER_THREAD + item; + if (i < top_k) { + selected[(uint64_t)t * top_k + i] = 0xffffffffu - (uint32_t)keys[item]; + } + } +} + +__global__ static void indexer_topk_1024_kernel( + uint32_t *selected, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + uint32_t t = blockIdx.x; + uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= 1024u) return; + __shared__ float vals[1024]; + __shared__ uint32_t idxs[1024]; + + const float *row = scores + (uint64_t)t * n_comp; + if (tid < n_comp) { + vals[tid] = row[tid]; + idxs[tid] = tid; + } else { + vals[tid] = -INFINITY; + idxs[tid] = UINT32_MAX; + } + __syncthreads(); + + for (uint32_t k = 2u; k <= 1024u; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + uint32_t other = tid ^ j; + if (other > tid && other < 1024u) { + const float av = vals[tid]; + const float bv = vals[other]; + const uint32_t ai = idxs[tid]; + const uint32_t bi = idxs[other]; + const bool desc_half = (tid & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[tid] = bv; + idxs[tid] = bi; + vals[other] = av; + idxs[other] = ai; + } + } + __syncthreads(); + } + } + + if (tid < top_k) selected[(uint64_t)t * top_k + tid] = idxs[tid]; +} + +template +__global__ static void indexer_topk_pow2_kernel( + uint32_t *selected, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + uint32_t t = blockIdx.x; + uint32_t tid = threadIdx.x; + if (t >= n_tokens) return; + __shared__ float vals[SORT_N]; + __shared__ uint32_t idxs[SORT_N]; + + const float *row = scores + (uint64_t)t * n_comp; + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + if (i < n_comp) { + vals[i] = row[i]; + idxs[i] = i; + } else { + vals[i] = -INFINITY; + idxs[i] = UINT32_MAX; + } + } + __syncthreads(); + + for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t other = i ^ j; + if (other > i && other < SORT_N) { + const float av = vals[i]; + const float bv = vals[other]; + const uint32_t ai = idxs[i]; + const uint32_t bi = idxs[other]; + const bool desc_half = (i & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[i] = bv; + idxs[i] = bi; + vals[other] = av; + idxs[other] = ai; + } + } + } + __syncthreads(); + } + } + + for (uint32_t i = tid; i < top_k; i += blockDim.x) { + selected[(uint64_t)t * top_k + i] = idxs[i]; + } +} + +template +__global__ static void indexer_topk_pow2_u16_kernel( + uint32_t *selected, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + uint32_t t = blockIdx.x; + uint32_t tid = threadIdx.x; + if (t >= n_tokens) return; + __shared__ float vals[SORT_N]; + __shared__ uint16_t idxs[SORT_N]; + + const float *row = scores + (uint64_t)t * n_comp; + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + if (i < n_comp) { + vals[i] = row[i]; + idxs[i] = (uint16_t)i; + } else { + vals[i] = -INFINITY; + idxs[i] = UINT16_MAX; + } + } + __syncthreads(); + + for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t other = i ^ j; + if (other > i && other < SORT_N) { + const float av = vals[i]; + const float bv = vals[other]; + const uint32_t ai = idxs[i]; + const uint32_t bi = idxs[other]; + const bool desc_half = (i & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[i] = bv; + idxs[i] = (uint16_t)bi; + vals[other] = av; + idxs[other] = (uint16_t)ai; + } + } + } + __syncthreads(); + } + } + + for (uint32_t i = tid; i < top_k; i += blockDim.x) { + selected[(uint64_t)t * top_k + i] = idxs[i]; + } +} + +template +__global__ static void indexer_topk_chunk_pow2_kernel( + uint32_t *candidates, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k, + uint32_t candidate_stride) { + uint32_t t = blockIdx.x; + uint32_t chunk = blockIdx.y; + uint32_t tid = threadIdx.x; + if (t >= n_tokens) return; + + const uint32_t chunk_start = chunk * SORT_N; + if (chunk_start >= n_comp) return; + const uint32_t chunk_n = n_comp - chunk_start < SORT_N ? n_comp - chunk_start : SORT_N; + __shared__ float vals[SORT_N]; + __shared__ uint32_t idxs[SORT_N]; + + const float *row = scores + (uint64_t)t * n_comp; + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + if (i < chunk_n) { + vals[i] = row[chunk_start + i]; + idxs[i] = chunk_start + i; + } else { + vals[i] = -INFINITY; + idxs[i] = UINT32_MAX; + } + } + __syncthreads(); + + for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t other = i ^ j; + if (other > i && other < SORT_N) { + const float av = vals[i]; + const float bv = vals[other]; + const uint32_t ai = idxs[i]; + const uint32_t bi = idxs[other]; + const bool desc_half = (i & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[i] = bv; + idxs[i] = bi; + vals[other] = av; + idxs[other] = ai; + } + } + } + __syncthreads(); + } + } + + uint32_t *out = candidates + (uint64_t)t * candidate_stride + chunk * top_k; + for (uint32_t i = tid; i < top_k; i += blockDim.x) { + out[i] = idxs[i]; + } +} + +template +__global__ static void indexer_topk_merge_pow2_kernel( + uint32_t *selected, + const uint32_t *candidates, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k, + uint32_t candidate_count, + uint32_t candidate_stride) { + uint32_t t = blockIdx.x; + uint32_t tid = threadIdx.x; + if (t >= n_tokens) return; + __shared__ float vals[SORT_N]; + __shared__ uint32_t idxs[SORT_N]; + + const float *row = scores + (uint64_t)t * n_comp; + const uint32_t *cand = candidates + (uint64_t)t * candidate_stride; + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t idx = UINT32_MAX; + float v = -INFINITY; + if (i < candidate_count) { + idx = cand[i]; + if (idx < n_comp) v = row[idx]; + } + vals[i] = v; + idxs[i] = idx; + } + __syncthreads(); + + for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t other = i ^ j; + if (other > i && other < SORT_N) { + const float av = vals[i]; + const float bv = vals[other]; + const uint32_t ai = idxs[i]; + const uint32_t bi = idxs[other]; + const bool desc_half = (i & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[i] = bv; + idxs[i] = bi; + vals[other] = av; + idxs[other] = ai; + } + } + } + __syncthreads(); + } + } + + for (uint32_t i = tid; i < top_k; i += blockDim.x) { + selected[(uint64_t)t * top_k + i] = idxs[i]; + } +} + +template +__global__ static void indexer_topk_tree_merge_pow2_kernel( + uint32_t *out, + const uint32_t *candidates, + const float *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k, + uint32_t n_sets, + uint32_t merge_group, + uint32_t candidate_stride, + uint32_t out_stride) { + uint32_t t = blockIdx.x; + uint32_t group = blockIdx.y; + uint32_t tid = threadIdx.x; + if (t >= n_tokens) return; + + const uint32_t set0 = group * merge_group; + if (set0 >= n_sets) return; + uint32_t set_count = n_sets - set0; + if (set_count > merge_group) set_count = merge_group; + const uint32_t candidate_count = set_count * top_k; + + __shared__ float vals[SORT_N]; + __shared__ uint32_t idxs[SORT_N]; + + const float *row = scores + (uint64_t)t * n_comp; + const uint32_t *cand = candidates + (uint64_t)t * candidate_stride + set0 * top_k; + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t idx = UINT32_MAX; + float v = -INFINITY; + if (i < candidate_count) { + idx = cand[i]; + if (idx < n_comp) v = row[idx]; + } + vals[i] = v; + idxs[i] = idx; + } + __syncthreads(); + + for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < SORT_N; i += blockDim.x) { + uint32_t other = i ^ j; + if (other > i && other < SORT_N) { + const float av = vals[i]; + const float bv = vals[other]; + const uint32_t ai = idxs[i]; + const uint32_t bi = idxs[other]; + const bool desc_half = (i & k) == 0u; + const bool swap = desc_half + ? topk_score_better(bv, bi, av, ai) + : topk_score_better(av, ai, bv, bi); + if (swap) { + vals[i] = bv; + idxs[i] = bi; + vals[other] = av; + idxs[other] = ai; + } + } + } + __syncthreads(); + } + } + + uint32_t *dst = out + (uint64_t)t * out_stride + group * top_k; + for (uint32_t i = tid; i < top_k; i += blockDim.x) { + dst[i] = idxs[i]; + } +} + +__global__ static void indexed_topk_sort_512_asc_kernel( + int32_t *dst, + const int32_t *src, + uint32_t n_tokens) { + const uint32_t t = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (t >= n_tokens || tid >= 512u) return; + __shared__ int32_t rows[512]; + + const int32_t *src_row = src + (uint64_t)t * 512u; + int32_t *dst_row = dst + (uint64_t)t * 512u; + rows[tid] = src_row[tid]; + __syncthreads(); + + for (uint32_t k = 2u; k <= 512u; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + const uint32_t other = tid ^ j; + if (other > tid && other < 512u) { + const int32_t a = rows[tid]; + const int32_t b = rows[other]; + const bool up = (tid & k) == 0u; + if ((up && a > b) || (!up && a < b)) { + rows[tid] = b; + rows[other] = a; + } + } + __syncthreads(); + } + } + + dst_row[tid] = rows[tid]; +} + +__global__ static void topk_mask_kernel(float *mask, const uint32_t *topk, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_comp; + if (gid >= n) return; + uint32_t t = gid / n_comp; + uint32_t c = gid - (uint64_t)t * n_comp; + float v = -INFINITY; + for (uint32_t k = 0; k < top_k; k++) { + if (topk[(uint64_t)t * top_k + k] == c) { + v = 0.0f; + break; + } + } + mask[gid] = v; +} + +extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n_vocab, uint32_t token, uint32_t n_embd, uint32_t n_hc) { + (void)n_vocab; + if (!out_hc || !model_map || weight_offset >= model_size) return 0; + uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0; + const int logical_tier = ds4_tensor_device_idx(out_hc); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "token_embd"); + if (!wptr) return 0; + uint32_t n = n_embd * n_hc; + embed_token_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc); + return cuda_ok(cudaGetLastError(), "embed token launch"); +} + +extern "C" int ds4_gpu_embed_tokens_hc_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *tokens_t, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!out_hc || !tokens_t || !model_map || + weight_offset > model_size || + (uint64_t)n_vocab * n_embd * sizeof(uint16_t) > model_size - weight_offset || + tokens_t->bytes < (uint64_t)n_tokens * sizeof(int32_t) || + out_hc->bytes < (uint64_t)n_tokens * n_hc * n_embd * sizeof(float)) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out_hc); + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, + (uint64_t)n_vocab * n_embd * sizeof(uint16_t), + logical_tier, + "token_embd"); + if (!wptr) return 0; + uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; + embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>( + (float *)out_hc->ptr, + (const int32_t *)tokens_t->ptr, + (const __half *)wptr, + n_vocab, n_tokens, n_embd, n_hc); + return cuda_ok(cudaGetLastError(), "embed tokens launch"); +} + +static int indexer_scores_launch( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale, + uint32_t causal) { + if (!scores || !q || !weights || !index_comp || + n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) || + index_comp->bytes < (uint64_t)n_comp * head_dim * sizeof(float) || + scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float)) { + return 0; + } + if (causal && ratio == 0) return 0; + if (n_tokens == 1u && head_dim == 128u && n_head == 64u && + getenv("DS4_CUDA_NO_INDEXER_DIRECT_ONE") == NULL) { + indexer_score_one_direct_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, pos0, ratio, + scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer score one direct launch"); + } + if (!g_quality_mode && head_dim == 128u && n_head == 64u && + getenv("DS4_CUDA_NO_INDEXER_WMMA") == NULL) { + if (getenv("DS4_CUDA_NO_INDEXER_WMMA128") == NULL) { + dim3 grid((n_comp + 127u) / 128u, (n_tokens + 15u) / 16u, 1); + indexer_scores_wmma128_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, n_tokens, pos0, n_head, + head_dim, ratio, scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer scores wmma128 launch"); + } else if (getenv("DS4_CUDA_NO_INDEXER_WMMA64") == NULL) { + dim3 grid((n_comp + 63u) / 64u, (n_tokens + 15u) / 16u, 1); + indexer_scores_wmma64_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, n_tokens, pos0, n_head, + head_dim, ratio, scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer scores wmma64 launch"); + } else if (getenv("DS4_CUDA_NO_INDEXER_WMMA32") == NULL) { + dim3 grid((n_comp + 31u) / 32u, (n_tokens + 15u) / 16u, 1); + indexer_scores_wmma32_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, n_tokens, pos0, n_head, + head_dim, ratio, scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer scores wmma32 launch"); + } else { + dim3 grid((n_comp + 15u) / 16u, (n_tokens + 15u) / 16u, 1); + indexer_scores_wmma_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, n_tokens, pos0, n_head, + head_dim, ratio, scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer scores wmma launch"); + } + } + dim3 grid(n_comp, n_tokens, 1); + indexer_scores_kernel<<>>((float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)index_comp->ptr, + n_comp, n_tokens, pos0, n_head, + head_dim, ratio, scale, causal ? 1 : 0); + return cuda_ok(cudaGetLastError(), "indexer scores launch"); +} + +extern "C" int ds4_gpu_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_head, + uint32_t head_dim, + float scale) { + return indexer_scores_launch(scores, q, weights, index_comp, n_comp, 1, 0, + n_head, head_dim, 1, scale, 0); +} + +extern "C" int ds4_gpu_indexer_scores_prefill_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, 0, + n_head, head_dim, ratio, scale, 1); +} + +extern "C" int ds4_gpu_indexer_scores_decode_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, pos0, + n_head, head_dim, ratio, scale, 1); +} + +extern "C" int ds4_gpu_dspark_markov_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits_row, + const void *model_map, + uint64_t model_size, + uint64_t w1_offset, + uint64_t w2_offset, + uint32_t prev_token, + uint32_t vocab, + uint32_t rank) { + if (!out_idx || !logits_row || !model_map || vocab == 0 || + rank == 0 || (rank & 31u) != 0u || rank > 256u || + out_idx->bytes < sizeof(unsigned long long) || + logits_row->bytes < (uint64_t)vocab * sizeof(float)) { + return 0; + } + const uint32_t rank_blocks = rank / 32u; + const uint64_t row_bytes = (uint64_t)rank_blocks * 34u; + if (w1_offset > model_size || + (uint64_t)prev_token * row_bytes + row_bytes > model_size - w1_offset || + w2_offset > model_size || + (uint64_t)vocab * row_bytes > model_size - w2_offset) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(logits_row); + const unsigned char *w1_row = (const unsigned char *)cuda_resolve_weight_ptr( + model_map, w1_offset + (uint64_t)prev_token * row_bytes, + row_bytes, logical_tier, "markov_w1_row"); + const unsigned char *w2 = (const unsigned char *)cuda_resolve_weight_ptr( + model_map, w2_offset, (uint64_t)vocab * row_bytes, + logical_tier, "markov_w2"); + if (!w1_row || !w2) return 0; + int dev_save = 0; + if (cudaGetDevice(&dev_save) != cudaSuccess) return 0; + if (logical_tier != dev_save && cudaSetDevice(logical_tier) != cudaSuccess) { + return 0; + } + int rc = cudaMemsetAsync(out_idx->ptr, 0, + sizeof(unsigned long long)) == cudaSuccess; + if (rc) { + dspark_markov_argmax_kernel<<<128, 256>>>( + (unsigned long long *)out_idx->ptr, + (const float *)logits_row->ptr, + w1_row, w2, vocab, rank_blocks); + rc = cuda_ok(cudaGetLastError(), "dspark markov argmax launch"); + } + if (logical_tier != dev_save) (void)cudaSetDevice(dev_save); + return rc; +} + +extern "C" int ds4_gpu_indexer_topk_tensor( + ds4_gpu_tensor *selected, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || + top_k > n_comp || + scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) { + return 0; + } + if (top_k == 1u && !g_cuda_no_top1) { + indexer_top1_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, + n_tokens); + return cuda_ok(cudaGetLastError(), "indexer top1 launch"); + } + if (top_k == 2048u && n_comp <= 4096u && + getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) { + indexer_topk_pow2_kernel<4096><<>>( + (uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 2048-wide launch"); + } + if (top_k == 2048u && n_comp > 4096u && + getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) { + const uint32_t chunk_n = 4096u; + const uint32_t merge_group = 2u; + const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n; + const uint32_t candidate_stride = n_chunks * top_k; + uint32_t n_sets = n_chunks; + uint64_t scratch_u32_per_token = candidate_stride; + while (n_sets > merge_group) { + n_sets = (n_sets + merge_group - 1u) / merge_group; + scratch_u32_per_token += (uint64_t)n_sets * top_k; + } + if (scratch_u32_per_token > + UINT64_MAX / n_tokens / sizeof(uint32_t)) { + return 0; + } + int exec_tier = ds4_tensor_device_idx(selected); + int current_device = -1; + if (cudaGetDevice(¤t_device) == cudaSuccess) { + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == current_device) { + exec_tier = t; + break; + } + } + } + const uint64_t tmp_bytes = + (uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t); + uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on( + exec_tier, tmp_bytes, "indexer topk 2048-wide tree"); + if (!scratch) return 0; + + uint32_t *cur = scratch; + n_sets = n_chunks; + uint32_t cur_stride = candidate_stride; + dim3 grid_chunks(n_tokens, n_chunks, 1); + indexer_topk_chunk_pow2_kernel<4096><<>>( + cur, (const float *)scores->ptr, + n_comp, n_tokens, top_k, candidate_stride); + if (!cuda_ok(cudaGetLastError(), + "indexer topk 2048-wide chunk launch")) { + return 0; + } + + while (n_sets > merge_group) { + const uint32_t next_sets = + (n_sets + merge_group - 1u) / merge_group; + const uint32_t next_stride = next_sets * top_k; + uint32_t *next = cur + (uint64_t)n_tokens * cur_stride; + dim3 grid_merge(n_tokens, next_sets, 1); + indexer_topk_tree_merge_pow2_kernel<4096><<>>( + next, cur, (const float *)scores->ptr, + n_comp, n_tokens, top_k, n_sets, merge_group, + cur_stride, next_stride); + if (!cuda_ok(cudaGetLastError(), + "indexer topk 2048-wide merge launch")) { + return 0; + } + cur = next; + n_sets = next_sets; + cur_stride = next_stride; + } + + indexer_topk_merge_pow2_kernel<4096><<>>( + (uint32_t *)selected->ptr, + cur, (const float *)scores->ptr, + n_comp, n_tokens, top_k, n_sets * top_k, cur_stride); + return cuda_ok(cudaGetLastError(), + "indexer topk 2048-wide final launch"); + } + if (top_k == 512u && n_comp <= 1024u && + getenv("DS4_CUDA_NO_TOPK1024") == NULL) { + indexer_topk_1024_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 1024 launch"); + } + if (top_k == 512u && n_comp <= 2048u && + getenv("DS4_CUDA_NO_TOPK2048") == NULL) { + indexer_topk_pow2_kernel<2048><<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 2048 launch"); + } + if (top_k == 512u && n_comp <= 4096u && + getenv("DS4_CUDA_NO_TOPK2048") == NULL) { + if (n_comp == 4096u) { + using TopkCubSort = cub::BlockRadixSort; + const int smem = (int)sizeof(typename TopkCubSort::TempStorage); + int dev = 0; + int max_optin_smem = 0; + cudaError_t attr_err = cudaGetDevice(&dev); + if (attr_err == cudaSuccess) { + attr_err = cudaDeviceGetAttribute(&max_optin_smem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + dev); + } + if (attr_err == cudaSuccess && max_optin_smem >= smem) { + attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + if (attr_err == cudaSuccess) { + indexer_topk_8192_cub_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 4096 cub launch"); + } + } + } + indexer_topk_pow2_kernel<4096><<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 4096 launch"); + } + if (top_k == 512u && n_comp <= 8192u && + getenv("DS4_CUDA_NO_TOPK2048") == NULL && + getenv("DS4_CUDA_NO_TOPK8192") == NULL) { + if (n_comp > 4096u) { + using TopkCubSort = cub::BlockRadixSort; + const int smem = (int)sizeof(typename TopkCubSort::TempStorage); + int dev = 0; + int max_optin_smem = 0; + cudaError_t attr_err = cudaGetDevice(&dev); + if (attr_err == cudaSuccess) { + attr_err = cudaDeviceGetAttribute(&max_optin_smem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + dev); + } + if (attr_err == cudaSuccess && max_optin_smem >= smem) { + attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + if (attr_err == cudaSuccess) { + indexer_topk_8192_cub_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 8192 cub launch"); + } + } + } + indexer_topk_pow2_u16_kernel<8192><<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 8192 launch"); + } + if (top_k == 512u && getenv("DS4_CUDA_NO_TOPK2048") == NULL && + getenv("DS4_CUDA_NO_TOPK_CHUNKED") == NULL) { + const uint32_t chunk_n = 4096u; + const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n; + const uint32_t candidate_stride = n_chunks * top_k; + uint32_t n_sets = n_chunks; + uint64_t scratch_u32_per_token = candidate_stride; + while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) { + n_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP; + scratch_u32_per_token += (uint64_t)n_sets * top_k; + } + if (scratch_u32_per_token > UINT64_MAX / n_tokens / sizeof(uint32_t)) return 0; + const uint64_t tmp_bytes = (uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t); + const int logical_tier = ds4_tensor_device_idx(selected); + uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "indexer topk tree"); + if (!scratch) return 0; + + uint32_t *cur = scratch; + n_sets = n_chunks; + uint32_t cur_stride = candidate_stride; + dim3 grid_chunks(n_tokens, n_chunks, 1); + indexer_topk_chunk_pow2_kernel<4096><<>>(cur, + (const float *)scores->ptr, + n_comp, + n_tokens, + top_k, + candidate_stride); + if (!cuda_ok(cudaGetLastError(), "indexer topk chunk launch")) return 0; + + while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) { + const uint32_t next_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP; + const uint32_t next_stride = next_sets * top_k; + uint32_t *next = cur + (uint64_t)n_tokens * cur_stride; + dim3 grid_merge(n_tokens, next_sets, 1); + indexer_topk_tree_merge_pow2_kernel<4096><<>>( + next, + cur, + (const float *)scores->ptr, + n_comp, + n_tokens, + top_k, + n_sets, + DS4_CUDA_TOPK_MERGE_GROUP, + cur_stride, + next_stride); + if (!cuda_ok(cudaGetLastError(), "indexer topk tree merge launch")) return 0; + cur = next; + n_sets = next_sets; + cur_stride = next_stride; + } + + indexer_topk_merge_pow2_kernel<4096><<>>((uint32_t *)selected->ptr, + cur, + (const float *)scores->ptr, + n_comp, + n_tokens, + top_k, + n_sets * top_k, + cur_stride); + return cuda_ok(cudaGetLastError(), "indexer topk tree final launch"); + } + indexer_topk_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk launch"); +} + +extern "C" int ds4_gpu_indexer_top1_value_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *values, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t index_offset) { + if (!selected || !values || !scores || n_comp == 0 || n_tokens == 0 || + scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * sizeof(uint32_t) || + values->bytes < (uint64_t)n_tokens * sizeof(float)) { + return 0; + } + indexer_top1_value_kernel<<>>((uint32_t *)selected->ptr, + (float *)values->ptr, + (const float *)scores->ptr, + n_comp, + n_tokens, + index_offset); + return cuda_ok(cudaGetLastError(), "indexer top1 value launch"); +} + +extern "C" int ds4_gpu_indexer_top2_value_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *values, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t index_offset) { + if (!selected || !values || !scores || n_comp < 2u || n_tokens == 0 || + scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * 2u * sizeof(uint32_t) || + values->bytes < (uint64_t)n_tokens * 2u * sizeof(float)) { + return 0; + } + indexer_top2_value_kernel<<>>((uint32_t *)selected->ptr, + (float *)values->ptr, + (const float *)scores->ptr, + n_comp, + n_tokens, + index_offset); + return cuda_ok(cudaGetLastError(), "indexer top2 value launch"); +} + +extern "C" int ds4_gpu_dsv4_topk_mask_tensor( + ds4_gpu_tensor *mask, + const ds4_gpu_tensor *topk, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0 || + mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) || + topk->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) { + return 0; + } + uint64_t n = (uint64_t)n_tokens * n_comp; + uint64_t nk = (uint64_t)n_tokens * top_k; + uint64_t blocks = ((n > nk ? n : nk) + 255) / 256; + topk_mask_kernel<<>>((float *)mask->ptr, + (const uint32_t *)topk->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "topk mask launch"); +} diff --git a/models/deepseek/cuda/dense_attention.inc b/models/deepseek/cuda/dense_attention.inc new file mode 100644 index 0000000000..ad374df521 --- /dev/null +++ b/models/deepseek/cuda/dense_attention.inc @@ -0,0 +1,5774 @@ +__global__ static void embed_token_hc_kernel(float *out, const unsigned short *w, uint32_t token, uint32_t n_embd, uint32_t n_hc) { + uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t n = n_embd * n_hc; + if (i >= n) return; + uint32_t e = i % n_embd; + out[i] = __half2float(reinterpret_cast(w)[(uint64_t)token * n_embd + e]); +} + +__global__ static void embed_tokens_hc_kernel( + float *out, + const int32_t *tokens, + const __half *w, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; + if (gid >= n) return; + uint32_t d = gid % n_embd; + uint64_t tmp = gid / n_embd; + uint32_t t = tmp / n_hc; + int32_t tok_i = tokens[t]; + uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i; + if (tok >= n_vocab) tok = 0; + out[gid] = __half2float(w[(uint64_t)tok * n_embd + d]); +} + +__global__ static void matmul_f16_kernel( + float *out, + const __half *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok) return; + + float sum = 0.0f; + const __half *wr = w + row * in_dim; + const float *xr = x + tok * in_dim; + for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) { + sum += __half2float(wr[i]) * xr[i]; + } + + __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(); + } + if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; +} + +__global__ static void matmul_f16_serial_kernel( + float *out, + const __half *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok || threadIdx.x != 0) return; + + float sum = 0.0f; + const __half *wr = w + row * in_dim; + const float *xr = x + tok * in_dim; + for (uint64_t i = 0; i < in_dim; i++) { + sum += __half2float(wr[i]) * xr[i]; + } + out[tok * out_dim + row] = sum; +} + +__global__ static void matmul_f16_ordered_chunks_kernel( + float *out, + const __half *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok) return; + + __shared__ float partial[32]; + const uint32_t tid = threadIdx.x; + float sum = 0.0f; + const uint64_t chunk = (in_dim + 31u) / 32u; + const uint64_t k0 = (uint64_t)tid * chunk; + uint64_t k1 = k0 + chunk; + if (k1 > in_dim) k1 = in_dim; + const __half *wr = w + row * in_dim; + const float *xr = x + tok * in_dim; + for (uint64_t i = k0; i < k1; i++) { + sum += __half2float(wr[i]) * xr[i]; + } + partial[tid] = sum; + __syncthreads(); + if (tid == 0) { + float total = 0.0f; + for (uint32_t i = 0; i < 32u; i++) total += partial[i]; + out[tok * out_dim + row] = total; + } +} + +__global__ static void matmul_f16_small_out_hx_ordered_chunks_kernel( + float *out, + const __half *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim) { + uint64_t row = (uint64_t)blockIdx.x; + if (row >= out_dim) return; + + __shared__ float partial[32]; + const uint32_t tid = threadIdx.x; + float sum = 0.0f; + const uint64_t chunk = (in_dim + 31u) / 32u; + const uint64_t k0 = (uint64_t)tid * chunk; + uint64_t k1 = k0 + chunk; + if (k1 > in_dim) k1 = in_dim; + const __half *wr = w + row * in_dim; + for (uint64_t i = k0; i < k1; i++) { + const float xv = __half2float(__float2half(x[i])); + sum += __half2float(wr[i]) * xv; + } + partial[tid] = sum; + __syncthreads(); + if (tid == 0) { + float total = 0.0f; + for (uint32_t i = 0; i < 32u; i++) total += partial[i]; + out[row] = total; + } +} + +__global__ static void matmul_f16_small_out_batch_kernel( + float *out, + const __half *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + const uint64_t tok = (uint64_t)blockIdx.x; + const uint32_t tid = threadIdx.x; + if (tok >= n_tok || out_dim > 32u || blockDim.x != 256u) return; + + float acc[32]; + #pragma unroll + for (uint32_t r = 0; r < 32u; r++) acc[r] = 0.0f; + + const float *xr = x + tok * in_dim; + for (uint64_t i = tid; i < in_dim; i += 256u) { + const float xv = xr[i]; + #pragma unroll + for (uint32_t r = 0; r < 32u; r++) { + if (r < out_dim) { + acc[r] += __half2float(w[(uint64_t)r * in_dim + i]) * xv; + } + } + } + + __shared__ float partial[32 * 256]; + #pragma unroll + for (uint32_t r = 0; r < 32u; r++) { + if (r < out_dim) partial[r * 256u + tid] = acc[r]; + } + __syncthreads(); + + for (uint32_t stride = 128u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + #pragma unroll + for (uint32_t r = 0; r < 32u; r++) { + if (r < out_dim) { + partial[r * 256u + tid] += partial[r * 256u + tid + stride]; + } + } + } + __syncthreads(); + } + + if (tid == 0) { + #pragma unroll + for (uint32_t r = 0; r < 32u; r++) { + if (r < out_dim) out[tok * out_dim + r] = partial[r * 256u]; + } + } +} + +__global__ static void matmul_f16_pair_ordered_chunks_kernel( + float *out0, + float *out1, + const __half *w0, + const __half *w1, + const float *x, + uint64_t in_dim, + uint64_t out0_dim, + uint64_t out1_dim) { + uint64_t row = (uint64_t)blockIdx.x; + if (row >= out0_dim && row >= out1_dim) return; + + __shared__ float partial0[32]; + __shared__ float partial1[32]; + const uint32_t tid = threadIdx.x; + float sum0 = 0.0f; + float sum1 = 0.0f; + const uint64_t chunk = (in_dim + 31u) / 32u; + const uint64_t k0 = (uint64_t)tid * chunk; + uint64_t k1 = k0 + chunk; + if (k1 > in_dim) k1 = in_dim; + const __half *wr0 = row < out0_dim ? w0 + row * in_dim : w0; + const __half *wr1 = row < out1_dim ? w1 + row * in_dim : w1; + for (uint64_t i = k0; i < k1; i++) { + const float xv = x[i]; + if (row < out0_dim) sum0 += __half2float(wr0[i]) * xv; + if (row < out1_dim) sum1 += __half2float(wr1[i]) * xv; + } + partial0[tid] = sum0; + partial1[tid] = sum1; + __syncthreads(); + if (tid == 0) { + float total0 = 0.0f; + float total1 = 0.0f; + for (uint32_t i = 0; i < 32u; i++) { + total0 += partial0[i]; + total1 += partial1[i]; + } + if (row < out0_dim) out0[row] = total0; + if (row < out1_dim) out1[row] = total1; + } +} + +__global__ static void matmul_f32_kernel( + float *out, + const float *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok) return; + + float sum = 0.0f; + const float *wr = w + row * in_dim; + const float *xr = x + tok * in_dim; + for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) { + sum += wr[i] * xr[i]; + } + + __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(); + } + if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; +} + +__global__ static void repeat_hc_kernel(float *out, const float *row, uint32_t n_embd, uint32_t n_hc) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_embd * n_hc; + if (i >= n) return; + out[i] = row[i % n_embd]; +} + +__global__ static void repeat_hc_rows_kernel(float *out, const float *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; + if (i >= n) return; + + uint64_t hc_row = (uint64_t)n_hc * n_embd; + uint64_t tok = i / hc_row; + uint64_t embd = i % n_embd; + out[i] = rows[tok * n_embd + embd]; +} + +__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_rows * n_slots * width; + if (i >= n) return; + + uint64_t col = i % width; + uint64_t slot = (i / width) % n_slots; + uint64_t row = i / ((uint64_t)n_slots * width); + out[i] = slots[((slot * slot_cap) + row) * width + col]; +} + +__global__ static void f32_to_f16_kernel(__half *out, const float *x, uint64_t n) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = __float2half(x[i]); +} + +__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 warp_max_f32(float v) { + for (int offset = 16; offset > 0; offset >>= 1) { + v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, offset)); + } + return v; +} + +__device__ static float dot4_f32(float4 a, float4 b) { + return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w; +} + +__device__ __forceinline__ static int32_t load_i8x4_i32_aligned(const int8_t *p) { + return *(const int32_t *)p; +} + +__device__ __forceinline__ static 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__ static 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__ static 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; +} + +__global__ static DS4_CUDA_UNUSED void matmul_q8_0_kernel( + float *out, + const unsigned char *w, + const float *x, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok) return; + const uint64_t blocks = (in_dim + 31) / 32; + const unsigned char *wr = w + row * blocks * 34; + const float *xr = x + tok * in_dim; + float acc = 0.0f; + + for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { + uint64_t i0 = b * 32; + uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32; + float amax = 0.0f; + for (uint64_t i = 0; i < bn; i++) amax = fmaxf(amax, fabsf(xr[i0 + i])); + float d = amax / 127.0f; + float id = d != 0.0f ? 1.0f / d : 0.0f; + const __half *scale_h = (const __half *)(wr + b * 34); + const int8_t *qs = (const int8_t *)(wr + b * 34 + 2); + int dot = 0; + for (uint64_t i = 0; i < bn; i++) { + int q = (int)lrintf(xr[i0 + i] * id); + q = q > 127 ? 127 : (q < -128 ? -128 : q); + dot += (int)qs[i] * q; + } + acc += __half2float(*scale_h) * d * (float)dot; + } + + __shared__ float partial[256]; + partial[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; +} + +__global__ static void quantize_q8_0_f32_kernel( + int8_t *xq, + float *xscale, + const float *x, + uint64_t in_dim, + uint64_t blocks) { + uint64_t b = blockIdx.x; + uint64_t tok = blockIdx.y; + if (b >= blocks) return; + uint64_t i0 = b * 32; + uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32; + const float *xr = x + tok * in_dim + i0; + + float a = 0.0f; + if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]); + __shared__ float vals[32]; + vals[threadIdx.x] = a; + __syncthreads(); + for (uint32_t stride = 16; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) vals[threadIdx.x] = fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]); + __syncthreads(); + } + const float d = vals[0] / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + if (threadIdx.x == 0) xscale[tok * blocks + b] = d; + int8_t *dst = xq + (tok * blocks + b) * 32; + if (threadIdx.x < bn) { + int v = (int)lrintf(xr[threadIdx.x] * id); + v = v > 127 ? 127 : (v < -128 ? -128 : v); + dst[threadIdx.x] = (int8_t)v; + } else { + dst[threadIdx.x] = 0; + } +} + +__global__ static void quantize_q8_0_group_slice_rows_kernel( + int8_t *xq, + float *xscale, + const float *x, + uint64_t group_dim, + uint64_t blocks, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt) { + const uint64_t b = blockIdx.x; + const uint64_t packed_row = blockIdx.y; + if (b >= blocks) return; + const uint64_t token = packed_row / group_cnt; + const uint64_t group = group0 + packed_row - token * group_cnt; + const uint64_t i0 = b * 32u; + const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u; + const float *xr = x + + (token * n_groups_total + group) * group_dim + i0; + + float a = 0.0f; + if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]); + __shared__ float vals[32]; + vals[threadIdx.x] = a; + __syncthreads(); + for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + vals[threadIdx.x] = + fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]); + } + __syncthreads(); + } + const float d = vals[0] / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + if (threadIdx.x == 0u) xscale[packed_row * blocks + b] = d; + int8_t *dst = xq + (packed_row * blocks + b) * 32u; + if (threadIdx.x < bn) { + int v = (int)lrintf(xr[threadIdx.x] * id); + v = v > 127 ? 127 : (v < -128 ? -128 : v); + dst[threadIdx.x] = (int8_t)v; + } else { + dst[threadIdx.x] = 0; + } +} + +__global__ static void matmul_q8_0_preq_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + int use_dp4a) { + uint64_t row = (uint64_t)blockIdx.x; + uint64_t tok = (uint64_t)blockIdx.y; + if (row >= out_dim || tok >= n_tok) return; + const unsigned char *wr = w + row * blocks * 34; + const int8_t *xqr = xq + tok * blocks * 32; + const float *xsr = xscale + tok * blocks; + float acc = 0.0f; + for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { + uint64_t i0 = b * 32; + 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 = xqr + b * 32; + int dot = dot_i8_block(qs, xqb, bn, use_dp4a); + acc += __half2float(*scale_h) * xsr[b] * (float)dot; + } + __shared__ float partial[256]; + partial[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0]; +} + +__global__ static void matmul_q8_0_preq_warp8_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks, + int use_dp4a) { + uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); + const uint64_t tok = (uint64_t)blockIdx.y; + uint32_t lane = threadIdx.x & 31u; + if (row >= out_dim) return; + const unsigned char *wr = w + row * blocks * 34; + const int8_t *xqr = xq + tok * blocks * 32u; + const float *xsr = xscale + tok * blocks; + float acc = 0.0f; + 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 __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) out[tok * out_dim + row] = acc; +} + +__device__ __forceinline__ static uint32_t q8_top1_float_ordered_key(float v) { + const uint32_t u = __float_as_uint(v); + return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); +} + +__device__ __forceinline__ static uint64_t q8_top1_pack_key(float v, uint32_t idx) { + return ((uint64_t)q8_top1_float_ordered_key(v) << 32u) | + (uint64_t)(0xffffffffu - idx); +} + +__device__ __forceinline__ static float q8_top1_unpack_value(uint32_t ordered) { + const uint32_t u = (ordered & 0x80000000u) + ? (ordered ^ 0x80000000u) + : ~ordered; + return __uint_as_float(u); +} + +__global__ static void matmul_q8_0_top1_preq_warp8_kernel( + unsigned long long *best_key, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks, + uint32_t index_offset, + int use_dp4a) { + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t lane = threadIdx.x & 31u; + const uint64_t row = (uint64_t)blockIdx.x * 8u + warp; + const bool valid = row < out_dim; + float acc = 0.0f; + + if (valid) { + const unsigned char *wr = w + row * blocks * 34; + 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 __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); + + __shared__ unsigned long long keys[8]; + if (lane == 0u) { + keys[warp] = valid + ? (unsigned long long)q8_top1_pack_key(acc, index_offset + (uint32_t)row) + : 0ull; + } + __syncthreads(); + + if (threadIdx.x == 0u) { + unsigned long long block_best = keys[0]; + #pragma unroll + for (uint32_t i = 1u; i < 8u; i++) { + if (keys[i] > block_best) block_best = keys[i]; + } + (void)atomicMax(best_key, block_best); + } +} + +__global__ static void matmul_q8_0_top1_unpack_kernel( + uint32_t *selected, + float *values, + const unsigned long long *best_key) { + if (threadIdx.x != 0u || blockIdx.x != 0u) return; + const uint64_t key = (uint64_t)best_key[0]; + const uint32_t ordered = (uint32_t)(key >> 32u); + const uint32_t idx = 0xffffffffu - (uint32_t)key; + selected[0] = idx; + values[0] = q8_top1_unpack_value(ordered); +} + +__global__ static void matmul_q8_0_kslice_preq_warp8_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t slice_dim, + uint64_t out_dim, + uint64_t full_blocks, + uint64_t block_start, + uint64_t slice_blocks, + int use_dp4a) { + const uint64_t row = + (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); + const uint64_t tok = blockIdx.y; + const uint32_t lane = threadIdx.x & 31u; + if (row >= out_dim) return; + out += tok * out_dim; + xq += tok * slice_blocks * 32u; + xscale += tok * slice_blocks; + const unsigned char *wr = + w + row * full_blocks * 34u + block_start * 34u; + float acc = 0.0f; + for (uint64_t b = lane; b < slice_blocks; b += 32u) { + uint64_t i0 = b * 32u; + uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u; + const __half *scale_h = (const __half *)(wr + b * 34u); + const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); + const int8_t *xqb = xq + b * 32u; + 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) out[row] = acc; +} + +__global__ static void matmul_q8_0_pair_preq_warp8_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); + const uint64_t tok = (uint64_t)blockIdx.y; + 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; + const int8_t *xqr = xq + tok * blocks * 32u; + const float *xsr = xscale + tok * blocks; + 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 = xqr + b * 32; + const float xs = xsr[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[tok * out0_dim + row] = acc0; + if (row < out1_dim) out1[tok * out1_dim + row] = acc1; + } +} + +__global__ static void shared_mid_q8_0_preq_warp8_exact_kernel( + float *mid, + const unsigned char *gate_w, + const unsigned char *up_w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks, + float clamp, + const int32_t *selected, + uint32_t expert_split, + bool home_rank, + 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; + if (selected) { + /* Complementary predicates select exactly one writer; ties stay on + * the home rank to avoid an unnecessary peer store. */ + uint32_t home_count = 0u; + uint32_t peer_count = 0u; + #pragma unroll + for (uint32_t i = 0; i < 6u; i++) { + const int32_t expert = selected[i]; + if (expert >= 0 && (uint32_t)expert < expert_split) { + home_count++; + } else if (expert >= 0 && + (uint32_t)expert < 2u * expert_split) { + peer_count++; + } + } + const bool assigned = home_rank + ? home_count <= peer_count : peer_count < home_count; + if (!assigned) return; + } + const unsigned char *gate_row = gate_w + row * blocks * 34u; + const unsigned char *up_row = up_w + row * blocks * 34u; + float gate = 0.0f; + float up = 0.0f; + for (uint64_t b = lane; b < blocks; b += 32u) { + const uint64_t i0 = b * 32u; + const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; + const int8_t *xqb = xq + b * 32u; + const float xs = xscale[b]; + const unsigned char *gb = gate_row + b * 34u; + const unsigned char *ub = up_row + b * 34u; + gate += __half2float(*(const __half *)gb) * xs * + (float)dot_i8_block((const int8_t *)(gb + 2u), xqb, bn, + use_dp4a); + up += __half2float(*(const __half *)ub) * xs * + (float)dot_i8_block((const int8_t *)(ub + 2u), xqb, bn, + use_dp4a); + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + gate = fminf(gate, clamp); + up = fminf(fmaxf(up, -clamp), clamp); + } + const float silu = gate / (1.0f + expf(-gate)); + mid[row] = silu * up * 1.0f; + } +} + +__global__ static void matmul_q8_0_pair_preq_batch_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 n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x; + const uint64_t tok = (uint64_t)blockIdx.y; + if (tok >= n_tok) return; + const int has0 = row < out0_dim; + const int has1 = row < out1_dim; + if (!has0 && !has1) return; + + const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL; + const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL; + const int8_t *xqr = xq + tok * blocks * 32u; + const float *xsr = xscale + tok * blocks; + float acc0 = 0.0f; + float acc1 = 0.0f; + + for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { + const uint64_t i0 = b * 32u; + const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; + const int8_t *xqb = xqr + b * 32u; + const float xs = xsr[b]; + if (has0) { + const __half *scale_h = (const __half *)(wr0 + b * 34u); + const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u); + const int dot = dot_i8_block(qs, xqb, bn, use_dp4a); + acc0 += __half2float(*scale_h) * xs * (float)dot; + } + if (has1) { + const __half *scale_h = (const __half *)(wr1 + b * 34u); + const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u); + const int dot = dot_i8_block(qs, xqb, bn, use_dp4a); + acc1 += __half2float(*scale_h) * xs * (float)dot; + } + } + + __shared__ float partial0[256]; + __shared__ float partial1[256]; + partial0[threadIdx.x] = acc0; + partial1[threadIdx.x] = acc1; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + partial0[threadIdx.x] += partial0[threadIdx.x + stride]; + partial1[threadIdx.x] += partial1[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + if (has0) out0[tok * out0_dim + row] = partial0[0]; + if (has1) out1[tok * out1_dim + row] = partial1[0]; + } +} + +__global__ static void matmul_q8_0_pair_preq_batch_tok2_exact_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 n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x; + const uint64_t tok0 = (uint64_t)blockIdx.y * 2u; + if (tok0 >= n_tok) return; + const int has0 = row < out0_dim; + const int has1 = row < out1_dim; + if (!has0 && !has1) return; + const int valid1 = tok0 + 1u < n_tok; + + const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL; + const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL; + const int8_t *xqr0 = xq + tok0 * blocks * 32u; + const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0; + const float *xsr0 = xscale + tok0 * blocks; + const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; + float acc00 = 0.0f; + float acc01 = 0.0f; + float acc10 = 0.0f; + float acc11 = 0.0f; + + for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { + const uint64_t i0 = b * 32u; + const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; + const int8_t *xqb0 = xqr0 + b * 32u; + const int8_t *xqb1 = xqr1 + b * 32u; + const float xs0 = xsr0[b]; + const float xs1 = valid1 ? xsr1[b] : 0.0f; + if (has0) { + const __half *scale_h = (const __half *)(wr0 + b * 34u); + const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u); + const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + int dot1 = 0; + if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + const float ws = __half2float(*scale_h); + acc00 += ws * xs0 * (float)dot0; + if (valid1) acc01 += ws * xs1 * (float)dot1; + } + if (has1) { + const __half *scale_h = (const __half *)(wr1 + b * 34u); + const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u); + const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + int dot1 = 0; + if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + const float ws = __half2float(*scale_h); + acc10 += ws * xs0 * (float)dot0; + if (valid1) acc11 += ws * xs1 * (float)dot1; + } + } + + __shared__ float partial00[256]; + __shared__ float partial01[256]; + __shared__ float partial10[256]; + __shared__ float partial11[256]; + partial00[threadIdx.x] = acc00; + partial01[threadIdx.x] = acc01; + partial10[threadIdx.x] = acc10; + partial11[threadIdx.x] = acc11; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + partial00[threadIdx.x] += partial00[threadIdx.x + stride]; + partial01[threadIdx.x] += partial01[threadIdx.x + stride]; + partial10[threadIdx.x] += partial10[threadIdx.x + stride]; + partial11[threadIdx.x] += partial11[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + if (has0) { + out0[tok0 * out0_dim + row] = partial00[0]; + if (valid1) out0[(tok0 + 1u) * out0_dim + row] = partial01[0]; + } + if (has1) { + out1[tok0 * out1_dim + row] = partial10[0]; + if (valid1) out1[(tok0 + 1u) * out1_dim + row] = partial11[0]; + } + } +} + +__device__ static float moe_owned_packed_combine_row( + const float *home_slots, + const float *peer_packed, + const int32_t *selected, + uint32_t row, + uint32_t out_dim, + uint32_t expert_split); + +__global__ static void matmul_q8_0_hc_expand_preq_warp8_kernel( + float *out_hc, + float *block_out, + const float *block_add, + const float *block_add2, + const float *owned_home_slots, + const float *owned_peer_packed, + const int32_t *owned_selected, + 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 has_add2, + int has_owned_slots, + uint32_t owned_expert_split, + 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_owned_slots) { + const float routed = moe_owned_packed_combine_row( + owned_home_slots, + owned_peer_packed, + owned_selected, + d, + (uint32_t)out_dim, + owned_expert_split); + block_v = __fadd_rn(block_v, routed); + } else if (has_add) { + float add_v = block_add[d]; + if (has_add2) add_v += block_add2[d]; + block_v += add_v; + } + 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; + } + } +} + +__global__ static void matmul_q8_0_kslice_hc_expand_add_preq_warp8_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 slice_dim, + uint64_t out_dim, + uint64_t full_blocks, + uint64_t block_start, + uint64_t slice_blocks, + uint32_t n_embd, + uint32_t n_hc, + 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 * full_blocks * 34u + block_start * 34u; + float acc = 0.0f; + for (uint64_t b = lane; b < slice_blocks; b += 32u) { + const uint64_t i0 = b * 32u; + const uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u; + const __half *scale_h = (const __half *)(wr + b * 34u); + const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); + const int8_t *xqb = xq + b * 32u; + const 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; + const float block_v = acc + 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; + } + } +} + +__global__ static void matmul_q8_0_preq_batch_warp8_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + 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; + if (row >= out_dim || tok >= n_tok) return; + + const unsigned char *wr = w + row * blocks * 34; + const int8_t *xqr = xq + tok * blocks * 32; + const float *xsr = xscale + tok * blocks; + 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 = 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) out[tok * out_dim + row] = acc; +} + +__global__ static void matmul_q8_0_preq_batch_warp8_tok2_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks, + 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 * 34u; + const int8_t *xqr0 = xq; + const int8_t *xqr1 = xq + blocks * 32u; + const float *xsr0 = xscale; + const float *xsr1 = xscale + blocks; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint64_t b = lane; b < blocks; b += 32u) { + const uint64_t i0 = b * 32u; + const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; + const __half *scale_h = (const __half *)(wr + b * 34u); + const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); + const int8_t *xqb0 = xqr0 + b * 32u; + const int8_t *xqb1 = xqr1 + b * 32u; + int dot0 = 0; + int dot1 = 0; + if (use_dp4a && bn == 32u) { +#pragma unroll + for (uint32_t i = 0; i < 32u; i += 4u) { + const int32_t w4 = load_i8x4_i32_unaligned(qs + i); + dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); + dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); + } + } else { + dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + } + const float ws = __half2float(*scale_h); + acc0 += ws * xsr0[b] * (float)dot0; + acc1 += ws * xsr1[b] * (float)dot1; + } + acc0 = warp_sum_f32(acc0); + acc1 = warp_sum_f32(acc1); + if (lane == 0) { + out[row] = acc0; + out[out_dim + row] = acc1; + } +} + +__global__ static void matmul_q8_0_preq_batch_warp8_tok4_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); + const uint64_t tok0 = (uint64_t)blockIdx.y * 4u; + const uint32_t lane = threadIdx.x & 31u; + if (row >= out_dim || tok0 >= n_tok) return; + + const unsigned char *wr = w + row * blocks * 34; + const int8_t *xqr0 = xq + tok0 * blocks * 32; + const int8_t *xqr1 = xqr0 + blocks * 32; + const int8_t *xqr2 = xqr1 + blocks * 32; + const int8_t *xqr3 = xqr2 + blocks * 32; + const float *xsr0 = xscale + tok0 * blocks; + const float *xsr1 = xsr0 + blocks; + const float *xsr2 = xsr1 + blocks; + const float *xsr3 = xsr2 + blocks; + const int valid1 = tok0 + 1u < n_tok; + const int valid2 = tok0 + 2u < n_tok; + const int valid3 = tok0 + 3u < n_tok; + + float acc0 = 0.0f; + float acc1 = 0.0f; + float acc2 = 0.0f; + float acc3 = 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 *xqb0 = xqr0 + b * 32; + const int8_t *xqb1 = xqr1 + b * 32; + const int8_t *xqb2 = xqr2 + b * 32; + const int8_t *xqb3 = xqr3 + b * 32; + int dot0 = 0; + int dot1 = 0; + int dot2 = 0; + int dot3 = 0; + if (use_dp4a && bn == 32u) { +#pragma unroll + for (uint32_t i = 0; i < 32u; i += 4u) { + const int32_t w4 = load_i8x4_i32_unaligned(qs + i); + dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); + if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); + if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2); + if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3); + } + } else { + dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a); + if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a); + } + const float ws = __half2float(*scale_h); + acc0 += ws * xsr0[b] * (float)dot0; + if (valid1) acc1 += ws * xsr1[b] * (float)dot1; + if (valid2) acc2 += ws * xsr2[b] * (float)dot2; + if (valid3) acc3 += ws * xsr3[b] * (float)dot3; + } + acc0 = warp_sum_f32(acc0); + acc1 = warp_sum_f32(acc1); + acc2 = warp_sum_f32(acc2); + acc3 = warp_sum_f32(acc3); + if (lane == 0) { + out[tok0 * out_dim + row] = acc0; + if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1; + if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2; + if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3; + } +} + +__global__ static void matmul_q8_0_preq_batch_warp8_tok8_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); + const uint64_t tok0 = (uint64_t)blockIdx.y * 8u; + const uint32_t lane = threadIdx.x & 31u; + if (row >= out_dim || tok0 >= n_tok) return; + + const unsigned char *wr = w + row * blocks * 34; + const uint64_t xq_stride = blocks * 32u; + const int8_t *xqr0 = xq + tok0 * xq_stride; + const int valid1 = tok0 + 1u < n_tok; + const int valid2 = tok0 + 2u < n_tok; + const int valid3 = tok0 + 3u < n_tok; + const int valid4 = tok0 + 4u < n_tok; + const int valid5 = tok0 + 5u < n_tok; + const int valid6 = tok0 + 6u < n_tok; + const int valid7 = tok0 + 7u < n_tok; + const int8_t *xqr1 = valid1 ? xqr0 + xq_stride : xqr0; + const int8_t *xqr2 = valid2 ? xqr1 + xq_stride : xqr0; + const int8_t *xqr3 = valid3 ? xqr2 + xq_stride : xqr0; + const int8_t *xqr4 = valid4 ? xqr3 + xq_stride : xqr0; + const int8_t *xqr5 = valid5 ? xqr4 + xq_stride : xqr0; + const int8_t *xqr6 = valid6 ? xqr5 + xq_stride : xqr0; + const int8_t *xqr7 = valid7 ? xqr6 + xq_stride : xqr0; + const float *xsr0 = xscale + tok0 * blocks; + const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; + const float *xsr2 = valid2 ? xsr1 + blocks : xsr0; + const float *xsr3 = valid3 ? xsr2 + blocks : xsr0; + const float *xsr4 = valid4 ? xsr3 + blocks : xsr0; + const float *xsr5 = valid5 ? xsr4 + blocks : xsr0; + const float *xsr6 = valid6 ? xsr5 + blocks : xsr0; + const float *xsr7 = valid7 ? xsr6 + blocks : xsr0; + + float acc0 = 0.0f; + float acc1 = 0.0f; + float acc2 = 0.0f; + float acc3 = 0.0f; + float acc4 = 0.0f; + float acc5 = 0.0f; + float acc6 = 0.0f; + float acc7 = 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 *xqb0 = xqr0 + b * 32; + const int8_t *xqb1 = xqr1 + b * 32; + const int8_t *xqb2 = xqr2 + b * 32; + const int8_t *xqb3 = xqr3 + b * 32; + const int8_t *xqb4 = xqr4 + b * 32; + const int8_t *xqb5 = xqr5 + b * 32; + const int8_t *xqb6 = xqr6 + b * 32; + const int8_t *xqb7 = xqr7 + b * 32; + int dot0 = 0; + int dot1 = 0; + int dot2 = 0; + int dot3 = 0; + int dot4 = 0; + int dot5 = 0; + int dot6 = 0; + int dot7 = 0; + if (use_dp4a && bn == 32u) { +#pragma unroll + for (uint32_t i = 0; i < 32u; i += 4u) { + const int32_t w4 = load_i8x4_i32_unaligned(qs + i); + dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0); + if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1); + if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2); + if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3); + if (valid4) dot4 = __dp4a(w4, load_i8x4_i32_aligned(xqb4 + i), dot4); + if (valid5) dot5 = __dp4a(w4, load_i8x4_i32_aligned(xqb5 + i), dot5); + if (valid6) dot6 = __dp4a(w4, load_i8x4_i32_aligned(xqb6 + i), dot6); + if (valid7) dot7 = __dp4a(w4, load_i8x4_i32_aligned(xqb7 + i), dot7); + } + } else { + dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a); + if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a); + if (valid4) dot4 = dot_i8_block(qs, xqb4, bn, use_dp4a); + if (valid5) dot5 = dot_i8_block(qs, xqb5, bn, use_dp4a); + if (valid6) dot6 = dot_i8_block(qs, xqb6, bn, use_dp4a); + if (valid7) dot7 = dot_i8_block(qs, xqb7, bn, use_dp4a); + } + const float ws = __half2float(*scale_h); + acc0 += ws * xsr0[b] * (float)dot0; + if (valid1) acc1 += ws * xsr1[b] * (float)dot1; + if (valid2) acc2 += ws * xsr2[b] * (float)dot2; + if (valid3) acc3 += ws * xsr3[b] * (float)dot3; + if (valid4) acc4 += ws * xsr4[b] * (float)dot4; + if (valid5) acc5 += ws * xsr5[b] * (float)dot5; + if (valid6) acc6 += ws * xsr6[b] * (float)dot6; + if (valid7) acc7 += ws * xsr7[b] * (float)dot7; + } + acc0 = warp_sum_f32(acc0); + acc1 = warp_sum_f32(acc1); + acc2 = warp_sum_f32(acc2); + acc3 = warp_sum_f32(acc3); + acc4 = warp_sum_f32(acc4); + acc5 = warp_sum_f32(acc5); + acc6 = warp_sum_f32(acc6); + acc7 = warp_sum_f32(acc7); + if (lane == 0) { + out[tok0 * out_dim + row] = acc0; + if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1; + if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2; + if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3; + if (valid4) out[(tok0 + 4u) * out_dim + row] = acc4; + if (valid5) out[(tok0 + 5u) * out_dim + row] = acc5; + if (valid6) out[(tok0 + 6u) * out_dim + row] = acc6; + if (valid7) out[(tok0 + 7u) * out_dim + row] = acc7; + } +} + +__global__ static void matmul_q8_0_preq_batch_tok2_exact_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x; + const uint64_t tok0 = (uint64_t)blockIdx.y * 2u; + if (row >= out_dim || tok0 >= n_tok) return; + const int valid1 = tok0 + 1u < n_tok; + const unsigned char *wr = w + row * blocks * 34u; + const int8_t *xqr0 = xq + tok0 * blocks * 32u; + const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0; + const float *xsr0 = xscale + tok0 * blocks; + const float *xsr1 = valid1 ? xsr0 + blocks : xsr0; + float acc0 = 0.0f; + float acc1 = 0.0f; + + for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) { + const uint64_t i0 = b * 32u; + const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u; + const __half *scale_h = (const __half *)(wr + b * 34u); + const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); + const int8_t *xqb0 = xqr0 + b * 32u; + const int8_t *xqb1 = xqr1 + b * 32u; + const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a); + int dot1 = 0; + if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a); + const float ws = __half2float(*scale_h); + acc0 += ws * xsr0[b] * (float)dot0; + if (valid1) acc1 += ws * xsr1[b] * (float)dot1; + } + + __shared__ float partial0[256]; + __shared__ float partial1[256]; + partial0[threadIdx.x] = acc0; + partial1[threadIdx.x] = acc1; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + partial0[threadIdx.x] += partial0[threadIdx.x + stride]; + partial1[threadIdx.x] += partial1[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + out[tok0 * out_dim + row] = partial0[0]; + if (valid1) out[(tok0 + 1u) * out_dim + row] = partial1[0]; + } +} + + +/* ---- INT8 tensor-core exact Q8_0 batch matmul -------------------------- + * Bit-identical replacement for the exact tok2/warp8-family batched Q8_0 + * kernels. Each output element's reduction is the reference's strided + * halving tree over T slots (T = reduction width: 32 for the warp kernels, + * cuda_q8_exact_threads(blocks) for the exact kernels; slots >= blocks hold + * +0.0f). The kernel decomposes that tree as: 32 streams at stride T/32 + * whose 32 terms per outer step j combine via an adjacent-pairwise static + * register stack taken in bit-reversed stream order (== the top five strided + * tree levels), plus per-(j&3) sequential accumulators and a fixed tail for + * the remaining levels. Fuzz-verified bitwise against both reference + * kernels across shapes, including blocks < T and ragged out_dim/n_tok. + * Rollback: DS4_CUDA_NO_Q8_MMA=1. */ +__device__ __forceinline__ static uint32_t ldu32_unaligned(const uint8_t *p) { + const uintptr_t addr = (uintptr_t)p; + const uint32_t *base = (const uint32_t *)(addr & ~(uintptr_t)3); + const uint32_t lo = base[0]; + const uint32_t hi = base[1]; + return __funnelshift_r(lo, hi, (uint32_t)(addr & 3u) * 8u); +} + +__device__ __forceinline__ static void mma_m16n8k32_s8( + int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3, + uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, + uint32_t b0, uint32_t b1) { +#if __CUDA_ARCH__ >= 800 + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3) + : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1)); +#else + (void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3; +#endif +} + +__device__ __forceinline__ static uint32_t bitrev5(uint32_t i) { + return ((i & 1u) << 4) | ((i & 2u) << 2) | (i & 4u) | ((i & 8u) >> 2) | ((i & 16u) >> 4); +} + +template +__global__ static void matmul_q8_0_mma_exact_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + uint64_t a_stride_blocks, /* activation row stride in blocks (>= blocks) */ + uint64_t out_stride) { /* output token stride in floats (>= out_dim) */ + extern __shared__ unsigned char q8mma_sh[]; + __half *sh_ws = (__half *)q8mma_sh; /* 64 rows x blocks */ + float *sh_xs = (float *)(q8mma_sh + 64u * blocks * 2u); /* 16 toks x blocks */ + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint64_t row_base = (uint64_t)blockIdx.x * 64u; + const uint64_t tok_base = (uint64_t)blockIdx.y * 16u; + + /* stage weight scales (64 rows) and activation scales (16 tokens) */ + for (uint32_t idx = threadIdx.x; idx < 64u * (uint32_t)blocks; idx += blockDim.x) { + const uint32_t rl = idx / (uint32_t)blocks; + const uint32_t b = idx - rl * (uint32_t)blocks; + uint64_t row = row_base + rl; + if (row >= out_dim) row = out_dim - 1u; + sh_ws[idx] = *(const __half *)(w + row * blocks * 34u + (uint64_t)b * 34u); + } + for (uint32_t idx = threadIdx.x; idx < 16u * (uint32_t)blocks; idx += blockDim.x) { + const uint32_t tl = idx / (uint32_t)blocks; + const uint32_t b = idx - tl * (uint32_t)blocks; + const uint64_t tok = tok_base + tl; + sh_xs[idx] = tok < n_tok ? xscale[tok * a_stride_blocks + b] : 0.0f; + } + __syncthreads(); + + const uint64_t row0 = row_base + (uint64_t)warp * 8u; + /* thread's C elements: rows n0,n0+1; tokens mt0, mt0+8 */ + const uint32_t n0 = (lane & 3u) * 2u; + const uint32_t mt0 = lane >> 2u; + const uint64_t tokA = tok_base + mt0; + const uint64_t tokB = tok_base + mt0 + 8u; + /* A source rows for loads (fragment layout): rows lane>>2 and (lane>>2)+8 */ + const uint64_t a_tok_lo = tok_base + (lane >> 2u); + const uint64_t a_tok_hi = a_tok_lo + 8u; + const int8_t *aq_lo = xq + (a_tok_lo < n_tok ? a_tok_lo : 0u) * a_stride_blocks * 32u; + const int8_t *aq_hi = xq + (a_tok_hi < n_tok ? a_tok_hi : 0u) * a_stride_blocks * 32u; + const bool a_lo_ok = a_tok_lo < n_tok; + const bool a_hi_ok = a_tok_hi < n_tok; + /* B source row for loads: row lane>>2 within the warp tile */ + uint64_t b_row = row0 + (lane >> 2u); + if (b_row >= out_dim) b_row = out_dim - 1u; + const unsigned char *b_wr = w + b_row * blocks * 34u; + + /* per-element (4) x per-(j&3) accumulators */ + float acc00 = 0.0f, acc01 = 0.0f, acc02 = 0.0f, acc03 = 0.0f; + float acc10 = 0.0f, acc11 = 0.0f, acc12 = 0.0f, acc13 = 0.0f; + float acc20 = 0.0f, acc21 = 0.0f, acc22 = 0.0f, acc23 = 0.0f; + float acc30 = 0.0f, acc31 = 0.0f, acc32 = 0.0f, acc33 = 0.0f; + + const uint32_t stride = T / 32u; + const uint32_t rl_ws0 = warp * 8u + n0; /* local row for ws of element cols */ + const uint32_t tl_xsA = mt0; /* local token rows for xs */ + const uint32_t tl_xsB = mt0 + 8u; + + for (uint32_t j = 0; j < stride; j++) { + /* adjacent-pairwise static stack over 32 terms in bitrev5 m order */ + float s0e0 = 0, s1e0 = 0, s2e0 = 0, s3e0 = 0, s4e0 = 0; + float s0e1 = 0, s1e1 = 0, s2e1 = 0, s3e1 = 0, s4e1 = 0; + float s0e2 = 0, s1e2 = 0, s2e2 = 0, s3e2 = 0, s4e2 = 0; + float s0e3 = 0, s1e3 = 0, s2e3 = 0, s3e3 = 0, s4e3 = 0; +#pragma unroll + for (uint32_t i = 0; i < 32u; i++) { + const uint32_t m = bitrev5(i); + const uint32_t s = j + m * stride; /* slot index */ + float t0 = 0.0f, t1 = 0.0f, t2 = 0.0f, t3 = 0.0f; + /* slot s sums blocks {s + k*T} sequentially (multi-term when + * blocks > T, exactly like the per-lane strided walk). */ + for (uint32_t b = s; b < blocks; b += T) { + const uint32_t koff = (lane & 3u) * 4u; + const int8_t *ablk_lo = aq_lo + b * 32u; + const int8_t *ablk_hi = aq_hi + b * 32u; + const uint32_t a0 = a_lo_ok ? *(const uint32_t *)(ablk_lo + koff) : 0u; + const uint32_t a1 = a_hi_ok ? *(const uint32_t *)(ablk_hi + koff) : 0u; + const uint32_t a2 = a_lo_ok ? *(const uint32_t *)(ablk_lo + 16u + koff) : 0u; + const uint32_t a3 = a_hi_ok ? *(const uint32_t *)(ablk_hi + 16u + koff) : 0u; + const uint8_t *bq = (const uint8_t *)(b_wr + (uint64_t)b * 34u + 2u); + const uint32_t b0 = ldu32_unaligned(bq + koff); + const uint32_t b1 = ldu32_unaligned(bq + 16u + koff); + int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; + mma_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); + /* term = ws * xs * dot, same expression as reference */ + const float ws0 = __half2float(sh_ws[rl_ws0 * (uint32_t)blocks + b]); + const float ws1 = __half2float(sh_ws[(rl_ws0 + 1u) * (uint32_t)blocks + b]); + const float xsA = sh_xs[tl_xsA * (uint32_t)blocks + b]; + const float xsB = sh_xs[tl_xsB * (uint32_t)blocks + b]; + t0 += ws0 * xsA * (float)c0; + t1 += ws1 * xsA * (float)c1; + t2 += ws0 * xsB * (float)c2; + t3 += ws1 * xsB * (float)c3; + } + /* static adjacent stack push (compile-time resolved) */ + if ((i & 1u) == 0u) { s0e0 = t0; s0e1 = t1; s0e2 = t2; s0e3 = t3; } + else { + t0 = s0e0 + t0; t1 = s0e1 + t1; t2 = s0e2 + t2; t3 = s0e3 + t3; + if ((i & 2u) == 0u) { s1e0 = t0; s1e1 = t1; s1e2 = t2; s1e3 = t3; } + else { + t0 = s1e0 + t0; t1 = s1e1 + t1; t2 = s1e2 + t2; t3 = s1e3 + t3; + if ((i & 4u) == 0u) { s2e0 = t0; s2e1 = t1; s2e2 = t2; s2e3 = t3; } + else { + t0 = s2e0 + t0; t1 = s2e1 + t1; t2 = s2e2 + t2; t3 = s2e3 + t3; + if ((i & 8u) == 0u) { s3e0 = t0; s3e1 = t1; s3e2 = t2; s3e3 = t3; } + else { + t0 = s3e0 + t0; t1 = s3e1 + t1; t2 = s3e2 + t2; t3 = s3e3 + t3; + if ((i & 16u) == 0u) { s4e0 = t0; s4e1 = t1; s4e2 = t2; s4e3 = t3; } + else { + t0 = s4e0 + t0; t1 = s4e1 + t1; t2 = s4e2 + t2; t3 = s4e3 + t3; + /* i == 31: t is the finished x_j */ + switch (j & 3u) { + case 0u: acc00 += t0; acc10 += t1; acc20 += t2; acc30 += t3; break; + case 1u: acc01 += t0; acc11 += t1; acc21 += t2; acc31 += t3; break; + case 2u: acc02 += t0; acc12 += t1; acc22 += t2; acc32 += t3; break; + default: acc03 += t0; acc13 += t1; acc23 += t2; acc33 += t3; break; + } + } + } + } + } + } + } + } + /* tail combine per T */ + float r0, r1, r2, r3; + if (T == 32u) { + r0 = acc00; r1 = acc10; r2 = acc20; r3 = acc30; + } else if (T == 64u) { + r0 = acc00 + acc01; r1 = acc10 + acc11; r2 = acc20 + acc21; r3 = acc30 + acc31; + } else { + r0 = (acc00 + acc02) + (acc01 + acc03); + r1 = (acc10 + acc12) + (acc11 + acc13); + r2 = (acc20 + acc22) + (acc21 + acc23); + r3 = (acc30 + acc32) + (acc31 + acc33); + } + /* writes */ + const uint64_t rowa = row0 + n0; + const uint64_t rowb = rowa + 1u; + if (tokA < n_tok) { + if (rowa < out_dim) out[tokA * out_stride + rowa] = r0; + if (rowb < out_dim) out[tokA * out_stride + rowb] = r1; + } + if (tokB < n_tok) { + if (rowa < out_dim) out[tokB * out_stride + rowa] = r2; + if (rowb < out_dim) out[tokB * out_stride + rowb] = r3; + } +} + +static int cuda_q4_mma_ok(void); +static int cuda_q8_mma_attr_ready[DS4_MAX_GPUS][4]; +static int cuda_q8_mma_try_launch( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t blocks, + uint64_t a_stride_blocks, + uint64_t out_stride, + uint32_t T) { + static int disabled = -1; + if (disabled < 0) disabled = getenv("DS4_CUDA_NO_Q8_MMA") != NULL ? 1 : 0; + if (disabled || !cuda_q4_mma_ok()) return 0; + if ((in_dim & 31u) != 0u || blocks > 256u || n_tok < 8u) return 0; + if (((uintptr_t)w & 1u) || ((uintptr_t)xq & 3u) || ((uintptr_t)xscale & 3u)) return 0; + const size_t shmem = (size_t)(64u * blocks * 2u + 16u * blocks * 4u); + int dev = 0; + cudaGetDevice(&dev); + if (dev < 0 || dev >= DS4_MAX_GPUS) return 0; + const int ti = T == 32u ? 0 : (T == 64u ? 1 : (T == 128u ? 2 : 3)); + dim3 grid(((unsigned)out_dim + 63u) / 64u, ((unsigned)n_tok + 15u) / 16u, 1); +#define DS4_Q8_MMA_LAUNCH(TT) \ + do { \ + if (!cuda_q8_mma_attr_ready[dev][ti]) { \ + cudaFuncAttributes fn_attr; \ + if (cudaFuncGetAttributes(&fn_attr, matmul_q8_0_mma_exact_kernel) != cudaSuccess || \ + fn_attr.binaryVersion < 80) { \ + disabled = 1; \ + return 0; \ + } \ + if (cudaFuncSetAttribute(matmul_q8_0_mma_exact_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + (int)(64u * 256u * 2u + 16u * 256u * 4u)) != cudaSuccess) { \ + disabled = 1; \ + return 0; \ + } \ + cuda_q8_mma_attr_ready[dev][ti] = 1; \ + } \ + matmul_q8_0_mma_exact_kernel<<>>( \ + out, w, xq, xscale, in_dim, out_dim, n_tok, blocks, \ + a_stride_blocks, out_stride); \ + } while (0) + if (T == 32u) DS4_Q8_MMA_LAUNCH(32u); + else if (T == 64u) DS4_Q8_MMA_LAUNCH(64u); + else if (T == 128u) DS4_Q8_MMA_LAUNCH(128u); + else DS4_Q8_MMA_LAUNCH(256u); +#undef DS4_Q8_MMA_LAUNCH + return cuda_ok(cudaGetLastError(), "matmul_q8_0 mma launch") ? 1 : -1; +} + + +__global__ static void dequant_q8_0_to_f16_kernel( + __half *out, + const unsigned char *w, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = in_dim * out_dim; + if (gid >= n) return; + uint64_t row = gid / in_dim; + uint64_t i = gid - row * in_dim; + uint64_t b = i / 32; + uint64_t j = i - b * 32; + const unsigned char *blk = w + (row * blocks + b) * 34; + const __half scale = *(const __half *)blk; + const int8_t q = *(const int8_t *)(blk + 2 + j); + out[gid] = __hmul(scale, __float2half((float)q)); +} + +__global__ static void dequant_q8_0_to_f32_kernel( + float *out, + const unsigned char *w, + uint64_t in_dim, + uint64_t out_dim, + uint64_t blocks) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = in_dim * out_dim; + if (gid >= n) return; + uint64_t row = gid / in_dim; + uint64_t i = gid - row * in_dim; + uint64_t b = i / 32; + uint64_t j = i - b * 32; + const unsigned char *blk = w + (row * blocks + b) * 34; + const float scale = __half2float(*(const __half *)blk); + const int8_t q = *(const int8_t *)(blk + 2 + j); + out[gid] = scale * (float)q; +} + +__global__ static void grouped_q8_0_a_preq_warp8_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__ static void grouped_q8_0_a_preq_warp8_tok2_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 uint32_t tid_in_tok = threadIdx.x & 255u; + const uint64_t row = (uint64_t)blockIdx.x * 8u + (tid_in_tok >> 5u); + const uint64_t tok = (uint64_t)blockIdx.y * 2u + (threadIdx.x >> 8u); + const uint32_t lane = threadIdx.x & 31u; + const uint64_t low_dim = (uint64_t)n_groups * rank; + + float acc = 0.0f; + if (row < low_dim && tok < n_tokens) { + 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 * 34u; + const uint64_t xrow = tok * (uint64_t)n_groups + group; + const int8_t *xqr = xq + xrow * blocks * 32u; + const float *xsr = xscale + xrow * blocks; + + for (uint64_t b = lane; b < blocks; b += 32u) { + const uint64_t i0 = b * 32u; + const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u; + const __half *scale_h = (const __half *)(wr + b * 34u); + const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u); + const int8_t *xqb = xqr + b * 32u; + const 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 && row < low_dim && tok < n_tokens) { + low[tok * low_dim + row] = acc; + } +} + +__global__ static void rms_norm_plain_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { + uint32_t row = blockIdx.x; + if (row >= rows) return; + const float *xr = x + (uint64_t)row * n; + float *orow = out + (uint64_t)row * n; + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < n; 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)n + eps); + for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { + orow[i] = xr[i] * scale; + } +} + +/* Latency-optimized RMS norm for the common n==4096 decode shape: one global + * read pass with register-batched loads, same per-thread accumulation order + * and shared-memory tree as rms_norm_plain_kernel (bit-identical, fuzz + * checked). */ +__global__ static void rms_norm_plain_fast4096_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { + uint32_t row = blockIdx.x; + if (row >= rows) return; + const float *xr = x + (uint64_t)row * n; + float *orow = out + (uint64_t)row * n; + float v[16]; +#pragma unroll + for (uint32_t j = 0; j < 16u; j++) v[j] = xr[threadIdx.x + j * 256u]; + float sum = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < 16u; j++) sum += v[j] * v[j]; + __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)n + eps); +#pragma unroll + for (uint32_t j = 0; j < 16u; j++) orow[threadIdx.x + j * 256u] = v[j] * scale; +} + +/* Batched-load RMS norm for larger rows (n multiple of 2048, e.g. the 16384 + * HC-concatenated decode rows). Two passes like the reference kernel, but + * eight independent loads are issued per accumulation group; the per-thread + * accumulation order (ascending i with stride 256) is unchanged, so results + * are bit-identical. */ +__global__ static void rms_norm_plain_batch8_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) { + uint32_t row = blockIdx.x; + if (row >= rows) return; + const float *xr = x + (uint64_t)row * n; + float *orow = out + (uint64_t)row * n; + float sum = 0.0f; +#pragma unroll 1 + for (uint32_t i = threadIdx.x; i < n; i += 2048u) { + const float v0 = xr[i]; + const float v1 = xr[i + 256u]; + const float v2 = xr[i + 512u]; + const float v3 = xr[i + 768u]; + const float v4 = xr[i + 1024u]; + const float v5 = xr[i + 1280u]; + const float v6 = xr[i + 1536u]; + const float v7 = xr[i + 1792u]; + sum += v0 * v0; + sum += v1 * v1; + sum += v2 * v2; + sum += v3 * v3; + sum += v4 * v4; + sum += v5 * v5; + sum += v6 * v6; + sum += v7 * v7; + } + __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)n + eps); +#pragma unroll 1 + for (uint32_t i = threadIdx.x; i < n; i += 2048u) { + const float v0 = xr[i]; + const float v1 = xr[i + 256u]; + const float v2 = xr[i + 512u]; + const float v3 = xr[i + 768u]; + const float v4 = xr[i + 1024u]; + const float v5 = xr[i + 1280u]; + const float v6 = xr[i + 1536u]; + const float v7 = xr[i + 1792u]; + orow[i] = v0 * scale; + orow[i + 256u] = v1 * scale; + orow[i + 512u] = v2 * scale; + orow[i + 768u] = v3 * scale; + orow[i + 1024u] = v4 * scale; + orow[i + 1280u] = v5 * scale; + orow[i + 1536u] = v6 * scale; + orow[i + 1792u] = v7 * scale; + } +} + +__global__ static void rms_norm_weight_kernel(float *out, const float *x, const float *w, uint32_t n, uint32_t rows, float eps) { + uint32_t row = blockIdx.x; + if (row >= rows) return; + const float *xr = x + (uint64_t)row * n; + float *orow = out + (uint64_t)row * n; + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < n; 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)n + eps); + for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { + orow[i] = xr[i] * scale * w[i]; + } +} + +__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; +} + +__device__ static float rope_yarn_ramp_dev(float low, float high, int i0); + +__global__ static void dsv4_qkv_rms_norm_rows_kv_rope_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, + uint32_t kv_n_head, + uint32_t kv_head_dim, + uint32_t n_rot, + uint32_t pos0, + 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) { + 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); + if (which == 0u) { + for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { + orow[i] = xr[i] * scale * w[i]; + } + return; + } + + const uint32_t n_nope = kv_head_dim - n_rot; + for (uint32_t h = 0; h < kv_n_head; h++) { + const uint32_t head_base = h * kv_head_dim; + for (uint32_t d = threadIdx.x; d < n_nope; d += blockDim.x) { + const uint32_t i = head_base + d; + orow[i] = xr[i] * scale * w[i]; + } + } + + 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); + } + const uint32_t pairs_per_head = n_rot / 2u; + const uint32_t total_pairs = kv_n_head * pairs_per_head; + for (uint32_t p = threadIdx.x; p < total_pairs; p += blockDim.x) { + const uint32_t h = p / pairs_per_head; + const uint32_t pair = p - h * pairs_per_head; + const uint32_t d = n_nope + pair * 2u; + const uint32_t i0 = h * kv_head_dim + d; + const uint32_t i = pair * 2u; + float theta_extrap = (float)(pos0 + row) * 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; + const float x0 = xr[i0] * scale * w[i0]; + const float x1 = xr[i0 + 1u] * scale * w[i0 + 1u]; + orow[i0] = x0 * c - x1 * s; + orow[i0 + 1u] = x0 * s + x1 * c; + } +} + +__global__ static void head_rms_norm_rope_tail_kernel( + float *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + 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; + } + + 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) * 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; + } +} + +__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)); +} + +__global__ static void rope_tail_kernel( + float *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + 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; + + 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 rope_tail_decode_rows_kernel( + float *x, + cuda_attention_decode_row_table rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + 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 gid = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t pairs = n_rows * n_head * (n_rot / 2u); + if (gid >= pairs) return; + const uint32_t pair = gid % (n_rot / 2u); + const uint32_t tmp = gid / (n_rot / 2u); + const uint32_t h = tmp % n_head; + const uint32_t row = tmp / n_head; + const uint32_t n_nope = head_dim - n_rot; + const uint32_t i = pair * 2u; + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + const 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 - 1u), corr1); + } + + const float theta_extrap = (float)rows.row[row].pos * + powf(freq_base, -((float)i) / (float)n_rot); + const float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = attn_factor; + if (ext_factor != 0.0f) { + const 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); + } + const float c = cosf(theta) * mscale; + float s = sinf(theta) * mscale; + if (inverse) s = -s; + + float *tail = x + ((uint64_t)row * n_head + h) * head_dim + n_nope; + const float x0 = tail[i]; + const float x1 = tail[i + 1u]; + tail[i] = x0 * c - x1 * s; + tail[i + 1u] = x0 * s + x1 * c; +} + +__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 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 float dsv4_e2m1fn_dequant_dev(float x) { + float sign = x < 0.0f ? -1.0f : 1.0f; + float ax = fminf(fabsf(x), 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 sign * dsv4_e2m1fn_value_dev(best); +} + +__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]; +} + +__device__ static float rope_yarn_ramp_cpu_equiv_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 DS4_CUDA_UNUSED void rope_tail_one_dev(float *x, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint32_t n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) { + uint32_t n_nope = head_dim - n_rot; + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + float denom = 2.0f * logf(freq_base); + corr0 = fmaxf(0.0f, floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom)); + corr1 = fminf((float)(n_rot - 1), ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom)); + } + for (uint32_t i = 0; i < n_rot; i += 2) { + float theta_extrap = (float)pos * 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 mix = rope_yarn_ramp_cpu_equiv_dev(corr0, corr1, (int)i) * ext_factor; + theta = theta_interp * (1.0f - mix) + theta_extrap * mix; + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + float c = cosf(theta) * mscale; + float s = sinf(theta) * mscale; + float x0 = x[n_nope + i]; + float x1 = x[n_nope + i + 1]; + x[n_nope + i] = x0 * c - x1 * s; + x[n_nope + i + 1] = x0 * s + x1 * c; + } +} + +__device__ static void fp8_kv_quantize_row( + float *xr, + uint32_t head_dim, + uint32_t n_rot, + float *scratch) { + uint32_t tid = threadIdx.x; + uint32_t n_nope = head_dim - n_rot; + 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 q = dsv4_e4m3fn_dequant_dev(fminf(448.0f, fmaxf(-448.0f, v / scale))) * scale; + xr[off + tid] = q; + } + __syncthreads(); + } +} + +__global__ static void fp8_kv_quantize_kernel( + float *x, + uint32_t n_tok, + uint32_t head_dim, + uint32_t n_rot) { + uint32_t row = blockIdx.x; + if (row >= n_tok) return; + __shared__ float scratch[64]; + fp8_kv_quantize_row( + x + (uint64_t)row * head_dim, head_dim, n_rot, scratch); +} + +__global__ static void fp8_kv_quantize_store_rows_kernel( + float *x, + cuda_attention_decode_row_table rows, + uint32_t n_rows, + uint32_t head_dim, + uint32_t n_rot) { + const uint32_t row = blockIdx.x; + if (row >= n_rows) return; + __shared__ float scratch[64]; + float *xr = x + (uint64_t)row * head_dim; + fp8_kv_quantize_row(xr, head_dim, n_rot, scratch); + + const ds4_gpu_attention_decode_row dsc = rows.row[row]; + float *raw = (float *)(uintptr_t)dsc.raw_kv; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + raw[(uint64_t)dsc.raw_start * head_dim + d] = + __half2float(__float2half(xr[d])); + } +} + +__global__ static void indexer_hadamard_fp4_kernel(float *x, uint32_t n_rows, uint32_t head_dim) { + 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))); + xr[tid] = dsv4_e2m1fn_dequant_dev(fminf(6.0f, fmaxf(-6.0f, v / scale))) * scale; +} + +__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) { + 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 = (pos0 + t) % raw_cap; + raw[(uint64_t)row * head_dim + d] = __half2float(__float2half(kv[(uint64_t)t * head_dim + d])); +} + +__global__ static void attention_prefill_raw_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + uint32_t n_tokens, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || h >= n_head) return; + uint32_t raw_count = t + 1 < window ? t + 1 : window; + uint32_t raw_start = t + 1 - raw_count; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + __shared__ float scores[256]; + __shared__ float partial[128]; + __shared__ float max_s; + __shared__ float denom; + float scale = rsqrtf((float)head_dim); + float local_max = sinks[h]; + __syncthreads(); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + const float *kv = raw_kv + (uint64_t)(raw_start + r) * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + if (threadIdx.x == 0) { + float den = expf(sinks[h] - max_s); + for (uint32_t r = 0; r < raw_count; r++) { + scores[r] = expf(scores[r] - max_s); + den += scores[r]; + } + denom = den; + } + __syncthreads(); + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r]; + } + oh[d] = acc / denom; + } +} + +__global__ static void attention_prefill_mixed_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || h >= n_head) return; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + uint32_t raw_start = (window != 0 && t + 1u > window) ? t + 1u - window : 0u; + uint32_t raw_count = t + 1u - raw_start; + uint32_t visible_comp = (t + 1u) / ratio; + if (visible_comp > n_comp) visible_comp = n_comp; + __shared__ float scores[512]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + float scale = rsqrtf((float)head_dim); + float local_max = sinks[h]; + uint32_t n_score = raw_count + visible_comp; + + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + const float *kvrow = raw_kv + (uint64_t)(raw_start + r) * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + for (uint32_t c = threadIdx.x; c < visible_comp; c += blockDim.x) { + float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + float s = -INFINITY; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)c * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale + add; + } + scores[raw_count + c] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r]; + for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; + oh[d] = acc / denom; + } +} + +__global__ static void attention_prefill_raw_softmax_kernel( + float *scores, + const float *sinks, + uint32_t n_tokens, + uint32_t window, + uint32_t n_keys) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens) return; + float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + float local_max = sinks[h]; + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { + bool valid = k <= t && (window == 0 || t - k < window); + float s = valid ? row[k] : -INFINITY; + row[k] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { + float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f; + row[k] = p; + den_local += p; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom; +} + +__global__ static void attention_prefill_mixed_softmax_kernel( + float *scores, + const float *sinks, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_keys) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || ratio == 0) return; + float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + float local_max = sinks[h]; + const uint32_t visible_comp = (t + 1u) / ratio; + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { + float s = -INFINITY; + if (k < n_tokens) { + if (k <= t && (window == 0 || t - k < window)) s = row[k]; + } else { + uint32_t c = k - n_tokens; + if (c < n_comp && c < visible_comp) { + float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + if (add > -1.0e20f) s = row[k] + add; + } + } + row[k] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) { + float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f; + row[k] = p; + den_local += p; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom; +} + +__global__ static void attention_prefill_pack_mixed_kv_kernel( + float *dst, + const float *raw_kv, + const float *comp_kv, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t head_dim) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)(n_tokens + n_comp) * head_dim; + if (gid >= n) return; + uint32_t d = gid % head_dim; + uint32_t r = gid / head_dim; + dst[gid] = r < n_tokens ? raw_kv[(uint64_t)r * head_dim + d] + : comp_kv[(uint64_t)(r - n_tokens) * head_dim + d]; +} + +__global__ static void attention_prefill_unpack_heads_kernel( + float *heads, + const float *tmp, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_head * head_dim; + if (gid >= n) return; + uint32_t d = gid % head_dim; + uint64_t q = gid / head_dim; + uint32_t h = q % n_head; + uint32_t t = q / n_head; + heads[gid] = tmp[((uint64_t)h * n_tokens + t) * head_dim + d]; +} + +__global__ static void attention_pack_group_heads_f16_kernel( + __half *dst, + const float *heads, + uint32_t n_tokens, + uint32_t n_groups, + uint32_t group_dim) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_groups * n_tokens * group_dim; + if (gid >= n) return; + uint32_t d = gid % group_dim; + uint64_t q = gid / group_dim; + uint32_t t = q % n_tokens; + uint32_t g = q / n_tokens; + dst[gid] = __float2half(heads[((uint64_t)t * n_groups + g) * group_dim + d]); +} + +__global__ static void attention_unpack_group_low_kernel( + float *low, + const float *tmp, + uint32_t n_tokens, + uint32_t n_groups, + uint32_t rank) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_groups * n_tokens * rank; + if (gid >= n) return; + uint32_t r = gid % rank; + uint64_t q = gid / rank; + uint32_t t = q % n_tokens; + uint32_t g = q / n_tokens; + uint32_t low_dim = n_groups * rank; + low[(uint64_t)t * low_dim + (uint64_t)g * rank + r] = tmp[gid]; +} + +__global__ static void attention_decode_mixed_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t score_lanes_single) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || h >= n_head) return; + const bool single_all = (n_tokens == 1u && ratio == 0u); + uint32_t qpos = pos0 + t; + uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; + const bool score_thread = threadIdx.x < score_threads; + float scale = rsqrtf((float)head_dim); + if (threadIdx.x == 0) { + raw_count = 0; + raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + if (score_thread) { + for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + } + __syncthreads(); + uint32_t n_score = raw_count + visible_comp; + float local_max = sinks[h]; + if (score_thread) { + if (visible_comp == 0 || (n_tokens == 1u && score_lanes_single == 0u)) { + for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { + const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + for (uint32_t c = threadIdx.x; c < visible_comp; c += score_threads) { + float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + float s = -INFINITY; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)c * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale + add; + } + scores[raw_count + c] = s; + local_max = fmaxf(local_max, s); + } + } else if (n_tokens == 1u && score_lanes_single == 4u) { + uint32_t qlane = threadIdx.x & 3u; + uint32_t qgroup = threadIdx.x >> 2u; + for (uint32_t row0 = 0; row0 < n_score; row0 += 64u) { + uint32_t row = row0 + qgroup; + if (row < n_score) { + float add = 0.0f; + const float *kvrow = NULL; + if (row < raw_count) { + kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim; + } else { + uint32_t c = row - raw_count; + add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim; + } + float s = -INFINITY; + if (kvrow) { + float dot = 0.0f; + for (uint32_t d = qlane; d < head_dim; d += 4u) dot += qh[d] * kvrow[d]; + const uint32_t mask = 0xfu << (threadIdx.x & 28u); + dot += __shfl_down_sync(mask, dot, 2, 4); + dot += __shfl_down_sync(mask, dot, 1, 4); + s = dot * scale + add; + } + if (qlane == 0) scores[row] = s; + } + } + __syncthreads(); + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + local_max = fmaxf(local_max, scores[i]); + } + } else { + uint32_t qlane = threadIdx.x & 7u; + uint32_t qgroup = threadIdx.x >> 3u; + for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { + uint32_t row = row0 + qgroup; + if (row < n_score) { + float add = 0.0f; + const float *kvrow = NULL; + if (row < raw_count) { + kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim; + } else { + uint32_t c = row - raw_count; + add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim; + } + float s = -INFINITY; + if (kvrow) { + float dot = 0.0f; + for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; + const uint32_t mask = 0xffu << (threadIdx.x & 24u); + for (uint32_t off = 4u; off > 0u; off >>= 1u) { + dot += __shfl_down_sync(mask, dot, off, 8); + } + s = dot * scale + add; + } + if (qlane == 0) scores[row] = s; + } + } + __syncthreads(); + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + local_max = fmaxf(local_max, scores[i]); + } + } + } + if (score_thread) partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + } + if (score_thread) partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + if (head_dim == 512u && blockDim.x >= 512u) { + uint32_t d = threadIdx.x; + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc += kv[d] * s; + } + for (uint32_t c = 0; c < visible_comp; c++) { + float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc += kv[d] * s; + } + oh[d] = acc / denom; + } else if (head_dim == 512u && blockDim.x == 256u) { + uint32_t d0 = threadIdx.x; + uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + for (uint32_t c = 0; c < visible_comp; c++) { + float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + oh[d0] = acc0 / denom; + oh[d1] = acc1 / denom; + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; + for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; + oh[d] = acc / denom; + } + } +} + +__global__ static void attention_decode_score_split_scores_kernel( + float *score_out, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + const uint32_t h = blockIdx.y; + const uint32_t j = blockIdx.z; + if (h >= n_head || j >= S) return; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return; + + const uint32_t qbase = n_score / S; + const uint32_t rem = n_score % S; + const uint32_t g0 = j * qbase + (j < rem ? j : rem); + const uint32_t cnt = qbase + (j < rem ? 1u : 0u); + const uint32_t g1 = g0 + cnt; + const float *qh = q + (uint64_t)h * head_dim; + float *row_scores = score_out + (uint64_t)h * n_score; + const float scale = rsqrtf((float)head_dim); + + for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { + float s = -INFINITY; + if (g < raw_count) { + const uint32_t raw_row = + (raw_start + raw_first_idx + g) % raw_cap; + const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale; + } else { + const uint32_t cidx = g - raw_count; + const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale + add; + } + } + row_scores[g] = s; + } +} + +__device__ __forceinline__ float ds4_dot_scalar_ldg( + const float *a, + const float *b, + uint32_t n) { + float dot = 0.0f; + for (uint32_t d = 0; d < n; d++) dot += __ldg(a + d) * __ldg(b + d); + return dot; +} + +__global__ static void attention_decode_score_split_scores_ldg_kernel( + float *score_out, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + const uint32_t h = blockIdx.y; + const uint32_t j = blockIdx.z; + if (h >= n_head || j >= S) return; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return; + + const uint32_t qbase = n_score / S; + const uint32_t rem = n_score % S; + const uint32_t g0 = j * qbase + (j < rem ? j : rem); + const uint32_t cnt = qbase + (j < rem ? 1u : 0u); + const uint32_t g1 = g0 + cnt; + const float *qh = q + (uint64_t)h * head_dim; + float *row_scores = score_out + (uint64_t)h * n_score; + const float scale = rsqrtf((float)head_dim); + + for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { + float s = -INFINITY; + if (g < raw_count) { + const uint32_t raw_row = + (raw_start + raw_first_idx + g) % raw_cap; + const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; + const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim); + s = dot * scale; + } else { + const uint32_t cidx = g - raw_count; + const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; + const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim); + s = dot * scale + add; + } + } + row_scores[g] = s; + } +} + +/* Head-tiled exact score kernel for head_dim==512. + * + * The reference score kernel assigns one (head, row-chunk) per block and lets + * every thread walk one KV row with a scalar sequential dot. Because MQA + * shares the same KV rows across all 64 heads, that reference layout re-reads + * every KV row once per head, and the per-thread row walk is fully + * uncoalesced (threads stride 2KB apart), which multiplies L2 traffic again. + * + * This kernel keeps the per-score arithmetic bit-identical (same ascending-d + * scalar accumulation `dot += q[d] * kv[d]`, same `dot * scale [+ add]` + * epilogue, same masked-row/raw-window classification) but stages a 16-row KV + * tile and a 16-head Q tile in shared memory with coalesced global loads, so + * each KV row is read from L2 once per 16 heads instead of once per head. + * Scores are independent outputs, so retiling the (head, row) space cannot + * change any output bit as long as each individual dot keeps its order. */ +#define DS4_SCORE_TILE_HEADS 16u +#define DS4_SCORE_TILE_ROWS 16u +#define DS4_SCORE_TILE_STRIDE 516u /* 512 + 4 floats: 16B-aligned rows, banks shifted by 4 */ + +__global__ static void attention_decode_score_split_scores_tile512_kernel( + float *score_out, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return; + + extern __shared__ float score_tile_shared[]; + float *sh_q = score_tile_shared; /* 16 x 516 */ + float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE; /* 16 x 516 */ + __shared__ float sh_add[DS4_SCORE_TILE_ROWS]; + + const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS; + const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS; + if (g_base >= n_score || h_base >= n_head) return; + + /* Cooperative Q tile load: 16 heads x 512 floats, float4 coalesced. */ + { + const float4 *q4 = (const float4 *)(q + (uint64_t)h_base * 512u); + const uint32_t tile_heads = + n_head - h_base < DS4_SCORE_TILE_HEADS ? n_head - h_base : DS4_SCORE_TILE_HEADS; + for (uint32_t idx = threadIdx.x; idx < tile_heads * 128u; idx += blockDim.x) { + const uint32_t hh = idx >> 7u; /* head within tile */ + const uint32_t dd = idx & 127u; /* float4 within row */ + const float4 v = q4[hh * 128u + dd]; + float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + /* Row classification + mask staging (thread per row). */ + if (threadIdx.x < DS4_SCORE_TILE_ROWS) { + const uint32_t g = g_base + threadIdx.x; + float add = -INFINITY; + if (g < n_score) { + if (g < raw_count) { + add = 0.0f; /* raw rows are always visible */ + } else { + const uint32_t cidx = g - raw_count; + add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; + } + } + sh_add[threadIdx.x] = add; + } + __syncthreads(); + /* Cooperative KV tile load: two rows at a time, float4 coalesced. + * Masked rows (add <= -1e20) are skipped; their scores never read KV. */ + { + const uint32_t rows_per_pass = blockDim.x >> 7u; /* 128 threads per row */ + const uint32_t rr0 = threadIdx.x >> 7u; + const uint32_t dd = threadIdx.x & 127u; + for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) { + const uint32_t g = g_base + r; + if (g >= n_score) continue; + const bool visible = g < raw_count || sh_add[r] > -1.0e20f; + if (!visible) continue; + const float4 *src; + if (g < raw_count) { + const uint32_t raw_row = (raw_start + raw_first_idx + g) % raw_cap; + src = (const float4 *)(raw_kv + (uint64_t)raw_row * 512u); + } else { + const uint32_t cidx = g - raw_count; + src = (const float4 *)(comp_kv + (uint64_t)cidx * 512u); + } + const float4 v = src[dd]; + float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + __syncthreads(); + + /* One score per thread: r = tid&15 (consecutive threads, coalesced score + * writes), h = tid>>4. The dot keeps the reference kernel's exact scalar + * ascending-d accumulation. */ + const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u); + const uint32_t h = h_base + (threadIdx.x >> 4u); + const uint32_t g = g_base + r; + if (h >= n_head || g >= n_score) return; + const float scale = rsqrtf((float)head_dim); + float *row_scores = score_out + (uint64_t)h * n_score; + const float *qh = sh_q + (uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE; + const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE; + float s = -INFINITY; + const bool need_dot = g < raw_count || sh_add[r] > -1.0e20f; + if (need_dot) { + /* The reference kernel's runtime-trip loop compiles to one sequential + * FFMA chain. Keep exactly that accumulation order here: batched loads + * for latency hiding, but a single explicit ascending fma chain. */ + float dot = 0.0f; +#pragma unroll 1 + for (uint32_t d = 0; d < 512u; d += 8u) { + const float a0 = qh[d + 0u], a1 = qh[d + 1u]; + const float a2 = qh[d + 2u], a3 = qh[d + 3u]; + const float a4 = qh[d + 4u], a5 = qh[d + 5u]; + const float a6 = qh[d + 6u], a7 = qh[d + 7u]; + const float b0 = kvrow[d + 0u], b1 = kvrow[d + 1u]; + const float b2 = kvrow[d + 2u], b3 = kvrow[d + 3u]; + const float b4 = kvrow[d + 4u], b5 = kvrow[d + 5u]; + const float b6 = kvrow[d + 6u], b7 = kvrow[d + 7u]; + dot = __fmaf_rn(a0, b0, dot); + dot = __fmaf_rn(a1, b1, dot); + dot = __fmaf_rn(a2, b2, dot); + dot = __fmaf_rn(a3, b3, dot); + dot = __fmaf_rn(a4, b4, dot); + dot = __fmaf_rn(a5, b5, dot); + dot = __fmaf_rn(a6, b6, dot); + dot = __fmaf_rn(a7, b7, dot); + } + if (g < raw_count) { + s = dot * scale; + } else { + /* The reference expression `dot * scale + add` contracts to one + * FFMA; keep that exact contraction explicit. */ + s = __fmaf_rn(dot, scale, sh_add[r]); + } + } + row_scores[g] = s; +} + +/* Multi-session form of the exact tiled score kernel. Each z-slice selects a + * private KV table entry, while every individual score keeps the same scalar + * ascending-d FMA chain as the one-session kernel. */ +__global__ static void attention_decode_score_split_scores_tile512_rows_kernel( + float *score_out, + const float *q, + cuda_attention_decode_row_table rows, + uint32_t n_rows, + uint32_t score_stride, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t row = blockIdx.z; + if (row >= n_rows) return; + const ds4_gpu_attention_decode_row dsc = rows.row[row]; + if (dsc.indexed) return; + + const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; + const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; + const bool single_all = dsc.ratio == 0u; + const uint32_t qpos = dsc.pos; + const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; + uint32_t visible_comp = single_all + ? dsc.n_comp + : (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u); + if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; + + uint32_t raw_count = 0u; + uint32_t raw_first_idx = 0u; + if (dsc.n_raw != 0u) { + const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; + if (single_all) { + raw_count = dsc.n_raw > 256u ? 256u : dsc.n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (dsc.window != 0u && qpos + 1u > dsc.window) { + const uint32_t wlo = qpos + 1u - dsc.window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + + extern __shared__ float score_tile_shared[]; + float *sh_q = score_tile_shared; + float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE; + + const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS; + const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS; + if (g_base >= n_score || h_base >= n_head) return; + + { + const float4 *q4 = (const float4 *)( + q + ((uint64_t)row * n_head + h_base) * head_dim); + const uint32_t tile_heads = + n_head - h_base < DS4_SCORE_TILE_HEADS + ? n_head - h_base : DS4_SCORE_TILE_HEADS; + for (uint32_t idx = threadIdx.x; + idx < tile_heads * 128u; + idx += blockDim.x) { + const uint32_t hh = idx >> 7u; + const uint32_t dd = idx & 127u; + const float4 v = q4[hh * 128u + dd]; + float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + __syncthreads(); + { + const uint32_t rows_per_pass = blockDim.x >> 7u; + const uint32_t rr0 = threadIdx.x >> 7u; + const uint32_t dd = threadIdx.x & 127u; + for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) { + const uint32_t g = g_base + r; + if (g >= n_score) continue; + const float4 *src; + if (g < raw_count) { + const uint32_t raw_row = + (dsc.raw_start + raw_first_idx + g) % dsc.raw_cap; + src = (const float4 *)(raw_kv + (uint64_t)raw_row * head_dim); + } else { + src = (const float4 *)(comp_kv + + (uint64_t)(g - raw_count) * head_dim); + } + const float4 v = src[dd]; + float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + __syncthreads(); + + const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u); + const uint32_t h = h_base + (threadIdx.x >> 4u); + const uint32_t g = g_base + r; + if (h >= n_head || g >= n_score) return; + const float scale = rsqrtf((float)head_dim); + float *row_scores = score_out + + ((uint64_t)row * n_head + h) * score_stride; + const float *qh = sh_q + + (uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE; + const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE; + float dot = 0.0f; +#pragma unroll 1 + for (uint32_t dd = 0; dd < 512u; dd += 8u) { + const float a0 = qh[dd + 0u], a1 = qh[dd + 1u]; + const float a2 = qh[dd + 2u], a3 = qh[dd + 3u]; + const float a4 = qh[dd + 4u], a5 = qh[dd + 5u]; + const float a6 = qh[dd + 6u], a7 = qh[dd + 7u]; + const float b0 = kvrow[dd + 0u], b1 = kvrow[dd + 1u]; + const float b2 = kvrow[dd + 2u], b3 = kvrow[dd + 3u]; + const float b4 = kvrow[dd + 4u], b5 = kvrow[dd + 5u]; + const float b6 = kvrow[dd + 6u], b7 = kvrow[dd + 7u]; + dot = __fmaf_rn(a0, b0, dot); + dot = __fmaf_rn(a1, b1, dot); + dot = __fmaf_rn(a2, b2, dot); + dot = __fmaf_rn(a3, b3, dot); + dot = __fmaf_rn(a4, b4, dot); + dot = __fmaf_rn(a5, b5, dot); + dot = __fmaf_rn(a6, b6, dot); + dot = __fmaf_rn(a7, b7, dot); + } + row_scores[g] = g < raw_count + ? dot * scale + : __fmaf_rn(dot, scale, 0.0f); +} + +__device__ __forceinline__ float ds4_dot512_float4_ordered( + const float *a, + const float *b) { + const float4 *a4 = (const float4 *)a; + const float4 *b4 = (const float4 *)b; + float dot = 0.0f; +#pragma unroll 1 + for (uint32_t i = 0; i < 128u; i++) { + const float4 av = a4[i]; + const float4 bv = b4[i]; + dot = __fadd_rn(dot, __fmul_rn(av.x, bv.x)); + dot = __fadd_rn(dot, __fmul_rn(av.y, bv.y)); + dot = __fadd_rn(dot, __fmul_rn(av.z, bv.z)); + dot = __fadd_rn(dot, __fmul_rn(av.w, bv.w)); + } + return dot; +} + +__device__ __forceinline__ float ds4_dot512_float4_plain( + const float *a, + const float *b) { + const float4 *a4 = (const float4 *)a; + const float4 *b4 = (const float4 *)b; + float dot = 0.0f; +#pragma unroll 1 + for (uint32_t i = 0; i < 128u; i++) { + const float4 av = a4[i]; + const float4 bv = b4[i]; + dot += av.x * bv.x; + dot += av.y * bv.y; + dot += av.z * bv.z; + dot += av.w * bv.w; + } + return dot; +} + +__global__ static void attention_decode_score_split_scores_vec4_kernel( + float *score_out, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t S) { + const uint32_t h = blockIdx.y; + const uint32_t j = blockIdx.z; + if (h >= n_head || j >= S) return; + const uint32_t head_dim = 512u; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return; + + const uint32_t qbase = n_score / S; + const uint32_t rem = n_score % S; + const uint32_t g0 = j * qbase + (j < rem ? j : rem); + const uint32_t cnt = qbase + (j < rem ? 1u : 0u); + const uint32_t g1 = g0 + cnt; + const float *qh = q + (uint64_t)h * head_dim; + float *row_scores = score_out + (uint64_t)h * n_score; + const float scale = rsqrtf((float)head_dim); + + for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { + float s = -INFINITY; + if (g < raw_count) { + const uint32_t raw_row = + (raw_start + raw_first_idx + g) % raw_cap; + const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; + const float dot = ds4_dot512_float4_ordered(qh, kvrow); + s = dot * scale; + } else { + const uint32_t cidx = g - raw_count; + const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; + const float dot = ds4_dot512_float4_ordered(qh, kvrow); + s = dot * scale + add; + } + } + row_scores[g] = s; + } +} + +__global__ static void attention_decode_score_split_scores_vec4_plain_kernel( + float *score_out, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t S) { + const uint32_t h = blockIdx.y; + const uint32_t j = blockIdx.z; + if (h >= n_head || j >= S) return; + const uint32_t head_dim = 512u; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u) return; + + const uint32_t qbase = n_score / S; + const uint32_t rem = n_score % S; + const uint32_t g0 = j * qbase + (j < rem ? j : rem); + const uint32_t cnt = qbase + (j < rem ? 1u : 0u); + const uint32_t g1 = g0 + cnt; + const float *qh = q + (uint64_t)h * head_dim; + float *row_scores = score_out + (uint64_t)h * n_score; + const float scale = rsqrtf((float)head_dim); + + for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) { + float s = -INFINITY; + if (g < raw_count) { + const uint32_t raw_row = + (raw_start + raw_first_idx + g) % raw_cap; + const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim; + const float dot = ds4_dot512_float4_plain(qh, kvrow); + s = dot * scale; + } else { + const uint32_t cidx = g - raw_count; + const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)cidx * head_dim; + const float dot = ds4_dot512_float4_plain(qh, kvrow); + s = dot * scale + add; + } + } + row_scores[g] = s; + } +} + +__global__ static void attention_decode_score_split_finalize_kernel( + float *heads, + const float *sinks, + const float *score_in, + const float *raw_kv, + const float *comp_kv, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t h = blockIdx.y; + if (h >= n_head) return; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count_s; + __shared__ uint32_t raw_first_idx_s; + + const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; + const bool score_thread = threadIdx.x < score_threads; + if (threadIdx.x == 0) { + raw_count_s = 0; + raw_first_idx_s = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count_s = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx_s = lo - first_raw_pos; + raw_count_s = hi - lo + 1u; + if (raw_count_s > 256u) raw_count_s = 256u; + } + } + } + } + __syncthreads(); + const uint32_t raw_count = raw_count_s; + const uint32_t raw_first_idx = raw_first_idx_s; + if (score_thread) { + for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + } + __syncthreads(); + const uint32_t n_score = raw_count + visible_comp; + const float *row_scores = score_in + (uint64_t)h * n_score; + float local_max = sinks[h]; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + const float s = row_scores[i]; + scores[i] = s; + local_max = fmaxf(local_max, s); + } + } + if (score_thread) partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = + fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + } + if (score_thread) partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + (uint64_t)h * head_dim; + if (head_dim == 512u && blockDim.x >= 512u) { + const uint32_t d = threadIdx.x; + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc += kv[d] * s; + } + for (uint32_t c = 0; c < visible_comp; c++) { + const float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc += kv[d] * s; + } + oh[d] = acc / denom; + } else if (head_dim == 512u && blockDim.x == 256u) { + const uint32_t d0 = threadIdx.x; + const uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + for (uint32_t c = 0; c < visible_comp; c++) { + const float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + oh[d0] = acc0 / denom; + oh[d1] = acc1 / denom; + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; + } + for (uint32_t c = 0; c < visible_comp; c++) { + acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c]; + } + oh[d] = acc / denom; + } + } +} + +__global__ static void attention_decode_score_split_finalize_rows_kernel( + float *heads, + const float *sinks, + const float *score_in, + cuda_attention_decode_row_table rows, + uint32_t n_rows, + uint32_t score_stride, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t row = blockIdx.x; + const uint32_t h = blockIdx.y; + if (row >= n_rows || h >= n_head) return; + const ds4_gpu_attention_decode_row dsc = rows.row[row]; + if (dsc.indexed) return; + const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; + const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; + const bool single_all = dsc.ratio == 0u; + const uint32_t qpos = dsc.pos; + const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; + uint32_t visible_comp = single_all + ? dsc.n_comp + : (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u); + if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; + + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count_s; + __shared__ uint32_t raw_first_idx_s; + + const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; + const bool score_thread = threadIdx.x < score_threads; + if (threadIdx.x == 0u) { + raw_count_s = 0u; + raw_first_idx_s = 0u; + if (dsc.n_raw != 0u) { + const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; + if (single_all) { + raw_count_s = dsc.n_raw > 256u ? 256u : dsc.n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (dsc.window != 0u && qpos + 1u > dsc.window) { + const uint32_t wlo = qpos + 1u - dsc.window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx_s = lo - first_raw_pos; + raw_count_s = hi - lo + 1u; + if (raw_count_s > 256u) raw_count_s = 256u; + } + } + } + } + __syncthreads(); + const uint32_t raw_count = raw_count_s; + const uint32_t raw_first_idx = raw_first_idx_s; + if (score_thread) { + for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { + raw_rows[r] = + (dsc.raw_start + raw_first_idx + r) % dsc.raw_cap; + } + } + __syncthreads(); + const uint32_t n_score = raw_count + visible_comp; + const float *row_scores = score_in + + ((uint64_t)row * n_head + h) * score_stride; + float local_max = sinks[h]; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + const float s = row_scores[i]; + scores[i] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + } + __syncthreads(); + for (uint32_t stride = score_threads >> 1u; + stride > 0u; + stride >>= 1u) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = + fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0u) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + } + __syncthreads(); + for (uint32_t stride = score_threads >> 1u; + stride > 0u; + stride >>= 1u) { + if (threadIdx.x < stride) { + partial[threadIdx.x] += partial[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0u) { + denom = partial[0] + expf(sinks[h] - max_s); + } + __syncthreads(); + + float *oh = heads + ((uint64_t)row * n_head + h) * head_dim; + if (head_dim == 512u && blockDim.x >= 512u) { + const uint32_t dim = threadIdx.x; + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc += kv[dim] * scores[r]; + } + for (uint32_t c = 0; c < visible_comp; c++) { + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc += kv[dim] * scores[raw_count + c]; + } + oh[dim] = acc / denom; + } else { + for (uint32_t dim = threadIdx.x; + dim < head_dim; + dim += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] * + scores[r]; + } + for (uint32_t c = 0; c < visible_comp; c++) { + acc += comp_kv[(uint64_t)c * head_dim + dim] * + scores[raw_count + c]; + } + oh[dim] = acc / denom; + } + } +} + +__global__ static void attention_decode_score_split_finalize_dim2_kernel( + float *heads, + const float *sinks, + const float *score_in, + const float *raw_kv, + const float *comp_kv, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t dim_half = blockIdx.x; + const uint32_t h = blockIdx.y; + if (h >= n_head || head_dim != 512u || dim_half >= 2u) return; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count_s; + __shared__ uint32_t raw_first_idx_s; + + const uint32_t score_threads = 256u; + if (threadIdx.x == 0) { + raw_count_s = 0; + raw_first_idx_s = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count_s = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx_s = lo - first_raw_pos; + raw_count_s = hi - lo + 1u; + if (raw_count_s > 256u) raw_count_s = 256u; + } + } + } + } + __syncthreads(); + const uint32_t raw_count = raw_count_s; + const uint32_t raw_first_idx = raw_first_idx_s; + for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + __syncthreads(); + + const uint32_t n_score = raw_count + visible_comp; + const float *row_scores = score_in + (uint64_t)h * n_score; + float local_max = sinks[h]; + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + const float s = row_scores[i]; + scores[i] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = + fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + + const uint32_t d = dim_half * 256u + threadIdx.x; + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc += kv[d] * s; + } + for (uint32_t c = 0; c < visible_comp; c++) { + const float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)c * head_dim; + acc += kv[d] * s; + } + heads[(uint64_t)h * head_dim + d] = acc / denom; +} + +__global__ static void attention_decode_global_softmax_kernel( + float *score_inout, + float *denom_out, + const float *sinks, + uint32_t n_score, + uint32_t n_head) { + const uint32_t h = blockIdx.x; + if (h >= n_head || n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return; + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom_s; + const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x; + const bool score_thread = threadIdx.x < score_threads; + float *row_scores = score_inout + (uint64_t)h * n_score; + + float local_max = sinks[h]; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + const float s = row_scores[i]; + scores[i] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + } + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = + fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + + float den_local = 0.0f; + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + const float e = expf(scores[i] - max_s); + scores[i] = e; + den_local += e; + } + partial[threadIdx.x] = den_local; + } + __syncthreads(); + for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom_s = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + + if (score_thread) { + for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) { + row_scores[i] = scores[i]; + } + } + if (threadIdx.x == 0) denom_out[h] = denom_s; +} + +__global__ static void attention_decode_split_value_kernel( + float *partials, + const float *score_exp, + const float *raw_kv, + const float *comp_kv, + uint32_t raw_count, + uint32_t raw_first_idx, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_score, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + const uint32_t h = blockIdx.y; + const uint32_t j = blockIdx.z; + if (h >= n_head || j >= S || n_score == 0u) return; + const uint32_t qbase = n_score / S; + const uint32_t rem = n_score % S; + const uint32_t g0 = j * qbase + (j < rem ? j : rem); + const uint32_t cnt = qbase + (j < rem ? 1u : 0u); + const uint32_t g1 = g0 + cnt; + const float *row_scores = score_exp + (uint64_t)h * n_score; + float *pout = partials + ((uint64_t)h * S + j) * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t g = g0; g < g1; g++) { + const float s = row_scores[g]; + if (g < raw_count) { + const uint32_t raw_row = + (raw_start + raw_first_idx + g) % raw_cap; + acc += raw_kv[(uint64_t)raw_row * head_dim + d] * s; + } else { + const uint32_t c = g - raw_count; + acc += comp_kv[(uint64_t)c * head_dim + d] * s; + } + } + pout[d] = acc; + } +} + +__global__ static void attention_decode_split_value_combine_kernel( + float *heads, + const float *partials, + const float *denom, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + const uint32_t h = blockIdx.y; + if (h >= n_head) return; + const float *base = partials + (uint64_t)h * S * head_dim; + const float den = denom[h]; + float *oh = heads + (uint64_t)h * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t j = 0; j < S; j++) { + acc += base[(uint64_t)j * head_dim + d]; + } + oh[d] = acc / den; + } +} + +typedef struct { + uint32_t n_rot; + uint32_t pos0; + uint32_t n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} cuda_attention_inv_rope_params; + +static void attention_decode_score_split_graph_destroy_one(int logical_tier) { + if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return; + cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier]; + if (c->exec) (void)cudaGraphExecDestroy(c->exec); + if (c->graph) (void)cudaGraphDestroy(c->graph); + memset(c, 0, sizeof(*c)); +} + +static int attention_decode_score_split_graph_launch( + int logical_tier, + float *heads, + const float *sinks, + float *scores, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t final_threads, + uint32_t S, + const cuda_attention_inv_rope_params *inv_rope) { + if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0; + cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier]; + const bool graph_inv_rope = + inv_rope && + head_dim == 512u && + inv_rope->n_rot != 0u && + inv_rope->n_rot <= head_dim && + (inv_rope->n_rot & 1u) == 0u; + const bool shape_match = + c->valid && + c->n_head == n_head && + c->head_dim == head_dim && + c->S == S && + c->final_threads == final_threads && + c->fuses_inv_rope == (graph_inv_rope ? 1 : 0) && + (!graph_inv_rope || c->n_rot == inv_rope->n_rot); + if (c->valid && !shape_match) { + attention_decode_score_split_graph_destroy_one(logical_tier); + c = &g_score_split_graph[logical_tier]; + } + + dim3 score_grid(1, n_head, S); + dim3 final_grid(1, n_head, 1); + dim3 score_block(256, 1, 1); + dim3 final_block(final_threads, 1, 1); + + void *score_args[] = { + &scores, &q, &raw_kv, &comp_kv, &comp_mask, &use_comp_mask, + &pos0, &n_raw, &raw_cap, &raw_start, &n_comp, &window, &ratio, + &n_head, &head_dim, &S + }; + cudaKernelNodeParams score_params; + memset(&score_params, 0, sizeof(score_params)); + score_params.func = (void *)attention_decode_score_split_scores_kernel; + score_params.gridDim = score_grid; + score_params.blockDim = score_block; + score_params.sharedMemBytes = 0; + score_params.kernelParams = score_args; + score_params.extra = NULL; + + void *final_args[] = { + &heads, &sinks, &scores, &raw_kv, &comp_kv, &pos0, &n_raw, + &raw_cap, &raw_start, &n_comp, &window, &ratio, &n_head, &head_dim + }; + cudaKernelNodeParams final_params; + memset(&final_params, 0, sizeof(final_params)); + final_params.func = (void *)attention_decode_score_split_finalize_kernel; + final_params.gridDim = final_grid; + final_params.blockDim = final_block; + final_params.sharedMemBytes = 0; + final_params.kernelParams = final_args; + final_params.extra = NULL; + + uint32_t rope_n_tok = 1u; + uint32_t rope_pos_stride = 1u; + int rope_inverse = 1; + uint32_t rope_n_rot = graph_inv_rope ? inv_rope->n_rot : 0u; + uint32_t rope_pos0 = graph_inv_rope ? inv_rope->pos0 : 0u; + uint32_t rope_n_ctx_orig = graph_inv_rope ? inv_rope->n_ctx_orig : 0u; + float rope_freq_base = graph_inv_rope ? inv_rope->freq_base : 0.0f; + float rope_freq_scale = graph_inv_rope ? inv_rope->freq_scale : 0.0f; + float rope_ext_factor = graph_inv_rope ? inv_rope->ext_factor : 0.0f; + float rope_attn_factor = graph_inv_rope ? inv_rope->attn_factor : 0.0f; + float rope_beta_fast = graph_inv_rope ? inv_rope->beta_fast : 0.0f; + float rope_beta_slow = graph_inv_rope ? inv_rope->beta_slow : 0.0f; + void *rope_args[] = { + &heads, &rope_n_tok, &n_head, &head_dim, &rope_n_rot, + &rope_pos0, &rope_pos_stride, &rope_n_ctx_orig, &rope_inverse, + &rope_freq_base, &rope_freq_scale, &rope_ext_factor, + &rope_attn_factor, &rope_beta_fast, &rope_beta_slow + }; + cudaKernelNodeParams rope_params; + memset(&rope_params, 0, sizeof(rope_params)); + if (graph_inv_rope) { + const uint32_t pairs = n_head * (rope_n_rot / 2u); + rope_params.func = (void *)rope_tail_kernel; + rope_params.gridDim = dim3((pairs + 255u) / 256u, 1, 1); + rope_params.blockDim = dim3(256, 1, 1); + rope_params.sharedMemBytes = 0; + rope_params.kernelParams = rope_args; + rope_params.extra = NULL; + } + + if (!c->valid) { + cudaError_t err = cudaGraphCreate(&c->graph, 0); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: attention score-split graph create failed: %s\n", + cudaGetErrorString(err)); + attention_decode_score_split_graph_destroy_one(logical_tier); + return -1; + } + err = cudaGraphAddKernelNode(&c->score_node, c->graph, NULL, 0, + &score_params); + if (err == cudaSuccess) { + err = cudaGraphAddKernelNode(&c->final_node, c->graph, + &c->score_node, 1, &final_params); + } + if (err == cudaSuccess && graph_inv_rope) { + err = cudaGraphAddKernelNode(&c->rope_node, c->graph, + &c->final_node, 1, &rope_params); + } + if (err == cudaSuccess) { + err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0); + } + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: attention score-split graph instantiate failed: %s\n", + cudaGetErrorString(err)); + attention_decode_score_split_graph_destroy_one(logical_tier); + return -1; + } + c->n_head = n_head; + c->head_dim = head_dim; + c->S = S; + c->final_threads = final_threads; + c->n_rot = graph_inv_rope ? inv_rope->n_rot : 0u; + c->fuses_inv_rope = graph_inv_rope ? 1 : 0; + c->valid = 1; + } else { + cudaError_t err = + cudaGraphExecKernelNodeSetParams(c->exec, c->score_node, + &score_params); + if (err == cudaSuccess) { + err = cudaGraphExecKernelNodeSetParams(c->exec, c->final_node, + &final_params); + } + if (err == cudaSuccess && graph_inv_rope) { + err = cudaGraphExecKernelNodeSetParams(c->exec, c->rope_node, + &rope_params); + } + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: attention score-split graph update failed: %s\n", + cudaGetErrorString(err)); + attention_decode_score_split_graph_destroy_one(logical_tier); + return -1; + } + } + + cudaError_t err = cudaGraphLaunch(c->exec, 0); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: attention score-split graph launch failed: %s\n", + cudaGetErrorString(err)); + attention_decode_score_split_graph_destroy_one(logical_tier); + return -1; + } + return 1; +} + +static int attention_decode_score_split_launch( + int logical_tier, + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t final_threads, + const cuda_attention_inv_rope_params *inv_rope) { + if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0)) return 0; + const int explicit_exact = + cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 0); + if (!cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1)) return 0; + if (!explicit_exact && cuda_splitkv_decode_requested()) return 0; + if (g_cuda_decode_score4 || g_cuda_decode_score8) return 0; + if (head_dim == 0u || n_head == 0u) return 0; + const bool single_all = (ratio == 0u); + const uint32_t qpos = pos0; + const uint32_t first_raw_pos = pos0 + 1u - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + uint32_t raw_count = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + const uint32_t n_score = raw_count + visible_comp; + if (n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return 0; + /* With the head-tiled score kernel the exact score-split path beats the + * one-block mixed kernel even for short score counts, so the gate that + * used to protect short contexts (512) now defaults to 1. */ + const uint32_t min_score = cuda_parse_u32_env_clamped( + "DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE", 1u, 0u, + DS4_CUDA_ATTENTION_SCORE_CAP, NULL); + if (n_score < min_score) return 0; + uint32_t chunk = cuda_parse_u32_env_clamped( + "DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK", DS4_CUDA_SPLITKV_CHUNK, + 1u, DS4_CUDA_ATTENTION_SCORE_CAP, NULL); + uint32_t s_floor = cuda_parse_u32_env_clamped( + "DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR", 6u, + 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); + uint32_t s_max = cuda_parse_u32_env_clamped( + "DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX", DS4_CUDA_SPLITKV_S_MAX, + 1u, DS4_CUDA_SPLITKV_S_MAX, NULL); + int exact_present = 0; + uint32_t S = cuda_parse_u32_env_clamped( + "DS4_CUDA_EXACT_SCORE_SPLIT_S", 0u, 1u, + DS4_CUDA_SPLITKV_S_MAX, &exact_present); + if (!exact_present) { + S = (n_score + chunk - 1u) / chunk; + if (S < s_floor) S = s_floor < n_score ? s_floor : n_score; + if (S > s_max) S = s_max; + } + if (S > n_score) S = n_score; + if (S <= 1u) return 0; + const bool graph_inv_rope = + g_cuda_exact_score_split_fuse_inv_rope && + inv_rope && + head_dim == 512u && + final_threads >= 512u && + inv_rope->n_rot != 0u && + inv_rope->n_rot <= 512u && + (inv_rope->n_rot & 1u) == 0u; + + const uint64_t score_count = (uint64_t)n_head * n_score; + float *scores = (float *)cuda_tmp_alloc_on(logical_tier, + score_count * sizeof(float), + "attention exact score split"); + if (!scores) return 0; + const bool use_ldg_scores = g_cuda_exact_score_split_ldg; + const bool use_vec4_plain_scores = + !use_ldg_scores && + g_cuda_exact_score_split_vec4_plain && + head_dim == 512u; + const bool use_vec4_scores = + !use_ldg_scores && + !use_vec4_plain_scores && + (g_cuda_exact_score_split_vec4 || g_decode_score_vec4) && + head_dim == 512u; + const bool use_dim2_finalize = + g_cuda_exact_score_split_dim2 && + head_dim == 512u && + final_threads >= 512u && + !graph_inv_rope; + if ((g_cuda_exact_score_split_graph || graph_inv_rope) && + !use_dim2_finalize && + !use_ldg_scores && + !use_vec4_plain_scores && + !use_vec4_scores) + { + int rc = attention_decode_score_split_graph_launch( + logical_tier, heads, sinks, scores, q, raw_kv, comp_kv, comp_mask, + use_comp_mask, pos0, n_raw, raw_cap, raw_start, n_comp, window, + ratio, n_head, head_dim, final_threads, S, + graph_inv_rope ? inv_rope : NULL); + if (rc == 1) return 1; + if (rc < 0) return -1; + } + if (graph_inv_rope) return 0; + static int score_tile_disabled = -1; + if (score_tile_disabled < 0) { + score_tile_disabled = getenv("DS4_CUDA_NO_SCORE_TILE") != NULL ? 1 : 0; + } + if (!score_tile_disabled && + head_dim == 512u && + !use_ldg_scores && + !use_vec4_plain_scores && + !use_vec4_scores) { + /* cudaFuncSetAttribute() applies to the current device only, so opt in + * to >48KB dynamic shared memory once per device. */ + static int tile_shmem_ready[DS4_MAX_GPUS] = {0}; + const size_t tile_shmem = + (size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) * + DS4_SCORE_TILE_STRIDE * sizeof(float); + int tile_dev = 0; + cudaGetDevice(&tile_dev); + if (tile_dev >= 0 && tile_dev < DS4_MAX_GPUS && + !tile_shmem_ready[tile_dev]) { + if (!cuda_ok(cudaFuncSetAttribute( + attention_decode_score_split_scores_tile512_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + (int)tile_shmem), + "attention score tile shared-memory opt-in")) { + score_tile_disabled = 1; + } + tile_shmem_ready[tile_dev] = 1; + } + if (score_tile_disabled) { + return 0; /* retry via the generic path on the next call */ + } + dim3 tile_grid((n_score + DS4_SCORE_TILE_ROWS - 1u) / DS4_SCORE_TILE_ROWS, + (n_head + DS4_SCORE_TILE_HEADS - 1u) / DS4_SCORE_TILE_HEADS, + 1); + attention_decode_score_split_scores_tile512_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), "attention exact score split tile launch")) return -1; + } else { + dim3 score_grid(1, n_head, S); + if (use_ldg_scores) { + attention_decode_score_split_scores_ldg_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim, S); + } else if (use_vec4_plain_scores) { + attention_decode_score_split_scores_vec4_plain_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, S); + } else if (use_vec4_scores) { + attention_decode_score_split_scores_vec4_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, S); + } else { + attention_decode_score_split_scores_kernel<<>>( + scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim, S); + } + if (!cuda_ok(cudaGetLastError(), "attention exact score split scores launch")) return -1; + } + if (use_dim2_finalize) { + dim3 final_grid(2, n_head, 1); + attention_decode_score_split_finalize_dim2_kernel<<>>( + heads, sinks, scores, raw_kv, comp_kv, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), "attention exact score split dim2 finalize launch")) return -1; + } else { + dim3 final_grid(1, n_head, 1); + attention_decode_score_split_finalize_kernel<<>>( + heads, sinks, scores, raw_kv, comp_kv, + pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, + n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), "attention exact score split finalize launch")) return -1; + } + return 1; +} + +/* ---- perf-02 split-KV / flash-decode (opt-in, default OFF) ---------------- + * + * attention_decode_splitkv_kernel computes a partial online-softmax over a + * contiguous chunk of the flattened logical row set [0, n_score) used by + * attention_decode_mixed_kernel (raw rows first, then compressed rows, same + * ascending ordering). Each block handles (t = blockIdx.x, h = blockIdx.y, + * chunk = blockIdx.z) and writes a partial (m_j, l_j, acc_j[head_dim]) WITHOUT + * the sink term. attention_decode_splitkv_combine_kernel merges the S partials + * per (t,h), folds the sink once, and writes the final normalized head output. + * + * The math is the standard flash-attention online-softmax rescale and is + * algebraically identical to attention_decode_mixed_kernel; it is NOT + * guaranteed bit-identical in FP32 (different expf inputs + add/mul grouping), + * hence default-OFF behind DS4_CUDA_SPLITKV_DECODE and the S==1 dispatch to the + * old kernel as the bit-exact anchor (handled in the launch helper). + * + * Partials scratch layout (per logical tier), contiguous floats: + * stride = head_dim + 2 + * base(t,h,j) = ((t*n_head + h)*S + j) * stride + * [0] = m_j (chunk running max; -INF if empty/all-masked) + * [1] = l_j (chunk denominator sum exp(s - m_j)) + * [2 .. 2+head_dim) = acc_j[head_dim] (chunk weighted value sum) + */ +__global__ static void attention_decode_splitkv_kernel( + float *partials, + const float *q, + const float *raw_kv, + const float *comp_kv, + const float *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + uint32_t j = blockIdx.z; + if (t >= n_tokens || h >= n_head || j >= S) return; + const bool single_all = (n_tokens == 1u && ratio == 0u); + uint32_t qpos = pos0 + t; + uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u); + if (visible_comp > n_comp) visible_comp = n_comp; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + /* scores buffer holds only this chunk's rows. The launch helper guarantees + * cnt <= DS4_CUDA_SPLITKV_SCORE_CAP, including env-tuned split counts. */ + __shared__ float scores[DS4_CUDA_SPLITKV_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float m_s; + __shared__ float l_s; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + float scale = rsqrtf((float)head_dim); + if (threadIdx.x == 0) { + raw_count = 0; + raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (single_all) { + raw_count = n_raw > 256u ? 256u : n_raw; + } else if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + uint32_t n_score = raw_count + visible_comp; + /* even split of [0, n_score) across S chunks: first (n_score % S) chunks + * get base+1, identical deterministic partition for every block. */ + uint32_t qbase = n_score / S; + uint32_t rem = n_score % S; + uint32_t g0 = j * qbase + (j < rem ? j : rem); + uint32_t cnt = qbase + (j < rem ? 1u : 0u); + uint32_t g1 = g0 + cnt; /* exclusive end of this chunk */ + /* Map raw rows that fall in this chunk into shared raw_rows[]. The chunk's + * raw portion is [raw_lo, raw_hi). cnt <= CHUNK and raw rows <= 256, so + * the slice fits raw_rows[256]. */ + uint32_t raw_lo = g0 < raw_count ? g0 : raw_count; + uint32_t raw_hi = g1 < raw_count ? g1 : raw_count; + for (uint32_t r = raw_lo + threadIdx.x; r < raw_hi; r += blockDim.x) { + raw_rows[r - raw_lo] = (raw_start + raw_first_idx + r) % raw_cap; + } + __syncthreads(); + float *pout = partials + (((uint64_t)t * n_head + h) * S + j) * (head_dim + 2u); + /* Pass 1: scores for this chunk's rows into shared scores[0..cnt). */ + float local_max = -INFINITY; + for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) { + uint32_t g = g0 + i; + float s; + if (g < raw_count) { + const float *kvrow = raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale; + } else { + uint32_t c = g - raw_count; + float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f; + s = -INFINITY; + if (add > -1.0e20f) { + const float *kvrow = comp_kv + (uint64_t)c * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale + add; + } + } + scores[i] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) m_s = partial[0]; + __syncthreads(); + float chunk_max = m_s; + /* All-masked / empty-chunk guard: never evaluate exp(-INF - -INF) -> NaN. + * Write zero partial (m=-INF, l=0, acc=0) and return. */ + if (!isfinite(chunk_max)) { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) pout[2u + d] = 0.0f; + if (threadIdx.x == 0) { pout[0] = -INFINITY; pout[1] = 0.0f; } + return; + } + /* Pass 2: exponentiate in place and reduce denominator. */ + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) { + float e = expf(scores[i] - chunk_max); + scores[i] = e; + den_local += e; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) l_s = partial[0]; + __syncthreads(); + /* Pass 3: weighted value accumulation over this chunk's rows (ascending g), + * preserving raw-then-comp ordering to match the reference accumulation. */ + if (head_dim == 512u && blockDim.x == 256u) { + uint32_t d0 = threadIdx.x; + uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t i = 0; i < cnt; i++) { + uint32_t g = g0 + i; + float s = scores[i]; + const float *kv = (g < raw_count) + ? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim + : comp_kv + (uint64_t)(g - raw_count) * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + pout[2u + d0] = acc0; + pout[2u + d1] = acc1; + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t i = 0; i < cnt; i++) { + uint32_t g = g0 + i; + float s = scores[i]; + const float *kv = (g < raw_count) + ? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim + : comp_kv + (uint64_t)(g - raw_count) * head_dim; + acc += kv[d] * s; + } + pout[2u + d] = acc; + } + } + if (threadIdx.x == 0) { + pout[0] = chunk_max; + pout[1] = l_s; + } +} + +__global__ static void attention_decode_splitkv_combine_kernel( + float *heads, + const float *sinks, + const float *partials, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t S) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || h >= n_head) return; + const float *base = partials + (((uint64_t)t * n_head + h) * S) * (head_dim + 2u); + uint32_t stride = head_dim + 2u; + __shared__ float M_s; + __shared__ float L_s; + if (threadIdx.x == 0) { + /* Global max M = max(sink, max_j m_j); sink placed first to match the + * reference (sink seeds local_max). */ + float M = sinks[h]; + for (uint32_t jj = 0; jj < S; jj++) { + float m_j = base[(uint64_t)jj * stride]; + M = fmaxf(M, m_j); /* -INF partials never raise M */ + } + M_s = M; + /* L = Σ_j exp(m_j - M) * l_j + exp(sink - M); sink term added last to + * mirror the reference's denom = Σ scores + expf(sink - max). Chunks + * with l_j == 0 / m_j == -INF contribute exactly 0 (guarded to avoid + * exp(-INF - finite) * 0 edge cases). */ + float L = 0.0f; + for (uint32_t jj = 0; jj < S; jj++) { + float m_j = base[(uint64_t)jj * stride]; + float l_j = base[(uint64_t)jj * stride + 1u]; + if (l_j != 0.0f && isfinite(m_j)) L += expf(m_j - M) * l_j; + } + L += expf(sinks[h] - M); + L_s = L; + } + __syncthreads(); + float M = M_s; + float L = L_s; + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float A = 0.0f; + for (uint32_t jj = 0; jj < S; jj++) { + float m_j = base[(uint64_t)jj * stride]; + float l_j = base[(uint64_t)jj * stride + 1u]; + if (l_j != 0.0f && isfinite(m_j)) { + A += expf(m_j - M) * base[(uint64_t)jj * stride + 2u + d]; + } + } + oh[d] = A / L; + } +} + +__device__ __forceinline__ void attention_compact_topk_stable( + uint32_t *comp_rows, + uint32_t *comp_count, + uint32_t *warp_offsets, + const int32_t *topk, + uint32_t top_k, + uint32_t visible_comp) { + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t n_warp = blockDim.x >> 5u; + if (threadIdx.x == 0u) *comp_count = 0u; + __syncthreads(); + + for (uint32_t base = 0u; base < 512u; base += blockDim.x) { + const uint32_t i = base + threadIdx.x; + const int32_t c = i < top_k ? topk[i] : -1; + const bool valid = c >= 0 && (uint32_t)c < visible_comp; + const uint32_t mask = __ballot_sync(0xffffffffu, valid); + if (lane == 0u) warp_offsets[warp] = __popc(mask); + __syncthreads(); + if (threadIdx.x == 0u) { + uint32_t out = *comp_count; + for (uint32_t w = 0u; w < n_warp; w++) { + const uint32_t count = warp_offsets[w]; + warp_offsets[w] = out; + out += count; + } + *comp_count = out; + } + __syncthreads(); + if (valid) { + const uint32_t lanes_before = lane == 0u + ? 0u : ((1u << lane) - 1u); + const uint32_t slot = warp_offsets[warp] + + __popc(mask & lanes_before); + if (slot < 512u) comp_rows[slot] = (uint32_t)c; + } + __syncthreads(); + } +} + +__global__ static void attention_indexed_mixed_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const int32_t *topk, + uint32_t n_tokens, + uint32_t pos0, + 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 n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tokens || h >= n_head) return; + uint32_t qpos = pos0 + t; + uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t visible_comp = n_comp; + if (ratio != 0) { + visible_comp = (qpos + 1u) / ratio; + if (visible_comp > n_comp) visible_comp = n_comp; + } + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + __shared__ float scores[768]; + __shared__ uint32_t raw_rows[256]; + __shared__ uint32_t comp_rows[512]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + __shared__ uint32_t comp_count; + __shared__ uint32_t comp_warp_offsets[8]; + float scale = rsqrtf((float)head_dim); + if (threadIdx.x == 0) { + raw_count = 0; + raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + attention_compact_topk_stable( + comp_rows, &comp_count, comp_warp_offsets, + topk + (uint64_t)t * top_k, top_k, visible_comp); + uint32_t n_score = raw_count + comp_count; + float local_max = sinks[h]; + if (comp_count == 0) { + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + } else { + uint32_t qlane = threadIdx.x & 7u; + uint32_t qgroup = threadIdx.x >> 3u; + for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { + uint32_t row = row0 + qgroup; + if (row < n_score) { + const float *kvrow = row < raw_count + ? raw_kv + (uint64_t)raw_rows[row] * head_dim + : comp_kv + (uint64_t)comp_rows[row - raw_count] * head_dim; + float dot = 0.0f; + for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d]; + const uint32_t mask = 0xffu << (threadIdx.x & 24u); + for (uint32_t off = 4u; off > 0u; off >>= 1u) { + dot += __shfl_down_sync(mask, dot, off, 8); + } + if (qlane == 0) scores[row] = dot * scale; + } + } + __syncthreads(); + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + local_max = fmaxf(local_max, scores[i]); + } + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + if (head_dim == 512u && blockDim.x == 256u) { + uint32_t d0 = threadIdx.x; + uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + for (uint32_t c = 0; c < comp_count; c++) { + float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + oh[d0] = acc0 / denom; + oh[d1] = acc1 / denom; + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; + for (uint32_t s = 0; s < comp_count; s++) acc += comp_kv[(uint64_t)comp_rows[s] * head_dim + d] * scores[raw_count + s]; + oh[d] = acc / denom; + } + } +} + +__global__ static void attention_indexed_mixed_decode_rows_kernel( + float *heads, + const float *sinks, + const float *q, + cuda_attention_decode_row_table rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t row = blockIdx.x; + const uint32_t h = blockIdx.y; + if (row >= n_rows || h >= n_head) return; + const ds4_gpu_attention_decode_row dsc = rows.row[row]; + if (!dsc.indexed) return; + const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv; + const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv; + const int32_t *topk = (const int32_t *)(uintptr_t)dsc.topk; + const uint32_t qpos = dsc.pos; + const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw; + uint32_t visible_comp = dsc.n_comp; + if (dsc.ratio != 0u) { + visible_comp = (qpos + 1u) / dsc.ratio; + if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp; + } + const float *qh = q + ((uint64_t)row * n_head + h) * head_dim; + __shared__ float scores[768]; + __shared__ uint32_t raw_rows[256]; + __shared__ uint32_t comp_rows[512]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + __shared__ uint32_t comp_count; + __shared__ uint32_t comp_warp_offsets[8]; + const float scale = rsqrtf((float)head_dim); + if (threadIdx.x == 0u) { + raw_count = 0u; + raw_first_idx = 0u; + if (dsc.n_raw != 0u) { + const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u; + if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (dsc.window != 0u && qpos + 1u > dsc.window) { + const uint32_t wlo = qpos + 1u - dsc.window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = + (dsc.raw_start + raw_first_idx + r) % dsc.raw_cap; + } + attention_compact_topk_stable( + comp_rows, &comp_count, comp_warp_offsets, + topk, dsc.top_k, visible_comp); + const uint32_t n_score = raw_count + comp_count; + float local_max = sinks[h]; + if (comp_count == 0u) { + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim; + float dot = 0.0f; + for (uint32_t dim = 0; dim < head_dim; dim++) { + dot += qh[dim] * kvrow[dim]; + } + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + } else { + const uint32_t qlane = threadIdx.x & 7u; + const uint32_t qgroup = threadIdx.x >> 3u; + for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) { + const uint32_t score_row = row0 + qgroup; + if (score_row < n_score) { + const float *kvrow = score_row < raw_count + ? raw_kv + (uint64_t)raw_rows[score_row] * head_dim + : comp_kv + + (uint64_t)comp_rows[score_row - raw_count] * head_dim; + float dot = 0.0f; + for (uint32_t dim = qlane; dim < head_dim; dim += 8u) { + dot += qh[dim] * kvrow[dim]; + } + const uint32_t mask = 0xffu << (threadIdx.x & 24u); + for (uint32_t off = 4u; off > 0u; off >>= 1u) { + dot += __shfl_down_sync(mask, dot, off, 8); + } + if (qlane == 0u) scores[score_row] = dot * scale; + } + } + __syncthreads(); + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + local_max = fmaxf(local_max, scores[i]); + } + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; + stride > 0u; + stride >>= 1u) { + if (threadIdx.x < stride) { + partial[threadIdx.x] = + fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + } + __syncthreads(); + } + if (threadIdx.x == 0u) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; + stride > 0u; + stride >>= 1u) { + if (threadIdx.x < stride) { + partial[threadIdx.x] += partial[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0u) { + denom = partial[0] + expf(sinks[h] - max_s); + } + __syncthreads(); + float *oh = heads + ((uint64_t)row * n_head + h) * head_dim; + if (head_dim == 512u && blockDim.x == 256u) { + const uint32_t d0 = threadIdx.x; + const uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float s = scores[r]; + const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + for (uint32_t c = 0; c < comp_count; c++) { + const float s = scores[raw_count + c]; + const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + oh[d0] = acc0 / denom; + oh[d1] = acc1 / denom; + } else { + for (uint32_t dim = threadIdx.x; + dim < head_dim; + dim += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] * + scores[r]; + } + for (uint32_t c = 0; c < comp_count; c++) { + acc += comp_kv[(uint64_t)comp_rows[c] * head_dim + dim] * + scores[raw_count + c]; + } + oh[dim] = acc / denom; + } + } +} + +__global__ static void attention_indexed_mixed_heads8_rb4_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const int32_t *topk, + uint32_t n_tokens, + uint32_t pos0, + 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 n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t head_group = blockIdx.y; + if (t >= n_tokens || head_dim != 512u) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t head = head_group * 8u + warp; + const bool valid_head = head < n_head; + + __shared__ uint32_t raw_rows[256]; + __shared__ uint32_t comp_rows[512]; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + __shared__ uint32_t comp_count; + __shared__ float4 kv_shared[4 * 128]; + __shared__ float scores[8 * 768]; + + uint32_t qpos = pos0 + t; + uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t visible_comp = n_comp; + if (ratio != 0) { + visible_comp = (qpos + 1u) / ratio; + if (visible_comp > n_comp) visible_comp = n_comp; + } + + if (threadIdx.x == 0) { + raw_count = 0; + raw_first_idx = 0; + comp_count = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + if (threadIdx.x == 0) { + for (uint32_t i = 0; i < top_k && comp_count < 512u; i++) { + int32_t c = topk[(uint64_t)t * top_k + i]; + if (c >= 0 && (uint32_t)c < visible_comp) comp_rows[comp_count++] = (uint32_t)c; + } + } + __syncthreads(); + + const uint32_t n_score = raw_count + comp_count; + const float scale = rsqrtf((float)head_dim); + const float4 *q4 = valid_head + ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) + : NULL; + float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 q1 = q0, q2 = q0, q3 = q0; + if (valid_head) { + q0 = q4[lane + 0u]; + q1 = q4[lane + 32u]; + q2 = q4[lane + 64u]; + q3 = q4[lane + 96u]; + } + + for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { + const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; + for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { + const uint32_t rr = off >> 7u; + const uint32_t c4 = off & 127u; + const uint32_t sr = row0 + rr; + const float4 *src = sr < raw_count + ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) + : (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim); + kv_shared[off] = src[c4]; + } + __syncthreads(); + if (valid_head) { + for (uint32_t rr = 0; rr < nr; rr++) { + const float4 *kv4 = kv_shared + rr * 128u; + float dot = dot4_f32(q0, kv4[lane + 0u]) + + dot4_f32(q1, kv4[lane + 32u]) + + dot4_f32(q2, kv4[lane + 64u]) + + dot4_f32(q3, kv4[lane + 96u]); + dot = warp_sum_f32(dot); + if (lane == 0) scores[warp * 768u + row0 + rr] = dot * scale; + } + } + __syncthreads(); + } + + float max_s = valid_head ? sinks[head] : -INFINITY; + if (valid_head) { + const float *score_row = scores + warp * 768u; + for (uint32_t i = lane; i < n_score; i += 32u) max_s = fmaxf(max_s, score_row[i]); + max_s = warp_max_f32(max_s); + max_s = __shfl_sync(0xffffffffu, max_s, 0); + } + float den = 0.0f; + if (valid_head) { + float *score_row = scores + warp * 768u; + for (uint32_t i = lane; i < n_score; i += 32u) { + float p = expf(score_row[i] - max_s); + score_row[i] = p; + den += p; + } + den = warp_sum_f32(den); + den += expf(sinks[head] - max_s); + den = __shfl_sync(0xffffffffu, den, 0); + } + + float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 o1 = o0, o2 = o0, o3 = o0; + for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { + const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; + for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { + const uint32_t rr = off >> 7u; + const uint32_t c4 = off & 127u; + const uint32_t sr = row0 + rr; + const float4 *src = sr < raw_count + ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) + : (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim); + kv_shared[off] = src[c4]; + } + __syncthreads(); + if (valid_head) { + const float *score_row = scores + warp * 768u; + for (uint32_t rr = 0; rr < nr; rr++) { + const float p = den == 0.0f ? 0.0f : score_row[row0 + rr] / den; + const float4 *kv4 = kv_shared + rr * 128u; + float4 k0 = kv4[lane + 0u]; + float4 k1 = kv4[lane + 32u]; + float4 k2 = kv4[lane + 64u]; + float4 k3 = kv4[lane + 96u]; + o0.x += k0.x * p; o0.y += k0.y * p; o0.z += k0.z * p; o0.w += k0.w * p; + o1.x += k1.x * p; o1.y += k1.y * p; o1.z += k1.z * p; o1.w += k1.w * p; + o2.x += k2.x * p; o2.y += k2.y * p; o2.z += k2.z * p; o2.w += k2.w * p; + o3.x += k3.x * p; o3.y += k3.y * p; o3.z += k3.z * p; o3.w += k3.w * p; + } + } + __syncthreads(); + } + if (valid_head) { + float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); + out4[lane + 0u] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; + } +} + +template +__global__ static void attention_indexed_mixed_heads8_online_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + const int32_t *topk, + uint32_t n_tokens, + uint32_t pos0, + 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 n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t head_group = blockIdx.y; + if (t >= n_tokens || head_dim != 512u) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t head = head_group * HEADS_PER_GROUP + warp; + const bool valid_head = head < n_head; + + __shared__ uint32_t raw_rows[256]; + __shared__ uint32_t raw_count; + __shared__ uint32_t raw_first_idx; + __shared__ float4 kv_shared[ROWS_PER_STAGE * 128]; + + uint32_t qpos = pos0 + t; + uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t visible_comp = n_comp; + if (ratio != 0) { + visible_comp = (qpos + 1u) / ratio; + if (visible_comp > n_comp) visible_comp = n_comp; + } + + if (threadIdx.x == 0) { + raw_count = 0; + raw_first_idx = 0; + if (n_raw != 0) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0 && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + } + __syncthreads(); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + __syncthreads(); + + uint32_t comp_count = top_k < visible_comp ? top_k : visible_comp; + if (comp_count > 512u) comp_count = 512u; + const uint32_t n_score = raw_count + comp_count; + const float scale = rsqrtf((float)head_dim); + const float4 *q4 = valid_head + ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) + : NULL; + float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 q1 = q0, q2 = q0, q3 = q0; + if (valid_head) { + q0 = q4[lane + 0u]; + q1 = q4[lane + 32u]; + q2 = q4[lane + 64u]; + q3 = q4[lane + 96u]; + } + + float max_s = -INFINITY; + float sum_s = 0.0f; + float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 o1 = o0, o2 = o0, o3 = o0; + + for (uint32_t row0 = 0; row0 < n_score; row0 += ROWS_PER_STAGE) { + const uint32_t nr = n_score - row0 < ROWS_PER_STAGE ? n_score - row0 : ROWS_PER_STAGE; + for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { + const uint32_t rr = off >> 7u; + const uint32_t c4 = off & 127u; + const uint32_t sr = row0 + rr; + const uint32_t comp_idx = sr < raw_count + ? 0u + : (uint32_t)topk[(uint64_t)t * top_k + (sr - raw_count)]; + const float4 *src = sr < raw_count + ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) + : (const float4 *)(comp_kv + (uint64_t)comp_idx * head_dim); + kv_shared[off] = src[c4]; + } + __syncthreads(); + if (valid_head) { + for (uint32_t rr = 0; rr < nr; rr++) { + const float4 *kv4 = kv_shared + rr * 128u; + float4 k0 = kv4[lane + 0u]; + float4 k1 = kv4[lane + 32u]; + float4 k2 = kv4[lane + 64u]; + float4 k3 = kv4[lane + 96u]; + float score = dot4_f32(q0, k0) + + dot4_f32(q1, k1) + + dot4_f32(q2, k2) + + dot4_f32(q3, k3); + score = warp_sum_f32(score) * scale; + score = __shfl_sync(0xffffffffu, score, 0); + + const float new_m = fmaxf(max_s, score); + const float old_scale = expf(max_s - new_m); + const float row_scale = expf(score - new_m); + sum_s = sum_s * old_scale + row_scale; + o0.x = o0.x * old_scale + k0.x * row_scale; + o0.y = o0.y * old_scale + k0.y * row_scale; + o0.z = o0.z * old_scale + k0.z * row_scale; + o0.w = o0.w * old_scale + k0.w * row_scale; + o1.x = o1.x * old_scale + k1.x * row_scale; + o1.y = o1.y * old_scale + k1.y * row_scale; + o1.z = o1.z * old_scale + k1.z * row_scale; + o1.w = o1.w * old_scale + k1.w * row_scale; + o2.x = o2.x * old_scale + k2.x * row_scale; + o2.y = o2.y * old_scale + k2.y * row_scale; + o2.z = o2.z * old_scale + k2.z * row_scale; + o2.w = o2.w * old_scale + k2.w * row_scale; + o3.x = o3.x * old_scale + k3.x * row_scale; + o3.y = o3.y * old_scale + k3.y * row_scale; + o3.z = o3.z * old_scale + k3.z * row_scale; + o3.w = o3.w * old_scale + k3.w * row_scale; + max_s = new_m; + } + } + __syncthreads(); + } + + if (valid_head) { + const float sink = sinks[head]; + const float new_m = fmaxf(max_s, sink); + const float old_scale = expf(max_s - new_m); + const float sink_scale = expf(sink - new_m); + sum_s = sum_s * old_scale + sink_scale; + o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; + o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; + o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; + o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; + + const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; + o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; + o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; + o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; + o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; + float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); + out4[lane + 0u] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; + } +} + +__global__ static void attention_static_mixed_heads8_online_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t head_group = blockIdx.y; + if (t >= n_tokens || head_dim != 512u) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t head = head_group * 8u + warp; + const bool valid_head = head < n_head; + + __shared__ float4 kv_shared[4 * 128]; + + const uint32_t raw_count = window != 0u && t + 1u > window ? window : t + 1u; + const uint32_t raw_start = t + 1u - raw_count; + uint32_t comp_count = 0; + if (n_comp != 0u && ratio != 0u) { + comp_count = (t + 1u) / ratio; + if (comp_count > n_comp) comp_count = n_comp; + } + const uint32_t n_score = raw_count + comp_count; + const float scale = rsqrtf((float)head_dim); + const float4 *q4 = valid_head + ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) + : NULL; + float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 q1 = q0, q2 = q0, q3 = q0; + if (valid_head) { + q0 = q4[lane + 0u]; + q1 = q4[lane + 32u]; + q2 = q4[lane + 64u]; + q3 = q4[lane + 96u]; + } + + float max_s = -INFINITY; + float sum_s = 0.0f; + float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 o1 = o0, o2 = o0, o3 = o0; + + for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { + const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; + for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { + const uint32_t rr = off >> 7u; + const uint32_t c4 = off & 127u; + const uint32_t sr = row0 + rr; + const float4 *src = sr < raw_count + ? (const float4 *)(raw_kv + (uint64_t)(raw_start + sr) * head_dim) + : (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim); + kv_shared[off] = src[c4]; + } + __syncthreads(); + if (valid_head) { + for (uint32_t rr = 0; rr < nr; rr++) { + const float4 *kv4 = kv_shared + rr * 128u; + float4 k0 = kv4[lane + 0u]; + float4 k1 = kv4[lane + 32u]; + float4 k2 = kv4[lane + 64u]; + float4 k3 = kv4[lane + 96u]; + float score = dot4_f32(q0, k0) + + dot4_f32(q1, k1) + + dot4_f32(q2, k2) + + dot4_f32(q3, k3); + score = warp_sum_f32(score) * scale; + score = __shfl_sync(0xffffffffu, score, 0); + + const float new_m = fmaxf(max_s, score); + const float old_scale = expf(max_s - new_m); + const float row_scale = expf(score - new_m); + sum_s = sum_s * old_scale + row_scale; + o0.x = o0.x * old_scale + k0.x * row_scale; + o0.y = o0.y * old_scale + k0.y * row_scale; + o0.z = o0.z * old_scale + k0.z * row_scale; + o0.w = o0.w * old_scale + k0.w * row_scale; + o1.x = o1.x * old_scale + k1.x * row_scale; + o1.y = o1.y * old_scale + k1.y * row_scale; + o1.z = o1.z * old_scale + k1.z * row_scale; + o1.w = o1.w * old_scale + k1.w * row_scale; + o2.x = o2.x * old_scale + k2.x * row_scale; + o2.y = o2.y * old_scale + k2.y * row_scale; + o2.z = o2.z * old_scale + k2.z * row_scale; + o2.w = o2.w * old_scale + k2.w * row_scale; + o3.x = o3.x * old_scale + k3.x * row_scale; + o3.y = o3.y * old_scale + k3.y * row_scale; + o3.z = o3.z * old_scale + k3.z * row_scale; + o3.w = o3.w * old_scale + k3.w * row_scale; + max_s = new_m; + } + } + __syncthreads(); + } + + if (valid_head) { + const float sink = sinks[head]; + const float new_m = fmaxf(max_s, sink); + const float old_scale = expf(max_s - new_m); + const float sink_scale = expf(sink - new_m); + sum_s = sum_s * old_scale + sink_scale; + o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; + o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; + o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; + o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; + + const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; + o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; + o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; + o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; + o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; + float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); + out4[lane + 0u] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; + } +} + +__global__ static void attention_decode_mixed_heads8_online_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + const float *comp_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + uint32_t t = blockIdx.x; + uint32_t head_group = blockIdx.y; + if (t >= n_tokens || head_dim != 512u) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t head = head_group * 8u + warp; + const bool valid_head = head < n_head; + + __shared__ uint32_t raw_rows[256]; + __shared__ uint32_t raw_count_s; + __shared__ uint32_t raw_first_idx_s; + __shared__ float4 kv_shared[4 * 128]; + + const uint32_t qpos = pos0 + t; + const uint32_t first_raw_pos = pos0 + n_tokens - n_raw; + uint32_t comp_count = 0; + if (n_comp != 0u) { + if (n_tokens == 1u && ratio == 0u) { + comp_count = n_comp; + } else if (ratio != 0u) { + comp_count = (qpos + 1u) / ratio; + if (comp_count > n_comp) comp_count = n_comp; + } + } + if (threadIdx.x == 0) { + uint32_t raw_count = 0; + uint32_t raw_first_idx = 0; + if (n_raw != 0u) { + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + if (qpos >= first_raw_pos) { + uint32_t lo = first_raw_pos; + if (window != 0u && qpos + 1u > window) { + const uint32_t wlo = qpos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos; + if (hi >= lo) { + raw_first_idx = lo - first_raw_pos; + raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + } + } + } + raw_count_s = raw_count; + raw_first_idx_s = raw_first_idx; + } + __syncthreads(); + const uint32_t raw_count = raw_count_s; + const uint32_t raw_first_idx = raw_first_idx_s; + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap; + } + __syncthreads(); + + const uint32_t n_score = raw_count + comp_count; + const float scale = rsqrtf((float)head_dim); + const float4 *q4 = valid_head + ? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim) + : NULL; + float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 q1 = q0, q2 = q0, q3 = q0; + if (valid_head) { + q0 = q4[lane + 0u]; + q1 = q4[lane + 32u]; + q2 = q4[lane + 64u]; + q3 = q4[lane + 96u]; + } + + float max_s = -INFINITY; + float sum_s = 0.0f; + float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 o1 = o0, o2 = o0, o3 = o0; + + for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) { + const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u; + for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) { + const uint32_t rr = off >> 7u; + const uint32_t c4 = off & 127u; + const uint32_t sr = row0 + rr; + const float4 *src = sr < raw_count + ? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim) + : (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim); + kv_shared[off] = src[c4]; + } + __syncthreads(); + if (valid_head) { + for (uint32_t rr = 0; rr < nr; rr++) { + const float4 *kv4 = kv_shared + rr * 128u; + float4 k0 = kv4[lane + 0u]; + float4 k1 = kv4[lane + 32u]; + float4 k2 = kv4[lane + 64u]; + float4 k3 = kv4[lane + 96u]; + float score = dot4_f32(q0, k0) + + dot4_f32(q1, k1) + + dot4_f32(q2, k2) + + dot4_f32(q3, k3); + score = warp_sum_f32(score) * scale; + score = __shfl_sync(0xffffffffu, score, 0); + + const float new_m = fmaxf(max_s, score); + const float old_scale = expf(max_s - new_m); + const float row_scale = expf(score - new_m); + sum_s = sum_s * old_scale + row_scale; + o0.x = o0.x * old_scale + k0.x * row_scale; + o0.y = o0.y * old_scale + k0.y * row_scale; + o0.z = o0.z * old_scale + k0.z * row_scale; + o0.w = o0.w * old_scale + k0.w * row_scale; + o1.x = o1.x * old_scale + k1.x * row_scale; + o1.y = o1.y * old_scale + k1.y * row_scale; + o1.z = o1.z * old_scale + k1.z * row_scale; + o1.w = o1.w * old_scale + k1.w * row_scale; + o2.x = o2.x * old_scale + k2.x * row_scale; + o2.y = o2.y * old_scale + k2.y * row_scale; + o2.z = o2.z * old_scale + k2.z * row_scale; + o2.w = o2.w * old_scale + k2.w * row_scale; + o3.x = o3.x * old_scale + k3.x * row_scale; + o3.y = o3.y * old_scale + k3.y * row_scale; + o3.z = o3.z * old_scale + k3.z * row_scale; + o3.w = o3.w * old_scale + k3.w * row_scale; + max_s = new_m; + } + } + __syncthreads(); + } + + if (valid_head) { + const float sink = sinks[head]; + const float new_m = fmaxf(max_s, sink); + const float old_scale = expf(max_s - new_m); + const float sink_scale = expf(sink - new_m); + sum_s = sum_s * old_scale + sink_scale; + o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale; + o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale; + o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale; + o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale; + + const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s; + o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s; + o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s; + o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s; + o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s; + float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim); + out4[lane + 0u] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; + } +} diff --git a/models/deepseek/cuda/hc.inc b/models/deepseek/cuda/hc.inc new file mode 100644 index 0000000000..54d2dcfc4c --- /dev/null +++ b/models/deepseek/cuda/hc.inc @@ -0,0 +1,420 @@ +extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) { + if (!out || !mix || !model_map || n_hc != 4) return 0; + const uint64_t mix_bytes = 24ull * sizeof(float); + if (scale_offset > model_size || model_size - scale_offset < 3ull * sizeof(float) || + base_offset > model_size || model_size - base_offset < mix_bytes || + mix->bytes < mix_bytes || out->bytes < mix_bytes) return 0; + const int logical_tier = ds4_tensor_device_idx(out); + const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale"); + const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base"); + if (!scale || !base) return 0; + uint32_t n_rows = (uint32_t)(mix->bytes / mix_bytes); + if (out->bytes / mix_bytes < n_rows) n_rows = (uint32_t)(out->bytes / mix_bytes); + hc_split_sinkhorn_kernel<<<(n_rows + 255) / 256, 256>>>( + (float *)out->ptr, (const float *)mix->ptr, + scale, + base, + n_rows, sinkhorn_iters, eps); + return cuda_ok(cudaGetLastError(), "hc_split_sinkhorn launch"); +} +extern "C" int ds4_gpu_hc_weighted_sum_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *weights, uint32_t n_embd, uint32_t n_hc) { + if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float))); + hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>( + (float *)out->ptr, (const float *)residual_hc->ptr, (const float *)weights->ptr, + n_embd, n_hc, n_tokens, n_hc); + return cuda_ok(cudaGetLastError(), "hc_weighted_sum launch"); +} +extern "C" int ds4_gpu_hc_weighted_sum_split_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { + if (!out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float))); + uint32_t stride = (uint32_t)(2u * n_hc + n_hc * n_hc); + hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>( + (float *)out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr, + n_embd, n_hc, n_tokens, stride); + return cuda_ok(cudaGetLastError(), "hc_weighted_sum_split launch"); +} +extern "C" int ds4_gpu_hc_split_weighted_sum_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps) { + if (!out || !split || !mix || !residual_hc || !model_map || + n_embd == 0 || n_hc != 4) { + return 0; + } + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 || + scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset) { + return 0; + } + uint64_t n_rows = out->bytes / out_row_bytes; + if (mix->bytes < n_rows * mix_bytes || + split->bytes < n_rows * mix_bytes || + residual_hc->bytes < n_rows * residual_row_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out); + const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale"); + const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base"); + if (!scale || !base) return 0; + hc_split_weighted_sum_fused_kernel<<<(uint32_t)n_rows, 256>>>( + (float *)out->ptr, + (float *)split->ptr, + (const float *)mix->ptr, + (const float *)residual_hc->ptr, + scale, + base, + n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps); + return cuda_ok(cudaGetLastError(), "hc split weighted sum launch"); +} +extern "C" int ds4_gpu_hc_split_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float norm_eps) { + if (getenv("DS4_CUDA_DISABLE_HC_SPLIT_NORM_FUSED") == NULL) { + if (!out || !norm_out || !split || !mix || !residual_hc || !model_map || + n_embd == 0 || n_hc != 4) { + return 0; + } + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 || + norm_out->bytes < out->bytes || + scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset || + norm_weight_offset > model_size || + (uint64_t)n_embd * sizeof(float) > model_size - norm_weight_offset) { + return 0; + } + uint64_t n_rows = out->bytes / out_row_bytes; + if (n_rows == 1) { + if (mix->bytes < n_rows * mix_bytes || + split->bytes < n_rows * mix_bytes || + residual_hc->bytes < n_rows * residual_row_bytes) { + return 0; + } + const int logical_tier = ds4_tensor_device_idx(out); + const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, + 3ull * sizeof(float), logical_tier, "hc_scale"); + const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, + mix_bytes, logical_tier, "hc_base"); + const float *norm_w = (const float *)cuda_resolve_weight_ptr(model_map, norm_weight_offset, + (uint64_t)n_embd * sizeof(float), logical_tier, "hc_norm_weight"); + if (!scale || !base || !norm_w) return 0; + hc_split_weighted_sum_norm_fused_kernel<<<(uint32_t)n_rows, 256>>>( + (float *)out->ptr, + (float *)norm_out->ptr, + (float *)split->ptr, + (const float *)mix->ptr, + (const float *)residual_hc->ptr, + scale, + base, + norm_w, + n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps, norm_eps); + return cuda_ok(cudaGetLastError(), "hc split weighted sum norm launch"); + } + } + /* Multi-row fallback: norm EVERY row (rms_norm_weight_tensor is the + * single-row entry and would leave rows 1..n-1 of norm_out untouched). */ + if (!out || n_embd == 0) return 0; + return ds4_gpu_hc_split_weighted_sum_tensor(out, split, mix, residual_hc, + model_map, model_size, + scale_offset, base_offset, + n_embd, n_hc, + sinkhorn_iters, eps) && + ds4_gpu_rms_norm_weight_rows_tensor( + norm_out, out, model_map, model_size, + norm_weight_offset, n_embd, + (uint32_t)(out->bytes / + ((uint64_t)n_embd * sizeof(float))), + norm_eps); +} +extern "C" int ds4_gpu_output_hc_weights_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *pre, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + float eps) { + if (!out || !pre || !model_map || n_hc == 0) return 0; + const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float); + if (row_bytes == 0 || out->bytes < row_bytes || out->bytes % row_bytes != 0 || + pre->bytes < out->bytes || + scale_offset > model_size || sizeof(float) > model_size - scale_offset || + base_offset > model_size || row_bytes > model_size - base_offset) { + return 0; + } + const uint64_t n_tokens = out->bytes / row_bytes; + const int logical_tier = ds4_tensor_device_idx(out); + const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, sizeof(float), logical_tier, "output_hc_scale"); + const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, row_bytes, logical_tier, "output_hc_base"); + if (!scale || !base) return 0; + uint64_t n = n_tokens * n_hc; + output_hc_weights_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, + (const float *)pre->ptr, + scale, + base, + n_hc, + (uint32_t)n_tokens, + eps); + return cuda_ok(cudaGetLastError(), "output hc weights launch"); +} +extern "C" int ds4_gpu_hc_expand_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) { + if (!out_hc || !block_out || !residual_hc || !post || !comb || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, + (const float *)block_out->ptr, + (const float *)block_out->ptr, + (const float *)block_out->ptr, + (const float *)residual_hc->ptr, + (const float *)post->ptr, + (const float *)comb->ptr, + n_embd, n_hc, n_tokens, + n_hc, n_hc * n_hc, 0, 0); + return cuda_ok(cudaGetLastError(), "hc_expand launch"); +} +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) { + if (!out_hc || !block_out || !block_add || !residual_hc || !post || !comb || + n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, + (const float *)block_out->ptr, + (const float *)block_add->ptr, + (const float *)block_out->ptr, + (const float *)residual_hc->ptr, + (const float *)post->ptr, + (const float *)comb->ptr, + n_embd, n_hc, n_tokens, + n_hc, n_hc * n_hc, 1, 0); + return cuda_ok(cudaGetLastError(), "hc_expand_add launch"); +} +extern "C" int ds4_gpu_hc_expand_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { + if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); + uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + const float *base = (const float *)split->ptr; + hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, + (const float *)block_out->ptr, + (const float *)block_out->ptr, + (const float *)block_out->ptr, + (const float *)residual_hc->ptr, + base + n_hc, + base + 2u * n_hc, + n_embd, n_hc, n_tokens, + mix_hc, mix_hc, 0, 0); + return cuda_ok(cudaGetLastError(), "hc_expand_split launch"); +} +extern "C" int ds4_gpu_hc_expand_add_split_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 *split, uint32_t n_embd, uint32_t n_hc) { + if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); + uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + const float *base = (const float *)split->ptr; + hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, + (const float *)block_out->ptr, + (const float *)block_add->ptr, + (const float *)block_out->ptr, + (const float *)residual_hc->ptr, + base + n_hc, + base + 2u * n_hc, + n_embd, n_hc, n_tokens, + mix_hc, mix_hc, 1, 0); + return cuda_ok(cudaGetLastError(), "hc_expand_add_split launch"); +} + +extern "C" int ds4_gpu_hc_expand_add2_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *block_add2, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) { + if (!out_hc || !block_out || !block_add || !block_add2 || + !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float))); + uint32_t mix_hc = 2u * n_hc + n_hc * n_hc; + uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd; + const float *base = (const float *)split->ptr; + hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr, + (const float *)block_out->ptr, + (const float *)block_add->ptr, + (const float *)block_add2->ptr, + (const float *)residual_hc->ptr, + base + n_hc, + base + 2u * n_hc, + n_embd, n_hc, n_tokens, + mix_hc, mix_hc, 1, 1); + return cuda_ok(cudaGetLastError(), "hc_expand_add2_split launch"); +} + +extern "C" int ds4_gpu_shared_down_hc_expand_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_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 *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { + return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out, + model_map, model_size, + weight_offset, + in_dim, out_dim, + shared_mid, + routed_out, + NULL, + NULL, NULL, NULL, 0, + residual_hc, + split, + n_embd, n_hc, + "shared_down_hc_expand"); + } + return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size, + weight_offset, in_dim, out_dim, + shared_mid, 1) && + ds4_gpu_hc_expand_add_split_tensor(out_hc, shared_out, routed_out, + residual_hc, split, n_embd, n_hc); +} + +extern "C" int ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_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 *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *routed_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { + return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out, + model_map, model_size, + weight_offset, + in_dim, out_dim, + shared_mid, + routed_out, + routed_add, + NULL, NULL, NULL, 0, + residual_hc, + split, + n_embd, n_hc, + "shared_down_hc_expand_add"); + } + return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size, + weight_offset, in_dim, out_dim, + shared_mid, 1) && + ds4_gpu_hc_expand_add2_split_tensor(out_hc, shared_out, routed_out, + routed_add, residual_hc, split, + n_embd, n_hc); +} + +extern "C" int ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_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 *shared_mid, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_packed, + const ds4_gpu_tensor *selected, + uint32_t expert_split, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") != NULL) return 0; + return cuda_matmul_q8_0_hc_expand_tensor_labeled( + out_hc, + shared_out, + model_map, + model_size, + weight_offset, + in_dim, + out_dim, + shared_mid, + NULL, + NULL, + home_slots, + peer_packed, + selected, + expert_split, + residual_hc, + split, + n_embd, + n_hc, + "shared_down_hc_expand_owned"); +} + +extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_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, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) { + return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, block_out, + model_map, model_size, + weight_offset, + in_dim, out_dim, + x, + NULL, + NULL, + NULL, NULL, NULL, 0, + residual_hc, + split, + n_embd, n_hc, + "q8_hc_expand"); + } + return ds4_gpu_matmul_q8_0_tensor(block_out, model_map, model_size, + weight_offset, in_dim, out_dim, x, 1) && + ds4_gpu_hc_expand_split_tensor(out_hc, block_out, residual_hc, + split, n_embd, n_hc); +} diff --git a/models/deepseek/cuda/moe.inc b/models/deepseek/cuda/moe.inc new file mode 100644 index 0000000000..02a2cebf16 --- /dev/null +++ b/models/deepseek/cuda/moe.inc @@ -0,0 +1,6223 @@ + +__device__ static float dev_f16_to_f32(uint16_t v) { + return __half2float(*reinterpret_cast(&v)); +} + +__device__ __forceinline__ static uint32_t dev_unpack_iq2_signs(uint32_t v) { + const uint32_t p = __popc(v) & 1u; + const uint32_t s = v ^ (p << 7u); + return s * 0x01010101u; +} + +__device__ __forceinline__ static int32_t dev_iq2_dp4a_8(uint64_t grid, uint32_t sign, const int8_t *q8, int32_t acc) { + const uint32_t signs = dev_unpack_iq2_signs(sign); + const int32_t sm0 = __vcmpne4(signs & 0x08040201u, 0); + const int32_t sm1 = __vcmpne4(signs & 0x80402010u, 0); + const int32_t g0 = __vsub4((int32_t)(uint32_t)grid ^ sm0, sm0); + const int32_t g1 = __vsub4((int32_t)(uint32_t)(grid >> 32) ^ sm1, sm1); + acc = __dp4a(g0, *(const int32_t *)(q8 + 0), acc); + acc = __dp4a(g1, *(const int32_t *)(q8 + 4), acc); + return acc; +} + +__device__ static int32_t dev_dot_q2_16(const uint8_t *q2, const int8_t *q8, int shift) { + int32_t sum = 0; + #pragma unroll + for (uint32_t i = 0; i < 16; i += 4) { + const int32_t v = (*(const int32_t *)(q2 + i) >> shift) & 0x03030303; + sum = __dp4a(v, *(const int32_t *)(q8 + i), sum); + } + return sum; +} + +__device__ static int32_t dev_dot_iq2_pair_16(uint8_t grid0, uint32_t sign0, uint8_t grid1, uint32_t sign1, const int8_t *q8) { + int32_t sum = 0; + sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid0], cuda_ksigns_iq2xs[sign0], q8, sum); + sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid1], cuda_ksigns_iq2xs[sign1], q8 + 8, sum); + return sum; +} + +__device__ __forceinline__ static void dev_iq2_i8x8_lut( + const uint64_t *grid, + const uint8_t *signs, + uint8_t grid_idx, + uint32_t sign_idx, + int32_t *w0, + int32_t *w1) { + const uint32_t s = dev_unpack_iq2_signs(signs[sign_idx]); + const int32_t sm0 = __vcmpne4(s & 0x08040201u, 0); + const int32_t sm1 = __vcmpne4(s & 0x80402010u, 0); + const uint64_t g = grid[grid_idx]; + *w0 = __vsub4((int32_t)(uint32_t)g ^ sm0, sm0); + *w1 = __vsub4((int32_t)(uint32_t)(g >> 32) ^ sm1, sm1); +} + +__device__ static float dev_dot_iq2_xxs_q8_K_block_lut( + const cuda_block_iq2_xxs *x, + const cuda_block_q8_K *y, + const uint64_t *grid, + const uint8_t *signs) { + const float xd = dev_f16_to_f32(x->d); + const uint16_t *q2 = x->qs; + const int8_t *q8 = y->qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u); + int32_t w[8]; + dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]); + int32_t sumi = 0; + sumi = __dp4a(w[0], *(const int32_t *)(q8 + ib32 * 32u + 0), sumi); + sumi = __dp4a(w[1], *(const int32_t *)(q8 + ib32 * 32u + 4), sumi); + sumi = __dp4a(w[2], *(const int32_t *)(q8 + ib32 * 32u + 8), sumi); + sumi = __dp4a(w[3], *(const int32_t *)(q8 + ib32 * 32u + 12), sumi); + sumi = __dp4a(w[4], *(const int32_t *)(q8 + ib32 * 32u + 16), sumi); + sumi = __dp4a(w[5], *(const int32_t *)(q8 + ib32 * 32u + 20), sumi); + sumi = __dp4a(w[6], *(const int32_t *)(q8 + ib32 * 32u + 24), sumi); + sumi = __dp4a(w[7], *(const int32_t *)(q8 + ib32 * 32u + 28), sumi); + bsum += sumi * ls; + } + return 0.125f * xd * y->d * (float)bsum; +} + +__device__ static float dev_dot_iq2_xxs_q8_K_block(const cuda_block_iq2_xxs *x, const cuda_block_q8_K *y) { + const float d = dev_f16_to_f32(x->d) * y->d; + const uint16_t *q2 = x->qs; + const int8_t *q8 = y->qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const uint32_t ls = 2u * (aux1 >> 28) + 1u; + const uint8_t a0 = (uint8_t)(aux0 & 0xffu); + const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); + const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); + const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); + int32_t sumi = 0; + sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8); + q8 += 16; + sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8); + q8 += 16; + bsum += sumi * (int32_t)ls; + } + return 0.125f * d * (float)bsum; +} + +__device__ static void dev_dot_iq2_xxs_q8_K_block8_deq_lut( + const cuda_block_iq2_xxs *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + const cuda_block_q8_K *y4, + const cuda_block_q8_K *y5, + const cuda_block_q8_K *y6, + const cuda_block_q8_K *y7, + uint32_t n, + float acc[8], + const uint64_t *grid, + const uint8_t *signs) { + const float xd = dev_f16_to_f32(x->d); + const uint16_t *q2 = x->qs; + int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const int8_t *q8[8] = { + y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL, + y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL, + }; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u); + int32_t w[8]; + dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]); + dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]); + for (uint32_t p = 0; p < n; p++) { + const int8_t *q = q8[p] + ib32 * 32; + int32_t sumi = 0; + sumi = __dp4a(w[0], *(const int32_t *)(q + 0), sumi); + sumi = __dp4a(w[1], *(const int32_t *)(q + 4), sumi); + sumi = __dp4a(w[2], *(const int32_t *)(q + 8), sumi); + sumi = __dp4a(w[3], *(const int32_t *)(q + 12), sumi); + sumi = __dp4a(w[4], *(const int32_t *)(q + 16), sumi); + sumi = __dp4a(w[5], *(const int32_t *)(q + 20), sumi); + sumi = __dp4a(w[6], *(const int32_t *)(q + 24), sumi); + sumi = __dp4a(w[7], *(const int32_t *)(q + 28), sumi); + bsum[p] += sumi * ls; + } + } + const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; + for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; +} + +__device__ static void dev_dot_iq2_xxs_q8_K_block4( + const cuda_block_iq2_xxs *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + uint32_t n, + float acc[4]) { + const float xd = dev_f16_to_f32(x->d); + const uint16_t *q2 = x->qs; + int32_t bsum[4] = {0, 0, 0, 0}; + const int8_t *q8[4] = { + y0 ? y0->qs : NULL, + y1 ? y1->qs : NULL, + y2 ? y2->qs : NULL, + y3 ? y3->qs : NULL, + }; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const uint32_t ls = 2u * (aux1 >> 28) + 1u; + const uint8_t a0 = (uint8_t)(aux0 & 0xffu); + const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); + const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); + const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); + for (uint32_t p = 0; p < n; p++) { + int32_t sumi = 0; + sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32); + sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16); + bsum[p] += sumi * (int32_t)ls; + } + } + const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 }; + for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; +} + +__device__ static DS4_CUDA_UNUSED void dev_dot_iq2_xxs_q8_K_block8( + const cuda_block_iq2_xxs *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + const cuda_block_q8_K *y4, + const cuda_block_q8_K *y5, + const cuda_block_q8_K *y6, + const cuda_block_q8_K *y7, + uint32_t n, + float acc[8]) { + const float xd = dev_f16_to_f32(x->d); + const uint16_t *q2 = x->qs; + int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const int8_t *q8[8] = { + y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL, + y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL, + }; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const uint32_t ls = 2u * (aux1 >> 28) + 1u; + const uint8_t a0 = (uint8_t)(aux0 & 0xffu); + const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu); + const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu); + const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu); + for (uint32_t p = 0; p < n; p++) { + int32_t sumi = 0; + sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32); + sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16); + bsum[p] += sumi * (int32_t)ls; + } + } + const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; + for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p]; +} + +__device__ static void dev_q4_K_get_scale_min( + uint32_t j, + const uint8_t *scales, + uint8_t *d_out, + uint8_t *m_out) { + if (j < 4u) { + *d_out = scales[j] & 63u; + *m_out = scales[j + 4u] & 63u; + } else { + *d_out = (scales[j + 4u] & 0x0fu) | ((scales[j - 4u] >> 6u) << 4u); + *m_out = (scales[j + 4u] >> 4u) | ((scales[j] >> 6u) << 4u); + } +} + +__device__ __forceinline__ static int32_t dev_dot_q4_32(const uint8_t *qs, const int8_t *q8, int shift) { + int32_t sum = 0; + #pragma unroll + for (uint32_t i = 0; i < 32u; i += 4u) { + const int32_t v = (*(const int32_t *)(qs + i) >> shift) & 0x0f0f0f0f; + sum = __dp4a(v, *(const int32_t *)(q8 + i), sum); + } + return sum; +} + +__device__ static float dev_dot_q4_K_q8_K_block(const cuda_block_q4_K *x, const cuda_block_q8_K *y) { + const float xd = dev_f16_to_f32(x->d); + const float xmin = dev_f16_to_f32(x->dmin); + int isum = 0; + int summs = 0; + #pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + uint8_t sc, m; + dev_q4_K_get_scale_min(j, x->scales, &sc, &m); + summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]); + const uint32_t byte_off = (j >> 1u) * 32u; + const int shift = (j & 1u) ? 4 : 0; + isum += (int)sc * dev_dot_q4_32(x->qs + byte_off, y->qs + j * 32u, shift); + } + return y->d * xd * (float)isum - y->d * xmin * (float)summs; +} + +/* Vector-load variant of dev_dot_q4_K_q8_K_block: loads the whole 144-byte + * Q4_K block with nine 16B loads (requires a 16B-aligned tensor base; block + * stride 144 and row strides are 16B multiples), then computes the exact same + * integer sums and float finish. Same values in the same order, so results + * are bit-identical; the wide loads just improve DRAM/memory-level + * parallelism for the bandwidth-bound decode matvecs. */ +__device__ __forceinline__ static void dev_dot_q4_K_q8_K_block_vec( + const cuda_block_q4_K *x, + const cuda_block_q8_K *y, + float *out_acc) { + const uint4 hdr = *(const uint4 *)x; /* d, dmin, scales[12] */ + uint4 qv[8]; +#pragma unroll + for (uint32_t i = 0; i < 8u; i++) qv[i] = ((const uint4 *)(x->qs))[i]; + const uint16_t xd_u = (uint16_t)(hdr.x & 0xffffu); + const uint16_t xmin_u = (uint16_t)(hdr.x >> 16u); + uint8_t scales[12]; + scales[0] = (uint8_t)(hdr.y); + scales[1] = (uint8_t)(hdr.y >> 8); + scales[2] = (uint8_t)(hdr.y >> 16); + scales[3] = (uint8_t)(hdr.y >> 24); + scales[4] = (uint8_t)(hdr.z); + scales[5] = (uint8_t)(hdr.z >> 8); + scales[6] = (uint8_t)(hdr.z >> 16); + scales[7] = (uint8_t)(hdr.z >> 24); + scales[8] = (uint8_t)(hdr.w); + scales[9] = (uint8_t)(hdr.w >> 8); + scales[10] = (uint8_t)(hdr.w >> 16); + scales[11] = (uint8_t)(hdr.w >> 24); + const float xd = dev_f16_to_f32(xd_u); + const float xmin = dev_f16_to_f32(xmin_u); + int isum = 0; + int summs = 0; + const int32_t *qw = (const int32_t *)qv; +#pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + uint8_t sc, m; + dev_q4_K_get_scale_min(j, scales, &sc, &m); + summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]); + const uint32_t word_off = (j >> 1u) * 8u; + const int shift = (j & 1u) ? 4 : 0; + int32_t sum = 0; +#pragma unroll + for (uint32_t i = 0; i < 8u; i++) { + const int32_t v = (qw[word_off + i] >> shift) & 0x0f0f0f0f; + sum = __dp4a(v, *(const int32_t *)(y->qs + j * 32u + i * 4u), sum); + } + isum += (int)sc * sum; + } + *out_acc += y->d * xd * (float)isum - y->d * xmin * (float)summs; +} + +__device__ static void dev_dot_q4_K_q8_K_block8( + const cuda_block_q4_K *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + const cuda_block_q8_K *y4, + const cuda_block_q8_K *y5, + const cuda_block_q8_K *y6, + const cuda_block_q8_K *y7, + uint32_t n, + float acc[8]) { + const float xd = dev_f16_to_f32(x->d); + const float xmin = dev_f16_to_f32(x->dmin); + const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; + int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + + #pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + uint8_t sc, m; + dev_q4_K_get_scale_min(j, x->scales, &sc, &m); + const uint32_t byte_off = (j >> 1u) * 32u; + const int shift = (j & 1u) ? 4 : 0; + for (uint32_t p = 0; p < n; p++) { + summs[p] += (int)m * (int)(ys[p]->bsums[2u * j] + ys[p]->bsums[2u * j + 1u]); + isum[p] += (int)sc * dev_dot_q4_32(x->qs + byte_off, ys[p]->qs + j * 32u, shift); + } + } + + for (uint32_t p = 0; p < n; p++) { + acc[p] += ys[p]->d * xd * (float)isum[p] - ys[p]->d * xmin * (float)summs[p]; + } +} + +__device__ static float dev_dot_q2_K_q8_K_block(const cuda_block_q2_K *x, const cuda_block_q8_K *y) { + const uint8_t *q2 = x->qs; + const int8_t *q8 = y->qs; + const uint8_t *sc = x->scales; + int summs = 0; + for (int j = 0; j < 16; j++) summs += y->bsums[j] * (sc[j] >> 4); + const float dall = y->d * dev_f16_to_f32(x->d); + const float dmin = y->d * dev_f16_to_f32(x->dmin); + int isum = 0; + int is = 0; + for (int k = 0; k < CUDA_QK_K / 128; k++) { + int shift = 0; + for (int j = 0; j < 4; j++) { + int d = sc[is++] & 0x0f; + isum += d * dev_dot_q2_16(q2, q8, shift); + d = sc[is++] & 0x0f; + isum += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); + shift += 2; + q8 += 32; + } + q2 += 32; + } + return dall * (float)isum - dmin * (float)summs; +} + +__device__ static void dev_dot_q2_K_q8_K_block4( + const cuda_block_q2_K *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + uint32_t n, + float acc[4]) { + const uint8_t *sc = x->scales; + const float xd = dev_f16_to_f32(x->d); + const float xmin = dev_f16_to_f32(x->dmin); + const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 }; + int isum[4] = {0, 0, 0, 0}; + int summs[4] = {0, 0, 0, 0}; + for (uint32_t p = 0; p < n; p++) { + for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4); + } + for (uint32_t p = 0; p < n; p++) { + const uint8_t *q2 = x->qs; + const int8_t *q8 = ys[p]->qs; + int is = 0; + for (int k = 0; k < CUDA_QK_K / 128; k++) { + int shift = 0; + for (int j = 0; j < 4; j++) { + int d = sc[is++] & 0x0f; + isum[p] += d * dev_dot_q2_16(q2, q8, shift); + d = sc[is++] & 0x0f; + isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); + shift += 2; + q8 += 32; + } + q2 += 32; + } + } + for (uint32_t p = 0; p < n; p++) { + const float yd = ys[p]->d; + acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p]; + } +} + +__device__ static void dev_dot_q2_K_q8_K_block8( + const cuda_block_q2_K *x, + const cuda_block_q8_K *y0, + const cuda_block_q8_K *y1, + const cuda_block_q8_K *y2, + const cuda_block_q8_K *y3, + const cuda_block_q8_K *y4, + const cuda_block_q8_K *y5, + const cuda_block_q8_K *y6, + const cuda_block_q8_K *y7, + uint32_t n, + float acc[8]) { + const uint8_t *sc = x->scales; + const float xd = dev_f16_to_f32(x->d); + const float xmin = dev_f16_to_f32(x->dmin); + const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 }; + int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (uint32_t p = 0; p < n; p++) { + for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4); + } + for (uint32_t p = 0; p < n; p++) { + const uint8_t *q2 = x->qs; + const int8_t *q8 = ys[p]->qs; + int is = 0; + for (int k = 0; k < CUDA_QK_K / 128; k++) { + int shift = 0; + for (int j = 0; j < 4; j++) { + int d = sc[is++] & 0x0f; + isum[p] += d * dev_dot_q2_16(q2, q8, shift); + d = sc[is++] & 0x0f; + isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift); + shift += 2; + q8 += 32; + } + q2 += 32; + } + } + for (uint32_t p = 0; p < n; p++) { + const float yd = ys[p]->d; + acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p]; + } +} + +__device__ static float half_warp_sum_f32(float v, uint32_t lane16) { + uint32_t mask = 0xffffu << (threadIdx.x & 16u); + for (int offset = 8; offset > 0; offset >>= 1) { + v += __shfl_down_sync(mask, v, offset, 16); + } + (void)lane16; + return v; +} + +__device__ static float quarter_warp_sum_f32(float v, uint32_t lane8) { + uint32_t mask = 0xffu << (threadIdx.x & 24u); + for (int offset = 4; offset > 0; offset >>= 1) { + v += __shfl_down_sync(mask, v, offset, 8); + } + (void)lane8; + return v; +} + +__global__ static void q8_K_quantize_kernel(cuda_block_q8_K *out, const float *x, uint32_t in_dim, uint32_t n_rows) { + uint32_t b = blockIdx.x; + uint32_t row = blockIdx.y; + if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; + const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; + cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b; + __shared__ float abs_part[256]; + __shared__ float val_part[256]; + __shared__ float maxv_s; + __shared__ float iscale_s; + uint32_t tid = threadIdx.x; + float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; + abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; + val_part[tid] = v; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { + abs_part[tid] = abs_part[tid + stride]; + val_part[tid] = val_part[tid + stride]; + } + __syncthreads(); + } + float amax = abs_part[0]; + if (amax == 0.0f) { + if (tid == 0) yb->d = 0.0f; + if (tid < CUDA_QK_K) yb->qs[tid] = 0; + if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; + return; + } + if (tid == 0) { + maxv_s = val_part[0]; + iscale_s = -127.0f / maxv_s; + } + __syncthreads(); + if (tid < CUDA_QK_K) { + int qv = (int)lrintf(iscale_s * xr[tid]); + if (qv > 127) qv = 127; + if (qv < -128) qv = -128; + yb->qs[tid] = (int8_t)qv; + } + __syncthreads(); + if (tid < CUDA_QK_K / 16) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; + yb->bsums[tid] = (int16_t)sum; + } + if (tid == 0) yb->d = 1.0f / iscale_s; +} + +/* Decode-only dual quantizer. The Q8_0 half mirrors + * quantize_q8_0_f32_kernel's 32-thread reduction and expression order, while + * the Q8_K half remains byte-for-byte the ordinary routed-MoE quantizer. */ +__global__ static void q8_K_q8_0_quantize_kernel( + cuda_block_q8_K *out, + int8_t *q8_0, + float *q8_0_scale, + const float *x, + uint32_t in_dim, + uint32_t n_rows) { + const uint32_t b = blockIdx.x; + const uint32_t row = blockIdx.y; + if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; + const float *xr = x + (uint64_t)row * in_dim + + (uint64_t)b * CUDA_QK_K; + cuda_block_q8_K *yb = out + + (uint64_t)row * (in_dim / CUDA_QK_K) + b; + __shared__ float abs_part[256]; + __shared__ float val_part[256]; + __shared__ float maxv_s; + __shared__ float iscale_s; + const uint32_t tid = threadIdx.x; + const uint32_t lane = tid & 31u; + const uint32_t warp = tid >> 5u; + const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; + + abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; + __syncthreads(); + for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { + if (lane < stride) { + abs_part[tid] = fmaxf(abs_part[tid], abs_part[tid + stride]); + } + __syncthreads(); + } + const uint32_t q8_blocks = in_dim / 32u; + const uint32_t q8_block = b * 8u + warp; + const float d = abs_part[warp * 32u] / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + if (lane == 0u) { + q8_0_scale[(uint64_t)row * q8_blocks + q8_block] = d; + } + int qv = (int)lrintf(v * id); + qv = qv > 127 ? 127 : (qv < -128 ? -128 : qv); + q8_0[((uint64_t)row * q8_blocks + q8_block) * 32u + lane] = + (int8_t)qv; + __syncthreads(); + + abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; + val_part[tid] = v; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { + abs_part[tid] = abs_part[tid + stride]; + val_part[tid] = val_part[tid + stride]; + } + __syncthreads(); + } + const float amax = abs_part[0]; + if (amax == 0.0f) { + if (tid == 0) yb->d = 0.0f; + if (tid < CUDA_QK_K) yb->qs[tid] = 0; + if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; + return; + } + if (tid == 0) { + maxv_s = val_part[0]; + iscale_s = -127.0f / maxv_s; + } + __syncthreads(); + if (tid < CUDA_QK_K) { + int kv = (int)lrintf(iscale_s * xr[tid]); + if (kv > 127) kv = 127; + if (kv < -128) kv = -128; + yb->qs[tid] = (int8_t)kv; + } + __syncthreads(); + if (tid < CUDA_QK_K / 16) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; + yb->bsums[tid] = (int16_t)sum; + } + if (tid == 0) yb->d = 1.0f / iscale_s; +} + +__device__ __forceinline__ static bool moe_owned_local_expert( + int32_t expert, + uint32_t expert_base, + uint32_t expert_count, + uint32_t *local_expert) { + if (expert < 0) return false; + const uint32_t e = (uint32_t)expert; + if (e < expert_base || e - expert_base >= expert_count) return false; + if (local_expert) *local_expert = e - expert_base; + return true; +} + +/* Quantize only selected slots owned by this expert-parallel rank. Rows keep + * their original slot index so the final rank-local reduction can visit slots + * in canonical order without compaction or a host synchronization. */ +__global__ static void q8_K_quantize_owned_kernel( + cuda_block_q8_K *out, + const float *x, + const int32_t *selected, + uint32_t in_dim, + uint32_t n_rows, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t b = blockIdx.x; + const uint32_t row = blockIdx.y; + if (row >= n_rows || b >= in_dim / CUDA_QK_K) return; + if (!moe_owned_local_expert(selected[row], expert_base, expert_count, NULL)) return; + + const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; + cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b; + __shared__ float abs_part[256]; + __shared__ float val_part[256]; + __shared__ float maxv_s; + __shared__ float iscale_s; + const uint32_t tid = threadIdx.x; + const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f; + abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f; + val_part[tid] = v; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { + abs_part[tid] = abs_part[tid + stride]; + val_part[tid] = val_part[tid + stride]; + } + __syncthreads(); + } + const float amax = abs_part[0]; + if (amax == 0.0f) { + if (tid == 0) yb->d = 0.0f; + if (tid < CUDA_QK_K) yb->qs[tid] = 0; + if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0; + return; + } + if (tid == 0) { + maxv_s = val_part[0]; + iscale_s = -127.0f / maxv_s; + } + __syncthreads(); + if (tid < CUDA_QK_K) { + int qv = (int)lrintf(iscale_s * xr[tid]); + if (qv > 127) qv = 127; + if (qv < -128) qv = -128; + yb->qs[tid] = (int8_t)qv; + } + __syncthreads(); + if (tid < CUDA_QK_K / 16) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i]; + yb->bsums[tid] = (int16_t)sum; + } + if (tid == 0) yb->d = 1.0f / iscale_s; +} + +__global__ static void moe_filter_owned_pairs_kernel( + int32_t *selected, + float *weights, + uint64_t pair_count, + uint32_t n_total_expert, + uint32_t expert_base, + uint32_t expert_count) { + const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (pair >= pair_count) return; + const int32_t expert_i = selected[pair]; + if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert && + (uint32_t)expert_i >= expert_base && + (uint32_t)expert_i - expert_base < expert_count) { + selected[pair] = expert_i - (int32_t)expert_base; + } else { + selected[pair] = -1; + weights[pair] = 0.0f; + } +} + +__global__ static void q8_K_quantize_sidecar_kernel( + cuda_block_q8_K *out, + const float *x, + const float *amax_sidecar, + uint32_t in_dim, + uint32_t n_rows) { + uint32_t b = blockIdx.x; + uint32_t row = blockIdx.y; + const uint32_t blocks = in_dim / CUDA_QK_K; + if (row >= n_rows || b >= blocks) return; + const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K; + const float *sc = amax_sidecar + ((uint64_t)row * blocks + b) * 32u; + cuda_block_q8_K *yb = out + (uint64_t)row * blocks + b; + __shared__ float abs_part[32]; + __shared__ float val_part[32]; + __shared__ float iscale_s; + const uint32_t tid = threadIdx.x; + if (tid < 32u) { + const float v = sc[tid]; + abs_part[tid] = fabsf(v); + val_part[tid] = v; + } + __syncthreads(); + for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) { + if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { + abs_part[tid] = abs_part[tid + stride]; + val_part[tid] = val_part[tid + stride]; + } + __syncthreads(); + } + const float amax = abs_part[0]; + if (amax == 0.0f) { + if (tid == 0u) yb->d = 0.0f; + if (tid < CUDA_QK_K) yb->qs[tid] = 0; + if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0; + return; + } + if (tid == 0u) iscale_s = -127.0f / val_part[0]; + __syncthreads(); + if (tid < CUDA_QK_K) { + int qv = (int)lrintf(iscale_s * xr[tid]); + if (qv > 127) qv = 127; + if (qv < -128) qv = -128; + yb->qs[tid] = (int8_t)qv; + } + __syncthreads(); + if (tid < CUDA_QK_K / 16u) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i]; + yb->bsums[tid] = (int16_t)sum; + } + if (tid == 0u) yb->d = 1.0f / iscale_s; +} + +__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t row = blockIdx.x; + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = threadIdx.x; b < xq_blocks; b += blockDim.x) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + __shared__ float partial_gate[256]; + __shared__ float partial_up[256]; + partial_gate[threadIdx.x] = gate; + partial_up[threadIdx.x] = up; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride]; + partial_up[threadIdx.x] += partial_up[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + gate = partial_gate[0]; + up = partial_up[0]; + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_warp8_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t lane = threadIdx.x & 31u; + uint32_t warp = threadIdx.x >> 5u; + uint32_t row = blockIdx.x * 8u + warp; + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_hwarp16_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t lane = threadIdx.x & 15u; + uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 16u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = half_warp_sum_f32(gate, lane); + up = half_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +// perf-04: launch-geometry tuning for the routed-MoE gate/up decode kernels +// (moe_gate_up_mid_qwarp32 / _decode_lut_qwarp32 / _decode_q4K_qwarp32). Each +// block processes MOE_DECODE_ROW_TILES tiles of 32 rows (row_lane in [0,32)). +// The historical value was 4 (128 rows/block -> ~96 blocks, occupancy ~16%, +// "grid too small to fill the device"). Lowering it issues correspondingly more +// blocks (e.g. 1 tile -> 32 rows/block -> ~4x more blocks -> ~384) to fill the +// SMs. The per-row arithmetic is identical regardless of this value, so output +// is bit-identical; only the qgrid.x divisor must match MOE_DECODE_ROWS_PER_BLOCK. +#ifndef MOE_DECODE_ROW_TILES +#define MOE_DECODE_ROW_TILES 1u +#endif +#define MOE_DECODE_ROWS_PER_BLOCK (32u * MOE_DECODE_ROW_TILES) + +__global__ static void moe_gate_up_mid_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t pair = blockIdx.y; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { + uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } + } +} + +__global__ static void moe_gate_up_mid_decode_lut_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t pair = blockIdx.y; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + xqb = sxq; + } + for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { + uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs); + up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } + } +} + +__global__ static void moe_gate_up_mid_decode_lut_owned_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t expert_base, + uint32_t expert_count, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t pair = blockIdx.y; + uint32_t expert = 0u; + if (!moe_owned_local_expert(selected[pair], expert_base, expert_count, + &expert)) return; + const cuda_block_q8_K *xqb = xq; + __shared__ cuda_block_q8_K sxq[16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + xqb = sxq; + } + for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { + uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs); + up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair]; + } + } +} + +__global__ static void moe_count_sorted_pairs_kernel( + uint32_t *counts, + const int32_t *selected, + uint32_t pair_count, + uint32_t n_total_expert) { + uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); + if (pair >= pair_count) return; + int32_t expert_i = selected[pair]; + if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return; + atomicAdd(counts + (uint32_t)expert_i, 1u); +} + +__global__ static void moe_prefix_sorted_pairs_kernel( + uint32_t *offsets, + uint32_t *cursors, + const uint32_t *counts, + uint32_t n_total_expert) { + if (threadIdx.x == 0) { + uint32_t sum = 0; + for (uint32_t e = 0; e < n_total_expert; e++) { + offsets[e] = sum; + cursors[e] = sum; + sum += counts[e]; + } + offsets[n_total_expert] = sum; + } +} + +__global__ static void moe_scatter_sorted_pairs_kernel( + uint32_t *sorted_pairs, + uint32_t *cursors, + const int32_t *selected, + uint32_t pair_count, + uint32_t n_total_expert) { + uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); + if (pair >= pair_count) return; + int32_t expert_i = selected[pair]; + if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return; + uint32_t pos = atomicAdd(cursors + (uint32_t)expert_i, 1u); + sorted_pairs[pos] = pair; +} + +__global__ static void moe_build_expert_tile_offsets_kernel( + uint32_t *tile_offsets, + uint32_t *tile_total, + const uint32_t *counts, + uint32_t block_m, + uint32_t n_total_expert) { + if (threadIdx.x == 0) { + uint32_t sum = 0; + for (uint32_t e = 0; e < n_total_expert; e++) { + tile_offsets[e] = sum; + sum += (counts[e] + block_m - 1u) / block_m; + } + tile_offsets[n_total_expert] = sum; + *tile_total = sum; + } +} + +__global__ static void moe_build_expert_tiles_kernel( + uint32_t *tile_experts, + uint32_t *tile_starts, + const uint32_t *tile_offsets, + const uint32_t *counts, + uint32_t block_m, + uint32_t n_total_expert) { + uint32_t e = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); + if (e >= n_total_expert) return; + uint32_t ntiles = (counts[e] + block_m - 1u) / block_m; + uint32_t off = tile_offsets[e]; + for (uint32_t t = 0; t < ntiles; t++) { + tile_experts[off + t] = e; + tile_starts[off + t] = t * block_m; + } +} + +/* Decode-sized routed batches spend more host time launching metadata kernels + * than doing the <= 96 pair / 128 expert setup. Build both tile lists in one + * deterministic block; the expensive expert kernels remain unchanged. */ +__global__ static void moe_prepare_sorted_tiles_small_kernel( + uint32_t *counts, + uint32_t *offsets, + uint32_t *cursors, + uint32_t *sorted_pairs, + uint32_t *tile_offsets, + uint32_t *tile_total, + uint32_t *tile_experts, + uint32_t *tile_starts, + uint32_t *tile16_offsets, + uint32_t *tile16_total, + uint32_t *tile16_experts, + uint32_t *tile16_starts, + const int32_t *selected, + uint32_t pair_count, + uint32_t n_total_expert, + uint32_t block_m, + bool build_tile16) { + if (blockIdx.x != 0) return; + const uint32_t tid = threadIdx.x; + __shared__ uint32_t local_counts[128]; + __shared__ int32_t local_selected[96]; + + for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { + local_counts[e] = 0u; + } + for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { + local_selected[pair] = selected[pair]; + } + __syncthreads(); + + for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { + const int32_t expert_i = local_selected[pair]; + if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) { + atomicAdd(local_counts + (uint32_t)expert_i, 1u); + } + } + __syncthreads(); + + for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { + counts[e] = local_counts[e]; + } + if (tid == 0u) { + uint32_t pair_sum = 0u; + uint32_t tile_sum = 0u; + uint32_t tile16_sum = 0u; + for (uint32_t e = 0; e < n_total_expert; e++) { + const uint32_t count = local_counts[e]; + offsets[e] = pair_sum; + pair_sum += count; + cursors[e] = pair_sum; + tile_offsets[e] = tile_sum; + tile_sum += (count + block_m - 1u) / block_m; + if (build_tile16) { + tile16_offsets[e] = tile16_sum; + tile16_sum += (count + 15u) / 16u; + } + } + offsets[n_total_expert] = pair_sum; + tile_offsets[n_total_expert] = tile_sum; + *tile_total = tile_sum; + if (build_tile16) { + tile16_offsets[n_total_expert] = tile16_sum; + *tile16_total = tile16_sum; + } + } + __syncthreads(); + + for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) { + const int32_t expert_i = local_selected[pair]; + if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) { + uint32_t rank = 0u; + for (uint32_t prev = 0; prev < pair; prev++) { + rank += local_selected[prev] == expert_i; + } + sorted_pairs[offsets[(uint32_t)expert_i] + rank] = pair; + } + } + for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) { + const uint32_t count = local_counts[e]; + const uint32_t ntiles = (count + block_m - 1u) / block_m; + const uint32_t tile_off = tile_offsets[e]; + for (uint32_t t = 0; t < ntiles; t++) { + tile_experts[tile_off + t] = e; + tile_starts[tile_off + t] = t * block_m; + } + if (build_tile16) { + const uint32_t ntiles16 = (count + 15u) / 16u; + const uint32_t tile16_off = tile16_offsets[e]; + for (uint32_t t = 0; t < ntiles16; t++) { + tile16_experts[tile16_off + t] = e; + tile16_starts[tile16_off + t] = t * 16u; + } + } + } +} + +__global__ static void moe_gate_up_mid_sorted_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t pair = sorted_pairs[blockIdx.y]; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_expert_tile8_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t group = threadIdx.x >> 3u; + uint32_t lane = threadIdx.x & 7u; + uint32_t pair_slot = group & 7u; + uint32_t row_lane = group >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_pair = tile_starts[tile] + pair_slot; + if (local_pair >= counts[expert]) return; + uint32_t sorted_idx = offsets[expert] + local_pair; + uint32_t pair = sorted_pairs[sorted_idx]; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + + for (uint32_t rr = 0; rr < 2u; rr++) { + uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } + } +} + +__global__ static void moe_gate_up_mid_expert_tile4_row32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[4][16]; + uint32_t pair[4] = {0, 0, 0, 0}; + uint32_t tok[4] = {0, 0, 0, 0}; + uint32_t slot[4] = {0, 0, 0, 0}; + const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 4u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= expert_mid_dim) return; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + float up[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + dev_dot_iq2_xxs_q8_K_block4(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, gate); + dev_dot_iq2_xxs_q8_K_block4(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, up); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } +} + +__global__ static void moe_gate_up_mid_expert_tile8_row32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= expert_mid_dim) return; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, + s_iq2_grid, s_iq2_signs); + dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, + s_iq2_grid, s_iq2_signs); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } +} + +__global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < 64u; rr++) { + uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, + s_iq2_grid, s_iq2_signs); + dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, + s_iq2_grid, s_iq2_signs); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } + } +} + +template +__global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { + uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate, + s_iq2_grid, s_iq2_signs); + dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up, + s_iq2_grid, s_iq2_signs); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } + } +} + +__global__ static void moe_gate_up_mid_sorted_p2_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t pair_count, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t pair_lane = (threadIdx.x >> 3u) & 1u; + uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); + uint32_t sorted_idx = blockIdx.y * 2u + pair_lane; + if (row >= expert_mid_dim || sorted_idx >= pair_count) return; + uint32_t pair = sorted_pairs[sorted_idx]; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b); + up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static DS4_CUDA_UNUSED void moe_down_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t row = blockIdx.x; + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = threadIdx.x; b < midq_blocks; b += blockDim.x) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + __shared__ float partial[256]; + partial[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0]; +} + +__global__ static DS4_CUDA_UNUSED void moe_down_warp8_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t lane = threadIdx.x & 31u; + uint32_t warp = threadIdx.x >> 5u; + uint32_t row = blockIdx.x * 8u + warp; + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 32u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = warp_sum_f32(acc); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +__global__ static DS4_CUDA_UNUSED void moe_down_hwarp16_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t lane = threadIdx.x & 15u; + uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 16u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = half_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +__global__ static void moe_down_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +__global__ static void moe_gate_up_mid_decode_q4K_qwarp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t pair = blockIdx.y; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { + uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 15u; + uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 16u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = half_warp_sum_f32(gate, lane); + up = half_warp_sum_f32(up, lane); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 15u; + uint32_t group = threadIdx.x >> 4u; + uint32_t pair = blockIdx.y; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + if (group >= 8u) return; + uint32_t row = blockIdx.x * 8u + group; + if (row >= expert_mid_dim) return; + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 16u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = half_warp_sum_f32(gate, lane); + up = half_warp_sum_f32(up, lane); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_warp32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 31u; + uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_kernel( + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t lane = threadIdx.x & 31u; + uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + /* Word-wise cooperative staging copy (same bytes, all lanes busy). */ + const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + uint32_t *dst = (uint32_t *)sxq; + const uint32_t *srcw = (const uint32_t *)xqb; + for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i]; + __syncthreads(); + xqb = sxq; + } + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base | + gate_row_bytes | gate_expert_bytes) & 15u) == 0u); + if (vec_ok) { + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate); + dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up); + } + } else { + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel( + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t expert_base, + uint32_t expert_count, + float clamp) { + uint32_t lane = threadIdx.x & 31u; + uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u); + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t expert = 0u; + if (!moe_owned_local_expert(selected[pair], expert_base, expert_count, + &expert)) return; + const cuda_block_q8_K *xqb = xq; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + uint32_t *dst = (uint32_t *)sxq; + const uint32_t *srcw = (const uint32_t *)xqb; + for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i]; + __syncthreads(); + xqb = sxq; + } + const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base | + gate_row_bytes | gate_expert_bytes) & 15u) == 0u); + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + if (vec_ok) { + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate); + dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up); + } + } else { + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair]; + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel( + float *mid_out, + float *amax_sidecar, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t row = blockIdx.x * 8u + warp; + const uint32_t pair = blockIdx.y; + const uint32_t tok = pair / n_expert; + const uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + __shared__ float tile_vals[8]; + __shared__ float tile_abs[8]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + + float midv = 0.0f; + const bool valid = row < expert_mid_dim; + if (valid) { + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + midv = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + mid_out[off] = midv; + } + } + if (lane == 0u) { + tile_vals[warp] = midv; + tile_abs[warp] = valid ? fabsf(midv) : 0.0f; + } + __syncthreads(); + if (threadIdx.x == 0u) { + float best_abs = tile_abs[0]; + float best_val = tile_vals[0]; + #pragma unroll + for (uint32_t i = 1u; i < 8u; i++) { + if (tile_abs[i] > best_abs) { + best_abs = tile_abs[i]; + best_val = tile_vals[i]; + } + } + const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; + const uint32_t qblock = blockIdx.x / 32u; + const uint32_t tile = blockIdx.x & 31u; + if (qblock < midq_blocks) { + amax_sidecar[((uint64_t)pair * midq_blocks + qblock) * 32u + tile] = best_val; + } + } +} + +__global__ static void moe_gate_up_mid_decode_q4K_warp32_row16_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t lane = threadIdx.x & 31u; + uint32_t warp = threadIdx.x >> 5u; + uint32_t row = blockIdx.x * 16u + warp; + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = warp_sum_f32(gate); + up = warp_sum_f32(up); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_gate_up_midq_decode_q4K_qwarp32_kernel( + float *mid_out, + cuda_block_q8_K *midq, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row_lane = threadIdx.x >> 3u; + const uint32_t qblock = blockIdx.x; + const uint32_t pair = blockIdx.y; + const uint32_t tok = pair / n_expert; + const uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const uint32_t expert = (uint32_t)expert_i; + const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks; + __shared__ cuda_block_q8_K sxq[16]; + __shared__ float vals[CUDA_QK_K]; + __shared__ float abs_part[CUDA_QK_K]; + __shared__ float val_part[CUDA_QK_K]; + __shared__ float iscale_s; + + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; + __syncthreads(); + xqb = sxq; + } + + const float w = weights[(uint64_t)tok * n_expert + slot]; + #pragma unroll + for (uint32_t rr = 0; rr < 8u; rr++) { + const uint32_t row_in_block = row_lane + rr * 32u; + const uint32_t row = qblock * CUDA_QK_K + row_in_block; + float midv = 0.0f; + if (row < expert_mid_dim) { + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b); + up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + midv = (gate / (1.0f + expf(-gate))) * up * w; + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + mid_out[off] = midv; + } + } + if (lane == 0u) vals[row_in_block] = midv; + } + __syncthreads(); + + cuda_block_q8_K *yb = midq + (uint64_t)pair * (expert_mid_dim / CUDA_QK_K) + qblock; + const uint32_t tid = threadIdx.x; + const float v = vals[tid]; + abs_part[tid] = fabsf(v); + val_part[tid] = v; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride && abs_part[tid + stride] > abs_part[tid]) { + abs_part[tid] = abs_part[tid + stride]; + val_part[tid] = val_part[tid + stride]; + } + __syncthreads(); + } + const float amax = abs_part[0]; + if (amax == 0.0f) { + if (tid == 0u) yb->d = 0.0f; + if (tid < CUDA_QK_K) yb->qs[tid] = 0; + if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0; + return; + } + if (tid == 0u) { + iscale_s = -127.0f / val_part[0]; + } + __syncthreads(); + int qv = (int)lrintf(iscale_s * v); + if (qv > 127) qv = 127; + if (qv < -128) qv = -128; + yb->qs[tid] = (int8_t)qv; + __syncthreads(); + if (tid < CUDA_QK_K / 16u) { + int sum = 0; + for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i]; + yb->bsums[tid] = (int16_t)sum; + } + if (tid == 0u) yb->d = 1.0f / iscale_s; +} + +template +__global__ static void moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + tok[np] = pair[np] / n_expert; + slot[np] = pair[np] - tok[np] * n_expert; + xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; + } + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) { + uint32_t p = i / xq_blocks; + uint32_t b = i - p * xq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { + uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + dev_dot_q4_K_q8_K_block8(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate); + dev_dot_q4_K_q8_K_block8(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up); + } + for (uint32_t p = 0; p < np; p++) { + gate[p] = quarter_warp_sum_f32(gate[p], lane); + up[p] = quarter_warp_sum_f32(up[p], lane); + if (lane == 0) { + if (clamp > 1.0e-6f) { + if (gate[p] > clamp) gate[p] = clamp; + if (up[p] > clamp) up[p] = clamp; + if (up[p] < -clamp) up[p] = -clamp; + } + const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate[p]; + up_out[off] = up[p]; + } + mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]]; + } + } + } +} + +__global__ static void moe_down_sum6_qwarp32_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + if (row >= out_dim) return; + float total = 0.0f; + #pragma unroll + for (uint32_t slot = 0; slot < 6u; slot++) { + int32_t expert_i = selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) total += acc; + } + if (lane == 0) out[row] = total; +} + +__global__ static void moe_down_owned_slots_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t slot = blockIdx.y; + if (row >= out_dim || slot >= 6u) return; + uint32_t expert = 0; + if (!moe_owned_local_expert(selected[slot], expert_base, + expert_count, &expert)) { + return; + } + const cuda_block_q2_K *wr = + (const cuda_block_q2_K *)(down_base + + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc; +} + +/* Map one of two packed operands for a three-slot reduction group. The only + * multi-slot operand is the peer-owned prefix (slots 0+1 within the group), + * which can be pre-added exactly because the reference reduction starts from + * +0. Every other peer slot remains a distinct operand in original order. */ +__device__ __forceinline__ static int moe_owned_packed_component( + const int32_t *selected, + uint32_t group, + uint32_t component, + uint32_t expert_base, + uint32_t expert_count, + bool *prefix_pair) { + const uint32_t slot0 = group * 3u; + uint32_t mask = 0u; + #pragma unroll + for (uint32_t i = 0; i < 3u; i++) { + if (moe_owned_local_expert(selected[slot0 + i], expert_base, + expert_count, NULL)) { + mask |= 1u << i; + } + } + *prefix_pair = false; + if ((mask & 3u) == 3u) { + if (component == 0u) { + *prefix_pair = true; + return (int)slot0; + } + return (mask & 4u) != 0u ? (int)(slot0 + 2u) : -1; + } + uint32_t ordinal = 0u; + #pragma unroll + for (uint32_t i = 0; i < 3u; i++) { + if ((mask & (1u << i)) == 0u) continue; + if (ordinal++ == component) return (int)(slot0 + i); + } + return -1; +} + +__global__ static void moe_down_owned_packed_qwarp32_kernel( + float *packed_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t packed_slot = blockIdx.y; + if (row >= out_dim || packed_slot >= 4u) return; + bool prefix_pair = false; + const int first_slot = moe_owned_packed_component( + selected, packed_slot / 2u, packed_slot & 1u, + expert_base, expert_count, &prefix_pair); + if (first_slot < 0) { + if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f; + return; + } + + float packed = 0.0f; + const uint32_t n_slots = prefix_pair ? 2u : 1u; + #pragma unroll + for (uint32_t i = 0; i < 2u; i++) { + if (i >= n_slots) break; + const uint32_t slot = (uint32_t)first_slot + i; + uint32_t expert = 0; + if (!moe_owned_local_expert(selected[slot], expert_base, + expert_count, &expert)) { + continue; + } + const cuda_block_q2_K *wr = + (const cuda_block_q2_K *)(down_base + + (uint64_t)expert * down_expert_bytes + + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) { + packed = prefix_pair ? __fadd_rn(packed, acc) : acc; + } + } + if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed; +} + +__global__ static void moe_down_sum3_qwarp32_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + if (row >= out_dim) return; + float total = 0.0f; + #pragma unroll + for (uint32_t slot = 0; slot < 3u; slot++) { + int32_t expert_i = selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) total += acc; + } + if (lane == 0) out[row] = total; +} + +__global__ static void moe_down_q4K_sum6_qwarp32_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + if (row >= out_dim) return; + const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u); + float total = 0.0f; + #pragma unroll + for (uint32_t slot = 0; slot < 6u; slot++) { + int32_t expert_i = selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + if (vec_ok) { + for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); + } else { + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) total += acc; + } + if (lane == 0) out[row] = total; +} + +__global__ static void moe_down_q4K_owned_slots_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t slot = blockIdx.y; + if (row >= out_dim || slot >= 6u) return; + uint32_t expert = 0; + if (!moe_owned_local_expert(selected[slot], expert_base, + expert_count, &expert)) { + return; + } + const cuda_block_q4_K *wr = + (const cuda_block_q4_K *)(down_base + + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | + down_expert_bytes) & 15u) == 0u); + float acc = 0.0f; + if (vec_ok) { + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); + } + } else { + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + } + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc; +} + +__global__ static void moe_down_q4K_owned_packed_qwarp32_kernel( + float *packed_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t packed_slot = blockIdx.y; + if (row >= out_dim || packed_slot >= 4u) return; + bool prefix_pair = false; + const int first_slot = moe_owned_packed_component( + selected, packed_slot / 2u, packed_slot & 1u, + expert_base, expert_count, &prefix_pair); + if (first_slot < 0) { + if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f; + return; + } + + const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | + down_expert_bytes) & 15u) == 0u); + float packed = 0.0f; + const uint32_t n_slots = prefix_pair ? 2u : 1u; + #pragma unroll + for (uint32_t i = 0; i < 2u; i++) { + if (i >= n_slots) break; + const uint32_t slot = (uint32_t)first_slot + i; + uint32_t expert = 0; + if (!moe_owned_local_expert(selected[slot], expert_base, + expert_count, &expert)) { + continue; + } + const cuda_block_q4_K *wr = + (const cuda_block_q4_K *)(down_base + + (uint64_t)expert * down_expert_bytes + + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + if (vec_ok) { + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); + } + } else { + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + } + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) { + packed = prefix_pair ? __fadd_rn(packed, acc) : acc; + } + } + if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed; +} + +__global__ static void moe_owned_slots_combine_fixed3_kernel( + float *out, + const float *home_slots, + const float *peer_slots, + const int32_t *selected, + uint32_t out_dim, + uint32_t expert_split) { + const uint32_t col = + (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); + const uint32_t row = blockIdx.y; + if (col >= out_dim) return; + out += (uint64_t)row * out_dim; + home_slots += (uint64_t)row * 6u * out_dim; + peer_slots += (uint64_t)row * 6u * out_dim; + selected += (uint64_t)row * 6u; + float slotv[6]; + #pragma unroll + for (uint32_t slot = 0; slot < 6u; slot++) { + const int32_t expert = selected[slot]; + if (expert < 0 || + (uint32_t)expert >= 2u * expert_split) { + slotv[slot] = 0.0f; + } else { + const bool on_home = (uint32_t)expert < expert_split; + const float *src = on_home ? home_slots : peer_slots; + slotv[slot] = src[(uint64_t)slot * out_dim + col]; + } + } + float home = __fadd_rn(0.0f, slotv[0]); + home = __fadd_rn(home, slotv[1]); + home = __fadd_rn(home, slotv[2]); + float peer = __fadd_rn(0.0f, slotv[3]); + peer = __fadd_rn(peer, slotv[4]); + peer = __fadd_rn(peer, slotv[5]); + out[col] = __fadd_rn(home, peer); +} + +__device__ static float moe_owned_packed_combine_row( + const float *home_slots, + const float *peer_packed, + const int32_t *selected, + uint32_t row, + uint32_t out_dim, + uint32_t expert_split) { + float groups[2]; + #pragma unroll + for (uint32_t group = 0; group < 2u; group++) { + const uint32_t slot0 = group * 3u; + uint32_t peer_mask = 0u; + uint32_t valid_mask = 0u; + #pragma unroll + for (uint32_t i = 0; i < 3u; i++) { + const int32_t expert = selected[slot0 + i]; + if (expert >= 0 && (uint32_t)expert < 2u * expert_split) { + valid_mask |= 1u << i; + } + if (expert >= 0 && (uint32_t)expert >= expert_split && + (uint32_t)expert < 2u * expert_split) { + peer_mask |= 1u << i; + } + } + const float *packed = peer_packed + + (uint64_t)group * 2u * out_dim + row; + float acc; + if ((peer_mask & 3u) == 3u) { + /* packed[0] is already (+0 + slot0) + slot1. */ + acc = packed[0]; + float slot2 = 0.0f; + if ((peer_mask & 4u) != 0u) { + slot2 = packed[out_dim]; + } else if ((valid_mask & 4u) != 0u) { + slot2 = home_slots[(uint64_t)(slot0 + 2u) * out_dim + row]; + } + acc = __fadd_rn(acc, slot2); + } else { + acc = 0.0f; + uint32_t peer_operand = 0u; + #pragma unroll + for (uint32_t i = 0; i < 3u; i++) { + float value; + if ((peer_mask & (1u << i)) != 0u) { + value = packed[(uint64_t)peer_operand * out_dim]; + peer_operand++; + } else if ((valid_mask & (1u << i)) != 0u) { + value = home_slots[(uint64_t)(slot0 + i) * out_dim + row]; + } else { + value = 0.0f; + } + acc = __fadd_rn(acc, value); + } + } + groups[group] = acc; + } + return __fadd_rn(groups[0], groups[1]); +} + +__global__ static void moe_owned_packed_combine_fixed3_kernel( + float *out, + const float *home_slots, + const float *peer_packed, + const int32_t *selected, + uint32_t out_dim, + uint32_t expert_split) { + const uint32_t row = + (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x); + if (row >= out_dim) return; + out[row] = moe_owned_packed_combine_row( + home_slots, peer_packed, selected, row, out_dim, expert_split); +} + +__global__ static void moe_down_q4K_sum3_qwarp32_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + if (row >= out_dim) return; + const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u); + float total = 0.0f; + #pragma unroll + for (uint32_t slot = 0; slot < 3u; slot++) { + int32_t expert_i = selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + if (vec_ok) { + for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc); + } else { + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) total += acc; + } + if (lane == 0) out[row] = total; +} + +__global__ static void moe_down_q4K_sum3_slotwarp_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim) { + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t slot = lane >> 3u; + const uint32_t qlane = lane & 7u; + const uint32_t row = blockIdx.x * 8u + warp; + if (row >= out_dim) return; + + float acc = 0.0f; + if (slot < 3u) { + int32_t expert_i = selected[slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q4_K *wr = + (const cuda_block_q4_K *)(down_base + + (uint64_t)(uint32_t)expert_i * down_expert_bytes + + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + for (uint32_t b = qlane; b < midq_blocks; b += 8u) { + acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, qlane); + } + + const float s1 = __shfl_sync(0xffffffffu, acc, 8); + const float s2 = __shfl_sync(0xffffffffu, acc, 16); + if (lane == 0u) { + const float s0 = acc; + out[row] = (s0 + s1) + s2; + } +} + +static void routed_moe_decode_graph_destroy_one(int logical_tier) { + if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return; + cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier]; + if (c->exec) (void)cudaGraphExecDestroy(c->exec); + if (c->graph) (void)cudaGraphDestroy(c->graph); + memset(c, 0, sizeof(*c)); +} + +static int routed_moe_decode_q4_graph_launch( + int logical_tier, + float *out, + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_w, + const char *up_w, + const char *down_w, + cuda_block_q8_K *xq, + cuda_block_q8_K *midq, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp, + const float *x) { + if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0; + if (n_expert != 3u && n_expert != 6u) return 0; + uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; + uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; + if (xq_blocks == 0u || midq_blocks == 0u) return 0; + + cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier]; + const bool shape_match = + c->valid && + c->n_expert == n_expert && + c->expert_in_dim == expert_in_dim && + c->expert_mid_dim == expert_mid_dim && + c->out_dim == out_dim; + if (c->valid && !shape_match) { + routed_moe_decode_graph_destroy_one(logical_tier); + c = &g_moe_decode_graph[logical_tier]; + } + + uint32_t x_rows = 1u; + uint32_t mid_rows = n_expert; + dim3 xq_grid(xq_blocks, 1, 1); + dim3 gate_grid((expert_mid_dim + 7u) / 8u, n_expert, 1); + dim3 midq_grid(midq_blocks, n_expert, 1); + dim3 down_grid((out_dim + 31u) / 32u, 1, 1); + dim3 block(256, 1, 1); + + void *xq_args[] = { &xq, &x, &expert_in_dim, &x_rows }; + cudaKernelNodeParams xq_params; + memset(&xq_params, 0, sizeof(xq_params)); + xq_params.func = (void *)q8_K_quantize_kernel; + xq_params.gridDim = xq_grid; + xq_params.blockDim = block; + xq_params.kernelParams = xq_args; + + void *gate_args[] = { + &gate_out, &up_out, &mid_out, &gate_w, &up_w, &xq, &selected, + &weights, &gate_expert_bytes, &gate_row_bytes, &xq_blocks, + &expert_mid_dim, &n_expert, &write_aux, &clamp + }; + cudaKernelNodeParams gate_params; + memset(&gate_params, 0, sizeof(gate_params)); + gate_params.func = (void *)moe_gate_up_mid_decode_q4K_warp32_kernel; + gate_params.gridDim = gate_grid; + gate_params.blockDim = block; + gate_params.kernelParams = gate_args; + + void *midq_args[] = { &midq, &mid_out, &expert_mid_dim, &mid_rows }; + cudaKernelNodeParams midq_params; + memset(&midq_params, 0, sizeof(midq_params)); + midq_params.func = (void *)q8_K_quantize_kernel; + midq_params.gridDim = midq_grid; + midq_params.blockDim = block; + midq_params.kernelParams = midq_args; + + void *down_args[] = { + &out, &down_w, &midq, &selected, &down_expert_bytes, + &down_row_bytes, &midq_blocks, &out_dim + }; + cudaKernelNodeParams down_params; + memset(&down_params, 0, sizeof(down_params)); + down_params.func = n_expert == 6u + ? (void *)moe_down_q4K_sum6_qwarp32_kernel + : (void *)moe_down_q4K_sum3_qwarp32_kernel; + down_params.gridDim = down_grid; + down_params.blockDim = block; + down_params.kernelParams = down_args; + + if (!c->valid) { + cudaError_t err = cudaGraphCreate(&c->graph, 0); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: routed MoE decode graph create failed: %s\n", + cudaGetErrorString(err)); + routed_moe_decode_graph_destroy_one(logical_tier); + return -1; + } + err = cudaGraphAddKernelNode(&c->xq_node, c->graph, NULL, 0, + &xq_params); + if (err == cudaSuccess) { + err = cudaGraphAddKernelNode(&c->gate_node, c->graph, + &c->xq_node, 1, &gate_params); + } + if (err == cudaSuccess) { + err = cudaGraphAddKernelNode(&c->midq_node, c->graph, + &c->gate_node, 1, &midq_params); + } + if (err == cudaSuccess) { + err = cudaGraphAddKernelNode(&c->down_node, c->graph, + &c->midq_node, 1, &down_params); + } + if (err == cudaSuccess) { + err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0); + } + if (err != cudaSuccess) { + fprintf(stderr, "ds4: routed MoE decode graph instantiate failed: %s\n", + cudaGetErrorString(err)); + routed_moe_decode_graph_destroy_one(logical_tier); + return -1; + } + c->n_expert = n_expert; + c->expert_in_dim = expert_in_dim; + c->expert_mid_dim = expert_mid_dim; + c->out_dim = out_dim; + c->valid = 1; + } else { + cudaError_t err = + cudaGraphExecKernelNodeSetParams(c->exec, c->xq_node, + &xq_params); + if (err == cudaSuccess) { + err = cudaGraphExecKernelNodeSetParams(c->exec, c->gate_node, + &gate_params); + } + if (err == cudaSuccess) { + err = cudaGraphExecKernelNodeSetParams(c->exec, c->midq_node, + &midq_params); + } + if (err == cudaSuccess) { + err = cudaGraphExecKernelNodeSetParams(c->exec, c->down_node, + &down_params); + } + if (err != cudaSuccess) { + fprintf(stderr, "ds4: routed MoE decode graph update failed: %s\n", + cudaGetErrorString(err)); + routed_moe_decode_graph_destroy_one(logical_tier); + return -1; + } + } + + cudaError_t err = cudaGraphLaunch(c->exec, 0); + if (err != cudaSuccess) { + fprintf(stderr, "ds4: routed MoE decode graph launch failed: %s\n", + cudaGetErrorString(err)); + routed_moe_decode_graph_destroy_one(logical_tier); + return -1; + } + return 1; +} + +/* Q4_K prefill (n_tokens > 1) down kernel. Mirrors moe_down_qwarp32_kernel + * geometry exactly; only the weight block type and dot helper differ. The + * pair = blockIdx.y indexing means the same grid shape (out_dim/32, n_tokens*n_expert) + * used by the IQ2 path applies here. The downstream moe_sum_kernel is + * weight-type-agnostic and sums these per-pair outputs into the final output. */ +__global__ static void moe_down_q4K_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +template +__global__ static void moe_down_q4K_expert_tile8_rowspan_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][8]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { + uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; + if (row >= out_dim) continue; + const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q4_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc); + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } +} + + +/* INT8 tensor-core (m8n8k16) exact MoE prefill tile kernels. + * + * Each warp computes an 8-token x 8-row tile. The Q4_K x Q8_K superblock dot + * keeps its integer sums (order-invariant, exact) but computes the 32-wide + * group dots on tensor cores; every output element keeps 8 float slot + * accumulators (slot[b & 7] += term_b, b ascending) and reduces them with the + * exact quarter_warp_sum_f32 grouping, so results are bit-identical to the + * scalar expert-tile kernels (fuzz-verified). Requires sm_75+, 16B-aligned + * expert tensors, and the staged activation-block counts (<=16 gate/up, + * <=8 down). Rollback: DS4_CUDA_MOE_NO_Q4_MMA=1. */ +__device__ __forceinline__ static void mma_m8n8k16_s8(int32_t &c0, int32_t &c1, uint32_t a, uint32_t b) { +#if __CUDA_ARCH__ >= 750 + asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%0,%1};" + : "+r"(c0), "+r"(c1) : "r"(a), "r"(b)); +#else + (void)a; (void)b; (void)c0; (void)c1; +#endif +} + +template +__global__ static void moe_gate_up_mid_q4K_tile8_mma_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][16]; + __shared__ uint32_t s_pair[8]; + __shared__ uint32_t s_tok[8]; + __shared__ uint32_t s_slot[8]; + __shared__ uint32_t s_np; + if (threadIdx.x == 0) { + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + uint32_t pr = sorted_pairs[offsets[expert] + local_pair]; + s_pair[np] = pr; + s_tok[np] = pr / n_expert; + s_slot[np] = pr - s_tok[np] * n_expert; + } + s_np = np; + } + __syncthreads(); + const uint32_t np = s_np; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < np * xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); i += blockDim.x) { + const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + uint32_t p = i / words_per_tok; + uint32_t w = i - p * words_per_tok; + ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w]; + } + __syncthreads(); + } + const uint32_t mtok = lane >> 2u; /* token row of this thread's C elems */ + const uint32_t n0 = (lane & 3u) * 2u; /* first C column (weight row) */ + /* 8 warps x 8 rows = 64 rows per pass */ + for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { + const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; + if (row0 >= expert_mid_dim) continue; + const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; + const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; + /* per-element slot accumulators (2 elements x 8 slots) */ + float sg0[8] = {0,0,0,0,0,0,0,0}, sg1[8] = {0,0,0,0,0,0,0,0}; + float su0[8] = {0,0,0,0,0,0,0,0}, su1[8] = {0,0,0,0,0,0,0,0}; + for (uint32_t b = 0; b < xq_blocks; b++) { + /* headers for this thread's two C columns */ + const uint4 ghdr0 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)n0 * gate_row_bytes) + b); + const uint4 ghdr1 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b); + const uint4 uhdr0 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)n0 * gate_row_bytes) + b); + const uint4 uhdr1 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b); + /* B-fragment source rows for loads: n_load = lane>>2. + * Batch all global loads for this superblock upfront so the + * memory system sees independent requests instead of a + * load->mma dependency chain. */ + const uint32_t *gqw = (const uint32_t *)(((const cuda_block_q4_K *)(grow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs); + const uint32_t *uqw = (const uint32_t *)(((const cuda_block_q4_K *)(urow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs); + const int8_t *aqs = sxq[mtok][b].qs; + uint32_t gw8[8], uw8[8]; +#pragma unroll + for (uint32_t k = 0; k < 8u; k++) { + gw8[k] = gqw[k * 4u + (lane & 3u)]; + uw8[k] = uqw[k * 4u + (lane & 3u)]; + } + int gi0 = 0, gi1 = 0, ui0 = 0, ui1 = 0; + int gs0 = 0, gs1 = 0, us0 = 0, us1 = 0; +#pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const int shift = (j & 1u) ? 4 : 0; + /* dot32 via two chained k16 mmas, per matrix */ + int32_t gc0 = 0, gc1 = 0, uc0 = 0, uc1 = 0; +#pragma unroll + for (uint32_t h = 0; h < 2u; h++) { + const uint32_t koff = h * 16u + (lane & 3u) * 4u; + const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff); + const uint32_t gw = (gw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; + const uint32_t uw = (uw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; + mma_m8n8k16_s8(gc0, gc1, a, gw); + mma_m8n8k16_s8(uc0, uc1, a, uw); + } + /* integer scale application for this thread's two columns */ + uint8_t sc, m; + dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr0.y, &sc, &m); + gi0 += (int)sc * gc0; + const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u]; + gs0 += (int)m * bs; + dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr1.y, &sc, &m); + gi1 += (int)sc * gc1; + gs1 += (int)m * bs; + dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr0.y, &sc, &m); + ui0 += (int)sc * uc0; + us0 += (int)m * bs; + dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr1.y, &sc, &m); + ui1 += (int)sc * uc1; + us1 += (int)m * bs; + } + /* float finish, exact dev_dot_q4_K_q8_K_block8 expression */ + const float yd = sxq[mtok][b].d; + const uint32_t sl = b & 7u; + sg0[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr0.x & 0xffffu)) * (float)gi0 - + yd * dev_f16_to_f32((uint16_t)(ghdr0.x >> 16u)) * (float)gs0; + sg1[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr1.x & 0xffffu)) * (float)gi1 - + yd * dev_f16_to_f32((uint16_t)(ghdr1.x >> 16u)) * (float)gs1; + su0[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr0.x & 0xffffu)) * (float)ui0 - + yd * dev_f16_to_f32((uint16_t)(uhdr0.x >> 16u)) * (float)us0; + su1[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr1.x & 0xffffu)) * (float)ui1 - + yd * dev_f16_to_f32((uint16_t)(uhdr1.x >> 16u)) * (float)us1; + } + /* quarter_warp_sum_f32 order: ((s0+s4)+(s2+s6)) + ((s1+s5)+(s3+s7)) */ + const uint32_t p = mtok; + if (p < np) { + const uint32_t rowa = row0 + n0; + const uint32_t rowb = row0 + n0 + 1u; + float gate2[2], up2[2]; + { + float a0 = sg0[0] + sg0[4], a1 = sg0[1] + sg0[5], a2 = sg0[2] + sg0[6], a3 = sg0[3] + sg0[7]; + gate2[0] = (a0 + a2) + (a1 + a3); + a0 = sg1[0] + sg1[4]; a1 = sg1[1] + sg1[5]; a2 = sg1[2] + sg1[6]; a3 = sg1[3] + sg1[7]; + gate2[1] = (a0 + a2) + (a1 + a3); + a0 = su0[0] + su0[4]; a1 = su0[1] + su0[5]; a2 = su0[2] + su0[6]; a3 = su0[3] + su0[7]; + up2[0] = (a0 + a2) + (a1 + a3); + a0 = su1[0] + su1[4]; a1 = su1[1] + su1[5]; a2 = su1[2] + su1[6]; a3 = su1[3] + su1[7]; + up2[1] = (a0 + a2) + (a1 + a3); + } +#pragma unroll + for (uint32_t e = 0; e < 2u; e++) { + const uint32_t row = e ? rowb : rowa; + if (row >= expert_mid_dim) continue; + float gate = gate2[e]; + float up = up2[e]; + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]]; + } + } + } +} + +template +__global__ static void moe_down_q4K_tile8_mma_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][8]; + __shared__ uint32_t s_pair[8]; + __shared__ uint32_t s_np; + if (threadIdx.x == 0) { + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + s_pair[np] = sorted_pairs[offsets[expert] + local_pair]; + } + s_np = np; + } + __syncthreads(); + const uint32_t np = s_np; + if (midq_blocks <= 8u) { + const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { + uint32_t p = i / words_per_tok; + uint32_t w = i - p * words_per_tok; + ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w]; + } + __syncthreads(); + } + const uint32_t mtok = lane >> 2u; + const uint32_t n0 = (lane & 3u) * 2u; + for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { + const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; + if (row0 >= out_dim) continue; + const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes; + float s0[8] = {0,0,0,0,0,0,0,0}, s1[8] = {0,0,0,0,0,0,0,0}; + for (uint32_t b = 0; b < midq_blocks; b++) { + const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b); + const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b); + const uint32_t *wqw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs); + const int8_t *aqs = sxq[mtok][b].qs; + uint32_t w8[8]; +#pragma unroll + for (uint32_t k = 0; k < 8u; k++) w8[k] = wqw[k * 4u + (lane & 3u)]; + int i0 = 0, i1 = 0, m0 = 0, m1 = 0; +#pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const int shift = (j & 1u) ? 4 : 0; + int32_t c0 = 0, c1 = 0; +#pragma unroll + for (uint32_t h = 0; h < 2u; h++) { + const uint32_t koff = h * 16u + (lane & 3u) * 4u; + const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff); + const uint32_t w = (w8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu; + mma_m8n8k16_s8(c0, c1, a, w); + } + uint8_t sc, m; + const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u]; + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc, &m); + i0 += (int)sc * c0; + m0 += (int)m * bs; + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc, &m); + i1 += (int)sc * c1; + m1 += (int)m * bs; + } + const float yd = sxq[mtok][b].d; + const uint32_t sl = b & 7u; + s0[sl] += yd * dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)) * (float)i0 - + yd * dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)) * (float)m0; + s1[sl] += yd * dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)) * (float)i1 - + yd * dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)) * (float)m1; + } + const uint32_t p = mtok; + if (p < np) { + float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; + const float r0 = (a0 + a2) + (a1 + a3); + a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; + const float r1 = (a0 + a2) + (a1 + a3); + if (row0 + n0 < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0] = r0; + if (row0 + n0 + 1u < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0 + 1u] = r1; + } + } +} + +/* 16-pair MoE expert tile kernels on sm_80+ m16n8k32 INT8 tensor cores. + * + * Same per-output math and reduction order as the 8-pair expert tile + * kernels (slot[b & 7] += term_b with b ascending, then the exact + * quarter_warp_sum_f32 grouping), so results are bit-identical; grouping + * 16 pairs per tile just halves how often each expert's weights are + * streamed from DRAM. Gate and up run as two passes over the superblocks + * to keep register pressure at the 8-pair kernel's level. */ + +__device__ __forceinline__ static void mma16_m16n8k32_s8( + int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3, + uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, + uint32_t b0, uint32_t b1) { +#if __CUDA_ARCH__ >= 800 + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3) + : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1)); +#else + (void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3; +#endif +} + +/* One matrix pass over all superblocks for this thread's 4 C elements + * (tokens mtokA/mtokB x rows n0/n0+1). Returns the quarter-tree-reduced + * values in r[4] with the exact reference ordering. */ +__device__ __forceinline__ static void moe_tile16_mma_pass( + const char *wrow, /* row0 base of this matrix */ + uint64_t row_bytes, + const cuda_block_q8_K (*sxq)[16], + uint32_t xq_blocks, + uint32_t lane, + float r[4]) { + const uint32_t mtokA = lane >> 2u; + const uint32_t mtokB = mtokA + 8u; + const uint32_t n0 = (lane & 3u) * 2u; + float s0[8] = {0,0,0,0,0,0,0,0}; + float s1[8] = {0,0,0,0,0,0,0,0}; + float s2[8] = {0,0,0,0,0,0,0,0}; + float s3[8] = {0,0,0,0,0,0,0,0}; + for (uint32_t b = 0; b < xq_blocks; b++) { + const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * row_bytes) + b); + const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * row_bytes) + b); + const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * row_bytes) + b)->qs); + uint32_t w8[8]; +#pragma unroll + for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)]; + const int8_t *aqsA = sxq[mtokA][b].qs; + const int8_t *aqsB = sxq[mtokB][b].qs; + int i0 = 0, i1 = 0, i2 = 0, i3 = 0; + int m0 = 0, m1 = 0, m2 = 0, m3 = 0; +#pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const int shift = (j & 1u) ? 4 : 0; + const uint32_t koff = (lane & 3u) * 4u; + const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff); + const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff); + const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff); + const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff); + const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu; + const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu; + int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; + mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); + uint8_t sc0, sm0, sc1, sm1; + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0); + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1); + const int bsA = (int)sxq[mtokA][b].bsums[2u * j] + (int)sxq[mtokA][b].bsums[2u * j + 1u]; + const int bsB = (int)sxq[mtokB][b].bsums[2u * j] + (int)sxq[mtokB][b].bsums[2u * j + 1u]; + i0 += (int)sc0 * c0; + i1 += (int)sc1 * c1; + i2 += (int)sc0 * c2; + i3 += (int)sc1 * c3; + m0 += (int)sm0 * bsA; + m1 += (int)sm1 * bsA; + m2 += (int)sm0 * bsB; + m3 += (int)sm1 * bsB; + } + const float ydA = sxq[mtokA][b].d; + const float ydB = sxq[mtokB][b].d; + const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)); + const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)); + const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)); + const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)); + const uint32_t sl = b & 7u; + s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0; + s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1; + s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2; + s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3; + } + { + float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; + r[0] = (a0 + a2) + (a1 + a3); + a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; + r[1] = (a0 + a2) + (a1 + a3); + a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7]; + r[2] = (a0 + a2) + (a1 + a3); + a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7]; + r[3] = (a0 + a2) + (a1 + a3); + } +} + +template +__global__ static void moe_gate_up_mid_q4K_tile16_mma_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t write_aux, + float clamp) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + extern __shared__ unsigned char t16_sh[]; + cuda_block_q8_K (*sxq)[16] = (cuda_block_q8_K (*)[16])t16_sh; /* [16][16] */ + __shared__ uint32_t s_pair[16]; + __shared__ uint32_t s_tok[16]; + __shared__ uint32_t s_slot[16]; + __shared__ uint32_t s_np; + if (threadIdx.x == 0) { + uint32_t np = 0; + for (; np < 16u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + uint32_t pr = sorted_pairs[offsets[expert] + local_pair]; + s_pair[np] = pr; + s_tok[np] = pr / n_expert; + s_slot[np] = pr - s_tok[np] * n_expert; + } + s_np = np; + } + __syncthreads(); + const uint32_t np = s_np; + if (xq_blocks <= 16u) { + const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { + uint32_t p = i / words_per_tok; + uint32_t w = i - p * words_per_tok; + ((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w]; + } + /* zero-fill missing pairs so the A fragments are defined */ + const uint32_t total_words = 16u * words_per_tok; + for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) { + ((uint32_t *)t16_sh)[i] = 0u; + } + __syncthreads(); + } + const uint32_t mtokA = lane >> 2u; + const uint32_t mtokB = mtokA + 8u; + const uint32_t n0 = (lane & 3u) * 2u; + for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { + const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; + if (row0 >= expert_mid_dim) continue; + const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; + const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes; + float gr[4], ur[4]; + moe_tile16_mma_pass(grow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, gr); + moe_tile16_mma_pass(urow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, ur); +#pragma unroll + for (uint32_t e = 0; e < 4u; e++) { + const uint32_t p = (e < 2u) ? mtokA : mtokB; + const uint32_t row = row0 + n0 + (e & 1u); + if (p >= np || row >= expert_mid_dim) continue; + float gate = gr[e]; + float up = ur[e]; + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]]; + } + } +} + +template +__global__ static void moe_down_q4K_tile16_mma_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t warp = threadIdx.x >> 5u; + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + extern __shared__ unsigned char t16_sh[]; + __shared__ uint32_t s_pair[16]; + __shared__ uint32_t s_np; + if (threadIdx.x == 0) { + uint32_t np = 0; + for (; np < 16u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + s_pair[np] = sorted_pairs[offsets[expert] + local_pair]; + } + s_np = np; + } + __syncthreads(); + const uint32_t np = s_np; + if (midq_blocks <= 16u) { + const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); + for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) { + uint32_t p = i / words_per_tok; + uint32_t w = i - p * words_per_tok; + ((uint32_t *)t16_sh)[i] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w]; + } + const uint32_t total_words = 16u * words_per_tok; + for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) { + ((uint32_t *)t16_sh)[i] = 0u; + } + __syncthreads(); + } + const uint32_t mtokA = lane >> 2u; + const uint32_t mtokB = mtokA + 8u; + const uint32_t n0 = (lane & 3u) * 2u; + for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) { + const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u; + if (row0 >= out_dim) continue; + const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes; + float s0[8] = {0,0,0,0,0,0,0,0}; + float s1[8] = {0,0,0,0,0,0,0,0}; + float s2[8] = {0,0,0,0,0,0,0,0}; + float s3[8] = {0,0,0,0,0,0,0,0}; + for (uint32_t b = 0; b < midq_blocks; b++) { + const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b); + const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b); + const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs); + uint32_t w8[8]; +#pragma unroll + for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)]; + /* activation rows: midq_blocks stride within the staged region */ + const int8_t *aqsA = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b)->qs; + const int8_t *aqsB = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b)->qs; + const cuda_block_q8_K *blkA = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b; + const cuda_block_q8_K *blkB = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b; + int i0 = 0, i1 = 0, i2 = 0, i3 = 0; + int m0 = 0, m1 = 0, m2 = 0, m3 = 0; +#pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const int shift = (j & 1u) ? 4 : 0; + const uint32_t koff = (lane & 3u) * 4u; + const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff); + const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff); + const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff); + const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff); + const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu; + const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu; + int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0; + mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); + uint8_t sc0, sm0, sc1, sm1; + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0); + dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1); + const int bsA = (int)blkA->bsums[2u * j] + (int)blkA->bsums[2u * j + 1u]; + const int bsB = (int)blkB->bsums[2u * j] + (int)blkB->bsums[2u * j + 1u]; + i0 += (int)sc0 * c0; + i1 += (int)sc1 * c1; + i2 += (int)sc0 * c2; + i3 += (int)sc1 * c3; + m0 += (int)sm0 * bsA; + m1 += (int)sm1 * bsA; + m2 += (int)sm0 * bsB; + m3 += (int)sm1 * bsB; + } + const float ydA = blkA->d; + const float ydB = blkB->d; + const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)); + const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)); + const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)); + const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)); + const uint32_t sl = b & 7u; + s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0; + s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1; + s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2; + s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3; + } + float rr4[4]; + { + float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7]; + rr4[0] = (a0 + a2) + (a1 + a3); + a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7]; + rr4[1] = (a0 + a2) + (a1 + a3); + a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7]; + rr4[2] = (a0 + a2) + (a1 + a3); + a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7]; + rr4[3] = (a0 + a2) + (a1 + a3); + } +#pragma unroll + for (uint32_t e = 0; e < 4u; e++) { + const uint32_t p = (e < 2u) ? mtokA : mtokB; + const uint32_t row = row0 + n0 + (e & 1u); + if (p >= np || row >= out_dim) continue; + down_out[(uint64_t)s_pair[p] * out_dim + row] = rr4[e]; + } + } +} + +static int cuda_q4_mma_tile16_shmem_ok(int which_down) { + /* Opt the tile16 kernels into >48KB dynamic shared memory, per device. */ + static int ready[DS4_MAX_GPUS][2]; + static int failed = 0; + if (failed) return 0; + int dev = 0; + cudaGetDevice(&dev); + if (dev < 0 || dev >= DS4_MAX_GPUS) return 0; + if (ready[dev][which_down]) return 1; + cudaFuncAttributes fn_attr; + cudaError_t err = which_down + ? cudaFuncGetAttributes(&fn_attr, moe_down_q4K_tile16_mma_kernel<512>) + : cudaFuncGetAttributes(&fn_attr, moe_gate_up_mid_q4K_tile16_mma_kernel<512>); + if (err != cudaSuccess || fn_attr.binaryVersion < 80) { + failed = 1; + return 0; + } + const int bytes = (int)(16u * 16u * sizeof(cuda_block_q8_K)); + if (which_down) { + err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<512>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (err == cudaSuccess) + err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<1024>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (err == cudaSuccess) + err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<2048>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + } else { + err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<512>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (err == cudaSuccess) + err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<1024>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (err == cudaSuccess) + err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<2048>, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + } + if (err != cudaSuccess) { + failed = 1; + return 0; + } + ready[dev][which_down] = 1; + return 1; +} + + + + +__global__ static void moe_down_sorted_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t pair = sorted_pairs[blockIdx.y]; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +__global__ static DS4_CUDA_UNUSED void moe_down_expert_tile8_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t group = threadIdx.x >> 3u; + uint32_t lane = threadIdx.x & 7u; + uint32_t pair_slot = group & 7u; + uint32_t row_lane = group >> 3u; + uint32_t expert = tile_experts[tile]; + uint32_t local_pair = tile_starts[tile] + pair_slot; + if (local_pair >= counts[expert]) return; + uint32_t sorted_idx = offsets[expert] + local_pair; + uint32_t pair = sorted_pairs[sorted_idx]; + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + + for (uint32_t rr = 0; rr < 2u; rr++) { + uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u; + if (row >= out_dim) continue; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; + } +} + +__global__ static void moe_down_expert_tile4_row32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t atomic_out) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[4][8]; + uint32_t pair[4] = {0, 0, 0, 0}; + const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 4u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= out_dim) return; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q2_K_q8_K_block4(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, acc); + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + if (atomic_out) { + uint32_t tok = pair[p] / n_expert; + atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); + } else { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } + } +} + +__global__ static void moe_down_expert_tile8_row32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t atomic_out) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + uint32_t local_start = tile_starts[tile]; + __shared__ cuda_block_q8_K sxq[8][8]; + uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + uint32_t np = 0; + for (; np < 8u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= out_dim) return; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc); + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + if (atomic_out) { + uint32_t tok = pair[p] / n_expert; + atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); + } else { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } + } +} + +__global__ static void moe_down_expert_tile16_row32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t atomic_out) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t local_start = tile_starts[tile]; + if (local_start & 8u) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + uint32_t expert = tile_experts[tile]; + __shared__ cuda_block_q8_K sxq[16][8]; + uint32_t pair[16] = {0}; + const cuda_block_q8_K *xqb[16] = {NULL}; + uint32_t np = 0; + for (; np < 16u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + if (row >= out_dim) return; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[16] = {0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); + if (np > 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, + xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, + xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, + xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); + } + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + if (atomic_out) { + uint32_t tok = pair[p] / n_expert; + atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); + } else { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } + } +} + +__global__ static void moe_down_expert_tile16_row2048_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t atomic_out) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t local_start = tile_starts[tile]; + if (local_start & 8u) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + __shared__ cuda_block_q8_K sxq[16][8]; + uint32_t pair[16] = {0}; + const cuda_block_q8_K *xqb[16] = {NULL}; + uint32_t np = 0; + for (; np < 16u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < 64u; rr++) { + uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u; + if (row >= out_dim) continue; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[16] = {0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); + if (np > 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, + xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, + xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, + xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); + } + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + if (atomic_out) { + uint32_t tok = pair[p] / n_expert; + atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); + } else { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } + } + } +} + +template +__global__ static void moe_down_expert_tile16_rowspan_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const uint32_t *offsets, + const uint32_t *counts, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t atomic_out) { + uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + uint32_t local_start = tile_starts[tile]; + if (local_start & 8u) return; + uint32_t lane = threadIdx.x & 7u; + uint32_t row_lane = threadIdx.x >> 3u; + uint32_t expert = tile_experts[tile]; + __shared__ cuda_block_q8_K sxq[16][8]; + uint32_t pair[16] = {0}; + const cuda_block_q8_K *xqb[16] = {NULL}; + uint32_t np = 0; + for (; np < 16u; np++) { + uint32_t local_pair = local_start + np; + if (local_pair >= counts[expert]) break; + pair[np] = sorted_pairs[offsets[expert] + local_pair]; + xqb[np] = midq + (uint64_t)pair[np] * midq_blocks; + } + if (midq_blocks <= 8u) { + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + uint32_t p = i / midq_blocks; + uint32_t b = i - p * midq_blocks; + sxq[p][b] = xqb[p][b]; + } + __syncthreads(); + for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; + } + for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { + uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u; + if (row >= out_dim) continue; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes); + float acc[16] = {0.0f}; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL, + xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, + xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL, + xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc); + if (np > 8u) { + dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL, + xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL, + xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL, + xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8); + } + } + for (uint32_t p = 0; p < np; p++) { + acc[p] = quarter_warp_sum_f32(acc[p], lane); + if (lane == 0) { + if (atomic_out) { + uint32_t tok = pair[p] / n_expert; + atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]); + } else { + down_out[(uint64_t)pair[p] * out_dim + row] = acc[p]; + } + } + } + } +} + +__global__ static void moe_down_sorted_p2_qwarp32_kernel( + float *down_out, + const char *down_base, + const cuda_block_q8_K *midq, + const uint32_t *sorted_pairs, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t pair_count) { + uint32_t lane = threadIdx.x & 7u; + uint32_t pair_lane = (threadIdx.x >> 3u) & 1u; + uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u); + uint32_t sorted_idx = blockIdx.y * 2u + pair_lane; + if (row >= out_dim || sorted_idx >= pair_count) return; + uint32_t pair = sorted_pairs[sorted_idx]; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b); + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; +} + +__global__ static void moe_sum_kernel(float *out, const float *down, uint32_t out_dim, uint32_t n_expert, uint32_t n_tokens) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * out_dim; + if (gid >= n) return; + uint32_t tok = gid / out_dim; + uint32_t row = gid - (uint64_t)tok * out_dim; + float acc = 0.0f; + for (uint32_t e = 0; e < n_expert; e++) acc += down[((uint64_t)tok * n_expert + e) * out_dim + row]; + out[gid] = acc; +} + +__global__ static void moe_sum_owned_kernel( + float *out, + const float *down, + const int32_t *selected, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t n = (uint64_t)n_tokens * out_dim; + if (gid >= n) return; + const uint32_t tok = (uint32_t)(gid / out_dim); + const uint32_t row = (uint32_t)(gid - (uint64_t)tok * out_dim); + float acc = 0.0f; + #pragma unroll + for (uint32_t slot = 0; slot < 6u; slot++) { + if (slot >= n_expert) break; + const uint64_t pair = (uint64_t)tok * n_expert + slot; + const float value = selected[pair] >= 0 + ? down[pair * out_dim + row] : 0.0f; + acc = __fadd_rn(acc, value); + } + out[gid] = acc; +} + +__device__ static float dev_iq2_xxs_dot_f32(const cuda_block_iq2_xxs *row, const float *x, uint32_t nb) { + float acc = 0.0f; + for (uint32_t b = 0; b < nb; b++) { + const cuda_block_iq2_xxs *xb = row + b; + const float d = dev_f16_to_f32(xb->d); + const uint16_t *q2 = xb->qs; + const float *xf = x + (uint64_t)b * CUDA_QK_K; + for (uint32_t ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux_g = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux_s = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const float dl = d * (0.5f + (float)(aux_s >> 28)) * 0.25f; + const uint8_t grids[4] = { + (uint8_t)(aux_g & 0xffu), + (uint8_t)((aux_g >> 8) & 0xffu), + (uint8_t)((aux_g >> 16) & 0xffu), + (uint8_t)((aux_g >> 24) & 0xffu), + }; + for (uint32_t half = 0; half < 2; half++) { + for (uint32_t g = 0; g < 2; g++) { + const uint32_t gi = half * 2 + g; + const uint64_t grid = cuda_iq2xxs_grid[grids[gi]]; + const uint8_t signs = cuda_ksigns_iq2xs[(aux_s >> (14u * half + 7u * g)) & 127u]; + for (uint32_t i = 0; i < 8; i++) { + float w = (float)((grid >> (8u * i)) & 0xffu); + if (signs & (1u << i)) w = -w; + acc += dl * w * xf[ib32 * 32u + half * 16u + g * 8u + i]; + } + } + } + } + } + return acc; +} + +__device__ static float dev_q2_K_dot_f32(const cuda_block_q2_K *row, const float *x, uint32_t nb) { + float acc = 0.0f; + for (uint32_t b = 0; b < nb; b++) { + const cuda_block_q2_K *xb = row + b; + const float d = dev_f16_to_f32(xb->d); + const float dmin = dev_f16_to_f32(xb->dmin); + for (uint32_t il = 0; il < 16; il++) { + const uint32_t chunk = il / 8u; + const uint32_t pair = il & 1u; + const uint32_t shift = ((il / 2u) & 3u) * 2u; + const uint8_t sc = xb->scales[il]; + const float dl = d * (float)(sc & 0x0fu); + const float ml = dmin * (float)(sc >> 4); + const uint8_t *q = xb->qs + 32u * chunk + 16u * pair; + const float *xf = x + (uint64_t)b * CUDA_QK_K + chunk * 128u + ((il % 8u) / 2u) * 32u + pair * 16u; + for (uint32_t i = 0; i < 16; i++) { + const float w = dl * (float)((q[i] >> shift) & 3u) - ml; + acc += w * xf[i]; + } + } + } + return acc; +} + +__global__ static void moe_gate_up_mid_f32_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const char *gate_base, + const char *up_base, + const float *x, + const int32_t *selected, + const float *weights, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t n_expert, + float clamp) { + uint32_t row = blockIdx.x; + uint32_t pair = blockIdx.y; + if (row >= expert_mid_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + uint32_t expert = (uint32_t)expert_i; + const uint32_t nb = expert_in_dim / CUDA_QK_K; + const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes); + const float *xr = x + (uint64_t)tok * expert_in_dim; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) { + gate += dev_iq2_xxs_dot_f32(gr + b, xr + (uint64_t)b * CUDA_QK_K, 1); + up += dev_iq2_xxs_dot_f32(ur + b, xr + (uint64_t)b * CUDA_QK_K, 1); + } + __shared__ float partial_gate[256]; + __shared__ float partial_up[256]; + partial_gate[threadIdx.x] = gate; + partial_up[threadIdx.x] = up; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride]; + partial_up[threadIdx.x] += partial_up[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + gate = partial_gate[0]; + up = partial_up[0]; + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + gate_out[off] = gate; + up_out[off] = up; + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot]; + } +} + +__global__ static void moe_down_f32_kernel( + float *down_out, + const char *down_base, + const float *mid, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_mid_dim, + uint32_t out_dim, + uint32_t n_expert) { + uint32_t row = blockIdx.x; + uint32_t pair = blockIdx.y; + if (row >= out_dim) return; + uint32_t tok = pair / n_expert; + uint32_t slot = pair - tok * n_expert; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const uint32_t nb = expert_mid_dim / CUDA_QK_K; + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes); + const float *xr = mid + (uint64_t)pair * expert_mid_dim; + float acc = 0.0f; + for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) acc += dev_q2_K_dot_f32(wr + b, xr + (uint64_t)b * CUDA_QK_K, 1); + __shared__ float partial[256]; + partial[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0]; +} + +static int routed_moe_launch( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *down, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + int allow_streaming, + int owned_filtered) { + if (!out || !gate || !up || !mid || !down || !model_map || !selected || !weights || !x || + n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || + expert_in_dim % CUDA_QK_K != 0 || expert_mid_dim % CUDA_QK_K != 0 || + gate_offset > model_size || up_offset > model_size || down_offset > model_size || + x->bytes < (uint64_t)n_tokens * expert_in_dim * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * n_expert * sizeof(int32_t) || + weights->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) || + gate->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || + up->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || + mid->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) || + down->bytes < (uint64_t)n_tokens * n_expert * out_dim * sizeof(float) || + out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) { + return 0; + } + const int q4k_path = (gate_type == 12u && down_type == 12u); + if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; + /* Q4_K routed-MoE dispatch: + * n_tokens == 1 and n_expert == 6: + * use_direct_down_sum + moe_gate_up_mid_decode_q4K_qwarp32 + * + moe_down_q4K_sum6_qwarp32. + * n_tokens == 1 and n_expert == 3: + * use the same direct path with moe_down_q4K_sum3_qwarp32. + * Decode TP relies on this for splitting the six selected + * experts into two groups. + * n_tokens == 1 and other n_expert: + * use the same per-pair gate/up kernel plus the generic + * q4K down + sum path. + * n_tokens > 1: default sorted-pairs expert-tile path groups token/expert + * pairs by expert and uses Q4_K tile8 gate/up + down kernels + * (`DS4_CUDA_MOE_NO_Q4_SORTED=1` restores the older + * token-indexed decode-style prefill kernels). */ + const uint64_t gate_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_bytes = (uint64_t)n_total_expert * down_expert_bytes; + if (gate_bytes > model_size - gate_offset || + gate_bytes > model_size - up_offset || + down_bytes > model_size - down_offset) { + return 0; + } + const uint64_t required_slot_count = (uint64_t)n_tokens * n_expert; + const int logical_tier = ds4_tensor_device_idx(out); + const int use_stream_selected_cache = + allow_streaming && + g_ssd_streaming_mode && + g_stream_selected_cache.valid && + g_stream_selected_cache.logical_tier == logical_tier && + g_stream_selected_cache.model_map == model_map && + g_stream_selected_cache.layer == layer_index && + g_stream_selected_cache.n_total_expert == n_total_expert && + g_stream_selected_cache.slot_count >= required_slot_count && + g_stream_selected_cache.gate_offset == gate_offset && + g_stream_selected_cache.up_offset == up_offset && + g_stream_selected_cache.down_offset == down_offset && + g_stream_selected_cache.gate_expert_bytes == gate_expert_bytes && + g_stream_selected_cache.down_expert_bytes == down_expert_bytes && + g_stream_selected_cache.gate_ptr && + g_stream_selected_cache.up_ptr && + g_stream_selected_cache.down_ptr && + g_stream_selected_cache.slot_selected_tensor.ptr && + g_stream_selected_cache.slot_selected_tensor.bytes >= + required_slot_count * sizeof(int32_t); + if (g_ssd_streaming_mode && allow_streaming && + !use_stream_selected_cache) { + fprintf(stderr, + "ds4: CUDA streaming selected experts are unavailable for layer %u\n", + layer_index); + return 0; + } + if (use_stream_selected_cache) { + selected = &g_stream_selected_cache.slot_selected_tensor; + } + const char *gate_w = use_stream_selected_cache ? + g_stream_selected_cache.gate_ptr : + cuda_resolve_weight_ptr(model_map, gate_offset, gate_bytes, + logical_tier, "moe_gate"); + const char *up_w = use_stream_selected_cache ? + g_stream_selected_cache.up_ptr : + cuda_resolve_weight_ptr(model_map, up_offset, gate_bytes, + logical_tier, "moe_up"); + const char *down_w = use_stream_selected_cache ? + g_stream_selected_cache.down_ptr : + cuda_resolve_weight_ptr(model_map, down_offset, down_bytes, + logical_tier, "moe_down"); + if (!gate_w || !up_w || !down_w) return 0; + + int ok = 1; + const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; + const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; + const uint64_t xq_count = (uint64_t)n_tokens * xq_blocks; + const uint64_t midq_count = (uint64_t)n_tokens * n_expert * midq_blocks; + const uint64_t xq_bytes = xq_count * sizeof(cuda_block_q8_K); + const uint64_t midq_bytes = midq_count * sizeof(cuda_block_q8_K); + if (down->bytes >= xq_bytes && gate->bytes >= midq_bytes) { + cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; + cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; + const uint32_t profile_moe = getenv("DS4_CUDA_MOE_PROFILE") != NULL; + cudaEvent_t prof_ev[7] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL}; + if (profile_moe) { + for (uint32_t i = 0; i < 7u; i++) { + if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) { + for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]); + memset(prof_ev, 0, sizeof(prof_ev)); + break; + } + } + if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0); + } + const uint32_t pair_count = n_tokens * n_expert; + const uint32_t use_q4_sorted_pairs = + q4k_path && n_tokens > 1u && + (owned_filtered || + (getenv("DS4_CUDA_MOE_NO_Q4_SORTED") == NULL && + getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL && + getenv("DS4_CUDA_MOE_TILE4") == NULL)); + const uint32_t use_sorted_pairs = + n_tokens > 1u && + (owned_filtered || !q4k_path || use_q4_sorted_pairs); + const uint32_t use_expert_tiles = + use_sorted_pairs && + (owned_filtered || getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL); + /* Small batches (DSpark stage chain / verify, n<=8) leave most of an + * 8-slot expert tile empty (1-2 rows per expert): tile4 halves the + * wasted dot-slots and measures ~2x faster there. Large prefill + * keeps tile8. Env overrides both ways. */ + const uint32_t expert_tile_m = + getenv("DS4_CUDA_MOE_TILE4") ? 4u : + (getenv("DS4_CUDA_MOE_TILE8") ? 8u : + (n_tokens <= 8u ? 4u : 8u)); + const uint32_t write_gate_up = getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL; + const uint32_t use_p2_sorted = + use_sorted_pairs && !owned_filtered && + getenv("DS4_CUDA_MOE_NO_P2") == NULL; + const uint32_t use_atomic_down = !q4k_path && use_expert_tiles && + (getenv("DS4_CUDA_MOE_ATOMIC_DOWN") != NULL || + (n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_ATOMIC_DOWN") == NULL)); + const uint32_t use_owned_sparse_buffers = owned_filtered && + getenv("DS4_CUDA_MOE_NO_OWNED_SPARSE_BUFFERS") == NULL; + const uint32_t use_gate_row2048 = use_expert_tiles && expert_tile_m == 8u && + (getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL || + getenv("DS4_CUDA_MOE_GATE_ROW256") != NULL || + getenv("DS4_CUDA_MOE_GATE_ROW128") != NULL || + (n_tokens >= 128u && + getenv("DS4_CUDA_MOE_NO_GATE_ROW2048") == NULL && + getenv("DS4_CUDA_MOE_NO_GATE_ROW256") == NULL && + getenv("DS4_CUDA_MOE_NO_GATE_ROW128") == NULL)); + const uint32_t use_q4_mma_tiles16 = q4k_path && use_expert_tiles && + expert_tile_m == 8u && cuda_q4_mma_ok() && + getenv("DS4_CUDA_MOE_NO_Q4_MMA_TILE16") == NULL; + const uint32_t use_down_tile16 = !q4k_path && use_atomic_down && expert_tile_m == 8u && + n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_DOWN_TILE16") == NULL; + const uint32_t use_small_sorted_prep = + owned_filtered && q4k_path && n_tokens <= 16u && pair_count <= 96u && + n_total_expert <= 128u && use_sorted_pairs && use_expert_tiles && + getenv("DS4_CUDA_MOE_NO_SMALL_SORTED_PREP") == NULL; + const uint32_t use_q4_down_rowspan = q4k_path && use_expert_tiles && expert_tile_m == 8u && + n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_Q4_DOWN_ROWSPAN") == NULL; + const uint32_t use_decode_lut_gate = + n_tokens == 1u && xq_blocks <= 16u && + getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") == NULL; + const uint32_t gate_row_span = + getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL ? 2048u : + getenv("DS4_CUDA_MOE_GATE_ROW1024") != NULL ? 1024u : 512u; + const uint32_t down_row_span = + getenv("DS4_CUDA_MOE_DOWN_ROW512") != NULL ? 512u : + getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL ? 2048u : + getenv("DS4_CUDA_MOE_DOWN_ROW1024") != NULL ? 1024u : 512u; + const uint32_t use_down_row2048 = !q4k_path && use_atomic_down && expert_tile_m == 8u && + (getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL || + getenv("DS4_CUDA_MOE_DOWN_ROW256") != NULL || + getenv("DS4_CUDA_MOE_DOWN_ROW128") != NULL || + getenv("DS4_CUDA_MOE_DOWN_ROW64") != NULL || + (use_down_tile16 && + getenv("DS4_CUDA_MOE_NO_DOWN_ROW2048") == NULL && + getenv("DS4_CUDA_MOE_NO_DOWN_ROW256") == NULL && + getenv("DS4_CUDA_MOE_NO_DOWN_ROW128") == NULL && + getenv("DS4_CUDA_MOE_NO_DOWN_ROW64") == NULL)); + const uint32_t use_direct_down_sum = + n_tokens == 1u && (n_expert == 6u || n_expert == 3u) && + getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6") == NULL; + const uint32_t use_direct_midq = + q4k_path && use_direct_down_sum && !write_gate_up && + getenv("DS4_CUDA_MOE_DIRECT_MIDQ") != NULL && + getenv("DS4_CUDA_MOE_NO_DIRECT_MIDQ") == NULL; + const uint32_t use_q4_gate_h16r8 = + q4k_path && !use_direct_midq && + getenv("DS4_CUDA_MOE_Q4_GATE_H16R8") != NULL && + getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16R8") == NULL; + const uint32_t use_q4_gate_h16 = + q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && + getenv("DS4_CUDA_MOE_Q4_GATE_H16") != NULL && + getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16") == NULL; + const uint32_t use_q4_gate_w32r16 = + q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 && + getenv("DS4_CUDA_MOE_Q4_GATE_W32R16") != NULL && + getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32R16") == NULL; + const uint32_t use_q4_gate_w32 = + q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 && + !use_q4_gate_w32r16 && + getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32") == NULL; + const uint32_t use_q4_gate_w32_noaux = + use_q4_gate_w32 && !write_gate_up && + getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32_NOAUX") == NULL; + const uint32_t use_q4_down_slot3 = + q4k_path && use_direct_down_sum && n_expert == 3u && + getenv("DS4_CUDA_MOE_Q4_DOWN_SLOT3") != NULL && + getenv("DS4_CUDA_MOE_NO_Q4_DOWN_SLOT3") == NULL; + const uint32_t use_q4_midq_sidecar = + q4k_path && use_direct_down_sum && use_q4_gate_w32_noaux && + !use_direct_midq && !write_gate_up && + (expert_mid_dim % CUDA_QK_K) == 0u && + getenv("DS4_CUDA_MOE_MIDQ_SIDECAR") != NULL && + getenv("DS4_CUDA_MOE_NO_MIDQ_SIDECAR") == NULL; + float *midq_sidecar = use_q4_midq_sidecar ? (float *)up->ptr : NULL; + if (g_cuda_moe_decode_graph && + !owned_filtered && + !profile_moe && + q4k_path && + n_tokens == 1u && + use_direct_down_sum && + use_q4_gate_w32 && + !use_q4_gate_w32r16 && + !use_q4_down_slot3 && + !use_direct_midq && + !use_q4_midq_sidecar && + (n_expert == 3u || n_expert == 6u)) { + int grc = routed_moe_decode_q4_graph_launch( + logical_tier, + (float *)out->ptr, + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + down_w, + xq, + midq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + n_expert, + write_gate_up, + clamp, + (const float *)x->ptr); + if (grc == 1) return 1; + if (grc < 0) return 0; + } + uint32_t *sorted_pairs = NULL; + uint32_t *sorted_offsets = NULL; + uint32_t *sorted_counts = NULL; + uint32_t *tile_total = NULL; + uint32_t *tile_experts = NULL; + uint32_t *tile_starts = NULL; + uint32_t *tile16_total = NULL; + uint32_t *tile16_experts = NULL; + uint32_t *tile16_starts = NULL; + uint32_t tile_capacity = 0; + uint32_t tile16_capacity = 0; + dim3 xq_grid(xq_blocks, n_tokens, 1); + q8_K_quantize_kernel<<>>(xq, (const float *)x->ptr, expert_in_dim, n_tokens); + ok = cuda_ok(cudaGetLastError(), "routed_moe x quantize launch"); + if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0); + if (ok && use_sorted_pairs) { + const uint64_t counts_bytes = (uint64_t)n_total_expert * sizeof(uint32_t); + const uint64_t offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t); + const uint64_t cursors_bytes = (uint64_t)n_total_expert * sizeof(uint32_t); + const uint64_t sorted_bytes = (uint64_t)pair_count * sizeof(uint32_t); + tile_capacity = (pair_count + expert_tile_m - 1u) / expert_tile_m + n_total_expert; + tile16_capacity = (use_down_tile16 || use_q4_mma_tiles16) ? ((pair_count + 15u) / 16u + n_total_expert) : 0u; + const uint64_t tile_offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t); + const uint64_t tile_total_bytes = sizeof(uint32_t); + const uint64_t tile_experts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t); + const uint64_t tile_starts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t); + const uint64_t tile16_offsets_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? (((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t)) : 0u; + const uint64_t tile16_total_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? sizeof(uint32_t) : 0u; + const uint64_t tile16_experts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t); + const uint64_t tile16_starts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t); + const uint64_t tile_offsets_off = counts_bytes + offsets_bytes + cursors_bytes + sorted_bytes; + const uint64_t tile_total_off = tile_offsets_off + tile_offsets_bytes; + const uint64_t tile_experts_off = tile_total_off + tile_total_bytes; + const uint64_t tile_starts_off = tile_experts_off + tile_experts_bytes; + const uint64_t tile16_offsets_off = tile_starts_off + tile_starts_bytes; + const uint64_t tile16_total_off = tile16_offsets_off + tile16_offsets_bytes; + const uint64_t tile16_experts_off = tile16_total_off + tile16_total_bytes; + const uint64_t tile16_starts_off = tile16_experts_off + tile16_experts_bytes; + const uint64_t scratch_bytes = tile16_starts_off + tile16_starts_bytes; + uint8_t *scratch = (uint8_t *)cuda_tmp_alloc_on(logical_tier, scratch_bytes, + "routed_moe sorted pairs"); + if (!scratch) { + ok = 0; + } else { + uint32_t *counts = (uint32_t *)scratch; + uint32_t *offsets = (uint32_t *)(scratch + counts_bytes); + uint32_t *cursors = (uint32_t *)(scratch + counts_bytes + offsets_bytes); + sorted_pairs = (uint32_t *)(scratch + counts_bytes + offsets_bytes + cursors_bytes); + sorted_offsets = offsets; + sorted_counts = counts; + uint32_t *tile_offsets = (uint32_t *)(scratch + tile_offsets_off); + tile_total = (uint32_t *)(scratch + tile_total_off); + tile_experts = (uint32_t *)(scratch + tile_experts_off); + tile_starts = (uint32_t *)(scratch + tile_starts_off); + uint32_t *tile16_offsets = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_offsets_off) : NULL; + tile16_total = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_total_off) : NULL; + tile16_experts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_experts_off) : NULL; + tile16_starts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_starts_off) : NULL; + if (use_small_sorted_prep) { + moe_prepare_sorted_tiles_small_kernel<<<1, 128>>>( + counts, offsets, cursors, sorted_pairs, + tile_offsets, tile_total, tile_experts, tile_starts, + tile16_offsets, tile16_total, tile16_experts, tile16_starts, + (const int32_t *)selected->ptr, pair_count, n_total_expert, + expert_tile_m, use_down_tile16 || use_q4_mma_tiles16); + ok = cuda_ok(cudaGetLastError(), + "routed_moe small sorted setup launch"); + } else { + ok = cuda_ok(cudaMemset(counts, 0, counts_bytes), + "routed_moe sorted counts clear"); + } + if (ok && !use_small_sorted_prep) { + moe_count_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( + counts, + (const int32_t *)selected->ptr, + pair_count, + n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe sorted count launch"); + } + if (ok && !use_small_sorted_prep) { + moe_prefix_sorted_pairs_kernel<<<1, 1>>>(offsets, cursors, counts, n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe sorted prefix launch"); + } + if (ok && !use_small_sorted_prep) { + moe_scatter_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( + sorted_pairs, + cursors, + (const int32_t *)selected->ptr, + pair_count, + n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe sorted scatter launch"); + } + if (ok && use_expert_tiles && !use_small_sorted_prep) { + moe_build_expert_tile_offsets_kernel<<<1, 1>>>(tile_offsets, tile_total, counts, expert_tile_m, n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile offsets launch"); + } + if (ok && use_expert_tiles && !use_small_sorted_prep) { + moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256>>>( + tile_experts, tile_starts, tile_offsets, counts, expert_tile_m, n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tiles launch"); + } + if (ok && use_expert_tiles && !use_small_sorted_prep && + (use_down_tile16 || use_q4_mma_tiles16)) { + moe_build_expert_tile_offsets_kernel<<<1, 1>>>(tile16_offsets, tile16_total, counts, 16u, n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 offsets launch"); + } + if (ok && use_expert_tiles && !use_small_sorted_prep && + (use_down_tile16 || use_q4_mma_tiles16)) { + moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256>>>( + tile16_experts, tile16_starts, tile16_offsets, counts, 16u, n_total_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 launch"); + } + } + } + if (prof_ev[2]) (void)cudaEventRecord(prof_ev[2], 0); + if (ok && owned_filtered && use_sorted_pairs && + !use_owned_sparse_buffers) { + const uint64_t mid_bytes = + (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float); + ok = cuda_ok(cudaMemset(mid->ptr, 0, (size_t)mid_bytes), + "owned routed_moe mid clear"); + } + if (ok) { + dim3 mgrid((expert_mid_dim + 31u) / 32u, n_tokens * n_expert, 1); + if (ok && sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && tile_total && tile_experts && tile_starts) { + if (q4k_path) { + const int use_q4_mma = cuda_q4_mma_ok() && + ((((uintptr_t)gate_w | (uintptr_t)up_w | + gate_row_bytes | gate_expert_bytes) & 15u) == 0u) && + xq_blocks <= 16u && (expert_mid_dim & 7u) == 0u; + const int use_q4_mma_t16 = use_q4_mma && use_q4_mma_tiles16 && + tile16_total && tile16_experts && tile16_starts && + xq_blocks == 16u && cuda_q4_mma_tile16_shmem_ok(0); + if (use_q4_mma_t16 && use_gate_row2048) { + const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert); + const size_t t16sh = 16u * 16u * sizeof(cuda_block_q8_K); + if (gate_row_span == 512u) { + dim3 tgrid((expert_mid_dim + 511u) / 512u, t16cap, 1); + moe_gate_up_mid_q4K_tile16_mma_kernel<512><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else if (gate_row_span == 1024u) { + dim3 tgrid((expert_mid_dim + 1023u) / 1024u, t16cap, 1); + moe_gate_up_mid_q4K_tile16_mma_kernel<1024><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else { + dim3 tgrid((expert_mid_dim + 2047u) / 2048u, t16cap, 1); + moe_gate_up_mid_q4K_tile16_mma_kernel<2048><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else if (use_q4_mma && use_gate_row2048) { + if (gate_row_span == 512u) { + dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); + moe_gate_up_mid_q4K_tile8_mma_kernel<512><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else if (gate_row_span == 1024u) { + dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); + moe_gate_up_mid_q4K_tile8_mma_kernel<1024><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else { + dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); + moe_gate_up_mid_q4K_tile8_mma_kernel<2048><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else if (use_gate_row2048) { + if (gate_row_span == 512u) { + dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); + moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<512><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else if (gate_row_span == 1024u) { + dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); + moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<1024><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else { + dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); + moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<2048><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else { + dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); + moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<32><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else if (use_gate_row2048) { + if (gate_row_span == 512u) { + dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1); + moe_gate_up_mid_expert_tile8_rowspan_kernel<512><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else if (gate_row_span == 1024u) { + dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1); + moe_gate_up_mid_expert_tile8_rowspan_kernel<1024><<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else { + dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1); + moe_gate_up_mid_expert_tile8_row2048_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else if (expert_tile_m == 8u) { + dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); + moe_gate_up_mid_expert_tile8_row32_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } else { + dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1); + moe_gate_up_mid_expert_tile4_row32_kernel<<>>( + (float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, + gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts, + tile_total, tile_experts, tile_starts, (const float *)weights->ptr, + gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert, + write_gate_up, clamp); + } + } else if (ok && sorted_pairs && use_p2_sorted) { + dim3 p2_mgrid((expert_mid_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1); + moe_gate_up_mid_sorted_p2_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + sorted_pairs, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + pair_count, + clamp); + } else if (ok && sorted_pairs) { + moe_gate_up_mid_sorted_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + sorted_pairs, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + clamp); + } else if (ok) { + dim3 qgrid((expert_mid_dim + MOE_DECODE_ROWS_PER_BLOCK - 1u) / MOE_DECODE_ROWS_PER_BLOCK, n_tokens * n_expert, 1); + if (q4k_path) { + /* Q4_K gate/up: the decode kernel is token-indexed via + * pair = blockIdx.y; tok = pair / n_expert, so the same + * launch covers both n_tokens == 1 (decode) and n_tokens > 1 + * (prefill). q4k_path is steered here by use_sorted_pairs = 0 + * cascading the IQ2 sorted/expert-tile branches off. */ + if (use_direct_midq) { + dim3 mqgrid(midq_blocks, n_tokens * n_expert, 1); + moe_gate_up_midq_decode_q4K_qwarp32_kernel<<>>( + (float *)mid->ptr, + midq, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + clamp); + } else if (use_q4_gate_h16r8) { + dim3 h8grid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1); + moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } else if (use_q4_gate_w32r16) { + dim3 w16grid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1); + moe_gate_up_mid_decode_q4K_warp32_row16_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } else if (use_q4_gate_w32) { + dim3 wgrid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1); + if (use_q4_gate_w32_noaux) { + if (use_q4_midq_sidecar) { + moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel<<>>( + (float *)mid->ptr, + midq_sidecar, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + clamp); + } else { + moe_gate_up_mid_decode_q4K_warp32_noaux_kernel<<>>( + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + clamp); + } + } else { + moe_gate_up_mid_decode_q4K_warp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } + } else if (use_q4_gate_h16) { + dim3 hgrid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1); + moe_gate_up_mid_decode_q4K_hwarp16_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } else { + moe_gate_up_mid_decode_q4K_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } + } else if (use_decode_lut_gate) { + moe_gate_up_mid_decode_lut_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + write_gate_up, + clamp); + } else { + moe_gate_up_mid_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + n_expert, + clamp); + } + } + ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch"); + } + if (prof_ev[3]) (void)cudaEventRecord(prof_ev[3], 0); + if (ok && !use_direct_midq) { + dim3 midq_grid(midq_blocks, n_tokens * n_expert, 1); + if (use_q4_midq_sidecar) { + q8_K_quantize_sidecar_kernel<<>>( + midq, + (const float *)mid->ptr, + midq_sidecar, + expert_mid_dim, + n_tokens * n_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe mid sidecar quantize launch"); + } else if (use_owned_sparse_buffers) { + q8_K_quantize_owned_kernel<<>>( + midq, + (const float *)mid->ptr, + (const int32_t *)selected->ptr, + expert_mid_dim, + n_tokens * n_expert, + 0u, + n_total_expert); + ok = cuda_ok(cudaGetLastError(), + "owned routed_moe active mid quantize launch"); + } else { + q8_K_quantize_kernel<<>>(midq, (const float *)mid->ptr, expert_mid_dim, n_tokens * n_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe mid quantize launch"); + } + } + if (prof_ev[4]) (void)cudaEventRecord(prof_ev[4], 0); + if (ok && owned_filtered && use_sorted_pairs && !use_atomic_down && + !use_owned_sparse_buffers) { + const uint64_t down_clear_bytes = + (uint64_t)n_tokens * n_expert * out_dim * sizeof(float); + ok = cuda_ok(cudaMemset(down->ptr, 0, (size_t)down_clear_bytes), + "owned routed_moe down clear"); + } + if (ok) { + dim3 dgrid((out_dim + 31u) / 32u, n_tokens * n_expert, 1); + uint32_t *down_tile_total = tile_total; + uint32_t *down_tile_experts = tile_experts; + uint32_t *down_tile_starts = tile_starts; + uint32_t down_tile_capacity = tile_capacity; + if (use_down_tile16 && tile16_total && tile16_experts && tile16_starts) { + down_tile_total = tile16_total; + down_tile_experts = tile16_experts; + down_tile_starts = tile16_starts; + down_tile_capacity = tile16_capacity; + } + if (use_direct_down_sum) { + dim3 sgrid((out_dim + 31u) / 32u, 1, 1); + if (q4k_path) { + if (n_expert == 6u) { + moe_down_q4K_sum6_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim); + } else { + if (use_q4_down_slot3) { + dim3 swgrid((out_dim + 7u) / 8u, 1, 1); + moe_down_q4K_sum3_slotwarp_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim); + } else { + moe_down_q4K_sum3_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim); + } + } + } else { + if (n_expert == 6u) { + moe_down_sum6_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim); + } else { + moe_down_sum3_qwarp32_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim); + } + } + } else if (use_atomic_down) { + uint64_t n = (uint64_t)n_tokens * out_dim; + zero_kernel<<<(n + 255u) / 256u, 256>>>((float *)out->ptr, n); + ok = cuda_ok(cudaGetLastError(), "routed_moe atomic zero launch"); + } + if (use_direct_down_sum) { + /* The direct decode kernel writes the final token row. */ + } else if (sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && + down_tile_total && down_tile_experts && down_tile_starts) { + if (q4k_path) { + const int use_q4_down_mma = cuda_q4_mma_ok() && + ((((uintptr_t)down_w | down_row_bytes | down_expert_bytes) & 15u) == 0u) && + midq_blocks <= 8u && (out_dim & 7u) == 0u; + const int use_q4_down_t16 = use_q4_down_mma && use_q4_mma_tiles16 && + tile16_total && tile16_experts && tile16_starts && + midq_blocks <= 16u && cuda_q4_mma_tile16_shmem_ok(1); + if (use_q4_down_t16 && use_q4_down_rowspan) { + const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert); + const size_t dt16sh = 16u * (size_t)midq_blocks * sizeof(cuda_block_q8_K); + if (down_row_span == 512u) { + dim3 tgrid((out_dim + 511u) / 512u, t16cap, 1); + moe_down_q4K_tile16_mma_kernel<512><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else if (down_row_span == 1024u) { + dim3 tgrid((out_dim + 1023u) / 1024u, t16cap, 1); + moe_down_q4K_tile16_mma_kernel<1024><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else { + dim3 tgrid((out_dim + 2047u) / 2048u, t16cap, 1); + moe_down_q4K_tile16_mma_kernel<2048><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } + } else if (use_q4_down_mma && use_q4_down_rowspan) { + if (down_row_span == 512u) { + dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); + moe_down_q4K_tile8_mma_kernel<512><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else if (down_row_span == 1024u) { + dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); + moe_down_q4K_tile8_mma_kernel<1024><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else { + dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); + moe_down_q4K_tile8_mma_kernel<2048><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } + } else if (use_q4_down_rowspan) { + if (down_row_span == 512u) { + dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); + moe_down_q4K_expert_tile8_rowspan_kernel<512><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else if (down_row_span == 1024u) { + dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); + moe_down_q4K_expert_tile8_rowspan_kernel<1024><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } else { + dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); + moe_down_q4K_expert_tile8_rowspan_kernel<2048><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } + } else { + dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); + moe_down_q4K_expert_tile8_rowspan_kernel<32><<>>( + (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert); + } + } else if (use_down_row2048) { + if (down_row_span == 512u) { + dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1); + moe_down_expert_tile16_rowspan_kernel<512><<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } else if (down_row_span == 1024u) { + dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1); + moe_down_expert_tile16_rowspan_kernel<1024><<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } else { + dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1); + moe_down_expert_tile16_row2048_kernel<<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } + } else if (use_down_tile16) { + dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); + moe_down_expert_tile16_row32_kernel<<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } else if (expert_tile_m == 8u) { + dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); + moe_down_expert_tile8_row32_kernel<<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } else { + dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1); + moe_down_expert_tile4_row32_kernel<<>>( + use_atomic_down ? (float *)out->ptr : (float *)down->ptr, + down_w, midq, sorted_pairs, sorted_offsets, sorted_counts, + down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, use_atomic_down); + } + } else if (sorted_pairs && use_p2_sorted) { + dim3 p2_dgrid((out_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1); + moe_down_sorted_p2_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + sorted_pairs, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert, + pair_count); + } else if (sorted_pairs) { + moe_down_sorted_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + sorted_pairs, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } else if (q4k_path) { + /* Q4_K prefill down. New kernel mirrors moe_down_qwarp32_kernel + * grid/geometry, swapping the weight block type to cuda_block_q4_K + * and the dot helper to dev_dot_q4_K_q8_K_block. Writes per-pair + * outputs into down->ptr; moe_sum_kernel below sums them across + * experts into out->ptr. */ + moe_down_q4K_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } else { + moe_down_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } + ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); + } + if (prof_ev[5]) (void)cudaEventRecord(prof_ev[5], 0); + if (ok && !use_atomic_down && !use_direct_down_sum) { + uint64_t n = (uint64_t)n_tokens * out_dim; + if (use_owned_sparse_buffers) { + moe_sum_owned_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, + (const float *)down->ptr, + (const int32_t *)selected->ptr, + out_dim, + n_expert, + n_tokens); + } else { + moe_sum_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, + (const float *)down->ptr, + out_dim, + n_expert, + n_tokens); + } + ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch"); + } + if (prof_ev[6]) { + (void)cudaEventRecord(prof_ev[6], 0); + if (cudaEventSynchronize(prof_ev[6]) == cudaSuccess) { + float ms_xq = 0.0f, ms_sort = 0.0f, ms_gate = 0.0f, ms_midq = 0.0f, ms_down = 0.0f, ms_sum = 0.0f, ms_total = 0.0f; + (void)cudaEventElapsedTime(&ms_xq, prof_ev[0], prof_ev[1]); + (void)cudaEventElapsedTime(&ms_sort, prof_ev[1], prof_ev[2]); + (void)cudaEventElapsedTime(&ms_gate, prof_ev[2], prof_ev[3]); + (void)cudaEventElapsedTime(&ms_midq, prof_ev[3], prof_ev[4]); + (void)cudaEventElapsedTime(&ms_down, prof_ev[4], prof_ev[5]); + (void)cudaEventElapsedTime(&ms_sum, prof_ev[5], prof_ev[6]); + (void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[6]); + fprintf(stderr, + "ds4: CUDA MoE profile tokens=%u pairs=%u xq=%.3f sort=%.3f gateup=%.3f midq=%.3f down=%.3f sum=%.3f total=%.3f ms\n", + n_tokens, pair_count, ms_xq, ms_sort, ms_gate, ms_midq, ms_down, ms_sum, ms_total); + } + for (uint32_t i = 0; i < 7u; i++) (void)cudaEventDestroy(prof_ev[i]); + } + return ok; + } + + if (ok) { + dim3 mgrid(expert_mid_dim, n_tokens * n_expert, 1); + moe_gate_up_mid_f32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + expert_in_dim, + expert_mid_dim, + n_expert, + clamp); + ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch"); + } + if (ok) { + dim3 dgrid(out_dim, n_tokens * n_expert, 1); + moe_down_f32_kernel<<>>( + (float *)down->ptr, + down_w, + (const float *)mid->ptr, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + expert_mid_dim, + out_dim, + n_expert); + ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); + } + if (ok) { + uint64_t n = (uint64_t)n_tokens * out_dim; + moe_sum_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)down->ptr, out_dim, n_expert, n_tokens); + ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch"); + } + return ok; +} + +extern "C" int ds4_gpu_routed_moe_one_owned_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *down, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t resident_expert_base, + uint32_t resident_expert_count, + float clamp, + const ds4_gpu_tensor *x, + ds4_gpu_tensor *down_output, + bool pack_fixed3, + ds4_gpu_tensor *shared_prequant) { + if (!out || !gate || !up || !mid || !down || !model_map || + !selected || !weights || !x || n_expert != 6u || + n_total_expert == 0u || resident_expert_count == 0u || + gate_expert_bytes == 0u || gate_row_bytes == 0u || + down_expert_bytes == 0u || down_row_bytes == 0u || + resident_expert_base >= n_total_expert || + resident_expert_count > n_total_expert - resident_expert_base || + expert_in_dim % CUDA_QK_K != 0u || + expert_mid_dim % CUDA_QK_K != 0u || + selected->bytes < 6u * sizeof(int32_t) || + weights->bytes < 6u * sizeof(float) || + x->bytes < (uint64_t)expert_in_dim * sizeof(float) || + mid->bytes < 6ull * expert_mid_dim * sizeof(float) || + out->bytes < (uint64_t)out_dim * sizeof(float)) { + return 0; + } + if (pack_fixed3 && resident_expert_base == 0u) return 0; + const bool q4k_path = gate_type == 12u && down_type == 12u; + if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; + if (q4k_path && getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL) { + fprintf(stderr, "ds4: CUDA owned Q4 decode does not support gate/up auxiliary output\n"); + return 0; + } + if (!q4k_path && getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") != NULL) { + fprintf(stderr, "ds4: CUDA owned IQ2 decode requires the LUT gate path\n"); + return 0; + } + const bool write_aux = + !q4k_path && getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL; + + if (resident_expert_base > UINT64_MAX / gate_expert_bytes || + resident_expert_count > UINT64_MAX / gate_expert_bytes || + resident_expert_base > UINT64_MAX / down_expert_bytes || + resident_expert_count > UINT64_MAX / down_expert_bytes) { + return 0; + } + const uint64_t gate_shift = (uint64_t)resident_expert_base * gate_expert_bytes; + const uint64_t down_shift = (uint64_t)resident_expert_base * down_expert_bytes; + const uint64_t gate_bytes = (uint64_t)resident_expert_count * gate_expert_bytes; + const uint64_t down_bytes = (uint64_t)resident_expert_count * down_expert_bytes; + if (gate_offset > model_size || gate_shift > model_size - gate_offset || + gate_bytes > model_size - gate_offset - gate_shift || + up_offset > model_size || gate_shift > model_size - up_offset || + gate_bytes > model_size - up_offset - gate_shift || + down_offset > model_size || down_shift > model_size - down_offset || + down_bytes > model_size - down_offset - down_shift) { + return 0; + } + + const int logical_tier = ds4_tensor_device_idx(out); + const char *gate_w = (const char *)cuda_resolve_weight_ptr( + model_map, gate_offset + gate_shift, gate_bytes, + logical_tier, "moe_owned_gate"); + const char *up_w = (const char *)cuda_resolve_weight_ptr( + model_map, up_offset + gate_shift, gate_bytes, + logical_tier, "moe_owned_up"); + const char *down_w = (const char *)cuda_resolve_weight_ptr( + model_map, down_offset + down_shift, down_bytes, + logical_tier, "moe_owned_down"); + if (!gate_w || !up_w || !down_w) return 0; + + const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K; + const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K; + const uint64_t xq_bytes = (uint64_t)xq_blocks * sizeof(cuda_block_q8_K); + const uint64_t midq_bytes = 6ull * midq_blocks * sizeof(cuda_block_q8_K); + const uint64_t down_output_bytes = + (uint64_t)(pack_fixed3 ? 4u : 6u) * out_dim * sizeof(float); + const uint64_t aux_bytes = 6ull * expert_mid_dim * sizeof(float); + const uint64_t shared_q8_blocks = expert_in_dim / 32u; + const uint64_t shared_q8_bytes = shared_q8_blocks * 32u; + const uint64_t shared_scale_offset = + (shared_q8_bytes + 15u) & ~15ull; + const uint64_t shared_prequant_bytes = + shared_scale_offset + shared_q8_blocks * sizeof(float); + if (down->bytes < xq_bytes || down->bytes < down_output_bytes || + (down_output && down_output->bytes < down_output_bytes) || + gate->bytes < midq_bytes || + (shared_prequant && + (shared_prequant->bytes < shared_prequant_bytes || + ds4_tensor_device_idx(shared_prequant) != logical_tier)) || + (write_aux && (gate->bytes < aux_bytes || up->bytes < aux_bytes))) { + return 0; + } + float *down_dst = (float *)(down_output ? down_output->ptr : down->ptr); + cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; + cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; + + dim3 xq_grid(xq_blocks, 1, 1); + if (shared_prequant) { + int8_t *shared_xq = (int8_t *)shared_prequant->ptr; + float *shared_scale = (float *)((char *)shared_prequant->ptr + + shared_scale_offset); + q8_K_q8_0_quantize_kernel<<>>( + xq, shared_xq, shared_scale, (const float *)x->ptr, + expert_in_dim, 1u); + } else { + q8_K_quantize_kernel<<>>( + xq, (const float *)x->ptr, expert_in_dim, 1u); + } + if (!cuda_ok(cudaGetLastError(), "owned routed_moe x quantize launch")) return 0; + + if (q4k_path) { + dim3 gate_grid((expert_mid_dim + 7u) / 8u, 6u, 1u); + moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel<<>>( + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + 6u, + resident_expert_base, + resident_expert_count, + clamp); + } else { + dim3 gate_grid((expert_mid_dim + 31u) / 32u, 6u, 1u); + moe_gate_up_mid_decode_lut_owned_qwarp32_kernel<<>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_w, + up_w, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + gate_expert_bytes, + gate_row_bytes, + xq_blocks, + expert_mid_dim, + 6u, + resident_expert_base, + resident_expert_count, + write_aux, + clamp); + } + if (!cuda_ok(cudaGetLastError(), "owned routed_moe gate/up launch")) return 0; + + dim3 midq_grid(midq_blocks, 6u, 1u); + q8_K_quantize_owned_kernel<<>>( + midq, + (const float *)mid->ptr, + (const int32_t *)selected->ptr, + expert_mid_dim, + 6u, + resident_expert_base, + resident_expert_count); + if (!cuda_ok(cudaGetLastError(), "owned routed_moe mid quantize launch")) return 0; + + dim3 down_grid((out_dim + 31u) / 32u, pack_fixed3 ? 4u : 6u, 1u); + if (q4k_path && pack_fixed3) { + moe_down_q4K_owned_packed_qwarp32_kernel<<>>( + down_dst, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); + } else if (q4k_path) { + moe_down_q4K_owned_slots_qwarp32_kernel<<>>( + down_dst, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); + } else if (pack_fixed3) { + moe_down_owned_packed_qwarp32_kernel<<>>( + down_dst, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); + } else { + moe_down_owned_slots_qwarp32_kernel<<>>( + down_dst, + down_w, + midq, + (const int32_t *)selected->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); + } + return cuda_ok(cudaGetLastError(), "owned routed_moe down launch"); +} + +extern "C" int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_slots, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split, + uint32_t rows) { + if (!out || !home_slots || !peer_slots || !selected || out_dim == 0u || + rows == 0u || rows > 65535u) { + return 0; + } + const uint64_t row_elems = (uint64_t)rows * out_dim; + if (row_elems > UINT64_MAX / (6u * sizeof(float))) return 0; + const uint64_t out_bytes = row_elems * sizeof(float); + const uint64_t slots_bytes = row_elems * 6u * sizeof(float); + const uint64_t selected_bytes = (uint64_t)rows * 6u * sizeof(int32_t); + if (out->bytes < out_bytes || home_slots->bytes < slots_bytes || + peer_slots->bytes < slots_bytes || selected->bytes < selected_bytes) { + return 0; + } + const dim3 grid((out_dim + 255u) / 256u, rows, 1u); + moe_owned_slots_combine_fixed3_kernel<<>>( + (float *)out->ptr, + (const float *)home_slots->ptr, + (const float *)peer_slots->ptr, + (const int32_t *)selected->ptr, + out_dim, + expert_split); + return cuda_ok(cudaGetLastError(), + "owned routed_moe slot rows combine launch"); +} + +extern "C" int ds4_gpu_routed_moe_owned_slots_combine_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_slots, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split) { + return ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( + out, home_slots, peer_slots, selected, + out_dim, expert_split, 1u); +} + +extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_packed, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split) { + const uint64_t home_bytes = 6ull * out_dim * sizeof(float); + const uint64_t peer_bytes = 4ull * out_dim * sizeof(float); + if (!out || !home_slots || !peer_packed || !selected || out_dim == 0u || + out->bytes < (uint64_t)out_dim * sizeof(float) || + home_slots->bytes < home_bytes || peer_packed->bytes < peer_bytes || + selected->bytes < 6u * sizeof(int32_t)) { + return 0; + } + moe_owned_packed_combine_fixed3_kernel<<< + (out_dim + 255u) / 256u, 256>>>( + (float *)out->ptr, + (const float *)home_slots->ptr, + (const float *)peer_packed->ptr, + (const int32_t *)selected->ptr, + out_dim, + expert_split); + return cuda_ok(cudaGetLastError(), + "owned routed_moe packed combine launch"); +} + +extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, + const ds4_gpu_tensor *add_in, + uint32_t layer_index, + bool force_resident) { + if (add_in) { + if (!ds4_gpu_add_tensor(out, out, add_in, + (uint32_t)(out->bytes / sizeof(float)))) return 0; + } + return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, + gate_offset, up_offset, down_offset, + gate_type, down_type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + expert_in_dim, expert_mid_dim, out_dim, + selected, weights, n_total_expert, n_expert, clamp, x, + layer_index, 1, force_resident ? 0 : 1, 0); +} +extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) { + (void)force_resident; + if (mid_is_f16) *mid_is_f16 = false; + return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, + gate_offset, up_offset, down_offset, + gate_type, down_type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + expert_in_dim, expert_mid_dim, out_dim, + selected, weights, n_total_expert, n_expert, clamp, x, + layer_index, n_tokens, 1, 0); +} + +extern "C" int ds4_gpu_routed_moe_batch_owned_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *down, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t resident_expert_base, + uint32_t resident_expert_count, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + bool *mid_is_f16) { + if (mid_is_f16) *mid_is_f16 = false; + if (!selected || !weights || n_tokens == 0u || n_expert == 0u || + n_total_expert == 0u || resident_expert_count == 0u || + gate_expert_bytes == 0u || gate_row_bytes == 0u || + down_expert_bytes == 0u || down_row_bytes == 0u || + resident_expert_base >= n_total_expert || + resident_expert_count > n_total_expert - resident_expert_base || + resident_expert_base > UINT64_MAX / gate_expert_bytes || + resident_expert_base > UINT64_MAX / down_expert_bytes) { + return 0; + } + const uint64_t pair_count = (uint64_t)n_tokens * n_expert; + if (pair_count > UINT32_MAX || + selected->bytes < pair_count * sizeof(int32_t) || + weights->bytes < pair_count * sizeof(float)) { + return 0; + } + const uint64_t gate_shift = + (uint64_t)resident_expert_base * gate_expert_bytes; + const uint64_t down_shift = + (uint64_t)resident_expert_base * down_expert_bytes; + if (gate_offset > model_size || gate_shift > model_size - gate_offset || + up_offset > model_size || gate_shift > model_size - up_offset || + down_offset > model_size || down_shift > model_size - down_offset) { + return 0; + } + moe_filter_owned_pairs_kernel<<<(pair_count + 255u) / 256u, 256>>>( + (int32_t *)selected->ptr, + (float *)weights->ptr, + pair_count, + n_total_expert, + resident_expert_base, + resident_expert_count); + if (!cuda_ok(cudaGetLastError(), "owned routed_moe pair filter launch")) return 0; + + return routed_moe_launch( + out, gate, up, mid, down, model_map, model_size, + gate_offset + gate_shift, + up_offset + gate_shift, + down_offset + down_shift, + gate_type, down_type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + expert_in_dim, expert_mid_dim, out_dim, + selected, weights, resident_expert_count, n_expert, + clamp, x, layer_index, n_tokens, 0, 1); +} diff --git a/models/deepseek/graph.inc b/models/deepseek/graph.inc new file mode 100644 index 0000000000..44e82da9ff --- /dev/null +++ b/models/deepseek/graph.inc @@ -0,0 +1,20333 @@ +/* + * DeepSeek V4 graph state, allocation, decode, prefill, and diagnostics. + * + * Included exactly once inside ds4.c's graph-backend conditional. This keeps + * all custom GPU orchestration concrete and statically linked. + */ + +/* + * Apple Metal stores the persistent attention-compressed KV cache in F16. The + * compressor still pools, normalizes, RoPEs, and FP8-rounds rows in F32 staging + * before writing the cache, while checkpoints and debug dumps expand back to + * F32 for the stable external format. This is a storage optimization rather + * than a semantic approximation: all Metal attention consumers already run the + * compressed K/V rows through F16 FlashAttention/indexed-attention paths. + */ +#if defined(__APPLE__) +#define DS4_GPU_ATTN_COMP_CACHE_F16 1 +#else +#define DS4_GPU_ATTN_COMP_CACHE_F16 0 +#endif + +#define DS4_GPU_GLM_COMPACT_CACHE_F16 DS4_GPU_ATTN_COMP_CACHE_F16 + +/* ========================================================================= + * Metal Release Graph State. + * ========================================================================= + * + * The release Metal executor owns one fixed set of tensors for single-token + * decode and another for batched prefill. The structure is DS4-specific: + * tensor names follow the model stages rather than generic graph nodes. + */ + +enum { DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS = 64 }; + +typedef struct { + /* Class P — per-tier replicated kernel scratch buffers. + * Each used tier has its own copy; active_tier names the slot the + * current dispatch step reads/writes. Single-tier paths leave + * active_tier == 0; multi-tier dispatch updates active_tier in B6. + * + * Decode hidden-state buffers. A generated token enters as an embedding + * in cur_hc and leaves as logits after all 43 layers update their + * raw/compressed/indexer caches. The hc_pre / hc_post / hc_comb views + * are derived from hc_split per tier (see metal_graph_alloc_raw_cap). */ + ds4_gpu_tensor *cur_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *flat_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *hc_mix_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *hc_split_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *hc_pre_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ + ds4_gpu_tensor *hc_post_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ + ds4_gpu_tensor *hc_comb_by_tier[DS4_MAX_GPUS]; /* views of hc_split */ + ds4_gpu_tensor *attn_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *attn_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *qr_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *qr_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *q_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *kv_raw_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *kv_by_tier[DS4_MAX_GPUS]; + int active_tier; + /* cached engine placement[] (length DS4_N_LAYER + 2) for the + * dispatch loops. NULL in single-tier mode — active_tier stays 0 and + * dispatch wrappers no-op the tier-switch + cross-device copy. The + * pointer aliases e->placement; the engine outlives the graph so this + * is safe. */ + const int *placement; + + /* Persistent KV state. Raw KV is a sliding-window ring per layer. Ratio-4 + * layers also keep an indexer-compressed cache; ratio-128 layers keep only + * the attention-compressed cache. The small state tensors are compressor + * frontiers for the next compressed row, so they must be snapshotted with + * the row counters whenever a checkpoint is saved or partially rewound. */ + ds4_gpu_tensor *layer_raw_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_comp_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_comp_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_raw_cache_tp[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_comp_cache_tp[DS4_MAX_LAYER]; + + /* Speculative decoding scratch. MTP is allowed to mutate graph state only + * if the target verifier can either commit it or restore the saved + * frontiers. The prefix1 buffers are the cheap partial-accept state for the + * common N=2 case. */ + ds4_gpu_tensor *spec_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_index_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_index_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_logits; + uint32_t layer_n_comp[DS4_MAX_LAYER]; + uint32_t layer_n_index_comp[DS4_MAX_LAYER]; + uint32_t spec_prefix1_n_comp[DS4_MAX_LAYER]; + uint32_t spec_prefix1_n_index_comp[DS4_MAX_LAYER]; + bool spec_capture_prefix1; + uint32_t raw_cap; + /* Maximum compressed-row capacity across layers. Shared work buffers use + * this worst-case size because ratio-4 indexer layers can still reach it. */ + uint32_t comp_cap; + /* Persistent compressed caches are per layer, so size them from the actual + * layer compression ratio instead of pessimistically using the ratio-4 cap + * for every ratio-128 layer. */ + uint32_t layer_comp_cap[DS4_MAX_LAYER]; + uint32_t attn_comp_stage_cap; + + /* Class P (per-layer work tensors). Each used tier has its + * own replica. They are reused in place by every layer instead of + * allocating a generic graph arena. This is why the code is verbose but + * predictable: each pointer names an actual DS4 stage. */ + ds4_gpu_tensor *comp_kv_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *comp_sc_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *attn_comp_stage_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *indexer_q_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *indexer_weights_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *indexer_scores_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *comp_mask_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *comp_selected_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *heads_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *attn_low_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *attn_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *after_attn_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *ffn_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *ffn_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *shared_gate_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *shared_up_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *shared_mid_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *shared_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *router_logits_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *router_probs_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *router_selected_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *router_weights_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *routed_gate_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *routed_up_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *routed_mid_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *routed_down_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *routed_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *ffn_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *after_ffn_hc_by_tier[DS4_MAX_GPUS]; + /* Class H — output-head buffers and logits live on the + * head tier only. head_tier is captured at metal_graph_alloc_raw_cap + * time from placement[DS4_N_LAYER + 1] (or 0 in single-tier / + * diagnostic paths). Non-head slots remain NULL. Readers go through + * the metal_graph_logits / metal_graph_output_* accessors below. */ + ds4_gpu_tensor *output_pre_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *output_weights_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *output_embd_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *output_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *logits_by_tier[DS4_MAX_GPUS]; + int head_tier; + + /* DSpark target features. The proposer consumes mean-over-HC rows from + * selected target layers; keeping them on-GPU avoids adding readbacks to + * the target path. */ + ds4_gpu_tensor *dspark_hc_mean_weights; + ds4_gpu_tensor *dspark_hc_mean_rows; + ds4_gpu_tensor *dspark_target_hidden; + ds4_gpu_tensor *dspark_target_hidden_batch; + ds4_gpu_tensor *dspark_stage0_packed; + ds4_gpu_tensor *dspark_stage0_proj; + ds4_gpu_tensor *dspark_main_x; + ds4_gpu_tensor *dspark_draft_tokens; + ds4_gpu_tensor *dspark_draft_hc; + ds4_gpu_tensor *dspark_target_hc; + ds4_gpu_tensor *dspark_stage_input_hc; + ds4_gpu_tensor *dspark_stage_output_hc; + ds4_gpu_tensor *dspark_position_ids; + ds4_gpu_tensor *dspark_raw_cache[DS4_DSPARK_MAX_STAGES]; + uint32_t dspark_cache_cap; + uint32_t dspark_cache_start; + uint32_t dspark_cache_token_start; + uint32_t dspark_cache_len; + uint32_t dspark_target_layer_count; + uint32_t dspark_block_size; + uint32_t dspark_target_layers[DS4_DSPARK_MAX_TARGET_LAYERS]; + uint32_t dspark_capture_mask; + uint32_t dspark_capture_checkpoint_len; + uint32_t dspark_capture_batch_mask; + uint32_t dspark_capture_batch_start; + uint32_t dspark_capture_batch_tokens; + bool dspark_capture_valid; + bool dspark_capture_batch_valid; + int dspark_exec_tier; + bool dspark_capture_enabled; + bool verify_small_batch_tp; + uint32_t pipeline_capture_chunk_start; + uint32_t pipeline_capture_chunk_len; + bool ssd_streaming; /* glm-branch SSD streaming; always false here */ + + /* Optional MTP model state. It has its own raw cache because the drafter + * runs on speculative future tokens; target KV state is updated only after + * verification accepts draft tokens. */ + ds4_gpu_tensor *mtp_embed; + ds4_gpu_tensor *mtp_enorm; + ds4_gpu_tensor *mtp_eproj; + ds4_gpu_tensor *mtp_eproj_hc; + ds4_gpu_tensor *mtp_hnorm_hc; + ds4_gpu_tensor *mtp_hproj_hc; + ds4_gpu_tensor *mtp_input_hc; + ds4_gpu_tensor *mtp_state_hc; + ds4_gpu_tensor *mtp_next_hc; + ds4_gpu_tensor *mtp_raw_cache; + uint32_t mtp_n_raw; + uint32_t prefill_cap; + uint32_t raw_window; + uint32_t batch_token_offset; + + /* Batched prefill tensors. Prefill is layer-major: a chunk of prompt + * tokens moves through layer 0, then layer 1, and so on, updating the same + * persistent caches used by decode. Keeping this separate from decode + * avoids a slow loop of one-token graph steps for long prompts. */ + /* Class E — embedding-tier-only prompt-token integer buffer. + * Captured at metal_graph_alloc_raw_cap time from placement[0] (or 0 in + * single-tier / diagnostic paths). Non-embedding slots stay NULL. Readers + * go through metal_graph_prefill_tokens() below. */ + ds4_gpu_tensor *prefill_tokens_by_tier[DS4_MAX_GPUS]; + int emb_tier; + /* Class P batch (chunked-prefill) scratch — per-tier + * replicated. The cur/next pair is ping-ponged per layer step on the + * layer's active tier; tier transitions copy the active buffer across + * boundaries via ds4_gpu_tensor_copy_xdev (handled in B6). */ + ds4_gpu_tensor *batch_cur_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_next_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_flat_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_hc_mix_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_hc_split_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_attn_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_attn_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_qr_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_qr_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_q_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_kv_raw_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_kv_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_comp_kv_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_comp_sc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_indexer_q_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_indexer_weights_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_heads_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_attn_low_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_attn_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_group_tmp_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_low_tmp_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_after_attn_hc_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_ffn_cur_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_ffn_norm_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_shared_gate_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_shared_up_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_shared_mid_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_shared_out_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_router_logits_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_router_probs_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_router_selected_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_router_weights_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_routed_gate_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_routed_up_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_routed_mid_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_routed_down_by_tier[DS4_MAX_GPUS]; + ds4_gpu_tensor *batch_routed_out_by_tier[DS4_MAX_GPUS]; + bool batch_routed_mid_is_f16; + ds4_gpu_tensor *batch_ffn_out_by_tier[DS4_MAX_GPUS]; + bool owns_prefill_workspace; + bool materialize_ffn_out; + /* Class P (replicated per tier — this is + * consumed in per-layer attn/FFN kernels, NOT embedding-only). Read-only + * after init; replicate by writing the same host directions buffer to + * every used tier's slot during session setup. */ + ds4_gpu_tensor *directional_steering_dirs_by_tier[DS4_MAX_GPUS]; + float directional_steering_attn_scale; + float directional_steering_ffn_scale; + bool cuda_tp_decode; + bool cuda_tp_attn; + bool cuda_tp_attn_peer_read; + bool cuda_tp_attn_heads; + bool cuda_tp_attn_cache_dup; + bool cuda_tp_moe; + bool cuda_tp_ep; + bool cuda_tp_ep_pack_exact; + bool cuda_tp_moe_delay_reduce; + bool cuda_tp_moe_copy3_handoff; + bool cuda_tp_moe_pack_handoff; + bool cuda_tp_moe_peer_read; + bool cuda_tp_moe_peer_router; + bool cuda_tp_shared; + bool cuda_tp_shared_fold; + bool cuda_tp_q; + bool cuda_tp_output; + bool cuda_tp_prefill_ffn; + bool cuda_tp_prefill_attn_output; + bool cuda_q_norm_rope_fuse; + bool cuda_qkv_kv_rope_fuse; + bool cuda_qkv_pair; + bool cuda_tp_attn_out_hc_fuse; + bool shared_gate_up_swiglu_fuse; + bool decode_stage_profile; + bool decode_index_stage_profile; + bool output_stage_profile; + ds4_gpu_tensor *tp_peer_tmp_by_tier[DS4_MAX_GPUS]; + uint32_t power_percent; + double prefill_layer_avg_sec[DS4_MAX_LAYER]; + double decode_token_avg_sec; + bool quality; + bool mtp_enabled; + /* Metal-only prefill helpers retained alongside the CUDA tiered workspace. */ + ds4_gpu_tensor *batch_q_half; + ds4_gpu_tensor *prefill_seed_router_selected; + uint32_t prefill_seed_tokens; + uint64_t prefill_selected_profile_rows; + uint64_t prefill_selected_profile_unique; + uint64_t prefill_selected_profile_selected_bytes; + uint64_t prefill_selected_profile_full_bytes; + uint32_t prefill_selected_profile_layers; + uint32_t prefill_selected_profile_min_unique; + uint32_t prefill_selected_profile_max_unique; + uint32_t streaming_preload_experts; + bool ssd_streaming_cold; + bool streaming_static_decode_map_current; + float *cpu_router_norm; + + /* Metal network tensor parallelism. These views alias engine-owned + * transport slabs except tp_logits_half, whose view object is session-owned. */ + uint32_t tp_world; + uint32_t tp_rank; + ds4_gpu_tensor **tp_out; + ds4_gpu_tensor **tp_in; + ds4_gpu_tensor **tp_batch_out; + ds4_gpu_tensor **tp_batch_in; + uint32_t tp_batch_rows; + ds4_gpu_tensor *tp_zero; + ds4_gpu_tensor *tp_logits_half; +} ds4_gpu_graph; + +/* Tensors that are temporary for chunked prefill and grouped multi-session + * decode. The batched server serializes every operation that uses them, so one + * engine-owned set can be aliased by all resident session graphs. */ +#define DS4_GPU_PREFILL_WORKSPACE_FIELDS(X) \ + X(prefill_tokens) \ + X(batch_ffn_out) \ + X(batch_routed_out) \ + X(batch_routed_down) \ + X(batch_routed_mid) \ + X(batch_routed_up) \ + X(batch_routed_gate) \ + X(batch_router_weights) \ + X(batch_router_selected) \ + X(batch_router_probs) \ + X(batch_router_logits) \ + X(batch_shared_out) \ + X(batch_shared_mid) \ + X(batch_shared_up) \ + X(batch_shared_gate) \ + X(batch_ffn_norm) \ + X(batch_ffn_cur) \ + X(batch_after_attn_hc) \ + X(batch_low_tmp) \ + X(batch_group_tmp) \ + X(batch_attn_out) \ + X(batch_attn_low) \ + X(batch_heads) \ + X(batch_indexer_weights) \ + X(batch_indexer_q) \ + X(batch_comp_sc) \ + X(batch_comp_kv) \ + X(batch_kv) \ + X(batch_kv_raw) \ + X(batch_q) \ + X(batch_qr_norm) \ + X(batch_qr) \ + X(batch_attn_norm) \ + X(batch_attn_cur) \ + X(batch_hc_split) \ + X(batch_hc_mix) \ + X(batch_flat_hc) \ + X(batch_next_hc) \ + X(batch_cur_hc) + +/* Class H accessors. All reader sites for the output-head + * tensors and the final logits route through these inlines, which read the + * head_tier slot captured at allocation time. Single-tier paths set + * head_tier == 0 and the slot is byte-identical to the legacy + * metal_graph_logits(g) / g->output_* pointers. Multi-tier paths set head_tier + * to placement[DS4_N_LAYER + 1]; other tier slots remain NULL. */ +static inline ds4_gpu_tensor *metal_graph_logits(const ds4_gpu_graph *g) { + return g->logits_by_tier[g->head_tier]; +} +static inline ds4_gpu_tensor *metal_graph_output_pre(const ds4_gpu_graph *g) { + return g->output_pre_by_tier[g->head_tier]; +} +static inline ds4_gpu_tensor *metal_graph_output_weights(const ds4_gpu_graph *g) { + return g->output_weights_by_tier[g->head_tier]; +} +static inline ds4_gpu_tensor *metal_graph_output_embd(const ds4_gpu_graph *g) { + return g->output_embd_by_tier[g->head_tier]; +} +static inline ds4_gpu_tensor *metal_graph_output_norm(const ds4_gpu_graph *g) { + return g->output_norm_by_tier[g->head_tier]; +} + +/* Class E accessor. The prompt-token integer buffer is + * consumed by the embedding kernel on the embedding tier only. Single-tier + * paths set emb_tier == 0 (byte-equivalent to the legacy single-tier + * pointer). Multi-tier paths set emb_tier = placement[0]. */ +static inline ds4_gpu_tensor *metal_graph_prefill_tokens(const ds4_gpu_graph *g) { + return g->prefill_tokens_by_tier[g->emb_tier]; +} + +/* Class P accessors. Each Class P kernel-scratch buffer is + * replicated across every tier the placement uses; the active_tier field + * names the slot the current dispatch step reads/writes. Single-tier paths + * leave active_tier == 0 (byte-equivalent to the legacy single-tier + * pointer). Multi-tier dispatch (wired up in B6) updates active_tier with + * the current layer's home tier before each kernel-dispatch wrapper runs. */ +#define DS4_GPU_GRAPH_CLASS_P_ACCESSOR(name) \ +static inline ds4_gpu_tensor *metal_graph_##name(const ds4_gpu_graph *g) { \ + return g->name##_by_tier[g->active_tier]; \ +} + +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(cur_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(flat_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_mix) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_split) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_pre) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_post) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(hc_comb) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(qr) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(qr_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(q) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(kv_raw) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(kv) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_kv_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_sc_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_comp_stage) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_q) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_weights) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(indexer_scores) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_mask) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(comp_selected) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(heads) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_low) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(attn_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(after_attn_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_gate) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_up) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_mid) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(shared_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_logits) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_probs) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_selected) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(router_weights) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_gate) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_up) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_mid) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_down) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(routed_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(ffn_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(after_ffn_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_cur_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_next_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_flat_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_hc_mix) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_hc_split) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_qr) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_qr_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_q) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_kv_raw) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_kv) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_comp_kv) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_comp_sc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_indexer_q) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_indexer_weights) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_heads) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_low) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_attn_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_group_tmp) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_low_tmp) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_after_attn_hc) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_cur) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_norm) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_gate) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_up) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_mid) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_shared_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_logits) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_probs) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_selected) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_router_weights) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_gate) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_up) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_mid) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_down) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_routed_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(batch_ffn_out) +DS4_GPU_GRAPH_CLASS_P_ACCESSOR(directional_steering_dirs) + +/* dispatch-loop helpers for multi-tier per-layer execution. + * + * Single-tier (g->placement == NULL): all helpers are no-ops; active_tier + * stays 0 from memset; behavior is byte-equivalent to legacy. + * + * Multi-tier: each helper switches g->active_tier to the requested tier + * BEFORE the next kernel-dispatch wrapper reads any Class P accessor. If + * the source-tier Class P cur_hc (or batch_cur_hc) differs from the new + * tier's, ds4_gpu_tensor_copy_xdev ferries the active hidden state across + * the boundary. copy_xdev returns 1 on success, 0 on failure. The + * destination tensor's device_id was stamped at alloc_on time and is + * immutable. + * + * For decode (one token at a time): metal_graph_set_active_tier_decode + * swaps to the requested tier and copies cur_hc across the boundary. + * + * For batch (chunked prefill): metal_graph_set_active_tier_batch swaps + * tier and copies batch_cur_hc across the boundary. The next/cur pair + * is maintained per tier — after a copy, batch_next_hc on the destination + * tier becomes the swap target for the next layer step on that tier. + * + * Helpers always invoke ds4_gpu_set_current_device(tier) so the next + * kernel-launch sees the correct CUDA device. */ + +/* ds4_gpu_set_current_device is declared in ds4_gpu_mgpu.h — single-tier + * (g_n_gpus <= 1) callers no-op. Returns 0 on success. */ + +#ifdef DS4_NO_GPU +static inline int ds4_gpu_set_current_device(int tier) { (void)tier; return 0; } +static inline int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint64_t bytes) { + (void)dst; (void)src; (void)bytes; return 1; +} +static inline 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) { + (void)dst0; (void)src0; (void)bytes0; + (void)dst1; (void)src1; (void)bytes1; + (void)dst2; (void)src2; (void)bytes2; + return 1; +} +static inline int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint64_t bytes) { + (void)dst; (void)src; (void)bytes; return 1; +} +static inline int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) { + (void)src; (void)dst_tier; return 1; +} +static inline int ds4_gpu_moe_handoff_pack_tensor( + ds4_gpu_tensor *packed, + const ds4_gpu_tensor *ffn_norm, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_expert) { + (void)packed; (void)ffn_norm; (void)selected; (void)weights; + (void)n_embd; (void)n_expert; return 1; +} +static inline int ds4_gpu_q8_cache_suppressed(void) { return 0; } +static inline void ds4_gpu_set_q8_cache_suppressed(int suppressed) { (void)suppressed; } +static inline int ds4_gpu_set_decode_fast_attention(int enabled) { + (void)enabled; + return 0; +} +static inline int ds4_gpu_set_decode_score_vec4(int enabled) { + (void)enabled; + return 0; +} +static inline 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) { + (void)q_out; (void)q; (void)model_map; (void)model_size; + (void)q_weight_offset; (void)q_n; (void)kv_out; (void)kv; + (void)kv_weight_offset; (void)kv_n; (void)rows; (void)kv_n_head; + (void)kv_head_dim; (void)n_rot; (void)pos0; (void)n_ctx_orig; + (void)inverse; (void)freq_base; (void)freq_scale; (void)ext_factor; + (void)attn_factor; (void)beta_fast; (void)beta_slow; (void)eps; + return 0; +} +#endif + +/* Returns true on success. Single-tier: no-op success. Multi-tier: + * sets the CUDA device, then if tier differs from current active_tier, + * copies cur_hc to the destination tier and updates active_tier. */ +static bool metal_graph_set_active_tier_decode(ds4_gpu_graph *g, int tier) { + if (!g->placement) { + /* Single-tier: just keep active_tier at 0; no device switch needed. */ + (void)tier; + return true; + } + if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + if (tier == g->active_tier) return true; + if (ds4_gpu_set_current_device(tier) != 0) return false; + /* Boundary hop: copy cur_hc from source-tier to destination-tier slot. */ + if (g->active_tier >= 0) { + ds4_gpu_tensor *src = g->cur_hc_by_tier[g->active_tier]; + ds4_gpu_tensor *dst = g->cur_hc_by_tier[tier]; + if (src && dst) { + const uint64_t hc_bytes = (uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float); + if (!ds4_gpu_tensor_copy_xdev(dst, src, hc_bytes)) return false; + } + } + g->active_tier = tier; + return true; +} + +/* Returns true on success. Same semantics as the decode helper but ferries + * batch_cur_hc (which contains chunk_tokens * hc_dim floats — variable per + * prefill call). The caller passes the chunk size in tokens; single-tier + * paths ignore the argument. */ +static bool metal_graph_set_active_tier_batch(ds4_gpu_graph *g, int tier, uint32_t chunk_tokens) { + if (!g->placement) { + (void)tier; + (void)chunk_tokens; + return true; + } + if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + if (tier == g->active_tier) return true; + if (ds4_gpu_set_current_device(tier) != 0) return false; + if (g->active_tier >= 0) { + ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[g->active_tier]; + ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[tier]; + if (src && dst) { + const uint64_t hc_bytes = + (uint64_t)chunk_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float); + if (!ds4_gpu_tensor_copy_xdev(dst, src, hc_bytes)) return false; + } + } + g->active_tier = tier; + return true; +} + +static bool metal_graph_set_active_tier_no_copy(ds4_gpu_graph *g, int tier) { + if (!g->placement) { + (void)tier; + return true; + } + if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + if (tier == g->active_tier) return true; + if (ds4_gpu_set_current_device(tier) != 0) return false; + g->active_tier = tier; + return true; +} + +/* Upstream: --power N GPU duty-cycle throttling helpers. The single-tier + * --power=100 path is a no-op; multi-tier inherits the same helpers via + * graph_power_note_prefill_layer / graph_power_note_decode_token which we + * call from the shared encode / decode loops. */ + +static bool graph_power_throttle_enabled(const ds4_gpu_graph *g) { + return g && g->power_percent > 0 && g->power_percent < 100; +} + +static double graph_power_update_avg(double avg, double sample) { + if (sample <= 0.0 || !isfinite(sample)) return avg; + if (avg <= 0.0 || !isfinite(avg)) return sample; + return avg * 0.875 + sample * 0.125; +} + +static void graph_power_sleep(double work_sec, uint32_t power_percent) { + if (power_percent == 0 || power_percent >= 100) return; + /* Target duty cycle: work / (work + sleep) = power / 100. + * At --power 50 this sleeps for one measured work interval; at 25 it + * sleeps for three. */ + const double sleep = work_sec * (100.0 - (double)power_percent) / + (double)power_percent; + sleep_sec(sleep); +} + +static void graph_power_note_prefill_layer(ds4_gpu_graph *g, + uint32_t il, + double elapsed_sec) { + if (!graph_power_throttle_enabled(g)) return; + if (il >= DS4_N_LAYER) return; + g->prefill_layer_avg_sec[il] = + graph_power_update_avg(g->prefill_layer_avg_sec[il], elapsed_sec); + graph_power_sleep(g->prefill_layer_avg_sec[il], g->power_percent); +} + +static void graph_power_note_decode_token(ds4_gpu_graph *g, double elapsed_sec) { + if (!graph_power_throttle_enabled(g)) return; + g->decode_token_avg_sec = + graph_power_update_avg(g->decode_token_avg_sec, elapsed_sec); + graph_power_sleep(g->decode_token_avg_sec, g->power_percent); +} + +static void metal_graph_copy_prefill_workspace_pointers( + ds4_gpu_graph *dst, + const ds4_gpu_graph *src) { +#define DS4_COPY_PREFILL_FIELD(name) \ + memcpy(dst->name##_by_tier, src->name##_by_tier, \ + sizeof(dst->name##_by_tier)); + DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_COPY_PREFILL_FIELD) +#undef DS4_COPY_PREFILL_FIELD + dst->batch_q_half = src->batch_q_half; + dst->prefill_seed_router_selected = src->prefill_seed_router_selected; +} + +static void metal_graph_transfer_prefill_workspace( + ds4_gpu_graph *dst, + ds4_gpu_graph *src) { + memset(dst, 0, sizeof(*dst)); + dst->prefill_cap = src->prefill_cap; + dst->emb_tier = src->emb_tier; + dst->owns_prefill_workspace = true; + metal_graph_copy_prefill_workspace_pointers(dst, src); + src->owns_prefill_workspace = false; +} + +static uint64_t metal_graph_prefill_workspace_bytes(const ds4_gpu_graph *g) { + uint64_t total = 0; + for (int t = 0; t < DS4_MAX_GPUS; t++) { +#define DS4_COUNT_PREFILL_FIELD(name) \ + total += ds4_gpu_tensor_bytes(g->name##_by_tier[t]); + DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_COUNT_PREFILL_FIELD) +#undef DS4_COUNT_PREFILL_FIELD + } + total += ds4_gpu_tensor_bytes(g->batch_q_half); + total += ds4_gpu_tensor_bytes(g->prefill_seed_router_selected); + return total; +} + +static void metal_graph_free_prefill_workspace(ds4_gpu_graph *g) { + if (!g || !g->owns_prefill_workspace) return; + for (int t = 0; t < DS4_MAX_GPUS; t++) { +#define DS4_FREE_PREFILL_FIELD(name) \ + ds4_gpu_tensor_free(g->name##_by_tier[t]); \ + g->name##_by_tier[t] = NULL; + DS4_GPU_PREFILL_WORKSPACE_FIELDS(DS4_FREE_PREFILL_FIELD) +#undef DS4_FREE_PREFILL_FIELD + } + ds4_gpu_tensor_free(g->batch_q_half); + ds4_gpu_tensor_free(g->prefill_seed_router_selected); + g->batch_q_half = NULL; + g->prefill_seed_router_selected = NULL; + g->owns_prefill_workspace = false; +} + +/* Release every Metal tensor owned by the whole-model graph runtime. */ +static void metal_graph_free(ds4_gpu_graph *g) { + /* free every Class P slot across all DS4_MAX_GPUS tier + * slots. Unallocated slots are NULL and ds4_gpu_tensor_free(NULL) is a + * no-op. The hc_pre / hc_post / hc_comb views must be freed BEFORE + * their parent hc_split — view destruction releases its own struct + * but does not touch the parent's memory. */ + metal_graph_free_prefill_workspace(g); + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(g->directional_steering_dirs_by_tier[t]); + g->directional_steering_dirs_by_tier[t] = NULL; + } + /* Class H free across all tier slots. Non-head slots are + * NULL and ds4_gpu_tensor_free(NULL) is a no-op. */ + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(g->logits_by_tier[t]); + g->logits_by_tier[t] = NULL; + } + ds4_gpu_tensor_free(g->mtp_raw_cache); + ds4_gpu_tensor_free(g->mtp_next_hc); + ds4_gpu_tensor_free(g->mtp_state_hc); + ds4_gpu_tensor_free(g->mtp_input_hc); + ds4_gpu_tensor_free(g->mtp_hproj_hc); + ds4_gpu_tensor_free(g->mtp_hnorm_hc); + ds4_gpu_tensor_free(g->mtp_eproj_hc); + ds4_gpu_tensor_free(g->mtp_eproj); + ds4_gpu_tensor_free(g->mtp_enorm); + ds4_gpu_tensor_free(g->mtp_embed); + ds4_gpu_tensor_free(g->spec_logits); + /* Class H output-head free across all tier slots. */ + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(g->output_norm_by_tier[t]); + g->output_norm_by_tier[t] = NULL; + ds4_gpu_tensor_free(g->output_embd_by_tier[t]); + g->output_embd_by_tier[t] = NULL; + ds4_gpu_tensor_free(g->output_weights_by_tier[t]); + g->output_weights_by_tier[t] = NULL; + ds4_gpu_tensor_free(g->output_pre_by_tier[t]); + g->output_pre_by_tier[t] = NULL; + } + /* Class P decode scratch + routed-FFN free across all + * tier slots. ffn_out is also a Class P field freed here. */ + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(g->after_ffn_hc_by_tier[t]); + ds4_gpu_tensor_free(g->ffn_out_by_tier[t]); + ds4_gpu_tensor_free(g->routed_out_by_tier[t]); + ds4_gpu_tensor_free(g->routed_down_by_tier[t]); + ds4_gpu_tensor_free(g->routed_mid_by_tier[t]); + ds4_gpu_tensor_free(g->routed_up_by_tier[t]); + ds4_gpu_tensor_free(g->routed_gate_by_tier[t]); + ds4_gpu_tensor_free(g->tp_peer_tmp_by_tier[t]); + ds4_gpu_tensor_free(g->router_weights_by_tier[t]); + ds4_gpu_tensor_free(g->router_selected_by_tier[t]); + ds4_gpu_tensor_free(g->router_probs_by_tier[t]); + ds4_gpu_tensor_free(g->router_logits_by_tier[t]); + ds4_gpu_tensor_free(g->shared_out_by_tier[t]); + ds4_gpu_tensor_free(g->shared_mid_by_tier[t]); + ds4_gpu_tensor_free(g->shared_up_by_tier[t]); + ds4_gpu_tensor_free(g->shared_gate_by_tier[t]); + ds4_gpu_tensor_free(g->ffn_norm_by_tier[t]); + ds4_gpu_tensor_free(g->ffn_cur_by_tier[t]); + ds4_gpu_tensor_free(g->after_attn_hc_by_tier[t]); + ds4_gpu_tensor_free(g->attn_out_by_tier[t]); + ds4_gpu_tensor_free(g->attn_low_by_tier[t]); + ds4_gpu_tensor_free(g->heads_by_tier[t]); + ds4_gpu_tensor_free(g->comp_sc_cur_by_tier[t]); + ds4_gpu_tensor_free(g->comp_kv_cur_by_tier[t]); + ds4_gpu_tensor_free(g->attn_comp_stage_by_tier[t]); + ds4_gpu_tensor_free(g->comp_mask_by_tier[t]); + ds4_gpu_tensor_free(g->comp_selected_by_tier[t]); + ds4_gpu_tensor_free(g->indexer_scores_by_tier[t]); + ds4_gpu_tensor_free(g->indexer_weights_by_tier[t]); + ds4_gpu_tensor_free(g->indexer_q_by_tier[t]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_raw_cache[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_raw_cache_tp[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_attn_comp_cache[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_attn_comp_cache_tp[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_attn_state_kv[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_attn_state_score[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_index_comp_cache[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_index_state_kv[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_index_state_score[il]); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + ds4_gpu_tensor_free(g->spec_attn_state_kv[il]); + ds4_gpu_tensor_free(g->spec_attn_state_score[il]); + ds4_gpu_tensor_free(g->spec_index_state_kv[il]); + ds4_gpu_tensor_free(g->spec_index_state_score[il]); + ds4_gpu_tensor_free(g->spec_prefix1_attn_state_kv[il]); + ds4_gpu_tensor_free(g->spec_prefix1_attn_state_score[il]); + ds4_gpu_tensor_free(g->spec_prefix1_index_state_kv[il]); + ds4_gpu_tensor_free(g->spec_prefix1_index_state_score[il]); + } + /* Class P decode-step scratch + decode HC group free across + * all tier slots. hc_pre / hc_post / hc_comb are VIEWS of hc_split — free + * them before hc_split so the view struct release happens with the parent + * still pointer-valid (view free does not touch parent memory). */ + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(g->kv_by_tier[t]); + ds4_gpu_tensor_free(g->kv_raw_by_tier[t]); + ds4_gpu_tensor_free(g->q_by_tier[t]); + ds4_gpu_tensor_free(g->qr_norm_by_tier[t]); + ds4_gpu_tensor_free(g->qr_by_tier[t]); + ds4_gpu_tensor_free(g->attn_norm_by_tier[t]); + ds4_gpu_tensor_free(g->attn_cur_by_tier[t]); + ds4_gpu_tensor_free(g->hc_comb_by_tier[t]); + ds4_gpu_tensor_free(g->hc_post_by_tier[t]); + ds4_gpu_tensor_free(g->hc_pre_by_tier[t]); + ds4_gpu_tensor_free(g->hc_split_by_tier[t]); + ds4_gpu_tensor_free(g->hc_mix_by_tier[t]); + ds4_gpu_tensor_free(g->flat_hc_by_tier[t]); + ds4_gpu_tensor_free(g->cur_hc_by_tier[t]); + } + ds4_gpu_tensor_free(g->dspark_position_ids); + ds4_gpu_tensor_free(g->dspark_stage_output_hc); + ds4_gpu_tensor_free(g->dspark_stage_input_hc); + ds4_gpu_tensor_free(g->dspark_target_hc); + ds4_gpu_tensor_free(g->dspark_draft_hc); + ds4_gpu_tensor_free(g->dspark_draft_tokens); + for (uint32_t stage = 0; stage < DS4_DSPARK_MAX_STAGES; stage++) { + ds4_gpu_tensor_free(g->dspark_raw_cache[stage]); + } + ds4_gpu_tensor_free(g->dspark_main_x); + ds4_gpu_tensor_free(g->dspark_stage0_proj); + ds4_gpu_tensor_free(g->dspark_stage0_packed); + ds4_gpu_tensor_free(g->dspark_target_hidden_batch); + ds4_gpu_tensor_free(g->dspark_target_hidden); + ds4_gpu_tensor_free(g->dspark_hc_mean_rows); + ds4_gpu_tensor_free(g->dspark_hc_mean_weights); + ds4_gpu_tensor_free(g->tp_logits_half); + free(g->cpu_router_norm); + memset(g, 0, sizeof(*g)); +} + +static bool metal_tensor_fill_f32(ds4_gpu_tensor *t, float v, uint64_t n) { + return ds4_gpu_tensor_fill_f32(t, v, n) != 0; +} + +/* ========================================================================= + * Directional Steering. + * ========================================================================= + * + * A steering file contains one normalized 4096-wide direction per layer. When + * enabled, the Metal graph edits selected block outputs in-place: + * + * y = y - scale * v * dot(v, y) + * + * Positive scales remove the represented direction from the activation. + * Negative scales add it. This is deliberately explicit and opt-in; with zero + * scales, the release graph does not allocate the direction tensor and follows + * the normal inference path. + */ + +/* directional_steering_dirs is Class P — replicated per tier. + * The same host directions buffer is written to every tier slot the engine's + * placement uses, then the load buffer is freed. Read-only after init, so + * the per-tier replicas stay byte-identical and never re-sync. */ +static bool metal_graph_load_directional_steering( + ds4_gpu_graph *g, + const char *path, + float attn_scale, + float ffn_scale) { + if (attn_scale == 0.0f && ffn_scale == 0.0f) return true; + + if (!path || !path[0]) { + fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); + return false; + } + + const uint64_t n = (uint64_t)DS4_N_LAYER * DS4_N_EMBD; + float *dirs = xmalloc((size_t)n * sizeof(dirs[0])); + bool ok = read_f32_binary_file(path, dirs, n); + if (ok) { + /* Replicate the directions buffer onto every Class P tier slot that + * has any other Class P scratch allocated (used_tier marker is the + * presence of g->cur_hc_by_tier[t]). Single-tier: only slot 0. */ + bool any = false; + for (int t = 0; ok && t < DS4_MAX_GPUS; t++) { + if (!g->cur_hc_by_tier[t]) continue; + g->directional_steering_dirs_by_tier[t] = + ds4_gpu_tensor_alloc_ptr_on(t, n * sizeof(dirs[0])); + ok = g->directional_steering_dirs_by_tier[t] != NULL && + ds4_gpu_tensor_write(g->directional_steering_dirs_by_tier[t], + 0, dirs, n * sizeof(dirs[0])) != 0; + if (ok) any = true; + } + if (ok && !any) { + /* No used tiers — graph not allocated yet. This shouldn't happen + * given the call site ordering, but bail cleanly. */ + ok = false; + } + } + free(dirs); + + if (!ok) { + fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); + return false; + } + g->directional_steering_attn_scale = attn_scale; + g->directional_steering_ffn_scale = ffn_scale; + fprintf(stderr, "ds4: directional steering enabled: %s attn=%g ffn=%g\n", + path, (double)attn_scale, (double)ffn_scale); + return true; +} + +static bool metal_graph_directional_steering_attn_enabled(const ds4_gpu_graph *g) { + return g && metal_graph_directional_steering_dirs(g) && + g->directional_steering_attn_scale != 0.0f; +} + +static bool metal_graph_directional_steering_ffn_enabled(const ds4_gpu_graph *g) { + return g && metal_graph_directional_steering_dirs(g) && + g->directional_steering_ffn_scale != 0.0f; +} + +static bool metal_graph_apply_directional_steering( + ds4_gpu_graph *g, + ds4_gpu_tensor *x, + uint32_t il, + uint32_t rows, + float scale) { + if (!g || !metal_graph_directional_steering_dirs(g) || scale == 0.0f) return true; + return ds4_gpu_directional_steering_project_tensor(x, + metal_graph_directional_steering_dirs(g), + il, + DS4_N_EMBD, + rows, + scale) != 0; +} + +static bool metal_graph_apply_directional_steering_attn( + ds4_gpu_graph *g, + ds4_gpu_tensor *x, + uint32_t il, + uint32_t rows) { + return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_attn_scale : 0.0f); +} + +static bool metal_graph_apply_directional_steering_ffn( + ds4_gpu_graph *g, + ds4_gpu_tensor *x, + uint32_t il, + uint32_t rows) { + return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f); +} + +static bool metal_graph_configure_dspark_capture( + ds4_gpu_graph *g, + const ds4_dspark_weights *dw) { + if (!g || !dw || dw->target_layer_count == 0) return true; + if (dw->target_layer_count > DS4_DSPARK_MAX_TARGET_LAYERS || + DS4_N_HC == 0 || + DS4_N_HC > DS4_MAX_HC) { + return false; + } + + g->dspark_hc_mean_weights = + ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); + g->dspark_hc_mean_rows = + ds4_gpu_tensor_alloc((uint64_t)g->prefill_cap * + DS4_N_HC * sizeof(float)); + g->dspark_target_hidden = + ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * + DS4_N_EMBD * sizeof(float)); + g->dspark_target_hidden_batch = + ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * + g->prefill_cap * + DS4_N_EMBD * sizeof(float)); + if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { + g->dspark_stage0_packed = + ds4_gpu_tensor_alloc(((uint64_t)dw->block_size + 1u) * + dw->target_layer_count * + DS4_N_EMBD * sizeof(float)); + } + g->dspark_stage0_proj = + ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + g->dspark_main_x = + ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + if (!g->dspark_hc_mean_weights || !g->dspark_hc_mean_rows || + !g->dspark_target_hidden || !g->dspark_target_hidden_batch || + !g->dspark_stage0_proj || !g->dspark_main_x) { + return false; + } + if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + g->dspark_draft_tokens = + ds4_gpu_tensor_alloc((uint64_t)dw->block_size * sizeof(int32_t)); + g->dspark_draft_hc = + ds4_gpu_tensor_alloc((uint64_t)dw->block_size * hc_dim * sizeof(float)); + g->dspark_target_hc = + ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->dspark_stage_input_hc = + ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * + hc_dim * sizeof(float)); + g->dspark_stage_output_hc = + ds4_gpu_tensor_alloc((uint64_t)dw->block_size * + hc_dim * sizeof(float)); + g->dspark_position_ids = + ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * + sizeof(int32_t)); + if (!g->dspark_draft_tokens || !g->dspark_draft_hc || + !g->dspark_target_hc || !g->dspark_stage_input_hc || + !g->dspark_stage_output_hc || !g->dspark_position_ids) { + return false; + } + if (dw->n_stages != 0 && g->raw_cap != 0) { + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + g->dspark_raw_cache[stage] = + ds4_gpu_tensor_alloc((uint64_t)g->raw_cap * + DS4_N_HEAD_DIM * sizeof(float)); + if (!g->dspark_raw_cache[stage]) return false; + } + g->dspark_cache_cap = g->raw_cap; + g->dspark_cache_start = 0; + g->dspark_cache_token_start = 0; + g->dspark_cache_len = 0; + } + g->dspark_block_size = dw->block_size; + } + + float mean[DS4_MAX_HC] = {0}; + const float inv_hc = 1.0f / (float)DS4_N_HC; + for (uint32_t i = 0; i < DS4_N_HC; i++) mean[i] = inv_hc; + if (ds4_gpu_tensor_write(g->dspark_hc_mean_weights, + 0, + mean, + (uint64_t)DS4_N_HC * sizeof(mean[0])) == 0) { + return false; + } + const uint64_t mean_rows_count = (uint64_t)g->prefill_cap * DS4_N_HC; + if (mean_rows_count == 0 || mean_rows_count > (uint64_t)SIZE_MAX / sizeof(float)) { + return false; + } + float *mean_rows = xmalloc((size_t)mean_rows_count * sizeof(mean_rows[0])); + for (uint64_t i = 0; i < mean_rows_count; i++) mean_rows[i] = inv_hc; + const bool mean_rows_ok = + ds4_gpu_tensor_write(g->dspark_hc_mean_rows, + 0, + mean_rows, + mean_rows_count * sizeof(mean_rows[0])) != 0; + free(mean_rows); + if (!mean_rows_ok) return false; + + g->dspark_target_layer_count = dw->target_layer_count; + memcpy(g->dspark_target_layers, + dw->target_layers, + (size_t)dw->target_layer_count * sizeof(g->dspark_target_layers[0])); + g->dspark_capture_mask = 0; + g->dspark_capture_checkpoint_len = 0; + g->dspark_capture_batch_mask = 0; + g->dspark_capture_batch_start = 0; + g->dspark_capture_batch_tokens = 0; + g->dspark_capture_valid = false; + g->dspark_capture_batch_valid = false; + g->dspark_capture_enabled = true; + return true; +} + +static uint64_t metal_graph_kv_cache_bytes_for_context(uint32_t ctx_size, uint32_t raw_cap) { + uint64_t bytes = (uint64_t)DS4_N_LAYER * + raw_cap * + DS4_N_HEAD_DIM * + sizeof(float); + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio == 0) continue; + const uint64_t comp_cap = (uint64_t)(ctx_size / ratio + 2u); + bytes += comp_cap * DS4_N_HEAD_DIM * + (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); + if (ratio == 4) { + bytes += comp_cap * DS4_N_INDEXER_HEAD_DIM * sizeof(float); + } + } + return bytes; +} + +static uint64_t metal_graph_context_bytes_for_kv_policy( + uint32_t ctx_size, + uint32_t raw_cap, + uint32_t prefill_cap, + uint64_t *kv_cache_bytes_out) { + uint32_t min_ratio = UINT32_MAX; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; + } + if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; + uint64_t comp_cap = (uint64_t)(ctx_size / min_ratio + 2u); + if (comp_cap < 2u) comp_cap = 2u; + const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context(ctx_size, raw_cap); + if (kv_cache_bytes_out) *kv_cache_bytes_out = kv_cache_bytes; + uint64_t bytes = kv_cache_bytes + + 2ull * comp_cap * prefill_cap * sizeof(float); + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + uint64_t attn_stage_cap = (uint64_t)(prefill_cap / min_ratio + 2u); + if (attn_stage_cap < 2u) attn_stage_cap = 2u; + bytes += attn_stage_cap * DS4_N_HEAD_DIM * sizeof(float); + } + return bytes; +} + +static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor_on( + bool managed, + int tier, + uint64_t bytes) { + if (g_n_gpus <= 1) { + return managed ? ds4_gpu_tensor_alloc_managed(bytes) + : ds4_gpu_tensor_alloc(bytes); + } + return managed ? ds4_gpu_tensor_alloc_managed_on(tier, bytes) + : ds4_gpu_tensor_alloc_ptr_on(tier, bytes); +} + +static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor(bool managed, uint64_t bytes) { + return metal_graph_alloc_kv_cache_tensor_on(managed, 0, bytes); +} + +/* ========================================================================= + * Metal Diagnostic Dump Hooks. + * ========================================================================= + * + * The release path calls these after important stages, but they are no-ops + * unless DS4_METAL_GRAPH_DUMP_PREFIX or DS4_ROCM_GRAPH_DUMP_PREFIX is set. + * Dumping synchronizes and restarts the command batch, so it is intentionally + * isolated here. + */ + +typedef struct { + int init; + const char *prefix; + const char *name; + int layer_set; + uint32_t layer; + int pos_set; + uint32_t pos; +} metal_graph_debug_config; + +static const metal_graph_debug_config *metal_graph_debug_get_config(void) { + static metal_graph_debug_config cfg; + if (!cfg.init) { + cfg.init = 1; + cfg.prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", + "DS4_METAL_GRAPH_DUMP_PREFIX"); + if (cfg.prefix && !cfg.prefix[0]) cfg.prefix = NULL; + cfg.name = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_NAME", + "DS4_METAL_GRAPH_DUMP_NAME"); + if (cfg.name && !cfg.name[0]) cfg.name = NULL; + + const char *layer_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_LAYER", + "DS4_METAL_GRAPH_DUMP_LAYER"); + if (layer_env && layer_env[0] && strcmp(layer_env, "all") != 0) { + cfg.layer_set = 1; + cfg.layer = (uint32_t)strtoul(layer_env, NULL, 10); + } + + const char *pos_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_POS", + "DS4_METAL_GRAPH_DUMP_POS"); + if (pos_env && pos_env[0]) { + cfg.pos_set = 1; + cfg.pos = (uint32_t)strtoul(pos_env, NULL, 10); + } + } + return &cfg; +} + +static const char *metal_graph_debug_prefix_for(const char *name, uint32_t il, uint32_t pos) { + const metal_graph_debug_config *cfg = metal_graph_debug_get_config(); + if (!cfg->prefix) return NULL; + if (cfg->name && strstr(cfg->name, name) == NULL) return NULL; + if (cfg->layer_set && cfg->layer != il) return NULL; + if (cfg->pos_set && cfg->pos != pos) return NULL; + return cfg->prefix; +} + +static bool metal_graph_debug_wants(const char *name, uint32_t il, uint32_t pos) { + return metal_graph_debug_prefix_for(name, il, pos) != NULL; +} + +static void metal_graph_debug_dump_tensor( + const char *name, + ds4_gpu_tensor *t, + uint64_t n_f32, + uint32_t il, + uint32_t pos) { + const char *prefix = metal_graph_debug_prefix_for(name, il, pos); + if (glm_graph_env_present("DS4_ROCM_GRAPH_DUMP_TRACE", + "DS4_METAL_GRAPH_DUMP_TRACE")) + fprintf(stderr, "ds4: dump? name=%s il=%u pos=%u t=%p n=%llu wants=%d\n", + name, il, pos, (void *)t, (unsigned long long)n_f32, + metal_graph_debug_wants(name, il, pos)); + if (!t || n_f32 == 0 || !metal_graph_debug_wants(name, il, pos)) return; + + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); + return; + } + + float *buf = xmalloc((size_t)n_f32 * sizeof(buf[0])); + if (ds4_gpu_tensor_read(t, 0, buf, n_f32 * sizeof(buf[0])) != 0) { + char path[1024]; + snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); + if (write_f32_binary_file(path, buf, n_f32)) { + fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); + } + } + free(buf); + + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); + } +} + +static void metal_graph_debug_dump_f16_tensor( + const char *name, + ds4_gpu_tensor *t, + uint64_t n_f16, + uint32_t il, + uint32_t pos) { + const char *prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", + "DS4_METAL_GRAPH_DUMP_PREFIX"); + if (!t || n_f16 == 0 || !metal_graph_debug_wants(name, il, pos)) return; + + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); + return; + } + + uint16_t *hbuf = xmalloc((size_t)n_f16 * sizeof(hbuf[0])); + float *fbuf = xmalloc((size_t)n_f16 * sizeof(fbuf[0])); + if (ds4_gpu_tensor_read(t, 0, hbuf, n_f16 * sizeof(hbuf[0])) != 0) { + for (uint64_t i = 0; i < n_f16; i++) fbuf[i] = f16_to_f32(hbuf[i]); + char path[1024]; + snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); + if (write_f32_binary_file(path, fbuf, n_f16)) { + fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); + } + } + free(fbuf); + free(hbuf); + + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); + } +} + +static void metal_graph_debug_dump_i32_tensor( + const char *name, + ds4_gpu_tensor *t, + uint64_t n_i32, + uint32_t il, + uint32_t pos) { + if (!t || n_i32 == 0) return; + const char *prefix = metal_graph_debug_prefix_for(name, il, pos); + if (!prefix) return; + + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); + return; + } + + int32_t *buf = xmalloc((size_t)n_i32 * sizeof(buf[0])); + if (ds4_gpu_tensor_read(t, 0, buf, n_i32 * sizeof(buf[0])) != 0) { + char path[1024]; + snprintf(path, sizeof(path), "%s_%s-%u_pos%u.i32", prefix, name, il, pos); + FILE *fp = fopen(path, "wb"); + if (fp) { + if (fwrite(buf, sizeof(buf[0]), (size_t)n_i32, fp) == (size_t)n_i32) { + fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); + } + fclose(fp); + } + } + free(buf); + + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); + } +} + +static bool metal_graph_needs_ffn_out(const ds4_gpu_graph *g, uint32_t il, uint32_t pos) { + return metal_graph_directional_steering_ffn_enabled(g) || + g->materialize_ffn_out || + metal_graph_debug_wants("ffn_out", il, pos); +} + +/* tier-aware lazy allocator. The Class P ffn_out scratch is + * created on demand the first time a layer that materializes ffn_out runs + * on a tier; subsequent visits to the same tier reuse the existing slot. + * Single-tier paths: active_tier == 0 always, behavior unchanged. */ +static bool metal_graph_ensure_ffn_out(ds4_gpu_graph *g) { + const int t = g->active_tier; + if (!g->ffn_out_by_tier[t]) { + g->ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( + t, (uint64_t)DS4_N_EMBD * sizeof(float)); + } + return g->ffn_out_by_tier[t] != NULL; +} + +static bool metal_graph_ensure_batch_ffn_out_on(ds4_gpu_graph *g, int t) { + if (t < 0 || t >= DS4_MAX_GPUS) return false; + if (!g->batch_ffn_out_by_tier[t]) { + g->batch_ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( + t, (uint64_t)g->prefill_cap * DS4_N_EMBD * sizeof(float)); + } + return g->batch_ffn_out_by_tier[t] != NULL; +} + +static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) { + return metal_graph_ensure_batch_ffn_out_on(g, g->active_tier); +} + +static bool metal_graph_tp_env_flag(const char *name, bool dflt) { + const char *env = getenv(name); + if (!env || !env[0]) return dflt; + return strcmp(env, "0") != 0; +} + +static bool metal_graph_cuda_tp_attn_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN", true); +#endif +} + +static bool metal_graph_cuda_tp_attn_peer_read_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_PEER_READ", true); +#endif +} + +static bool metal_graph_cuda_tp_attn_heads_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_HEADS", false); +#endif +} + +static bool metal_graph_cuda_tp_attn_cache_dup_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_CACHE_DUP", false); +#endif +} + +static bool metal_graph_cuda_tp_moe_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE", true); +#endif +} + +static bool metal_graph_cuda_tp_ep_pack_exact_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_PACK_EXACT", true); +#endif +} + +static bool metal_graph_cuda_tp_ep_direct_return_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DIRECT_RETURN", true); +#endif +} + +static bool metal_graph_cuda_tp_ep_delay_reduce_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DELAY_REDUCE", true); +#endif +} + +static bool metal_graph_cuda_tp_ep_fused_hc_reduce_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_HC_REDUCE", true); +#endif +} + +static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_fused_shared_mid_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_SHARED_MID", true); +#endif +} + +static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_balanced_shared_mid_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag( + "DS4_CUDA_TP_EP_BALANCED_SHARED_MID", true); +#endif +} + +static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_dual_prequant_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag( + "DS4_CUDA_TP_EP_DUAL_PREQUANT", true); +#endif +} + +static bool metal_graph_cuda_tp_moe_delay_reduce_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_DELAY_REDUCE", true); +#endif +} + +static bool metal_graph_cuda_tp_moe_pack_handoff_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PACK", false); +#endif +} + +static bool metal_graph_cuda_tp_moe_copy3_handoff_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_COPY3", false); +#endif +} + +static bool metal_graph_cuda_tp_moe_peer_read_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_READ", false); +#endif +} + +static bool metal_graph_cuda_tp_moe_peer_router_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_ROUTER", false); +#endif +} + +static bool metal_graph_cuda_tp_shared_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED", false); +#endif +} + +static bool metal_graph_cuda_tp_shared_fold_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED_FOLD", true); +#endif +} + +static bool metal_graph_cuda_tp_q_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_Q", false); +#endif +} + +static bool metal_graph_cuda_tp_output_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_OUTPUT", true); +#endif +} + +static bool metal_graph_cuda_greedy_split_top1_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLIT_TOP1"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLIT_TOP1", false); +#endif +} + +static bool metal_graph_cuda_output_fused_top1_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_OUTPUT_FUSED_TOP1", false); +#endif +} + +static bool metal_graph_cuda_verify_decode2_split_top1_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_VERIFY_DECODE2_SPLIT_TOP1"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_VERIFY_DECODE2_SPLIT_TOP1", false); +#endif +} + +static bool metal_graph_cuda_greedy_splitkv_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV", false); +#endif +} + +static bool metal_graph_cuda_greedy_vec4_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_VEC4", false); +#endif +} + +static bool metal_graph_cuda_splitkv_spec_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC", false); +#endif +} + +static bool metal_graph_cuda_splitkv_spec_toponly_row0_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_TOPONLY_ROW0"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_TOPONLY_ROW0", false); +#endif +} + +static bool metal_graph_cuda_splitkv_spec_batch_verify_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_BATCH_VERIFY"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_BATCH_VERIFY", false); +#endif +} + +static float metal_graph_cuda_greedy_vec4_margin_threshold(void) { +#if defined(__APPLE__) + return 0.0f; +#else + const char *env = getenv("DS4_CUDA_GREEDY_VEC4_MARGIN"); + if (env && env[0]) { + char *end = NULL; + double v = strtod(env, &end); + while (end && isspace((unsigned char)*end)) end++; + if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { + return (float)v; + } + fprintf(stderr, + "ds4: invalid DS4_CUDA_GREEDY_VEC4_MARGIN=%s; using 0.25\n", + env); + } + return 0.25f; +#endif +} + +static bool metal_graph_cuda_greedy_vec4_fallback_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4_FALLBACK"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_cuda_greedy_vec4_margin_threshold() > 0.0f; +#endif +} + +static float metal_graph_cuda_greedy_splitkv_margin_threshold(void) { +#if defined(__APPLE__) + return 0.0f; +#else + const char *env = getenv("DS4_CUDA_GREEDY_SPLITKV_MARGIN"); + if (env && env[0]) { + char *end = NULL; + double v = strtod(env, &end); + while (end && isspace((unsigned char)*end)) end++; + if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { + return (float)v; + } + fprintf(stderr, + "ds4: invalid DS4_CUDA_GREEDY_SPLITKV_MARGIN=%s; using 0.25\n", + env); + } + return 0.25f; +#endif +} + +static bool metal_graph_cuda_greedy_splitkv_fallback_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_FALLBACK"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_cuda_greedy_splitkv_margin_threshold() > 0.0f; +#endif +} + +static bool metal_graph_cuda_greedy_splitkv_top2_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_TOP2"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TOP2", true); +#endif +} + +static bool metal_graph_cuda_greedy_splitkv_trust_replay_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TRUST_REPLAY", false); +#endif +} + +static bool metal_graph_cuda_greedy_splitkv_pair_replay_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_PAIR_REPLAY"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_PAIR_REPLAY", false); +#endif +} + +static DS4_MAYBE_UNUSED uint32_t metal_graph_cuda_greedy_max_segment(const char *name) { + const char *env = getenv(name); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + while (end && isspace((unsigned char)*end)) end++; + if (end != env && end && *end == '\0' && v <= INT32_MAX) { + return (uint32_t)v; + } + fprintf(stderr, + "ds4: invalid %s=%s; expected 0..%d, using disabled\n", + name, + env, + INT32_MAX); + } + return 0; +} + +static uint32_t metal_graph_cuda_greedy_splitkv_max_segment(void) { +#if defined(__APPLE__) + return 0; +#else + return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_SPLITKV_MAX_SEGMENT"); +#endif +} + +static uint32_t metal_graph_cuda_greedy_vec4_max_segment(void) { +#if defined(__APPLE__) + return 0; +#else + return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_VEC4_MAX_SEGMENT"); +#endif +} + +static uint32_t metal_graph_cuda_greedy_splitkv_min_score(void) { + const char *env = getenv("DS4_CUDA_SPLITKV_MIN_SCORE"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + while (end && isspace((unsigned char)*end)) end++; + if (end != env && end && *end == '\0' && v <= UINT32_MAX) { + return (uint32_t)v; + } + } + return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_DECODE", false) ? 0u : 512u; +} + +static bool metal_graph_cuda_q_norm_rope_fuse_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_Q_NORM_ROPE_FUSE", true); +#endif +} + +static bool metal_graph_cuda_qkv_kv_rope_fuse_requested(void) { +#if defined(__APPLE__) + return false; +#else + const char *no = getenv("DS4_CUDA_NO_QKV_KV_ROPE_FUSE"); + if (no && no[0] && strcmp(no, "0") != 0) return false; + if (getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL) return false; + return metal_graph_tp_env_flag("DS4_CUDA_QKV_KV_ROPE_FUSE", true); +#endif +} + +static bool metal_graph_cuda_tp_prefill_ffn_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_FFN", true); +#endif +} + +static bool metal_graph_cuda_tp_prefill_attn_output_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_ATTN_OUTPUT", true); +#endif +} + +static bool metal_graph_cuda_prefill_pipeline_requested(const ds4_gpu_graph *g) { +#if defined(__APPLE__) + (void)g; + return false; +#else + const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE"); + if (env && env[0]) return strcmp(env, "0") != 0; + return g && g->cuda_tp_decode; +#endif +} + +static bool metal_graph_cuda_prefill_pipeline_q8_cache_requested(void) { +#if defined(__APPLE__) + return false; +#else + return metal_graph_tp_env_flag("DS4_CUDA_PREFILL_PIPELINE_Q8_CACHE", false); +#endif +} + +static uint32_t metal_graph_cuda_prefill_pipeline_microbatch(void) { + const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && v > 0 && v <= UINT32_MAX) return (uint32_t)v; + } + return 512; +} + +static int metal_graph_cuda_tp_partner_tier(int tier) { + if (g_n_gpus < 2 || (g_n_gpus & 1) != 0) return -1; + const int half = g_n_gpus / 2; + if (tier < 0 || tier >= half) return -1; + return tier + half; +} + +static uint32_t metal_graph_cuda_tp_output_tiers( + const ds4_gpu_graph *g, + int tiers[DS4_MAX_GPUS]) { + if (!g) return 0; + return metal_graph_cuda_tp_output_tiers_for_head(g->head_tier, + g->cuda_tp_output, + g_n_gpus, + tiers); +} + +static uint64_t metal_graph_q8_0_row_bytes(uint64_t in_dim) { + return ((in_dim + 31u) / 32u) * 34u; +} + +/* ========================================================================= + * Metal Release Graph Allocation. + * ========================================================================= */ + +/* Allocate the Metal graph state for a chosen raw-cache capacity. The model + * weights are not copied here; tensors reference the mapped GGUF. + * + * tier-aware per-layer allocation. + * placement: when non-NULL, an array of DS4_N_LAYER + 2 logical tiers + * (embedding, per-layer..., head). The per-layer KV / state allocations + * in this function use placement[il + 1] as the home tier for each + * layer il. When NULL (single-tier callers, diagnostic paths), all + * per-layer allocations land on tier 0 — byte-equivalent to legacy. + * + * Single-tier (g_n_gpus <= 1) is byte-equivalent regardless of placement, + * because metal_graph_alloc_kv_cache_tensor_on short-circuits to the + * legacy 1-arg helpers when g_n_gpus <= 1. */ +static bool metal_graph_alloc_raw_cap( + ds4_gpu_graph *g, + const ds4_weights *weights, + const ds4_layer_weights *layer, + uint32_t raw_cap, + uint32_t ctx_size, + uint32_t prefill_cap, + bool enable_mtp, + const int *placement, + bool cuda_tensor_parallel, + const ds4_gpu_graph *shared_prefill_workspace) { + const int saved_dspark_exec_tier = g->dspark_exec_tier; + memset(g, 0, sizeof(*g)); + g->dspark_exec_tier = saved_dspark_exec_tier; + g->owns_prefill_workspace = shared_prefill_workspace == NULL; + g->cpu_router_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])); + g->active_tier = placement ? -1 : 0; + /* cache placement on the graph so the dispatch loops can + * walk it without threading the engine pointer through every + * kernel-dispatch wrapper. NULL in single-tier callers (placement + * was already NULL on entry). */ + g->placement = placement; + g->cuda_tp_decode = placement && cuda_tensor_parallel; + g->cuda_tp_attn = g->cuda_tp_decode && metal_graph_cuda_tp_attn_requested(); + g->cuda_tp_attn_peer_read = metal_graph_cuda_tp_attn_peer_read_requested(); + g->cuda_tp_attn_heads = g->cuda_tp_decode && metal_graph_cuda_tp_attn_heads_requested(); + g->cuda_tp_attn_cache_dup = g->cuda_tp_attn_heads && + metal_graph_cuda_tp_attn_cache_dup_requested(); + g->cuda_tp_moe = g->cuda_tp_decode && metal_graph_cuda_tp_moe_requested(); + g->cuda_tp_ep = g->cuda_tp_moe && cuda_tensor_parallel; + g->cuda_tp_ep_pack_exact = + g->cuda_tp_ep && metal_graph_cuda_tp_ep_pack_exact_requested(); + g->cuda_tp_moe_delay_reduce = metal_graph_cuda_tp_moe_delay_reduce_requested(); + g->cuda_tp_moe_copy3_handoff = metal_graph_cuda_tp_moe_copy3_handoff_requested(); + g->cuda_tp_moe_pack_handoff = metal_graph_cuda_tp_moe_pack_handoff_requested(); + g->cuda_tp_moe_peer_read = metal_graph_cuda_tp_moe_peer_read_requested(); + g->cuda_tp_moe_peer_router = metal_graph_cuda_tp_moe_peer_router_requested(); + g->cuda_tp_shared = g->cuda_tp_decode && metal_graph_cuda_tp_shared_requested(); + g->cuda_tp_shared_fold = metal_graph_cuda_tp_shared_fold_requested(); + g->cuda_tp_q = g->cuda_tp_decode && metal_graph_cuda_tp_q_requested(); + g->cuda_tp_output = g->cuda_tp_decode && metal_graph_cuda_tp_output_requested(); + g->cuda_tp_prefill_ffn = g->cuda_tp_decode && metal_graph_cuda_tp_prefill_ffn_requested(); + g->cuda_tp_prefill_attn_output = + g->cuda_tp_decode && metal_graph_cuda_tp_prefill_attn_output_requested(); + g->cuda_q_norm_rope_fuse = metal_graph_cuda_q_norm_rope_fuse_requested(); + g->cuda_qkv_kv_rope_fuse = metal_graph_cuda_qkv_kv_rope_fuse_requested(); + g->cuda_qkv_pair = getenv("DS4_CUDA_NO_QKV_PAIR") == NULL; + g->cuda_tp_attn_out_hc_fuse = + getenv("DS4_CUDA_TP_ATTN_OUT_HC_FUSE") != NULL && + getenv("DS4_CUDA_NO_TP_ATTN_OUT_HC_FUSE") == NULL; + g->shared_gate_up_swiglu_fuse = + getenv("DS4_METAL_DISABLE_SHARED_GATE_UP_SWIGLU_FUSION") == NULL; + g->decode_stage_profile = getenv("DS4_METAL_DECODE_STAGE_PROFILE") != NULL; + g->decode_index_stage_profile = getenv("DS4_METAL_INDEXER_STAGE_PROFILE") != NULL; + g->output_stage_profile = getenv("DS4_METAL_OUTPUT_STAGE_PROFILE") != NULL; + const bool enable_splitkv_spec = metal_graph_cuda_splitkv_spec_requested(); + const bool enable_splitkv_batch_verify = + enable_splitkv_spec && metal_graph_cuda_splitkv_spec_batch_verify_requested(); + const bool enable_spec_logits = enable_mtp || enable_splitkv_batch_verify; + const bool enable_prefix1_snapshot = enable_mtp || enable_splitkv_spec; + const bool enable_frontier_snapshot = + enable_mtp || + enable_splitkv_spec || + (metal_graph_cuda_greedy_splitkv_requested() && + metal_graph_cuda_greedy_splitkv_fallback_requested()) || + (metal_graph_cuda_greedy_vec4_requested() && + metal_graph_cuda_greedy_vec4_fallback_requested()); + if (g->cuda_tp_decode && metal_graph_cuda_tp_partner_tier(0) < 0) { + fprintf(stderr, + "ds4: CUDA tensor parallelism requires an even multi-GPU placement; " + "have %d GPU tiers\n", + g_n_gpus); + metal_graph_free(g); + return false; + } + if (g->cuda_tp_ep && + (g_ds4_shape.family != DS4_MODEL_FAMILY_DEEPSEEK4 || + (DS4_N_EXPERT & 1u) != 0u)) { + fprintf(stderr, + "ds4: CUDA tensor parallelism requires an even-expert DeepSeek model\n"); + metal_graph_free(g); + return false; + } + if (g->cuda_tp_ep) { + fprintf(stderr, + "ds4: CUDA routed MoE expert ownership enabled " + "(half-resident decode and prefill)\n"); + } + g->mtp_enabled = enable_mtp; + if (raw_cap == 0) raw_cap = 1; + if (ctx_size == 0) ctx_size = raw_cap; + if (prefill_cap == 0) prefill_cap = 1; + uint32_t raw_window = DS4_N_SWA; + if (raw_window > ctx_size) raw_window = ctx_size; + if (raw_window == 0) raw_window = 1; + if (raw_cap < raw_window) raw_cap = raw_window; + if (raw_cap > ctx_size) raw_cap = ctx_size; + if (raw_cap == 0) raw_cap = 1; + g->raw_cap = raw_cap; + g->raw_window = raw_window; + g->prefill_cap = prefill_cap; + uint32_t min_ratio = UINT32_MAX; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) continue; + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; + } + if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; + g->comp_cap = ctx_size / min_ratio + 2u; + if (g->comp_cap < 2u) g->comp_cap = 2u; + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + g->attn_comp_stage_cap = prefill_cap / min_ratio + 2u; + if (g->attn_comp_stage_cap < 2u) g->attn_comp_stage_cap = 2u; + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) { + g->layer_comp_cap[il] = 0; + continue; + } + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio == 0) { + g->layer_comp_cap[il] = 0; + } else { + g->layer_comp_cap[il] = ctx_size / ratio + 2u; + if (g->layer_comp_cap[il] < 2u) g->layer_comp_cap[il] = 2u; + } + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; + const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); + const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; + const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; + /* Distributed coordinators do not normally own the output head. The + * logits workspace still has a fixed model-vocabulary shape, while the + * actual head is encoded only on a node that bound its tensors. */ + const uint64_t vocab_dim = + weights->output ? weights->output->dim[1] : DS4_N_VOCAB; + const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM + ? DS4_N_HEAD_DIM + : DS4_N_INDEXER_HEAD_DIM); + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + const uint64_t pc = prefill_cap; + uint64_t kv_cache_bytes = 0; + const uint64_t context_bytes = + metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes); + const bool managed_kv_cache = + ds4_gpu_should_use_managed_kv_cache(kv_cache_bytes, context_bytes) != 0; + if (managed_kv_cache) { + /* + * CUDA device allocations are fastest, but a million-token KV cache is + * large enough to starve DGX Spark's unified CPU/GPU memory once the + * model cache and driver allocations are present. For this one + * long-lived cache class, managed memory restores the old demand-paged + * behavior. It can be slower, but it keeps oversized contexts from + * turning memory pressure into a machine-wide lockup. + */ + fprintf(stderr, + "ds4: CUDA using managed KV cache for ctx=%u " + "(kv cache %.2f GiB, context buffers %.2f GiB); " + "this may degrade performance but is needed for very large contexts\n", + ctx_size, + (double)kv_cache_bytes / 1073741824.0, + (double)context_bytes / 1073741824.0); + } + + /* Class P decode HC scratch — replicated across every tier + * the placement uses (per-tier kernel-scratch). Single-tier path + * (placement == NULL) collapses to tier 0 only; _ptr_on(0, ...) short- + * circuits to legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte- + * equivalent. The hc_pre/hc_post/hc_comb buffers are VIEWS of hc_split + * and therefore allocated per tier alongside their parent. */ + bool used_tier[DS4_MAX_GPUS] = {0}; + used_tier[0] = true; /* single-tier baseline always uses tier 0 */ + if (placement) { + for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER + 2u; i++) { + const int p = placement[i]; + if (p >= 0 && p < DS4_MAX_GPUS) used_tier[p] = true; + } + } + if (g->cuda_tp_decode) { + const int half = g_n_gpus / 2; + for (int t = half; t < g_n_gpus; t++) { + if (used_tier[t]) { + fprintf(stderr, + "ds4: CUDA tensor parallelism expects layer homes in lower-half " + "tiers; placement already uses tier %d\n", + t); + metal_graph_free(g); + return false; + } + } + for (int t = 0; t < half; t++) { + if (used_tier[t]) used_tier[t + half] = true; + } + fprintf(stderr, + "ds4: CUDA decode TP enabled: pairing lower-half tiers with " + "upper-half tiers\n"); + } + for (int t = 0; t < DS4_MAX_GPUS; t++) { + if (!used_tier[t]) continue; + g->cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); + g->flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); + g->hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); + g->hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); + g->hc_pre_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], + 0, + (uint64_t)DS4_N_HC * sizeof(float)); + g->hc_post_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], + (uint64_t)DS4_N_HC * sizeof(float), + (uint64_t)DS4_N_HC * sizeof(float)); + g->hc_comb_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], + 2ull * DS4_N_HC * sizeof(float), + (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); + g->attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); + g->qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); + g->q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_dim * sizeof(float)); + g->kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + g->kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + } + bool state_init_ok = true; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + /* A distributed process owns only its bound layer slice. Persistent + * KV state must follow that ownership just like the model tensors; + * allocating every model layer here defeats split-model residency. */ + if (!weights_layer_has_required(&weights->layer[il], il)) continue; + /* per-layer Class L allocations land on the layer's + * home tier. placement is NULL on single-tier / diagnostic paths + * (all-tier-0); non-NULL on the engine path that opted into + * multi-tier. layer_tier == 0 in single-tier mode is the + * byte-equivalent path through metal_graph_alloc_kv_cache_tensor_on + * and ds4_gpu_tensor_alloc_ptr_on. */ + const int layer_tier = placement ? placement[il + 1] : 0; + g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor_on( + managed_kv_cache, + layer_tier, + (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); + const int layer_tp_partner = g->cuda_tp_attn_cache_dup + ? metal_graph_cuda_tp_partner_tier(layer_tier) : -1; + if (layer_tp_partner >= 0) { + g->layer_raw_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( + managed_kv_cache, + layer_tp_partner, + (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); + } + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0) { + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; + const uint64_t attn_rows = (uint64_t)coff * ratio; + g->layer_attn_comp_cache[il] = metal_graph_alloc_kv_cache_tensor_on( + managed_kv_cache, + layer_tier, + (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * + (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); + if (layer_tp_partner >= 0) { + g->layer_attn_comp_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( + managed_kv_cache, + layer_tp_partner, + (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * + (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); + } + g->layer_attn_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + g->layer_attn_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + if (enable_frontier_snapshot) { + g->spec_attn_state_kv[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + g->spec_attn_state_score[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + if (enable_prefix1_snapshot) { + g->spec_prefix1_attn_state_kv[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + g->spec_prefix1_attn_state_score[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); + } + } + if (g->layer_attn_state_kv[il]) { + state_init_ok = state_init_ok && + metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows); + } + if (g->layer_attn_state_score[il]) { + state_init_ok = state_init_ok && + metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows); + } + + if (ratio == 4) { + const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM; + const uint64_t index_rows = (uint64_t)coff * ratio; + g->layer_index_comp_cache[il] = metal_graph_alloc_kv_cache_tensor_on( + managed_kv_cache, + layer_tier, + (uint64_t)g->layer_comp_cap[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + g->layer_index_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + g->layer_index_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + if (enable_frontier_snapshot) { + g->spec_index_state_kv[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + g->spec_index_state_score[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + if (enable_prefix1_snapshot) { + g->spec_prefix1_index_state_kv[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + g->spec_prefix1_index_state_score[il] = + ds4_gpu_tensor_alloc_ptr_on(layer_tier, index_width * index_rows * sizeof(float)); + } + } + if (g->layer_index_state_kv[il]) { + state_init_ok = state_init_ok && + metal_tensor_fill_f32(g->layer_index_state_kv[il], 0.0f, index_width * index_rows); + } + if (g->layer_index_state_score[il]) { + state_init_ok = state_init_ok && + metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows); + } + } + } + } + /* Class P per-layer decode scratch + routed-expert state — + * replicated across every used tier. ffn_out is lazily allocated by + * metal_graph_ensure_ffn_out (per-tier on first touch). */ + for (int t = 0; t < DS4_MAX_GPUS; t++) { + if (!used_tier[t]) continue; + g->comp_kv_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, comp_width_max * sizeof(float)); + g->comp_sc_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, comp_width_max * sizeof(float)); + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + /* Upstream's F16-compressed attn staging buffer. Only allocated when + * the F16-cache mode is enabled (the non-F16 path stages in-place). */ + g->attn_comp_stage_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)g->attn_comp_stage_cap * DS4_N_HEAD_DIM * sizeof(float)); + } + g->indexer_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, indexer_q_dim * sizeof(float)); + g->indexer_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float)); + g->indexer_scores_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)g->comp_cap * pc * sizeof(float)); + g->comp_mask_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)g->comp_cap * pc * sizeof(float)); + g->comp_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)(DS4_N_INDEXER_TOP_K ? DS4_N_INDEXER_TOP_K : 1u) * pc * sizeof(uint32_t)); + g->heads_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_dim * sizeof(float)); + g->attn_low_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, low_dim * sizeof(float)); + g->attn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->after_attn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); + g->ffn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->ffn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->shared_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); + g->shared_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); + g->shared_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, shared_dim * sizeof(float)); + g->shared_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->router_logits_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT * sizeof(float)); + g->router_probs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT * sizeof(float)); + g->router_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT_USED * sizeof(int)); + g->router_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, DS4_N_EXPERT_USED * sizeof(float)); + g->routed_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->routed_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->routed_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->routed_down_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); + g->routed_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); + if (g->cuda_tp_decode) { + g->tp_peer_tmp_by_tier[t] = + ds4_gpu_tensor_alloc_ptr_on(t, DS4_CUDA_TP_PEER_TMP_BYTES); + } + g->after_ffn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); + } + /* Class H — head_tier captured from placement[DS4_N_LAYER + 1] + * (or 0 in single-tier / diagnostic paths). Output-head tensors and the + * final logits buffer allocate on head_tier only; other tier slots stay + * NULL. The _ptr_on(0, ...) path short-circuits to the legacy + * ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte-equivalent. */ + g->head_tier = placement ? placement[DS4_N_LAYER + 1] : 0; + int output_tp_tiers[DS4_MAX_GPUS] = {0}; + const uint32_t output_tp_ways = g->cuda_tp_output + ? metal_graph_cuda_tp_output_tiers(g, output_tp_tiers) : 0; + if (g->cuda_tp_output && output_tp_ways < 2u) { + fprintf(stderr, + "ds4: CUDA output TP requires output head tier %d to be in " + "the lower half of the CUDA placement\n", + g->head_tier); + metal_graph_free(g); + return false; + } + uint64_t output_logits_elems = vocab_dim; + if (enable_spec_logits && output_tp_ways >= 2u) { + const uint64_t max_shard_vocab = + (vocab_dim + output_tp_ways - 1u) / output_tp_ways; + const uint64_t spec_shard_elems = + (uint64_t)DS4_DSPARK_MAX_BLOCK_SIZE * max_shard_vocab; + if (spec_shard_elems > output_logits_elems) { + output_logits_elems = spec_shard_elems; + } + } + g->output_pre_by_tier[g->head_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_HC * sizeof(float)); + g->output_weights_by_tier[g->head_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_HC * sizeof(float)); + g->output_embd_by_tier[g->head_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->output_norm_by_tier[g->head_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->head_tier, (uint64_t)DS4_N_EMBD * sizeof(float)); + g->logits_by_tier[g->head_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->head_tier, + output_logits_elems * sizeof(float)); + for (uint32_t i = 1; i < output_tp_ways; i++) { + const int t = output_tp_tiers[i]; + if (t < 0 || t >= DS4_MAX_GPUS || t == g->head_tier) continue; + g->output_norm_by_tier[t] = + ds4_gpu_tensor_alloc_ptr_on(t, + (uint64_t)DS4_N_EMBD * sizeof(float)); + g->logits_by_tier[t] = + ds4_gpu_tensor_alloc_ptr_on(t, + output_logits_elems * sizeof(float)); + } + /* + * MTP is deliberately outside the normal graph footprint. A session that + * does not opt in with --mtp must allocate and execute exactly the same + * buffers as the plain decoder: no support-model mapping, no draft logits, + * and no MTP scratch hidden behind otherwise unused tensors. + */ + if (enable_mtp) { + g->mtp_embed = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + g->mtp_enorm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + g->mtp_eproj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + g->mtp_eproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_hnorm_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_hproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor( + managed_kv_cache, + (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); + g->mtp_n_raw = 0; + } + if (enable_spec_logits) { + const int spec_tier = + g->dspark_exec_tier > 0 && g->dspark_exec_tier < DS4_MAX_GPUS + ? g->dspark_exec_tier : 0; + g->spec_logits = spec_tier + ? ds4_gpu_tensor_alloc_ptr_on(spec_tier, (uint64_t)16 * DS4_N_VOCAB * sizeof(float)) + : ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); + } + + /* Class E — emb_tier captured from placement[0] (or 0 in + * single-tier / diagnostic paths). _ptr_on(0, ...) short-circuits to the + * legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte-equivalent. */ + g->emb_tier = placement ? placement[0] : 0; + /* Class P chunked-prefill batch scratch — replicated across + * every used tier. The cur/next pair (batch_cur_hc / batch_next_hc) is + * ping-ponged per layer step on each tier; tier transitions copy via + * ds4_gpu_tensor_copy_xdev (handled in B6). batch_ffn_out is lazily + * allocated by metal_graph_ensure_batch_ffn_out (per-tier on first touch) + * and included in the CUDA scratch estimate because TP prefill can use it + * as the combined routed+shared FFN buffer. */ + if (shared_prefill_workspace) { + if (shared_prefill_workspace->prefill_cap < prefill_cap || + shared_prefill_workspace->emb_tier != g->emb_tier) { + fprintf(stderr, + "ds4: shared prefill workspace is incompatible " + "(capacity %u/%u, embedding tier %d/%d)\n", + shared_prefill_workspace->prefill_cap, + prefill_cap, + shared_prefill_workspace->emb_tier, + g->emb_tier); + } else { + metal_graph_copy_prefill_workspace_pointers( + g, shared_prefill_workspace); + } + } else { + g->prefill_tokens_by_tier[g->emb_tier] = + ds4_gpu_tensor_alloc_ptr_on(g->emb_tier, pc * sizeof(int32_t)); + for (int t = 0; t < DS4_MAX_GPUS; t++) { + if (!used_tier[t]) continue; + g->batch_cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); + g->batch_next_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); + g->batch_flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); + g->batch_hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); + g->batch_hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); + g->batch_attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); + g->batch_qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); + g->batch_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); + g->batch_kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); + g->batch_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); + g->batch_comp_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); + g->batch_comp_sc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); + g->batch_indexer_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * indexer_q_dim * sizeof(float)); + g->batch_indexer_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_INDEXER_HEAD * sizeof(float)); + g->batch_heads_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); + g->batch_attn_low_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * low_dim * sizeof(float)); + g->batch_attn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_group_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * group_dim * sizeof(float)); + g->batch_low_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_LORA_O * sizeof(float)); + g->batch_after_attn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); + g->batch_ffn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_ffn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_shared_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); + g->batch_shared_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); + g->batch_shared_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); + g->batch_shared_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + g->batch_router_logits_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); + g->batch_router_probs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); + g->batch_router_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(int)); + g->batch_router_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(float)); + g->batch_routed_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->batch_routed_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->batch_routed_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + g->batch_routed_down_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); + g->batch_routed_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); + } + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + g->batch_q_half = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(uint16_t)); + } + g->prefill_seed_router_selected = ds4_gpu_tensor_alloc( + (uint64_t)DS4_N_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * sizeof(int32_t)); + } + + bool layer_cache_ok = true; + for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) continue; + layer_cache_ok = g->layer_raw_cache[il] != NULL; + if (layer_cache_ok && g->cuda_tp_attn_cache_dup) { + layer_cache_ok = g->layer_raw_cache_tp[il] != NULL; + } + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (layer_cache_ok && ratio != 0) { + layer_cache_ok = g->layer_attn_comp_cache[il] != NULL && + (!g->cuda_tp_attn_cache_dup || + g->layer_attn_comp_cache_tp[il] != NULL) && + g->layer_attn_state_kv[il] != NULL && + g->layer_attn_state_score[il] != NULL && + (!enable_frontier_snapshot || + (g->spec_attn_state_kv[il] != NULL && + g->spec_attn_state_score[il] != NULL)) && + (!enable_prefix1_snapshot || + (g->spec_prefix1_attn_state_kv[il] != NULL && + g->spec_prefix1_attn_state_score[il] != NULL)); + } + if (layer_cache_ok && ratio == 4) { + layer_cache_ok = g->layer_index_comp_cache[il] != NULL && + g->layer_index_state_kv[il] != NULL && + g->layer_index_state_score[il] != NULL && + (!enable_frontier_snapshot || + (g->spec_index_state_kv[il] != NULL && + g->spec_index_state_score[il] != NULL)) && + (!enable_prefix1_snapshot || + (g->spec_prefix1_index_state_kv[il] != NULL && + g->spec_prefix1_index_state_score[il] != NULL)); + } + } + + /* Class P validation — check every used tier's slot. */ + bool class_p_ok = true; + for (int t = 0; class_p_ok && t < DS4_MAX_GPUS; t++) { + if (!used_tier[t]) continue; + class_p_ok = + g->cur_hc_by_tier[t] && g->flat_hc_by_tier[t] && g->hc_mix_by_tier[t] && g->hc_split_by_tier[t] && + g->hc_pre_by_tier[t] && g->hc_post_by_tier[t] && g->hc_comb_by_tier[t] && + g->attn_cur_by_tier[t] && g->attn_norm_by_tier[t] && g->qr_by_tier[t] && g->qr_norm_by_tier[t] && + g->q_by_tier[t] && g->kv_raw_by_tier[t] && g->kv_by_tier[t] && + g->comp_kv_cur_by_tier[t] && g->comp_sc_cur_by_tier[t] && + (!DS4_GPU_ATTN_COMP_CACHE_F16 || g->attn_comp_stage_by_tier[t]) && + g->indexer_q_by_tier[t] && g->indexer_weights_by_tier[t] && g->indexer_scores_by_tier[t] && + g->comp_mask_by_tier[t] && g->comp_selected_by_tier[t] && + g->heads_by_tier[t] && g->attn_low_by_tier[t] && g->attn_out_by_tier[t] && + g->after_attn_hc_by_tier[t] && g->ffn_cur_by_tier[t] && g->ffn_norm_by_tier[t] && + g->shared_gate_by_tier[t] && g->shared_up_by_tier[t] && g->shared_mid_by_tier[t] && + g->shared_out_by_tier[t] && + g->router_logits_by_tier[t] && g->router_probs_by_tier[t] && + g->router_selected_by_tier[t] && g->router_weights_by_tier[t] && + g->routed_gate_by_tier[t] && g->routed_up_by_tier[t] && g->routed_mid_by_tier[t] && + g->routed_down_by_tier[t] && g->routed_out_by_tier[t] && + (!g->cuda_tp_decode || g->tp_peer_tmp_by_tier[t]) && + g->after_ffn_hc_by_tier[t] && + g->batch_cur_hc_by_tier[t] && g->batch_next_hc_by_tier[t] && g->batch_flat_hc_by_tier[t] && + g->batch_hc_mix_by_tier[t] && g->batch_hc_split_by_tier[t] && + g->batch_attn_cur_by_tier[t] && g->batch_attn_norm_by_tier[t] && + g->batch_qr_by_tier[t] && g->batch_qr_norm_by_tier[t] && g->batch_q_by_tier[t] && + g->batch_kv_raw_by_tier[t] && g->batch_kv_by_tier[t] && + g->batch_comp_kv_by_tier[t] && g->batch_comp_sc_by_tier[t] && + g->batch_indexer_q_by_tier[t] && g->batch_indexer_weights_by_tier[t] && + g->batch_heads_by_tier[t] && g->batch_attn_low_by_tier[t] && g->batch_attn_out_by_tier[t] && + g->batch_group_tmp_by_tier[t] && g->batch_low_tmp_by_tier[t] && g->batch_after_attn_hc_by_tier[t] && + g->batch_ffn_cur_by_tier[t] && g->batch_ffn_norm_by_tier[t] && + g->batch_shared_gate_by_tier[t] && g->batch_shared_up_by_tier[t] && + g->batch_shared_mid_by_tier[t] && g->batch_shared_out_by_tier[t] && + g->batch_router_logits_by_tier[t] && g->batch_router_probs_by_tier[t] && + g->batch_router_selected_by_tier[t] && g->batch_router_weights_by_tier[t] && + g->batch_routed_gate_by_tier[t] && g->batch_routed_up_by_tier[t] && + g->batch_routed_mid_by_tier[t] && g->batch_routed_down_by_tier[t] && + g->batch_routed_out_by_tier[t]; + } + bool output_tp_ok = true; + for (uint32_t i = 1; i < output_tp_ways; i++) { + const int t = output_tp_tiers[i]; + if (t < 0 || t >= DS4_MAX_GPUS || t == g->head_tier) continue; + output_tp_ok = output_tp_ok && + g->output_norm_by_tier[t] != NULL && + g->logits_by_tier[t] != NULL; + } + const bool ok = state_init_ok && layer_cache_ok && class_p_ok && + /* Class H — validate the head_tier slot + * (single-tier: head_tier == 0, byte-equivalent). */ + metal_graph_output_pre(g) && metal_graph_output_weights(g) && + metal_graph_output_embd(g) && metal_graph_output_norm(g) && + metal_graph_logits(g) && output_tp_ok && + (!enable_mtp || + (g->mtp_embed && g->mtp_enorm && g->mtp_eproj && + g->mtp_eproj_hc && g->mtp_hnorm_hc && g->mtp_hproj_hc && + g->mtp_input_hc && g->mtp_state_hc && g->mtp_next_hc && + g->mtp_raw_cache)) && + (!enable_spec_logits || g->spec_logits) && + /* Class E — validate the emb_tier slot. */ + metal_graph_prefill_tokens(g) && + g->cpu_router_norm && + (!DS4_GPU_ATTN_COMP_CACHE_F16 || g->batch_q_half) && + g->prefill_seed_router_selected; + if (!ok) metal_graph_free(g); + return ok; +} + +static bool metal_graph_alloc( + ds4_gpu_graph *g, + const ds4_weights *weights, + const ds4_layer_weights *layer) { + /* single-tier convenience wrapper; placement=NULL routes + * all per-layer allocations to tier 0. */ + return metal_graph_alloc_raw_cap(g, weights, layer, DS4_N_SWA, DS4_N_SWA, + 1, false, NULL, false, NULL); +} + +static bool metal_graph_install_model_spans( + const ds4_model *model, + const ds4_model_map_span_vec *spans, + const char *label) { + if (!model || !spans || spans->len == 0) return false; + + uint64_t *offsets = xmalloc((size_t)spans->len * sizeof(offsets[0])); + uint64_t *sizes = xmalloc((size_t)spans->len * sizeof(sizes[0])); + for (uint32_t i = 0; i < spans->len; i++) { + offsets[i] = spans->v[i].off; + sizes[i] = spans->v[i].end - spans->v[i].off; + } + + const bool ok = ds4_gpu_set_model_map_spans(model->map, + model->size, + offsets, + sizes, + spans->len, + spans->max_tensor_bytes) != 0; + if (!ok) { + fprintf(stderr, + "ds4: Metal SSD streaming failed to map %s model spans\n", + label ? label : "requested"); + } + free(offsets); + free(sizes); + return ok; +} + +static bool metal_graph_stream_readahead_enabled(void) { + return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_READAHEAD", + "DS4_METAL_ENABLE_STREAMING_READAHEAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_READAHEAD"); +} + +static bool metal_graph_stream_madvise_willneed_enabled(void) { + return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_MADVISE_WILLNEED", + "DS4_METAL_ENABLE_STREAMING_MADVISE_WILLNEED") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_MADVISE_WILLNEED", + "DS4_METAL_DISABLE_STREAMING_MADVISE_WILLNEED"); +} + +static bool metal_graph_stream_decode_static_map_enabled(void) { + if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_STATIC_DECODE_MAP", + "DS4_METAL_DISABLE_STREAMING_STATIC_DECODE_MAP")) { + return false; + } +#ifdef DS4_ROCM_BUILD + return glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_STATIC_DECODE_MAP", + "DS4_METAL_ENABLE_STREAMING_STATIC_DECODE_MAP"); +#else + return true; +#endif +} + +static bool metal_graph_stream_decode_static_map_state_cache_enabled(void) { + return !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_STATIC_MAP_STATE_CACHE", + "DS4_METAL_DISABLE_STREAMING_STATIC_MAP_STATE_CACHE"); +} + +static bool metal_graph_stream_decode_layer_batch_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + !g_expert_profile.active && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_LAYER_BATCH", + "DS4_METAL_DISABLE_STREAMING_LAYER_BATCH") && + (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE", + "DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE", + "DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE")) && + !glm_graph_env_present("DS4_ROCM_DECODE_STAGE_PROFILE", + "DS4_METAL_DECODE_STAGE_PROFILE") && + !glm_graph_env_present("DS4_ROCM_GRAPH_DUMP_PREFIX", + "DS4_METAL_GRAPH_DUMP_PREFIX"); +} + +static void metal_graph_stream_readahead_range_impl( + const ds4_model *model, + uint64_t offset, + uint64_t size, + bool enabled) { + if (!enabled || + !model || + model->fd < 0 || + !model->map || + offset > model->size || + size == 0 || + size > model->size - offset) { + return; + } + +#if defined(F_RDADVISE) + uint64_t pos = offset; + uint64_t rem = size; + while (rem > 0) { + const uint64_t chunk64 = + rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem; + if (pos > (uint64_t)LLONG_MAX) break; + + struct radvisory ra; + ra.ra_offset = (off_t)pos; + ra.ra_count = (int)chunk64; + (void)fcntl(model->fd, F_RDADVISE, &ra); + + pos += chunk64; + rem -= chunk64; + } +#else + (void)model; + (void)offset; + (void)size; +#endif +} + +static bool metal_graph_stream_madvise_willneed_range_impl( + const ds4_model *model, + uint64_t offset, + uint64_t size, + bool enabled, + uint64_t *advised) { + if (!enabled || + !model || + !model->map || + offset > model->size || + size == 0 || + size > model->size - offset) { + return !enabled; + } + +#if defined(POSIX_MADV_WILLNEED) + const uint64_t page = (uint64_t)getpagesize(); + if (page == 0) return false; + const uint64_t page_offset = offset & ~(page - 1u); + const uint64_t leading = offset - page_offset; + if (size > UINT64_MAX - leading || + leading + size > UINT64_MAX - (page - 1u)) { + return false; + } + uint64_t advise_bytes = align_up(leading + size, page); + if (advise_bytes > model->size - page_offset) { + advise_bytes = model->size - page_offset; + } + if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) { + return false; + } + uint8_t *base = (uint8_t *)model->map; + const int rc = posix_madvise((void *)(base + page_offset), + (size_t)advise_bytes, + POSIX_MADV_WILLNEED); + if (rc != 0) return false; + if (advised) { + if (*advised > UINT64_MAX - advise_bytes) { + *advised = UINT64_MAX; + } else { + *advised += advise_bytes; + } + } + return true; +#else + (void)model; + (void)offset; + (void)size; + (void)advised; + return true; +#endif +} + +static void metal_graph_stream_readahead_range( + const ds4_model *model, + uint64_t offset, + uint64_t size) { + metal_graph_stream_readahead_range_impl(model, + offset, + size, + metal_graph_stream_readahead_enabled()); + metal_graph_stream_madvise_willneed_range_impl( + model, + offset, + size, + metal_graph_stream_madvise_willneed_enabled(), + NULL); +} + +static void metal_graph_stream_readahead_spans( + const ds4_model *model, + const ds4_model_map_span_vec *spans) { + if (!spans) return; + for (uint32_t i = 0; i < spans->len; i++) { + metal_graph_stream_readahead_range(model, + spans->v[i].off, + spans->v[i].end - spans->v[i].off); + } +} + +typedef struct { + uint64_t off; + uint64_t size; +} metal_graph_stream_pagein_range; + +typedef struct { + pthread_t thread; + const ds4_model *model; + metal_graph_stream_pagein_range *ranges; + pthread_t *threads; + struct metal_graph_stream_pagein_worker *workers; + uint32_t n_ranges; + uint32_t n_threads; + uint32_t layer; + uint32_t n_tokens; + uint32_t unique; + uint64_t bytes; + uint64_t touched; + double read_ms; + double thread_ms; + bool profile; + bool madvise_only; + bool pread_only; + bool readahead_only; + bool started; + bool ok; + uint8_t sink; +} metal_graph_stream_pagein_job; + +typedef struct metal_graph_stream_pagein_worker { + metal_graph_stream_pagein_job *job; + uint32_t first; + uint32_t stride; + uint64_t touched; + double thread_ms; + bool ok; + uint8_t sink; +} metal_graph_stream_pagein_worker; + +static bool metal_graph_stream_prefill_selected_pagein_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_PAGEIN", + "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_PAGEIN") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_PAGEIN", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_PAGEIN"); +} + +static bool metal_graph_stream_prefill_selected_madvise_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_MADVISE", + "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_MADVISE") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_MADVISE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_MADVISE"); +} + +static bool metal_graph_stream_prefill_layer_pagein_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN"); +} + +static bool metal_graph_stream_prefill_layer_readahead_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); +} + +static bool metal_graph_stream_prefill_layer_pread_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); +} + +static bool metal_graph_stream_prefill_layer_madvise_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE"); +} + +static uint32_t metal_graph_stream_prefill_batch_selected_addr_auto_max(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX", + "DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } +#ifdef DS4_ROCM_BUILD + if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO || + DS4_MODEL_VARIANT == DS4_VARIANT_FLASH || + DS4_MODEL_VARIANT == DS4_VARIANT_GLM52) return UINT32_MAX; +#endif + if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO) return 800u; + if (DS4_MODEL_VARIANT == DS4_VARIANT_FLASH) return 760u; + return 0; +} + +static uint32_t metal_graph_stream_prefill_batch_selected_addr_auto_min(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN", + "DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } +#ifdef DS4_ROCM_BUILD + if (DS4_MODEL_VARIANT == DS4_VARIANT_GLM52) return 2u; +#endif + if (DS4_MODEL_VARIANT == DS4_VARIANT_PRO || + DS4_MODEL_VARIANT == DS4_VARIANT_FLASH) return 2u; + return 0; +} + +static bool metal_graph_stream_prefill_batch_selected_addr_enabled( + const ds4_gpu_graph *g, + const ds4_weights *weights, + uint32_t n_tokens) { + if (!g || + !g->ssd_streaming || + g->quality || + !weights || + n_tokens <= 1 || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR", + "DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_EXPERT_ADDR_TABLE", + "DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") || + glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") || + glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") || + DS4_N_LAYER == 0) { + return false; + } + const uint32_t routed_il = + DS4_N_LEADING_DENSE < DS4_N_LAYER ? DS4_N_LEADING_DENSE : 0u; + const ds4_layer_weights *layer = &weights->layer[routed_il]; + if (!layer->ffn_gate_exps || !layer->ffn_up_exps || !layer->ffn_down_exps) { + return false; + } +#ifdef DS4_ROCM_BUILD + const bool selected_iq2 = + glm_stream_selected_expert_cache_supported(layer, routed_il); + const bool selected_q2 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q2_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q2_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && + glm_stream_expert_cache_addr_layout_supported(weights, layer, routed_il); + if (!selected_iq2 && !selected_q2) return false; +#else + if (DS4_N_EXPERT_USED != 6 || + layer->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || + layer->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || + layer->ffn_down_exps->type != DS4_TENSOR_Q2_K) { + return false; + } +#endif + + const uint32_t cache_configured = + ds4_gpu_stream_expert_cache_configured_count(); +#ifdef DS4_ROCM_BUILD + if (cache_configured == 0) { + return false; + } +#else + if (cache_configured < DS4_N_EXPERT) { + return false; + } +#endif + + if (glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR", + "DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR")) { + return true; + } + + const uint32_t max_tokens = + metal_graph_stream_prefill_batch_selected_addr_auto_max(); + const uint32_t min_tokens = + metal_graph_stream_prefill_batch_selected_addr_auto_min(); + return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens; +} + +static bool metal_graph_cuda_stream_prefill_batch_selected_addr_enabled( + const ds4_gpu_graph *g, + const ds4_weights *weights, + uint32_t n_tokens) { +#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) + if (!g || + !g->ssd_streaming || + g->quality || + !weights || + n_tokens <= 1 || + getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || + getenv("DS4_CUDA_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL || + DS4_N_LAYER == 0 || + DS4_N_EXPERT < 128 || + DS4_N_EXPERT_USED != 6) { + return false; + } + + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const ds4_layer_weights *layer = &weights->layer[il]; + if (!layer->ffn_gate_exps || !layer->ffn_up_exps || + !layer->ffn_down_exps) { + continue; + } + const bool q4 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q4_K; + const bool iq2 = + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K; + if (q4 || iq2) return true; + } + return false; +#else + (void)g; + (void)weights; + (void)n_tokens; + return false; +#endif +} + +#ifdef DS4_ROCM_BUILD +enum { DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 1024 }; +enum { DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS = 8 }; + +typedef struct rocm_graph_stream_layer_expert_load { + pthread_t thread; + bool active; + bool ok; + const ds4_model *model; + const ds4_layer_weights *layer; + uint32_t il; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; +} rocm_graph_stream_layer_expert_load; + +static bool rocm_graph_stream_prefill_full_layer_enabled( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens) { + return g && + g->ssd_streaming && + !g->quality && + layer && + n_tokens >= DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS && + glm_stream_resident_decode_layer_supported(layer, il); +} + +static uint32_t rocm_graph_stream_prefill_full_layer_seed_tokens(void) { + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_count(); + const uint64_t entries_per_token = + (uint64_t)DS4_N_LAYER * (uint64_t)DS4_N_EXPERT_USED; + if (entries_per_token == 0) return 1; + uint32_t seed_tokens = budget == 0 ? 1 : (uint32_t)(budget / entries_per_token); + if (seed_tokens < 1) seed_tokens = 1; + if (seed_tokens > DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS) { + seed_tokens = DS4_ROCM_STREAM_PREFILL_FULL_LAYER_MAX_SEED_TOKENS; + } + return seed_tokens; +} + +static bool rocm_graph_stream_layer_expert_bytes( + const ds4_layer_weights *layer, + uint64_t *gate_expert_bytes, + uint64_t *down_expert_bytes) { + return streaming_layer_gate_down_expert_bytes(layer, + gate_expert_bytes, + down_expert_bytes); +} + +static bool rocm_graph_stream_layer_expert_load_sync( + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + return model && + layer && + ds4_gpu_stream_expert_cache_load_layer(&table) != 0; +} + +static void *rocm_graph_stream_layer_expert_load_thread_main(void *arg) { + rocm_graph_stream_layer_expert_load *job = arg; + if (!job) return NULL; + job->ok = rocm_graph_stream_layer_expert_load_sync(job->model, + job->layer, + job->il, + job->gate_expert_bytes, + job->down_expert_bytes); + return NULL; +} + +static bool rocm_graph_stream_layer_expert_load_join( + rocm_graph_stream_layer_expert_load *job) { + if (!job || !job->active) return true; + const int rc = pthread_join(job->thread, NULL); + const bool ok = rc == 0 && job->ok; + if (rc != 0) { + fprintf(stderr, + "ds4: ROCm streaming full-layer expert load join failed: %s\n", + strerror(rc)); + } + memset(job, 0, sizeof(*job)); + return ok; +} + +static bool rocm_graph_stream_layer_expert_load_start( + rocm_graph_stream_layer_expert_load *job, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!job || job->active || !model || !layer) return false; + memset(job, 0, sizeof(*job)); + job->model = model; + job->layer = layer; + job->il = il; + job->gate_expert_bytes = gate_expert_bytes; + job->down_expert_bytes = down_expert_bytes; + const int rc = pthread_create(&job->thread, + NULL, + rocm_graph_stream_layer_expert_load_thread_main, + job); + if (rc != 0) { + fprintf(stderr, + "ds4: failed to start ROCm streaming full-layer expert load " + "thread for layer %u: %s\n", + il, + strerror(rc)); + memset(job, 0, sizeof(*job)); + return false; + } + job->active = true; + return true; +} + +static bool rocm_graph_stream_layer_expert_load_start_next( + rocm_graph_stream_layer_expert_load *job, + const ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il, + uint32_t n_tokens) { + if (!job || + !model || + !weights || + il >= DS4_N_LAYER || + !rocm_graph_stream_prefill_full_layer_enabled(g, + &weights->layer[il], + il, + n_tokens)) { + return true; + } + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!rocm_graph_stream_layer_expert_bytes(&weights->layer[il], + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + return rocm_graph_stream_layer_expert_load_start(job, + model, + &weights->layer[il], + il, + gate_expert_bytes, + down_expert_bytes); +} + +static bool rocm_graph_stream_layer_expert_load_ready( + rocm_graph_stream_layer_expert_load *job, + const ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il, + uint32_t n_tokens) { + if (!model || !weights || il >= DS4_N_LAYER) return false; + if (!rocm_graph_stream_prefill_full_layer_enabled(g, + &weights->layer[il], + il, + n_tokens)) { + return true; + } + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!rocm_graph_stream_layer_expert_bytes(&weights->layer[il], + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + if (job && job->active) { + if (job->il != il) { + fprintf(stderr, + "ds4: ROCm streaming full-layer expert load expected layer " + "%u but pending job is layer %u\n", + il, + job->il); + return false; + } + return rocm_graph_stream_layer_expert_load_join(job); + } + return rocm_graph_stream_layer_expert_load_sync(model, + &weights->layer[il], + il, + gate_expert_bytes, + down_expert_bytes); +} + +static bool rocm_graph_stream_seed_full_layer_selected( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens) { + if (!rocm_graph_stream_prefill_full_layer_enabled(g, layer, il, n_tokens)) { + return true; + } + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!rocm_graph_stream_layer_expert_bytes(layer, + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + if (ds4_gpu_stream_expert_cache_seed_from_layer_selected( + &table, + metal_graph_batch_router_selected(g), + n_tokens, + rocm_graph_stream_prefill_full_layer_seed_tokens(), + DS4_N_EXPERT_USED) == 0) { + static bool warned = false; + if (!warned) { + fprintf(stderr, + "ds4: ROCm streaming full-layer prefill seed skipped; " + "decode may start with a colder expert cache\n"); + warned = true; + } + } + return true; +} +#endif + +static bool metal_graph_stream_prefill_selected_profile_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_PROFILE", + "DS4_METAL_STREAMING_PREFILL_SELECTED_PROFILE") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_PROFILE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_PROFILE"); +} + +static void metal_graph_stream_prefill_selected_profile_reset( + ds4_gpu_graph *g) { + if (!g) return; + g->prefill_selected_profile_rows = 0; + g->prefill_selected_profile_unique = 0; + g->prefill_selected_profile_selected_bytes = 0; + g->prefill_selected_profile_full_bytes = 0; + g->prefill_selected_profile_layers = 0; + g->prefill_selected_profile_min_unique = UINT32_MAX; + g->prefill_selected_profile_max_unique = 0; +} + +static uint64_t metal_graph_stream_prefill_selected_profile_add_bytes( + uint64_t a, + uint64_t b) { + return a > UINT64_MAX - b ? UINT64_MAX : a + b; +} + +static bool metal_graph_selected_profile_layer_impl( + ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + const char *label) { + if (!layer || !metal_graph_batch_router_selected(g) || n_tokens == 0 || + DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || + DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return false; + } + + const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; + int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); + const bool read_ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), + 0, + selected, + n_ids * sizeof(selected[0])) != 0; + if (!read_ok) { + free(selected); + return false; + } + + bool seen[DS4_MAX_EXPERT] = { false }; + uint32_t unique = 0; + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t expert = selected[i]; + if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: Metal streaming prefill selected profile expert id %d is outside 0..%u at layer %u\n", + expert, + (uint32_t)DS4_N_EXPERT, + il); + free(selected); + return false; + } + if (!seen[expert]) { + seen[expert] = true; + unique++; + } + } + free(selected); + + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: Metal streaming prefill selected profile byte size overflow at layer %u\n", il); + return false; + } + const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; + const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes || + gate_expert_bytes + gate_expert_bytes > UINT64_MAX - down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming prefill selected profile byte size overflow at layer %u\n", il); + return false; + } + const uint64_t per_expert_bytes = gate_expert_bytes + gate_expert_bytes + + down_expert_bytes; + const uint64_t selected_bytes = + unique > UINT64_MAX / per_expert_bytes ? + UINT64_MAX : (uint64_t)unique * per_expert_bytes; + const uint64_t full_bytes = + (uint64_t)DS4_N_EXPERT > UINT64_MAX / per_expert_bytes ? + UINT64_MAX : (uint64_t)DS4_N_EXPERT * per_expert_bytes; + const double ratio = full_bytes == 0 ? 0.0 : + (double)selected_bytes / (double)full_bytes; + + g->prefill_selected_profile_layers++; + g->prefill_selected_profile_rows = + metal_graph_stream_prefill_selected_profile_add_bytes( + g->prefill_selected_profile_rows, + n_ids); + g->prefill_selected_profile_unique = + metal_graph_stream_prefill_selected_profile_add_bytes( + g->prefill_selected_profile_unique, + unique); + g->prefill_selected_profile_selected_bytes = + metal_graph_stream_prefill_selected_profile_add_bytes( + g->prefill_selected_profile_selected_bytes, + selected_bytes); + g->prefill_selected_profile_full_bytes = + metal_graph_stream_prefill_selected_profile_add_bytes( + g->prefill_selected_profile_full_bytes, + full_bytes); + if (unique < g->prefill_selected_profile_min_unique) { + g->prefill_selected_profile_min_unique = unique; + } + if (unique > g->prefill_selected_profile_max_unique) { + g->prefill_selected_profile_max_unique = unique; + } + + fprintf(stderr, + "ds4: %s layer=%u " + "tokens=%u unique=%u/%u selected=%.2f GiB full=%.2f GiB ratio=%.3f\n", + label ? label : "selected expert profile", + il, + n_tokens, + unique, + (uint32_t)DS4_N_EXPERT, + (double)selected_bytes / (1024.0 * 1024.0 * 1024.0), + (double)full_bytes / (1024.0 * 1024.0 * 1024.0), + ratio); + return true; +} + +static bool metal_graph_stream_prefill_selected_profile_layer( + ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens) { + if (!metal_graph_stream_prefill_selected_profile_enabled(g)) return true; + return metal_graph_selected_profile_layer_impl( + g, + layer, + il, + n_tokens, + "Metal streaming prefill selected profile"); +} + +static void metal_graph_selected_profile_summary_impl( + const ds4_gpu_graph *g, + const char *label) { + if (!g || g->prefill_selected_profile_layers == 0) { + return; + } + const double layers = (double)g->prefill_selected_profile_layers; + const double avg_unique = (double)g->prefill_selected_profile_unique / layers; + const double ratio = g->prefill_selected_profile_full_bytes == 0 ? 0.0 : + (double)g->prefill_selected_profile_selected_bytes / + (double)g->prefill_selected_profile_full_bytes; + fprintf(stderr, + "ds4: %s summary " + "layers=%u avg_unique=%.1f min_unique=%u max_unique=%u " + "selected=%.2f GiB full=%.2f GiB ratio=%.3f rows=%" PRIu64 "\n", + label ? label : "selected expert profile", + g->prefill_selected_profile_layers, + avg_unique, + g->prefill_selected_profile_min_unique == UINT32_MAX ? + 0 : g->prefill_selected_profile_min_unique, + g->prefill_selected_profile_max_unique, + (double)g->prefill_selected_profile_selected_bytes / + (1024.0 * 1024.0 * 1024.0), + (double)g->prefill_selected_profile_full_bytes / + (1024.0 * 1024.0 * 1024.0), + ratio, + g->prefill_selected_profile_rows); +} + +static void metal_graph_stream_prefill_selected_profile_summary( + const ds4_gpu_graph *g) { + if (!metal_graph_stream_prefill_selected_profile_enabled(g)) return; + metal_graph_selected_profile_summary_impl( + g, + "Metal streaming prefill selected profile"); +} + +static bool metal_graph_stream_pagein_touch_range( + const ds4_model *model, + uint64_t offset, + uint64_t size, + uint64_t *touched, + uint8_t *sink) { + if (!model || + !model->map || + model->size == 0 || + offset > model->size || + size == 0 || + size > model->size - offset) { + return false; + } + + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t page_offset = offset & ~(page - 1u); + const uint64_t leading = offset - page_offset; + if (size > UINT64_MAX - leading || + leading + size > UINT64_MAX - (page - 1u)) { + return false; + } + uint64_t touch_bytes = align_up(leading + size, page); + if (touch_bytes > model->size - page_offset) { + touch_bytes = model->size - page_offset; + } + if (touch_bytes == 0 || touch_bytes > (uint64_t)SIZE_MAX) { + return false; + } + + const uint8_t *base = (const uint8_t *)model->map; + const volatile uint8_t *p = + (const volatile uint8_t *)(base + page_offset); + +#if defined(POSIX_MADV_WILLNEED) + (void)posix_madvise((void *)(base + page_offset), + (size_t)touch_bytes, + POSIX_MADV_WILLNEED); +#endif + + uint8_t s = sink ? *sink : 0; + for (uint64_t off = 0; off < touch_bytes; off += page) { + s ^= p[off]; + } + s ^= p[touch_bytes - 1u]; + if (sink) *sink = s; + if (touched) *touched += touch_bytes; + return true; +} + +static bool metal_graph_stream_pread_range( + const ds4_model *model, + uint64_t offset, + uint64_t size, + uint64_t *read_bytes, + uint8_t *sink) { + if (!model || + model->fd < 0 || + offset > model->size || + size == 0 || + size > model->size - offset) { + return false; + } + if (offset > (uint64_t)LLONG_MAX) return false; + + const size_t chunk = 1024u * 1024u; + uint8_t *buf = xmalloc(chunk); + uint64_t pos = offset; + uint64_t rem = size; + uint8_t s = sink ? *sink : 0; + bool ok = true; + while (rem != 0) { + const size_t want = rem > (uint64_t)chunk ? chunk : (size_t)rem; + ssize_t nread; + do { + nread = pread(model->fd, buf, want, (off_t)pos); + } while (nread < 0 && errno == EINTR); + if (nread <= 0) { + ok = false; + break; + } + s ^= buf[0]; + s ^= buf[(size_t)nread - 1u]; + pos += (uint64_t)nread; + rem -= (uint64_t)nread; + if (read_bytes) { + *read_bytes = *read_bytes > UINT64_MAX - (uint64_t)nread ? + UINT64_MAX : *read_bytes + (uint64_t)nread; + } + } + if (sink) *sink = s; + free(buf); + return ok; +} + +static bool metal_graph_stream_prepare_range( + const metal_graph_stream_pagein_job *job, + uint64_t offset, + uint64_t size, + uint64_t *touched, + uint8_t *sink) { + if (!job) return false; + if (job->pread_only) { + return metal_graph_stream_pread_range(job->model, + offset, + size, + touched, + sink); + } + if (job->readahead_only) { + metal_graph_stream_readahead_range_impl(job->model, + offset, + size, + true); + if (touched) { + *touched = *touched > UINT64_MAX - size ? + UINT64_MAX : *touched + size; + } + return true; + } + if (job->madvise_only) { + return metal_graph_stream_madvise_willneed_range_impl(job->model, + offset, + size, + true, + touched); + } + return metal_graph_stream_pagein_touch_range(job->model, + offset, + size, + touched, + sink); +} + +static void *metal_graph_stream_pagein_thread_main(void *arg) { + metal_graph_stream_pagein_job *job = arg; + const double t0 = job->profile ? now_sec() : 0.0; + job->ok = true; + for (uint32_t i = 0; i < job->n_ranges; i++) { + const bool ok = metal_graph_stream_prepare_range(job, + job->ranges[i].off, + job->ranges[i].size, + &job->touched, + &job->sink); + if (!ok) { + job->ok = false; + break; + } + } + if (job->profile) { + job->thread_ms = (now_sec() - t0) * 1000.0; + } + return NULL; +} + +static void *metal_graph_stream_pagein_worker_main(void *arg) { + metal_graph_stream_pagein_worker *worker = arg; + metal_graph_stream_pagein_job *job = worker ? worker->job : NULL; + const double t0 = job && job->profile ? now_sec() : 0.0; + worker->ok = true; + if (!job || worker->stride == 0) { + worker->ok = false; + return NULL; + } + for (uint32_t i = worker->first; i < job->n_ranges; i += worker->stride) { + const bool ok = metal_graph_stream_prepare_range(job, + job->ranges[i].off, + job->ranges[i].size, + &worker->touched, + &worker->sink); + if (!ok) { + worker->ok = false; + break; + } + } + if (job->profile) { + worker->thread_ms = (now_sec() - t0) * 1000.0; + } + return NULL; +} + +static uint32_t metal_graph_stream_prefill_layer_pagein_threads(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_THREADS", + "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_THREADS"); + if (!env || !env[0]) { + env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_THREADS", + "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_THREADS"); + } + if (!env || !env[0]) return 8; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0' || v == 0) return 1; + return v > 16 ? 16u : (uint32_t)v; +} + +static uint32_t metal_graph_stream_prefill_selected_prepare_threads( + bool madvise_only) { + if (!madvise_only) return 1; + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_SELECTED_PREPARE_THREADS", + "DS4_METAL_STREAMING_PREFILL_SELECTED_PREPARE_THREADS"); + if (!env || !env[0]) { + env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_SELECTED_MADVISE_THREADS", + "DS4_METAL_STREAMING_PREFILL_SELECTED_MADVISE_THREADS"); + } + if (!env || !env[0]) return metal_graph_stream_prefill_layer_pagein_threads(); + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0' || v == 0) return 1; + return v > 16 ? 16u : (uint32_t)v; +} + +static uint32_t metal_graph_stream_prefill_selected_prepare_gap(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_SELECTED_PREPARE_GAP", + "DS4_METAL_STREAMING_PREFILL_SELECTED_PREPARE_GAP"); + if (!env || !env[0]) return 0; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0') return 0; + return v > 8 ? 8u : (uint32_t)v; +} + +static bool metal_graph_stream_prefill_layer_pagein_overlap_enabled(void) { + return !glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_NO_OVERLAP", + "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_NO_OVERLAP") && + !glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_NO_OVERLAP", + "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_NO_OVERLAP") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE_OVERLAP", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE_OVERLAP") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN_OVERLAP", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN_OVERLAP"); +} + +enum { DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD = 4 }; + +static uint32_t metal_graph_stream_prefill_layer_prepare_ahead(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_LAYER_PREPARE_AHEAD", + "DS4_METAL_STREAMING_PREFILL_LAYER_PREPARE_AHEAD"); + if (!env || !env[0]) return 1; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0' || v == 0) return 1; + if (v > DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD) { + return DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD; + } + return (uint32_t)v; +} + +static bool metal_graph_stream_prefill_selected_pagein_start( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + metal_graph_stream_pagein_job *job) { + if (!job) return false; + memset(job, 0, sizeof(*job)); + job->ok = true; + job->profile = + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_PAGEIN_PROFILE", + "DS4_METAL_STREAMING_PREFILL_SELECTED_PAGEIN_PROFILE") || + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_MADVISE_PROFILE", + "DS4_METAL_STREAMING_PREFILL_SELECTED_MADVISE_PROFILE"); + job->layer = il; + job->n_tokens = n_tokens; + + const bool madvise_only = + metal_graph_stream_prefill_selected_madvise_enabled(g); + job->madvise_only = madvise_only; + if (!metal_graph_stream_prefill_selected_pagein_enabled(g) && + !madvise_only) return true; + if (!model || !layer || !metal_graph_batch_router_selected(g) || n_tokens == 0) { + return false; + } + + const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; + int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); + + const double t_read0 = job->profile ? now_sec() : 0.0; + bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), + 0, + selected, + n_ids * sizeof(selected[0])) != 0; + if (job->profile) { + job->read_ms = (now_sec() - t_read0) * 1000.0; + } + + bool seen[DS4_MAX_EXPERT] = { false }; + if (ok) { + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t expert = selected[i]; + if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: Metal streaming prefill selected page-in expert id %d is outside 0..%u at layer %u\n", + expert, + (uint32_t)DS4_N_EXPERT, + il); + ok = false; + break; + } + if (seen[expert]) continue; + seen[expert] = true; + job->unique++; + } + } + free(selected); + + metal_graph_stream_pagein_range *ranges = NULL; + uint32_t n_ranges = 0; + if (ok && job->unique != 0) { + ranges = xmalloc((size_t)DS4_N_EXPERT * 3u * sizeof(ranges[0])); + const uint32_t gap = madvise_only ? + metal_graph_stream_prefill_selected_prepare_gap() : 0; + uint32_t e = 0; + while (e < DS4_N_EXPERT) { + while (e < DS4_N_EXPERT && !seen[e]) e++; + if (e >= DS4_N_EXPERT) break; + const uint32_t first = e; + uint32_t last = e; + uint32_t skipped = 0; + e++; + while (e < DS4_N_EXPERT) { + if (seen[e]) { + last = e; + skipped = 0; + } else if (skipped < gap) { + skipped++; + } else { + break; + } + e++; + } + + const uint64_t first_id = first; + const uint64_t n_experts = (uint64_t)last - (uint64_t)first + 1u; + if (first_id > UINT64_MAX / gate_expert_bytes || + first_id > UINT64_MAX / down_expert_bytes || + n_experts > UINT64_MAX / gate_expert_bytes || + n_experts > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming prefill selected page-in offset overflow\n"); + ok = false; + break; + } + const uint64_t gate_rel = first_id * gate_expert_bytes; + const uint64_t down_rel = first_id * down_expert_bytes; + const uint64_t gate_bytes = n_experts * gate_expert_bytes; + const uint64_t down_bytes = n_experts * down_expert_bytes; + if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || + gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || + down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { + fprintf(stderr, "ds4: Metal streaming prefill selected page-in offset overflow\n"); + ok = false; + break; + } + ranges[n_ranges++] = (metal_graph_stream_pagein_range){ + layer->ffn_gate_exps->abs_offset + gate_rel, + gate_bytes, + }; + ranges[n_ranges++] = (metal_graph_stream_pagein_range){ + layer->ffn_up_exps->abs_offset + gate_rel, + gate_bytes, + }; + ranges[n_ranges++] = (metal_graph_stream_pagein_range){ + layer->ffn_down_exps->abs_offset + down_rel, + down_bytes, + }; + uint64_t run_bytes = UINT64_MAX; + if (gate_bytes <= (UINT64_MAX - down_bytes) / 2ull) { + run_bytes = gate_bytes * 2ull + down_bytes; + } + if (run_bytes == UINT64_MAX || + job->bytes > UINT64_MAX - run_bytes) { + job->bytes = UINT64_MAX; + } else { + job->bytes += run_bytes; + } + } + } + + if (!ok || n_ranges == 0) { + free(ranges); + return ok; + } + + job->model = model; + job->ranges = ranges; + job->n_ranges = n_ranges; + job->n_threads = metal_graph_stream_prefill_selected_prepare_threads(madvise_only); + if (job->n_threads <= 1) { + const int rc = pthread_create(&job->thread, + NULL, + metal_graph_stream_pagein_thread_main, + job); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill selected page-in thread failed: %s\n", + strerror(rc)); + free(ranges); + memset(job, 0, sizeof(*job)); + return false; + } + } else { + job->threads = xcalloc(job->n_threads, sizeof(job->threads[0])); + job->workers = xcalloc(job->n_threads, sizeof(job->workers[0])); + for (uint32_t t = 0; t < job->n_threads; t++) { + job->workers[t].job = job; + job->workers[t].first = t; + job->workers[t].stride = job->n_threads; + const int rc = pthread_create(&job->threads[t], + NULL, + metal_graph_stream_pagein_worker_main, + &job->workers[t]); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill selected page-in worker failed: %s\n", + strerror(rc)); + for (uint32_t j = 0; j < t; j++) { + (void)pthread_join(job->threads[j], NULL); + } + free(job->workers); + free(job->threads); + free(ranges); + memset(job, 0, sizeof(*job)); + return false; + } + } + } + job->started = true; + return true; +} + +static bool metal_graph_stream_prefill_selected_pagein_join( + metal_graph_stream_pagein_job *job) { + if (!job || !job->started) return true; + const double t0 = job->profile ? now_sec() : 0.0; + int rc = 0; + bool ok = true; + if (job->n_threads <= 1) { + rc = pthread_join(job->thread, NULL); + ok = rc == 0 && job->ok; + } else { + job->touched = 0; + job->thread_ms = 0.0; + job->sink = 0; + for (uint32_t t = 0; t < job->n_threads; t++) { + const int trc = pthread_join(job->threads[t], NULL); + if (trc != 0 && rc == 0) rc = trc; + if (trc != 0 || !job->workers[t].ok) ok = false; + if (job->touched > UINT64_MAX - job->workers[t].touched) { + job->touched = UINT64_MAX; + } else { + job->touched += job->workers[t].touched; + } + if (job->workers[t].thread_ms > job->thread_ms) { + job->thread_ms = job->workers[t].thread_ms; + } + job->sink ^= job->workers[t].sink; + } + } + const double wait_ms = job->profile ? (now_sec() - t0) * 1000.0 : 0.0; + if (job->profile) { + const char *kind = job->madvise_only ? "madvise" : "page-in"; + const char *bytes_label = job->madvise_only ? "advised" : "touched"; + fprintf(stderr, + "ds4: Metal streaming prefill selected %s layer=%u " + "tokens=%u unique=%u ranges=%u bytes=%.2f GiB " + "read=%.3f ms wait=%.3f ms thread=%.3f ms %s=%.2f GiB ok=%d\n", + kind, + job->layer, + job->n_tokens, + job->unique, + job->n_ranges, + (double)job->bytes / (1024.0 * 1024.0 * 1024.0), + job->read_ms, + wait_ms, + job->thread_ms, + bytes_label, + (double)job->touched / (1024.0 * 1024.0 * 1024.0), + ok ? 1 : 0); + } + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill selected page-in join failed: %s\n", + strerror(rc)); + } + free(job->workers); + free(job->threads); + free(job->ranges); + memset(job, 0, sizeof(*job)); + return ok; +} + +static bool metal_graph_stream_prefill_layer_pagein_start( + const ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il, + uint32_t n_tokens, + bool madvise_only, + bool pread_only, + bool readahead_only, + bool decode_only, + metal_graph_stream_pagein_job *job) { + if (!job) return false; + memset(job, 0, sizeof(*job)); + job->ok = true; + job->madvise_only = madvise_only; + job->pread_only = pread_only; + job->readahead_only = readahead_only; + job->profile = + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PAGEIN_PROFILE", + "DS4_METAL_STREAMING_PREFILL_LAYER_PAGEIN_PROFILE") || + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_PREAD_PROFILE", + "DS4_METAL_STREAMING_PREFILL_LAYER_PREAD_PROFILE") || + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_MADVISE_PROFILE", + "DS4_METAL_STREAMING_PREFILL_LAYER_MADVISE_PROFILE") || + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_LAYER_READAHEAD_PROFILE", + "DS4_METAL_STREAMING_PREFILL_LAYER_READAHEAD_PROFILE"); + job->layer = il; + job->n_tokens = n_tokens; + + if (pread_only) { + if (g) { + if (!metal_graph_stream_prefill_layer_pread_enabled(g)) return true; + } else if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE")) { + return true; + } + } else if (readahead_only) { + if (g) { + if (!metal_graph_stream_prefill_layer_readahead_enabled(g)) return true; + } else if (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE")) { + return true; + } + } else if (madvise_only) { + if (g) { + if (!metal_graph_stream_prefill_layer_madvise_enabled(g)) return true; + } else if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE")) { + return true; + } + } else { + if (g) { + if (!metal_graph_stream_prefill_layer_pagein_enabled(g)) return true; + } else if (!glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") || + glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN")) { + return true; + } + } + if (!model || !weights || il >= DS4_N_LAYER) return false; + + const uint32_t n_threads = + metal_graph_stream_prefill_layer_pagein_threads(); + ds4_model_map_span_vec spans; + const bool spans_ok = decode_only ? + weights_model_map_decode_layer_spans(weights, il, &spans) : + weights_model_map_spans(weights, il, il, false, &spans); + if (!spans_ok) return false; + metal_graph_stream_pagein_range *ranges = + xmalloc((size_t)spans.len * n_threads * sizeof(ranges[0])); + uint32_t n_ranges = 0; + const uint64_t page = (uint64_t)getpagesize(); + for (uint32_t i = 0; i < spans.len; i++) { + const uint64_t size = spans.v[i].end - spans.v[i].off; + uint64_t consumed = 0; + uint64_t chunk = size / n_threads; + if (chunk > page) chunk = (chunk / page) * page; + if (chunk == 0) chunk = size; + for (uint32_t t = 0; t < n_threads && consumed < size; t++) { + uint64_t this_size = + (t + 1u == n_threads || size - consumed <= chunk) ? + size - consumed : chunk; + ranges[n_ranges++] = (metal_graph_stream_pagein_range){ + spans.v[i].off + consumed, + this_size, + }; + consumed += this_size; + } + if (job->bytes > UINT64_MAX - size) { + job->bytes = UINT64_MAX; + } else { + job->bytes += size; + } + } + job->unique = spans.len; + free(spans.v); + + job->model = model; + job->ranges = ranges; + job->n_ranges = n_ranges; + job->n_threads = n_threads; + if (n_threads == 1) { + const int rc = pthread_create(&job->thread, + NULL, + metal_graph_stream_pagein_thread_main, + job); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill layer page-in thread failed: %s\n", + strerror(rc)); + free(ranges); + memset(job, 0, sizeof(*job)); + return false; + } + } else { + job->threads = xcalloc(n_threads, sizeof(job->threads[0])); + job->workers = xcalloc(n_threads, sizeof(job->workers[0])); + for (uint32_t t = 0; t < n_threads; t++) { + job->workers[t].job = job; + job->workers[t].first = t; + job->workers[t].stride = n_threads; + const int rc = pthread_create(&job->threads[t], + NULL, + metal_graph_stream_pagein_worker_main, + &job->workers[t]); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill layer page-in worker failed: %s\n", + strerror(rc)); + for (uint32_t j = 0; j < t; j++) { + (void)pthread_join(job->threads[j], NULL); + } + free(job->workers); + free(job->threads); + free(ranges); + memset(job, 0, sizeof(*job)); + return false; + } + } + } + job->started = true; + return true; +} + +static bool metal_graph_stream_prefill_layer_pagein_join( + metal_graph_stream_pagein_job *job) { + if (!job || !job->started) return true; + const double t0 = job->profile ? now_sec() : 0.0; + int rc = 0; + bool ok = true; + if (job->n_threads <= 1) { + rc = pthread_join(job->thread, NULL); + ok = rc == 0 && job->ok; + } else { + job->touched = 0; + job->thread_ms = 0.0; + job->sink = 0; + for (uint32_t t = 0; t < job->n_threads; t++) { + const int trc = pthread_join(job->threads[t], NULL); + if (trc != 0 && rc == 0) rc = trc; + if (trc != 0 || !job->workers[t].ok) ok = false; + if (job->touched > UINT64_MAX - job->workers[t].touched) { + job->touched = UINT64_MAX; + } else { + job->touched += job->workers[t].touched; + } + if (job->workers[t].thread_ms > job->thread_ms) { + job->thread_ms = job->workers[t].thread_ms; + } + job->sink ^= job->workers[t].sink; + } + } + const double wait_ms = job->profile ? (now_sec() - t0) * 1000.0 : 0.0; + if (job->profile) { + const char *kind = job->pread_only ? "pread" : + job->readahead_only ? "readahead" : + job->madvise_only ? "madvise" : "page-in"; + const char *bytes_label = job->pread_only ? "read" : + job->readahead_only ? "requested" : + job->madvise_only ? "advised" : "touched"; + fprintf(stderr, + "ds4: Metal streaming prefill layer %s layer=%u " + "tokens=%u threads=%u ranges=%u bytes=%.2f GiB wait=%.3f ms " + "thread=%.3f ms %s=%.2f GiB ok=%d\n", + kind, + job->layer, + job->n_tokens, + job->n_threads ? job->n_threads : 1u, + job->n_ranges, + (double)job->bytes / (1024.0 * 1024.0 * 1024.0), + wait_ms, + job->thread_ms, + bytes_label, + (double)job->touched / (1024.0 * 1024.0 * 1024.0), + ok ? 1 : 0); + } + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming prefill layer page-in join failed: %s\n", + strerror(rc)); + } + free(job->workers); + free(job->threads); + free(job->ranges); + memset(job, 0, sizeof(*job)); + return ok; +} + +typedef struct { + metal_graph_stream_pagein_job job; + uint32_t layer; + bool active; +} metal_graph_stream_prepare_slot; + +static metal_graph_stream_prepare_slot *metal_graph_stream_prepare_slot_find( + metal_graph_stream_prepare_slot *slots, + uint32_t n_slots, + uint32_t layer) { + for (uint32_t i = 0; i < n_slots; i++) { + if (slots[i].active && slots[i].layer == layer) return &slots[i]; + } + return NULL; +} + +static metal_graph_stream_prepare_slot *metal_graph_stream_prepare_slot_free( + metal_graph_stream_prepare_slot *slots, + uint32_t n_slots) { + for (uint32_t i = 0; i < n_slots; i++) { + if (!slots[i].active) return &slots[i]; + } + return NULL; +} + +static bool metal_graph_stream_prepare_start_if_needed( + const ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t layer, + uint32_t n_tokens, + bool madvise_only, + bool pread_only, + bool readahead_only, + bool decode_only, + metal_graph_stream_prepare_slot *slots, + uint32_t n_slots) { + if (layer >= DS4_N_LAYER) return true; + if (metal_graph_stream_prepare_slot_find(slots, n_slots, layer)) { + return true; + } + metal_graph_stream_prepare_slot *slot = + metal_graph_stream_prepare_slot_free(slots, n_slots); + if (!slot) { + fprintf(stderr, + "ds4: Metal streaming prefill prepare queue is full before layer %u\n", + layer); + return false; + } + memset(slot, 0, sizeof(*slot)); + slot->layer = layer; + if (!metal_graph_stream_prefill_layer_pagein_start(g, + model, + weights, + layer, + n_tokens, + madvise_only, + pread_only, + readahead_only, + decode_only, + &slot->job)) { + memset(slot, 0, sizeof(*slot)); + return false; + } + slot->active = slot->job.started; + return true; +} + +static bool metal_graph_stream_prepare_join_layer( + const ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t layer, + uint32_t n_tokens, + bool madvise_only, + bool pread_only, + bool readahead_only, + bool decode_only, + metal_graph_stream_prepare_slot *slots, + uint32_t n_slots) { + metal_graph_stream_prepare_slot *slot = + metal_graph_stream_prepare_slot_find(slots, n_slots, layer); + if (!slot) { + metal_graph_stream_pagein_job job; + memset(&job, 0, sizeof(job)); + if (!metal_graph_stream_prefill_layer_pagein_start(g, + model, + weights, + layer, + n_tokens, + madvise_only, + pread_only, + readahead_only, + decode_only, + &job)) { + return false; + } + return metal_graph_stream_prefill_layer_pagein_join(&job); + } + const bool ok = metal_graph_stream_prefill_layer_pagein_join(&slot->job); + memset(slot, 0, sizeof(*slot)); + return ok; +} + +static bool metal_graph_stream_prepare_join_all( + metal_graph_stream_prepare_slot *slots, + uint32_t n_slots) { + bool ok = true; + for (uint32_t i = 0; i < n_slots; i++) { + if (!slots[i].active) continue; + if (!metal_graph_stream_prefill_layer_pagein_join(&slots[i].job)) { + ok = false; + } + memset(&slots[i], 0, sizeof(slots[i])); + } + return ok; +} + +static void metal_graph_stream_readahead_layer( + const ds4_model *model, + const ds4_weights *weights, + uint32_t il) { + ds4_model_map_span_vec spans; + if (!weights_model_map_spans(weights, il, il, false, &spans)) return; + metal_graph_stream_readahead_spans(model, &spans); + free(spans.v); +} + +static void metal_graph_stream_readahead_layer_decode( + const ds4_model *model, + const ds4_weights *weights, + uint32_t il) { + ds4_model_map_span_vec spans; + if (!weights_model_map_decode_layer_spans(weights, il, &spans)) return; + metal_graph_stream_readahead_spans(model, &spans); + free(spans.v); +} + +static void metal_graph_stream_readahead_output( + const ds4_model *model, + const ds4_weights *weights) { + ds4_model_map_span_vec spans; + if (!weights_model_map_output_spans(weights, &spans)) return; + metal_graph_stream_readahead_spans(model, &spans); + free(spans.v); +} + +static bool metal_graph_stream_prefill_selected_readahead_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + (glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD", + "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD") || + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", + "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED")) && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD"); +} + +static bool metal_graph_stream_prefill_selected_readahead_shared_enabled( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", + "DS4_METAL_ENABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD_SHARED") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_READAHEAD"); +} + +static uint32_t metal_graph_stream_prefill_selected_readahead_gap(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_PREFILL_SELECTED_READAHEAD_GAP", + "DS4_METAL_STREAMING_PREFILL_SELECTED_READAHEAD_GAP"); + if (!env || !env[0]) return 0; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0') return 0; + return v > 8 ? 8u : (uint32_t)v; +} + +static bool metal_graph_stream_readahead_selected_run( + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t first, + uint32_t last, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint64_t *hint_bytes) { + if (!model || !layer || first > last || last >= DS4_N_EXPERT) return false; + + const uint64_t first_id = first; + const uint64_t n_experts = (uint64_t)last - (uint64_t)first + 1u; + if (first_id > UINT64_MAX / gate_expert_bytes || + first_id > UINT64_MAX / down_expert_bytes || + n_experts > UINT64_MAX / gate_expert_bytes || + n_experts > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming prefill selected expert readahead overflow\n"); + return false; + } + + const uint64_t gate_rel = first_id * gate_expert_bytes; + const uint64_t down_rel = first_id * down_expert_bytes; + const uint64_t gate_bytes = n_experts * gate_expert_bytes; + const uint64_t down_bytes = n_experts * down_expert_bytes; + if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || + gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || + down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { + fprintf(stderr, "ds4: Metal streaming prefill selected expert readahead overflow\n"); + return false; + } + + metal_graph_stream_readahead_range_impl(model, + layer->ffn_gate_exps->abs_offset + gate_rel, + gate_bytes, + true); + metal_graph_stream_readahead_range_impl(model, + layer->ffn_up_exps->abs_offset + gate_rel, + gate_bytes, + true); + metal_graph_stream_readahead_range_impl(model, + layer->ffn_down_exps->abs_offset + down_rel, + down_bytes, + true); + if (hint_bytes) { + if (*hint_bytes > UINT64_MAX - gate_bytes || + *hint_bytes + gate_bytes > UINT64_MAX - gate_bytes || + *hint_bytes + gate_bytes * 2u > UINT64_MAX - down_bytes) { + *hint_bytes = UINT64_MAX; + } else { + *hint_bytes += gate_bytes * 2u + down_bytes; + } + } + return true; +} + +static bool metal_graph_stream_readahead_selected_experts_from_gpu( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!metal_graph_stream_prefill_selected_readahead_enabled(g)) return true; + if (!model || !layer || !g || !metal_graph_batch_router_selected(g) || n_tokens == 0) { + return false; + } + if (sizeof(int) != sizeof(int32_t) || DS4_N_EXPERT > DS4_MAX_EXPERT) { + return false; + } + + const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; + int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); + + const bool profile = + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_SELECTED_READAHEAD_PROFILE", + "DS4_METAL_STREAMING_PREFILL_SELECTED_READAHEAD_PROFILE"); + const double t0 = profile ? now_sec() : 0.0; + bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), + 0, + selected, + n_ids * sizeof(selected[0])) != 0; + bool seen[DS4_MAX_EXPERT] = { false }; + uint32_t unique = 0; + uint32_t ranges = 0; + uint64_t hint_bytes = 0; + const uint32_t gap = metal_graph_stream_prefill_selected_readahead_gap(); + if (ok) { + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t expert = selected[i]; + if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: Metal streaming prefill selected expert id %d is outside 0..%u at layer %u\n", + expert, + (uint32_t)DS4_N_EXPERT, + il); + ok = false; + break; + } + if (seen[expert]) continue; + seen[expert] = true; + unique++; + } + } + + if (ok) { + uint32_t e = 0; + while (e < DS4_N_EXPERT) { + while (e < DS4_N_EXPERT && !seen[e]) e++; + if (e >= DS4_N_EXPERT) break; + const uint32_t first = e; + uint32_t last = e; + uint32_t skipped = 0; + e++; + while (e < DS4_N_EXPERT) { + if (seen[e]) { + last = e; + skipped = 0; + } else if (skipped < gap) { + skipped++; + } else { + break; + } + e++; + } + + if (!metal_graph_stream_readahead_selected_run(model, + layer, + first, + last, + gate_expert_bytes, + down_expert_bytes, + &hint_bytes)) { + ok = false; + break; + } + ranges++; + } + } + if (profile) { + fprintf(stderr, + "ds4: Metal streaming prefill selected readahead layer=%u " + "tokens=%u unique=%u ranges=%u gap=%u hint=%.2f GiB time=%.3f ms\n", + il, + n_tokens, + unique, + ranges, + gap, + (double)hint_bytes / (1024.0 * 1024.0 * 1024.0), + (now_sec() - t0) * 1000.0); + } + free(selected); + return ok; +} + +static bool metal_graph_stream_map_token( + const ds4_model *model, + const ds4_weights *weights) { + ds4_model_map_span_vec spans; + if (!weights_model_map_token_spans(weights, &spans)) { + fprintf(stderr, "ds4: Metal SSD streaming could not build token embedding span\n"); + return false; + } + const bool ok = metal_graph_install_model_spans(model, &spans, "token embedding"); + free(spans.v); + return ok; +} + +static bool metal_graph_stream_map_decode_static_all( + const ds4_model *model, + const ds4_weights *weights) { + ds4_model_map_span_vec spans; + if (!weights_model_map_decode_static_spans(weights, true, true, &spans)) { + fprintf(stderr, "ds4: Metal SSD streaming could not build static decode spans\n"); + return false; + } + const bool ok = metal_graph_install_model_spans(model, &spans, "static decode"); + free(spans.v); + return ok; +} + +static bool metal_graph_stream_map_layer( + const ds4_model *model, + const ds4_weights *weights, + uint32_t il) { + ds4_model_map_span_vec spans; + if (!weights_model_map_spans(weights, il, il, false, &spans)) { + fprintf(stderr, "ds4: Metal SSD streaming could not build layer %u spans\n", il); + return false; + } + const bool ok = metal_graph_install_model_spans(model, &spans, "layer"); + free(spans.v); + return ok; +} + +static bool metal_graph_stream_map_layer_decode( + const ds4_model *model, + const ds4_weights *weights, + uint32_t il) { + ds4_model_map_span_vec spans; + if (!weights_model_map_decode_layer_spans(weights, il, &spans)) { + fprintf(stderr, "ds4: Metal SSD streaming could not build decode layer %u spans\n", il); + return false; + } + const bool ok = metal_graph_install_model_spans(model, &spans, "decode layer"); + free(spans.v); + return ok; +} + +static bool metal_graph_stream_map_output( + const ds4_model *model, + const ds4_weights *weights) { + ds4_model_map_span_vec spans; + if (!weights_model_map_output_spans(weights, &spans)) { + fprintf(stderr, "ds4: Metal SSD streaming could not build output head spans\n"); + return false; + } + const bool ok = metal_graph_install_model_spans(model, &spans, "output head"); + free(spans.v); + return ok; +} + +static uint32_t metal_graph_raw_span_for_batch( + const ds4_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens) { + if (!g || g->raw_cap == 0 || n_tokens == 0) return 0; + + const uint32_t window = g->raw_window ? g->raw_window : DS4_N_SWA; + const uint32_t last_pos = pos0 + n_tokens - 1u; + uint64_t needed = (uint64_t)n_tokens; + if (window != 0) { + needed += n_tokens == 1 ? (uint64_t)window - 1u : (uint64_t)window; + } + uint64_t available = (uint64_t)last_pos + 1u; + if (needed > available) needed = available; + if (needed > g->raw_cap) needed = g->raw_cap; + return (uint32_t)needed; +} + +static uint32_t metal_graph_raw_start_for_span( + const ds4_gpu_graph *g, + uint32_t last_pos, + uint32_t n_raw) { + if (!g || g->raw_cap == 0 || n_raw == 0) return 0; + const uint32_t first_raw_pos = last_pos + 1u - n_raw; + return first_raw_pos % g->raw_cap; +} + +static uint32_t metal_graph_decode_raw_score_count( + const ds4_gpu_graph *g, + uint32_t pos, + uint32_t n_raw, + uint32_t ratio) { + if (!g || n_raw == 0) return 0; + if (ratio == 0) return n_raw > 256u ? 256u : n_raw; + + const uint32_t first_raw_pos = pos + 1u - n_raw; + const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u; + uint32_t lo = first_raw_pos; + const uint32_t window = g->raw_window ? g->raw_window : DS4_N_SWA; + if (window != 0 && pos + 1u > window) { + const uint32_t wlo = pos + 1u - window; + if (wlo > lo) lo = wlo; + } + const uint32_t hi = pos < raw_last_pos ? pos : raw_last_pos; + if (hi < lo) return 0; + uint32_t raw_count = hi - lo + 1u; + if (raw_count > 256u) raw_count = 256u; + return raw_count; +} + +static bool metal_graph_cuda_splitkv_score_may_engage( + const ds4_gpu_graph *g, + uint32_t pos) { + if (!g) return false; + + const uint32_t min_score = metal_graph_cuda_greedy_splitkv_min_score(); + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + uint32_t visible_comp = 0; + const uint32_t n_comp = g->layer_n_comp[il]; + if (n_comp != 0) { + visible_comp = ratio == 0 ? n_comp : (pos + 1u) / ratio; + if (visible_comp > n_comp) visible_comp = n_comp; + } + const uint32_t raw_count = + metal_graph_decode_raw_score_count(g, pos, n_raw, ratio); + const uint32_t n_score = raw_count + visible_comp; + if (n_score > 1u && n_score >= min_score) return true; + } + return false; +} + +static bool metal_graph_cuda_greedy_splitkv_may_engage( + const ds4_gpu_graph *g, + uint32_t pos) { + if (!metal_graph_cuda_greedy_splitkv_requested()) return false; + return metal_graph_cuda_splitkv_score_may_engage(g, pos); +} + +/* Capture the verifier prefix after the first speculative token. + * + * Exact MTP speculation is only profitable if partial accepts are cheap. The + * target verifier computes two draft tokens together; if only the first token + * is accepted, replaying a one-token verifier throws away most of the gain. + * For compressed-attention layers the mutable frontier is just the small + * compressor state plus append counters, so we save that prefix-1 state while + * the N=2 verifier is already stepping the compressor token by token. + * + * Raw SWA rows are not captured here. This graph uses a raw ring larger than + * the 128-token logical SWA window, so writing speculative future rows does + * not evict visible raw rows. If the raw cache is ever reduced to a strict + * 128-row ring, speculative raw rows must become shadow rows and be copied + * into the ring only on commit. */ +static bool metal_graph_capture_prefix1_attn_state(ds4_gpu_graph *g, uint32_t il) { + if (!g->spec_capture_prefix1 || !g->spec_prefix1_attn_state_kv[il]) return true; + const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); + g->spec_prefix1_n_comp[il] = g->layer_n_comp[il]; + return ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_kv[il], 0, + g->layer_attn_state_kv[il], 0, bytes) != 0 && + ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_score[il], 0, + g->layer_attn_state_score[il], 0, bytes) != 0; +} + +static bool metal_graph_capture_prefix1_index_state(ds4_gpu_graph *g, uint32_t il) { + if (!g->spec_capture_prefix1 || !g->spec_prefix1_index_state_kv[il]) return true; + const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); + g->spec_prefix1_n_index_comp[il] = g->layer_n_index_comp[il]; + return ds4_gpu_tensor_copy(g->spec_prefix1_index_state_kv[il], 0, + g->layer_index_state_kv[il], 0, bytes) != 0 && + ds4_gpu_tensor_copy(g->spec_prefix1_index_state_score[il], 0, + g->layer_index_state_score[il], 0, bytes) != 0; +} + +static uint32_t metal_graph_decode_indexer_sparse_threshold(const ds4_gpu_graph *g) { + (void)g; + static int parsed = -1; + static uint32_t cached = 0; + if (parsed < 0) { + parsed = 0; +#ifndef DS4_ROCM_BUILD + const char *env = getenv("DS4_METAL_DECODE_INDEXER_SPARSE_THRESHOLD"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + while (end && isspace((unsigned char)*end)) end++; + if (end != env && end && *end == '\0' && + (v == 64ul || v == 128ul || v == 256ul || v == 512ul || + v == 1024ul || v == 2048ul || v == 4096ul)) { + cached = (uint32_t)v; + parsed = 1; + } else { + fprintf(stderr, + "ds4: invalid DS4_METAL_DECODE_INDEXER_SPARSE_THRESHOLD=%s; " + "expected 64, 128, 256, 512, 1024, 2048, or 4096\n", + env); + } + } +#endif + } + if (parsed > 0) return cached; + + /* Keep dense attention longer than the legacy 512-row window by default. + * Around the 2K frontier the sparse path's score/top-k setup dominates + * the smaller attention scan, while larger contexts benefit from sparse + * indexed attention. This threshold changes only the implementation used + * to consume the compressed rows; it must not lower the 512-row indexer + * selection defined by DS4_N_INDEXER_TOP_K. */ + return 1024u; +} + +/* ========================================================================= + * Metal Decode Release Helpers and Reference Fallbacks. + * ========================================================================= + * + * The normal generation path uses the fused helpers below. The older unfused + * kernels remain available as diagnostic reference paths selected only by the + * DS4_METAL_DISABLE_*_FUSION environment switches. + */ + +static bool metal_graph_env_flag(const char *name, int *cache) { + if (*cache == -1) { +#ifdef DS4_ROCM_BUILD + (void)name; + *cache = 0; +#else + const char *env = getenv(name); + *cache = env && env[0] && strcmp(env, "0") != 0; +#endif + } + return *cache != 0; +} + +static bool metal_graph_use_reference_hc_decode(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_HC_FUSION", &cache); +} + +static bool metal_graph_use_reference_kv_decode(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_KV_FUSION", &cache); +} + +static bool metal_graph_use_reference_qkv_norm(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_NORM_FUSION", &cache); +} + +static bool metal_graph_use_reference_qkv_pair_proj(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_PAIR_PROJ", &cache); +} + +static bool metal_graph_use_reference_compressor_pair_proj(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ", &cache); +} + +static bool metal_graph_use_reference_hc_norm_decode(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_HC_NORM_FUSION", &cache); +} + +static bool metal_graph_enable_batch_hc_norm_fusion(void) { + static int cache = -1; + if (metal_graph_use_reference_hc_norm_decode()) return false; + if (cache == -1) { + const char *disable = getenv("DS4_METAL_DISABLE_BATCH_HC_NORM_FUSION"); + if (disable && disable[0] && strcmp(disable, "0") != 0) { + cache = 0; + } else { + const char *legacy_enable = + getenv("DS4_METAL_ENABLE_BATCH_HC_NORM_FUSION"); + cache = (!legacy_enable || !legacy_enable[0] || + strcmp(legacy_enable, "0") != 0) ? 1 : 0; + } + } + return cache != 0; +} + +static bool metal_graph_use_reference_shared_down_hc(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_SHARED_DOWN_HC_FUSION", &cache); +} + +static bool metal_graph_use_reference_attn_out_hc(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_ATTN_OUT_HC_FUSION", &cache); +} + +static bool metal_graph_decode_hc_pre( + ds4_gpu_tensor *out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const ds4_model *model, + uint64_t scale_offset, + uint64_t base_offset) { + if (metal_graph_use_reference_hc_decode()) { + return ds4_gpu_hc_split_sinkhorn_tensor(split, + mix, + model->map, + model->size, + scale_offset, + base_offset, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0 && + ds4_gpu_hc_weighted_sum_tensor(out, + residual_hc, + split, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + + return ds4_gpu_hc_split_weighted_sum_tensor(out, + split, + mix, + residual_hc, + model->map, + model->size, + scale_offset, + base_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; +} + +static bool metal_graph_hc_norm_fusion_check_enabled(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_HC_NORM_FUSION_CHECK", &cache); +} + +static float metal_graph_hc_norm_fusion_check_tolerance(void) { + static int initialized; + static float tolerance; + if (initialized) return tolerance; + tolerance = 2.0e-4f; +#ifndef DS4_ROCM_BUILD + const char *env = getenv("DS4_METAL_HC_NORM_FUSION_CHECK_TOL"); + if (env && env[0]) { + char *end = NULL; + const float v = strtof(env, &end); + if (end != env && isfinite(v) && v > 0.0f) tolerance = v; + } +#endif + initialized = 1; + return tolerance; +} + +static bool metal_graph_check_hc_norm_fusion( + const char *label, + ds4_gpu_tensor *fused_out, + ds4_gpu_tensor *fused_norm, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const ds4_model *model, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t il, + uint32_t pos) { + if (!metal_graph_hc_norm_fusion_check_enabled()) return true; + if (!fused_out || !fused_norm || !mix || !residual_hc || !model) return false; + + const uint64_t n_embd = DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + ds4_gpu_tensor *ref_split = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); + ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(n_embd * sizeof(float)); + ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(n_embd * sizeof(float)); + bool ok = ref_split && ref_out && ref_norm; + + if (ok) { + ok = ds4_gpu_hc_split_sinkhorn_tensor(ref_split, + mix, + model->map, + model->size, + scale_offset, + base_offset, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0 && + ds4_gpu_hc_weighted_sum_tensor(ref_out, + residual_hc, + ref_split, + DS4_N_EMBD, + DS4_N_HC) != 0 && + ds4_gpu_rms_norm_weight_tensor(ref_norm, + ref_out, + model->map, + model->size, + norm_weight_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } + + if (ok) ok = ds4_gpu_end_commands() != 0; + + float *fused_out_cpu = NULL; + float *ref_out_cpu = NULL; + float *fused_norm_cpu = NULL; + float *ref_norm_cpu = NULL; + if (ok) { + fused_out_cpu = xmalloc((size_t)n_embd * sizeof(float)); + ref_out_cpu = xmalloc((size_t)n_embd * sizeof(float)); + fused_norm_cpu = xmalloc((size_t)n_embd * sizeof(float)); + ref_norm_cpu = xmalloc((size_t)n_embd * sizeof(float)); + ok = ds4_gpu_tensor_read(fused_out, 0, fused_out_cpu, n_embd * sizeof(float)) != 0 && + ds4_gpu_tensor_read(ref_out, 0, ref_out_cpu, n_embd * sizeof(float)) != 0 && + ds4_gpu_tensor_read(fused_norm, 0, fused_norm_cpu, n_embd * sizeof(float)) != 0 && + ds4_gpu_tensor_read(ref_norm, 0, ref_norm_cpu, n_embd * sizeof(float)) != 0; + } + + if (ok) { + const float out_max = max_abs_diff(fused_out_cpu, ref_out_cpu, n_embd); + const float out_rms = rms_abs_diff(fused_out_cpu, ref_out_cpu, n_embd); + const float norm_max = max_abs_diff(fused_norm_cpu, ref_norm_cpu, n_embd); + const float norm_rms = rms_abs_diff(fused_norm_cpu, ref_norm_cpu, n_embd); + const float tol = metal_graph_hc_norm_fusion_check_tolerance(); + fprintf(stderr, + "ds4: Metal HC norm fusion check %s layer=%u pos=%u " + "out_max=%g out_rms=%g norm_max=%g norm_rms=%g tol=%g\n", + label ? label : "hc", + il, + pos, + out_max, + out_rms, + norm_max, + norm_rms, + tol); + if (out_max > tol || norm_max > tol) { + fprintf(stderr, + "ds4: Metal HC norm fusion check failed for %s layer=%u pos=%u\n", + label ? label : "hc", + il, + pos); + ok = false; + } + } + + free(fused_out_cpu); + free(ref_out_cpu); + free(fused_norm_cpu); + free(ref_norm_cpu); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(ref_split); + + const bool restart_ok = ds4_gpu_begin_commands() != 0; + return ok && restart_ok; +} + +static bool metal_graph_decode_kv_store( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t raw_row) { + if (metal_graph_use_reference_kv_decode()) { + return ds4_gpu_dsv4_fp8_kv_quantize_tensor(kv, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0 && + ds4_gpu_store_raw_kv_tensor(raw_cache, kv, raw_cap, raw_row, DS4_N_HEAD_DIM) != 0; + } + + return ds4_gpu_kv_fp8_store_raw_tensor(kv, + raw_cache, + raw_cap, + raw_row, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; +} + +static uint64_t metal_graph_attn_comp_cache_row_bytes(void) { + return (uint64_t)DS4_N_HEAD_DIM * + (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); +} + +static uint32_t metal_graph_attn_comp_cache_is_f16(void) { + return DS4_GPU_ATTN_COMP_CACHE_F16 ? 1u : 0u; +} + +static bool metal_graph_store_attn_comp_stage( + ds4_gpu_graph *g, + uint32_t il, + uint32_t first_row, + uint32_t rows) { + if (!g || il >= DS4_N_LAYER) return false; + if (rows == 0) return true; + if (!g->layer_attn_comp_cache[il] || !metal_graph_attn_comp_stage(g)) return false; + if (rows > g->attn_comp_stage_cap || first_row > g->layer_comp_cap[il] || + rows > g->layer_comp_cap[il] - first_row) { + return false; + } + + const uint64_t count = (uint64_t)rows * DS4_N_HEAD_DIM; + const uint64_t dst_offset = (uint64_t)first_row * + metal_graph_attn_comp_cache_row_bytes(); + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + return ds4_gpu_tensor_copy_f32_to_f16(g->layer_attn_comp_cache[il], + dst_offset, + metal_graph_attn_comp_stage(g), + 0, + count) != 0; + } + + return ds4_gpu_tensor_copy(g->layer_attn_comp_cache[il], + dst_offset, + metal_graph_attn_comp_stage(g), + 0, + count * sizeof(float)) != 0; +} + +static ds4_gpu_tensor *metal_graph_attn_comp_update_target( + ds4_gpu_graph *g, + uint32_t il) { + return DS4_GPU_ATTN_COMP_CACHE_F16 + ? metal_graph_attn_comp_stage(g) + : g->layer_attn_comp_cache[il]; +} + +static uint32_t metal_graph_attn_comp_update_row(uint32_t row) { + return DS4_GPU_ATTN_COMP_CACHE_F16 ? 0u : row; +} + +static bool metal_graph_commit_attn_comp_stage( + ds4_gpu_graph *g, + uint32_t il, + uint32_t first_row, + uint32_t rows) { + if (!DS4_GPU_ATTN_COMP_CACHE_F16) return true; + return metal_graph_store_attn_comp_stage(g, il, first_row, rows); +} + +static ds4_gpu_tensor *metal_graph_attn_comp_row_view( + ds4_gpu_graph *g, + uint32_t il, + uint32_t row) { + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + return ds4_gpu_tensor_view(metal_graph_attn_comp_stage(g), + 0, + (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + } + return ds4_gpu_tensor_view(g->layer_attn_comp_cache[il], + (uint64_t)row * DS4_N_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); +} + +static ds4_gpu_tensor *metal_graph_attn_comp_prefill_target( + ds4_gpu_graph *g, + uint32_t il, + uint32_t first_row, + uint32_t rows) { + if (DS4_GPU_ATTN_COMP_CACHE_F16) return metal_graph_attn_comp_stage(g); + const uint32_t view_rows = rows ? rows : 1u; + return ds4_gpu_tensor_view(g->layer_attn_comp_cache[il], + (uint64_t)first_row * DS4_N_HEAD_DIM * sizeof(float), + (uint64_t)view_rows * DS4_N_HEAD_DIM * sizeof(float)); +} + +static void metal_graph_attn_comp_prefill_target_free(ds4_gpu_tensor *t) { + if (!DS4_GPU_ATTN_COMP_CACHE_F16) ds4_gpu_tensor_free(t); +} + +static bool metal_graph_cuda_tp_attn_cache_dup_layer_ready( + const ds4_gpu_graph *g, + uint32_t il) { + if (!g || il >= DS4_N_LAYER || !g->cuda_tp_attn_cache_dup) return false; + if (!g->placement || !g->layer_raw_cache[il] || !g->layer_raw_cache_tp[il]) { + return false; + } + const int layer_tier = g->placement[il + 1]; + if (metal_graph_cuda_tp_partner_tier(layer_tier) < 0) return false; + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0 && + (!g->layer_attn_comp_cache[il] || + !g->layer_attn_comp_cache_tp[il])) { + return false; + } + return true; +} + +static bool metal_graph_cuda_tp_attn_cache_copy_row( + ds4_gpu_tensor *dst_base, + const ds4_gpu_tensor *src_base, + uint64_t offset, + uint64_t bytes) { + if (bytes == 0) return true; + ds4_gpu_tensor *dst = ds4_gpu_tensor_view(dst_base, offset, bytes); + ds4_gpu_tensor *src = ds4_gpu_tensor_view(src_base, offset, bytes); + bool ok = dst && src && ds4_gpu_tensor_copy_xdev(dst, src, bytes) != 0; + ds4_gpu_tensor_free(src); + ds4_gpu_tensor_free(dst); + return ok; +} + +static bool metal_graph_cuda_tp_attn_cache_sync_raw_row( + ds4_gpu_graph *g, + uint32_t il, + uint32_t raw_row) { + if (!metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il)) return true; + if (raw_row >= g->raw_cap) return false; + const uint64_t row_bytes = (uint64_t)DS4_N_HEAD_DIM * sizeof(float); + return metal_graph_cuda_tp_attn_cache_copy_row( + g->layer_raw_cache_tp[il], + g->layer_raw_cache[il], + (uint64_t)raw_row * row_bytes, + row_bytes); +} + +static bool metal_graph_cuda_tp_attn_cache_sync_all(ds4_gpu_graph *g) { + if (!g || !g->cuda_tp_attn_cache_dup) return true; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il)) return false; + const uint64_t raw_bytes = + (uint64_t)g->raw_cap * DS4_N_HEAD_DIM * sizeof(float); + if (!ds4_gpu_tensor_copy_xdev(g->layer_raw_cache_tp[il], + g->layer_raw_cache[il], + raw_bytes)) { + return false; + } + const uint32_t ratio = ds4_layer_compress_ratio(il); + const uint32_t n_comp = g->layer_n_comp[il]; + if (ratio != 0 && n_comp != 0) { + const uint64_t comp_bytes = + (uint64_t)n_comp * metal_graph_attn_comp_cache_row_bytes(); + if (!ds4_gpu_tensor_copy_xdev(g->layer_attn_comp_cache_tp[il], + g->layer_attn_comp_cache[il], + comp_bytes)) { + return false; + } + } + } + return true; +} + +/* Encode one DS4 decode layer on Metal. This is the release single-token + * layer path; diagnostics reuse it so they compare exactly what generation + * runs. */ +static bool metal_graph_indexer_stage_profile_boundary( + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_comp, + double *stage_t0); +static bool metal_graph_layer_stage_profile_boundary( + const char *part, + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + double *stage_t0); +static bool metal_graph_decode_stage_profile_enabled(uint32_t il); +static bool metal_graph_matmul_plain_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); +static bool metal_graph_matmul_dense_quant_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); +static bool metal_graph_dense_quant_row_bytes( + const ds4_tensor *w, + uint64_t in_dim, + uint64_t *row_bytes); +static bool metal_graph_matmul_dense_quant_abs( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); +static bool metal_graph_matmul_dense_quant_kslice( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + 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); +static bool metal_graph_attention_output_dense_quant_low( + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + uint64_t group_dim, + uint64_t rank, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads); +static bool metal_graph_attention_output_dense_quant_tp( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + const ds4_tensor *out_b, + 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); +static bool metal_graph_attention_output_dense_quant_batch( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + const ds4_tensor *out_b, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens); + +static bool metal_graph_use_pro_q4_cpu_router(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_PRO_Q4_CPU_ROUTER", &cache); +} + +static bool metal_graph_use_streaming_iq2_cpu_router(void) { + return getenv("DS4_METAL_ENABLE_STREAMING_IQ2_CPU_ROUTER") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_IQ2_CPU_ROUTER") == NULL; +} + +static bool metal_graph_use_q4_selected_shared_overlap(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_Q4_SELECTED_OVERLAP_SHARED", &cache); +} + +static bool metal_graph_use_cuda_selected_shared_overlap(const ds4_gpu_graph *g) { +#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) + return g && + g->ssd_streaming && + getenv("DS4_CUDA_DISABLE_STREAMING_SELECTED_SHARED_OVERLAP") == NULL; +#else + (void)g; + return false; +#endif +} + +static bool metal_graph_q4_non_streaming_opt_in_enabled(void) { + return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL; +} + +static bool metal_graph_q4_selected_paths_allowed(const ds4_gpu_graph *g) { + if (!g) return false; + if (g->ssd_streaming) return true; + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) return false; + return metal_graph_q4_non_streaming_opt_in_enabled(); +} + +static bool metal_graph_use_iq2_selected_shared_overlap(const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_SHARED_OVERLAP") == NULL && + getenv("DS4_METAL_DISABLE_IQ2_SELECTED_SHARED_OVERLAP") == NULL; +} + +static bool metal_graph_use_iq2_selected_async_load(const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && +#ifndef DS4_ROCM_BUILD + getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_LOAD") == NULL; +#else + true; +#endif +} + +static bool metal_graph_use_iq2_selected_async_early_commit( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && +#ifndef DS4_ROCM_BUILD + getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_EARLY_COMMIT") == NULL; +#else + false; +#endif +} + +static bool metal_graph_use_pro_q4_expert_table_auto(const ds4_gpu_graph *g) { + if (getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") != NULL) { + return false; + } + if (!g || (!g->ssd_streaming && + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL)) { + return false; + } +#ifndef DS4_NO_GPU + return ds4_gpu_pro_q4_expert_table_auto_available() != 0; +#else + return false; +#endif +} + +static bool metal_graph_decode_cpu_router_applicable( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer) { + const bool pro_q4 = + DS4_MODEL_VARIANT == DS4_VARIANT_PRO && + metal_graph_use_pro_q4_cpu_router() && + layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q4_K; + const bool streaming_iq2 = + g && + g->ssd_streaming && + !g->quality && + metal_graph_use_streaming_iq2_cpu_router() && + layer->ffn_gate_tid2eid == NULL && + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && + DS4_N_EXPERT_USED == 6 && + DS4_N_EXPERT >= 128 && + !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && + !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && + !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", + "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); + return pro_q4 || streaming_iq2; +} + +static bool metal_graph_decode_pro_q4_expert_table_expected( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + return !g->quality && + DS4_MODEL_VARIANT == DS4_VARIANT_PRO && + metal_graph_q4_selected_paths_allowed(g) && + layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && + DS4_N_EXPERT == 384 && + DS4_N_EXPERT_USED == 6 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && + !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && + (metal_graph_use_pro_q4_expert_table_auto(g) || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL) && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL; +} + +static bool metal_graph_decode_q4_selected_slots_expected( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + if (metal_graph_decode_pro_q4_expert_table_expected(g, layer, + gate_tensor_bytes, + down_tensor_bytes)) { + return false; + } + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + return !g->quality && + metal_graph_q4_selected_paths_allowed(g) && + layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && + DS4_N_EXPERT_USED == 6 && + DS4_N_EXPERT >= 128 && + (g->ssd_streaming || + (gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes)) && + !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && + !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && + !glm_graph_env_present("DS4_ROCM_DISABLE_Q4_SELECTED_EXPERT_VIEWS", + "DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS"); +} + +static bool metal_graph_decode_iq2_selected_slots_expected( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer) { + return g && + g->ssd_streaming && + !g->quality && + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && + DS4_N_EXPERT_USED == 6 && + DS4_N_EXPERT >= 128 && + !glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", + "DS4_METAL_MOE_WRITE_CLAMPED_ACT") && + !glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", + "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") && + !glm_graph_env_present("DS4_ROCM_DISABLE_IQ2_SELECTED_EXPERT_VIEWS", + "DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS"); +} + +static bool metal_graph_streaming_expert_cache_seed_layer_expected( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer) { + if (!g || + !g->ssd_streaming || + !layer || + !layer->ffn_gate_exps || + !layer->ffn_up_exps || + !layer->ffn_down_exps) { + return false; + } + if (metal_graph_decode_iq2_selected_slots_expected(g, layer)) return true; + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + !g->quality && + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS && + DS4_N_EXPERT_USED != 0 && + DS4_N_EXPERT_USED <= 8 && + DS4_N_EXPERT >= 128 && + !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", + "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { + return true; + } + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || + g->quality || + layer->ffn_gate_exps->type != layer->ffn_up_exps->type || + layer->ffn_gate_exps->type != layer->ffn_down_exps->type || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > 8 || + DS4_N_EXPERT < 128 || + glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", + "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { + return false; + } + const uint32_t type = layer->ffn_gate_exps->type; + return type == DS4_TENSOR_Q2_K || type == DS4_TENSOR_Q4_K; +} + +static bool metal_graph_decode_cuda_selected_slots_expected( + const ds4_gpu_graph *g, + const ds4_layer_weights *layer) { +#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) + if (!g || + !g->ssd_streaming || + g->quality || + !layer || + !layer->ffn_gate_exps || + !layer->ffn_up_exps || + !layer->ffn_down_exps || + DS4_N_EXPERT_USED != 6 || + DS4_N_EXPERT < 128 || + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) { + return false; + } + const bool q4 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_up_exps->type == DS4_TENSOR_Q4_K && + layer->ffn_down_exps->type == DS4_TENSOR_Q4_K && + getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL; + const bool iq2 = + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K && + getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; + return q4 || iq2; +#else + (void)g; + (void)layer; + return false; +#endif +} + +static uint32_t metal_graph_streaming_prefill_cache_seed_k(const ds4_gpu_graph *g) { + const bool enabled = + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_CACHE_SEED", + "DS4_METAL_ENABLE_STREAMING_PREFILL_CACHE_SEED"); + if (!g || + !g->ssd_streaming || + !enabled) { + return 0; + } + + uint32_t k = 1; + const char *env = glm_graph_env_value("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_K", + "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_K"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0') { + if (v == 0) return 0; + k = v > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS ? + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS : (uint32_t)v; + } + } + return k; +} + +static bool metal_graph_streaming_prefill_cache_seed_enabled(const ds4_gpu_graph *g) { + return metal_graph_streaming_prefill_cache_seed_k(g) != 0; +} + +static bool metal_graph_streaming_expert_hotlist_enabled(const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + !g->ssd_streaming_cold && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_EXPERT_HOTLIST", + "DS4_METAL_DISABLE_STREAMING_EXPERT_HOTLIST"); +} + +static bool metal_graph_streaming_expert_hotlist_add( + uint32_t layer, + uint32_t expert, + uint32_t priority, + int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t counts[DS4_MAX_LAYER], + bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t *loaded) { + if (layer >= DS4_N_LAYER || expert >= DS4_N_EXPERT) return true; + if (layer >= DS4_MAX_LAYER || expert >= DS4_MAX_EXPERT) return true; + if (seen[layer][expert]) return true; + if (counts[layer] >= DS4_MAX_EXPERT) return false; + seen[layer][expert] = true; + if (priority == 0) priority = 1; + priorities[layer][counts[layer]] = priority; + experts[layer][counts[layer]++] = (int32_t)expert; + (*loaded)++; + return true; +} + +static bool metal_graph_streaming_expert_hotlist_load_file( + const char *path, + uint32_t max_entries, + int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t counts[DS4_MAX_LAYER], + bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t *loaded_out) { + if (!path || !path[0] || max_entries == 0 || + !experts || !priorities || !counts || !seen || !loaded_out) { + return false; + } + FILE *fp = fopen(path, "rb"); + if (!fp) { + fprintf(stderr, + "ds4: failed to open streaming expert hotlist %s: %s\n", + path, + strerror(errno)); + return false; + } + + char line[256]; + uint64_t lineno = 0; + uint32_t loaded = 0; + while (fgets(line, sizeof(line), fp)) { + if (loaded >= max_entries) break; + lineno++; + char *p = line; + while (*p && isspace((unsigned char)*p)) p++; + if (*p == '\0' || *p == '#') continue; + + errno = 0; + char *end = NULL; + unsigned long layer = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long expert = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long long hits = strtoull(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + if (hits == 0) continue; + const uint32_t priority = + hits > UINT32_MAX ? UINT32_MAX : (uint32_t)hits; + if (!metal_graph_streaming_expert_hotlist_add((uint32_t)layer, + (uint32_t)expert, + priority, + experts, + priorities, + counts, + seen, + &loaded)) { + goto bad_line; + } + continue; + +bad_line: + fprintf(stderr, + "ds4: invalid streaming expert hotlist line %" PRIu64 " in %s\n", + lineno, + path); + fclose(fp); + return false; + } + if (ferror(fp)) { + fprintf(stderr, + "ds4: failed to read streaming expert hotlist %s: %s\n", + path, + strerror(errno)); + fclose(fp); + return false; + } + fclose(fp); + + if (loaded == 0) { + fprintf(stderr, "ds4: streaming expert hotlist %s had no usable nonzero entries\n", path); + } + *loaded_out = loaded; + return true; +} + +static bool metal_graph_streaming_expert_hotlist_load_default( + uint32_t max_entries, + int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t counts[DS4_MAX_LAYER], + bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT], + uint32_t *loaded_out) { + if (max_entries == 0 || !experts || !priorities || !counts || !seen || !loaded_out) { + return false; + } + const uint16_t (*hotlist)[2] = NULL; + uint32_t hotlist_count = 0; + if (g_ds4_shape.variant == DS4_VARIANT_PRO) { + hotlist = ds4_default_streaming_hotlist_pro; + hotlist_count = ds4_default_streaming_hotlist_pro_count; + } else if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { + hotlist = ds4_default_streaming_hotlist_flash; + hotlist_count = ds4_default_streaming_hotlist_flash_count; + } else if (g_ds4_shape.variant == DS4_VARIANT_GLM52) { + hotlist = ds4_default_streaming_hotlist_glm52; + hotlist_count = ds4_default_streaming_hotlist_glm52_count; + } else { + *loaded_out = 0; + return true; + } + uint32_t loaded = 0; + for (uint32_t i = 0; + i < hotlist_count && loaded < max_entries; + i++) { + if (!metal_graph_streaming_expert_hotlist_add( + hotlist[i][0], + hotlist[i][1], + max_entries - loaded, + experts, + priorities, + counts, + seen, + &loaded)) { + return false; + } + } + *loaded_out = loaded; + return true; +} + +static uint32_t metal_graph_streaming_expert_preload_count( + const ds4_gpu_graph *g, + uint32_t cache_budget) { + if (!g || cache_budget == 0) return 0; + uint32_t preload = g->streaming_preload_experts; + if (preload == 0) { + preload = cache_budget; + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_EXPERT_AUTO_PRELOAD_CAP", + "DS4_METAL_STREAMING_EXPERT_AUTO_PRELOAD_CAP"); +#ifdef DS4_ROCM_BUILD + if (g_ds4_shape.variant == DS4_VARIANT_GLM52 && + (!env || !env[0])) { + return 0; + } +#endif + /* Auto mode is a hot seed, not a request to synchronously fill the + * whole cache. Large Flash caches can otherwise spend startup doing + * thousands of preads into shared Metal buffers and trip the system + * watchdog before decode begins. ROCm GLM52 uses indexed batch prefill + * by default, which already populates the cache; explicit CLI preload + * counts and auto-preload env caps bypass that default. */ + uint32_t cap = 4096; + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0') { + cap = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v; + } + } + if (cap != 0 && preload > cap) preload = cap; + } + if (preload > cache_budget) preload = cache_budget; + const uint64_t max_possible = (uint64_t)DS4_N_LAYER * DS4_N_EXPERT; + if ((uint64_t)preload > max_possible) preload = (uint32_t)max_possible; + return preload; +} + +static bool metal_graph_decode_set_hash_selected_override( + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t token, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes, + const ds4_gpu_graph *g) { + if (!layer->ffn_gate_tid2eid) return true; + + const bool q4_selected = + metal_graph_decode_q4_selected_slots_expected(g, + layer, + gate_tensor_bytes, + down_tensor_bytes); + const bool iq2_selected = + metal_graph_decode_iq2_selected_slots_expected(g, layer); + if (!q4_selected && !iq2_selected) { + return true; + } + + int selected[DS4_MAX_EXPERT_USED]; + int32_t selected_i32[DS4_MAX_EXPERT_USED]; + layer_hash_selected_experts(selected, model, layer, (int)token); + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + selected_i32[i] = (int32_t)selected[i]; + } + if (g && g->ssd_streaming) { + if (DS4_N_EXPERT == 0 || + gate_tensor_bytes % DS4_N_EXPERT != 0 || + down_tensor_bytes % DS4_N_EXPERT != 0) { + return false; + } + const uint64_t gate_expert_bytes = gate_tensor_bytes / DS4_N_EXPERT; + const uint64_t down_expert_bytes = down_tensor_bytes / DS4_N_EXPERT; + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + if (ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_i32, + DS4_N_EXPERT_USED) == 0) { + return false; + } + } + return ds4_gpu_routed_moe_set_selected_override(selected_i32, DS4_N_EXPERT_USED) != 0; +} + +static bool metal_graph_decode_cpu_router( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t token) { + const bool profile = + getenv("DS4_METAL_PRO_Q4_CPU_ROUTER_PROFILE") != NULL || + getenv("DS4_METAL_STREAMING_IQ2_CPU_ROUTER_PROFILE") != NULL; + const double t0 = profile ? now_sec() : 0.0; + if (ds4_gpu_end_commands() == 0) return false; + const double t_sync = profile ? now_sec() : 0.0; + if (ds4_gpu_tensor_read(metal_graph_ffn_norm(g), + 0, + g->cpu_router_norm, + (uint64_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])) == 0) { + return false; + } + const double t_read = profile ? now_sec() : 0.0; + + float logits[DS4_MAX_EXPERT]; + float probs[DS4_MAX_EXPERT]; + int selected[DS4_MAX_EXPERT_USED]; + int32_t selected_i32[DS4_MAX_EXPERT_USED]; + float weights[DS4_MAX_EXPERT_USED]; + + matvec_any(logits, model, layer->ffn_gate_inp, g->cpu_router_norm); + for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { + probs[i] = sqrtf(softplus_stable(logits[i])); + } + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(selected, model, layer, (int)token); + layer_hash_router_weights_from_probs(weights, probs, selected); + } else { + layer_topk_selected_experts_from_probs(selected, weights, model, layer, probs); + } + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + selected_i32[i] = (int32_t)selected[i]; + } + const double t_cpu = profile ? now_sec() : 0.0; + + if (ds4_gpu_tensor_write(metal_graph_router_logits(g), + 0, + logits, + (uint64_t)DS4_N_EXPERT * sizeof(logits[0])) == 0 || + ds4_gpu_tensor_write(metal_graph_router_probs(g), + 0, + probs, + (uint64_t)DS4_N_EXPERT * sizeof(probs[0])) == 0 || + ds4_gpu_tensor_write(metal_graph_router_selected(g), + 0, + selected_i32, + (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_i32[0])) == 0 || + ds4_gpu_tensor_write(metal_graph_router_weights(g), + 0, + weights, + (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) == 0) { + return false; + } + const double t_write = profile ? now_sec() : 0.0; + if (ds4_gpu_begin_commands() == 0) return false; + if (ds4_gpu_routed_moe_set_selected_override(selected_i32, DS4_N_EXPERT_USED) == 0) return false; + + if (profile) { + fprintf(stderr, + "ds4: Metal CPU router layer=%u gate=%s down=%s sync=%.3f ms read=%.3f ms cpu=%.3f ms write=%.3f ms total=%.3f ms\n", + il, + tensor_type_name(layer->ffn_gate_exps->type), + tensor_type_name(layer->ffn_down_exps->type), + (t_sync - t0) * 1000.0, + (t_read - t_sync) * 1000.0, + (t_cpu - t_read) * 1000.0, + (t_write - t_cpu) * 1000.0, + (t_write - t0) * 1000.0); + } + return true; +} + +static bool metal_graph_use_iq2_selected_readahead_shared_delay( + const ds4_gpu_graph *g) { + return g && + g->ssd_streaming && + getenv("DS4_METAL_ENABLE_STREAMING_SELECTED_READAHEAD_SHARED_DELAY") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_SELECTED_READAHEAD_SHARED_DELAY") == NULL; +} + +static bool metal_graph_decode_selected_readahead_override( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!g || !model || !layer || !metal_graph_router_selected(g) || + DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || + DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return false; + } + + const bool profile = + getenv("DS4_METAL_STREAMING_SELECTED_READAHEAD_PROFILE") != NULL; + const double t0 = profile ? now_sec() : 0.0; + if (ds4_gpu_end_commands() == 0) return false; + const double t_sync = profile ? now_sec() : 0.0; + + int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; + if (ds4_gpu_tensor_read(metal_graph_router_selected(g), + 0, + selected_ids, + (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) == 0) { + return false; + } + const double t_read = profile ? now_sec() : 0.0; + + bool seen[DS4_MAX_EXPERT] = {0}; + uint32_t unique = 0; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: Metal streaming selected readahead expert id %d is outside 0..%u at layer %u\n", + selected_ids[i], + DS4_N_EXPERT, + il); + return false; + } + const uint32_t expert = (uint32_t)selected_ids[i]; + if (seen[expert]) continue; + seen[expert] = true; + unique++; + + const uint64_t expert_id = (uint64_t)expert; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming selected readahead offset overflow\n"); + return false; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - layer->ffn_gate_exps->abs_offset || + gate_rel > UINT64_MAX - layer->ffn_up_exps->abs_offset || + down_rel > UINT64_MAX - layer->ffn_down_exps->abs_offset) { + fprintf(stderr, "ds4: Metal streaming selected readahead offset overflow\n"); + return false; + } + + metal_graph_stream_readahead_range_impl(model, + layer->ffn_gate_exps->abs_offset + gate_rel, + gate_expert_bytes, + true); + metal_graph_stream_readahead_range_impl(model, + layer->ffn_up_exps->abs_offset + gate_rel, + gate_expert_bytes, + true); + metal_graph_stream_readahead_range_impl(model, + layer->ffn_down_exps->abs_offset + down_rel, + down_expert_bytes, + true); + } + const double t_hint = profile ? now_sec() : 0.0; + + if (ds4_gpu_routed_moe_set_selected_override(selected_ids, + DS4_N_EXPERT_USED) == 0) { + return false; + } + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + if (ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_ids, + DS4_N_EXPERT_USED) == 0) { + return false; + } + if (ds4_gpu_begin_commands() == 0) return false; + const double t_done = profile ? now_sec() : 0.0; + + if (profile) { + fprintf(stderr, + "ds4: Metal streaming selected readahead layer=%u unique=%u sync=%.3f ms read=%.3f ms hint=%.3f ms resume=%.3f ms total=%.3f ms\n", + il, + unique, + (t_sync - t0) * 1000.0, + (t_read - t_sync) * 1000.0, + (t_hint - t_read) * 1000.0, + (t_done - t_hint) * 1000.0, + (t_done - t0) * 1000.0); + } + return true; +} + +static bool metal_graph_decode_cuda_selected_load( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { +#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) + if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || + !model || + !metal_graph_router_selected(g) || + DS4_N_EXPERT == 0 || + DS4_N_EXPERT > DS4_MAX_EXPERT || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return false; + } + + const bool profile = + getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_PROFILE") != NULL; + const double t0 = profile ? now_sec() : 0.0; + + if (ds4_gpu_end_commands() == 0) return false; + const double t_sync = profile ? now_sec() : 0.0; + + int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; + bool ok = ds4_gpu_tensor_read(metal_graph_router_selected(g), + 0, + selected_ids, + (uint64_t)DS4_N_EXPERT_USED * + sizeof(selected_ids[0])) != 0; + const double t_read = profile ? now_sec() : 0.0; + + if (ok) { + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + ok = ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_ids, + DS4_N_EXPERT_USED) != 0; + } + const double t_load = profile ? now_sec() : 0.0; + + if (ds4_gpu_begin_commands() == 0) ok = false; + const double t_done = profile ? now_sec() : 0.0; + + if (profile) { + fprintf(stderr, + "ds4: CUDA streaming selected load layer=%u sync=%.3f ms read=%.3f ms load=%.3f ms resume=%.3f ms total=%.3f ms\n", + il, + (t_sync - t0) * 1000.0, + (t_read - t_sync) * 1000.0, + (t_load - t_read) * 1000.0, + (t_done - t_load) * 1000.0, + (t_done - t0) * 1000.0); + } + return ok; +#else + (void)g; + (void)model; + (void)layer; + (void)il; + (void)gate_expert_bytes; + (void)down_expert_bytes; + return false; +#endif +} + +static bool metal_graph_cuda_stream_prefill_batch_selected_load( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { +#if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) + if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || + !model || + !metal_graph_batch_router_selected(g) || + n_tokens <= 1 || + DS4_N_EXPERT == 0 || + DS4_N_EXPERT_USED == 0 || + getenv("DS4_CUDA_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_LOAD") != NULL) { + return true; + } + + if ((uint64_t)n_tokens > UINT64_MAX / (uint64_t)DS4_N_EXPERT_USED) { + fprintf(stderr, "ds4: CUDA streaming prefill selected-id count overflow at layer %u\n", il); + return false; + } + const uint64_t n_ids64 = (uint64_t)n_tokens * DS4_N_EXPERT_USED; + if (n_ids64 == 0 || n_ids64 > SIZE_MAX / sizeof(int32_t)) { + fprintf(stderr, "ds4: CUDA streaming prefill selected-id byte size overflow at layer %u\n", il); + return false; + } + + const bool profile = + getenv("DS4_CUDA_STREAMING_PREFILL_BATCH_SELECTED_PROFILE") != NULL; + const double t0 = profile ? now_sec() : 0.0; + + if (ds4_gpu_end_commands() == 0) return false; + const double t_sync = profile ? now_sec() : 0.0; + + int32_t *selected_ids = xmalloc((size_t)n_ids64 * sizeof(selected_ids[0])); + bool ok = ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), + 0, + selected_ids, + n_ids64 * sizeof(selected_ids[0])) != 0; + const double t_read = profile ? now_sec() : 0.0; + if (ok) { + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + ok = ds4_gpu_stream_expert_cache_prepare_selected_batch( + &table, + selected_ids, + n_tokens, + DS4_N_EXPERT_USED) != 0; + } + free(selected_ids); + const double t_load = profile ? now_sec() : 0.0; + + if (ds4_gpu_begin_commands() == 0) ok = false; + const double t_done = profile ? now_sec() : 0.0; + + if (profile) { + fprintf(stderr, + "ds4: CUDA streaming prefill batch selected load layer=%u tokens=%u sync=%.3f ms read=%.3f ms load=%.3f ms resume=%.3f ms total=%.3f ms\n", + il, + n_tokens, + (t_sync - t0) * 1000.0, + (t_read - t_sync) * 1000.0, + (t_load - t_read) * 1000.0, + (t_done - t_load) * 1000.0, + (t_done - t0) * 1000.0); + } + return ok; +#else + (void)g; + (void)model; + (void)layer; + (void)il; + (void)n_tokens; + (void)gate_expert_bytes; + (void)down_expert_bytes; + return true; +#endif +} + +typedef struct metal_graph_selected_async_load { + bool active; + bool ok; + /* Selected ids remain usable for a synchronous retry if the service + * thread cannot stage the cache load without waiting on GPU work. */ + bool ids_ok; + ds4_gpu_tensor *router_selected; + const ds4_model *model; + const ds4_layer_weights *layer; + uint32_t il; + uint64_t event_value; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t selected_ids[DS4_MAX_EXPERT_USED]; +} metal_graph_selected_async_load; + +static pthread_mutex_t g_metal_graph_selected_async_load_mutex = + PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_metal_graph_selected_async_load_cond = + PTHREAD_COND_INITIALIZER; +static pthread_cond_t g_metal_graph_selected_async_load_done_cond = + PTHREAD_COND_INITIALIZER; +static pthread_t g_metal_graph_selected_async_load_thread; +static bool g_metal_graph_selected_async_load_thread_started = false; +static bool g_metal_graph_selected_async_load_has_job = false; +static bool g_metal_graph_selected_async_load_done = false; +static metal_graph_selected_async_load g_metal_graph_selected_async_load_job; + +static void metal_graph_selected_async_load_run( + metal_graph_selected_async_load *job) { + job->ok = false; + + if (!job->router_selected || !job->model || !job->layer || + DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || + DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return; + } + if (job->event_value != 0) { +#ifdef DS4_ROCM_BUILD + if (ds4_gpu_tensor_read_after_selected_event( + job->router_selected, + 0, + job->selected_ids, + (uint64_t)DS4_N_EXPERT_USED * + sizeof(job->selected_ids[0]), + job->event_value, + "selected-id async expert load") == 0) { + return; + } +#else + if (ds4_gpu_wait_selected_readback_ready(job->event_value, + "selected-id async expert load") == 0) { + return; + } + if (ds4_gpu_tensor_read(job->router_selected, + 0, + job->selected_ids, + (uint64_t)DS4_N_EXPERT_USED * + sizeof(job->selected_ids[0])) == 0) { + return; + } +#endif + } + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + if (job->selected_ids[i] < 0 || + (uint32_t)job->selected_ids[i] >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: Metal streaming async selected expert id %d is outside 0..%u at layer %u\n", + job->selected_ids[i], + DS4_N_EXPERT, + job->il); + return; + } + } + job->ids_ok = true; + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(job->model, + job->layer, + job->il, + job->gate_expert_bytes, + job->down_expert_bytes); + if (ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + job->selected_ids, + DS4_N_EXPERT_USED) == 0) { + return; + } + + job->ok = true; +} + +static void *metal_graph_selected_async_load_worker_main(void *arg) { + (void)arg; +#ifdef __APPLE__ + /* The Metal cache paths must never wait on command buffers from this + * thread while the main thread is encoding; register it so those waits + * turn into load failures that the caller retries synchronously. */ + ds4_gpu_stream_expert_cache_note_service_thread(); +#endif + for (;;) { + pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); + while (!g_metal_graph_selected_async_load_has_job) { + pthread_cond_wait(&g_metal_graph_selected_async_load_cond, + &g_metal_graph_selected_async_load_mutex); + } + metal_graph_selected_async_load job = + g_metal_graph_selected_async_load_job; + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + + metal_graph_selected_async_load_run(&job); + + pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); + g_metal_graph_selected_async_load_job = job; + g_metal_graph_selected_async_load_has_job = false; + g_metal_graph_selected_async_load_done = true; + pthread_cond_signal(&g_metal_graph_selected_async_load_done_cond); + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + } + return NULL; +} + +static bool metal_graph_selected_async_load_ensure_worker(void) { + pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); + if (g_metal_graph_selected_async_load_thread_started) { + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + return true; + } + const int rc = pthread_create(&g_metal_graph_selected_async_load_thread, + NULL, + metal_graph_selected_async_load_worker_main, + NULL); + if (rc != 0) { + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + fprintf(stderr, + "ds4: failed to start Metal streaming async selected load worker: %s\n", + strerror(rc)); + return false; + } + g_metal_graph_selected_async_load_thread_started = true; + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + return true; +} + +static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start_tensor( + metal_graph_selected_async_load *job, + ds4_gpu_tensor *router_selected, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t event_value, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!job || !router_selected || event_value == 0) return false; + if (!metal_graph_selected_async_load_ensure_worker()) return false; + memset(job, 0, sizeof(*job)); + job->router_selected = router_selected; + job->model = model; + job->layer = layer; + job->il = il; + job->event_value = event_value; + job->gate_expert_bytes = gate_expert_bytes; + job->down_expert_bytes = down_expert_bytes; + + pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); + if (g_metal_graph_selected_async_load_has_job || + g_metal_graph_selected_async_load_done) { + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + return false; + } + g_metal_graph_selected_async_load_job = *job; + g_metal_graph_selected_async_load_job.ok = false; + g_metal_graph_selected_async_load_has_job = true; + pthread_cond_signal(&g_metal_graph_selected_async_load_cond); + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + job->active = true; + return true; +} + +static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start( + metal_graph_selected_async_load *job, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint64_t event_value, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return metal_graph_selected_async_load_start_tensor( + job, + g ? metal_graph_router_selected(g) : NULL, + model, + layer, + il, + event_value, + gate_expert_bytes, + down_expert_bytes); +} + +static bool metal_graph_selected_async_load_finish( + metal_graph_selected_async_load *job) { + if (!job || !job->active) return false; + pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); + while (!g_metal_graph_selected_async_load_done) { + pthread_cond_wait(&g_metal_graph_selected_async_load_done_cond, + &g_metal_graph_selected_async_load_mutex); + } + *job = g_metal_graph_selected_async_load_job; + g_metal_graph_selected_async_load_done = false; + pthread_mutex_unlock(&g_metal_graph_selected_async_load_mutex); + job->active = false; + if (!job->ok) return false; + return ds4_gpu_routed_moe_set_selected_override(job->selected_ids, + DS4_N_EXPERT_USED) != 0; +} + +#ifdef DS4_ROCM_BUILD +typedef struct rocm_graph_batch_selected_async_load { + bool active; + bool ok; + const ds4_gpu_tensor *selected; + const ds4_model *model; + const ds4_layer_weights *layer; + uint32_t il; + uint32_t n_tokens; + uint64_t event_value; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t *selected_ids; +} rocm_graph_batch_selected_async_load; + +static pthread_mutex_t g_rocm_graph_batch_selected_async_load_mutex = + PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_rocm_graph_batch_selected_async_load_cond = + PTHREAD_COND_INITIALIZER; +static pthread_cond_t g_rocm_graph_batch_selected_async_load_done_cond = + PTHREAD_COND_INITIALIZER; +static pthread_t g_rocm_graph_batch_selected_async_load_thread; +static bool g_rocm_graph_batch_selected_async_load_thread_started = false; +static bool g_rocm_graph_batch_selected_async_load_has_job = false; +static bool g_rocm_graph_batch_selected_async_load_done = false; +static rocm_graph_batch_selected_async_load + g_rocm_graph_batch_selected_async_load_job; + +static void rocm_graph_batch_selected_async_load_run( + rocm_graph_batch_selected_async_load *job) { + job->ok = false; + if (!job->selected || !job->model || !job->layer || !job->selected_ids || + job->n_tokens <= 1 || + DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || + DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return; + } + if (DS4_N_EXPERT_USED != 0 && + job->n_tokens > UINT64_MAX / DS4_N_EXPERT_USED) { + return; + } + const uint64_t n_ids = (uint64_t)job->n_tokens * DS4_N_EXPERT_USED; + if (n_ids > SIZE_MAX / sizeof(job->selected_ids[0])) return; + if (ds4_gpu_tensor_read_after_selected_event( + job->selected, + 0, + job->selected_ids, + n_ids * sizeof(job->selected_ids[0]), + job->event_value, + "prefill selected-id async expert load") == 0) { + return; + } + for (uint64_t i = 0; i < n_ids; i++) { + if (job->selected_ids[i] < 0 || + (uint32_t)job->selected_ids[i] >= DS4_N_EXPERT) { + fprintf(stderr, + "ds4: ROCm streaming async batch selected expert id %d " + "is outside 0..%u at layer %u\n", + job->selected_ids[i], + DS4_N_EXPERT, + job->il); + return; + } + } + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(job->model, + job->layer, + job->il, + job->gate_expert_bytes, + job->down_expert_bytes); + if (ds4_gpu_stream_expert_cache_prepare_selected_batch( + &table, + job->selected_ids, + job->n_tokens, + DS4_N_EXPERT_USED) == 0) { + return; + } + job->ok = true; +} + +static void *rocm_graph_batch_selected_async_load_worker_main(void *arg) { + (void)arg; + for (;;) { + pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); + while (!g_rocm_graph_batch_selected_async_load_has_job) { + pthread_cond_wait(&g_rocm_graph_batch_selected_async_load_cond, + &g_rocm_graph_batch_selected_async_load_mutex); + } + rocm_graph_batch_selected_async_load job = + g_rocm_graph_batch_selected_async_load_job; + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + + rocm_graph_batch_selected_async_load_run(&job); + + pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); + g_rocm_graph_batch_selected_async_load_job = job; + g_rocm_graph_batch_selected_async_load_has_job = false; + g_rocm_graph_batch_selected_async_load_done = true; + pthread_cond_signal(&g_rocm_graph_batch_selected_async_load_done_cond); + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + } + return NULL; +} + +static bool rocm_graph_batch_selected_async_load_ensure_worker(void) { + pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); + if (g_rocm_graph_batch_selected_async_load_thread_started) { + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + return true; + } + const int rc = pthread_create(&g_rocm_graph_batch_selected_async_load_thread, + NULL, + rocm_graph_batch_selected_async_load_worker_main, + NULL); + if (rc != 0) { + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + fprintf(stderr, + "ds4: failed to start ROCm streaming async batch selected " + "load worker: %s\n", + strerror(rc)); + return false; + } + g_rocm_graph_batch_selected_async_load_thread_started = true; + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + return true; +} + +static bool rocm_graph_batch_selected_async_load_start( + rocm_graph_batch_selected_async_load *job, + const ds4_gpu_tensor *selected, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + uint64_t event_value, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!job || !selected || event_value == 0 || n_tokens <= 1) return false; + if (!rocm_graph_batch_selected_async_load_ensure_worker()) return false; + if (DS4_N_EXPERT_USED != 0 && + n_tokens > UINT64_MAX / DS4_N_EXPERT_USED) { + return false; + } + const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; + memset(job, 0, sizeof(*job)); + job->selected_ids = xmalloc((size_t)n_ids * sizeof(job->selected_ids[0])); + job->selected = selected; + job->model = model; + job->layer = layer; + job->il = il; + job->n_tokens = n_tokens; + job->event_value = event_value; + job->gate_expert_bytes = gate_expert_bytes; + job->down_expert_bytes = down_expert_bytes; + + pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); + if (g_rocm_graph_batch_selected_async_load_has_job || + g_rocm_graph_batch_selected_async_load_done) { + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + free(job->selected_ids); + memset(job, 0, sizeof(*job)); + return false; + } + g_rocm_graph_batch_selected_async_load_job = *job; + g_rocm_graph_batch_selected_async_load_job.ok = false; + g_rocm_graph_batch_selected_async_load_has_job = true; + pthread_cond_signal(&g_rocm_graph_batch_selected_async_load_cond); + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + job->active = true; + return true; +} + +static bool rocm_graph_batch_selected_async_load_finish( + rocm_graph_batch_selected_async_load *job) { + if (!job || !job->active) return false; + pthread_mutex_lock(&g_rocm_graph_batch_selected_async_load_mutex); + while (!g_rocm_graph_batch_selected_async_load_done) { + pthread_cond_wait(&g_rocm_graph_batch_selected_async_load_done_cond, + &g_rocm_graph_batch_selected_async_load_mutex); + } + *job = g_rocm_graph_batch_selected_async_load_job; + g_rocm_graph_batch_selected_async_load_done = false; + pthread_mutex_unlock(&g_rocm_graph_batch_selected_async_load_mutex); + const bool ok = job->ok; + free(job->selected_ids); + memset(job, 0, sizeof(*job)); + return ok; +} +#endif + +static bool metal_graph_profile_router_selection( + ds4_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos) { + if (!g_expert_profile.active) return true; + if (!g || !layer || !metal_graph_router_selected(g) || !metal_graph_router_weights(g)) return false; + + if (ds4_gpu_end_commands() == 0) { + fprintf(stderr, + "ds4: failed to end Metal command batch for expert profile readback\n"); + return false; + } + + int32_t selected[DS4_MAX_EXPERT_USED] = {0}; + float weights[DS4_MAX_EXPERT_USED] = {0}; + const bool read_ok = + ds4_gpu_tensor_read(metal_graph_router_selected(g), + 0, + selected, + (uint64_t)DS4_N_EXPERT_USED * sizeof(selected[0])) != 0 && + ds4_gpu_tensor_read(metal_graph_router_weights(g), + 0, + weights, + (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) != 0; + + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: failed to resume Metal command batch after expert profile readback\n"); + return false; + } + if (!read_ok) { + fprintf(stderr, "ds4: failed to read Metal router tensors for expert profile\n"); + return false; + } + + ds4_expert_profile_record(il, + pos, + selected, + weights, + layer->ffn_gate_tid2eid != NULL); + return true; +} + +/* Diagnostic skip-ablation for TP profiling (DS4_TP_ABLATE=chain[,chain]): + * drops whole encode chains so their true in-situ cost shows up as a t/s + * delta. Output is semantically wrong while enabled; both ranks must set + * the same value. Chains: hcpre, router, kv, compidx. */ +static bool metal_graph_tp_ablate(const char *chain) { + static const char *env = NULL; + static int init = 0; + if (!init) { + env = getenv("DS4_TP_ABLATE"); + init = 1; + } + return env && strstr(env, chain) != NULL; +} + +static bool metal_graph_borrow_tensor_view( + ds4_gpu_tensor *view, + const ds4_gpu_tensor *base, + uint64_t offset, + uint64_t bytes) { + if (!view || !base || offset > base->bytes || bytes > base->bytes - offset) { + return false; + } + memset(view, 0, sizeof(*view)); + view->ptr = (char *)base->ptr + offset; + view->bytes = bytes; + view->owner = 0; + view->device_id = base->device_id; + return true; +} + +static bool metal_graph_cuda_tp_ep_finish_reduce( + ds4_gpu_graph *g, + int home_tier, + int partner_tier, + bool direct_return, + uint64_t return_bytes, + bool combine) { + bool ok; + if (direct_return) { + ok = ds4_gpu_tensor_wait_xdev_default( + g->routed_down_by_tier[partner_tier], home_tier) != 0; + } else { + ok = ds4_gpu_tensor_copy_xdev_default( + g->tp_peer_tmp_by_tier[home_tier], + g->routed_down_by_tier[partner_tier], + return_bytes) != 0; + } + if (!combine) return ok; + if (ok && g->cuda_tp_ep_pack_exact) { + ok = ds4_gpu_routed_moe_owned_packed_combine_tensor( + metal_graph_routed_out(g), + metal_graph_routed_down(g), + g->tp_peer_tmp_by_tier[home_tier], + metal_graph_router_selected(g), + DS4_N_EMBD, + DS4_N_EXPERT / 2u) != 0; + } else if (ok) { + ok = ds4_gpu_routed_moe_owned_slots_combine_tensor( + metal_graph_routed_out(g), + metal_graph_routed_down(g), + g->tp_peer_tmp_by_tier[home_tier], + metal_graph_router_selected(g), + DS4_N_EMBD, + DS4_N_EXPERT / 2u) != 0; + } + return ok; +} + +typedef enum { + METAL_DECODE_LAYER_FULL = 0, + METAL_DECODE_LAYER_TO_FFN, + METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN, + METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN, + METAL_DECODE_LAYER_TO_QKV, + METAL_DECODE_LAYER_FROM_QKV_TO_ATTN, + METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID, + METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN, + METAL_DECODE_LAYER_FROM_ATTN_TO_FFN, + METAL_DECODE_LAYER_TO_ROUTER, + METAL_DECODE_LAYER_TO_SHARED_MID, + METAL_DECODE_LAYER_FROM_ROUTER, +} metal_decode_layer_phase; + +static bool metal_graph_encode_decode_layer_phase( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t raw_row, + uint32_t n_raw, + int token, + metal_decode_layer_phase phase) { + /* switch to this layer's home tier before any Class P + * accessor reads. Single-tier (placement == NULL): no-op. */ + if (g->placement) { + const int this_tier = g->placement[il + 1]; + if (!metal_graph_set_active_tier_decode(g, this_tier)) return false; + } + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint32_t n_groups = DS4_N_OUT_GROUP; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = DS4_N_LORA_O; + const uint32_t shared_dim = (uint32_t)layer->ffn_gate_shexp->dim[1]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + const uint64_t down_expert_bytes = routed_out_dim * down_row_bytes; + const bool compressed = ds4_layer_compress_ratio(il) != 0; + const float freq_base = layer_rope_freq_base(il); + const float freq_scale = layer_rope_freq_scale(il); + const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; + float attn_factor = 1.0f; + if (ext_factor != 0.0f && freq_scale > 0.0f) { + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); + const int cuda_tp_home_tier = g->active_tier; + const int cuda_tp_partner_tier = g->cuda_tp_decode + ? metal_graph_cuda_tp_partner_tier(cuda_tp_home_tier) : -1; + const bool tp_split_attn = g->tp_world == 2; + const uint32_t tp_heads = tp_split_attn ? + (uint32_t)DS4_N_HEAD / 2u : (uint32_t)DS4_N_HEAD; + const uint32_t tp_head0 = tp_split_attn ? g->tp_rank * tp_heads : 0; + + bool ok = true; + const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); + double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; +#define DS4_METAL_PROFILE_DECODE_STAGE(name) do { \ + if (ok && decode_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary("decode", (name), il, pos, 1, &decode_stage_t0); \ + } \ + } while (0) + const bool tp_ablate_hcpre = metal_graph_tp_ablate("hcpre"); + if (phase != METAL_DECODE_LAYER_FROM_ROUTER) { + const bool fuse_hc_norm = + DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + !metal_graph_use_reference_hc_norm_decode(); + const bool stop_before_attn = + phase == METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN || + phase == METAL_DECODE_LAYER_FROM_QKV_TO_ATTN || + phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; + const bool resume_after_qkv = + phase == METAL_DECODE_LAYER_FROM_QKV_TO_ATTN || + phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; + const bool resume_after_qa_kv_raw = + phase == METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID; + const bool resume_after_kv_store = + phase == METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN; + const bool resume_after_attn = + phase == METAL_DECODE_LAYER_FROM_ATTN_TO_FFN; + bool cuda_tp_attn_heads_active = false; + bool attn_inv_rope_done = resume_after_attn; + if (phase != METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN && + phase != METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN && + phase != METAL_DECODE_LAYER_FROM_QKV_TO_ATTN && + phase != METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID && + phase != METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN && + !resume_after_attn) { + if (ok && !tp_ablate_hcpre) { + ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_hc_mix(g), model, layer->hc_attn_fn, + hc_dim, mix_hc, metal_graph_flat_hc(g), 1); + } + if (ok && fuse_hc_norm) { + ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(metal_graph_attn_cur(g), + metal_graph_attn_norm(g), + metal_graph_hc_split(g), + metal_graph_hc_mix(g), + metal_graph_cur_hc(g), + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + if (ok) { + ok = metal_graph_check_hc_norm_fusion("attn", + metal_graph_attn_cur(g), + metal_graph_attn_norm(g), + metal_graph_hc_mix(g), + metal_graph_cur_hc(g), + model, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + layer->attn_norm->abs_offset, + il, + pos); + } + } else if (ok) { + ok = metal_graph_decode_hc_pre(metal_graph_attn_cur(g), + metal_graph_hc_split(g), + metal_graph_hc_mix(g), + metal_graph_cur_hc(g), + model, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset); + } + DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_pre"); + if (ok) { + metal_graph_debug_dump_tensor("hc_attn_pre_mixes", metal_graph_hc_mix(g), mix_hc, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_weights", metal_graph_hc_pre(g), DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_post_weights", metal_graph_hc_post(g), DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_comb", metal_graph_hc_comb(g), (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("hc_attn_pre", metal_graph_attn_cur(g), DS4_N_EMBD, il, pos); + } + if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_attn_norm(g), metal_graph_attn_cur(g), + model->map, model->size, + layer->attn_norm->abs_offset, + DS4_N_EMBD, DS4_RMS_EPS) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("attn_norm"); + if (ok) { + metal_graph_debug_dump_tensor("attn_norm", metal_graph_attn_norm(g), DS4_N_EMBD, il, pos); + } + if (phase == METAL_DECODE_LAYER_TO_QKV) return ok; + } + if (!resume_after_attn) { + if (!resume_after_qkv) { + bool qkv_pair_projected = resume_after_qa_kv_raw; + if (!resume_after_qa_kv_raw && ok && qkv_rms_fused && + g->cuda_qkv_pair && !metal_graph_use_reference_qkv_pair_proj()) { + qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor( + metal_graph_qr(g), + metal_graph_kv_raw(g), + model->map, + model->size, + layer->attn_q_a->abs_offset, + layer->attn_kv->abs_offset, + DS4_N_EMBD, + q_rank, + DS4_N_HEAD_DIM, + metal_graph_attn_norm(g), + 1) != 0; + } + if (!resume_after_qa_kv_raw && ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(metal_graph_qr(g), model->map, model->size, + layer->attn_q_a->abs_offset, + DS4_N_EMBD, q_rank, + metal_graph_attn_norm(g), 1) != 0; + if (ok) { + metal_graph_debug_dump_tensor("q_lora", metal_graph_qr(g), q_rank, il, pos); + } + const bool kvnorm_dump = metal_graph_debug_wants("KVnorm", il, pos); + bool kv_rope_fused = false; + if (qkv_rms_fused) { + if (!resume_after_qa_kv_raw && ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(metal_graph_kv_raw(g), model->map, model->size, + layer->attn_kv->abs_offset, + DS4_N_EMBD, DS4_N_HEAD_DIM, + metal_graph_attn_norm(g), 1) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVraw", metal_graph_kv_raw(g), DS4_N_HEAD_DIM, il, pos); + } + if (ok && g->cuda_qkv_kv_rope_fuse && !kvnorm_dump && DS4_N_HEAD_KV == 1u) { + ok = ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor( + metal_graph_qr_norm(g), + metal_graph_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + metal_graph_kv(g), + metal_graph_kv_raw(g), + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + 1, + DS4_N_HEAD_KV, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + kv_rope_fused = ok; + } else if (ok) { + ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(metal_graph_qr_norm(g), + metal_graph_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + metal_graph_kv(g), + metal_graph_kv_raw(g), + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + 1, + DS4_RMS_EPS) != 0; + } + } else { + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_qr_norm(g), metal_graph_qr(g), + model->map, model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, DS4_RMS_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("q_lora_norm", metal_graph_qr_norm(g), q_rank, il, pos); + } + if (qkv_rms_fused && ok && !kv_rope_fused) { + metal_graph_debug_dump_tensor("KVnorm", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + /* Phase B head slice: under the real TP split this rank computes only + * its heads [tp_head0, tp_head0 + tp_heads) end to end — q_b rows, the + * per-head norm/rope, the attention core and its owned output groups. + * q and heads hold the owned half compactly at the buffer base; the + * head range lines up with the output-group split (32 heads = 4 of the + * 8 groups). */ + uint64_t tp_q_row_bytes = 0; + if (ok) ok = metal_graph_dense_quant_row_bytes(layer->attn_q_b, + q_rank, + &tp_q_row_bytes); + const uint64_t tp_q_rows_off = + (uint64_t)tp_head0 * DS4_N_HEAD_DIM * tp_q_row_bytes; + if (ok) ok = metal_graph_matmul_dense_quant_abs(metal_graph_q(g), + model, + layer->attn_q_b, + layer->attn_q_b->abs_offset + tp_q_rows_off, + q_rank, + (uint64_t)tp_heads * DS4_N_HEAD_DIM, + metal_graph_qr_norm(g), + 1); + if (ok) { + metal_graph_debug_dump_tensor("Qraw", metal_graph_q(g), q_dim, il, pos); + } + const bool decode_q_norm_debug = metal_graph_debug_wants("Qnorm", il, pos); + bool decode_q_norm_rope_fused = false; + if (ok && !decode_q_norm_debug) { + decode_q_norm_rope_fused = + ds4_gpu_head_rms_norm_rope_tail_tensor(metal_graph_q(g), + 1, + tp_heads, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } + if (!decode_q_norm_rope_fused) { + if (ok) ok = ds4_gpu_head_rms_norm_tensor(metal_graph_q(g), 1, tp_heads, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) { + metal_graph_debug_dump_tensor("Qnorm", metal_graph_q(g), q_dim, il, pos); + } + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_q(g), 1, tp_heads, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, freq_base, freq_scale, ext_factor, attn_factor, + DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("q_path"); + if (ok) { + metal_graph_debug_dump_tensor("Qcur", metal_graph_q(g), q_dim, il, pos); + } + if (!qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_attn_norm(g), + 1); + if (ok) { + metal_graph_debug_dump_tensor("KVraw", metal_graph_kv_raw(g), DS4_N_HEAD_DIM, il, pos); + } + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_kv(g), metal_graph_kv_raw(g), + model->map, model->size, + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVnorm", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + const bool tp_ablate_kv = metal_graph_tp_ablate("kv"); + if (ok && !tp_ablate_kv && !kv_rope_fused) { + ok = ds4_gpu_rope_tail_tensor(metal_graph_kv(g), 1, + DS4_N_HEAD_KV, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, freq_base, freq_scale, + ext_factor, attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("KVrope", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + if (!resume_after_kv_store) { + /* The common no-debug path may fuse KV RMS with RoPE above. KV + * storage starts here after metal_graph_kv(g) contains the RoPE row. */ + if (ok) ok = metal_graph_decode_kv_store(metal_graph_kv(g), raw_cache, raw_cap, raw_row); + if (ok) ok = metal_graph_cuda_tp_attn_cache_sync_raw_row(g, il, raw_row); + DS4_METAL_PROFILE_DECODE_STAGE("kv_path"); + if (ok) { + metal_graph_debug_dump_tensor("KVcur", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + + uint32_t n_comp = 0; + ds4_gpu_tensor *comp_cache = NULL; + ds4_gpu_tensor *comp_selected = NULL; + uint32_t n_selected = 0; + double decode_index_stage_t0 = 0.0; + const bool decode_index_stage_profile = g->decode_index_stage_profile; + if (ok && compressed) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_width = coff * DS4_N_HEAD_DIM; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (!layer->attn_compressor_kv || !layer->attn_compressor_gate || + !layer->attn_compressor_ape || !layer->attn_compressor_norm || + layer->attn_compressor_kv->type != DS4_TENSOR_F16 || + layer->attn_compressor_gate->type != DS4_TENSOR_F16 || + layer->attn_compressor_kv->dim[0] != DS4_N_EMBD || + layer->attn_compressor_gate->dim[0] != DS4_N_EMBD || + layer->attn_compressor_kv->dim[1] != comp_width || + layer->attn_compressor_gate->dim[1] != comp_width) { + fprintf(stderr, "ds4: Metal graph compressor expects paired F16 compressor projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + bool comp_state_already_stored = false; + if (ok && !metal_graph_use_reference_compressor_pair_proj()) { + const int fused_store = + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + layer->attn_compressor_gate->abs_offset, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + DS4_N_EMBD, + comp_width, + metal_graph_attn_norm(g), + ratio, + pos); + if (fused_store < 0) { + ok = false; + } else if (fused_store > 0) { + comp_state_already_stored = true; + } else { + ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, + comp_width, + metal_graph_attn_norm(g), + 1) != 0; + } + } else { + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size, + layer->attn_compressor_kv->abs_offset, + DS4_N_EMBD, comp_width, + metal_graph_attn_norm(g), 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, comp_width, + metal_graph_attn_norm(g), 1) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("compressor_proj"); + const uint32_t comp_row = g->layer_n_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_tensor(metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_attn_comp_update_target(g, il), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos, + metal_graph_attn_comp_update_row(comp_row), + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + comp_state_already_stored) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("compressor_update"); + if (ok && emit) { + ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); + if (!comp_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVcompress", comp_row_view, DS4_N_HEAD_DIM, il, pos); + } + } + ds4_gpu_tensor_free(comp_row_view); + DS4_METAL_PROFILE_DECODE_STAGE("compressor_quantize"); + if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); + DS4_METAL_PROFILE_DECODE_STAGE("compressor_commit"); + } + if (ok && emit) g->layer_n_comp[il]++; + + if (ok && ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || + !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || + layer->indexer_compressor_kv->type != DS4_TENSOR_F16 || + layer->indexer_compressor_gate->type != DS4_TENSOR_F16 || + layer->indexer_compressor_kv->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_gate->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_kv->dim[1] != index_width || + layer->indexer_compressor_gate->dim[1] != index_width) { + fprintf(stderr, "ds4: Metal graph indexer compressor expects paired F16 projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + bool index_state_already_stored = false; + if (ok && !metal_graph_use_reference_compressor_pair_proj()) { + const int fused_store = + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + layer->indexer_compressor_gate->abs_offset, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + DS4_N_EMBD, + index_width, + metal_graph_attn_norm(g), + ratio, + pos); + if (fused_store < 0) { + ok = false; + } else if (fused_store > 0) { + index_state_already_stored = true; + } else { + ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, + index_width, + metal_graph_attn_norm(g), + 1) != 0; + } + } else { + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size, + layer->indexer_compressor_kv->abs_offset, + DS4_N_EMBD, index_width, + metal_graph_attn_norm(g), 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, index_width, + metal_graph_attn_norm(g), 1) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_proj"); + const uint32_t index_row = g->layer_n_index_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_tensor(metal_graph_comp_kv_cur(g), + metal_graph_comp_sc_cur(g), + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + g->layer_index_comp_cache[il], + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos, + index_row, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + index_state_already_stored) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_update"); + if (ok && emit) { +#if defined(__APPLE__) + ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + if (!index_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + ds4_gpu_tensor_free(index_row_view); +#else + ds4_gpu_tensor index_row_view; + if (!metal_graph_borrow_tensor_view( + &index_row_view, + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float))) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(&index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } +#endif + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_qat"); + } + if (ok && emit) g->layer_n_index_comp[il]++; + const uint32_t decode_sparse_threshold = + metal_graph_decode_indexer_sparse_threshold(g); + if (ok && + g->layer_n_comp[il] > decode_sparse_threshold && + g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K) { + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_attn_q_b || + !tensor_type_is_f16_or_q8_0(layer->indexer_attn_q_b->type) || + layer->indexer_attn_q_b->dim[0] != q_rank || + layer->indexer_attn_q_b->dim[1] != indexer_q_dim) { + fprintf(stderr, "ds4: Metal graph indexer q projection expects F16 or Q8_0 weights\n"); + ok = false; + } + if (ok && (!layer->indexer_proj || + layer->indexer_proj->type != DS4_TENSOR_F16 || + layer->indexer_proj->dim[0] != DS4_N_EMBD || + layer->indexer_proj->dim[1] != DS4_N_INDEXER_HEAD)) { + fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); + ok = false; + } + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_indexer_q(g), + model, + layer->indexer_attn_q_b, + q_rank, + indexer_q_dim, + metal_graph_qr_norm(g), + 1); + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_indexer_q(g), 1, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_indexer_q(g), + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_indexer_weights(g), model->map, model->size, + layer->indexer_proj->abs_offset, + DS4_N_EMBD, DS4_N_INDEXER_HEAD, + metal_graph_attn_norm(g), 1) != 0; + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), + metal_graph_indexer_q(g), + metal_graph_indexer_weights(g), + g->layer_index_comp_cache[il], + g->layer_n_index_comp[il], + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + index_scale) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_score", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + g->layer_n_index_comp[il], + 1, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_topk", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + /* Decode used to materialize a dense compressed-row mask and + * call the generic gathered FlashAttention wrapper below. + * That wrapper scans every compressed row and rejects long + * contexts once raw+compressed rows exceed 8192. Ratio-4 DS4 + * attention is sparse after indexer top-k, so use the private + * indexed attention kernel instead: it scans only SWA raw rows + * plus the selected compressed rows, matching prefill and + * avoiding the long-context decode failure. */ + if (ok) { + comp_selected = metal_graph_comp_selected(g); + /* + * Contract: the indexer top-k is fixed by the model config + * and must remain the full 512 rows. Do not reduce this for + * throughput benchmarks. + * + * Why: the indexer is not just an implementation detail. It + * decides which compressed memory rows are visible to the + * attention kernel. If we keep only 128/256 rows, the later + * indexed-attention math may be perfectly computed, but it is + * computed over the wrong candidate set: rows ranked 257-512 + * are removed before softmax/PV can use them. Those rows may + * carry weak-but-necessary evidence for retrieval, name/number + * recall, or long-context disambiguation. The error is + * therefore semantic/algorithmic, not the acceptable kind of + * local numerical drift caused by a different reduction order + * or Tensor/NAX precision. + * + * Short prompt tests, first-token agreement, or even a small + * official-vector set can miss this because many prompts do + * not need the tail of the 512 selected compressed rows. The + * failure appears only when the model needs information that + * fell below the reduced cutoff. Optimizations belong inside + * the score/top-k/attention implementation while preserving + * DS4_N_INDEXER_TOP_K. + */ + n_selected = DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] + ? DS4_N_INDEXER_TOP_K + : g->layer_n_index_comp[il]; + } + } + } + + n_comp = g->layer_n_comp[il]; + comp_cache = g->layer_attn_comp_cache[il]; + } + DS4_METAL_PROFILE_DECODE_STAGE("compressor_indexer"); + + if (stop_before_attn) return ok; + if (ok) { + const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); + const bool indexed_attention = n_comp != 0 && comp_selected != NULL && n_selected != 0; + const bool cuda_tp_attn_heads_requested = g->cuda_tp_attn_heads; + const uint32_t cuda_tp_heads = DS4_N_HEAD / 2u; + const bool cuda_tp_attn_local_cache = + metal_graph_cuda_tp_attn_cache_dup_layer_ready(g, il); + cuda_tp_attn_heads_active = + cuda_tp_attn_heads_requested && + cuda_tp_partner_tier >= 0 && + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] && + (DS4_N_HEAD % 2u) == 0u && + (n_groups % 2u) == 0u && + g->q_by_tier[cuda_tp_partner_tier] && + g->heads_by_tier[cuda_tp_partner_tier] && + !metal_graph_debug_wants("kqv_out", il, pos) && + !metal_graph_debug_wants("kqv_back", il, pos); + if (cuda_tp_attn_heads_requested && !cuda_tp_attn_heads_active) { + fprintf(stderr, + "ds4: CUDA decode TP cannot split attention heads for tier %d " + "(partner=%d heads=%u peer=%d)\n", + cuda_tp_home_tier, + cuda_tp_partner_tier, + DS4_N_HEAD, + cuda_tp_partner_tier >= 0 ? + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] : 0); + ok = false; + } + if (ok && cuda_tp_attn_heads_active) { + const uint64_t tp_head_bytes = (uint64_t)cuda_tp_heads * DS4_N_HEAD_DIM * sizeof(float); + ds4_gpu_tensor q_peer_src; + ds4_gpu_tensor q_peer_dst; + ds4_gpu_tensor peer_heads_tail; + ok = metal_graph_borrow_tensor_view(&q_peer_src, + metal_graph_q(g), + tp_head_bytes, + tp_head_bytes) && + metal_graph_borrow_tensor_view(&q_peer_dst, + g->q_by_tier[cuda_tp_partner_tier], + 0, + tp_head_bytes) && + metal_graph_borrow_tensor_view(&peer_heads_tail, + g->heads_by_tier[cuda_tp_partner_tier], + tp_head_bytes, + tp_head_bytes); + if (ok) { + ok = ds4_gpu_tensor_copy_xdev(&q_peer_dst, &q_peer_src, + tp_head_bytes) != 0; + } + ds4_gpu_tensor *peer_raw_cache = cuda_tp_attn_local_cache + ? g->layer_raw_cache_tp[il] : raw_cache; + ds4_gpu_tensor *peer_comp_cache = cuda_tp_attn_local_cache + ? g->layer_attn_comp_cache_tp[il] : comp_cache; + ds4_gpu_tensor *peer_selected = comp_selected; + if (ok && indexed_attention && cuda_tp_attn_local_cache) { + peer_selected = g->comp_selected_by_tier[cuda_tp_partner_tier]; + ok = peer_selected && + ds4_gpu_tensor_copy_xdev(peer_selected, + comp_selected, + (uint64_t)n_selected * sizeof(int32_t)) != 0; + } + if (ok && !cuda_tp_attn_local_cache) { + ok = ds4_gpu_tensor_wait_xdev(raw_cache, cuda_tp_partner_tier) != 0; + } + if (ok) ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; + if (ok && indexed_attention) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + &peer_heads_tail, + model->map, + model->size, + layer->attn_sinks->abs_offset + (uint64_t)cuda_tp_heads * sizeof(float), + &q_peer_dst, + peer_raw_cache, + peer_comp_cache, + metal_graph_attn_comp_cache_is_f16(), + peer_selected, + 1, + pos, + n_raw, + raw_cap, + raw_start, + n_comp, + n_selected, + g->raw_window, + ds4_layer_compress_ratio(il), + cuda_tp_heads, + DS4_N_HEAD_DIM) != 0; + } else if (ok) { + ok = ds4_gpu_attention_decode_heads_tensor( + &peer_heads_tail, + model->map, + model->size, + layer->attn_sinks->abs_offset + (uint64_t)cuda_tp_heads * sizeof(float), + &q_peer_dst, + peer_raw_cache, + n_raw, + raw_cap, + raw_start, + n_comp ? peer_comp_cache : NULL, + metal_graph_attn_comp_cache_is_f16(), + n_comp, + NULL, + 0, + cuda_tp_heads, + DS4_N_HEAD_DIM) != 0; + } + if (ok) { + ok = ds4_gpu_rope_tail_tensor(&peer_heads_tail, + 1, cuda_tp_heads, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + if (ok) ok = ds4_gpu_set_current_device(cuda_tp_home_tier) == 0; + if (ok && indexed_attention) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + metal_graph_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_q(g), + raw_cache, + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + comp_selected, + 1, + pos, + n_raw, + raw_cap, + raw_start, + n_comp, + n_selected, + g->raw_window, + ds4_layer_compress_ratio(il), + cuda_tp_heads, + DS4_N_HEAD_DIM) != 0; + } else if (ok) { + ok = ds4_gpu_attention_decode_heads_tensor(metal_graph_heads(g), + model->map, model->size, + layer->attn_sinks->abs_offset, + metal_graph_q(g), raw_cache, n_raw, + raw_cap, + raw_start, + n_comp ? comp_cache : NULL, + metal_graph_attn_comp_cache_is_f16(), + n_comp, + NULL, + 0, + cuda_tp_heads, DS4_N_HEAD_DIM) != 0; + } + if (ok) { + ok = ds4_gpu_rope_tail_tensor(metal_graph_heads(g), + 1, cuda_tp_heads, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + if (ok && indexed_attention && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_attention", + il, + pos, + 1, + n_comp, + &decode_index_stage_t0); + } + } else if (ok && indexed_attention) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + metal_graph_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset + (uint64_t)tp_head0 * (layer->attn_sinks->bytes / DS4_N_HEAD), + metal_graph_q(g), + raw_cache, + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + comp_selected, + 1, + pos, + n_raw, + raw_cap, + raw_start, + n_comp, + n_selected, + g->raw_window, + ds4_layer_compress_ratio(il), + tp_heads, + DS4_N_HEAD_DIM) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_attention", + il, + pos, + 1, + n_comp, + &decode_index_stage_t0); + } + } else { + ok = ds4_gpu_attention_decode_heads_tensor(metal_graph_heads(g), + model->map, model->size, + layer->attn_sinks->abs_offset + (uint64_t)tp_head0 * (layer->attn_sinks->bytes / DS4_N_HEAD), + metal_graph_q(g), raw_cache, n_raw, + raw_cap, + raw_start, + n_comp ? comp_cache : NULL, + metal_graph_attn_comp_cache_is_f16(), + n_comp, + NULL, + 0, + tp_heads, DS4_N_HEAD_DIM) != 0; + } + } + } + if (ok && !cuda_tp_attn_heads_active && !attn_inv_rope_done) { + ok = ds4_gpu_rope_tail_tensor(metal_graph_heads(g), + 1, tp_heads, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("attn_inv_rope"); + if (ok && !cuda_tp_attn_heads_active) { + metal_graph_debug_dump_tensor("kqv_back", metal_graph_heads(g), q_dim, il, pos); + } + ds4_gpu_tensor *cuda_tp_attn_peer = NULL; + bool cuda_tp_attn_hc_fused = false; + const bool cuda_tp_attn_requested = g->cuda_tp_attn; + const bool cuda_tp_attn = + cuda_tp_attn_requested && + !metal_graph_directional_steering_attn_enabled(g) && + cuda_tp_partner_tier >= 0 && + (n_groups % 2u) == 0u; + ds4_gpu_tensor *tp_attn_a = NULL; /* rank partials consumed directly */ + ds4_gpu_tensor *tp_attn_b = NULL; /* by the HC expand */ + const bool fuse_attn_out_hc = + !cuda_tp_attn && + g->tp_world < 2 && + layer->attn_output_a->type == DS4_TENSOR_Q8_0 && + layer->attn_output_b->type == DS4_TENSOR_Q8_0 && + !metal_graph_directional_steering_attn_enabled(g) && + !metal_graph_use_reference_attn_out_hc(); + const bool fuse_tp_attn_out_hc = + cuda_tp_attn && + !metal_graph_use_reference_attn_out_hc() && + g->cuda_tp_attn_out_hc_fuse; + if (ok && cuda_tp_attn_requested && !cuda_tp_attn) { + fprintf(stderr, + "ds4: CUDA decode TP cannot split attention output for tier %d " + "(partner=%d groups=%u)\n", + cuda_tp_home_tier, cuda_tp_partner_tier, n_groups); + ok = false; + } + if (ok && cuda_tp_attn) { + const uint32_t tp_groups = n_groups / 2u; + const uint64_t tp_heads_bytes = (uint64_t)tp_groups * group_dim * sizeof(float); + const uint64_t tp_heads_off = (uint64_t)tp_groups * group_dim * sizeof(float); + const bool cuda_tp_attn_peer_read = + !cuda_tp_attn_heads_active && + g->cuda_tp_attn_peer_read && + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; + ds4_gpu_tensor peer_heads_dst_view; + ds4_gpu_tensor peer_heads_src_view; + ds4_gpu_tensor *peer_heads_dst = NULL; + ds4_gpu_tensor *peer_heads_src = NULL; + if (cuda_tp_attn_heads_active) { + peer_heads_src = g->heads_by_tier[cuda_tp_partner_tier]; + } else if (!cuda_tp_attn_peer_read) { + ok = metal_graph_borrow_tensor_view( + &peer_heads_dst_view, + g->heads_by_tier[cuda_tp_partner_tier], + tp_heads_off, + tp_heads_bytes); + if (ok) peer_heads_dst = &peer_heads_dst_view; + } + if (ok) { + if (cuda_tp_attn_heads_active) { + peer_heads_src = g->heads_by_tier[cuda_tp_partner_tier]; + } else if (cuda_tp_attn_peer_read) { + peer_heads_src = metal_graph_heads(g); + } else { + ok = metal_graph_borrow_tensor_view(&peer_heads_src_view, + metal_graph_heads(g), + tp_heads_off, + tp_heads_bytes); + if (ok) peer_heads_src = &peer_heads_src_view; + } + } + ds4_gpu_tensor *peer_heads = cuda_tp_attn_peer_read ? + metal_graph_heads(g) : g->heads_by_tier[cuda_tp_partner_tier]; + ok = (cuda_tp_attn_heads_active || cuda_tp_attn_peer_read || peer_heads_dst) && + peer_heads_src && peer_heads && + g->attn_low_by_tier[cuda_tp_partner_tier] && + g->attn_out_by_tier[cuda_tp_partner_tier] && + g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; + if (ok && !cuda_tp_attn_heads_active && !cuda_tp_attn_peer_read) { + ok = ds4_gpu_tensor_copy_xdev(peer_heads_dst, peer_heads_src, + tp_heads_bytes) != 0; + } else if (ok && cuda_tp_attn_peer_read) { + ok = ds4_gpu_tensor_wait_xdev(peer_heads_src, cuda_tp_partner_tier) != 0; + } + if (ok) ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; + if (ok) { + ok = ds4_gpu_attention_output_q8_tp_tensor( + g->attn_out_by_tier[cuda_tp_partner_tier], + g->attn_low_by_tier[cuda_tp_partner_tier], + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, + rank, + n_groups, + tp_groups, + tp_groups, + DS4_N_EMBD, + peer_heads) != 0; + } + if (ok) ok = ds4_gpu_set_current_device(cuda_tp_home_tier) == 0; + if (ok && fuse_tp_attn_out_hc) { + ok = ds4_gpu_attention_output_low_q8_tensor( + metal_graph_attn_low(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + group_dim, + rank, + tp_groups, + metal_graph_heads(g)) != 0; + } else if (ok) { + ok = ds4_gpu_attention_output_q8_tp_tensor( + metal_graph_attn_out(g), + metal_graph_attn_low(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, + rank, + n_groups, + 0, + tp_groups, + DS4_N_EMBD, + metal_graph_heads(g)) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy_xdev(g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + g->attn_out_by_tier[cuda_tp_partner_tier], + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + if (ok) cuda_tp_attn_peer = g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; + } + if (ok && fuse_tp_attn_out_hc) { + ok = ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( + metal_graph_after_attn_hc(g), + metal_graph_attn_out(g), + model->map, + model->size, + layer->attn_output_b->abs_offset, + (uint64_t)n_groups * rank, + DS4_N_EMBD, + 0, + (uint64_t)tp_groups * rank, + metal_graph_attn_low(g), + cuda_tp_attn_peer, + metal_graph_cur_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) cuda_tp_attn_hc_fused = true; + } + } else if (ok && fuse_attn_out_hc) { + ok = ds4_gpu_attention_output_low_q8_tensor(metal_graph_attn_low(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + group_dim, + rank, + n_groups, + metal_graph_heads(g)) != 0; + if (ok) { + ok = ds4_gpu_matmul_q8_0_hc_expand_tensor(metal_graph_after_attn_hc(g), + metal_graph_attn_out(g), + model->map, + model->size, + layer->attn_output_b->abs_offset, + (uint64_t)n_groups * rank, + DS4_N_EMBD, + metal_graph_attn_low(g), + metal_graph_cur_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + } else if (ok && g->tp_world == 2) { + /* Group-sliced attention output: this rank computes its half of the + * output groups and the matching k-window of the expand projection, + * leaving a partial block output in the gate slot. */ + const uint32_t tp_groups = n_groups / 2; + ok = metal_graph_attention_output_dense_quant_tp( + g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN], + metal_graph_attn_low(g), + g, + model, + layer->attn_output_a, + layer->attn_output_b, + group_dim, rank, + n_groups, + g->tp_rank * tp_groups, tp_groups, + DS4_N_EMBD, + metal_graph_heads(g)); + } else if (ok && layer->attn_output_a->type != DS4_TENSOR_Q8_0) { + ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? + g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); + ok = metal_graph_attention_output_dense_quant_low(metal_graph_attn_low(g), + g, + model, + layer->attn_output_a, + group_dim, + rank, + 0, + n_groups, + metal_graph_heads(g)); + if (ok) ok = metal_graph_matmul_dense_quant_tensor(attn_out_dst, + model, + layer->attn_output_b, + (uint64_t)n_groups * rank, + DS4_N_EMBD, + metal_graph_attn_low(g), + 1); + } else if (ok) { + ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? + g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); + ok = ds4_gpu_attention_output_q8_batch_tensor(attn_out_dst, + metal_graph_attn_low(g), + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, rank, + n_groups, DS4_N_EMBD, + metal_graph_heads(g), 1) != 0; + } + if (ok && g->tp_world == 2) { + /* Gate ATTN: exchange the attention block output with the peer and + * rebuild the canonical sum (rank0 first, then rank1) in attn_out + * on both ranks — identical expression on both machines keeps them + * bit-exact. */ + const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN; + ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_ATTN) != 0; + if (ok) { + ds4_gpu_tensor *first = g->tp_rank == 0 ? g->tp_out[slot] : g->tp_in[slot]; + if (metal_graph_directional_steering_attn_enabled(g)) { + ds4_gpu_tensor *second = g->tp_rank == 0 ? g->tp_in[slot] : g->tp_out[slot]; + ok = ds4_gpu_add_tensor(metal_graph_attn_out(g), first, second, DS4_N_EMBD) != 0; + } else { + /* Combine folded into the HC expand below; attn_out is not + * materialized on this path. */ + tp_attn_a = first; + tp_attn_b = g->tp_rank == 0 ? g->tp_in[slot] : g->tp_out[slot]; + } + } + } + DS4_METAL_PROFILE_DECODE_STAGE("attn_output"); + if (ok) { + metal_graph_debug_dump_tensor("attn_low", metal_graph_attn_low(g), (uint64_t)n_groups * rank, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("attn_out", metal_graph_attn_out(g), DS4_N_EMBD, il, pos); + } + if (ok && metal_graph_directional_steering_attn_enabled(g)) { + ok = metal_graph_apply_directional_steering_attn(g, metal_graph_attn_out(g), il, 1); + } + if (ok && !fuse_attn_out_hc && !cuda_tp_attn_hc_fused) { + if (tp_attn_a) { + ok = ds4_gpu_hc_expand_add_tensor(metal_graph_after_attn_hc(g), tp_attn_a, tp_attn_b, + metal_graph_cur_hc(g), metal_graph_hc_post(g), metal_graph_hc_comb(g), + DS4_N_EMBD, DS4_N_HC) != 0; + } else if (cuda_tp_attn_peer) { + ok = ds4_gpu_hc_expand_add_tensor( + metal_graph_after_attn_hc(g), + metal_graph_attn_out(g), + cuda_tp_attn_peer, + metal_graph_cur_hc(g), + metal_graph_hc_post(g), + metal_graph_hc_comb(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else { + ok = ds4_gpu_hc_expand_tensor(metal_graph_after_attn_hc(g), metal_graph_attn_out(g), metal_graph_cur_hc(g), + metal_graph_hc_post(g), metal_graph_hc_comb(g), DS4_N_EMBD, DS4_N_HC) != 0; + } + } + DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_post"); + if (ok) { + metal_graph_debug_dump_tensor("hc_attn_post", metal_graph_after_attn_hc(g), hc_dim, il, pos); + } + if (ok && !tp_ablate_hcpre) { + ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_after_attn_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_hc_mix(g), model, layer->hc_ffn_fn, + hc_dim, mix_hc, metal_graph_flat_hc(g), 1); + } + if (ok && fuse_hc_norm) { + ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(metal_graph_ffn_cur(g), + metal_graph_ffn_norm(g), + metal_graph_hc_split(g), + metal_graph_hc_mix(g), + metal_graph_after_attn_hc(g), + model->map, + model->size, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset, + layer->ffn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + if (ok) { + ok = metal_graph_check_hc_norm_fusion("ffn", + metal_graph_ffn_cur(g), + metal_graph_ffn_norm(g), + metal_graph_hc_mix(g), + metal_graph_after_attn_hc(g), + model, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset, + layer->ffn_norm->abs_offset, + il, + pos); + } + } else if (ok) { + ok = metal_graph_decode_hc_pre(metal_graph_ffn_cur(g), + metal_graph_hc_split(g), + metal_graph_hc_mix(g), + metal_graph_after_attn_hc(g), + model, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset); + } + DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_pre"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_pre_mixes", metal_graph_hc_mix(g), mix_hc, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_weights", metal_graph_hc_pre(g), DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_post_weights", metal_graph_hc_post(g), DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_comb", metal_graph_hc_comb(g), (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_pre", metal_graph_ffn_cur(g), DS4_N_EMBD, il, pos); + } + if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_ffn_norm(g), metal_graph_ffn_cur(g), + model->map, model->size, + layer->ffn_norm->abs_offset, + DS4_N_EMBD, DS4_RMS_EPS) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("ffn_norm"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_norm", metal_graph_ffn_norm(g), DS4_N_EMBD, il, pos); + } + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t gate_expert_bytes DS4_MAYBE_UNUSED = expert_mid_dim * gate_row_bytes; + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + const uint64_t down_expert_bytes DS4_MAYBE_UNUSED = routed_out_dim * down_row_bytes; + if (ok && metal_graph_decode_cpu_router_applicable(g, layer)) { + ok = metal_graph_decode_cpu_router(g, model, layer, il, (uint32_t)token); + } else { + if (ok && !metal_graph_tp_ablate("router")) { + ok = metal_graph_matmul_plain_tensor(metal_graph_router_logits(g), model, layer->ffn_gate_inp, + DS4_N_EMBD, DS4_N_EXPERT, metal_graph_ffn_norm(g), 1); + if (ok) ok = ds4_gpu_router_select_tensor(metal_graph_router_selected(g), metal_graph_router_weights(g), metal_graph_router_probs(g), + model->map, model->size, + layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, + layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, + layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, + (uint32_t)token, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE, + 0, + 0, + layer->ffn_exp_probs_b != NULL, + layer->ffn_gate_tid2eid != NULL, + metal_graph_router_logits(g)) != 0; + } + if (ok) ok = metal_graph_decode_set_hash_selected_override(model, + layer, + il, + (uint32_t)token, + layer->ffn_gate_exps->bytes, + layer->ffn_down_exps->bytes, + g); + } + DS4_METAL_PROFILE_DECODE_STAGE("router"); + if (ok) ok = metal_graph_profile_router_selection(g, layer, il, pos); + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_logits", metal_graph_router_logits(g), DS4_N_EXPERT, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_probs", metal_graph_router_probs(g), DS4_N_EXPERT, il, pos); + metal_graph_debug_dump_i32_tensor("ffn_moe_topk", metal_graph_router_selected(g), DS4_N_EXPERT_USED, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", metal_graph_router_weights(g), DS4_N_EXPERT_USED, il, pos); + } + if (phase == METAL_DECODE_LAYER_TO_ROUTER) return ok; + } + const bool external_routed = phase == METAL_DECODE_LAYER_FROM_ROUTER; + const bool fuse_shared_gate_up = + !g->quality && + g->tp_world < 2 && + layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + g->shared_gate_up_swiglu_fuse; + const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos); + const bool cuda_tp_shared_requested = g->cuda_tp_shared; + const bool cuda_tp_moe_requested = !external_routed && g->cuda_tp_moe; + const uint64_t shared_tp_local = shared_dim / 2u; + const uint64_t shared_tp_peer = shared_dim - shared_tp_local; + const uint64_t shared_q8_blocks = ((uint64_t)DS4_N_EMBD + 31u) / 32u; + const uint64_t shared_q8_x_bytes = shared_q8_blocks * 32u; + const uint64_t shared_q8_scale_offset = + (shared_q8_x_bytes + 15u) & ~15ull; + const uint64_t shared_q8_prequant_bytes DS4_MAYBE_UNUSED = + shared_q8_scale_offset + shared_q8_blocks * sizeof(float); + const bool cuda_tp_shared = + cuda_tp_shared_requested && + fuse_shared_gate_up && + !metal_graph_use_reference_shared_down_hc() && + cuda_tp_partner_tier >= 0 && + shared_tp_local != 0 && + shared_tp_peer != 0 && + (shared_tp_local % 32u) == 0 && + (shared_tp_peer % 32u) == 0 && + g->ffn_norm_by_tier[cuda_tp_partner_tier] && + g->shared_gate_by_tier[cuda_tp_partner_tier] && + g->shared_up_by_tier[cuda_tp_partner_tier] && + g->shared_mid_by_tier[cuda_tp_partner_tier] && + g->shared_out_by_tier[cuda_tp_partner_tier] && + g->tp_peer_tmp_by_tier[cuda_tp_home_tier]; + const bool cuda_tp_shared_fold = + cuda_tp_shared && + g->cuda_tp_shared_fold && + !g->cuda_tp_ep && + cuda_tp_moe_requested && + !keep_ffn_out && + !metal_graph_directional_steering_ffn_enabled(g) && + !metal_graph_debug_wants("ffn_moe_out", il, pos) && + !metal_graph_debug_wants("ffn_shexp", il, pos); + const bool fuse_shared_down_hc = + g->tp_world < 2 && + layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && + !cuda_tp_shared && + !keep_ffn_out && + !metal_graph_use_reference_shared_down_hc(); + const bool cuda_tp_moe = + cuda_tp_moe_requested && + cuda_tp_partner_tier >= 0 && + (DS4_N_EXPERT_USED % 2u) == 0u && + g->ffn_norm_by_tier[cuda_tp_partner_tier] && + g->router_selected_by_tier[cuda_tp_partner_tier] && + g->router_weights_by_tier[cuda_tp_partner_tier] && + g->routed_gate_by_tier[cuda_tp_partner_tier] && + g->routed_up_by_tier[cuda_tp_partner_tier] && + g->routed_mid_by_tier[cuda_tp_partner_tier] && + g->routed_down_by_tier[cuda_tp_partner_tier] && + g->routed_out_by_tier[cuda_tp_partner_tier] && + g->tp_peer_tmp_by_tier[cuda_tp_home_tier] && + g->tp_peer_tmp_by_tier[cuda_tp_partner_tier]; + const bool cuda_tp_ep = cuda_tp_moe && g->cuda_tp_ep; + const bool cuda_tp_moe_delay_reduce = + cuda_tp_moe && + g->cuda_tp_moe_delay_reduce && + fuse_shared_down_hc && + !metal_graph_debug_wants("ffn_moe_out", il, pos); + bool cuda_tp_moe_peer_tmp = false; + bool cuda_tp_moe_peer_copy_deferred = false; + bool cuda_tp_ep_reduce_deferred = false; + bool cuda_tp_ep_fused_hc_reduce = false; + bool cuda_tp_ep_direct_return = false; + bool cuda_tp_ep_balanced_shared_mid DS4_MAYBE_UNUSED = false; + bool cuda_tp_ep_dual_prequant = false; + uint64_t cuda_tp_ep_return_bytes = 0; + bool cuda_tp_shared_fold_peer_tmp = false; + if (ok && cuda_tp_moe_requested && !cuda_tp_moe) { + fprintf(stderr, + "ds4: CUDA decode TP cannot split routed MoE for tier %d " + "(partner=%d experts=%u)\n", + cuda_tp_home_tier, cuda_tp_partner_tier, DS4_N_EXPERT_USED); + ok = false; + } + if (ok && cuda_tp_moe) { + const uint32_t tp_experts = cuda_tp_ep + ? DS4_N_EXPERT_USED : DS4_N_EXPERT_USED / 2u; + const uint64_t tp_selected_bytes = (uint64_t)tp_experts * sizeof(int32_t); + const uint64_t tp_weights_bytes = (uint64_t)tp_experts * sizeof(float); + const uint64_t peer_selected_offset = cuda_tp_ep ? 0 : tp_selected_bytes; + const uint64_t peer_weights_offset = cuda_tp_ep ? 0 : tp_weights_bytes; + ds4_gpu_tensor local_selected; + ds4_gpu_tensor local_weights; + ds4_gpu_tensor peer_selected_src; + ds4_gpu_tensor peer_weights_src; + ds4_gpu_tensor packed_peer_ffn_norm; + ds4_gpu_tensor packed_peer_selected; + ds4_gpu_tensor packed_peer_weights; + ds4_gpu_tensor direct_peer_down; + ds4_gpu_tensor *peer_down_output = NULL; + cuda_tp_ep_return_bytes = + (uint64_t)(g->cuda_tp_ep_pack_exact ? 4u : DS4_N_EXPERT_USED) * + DS4_N_EMBD * sizeof(float); + cuda_tp_ep_direct_return = + cuda_tp_ep && + metal_graph_cuda_tp_ep_direct_return_requested() && + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier] && + metal_graph_borrow_tensor_view( + &direct_peer_down, + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + 0, + cuda_tp_ep_return_bytes); + if (cuda_tp_ep_direct_return) peer_down_output = &direct_peer_down; +#if !defined(__APPLE__) + /* Run the shared gate/up projection on the less-loaded EP rank. The + * two kernels use complementary predicates over the same top-k IDs; + * a partner result is ordered by the existing direct-return event. */ + cuda_tp_ep_balanced_shared_mid = + cuda_tp_ep_direct_return && + cuda_tp_moe_delay_reduce && + metal_graph_cuda_tp_ep_delay_reduce_requested() && + metal_graph_cuda_tp_ep_fused_shared_mid_requested() && + metal_graph_cuda_tp_ep_balanced_shared_mid_requested() && + fuse_shared_gate_up && + !cuda_tp_shared_requested && + !g->cuda_tp_moe_peer_read && + !g->cuda_tp_moe_peer_router && + !g->decode_stage_profile; + cuda_tp_ep_dual_prequant = + cuda_tp_ep_balanced_shared_mid && + metal_graph_cuda_tp_ep_dual_prequant_requested() && + metal_graph_shared_gate(g) && + metal_graph_shared_gate(g)->bytes >= shared_q8_prequant_bytes && + g->shared_gate_by_tier[cuda_tp_partner_tier] && + g->shared_gate_by_tier[cuda_tp_partner_tier]->bytes >= + shared_q8_prequant_bytes; +#endif + const bool cuda_tp_moe_peer_read = + g->cuda_tp_moe_peer_read && + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; + const bool cuda_tp_moe_peer_router = + !cuda_tp_moe_peer_read && + g->cuda_tp_moe_peer_router && + g_gpu_peer_ok[cuda_tp_partner_tier][cuda_tp_home_tier]; + const ds4_gpu_tensor *peer_selected = g->router_selected_by_tier[cuda_tp_partner_tier]; + const ds4_gpu_tensor *peer_weights = g->router_weights_by_tier[cuda_tp_partner_tier]; + const ds4_gpu_tensor *peer_ffn_norm = g->ffn_norm_by_tier[cuda_tp_partner_tier]; + ok = metal_graph_borrow_tensor_view(&local_selected, + metal_graph_router_selected(g), + 0, + tp_selected_bytes) && + metal_graph_borrow_tensor_view(&local_weights, + metal_graph_router_weights(g), + 0, + tp_weights_bytes) && + metal_graph_borrow_tensor_view(&peer_selected_src, + metal_graph_router_selected(g), + peer_selected_offset, + tp_selected_bytes) && + metal_graph_borrow_tensor_view(&peer_weights_src, + metal_graph_router_weights(g), + peer_weights_offset, + tp_weights_bytes); + if (ok && cuda_tp_moe_peer_read) { + ok = ds4_gpu_tensor_wait_xdev(metal_graph_router_weights(g), + cuda_tp_partner_tier) != 0; + peer_selected = &peer_selected_src; + peer_weights = &peer_weights_src; + peer_ffn_norm = metal_graph_ffn_norm(g); + } else if (ok && cuda_tp_moe_peer_router) { + ok = ds4_gpu_tensor_copy_xdev(g->ffn_norm_by_tier[cuda_tp_partner_tier], + metal_graph_ffn_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + peer_selected = &peer_selected_src; + peer_weights = &peer_weights_src; + peer_ffn_norm = g->ffn_norm_by_tier[cuda_tp_partner_tier]; + } else if (ok && g->cuda_tp_moe_pack_handoff) { + const uint64_t norm_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + const uint64_t packed_selected_off = norm_bytes; + const uint64_t packed_weights_off = packed_selected_off + tp_selected_bytes; + const uint64_t packed_bytes = packed_weights_off + tp_weights_bytes; + ok = metal_graph_borrow_tensor_view(&packed_peer_ffn_norm, + g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], + 0, + norm_bytes) && + metal_graph_borrow_tensor_view(&packed_peer_selected, + g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], + packed_selected_off, + tp_selected_bytes) && + metal_graph_borrow_tensor_view(&packed_peer_weights, + g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], + packed_weights_off, + tp_weights_bytes); + if (ok) { + ok = ds4_gpu_moe_handoff_pack_tensor( + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + metal_graph_ffn_norm(g), + &peer_selected_src, + &peer_weights_src, + DS4_N_EMBD, + tp_experts) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy_xdev( + g->tp_peer_tmp_by_tier[cuda_tp_partner_tier], + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + packed_bytes) != 0; + } + peer_selected = &packed_peer_selected; + peer_weights = &packed_peer_weights; + peer_ffn_norm = &packed_peer_ffn_norm; + } else if (ok && g->cuda_tp_moe_copy3_handoff) { + ok = ds4_gpu_tensor_copy_xdev3( + g->ffn_norm_by_tier[cuda_tp_partner_tier], + metal_graph_ffn_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float), + g->router_selected_by_tier[cuda_tp_partner_tier], + &peer_selected_src, + tp_selected_bytes, + g->router_weights_by_tier[cuda_tp_partner_tier], + &peer_weights_src, + tp_weights_bytes) != 0; + } else if (ok) { + ok = ds4_gpu_tensor_copy_xdev(g->ffn_norm_by_tier[cuda_tp_partner_tier], + metal_graph_ffn_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_copy_xdev(g->router_selected_by_tier[cuda_tp_partner_tier], + &peer_selected_src, + tp_selected_bytes) != 0 && + ds4_gpu_tensor_copy_xdev(g->router_weights_by_tier[cuda_tp_partner_tier], + &peer_weights_src, + tp_weights_bytes) != 0; + } + + bool switched_to_partner = false; + if (ok) { + ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; + switched_to_partner = ok; + } + if (ok) { + if (cuda_tp_ep) { + ok = ds4_gpu_routed_moe_one_owned_tensor( + g->routed_out_by_tier[cuda_tp_partner_tier], + g->routed_gate_by_tier[cuda_tp_partner_tier], + g->routed_up_by_tier[cuda_tp_partner_tier], + g->routed_mid_by_tier[cuda_tp_partner_tier], + g->routed_down_by_tier[cuda_tp_partner_tier], + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + peer_selected, + peer_weights, + DS4_N_EXPERT, + tp_experts, + DS4_N_EXPERT / 2u, + DS4_N_EXPERT - DS4_N_EXPERT / 2u, + DS4_SWIGLU_CLAMP_EXP, + peer_ffn_norm, + peer_down_output, + g->cuda_tp_ep_pack_exact, + cuda_tp_ep_dual_prequant + ? g->shared_gate_by_tier[cuda_tp_partner_tier] + : NULL) != 0; + } else { + ok = ds4_gpu_routed_moe_one_tensor( + g->routed_out_by_tier[cuda_tp_partner_tier], + g->routed_gate_by_tier[cuda_tp_partner_tier], + g->routed_up_by_tier[cuda_tp_partner_tier], + g->routed_mid_by_tier[cuda_tp_partner_tier], + g->routed_down_by_tier[cuda_tp_partner_tier], + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + peer_selected, + peer_weights, + DS4_N_EXPERT, + tp_experts, + DS4_SWIGLU_CLAMP_EXP, + peer_ffn_norm, NULL, 0, false) != 0; + } + } +#if !defined(__APPLE__) + if (ok && cuda_tp_ep_balanced_shared_mid) { + ok = ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + peer_ffn_norm, + DS4_SWIGLU_CLAMP_EXP, + peer_selected, + cuda_tp_ep_dual_prequant + ? g->shared_gate_by_tier[cuda_tp_partner_tier] + : NULL, + DS4_N_EXPERT / 2u, + false) != 0; + } +#endif + if (switched_to_partner && ds4_gpu_set_current_device(cuda_tp_home_tier) != 0) { + ok = false; + } + if (ok) { + if (cuda_tp_ep) { + ok = ds4_gpu_routed_moe_one_owned_tensor( + metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + &local_selected, + &local_weights, + DS4_N_EXPERT, + tp_experts, + 0, + DS4_N_EXPERT / 2u, + DS4_SWIGLU_CLAMP_EXP, + metal_graph_ffn_norm(g), + NULL, + false, + cuda_tp_ep_dual_prequant + ? metal_graph_shared_gate(g) + : NULL) != 0; + } else { + ok = ds4_gpu_routed_moe_one_tensor( + metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + &local_selected, + &local_weights, + DS4_N_EXPERT, + tp_experts, + DS4_SWIGLU_CLAMP_EXP, + metal_graph_ffn_norm(g), NULL, 0, false) != 0; + } + } +#if !defined(__APPLE__) + if (ok && cuda_tp_ep_balanced_shared_mid) { + ok = ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP, + &local_selected, + cuda_tp_ep_dual_prequant + ? metal_graph_shared_gate(g) + : NULL, + DS4_N_EXPERT / 2u, + true) != 0; + } +#endif + if (ok) { + if (cuda_tp_ep) { + if (cuda_tp_moe_delay_reduce && + metal_graph_cuda_tp_ep_delay_reduce_requested() && + !g->decode_stage_profile) { + cuda_tp_ep_reduce_deferred = true; + cuda_tp_ep_fused_hc_reduce = + g->cuda_tp_ep_pack_exact && + metal_graph_cuda_tp_ep_fused_hc_reduce_requested(); + } else { + ok = metal_graph_cuda_tp_ep_finish_reduce( + g, + cuda_tp_home_tier, + cuda_tp_partner_tier, + cuda_tp_ep_direct_return, + cuda_tp_ep_return_bytes, + true); + } + } else if (cuda_tp_moe_delay_reduce && + !cuda_tp_shared && !cuda_tp_shared_fold) { + /* Defer the peer routed-half copy until after the shared + * expert gate/up launch below: that work depends only on + * ffn_norm, so the home stream computes it while waiting + * for the partner instead of idling at the copy. */ + cuda_tp_moe_peer_copy_deferred = true; + } else if (cuda_tp_moe_delay_reduce) { + ok = ds4_gpu_tensor_copy_xdev( + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + g->routed_out_by_tier[cuda_tp_partner_tier], + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + cuda_tp_moe_peer_tmp = ok; + } else if (!cuda_tp_shared_fold) { + ok = ds4_gpu_add_xdev_tensor(metal_graph_routed_out(g), + metal_graph_routed_out(g), + g->routed_out_by_tier[cuda_tp_partner_tier], + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + DS4_N_EMBD) != 0; + } + } + } + /* Real TP split slices the shared expert by intermediate lanes, which + * needs the unfused gate/up/swiglu/down sequence. */ + const bool tp_split_shared = g->tp_world == 2; + const bool q4_selected_shared_overlap = + metal_graph_use_q4_selected_shared_overlap() && + metal_graph_decode_q4_selected_slots_expected(g, + layer, + layer->ffn_gate_exps->bytes, + layer->ffn_down_exps->bytes); + const bool iq2_selected_shared_overlap = + metal_graph_use_iq2_selected_shared_overlap(g) && + metal_graph_decode_iq2_selected_slots_expected(g, layer); + const bool cuda_selected_shared_overlap = + metal_graph_use_cuda_selected_shared_overlap(g) && + metal_graph_decode_cuda_selected_slots_expected(g, layer); + const bool overlap_selected_shared = + ok && + g->tp_world < 2 && + !decode_stage_profile && + !metal_graph_decode_cpu_router_applicable(g, layer) && + layer->ffn_gate_tid2eid == NULL && + getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL && + (q4_selected_shared_overlap || + iq2_selected_shared_overlap || + cuda_selected_shared_overlap); + const bool async_selected_load = + overlap_selected_shared && + ((iq2_selected_shared_overlap && + metal_graph_use_iq2_selected_async_load(g)) || + cuda_selected_shared_overlap); + const bool selected_readahead_shared_delay = + ok && + g->tp_world < 2 && + !overlap_selected_shared && + !decode_stage_profile && + metal_graph_use_iq2_selected_readahead_shared_delay(g) && + metal_graph_decode_iq2_selected_slots_expected(g, layer) && + !metal_graph_decode_cpu_router_applicable(g, layer) && + layer->ffn_gate_tid2eid == NULL && + getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; + const bool cuda_stream_selected_load = + ok && + !overlap_selected_shared && + !selected_readahead_shared_delay && + g->ssd_streaming && + metal_graph_decode_cuda_selected_slots_expected(g, layer) && + layer->ffn_gate_tid2eid == NULL && + getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; + if (cuda_stream_selected_load) { + ok = metal_graph_decode_cuda_selected_load(g, + model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + } + if (selected_readahead_shared_delay) { + if (ok) { + ok = metal_graph_decode_selected_readahead_override(g, + model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + } + if (ok && fuse_shared_gate_up) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), + metal_graph_shared_up(g), + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP) != 0; + } else if (ok) { + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), + model, + layer->ffn_gate_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), + model, + layer->ffn_up_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), + shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); + if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + metal_graph_router_selected(g), metal_graph_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), + NULL, + il, + false) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); + } + if (ok && fuse_shared_down_hc) { + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok) { + ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 1); + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); + } + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_ffn_out(g) && + ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; + } + if (ok && keep_ffn_out) { + metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), + metal_graph_ffn_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_post(g), + metal_graph_hc_comb(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok && !fuse_shared_down_hc) { + ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); + } + return ok; + } + if (overlap_selected_shared) { + uint64_t selected_event = 0; + if (ok) ok = ds4_gpu_signal_selected_readback_ready(&selected_event) != 0; + metal_graph_selected_async_load async_load = {0}; + bool async_load_started = false; + const bool async_early_commit = + async_selected_load && + metal_graph_use_iq2_selected_async_early_commit(g); + if (ok && async_selected_load) { + ok = metal_graph_selected_async_load_start(&async_load, + g, + model, + layer, + il, + selected_event, + gate_expert_bytes, + down_expert_bytes); + async_load_started = ok; + } + if (ok && async_early_commit) { + ok = ds4_gpu_flush_commands() != 0; + } + if (ok && fuse_shared_gate_up) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), + metal_graph_shared_up(g), + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP) != 0; + } else if (ok) { + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), + model, + layer->ffn_gate_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), + model, + layer->ffn_up_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), + shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); + if (ok && !fuse_shared_down_hc) { + ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 1); + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); + if (async_load_started) { + const bool flush_ok = ds4_gpu_flush_commands() != 0; + bool finish_ok = + metal_graph_selected_async_load_finish(&async_load); + if (!finish_ok && async_load.ids_ok) { + /* The worker read valid ids but could not stage the load + * (it is not allowed to wait on in-flight cache entries). + * This thread is, so retry the same load synchronously. */ + const ds4_gpu_stream_expert_table retry_table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + finish_ok = + ds4_gpu_stream_expert_cache_begin_selected_load( + &retry_table, + async_load.selected_ids, + DS4_N_EXPERT_USED) != 0 && + ds4_gpu_routed_moe_set_selected_override( + async_load.selected_ids, + DS4_N_EXPERT_USED) != 0; + } + ok = ok && flush_ok && finish_ok; + } else if (ok) { + ok = ds4_gpu_commit_and_wait_selected_readback(selected_event, + "selected-id shared-overlap") != 0; + } + if (ok && !async_load_started) { + int32_t selected_ids[DS4_MAX_EXPERT_USED]; + ok = ds4_gpu_tensor_read(metal_graph_router_selected(g), + 0, + selected_ids, + (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) != 0 && + ds4_gpu_routed_moe_set_selected_override(selected_ids, + DS4_N_EXPERT_USED) != 0; + if (ok) { + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + ok = ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_ids, + DS4_N_EXPERT_USED) != 0; + } + } + if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + metal_graph_router_selected(g), metal_graph_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), + NULL, + il, + false) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); + } + if (ok && fuse_shared_down_hc) { + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); + } + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_ffn_out(g) && + ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; + } + if (ok && keep_ffn_out) { + metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), + metal_graph_ffn_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_post(g), + metal_graph_hc_comb(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok && !fuse_shared_down_hc) { + ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); + } + return ok; + } + /* Under the TP split the routed experts run after the shared expert so + * the sum6 kernel can fold the shared partial and write the slab slot + * directly (no separate local add). */ + const bool tp_fold_ffn = tp_split_shared && + !keep_ffn_out && + !metal_graph_directional_steering_ffn_enabled(g); + if (ok && !tp_fold_ffn && !cuda_tp_moe) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + metal_graph_router_selected(g), metal_graph_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), + NULL, + il, + false) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), + (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); + } + if (phase == METAL_DECODE_LAYER_TO_SHARED_MID || + phase == METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID) { + return ok; + } + if (ok && tp_split_shared) { + /* Shared expert lane slice: the fused gate/up/swiglu kernel covers + * this rank's half of the intermediate (row slicing is pure offset + * math), compact at the buffer base; the down k-slice below turns + * it into a partial output. */ + const uint32_t tp_half = shared_dim / 2; + uint64_t shexp_row_bytes = 0; + ok = metal_graph_dense_quant_row_bytes(layer->ffn_gate_shexp, + DS4_N_EMBD, + &shexp_row_bytes) && + layer->ffn_gate_shexp->type == layer->ffn_up_shexp->type; + const uint64_t tp_lane_off = (uint64_t)g->tp_rank * tp_half * shexp_row_bytes; + ok = ok && (tp_half % 32u) == 0; + if (!ok) { + fprintf(stderr, "ds4: TP shared expert width %u is not sliceable\n", shared_dim); + } + if (ok && layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), + metal_graph_shared_up(g), + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset + tp_lane_off, + layer->ffn_up_shexp->abs_offset + tp_lane_off, + DS4_N_EMBD, + tp_half, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP) != 0; + } else if (ok) { + ok = metal_graph_matmul_dense_quant_abs(metal_graph_shared_gate(g), + model, + layer->ffn_gate_shexp, + layer->ffn_gate_shexp->abs_offset + tp_lane_off, + DS4_N_EMBD, + tp_half, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = metal_graph_matmul_dense_quant_abs(metal_graph_shared_up(g), + model, + layer->ffn_up_shexp, + layer->ffn_up_shexp->abs_offset + tp_lane_off, + DS4_N_EMBD, + tp_half, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), + tp_half, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; + } + } else if (ok && fuse_shared_gate_up) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), + metal_graph_shared_up(g), + metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP) != 0; + } else { + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), + model, + layer->ffn_gate_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), + model, + layer->ffn_up_shexp, + DS4_N_EMBD, + shared_dim, + metal_graph_ffn_norm(g), + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), + shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); + if (ok && cuda_tp_ep_reduce_deferred) { + ok = metal_graph_cuda_tp_ep_finish_reduce( + g, + cuda_tp_home_tier, + cuda_tp_partner_tier, + cuda_tp_ep_direct_return, + cuda_tp_ep_return_bytes, + !cuda_tp_ep_fused_hc_reduce); + } + if (ok && cuda_tp_moe_peer_copy_deferred) { + ok = ds4_gpu_tensor_copy_xdev( + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + g->routed_out_by_tier[cuda_tp_partner_tier], + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + cuda_tp_moe_peer_tmp = ok; + } + if (ok && cuda_tp_shared_fold) { + bool switched_to_partner = false; + ok = ds4_gpu_set_current_device(cuda_tp_partner_tier) == 0; + switched_to_partner = ok; + if (ok) { + ok = ds4_gpu_add_tensor( + g->routed_out_by_tier[cuda_tp_partner_tier], + g->routed_out_by_tier[cuda_tp_partner_tier], + g->shared_out_by_tier[cuda_tp_partner_tier], + DS4_N_EMBD) != 0; + } + if (switched_to_partner && + ds4_gpu_set_current_device(cuda_tp_home_tier) != 0) { + ok = false; + } + if (ok) { + ok = ds4_gpu_add_tensor(metal_graph_routed_out(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + DS4_N_EMBD) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy_xdev( + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + g->routed_out_by_tier[cuda_tp_partner_tier], + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + cuda_tp_shared_fold_peer_tmp = ok; + } + if (ok) { + ok = cuda_tp_shared_fold_peer_tmp && + ds4_gpu_hc_expand_add_split_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_routed_out(g), + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + } else if (ok && cuda_tp_shared) { + /* shared_out already contains the reduced local and partner partials. */ + } else if (ok && cuda_tp_ep_fused_hc_reduce) { + ok = ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_down(g), + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + metal_graph_router_selected(g), + DS4_N_EXPERT / 2u, + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok && fuse_shared_down_hc) { + if (cuda_tp_moe_peer_tmp) { + ok = ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_out(g), + g->tp_peer_tmp_by_tier[cuda_tp_home_tier], + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else { + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + } else if (ok && tp_split_shared) { + ok = metal_graph_matmul_dense_quant_kslice(metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + (uint64_t)g->tp_rank * (shared_dim / 2), + shared_dim / 2, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 0); + } else if (ok) { + ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 1); + } + DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); + if (ok) { + metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); + } + if (ok && tp_fold_ffn) { + ok = ds4_gpu_routed_moe_one_tensor( + g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN], + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, gate_row_bytes, + down_expert_bytes, down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + metal_graph_router_selected(g), metal_graph_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), + metal_graph_shared_out(g), + il, + false) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("routed_moe_folded"); + } + ds4_gpu_tensor *tp_ffn_a = NULL; /* rank0/rank1 partials consumed */ + ds4_gpu_tensor *tp_ffn_b = NULL; /* directly by the HC expand */ + if (ok && g->tp_world == 2) { + /* Gate FFN: local partial = shared expert + owned routed experts. + * The HC expand below already sums two block vectors, so after the + * exchange the two rank partials feed it directly (canonical rank + * order) with no separate combine dispatch. The paths that need + * the materialized sum (ffn_out consumers) still builds it in + * routed_out. */ + const uint32_t tp_slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; + if (!tp_fold_ffn) { + ok = ds4_gpu_add_tensor(g->tp_out[tp_slot], metal_graph_shared_out(g), metal_graph_routed_out(g), + DS4_N_EMBD) != 0; + } + if (ok) ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_FFN) != 0; + if (ok) { + ds4_gpu_tensor *first = g->tp_rank == 0 ? g->tp_out[tp_slot] : g->tp_in[tp_slot]; + ds4_gpu_tensor *second = g->tp_rank == 0 ? g->tp_in[tp_slot] : g->tp_out[tp_slot]; + if (keep_ffn_out || metal_graph_directional_steering_ffn_enabled(g)) { + ok = ds4_gpu_add_tensor(metal_graph_routed_out(g), first, second, DS4_N_EMBD) != 0; + } else { + tp_ffn_a = first; + tp_ffn_b = second; + } + } + } + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_ffn_out(g) && + ds4_gpu_add_tensor(metal_graph_ffn_out(g), + g->tp_world == 2 ? g->tp_zero : metal_graph_shared_out(g), + metal_graph_routed_out(g), DS4_N_EMBD) != 0; + } + if (ok && keep_ffn_out) { + metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), + metal_graph_ffn_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_post(g), + metal_graph_hc_comb(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok && !cuda_tp_shared_fold && !fuse_shared_down_hc) { + ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), + tp_ffn_a ? tp_ffn_a : metal_graph_routed_out(g), + tp_ffn_a ? tp_ffn_b : + (g->tp_world == 2 ? g->tp_zero : metal_graph_shared_out(g)), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); +#undef DS4_METAL_PROFILE_DECODE_STAGE + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); + } + return ok; +} + +static bool metal_graph_encode_decode_layer( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t raw_row, + uint32_t n_raw, + int token) { + return metal_graph_encode_decode_layer_phase( + g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, + token, METAL_DECODE_LAYER_FULL); +} + +static bool metal_graph_output_logits_head_matmul( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + ds4_gpu_tensor *norm_full, + ds4_gpu_tensor *dst_logits, + uint32_t n_tokens, + uint64_t vocab_dim); + +/* Encode the final HC collapse, output norm, and vocab projection on Metal. */ +static bool metal_graph_encode_output_head( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint64_t vocab_dim) { + /* switch to head_tier before the output-head pipeline. + * Single-tier (placement == NULL): no-op (head_tier == 0 == active_tier). + * Note: head_tier was captured in metal_graph_alloc_raw_cap; this + * helper consults it directly (and also covers the case where the + * preceding decode layer ran on a different tier — copy_xdev ferries + * the active cur_hc across the boundary). */ + if (g->placement) { + if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; + } + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const bool output_stage_profile = g->output_stage_profile; + double output_stage_t0 = output_stage_profile ? now_sec() : 0.0; +#define DS4_METAL_PROFILE_OUTPUT_STAGE(name) do { \ + if (ok && output_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary("output", (name), DS4_N_LAYER, 0, 1, &output_stage_t0); \ + } \ + } while (0) + bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + DS4_METAL_PROFILE_OUTPUT_STAGE("hc_flat_norm"); + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), + model->map, + model->size, + weights->output_hc_fn->abs_offset, + hc_dim, + DS4_N_HC, + metal_graph_flat_hc(g), + 1) != 0; + DS4_METAL_PROFILE_OUTPUT_STAGE("hc_pre"); + if (ok) { + metal_graph_debug_dump_tensor("result_hc_pre", metal_graph_output_pre(g), DS4_N_HC, DS4_N_LAYER, 0); + } + if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), + metal_graph_output_pre(g), + model->map, + model->size, + weights->output_hc_scale->abs_offset, + weights->output_hc_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weights"); + if (ok) { + metal_graph_debug_dump_tensor("result_hc_weights", metal_graph_output_weights(g), DS4_N_HC, DS4_N_LAYER, 0); + } + bool output_sum_norm_fused = false; +#if defined(__APPLE__) + if (ok) { + output_sum_norm_fused = + ds4_gpu_hc_weighted_sum_norm_tensor( + metal_graph_output_embd(g), + metal_graph_output_norm(g), + metal_graph_cur_hc(g), + metal_graph_output_weights(g), + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_RMS_EPS) != 0; + if (!output_sum_norm_fused && + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL) { + ok = false; + } + } +#endif + if (ok && !output_sum_norm_fused) { + ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), + metal_graph_cur_hc(g), + metal_graph_output_weights(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weighted_sum"); + if (ok) { + metal_graph_debug_dump_tensor("result_hc", metal_graph_output_embd(g), DS4_N_EMBD, DS4_N_LAYER, 0); + } + if (ok && !output_sum_norm_fused) { + ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), + metal_graph_output_embd(g), + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } + DS4_METAL_PROFILE_OUTPUT_STAGE("output_norm"); + if (ok) { + metal_graph_debug_dump_tensor("result_norm", metal_graph_output_norm(g), DS4_N_EMBD, DS4_N_LAYER, 0); + } + if (ok && g->tp_world == 2 && g->tp_logits_half) { + /* Vocab-split: this rank computes its half of the head rows into + * its logits view; the halves are bit-identical to the full head + * (same kernel, same rows) and the worker ships its half to the + * leader after the eval. */ + const uint64_t tp_vhalf = vocab_dim / 2u; + uint64_t head_row_bytes = 0; + ok = metal_graph_dense_quant_row_bytes(weights->output, + DS4_N_EMBD, + &head_row_bytes); + if (ok) ok = metal_graph_matmul_dense_quant_abs(g->tp_logits_half, + model, + weights->output, + weights->output->abs_offset + + (uint64_t)g->tp_rank * tp_vhalf * head_row_bytes, + DS4_N_EMBD, + tp_vhalf, + metal_graph_output_norm(g), + 1); + } else if (ok && g->cuda_tp_ep && g->cuda_tp_output) { + ok = metal_graph_output_logits_head_matmul( + g, model, weights, metal_graph_output_norm(g), + metal_graph_logits(g), 1, vocab_dim); + } else if (ok) { + ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), + model, + weights->output, + DS4_N_EMBD, + vocab_dim, + metal_graph_output_norm(g), + 1); + } + if (ok) { + metal_graph_debug_dump_tensor("result_output", metal_graph_logits(g), vocab_dim, DS4_N_LAYER, 0); + } +#undef DS4_METAL_PROFILE_OUTPUT_STAGE + return ok; +} + +/* Greedy-only output head: compute one local top-1 candidate per output TP + * split and leave the full split logits on their owning tiers. This avoids + * gathering the whole vocabulary row back to the head tier when the caller only + * needs the next argmax token. */ +static bool metal_graph_encode_output_head_split_top1( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint64_t vocab_dim, + int cuda_tp_output_tiers[DS4_MAX_GPUS], + uint32_t *cuda_tp_output_ways_out) { + if (!g || !model || !weights || + !cuda_tp_output_tiers || !cuda_tp_output_ways_out || + vocab_dim > UINT32_MAX) { + return false; + } + *cuda_tp_output_ways_out = 0; + + if (g->placement) { + if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; + } + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), + metal_graph_cur_hc(g), + (uint32_t)hc_dim, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), + model->map, + model->size, + weights->output_hc_fn->abs_offset, + hc_dim, + DS4_N_HC, + metal_graph_flat_hc(g), + 1) != 0; + if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), + metal_graph_output_pre(g), + model->map, + model->size, + weights->output_hc_scale->abs_offset, + weights->output_hc_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), + metal_graph_cur_hc(g), + metal_graph_output_weights(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), + metal_graph_output_embd(g), + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + if (!ok) return false; + + const uint32_t cuda_tp_output_ways = + g->cuda_tp_output ? metal_graph_cuda_tp_output_tiers(g, cuda_tp_output_tiers) : 0; + const bool cuda_tp_output = + g->cuda_tp_output && + cuda_tp_output_ways >= 2u && + weights->output->type == DS4_TENSOR_Q8_0 && + weights->output->ndim == 2 && + weights->output->dim[0] == DS4_N_EMBD && + weights->output->dim[1] == vocab_dim && + vocab_dim >= 2; + if (!cuda_tp_output) return false; + + for (uint32_t i = 0; i < cuda_tp_output_ways; i++) { + const int t = cuda_tp_output_tiers[i]; + if (t < 0 || t >= DS4_MAX_GPUS || + !g->output_norm_by_tier[t] || + !g->logits_by_tier[t] || + !g->comp_selected_by_tier[t] || + !g->comp_mask_by_tier[t]) { + return false; + } + } + + const bool fused_top1 = metal_graph_cuda_output_fused_top1_requested(); + const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); + uint64_t split_start[DS4_MAX_GPUS] = {0}; + uint64_t split_count[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor split_logits[DS4_MAX_GPUS]; + memset(split_logits, 0, sizeof(split_logits)); + + for (uint32_t i = 0; ok && i < cuda_tp_output_ways; i++) { + const int t = cuda_tp_output_tiers[i]; + split_start[i] = (vocab_dim * (uint64_t)i) / cuda_tp_output_ways; + const uint64_t split_end = + (vocab_dim * (uint64_t)(i + 1u)) / cuda_tp_output_ways; + split_count[i] = split_end - split_start[i]; + if (split_count[i] == 0 || split_count[i] > UINT32_MAX) { + ok = false; + break; + } + if (fused_top1) { + if (t != g->head_tier) { + ok = ds4_gpu_tensor_copy_xdev( + g->output_norm_by_tier[t], + metal_graph_output_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } + } else if (t == g->head_tier) { + ok = metal_graph_borrow_tensor_view(&split_logits[i], + metal_graph_logits(g), + split_start[i] * sizeof(float), + split_count[i] * sizeof(float)); + } else { + ok = metal_graph_borrow_tensor_view(&split_logits[i], + g->logits_by_tier[t], + 0, + split_count[i] * sizeof(float)) && + ds4_gpu_tensor_copy_xdev( + g->output_norm_by_tier[t], + metal_graph_output_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } + } + + for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { + const int t = cuda_tp_output_tiers[i]; + ok = ds4_gpu_set_current_device(t) == 0; + if (ok && fused_top1) { + ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[t], + g->comp_mask_by_tier[t], + model->map, + model->size, + weights->output->abs_offset + + split_start[i] * row_bytes, + DS4_N_EMBD, + split_count[i], + g->output_norm_by_tier[t], + (uint32_t)split_start[i]) != 0; + } else if (ok) { + ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[i], + model->map, + model->size, + weights->output->abs_offset + + split_start[i] * row_bytes, + DS4_N_EMBD, + split_count[i], + g->output_norm_by_tier[t], + 1) != 0; + } + if (ok && !fused_top1) { + ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[t], + g->comp_mask_by_tier[t], + &split_logits[i], + (uint32_t)split_count[i], + 1, + (uint32_t)split_start[i]) != 0; + } + } + if (ok) ok = ds4_gpu_set_current_device(g->head_tier) == 0; + if (ok && fused_top1) { + ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[g->head_tier], + g->comp_mask_by_tier[g->head_tier], + model->map, + model->size, + weights->output->abs_offset, + DS4_N_EMBD, + split_count[0], + metal_graph_output_norm(g), + (uint32_t)split_start[0]) != 0; + } else if (ok) { + ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[0], + model->map, + model->size, + weights->output->abs_offset, + DS4_N_EMBD, + split_count[0], + metal_graph_output_norm(g), + 1) != 0; + } + if (ok && !fused_top1) { + ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[g->head_tier], + g->comp_mask_by_tier[g->head_tier], + &split_logits[0], + (uint32_t)split_count[0], + 1, + (uint32_t)split_start[0]) != 0; + } + for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { + const int t = cuda_tp_output_tiers[i]; + ds4_gpu_tensor head_id_dst; + ds4_gpu_tensor head_value_dst; + ok = metal_graph_borrow_tensor_view(&head_id_dst, + g->comp_selected_by_tier[g->head_tier], + (uint64_t)i * sizeof(uint32_t), + sizeof(uint32_t)) && + metal_graph_borrow_tensor_view(&head_value_dst, + g->comp_mask_by_tier[g->head_tier], + (uint64_t)i * sizeof(float), + sizeof(float)) && + ds4_gpu_tensor_copy_xdev3(&head_id_dst, + g->comp_selected_by_tier[t], + sizeof(uint32_t), + &head_value_dst, + g->comp_mask_by_tier[t], + sizeof(float), + NULL, + NULL, + 0) != 0; + } + if (ok) { + ok = ds4_gpu_set_current_device(g->head_tier) == 0; + *cuda_tp_output_ways_out = cuda_tp_output_ways; + } + return ok; +} + +static bool metal_graph_read_output_split_top1( + ds4_gpu_graph *g, + uint32_t output_ways, + int *top_id) { + if (!g || !top_id || output_ways == 0 || output_ways > DS4_MAX_GPUS) { + return false; + } + bool have_best = false; + uint32_t best_id = 0; + float best_value = 0.0f; + uint32_t cand_ids[DS4_MAX_GPUS] = {0}; + float cand_values[DS4_MAX_GPUS] = {0.0f}; + bool ok = ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], + 0, + cand_ids, + (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && + ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], + 0, + cand_values, + (uint64_t)output_ways * sizeof(cand_values[0])) != 0; + for (uint32_t i = 0; ok && i < output_ways; i++) { + const uint32_t cand_id = cand_ids[i]; + const float cand_value = cand_values[i]; + if (!have_best || + cand_value > best_value || + (cand_value == best_value && cand_id < best_id)) { + have_best = true; + best_id = cand_id; + best_value = cand_value; + } + } + ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; + if (ok) *top_id = (int)best_id; + return ok; +} + +/* Batched output head for speculative verification. + * + * A target verifier only needs top-1 ids for intermediate draft rows and full + * logits for the last accepted row. Running the normal one-row output head in + * a loop serializes the HC collapse, output norm, and Q8 vocab projection. For + * tiny MTP suffixes we instead process all rows together and let the GPU reduce + * each row to a top id; the CPU reads back just those ids plus the last row's + * logits needed to continue the exact target stream. */ +/* Shared vocab-head matmul: pads small batches to 8 rows for the exact-mma Q8 + * kernel and shards the vocabulary across output-TP tiers. */ +static bool metal_graph_output_logits_head_matmul( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + ds4_gpu_tensor *norm_full, + ds4_gpu_tensor *dst_logits, + uint32_t n_tokens, + uint64_t vocab_dim) { + if (!g || !model || !weights || !norm_full || n_tokens == 0 || + !dst_logits || + ds4_gpu_tensor_bytes(dst_logits) < + (uint64_t)n_tokens * vocab_dim * sizeof(float)) { + return false; + } + const uint32_t head_rows = + (n_tokens > 1 && n_tokens < 8 && + ds4_gpu_tensor_bytes(dst_logits) >= 8u * vocab_dim * sizeof(float) && + ds4_gpu_tensor_bytes(norm_full) >= + 8u * DS4_N_EMBD * sizeof(float)) ? 8u : n_tokens; + ds4_gpu_tensor *output_norm = + ds4_gpu_tensor_view(norm_full, + 0, + (uint64_t)head_rows * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *logits = + ds4_gpu_tensor_view(dst_logits, + 0, + (uint64_t)head_rows * vocab_dim * sizeof(float)); + bool ok = output_norm && logits; + if (ok && head_rows > n_tokens) { + ds4_gpu_tensor *pad = + ds4_gpu_tensor_view(norm_full, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float), + (uint64_t)(head_rows - n_tokens) * DS4_N_EMBD * + sizeof(float)); + ok = pad && + ds4_gpu_tensor_fill_f32(pad, + 0.0f, + (uint64_t)(head_rows - n_tokens) * + DS4_N_EMBD) != 0; + ds4_gpu_tensor_free(pad); + } + /* Output TP for the speculative batch, mirroring the decode head: each + * device matmuls its VRAM-resident vocab shard. Shard outputs land + * compactly in logits_by_tier[t] ([head_rows x split]) and are gathered + * into spec_logits rows. */ + int tp_tiers[DS4_MAX_GPUS] = {0}; + const uint32_t tp_ways = (ok && g->cuda_tp_output) + ? metal_graph_cuda_tp_output_tiers(g, tp_tiers) : 0; + bool tp_ok = ok && tp_ways >= 2u && + weights->output->type == DS4_TENSOR_Q8_0 && + weights->output->ndim == 2 && + weights->output->dim[0] == DS4_N_EMBD && + weights->output->dim[1] == vocab_dim && + head_rows <= DS4_DSPARK_MAX_BLOCK_SIZE && + getenv("DS4_DSPARK_VERIFY_HEAD_NO_TP") == NULL; + for (uint32_t i = 0; tp_ok && i < tp_ways; i++) { + const int t = tp_tiers[i]; + tp_ok = t >= 0 && t < DS4_MAX_GPUS && + g->logits_by_tier[t] && + ds4_gpu_tensor_bytes(g->logits_by_tier[t]) >= + (uint64_t)head_rows * + ((vocab_dim + tp_ways - 1u) / tp_ways) * + sizeof(float) && + (t == g->active_tier || + (g->batch_ffn_norm_by_tier[t] && + ds4_gpu_tensor_bytes(g->batch_ffn_norm_by_tier[t]) >= + (uint64_t)head_rows * DS4_N_EMBD * sizeof(float))); + } + if (tp_ok) { + const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); + const int home_tier = g->active_tier; + uint64_t split_start[DS4_MAX_GPUS] = {0}; + uint64_t split_count[DS4_MAX_GPUS] = {0}; + for (uint32_t i = 0; ok && i < tp_ways; i++) { + const int t = tp_tiers[i]; + split_start[i] = (vocab_dim * (uint64_t)i) / tp_ways; + const uint64_t split_end = + (vocab_dim * (uint64_t)(i + 1u)) / tp_ways; + split_count[i] = split_end - split_start[i]; + if (split_count[i] == 0) { ok = false; break; } + if (t != home_tier) { + ok = ds4_gpu_tensor_copy_xdev( + g->batch_ffn_norm_by_tier[t], + output_norm, + (uint64_t)head_rows * DS4_N_EMBD * + sizeof(float)) != 0; + } + } + for (uint32_t i = 0; ok && i < tp_ways; i++) { + const int t = tp_tiers[i]; + ok = ds4_gpu_set_current_device(t) == 0; + if (!ok) break; + ds4_gpu_tensor *shard_out = + ds4_gpu_tensor_view(g->logits_by_tier[t], + 0, + (uint64_t)head_rows * split_count[i] * + sizeof(float)); + ds4_gpu_tensor *shard_in = t == home_tier ? + NULL : + ds4_gpu_tensor_view(g->batch_ffn_norm_by_tier[t], + 0, + (uint64_t)head_rows * DS4_N_EMBD * + sizeof(float)); + ok = shard_out && (t == home_tier || shard_in) && + ds4_gpu_matmul_q8_0_tensor(shard_out, + model->map, + model->size, + weights->output->abs_offset + + split_start[i] * row_bytes, + DS4_N_EMBD, + split_count[i], + t == home_tier ? output_norm : + shard_in, + head_rows) != 0; + ds4_gpu_tensor_free(shard_in); + ds4_gpu_tensor_free(shard_out); + } + if (ok) ok = ds4_gpu_set_current_device(home_tier) == 0; + for (uint32_t i = 0; ok && i < tp_ways; i++) { + const int t = tp_tiers[i]; + for (uint32_t r = 0; ok && r < n_tokens; r++) { + ds4_gpu_tensor *dst = + ds4_gpu_tensor_view(dst_logits, + ((uint64_t)r * vocab_dim + + split_start[i]) * sizeof(float), + split_count[i] * sizeof(float)); + ds4_gpu_tensor *src = + ds4_gpu_tensor_view(g->logits_by_tier[t], + (uint64_t)r * split_count[i] * + sizeof(float), + split_count[i] * sizeof(float)); + ok = dst && src && + ds4_gpu_tensor_copy_xdev(dst, + src, + split_count[i] * + sizeof(float)) != 0; + ds4_gpu_tensor_free(src); + ds4_gpu_tensor_free(dst); + } + } + } else if (ok && !(g->cuda_tp_ep && g->cuda_tp_output)) { + ok = ds4_gpu_matmul_q8_0_tensor(logits, + model->map, + model->size, + weights->output->abs_offset, + DS4_N_EMBD, + vocab_dim, + output_norm, + head_rows) != 0; + } else if (ok) { + /* The expert-parallel cache stores only output vocabulary shards, so + * a single-device full-head fallback would access uncached weights. */ + ok = false; + } + ds4_gpu_tensor_free(logits); + ds4_gpu_tensor_free(output_norm); + return ok; +} + +static bool metal_graph_encode_output_head_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t n_tokens, + uint64_t vocab_dim) { + if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + ds4_gpu_tensor *output_pre = NULL; + ds4_gpu_tensor *output_weights = NULL; + ds4_gpu_tensor *output_embd = NULL; + ds4_gpu_tensor *output_norm = NULL; + + bool ok = true; + output_pre = ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), + 0, + (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); + output_weights = ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + 0, + (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); + output_embd = ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), + 0, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + output_norm = ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), + 0, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + ok = output_pre && output_weights && output_embd && output_norm; + + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), + metal_graph_batch_cur_hc(g), + (uint32_t)hc_dim, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(output_pre, + model->map, + model->size, + weights->output_hc_fn->abs_offset, + hc_dim, + DS4_N_HC, + metal_graph_batch_flat_hc(g), + n_tokens) != 0; + if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, + output_pre, + model->map, + model->size, + weights->output_hc_scale->abs_offset, + weights->output_hc_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, + metal_graph_batch_cur_hc(g), + output_weights, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, + output_embd, + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_output_logits_head_matmul( + g, model, weights, metal_graph_batch_ffn_norm(g), + g->spec_logits, n_tokens, vocab_dim); + + ds4_gpu_tensor_free(output_norm); + ds4_gpu_tensor_free(output_embd); + ds4_gpu_tensor_free(output_weights); + ds4_gpu_tensor_free(output_pre); + return ok; +} + +static bool metal_graph_matmul_plain_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (w->type == DS4_TENSOR_F16) { + return ds4_gpu_matmul_f16_tensor(out, model->map, model->size, + w->abs_offset, in_dim, out_dim, x, n_tok) != 0; + } + if (w->type == DS4_TENSOR_F32) { + return ds4_gpu_matmul_f32_tensor(out, model->map, model->size, + w->abs_offset, in_dim, out_dim, x, n_tok) != 0; + } + if (w->type == DS4_TENSOR_Q8_0) { + return ds4_gpu_matmul_q8_0_tensor(out, model->map, model->size, + w->abs_offset, in_dim, out_dim, x, n_tok) != 0; + } + if (tensor_type_is_dense_quant(w->type)) { + return ds4_gpu_matmul_quant_tensor(out, + model->map, + model->size, + w->abs_offset, + w->type, + in_dim, + out_dim, + x, + n_tok) != 0; + } + fprintf(stderr, "ds4: Metal plain matmul does not support %s\n", tensor_type_name(w->type)); + return false; +} + +static bool metal_graph_dense_quant_row_bytes( + const ds4_tensor *w, + uint64_t in_dim, + uint64_t *row_bytes) { + if (row_bytes) *row_bytes = 0; + if (!w || !row_bytes || !tensor_type_is_dense_quant(w->type)) return false; + return tensor_nbytes(w->type, in_dim, row_bytes); +} + +static bool metal_graph_matmul_dense_quant_abs( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!w || !tensor_type_is_dense_quant(w->type)) return false; + return ds4_gpu_matmul_quant_tensor(out, + model->map, + model->size, + weight_offset, + w->type, + in_dim, + out_dim, + x, + n_tok) != 0; +} + +static bool metal_graph_matmul_dense_quant_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!w) return false; + return metal_graph_matmul_dense_quant_abs(out, + model, + w, + w->abs_offset, + in_dim, + out_dim, + x, + n_tok); +} + +static bool metal_graph_matmul_dense_quant_kslice( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + 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) { + if (!w || !tensor_type_is_dense_quant(w->type)) return false; + return ds4_gpu_matmul_quant_kslice_tensor(out, + model->map, + model->size, + w->abs_offset, + w->type, + full_in_dim, + k_off, + k_cnt, + out_dim, + x, + x_elem_off) != 0; +} + +static bool metal_graph_attention_output_dense_quant_low( + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + uint64_t group_dim, + uint64_t rank, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads) { + (void)g; + if (!low || !model || !out_a || !heads || + group_dim == 0 || rank == 0 || group_cnt == 0) { + return false; + } + if (out_a->type == DS4_TENSOR_Q8_0 && group0 == 0) { + return ds4_gpu_attention_output_low_q8_tensor(low, + model->map, + model->size, + out_a->abs_offset, + group_dim, + rank, + group_cnt, + heads) != 0; + } + if (out_a->type == DS4_TENSOR_Q4_K) { + return ds4_gpu_attention_output_low_q4_K_slice_tensor(low, + model->map, + model->size, + out_a->abs_offset, + group_dim, + rank, + group0, + group_cnt, + heads) != 0; + } + uint64_t row_bytes = 0; + if (!metal_graph_dense_quant_row_bytes(out_a, group_dim, &row_bytes)) return false; + const uint64_t group_weight_bytes = rank * row_bytes; + bool ok = true; + for (uint32_t i = 0; ok && i < group_cnt; i++) { + ds4_gpu_tensor *head_view = ds4_gpu_tensor_view( + heads, + (uint64_t)i * group_dim * sizeof(float), + group_dim * sizeof(float)); + ds4_gpu_tensor *low_view = ds4_gpu_tensor_view( + low, + (uint64_t)i * rank * sizeof(float), + rank * sizeof(float)); + ok = head_view && low_view && + metal_graph_matmul_dense_quant_abs(low_view, + model, + out_a, + out_a->abs_offset + + (uint64_t)(group0 + i) * group_weight_bytes, + group_dim, + rank, + head_view, + 1); + ds4_gpu_tensor_free(low_view); + ds4_gpu_tensor_free(head_view); + } + return ok; +} + +static bool metal_graph_attention_output_dense_quant_tp( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + const ds4_tensor *out_b, + 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) { + if (!out || !low || !g || !model || !out_a || !out_b || !heads || + group0 + group_cnt > n_groups_total) { + return false; + } + if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { + return ds4_gpu_attention_output_q8_tp_tensor(out, + low, + model->map, + model->size, + out_a->abs_offset, + out_b->abs_offset, + group_dim, + rank, + n_groups_total, + group0, + group_cnt, + out_dim, + heads) != 0; + } + if (!metal_graph_attention_output_dense_quant_low(low, + g, + model, + out_a, + group_dim, + rank, + group0, + group_cnt, + heads)) { + return false; + } + return metal_graph_matmul_dense_quant_kslice(out, + model, + out_b, + (uint64_t)n_groups_total * rank, + (uint64_t)group0 * rank, + (uint64_t)group_cnt * rank, + out_dim, + low, + 0); +} + +static bool metal_graph_attention_output_dense_quant_batch( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_tensor *out_a, + const ds4_tensor *out_b, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!out || !low || !g || !model || !out_a || !out_b || !heads || + n_groups == 0 || n_tokens == 0) { + return false; + } + if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { + return ds4_gpu_attention_output_q8_batch_tensor(out, + low, + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + model->map, + model->size, + out_a->abs_offset, + out_b->abs_offset, + group_dim, + rank, + n_groups, + out_dim, + heads, + n_tokens) != 0; + } + if (out_a->type == DS4_TENSOR_Q4_K && n_tokens >= 32u) { + if (ds4_gpu_attention_output_q4_K_batch_tensor(out, + low, + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + model->map, + model->size, + out_a->abs_offset, + out_b->abs_offset, + out_b->type, + group_dim, + rank, + n_groups, + out_dim, + heads, + n_tokens) != 0) { + return true; + } + } + + const uint64_t heads_row_elems = (uint64_t)n_groups * group_dim; + const uint64_t low_row_elems = (uint64_t)n_groups * rank; + bool ok = true; + for (uint32_t t = 0; ok && t < n_tokens; t++) { + ds4_gpu_tensor *heads_row = ds4_gpu_tensor_view( + heads, + (uint64_t)t * heads_row_elems * sizeof(float), + heads_row_elems * sizeof(float)); + ds4_gpu_tensor *low_row = ds4_gpu_tensor_view( + low, + (uint64_t)t * low_row_elems * sizeof(float), + low_row_elems * sizeof(float)); + ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( + out, + (uint64_t)t * out_dim * sizeof(float), + out_dim * sizeof(float)); + ok = heads_row && low_row && out_row && + metal_graph_attention_output_dense_quant_low(low_row, + g, + model, + out_a, + group_dim, + rank, + 0, + n_groups, + heads_row); + if (ok) ok = metal_graph_matmul_dense_quant_tensor(out_row, + model, + out_b, + low_row_elems, + out_dim, + low_row, + 1); + ds4_gpu_tensor_free(out_row); + ds4_gpu_tensor_free(low_row); + ds4_gpu_tensor_free(heads_row); + } + return ok; +} + +static bool metal_graph_matmul_q8_0_named_tensor( + const char *module, + uint32_t il, + uint32_t pos0, + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + (void)module; + (void)il; + (void)pos0; + return metal_graph_matmul_dense_quant_tensor(out, + model, + w, + in_dim, + out_dim, + x, + n_tok); +} + +static bool metal_graph_encode_output_head_mtp( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_model *mtp_model, + const ds4_mtp_weights *mtp, + uint64_t vocab_dim) { + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_output_pre(g), mtp_model, mtp->hc_head_fn, + hc_dim, DS4_N_HC, metal_graph_flat_hc(g), 1); + if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), + metal_graph_output_pre(g), + mtp_model->map, + mtp_model->size, + mtp->hc_head_scale->abs_offset, + mtp->hc_head_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), + metal_graph_cur_hc(g), + metal_graph_output_weights(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), + metal_graph_output_embd(g), + mtp_model->map, + mtp_model->size, + mtp->norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), + base_model, + base_weights->output, + DS4_N_EMBD, + vocab_dim, + metal_graph_output_norm(g), + 1); + return ok; +} + +/* ========================================================================= + * Metal Diagnostic Comparisons. + * ========================================================================= + * + * These routines deliberately allocate CPU-side reference buffers and read + * Metal tensors back. They are not part of generation; command-line tests use + * them to localize drift against the C reference pipeline. + */ + +static void metal_graph_trace_layer_stages( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + const float *cpu_in_hc, + uint32_t il, + int token) { + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t shared_in_dim = layer->ffn_gate_shexp->dim[0]; + const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; + const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; + + float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); + float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); + float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); + float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); + float *cpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); + float *cpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); + float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float post[4]; + float comb[16]; + float ffn_post[4]; + float ffn_comb[16]; + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + const uint64_t shared_blocks = (shared_in_dim + 31) / 32; + int8_t *shared_xq = xmalloc((size_t)shared_blocks * 32); + float *shared_xscale = xmalloc((size_t)shared_blocks * sizeof(float)); + float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); + block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); + block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); + int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; + float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; + int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; + float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; + + hc_pre_from_state_one(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + cpu_in_hc, cpu_attn_cur, post, comb); + layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); + layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); + layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); + rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); + rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); + dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); + f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); + layer_attention_one(cpu_heads, model, layer, cpu_q, cpu_kv); + rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, true); + layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); + hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_in_hc, post, comb, DS4_N_EMBD, DS4_N_HC); + hc_pre_from_state_one(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + cpu_after_attn_hc, cpu_ffn_cur, ffn_post, ffn_comb); + rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); + quantize_q8_0_activation(cpu_ffn_norm, shared_xq, shared_xscale, shared_in_dim); + matvec_q8_0_pair_prequant(cpu_shared_gate, + cpu_shared_up, + model, + layer->ffn_gate_shexp, + layer->ffn_up_shexp, + shared_xq, + shared_xscale); + swiglu(cpu_shared_mid, cpu_shared_gate, cpu_shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP); + matvec_q8_0(cpu_shared, model, layer->ffn_down_shexp, cpu_shared_mid); + layer_routed_moe_one_prealloc(cpu_routed, + model, + layer, + cpu_ffn_norm, + il, + token, + DS4_SWIGLU_CLAMP_EXP, + routed_mid_all, + routed_xq, + routed_midq, + routed_q8_xq, + routed_q8_xscale, + routed_q8_midq, + routed_q8_midscale); + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(selected, model, layer, token); + layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); + } else { + layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); + } + for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; + hc_post_one(cpu_after_ffn_hc, cpu_ffn_out, cpu_after_attn_hc, ffn_post, ffn_comb, DS4_N_EMBD, DS4_N_HC); + + float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); + float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); + float *gpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); + float *gpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); + float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); + float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + int gpu_selected[DS4_MAX_EXPERT_USED]; + float gpu_expert_weight[DS4_MAX_EXPERT_USED]; + + bool ok = ds4_gpu_tensor_read(metal_graph_attn_cur(g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_attn_norm(g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_q(g), 0, gpu_q, q_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_kv(g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_attn_out(g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_after_attn_hc(g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_cur(g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_norm(g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_shared_gate(g), 0, gpu_shared_gate, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_shared_up(g), 0, gpu_shared_up, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_shared_mid(g), 0, gpu_shared_mid, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_shared_out(g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_router_selected(g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && + ds4_gpu_tensor_read(metal_graph_router_weights(g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && + ds4_gpu_tensor_read(metal_graph_routed_mid(g), 0, gpu_routed_mid_all, (uint64_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_routed_out(g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_out(g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_cur_hc(g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0; + + if (ok) { + fprintf(stderr, + "ds4: Metal stage layer %u attn_cur=%g/%g attn_norm=%g/%g q=%g/%g kv=%g/%g attn_out=%g/%g after_attn_hc=%g/%g ffn_cur=%g/%g ffn_norm=%g/%g shared=%g/%g router_w=%g routed=%g/%g ffn_out=%g/%g after_ffn_hc=%g/%g\n", + il, + max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), rms_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), + max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), rms_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), + max_abs_diff(cpu_q, gpu_q, q_dim), rms_abs_diff(cpu_q, gpu_q, q_dim), + max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), rms_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), + max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), rms_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), + max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), rms_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), + max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), rms_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), + max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), rms_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), + max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), + max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), + max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), + max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), rms_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), + max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), rms_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim)); + fprintf(stderr, + "ds4: Metal shared layer %u gate=%g/%g up=%g/%g mid=%g/%g down=%g/%g\n", + il, + max_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), rms_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), + max_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), rms_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), + max_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), rms_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), + max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD)); + fprintf(stderr, + "ds4: Metal routed layer %u mid=%g/%g out=%g/%g\n", + il, + max_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), + rms_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), + max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), + rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD)); + if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { + fprintf(stderr, + "ds4: Metal stage layer %u router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", + il, + selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], + gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); + } + } + + free(gpu_after_ffn_hc); + free(gpu_ffn_out); + free(gpu_routed); + free(gpu_routed_mid_all); + free(gpu_shared); + free(gpu_shared_mid); + free(gpu_shared_up); + free(gpu_shared_gate); + free(gpu_ffn_norm); + free(gpu_ffn_cur); + free(gpu_after_attn_hc); + free(gpu_attn_out); + free(gpu_kv); + free(gpu_q); + free(gpu_attn_norm); + free(gpu_attn_cur); + free(routed_q8_midscale); + free(routed_q8_midq); + free(routed_q8_xscale); + free(routed_q8_xq); + free(routed_midq); + free(routed_xq); + free(routed_mid_all); + free(shared_xscale); + free(shared_xq); + free(cpu_after_ffn_hc); + free(cpu_ffn_out); + free(cpu_routed); + free(cpu_shared); + free(cpu_shared_mid); + free(cpu_shared_up); + free(cpu_shared_gate); + free(cpu_ffn_norm); + free(cpu_ffn_cur); + free(cpu_after_attn_hc); + free(cpu_attn_out); + free(cpu_heads); + free(cpu_kv); + free(cpu_qr_norm); + free(cpu_q); + free(cpu_attn_norm); + free(cpu_attn_cur); +} + +static int metal_graph_decode_test( + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + bool quality) { + if (prompt->len <= 0) { + fprintf(stderr, "ds4: Metal graph test needs a non-empty prompt\n"); + return 1; + } + + const int token = prompt->v[0]; + const ds4_layer_weights *layer = &weights->layer[0]; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const uint64_t vocab_dim = weights->output->dim[1]; + const bool routed_q8_0 = + layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; + const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; + const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; + + float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); + float *cpu_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); + float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); + float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); + float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); + float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_ffn_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); + float *cpu_ffn_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); + float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); + float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); + float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + float *gpu_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); + float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); + int gpu_selected[DS4_MAX_EXPERT_USED]; + float gpu_expert_weight[DS4_MAX_EXPERT_USED]; + float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); + block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); + block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); + int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; + float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; + int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; + float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + + embed_token_f16(model, weights, token, plain); + hc_from_plain_embedding(cpu_hc, plain, DS4_N_EMBD, DS4_N_HC); + hc_pre_from_state_one(model, + layer->hc_attn_fn, + layer->hc_attn_scale, + layer->hc_attn_base, + cpu_hc, cpu_attn_cur, cpu_post, cpu_comb); + layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); + layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); + layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); + rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); + rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); + dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); + f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); + layer_attention_rows_one(cpu_heads, model, layer, cpu_q, cpu_kv, 1); + rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, true); + layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); + hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_hc, cpu_post, cpu_comb, DS4_N_EMBD, DS4_N_HC); + hc_pre_from_state_one(model, + layer->hc_ffn_fn, + layer->hc_ffn_scale, + layer->hc_ffn_base, + cpu_after_attn_hc, cpu_ffn_cur, cpu_ffn_post, cpu_ffn_comb); + rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); + layer_shared_ffn_one(cpu_shared, model, layer, cpu_ffn_norm); + layer_routed_moe_one_prealloc(cpu_routed, + model, + layer, + cpu_ffn_norm, + 0, + token, + DS4_SWIGLU_CLAMP_EXP, + routed_mid_all, + routed_xq, + routed_midq, + routed_q8_xq, + routed_q8_xscale, + routed_q8_midq, + routed_q8_midscale); + if (layer->ffn_gate_tid2eid) { + layer_hash_selected_experts(selected, model, layer, token); + layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); + } else { + layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); + } + for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; + hc_post_one(cpu_after_ffn_hc, + cpu_ffn_out, + cpu_after_attn_hc, + cpu_ffn_post, + cpu_ffn_comb, + DS4_N_EMBD, + DS4_N_HC); + output_logits_one(cpu_logits, model, weights, cpu_after_ffn_hc); + + ds4_gpu_graph g; + bool ok = metal_graph_alloc(&g, weights, layer); + g.quality = quality; + g.materialize_ffn_out = true; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = metal_graph_encode_decode_layer(&g, + model, + layer, + 0, + 0, + g.layer_raw_cache[0], + g.raw_cap, + 0, + 1, + token); + if (ok) { + /* Single-tier diagnostic: swap the active-tier slots so the head + * pipeline reads the embedded hidden state from cur_hc. */ + ds4_gpu_tensor *embedded_hc = g.cur_hc_by_tier[g.active_tier]; + g.cur_hc_by_tier[g.active_tier] = g.after_ffn_hc_by_tier[g.active_tier]; + g.after_ffn_hc_by_tier[g.active_tier] = embedded_hc; + } + if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); + if (ok) ok = ds4_gpu_end_commands() != 0; + + if (ok) { + ok = ds4_gpu_tensor_read(metal_graph_after_ffn_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_attn_cur(&g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_attn_norm(&g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_q(&g), 0, gpu_q, q_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_kv(&g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.layer_raw_cache[0], 0, gpu_raw, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_attn_out(&g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_after_attn_hc(&g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_cur(&g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_norm(&g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_shared_out(&g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_router_selected(&g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && + ds4_gpu_tensor_read(metal_graph_router_weights(&g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && + ds4_gpu_tensor_read(metal_graph_routed_out(&g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_ffn_out(&g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; + } + + if (ok) { + fprintf(stderr, + "ds4: Metal graph test layer0 diffs: embed_hc=%g hc_pre=%g attn_norm=%g q_rope=%g kv_rope=%g raw_cache=%g attn_out=%g after_attn_hc=%g ffn_cur=%g ffn_norm=%g shared=%g router_w=%g routed=%g ffn_out=%g after_ffn_hc=%g logits=%g\n", + max_abs_diff(cpu_hc, gpu_hc, hc_dim), + max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), + max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), + max_abs_diff(cpu_q, gpu_q, q_dim), + max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), + max_abs_diff(cpu_kv, gpu_raw, DS4_N_HEAD_DIM), + max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), + max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), + max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), + max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), + max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), + max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), + max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), + max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), + max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), + max_abs_diff(cpu_logits, gpu_logits, vocab_dim)); + if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { + fprintf(stderr, + "ds4: Metal graph router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", + selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], + gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); + } + print_vec_stats("metal graph q", gpu_q, q_dim); + print_vec_stats("metal graph kv", gpu_kv, DS4_N_HEAD_DIM); + print_vec_stats("metal graph routed", gpu_routed, DS4_N_EMBD); + } else { + fprintf(stderr, "ds4: Metal graph test failed while encoding first decode stages\n"); + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after graph test failure also failed\n"); + } + } + + metal_graph_free(&g); + free(routed_q8_midscale); + free(routed_q8_midq); + free(routed_q8_xscale); + free(routed_q8_xq); + free(routed_midq); + free(routed_xq); + free(routed_mid_all); + free(gpu_logits); + free(gpu_after_ffn_hc); + free(gpu_ffn_out); + free(gpu_routed); + free(gpu_shared); + free(gpu_ffn_norm); + free(gpu_ffn_cur); + free(gpu_after_attn_hc); + free(gpu_attn_out); + free(gpu_raw); + free(gpu_kv); + free(gpu_q); + free(gpu_attn_norm); + free(gpu_attn_cur); + free(gpu_hc); + free(cpu_kv); + free(cpu_q); + free(cpu_attn_out); + free(cpu_heads); + free(cpu_ffn_norm); + free(cpu_routed); + free(cpu_logits); + free(cpu_after_ffn_hc); + free(cpu_ffn_out); + free(cpu_shared); + free(cpu_ffn_comb); + free(cpu_ffn_post); + free(cpu_ffn_cur); + free(cpu_after_attn_hc); + free(cpu_qr_norm); + free(cpu_attn_norm); + free(cpu_comb); + free(cpu_post); + free(cpu_attn_cur); + free(cpu_hc); + free(plain); + return ok ? 0 : 1; +} + +static int metal_graph_first_token_full_test( + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + bool quality) { + if (prompt->len <= 0) { + fprintf(stderr, "ds4: full Metal graph test needs a non-empty prompt\n"); + return 1; + } + + const int token = prompt->v[0]; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t vocab_dim = weights->output->dim[1]; + float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); + float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); + + forward_first_token_cpu(cpu_hc, model, weights, token); + output_logits_one(cpu_logits, model, weights, cpu_hc); + + ds4_gpu_graph g; + bool ok = metal_graph_alloc(&g, weights, &weights->layer[0]); + g.quality = quality; + const bool trace_layers = getenv("DS4_METAL_GRAPH_TRACE_LAYERS") != NULL; + if (trace_layers && ok) { + g.materialize_ffn_out = true; + const bool teacher_force = getenv("DS4_METAL_GRAPH_TEACHER_FORCE") != NULL; + const char *stage_layer_env = getenv("DS4_METAL_GRAPH_TRACE_STAGE_LAYER"); + const long stage_layer = stage_layer_env ? strtol(stage_layer_env, NULL, 10) : -1; + float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); + float *cpu_cur = xmalloc((size_t)hc_dim * sizeof(float)); + float *cpu_next = xmalloc((size_t)hc_dim * sizeof(float)); + + embed_token_f16(model, weights, token, plain); + hc_from_plain_embedding(cpu_cur, plain, DS4_N_EMBD, DS4_N_HC); + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_end_commands() != 0; + + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + if (teacher_force) { + ok = ds4_gpu_tensor_write(metal_graph_cur_hc(&g), 0, cpu_cur, hc_dim * sizeof(float)) != 0; + } + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], + il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); + ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); + g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); + g.after_ffn_hc_by_tier[g.active_tier] = tmp; + if (ok) ok = ds4_gpu_end_commands() != 0; + + layer_forward_self_one(cpu_next, model, &weights->layer[il], cpu_cur, il, 0, token); + if (ok) ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0; + if (ok) { + fprintf(stderr, + "ds4: Metal full graph layer %u%s hc_max=%g hc_rms=%g\n", + il, + teacher_force ? " teacher" : "", + max_abs_diff(cpu_next, gpu_hc, hc_dim), + rms_abs_diff(cpu_next, gpu_hc, hc_dim)); + if (stage_layer == (long)il) { + metal_graph_trace_layer_stages(&g, model, &weights->layer[il], cpu_cur, il, token); + } + } + float *ctmp = cpu_cur; + cpu_cur = cpu_next; + cpu_next = ctmp; + } + + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); + if (ok) ok = ds4_gpu_end_commands() != 0; + + free(cpu_next); + free(cpu_cur); + free(plain); + } else { + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], + il, 0, g.layer_raw_cache[il], + g.raw_cap, 0, 1, token); + ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); + g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); + g.after_ffn_hc_by_tier[g.active_tier] = tmp; + } + + if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); + if (ok) ok = ds4_gpu_end_commands() != 0; + } + + if (ok) { + ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; + } + + if (ok) { + const uint64_t cpu_top = argmax_f32(cpu_logits, vocab_dim); + const uint64_t gpu_top = argmax_f32(gpu_logits, vocab_dim); + fprintf(stderr, + "ds4: Metal full first-token graph diffs: final_hc_max=%g final_hc_rms=%g logits_max=%g logits_rms=%g cpu_top=%llu gpu_top=%llu cpu_top_logit=%g gpu_top_logit=%g\n", + max_abs_diff(cpu_hc, gpu_hc, hc_dim), + rms_abs_diff(cpu_hc, gpu_hc, hc_dim), + max_abs_diff(cpu_logits, gpu_logits, vocab_dim), + rms_abs_diff(cpu_logits, gpu_logits, vocab_dim), + (unsigned long long)cpu_top, + (unsigned long long)gpu_top, + cpu_logits[cpu_top], + gpu_logits[gpu_top]); + } else { + fprintf(stderr, "ds4: Metal full first-token graph test failed\n"); + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after full graph failure also failed\n"); + } + } + + metal_graph_free(&g); + free(gpu_logits); + free(cpu_logits); + free(gpu_hc); + free(cpu_hc); + return ok ? 0 : 1; +} + +/* ========================================================================= + * Metal Release Decode and Prefill. + * ========================================================================= + * + * Everything below is the user-facing Metal backend. It uses the same layer + * encoder as diagnostics, but diagnostics are not required for normal command + * flow and their CPU reads stay outside these generation entry points. + */ + +static uint32_t metal_graph_token_split_after_layers(void) { + uint32_t split_after_layers = 4; +#ifndef DS4_ROCM_BUILD + const char *split_env = getenv("DS4_METAL_GRAPH_TOKEN_SPLIT_LAYERS"); + if (split_env && split_env[0]) { + char *end = NULL; + unsigned long v = strtoul(split_env, &end, 10); + if (end != split_env && v <= DS4_N_LAYER) split_after_layers = (uint32_t)v; + } +#endif + return split_after_layers; +} + +static int metal_graph_dspark_target_slot( + const ds4_gpu_graph *g, + uint32_t il) { + if (!g || !g->dspark_capture_enabled) return -1; + for (uint32_t i = 0; i < g->dspark_target_layer_count; i++) { + if (g->dspark_target_layers[i] == il) return (int)i; + } + return -1; +} + +static uint32_t metal_graph_dspark_capture_complete_mask( + const ds4_gpu_graph *g) { + if (!g || g->dspark_target_layer_count == 0) return 0; + return g->dspark_target_layer_count >= 32u ? + UINT32_MAX : ((1u << g->dspark_target_layer_count) - 1u); +} + +static void metal_graph_dspark_capture_note_slot(ds4_gpu_graph *g, + uint32_t slot) { + if (!g || slot >= g->dspark_target_layer_count) return; + g->dspark_capture_mask |= 1u << slot; + g->dspark_capture_valid = + g->dspark_capture_mask == metal_graph_dspark_capture_complete_mask(g); +} + +static void metal_graph_dspark_capture_row_invalidate(ds4_gpu_graph *g) { + if (!g || !g->dspark_capture_enabled) return; + g->dspark_capture_mask = 0; + g->dspark_capture_checkpoint_len = 0; + g->dspark_capture_valid = false; +} + +static void metal_graph_dspark_capture_batch_invalidate(ds4_gpu_graph *g) { + if (!g || !g->dspark_capture_enabled) return; + g->dspark_capture_batch_mask = 0; + g->dspark_capture_batch_start = 0; + g->dspark_capture_batch_tokens = 0; + g->dspark_capture_batch_valid = false; +} + +static void metal_graph_dspark_capture_invalidate(ds4_gpu_graph *g) { + metal_graph_dspark_capture_row_invalidate(g); + metal_graph_dspark_capture_batch_invalidate(g); +} + +static void metal_graph_dspark_cache_reset(ds4_gpu_graph *g) { + if (!g) return; + g->dspark_cache_start = 0; + g->dspark_cache_token_start = 0; + g->dspark_cache_len = 0; +} + +static bool metal_graph_dspark_cache_window_valid( + const ds4_gpu_graph *g, + uint32_t token_start, + uint32_t raw_start, + uint32_t len) { + if (!g || len > g->dspark_cache_cap) return false; + if (len == 0) return true; + if (g->dspark_cache_cap == 0 || + raw_start >= g->dspark_cache_cap || + token_start > UINT32_MAX - len || + raw_start != token_start % g->dspark_cache_cap) { + return false; + } + return true; +} + +static bool metal_graph_dspark_cache_current_window_valid( + const ds4_gpu_graph *g) { + if (!g) return false; + return metal_graph_dspark_cache_window_valid(g, + g->dspark_cache_token_start, + g->dspark_cache_start, + g->dspark_cache_len); +} + +static bool metal_graph_dspark_cache_set_window(ds4_gpu_graph *g, + uint32_t token_start, + uint32_t len) { + if (!g || len > g->dspark_cache_cap) return false; + const uint32_t raw_start = + len && g->dspark_cache_cap ? token_start % g->dspark_cache_cap : 0; + if (!metal_graph_dspark_cache_window_valid(g, + len ? token_start : 0, + raw_start, + len)) { + return false; + } + g->dspark_cache_start = raw_start; + g->dspark_cache_token_start = len ? token_start : 0; + g->dspark_cache_len = len; + return true; +} + +static bool metal_graph_dspark_cache_crop_to_prefix(ds4_gpu_graph *g, + uint32_t prefix_len) { + if (!g) return false; + if (g->dspark_cache_len == 0) return true; + if (!metal_graph_dspark_cache_current_window_valid(g)) return false; + + const uint32_t start = g->dspark_cache_token_start; + const uint32_t end = start + g->dspark_cache_len; + if (prefix_len <= start || prefix_len > end) { + metal_graph_dspark_cache_reset(g); + return true; + } + g->dspark_cache_len = prefix_len - start; + return true; +} + +static bool metal_graph_dspark_cache_ends_at(const ds4_gpu_graph *g, + uint32_t pos) { + if (!metal_graph_dspark_cache_current_window_valid(g)) return false; + if (g->dspark_cache_len == 0) return true; + return g->dspark_cache_token_start <= UINT32_MAX - g->dspark_cache_len && + g->dspark_cache_token_start + g->dspark_cache_len == pos; +} + +static bool metal_graph_dspark_cache_claim_appended_row(ds4_gpu_graph *g, + uint32_t pos) { + if (!g || g->dspark_cache_len == 0 || + !metal_graph_dspark_cache_ends_at(g, pos)) return false; + g->dspark_cache_len += 1u; + if (g->dspark_cache_len > g->dspark_cache_cap) { + const uint32_t excess = g->dspark_cache_len - g->dspark_cache_cap; + g->dspark_cache_token_start += excess; + g->dspark_cache_len = g->dspark_cache_cap; + g->dspark_cache_start = + g->dspark_cache_token_start % g->dspark_cache_cap; + } + return true; +} + +bool ds4_test_dspark_cache_window_crop(void) { + ds4_gpu_graph g; + memset(&g, 0, sizeof(g)); + g.dspark_cache_cap = 8; + + if (!metal_graph_dspark_cache_set_window(&g, 10, 5)) return false; + if (g.dspark_cache_token_start != 10 || + g.dspark_cache_start != 2 || + g.dspark_cache_len != 5) return false; + if (!metal_graph_dspark_cache_ends_at(&g, 15)) return false; + if (metal_graph_dspark_cache_ends_at(&g, 14)) return false; + if (metal_graph_dspark_cache_window_valid(&g, 10, 3, 5)) return false; + + if (!metal_graph_dspark_cache_crop_to_prefix(&g, 13)) return false; + if (g.dspark_cache_token_start != 10 || + g.dspark_cache_start != 2 || + g.dspark_cache_len != 3) return false; + if (!metal_graph_dspark_cache_ends_at(&g, 13)) return false; + if (!metal_graph_dspark_cache_claim_appended_row(&g, 13)) return false; + if (g.dspark_cache_token_start != 10 || + g.dspark_cache_start != 2 || + g.dspark_cache_len != 4) return false; + if (!metal_graph_dspark_cache_ends_at(&g, 14)) return false; + if (metal_graph_dspark_cache_claim_appended_row(&g, 13)) return false; + + if (!metal_graph_dspark_cache_crop_to_prefix(&g, 20)) return false; + if (g.dspark_cache_token_start != 0 || + g.dspark_cache_start != 0 || + g.dspark_cache_len != 0) return false; + if (!metal_graph_dspark_cache_ends_at(&g, 20)) return false; + + if (metal_graph_dspark_cache_set_window(&g, UINT32_MAX - 1u, 2)) { + return false; + } + return true; +} + +static void metal_graph_dspark_capture_begin(ds4_gpu_graph *g) { + metal_graph_dspark_capture_row_invalidate(g); +} + +static void metal_graph_dspark_capture_begin_prefill(ds4_gpu_graph *g) { + metal_graph_dspark_capture_invalidate(g); +} + +static bool metal_graph_dspark_capture_hc( + ds4_gpu_graph *g, + const ds4_gpu_tensor *hc, + uint32_t slot) { + if (!g || !hc || !g->dspark_target_hidden || + !g->dspark_hc_mean_weights || + slot >= g->dspark_target_layer_count) { + return false; + } + + ds4_gpu_tensor *dst = + ds4_gpu_tensor_view(g->dspark_target_hidden, + (uint64_t)slot * DS4_N_EMBD * sizeof(float), + (uint64_t)DS4_N_EMBD * sizeof(float)); + if (!dst) return false; + const bool ok = ds4_gpu_hc_weighted_sum_tensor(dst, + hc, + g->dspark_hc_mean_weights, + DS4_N_EMBD, + DS4_N_HC) != 0; + ds4_gpu_tensor_free(dst); + if (ok) metal_graph_dspark_capture_note_slot(g, slot); + return ok; +} + +static bool metal_graph_dspark_capture_batch_note_slot( + ds4_gpu_graph *g, + uint32_t slot, + uint32_t start, + uint32_t n_tokens) { + if (!g || slot >= g->dspark_target_layer_count || n_tokens == 0) { + return false; + } + if (g->dspark_capture_batch_mask == 0) { + g->dspark_capture_batch_start = start; + g->dspark_capture_batch_tokens = n_tokens; + } else if (g->dspark_capture_batch_start != start || + g->dspark_capture_batch_tokens != n_tokens) { + metal_graph_dspark_capture_batch_invalidate(g); + return false; + } + g->dspark_capture_batch_mask |= 1u << slot; + g->dspark_capture_batch_valid = + g->dspark_capture_batch_mask == + metal_graph_dspark_capture_complete_mask(g); + return true; +} + +static bool metal_graph_dspark_capture_decode_layer( + ds4_gpu_graph *g, + uint32_t il) { + const int slot = metal_graph_dspark_target_slot(g, il); + if (slot < 0) return true; + return metal_graph_dspark_capture_hc(g, metal_graph_cur_hc(g), (uint32_t)slot); +} + +static bool metal_graph_dspark_capture_prefill_layer( + ds4_gpu_graph *g, + uint32_t il, + uint32_t start, + uint32_t n_tokens) { + const int slot = metal_graph_dspark_target_slot(g, il); + if (slot < 0) return true; + if (n_tokens == 0) return false; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + if (g->dspark_target_hidden_batch && + g->dspark_hc_mean_rows && + n_tokens <= g->prefill_cap) { + ds4_gpu_tensor *batch_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden_batch, + ((uint64_t)slot * g->prefill_cap * + DS4_N_EMBD) * sizeof(float), + (uint64_t)n_tokens * embd_bytes); + ds4_gpu_tensor *last_src = + batch_dst ? + ds4_gpu_tensor_view(batch_dst, + (uint64_t)(n_tokens - 1u) * embd_bytes, + embd_bytes) : NULL; + ds4_gpu_tensor *last_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden, + (uint64_t)slot * embd_bytes, + embd_bytes); + bool ok = batch_dst && last_src && last_dst && + ds4_gpu_hc_weighted_sum_tensor(batch_dst, + metal_graph_batch_cur_hc(g), + g->dspark_hc_mean_rows, + DS4_N_EMBD, + DS4_N_HC) != 0 && + ds4_gpu_tensor_copy(last_dst, + 0, + last_src, + 0, + embd_bytes) != 0; + ds4_gpu_tensor_free(last_dst); + ds4_gpu_tensor_free(last_src); + ds4_gpu_tensor_free(batch_dst); + if (ok) { + metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); + ok = metal_graph_dspark_capture_batch_note_slot(g, + (uint32_t)slot, + start, + n_tokens); + } + return ok; + } + + ds4_gpu_tensor *last_hc = + ds4_gpu_tensor_view(metal_graph_batch_cur_hc(g), + (uint64_t)(n_tokens - 1u) * hc_dim * sizeof(float), + hc_dim * sizeof(float)); + if (!last_hc) return false; + const bool ok = metal_graph_dspark_capture_hc(g, last_hc, (uint32_t)slot); + ds4_gpu_tensor_free(last_hc); + return ok; +} + +static bool metal_graph_dspark_capture_prefill_rows( + ds4_gpu_graph *g, + uint32_t il, + uint32_t chunk_start, + uint32_t chunk_len, + uint32_t pos0, + uint32_t n_tokens) { + const int slot = metal_graph_dspark_target_slot(g, il); + if (slot < 0) return true; + if (!g->dspark_target_hidden_batch || + !g->dspark_target_hidden || + !g->dspark_hc_mean_rows || + n_tokens == 0 || + chunk_len == 0 || + chunk_len > g->prefill_cap || + pos0 < chunk_start) { + return true; + } + const uint32_t row0 = pos0 - chunk_start; + if (row0 > chunk_len || n_tokens > chunk_len - row0) return true; + + const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + ds4_gpu_tensor *batch_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden_batch, + (((uint64_t)(uint32_t)slot * g->prefill_cap + + row0) * DS4_N_EMBD) * sizeof(float), + (uint64_t)n_tokens * embd_bytes); + bool ok = batch_dst && + ds4_gpu_hc_weighted_sum_tensor(batch_dst, + metal_graph_batch_cur_hc(g), + g->dspark_hc_mean_rows, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (!ok) fprintf(stderr, "ds4: pipeline capture rows FAIL il=%u row0=%u n=%u dst=%d\n", + il, row0, n_tokens, batch_dst != NULL); + if (ok && row0 + n_tokens == chunk_len) { + ds4_gpu_tensor *last_src = + ds4_gpu_tensor_view(batch_dst, + (uint64_t)(n_tokens - 1u) * embd_bytes, + embd_bytes); + ds4_gpu_tensor *last_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden, + (uint64_t)(uint32_t)slot * embd_bytes, + embd_bytes); + ok = last_src && last_dst && + ds4_gpu_tensor_copy(last_dst, 0, last_src, 0, embd_bytes) != 0; + ds4_gpu_tensor_free(last_dst); + ds4_gpu_tensor_free(last_src); + if (ok) { + metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); + ok = metal_graph_dspark_capture_batch_note_slot(g, + (uint32_t)slot, + chunk_start, + chunk_len); + } + } + ds4_gpu_tensor_free(batch_dst); + return ok; +} + +static bool metal_graph_dspark_capture_verified_suffix_begin( + ds4_gpu_graph *g, + uint32_t start, + uint32_t n_tokens, + bool commands_open) { + if (!g || !g->dspark_capture_enabled || + !g->dspark_target_hidden || + !g->dspark_target_hidden_batch || + start == 0 || + n_tokens == 0 || + n_tokens + 1u < n_tokens || + n_tokens + 1u > g->prefill_cap || + !g->dspark_capture_valid || + g->dspark_capture_checkpoint_len != start) { + metal_graph_dspark_capture_invalidate(g); + return false; + } + + const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + metal_graph_dspark_capture_batch_invalidate(g); + bool ok = commands_open || ds4_gpu_begin_commands() != 0; + for (uint32_t slot = 0; ok && slot < g->dspark_target_layer_count; slot++) { + ds4_gpu_tensor *dst = + ds4_gpu_tensor_view(g->dspark_target_hidden_batch, + ((uint64_t)slot * g->prefill_cap * + DS4_N_EMBD) * sizeof(float), + embd_bytes); + ds4_gpu_tensor *src = + ds4_gpu_tensor_view(g->dspark_target_hidden, + (uint64_t)slot * embd_bytes, + embd_bytes); + ok = dst && src && + ds4_gpu_tensor_copy(dst, 0, src, 0, embd_bytes) != 0; + ds4_gpu_tensor_free(src); + ds4_gpu_tensor_free(dst); + } + if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; + else if (!ok && !commands_open) (void)ds4_gpu_synchronize(); + if (!ok) { + metal_graph_dspark_capture_invalidate(g); + return false; + } + + metal_graph_dspark_capture_row_invalidate(g); + return true; +} + +static bool metal_graph_dspark_capture_verified_suffix_layer( + ds4_gpu_graph *g, + uint32_t il, + uint32_t start, + uint32_t n_tokens) { + const int slot = metal_graph_dspark_target_slot(g, il); + if (slot < 0) return true; + if (!g || !g->dspark_target_hidden_batch || + !g->dspark_target_hidden || + !g->dspark_hc_mean_rows || + start == 0 || + n_tokens == 0 || + n_tokens + 1u < n_tokens || + n_tokens + 1u > g->prefill_cap) { + return false; + } + + const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + ds4_gpu_tensor *batch_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden_batch, + (((uint64_t)(uint32_t)slot * g->prefill_cap + + 1u) * DS4_N_EMBD) * sizeof(float), + (uint64_t)n_tokens * embd_bytes); + ds4_gpu_tensor *last_src = + batch_dst ? + ds4_gpu_tensor_view(batch_dst, + (uint64_t)(n_tokens - 1u) * embd_bytes, + embd_bytes) : NULL; + ds4_gpu_tensor *last_dst = + ds4_gpu_tensor_view(g->dspark_target_hidden, + (uint64_t)(uint32_t)slot * embd_bytes, + embd_bytes); + bool ok = batch_dst && last_src && last_dst && + ds4_gpu_hc_weighted_sum_tensor(batch_dst, + metal_graph_batch_cur_hc(g), + g->dspark_hc_mean_rows, + DS4_N_EMBD, + DS4_N_HC) != 0 && + ds4_gpu_tensor_copy(last_dst, + 0, + last_src, + 0, + embd_bytes) != 0; + ds4_gpu_tensor_free(last_dst); + ds4_gpu_tensor_free(last_src); + ds4_gpu_tensor_free(batch_dst); + if (ok) { + metal_graph_dspark_capture_note_slot(g, (uint32_t)slot); + ok = metal_graph_dspark_capture_batch_note_slot(g, + (uint32_t)slot, + start - 1u, + n_tokens + 1u); + } + return ok; +} + +/* Encode a full single-token decode step on Metal. This is the generation + * hot path: update caches, run all layers, then produce logits. */ +static bool metal_graph_encode_token_raw_swa( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + uint32_t pos, + bool need_logits, + bool allow_split_flush) { + if (g->raw_cap == 0) { + fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); + return false; + } + /* Under the vocab split both ranks materialize their logits half. */ + if (g->tp_world == 2 && g->tp_rank == 1 && + !g->tp_logits_half) need_logits = false; + const uint32_t raw_row = pos % g->raw_cap; + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); + metal_graph_dspark_capture_begin(g); + + /* write the embedded token on the embedding tier. Single- + * tier: emb_tier == 0 == active_tier; no-op. Multi-tier: switch to + * emb_tier (no cross-device copy needed — embed writes from scratch). */ + if (g->placement) { + if (!metal_graph_set_active_tier_decode(g, g->emb_tier)) return false; + } + bool ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + + /* + * Start executing the prefix of the decode graph while the CPU is still + * encoding the rest. The split point is layer-based because this executor is + * a fixed DS4 tape, not a dynamic node graph; four layers is the measured + * point where the prefix is large enough to hide useful work without + * starving the second command buffer. + */ + const uint32_t split_after_layers = metal_graph_token_split_after_layers(); + + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + ok = metal_graph_encode_decode_layer(g, + model, + &weights->layer[il], + il, + pos, + g->layer_raw_cache[il], + g->raw_cap, + raw_row, + n_raw, + token); + ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); + g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); + g->after_ffn_hc_by_tier[g->active_tier] = tmp; + if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); + /* A TP gate uses one monotonic shared event for the whole token. A + * later command buffer may signal a higher value while the prefix is + * blocked at an earlier gate, making the transport consume a slab + * slot before its payload is ready. Keep each TP token in one command + * buffer; non-TP decode retains the encode/execute overlap. */ + if (ok && allow_split_flush && g->tp_world != 2 && + split_after_layers != 0 && il + 1u == split_after_layers) { + ok = ds4_gpu_flush_commands() != 0; + } + } + + if (ok && need_logits) { + ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + } + return ok; +} + +static ds4_gpu_tensor *metal_graph_tensor_row_view( + ds4_gpu_tensor *base, + uint32_t row, + uint64_t row_values) { + return ds4_gpu_tensor_view(base, + (uint64_t)row * row_values * sizeof(float), + row_values * sizeof(float)); +} + +/* Upload prompt token ids for kernels that need token-aware hash routing. */ +static bool metal_graph_upload_prompt_tokens( + ds4_gpu_tensor *out_tokens, + const token_vec *prompt, + uint32_t pos0, + uint32_t n_tokens) { + if (!out_tokens || pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) { + return false; + } + + int32_t *tokens = xmalloc((size_t)n_tokens * sizeof(tokens[0])); + for (uint32_t i = 0; i < n_tokens; i++) tokens[i] = prompt->v[pos0 + i]; + + const bool ok = ds4_gpu_tensor_write(out_tokens, + 0, + tokens, + (uint64_t)n_tokens * sizeof(tokens[0])) != 0; + free(tokens); + return ok; +} + +/* Rebuild ratio-4 compressor state after chunked prefill so a following decode + * token sees the same rolling compression window. */ +static bool metal_graph_refresh_ratio4_compressor_state( + ds4_gpu_graph *g, + const ds4_model *model, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_tensor *kv_weight, + const ds4_tensor *score_weight, + const ds4_tensor *ape, + uint32_t head_dim, + uint32_t width, + uint32_t pos0, + uint32_t n_tokens) { + if (n_tokens < 4) { + return true; + } + if (!g || !model || !state_kv || !state_score || !kv_weight || !score_weight || !ape || + head_dim == 0 || width == 0) { + return false; + } + + /* + * The recurrent ratio-4 state is intentionally rebuilt from the last + * four tokens using the small-batch projection kernel. The full-chunk + * projection is already available, but it uses the matrix-matrix path; + * mixing those two accumulation orders changes a few FP8 rounding + * decisions in later chunks. + */ + ds4_gpu_tensor *tail_hc = ds4_gpu_tensor_view( + metal_graph_batch_attn_norm(g), + (uint64_t)(n_tokens - 4u) * DS4_N_EMBD * sizeof(float), + 4ull * DS4_N_EMBD * sizeof(float)); + bool ok = tail_hc != NULL; + if (!ok) { + fprintf(stderr, "ds4: ratio-4 compressor tail view creation failed\n"); + } + if (ok) { +#if defined(__APPLE__) + ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), + model->map, + model->size, + kv_weight->abs_offset, + DS4_N_EMBD, + width, + tail_hc, + 4) != 0; + if (ok) { + ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), + model->map, + model->size, + score_weight->abs_offset, + DS4_N_EMBD, + width, + tail_hc, + 4) != 0; + } +#else + ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + kv_weight->abs_offset, + score_weight->abs_offset, + DS4_N_EMBD, + width, + tail_hc, + 4) != 0; +#endif + if (!ok) { + fprintf(stderr, "ds4: ratio-4 compressor tail projection failed\n"); + } + } + if (ok) { + ok = ds4_gpu_compressor_prefill_state_ratio4_tensor(state_kv, + state_score, + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + ape->abs_offset, + ape->type, + head_dim, + pos0 + n_tokens - 4u) != 0; + if (!ok) { + fprintf(stderr, "ds4: ratio-4 compressor state refresh failed\n"); + } + } + ds4_gpu_tensor_free(tail_hc); + return ok; +} + +/* CPU fallback for seeding batched HC state from token embeddings. It is still + * useful for tiny speculative verifier batches where a separate GPU embedding + * command buffer costs more than the small host write. */ +static bool metal_graph_upload_prompt_embeddings_hc_cpu( + ds4_gpu_tensor *out_hc, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t pos0, + uint32_t n_tokens) { + if (pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) return false; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t total = (uint64_t)n_tokens * hc_dim; + float *hc = xmalloc((size_t)total * sizeof(hc[0])); + float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); + + for (uint32_t t = 0; t < n_tokens; t++) { + embed_token_f16(model, weights, prompt->v[pos0 + t], plain); + float *dst = hc + (uint64_t)t * hc_dim; + for (uint32_t h = 0; h < DS4_N_HC; h++) { + memcpy(dst + (uint64_t)h * DS4_N_EMBD, + plain, + (size_t)DS4_N_EMBD * sizeof(plain[0])); + } + } + + const bool ok = ds4_gpu_tensor_write(out_hc, 0, hc, total * sizeof(hc[0])) != 0; + free(plain); + free(hc); + return ok; +} + +/* Seed the batched HC state from token ids: every HC stream starts as the same + * 4096-wide embedding. Long prefill chunks use the Metal get-rows/repeat + * kernel so the CPU does not build and upload a large [token, HC, dim] tensor. */ +static bool metal_graph_upload_prompt_embeddings_hc( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *tokens, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t pos0, + uint32_t n_tokens) { + if (pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) return false; + + uint32_t gpu_min = 512; +#ifndef DS4_ROCM_BUILD + const char *gpu_min_env = getenv("DS4_METAL_GPU_BATCH_EMBED_MIN"); + if (gpu_min_env && gpu_min_env[0]) { + char *end = NULL; + unsigned long v = strtoul(gpu_min_env, &end, 10); + if (end != gpu_min_env && v <= UINT32_MAX) gpu_min = (uint32_t)v; + } +#endif + + if (tokens && n_tokens >= gpu_min) { + return ds4_gpu_embed_tokens_hc_tensor(out_hc, + tokens, + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + n_tokens, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + + return metal_graph_upload_prompt_embeddings_hc_cpu(out_hc, + model, + weights, + prompt, + pos0, + n_tokens); +} + +static bool metal_graph_hc_rms_scale_project( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_scratch, + const ds4_model *model, + const ds4_tensor *weight, + const ds4_gpu_tensor *x, + uint64_t in_dim, + uint32_t n_tokens) { + if (!out || !norm_scratch || !model || !weight || !x || + in_dim > UINT32_MAX) { + return false; + } +#if defined(__APPLE__) + return ds4_gpu_hc_rms_scale_project_f16_tensor( + out, + norm_scratch, + model->map, + model->size, + weight->abs_offset, + (uint32_t)in_dim, + 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, + x, + n_tokens, + DS4_RMS_EPS) != 0; +#else + bool ok = ds4_gpu_rms_norm_plain_rows_tensor( + norm_scratch, + x, + (uint32_t)in_dim, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) { + ok = ds4_gpu_matmul_f16_tensor( + out, + model->map, + model->size, + weight->abs_offset, + in_dim, + 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, + norm_scratch, + n_tokens) != 0; + } + return ok; +#endif +} + +static bool metal_graph_warmup_prefill_kernels( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t n_tokens) { + static bool warmed = false; + if (g && g->ssd_streaming) return true; + if (warmed) return true; +#ifndef DS4_ROCM_BUILD + if (getenv("DS4_METAL_NO_PREFILL_KERNEL_WARMUP") != NULL) return true; +#endif + + /* + * The first batched F16 matmul can pay Metal's one-time pipeline execution + * cost. Run the same HC attention projection on scratch storage before the + * measured prefill. The output is overwritten by the real graph. + */ + if (n_tokens <= 8) return true; + + /* (B6 fix, ): warm-up uses layer-0's hc_attn_fn + * weight, which in multi-tier is resolved on placement[1]'s tier. + * Switch active_tier so the F16 matmul reads/writes the correct + * Class P scratch and resolves the weight on the right device. + * Single-tier (g->placement == NULL): no-op. */ + if (g->placement) { + if (!metal_graph_set_active_tier_batch(g, g->placement[1], n_tokens)) return false; + } + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = metal_graph_hc_rms_scale_project( + metal_graph_batch_hc_mix(g), + metal_graph_batch_flat_hc(g), + model, + weights->layer[0].hc_attn_fn, + metal_graph_batch_cur_hc(g), + hc_dim, + n_tokens); + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) { + fprintf(stderr, "ds4: Metal prefill kernel warmup failed\n"); + return false; + } + + warmed = true; + return true; +} + +/* Encode the batched prefill attention half for one layer. It mirrors the CPU + * layer-major path: HC pre/norm, Q/KV, cache/compression, prefix attention. */ +static bool metal_graph_indexer_stage_profile_boundary( + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_comp, + double *stage_t0) { + if (ds4_gpu_end_commands() == 0) return false; + const double now = now_sec(); + if (stage != NULL) { + fprintf(stderr, + "ds4: metal indexer stage layer=%u pos=%u tokens=%u comp=%u %s=%.3f ms\n", + il, + pos0, + n_tokens, + n_comp, + stage, + (now - *stage_t0) * 1000.0); + } + *stage_t0 = now; + return ds4_gpu_begin_commands() != 0; +} + +static bool metal_graph_env_value_eq(const char *v, + size_t n, + const char *literal) { + const size_t m = strlen(literal); + if (n != m) return false; + for (size_t i = 0; i < n; i++) { + if (tolower((unsigned char)v[i]) != + tolower((unsigned char)literal[i])) { + return false; + } + } + return true; +} + +static const char *metal_graph_env_trim(const char *v, size_t *len_out) { + if (!v) { + if (len_out) *len_out = 0; + return NULL; + } + while (isspace((unsigned char)*v)) v++; + size_t n = strlen(v); + while (n > 0 && isspace((unsigned char)v[n - 1])) n--; + if (len_out) *len_out = n; + return v; +} + +static bool metal_graph_profile_layer_value_match(const char *layer_env, + uint32_t il) { + size_t n = 0; + layer_env = metal_graph_env_trim(layer_env, &n); + if (!layer_env || n == 0) return true; + + char *end = NULL; + const unsigned long layer = strtoul(layer_env, &end, 10); + return end != layer_env && + (size_t)(end - layer_env) == n && + layer <= UINT32_MAX && + (uint32_t)layer == il; +} + +static bool metal_graph_stage_profile_enabled_for_layer( + const char *flag_env_name, + const char *layer_env_name, + uint32_t il) { + size_t flag_len = 0; + const char *flag = metal_graph_env_trim(getenv(flag_env_name), &flag_len); + if (!flag) return false; + + const char *layer_env = getenv(layer_env_name); + const bool has_layer_filter = layer_env && layer_env[0]; + + if (flag_len != 0) { + if (metal_graph_env_value_eq(flag, flag_len, "0") || + metal_graph_env_value_eq(flag, flag_len, "false") || + metal_graph_env_value_eq(flag, flag_len, "no") || + metal_graph_env_value_eq(flag, flag_len, "off")) { + return false; + } + if (!has_layer_filter && + !metal_graph_env_value_eq(flag, flag_len, "1") && + !metal_graph_env_value_eq(flag, flag_len, "true") && + !metal_graph_env_value_eq(flag, flag_len, "yes") && + !metal_graph_env_value_eq(flag, flag_len, "on") && + !metal_graph_env_value_eq(flag, flag_len, "all")) { + return metal_graph_profile_layer_value_match(flag, il); + } + } + + return metal_graph_profile_layer_value_match(layer_env, il); +} + +static bool metal_graph_layer_stage_profile_enabled(uint32_t il) { + return metal_graph_stage_profile_enabled_for_layer( + "DS4_ROCM_LAYER_STAGE_PROFILE", + "DS4_ROCM_LAYER_STAGE_PROFILE_LAYER", + il) || + metal_graph_stage_profile_enabled_for_layer( + "DS4_METAL_LAYER_STAGE_PROFILE", + "DS4_METAL_LAYER_STAGE_PROFILE_LAYER", + il); +} + +static bool metal_graph_decode_stage_profile_enabled(uint32_t il) { + return metal_graph_stage_profile_enabled_for_layer( + "DS4_ROCM_DECODE_STAGE_PROFILE", + "DS4_ROCM_DECODE_STAGE_PROFILE_LAYER", + il) || + metal_graph_stage_profile_enabled_for_layer( + "DS4_METAL_DECODE_STAGE_PROFILE", + "DS4_METAL_DECODE_STAGE_PROFILE_LAYER", + il); +} + +static bool metal_graph_layer_stage_profile_start(uint32_t il) { + if (!metal_graph_layer_stage_profile_enabled(il)) return true; + if (ds4_gpu_end_commands() == 0) return false; + return ds4_gpu_begin_commands() != 0; +} + +/* Optional prefill stage profiler. It intentionally ends the current Metal + * command buffer and waits, so the printed number includes encoding plus GPU + * execution for the stage just emitted. This is disabled by default because it + * adds synchronization points and changes scheduling. */ +static bool metal_graph_layer_stage_profile_boundary( + const char *part, + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + double *stage_t0) { + if (ds4_gpu_end_commands() == 0) return false; + const double now = now_sec(); + if (stage != NULL) { + fprintf(stderr, + "ds4: metal layer stage part=%s layer=%u pos=%u tokens=%u %s=%.3f ms\n", + part, + il, + pos0, + n_tokens, + stage, + (now - *stage_t0) * 1000.0); + } + *stage_t0 = now; + return ds4_gpu_begin_commands() != 0; +} + +static bool metal_graph_q_stage_profile_boundary( + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + double *stage_t0) { + if (ds4_gpu_end_commands() == 0) return false; + const double now = now_sec(); + fprintf(stderr, + "ds4: metal Q path stage layer=%u pos=%u tokens=%u %s=%.3f ms\n", + il, + pos0, + n_tokens, + stage, + (now - *stage_t0) * 1000.0); + *stage_t0 = now; + return ds4_gpu_begin_commands() != 0; +} + +static ds4_gpu_tensor *metal_graph_tensor_row_range_view( + ds4_gpu_tensor *base, + uint32_t row0, + uint32_t rows, + uint64_t row_values) { + return ds4_gpu_tensor_view(base, + (uint64_t)row0 * row_values * sizeof(float), + (uint64_t)rows * row_values * sizeof(float)); +} + +/* TP prefill threshold for row-splitting the replicated shared expert. + * Routed experts remain ownership-split at every batch size. */ +static uint32_t metal_graph_tp_prefill_split_min(void) { + static int cached = -1; + if (cached < 0) { + cached = 32; + const char *env = getenv("DS4_TP_PREFILL_SPLIT_MIN"); + if (env && env[0]) cached = atoi(env); + if (cached < 2) cached = 2; + } + return (uint32_t)cached; +} + +/* Opt-in sub-chunk gate pipelining for the TP prefill row swaps. Must be + * set on BOTH ranks (it changes the per-layer gate count; asymmetric + * settings deadlock the big gates). Default off: measured net-negative + * on the M5 Max pair, see the pipelined blocks for the numbers. */ +static bool metal_graph_tp_subgate_pipeline(void) { + static int cached = -1; + if (cached < 0) { + const char *env = getenv("DS4_TP_SUBGATE_PIPELINE"); + cached = env && env[0] && atoi(env) != 0; + } + return cached != 0; +} + +static bool metal_graph_encode_layer_attention_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens) { + if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint32_t n_groups = DS4_N_OUT_GROUP; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = DS4_N_LORA_O; + const uint32_t ratio = ds4_layer_compress_ratio(il); + const bool compressed = ratio != 0; + const bool zero_prefix = pos0 == 0; + /* TP attention row split for large zero-prefix chunks: q_a and the KV + * path stay full (both ranks need every row's KV, and the compressor/ + * indexer keep updating their state from full rows), q_b onward runs on + * this rank's half of the chunk rows, and the computed row halves of + * batch_attn_out are swapped in place through one big gate per layer. + * Three chunk shapes split: full-raw (uncompressed layer at pos0 == 0, + * or a compressed layer whose chunk is too short to emit compressed + * keys, i.e. raw_prefix_tokens == n_tokens), static-mixed (compressed + * layer whose whole chunk attends through the one-shot mixed kernel + * over the full raw keys plus n_tokens / ratio compressed keys, without + * indexer top-k), and indexed (ratio-4 layer with indexer top-k, whose + * per-token score/top-k selection stays replicated while the attention + * consumption splits by rows). Every condition derives from + * pos0/n_tokens/ratio/model shape so both ranks stay in lockstep. */ + const bool tp_attn_full_raw = zero_prefix && + (ratio == 0 || (n_tokens < ratio && n_tokens <= g->raw_cap)); + const uint32_t tp_attn_n_comp = ratio != 0 ? n_tokens / ratio : 0; + const bool tp_attn_static_mixed = zero_prefix && ratio != 0 && + tp_attn_n_comp != 0 && + !(ratio == 4 && tp_attn_n_comp > DS4_N_INDEXER_TOP_K); + const bool tp_attn_indexed = zero_prefix && ratio == 4 && + tp_attn_n_comp > DS4_N_INDEXER_TOP_K; + const bool tp_row_split_attn = + g->tp_world == 2 && + g->tp_batch_rows != n_tokens && + (tp_attn_full_raw || tp_attn_static_mixed || tp_attn_indexed) && + !metal_graph_directional_steering_attn_enabled(g) && + n_tokens >= metal_graph_tp_prefill_split_min(); + const uint32_t tp_half_rows = (n_tokens + 1u) / 2u; + const uint32_t tp_row0 = (tp_row_split_attn && g->tp_rank != 0) ? tp_half_rows : 0; + const uint32_t tp_rows = tp_row_split_attn ? + (g->tp_rank == 0 ? tp_half_rows : n_tokens - tp_half_rows) : n_tokens; + const bool index_stage_profile = + glm_graph_env_present("DS4_ROCM_INDEXER_STAGE_PROFILE", + "DS4_METAL_INDEXER_STAGE_PROFILE"); + const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); + const bool q_stage_profile = + glm_graph_env_present("DS4_ROCM_Q_STAGE_PROFILE", + "DS4_METAL_Q_STAGE_PROFILE"); + double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; + double q_stage_t0 = q_stage_profile ? now_sec() : 0.0; +#define DS4_METAL_PROFILE_ATTN_STAGE(name) do { \ + if (ok && layer_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary("attn", (name), il, pos0, n_tokens, &layer_stage_t0); \ + } \ + } while (0) +#define DS4_METAL_PROFILE_Q_STAGE(name) do { \ + if (ok && q_stage_profile) { \ + ok = metal_graph_q_stage_profile_boundary((name), il, pos0, n_tokens, &q_stage_t0); \ + } \ + } while (0) + const float freq_base = layer_rope_freq_base(il); + const float freq_scale = layer_rope_freq_scale(il); + const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; + float attn_factor = 1.0f; + if (ext_factor != 0.0f && freq_scale > 0.0f) { + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + enum { stack_count_cap = 16 }; + uint32_t comp_counts_stack[stack_count_cap]; + uint32_t index_counts_stack[stack_count_cap]; + uint32_t *comp_counts = NULL; + uint32_t *index_counts = NULL; + if (compressed) { + if (n_tokens <= stack_count_cap) { + memset(comp_counts_stack, 0, + (size_t)n_tokens * sizeof(comp_counts_stack[0])); + comp_counts = comp_counts_stack; + } else { + comp_counts = xcalloc(n_tokens, sizeof(comp_counts[0])); + } + } + if (ratio == 4) { + if (n_tokens <= stack_count_cap) { + memset(index_counts_stack, 0, + (size_t)n_tokens * sizeof(index_counts_stack[0])); + index_counts = index_counts_stack; + } else { + index_counts = xcalloc(n_tokens, sizeof(index_counts[0])); + } + } + const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); + ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_mix(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_split(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *attn_cur_view = ds4_gpu_tensor_view( + metal_graph_batch_attn_cur(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( + metal_graph_batch_after_attn_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); + bool ok = hc_mix_view && hc_split_view && attn_cur_view && after_attn_hc_view; + const bool fuse_hc_norm = n_tokens > 1 && + DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + metal_graph_enable_batch_hc_norm_fusion(); + if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, + metal_graph_batch_flat_hc(g), + model, + layer->hc_attn_fn, + metal_graph_batch_cur_hc(g), + hc_dim, + n_tokens); + if (metal_graph_use_reference_hc_decode()) { + if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, + hc_mix_view, + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(attn_cur_view, + metal_graph_batch_cur_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (fuse_hc_norm) { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, + metal_graph_batch_attn_norm(g), + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("hc_attn_pre", metal_graph_batch_attn_cur(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + DS4_METAL_PROFILE_ATTN_STAGE("hc_pre"); + if (ok && !fuse_hc_norm) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), + metal_graph_batch_attn_cur(g), + model->map, + model->size, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("attn_norm", metal_graph_batch_attn_norm(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + DS4_METAL_PROFILE_ATTN_STAGE("norm"); + DS4_METAL_PROFILE_Q_STAGE("pre_q"); + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_a", + il, + pos0, + metal_graph_batch_qr(g), + model, + layer->attn_q_a, + DS4_N_EMBD, + q_rank, + metal_graph_batch_attn_norm(g), + n_tokens); + if (ok) { + metal_graph_debug_dump_tensor("q_lora", metal_graph_batch_qr(g), + (uint64_t)n_tokens * q_rank, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("q_a"); + if (qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", + il, + pos0, + metal_graph_batch_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + n_tokens); + if (ok) { + metal_graph_debug_dump_tensor("KVraw", metal_graph_batch_kv_raw(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(metal_graph_batch_qr_norm(g), + metal_graph_batch_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + n_tokens, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_qr_norm(g), + metal_graph_batch_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("q_lora_norm", metal_graph_batch_qr_norm(g), + (uint64_t)n_tokens * q_rank, il, pos0); + } + if (qkv_rms_fused && ok) { + metal_graph_debug_dump_tensor("KVnorm", metal_graph_batch_kv(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("q_a_norm"); + const bool q_path_debug = + metal_graph_debug_wants("Qraw", il, pos0) || + metal_graph_debug_wants("Qnorm", il, pos0); + /* Under the TP row split everything from q_b to the output projection + * runs on this rank's rows only, through row-range views of the batch + * tensors (batch_q_half is F16, so its view is built directly). */ + ds4_gpu_tensor *tp_q = tp_row_split_attn ? + metal_graph_tensor_row_range_view(metal_graph_batch_q(g), tp_row0, tp_rows, q_dim) : NULL; + ds4_gpu_tensor *tp_q_half = tp_row_split_attn ? + ds4_gpu_tensor_view(g->batch_q_half, + (uint64_t)tp_row0 * q_dim * sizeof(uint16_t), + (uint64_t)tp_rows * q_dim * sizeof(uint16_t)) : NULL; + ds4_gpu_tensor *tp_qr_norm = tp_row_split_attn ? + metal_graph_tensor_row_range_view(metal_graph_batch_qr_norm(g), tp_row0, tp_rows, q_rank) : NULL; + ds4_gpu_tensor *tp_heads = tp_row_split_attn ? + metal_graph_tensor_row_range_view(metal_graph_batch_heads(g), tp_row0, tp_rows, q_dim) : NULL; + ds4_gpu_tensor *tp_attn_out = tp_row_split_attn ? + metal_graph_tensor_row_range_view(metal_graph_batch_attn_out(g), tp_row0, tp_rows, + DS4_N_EMBD) : NULL; + if (tp_row_split_attn && + (!tp_q || !tp_q_half || !tp_qr_norm || !tp_heads || !tp_attn_out)) { + ok = false; + } + bool q_b_f16_out = false; + if (ok && !q_path_debug && layer->attn_q_b->type == DS4_TENSOR_Q8_0) { + q_b_f16_out = ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(tp_q ? tp_q : metal_graph_batch_q(g), + tp_q_half ? tp_q_half : g->batch_q_half, + model->map, + model->size, + layer->attn_q_b->abs_offset, + q_rank, + q_dim, + tp_qr_norm ? tp_qr_norm : metal_graph_batch_qr_norm(g), + tp_rows, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0 + tp_row0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } + if (q_b_f16_out) { + DS4_METAL_PROFILE_Q_STAGE("q_b"); + DS4_METAL_PROFILE_Q_STAGE("head_norm"); + if (ok) { + metal_graph_debug_dump_tensor("Qcur", metal_graph_batch_q(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("rope"); + } else { + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_b", + il, + pos0, + tp_q ? tp_q : metal_graph_batch_q(g), + model, + layer->attn_q_b, + q_rank, + q_dim, + tp_qr_norm ? tp_qr_norm : metal_graph_batch_qr_norm(g), + tp_rows); + if (ok) { + metal_graph_debug_dump_tensor("Qraw", metal_graph_batch_q(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("q_b"); + if (ok) ok = ds4_gpu_head_rms_norm_tensor(tp_q ? tp_q : metal_graph_batch_q(g), + tp_rows, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_RMS_EPS) != 0; + if (ok) { + metal_graph_debug_dump_tensor("Qnorm", metal_graph_batch_q(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("head_norm"); + if (ok) ok = ds4_gpu_rope_tail_tensor(tp_q ? tp_q : metal_graph_batch_q(g), + tp_rows, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0 + tp_row0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) { + metal_graph_debug_dump_tensor("Qcur", metal_graph_batch_q(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + DS4_METAL_PROFILE_Q_STAGE("rope"); + } + DS4_METAL_PROFILE_ATTN_STAGE("q_path"); + if (!qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", + il, + pos0, + metal_graph_batch_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + n_tokens); + if (ok) { + metal_graph_debug_dump_tensor("KVraw", metal_graph_batch_kv_raw(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + model->map, + model->size, + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVnorm", metal_graph_batch_kv(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + } + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), + n_tokens, + DS4_N_HEAD_KV, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVrope", metal_graph_batch_kv(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), + n_tokens, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVcur", metal_graph_batch_kv(g), + (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); + } + DS4_METAL_PROFILE_ATTN_STAGE("kv_path"); + /* + * Static graph order is q, kv, cpy_k(raw SWA), then attention. For a + * zero-prefix batch it is safe to store the whole batch at once: attention + * reads the contiguous batch KV, and the ring only has to end with the last + * SWA rows for later chunks/decode. For nonzero chunks the physical ring is + * sized to hold the current chunk plus the previous SWA window, while the + * attention mask still enforces the 128-token logical window. + */ + if (ok && zero_prefix) ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + metal_graph_batch_kv(g), + g->raw_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u raw KV batch store failed\n", il); + } + const bool raw_batch_attention = zero_prefix && ratio == 0; + bool batch_attention_done = false; + + if (ok && raw_batch_attention) { + if (tp_row_split_attn) { + ok = ds4_gpu_attention_prefill_raw_heads_range_tensor(tp_heads, + model->map, + model->size, + layer->attn_sinks->abs_offset, + tp_q, + metal_graph_batch_kv(g), + tp_row0, + tp_rows, + n_tokens, + g->raw_window, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } else { + ok = ds4_gpu_attention_prefill_raw_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + metal_graph_batch_kv(g), + n_tokens, + g->raw_window, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + if (ok) batch_attention_done = true; + } else if (ok && !zero_prefix && ratio == 0 && n_tokens <= g->raw_cap) { + /* + * The ubatch path stores the whole batch in the SWA cache, then runs + * one batched attention kernel with an absolute-position causal/window + * mask. This avoids mixing prefill with the different single-token + * attention path. + */ + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos0, n_tokens); + /* Nonzero prompt chunks read the SWA cache as a ring. FlashAttention + * receives a linearized window starting at raw_start, not physical row + * zero; otherwise wrapped chunks silently miss recent raw keys. */ + const uint32_t raw_start = metal_graph_raw_start_for_span(g, + pos0 + n_tokens - 1u, + n_raw); + ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + metal_graph_batch_kv(g), + g->raw_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + if (ok) { + metal_graph_debug_dump_tensor("raw_cache", + g->layer_raw_cache[il], + (uint64_t)n_raw * DS4_N_HEAD_DIM, + il, + pos0); + } + if (ok) { + ok = ds4_gpu_attention_decode_raw_batch_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->layer_raw_cache[il], + n_tokens, + pos0, + n_raw, + g->raw_cap, + raw_start, + g->raw_window, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + if (ok) batch_attention_done = true; + } else if (ok && ratio != 0) { + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_width = coff * DS4_N_HEAD_DIM; + const bool have_attn_comp = layer->attn_compressor_kv && layer->attn_compressor_gate && + layer->attn_compressor_ape && layer->attn_compressor_norm; + if (!have_attn_comp) { + fprintf(stderr, "ds4: Metal layer-major prefill needs attention compressor weights\n"); + ok = false; + } + if (ok) { + ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + DS4_N_EMBD, + comp_width, + metal_graph_batch_attn_norm(g), + n_tokens) != 0; + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u attention compressor KV projection failed\n", il); + } + if (ok) { + ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, + comp_width, + metal_graph_batch_attn_norm(g), + n_tokens) != 0; + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u attention compressor score projection failed\n", il); + } + } + } + if (ok) metal_graph_debug_dump_tensor("attn_comp_kv_raw", + metal_graph_batch_comp_kv(g), + (uint64_t)comp_width * n_tokens, + il, + pos0); + if (ok) metal_graph_debug_dump_tensor("attn_comp_score_raw", + metal_graph_batch_comp_sc(g), + (uint64_t)comp_width * n_tokens, + il, + pos0); + uint32_t n_comp = g->layer_n_comp[il]; + if (zero_prefix) { + n_comp = n_tokens / ratio; + if (ok && n_comp > g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal layer-major compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + if (ok && DS4_GPU_ATTN_COMP_CACHE_F16 && n_comp > g->attn_comp_stage_cap) { + fprintf(stderr, "ds4: Metal graph compressed KV staging capacity exceeded at layer %u\n", il); + ok = false; + } + ds4_gpu_tensor *attn_comp_target = NULL; + if (ok) { + attn_comp_target = metal_graph_attn_comp_prefill_target(g, il, 0, n_comp); + if (!attn_comp_target) { + fprintf(stderr, "ds4: gpu layer %u attention compressor target creation failed\n", il); + ok = false; + } + if (ok) ok = ds4_gpu_compressor_prefill_tensor(attn_comp_target, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos0, + n_tokens, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u attention compressor prefill failed\n", il); + } + DS4_METAL_PROFILE_ATTN_STAGE("compressor_prefill"); + if (ok && n_comp != 0) { + ok = metal_graph_commit_attn_comp_stage(g, il, 0, n_comp); + } + DS4_METAL_PROFILE_ATTN_STAGE("compressor_commit"); + if (ok && ratio == 4) { + ok = metal_graph_refresh_ratio4_compressor_state(g, + model, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + layer->attn_compressor_kv, + layer->attn_compressor_gate, + layer->attn_compressor_ape, + DS4_N_HEAD_DIM, + comp_width, + pos0, + n_tokens); + } + DS4_METAL_PROFILE_ATTN_STAGE("compressor_refresh"); + } + if (ok) { + g->layer_n_comp[il] = n_comp; + for (uint32_t t = 0; t < n_tokens; t++) { + comp_counts[t] = (pos0 + t + 1u) / ratio; + } + if (n_comp != 0) { + metal_graph_debug_dump_tensor("KVcompress", + attn_comp_target, + (uint64_t)n_comp * DS4_N_HEAD_DIM, + il, + pos0); + } + metal_graph_debug_dump_tensor("attn_state_kv", + g->layer_attn_state_kv[il], + (uint64_t)comp_width * coff * ratio, + il, + pos0); + metal_graph_debug_dump_tensor("attn_state_score", + g->layer_attn_state_score[il], + (uint64_t)comp_width * coff * ratio, + il, + pos0); + } + metal_graph_attn_comp_prefill_target_free(attn_comp_target); + } else { + const bool aligned_chunk = + getenv("DS4_CUDA_NO_COMPRESSOR_PREFILL_BATCH") == NULL && + (pos0 % ratio) == 0u && (n_tokens % ratio) == 0u; + if (aligned_chunk) { + const uint32_t comp_before = g->layer_n_comp[il]; + const uint32_t comp_chunk = n_tokens / ratio; + if (comp_before + comp_chunk > g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + if (ok && DS4_GPU_ATTN_COMP_CACHE_F16 && comp_chunk > g->attn_comp_stage_cap) { + fprintf(stderr, "ds4: Metal graph compressed KV staging capacity exceeded at layer %u\n", il); + ok = false; + } + ds4_gpu_tensor *attn_comp_target = + ok ? metal_graph_attn_comp_prefill_target(g, il, comp_before, comp_chunk) : NULL; + if (ok && !attn_comp_target) ok = false; + if (ok && ratio == 4) { + ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( + attn_comp_target, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + pos0, + n_tokens, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } else if (ok) { + ok = ds4_gpu_compressor_prefill_tensor( + attn_comp_target, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos0, + n_tokens, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } + if (ok && comp_chunk != 0) { + ok = metal_graph_commit_attn_comp_stage(g, il, comp_before, comp_chunk); + } + if (ok && ratio == 4) { + ok = metal_graph_refresh_ratio4_compressor_state(g, + model, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + layer->attn_compressor_kv, + layer->attn_compressor_gate, + layer->attn_compressor_ape, + DS4_N_HEAD_DIM, + comp_width, + pos0, + n_tokens); + } + if (ok) { + g->layer_n_comp[il] = comp_before + comp_chunk; + if (comp_counts) { + for (uint32_t t = 0; t < n_tokens; t++) { + comp_counts[t] = (pos0 + t + 1u) / ratio; + } + } + metal_graph_debug_dump_tensor("KVcompress", + attn_comp_target, + (uint64_t)comp_chunk * DS4_N_HEAD_DIM, + il, + pos0); + metal_graph_debug_dump_tensor("attn_state_kv", + g->layer_attn_state_kv[il], + (uint64_t)comp_width * coff * ratio, + il, + pos0); + metal_graph_debug_dump_tensor("attn_state_score", + g->layer_attn_state_score[il], + (uint64_t)comp_width * coff * ratio, + il, + pos0); + } + metal_graph_attn_comp_prefill_target_free(attn_comp_target); + } else { + for (uint32_t t = 0; ok && t < n_tokens; t++) { + const uint32_t pos = pos0 + t; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + break; + } + ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(metal_graph_batch_comp_kv(g), t, comp_width); + ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(metal_graph_batch_comp_sc(g), t, comp_width); + const uint32_t comp_row = g->layer_n_comp[il]; + ok = kv_view && sc_view && + ds4_gpu_compressor_update_tensor(kv_view, + sc_view, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_attn_comp_update_target(g, il), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos, + metal_graph_attn_comp_update_row(comp_row), + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + false) != 0; + if (ok && emit) { + ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); + ok = comp_row_view && + ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVcompress", + comp_row_view, + DS4_N_HEAD_DIM, + il, + pos); + } + ds4_gpu_tensor_free(comp_row_view); + if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); + } + if (ok && emit) g->layer_n_comp[il]++; + if (comp_counts) comp_counts[t] = g->layer_n_comp[il]; + if (ok && t == 0) ok = metal_graph_capture_prefix1_attn_state(g, il); + ds4_gpu_tensor_free(sc_view); + ds4_gpu_tensor_free(kv_view); + } + } + n_comp = g->layer_n_comp[il]; + } + DS4_METAL_PROFILE_ATTN_STAGE("compressor"); + + if (ok && ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || + !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || + !layer->indexer_attn_q_b || !layer->indexer_proj) { + fprintf(stderr, "ds4: Metal layer-major prefill needs indexer weights\n"); + ok = false; + } + if (ok) { + ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + DS4_N_EMBD, + index_width, + metal_graph_batch_attn_norm(g), + n_tokens) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, + index_width, + metal_graph_batch_attn_norm(g), + n_tokens) != 0; + } + if (ok) metal_graph_debug_dump_tensor("indexer_comp_kv_raw", + metal_graph_batch_comp_kv(g), + (uint64_t)index_width * n_tokens, + il, + pos0); + if (ok) metal_graph_debug_dump_tensor("indexer_comp_score_raw", + metal_graph_batch_comp_sc(g), + (uint64_t)index_width * n_tokens, + il, + pos0); + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_indexer_q(g), + model, + layer->indexer_attn_q_b, + q_rank, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, + metal_graph_batch_qr_norm(g), + n_tokens); + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_indexer_q(g), + n_tokens, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_batch_indexer_q(g), + n_tokens * DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_indexer_weights(g), + model->map, + model->size, + layer->indexer_proj->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD, + metal_graph_batch_attn_norm(g), + n_tokens) != 0; + if (zero_prefix) { + if (ok && n_comp > g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal layer-major indexer cache capacity exceeded at layer %u\n", il); + ok = false; + } + if (ok) { + ok = ds4_gpu_compressor_prefill_tensor(g->layer_index_comp_cache[il], + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos0, + n_tokens, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } + if (ok && n_comp != 0) { + ok = ds4_gpu_dsv4_indexer_qat_tensor(g->layer_index_comp_cache[il], + n_comp, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + if (ok) { + ok = metal_graph_refresh_ratio4_compressor_state(g, + model, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + layer->indexer_compressor_kv, + layer->indexer_compressor_gate, + layer->indexer_compressor_ape, + DS4_N_INDEXER_HEAD_DIM, + index_width, + pos0, + n_tokens); + } + if (ok) { + g->layer_n_index_comp[il] = n_comp; + for (uint32_t t = 0; t < n_tokens; t++) { + index_counts[t] = (pos0 + t + 1u) / ratio; + } + if (n_comp != 0) { + metal_graph_debug_dump_tensor("indexer_KVcompress", + g->layer_index_comp_cache[il], + (uint64_t)n_comp * DS4_N_INDEXER_HEAD_DIM, + il, + pos0); + } + metal_graph_debug_dump_tensor("indexer_state_kv", + g->layer_index_state_kv[il], + (uint64_t)index_width * coff * ratio, + il, + pos0); + metal_graph_debug_dump_tensor("indexer_state_score", + g->layer_index_state_score[il], + (uint64_t)index_width * coff * ratio, + il, + pos0); + } + } else { + const bool aligned_chunk = + getenv("DS4_CUDA_NO_COMPRESSOR_PREFILL_BATCH") == NULL && + (pos0 % ratio) == 0u && (n_tokens % ratio) == 0u; + if (aligned_chunk) { + const uint32_t index_before = g->layer_n_index_comp[il]; + const uint32_t index_chunk = n_tokens / ratio; + if (index_before + index_chunk > g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + ds4_gpu_tensor *index_view = NULL; + if (ok) { + index_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_before * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)index_chunk * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + ok = index_view != NULL; + } + if (ok) { + ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( + index_view, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + metal_graph_batch_comp_kv(g), + metal_graph_batch_comp_sc(g), + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + pos0, + n_tokens, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + } + if (ok && index_chunk != 0) { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_view, + index_chunk, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + if (ok) { + ok = metal_graph_refresh_ratio4_compressor_state(g, + model, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + layer->indexer_compressor_kv, + layer->indexer_compressor_gate, + layer->indexer_compressor_ape, + DS4_N_INDEXER_HEAD_DIM, + index_width, + pos0, + n_tokens); + } + if (ok) { + g->layer_n_index_comp[il] = index_before + index_chunk; + if (index_counts) { + for (uint32_t t = 0; t < n_tokens; t++) { + index_counts[t] = (pos0 + t + 1u) / ratio; + } + } + metal_graph_debug_dump_tensor("indexer_KVcompress", + index_view, + (uint64_t)index_chunk * DS4_N_INDEXER_HEAD_DIM, + il, + pos0); + metal_graph_debug_dump_tensor("indexer_state_kv", + g->layer_index_state_kv[il], + (uint64_t)index_width * coff * ratio, + il, + pos0); + metal_graph_debug_dump_tensor("indexer_state_score", + g->layer_index_state_score[il], + (uint64_t)index_width * coff * ratio, + il, + pos0); + } + ds4_gpu_tensor_free(index_view); + } else { + for (uint32_t t = 0; ok && t < n_tokens; t++) { + const uint32_t pos = pos0 + t; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + break; + } + ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(metal_graph_batch_comp_kv(g), t, index_width); + ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(metal_graph_batch_comp_sc(g), t, index_width); + const uint32_t index_row = g->layer_n_index_comp[il]; + ok = kv_view && sc_view && + ds4_gpu_compressor_update_tensor(kv_view, + sc_view, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + g->layer_index_comp_cache[il], + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos, + index_row, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + false) != 0; + if (ok && emit) { + ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + if (!index_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + ds4_gpu_tensor_free(index_row_view); + } + } + if (ok && emit) g->layer_n_index_comp[il]++; + if (index_counts) index_counts[t] = g->layer_n_index_comp[il]; + if (ok && t == 0) ok = metal_graph_capture_prefix1_index_state(g, il); + ds4_gpu_tensor_free(sc_view); + ds4_gpu_tensor_free(kv_view); + } + } + } + } + if (ratio == 4) DS4_METAL_PROFILE_ATTN_STAGE("indexer_setup"); + + if (ok && !zero_prefix && n_tokens <= g->raw_cap) { + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos0, n_tokens); + /* See the raw-only branch above: batched mixed attention also + * consumes a logical raw window, linearized out of the ring. */ + const uint32_t raw_start = metal_graph_raw_start_for_span(g, + pos0 + n_tokens - 1u, + n_raw); + uint32_t use_comp_mask = 0; + bool use_indexed_comp = false; + double index_stage_t0 = 0.0; + + ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + metal_graph_batch_kv(g), + g->raw_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + if (ok && ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K) { + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + if (index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + ok = ds4_gpu_indexer_scores_decode_batch_tensor(metal_graph_indexer_scores(g), + metal_graph_batch_indexer_q(g), + metal_graph_batch_indexer_weights(g), + g->layer_index_comp_cache[il], + n_comp, + n_tokens, + pos0, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + ratio, + index_scale) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("score", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + if (ok) { + metal_graph_debug_dump_tensor("indexer_scores", + metal_graph_indexer_scores(g), + (uint64_t)n_comp * n_tokens, + il, + pos0); + } + if (ok) { + ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + n_comp, + n_tokens, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("topk", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + if (ok) { + metal_graph_debug_dump_i32_tensor("indexer_topk", + metal_graph_comp_selected(g), + (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, + il, + pos0); + } + } + if (ok) { + use_indexed_comp = true; + } + use_comp_mask = 1; + } + if (ok) { + if (use_indexed_comp) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->layer_raw_cache[il], + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + metal_graph_comp_selected(g), + n_tokens, + pos0, + n_raw, + g->raw_cap, + raw_start, + n_comp, + DS4_N_INDEXER_TOP_K, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("attention", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + } else { + ok = ds4_gpu_attention_decode_mixed_batch_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->layer_raw_cache[il], + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + use_comp_mask ? metal_graph_comp_mask(g) : NULL, + use_comp_mask, + n_tokens, + pos0, + n_raw, + g->raw_cap, + raw_start, + n_comp, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + } + if (ok) batch_attention_done = true; + } + + const bool topk_prefill_needed = ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K; + if (ok && zero_prefix && topk_prefill_needed && n_comp != 0) { + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + double index_stage_t0 = 0.0; + if (index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + ok = ds4_gpu_indexer_scores_prefill_tensor(metal_graph_indexer_scores(g), + metal_graph_batch_indexer_q(g), + metal_graph_batch_indexer_weights(g), + g->layer_index_comp_cache[il], + n_comp, + n_tokens, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + ratio, + index_scale) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("score", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + if (ok) { + metal_graph_debug_dump_tensor("indexer_scores", + metal_graph_indexer_scores(g), + (uint64_t)n_comp * n_tokens, + il, + pos0); + } + if (ok) { + ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + n_comp, + n_tokens, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("topk", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + if (ok) { + metal_graph_debug_dump_i32_tensor("indexer_topk", + metal_graph_comp_selected(g), + (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, + il, + pos0); + } + } + if (ok && tp_row_split_attn) { + /* Score/top-k selection above ran replicated over all rows; + * only the attention consumption splits. Passing the row + * offset through pos0 and clamping n_raw to the rows this + * rank can see keeps the kernel's first_raw_pos at the + * chunk origin, so the raw ring mapping is unchanged. */ + ds4_gpu_tensor *tp_topk = metal_graph_tensor_row_range_view( + metal_graph_comp_selected(g), tp_row0, tp_rows, DS4_N_INDEXER_TOP_K); + ok = tp_topk && + ds4_gpu_attention_indexed_mixed_batch_heads_tensor(tp_heads, + model->map, + model->size, + layer->attn_sinks->abs_offset, + tp_q, + g->layer_raw_cache[il], + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + tp_topk, + tp_rows, + pos0 + tp_row0, + tp_row0 + tp_rows, + g->raw_cap, + 0, + n_comp, + DS4_N_INDEXER_TOP_K, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + ds4_gpu_tensor_free(tp_topk); + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("attention", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + } else if (ok) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->layer_raw_cache[il], + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + metal_graph_comp_selected(g), + n_tokens, + pos0, + n_tokens, + g->raw_cap, + 0, + n_comp, + DS4_N_INDEXER_TOP_K, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + if (ok && index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("attention", + il, + pos0, + n_tokens, + n_comp, + &index_stage_t0); + } + } + if (ok) batch_attention_done = true; + } + if (ok && zero_prefix && !topk_prefill_needed && n_comp != 0) { + if (tp_row_split_attn) { + ok = ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(tp_heads, + model->map, + model->size, + layer->attn_sinks->abs_offset, + tp_q, + metal_graph_batch_kv(g), + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + tp_row0, + tp_rows, + n_tokens, + n_comp, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } else { + ok = ds4_gpu_attention_prefill_static_mixed_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + metal_graph_batch_kv(g), + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + n_tokens, + n_comp, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + if (ok) batch_attention_done = true; + } + } + + if (ok && !raw_batch_attention && !batch_attention_done) { + uint32_t raw_prefix_tokens = 0; + if (zero_prefix && ratio != 0 && n_tokens <= g->raw_cap && comp_counts != NULL) { + while (raw_prefix_tokens < n_tokens && comp_counts[raw_prefix_tokens] == 0u) { + raw_prefix_tokens++; + } + } + + if (raw_prefix_tokens != 0) { + if (tp_row_split_attn && raw_prefix_tokens == n_tokens) { + /* tp_attn_full_raw guarantees the whole chunk stays raw + * (n_tokens < ratio), so the split covers every row. */ + ok = ds4_gpu_attention_prefill_raw_heads_range_tensor(tp_heads, + model->map, + model->size, + layer->attn_sinks->abs_offset, + tp_q, + metal_graph_batch_kv(g), + tp_row0, + tp_rows, + n_tokens, + g->raw_window, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } else { + ok = ds4_gpu_attention_prefill_raw_heads_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + metal_graph_batch_kv(g), + raw_prefix_tokens, + g->raw_window, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + } + if (raw_prefix_tokens < n_tokens) { + for (uint32_t t = raw_prefix_tokens; ok && t < n_tokens; t++) { + const uint32_t pos = pos0 + t; + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); + const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); + const uint32_t cur_comp = comp_counts ? comp_counts[t] : 0u; + const uint32_t cur_index = index_counts ? index_counts[t] : 0u; + uint32_t n_selected = 0; + ds4_gpu_tensor *comp_mask = NULL; + + if (ratio == 4 && cur_comp > DS4_N_INDEXER_TOP_K) { + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + ds4_gpu_tensor *indexer_q_view = metal_graph_tensor_row_view( + metal_graph_batch_indexer_q(g), t, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); + ds4_gpu_tensor *indexer_w_view = metal_graph_tensor_row_view( + metal_graph_batch_indexer_weights(g), t, DS4_N_INDEXER_HEAD); + ok = indexer_q_view && indexer_w_view && + ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), + indexer_q_view, + indexer_w_view, + g->layer_index_comp_cache[il], + cur_index, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + index_scale) != 0 && + ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + cur_index, + 1, + DS4_N_INDEXER_TOP_K) != 0 && + ds4_gpu_dsv4_topk_mask_tensor(metal_graph_comp_mask(g), + metal_graph_comp_selected(g), + cur_index, + 1, + DS4_N_INDEXER_TOP_K) != 0; + ds4_gpu_tensor_free(indexer_w_view); + ds4_gpu_tensor_free(indexer_q_view); + if (ok) { + comp_mask = metal_graph_comp_mask(g); + n_selected = DS4_N_INDEXER_TOP_K < cur_index + ? DS4_N_INDEXER_TOP_K + : cur_index; + } + } + + ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(metal_graph_batch_q(g), t, q_dim); + ds4_gpu_tensor *kv_cache_view = metal_graph_tensor_row_view(metal_graph_batch_kv(g), t, DS4_N_HEAD_DIM); + ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(metal_graph_batch_heads(g), t, q_dim); + ok = ok && q_view && kv_cache_view && heads_view; + if (ok && !zero_prefix) { + ok = ds4_gpu_store_raw_kv_tensor(g->layer_raw_cache[il], + kv_cache_view, + g->raw_cap, + pos % g->raw_cap, + DS4_N_HEAD_DIM) != 0; + } + if (ok && comp_mask != NULL && n_selected != 0) { + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(heads_view, + model->map, + model->size, + layer->attn_sinks->abs_offset, + q_view, + g->layer_raw_cache[il], + g->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + metal_graph_comp_selected(g), + 1, + pos, + n_raw, + g->raw_cap, + raw_start, + cur_comp, + n_selected, + g->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } else if (ok) { + ok = ds4_gpu_attention_decode_heads_tensor(heads_view, + model->map, + model->size, + layer->attn_sinks->abs_offset, + q_view, + g->layer_raw_cache[il], + n_raw, + g->raw_cap, + raw_start, + cur_comp ? g->layer_attn_comp_cache[il] : NULL, + metal_graph_attn_comp_cache_is_f16(), + cur_comp, + comp_mask, + n_selected, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + ds4_gpu_tensor_free(heads_view); + ds4_gpu_tensor_free(kv_cache_view); + ds4_gpu_tensor_free(q_view); + } + } + } + DS4_METAL_PROFILE_ATTN_STAGE("attention"); + + if (ok) { + metal_graph_debug_dump_tensor("kqv_out", metal_graph_batch_heads(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + if (ok) ok = ds4_gpu_rope_tail_tensor(tp_heads ? tp_heads : metal_graph_batch_heads(g), + tp_rows, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0 + tp_row0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) { + metal_graph_debug_dump_tensor("kqv_back", metal_graph_batch_heads(g), + (uint64_t)n_tokens * q_dim, il, pos0); + } + DS4_METAL_PROFILE_ATTN_STAGE("inv_rope"); + const bool attn_out_debug = + metal_graph_debug_wants("attn_low", il, pos0) || + metal_graph_debug_wants("attn_out", il, pos0); + bool attn_out_f16 = false; + if (ok && + !attn_out_debug && + !tp_row_split_attn && + layer->attn_output_a->type == DS4_TENSOR_Q8_0 && + layer->attn_output_b->type == DS4_TENSOR_Q8_0 && + !metal_graph_directional_steering_attn_enabled(g)) { + attn_out_f16 = ds4_gpu_attention_output_q8_batch_f16_tensor(g->batch_q_half, + metal_graph_batch_attn_low(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + metal_graph_batch_heads(g), + n_tokens) != 0; + } + uint64_t tp_attn_gate_seq = 0; + /* Opt-in sub-chunk gate pipelining (see metal_graph_tp_subgate_pipeline; + * measured net-negative on the M5 Max pair, kept for slower wires). + * Kernel-path constraint: the output projection picks its kernel by row + * count (direct low below 32 rows, the 64-token-tile TensorOps path at + * multiples of 64, ids-cache fallback otherwise), and the paths are not + * bit-identical per row. Parity with the single-node reference + * therefore requires every sub-call to land on the same path as the + * full-chunk call: n_tokens % 256 == 0 puts the chunk, the rank halves, + * and the quarter sub-calls all on the TensorOps path. Other sizes + * keep the proven single-gate swap. */ + const bool tp_attn_pipeline = + tp_row_split_attn && (n_tokens % 256u) == 0u && + metal_graph_tp_subgate_pipeline(); + if (!attn_out_f16) { + if (ok && tp_attn_pipeline) { + /* Sub-chunk pipelined swap: the output projection runs in two + * sub-halves of this rank's rows and each sub-half's row swap + * is kicked as soon as its rows land in batch_attn_out, so the + * first wire exchange overlaps the second sub-half's compute. + * The two kicks use opposite flag-slot parities; the wait below + * (before the HC post expand) covers both. */ + const uint32_t tp_sub1 = (tp_half_rows + 1u) / 2u; + const uint32_t tp_c1 = tp_rows < tp_sub1 ? tp_rows : tp_sub1; + const uint32_t tp_own_base = g->tp_rank == 0 ? 0u : tp_half_rows; + const uint32_t tp_peer_base = g->tp_rank == 0 ? tp_half_rows : 0u; + for (uint32_t sub = 0; ok && sub < 2u; sub++) { + const uint32_t coff = sub == 0 ? 0u : tp_c1; + const uint32_t crows = sub == 0 ? tp_c1 : tp_rows - tp_c1; + const uint32_t soff = sub == 0 ? 0u : tp_sub1; + const uint32_t srows = sub == 0 ? tp_sub1 : tp_half_rows - tp_sub1; + if (crows != 0) { + ds4_gpu_tensor *sub_heads = metal_graph_tensor_row_range_view( + metal_graph_batch_heads(g), tp_row0 + coff, crows, q_dim); + ds4_gpu_tensor *sub_out = metal_graph_tensor_row_range_view( + metal_graph_batch_attn_out(g), tp_row0 + coff, crows, DS4_N_EMBD); + ok = sub_heads && sub_out && + metal_graph_attention_output_dense_quant_batch(sub_out, + metal_graph_batch_attn_low(g), + g, + model, + layer->attn_output_a, + layer->attn_output_b, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + sub_heads, + crows); + ds4_gpu_tensor_free(sub_out); + ds4_gpu_tensor_free(sub_heads); + } + if (ok && srows != 0) { + ds4_gpu_tensor *send_sub = metal_graph_tensor_row_range_view( + metal_graph_batch_attn_out(g), tp_own_base + soff, srows, DS4_N_EMBD); + ds4_gpu_tensor *recv_sub = metal_graph_tensor_row_range_view( + metal_graph_batch_attn_out(g), tp_peer_base + soff, srows, DS4_N_EMBD); + uint64_t seq = 0; + if (send_sub && recv_sub) { + seq = ds4_gpu_tp_big_gate_kick(il, n_tokens, + send_sub, recv_sub, + (uint64_t)srows * DS4_N_EMBD * sizeof(float)); + } + ok = seq != 0; + if (ok) tp_attn_gate_seq = seq; + ds4_gpu_tensor_free(recv_sub); + ds4_gpu_tensor_free(send_sub); + } + } + } else if (ok) { + ok = metal_graph_attention_output_dense_quant_batch(tp_attn_out ? tp_attn_out : metal_graph_batch_attn_out(g), + metal_graph_batch_attn_low(g), + g, + model, + layer->attn_output_a, + layer->attn_output_b, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + tp_heads ? tp_heads : metal_graph_batch_heads(g), + tp_rows); + } + if (ok) { + metal_graph_debug_dump_tensor("attn_low", metal_graph_batch_attn_low(g), + (uint64_t)n_tokens * n_groups * rank, + il, + pos0); + } + if (ok) { + metal_graph_debug_dump_tensor("attn_out", metal_graph_batch_attn_out(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + } + DS4_METAL_PROFILE_ATTN_STAGE("output_proj"); + if (ok && tp_row_split_attn) { + /* Release point for the pipelined row swaps of batch_attn_out: both + * ranks reach the HC post expand with identical full tensors. */ + if (tp_attn_pipeline) { + ok = tp_attn_gate_seq != 0 && + ds4_gpu_tp_big_gate_wait(tp_attn_gate_seq) != 0; + } else { + const uint64_t half_bytes = + (uint64_t)tp_half_rows * DS4_N_EMBD * sizeof(float); + ds4_gpu_tensor *send_half = metal_graph_tensor_row_range_view( + metal_graph_batch_attn_out(g), + g->tp_rank == 0 ? 0 : tp_half_rows, tp_half_rows, DS4_N_EMBD); + ds4_gpu_tensor *recv_half = metal_graph_tensor_row_range_view( + metal_graph_batch_attn_out(g), + g->tp_rank == 0 ? tp_half_rows : 0, tp_half_rows, DS4_N_EMBD); + ok = send_half && recv_half && + ds4_gpu_tp_big_gate_encode(il, n_tokens, + send_half, recv_half, + half_bytes) != 0; + ds4_gpu_tensor_free(send_half); + ds4_gpu_tensor_free(recv_half); + } + if (!ok) fprintf(stderr, "ds4: TP prefill attention row gate failed (layer %u)\n", il); + } + if (ok && !attn_out_f16 && metal_graph_directional_steering_attn_enabled(g)) { + ok = metal_graph_apply_directional_steering_attn(g, metal_graph_batch_attn_out(g), il, n_tokens); + } + if (ok && attn_out_f16) { + ok = ds4_gpu_hc_expand_split_half_tensor(after_attn_hc_view, + g->batch_q_half, + metal_graph_batch_cur_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok) { + ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, + metal_graph_batch_attn_out(g), + metal_graph_batch_cur_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("hc_attn_post", metal_graph_batch_after_attn_hc(g), + (uint64_t)n_tokens * hc_dim, il, pos0); + } + DS4_METAL_PROFILE_ATTN_STAGE("hc_post"); + ds4_gpu_tensor_free(tp_attn_out); + ds4_gpu_tensor_free(tp_heads); + ds4_gpu_tensor_free(tp_qr_norm); + ds4_gpu_tensor_free(tp_q_half); + ds4_gpu_tensor_free(tp_q); + ds4_gpu_tensor_free(after_attn_hc_view); + ds4_gpu_tensor_free(attn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); + if (index_counts != index_counts_stack) free(index_counts); + if (comp_counts != comp_counts_stack) free(comp_counts); +#undef DS4_METAL_PROFILE_ATTN_STAGE +#undef DS4_METAL_PROFILE_Q_STAGE + return ok; +} + +static bool metal_graph_encode_mixed_routed_rows( + ds4_gpu_graph *g, + ds4_decode_item *decode_items, + int decode_count, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t prefill_rows); + +/* Encode the batched prefill FFN half: HC pre/norm, shared expert, routed + * experts, sum, and HC post. A non-empty decode tail has already been + * prepared in rows [n_tokens, n_tokens + decode_count); only the routed + * expert dispatch is shared between the two arithmetic paths. */ +static bool metal_graph_encode_layer_ffn_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + ds4_decode_item *decode_items, + int decode_count) { + if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; + if (decode_count < 0 || + (decode_count > 0 && + (!decode_items || (uint64_t)n_tokens + (uint32_t)decode_count > g->prefill_cap))) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + const uint64_t down_expert_bytes = routed_out_dim * down_row_bytes; + const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); + double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; +#define DS4_METAL_PROFILE_FFN_STAGE(name) do { \ + if (ok && layer_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary("ffn", (name), il, pos0, n_tokens, &layer_stage_t0); \ + } \ + } while (0) + + ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_mix(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_split(g), 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *ffn_cur_view = ds4_gpu_tensor_view( + metal_graph_batch_ffn_cur(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view( + metal_graph_batch_next_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); + bool ok = hc_mix_view && hc_split_view && ffn_cur_view && next_hc_view; + const bool fuse_hc_norm = n_tokens > 1 && + DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + metal_graph_enable_batch_hc_norm_fusion(); + if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, + metal_graph_batch_flat_hc(g), + model, + layer->hc_ffn_fn, + metal_graph_batch_after_attn_hc(g), + hc_dim, + n_tokens); + if (metal_graph_use_reference_hc_decode()) { + if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, + hc_mix_view, + model->map, + model->size, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(ffn_cur_view, + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (fuse_hc_norm) { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(ffn_cur_view, + metal_graph_batch_ffn_norm(g), + hc_split_view, + hc_mix_view, + metal_graph_batch_after_attn_hc(g), + model->map, + model->size, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset, + layer->ffn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(ffn_cur_view, + hc_split_view, + hc_mix_view, + metal_graph_batch_after_attn_hc(g), + model->map, + model->size, + layer->hc_ffn_scale->abs_offset, + layer->hc_ffn_base->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_pre", metal_graph_batch_ffn_cur(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + DS4_METAL_PROFILE_FFN_STAGE("hc_pre"); + if (ok && !fuse_hc_norm) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_ffn_norm(g), + metal_graph_batch_ffn_cur(g), + model->map, + model->size, + layer->ffn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_norm", metal_graph_batch_ffn_norm(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + DS4_METAL_PROFILE_FFN_STAGE("norm"); + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_router_logits(g), + model, + layer->ffn_gate_inp, + DS4_N_EMBD, + DS4_N_EXPERT, + metal_graph_batch_ffn_norm(g), + n_tokens); + + ds4_gpu_tensor *router_tokens = NULL; + if (ok) { + router_tokens = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), + (uint64_t)g->batch_token_offset * sizeof(int32_t), + (uint64_t)n_tokens * sizeof(int32_t)); + ok = router_tokens != NULL; + } + if (ok) ok = ds4_gpu_router_select_batch_tensor(metal_graph_batch_router_selected(g), + metal_graph_batch_router_weights(g), + metal_graph_batch_router_probs(g), + model->map, + model->size, + layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, + layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, + layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, + 0, + 0, + layer->ffn_exp_probs_b != NULL, + layer->ffn_gate_tid2eid != NULL, + metal_graph_batch_router_logits(g), + metal_graph_prefill_tokens(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE, + n_tokens) != 0; + ds4_gpu_tensor_free(router_tokens); + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_logits", metal_graph_batch_router_logits(g), + (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); + metal_graph_debug_dump_tensor("ffn_moe_probs", metal_graph_batch_router_probs(g), + (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); + metal_graph_debug_dump_i32_tensor("ffn_moe_topk", metal_graph_batch_router_selected(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); + metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", metal_graph_batch_router_weights(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); + } + DS4_METAL_PROFILE_FFN_STAGE("router"); + + if (ok) { + ok = metal_graph_cuda_stream_prefill_batch_selected_load(g, + model, + layer, + il, + n_tokens, + gate_expert_bytes, + down_expert_bytes); + } + +#ifdef DS4_ROCM_BUILD + rocm_graph_batch_selected_async_load rocm_batch_selected_async = {0}; + bool rocm_batch_selected_async_started = false; + const bool rocm_batch_selected_shared_overlap = + ok && + g->ssd_streaming && + !g->quality && + n_tokens > 1 && + DS4_N_EXPERT_USED == 6 && + !rocm_graph_stream_prefill_full_layer_enabled(g, layer, il, n_tokens) && + layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + layer->ffn_down_exps->type == DS4_TENSOR_Q2_K; + if (rocm_batch_selected_shared_overlap) { + uint64_t selected_event = 0; + if (ds4_gpu_signal_selected_readback_ready(&selected_event) == 0) { + ok = false; + } else { + ok = rocm_graph_batch_selected_async_load_start( + &rocm_batch_selected_async, + metal_graph_batch_router_selected(g), + model, + layer, + il, + n_tokens, + selected_event, + gate_expert_bytes, + down_expert_bytes); + rocm_batch_selected_async_started = ok; + } + } +#endif + + const bool selected_readahead_shared = + metal_graph_stream_prefill_selected_readahead_shared_enabled(g) +#ifdef DS4_ROCM_BUILD + && !rocm_batch_selected_async_started +#endif + ; + if (ok && + metal_graph_stream_prefill_selected_readahead_enabled(g) && +#ifdef DS4_ROCM_BUILD + !rocm_batch_selected_async_started && +#endif + !selected_readahead_shared) { + if (ds4_gpu_end_commands() == 0) { + ok = false; + } else { + ok = metal_graph_stream_readahead_selected_experts_from_gpu(g, + model, + layer, + il, + n_tokens, + gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_begin_commands() != 0; + } + } + + const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos0); + bool shared_down_f16 = false; + +#define DS4_METAL_TRY_SHARED_DOWN_F16() do { \ + if (ok && !tp_row_split_ffn && !keep_ffn_out && \ + !metal_graph_debug_wants("ffn_shexp", il, pos0)) { \ + shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->batch_q_half, \ + model->map, \ + model->size, \ + layer->ffn_down_shexp->abs_offset, \ + shared_dim, \ + DS4_N_EMBD, \ + metal_graph_batch_shared_mid(g), \ + n_tokens) != 0; \ + } \ + } while (0) + +#define DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT() do { \ + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_gate", \ + il, \ + pos0, \ + metal_graph_batch_shared_gate(g), \ + model, \ + layer->ffn_gate_shexp, \ + DS4_N_EMBD, \ + shared_dim, \ + tp_ffn_x ? tp_ffn_x : metal_graph_batch_ffn_norm(g), \ + tp_rows); \ + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_up", \ + il, \ + pos0, \ + metal_graph_batch_shared_up(g), \ + model, \ + layer->ffn_up_shexp, \ + DS4_N_EMBD, \ + shared_dim, \ + tp_ffn_x ? tp_ffn_x : metal_graph_batch_ffn_norm(g), \ + tp_rows); \ + DS4_METAL_PROFILE_FFN_STAGE("shared_gate_up"); \ + if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_batch_shared_mid(g), \ + metal_graph_batch_shared_gate(g), \ + metal_graph_batch_shared_up(g), \ + (uint32_t)((uint64_t)tp_rows * shared_dim), \ + DS4_SWIGLU_CLAMP_EXP, \ + 1.0f) != 0; \ + DS4_METAL_TRY_SHARED_DOWN_F16(); \ + if (ok && !shared_down_f16) ok = metal_graph_matmul_q8_0_named_tensor("shared_down", \ + il, \ + pos0, \ + metal_graph_batch_shared_out(g), \ + model, \ + layer->ffn_down_shexp, \ + shared_dim, \ + DS4_N_EMBD, \ + metal_graph_batch_shared_mid(g), \ + tp_rows); \ + DS4_METAL_PROFILE_FFN_STAGE("shared_down"); \ + if (ok && !shared_down_f16) { \ + metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_batch_shared_out(g), \ + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); \ + } \ + } while (0) + + bool shared_done = false; + /* With 50/50 expert residency, every rank evaluates every prompt row + * against its local expert half. For large chunks the replicated shared + * expert remains row-split; its rows are folded into the local routed + * partial before the one all-reduce-style bulk exchange. */ + const bool tp_split_ffn = g->tp_world == 2; + const bool tp_row_split_ffn = + tp_split_ffn && g->tp_batch_rows != n_tokens && !keep_ffn_out && + !metal_graph_directional_steering_ffn_enabled(g) && + n_tokens >= metal_graph_tp_prefill_split_min(); + const uint32_t tp_half_rows = (n_tokens + 1u) / 2u; + const uint32_t tp_row0 = (tp_row_split_ffn && g->tp_rank != 0) ? tp_half_rows : 0; + const uint32_t tp_rows = tp_row_split_ffn ? + (g->tp_rank == 0 ? tp_half_rows : n_tokens - tp_half_rows) : n_tokens; + ds4_gpu_tensor *tp_ffn_x = tp_row_split_ffn ? + metal_graph_tensor_row_range_view(metal_graph_batch_ffn_norm(g), tp_row0, tp_rows, + DS4_N_EMBD) : NULL; + if (tp_row_split_ffn && !tp_ffn_x) ok = false; + if (ok && selected_readahead_shared) { + if (ds4_gpu_end_commands() == 0) { + ok = false; + } else { + ok = metal_graph_stream_readahead_selected_experts_from_gpu(g, + model, + layer, + il, + n_tokens, + gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_begin_commands() != 0; + } + if (ok) { + DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); + shared_done = ok; + } + if (ok) { + if (ds4_gpu_end_commands() == 0) { + ok = false; + } else { + ok = ds4_gpu_begin_commands() != 0; + } + } + } + + if (ok && + !shared_done && + (metal_graph_stream_prefill_selected_pagein_enabled(g) || + metal_graph_stream_prefill_selected_madvise_enabled(g))) { + metal_graph_stream_pagein_job pagein_job; + memset(&pagein_job, 0, sizeof(pagein_job)); + bool pagein_commands_open = false; + if (ds4_gpu_end_commands() == 0) { + ok = false; + } else { + ok = metal_graph_stream_prefill_selected_pagein_start(g, + model, + layer, + il, + n_tokens, + gate_expert_bytes, + down_expert_bytes, + &pagein_job); + } + if (ok) { + if (ds4_gpu_begin_commands() == 0) { + ok = false; + } else { + pagein_commands_open = true; + } + } + if (ok) { + DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); + shared_done = ok; + } + if (pagein_commands_open) { + if (ds4_gpu_end_commands() == 0) ok = false; + } + if (!metal_graph_stream_prefill_selected_pagein_join(&pagein_job)) { + ok = false; + } + if (ok) ok = ds4_gpu_begin_commands() != 0; + } + +#ifdef DS4_ROCM_BUILD + if (rocm_batch_selected_async_started) { + if (ok && !shared_done) { + DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); + shared_done = ok; + } + const bool finish_ok = + rocm_graph_batch_selected_async_load_finish(&rocm_batch_selected_async); + ok = ok && finish_ok; + } +#endif + + const bool tp_split_batch_moe = + g->tp_batch_rows == n_tokens && n_tokens > 0 && + g->tp_world == 2 && + g->tp_batch_out && g->tp_batch_in; + const bool cuda_tp_owned_batch_moe = + g->cuda_tp_ep && g->cuda_tp_prefill_ffn; + if (ok && cuda_tp_owned_batch_moe) { + ok = metal_graph_encode_mixed_routed_rows( + g, decode_items, decode_count, model, layer, il, n_tokens); + } else if (ok && tp_split_batch_moe) { + /* Verify-block expert split: run the contiguous-half split + * single-token routed kernels per row into the slab batch-out + * rows, exchange all rows with one gate, then materialize the + * combined routed output. The add is commutative, so both ranks + * compute bit-identical sums and stay in lockstep. */ + const uint64_t vec_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + for (uint32_t r = 0; ok && r < n_tokens; r++) { + ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( + g->tp_batch_out[il], (uint64_t)r * vec_bytes, vec_bytes); + ds4_gpu_tensor *x_row = ds4_gpu_tensor_view( + metal_graph_batch_ffn_norm(g), (uint64_t)r * vec_bytes, vec_bytes); + ds4_gpu_tensor *sel_row = ds4_gpu_tensor_view( + metal_graph_batch_router_selected(g), + (uint64_t)r * DS4_N_EXPERT_USED * sizeof(int32_t), + (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + ds4_gpu_tensor *w_row = ds4_gpu_tensor_view( + metal_graph_batch_router_weights(g), + (uint64_t)r * DS4_N_EXPERT_USED * sizeof(float), + (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); + ok = out_row && x_row && sel_row && w_row && + ds4_gpu_routed_moe_one_tensor(out_row, + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + sel_row, w_row, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_SWIGLU_CLAMP_EXP, + x_row, + NULL, + il, + false) != 0; + ds4_gpu_tensor_free(w_row); + ds4_gpu_tensor_free(sel_row); + ds4_gpu_tensor_free(x_row); + ds4_gpu_tensor_free(out_row); + } + if (ok) ok = ds4_gpu_tp_batch_gate_encode(il, n_tokens) != 0; + if (ok) { + ok = ds4_gpu_add_tensor(metal_graph_batch_routed_out(g), + g->tp_batch_out[il], + g->tp_batch_in[il], + (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; + } + } else if (ok) { + ok = ds4_gpu_routed_moe_batch_tensor(metal_graph_batch_routed_out(g), + metal_graph_batch_routed_gate(g), + metal_graph_batch_routed_up(g), + metal_graph_batch_routed_mid(g), + metal_graph_batch_routed_down(g), + model->map, + model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + metal_graph_batch_router_selected(g), + metal_graph_batch_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_SWIGLU_CLAMP_EXP, + metal_graph_batch_ffn_norm(g), + il, + n_tokens, + &g->batch_routed_mid_is_f16, + false) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_batch_routed_gate(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_batch_routed_up(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); + } + if (ok) { + const uint64_t routed_mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim; + if (g->batch_routed_mid_is_f16) { + metal_graph_debug_dump_f16_tensor("ffn_moe_weighted_swiglu", metal_graph_batch_routed_mid(g), + routed_mid_elems, il, pos0); + } else { + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_batch_routed_mid(g), + routed_mid_elems, il, pos0); + } + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_batch_routed_down(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos0); + } + if (ok) { + metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_batch_routed_out(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + DS4_METAL_PROFILE_FFN_STAGE("routed_moe"); + if (!shared_done) { + DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); + } +#undef DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT +#undef DS4_METAL_TRY_SHARED_DOWN_F16 + + if (ok && tp_row_split_ffn) { + /* Each shared-expert row must appear exactly once in the all-reduce. + * Fold this rank's shared rows into its full-row routed partial; the + * peer does the same for the complementary rows. */ + ds4_gpu_tensor *own_rows = + metal_graph_tensor_row_range_view(metal_graph_batch_routed_out(g), tp_row0, + tp_rows, DS4_N_EMBD); + ok = own_rows && + ds4_gpu_add_tensor(own_rows, own_rows, metal_graph_batch_shared_out(g), + (uint32_t)((uint64_t)tp_rows * DS4_N_EMBD)) != 0; + ds4_gpu_tensor_free(own_rows); + } + + if (ok && tp_split_ffn && !tp_split_batch_moe) { + /* All rows contain this rank's routed-expert partial. Exchange that + * matrix in one bulk gate, then add in canonical rank order. The + * batch verify path above already performed the equivalent slab gate. */ + const uint64_t bytes = + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); + ok = metal_graph_ensure_batch_ffn_out(g) && + ds4_gpu_tp_big_gate_encode(il, n_tokens, + metal_graph_batch_routed_out(g), + metal_graph_batch_ffn_out(g), + bytes) != 0; + if (ok) { + ds4_gpu_tensor *first = g->tp_rank == 0 ? + metal_graph_batch_routed_out(g) : metal_graph_batch_ffn_out(g); + ds4_gpu_tensor *second = g->tp_rank == 0 ? + metal_graph_batch_ffn_out(g) : metal_graph_batch_routed_out(g); + ok = ds4_gpu_add_tensor(metal_graph_batch_routed_out(g), first, second, + (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; + } + if (!ok) { + fprintf(stderr, "ds4: TP prefill FFN all-reduce failed (layer %u)\n", il); + } + } + + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_batch_ffn_out(g) && + ds4_gpu_add_tensor(metal_graph_batch_ffn_out(g), + metal_graph_batch_shared_out(g), + metal_graph_batch_routed_out(g), + (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; + } + if (ok && keep_ffn_out) { + metal_graph_debug_dump_tensor("ffn_out", metal_graph_batch_ffn_out(g), + (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_batch_ffn_out(g), il, n_tokens); + } + if (ok && metal_graph_directional_steering_ffn_enabled(g)) { + ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, + metal_graph_batch_ffn_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + else if (ok && shared_down_f16) { + ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, + metal_graph_batch_routed_out(g), + g->batch_q_half, + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + else if (ok && tp_row_split_ffn) { + /* Shared expert already folded into the exchanged routed rows. */ + ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, + metal_graph_batch_routed_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + else if (ok) { + ok = ds4_gpu_hc_expand_add_split_tensor(next_hc_view, + metal_graph_batch_routed_out(g), + metal_graph_batch_shared_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + DS4_METAL_PROFILE_FFN_STAGE("hc_post"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_batch_next_hc(g), + (uint64_t)n_tokens * hc_dim, il, pos0); + } + DS4_METAL_PROFILE_FFN_STAGE("hc_post"); + ds4_gpu_tensor_free(tp_ffn_x); + ds4_gpu_tensor_free(next_hc_view); + ds4_gpu_tensor_free(ffn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); +#undef DS4_METAL_PROFILE_FFN_STAGE + return ok; +} + +/* Encode one complete layer for prefill by chaining attention and FFN batches. */ +static bool metal_graph_encode_layer_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens) { + if (g->placement) { + const int this_tier = g->placement[il + 1u]; + if (!metal_graph_set_active_tier_batch(g, this_tier, n_tokens)) { + return false; + } + } + bool ok = metal_graph_layer_stage_profile_start(il); + if (ok) { + ok = metal_graph_encode_layer_attention_batch(g, model, layer, il, pos0, n_tokens); + } + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u attention batch encode failed\n", il); + } + if (ok) { + ok = metal_graph_encode_layer_ffn_batch(g, model, layer, il, pos0, + n_tokens, NULL, 0); + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u ffn batch encode failed\n", il); + } + } + if (ok) { + ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); + g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); + g->batch_next_hc_by_tier[g->active_tier] = tmp; + } + return ok; +} + +static bool metal_graph_eval_token_raw_swa_streaming( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + uint32_t pos, + float *logits) { + if (g->raw_cap == 0) { + fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); + return false; + } + + const bool profile = + glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", + "DS4_METAL_GRAPH_TOKEN_PROFILE"); + const bool throttle = graph_power_throttle_enabled(g); + const double t0 = (profile || throttle) ? now_sec() : 0.0; + const uint32_t raw_row = pos % g->raw_cap; + const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); + metal_graph_dspark_capture_begin(g); + + const bool static_decode_map = metal_graph_stream_decode_static_map_enabled(); + const bool static_map_state_cache = + static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); + const bool batch_static_decode = + static_decode_map && metal_graph_stream_decode_layer_batch_enabled(g); + bool ok = true; + if (static_decode_map) { + if (!static_map_state_cache || !g->streaming_static_decode_map_current) { + ok = metal_graph_stream_map_decode_static_all(model, weights); + if (ok) g->streaming_static_decode_map_current = static_map_state_cache; + } + } else { + g->streaming_static_decode_map_current = false; + ok = metal_graph_stream_map_token(model, weights); + } + if (ok && !static_decode_map && DS4_N_LAYER > 0) { + metal_graph_stream_readahead_layer_decode(model, weights, 0); + } + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (batch_static_decode) { + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + ok = metal_graph_encode_decode_layer(g, + model, + &weights->layer[il], + il, + pos, + g->layer_raw_cache[il], + g->raw_cap, + raw_row, + n_raw, + token); + if (ok) { + ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); + g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); + g->after_ffn_hc_by_tier[g->active_tier] = tmp; + ok = metal_graph_dspark_capture_decode_layer(g, il); + } + } + if (ok && logits) { + ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + } + const double t_encoded = (profile || throttle) ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = (profile || throttle) ? now_sec() : 0.0; + if (ok && logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + const double t_read = (profile || throttle) ? now_sec() : 0.0; + if (profile) { + fprintf(stderr, + "ds4: metal SSD streaming batched token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", + pos, + (t_encoded - t0) * 1000.0, + (t_done - t_encoded) * 1000.0, + (t_read - t_done) * 1000.0, + (t_read - t0) * 1000.0, + logits != NULL); + } + if (ok && throttle) { + graph_power_note_decode_token(g, t_read - t0); + } + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after batched SSD streaming graph eval failure also failed\n"); + } + } + return ok; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + + double encode_s = 0.0; + double execute_s = 0.0; + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + const double tl0 = profile ? now_sec() : 0.0; + if (!static_decode_map && !metal_graph_stream_map_layer_decode(model, weights, il)) { + ok = false; + break; + } + if (!static_decode_map && il + 1 < DS4_N_LAYER) { + metal_graph_stream_readahead_layer_decode(model, weights, il + 1); + } else if (!static_decode_map && logits) { + metal_graph_stream_readahead_output(model, weights); + } + if (ok) ok = ds4_gpu_begin_commands() != 0; + bool encoded_layer = false; + if (ok) { + ok = metal_graph_encode_decode_layer(g, + model, + &weights->layer[il], + il, + pos, + g->layer_raw_cache[il], + g->raw_cap, + raw_row, + n_raw, + token); + encoded_layer = true; + } + if (encoded_layer) { + ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); + g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); + g->after_ffn_hc_by_tier[g->active_tier] = tmp; + if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); + } + const double tl_encoded = profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double tl_done = profile ? now_sec() : 0.0; + if (profile) { + encode_s += tl_encoded - tl0; + execute_s += tl_done - tl_encoded; + } + } + + if (ok && logits && !static_decode_map) ok = metal_graph_stream_map_output(model, weights); + const double t_head0 = profile ? now_sec() : 0.0; + if (ok && logits) ok = ds4_gpu_begin_commands() != 0; + if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + const double t_head_encoded = profile ? now_sec() : 0.0; + if (ok && logits) ok = ds4_gpu_end_commands() != 0; + const double t_done = (profile || throttle) ? now_sec() : 0.0; + if (ok && logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + const double t_read = (profile || throttle) ? now_sec() : 0.0; + + if (profile) { + if (logits) { + encode_s += t_head_encoded - t_head0; + execute_s += t_done - t_head_encoded; + } + fprintf(stderr, + "ds4: metal SSD streaming token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", + pos, + encode_s * 1000.0, + execute_s * 1000.0, + (t_read - t_done) * 1000.0, + (t_read - t0) * 1000.0, + logits != NULL); + } + if (ok) graph_power_note_decode_token(g, t_read - t0); + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after SSD streaming graph eval failure also failed\n"); + } + } + return ok; +} + +/* Execute one Metal decode token and read back logits. */ +static bool metal_graph_eval_token_raw_swa( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + uint32_t pos, + float *logits) { + if (g && g->ssd_streaming) { + return metal_graph_eval_token_raw_swa_streaming(g, model, weights, token, pos, logits); + } + + const bool profile = + glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", + "DS4_METAL_GRAPH_TOKEN_PROFILE"); + const bool throttle = graph_power_throttle_enabled(g); + const double t0 = (profile || throttle) ? now_sec() : 0.0; + + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, logits != NULL, true); + const double t_encoded = (profile || throttle) ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = (profile || throttle) ? now_sec() : 0.0; + + if (ok && logits && g->tp_world == 2 && g->tp_logits_half) { + const uint64_t tp_vhalf = (uint64_t)DS4_N_VOCAB / 2u; + const uint64_t off = (uint64_t)g->tp_rank * tp_vhalf * sizeof(float); + ok = ds4_gpu_tensor_read(metal_graph_logits(g), off, logits + g->tp_rank * tp_vhalf, + tp_vhalf * sizeof(float)) != 0; + } else if (ok && logits && !(g->tp_world == 2 && g->tp_rank == 1)) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + const double t_read = (profile || throttle) ? now_sec() : 0.0; + if (profile) { + fprintf(stderr, + "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", + pos, + (t_encoded - t0) * 1000.0, + (t_done - t_encoded) * 1000.0, + (t_read - t_done) * 1000.0, + (t_read - t0) * 1000.0, + logits != NULL); + } + if (ok) graph_power_note_decode_token(g, t_read - t0); + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after graph eval failure also failed\n"); + } + } + return ok; +} + +static bool metal_graph_streaming_decode_prefill_wide_default( + const ds4_weights *weights) { + return DS4_MODEL_VARIANT == DS4_VARIANT_FLASH && + weights && + DS4_N_LAYER > 0 && + weights->layer[0].ffn_gate_exps->type == DS4_TENSOR_Q4_K && + weights->layer[0].ffn_up_exps->type == DS4_TENSOR_Q4_K && + weights->layer[0].ffn_down_exps->type == DS4_TENSOR_Q4_K; +} + +static uint32_t metal_graph_streaming_decode_prefill_max_tokens( + const ds4_gpu_graph *g, + const ds4_weights *weights) { + (void)g; + if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_DECODE_PREFILL", + "DS4_METAL_DISABLE_STREAMING_DECODE_PREFILL")) { + return 0; + } + + const char *env = glm_graph_env_value( + "DS4_ROCM_STREAMING_DECODE_PREFILL_MAX", + "DS4_METAL_STREAMING_DECODE_PREFILL_MAX"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } + + if (DS4_MODEL_VARIANT != DS4_VARIANT_PRO && + DS4_MODEL_VARIANT != DS4_VARIANT_FLASH) { + return 0u; + } + return metal_graph_streaming_decode_prefill_wide_default(weights) ? 64u : 18u; +} + +static bool metal_graph_use_streaming_decode_prefill( + const ds4_gpu_graph *g, + const ds4_weights *weights, + uint32_t n_tokens) { + const uint32_t max_tokens = + metal_graph_streaming_decode_prefill_max_tokens(g, weights); + return g && + g->ssd_streaming && + !g->quality && + n_tokens != 0 && + max_tokens != 0 && + n_tokens <= max_tokens; +} + +static bool metal_graph_use_streaming_decode_prefill_range( + const ds4_gpu_graph *g, + const ds4_weights *weights, + uint32_t start, + uint32_t n_tokens) { + /* + * Short streamed prefill is latency-sensitive. Use the decode-style path + * by default for SSD streaming, while keeping a cold-only escape hatch for + * strict-vector tests that need canonical layer-major prefill semantics. + */ + if (start == 0) { + if (glm_graph_env_present( + "DS4_ROCM_DISABLE_STREAMING_COLD_DECODE_PREFILL", + "DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL")) { + return false; + } + } + return metal_graph_use_streaming_decode_prefill(g, weights, n_tokens); +} + +static bool metal_graph_prefill_decode_streaming_range( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + float *logits, + bool show_progress, + ds4_session_progress_fn progress, + void *progress_ud, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + ds4_session_cancel_fn cancel, + void *cancel_ud, + bool *cancelled) { + if (!metal_graph_use_streaming_decode_prefill(g, weights, n_tokens)) return false; + if (!prompt || start > (uint32_t)prompt->len || + n_tokens > (uint32_t)prompt->len - start) return false; + if (start == 0) { + ds4_gpu_stream_expert_cache_reset_route_hotness(); + } + + const bool profile = + glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", + "DS4_METAL_GRAPH_PREFILL_PROFILE"); + const double t0 = profile ? now_sec() : 0.0; + + /* + * `prefill_chunk` is not just UI progress: ds4_session_sync() wraps it to + * advance the live checkpoint, and ds4-server may save that checkpoint. + * Decode-style prefill only reads logits for the final token, so report one + * cacheable chunk at the end. `prefill_display` remains per-token UI only. + */ + if (progress) progress(progress_ud, "prefill_chunk", (int)start, prompt->len); + if (display_progress) { + display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); + } + + for (uint32_t i = 0; i < n_tokens; i++) { + if (cancel && cancel(cancel_ud)) { + if (cancelled) *cancelled = true; + return true; + } + const uint32_t pos = start + i; + const bool last = i + 1u == n_tokens; + float *token_logits = (last && logits) ? logits : NULL; + if (!metal_graph_eval_token_raw_swa(g, + model, + weights, + prompt->v[pos], + pos, + token_logits)) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after decode-style streaming prefill failure also failed\n"); + } + return false; + } + + if (last && progress && logits) { + progress(progress_ud, "prefill_chunk", (int)(pos + 1u), prompt->len); + } + if (display_progress) { + display_progress(display_progress_ud, "prefill_display", (int)(pos + 1u), prompt->len); + } + if (cancel && cancel(cancel_ud)) { + if (cancelled) *cancelled = true; + return true; + } + if (show_progress) { + fprintf(stderr, "ds4: gpu streaming prefill token %u/%u\r", + i + 1u, + n_tokens); + fflush(stderr); + } + } + if (show_progress) fputc('\n', stderr); + + if (profile) { + const double t1 = now_sec(); + fprintf(stderr, + "ds4: gpu decode-style streaming prefill start=%u tokens=%u total=%.3f ms\n", + start, + n_tokens, + (t1 - t0) * 1000.0); + } + return true; +} + +static bool metal_graph_capture_prefill_seed_router_selected( + ds4_gpu_graph *g, + uint32_t il, + uint32_t n_tokens) { + uint32_t k = metal_graph_streaming_prefill_cache_seed_k(g); + if (k == 0) return true; + if (k > n_tokens) k = n_tokens; + g->prefill_seed_tokens = k; + if (!g->prefill_seed_router_selected || !metal_graph_batch_router_selected(g) || + il >= DS4_N_LAYER || n_tokens == 0 || sizeof(int) != sizeof(int32_t)) { + return false; + } + + const uint64_t bytes = (uint64_t)k * DS4_N_EXPERT_USED * sizeof(int32_t); + const uint64_t src_off = (uint64_t)(n_tokens - k) * + DS4_N_EXPERT_USED * sizeof(int); + const uint64_t dst_off = (uint64_t)il * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * sizeof(int32_t); + return ds4_gpu_tensor_copy(g->prefill_seed_router_selected, + dst_off, + metal_graph_batch_router_selected(g), + src_off, + bytes) != 0; +} + +static bool metal_graph_seed_streaming_expert_cache_from_prefill( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; + if (!metal_graph_streaming_prefill_cache_seed_enabled(g)) return true; + if (!model || !weights || !g->prefill_seed_router_selected || seed_tokens == 0) { + return false; + } + + int32_t selected[DS4_MAX_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_MAX_EXPERT_USED]; + const uint64_t bytes = (uint64_t)DS4_N_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * sizeof(selected[0]); + if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, + 0, + selected, + bytes) == 0) { + return false; + } + + const bool profile = + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_PROFILE", + "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE"); + const double t0 = profile ? now_sec() : 0.0; + uint32_t seeded_layers = 0; + uint32_t seeded_rows = 0; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const ds4_layer_weights *layer = &weights->layer[il]; + if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { + continue; + } + + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: Metal prefill expert-cache seed byte size overflow at layer %u\n", il); + return false; + } + const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; + const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + for (uint32_t row = 0; row < seed_tokens; row++) { + const size_t sel_off = ((size_t)il * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS + + row) * DS4_N_EXPERT_USED; + if (ds4_gpu_stream_expert_cache_seed_selected( + &table, + selected + sel_off, + DS4_N_EXPERT_USED) == 0) { + return false; + } + seeded_rows++; + } + seeded_layers++; + } + if (profile) { + fprintf(stderr, + "ds4: Metal streaming prefill expert-cache seed k=%u layers=%u rows=%u time=%.3f ms\n", + seed_tokens, + seeded_layers, + seeded_rows, + (now_sec() - t0) * 1000.0); + } + return true; +} + +static bool metal_graph_seed_streaming_expert_cache_from_hotlist( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + if (!metal_graph_streaming_expert_hotlist_enabled(g)) return true; + if (!model || !weights) return false; + + uint32_t cache_budget = 0; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const ds4_layer_weights *layer = &weights->layer[il]; + if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { + continue; + } + + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: streaming expert hotlist budget byte size overflow at layer %u\n", il); + return false; + } + const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; + const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + cache_budget = ds4_gpu_stream_expert_cache_budget_for_expert_size( + gate_expert_bytes, + down_expert_bytes); + break; + } + if (cache_budget == 0) return true; + const uint32_t preload_count = + metal_graph_streaming_expert_preload_count(g, cache_budget); + if (preload_count == 0) return true; + const uint32_t current_count = + ds4_gpu_stream_expert_cache_current_count(); + const char *path = glm_graph_env_value("DS4_ROCM_STREAMING_EXPERT_HOTLIST", + "DS4_METAL_STREAMING_EXPERT_HOTLIST"); + const bool from_file = path && path[0]; + const bool refresh_builtin_glm = + !from_file && g_ds4_shape.variant == DS4_VARIANT_GLM52; + const bool profile = + glm_graph_env_present("DS4_ROCM_STREAMING_EXPERT_HOTLIST_PROFILE", + "DS4_METAL_STREAMING_EXPERT_HOTLIST_PROFILE"); + if (!from_file && !refresh_builtin_glm && current_count >= preload_count) { + if (profile) { + fprintf(stderr, + "ds4: streaming expert hotlist seed skipped preload=%u current=%u\n", + preload_count, + current_count); + } + return true; + } + + int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT]; + uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT]; + uint32_t counts[DS4_MAX_LAYER]; + bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT]; + memset(experts, 0, sizeof(experts)); + memset(priorities, 0, sizeof(priorities)); + memset(counts, 0, sizeof(counts)); + memset(seen, 0, sizeof(seen)); + + uint32_t loaded = 0; + if (from_file) { + if (!metal_graph_streaming_expert_hotlist_load_file(path, + preload_count, + experts, + priorities, + counts, + seen, + &loaded)) { + return false; + } + } else if (!metal_graph_streaming_expert_hotlist_load_default(preload_count, + experts, + priorities, + counts, + seen, + &loaded)) { + return false; + } + if (loaded == 0) return true; + + const double t0 = profile ? now_sec() : 0.0; + uint32_t seeded_layers = 0; + uint32_t seeded_experts = 0; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t n = counts[il]; + if (n == 0) continue; + const ds4_layer_weights *layer = &weights->layer[il]; + if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { + continue; + } + + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || + layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: streaming expert hotlist seed byte size overflow at layer %u\n", il); + return false; + } + const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; + const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + if (ds4_gpu_stream_expert_cache_seed_experts( + &table, + experts[il], + priorities[il], + n) == 0) { + return false; + } + seeded_layers++; + seeded_experts += n; + } + if (profile) { + const char *source_name = NULL; + if (from_file) { + source_name = path; + } else if (g_ds4_shape.variant == DS4_VARIANT_GLM52) { + source_name = "built-in-glm52"; + } else if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { + source_name = "built-in-flash"; + } else if (g_ds4_shape.variant == DS4_VARIANT_PRO) { + source_name = "built-in-pro"; + } else { + source_name = "built-in"; + } + fprintf(stderr, + "ds4: streaming expert hotlist seed source=%s preload=%u loaded=%u layers=%u experts=%u time=%.3f ms\n", + source_name, + preload_count, + loaded, + seeded_layers, + seeded_experts, + (now_sec() - t0) * 1000.0); + } + return true; +} + +typedef struct { + int id0; + int id1; + float value0; + float value1; + bool valid; + bool fast_attention; +} metal_graph_top2_result; + +/* Greedy verifier helper. Speculative decoding only needs the target model's + * top token after most accepted draft rows; the full vocabulary row is needed + * once, for the final committed state that normal sampling will continue from. + * Keeping intermediate rows device-resident avoids turning verification into a + * sequence of large CPU readbacks. */ +static bool metal_graph_eval_token_raw_swa_top( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + uint32_t pos, + int *top_id, + float *logits, + bool allow_split_top1, + metal_graph_top2_result *top2, + bool force_fast_attention) { + if (!top_id) return false; + if (top2) memset(top2, 0, sizeof(*top2)); + + const bool fast_attention = + allow_split_top1 && + logits == NULL && + (force_fast_attention || metal_graph_cuda_greedy_splitkv_requested()); + if (top2) top2->fast_attention = fast_attention; + const int old_fast_attention = + ds4_gpu_set_decode_fast_attention(fast_attention ? 1 : 0); + const bool profile = getenv("DS4_METAL_GRAPH_TOKEN_PROFILE") != NULL; + const double t0 = profile ? now_sec() : 0.0; + const bool split_top1 = + allow_split_top1 && + logits == NULL && + top2 == NULL && + g->cuda_tp_output && + metal_graph_cuda_greedy_split_top1_requested(); + if (split_top1) { + int output_tiers[DS4_MAX_GPUS] = {0}; + uint32_t output_ways = 0; + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, + token, pos, false, true); + if (ok) ok = metal_graph_encode_output_head_split_top1(g, + model, + weights, + weights->output->dim[1], + output_tiers, + &output_ways); + const double t_encoded = profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = profile ? now_sec() : 0.0; + if (ok) { + bool have_best = false; + uint32_t best_id = 0; + float best_value = 0.0f; + uint32_t cand_ids[DS4_MAX_GPUS] = {0}; + float cand_values[DS4_MAX_GPUS] = {0.0f}; + ok = output_ways <= DS4_MAX_GPUS && + ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], + 0, + cand_ids, + (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && + ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], + 0, + cand_values, + (uint64_t)output_ways * sizeof(cand_values[0])) != 0; + for (uint32_t i = 0; ok && i < output_ways; i++) { + const uint32_t cand_id = cand_ids[i]; + const float cand_value = cand_values[i]; + if (ok && + (!have_best || + cand_value > best_value || + (cand_value == best_value && cand_id < best_id))) { + have_best = true; + best_id = cand_id; + best_value = cand_value; + } + } + ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; + if (ok) *top_id = (int)best_id; + } + const double t_read = profile ? now_sec() : 0.0; + if (profile) { + fprintf(stderr, + "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=0 top=1 split_top1=1\n", + pos, + (t_encoded - t0) * 1000.0, + (t_done - t_encoded) * 1000.0, + (t_read - t_done) * 1000.0, + (t_read - t0) * 1000.0); + } + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after split-top graph eval failure also failed\n"); + } + } + (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); + return ok; + } + + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, + token, pos, true, true); + if (ok) { + ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), + metal_graph_logits(g), + DS4_N_VOCAB) != 0; + } + const double t_encoded = profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = profile ? now_sec() : 0.0; + if (ok && top2) { + uint32_t ids[2] = {0, 0}; + float values[2] = {0.0f, 0.0f}; + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), + 0, + ids, + sizeof(ids)) != 0 && + ds4_gpu_tensor_read(metal_graph_comp_mask(g), + 0, + values, + sizeof(values)) != 0; + if (ok && ids[0] <= (uint32_t)INT32_MAX && ids[1] <= (uint32_t)INT32_MAX) { + top2->id0 = (int)ids[0]; + top2->id1 = (int)ids[1]; + top2->value0 = values[0]; + top2->value1 = values[1]; + top2->valid = isfinite(values[0]) && isfinite(values[1]); + *top_id = top2->id0; + } else { + ok = false; + } + } else if (ok) { + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top_id, sizeof(*top_id)) != 0; + } + if (ok && logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + const double t_read = profile ? now_sec() : 0.0; + if (profile) { + fprintf(stderr, + "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d top=1 split_top1=0\n", + pos, + (t_encoded - t0) * 1000.0, + (t_done - t_encoded) * 1000.0, + (t_read - t_done) * 1000.0, + (t_read - t0) * 1000.0, + logits != NULL); + } + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after top-only graph eval failure also failed\n"); + } + } + (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); + return ok; +} + +static bool dspark_stage0_weights_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw) { + if (!g || !dw || dw->n_stages == 0 || dw->target_layer_count == 0 || + dw->target_layer_count != g->dspark_target_layer_count || + !g->dspark_target_hidden || !g->dspark_stage0_proj || + !g->dspark_main_x) { + return false; + } + + const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; + const ds4_tensor *main_proj = stage0->main_proj; + const ds4_tensor *main_norm = stage0->main_norm; + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + return main_proj && + main_norm && + dspark_tensor_type_matches(main_proj->type, DS4_DSPARK_LAYOUT_DENSE) && + main_norm->type == DS4_TENSOR_F32 && + main_proj->ndim == 2 && + main_proj->dim[0] == in_dim && + main_proj->dim[1] == DS4_N_EMBD && + main_norm->ndim == 1 && + main_norm->dim[0] == DS4_N_EMBD; +} + +static bool metal_graph_eval_dspark_stage0( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw) { + if (!g || !dspark_model || !dw || !dspark_stage0_weights_ready(g, dw)) { + return false; + } + + const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, + dspark_model, + stage0->main_proj, + in_dim, + DS4_N_EMBD, + g->dspark_target_hidden, + 1); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, + g->dspark_stage0_proj, + dspark_model->map, + dspark_model->size, + stage0->main_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + return ok; +} + +static bool dspark_stage0_batch_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw, + uint32_t n_tokens) { + if (!dspark_stage0_weights_ready(g, dw) || + n_tokens == 0 || + n_tokens > g->prefill_cap || + !g->dspark_target_hidden_batch || + !metal_graph_batch_ffn_cur(g) || + !metal_graph_batch_ffn_norm(g) || + !metal_graph_batch_cur_hc(g)) { + return false; + } + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + return ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) >= + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) >= + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_cur_hc(g)) >= + (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float) && + in_dim <= SIZE_MAX / sizeof(float); +} + +static bool metal_graph_pack_dspark_target_hidden_batch( + ds4_gpu_graph *g, + const ds4_dspark_weights *dw, + ds4_gpu_tensor *packed, + uint32_t n_tokens) { + if (!g || !dw || !packed || n_tokens == 0 || + n_tokens > g->prefill_cap || + dw->target_layer_count != g->dspark_target_layer_count) { + return false; + } + + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + const uint64_t packed_count = (uint64_t)n_tokens * in_dim; + if (packed_count == 0 || + packed_count > (uint64_t)SIZE_MAX / sizeof(float)) { + return false; + } + return ds4_gpu_pack_slot_rows_f32_tensor(packed, + g->dspark_target_hidden_batch, + n_tokens, + DS4_N_EMBD, + dw->target_layer_count, + g->prefill_cap) != 0; +} + +static bool metal_graph_eval_dspark_stage0_batch( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t n_tokens, + bool commands_open) { + if (!g || !dspark_model || !dw || + !dspark_stage0_batch_ready(g, dw, n_tokens)) { + return false; + } + + const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + const uint64_t packed_bytes = (uint64_t)n_tokens * in_dim * sizeof(float); + bool packed_owned = false; + ds4_gpu_tensor *packed = NULL; + if (g->dspark_stage0_packed && + ds4_gpu_tensor_bytes(g->dspark_stage0_packed) >= packed_bytes) { + packed = g->dspark_stage0_packed; + } else { + packed = ds4_gpu_tensor_alloc(packed_bytes); + packed_owned = true; + } + if (!packed) return false; + + bool ok = metal_graph_pack_dspark_target_hidden_batch(g, dw, packed, n_tokens); + if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = metal_graph_matmul_plain_tensor(metal_graph_batch_ffn_cur(g), + dspark_model, + stage0->main_proj, + in_dim, + DS4_N_EMBD, + packed, + n_tokens); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_ffn_norm(g), + metal_graph_batch_ffn_cur(g), + dspark_model->map, + dspark_model->size, + stage0->main_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; + if (!ok && !commands_open) (void)ds4_gpu_synchronize(); + if (packed_owned) ds4_gpu_tensor_free(packed); + return ok; +} + +static bool dspark_draft_block_ready( + const ds4_gpu_graph *g, + const ds4_weights *base_weights, + const ds4_dspark_weights *dw, + int token) { + if (!g || !base_weights || !dw || !base_weights->token_embd || + !g->dspark_draft_tokens || !g->dspark_draft_hc || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + g->dspark_block_size != dw->block_size || + !dw->has_noise_token_id) { + return false; + } + const uint32_t n_vocab = (uint32_t)base_weights->token_embd->dim[1]; + return token >= 0 && + (uint32_t)token < n_vocab && + dw->noise_token_id < n_vocab; +} + +static bool dspark_stage_input_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw) { + if (!g || !dw || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + g->dspark_block_size != dw->block_size || + !g->dspark_main_x || !g->dspark_draft_hc || + !g->dspark_target_hc || !g->dspark_stage_input_hc || + !g->dspark_position_ids) { + return false; + } + if (dw->block_size == UINT32_MAX) return false; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t rows = (uint64_t)dw->block_size + 1u; + return ds4_gpu_tensor_bytes(g->dspark_target_hc) >= + hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(g->dspark_stage_input_hc) >= + rows * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(g->dspark_position_ids) >= + rows * sizeof(int32_t); +} + +static bool dspark_stage_cache_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw) { + if (!g || !dw || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + g->dspark_cache_cap == 0 || + !metal_graph_dspark_cache_current_window_valid(g)) { + return false; + } + const uint64_t bytes = + (uint64_t)g->dspark_cache_cap * DS4_N_HEAD_DIM * sizeof(float); + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + if (!g->dspark_raw_cache[stage] || + ds4_gpu_tensor_bytes(g->dspark_raw_cache[stage]) < bytes) { + return false; + } + } + return true; +} + +static bool dspark_noncausal_attention_probe_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw) { + if (!g || !dw || + dw->n_stages == 0 || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + g->prefill_cap < dw->block_size + 1u || + !metal_graph_batch_q(g) || !metal_graph_batch_heads(g) || + !g->dspark_raw_cache[0]) { + return false; + } + const ds4_layer_weights *block = &dw->stage[0].block; + if (!block->attn_sinks) return false; + + const uint64_t rows = (uint64_t)dw->block_size + 1u; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + return ds4_gpu_tensor_bytes(metal_graph_batch_q(g)) >= + rows * q_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_heads(g)) >= + rows * q_dim * sizeof(float) && + ds4_gpu_tensor_bytes(g->dspark_raw_cache[0]) >= + rows * DS4_N_HEAD_DIM * sizeof(float); +} + +static bool metal_graph_probe_dspark_noncausal_attention( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw) { + if (!g || !dspark_model || !dspark_noncausal_attention_probe_ready(g, dw)) { + return false; + } + + const uint32_t rows = dw->block_size + 1u; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const ds4_layer_weights *block = &dw->stage[0].block; + bool ok = ds4_gpu_tensor_fill_f32(metal_graph_batch_q(g), + 0.0f, + (uint64_t)rows * q_dim) != 0 && + ds4_gpu_tensor_fill_f32(g->dspark_raw_cache[0], + 0.0f, + (uint64_t)rows * DS4_N_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + metal_graph_batch_heads(g), + dspark_model->map, + dspark_model->size, + block->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->dspark_raw_cache[0], + rows, + rows, + rows, + 0, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + return ok; +} + +static bool metal_graph_prepare_dspark_setup_block( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_dspark_weights *dw, + int token, + uint32_t pos) { + if (!g || !base_model || + !dspark_draft_block_ready(g, base_weights, dw, token) || + !dspark_stage_input_ready(g, dw)) { + return false; + } + if (pos > (uint32_t)INT32_MAX || + dw->block_size > (uint32_t)INT32_MAX || + pos > (uint32_t)INT32_MAX - dw->block_size) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t hc_bytes = hc_dim * sizeof(float); + int32_t positions[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; + positions[0] = (int32_t)pos; + for (uint32_t i = 0; i < dw->block_size; i++) { + positions[i + 1u] = (int32_t)(pos + i); + } + int32_t ids[DS4_DSPARK_MAX_BLOCK_SIZE]; + ids[0] = (int32_t)token; + for (uint32_t i = 1; i < dw->block_size; i++) { + ids[i] = (int32_t)dw->noise_token_id; + } + + bool ok = ds4_gpu_tensor_write(g->dspark_draft_tokens, + 0, + ids, + (uint64_t)dw->block_size * sizeof(ids[0])) != 0 && + ds4_gpu_tensor_write(g->dspark_position_ids, + 0, + positions, + ((uint64_t)dw->block_size + 1u) * + sizeof(positions[0])) != 0; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = ds4_gpu_embed_tokens_hc_tensor(g->dspark_draft_hc, + g->dspark_draft_tokens, + base_model->map, + base_model->size, + base_weights->token_embd->abs_offset, + (uint32_t)base_weights->token_embd->dim[1], + dw->block_size, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (ok) { + ok = ds4_gpu_repeat_hc_tensor(g->dspark_target_hc, + g->dspark_main_x, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, + 0, + g->dspark_target_hc, + 0, + hc_bytes) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, + hc_bytes, + g->dspark_draft_hc, + 0, + (uint64_t)dw->block_size * hc_bytes) != 0; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + return ok; +} + +static bool metal_graph_prepare_dspark_stage0_setup_block( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + int token, + uint32_t pos) { + if (!g || !base_model || !dspark_model || + !dspark_stage0_weights_ready(g, dw) || + !dspark_draft_block_ready(g, base_weights, dw, token) || + !dspark_stage_input_ready(g, dw)) { + return false; + } + if (pos > (uint32_t)INT32_MAX || + dw->block_size > (uint32_t)INT32_MAX || + pos > (uint32_t)INT32_MAX - dw->block_size) { + return false; + } + + const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t hc_bytes = hc_dim * sizeof(float); + int32_t positions[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; + positions[0] = (int32_t)pos; + for (uint32_t i = 0; i < dw->block_size; i++) { + positions[i + 1u] = (int32_t)(pos + i); + } + int32_t ids[DS4_DSPARK_MAX_BLOCK_SIZE]; + ids[0] = (int32_t)token; + for (uint32_t i = 1; i < dw->block_size; i++) { + ids[i] = (int32_t)dw->noise_token_id; + } + + /* DS4_DSPARK_PROP_PROFILE=1: break the setup block into phases to + * localize the TP-only prop_setup inflation (26ms vs 1.3ms single). */ + const bool prop_profile = getenv("DS4_DSPARK_PROP_PROFILE") != NULL; + const double pp_t0 = prop_profile ? now_sec() : 0.0; + bool ok = ds4_gpu_tensor_write(g->dspark_draft_tokens, + 0, + ids, + (uint64_t)dw->block_size * sizeof(ids[0])) != 0 && + ds4_gpu_tensor_write(g->dspark_position_ids, + 0, + positions, + ((uint64_t)dw->block_size + 1u) * + sizeof(positions[0])) != 0; + const double pp_t1 = prop_profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_begin_commands() != 0; + const double pp_t2 = prop_profile ? now_sec() : 0.0; + if (ok) { + ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, + dspark_model, + stage0->main_proj, + in_dim, + DS4_N_EMBD, + g->dspark_target_hidden, + 1); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, + g->dspark_stage0_proj, + dspark_model->map, + dspark_model->size, + stage0->main_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } + if (ok) { + ok = ds4_gpu_embed_tokens_hc_tensor(g->dspark_draft_hc, + g->dspark_draft_tokens, + base_model->map, + base_model->size, + base_weights->token_embd->abs_offset, + (uint32_t)base_weights->token_embd->dim[1], + dw->block_size, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (ok) { + ok = ds4_gpu_repeat_hc_tensor(g->dspark_target_hc, + g->dspark_main_x, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, + 0, + g->dspark_target_hc, + 0, + hc_bytes) != 0; + } + if (ok) { + ok = ds4_gpu_tensor_copy(g->dspark_stage_input_hc, + hc_bytes, + g->dspark_draft_hc, + 0, + (uint64_t)dw->block_size * hc_bytes) != 0; + } + const double pp_t3 = prop_profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + if (prop_profile) { + const double pp_t4 = now_sec(); + fprintf(stderr, + "ds4: DSpark prop-setup phases: writes=%.3fms begin=%.3fms " + "encode=%.3fms end/wait=%.3fms\n", + (pp_t1 - pp_t0) * 1000.0, + (pp_t2 - pp_t1) * 1000.0, + (pp_t3 - pp_t2) * 1000.0, + (pp_t4 - pp_t3) * 1000.0); + } + if (!ok) (void)ds4_gpu_synchronize(); + return ok; +} + +static bool dspark_stage_block_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw, + uint32_t stage) { + if (!g || !dw || + stage >= dw->n_stages || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + g->prefill_cap < dw->block_size + 1u || + !g->dspark_stage_output_hc || + !dspark_stage_input_ready(g, dw) || + !dspark_stage_cache_ready(g, dw)) { + return false; + } + + const ds4_layer_weights *l = &dw->stage[stage].block; + if (!l->hc_attn_fn || !l->hc_attn_scale || !l->hc_attn_base || + !l->attn_norm || !l->attn_q_a || !l->attn_q_a_norm || + !l->attn_q_b || !l->attn_kv || !l->attn_kv_a_norm || + !l->attn_sinks || !l->attn_output_a || !l->attn_output_b || + !l->hc_ffn_fn || !l->hc_ffn_scale || !l->hc_ffn_base || + !l->ffn_norm || !l->ffn_gate_inp || !l->ffn_exp_probs_b || + !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || + !l->ffn_gate_shexp || !l->ffn_up_shexp || !l->ffn_down_shexp) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t rows = (uint64_t)dw->block_size + 1u; + const uint64_t draft = dw->block_size; + const uint64_t out_low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; + const uint64_t group_dim = + (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); + + return + dspark_tensor_type_matches(l->hc_attn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && + l->hc_attn_scale->type == DS4_TENSOR_F32 && + l->hc_attn_base->type == DS4_TENSOR_F32 && + l->attn_norm->type == DS4_TENSOR_F32 && + dspark_tensor_type_matches(l->attn_q_a->type, DS4_DSPARK_LAYOUT_DENSE) && + l->attn_q_a_norm->type == DS4_TENSOR_F32 && + dspark_tensor_type_matches(l->attn_q_b->type, DS4_DSPARK_LAYOUT_DENSE) && + dspark_tensor_type_matches(l->attn_kv->type, DS4_DSPARK_LAYOUT_DENSE) && + l->attn_kv_a_norm->type == DS4_TENSOR_F32 && + l->attn_sinks->type == DS4_TENSOR_F32 && + l->attn_output_a->type == DS4_TENSOR_Q8_0 && + l->attn_output_b->type == DS4_TENSOR_Q8_0 && + dspark_tensor_type_matches(l->hc_ffn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && + l->hc_ffn_scale->type == DS4_TENSOR_F32 && + l->hc_ffn_base->type == DS4_TENSOR_F32 && + l->ffn_norm->type == DS4_TENSOR_F32 && + dspark_tensor_type_matches(l->ffn_gate_inp->type, DS4_DSPARK_LAYOUT_DENSE) && + l->ffn_exp_probs_b->type == DS4_TENSOR_F32 && + tensor_is_routed_expert_type(l->ffn_gate_exps->type) && + l->ffn_gate_exps->type == l->ffn_up_exps->type && + tensor_is_routed_expert_type(l->ffn_down_exps->type) && + l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && + l->hc_attn_fn->ndim == 2 && + l->hc_attn_fn->dim[0] == hc_dim && + l->hc_attn_fn->dim[1] == mix_hc && + l->attn_q_a->ndim == 2 && + l->attn_q_a->dim[0] == DS4_N_EMBD && + l->attn_q_a->dim[1] == DS4_N_LORA_Q && + l->attn_q_b->ndim == 2 && + l->attn_q_b->dim[0] == DS4_N_LORA_Q && + l->attn_q_b->dim[1] == q_dim && + l->attn_kv->ndim == 2 && + l->attn_kv->dim[0] == DS4_N_EMBD && + l->attn_kv->dim[1] == DS4_N_HEAD_DIM && + l->attn_output_a->ndim == 2 && + l->attn_output_a->dim[0] == group_dim && + l->attn_output_a->dim[1] == out_low_dim && + l->attn_output_b->ndim == 2 && + l->attn_output_b->dim[0] == out_low_dim && + l->attn_output_b->dim[1] == DS4_N_EMBD && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= rows * mix_hc * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= rows * mix_hc * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= rows * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_attn_cur(g)) >= rows * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_attn_norm(g)) >= rows * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_qr(g)) >= draft * DS4_N_LORA_Q * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_qr_norm(g)) >= draft * DS4_N_LORA_Q * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_q(g)) >= draft * q_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_kv_raw(g)) >= rows * DS4_N_HEAD_DIM * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_kv(g)) >= rows * DS4_N_HEAD_DIM * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_heads(g)) >= draft * q_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_attn_out(g)) >= draft * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_after_attn_hc(g)) >= draft * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_next_hc(g)) >= draft * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_prefill_tokens(g)) >= draft * sizeof(int32_t) && + ds4_gpu_tensor_bytes(g->dspark_stage_output_hc) >= draft * hc_dim * sizeof(float); +} + +static bool dspark_stage_target_cache_seed_ready( + const ds4_gpu_graph *g, + const ds4_dspark_weights *dw, + uint32_t stage, + uint32_t n_tokens) { + if (!g || !dw || + stage >= dw->n_stages || + n_tokens == 0 || + n_tokens > g->prefill_cap || + !dspark_stage_cache_ready(g, dw) || + !metal_graph_batch_cur_hc(g) || + !metal_graph_batch_hc_mix(g) || + !metal_graph_batch_hc_split(g) || + !metal_graph_batch_flat_hc(g) || + !metal_graph_batch_attn_cur(g) || + !metal_graph_batch_attn_norm(g) || + !metal_graph_batch_kv_raw(g) || + !metal_graph_batch_kv(g)) { + return false; + } + + const ds4_layer_weights *l = &dw->stage[stage].block; + if (!l->hc_attn_fn || !l->hc_attn_scale || !l->hc_attn_base || + !l->attn_norm || !l->attn_kv || !l->attn_kv_a_norm) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + return + dspark_tensor_type_matches(l->hc_attn_fn->type, DS4_DSPARK_LAYOUT_PLAIN) && + l->hc_attn_scale->type == DS4_TENSOR_F32 && + l->hc_attn_base->type == DS4_TENSOR_F32 && + l->attn_norm->type == DS4_TENSOR_F32 && + dspark_tensor_type_matches(l->attn_kv->type, DS4_DSPARK_LAYOUT_DENSE) && + l->attn_kv_a_norm->type == DS4_TENSOR_F32 && + l->hc_attn_fn->ndim == 2 && + l->hc_attn_fn->dim[0] == hc_dim && + l->hc_attn_fn->dim[1] == mix_hc && + l->attn_kv->ndim == 2 && + l->attn_kv->dim[0] == DS4_N_EMBD && + l->attn_kv->dim[1] == DS4_N_HEAD_DIM && + ds4_gpu_tensor_bytes(metal_graph_batch_cur_hc(g)) >= + (uint64_t)n_tokens * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= + (uint64_t)n_tokens * mix_hc * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= + (uint64_t)n_tokens * mix_hc * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= + (uint64_t)n_tokens * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_attn_cur(g)) >= + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_attn_norm(g)) >= + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_kv_raw(g)) >= + (uint64_t)n_tokens * DS4_N_HEAD_DIM * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_kv(g)) >= + (uint64_t)n_tokens * DS4_N_HEAD_DIM * sizeof(float); +} + +static bool metal_graph_seed_dspark_stage_target_cache( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t stage, + uint32_t pos0, + uint32_t n_tokens, + bool commands_open) { + if (!g || !dspark_model || !dw || + !dspark_stage_target_cache_seed_ready(g, dw, stage, n_tokens) || + n_tokens > g->dspark_cache_cap) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const ds4_layer_weights *block = &dw->stage[stage].block; + const bool fuse_hc_norm = DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + metal_graph_enable_batch_hc_norm_fusion(); + + ds4_gpu_tensor *hc_mix_view = + ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), + 0, + (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *hc_split_view = + ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + 0, + (uint64_t)n_tokens * mix_hc * sizeof(float)); + ds4_gpu_tensor *attn_cur_view = + ds4_gpu_tensor_view(metal_graph_batch_attn_cur(g), + 0, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + bool ok = hc_mix_view && hc_split_view && attn_cur_view; + + const float freq_base = DS4_ROPE_FREQ_BASE; + const float freq_scale = 1.0f; + const float ext_factor = 0.0f; + const float attn_factor = 1.0f; + + if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), + metal_graph_batch_cur_hc(g), + (uint32_t)hc_dim, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(hc_mix_view, + dspark_model, + block->hc_attn_fn, + hc_dim, + mix_hc, + metal_graph_batch_flat_hc(g), + n_tokens); + if (fuse_hc_norm) { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, + metal_graph_batch_attn_norm(g), + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + dspark_model->map, + dspark_model->size, + block->hc_attn_scale->abs_offset, + block->hc_attn_base->abs_offset, + block->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + dspark_model->map, + dspark_model->size, + block->hc_attn_scale->abs_offset, + block->hc_attn_base->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), + metal_graph_batch_attn_cur(g), + dspark_model->map, + dspark_model->size, + block->attn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_kv_raw(g), + dspark_model, + block->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + n_tokens); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + dspark_model->map, + dspark_model->size, + block->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), + n_tokens, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), + n_tokens, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; + if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor(g->dspark_raw_cache[stage], + metal_graph_batch_kv(g), + g->dspark_cache_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; + + ds4_gpu_tensor_free(attn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); + if (!ok && !commands_open) (void)ds4_gpu_synchronize(); + return ok; +} + +static bool metal_graph_seed_dspark_initial_cache_from_prefill( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t batch_start, + uint32_t n_tokens, + uint32_t *seeded_rows) { + if (seeded_rows) *seeded_rows = 0; + if (!g || !dspark_model || !dw || + n_tokens == 0 || + n_tokens > g->prefill_cap || + n_tokens > g->dspark_cache_cap || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + !dspark_stage0_batch_ready(g, dw, n_tokens) || + !dspark_stage_cache_ready(g, dw)) { + return false; + } + + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = metal_graph_eval_dspark_stage0_batch(g, + dspark_model, + dw, + n_tokens, + true); + } + if (ok) { + ok = ds4_gpu_repeat_hc_rows_tensor(metal_graph_batch_cur_hc(g), + metal_graph_batch_ffn_norm(g), + n_tokens, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { + ok = metal_graph_seed_dspark_stage_target_cache(g, + dspark_model, + dw, + stage, + batch_start, + n_tokens, + true); + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) { + (void)ds4_gpu_synchronize(); + return false; + } + if (!metal_graph_dspark_cache_set_window(g, batch_start, n_tokens)) { + return false; + } + if (seeded_rows) *seeded_rows = n_tokens; + return true; +} + +static bool metal_graph_encode_dspark_next_stage_draft_input_from( + ds4_gpu_graph *g, + const ds4_dspark_weights *dw, + const ds4_gpu_tensor *draft_hc) { + if (!g || !dw || !dspark_stage_input_ready(g, dw) || + !draft_hc) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t hc_bytes = hc_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(draft_hc) < + (uint64_t)dw->block_size * hc_bytes) { + return false; + } + + return ds4_gpu_tensor_copy(g->dspark_stage_input_hc, + hc_bytes, + draft_hc, + 0, + (uint64_t)dw->block_size * hc_bytes) != 0; +} + +static bool metal_graph_profile_layer_env_match(const char *env_name, uint32_t il) { + const char *layer_env = getenv(env_name); + if (!layer_env || !layer_env[0]) return true; + + char *end = NULL; + const unsigned long layer = strtoul(layer_env, &end, 10); + return end != layer_env && + *end == '\0' && + layer <= UINT32_MAX && + (uint32_t)layer == il; +} + +static bool metal_graph_dspark_stage_profile_enabled(uint32_t stage) { + return getenv("DS4_DSPARK_STAGE_PROFILE") != NULL && + metal_graph_profile_layer_env_match("DS4_DSPARK_STAGE_PROFILE_STAGE", + stage); +} + +static bool metal_graph_dspark_stage_profile_boundary( + const char *part, + uint32_t stage, + uint32_t pos, + uint32_t rows, + double *stage_t0) { + if (ds4_gpu_end_commands() == 0) return false; + const double now = now_sec(); + fprintf(stderr, + "ds4: DSpark stage profile stage=%u pos=%u rows=%u %s=%.3f ms\n", + stage, + pos, + rows, + part, + (now - *stage_t0) * 1000.0); + *stage_t0 = now; + return ds4_gpu_begin_commands() != 0; +} + +static bool metal_graph_eval_dspark_stage_block( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t stage, + uint32_t pos, + uint32_t support_len, + uint32_t raw_start, + bool prepare_next_stage_input, + bool commands_open) { + if (!g || !dspark_model || !dw || + !dspark_stage_block_ready(g, dw, stage)) { + return false; + } + + const uint32_t draft = dw->block_size; + const uint32_t rows = draft + 1u; + if (support_len > g->dspark_cache_cap || + rows > g->dspark_cache_cap - support_len || + (support_len != 0 && raw_start >= g->dspark_cache_cap)) { + return false; + } + const uint32_t visible_rows = support_len + rows; + const uint32_t attention_raw_start = + support_len ? raw_start : (pos % g->dspark_cache_cap); + const uint32_t append_pos = support_len ? + (uint32_t)(((uint64_t)raw_start + support_len) % + g->dspark_cache_cap) : + attention_raw_start; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t group_dim = + (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); + const ds4_layer_weights *block = &dw->stage[stage].block; + const bool fuse_hc_norm = DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + metal_graph_enable_batch_hc_norm_fusion(); + + ds4_gpu_tensor *hc_mix_view = + ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), + 0, + (uint64_t)rows * mix_hc * sizeof(float)); + ds4_gpu_tensor *hc_split_view = + ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + 0, + (uint64_t)rows * mix_hc * sizeof(float)); + ds4_gpu_tensor *attn_cur_view = + ds4_gpu_tensor_view(metal_graph_batch_attn_cur(g), + 0, + (uint64_t)rows * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *draft_attn_norm_view = + ds4_gpu_tensor_view(metal_graph_batch_attn_norm(g), + (uint64_t)DS4_N_EMBD * sizeof(float), + (uint64_t)draft * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *draft_hc_view = + ds4_gpu_tensor_view(g->dspark_stage_input_hc, + hc_dim * sizeof(float), + (uint64_t)draft * hc_dim * sizeof(float)); + ds4_gpu_tensor *draft_hc_split_view = + ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + mix_hc * sizeof(float), + (uint64_t)draft * mix_hc * sizeof(float)); + ds4_gpu_tensor *kv_target_view = + ds4_gpu_tensor_view(metal_graph_batch_kv(g), + 0, + (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + ds4_gpu_tensor *kv_draft_view = + ds4_gpu_tensor_view(metal_graph_batch_kv(g), + (uint64_t)DS4_N_HEAD_DIM * sizeof(float), + (uint64_t)draft * DS4_N_HEAD_DIM * sizeof(float)); + ds4_gpu_tensor *after_attn_hc_view = + ds4_gpu_tensor_view(metal_graph_batch_after_attn_hc(g), + 0, + (uint64_t)draft * hc_dim * sizeof(float)); + + bool ok = hc_mix_view && hc_split_view && attn_cur_view && + draft_attn_norm_view && draft_hc_view && + draft_hc_split_view && kv_target_view && kv_draft_view && + after_attn_hc_view; + const bool saved_streaming = g->ssd_streaming; + g->ssd_streaming = false; + + const float freq_base = DS4_ROPE_FREQ_BASE; + const float freq_scale = 1.0f; + const float ext_factor = 0.0f; + const float attn_factor = 1.0f; + const bool stage_profile = + metal_graph_dspark_stage_profile_enabled(stage); + double stage_t0 = stage_profile ? now_sec() : 0.0; +#define DS4_DSPARK_PROFILE_STAGE(part_) do { \ + if (ok && stage_profile) { \ + ok = metal_graph_dspark_stage_profile_boundary((part_), \ + stage, \ + pos, \ + rows, \ + &stage_t0); \ + } \ + } while (0) + + if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), + g->dspark_stage_input_hc, + (uint32_t)hc_dim, + rows, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(hc_mix_view, + dspark_model, + block->hc_attn_fn, + hc_dim, + mix_hc, + metal_graph_batch_flat_hc(g), + rows); + DS4_DSPARK_PROFILE_STAGE("attn_hc_pre"); + if (fuse_hc_norm) { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, + metal_graph_batch_attn_norm(g), + hc_split_view, + hc_mix_view, + g->dspark_stage_input_hc, + dspark_model->map, + dspark_model->size, + block->hc_attn_scale->abs_offset, + block->hc_attn_base->abs_offset, + block->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, + hc_split_view, + hc_mix_view, + g->dspark_stage_input_hc, + dspark_model->map, + dspark_model->size, + block->hc_attn_scale->abs_offset, + block->hc_attn_base->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), + metal_graph_batch_attn_cur(g), + dspark_model->map, + dspark_model->size, + block->attn_norm->abs_offset, + DS4_N_EMBD, + rows, + DS4_RMS_EPS) != 0; + } + DS4_DSPARK_PROFILE_STAGE("attn_norm"); + + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_qr(g), + dspark_model, + block->attn_q_a, + DS4_N_EMBD, + DS4_N_LORA_Q, + draft_attn_norm_view, + draft); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_qr_norm(g), + metal_graph_batch_qr(g), + dspark_model->map, + dspark_model->size, + block->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + draft, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_q(g), + dspark_model, + block->attn_q_b, + DS4_N_LORA_Q, + q_dim, + metal_graph_batch_qr_norm(g), + draft); + if (ok) ok = ds4_gpu_head_rms_norm_tensor(metal_graph_batch_q(g), + draft, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_q(g), + draft, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_DSPARK_PROFILE_STAGE("q_path"); + + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_kv_raw(g), + dspark_model, + block->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + rows); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + dspark_model->map, + dspark_model->size, + block->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + rows, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(kv_target_view, + 1, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(kv_draft_view, + draft, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), + rows, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; + if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor(g->dspark_raw_cache[stage], + metal_graph_batch_kv(g), + g->dspark_cache_cap, + append_pos, + rows, + DS4_N_HEAD_DIM) != 0; + DS4_DSPARK_PROFILE_STAGE("kv_path"); + + if (ok) ok = ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + metal_graph_batch_heads(g), + dspark_model->map, + dspark_model->size, + block->attn_sinks->abs_offset, + metal_graph_batch_q(g), + g->dspark_raw_cache[stage], + draft, + visible_rows, + g->dspark_cache_cap, + attention_raw_start, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_heads(g), + draft, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_DSPARK_PROFILE_STAGE("attention"); + if (ok) ok = ds4_gpu_attention_output_q8_batch_tensor( + metal_graph_batch_attn_out(g), + metal_graph_batch_attn_low(g), + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + dspark_model->map, + dspark_model->size, + block->attn_output_a->abs_offset, + block->attn_output_b->abs_offset, + group_dim, + DS4_N_LORA_O, + DS4_N_OUT_GROUP, + DS4_N_EMBD, + metal_graph_batch_heads(g), + draft) != 0; + if (ok) ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, + metal_graph_batch_attn_out(g), + draft_hc_view, + draft_hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + DS4_DSPARK_PROFILE_STAGE("attn_output_hc"); + + if (ok) ok = metal_graph_encode_layer_ffn_batch(g, + dspark_model, + block, + stage, + pos, + draft, + NULL, + 0); + DS4_DSPARK_PROFILE_STAGE("ffn"); + if (ok && + !prepare_next_stage_input && + getenv("DS4_DSPARK_DISABLE_FINAL_OUTPUT_ALIAS") != NULL) { + ok = ds4_gpu_tensor_copy(g->dspark_stage_output_hc, + 0, + metal_graph_batch_next_hc(g), + 0, + (uint64_t)draft * hc_dim * sizeof(float)) != 0; + } + DS4_DSPARK_PROFILE_STAGE("copy_output"); + if (ok && prepare_next_stage_input) { + ok = metal_graph_encode_dspark_next_stage_draft_input_from( + g, dw, metal_graph_batch_next_hc(g)); + } + DS4_DSPARK_PROFILE_STAGE("next_input"); + if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; + g->ssd_streaming = saved_streaming; + + ds4_gpu_tensor_free(after_attn_hc_view); + ds4_gpu_tensor_free(kv_draft_view); + ds4_gpu_tensor_free(kv_target_view); + ds4_gpu_tensor_free(draft_hc_split_view); + ds4_gpu_tensor_free(draft_hc_view); + ds4_gpu_tensor_free(draft_attn_norm_view); + ds4_gpu_tensor_free(attn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); + if (!ok && !commands_open) (void)ds4_gpu_synchronize(); +#undef DS4_DSPARK_PROFILE_STAGE + return ok; +} + +static bool metal_graph_eval_dspark_stage_chain( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t pos, + uint32_t *completed_stages, + uint32_t *cache_start_out, + uint32_t *cache_rows_out) { + if (completed_stages) *completed_stages = 0; + if (cache_start_out) *cache_start_out = 0; + if (cache_rows_out) *cache_rows_out = 0; + if (!g || !dspark_model || !dw || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + !dspark_stage_input_ready(g, dw) || + !dspark_stage_cache_ready(g, dw) || + !metal_graph_prefill_tokens(g) || + !g->dspark_draft_tokens) { + return false; + } + + const uint32_t rows = dw->block_size + 1u; + const uint32_t support_len = g->dspark_cache_len; + const uint32_t raw_start = support_len ? g->dspark_cache_start : 0; + if (support_len > g->dspark_cache_cap || + rows > g->dspark_cache_cap - support_len || + (support_len != 0 && raw_start >= g->dspark_cache_cap) || + !metal_graph_dspark_cache_ends_at(g, pos)) { + return false; + } + if (cache_start_out) { + *cache_start_out = support_len ? raw_start : + (pos % g->dspark_cache_cap); + } + if (cache_rows_out) *cache_rows_out = support_len + rows; + + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + if (!dspark_stage_block_ready(g, dw, stage)) return false; + } + + /* The support model runs only on the coordinator. Its generic layer + * helpers share the base graph object, so temporarily disarm TP or they + * would encode expert gates that the worker can never reach. The base + * model's later verification restores and uses the normal 50/50 split. */ + const uint32_t saved_tp_world = g->tp_world; + const uint32_t saved_tp_batch_rows = g->tp_batch_rows; + g->tp_world = 0; + g->tp_batch_rows = 0; + const bool suspended_expert_sharding = saved_tp_world == 2; + if (suspended_expert_sharding) { + ds4_gpu_tp_suspend_expert_sharding(1); + } + + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) { + ok = ds4_gpu_tensor_copy(metal_graph_prefill_tokens(g), + 0, + g->dspark_draft_tokens, + 0, + (uint64_t)dw->block_size * sizeof(int32_t)) != 0; + } + for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { + const bool stage_ok = + metal_graph_eval_dspark_stage_block(g, + dspark_model, + dw, + stage, + pos, + support_len, + raw_start, + stage + 1u < dw->n_stages, + true); + if (!stage_ok) { + ok = false; + break; + } + if (completed_stages) *completed_stages = stage + 1u; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (suspended_expert_sharding) { + ds4_gpu_tp_suspend_expert_sharding(0); + } + g->tp_world = saved_tp_world; + g->tp_batch_rows = saved_tp_batch_rows; + if (!ok) { + (void)ds4_gpu_synchronize(); + return false; + } + return true; +} + +/* Keep the support KV ring aligned while the scheduler skips proposals. */ +static bool metal_graph_dspark_ring_maintain( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + uint32_t pos) { + if (!g || !dspark_model || !dw || + !g->dspark_capture_valid || + g->dspark_cache_len == 0 || + !metal_graph_dspark_cache_ends_at(g, pos) || + !dspark_stage0_weights_ready(g, dw) || + !dspark_stage_cache_ready(g, dw) || + !metal_graph_batch_kv_raw(g) || !metal_graph_batch_kv(g)) { + return false; + } + for (uint32_t stage = 0; stage < dw->n_stages; stage++) { + if (!dspark_stage_block_ready(g, dw, stage)) return false; + } + + const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; + const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; + ds4_gpu_tensor *kv_raw_view = + ds4_gpu_tensor_view(metal_graph_batch_kv_raw(g), + 0, + (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + ds4_gpu_tensor *kv_view = + ds4_gpu_tensor_view(metal_graph_batch_kv(g), + 0, + (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + bool ok = kv_raw_view && kv_view && ds4_gpu_begin_commands() != 0; + if (ok) { + ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, + dspark_model, + stage0->main_proj, + in_dim, + DS4_N_EMBD, + g->dspark_target_hidden, + 1); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, + g->dspark_stage0_proj, + dspark_model->map, + dspark_model->size, + stage0->main_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } + for (uint32_t stage = 0; ok && stage < dw->n_stages; stage++) { + const ds4_layer_weights *block = &dw->stage[stage].block; + ok = metal_graph_matmul_plain_tensor(kv_raw_view, + dspark_model, + block->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + g->dspark_main_x, + 1); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor( + kv_view, + kv_raw_view, + dspark_model->map, + dspark_model->size, + block->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + 1, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_rope_tail_tensor(kv_view, + 1, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + false, + DS4_ROPE_FREQ_BASE, + 1.0f, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(kv_view, + 1, + DS4_N_HEAD_DIM, + DS4_N_ROT) != 0; + if (ok) ok = ds4_gpu_store_raw_kv_batch_tensor( + g->dspark_raw_cache[stage], + kv_view, + g->dspark_cache_cap, + pos, + 1, + DS4_N_HEAD_DIM) != 0; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + ds4_gpu_tensor_free(kv_view); + ds4_gpu_tensor_free(kv_raw_view); + if (ok) (void)metal_graph_dspark_cache_claim_appended_row(g, pos); + return ok; +} + +static ds4_gpu_tensor *metal_graph_dspark_final_output_hc(const ds4_gpu_graph *g) { + if (!g) return NULL; + if (getenv("DS4_DSPARK_DISABLE_FINAL_OUTPUT_ALIAS") == NULL && + metal_graph_batch_next_hc(g)) { + return metal_graph_batch_next_hc(g); + } + return g->dspark_stage_output_hc; +} + +static bool dspark_final_head_ready( + const ds4_gpu_graph *g, + const ds4_weights *base_weights, + const ds4_dspark_weights *dw) { + if (!g || !base_weights || !dw || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + !base_weights->output || + !metal_graph_dspark_final_output_hc(g) || + !metal_graph_batch_hc_mix(g) || + !metal_graph_batch_hc_split(g) || + !metal_graph_batch_flat_hc(g) || + !metal_graph_batch_ffn_cur(g) || + !metal_graph_batch_ffn_norm(g) || + !g->spec_logits) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t draft = dw->block_size; + const uint64_t vocab_dim = base_weights->output->dim[1]; + if (!final->norm || + !final->hc_head_base || + !final->hc_head_fn || + !final->hc_head_scale || + final->norm->type != DS4_TENSOR_F32 || + final->hc_head_base->type != DS4_TENSOR_F32 || + !dspark_tensor_type_matches(final->hc_head_fn->type, + DS4_DSPARK_LAYOUT_PLAIN) || + final->hc_head_scale->type != DS4_TENSOR_F32 || + !tensor_type_is_dense_quant(base_weights->output->type)) { + return false; + } + + return final->norm->ndim == 1 && + final->norm->dim[0] == DS4_N_EMBD && + final->hc_head_base->ndim == 1 && + final->hc_head_base->dim[0] == DS4_N_HC && + final->hc_head_fn->ndim == 2 && + final->hc_head_fn->dim[0] == hc_dim && + final->hc_head_fn->dim[1] == DS4_N_HC && + final->hc_head_scale->ndim == 1 && + final->hc_head_scale->dim[0] == 1 && + base_weights->output->ndim == 2 && + base_weights->output->dim[0] == DS4_N_EMBD && + vocab_dim == DS4_N_VOCAB && + ds4_gpu_tensor_bytes( + metal_graph_dspark_final_output_hc(g)) >= + draft * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) >= + draft * hc_dim * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) >= + draft * DS4_N_HC * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) >= + draft * DS4_N_HC * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) >= + draft * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) >= + draft * DS4_N_EMBD * sizeof(float) && + ds4_gpu_tensor_bytes(g->spec_logits) >= + draft * vocab_dim * sizeof(float); +} + +static bool metal_graph_eval_dspark_base_logits( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw) { + if (!g || !base_model || !base_weights || !dspark_model || !dw || + !dspark_final_head_ready(g, base_weights, dw)) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + const uint32_t draft = dw->block_size; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t vocab_dim = base_weights->output->dim[1]; + ds4_gpu_tensor *stage_output_hc = metal_graph_dspark_final_output_hc(g); + ds4_gpu_tensor *output_pre = + ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), + 0, + (uint64_t)draft * DS4_N_HC * sizeof(float)); + ds4_gpu_tensor *output_weights = + ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + 0, + (uint64_t)draft * DS4_N_HC * sizeof(float)); + ds4_gpu_tensor *output_embd = + ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), + 0, + (uint64_t)draft * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *output_norm = + ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), + 0, + (uint64_t)draft * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *logits = + ds4_gpu_tensor_view(g->spec_logits, + 0, + (uint64_t)draft * vocab_dim * sizeof(float)); + + bool ok = stage_output_hc && output_pre && output_weights && output_embd && + output_norm && logits; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), + stage_output_hc, + (uint32_t)hc_dim, + draft, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(output_pre, + dspark_model, + final->hc_head_fn, + hc_dim, + DS4_N_HC, + metal_graph_batch_flat_hc(g), + draft); + if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, + output_pre, + dspark_model->map, + dspark_model->size, + final->hc_head_scale->abs_offset, + final->hc_head_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, + stage_output_hc, + output_weights, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, + output_embd, + dspark_model->map, + dspark_model->size, + final->norm->abs_offset, + DS4_N_EMBD, + draft, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(logits, + base_model, + base_weights->output, + DS4_N_EMBD, + vocab_dim, + output_norm, + draft); + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + + ds4_gpu_tensor_free(logits); + ds4_gpu_tensor_free(output_norm); + ds4_gpu_tensor_free(output_embd); + ds4_gpu_tensor_free(output_weights); + ds4_gpu_tensor_free(output_pre); + return ok; +} + +static bool metal_graph_eval_dspark_final_hidden( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw) { + if (!g || !dspark_model || !dw || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + !metal_graph_dspark_final_output_hc(g) || + !metal_graph_batch_hc_mix(g) || + !metal_graph_batch_hc_split(g) || + !metal_graph_batch_flat_hc(g) || + !metal_graph_batch_ffn_cur(g) || + !metal_graph_batch_ffn_norm(g)) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + const uint32_t draft = dw->block_size; + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + if (!final->norm || + !final->hc_head_base || + !final->hc_head_fn || + !final->hc_head_scale || + final->norm->type != DS4_TENSOR_F32 || + final->hc_head_base->type != DS4_TENSOR_F32 || + !dspark_tensor_type_matches(final->hc_head_fn->type, + DS4_DSPARK_LAYOUT_PLAIN) || + final->hc_head_scale->type != DS4_TENSOR_F32 || + final->norm->ndim != 1 || + final->norm->dim[0] != DS4_N_EMBD || + final->hc_head_base->ndim != 1 || + final->hc_head_base->dim[0] != DS4_N_HC || + final->hc_head_fn->ndim != 2 || + final->hc_head_fn->dim[0] != hc_dim || + final->hc_head_fn->dim[1] != DS4_N_HC || + final->hc_head_scale->ndim != 1 || + final->hc_head_scale->dim[0] != 1 || + ds4_gpu_tensor_bytes(metal_graph_dspark_final_output_hc(g)) < + (uint64_t)draft * hc_dim * sizeof(float) || + ds4_gpu_tensor_bytes(metal_graph_batch_flat_hc(g)) < + (uint64_t)draft * hc_dim * sizeof(float) || + ds4_gpu_tensor_bytes(metal_graph_batch_hc_mix(g)) < + (uint64_t)draft * DS4_N_HC * sizeof(float) || + ds4_gpu_tensor_bytes(metal_graph_batch_hc_split(g)) < + (uint64_t)draft * DS4_N_HC * sizeof(float) || + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) < + (uint64_t)draft * DS4_N_EMBD * sizeof(float) || + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < + (uint64_t)draft * DS4_N_EMBD * sizeof(float)) { + return false; + } + + ds4_gpu_tensor *output_pre = + ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), + 0, + (uint64_t)draft * DS4_N_HC * sizeof(float)); + ds4_gpu_tensor *output_weights = + ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), + 0, + (uint64_t)draft * DS4_N_HC * sizeof(float)); + ds4_gpu_tensor *output_embd = + ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), + 0, + (uint64_t)draft * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *output_norm = + ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), + 0, + (uint64_t)draft * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *stage_output_hc = metal_graph_dspark_final_output_hc(g); + + bool ok = stage_output_hc && output_pre && output_weights && + output_embd && output_norm; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), + stage_output_hc, + (uint32_t)hc_dim, + draft, + DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(output_pre, + dspark_model, + final->hc_head_fn, + hc_dim, + DS4_N_HC, + metal_graph_batch_flat_hc(g), + draft); + if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, + output_pre, + dspark_model->map, + dspark_model->size, + final->hc_head_scale->abs_offset, + final->hc_head_base->abs_offset, + DS4_N_HC, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, + stage_output_hc, + output_weights, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, + output_embd, + dspark_model->map, + dspark_model->size, + final->norm->abs_offset, + DS4_N_EMBD, + draft, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + + ds4_gpu_tensor_free(output_norm); + ds4_gpu_tensor_free(output_embd); + ds4_gpu_tensor_free(output_weights); + ds4_gpu_tensor_free(output_pre); + return ok; +} + +static bool metal_graph_eval_dspark_base_logits_from_hidden( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_dspark_weights *dw) { + if (!g || !base_model || !base_weights || !dw || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + !base_weights->output || + !tensor_type_is_dense_quant(base_weights->output->type) || + base_weights->output->ndim != 2 || + base_weights->output->dim[0] != DS4_N_EMBD || + base_weights->output->dim[1] != DS4_N_VOCAB || + !metal_graph_batch_ffn_norm(g) || + !g->spec_logits || + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < + (uint64_t)dw->block_size * DS4_N_EMBD * sizeof(float) || + ds4_gpu_tensor_bytes(g->spec_logits) < + (uint64_t)dw->block_size * DS4_N_VOCAB * sizeof(float)) { + return false; + } + + ds4_gpu_tensor *output_norm = + ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), + 0, + (uint64_t)dw->block_size * + DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *logits = + ds4_gpu_tensor_view(g->spec_logits, + 0, + (uint64_t)dw->block_size * + DS4_N_VOCAB * sizeof(float)); + bool ok = output_norm && logits; + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(logits, + base_model, + base_weights->output, + DS4_N_EMBD, + DS4_N_VOCAB, + output_norm, + dw->block_size); + if (ok) ok = ds4_gpu_end_commands() != 0; + if (!ok) (void)ds4_gpu_synchronize(); + + ds4_gpu_tensor_free(logits); + ds4_gpu_tensor_free(output_norm); + return ok; +} + +static bool dspark_markov_probe_ready( + const ds4_dspark_weights *dw) { + if (!dw || + dw->n_stages == 0 || + dw->n_stages > DS4_DSPARK_MAX_STAGES || + dw->block_size == 0 || + dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || + dw->markov_rank == 0) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + if (!final->markov_w1 || + !final->markov_w2 || + !dspark_tensor_type_matches(final->markov_w1->type, + DS4_DSPARK_LAYOUT_DENSE) || + !dspark_tensor_type_matches(final->markov_w2->type, + DS4_DSPARK_LAYOUT_DENSE)) { + return false; + } + + return final->markov_w1->ndim == 2 && + final->markov_w1->dim[0] == dw->markov_rank && + final->markov_w1->dim[1] == DS4_N_VOCAB && + final->markov_w2->ndim == 2 && + final->markov_w2->dim[0] == dw->markov_rank && + final->markov_w2->dim[1] == DS4_N_VOCAB; +} + +static bool dspark_dense_row_to_f32( + float *out, + const ds4_model *model, + const ds4_tensor *t, + uint32_t row) { + if (!out || !model || !t || t->ndim != 2 || row >= t->dim[1]) { + return false; + } + + const uint64_t width = t->dim[0]; + if (t->type == DS4_TENSOR_F32) { + const float *base = tensor_data(model, t); + memcpy(out, base + (uint64_t)row * width, width * sizeof(out[0])); + return true; + } + if (t->type == DS4_TENSOR_F16) { + const uint16_t *base = tensor_data(model, t); + const uint16_t *src = base + (uint64_t)row * width; + for (uint64_t i = 0; i < width; i++) out[i] = f16_to_f32(src[i]); + return true; + } + if (t->type == DS4_TENSOR_Q8_0) { + const uint64_t blocks = (width + 31u) / 32u; + const uint8_t *src = + (const uint8_t *)tensor_data(model, t) + + (uint64_t)row * blocks * 34u; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t scale_bits; + memcpy(&scale_bits, src + b * 34u, sizeof(scale_bits)); + const float scale = f16_to_f32(scale_bits); + const int8_t *qs = (const int8_t *)(src + b * 34u + 2u); + const uint64_t i0 = b * 32u; + const uint64_t n = width - i0 < 32u ? width - i0 : 32u; + for (uint64_t i = 0; i < n; i++) { + out[i0 + i] = scale * (float)qs[i]; + } + } + return true; + } + return false; +} + +static uint32_t dspark_argmax_f32(const float *x, uint32_t n) { + uint32_t best = 0; + float best_v = x[0]; + for (uint32_t i = 1; i < n; i++) { + if (x[i] > best_v) { + best_v = x[i]; + best = i; + } + } + return best; +} + +typedef struct { + const uint8_t *data; + const int8_t *xq; + const float *xscale; + const float *logits; + uint64_t in_dim; + uint64_t blocks; + uint64_t rows_per_slot; + uint32_t best_idx[DS4_MAX_THREADS]; + float best_val[DS4_MAX_THREADS]; +} dspark_markov_q8_0_argmax_ctx; + +static void dspark_markov_q8_0_argmax_worker( + void *vctx, + uint64_t row0, + uint64_t row1) { + dspark_markov_q8_0_argmax_ctx *ctx = vctx; + uint64_t slot = ctx->rows_per_slot ? row0 / ctx->rows_per_slot : 0; + if (slot >= DS4_MAX_THREADS) slot = DS4_MAX_THREADS - 1u; + + float best_v = -FLT_MAX; + uint32_t best = (uint32_t)row0; + for (uint64_t row = row0; row < row1; row++) { + const uint8_t *wrow = ctx->data + row * ctx->blocks * 34u; + const float score = + ctx->logits[row] + + dot_q8_0_row(wrow, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks); + if (score > best_v) { + best_v = score; + best = (uint32_t)row; + } + } + + ctx->best_idx[slot] = best; + ctx->best_val[slot] = best_v; +} + +static bool dspark_markov_q8_0_argmax( + uint32_t *token_out, + const ds4_model *model, + const ds4_tensor *w, + const float *state, + const float *logits) { + if (!token_out || + !model || + !w || + !state || + !logits || + w->type != DS4_TENSOR_Q8_0 || + w->ndim != 2 || + w->dim[1] > UINT32_MAX) { + return false; + } + + const uint64_t in_dim = w->dim[0]; + const uint64_t out_dim = w->dim[1]; + const uint64_t blocks = (in_dim + 31u) / 32u; + if (out_dim == 0 || + blocks == 0 || + blocks > (uint64_t)SIZE_MAX / 32u || + blocks > (uint64_t)SIZE_MAX / sizeof(float)) { + return false; + } + + enum { DSPARK_MARKOV_ARGMAX_STACK_BLOCKS = 32 }; + int8_t xq_stack[DSPARK_MARKOV_ARGMAX_STACK_BLOCKS * 32u]; + float xscale_stack[DSPARK_MARKOV_ARGMAX_STACK_BLOCKS]; + const bool use_stack = blocks <= DSPARK_MARKOV_ARGMAX_STACK_BLOCKS; + int8_t *xq = use_stack ? xq_stack : xmalloc((size_t)blocks * 32u); + float *xscale = use_stack ? xscale_stack : + xmalloc((size_t)blocks * sizeof(xscale[0])); + quantize_q8_0_activation(state, xq, xscale, in_dim); + + ds4_threads_init(); + const uint32_t n_slots = + g_pool.n_threads == 0 ? 1u : g_pool.n_threads; + const uint64_t rows_per_slot = (out_dim + n_slots - 1u) / n_slots; + dspark_markov_q8_0_argmax_ctx ctx = { + .data = tensor_data(model, w), + .xq = xq, + .xscale = xscale, + .logits = logits, + .in_dim = in_dim, + .blocks = blocks, + .rows_per_slot = rows_per_slot, + }; + for (uint32_t i = 0; i < DS4_MAX_THREADS; i++) { + ctx.best_idx[i] = 0; + ctx.best_val[i] = -FLT_MAX; + } + + ds4_parallel_for(out_dim, dspark_markov_q8_0_argmax_worker, &ctx); + + uint32_t best = 0; + float best_v = -FLT_MAX; + for (uint32_t slot = 0; slot < n_slots && slot < DS4_MAX_THREADS; slot++) { + const uint64_t row0 = (uint64_t)slot * rows_per_slot; + if (row0 >= out_dim) break; + if (ctx.best_val[slot] > best_v) { + best_v = ctx.best_val[slot]; + best = ctx.best_idx[slot]; + } + } + + if (!use_stack) { + free(xscale); + free(xq); + } + *token_out = best; + return true; +} + +/* Exact target verification preserves correctness when this diagnostic mode + * proposes directly from the support model's base logits. */ +static bool dspark_markov_bias_disabled(void) { + static int cached = -1; + if (cached < 0) { + const char *env = getenv("DS4_DSPARK_NO_MARKOV"); + cached = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; + } + return cached == 1; +} + +static bool dspark_disable_fused_cpu_markov_argmax(void) { + static int cache = -1; + if (cache < 0) { + const char *env = getenv("DS4_DSPARK_DISABLE_FUSED_CPU_MARKOV_ARGMAX"); + cache = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; + } + return cache != 0; +} + +static bool dspark_disable_reuse_confidence0_markov(void) { + static int cache = -1; + if (cache < 0) { + const char *env = getenv("DS4_DSPARK_DISABLE_REUSE_CONFIDENCE0_MARKOV"); + cache = (env && env[0] && strcmp(env, "0") != 0) ? 1 : 0; + } + return cache != 0; +} + +static bool dspark_apply_markov_greedy_probe( + float *logits, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + int first_prev_token, + float *markov_state, + float *markov_bias, + int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], + uint32_t *proposal_len) { + if (proposal_len) *proposal_len = 0; + if (!logits || + !dspark_model || + !dw || + !markov_state || + !markov_bias || + !proposal || + first_prev_token < 0 || + (uint32_t)first_prev_token >= DS4_N_VOCAB || + !dspark_markov_probe_ready(dw)) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + const bool no_bias = dspark_markov_bias_disabled(); + int32_t prev_token = first_prev_token; + for (uint32_t draft = 0; draft < dw->block_size; draft++) { + float *row = logits + (uint64_t)draft * DS4_N_VOCAB; + if (!no_bias) { + if (!dspark_dense_row_to_f32(markov_state, + dspark_model, + final->markov_w1, + (uint32_t)prev_token)) { + return false; + } + matvec_any(markov_bias, dspark_model, final->markov_w2, markov_state); + for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { + row[i] += markov_bias[i]; + } + } + const uint32_t token = dspark_argmax_f32(row, DS4_N_VOCAB); + proposal[draft] = (int32_t)token; + prev_token = (int32_t)token; + } + + if (proposal_len) *proposal_len = dw->block_size; + return true; +} + +static bool dspark_confidence_probe_ready( + const ds4_dspark_weights *dw) { + if (!dspark_markov_probe_ready(dw)) return false; + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + if (!final->confidence_proj || + !dspark_tensor_type_matches(final->confidence_proj->type, + DS4_DSPARK_LAYOUT_DENSE)) { + return false; + } + return final->confidence_proj->ndim == 2 && + final->confidence_proj->dim[0] == + (uint64_t)DS4_N_EMBD + dw->markov_rank && + final->confidence_proj->dim[1] == 1; +} + +static bool dspark_eval_confidence_probe( + float *confidence_logits, + const float *hidden_rows, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + int first_prev_token, + const int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], + float *markov_state, + float *features, + uint32_t *confidence_len) { + if (confidence_len) *confidence_len = 0; + if (!confidence_logits || + !hidden_rows || + !dspark_model || + !dw || + !proposal || + !markov_state || + !features || + first_prev_token < 0 || + (uint32_t)first_prev_token >= DS4_N_VOCAB || + !dspark_confidence_probe_ready(dw)) { + return false; + } + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + int32_t prev_token = first_prev_token; + for (uint32_t draft = 0; draft < dw->block_size; draft++) { + if (prev_token < 0 || (uint32_t)prev_token >= DS4_N_VOCAB) { + return false; + } + if (!dspark_dense_row_to_f32(markov_state, + dspark_model, + final->markov_w1, + (uint32_t)prev_token)) { + return false; + } + memcpy(features, + hidden_rows + (uint64_t)draft * DS4_N_EMBD, + (uint64_t)DS4_N_EMBD * sizeof(features[0])); + memcpy(features + DS4_N_EMBD, + markov_state, + (uint64_t)dw->markov_rank * sizeof(features[0])); + matvec_any(confidence_logits + draft, + dspark_model, + final->confidence_proj, + features); + prev_token = proposal[draft]; + } + + if (confidence_len) *confidence_len = dw->block_size; + return true; +} + +static bool dspark_apply_markov_confidence_lazy_runtime( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + int first_prev_token, + float confidence_threshold, + float *logits, + float *markov_bias, + float *features, + size_t features_cap, + int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], + uint32_t *proposal_len, + uint32_t *confidence_len, + uint32_t *confidence_prefix_len, + bool reuse_first_confidence, + float *confidence0) { + if (proposal_len) *proposal_len = 0; + if (confidence_len) *confidence_len = 0; + if (confidence_prefix_len) *confidence_prefix_len = 0; + if (confidence0 && !reuse_first_confidence) *confidence0 = 0.0f; + if (!g || + !g->spec_logits || + !metal_graph_batch_ffn_norm(g) || + !dspark_model || + !dw || + !logits || + !markov_bias || + !features || + !proposal || + confidence_threshold <= 0.0f || + first_prev_token < 0 || + (uint32_t)first_prev_token >= DS4_N_VOCAB || + (reuse_first_confidence && !confidence0) || + !dspark_markov_probe_ready(dw) || + !dspark_confidence_probe_ready(dw)) { + return false; + } + + const uint64_t logits_bytes = + (uint64_t)DS4_N_VOCAB * sizeof(float); + const uint64_t hidden_bytes = + (uint64_t)DS4_N_EMBD * sizeof(float); + const uint64_t feature_count = + (uint64_t)DS4_N_EMBD + (uint64_t)dw->markov_rank; + if (feature_count > features_cap) return false; + float *markov_state = features + DS4_N_EMBD; + bool ok = true; + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + int32_t prev_token = first_prev_token; + uint32_t produced = 0; + uint32_t confident = 0; + for (uint32_t draft = 0; ok && draft < dw->block_size; draft++) { + if (prev_token < 0 || (uint32_t)prev_token >= DS4_N_VOCAB) { + ok = false; + break; + } + float confidence_logit = 0.0f; + if (draft == 0 && reuse_first_confidence) { + confidence_logit = *confidence0; + } else { + ok = dspark_dense_row_to_f32(markov_state, + dspark_model, + final->markov_w1, + (uint32_t)prev_token); + if (!ok) break; + + ok = ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), + (uint64_t)draft * hidden_bytes, + features, + hidden_bytes) != 0; + if (!ok) break; + matvec_any(&confidence_logit, + dspark_model, + final->confidence_proj, + features); + } + if (draft == 0 && confidence0) *confidence0 = confidence_logit; + if (confidence_len) *confidence_len = draft + 1u; + if (sigmoid_stable(confidence_logit) < confidence_threshold) { + ok = true; + break; + } + + int32_t token = -1; +#ifndef __APPLE__ + /* CUDA can apply the Markov bias and argmax without reading back the + * full logits row. Metal currently falls through to the CPU path. */ + if (ok && !dspark_markov_bias_disabled() && + getenv("DS4_DSPARK_NO_GPU_MARKOV") == NULL && + g->dspark_draft_tokens && + dw->markov_rank != 0 && (dw->markov_rank & 31u) == 0 && + final->markov_w1->type == DS4_TENSOR_Q8_0 && + final->markov_w2->type == DS4_TENSOR_Q8_0) { + ds4_gpu_tensor *row_view = + ds4_gpu_tensor_view(g->spec_logits, + (uint64_t)draft * logits_bytes, + logits_bytes); + uint64_t gpu_key = 0; + bool gpu_ok = row_view && + ds4_gpu_dspark_markov_argmax_tensor( + g->dspark_draft_tokens, + row_view, + dspark_model->map, + dspark_model->size, + final->markov_w1->abs_offset, + final->markov_w2->abs_offset, + (uint32_t)prev_token, + DS4_N_VOCAB, + dw->markov_rank) != 0 && + ds4_gpu_tensor_read(g->dspark_draft_tokens, + 0, + &gpu_key, + sizeof(gpu_key)) != 0; + ds4_gpu_tensor_free(row_view); + const uint32_t gpu_token = ~(uint32_t)(gpu_key & 0xffffffffu); + if (gpu_ok && gpu_key != 0 && gpu_token < DS4_N_VOCAB) { + token = (int32_t)gpu_token; + proposal[draft] = token; + produced = draft + 1u; + confident = produced; + prev_token = token; + continue; + } + } +#endif + if (ok) { + ok = ds4_gpu_tensor_read(g->spec_logits, + (uint64_t)draft * logits_bytes, + logits, + logits_bytes) != 0; + if (ok) { + uint32_t fused_token = 0; + if (dspark_markov_bias_disabled()) { + token = (int32_t)dspark_argmax_f32(logits, DS4_N_VOCAB); + } else if (!dspark_disable_fused_cpu_markov_argmax() && + dspark_markov_q8_0_argmax(&fused_token, + dspark_model, + final->markov_w2, + markov_state, + logits)) { + token = (int32_t)fused_token; + } else { + matvec_any(markov_bias, dspark_model, final->markov_w2, markov_state); + for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { + logits[i] += markov_bias[i]; + } + token = (int32_t)dspark_argmax_f32(logits, DS4_N_VOCAB); + } + } + } + if (!ok || token < 0 || (uint32_t)token >= DS4_N_VOCAB) { + ok = false; + break; + } + proposal[draft] = token; + produced = draft + 1u; + confident = produced; + prev_token = token; + } + + if (ok) { + if (proposal_len) *proposal_len = produced; + if (confidence_prefix_len) *confidence_prefix_len = confident; + } + return ok; +} + +static bool dspark_eval_confidence0_runtime( + ds4_gpu_graph *g, + const ds4_model *dspark_model, + const ds4_dspark_weights *dw, + int first_prev_token, + float *features, + size_t features_cap, + float *confidence0) { + if (confidence0) *confidence0 = 0.0f; + if (!confidence0 || + !g || + !metal_graph_batch_ffn_norm(g) || + !dspark_model || + !dw || + !features || + first_prev_token < 0 || + (uint32_t)first_prev_token >= DS4_N_VOCAB || + !dspark_confidence_probe_ready(dw)) { + return false; + } + + const uint64_t hidden_bytes = + (uint64_t)DS4_N_EMBD * sizeof(float); + const uint64_t feature_count = + (uint64_t)DS4_N_EMBD + (uint64_t)dw->markov_rank; + if (feature_count > features_cap || + ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < hidden_bytes) { + return false; + } + + float *markov_state = features + DS4_N_EMBD; + bool ok = true; + + const ds4_dspark_stage_weights *final = + &dw->stage[dw->n_stages - 1u]; + if (ok) { + ok = dspark_dense_row_to_f32(markov_state, + dspark_model, + final->markov_w1, + (uint32_t)first_prev_token); + } + if (ok) { + ok = ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), + 0, + features, + hidden_bytes) != 0; + } + if (ok) { + matvec_any(confidence0, dspark_model, final->confidence_proj, features); + } + + return ok; +} + +static uint32_t dspark_confident_prefix_len( + const float *confidence_logits, + uint32_t confidence_len, + float threshold) { + if (!confidence_logits || confidence_len == 0 || threshold <= 0.0f) { + return confidence_len; + } + for (uint32_t i = 0; i < confidence_len; i++) { + if (sigmoid_stable(confidence_logits[i]) < threshold) return i; + } + return confidence_len; +} + +static bool metal_graph_eval_mtp_draft_from_hc( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_model *mtp_model, + const ds4_mtp_weights *mtp, + ds4_gpu_tensor *prev_hc, + ds4_gpu_tensor *out_hc, + int token, + uint32_t pos, + float *logits, + int *top_id) { + if (!mtp || !mtp->block.attn_q_a || !g->mtp_raw_cache || !prev_hc || !out_hc) return false; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint32_t raw_row = pos % g->raw_cap; + uint32_t n_raw = g->mtp_n_raw + 1u; + if (n_raw > g->raw_window) n_raw = g->raw_window; + if (n_raw > g->raw_cap) n_raw = g->raw_cap; + + ds4_gpu_tensor *saved_cur = metal_graph_cur_hc(g); + ds4_gpu_tensor *saved_after = metal_graph_after_ffn_hc(g); + const uint32_t saved_tp_world = g->tp_world; + const uint32_t saved_tp_batch_rows = g->tp_batch_rows; + g->tp_world = 0; + g->tp_batch_rows = 0; + const bool suspended_expert_sharding = saved_tp_world == 2; + if (suspended_expert_sharding) { + ds4_gpu_tp_suspend_expert_sharding(1); + } + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = ds4_gpu_embed_token_hc_tensor(g->mtp_embed, + base_model->map, + base_model->size, + base_weights->token_embd->abs_offset, + (uint32_t)base_weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + 1) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->mtp_enorm, + g->mtp_embed, + mtp_model->map, + mtp_model->size, + mtp->enorm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_eproj, + mtp_model->map, + mtp_model->size, + mtp->e_proj->abs_offset, + DS4_N_EMBD, + DS4_N_EMBD, + g->mtp_enorm, + 1) != 0; + if (ok) ok = ds4_gpu_repeat_hc_tensor(g->mtp_eproj_hc, + g->mtp_eproj, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->mtp_hnorm_hc, + prev_hc, + mtp_model->map, + mtp_model->size, + mtp->hnorm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_hproj_hc, + mtp_model->map, + mtp_model->size, + mtp->h_proj->abs_offset, + DS4_N_EMBD, + DS4_N_EMBD, + g->mtp_hnorm_hc, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_add_tensor(g->mtp_input_hc, + g->mtp_eproj_hc, + g->mtp_hproj_hc, + (uint32_t)hc_dim) != 0; + if (ok) { + g->cur_hc_by_tier[g->active_tier] = g->mtp_input_hc; + g->after_ffn_hc_by_tier[g->active_tier] = out_hc; + ok = metal_graph_encode_decode_layer(g, + mtp_model, + &mtp->block, + 1, + pos, + g->mtp_raw_cache, + g->raw_cap, + raw_row, + n_raw, + token); + } + if (ok) g->cur_hc_by_tier[g->active_tier] = out_hc; + if (ok) ok = metal_graph_encode_output_head_mtp(g, + base_model, + base_weights, + mtp_model, + mtp, + base_weights->output->dim[1]); + if (ok && top_id) { + ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), + metal_graph_logits(g), + DS4_N_VOCAB) != 0; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + if (suspended_expert_sharding) { + ds4_gpu_tp_suspend_expert_sharding(0); + } + g->cur_hc_by_tier[g->active_tier] = saved_cur; + g->after_ffn_hc_by_tier[g->active_tier] = saved_after; + + if (ok && logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + if (ok && top_id) { + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top_id, sizeof(*top_id)) != 0; + } + if (ok && g->mtp_n_raw < g->raw_window) g->mtp_n_raw++; + g->tp_world = saved_tp_world; + g->tp_batch_rows = saved_tp_batch_rows; + if (!ok) { + (void)ds4_gpu_synchronize(); + g->cur_hc_by_tier[g->active_tier] = saved_cur; + g->after_ffn_hc_by_tier[g->active_tier] = saved_after; + } + return ok; +} + +static bool metal_graph_eval_mtp_draft( + ds4_gpu_graph *g, + const ds4_model *base_model, + const ds4_weights *base_weights, + const ds4_model *mtp_model, + const ds4_mtp_weights *mtp, + int token, + uint32_t pos, + float *logits, + int *top_id) { + return metal_graph_eval_mtp_draft_from_hc(g, + base_model, + base_weights, + mtp_model, + mtp, + metal_graph_cur_hc(g), + g->mtp_state_hc, + token, + pos, + logits, + top_id); +} + +/* ========================================================================= + * Imatrix Collection. + * ========================================================================= + * + * The 2-bit DS4 quants care most about routed MoE experts. For expert gate + * and up matrices the matmul input is the FFN-normalized activation row. For + * expert down matrices the matmul input is the routed SwiGLU row after route + * weighting. During Metal prefill those tensors are already materialized as + * `batch_ffn_norm`, `batch_router_selected`, and `batch_routed_mid`, so the + * collector observes the exact release graph without changing inference math. + * + * The output is llama.cpp's legacy imatrix `.dat` format. Entries are packed + * by expert: one tensor entry contains `n_expert * n_columns` floats and the + * quantizer slices the vector for each expert. + */ +typedef struct { + float *gate_up_sum2; /* [active layer][active expert][hidden] */ + float *down_sum2; /* [active layer][active expert][expert FFN] */ + uint32_t gate_up_count[DS4_MAX_LAYER][DS4_MAX_EXPERT]; + uint32_t down_count[DS4_MAX_LAYER][DS4_MAX_EXPERT]; + float *ffn_norm_buf; + float *routed_mid_buf; + uint16_t *routed_mid_f16_buf; + int *selected_buf; + float *sq_tmp; + uint32_t cap_tokens; + uint64_t observed_tokens; + uint64_t observed_routes; + uint32_t chunks; + const char *dataset_path; +} ds4_imatrix_collector; + +static bool imatrix_collector_init(ds4_imatrix_collector *c, uint32_t cap_tokens, const char *dataset_path) { + memset(c, 0, sizeof(*c)); + c->cap_tokens = cap_tokens ? cap_tokens : 1u; + c->dataset_path = dataset_path; + const size_t gate_n = (size_t)DS4_N_LAYER * DS4_N_EXPERT * DS4_N_EMBD; + const size_t down_n = (size_t)DS4_N_LAYER * DS4_N_EXPERT * DS4_N_FF_EXP; + c->gate_up_sum2 = xcalloc(gate_n, sizeof(c->gate_up_sum2[0])); + c->down_sum2 = xcalloc(down_n, sizeof(c->down_sum2[0])); + c->ffn_norm_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EMBD * sizeof(c->ffn_norm_buf[0])); + c->routed_mid_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(c->routed_mid_buf[0])); + c->routed_mid_f16_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(c->routed_mid_f16_buf[0])); + c->selected_buf = xmalloc((size_t)c->cap_tokens * DS4_N_EXPERT_USED * sizeof(c->selected_buf[0])); + c->sq_tmp = xmalloc((size_t)DS4_N_EMBD * sizeof(c->sq_tmp[0])); + return c->gate_up_sum2 && c->down_sum2 && c->ffn_norm_buf && + c->routed_mid_buf && c->routed_mid_f16_buf && c->selected_buf && c->sq_tmp; +} + +static void imatrix_collector_free(ds4_imatrix_collector *c) { + if (!c) return; + free(c->gate_up_sum2); + free(c->down_sum2); + free(c->ffn_norm_buf); + free(c->routed_mid_buf); + free(c->routed_mid_f16_buf); + free(c->selected_buf); + free(c->sq_tmp); + memset(c, 0, sizeof(*c)); +} + +static float *imatrix_gate_up_ptr(ds4_imatrix_collector *c, uint32_t il, uint32_t expert) { + return c->gate_up_sum2 + ((size_t)il * DS4_N_EXPERT + expert) * DS4_N_EMBD; +} + +static float *imatrix_down_ptr(ds4_imatrix_collector *c, uint32_t il, uint32_t expert) { + return c->down_sum2 + ((size_t)il * DS4_N_EXPERT + expert) * DS4_N_FF_EXP; +} + +static bool imatrix_collect_layer_batch( + ds4_imatrix_collector *c, + ds4_gpu_graph *g, + uint32_t il, + uint32_t n_tokens) { + if (!c || n_tokens == 0) return true; + if (n_tokens > c->cap_tokens) return false; + + const uint64_t norm_bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); + const uint64_t mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP; + const uint64_t mid_bytes = mid_elems * (g->batch_routed_mid_is_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t sel_bytes = (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int); + void *mid_dst = g->batch_routed_mid_is_f16 + ? (void *)c->routed_mid_f16_buf + : (void *)c->routed_mid_buf; + if (ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(g), 0, c->ffn_norm_buf, norm_bytes) == 0 || + ds4_gpu_tensor_read(metal_graph_batch_routed_mid(g), 0, mid_dst, mid_bytes) == 0 || + ds4_gpu_tensor_read(metal_graph_batch_router_selected(g), 0, c->selected_buf, sel_bytes) == 0) + { + return false; + } + + for (uint32_t t = 0; t < n_tokens; t++) { + const float *x = c->ffn_norm_buf + (size_t)t * DS4_N_EMBD; + for (uint32_t i = 0; i < DS4_N_EMBD; i++) c->sq_tmp[i] = x[i] * x[i]; + + for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { + const int expert = c->selected_buf[(size_t)t * DS4_N_EXPERT_USED + slot]; + if (expert < 0 || (uint32_t)expert >= DS4_N_EXPERT) continue; + + float *gate_up = imatrix_gate_up_ptr(c, il, (uint32_t)expert); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) gate_up[i] += c->sq_tmp[i]; + c->gate_up_count[il][expert]++; + + float *down = imatrix_down_ptr(c, il, (uint32_t)expert); + const size_t mid_off = ((size_t)t * DS4_N_EXPERT_USED + slot) * DS4_N_FF_EXP; + if (g->batch_routed_mid_is_f16) { + const uint16_t *mid = c->routed_mid_f16_buf + mid_off; + for (uint32_t i = 0; i < DS4_N_FF_EXP; i++) { + const float v = f16_to_f32(mid[i]); + down[i] += v * v; + } + } else { + const float *mid = c->routed_mid_buf + mid_off; + for (uint32_t i = 0; i < DS4_N_FF_EXP; i++) down[i] += mid[i] * mid[i]; + } + c->down_count[il][expert]++; + c->observed_routes++; + } + } + c->observed_tokens += n_tokens; + c->chunks++; + return true; +} + +static void imatrix_write_i32(FILE *fp, int32_t v) { + if (fwrite(&v, sizeof(v), 1, fp) != 1) ds4_die("failed to write imatrix"); +} + +static void imatrix_write_entry( + FILE *fp, + const char *name, + const float *sum2, + const uint32_t *counts, + uint32_t n_expert, + uint32_t n_col) { + const int32_t len = (int32_t)strlen(name); + const int32_t ncall = 1; + const int32_t nval = (int32_t)((uint64_t)n_expert * n_col); + imatrix_write_i32(fp, len); + if (fwrite(name, 1, (size_t)len, fp) != (size_t)len) ds4_die("failed to write imatrix name"); + imatrix_write_i32(fp, ncall); + imatrix_write_i32(fp, nval); + + float *tmp = xmalloc((size_t)n_col * sizeof(tmp[0])); + for (uint32_t e = 0; e < n_expert; e++) { + const uint32_t count = counts[e]; + const float *src = sum2 + (size_t)e * n_col; + if (count == 0) { + for (uint32_t i = 0; i < n_col; i++) tmp[i] = 1.0f; + } else { + const float inv = 1.0f / (float)count; + for (uint32_t i = 0; i < n_col; i++) tmp[i] = src[i] * inv; + } + if (fwrite(tmp, sizeof(tmp[0]), n_col, fp) != n_col) ds4_die("failed to write imatrix values"); + } + free(tmp); +} + +static bool imatrix_collector_save( + const ds4_imatrix_collector *c, + const ds4_weights *weights, + const char *path) { + FILE *fp = fopen(path, "wb"); + if (!fp) { + fprintf(stderr, "ds4: failed to open imatrix output %s: %s\n", path, strerror(errno)); + return false; + } + + const int32_t entries = (int32_t)(DS4_N_LAYER * 3); + imatrix_write_i32(fp, entries); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const ds4_layer_weights *layer = &weights->layer[il]; + char name[256]; + snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_gate_exps->name.len, layer->ffn_gate_exps->name.ptr); + imatrix_write_entry(fp, name, + c->gate_up_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_EMBD, + c->gate_up_count[il], + DS4_N_EXPERT, + DS4_N_EMBD); + snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_up_exps->name.len, layer->ffn_up_exps->name.ptr); + imatrix_write_entry(fp, name, + c->gate_up_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_EMBD, + c->gate_up_count[il], + DS4_N_EXPERT, + DS4_N_EMBD); + snprintf(name, sizeof(name), "%.*s", (int)layer->ffn_down_exps->name.len, layer->ffn_down_exps->name.ptr); + imatrix_write_entry(fp, name, + c->down_sum2 + (size_t)il * DS4_N_EXPERT * DS4_N_FF_EXP, + c->down_count[il], + DS4_N_EXPERT, + DS4_N_FF_EXP); + } + + const int32_t chunks = (int32_t)c->chunks; + imatrix_write_i32(fp, chunks); + const char *dataset = c->dataset_path ? c->dataset_path : ""; + const int32_t dataset_len = (int32_t)strlen(dataset); + imatrix_write_i32(fp, dataset_len); + if (dataset_len && fwrite(dataset, 1, (size_t)dataset_len, fp) != (size_t)dataset_len) { + ds4_die("failed to write imatrix dataset name"); + } + + if (fclose(fp) != 0) { + fprintf(stderr, "ds4: failed to close imatrix output %s: %s\n", path, strerror(errno)); + return false; + } + return true; +} + +static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { + memset(g->layer_n_comp, 0, sizeof(g->layer_n_comp)); + memset(g->layer_n_index_comp, 0, sizeof(g->layer_n_index_comp)); + g->mtp_n_raw = 0; + metal_graph_dspark_cache_reset(g); + metal_graph_dspark_capture_invalidate(g); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!g->layer_raw_cache[il]) continue; + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio == 0) continue; + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; + const uint64_t attn_rows = (uint64_t)coff * ratio; + if (!metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows)) return false; + if (!metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows)) return false; + if (ratio == 4) { + const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM; + const uint64_t index_rows = (uint64_t)coff * ratio; + if (!metal_tensor_fill_f32(g->layer_index_state_kv[il], 0.0f, index_width * index_rows)) return false; + if (!metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows)) return false; + } + } + return true; +} + +/* Execute graph-backend prefill in layer-major order so intermediate + * activations stay on the GPU and cache state is built exactly once. */ +static void gpu_graph_report_prefill_display_progress( + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t start, + uint32_t n_tokens, + uint32_t layer_done, + int total) { + if (!display_progress) return; + if (layer_done > (uint32_t)DS4_N_LAYER) layer_done = (uint32_t)DS4_N_LAYER; + uint64_t done = (uint64_t)n_tokens * layer_done / (uint32_t)DS4_N_LAYER; + if (layer_done == (uint32_t)DS4_N_LAYER) done = n_tokens; + display_progress(display_progress_ud, "prefill_display", + (int)(start + (uint32_t)done), total); +} + +typedef struct { + int tier; + uint32_t first_layer; + uint32_t end_layer; +} metal_graph_prefill_stage; + +static bool metal_graph_build_prefill_stages( + const ds4_gpu_graph *g, + metal_graph_prefill_stage *stages, + uint32_t *n_stages) { + if (!g || !g->placement || !stages || !n_stages) return false; + uint32_t ns = 0; + int prev_tier = -1; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const int tier = g->placement[il + 1]; + if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + if (il == 0 || tier != prev_tier) { + if (ns >= DS4_MAX_GPUS) return false; + stages[ns].tier = tier; + stages[ns].first_layer = il; + stages[ns].end_layer = il + 1u; + ns++; + prev_tier = tier; + } else { + stages[ns - 1u].end_layer = il + 1u; + } + } + *n_stages = ns; + return ns != 0; +} + +static bool metal_graph_encode_prefill_stage_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const metal_graph_prefill_stage *stage, + uint32_t pos0, + uint32_t n_tokens) { + if (!g || !model || !weights || !stage || n_tokens == 0) return false; + if (!metal_graph_set_active_tier_no_copy(g, stage->tier)) return false; + for (uint32_t il = stage->first_layer; il < stage->end_layer; il++) { + if (g->placement && g->placement[il + 1] != stage->tier) return false; + if (!metal_graph_encode_layer_batch(g, + model, + &weights->layer[il], + il, + pos0, + n_tokens)) { + return false; + } + if (g->pipeline_capture_chunk_len != 0 && + !metal_graph_dspark_capture_prefill_rows( + g, il, + g->pipeline_capture_chunk_start, + g->pipeline_capture_chunk_len, + pos0, + n_tokens)) { + return false; + } + } + return true; +} + +static bool metal_graph_prefill_pipeline_stage_major( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + float *logits, + bool show_progress, + ds4_session_progress_fn display_progress, + void *display_progress_ud) { + if (!g || !model || !weights || !prompt || !g->placement || + n_tokens == 0 || n_tokens > g->prefill_cap || + start > (uint32_t)prompt->len || + n_tokens > (uint32_t)prompt->len - start) { + return false; + } + + metal_graph_prefill_stage stages[DS4_MAX_GPUS]; + uint32_t n_stages = 0; + g->pipeline_capture_chunk_start = start; + g->pipeline_capture_chunk_len = + g->dspark_capture_enabled ? n_tokens : 0; + if (!metal_graph_build_prefill_stages(g, stages, &n_stages) || n_stages < 2) { + return false; + } + if (stages[0].tier != g->emb_tier) { + return false; + } + + uint32_t mb_cap = metal_graph_cuda_prefill_pipeline_microbatch(); + if (mb_cap == 0 || mb_cap >= n_tokens) return false; + if (mb_cap > g->prefill_cap) mb_cap = g->prefill_cap; + const uint32_t n_mb = (n_tokens + mb_cap - 1u) / mb_cap; + if (n_mb < 2) return false; + + if (display_progress) + display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + bool ok = true; + const double t0 = getenv("DS4_METAL_GRAPH_PREFILL_PROFILE") ? now_sec() : 0.0; + + const bool sequential = getenv("DS4_CUDA_PREFILL_PIPELINE_SEQUENTIAL") != NULL; + const bool suppress_q8_cache = + !metal_graph_cuda_prefill_pipeline_q8_cache_requested(); + const int saved_q8_cache_suppressed = ds4_gpu_q8_cache_suppressed(); + if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(1); + ok = ds4_gpu_begin_commands() != 0; + if (sequential) { + for (uint32_t mb_i = 0; ok && mb_i < n_mb; mb_i++) { + const uint32_t mb_off = mb_i * mb_cap; + uint32_t mb_len = n_tokens - mb_off; + if (mb_len > mb_cap) mb_len = mb_cap; + const uint32_t pos0 = start + mb_off; + + for (uint32_t stage_i = 0; ok && stage_i < n_stages; stage_i++) { + g->batch_token_offset = mb_off; + if (stage_i == 0) { + ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); + ds4_gpu_tensor *tokens_view = NULL; + if (ok) { + tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), + (uint64_t)mb_off * sizeof(int32_t), + (uint64_t)mb_len * sizeof(int32_t)); + ok = tokens_view != NULL; + } + if (ok) { + ok = metal_graph_upload_prompt_embeddings_hc( + g->batch_cur_hc_by_tier[stages[0].tier], + tokens_view, + model, + weights, + prompt, + pos0, + mb_len); + } + ds4_gpu_tensor_free(tokens_view); + } + if (ok) { + ok = metal_graph_encode_prefill_stage_batch(g, + model, + weights, + &stages[stage_i], + pos0, + mb_len); + } + if (ok && stage_i + 1u < n_stages) { + ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; + ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; + if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { + ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && + ds4_gpu_synchronize() != 0; + } + ok = src && dst && + ds4_gpu_tensor_copy_xdev_ordered(dst, + src, + (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok && display_progress) { + uint32_t done = mb_off + mb_len; + if (done > n_tokens) done = n_tokens; + display_progress(display_progress_ud, + "prefill_display", + (int)(start + done), + prompt->len); + } + if (show_progress) { + fprintf(stderr, "ds4: gpu sequential pipeline prefill microbatch %u/%u\r", + mb_i + 1u, + n_mb); + fflush(stderr); + } + if (ok && mb_i + 1u < n_mb) ok = ds4_gpu_begin_commands() != 0; + } + } else { + for (uint32_t wave = 0; ok && wave < n_mb + n_stages - 1u; wave++) { + uint32_t smax = wave < n_stages ? wave : n_stages - 1u; + for (int si = (int)smax; ok && si >= 0; si--) { + const uint32_t stage_i = (uint32_t)si; + const uint32_t mb_i = wave - stage_i; + if (mb_i >= n_mb) continue; + const uint32_t mb_off = mb_i * mb_cap; + uint32_t mb_len = n_tokens - mb_off; + if (mb_len > mb_cap) mb_len = mb_cap; + const uint32_t pos0 = start + mb_off; + + g->batch_token_offset = mb_off; + if (stage_i == 0) { + ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); + ds4_gpu_tensor *tokens_view = NULL; + if (ok) { + tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), + (uint64_t)mb_off * sizeof(int32_t), + (uint64_t)mb_len * sizeof(int32_t)); + ok = tokens_view != NULL; + } + if (ok) { + ok = metal_graph_upload_prompt_embeddings_hc( + g->batch_cur_hc_by_tier[stages[0].tier], + tokens_view, + model, + weights, + prompt, + pos0, + mb_len); + } + ds4_gpu_tensor_free(tokens_view); + } + + if (ok) { + ok = metal_graph_encode_prefill_stage_batch(g, + model, + weights, + &stages[stage_i], + pos0, + mb_len); + } + if (ok && stage_i + 1u < n_stages) { + ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; + ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; + if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { + ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && + ds4_gpu_synchronize() != 0; + } + ok = src && dst && + ds4_gpu_tensor_copy_xdev_ordered(dst, + src, + (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; + } + if (ok && display_progress && stage_i + 1u == n_stages) { + uint32_t done = mb_off + mb_len; + if (done > n_tokens) done = n_tokens; + display_progress(display_progress_ud, + "prefill_display", + (int)(start + done), + prompt->len); + } + } + if (show_progress) { + fprintf(stderr, "ds4: gpu pipeline prefill wave %u/%u\r", + wave + 1u, + n_mb + n_stages - 1u); + fflush(stderr); + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + } + if (show_progress) fputc('\n', stderr); + g->batch_token_offset = 0; + if (!ok) { + if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); + return false; + } + + const uint32_t final_len = n_tokens - (n_mb - 1u) * mb_cap; + const int src_tier = stages[n_stages - 1u].tier; + if (!metal_graph_set_active_tier_no_copy(g, src_tier)) { + if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); + return false; + } + + ds4_gpu_tensor *saved_cur = g->cur_hc_by_tier[src_tier]; + ds4_gpu_tensor *last_hc = NULL; + if (logits) { + last_hc = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[src_tier], + final_len - 1u, + hc_dim); + ok = last_hc != NULL; + } + if (ok && logits) { + g->cur_hc_by_tier[src_tier] = last_hc; + ok = ds4_gpu_begin_commands() != 0; + } + if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + if (ok && logits) ok = ds4_gpu_end_commands() != 0; + else if (!ok) (void)ds4_gpu_synchronize(); + g->cur_hc_by_tier[src_tier] = saved_cur; + ds4_gpu_tensor_free(last_hc); + if (g->placement && g->active_tier != src_tier) { + ok = metal_graph_set_active_tier_no_copy(g, src_tier); + } + if (ok && logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), + 0, + logits, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + if (ok && display_progress) + display_progress(display_progress_ud, "prefill_display", + (int)(start + n_tokens), prompt->len); + if (ok && t0 != 0.0) { + const double t1 = now_sec(); + fprintf(stderr, + "ds4: gpu pipeline prefill total tokens=%u stages=%u mb=%u total=%.3f ms\n", + n_tokens, + n_stages, + mb_cap, + (t1 - t0) * 1000.0); + } + if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); + return ok; +} + +static bool metal_graph_prefill_layer_major( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + float *logits, + bool show_progress, + ds4_imatrix_collector *imatrix, + ds4_session_progress_fn display_progress, + void *display_progress_ud) { + if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; + if (start > (uint32_t)prompt->len) return false; + if (n_tokens > (uint32_t)prompt->len - start) return false; + + if (display_progress) + display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); + + bool ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), prompt, start, n_tokens); + if (!ok) return false; + +#ifdef DS4_ROCM_BUILD + if (g->ssd_streaming && + DS4_MODEL_VARIANT == DS4_VARIANT_PRO && + n_tokens >= 1024u) { + ds4_gpu_stream_expert_cache_release_resident(); + } +#endif + + if (!metal_graph_warmup_prefill_kernels(g, model, weights, n_tokens)) return false; + if (g->placement && + !metal_graph_set_active_tier_no_copy(g, g->emb_tier)) { + return false; + } + metal_graph_dspark_capture_begin_prefill(g); + + const bool split_profile = + glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_SPLIT_PROFILE", + "DS4_METAL_GRAPH_PREFILL_SPLIT_PROFILE"); + /* + * A full long-prompt prefill can keep the GPU busy for a long time. Split + * non-tiny prefills when a frontend asked for display progress: completed + * layer command buffers are real scheduling/keepalive points, while + * callbacks emitted while encoding one huge command buffer would only be + * cosmetic. + */ + const bool throttle = graph_power_throttle_enabled(g); + const bool callback_split = display_progress != NULL && n_tokens >= 32; + const bool split_commands = g->ssd_streaming || + split_profile || throttle || callback_split || + n_tokens > 2048 || imatrix != NULL; + const bool profile = + glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", + "DS4_METAL_GRAPH_PREFILL_PROFILE") || + split_profile; + const double t0 = profile ? now_sec() : 0.0; + double encode_s = 0.0; + double execute_s = 0.0; + + const uint32_t pipeline_mb = metal_graph_cuda_prefill_pipeline_microbatch(); + if (!split_commands && + !profile && + imatrix == NULL && + metal_graph_cuda_prefill_pipeline_requested(g) && + pipeline_mb != 0 && + pipeline_mb < n_tokens) { + return metal_graph_prefill_pipeline_stage_major(g, + model, + weights, + prompt, + start, + n_tokens, + logits, + show_progress, + display_progress, + display_progress_ud); + } + + if (!split_commands) { + ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), + metal_graph_prefill_tokens(g), + model, + weights, + prompt, + start, + n_tokens); + if (ok) ok = ds4_gpu_begin_commands() != 0; + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + ok = metal_graph_encode_layer_batch(g, + model, + &weights->layer[il], + il, + start, + n_tokens); + if (!ok) { + fprintf(stderr, "ds4: gpu whole-prefill layer %u encode failed\n", il); + } + if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, + il, + start, + n_tokens); + if (show_progress) { + fprintf(stderr, "ds4: gpu prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); + fflush(stderr); + } + } + if (show_progress) fputc('\n', stderr); + if (display_progress) + display_progress(display_progress_ud, "prefill_display", + (int)(start + n_tokens), prompt->len); + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + uint32_t output_row = (uint32_t)n_tokens - 1u; + const char *output_row_env = glm_graph_env_value( + "DS4_ROCM_GRAPH_OUTPUT_ROW", + "DS4_METAL_GRAPH_OUTPUT_ROW"); + if (output_row_env && output_row_env[0]) { + char *end = NULL; + unsigned long v = strtoul(output_row_env, &end, 10); + if (end != output_row_env && v < (unsigned long)n_tokens) { + output_row = (uint32_t)v; + } + } + const int src_tier = g->active_tier; + ds4_gpu_tensor *saved_cur = g->cur_hc_by_tier[src_tier]; + ds4_gpu_tensor *last_hc = NULL; + if (ok && logits) { + last_hc = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), output_row, hc_dim); + ok = last_hc != NULL; + } + if (ok && logits) { + g->cur_hc_by_tier[src_tier] = last_hc; + ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + g->cur_hc_by_tier[src_tier] = saved_cur; + } + + const double t_encoded = profile ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = profile ? now_sec() : 0.0; + g->cur_hc_by_tier[src_tier] = saved_cur; + if (last_hc) ds4_gpu_tensor_free(last_hc); + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after whole-prefill graph failure also failed\n"); + } + return false; + } +#ifdef __APPLE__ + ds4_gpu_release_zero_prefix_prefill_mask_cache(); +#endif + + const double t_before_read = profile ? now_sec() : 0.0; + if (logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + if (profile) { + const double t_read = now_sec(); + fprintf(stderr, + "ds4: gpu graph prefill total tokens=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms\n", + n_tokens, + (t_encoded - t0) * 1000.0, + (t_done - t_encoded) * 1000.0, + (t_read - t_before_read) * 1000.0, + (t_read - t0) * 1000.0); + } + return ok; + } + + if (g->ssd_streaming) { + g->streaming_static_decode_map_current = false; + if (!metal_graph_stream_map_token(model, weights)) return false; + } + metal_graph_stream_prefill_selected_profile_reset(g); + metal_graph_stream_prepare_slot layer_prepare_slots[DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD]; + memset(layer_prepare_slots, 0, sizeof(layer_prepare_slots)); + const bool layer_pagein = + metal_graph_stream_prefill_layer_pagein_enabled(g); + const bool layer_readahead = + !layer_pagein && + metal_graph_stream_prefill_layer_readahead_enabled(g); + const bool layer_pread = + !layer_pagein && !layer_readahead && + metal_graph_stream_prefill_layer_pread_enabled(g); + const bool layer_madvise = + !layer_pagein && !layer_pread && !layer_readahead && + metal_graph_stream_prefill_layer_madvise_enabled(g); + const bool layer_prepare = + layer_pagein || layer_pread || layer_readahead || layer_madvise; + const bool layer_prepare_overlap = + layer_prepare && metal_graph_stream_prefill_layer_pagein_overlap_enabled(); + const uint32_t layer_prepare_ahead = + layer_prepare && layer_prepare_overlap ? + metal_graph_stream_prefill_layer_prepare_ahead() : 1u; + const bool batch_selected_addr = + metal_graph_stream_prefill_batch_selected_addr_enabled(g, weights, n_tokens) || + metal_graph_cuda_stream_prefill_batch_selected_addr_enabled(g, weights, n_tokens); +#ifdef DS4_ROCM_BUILD + rocm_graph_stream_layer_expert_load rocm_full_layer_load; + memset(&rocm_full_layer_load, 0, sizeof(rocm_full_layer_load)); +#endif + if (g->ssd_streaming && DS4_N_LAYER > 0) { + if (layer_prepare) { + if (!metal_graph_stream_prepare_start_if_needed(g, + model, + weights, + 0, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + batch_selected_addr, + layer_prepare_slots, + layer_prepare_ahead)) { + return false; + } + } else { + if (batch_selected_addr) { + metal_graph_stream_readahead_layer_decode(model, weights, 0); + } else { + metal_graph_stream_readahead_layer(model, weights, 0); + } + } + } +#ifdef DS4_ROCM_BUILD + if (g->ssd_streaming && DS4_N_LAYER > 0 && + !rocm_graph_stream_layer_expert_load_start_next(&rocm_full_layer_load, + g, + model, + weights, + 0, + n_tokens)) { + return false; + } +#endif + + double t_layer0 = (profile || throttle) ? now_sec() : 0.0; + ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), + metal_graph_prefill_tokens(g), + model, + weights, + prompt, + start, + n_tokens); + const double t_embed_encoded = (profile || throttle) ? now_sec() : 0.0; + const double t_embed_done = (profile || throttle) ? now_sec() : 0.0; + if (profile) { + encode_s += t_embed_encoded - t_layer0; + execute_s += t_embed_done - t_embed_encoded; + if (split_profile) { + fprintf(stderr, + "ds4: metal layer-major prefill embed encode=%.3f ms execute=%.3f ms\n", + (t_embed_encoded - t_layer0) * 1000.0, + (t_embed_done - t_embed_encoded) * 1000.0); + } + } + if (!ok) { +#ifdef DS4_ROCM_BUILD + (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); +#endif + if (layer_prepare) { + (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, + layer_prepare_ahead); + } + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after layer-major prefill embed failure also failed\n"); + } + return false; + } + + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + double layer_elapsed = 0.0; + if (layer_prepare && + !metal_graph_stream_prepare_join_layer(g, + model, + weights, + il, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + batch_selected_addr, + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + break; + } +#ifdef DS4_ROCM_BUILD + const bool rocm_full_layer_stream_prefill = + rocm_graph_stream_prefill_full_layer_enabled(g, + &weights->layer[il], + il, + n_tokens); + if (rocm_full_layer_stream_prefill && + !rocm_graph_stream_layer_expert_load_ready(&rocm_full_layer_load, + g, + model, + weights, + il, + n_tokens)) { + ok = false; + break; + } + if (rocm_full_layer_stream_prefill && + !rocm_graph_stream_layer_expert_load_start_next(&rocm_full_layer_load, + g, + model, + weights, + il + 1u, + n_tokens)) { + ok = false; + break; + } +#endif + if (g->ssd_streaming) { + g->streaming_static_decode_map_current = false; + bool decode_only_map = batch_selected_addr; +#ifdef DS4_ROCM_BUILD + decode_only_map = decode_only_map || rocm_full_layer_stream_prefill; +#endif + const bool map_ok = decode_only_map ? + metal_graph_stream_map_layer_decode(model, weights, il) : + metal_graph_stream_map_layer(model, weights, il); + if (!map_ok) { + ok = false; + break; + } + } + if (g->ssd_streaming) { + if (layer_prepare && layer_prepare_overlap) { + bool started_future = false; + for (uint32_t ahead = 1; ahead <= layer_prepare_ahead; ahead++) { + if (il + ahead >= DS4_N_LAYER) break; + started_future = true; + if (!metal_graph_stream_prepare_start_if_needed(g, + model, + weights, + il + ahead, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + batch_selected_addr, + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + break; + } + } + if (!ok) break; + if (!started_future && logits) { + metal_graph_stream_readahead_output(model, weights); + } + } else if (!layer_prepare && il + 1 < DS4_N_LAYER) { + if (batch_selected_addr) { + metal_graph_stream_readahead_layer_decode(model, weights, il + 1); + } else { + metal_graph_stream_readahead_layer(model, weights, il + 1); + } + } else if (logits) { + metal_graph_stream_readahead_output(model, weights); + } + } + if (split_profile) { + /* (B6 fix): split-profile diagnostic bypasses the + * metal_graph_encode_layer_batch wrapper that normally does + * the per-layer tier switch. Replicate the switch here so the + * diagnostic / profile mode stays multi-tier-correct. + * Single-tier (g->placement == NULL): no-op. */ + if (g->placement) { + const int this_tier = g->placement[il + 1]; + if (!metal_graph_set_active_tier_batch(g, this_tier, (uint32_t)n_tokens)) { + ok = false; + break; + } + } + const double t_attn0 = now_sec(); + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_layer_attention_batch(g, + model, + &weights->layer[il], + il, + start, + n_tokens); + if (!ok) { + fprintf(stderr, "ds4: gpu layer-major prefill layer %u attention encode failed\n", il); + } + const double t_attn_encoded = now_sec(); + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_attn_done = now_sec(); + + const double t_ffn0 = now_sec(); + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_layer_ffn_batch(g, + model, + &weights->layer[il], + il, + start, + n_tokens, + NULL, + 0); + if (!ok) { + fprintf(stderr, "ds4: gpu layer-major prefill layer %u ffn encode failed\n", il); + } + if (ok) { + ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); + g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); + g->batch_next_hc_by_tier[g->active_tier] = tmp; + } + if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, + il, + start, + n_tokens); + if (ok) ok = metal_graph_capture_prefill_seed_router_selected(g, + il, + n_tokens); + const double t_ffn_encoded = now_sec(); + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_ffn_done = now_sec(); +#ifdef DS4_ROCM_BUILD + if (ok) { + ok = rocm_graph_stream_seed_full_layer_selected(g, + model, + &weights->layer[il], + il, + n_tokens); + } +#endif + if (ok) { + ok = metal_graph_stream_prefill_selected_profile_layer( + g, + &weights->layer[il], + il, + n_tokens); + } + if (ok && imatrix) ok = imatrix_collect_layer_batch(imatrix, g, il, (uint32_t)n_tokens); + layer_elapsed = (t_attn_done - t_attn0) + (t_ffn_done - t_ffn0); + + encode_s += (t_attn_encoded - t_attn0) + (t_ffn_encoded - t_ffn0); + execute_s += (t_attn_done - t_attn_encoded) + (t_ffn_done - t_ffn_encoded); + fprintf(stderr, + "ds4: metal layer-major prefill layer %u attn encode=%.3f execute=%.3f ms ffn encode=%.3f execute=%.3f ms\n", + il, + (t_attn_encoded - t_attn0) * 1000.0, + (t_attn_done - t_attn_encoded) * 1000.0, + (t_ffn_encoded - t_ffn0) * 1000.0, + (t_ffn_done - t_ffn_encoded) * 1000.0); + } else { + const double t_chunk0 = (profile || throttle) ? now_sec() : 0.0; + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_layer_batch(g, + model, + &weights->layer[il], + il, + start, + n_tokens); + if (!ok) { + fprintf(stderr, "ds4: gpu layer-major prefill layer %u encode failed\n", il); + } + if (ok) ok = metal_graph_dspark_capture_prefill_layer(g, + il, + start, + n_tokens); + if (ok) ok = metal_graph_capture_prefill_seed_router_selected(g, + il, + n_tokens); + const double t_encoded = (profile || throttle) ? now_sec() : 0.0; + if (ok) ok = ds4_gpu_end_commands() != 0; + const double t_done = (profile || throttle) ? now_sec() : 0.0; +#ifdef DS4_ROCM_BUILD + if (ok) { + ok = rocm_graph_stream_seed_full_layer_selected(g, + model, + &weights->layer[il], + il, + n_tokens); + } +#endif + if (ok) { + ok = metal_graph_stream_prefill_selected_profile_layer( + g, + &weights->layer[il], + il, + n_tokens); + } + if (ok && imatrix) ok = imatrix_collect_layer_batch(imatrix, g, il, (uint32_t)n_tokens); + layer_elapsed = t_done - t_chunk0; + if (profile) { + encode_s += t_encoded - t_chunk0; + execute_s += t_done - t_encoded; + fprintf(stderr, + "ds4: gpu layer-major prefill layer %u encode=%.3f ms execute=%.3f ms\n", + il, + (t_encoded - t_chunk0) * 1000.0, + (t_done - t_encoded) * 1000.0); + } + } + if (ok && + g->ssd_streaming && + layer_prepare && + !layer_prepare_overlap) { + if (il + 1 < DS4_N_LAYER) { + if (!metal_graph_stream_prepare_start_if_needed(g, + model, + weights, + il + 1, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + batch_selected_addr, + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + } + } else if (logits) { + metal_graph_stream_readahead_output(model, weights); + } + } + if (!ok) { +#ifdef DS4_ROCM_BUILD + (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); +#endif + if (layer_prepare) { + (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, + layer_prepare_ahead); + } + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after layer-major prefill failure also failed\n"); + } + return false; + } + graph_power_note_prefill_layer(g, il, layer_elapsed); + gpu_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + start, + n_tokens, + il + 1, + prompt->len); + if (show_progress) { + fprintf(stderr, "ds4: gpu prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); + fflush(stderr); + } + } + if (!ok) { +#ifdef DS4_ROCM_BUILD + (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); +#endif + if (layer_prepare) { + (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, + layer_prepare_ahead); + } + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after layer-major prefill failure also failed\n"); + } + return false; + } +#ifdef __APPLE__ + /* Zero-prefix masks are shared across the 43 per-layer command batches, + * then become dead weight. Release them before the output head and later + * replay/decode chunks so the prefill win does not add residency pressure. */ + ds4_gpu_release_zero_prefix_prefill_mask_cache(); +#endif + if (show_progress) fputc('\n', stderr); + metal_graph_stream_prefill_selected_profile_summary(g); +#ifdef DS4_ROCM_BUILD + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); + if (g->ssd_streaming) ds4_gpu_release_q8_f16_cache(); +#endif + if (!metal_graph_seed_streaming_expert_cache_from_hotlist(g, model, weights)) { + return false; + } + if (!metal_graph_seed_streaming_expert_cache_from_prefill(g, model, weights)) { + return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + uint32_t output_row = (uint32_t)n_tokens - 1u; + const char *output_row_env = glm_graph_env_value( + "DS4_ROCM_GRAPH_OUTPUT_ROW", + "DS4_METAL_GRAPH_OUTPUT_ROW"); + if (output_row_env && output_row_env[0]) { + char *end = NULL; + unsigned long v = strtoul(output_row_env, &end, 10); + if (end != output_row_env && v < (unsigned long)n_tokens) { + output_row = (uint32_t)v; + } + } + ds4_gpu_tensor *saved_cur = metal_graph_cur_hc(g); + ds4_gpu_tensor *last_hc = NULL; + + const double t_head0 = profile ? now_sec() : 0.0; + if (logits) { + last_hc = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), + output_row, + hc_dim); + ok = last_hc != NULL; + } + if (ok && logits && g->ssd_streaming) { + const bool static_decode_map = + metal_graph_stream_decode_static_map_enabled(); + const bool static_map_state_cache = + static_decode_map && + metal_graph_stream_decode_static_map_state_cache_enabled(); + g->streaming_static_decode_map_current = false; + if (static_map_state_cache) { + ok = metal_graph_stream_map_decode_static_all(model, weights); + if (ok) g->streaming_static_decode_map_current = true; + } else { + ok = metal_graph_stream_map_output(model, weights); + } + } + if (ok && logits) { + g->cur_hc_by_tier[g->active_tier] = last_hc; + ok = ds4_gpu_begin_commands() != 0; + } + if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + const double t_head_encoded = profile ? now_sec() : 0.0; + if (ok && logits) ok = ds4_gpu_end_commands() != 0; + const double t_head_done = profile ? now_sec() : 0.0; + g->cur_hc_by_tier[g->active_tier] = saved_cur; + if (last_hc) ds4_gpu_tensor_free(last_hc); + if (!ok) return false; + + const double t_before_read = profile ? now_sec() : 0.0; + if (logits) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + if (profile) { + const double t_read = now_sec(); + encode_s += t_head_encoded - t_head0; + execute_s += t_head_done - t_head_encoded; + if (split_profile) { + fprintf(stderr, + "ds4: gpu layer-major prefill head encode=%.3f ms execute=%.3f ms\n", + (t_head_encoded - t_head0) * 1000.0, + (t_head_done - t_head_encoded) * 1000.0); + } + fprintf(stderr, + "ds4: gpu layer-major prefill total tokens=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms\n", + n_tokens, + encode_s * 1000.0, + execute_s * 1000.0, + (t_read - t_before_read) * 1000.0, + (t_read - t0) * 1000.0); + } + return ok; +} + +static bool metal_graph_prefill_raw_swa( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + int n_tokens, + float *logits, + bool show_progress, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + ds4_session_cancel_fn cancel, + void *cancel_ud, + bool *cancelled) { + if (n_tokens <= 0 || n_tokens > prompt->len) return false; + if ((uint32_t)n_tokens > g->prefill_cap) return false; + if (metal_graph_use_streaming_decode_prefill_range(g, weights, 0, + (uint32_t)n_tokens)) { + return metal_graph_prefill_decode_streaming_range(g, + model, + weights, + prompt, + 0, + (uint32_t)n_tokens, + logits, + show_progress, + NULL, + NULL, + display_progress, + display_progress_ud, + cancel, + cancel_ud, + cancelled); + } + /* The layer-major fallback below may submit the whole short prefill as one + * Metal command buffer. Once that command is in flight there is no useful + * safe prefix to expose: by the time cancellation can be observed again, + * the prompt has already been fully read and the KV is valid. Let the + * caller observe the pending interrupt at generation time instead. */ + (void)cancel; + (void)cancel_ud; + (void)cancelled; + return metal_graph_prefill_layer_major(g, + model, + weights, + prompt, + 0, + (uint32_t)n_tokens, + logits, + show_progress, + NULL, + display_progress, + display_progress_ud); +} + +/* Prefill a contiguous token range in fixed-size chunks. + * + * The common case starts at token zero, but server sessions also use this to + * extend an existing KV cache with a long suffix. Resumed chunks are aligned + * to the same absolute prefill-cap boundaries used by a cold full prompt, so + * compression windows and row finalization follow the same schedule after the + * cached prefix. + */ +static bool metal_graph_prefill_chunked_range( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + float *logits, + bool show_progress, + ds4_session_progress_fn progress, + void *progress_ud, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + ds4_imatrix_collector *imatrix, + ds4_session_cancel_fn cancel, + void *cancel_ud, + bool *cancelled) { + if (n_tokens == 0 || g->prefill_cap == 0) return false; + if (start > (uint32_t)prompt->len) return false; + if (n_tokens > (uint32_t)prompt->len - start) return false; + if (g->ssd_streaming && start == 0) { + ds4_gpu_stream_expert_cache_reset_route_hotness(); + } + if (!imatrix && + metal_graph_use_streaming_decode_prefill_range(g, weights, + start, n_tokens)) { + return metal_graph_prefill_decode_streaming_range(g, + model, + weights, + prompt, + start, + n_tokens, + logits, + show_progress, + progress, + progress_ud, + display_progress, + display_progress_ud, + cancel, + cancel_ud, + cancelled); + } + + uint32_t chunk_cap = g->prefill_cap; + if (start != 0 && chunk_cap > g->raw_cap) chunk_cap = g->raw_cap; + if (chunk_cap == 0) return false; + + const bool profile = + glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", + "DS4_METAL_GRAPH_PREFILL_PROFILE"); + const double t0 = profile ? now_sec() : 0.0; + const uint32_t end = start + n_tokens; + + if (progress) { + progress(progress_ud, "prefill_chunk", (int)start, prompt->len); + } + if (display_progress) { + display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); + } + + for (uint32_t pos0 = start; pos0 < end; ) { + if (cancel && cancel(cancel_ud)) { + if (cancelled) *cancelled = true; + return true; + } + const uint32_t remaining = end - pos0; + uint32_t local_cap = chunk_cap; + if (start != 0 && g->prefill_cap != 0) { + const uint32_t mod = pos0 % g->prefill_cap; + if (mod != 0) { + const uint32_t to_boundary = g->prefill_cap - mod; + if (to_boundary < local_cap) local_cap = to_boundary; + } + } + const uint32_t chunk = remaining < local_cap ? remaining : local_cap; + const uint32_t chunk_end = pos0 + chunk; + float *chunk_logits = (progress || chunk_end == end) ? logits : NULL; + bool ok = metal_graph_prefill_layer_major(g, + model, + weights, + prompt, + pos0, + chunk, + chunk_logits, + show_progress, + imatrix, + display_progress, + display_progress_ud); + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after chunked prefill failure also failed\n"); + } + return false; + } + if (progress) { + progress(progress_ud, "prefill_chunk", (int)chunk_end, prompt->len); + } + if (display_progress) { + display_progress(display_progress_ud, "prefill_display", (int)chunk_end, prompt->len); + } + if (cancel && cancel(cancel_ud)) { + if (cancelled) *cancelled = true; + return true; + } + pos0 = chunk_end; + } + if (show_progress) fputc('\n', stderr); + if (profile) { + const double t_read = now_sec(); + fprintf(stderr, + "ds4: gpu chunked prefill start=%u tokens=%u chunk=%u total=%.3f ms\n", + start, + n_tokens, + chunk_cap, + (t_read - t0) * 1000.0); + } + return true; +} + +/* Long prompts are prefetched in fixed-size chunks. Chunks bound transient + * attention buffers while preserving the same final KV/cache state. */ +static bool metal_graph_prefill_chunked( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + int n_tokens, + float *logits, + bool show_progress, + ds4_session_progress_fn progress, + void *progress_ud, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + ds4_session_cancel_fn cancel, + void *cancel_ud, + bool *cancelled) { + if (n_tokens <= 0) return false; + return metal_graph_prefill_chunked_range(g, + model, + weights, + prompt, + 0, + (uint32_t)n_tokens, + logits, + show_progress, + progress, + progress_ud, + display_progress, + display_progress_ud, + NULL, + cancel, + cancel_ud, + cancelled); +} + +typedef struct ds4_verify_suffix_timing { + double upload_ms; + double layer_ms; + double head_ms; + double read_ms; + bool fused_head; +} ds4_verify_suffix_timing; + +static bool metal_graph_dspark_verify_selected_profile_enabled(void) { + return getenv("DS4_DSPARK_VERIFY_SELECTED_PROFILE") != NULL && + getenv("DS4_DSPARK_DISABLE_VERIFY_SELECTED_PROFILE") == NULL; +} + +/* Layer-major speculative target verifier for tiny MTP suffixes. + * + * This is the first production-shaped verifier attempt: unlike repeated decode + * it runs the target model layer-by-layer for the whole speculative suffix, and + * unlike the diagnostic path it does not read back full logits for every row. + * The verifier returns the row top-1 ids needed for acceptance. The caller + * then reads exactly one logits row: the row that becomes the new continuation + * state. It still reuses the existing batch layer kernels, so it is not yet + * the final hand-written N=2/N=4 decode microbatch, but it exercises the right + * verifier contract and removes the obvious diagnostic overheads first. */ +static bool metal_graph_verify_suffix_tops_impl( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + bool capture_prefix1, + bool capture_dspark_hidden, + int *row_tops, + float *row_logits, + ds4_verify_suffix_timing *timing) { + if (timing) memset(timing, 0, sizeof(*timing)); + if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; + if (start > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - start) return false; + const uint32_t top_rows = n_tokens > 1 ? n_tokens - 1 : 0; + if (top_rows && !row_tops) return false; + + const double upload_t0 = timing ? now_sec() : 0.0; + bool ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), prompt, start, n_tokens); + if (ok) ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), + metal_graph_prefill_tokens(g), + model, + weights, + prompt, + start, + n_tokens); + if (!ok) return false; + + const bool saved_capture = g->spec_capture_prefix1; + g->spec_capture_prefix1 = capture_prefix1 && n_tokens == 2; + const char *split_head_env = getenv("DS4_DSPARK_VERIFY_SPLIT_HEAD"); + const bool fuse_head = + !split_head_env || !split_head_env[0] || + strcmp(split_head_env, "0") == 0; + if (timing) timing->fused_head = fuse_head; + if (timing) timing->upload_ms += (now_sec() - upload_t0) * 1000.0; + + const bool selected_profile = + metal_graph_dspark_verify_selected_profile_enabled(); + if (selected_profile) { + metal_graph_stream_prefill_selected_profile_reset(g); + } + + /* Under TP, verify every speculative block against the two resident + * expert halves. Both ranks encode identically, so one batch gate per + * layer reconstructs the routed result while preserving their KV state. */ + g->tp_batch_rows = (g->tp_world == 2 && + g->tp_batch_out != NULL && g->tp_batch_in != NULL && + n_tokens <= (uint32_t)DS4_TP_BATCH_MAX_ROWS) + ? n_tokens : 0; + const double layer_t0 = timing ? now_sec() : 0.0; + ok = ds4_gpu_begin_commands() != 0; + const bool dspark_capture_active = + ok && + capture_dspark_hidden && + metal_graph_dspark_capture_verified_suffix_begin(g, + start, + n_tokens, + true); + static int verify_profile_left = -1; + if (verify_profile_left < 0) { + verify_profile_left = + getenv("DS4_DSPARK_VERIFY_PROFILE") != NULL ? 1 : 0; + } + const bool verify_profile = verify_profile_left > 0 && ok; + if (verify_profile) verify_profile_left--; + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + if (verify_profile) { + ok = ds4_gpu_end_commands() != 0; + if (ok) (void)ds4_gpu_synchronize(); + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (!ok) break; + } + ok = metal_graph_encode_layer_batch(g, + model, + &weights->layer[il], + il, + start, + n_tokens); + if (ok && dspark_capture_active) { + ok = metal_graph_dspark_capture_verified_suffix_layer(g, + il, + start, + n_tokens); + } + if (ok && selected_profile) { + ok = ds4_gpu_end_commands() != 0 && + metal_graph_selected_profile_layer_impl( + g, + &weights->layer[il], + il, + n_tokens, + "DSpark verifier selected profile") && + ds4_gpu_begin_commands() != 0; + } + } + g->tp_batch_rows = 0; + if (ok && fuse_head) { + ok = metal_graph_encode_output_head_batch(g, + model, + weights, + n_tokens, + weights->output->dim[1]); + } + if (ok && fuse_head) { + if (top_rows == 1) { + ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), + g->spec_logits, + DS4_N_VOCAB) != 0; + } else if (top_rows) { + ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + g->spec_logits, + DS4_N_VOCAB, + top_rows, + 1) != 0; + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + g->spec_capture_prefix1 = saved_capture; + if (!ok && dspark_capture_active) { + metal_graph_dspark_capture_invalidate(g); + } + if (timing) timing->layer_ms += (now_sec() - layer_t0) * 1000.0; + if (!ok) return false; + if (selected_profile) { + metal_graph_selected_profile_summary_impl( + g, + "DSpark verifier selected profile"); + } + + if (!fuse_head) { + const double head_t0 = timing ? now_sec() : 0.0; + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_output_head_batch(g, + model, + weights, + n_tokens, + weights->output->dim[1]); + if (ok) { + if (top_rows == 1) { + /* Common K=2 verify case: top_k=1 over n_vocab → use the dedicated + * argmax kernel (single-block tree-reduce) instead of the legacy + * indexer_topk_kernel's single-thread O(n_vocab * top_k) fall-through. */ + ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), + g->spec_logits, + DS4_N_VOCAB) != 0; + } else if (top_rows) { + /* top-1 of each of the top_rows rows: n_tokens=top_rows, top_k=1. + * The order is transposed vs the indexer-score callers; a swap + * silently scores row 0's runner-ups instead of each row. */ + ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + g->spec_logits, + DS4_N_VOCAB, + top_rows, + 1) != 0; + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (timing) timing->head_ms += (now_sec() - head_t0) * 1000.0; + } + const double read_t0 = timing ? now_sec() : 0.0; + if (ok && top_rows) { + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), + 0, + row_tops, + (uint64_t)top_rows * sizeof(row_tops[0])) != 0; + if (ok && getenv("DS4_DSPARK_VERIFY_TOPS_CHECK") != NULL) { + float *chk = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); + for (uint32_t r = 0; r < top_rows; r++) { + if (ds4_gpu_tensor_read(g->spec_logits, + (uint64_t)r * DS4_N_VOCAB * + sizeof(float), + chk, + (uint64_t)DS4_N_VOCAB * + sizeof(float)) == 0) break; + uint32_t am = 0; + for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { + if (chk[i] > chk[am]) am = i; + } + fprintf(stderr, + "ds4: verify tops-check row=%u gpu_top=%d cpu_argmax=%u " + "cpu_max=%.3f\n", + r, row_tops[r], am, chk[am]); + } + free(chk); + } + } + if (ok && row_logits) { + ok = ds4_gpu_tensor_read(g->spec_logits, + 0, + row_logits, + (uint64_t)n_tokens * DS4_N_VOCAB * sizeof(row_logits[0])) != 0; + } + if (timing) timing->read_ms += (now_sec() - read_t0) * 1000.0; + return ok; +} + +/* The verify block keeps the GPU genuinely busy, so the TP DVFS + * keep-alive is a pure parasite for its duration — pause it. */ +static bool metal_graph_verify_suffix_tops( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + uint32_t start, + uint32_t n_tokens, + bool capture_prefix1, + bool capture_dspark_hidden, + int *row_tops, + float *row_logits, + ds4_verify_suffix_timing *timing) { + ds4_gpu_tp_keepalive_pause(1); + const bool ok = metal_graph_verify_suffix_tops_impl(g, model, weights, + prompt, start, + n_tokens, + capture_prefix1, + capture_dspark_hidden, + row_tops, row_logits, + timing); + ds4_gpu_tp_keepalive_pause(0); + return ok; +} + +static bool metal_graph_read_spec_logits_row(ds4_gpu_graph *g, uint32_t row, float *logits) { + if (!g || !g->spec_logits || !logits || row >= g->prefill_cap) return false; + const uint64_t row_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); + return ds4_gpu_tensor_read(g->spec_logits, + (uint64_t)row * row_bytes, + logits, + row_bytes) != 0; +} + +/* Exact N=2 target verifier for MTP. + * + * The generic batch prefill path is fast, but it is not a safe substitute for + * autoregressive decode: small row-wise differences in HC/MoE/output kernels + * are enough to flip future greedy tokens. This verifier keeps the exact + * decode kernels and cache update order, but encodes the two proposed tokens + * layer-by-layer in one command stream. It returns the exact target top after + * token0, and exact logits after token1. */ +static bool metal_graph_verify_decode2_exact( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token0, + int token1, + uint32_t start, + int *top0, + int *top1, + float *logits0, + float *logits1) { + if (!g || !top0 || (!top1 && !logits1) || g->raw_cap == 0) return false; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t hc_bytes = hc_dim * sizeof(float); + ds4_gpu_tensor *cur0_by_tier[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor *cur1_by_tier[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor *next0_by_tier[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor *next1_by_tier[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor *saved_cur_by_tier[DS4_MAX_GPUS] = {0}; + ds4_gpu_tensor *saved_after_by_tier[DS4_MAX_GPUS] = {0}; + const int saved_active_tier = g->active_tier; + const bool saved_capture = g->spec_capture_prefix1; + + bool ok = true; + for (int t = 0; t < DS4_MAX_GPUS; t++) { + saved_cur_by_tier[t] = g->cur_hc_by_tier[t]; + saved_after_by_tier[t] = g->after_ffn_hc_by_tier[t]; + if (!g->batch_cur_hc_by_tier[t] && !g->batch_next_hc_by_tier[t]) continue; + if (!g->batch_cur_hc_by_tier[t] || !g->batch_next_hc_by_tier[t]) { + ok = false; + break; + } + cur0_by_tier[t] = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[t], 0, hc_dim); + cur1_by_tier[t] = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[t], 1, hc_dim); + next0_by_tier[t] = metal_graph_tensor_row_view(g->batch_next_hc_by_tier[t], 0, hc_dim); + next1_by_tier[t] = metal_graph_tensor_row_view(g->batch_next_hc_by_tier[t], 1, hc_dim); + if (!cur0_by_tier[t] || !cur1_by_tier[t] || + !next0_by_tier[t] || !next1_by_tier[t]) { + ok = false; + break; + } + } + + int cur_tier = g->emb_tier; + if (cur_tier < 0 || cur_tier >= DS4_MAX_GPUS || + !cur0_by_tier[cur_tier] || !cur1_by_tier[cur_tier]) { + ok = false; + } + if (ok) ok = metal_graph_set_active_tier_no_copy(g, cur_tier); + if (ok) ok = ds4_gpu_embed_token_hc_tensor(cur0_by_tier[cur_tier], + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token0, + DS4_N_EMBD, + DS4_N_HC) != 0; + if (ok) ok = ds4_gpu_embed_token_hc_tensor(cur1_by_tier[cur_tier], + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token1, + DS4_N_EMBD, + DS4_N_HC) != 0; + + g->spec_capture_prefix1 = true; + if (ok) ok = ds4_gpu_begin_commands() != 0; + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + const uint32_t pos0 = start; + const uint32_t pos1 = start + 1u; + const int this_tier = g->placement ? g->placement[il + 1] : cur_tier; + if (this_tier < 0 || this_tier >= DS4_MAX_GPUS || + !cur0_by_tier[this_tier] || !cur1_by_tier[this_tier]) { + ok = false; + break; + } + if (this_tier != cur_tier) { + ok = ds4_gpu_tensor_copy_xdev(cur0_by_tier[this_tier], + cur0_by_tier[cur_tier], + hc_bytes) != 0 && + ds4_gpu_tensor_copy_xdev(cur1_by_tier[this_tier], + cur1_by_tier[cur_tier], + hc_bytes) != 0; + if (!ok) break; + cur_tier = this_tier; + } + ok = metal_graph_set_active_tier_no_copy(g, this_tier); + if (!ok) break; + + g->cur_hc_by_tier[this_tier] = cur0_by_tier[this_tier]; + g->after_ffn_hc_by_tier[this_tier] = next0_by_tier[this_tier]; + ok = metal_graph_encode_decode_layer(g, + model, + &weights->layer[il], + il, + pos0, + g->layer_raw_cache[il], + g->raw_cap, + pos0 % g->raw_cap, + metal_graph_raw_span_for_batch(g, pos0, 1), + token0); + if (!ok) break; + ok = metal_graph_capture_prefix1_attn_state(g, il) && + metal_graph_capture_prefix1_index_state(g, il); + if (!ok) break; + + g->cur_hc_by_tier[this_tier] = cur1_by_tier[this_tier]; + g->after_ffn_hc_by_tier[this_tier] = next1_by_tier[this_tier]; + ok = metal_graph_encode_decode_layer(g, + model, + &weights->layer[il], + il, + pos1, + g->layer_raw_cache[il], + g->raw_cap, + pos1 % g->raw_cap, + metal_graph_raw_span_for_batch(g, pos1, 1), + token1); + if (!ok) break; + + ds4_gpu_tensor *tmp = cur0_by_tier[this_tier]; + cur0_by_tier[this_tier] = next0_by_tier[this_tier]; + next0_by_tier[this_tier] = tmp; + tmp = cur1_by_tier[this_tier]; + cur1_by_tier[this_tier] = next1_by_tier[this_tier]; + next1_by_tier[this_tier] = tmp; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + g->spec_capture_prefix1 = saved_capture; + + if (ok) { + ok = metal_graph_set_active_tier_no_copy(g, cur_tier); + } + if (ok) { + const bool split_top1 = + logits0 == NULL && + g->cuda_tp_output && + metal_graph_cuda_verify_decode2_split_top1_requested(); + uint32_t output_ways = 0; + g->cur_hc_by_tier[cur_tier] = cur0_by_tier[cur_tier]; + ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + if (ok) ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), + metal_graph_logits(g), + DS4_N_VOCAB) != 0; + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok && split_top1) { + ok = metal_graph_read_output_split_top1(g, output_ways, top0); + } else if (ok) { + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top0, sizeof(*top0)) != 0; + } + if (ok && logits0) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), + 0, + logits0, + (uint64_t)DS4_N_VOCAB * sizeof(logits0[0])) != 0; + } + } + + if (ok) { + ok = metal_graph_set_active_tier_no_copy(g, cur_tier); + } + if (ok) { + const bool split_top1 = + logits1 == NULL && + top1 != NULL && + g->cuda_tp_output && + metal_graph_cuda_verify_decode2_split_top1_requested(); + int output_tiers[DS4_MAX_GPUS] = {0}; + uint32_t output_ways = 0; + g->cur_hc_by_tier[cur_tier] = cur1_by_tier[cur_tier]; + ok = ds4_gpu_begin_commands() != 0; + if (ok && split_top1) { + ok = metal_graph_encode_output_head_split_top1(g, + model, + weights, + weights->output->dim[1], + output_tiers, + &output_ways); + } else if (ok) { + ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + if (ok && top1) { + ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_logits(g), + DS4_N_VOCAB, + 1, + 1) != 0; + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok && split_top1) { + ok = metal_graph_read_output_split_top1(g, output_ways, top1); + } else if (ok && top1) { + ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top1, sizeof(*top1)) != 0; + } + if (ok) { + if (logits1) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), + 0, + logits1, + (uint64_t)DS4_N_VOCAB * sizeof(logits1[0])) != 0; + } + } + } + g->spec_capture_prefix1 = saved_capture; + for (int t = 0; t < DS4_MAX_GPUS; t++) { + g->cur_hc_by_tier[t] = saved_cur_by_tier[t]; + g->after_ffn_hc_by_tier[t] = saved_after_by_tier[t]; + } + if (g->placement) { + if (saved_active_tier >= 0) { + (void)metal_graph_set_active_tier_no_copy(g, saved_active_tier); + } else { + g->active_tier = saved_active_tier; + } + } + for (int t = 0; t < DS4_MAX_GPUS; t++) { + ds4_gpu_tensor_free(next1_by_tier[t]); + ds4_gpu_tensor_free(next0_by_tier[t]); + ds4_gpu_tensor_free(cur1_by_tier[t]); + ds4_gpu_tensor_free(cur0_by_tier[t]); + } + return ok; +} + +/* Pick a raw SWA cache size for Metal. During batched prefill it must cover + * the previous window plus the current ubatch. */ +static uint32_t metal_graph_raw_cap_for_context(int ctx_size, uint32_t prefill_cap) { + uint32_t raw_window = DS4_N_SWA; + if (raw_window > (uint32_t)ctx_size) raw_window = (uint32_t)ctx_size; + if (raw_window == 0) raw_window = 1; + + /* + * During batched prefill the SWA cache must hold the current ubatch plus + * the previous logical window. The cache is padded to a 256-row multiple + * so the physical row order and FlashAttention block grouping match the + * model path we compare against. + */ + uint64_t wanted = (uint64_t)raw_window + prefill_cap; + if (wanted > (uint32_t)ctx_size) wanted = (uint32_t)ctx_size; + if (wanted == 0) wanted = 1; + wanted = align_up(wanted, 256u); + if (wanted > 8192u) wanted = 8192u; + uint32_t raw_cap = (uint32_t)wanted; + if (raw_cap < raw_window) raw_cap = raw_window; + +#ifndef DS4_ROCM_BUILD + const char *env = getenv("DS4_METAL_GRAPH_RAW_CAP"); + if (env && env[0]) { + char *endp = NULL; + const long v = strtol(env, &endp, 10); + if (endp != env && v > 0) { + raw_cap = (uint32_t)v; + if (raw_cap > (uint32_t)ctx_size) raw_cap = (uint32_t)ctx_size; + if (raw_cap > 8192u) raw_cap = 8192u; + if (raw_cap < raw_window) raw_cap = raw_window; + } + } +#endif + + return raw_cap; +} + +/* Choose the prefill ubatch size. Whole-batch is fastest for normal prompts. + * Long Flash prompts default to 4096-token chunks; PRO defaults to 8192. */ +static uint32_t metal_graph_prefill_cap_for_prompt(int prompt_len, + uint32_t prefill_chunk) { + return ds4_prefill_cap_for_prompt(prompt_len, prefill_chunk); +} + +/* When a server request shares a large prefix with the live checkpoint, extend + * the KV cache with batched prefill instead of single-token decode. On an M3 + * Max, prefill is faster from 2-token suffixes upward; keep the default at 4 + * as a conservative crossover. The env knob remains useful for retuning. */ +static uint32_t metal_graph_resume_prefill_min_tokens(void) { +#ifndef DS4_ROCM_BUILD + const char *env = getenv("DS4_METAL_RESUME_PREFILL_MIN"); + if (env && env[0]) { + char *endp = NULL; + const long v = strtol(env, &endp, 10); + if (endp != env) { + if (v <= 0) return UINT32_MAX; + return (uint32_t)v; + } + } +#endif + return 4u; +} + +static uint32_t glm_graph_resume_prefill_min_tokens(void) { +#ifndef DS4_ROCM_BUILD + const char *env = getenv("DS4_GLM_RESUME_PREFILL_MIN"); + if (env && env[0]) { + char *endp = NULL; + const long v = strtol(env, &endp, 10); + if (endp != env) { + if (v <= 0) return UINT32_MAX; + return (uint32_t)v; + } + } +#endif + return 4u; +} + +#define DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT 4096u +#define DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT 8192u +#define DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT 2048u +#define DS4_GLM_METAL_DISPLAY_PROGRESS_LAYER_TOKENS 32u +#define DS4_GLM_METAL_SMALL_PREFILL_STAGE_SYNC_TOKENS 0u +#define DS4_GLM_METAL_LONG_CONTEXT_THRESHOLD 65536u +#define DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT 4096u +#define DS4_GLM_METAL_INDEXED_PREFILL_CHUNK_TOKENS 4096u +#define DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB 256u + +static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, + bool ssd_streaming); +static uint32_t glm_graph_indexed_prefill_chunk_tokens( + uint32_t full_attention_cap, + uint32_t compact_cap); +static uint32_t glm_graph_indexed_prefill_score_tokens( + uint32_t indexed_prefill_cap, + uint32_t compact_cap); + +static uint64_t glm_graph_compact_cache_elem_bytes(void) { + return DS4_GPU_GLM_COMPACT_CACHE_F16 ? sizeof(uint16_t) : sizeof(float); +} + +static uint32_t glm_graph_compact_cache_is_f16(void) { + return DS4_GPU_GLM_COMPACT_CACHE_F16 ? 1u : 0u; +} + +static bool glm_graph_expanded_kv_cache_enabled(bool ssd_streaming) { + (void)ssd_streaming; + return false; +} + +static bool glm_graph_layer_uses_full_indexer(uint32_t il) { + if (il < DS4_N_LEADING_DENSE) return true; + return il >= 6u && ((il - 6u) % 4u) == 0u; +} + +static uint32_t glm_graph_normal_layer_count(void) { + if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT || DS4_N_LAYER > DS4_MAX_LAYER) { + return 0; + } + return DS4_N_LAYER - DS4_N_NEXTN_PREDICT; +} + +static uint32_t glm_graph_full_indexer_layer_count_range(uint32_t layer_start, + uint32_t layer_end) { + if (layer_start > layer_end) return 0; + uint32_t n = 0; + for (uint32_t il = layer_start; il <= layer_end; il++) { + if (glm_graph_layer_uses_full_indexer(il)) n++; + } + return n; +} + +static uint64_t glm_graph_full_kv_cache_elem_bytes(void) { + return sizeof(uint16_t); +} + +static uint32_t glm_graph_indexer_top_k_limit(void) { + return DS4_N_INDEXER_TOP_K; +} + +static uint32_t glm_tp_head_split_min(void) { + static int cached = -1; + if (cached < 0) { + cached = 64; + const char *env = getenv("DS4_GLM_TP_HEAD_SPLIT_MIN"); + if (env && env[0]) cached = atoi(env); + if (cached < 0) cached = 0; + } + return (uint32_t)cached; +} + +/* Correctness isolation: dump a hidden row (pre-output-norm), overwriting + * on each call — run with -n 0 so the file ends as the final prompt row. */ +static void glm_debug_dump_hidden_row(const ds4_gpu_tensor *t, uint32_t row) { + const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); + if (!path || !path[0] || !t) return; + float *buf = malloc((size_t)DS4_N_EMBD * sizeof(float)); + if (!buf) return; + if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, + (uint64_t)row * DS4_N_EMBD * sizeof(float), + buf, + (uint64_t)DS4_N_EMBD * sizeof(float))) { + FILE *f = fopen(path, "wb"); + if (f) { + fwrite(buf, sizeof(float), (size_t)DS4_N_EMBD, f); + fclose(f); + } + } + free(buf); +} + +/* Layer bisect: which layer's output hidden to dump (-1 = final/off, + * -2 = every layer, one file per layer). */ +static int glm_debug_hidden_dump_layer(void) { + const char *v = getenv("DS4_GLM_HIDDEN_DUMP_LAYER"); + if (!v || !v[0]) return -1; + if (strcmp(v, "all") == 0) return -2; + return atoi(v); +} + +static bool glm_debug_hidden_dump_layer_match(uint32_t il) { + const int dl = glm_debug_hidden_dump_layer(); + return dl == -2 || dl == (int)il; +} + +static void glm_debug_dump_raw_layer(const ds4_gpu_tensor *t, + const char *tag, + uint64_t bytes, + uint32_t il, + int pos) { + const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); + if (!path || !path[0] || !t) return; + char full[1024]; + if (pos >= 0) + snprintf(full, sizeof(full), "%s.%s.L%02u.T%02u", path, tag, il, pos); + else + snprintf(full, sizeof(full), "%s.%s.L%02u", path, tag, il); + void *buf = malloc(bytes); + if (!buf) return; + if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, 0, buf, bytes)) { + FILE *f = fopen(full, "wb"); + if (f) { fwrite(buf, 1, bytes, f); fclose(f); } + } + free(buf); +} + +static void glm_debug_dump_hidden_layer(const ds4_gpu_tensor *t, + uint32_t row, + uint32_t il, + uint32_t pos) { + const char *path = getenv("DS4_GLM_HIDDEN_DUMP"); + if (!path || !path[0] || !t) return; + char full[1024]; + snprintf(full, sizeof(full), "%s.L%02u.T%02u", path, il, pos); + float *buf = malloc((size_t)DS4_N_EMBD * sizeof(float)); + if (!buf) return; + if (ds4_gpu_tensor_read((ds4_gpu_tensor *)t, + (uint64_t)row * DS4_N_EMBD * sizeof(float), + buf, + (uint64_t)DS4_N_EMBD * sizeof(float))) { + FILE *f = fopen(full, "wb"); + if (f) { + fwrite(buf, sizeof(float), (size_t)DS4_N_EMBD, f); + fclose(f); + } + } + free(buf); +} + +/* Correctness isolation: dump the post-prefill logits vector once. */ +static void glm_debug_dump_prefill_logits(const float *logits) { + static int dumped; + const char *path = getenv("DS4_GLM_LOGIT_DUMP"); + if (!path || !path[0] || dumped || !logits) return; + FILE *f = fopen(path, "wb"); + if (!f) return; + fwrite(logits, sizeof(float), (size_t)DS4_N_VOCAB, f); + fclose(f); + dumped = 1; + fprintf(stderr, "ds4: prefill logits dumped to %s\n", path); +} + +static bool glm_graph_indexed_prefill_trace_enabled(void) { + return false; +} + +static bool glm_graph_indexed_prefill_trace_all(void) { + return false; +} + +static uint32_t glm_graph_indexed_prefill_trace_slow_ms(void) { + return 100u; +} + +static uint32_t glm_graph_indexed_prefill_drain_interval(void) { + return 16u; +} + +static bool glm_graph_full_prefill_trace_enabled(void) { + return false; +} + +static bool glm_graph_full_prefill_trace_all(void) { + return false; +} + +static uint32_t glm_graph_full_prefill_trace_slow_ms(void) { + return 100u; +} + +static uint32_t glm_graph_full_prefill_drain_interval(void) { + return 16u; +} + +static void glm_graph_full_prefill_tracef(const char *fmt, ...) { + va_list ap; + va_start(ap, fmt); + fprintf(stderr, "ds4: GLM full prefill trace "); + vfprintf(stderr, fmt, ap); + fputc('\n', stderr); + fflush(stderr); + va_end(ap); +} + +static void glm_graph_indexed_prefill_tracef(const char *fmt, ...) { + va_list ap; + va_start(ap, fmt); + fprintf(stderr, "ds4: GLM indexed prefill trace "); + vfprintf(stderr, fmt, ap); + fputc('\n', stderr); + fflush(stderr); + va_end(ap); +} + +static uint32_t glm_graph_compact_cache_initial_cap( + uint32_t ctx_size, + uint32_t full_attention_cap) { + if (ctx_size == 0) return 0; + if (ctx_size <= full_attention_cap) return ctx_size; + + uint32_t cap = ctx_size; + if (cap == 0) cap = full_attention_cap ? full_attention_cap : 1u; + if (cap > ctx_size) cap = ctx_size; + return cap; +} + +static uint64_t glm_graph_compact_cache_bytes_for_cap( + uint32_t normal_layers, + uint32_t indexer_layers, + uint32_t compact_cap) { + if (compact_cap == 0) return 0; + const uint64_t elem = glm_graph_compact_cache_elem_bytes(); + uint64_t total = + (uint64_t)normal_layers * + compact_cap * + ((uint64_t)DS4_N_KV_LORA + DS4_N_ROT) * + elem; + total += + (uint64_t)indexer_layers * + compact_cap * + DS4_N_INDEXER_HEAD_DIM * + elem; + return total; +} + +static uint64_t glm_graph_indexed_scratch_bytes_for_cap( + uint32_t full_attention_cap, + uint32_t compact_cap) { + if (compact_cap == 0) return 0; + const uint64_t indexed_rows = + glm_graph_indexed_prefill_chunk_tokens(full_attention_cap, compact_cap); + const uint64_t indexed_score_rows = + glm_graph_indexed_prefill_score_tokens((uint32_t)indexed_rows, + compact_cap); + const uint64_t indexer_q_elems = + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + const uint64_t qk_low_elems = + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA; + uint64_t bytes = (uint64_t)compact_cap * sizeof(float); + const uint64_t indexer_top_k = glm_graph_indexer_top_k_limit(); + bytes += indexed_score_rows * (uint64_t)compact_cap * sizeof(float); + bytes += indexed_rows * indexer_q_elems * sizeof(float); + bytes += indexed_rows * DS4_N_INDEXER_HEAD * sizeof(float); + bytes += indexed_rows * indexer_top_k * sizeof(uint32_t); /* batch_indexer_selected */ + bytes += indexed_rows * qk_low_elems * sizeof(float); /* batch_qk_low */ + bytes += indexed_rows * qk_low_elems * sizeof(float); /* batch_attn_lora */ + return bytes; +} + +static uint64_t glm_graph_workspace_add_bytes( + uint64_t total, + uint64_t count, + uint64_t elem_bytes) { + return ds4_add_sat_u64(total, ds4_mul_sat_u64(count, elem_bytes)); +} + +static uint32_t glm_graph_indexed_decode_split_blocks(void); + +static uint64_t glm_graph_workspace_bytes_for_cap( + uint32_t full_attention_cap, + uint32_t compact_cap, + bool ssd_streaming) { + const bool expanded_kv = + glm_graph_expanded_kv_cache_enabled(ssd_streaming); + const uint64_t indexed_rows = + compact_cap != 0 ? + glm_graph_indexed_prefill_chunk_tokens(full_attention_cap, + compact_cap) : + 0; + const uint64_t batch_rows = + expanded_kv || indexed_rows == 0 ? full_attention_cap : indexed_rows; + const uint64_t indexer_top_k = glm_graph_indexer_top_k_limit(); + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; + const uint64_t q_nope = + DS4_N_KEY_MLA > DS4_N_ROT ? (uint64_t)DS4_N_KEY_MLA - DS4_N_ROT : 0; + const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; + const uint64_t kv_raw_dim = (uint64_t)DS4_N_KV_LORA + DS4_N_ROT; + uint64_t dense_hidden_max = + DS4_N_FF_DENSE > DS4_N_FF_EXP ? DS4_N_FF_DENSE : DS4_N_FF_EXP; + if (dense_hidden_max == 0) dense_hidden_max = DS4_N_FF_EXP; + const uint64_t sparse_mid_elems = + (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; + const uint64_t ffn_mid_elems = + dense_hidden_max > sparse_mid_elems ? + dense_hidden_max : + sparse_mid_elems; + const uint64_t qk_low_elems = + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA; + const uint64_t split_attn_blocks = + glm_graph_indexed_decode_split_blocks(); + + uint64_t bytes = + glm_graph_indexed_scratch_bytes_for_cap(full_attention_cap, + compact_cap); + + bytes = glm_graph_workspace_add_bytes(bytes, 3u, DS4_N_EMBD * sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, 2u, DS4_N_LORA_Q * sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, q_dim, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_INDEXER_HEAD_DIM, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + (uint64_t)DS4_N_INDEXER_HEAD * + DS4_N_INDEXER_HEAD_DIM, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_INDEXER_HEAD, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, indexer_top_k, sizeof(uint32_t)); + bytes = glm_graph_workspace_add_bytes(bytes, qk_low_elems, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + split_attn_blocks * qk_low_elems, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + split_attn_blocks * + DS4_N_HEAD * 2u, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, kv_raw_dim, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_KV_LORA, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + DS4_N_HEAD * q_nope, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, 2u * heads_dim, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, 6u, DS4_N_EMBD * sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, 2u * dense_hidden_max, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, ffn_mid_elems, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT * 2u, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT_USED, sizeof(int32_t)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_EXPERT_USED, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_VOCAB, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + (uint64_t)DS4_N_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED, + sizeof(int32_t)); + + bytes = glm_graph_workspace_add_bytes(bytes, batch_rows, sizeof(int32_t)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_EXPERT * 2u, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_EXPERT_USED, + sizeof(int32_t)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_EXPERT_USED, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_EMBD * 7u, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_LORA_Q * 2u, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * q_dim, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_INDEXER_HEAD_DIM, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * kv_raw_dim, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_KV_LORA, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_HEAD * q_nope, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * heads_dim * 2u, sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * dense_hidden_max * 2u, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * DS4_N_FF_EXP, + sizeof(float)); + bytes = glm_graph_workspace_add_bytes(bytes, + batch_rows * ffn_mid_elems, + sizeof(float)); + if (indexed_rows != 0) { + bytes = glm_graph_workspace_add_bytes(bytes, + indexed_rows * qk_low_elems, + sizeof(float)); + } + return bytes; +} + +static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap_slice( + uint32_t ctx, + uint32_t work_ctx, + uint32_t compact_cap, + bool ssd_streaming, + uint32_t layer_start, + uint32_t layer_end) { + ds4_context_memory m = {0}; + const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0 || layer_start >= normal_layers || + layer_end < layer_start) { + return m; + } + if (layer_end >= normal_layers) layer_end = normal_layers - 1u; + const uint32_t layer_count = layer_end - layer_start + 1u; + if (compact_cap > ctx) compact_cap = ctx; + + const bool expanded_kv = glm_graph_expanded_kv_cache_enabled(ssd_streaming); + const uint32_t indexed_rows = + compact_cap != 0 ? + glm_graph_indexed_prefill_chunk_tokens(work_ctx, compact_cap) : + 0; + const uint32_t batch_rows = + expanded_kv || indexed_rows == 0 ? work_ctx : indexed_rows; + + m.prefill_cap = batch_rows; + m.raw_cap = expanded_kv ? work_ctx : 0; + if (expanded_kv) { + m.raw_bytes = (uint64_t)layer_count * + work_ctx * + ((uint64_t)DS4_N_HEAD * (DS4_N_KEY_MLA + DS4_N_VALUE_MLA)) * + glm_graph_full_kv_cache_elem_bytes(); + } + m.scratch_bytes = + glm_graph_workspace_bytes_for_cap(work_ctx, + compact_cap, + ssd_streaming); + if (compact_cap != 0) { + m.comp_cap = compact_cap; + m.compressed_bytes = + glm_graph_compact_cache_bytes_for_cap( + layer_count, + glm_graph_full_indexer_layer_count_range(layer_start, + layer_end), + compact_cap); + } + m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; + return m; +} + +static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( + uint32_t ctx, + uint32_t work_ctx, + uint32_t compact_cap, + bool ssd_streaming) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0) { + const ds4_context_memory empty = {0}; + return empty; + } + return glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx, + work_ctx, + compact_cap, + ssd_streaming, + 0, + normal_layers - 1u); +} + +ds4_context_memory ds4_context_memory_estimate_with_prefill_mode( + ds4_backend backend, + int ctx_size, + uint32_t prefill_chunk, + bool ssd_streaming) { + ds4_context_memory m = {0}; + uint32_t ctx = ctx_size > 0 ? (uint32_t)ctx_size : 1u; + + if (ds4_backend_uses_graph(backend)) { + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx, ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx, work_ctx); + m = glm_graph_context_memory_estimate_for_compact_cap(ctx, + work_ctx, + compact_cap, + ssd_streaming); + return m; + } + m.prefill_cap = metal_graph_prefill_cap_for_prompt((int)ctx, + prefill_chunk); + m.raw_cap = metal_graph_raw_cap_for_context((int)ctx, m.prefill_cap); + + uint32_t min_ratio = UINT32_MAX; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; + } + if (min_ratio == UINT32_MAX) min_ratio = ctx; + m.comp_cap = ctx / min_ratio + 2u; + if (m.comp_cap < 2u) m.comp_cap = 2u; + + m.raw_bytes = (uint64_t)DS4_N_LAYER * + m.raw_cap * + DS4_N_HEAD_DIM * + sizeof(float); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio == 0) continue; + const uint32_t layer_comp_cap = ctx / ratio + 2u; + m.compressed_bytes += (uint64_t)layer_comp_cap * + DS4_N_HEAD_DIM * + (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); + if (ratio == 4) { + m.compressed_bytes += (uint64_t)layer_comp_cap * + DS4_N_INDEXER_HEAD_DIM * + sizeof(float); + } + } + uint64_t attn_stage_cap = (uint64_t)(m.prefill_cap / min_ratio + 2u); + if (attn_stage_cap < 2u) attn_stage_cap = 2u; + m.scratch_bytes = 2ull * + m.comp_cap * + m.prefill_cap * + sizeof(float) + + attn_stage_cap * DS4_N_HEAD_DIM * sizeof(float); + } else { + m.raw_cap = ds4_default_raw_cap(ctx); + m.raw_bytes = (uint64_t)DS4_N_LAYER * + m.raw_cap * + DS4_N_HEAD_DIM * + sizeof(float); + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + if (ratio == 0) continue; + const uint32_t comp_cap = ctx / ratio + 2u; + if (ratio == 4) m.comp_cap = comp_cap; + m.compressed_bytes += (uint64_t)comp_cap * + DS4_N_HEAD_DIM * + sizeof(float); + if (ratio == 4) { + m.compressed_bytes += (uint64_t)comp_cap * + DS4_N_INDEXER_HEAD_DIM * + sizeof(float); + } + } + if (m.comp_cap == 0) m.comp_cap = ctx / 4u + 2u; + m.scratch_bytes = ((uint64_t)(m.raw_cap + m.comp_cap) * sizeof(float)) + + ((uint64_t)m.comp_cap * sizeof(float)) + + ((uint64_t)m.comp_cap * sizeof(bool)); + } + + m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; + return m; +} + +ds4_context_memory ds4_context_memory_estimate_with_prefill( + ds4_backend backend, + int ctx_size, + uint32_t prefill_chunk) { + return ds4_context_memory_estimate_with_prefill_mode(backend, + ctx_size, + prefill_chunk, + false); +} + +ds4_context_memory ds4_context_memory_estimate(ds4_backend backend, + int ctx_size) { + return ds4_context_memory_estimate_with_prefill(backend, ctx_size, 0); +} + +static int metal_graph_prompt_logits_test( + const ds4_model *model, + const ds4_weights *weights, + const token_vec *prompt, + int ctx_size) { + int n_test = prompt->len; + const char *n_test_env = getenv("DS4_METAL_GRAPH_PROMPT_TOKENS"); + if (n_test_env && n_test_env[0]) { + char *endp = NULL; + const long v = strtol(n_test_env, &endp, 10); + if (endp != n_test_env && v > 0 && v <= prompt->len) n_test = (int)v; + } + + if (n_test <= 0 || n_test > ctx_size) { + fprintf(stderr, "ds4: Metal graph prompt test needs 1..%d prompt tokens\n", ctx_size); + return 1; + } + + const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, (uint32_t)n_test); + + ds4_gpu_graph g; + /* diagnostic single-tier callsite; placement=NULL. */ + bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], + raw_cap, (uint32_t)ctx_size, + (uint32_t)n_test, false, NULL, false, NULL); + if (!ok) { + metal_graph_free(&g); + fprintf(stderr, "ds4: failed to initialize Metal graph prompt test runtime\n"); + return 1; + } + const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; + if (memory_report) ds4_gpu_print_memory_report("after graph alloc"); + + ds4_kv_cache cpu_cache; + kv_cache_init(&cpu_cache, (uint32_t)ctx_size, raw_cap); + float *cpu_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); + float *gpu_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); + float *oracle_logits = NULL; + + const char *oracle_path = getenv("DS4_ORACLE_LOGITS"); + if (oracle_path && oracle_path[0]) { + oracle_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); + if (!read_f32_binary_file(oracle_path, oracle_logits, DS4_N_VOCAB)) { + free(oracle_logits); + oracle_logits = NULL; + } + } + + for (int t = 0; t < n_test; t++) { + const bool last = t == n_test - 1; + forward_token_raw_swa_cpu(last ? cpu_logits : NULL, + model, + weights, + &cpu_cache, + prompt->v[t], + (uint32_t)t); + } + ok = metal_graph_prefill_raw_swa(&g, model, weights, prompt, n_test, + gpu_logits, true, NULL, NULL, + NULL, NULL, NULL); + if (memory_report) ds4_gpu_print_memory_report("after prompt graph"); + + if (ok) { + const char *dump_gpu = getenv("DS4_METAL_GRAPH_DUMP_LOGITS"); + if (dump_gpu && dump_gpu[0]) { + if (write_f32_binary_file(dump_gpu, gpu_logits, DS4_N_VOCAB)) { + fprintf(stderr, "ds4: wrote Metal graph logits to %s\n", dump_gpu); + } + } + const char *dump_cpu = getenv("DS4_CPU_DUMP_LOGITS"); + if (dump_cpu && dump_cpu[0]) { + if (write_f32_binary_file(dump_cpu, cpu_logits, DS4_N_VOCAB)) { + fprintf(stderr, "ds4: wrote CPU logits to %s\n", dump_cpu); + } + } + if (getenv("DS4_METAL_GRAPH_TRACE_CACHE") != NULL || + getenv("DS4_METAL_GRAPH_TRACE_COMP") != NULL) { + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + const uint32_t n_raw = cpu_cache.layer[il].n_raw; + if (n_raw != 0) { + const uint64_t raw_phys_n = (uint64_t)raw_cap * DS4_N_HEAD_DIM; + const uint64_t raw_logical_n = (uint64_t)n_raw * DS4_N_HEAD_DIM; + const uint32_t raw_start = n_raw < raw_cap ? 0u : ((uint32_t)n_test % raw_cap); + float *gpu_raw_phys = xmalloc((size_t)raw_phys_n * sizeof(float)); + float *gpu_raw_logical = xmalloc((size_t)raw_logical_n * sizeof(float)); + if (ds4_gpu_tensor_read(g.layer_raw_cache[il], 0, gpu_raw_phys, raw_phys_n * sizeof(float)) != 0) { + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t phys = (raw_start + r) % raw_cap; + memcpy(gpu_raw_logical + (uint64_t)r * DS4_N_HEAD_DIM, + gpu_raw_phys + (uint64_t)phys * DS4_N_HEAD_DIM, + (size_t)DS4_N_HEAD_DIM * sizeof(float)); + } + fprintf(stderr, + "ds4: cache trace layer %u raw_n=%u raw_start=%u raw_max=%g raw_rms=%g\n", + il, n_raw, raw_start, + max_abs_diff(cpu_cache.layer[il].raw_kv, gpu_raw_logical, raw_logical_n), + rms_abs_diff(cpu_cache.layer[il].raw_kv, gpu_raw_logical, raw_logical_n)); + } + free(gpu_raw_logical); + free(gpu_raw_phys); + } + + const uint32_t n_comp = cpu_cache.layer[il].n_comp; + if (n_comp == 0) continue; + const uint64_t n = (uint64_t)n_comp * DS4_N_HEAD_DIM; + float *gpu_comp = xmalloc((size_t)n * sizeof(float)); + bool comp_read = false; + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + uint16_t *gpu_comp_h = xmalloc((size_t)n * sizeof(uint16_t)); + if (ds4_gpu_tensor_read(g.layer_attn_comp_cache[il], 0, + gpu_comp_h, n * sizeof(uint16_t)) != 0) { + for (uint64_t i = 0; i < n; i++) gpu_comp[i] = f16_to_f32(gpu_comp_h[i]); + comp_read = true; + } + free(gpu_comp_h); + } else { + comp_read = ds4_gpu_tensor_read(g.layer_attn_comp_cache[il], 0, + gpu_comp, n * sizeof(float)) != 0; + } + if (comp_read) { + fprintf(stderr, + "ds4: comp trace layer %u n=%u attn_max=%g attn_rms=%g\n", + il, n_comp, + max_abs_diff(cpu_cache.layer[il].attn_comp_kv, gpu_comp, n), + rms_abs_diff(cpu_cache.layer[il].attn_comp_kv, gpu_comp, n)); + } + free(gpu_comp); + + const uint32_t n_index = cpu_cache.layer[il].n_index_comp; + if (n_index != 0 && g.layer_index_comp_cache[il]) { + const uint64_t ni = (uint64_t)n_index * DS4_N_INDEXER_HEAD_DIM; + float *gpu_index = xmalloc((size_t)ni * sizeof(float)); + if (ds4_gpu_tensor_read(g.layer_index_comp_cache[il], 0, gpu_index, ni * sizeof(float)) != 0) { + fprintf(stderr, + "ds4: comp trace layer %u n=%u index_max=%g index_rms=%g\n", + il, n_index, + max_abs_diff(cpu_cache.layer[il].index_comp_kv, gpu_index, ni), + rms_abs_diff(cpu_cache.layer[il].index_comp_kv, gpu_index, ni)); + } + free(gpu_index); + } + } + } + const uint64_t cpu_top = argmax_f32(cpu_logits, DS4_N_VOCAB); + const uint64_t gpu_top = argmax_f32(gpu_logits, DS4_N_VOCAB); + fprintf(stderr, + "ds4: Metal prompt graph logits: tokens=%d logits_max=%g logits_rms=%g cpu_top=%llu gpu_top=%llu cpu_top_logit=%g gpu_top_logit=%g\n", + n_test, + max_abs_diff(cpu_logits, gpu_logits, DS4_N_VOCAB), + rms_abs_diff(cpu_logits, gpu_logits, DS4_N_VOCAB), + (unsigned long long)cpu_top, + (unsigned long long)gpu_top, + cpu_logits[cpu_top], + gpu_logits[gpu_top]); + if (oracle_logits) { + const uint64_t oracle_top = argmax_f32(oracle_logits, DS4_N_VOCAB); + fprintf(stderr, + "ds4: oracle logits: tokens=%d oracle_top=%llu oracle_top_logit=%g cpu_max=%g cpu_rms=%g metal_max=%g metal_rms=%g\n", + n_test, + (unsigned long long)oracle_top, + oracle_logits[oracle_top], + max_abs_diff(cpu_logits, oracle_logits, DS4_N_VOCAB), + rms_abs_diff(cpu_logits, oracle_logits, DS4_N_VOCAB), + max_abs_diff(gpu_logits, oracle_logits, DS4_N_VOCAB), + rms_abs_diff(gpu_logits, oracle_logits, DS4_N_VOCAB)); + } + } else { + fprintf(stderr, "ds4: Metal prompt graph logits test failed\n"); + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after prompt graph failure also failed\n"); + } + } + + free(gpu_logits); + free(cpu_logits); + free(oracle_logits); + kv_cache_free(&cpu_cache); + metal_graph_free(&g); + return ok ? 0 : 1; +} diff --git a/models/deepseek/metal/host/attention.inc b/models/deepseek/metal/host/attention.inc new file mode 100644 index 0000000000..f194aea4ff --- /dev/null +++ b/models/deepseek/metal/host/attention.inc @@ -0,0 +1,6927 @@ +int ds4_gpu_store_raw_kv_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!raw_cache || !kv || raw_cap == 0 || row >= raw_cap || head_dim == 0 || raw_cap > INT32_MAX) return 0; + + @autoreleasepool { + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal raw KV store received undersized destination buffer\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t row_i32 = (int32_t)row; + if (!ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, head_dim) || + !ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, + g_raw_store_round_buffer, + 0, + &row_i32, + 1, + raw_cap, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV DS4 set_rows store")) return 0; + } + + return 1; +} + +/* Release decode fused KV finalizer. Reference paths are selected by the C + * graph driver; this Objective-C entry point always means "use the fused + * Metal kernel." */ +int ds4_gpu_kv_fp8_store_raw_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim, + uint32_t n_rot) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 || + n_rot > head_dim || raw_cap > INT32_MAX) { + return 0; + } + + @autoreleasepool { + id kvbuf = ds4_gpu_tensor_buffer(kv); + id rawbuf = ds4_gpu_tensor_buffer(raw_cache); + const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (!kvbuf || !rawbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal fused KV FP8/raw-store received undersized buffers\n"); + return 0; + } + + ds4_gpu_dsv4_kv_fp8_store_args args = { + .head_dim = (int32_t)head_dim, + .n_rot = (int32_t)n_rot, + .raw_row = (int32_t)row, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_kv_fp8_store_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:2]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "KV FP8/raw-store fused")) return 0; + } + + return 1; +} + +int ds4_gpu_store_raw_kv_batch_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t pos0, + uint32_t n_tokens, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!raw_cache || !kv || raw_cap == 0 || n_tokens == 0 || head_dim == 0 || raw_cap > INT32_MAX) return 0; + + @autoreleasepool { + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal raw KV batch store received undersized destination buffer\n"); + return 0; + } + + int32_t rows_stack[512]; + int32_t *rows = rows_stack; + if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { + rows = malloc((size_t)n_tokens * sizeof(*rows)); + if (!rows) { + fprintf(stderr, "ds4: failed to allocate raw KV set_rows index list\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tokens; t++) { + rows[t] = (int32_t)((pos0 + t) % raw_cap); + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { + if (rows != rows_stack) free(rows); + return 0; + } + + const uint64_t n = (uint64_t)n_tokens * head_dim; + const int ok = n <= UINT32_MAX && + ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, (uint32_t)n) && + ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, + g_raw_store_round_buffer, + 0, + rows, + n_tokens, + raw_cap, + head_dim); + if (rows != rows_stack) free(rows); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV batch DS4 set_rows store")) return 0; + } + + return 1; +} + +static int ds4_gpu_encode_compressor_score_with_ape( + id cb, + id score_src, + NSUInteger score_src_offset, + id score_dst, + NSUInteger score_dst_offset, + id apebuf, + NSUInteger ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens) { + if (!cb || !score_src || !score_dst || !apebuf || + width == 0 || ratio == 0 || n_tokens == 0 || + (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + const uint64_t total_elems64 = (uint64_t)n_tokens * width; + if (total_elems64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal compressor APE add received too many elements\n"); + return 0; + } + const uint32_t total_elems = (uint32_t)total_elems64; + + const bool force_fused = + getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL; + const bool use_fused = + (ds4_gpu_device_name_contains("M3") || force_fused) && + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL; + if (use_fused) { + id pipeline = ds4_gpu_get_pipeline( + ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16" + : "kernel_dsv4_compressor_score_ape_f32"); + if (pipeline) { + ds4_gpu_dsv4_compressor_score_ape_args args = { + .width = width, + .ratio = ratio, + .pos0 = pos0, + .n_tokens = n_tokens, + }; + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth == 0) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:score_src offset:score_src_offset atIndex:1]; + [enc setBuffer:apebuf offset:ape_offset atIndex:2]; + [enc setBuffer:score_dst offset:score_dst_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)total_elems + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + if (force_fused) return 0; + } + + const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, + &g_compressor_store_ape_bytes, + scratch_bytes, + "ds4_compressor_store_ape")) { + return 0; + } + + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + uint32_t copied_rows = 0; + uint32_t pos_mod = pos0 % ratio; + while (copied_rows < n_tokens) { + uint32_t seg_rows = ratio - pos_mod; + if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; + const uint32_t seg_elems = seg_rows * width; + const NSUInteger src_off = ape_offset + (NSUInteger)pos_mod * width * elem_ape; + const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); + int ok; + if (ape_type == 1u) { + ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } else { + ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } + if (!ok) return 0; + copied_rows += seg_rows; + pos_mod = 0; + } + + return ds4_gpu_encode_add_f32_1d(cb, + score_src, + score_src_offset, + g_compressor_store_ape_buffer, + 0, + score_dst, + score_dst_offset, + total_elems); +} + +static int ds4_gpu_encode_compressor_set_rows_projected( + id cb, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id apebuf, + NSUInteger ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos0, + const int32_t *rows, + uint32_t n_rows, + uint32_t state_rows) { + if (!cb || !state_kv || !state_score || !kvbuf || !scorebuf || + !apebuf || !rows || width == 0 || n_rows == 0 || state_rows == 0) { + return 0; + } + + const NSUInteger score_scratch_bytes = (NSUInteger)n_rows * width * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_scratch_bytes, + "ds4_compressor_store_score")) { + return 0; + } + + return ds4_gpu_encode_compressor_score_with_ape(cb, + scorebuf, + score_offset, + g_compressor_store_score_buffer, + 0, + apebuf, + ape_offset, + ape_type, + width, + ratio, + pos0, + n_rows) && + ds4_gpu_encode_set_rows_f32_i32(cb, + state_kv, + kvbuf, + kv_offset, + rows, + n_rows, + state_rows, + width) && + ds4_gpu_encode_set_rows_f32_i32(cb, + state_score, + g_compressor_store_score_buffer, + 0, + rows, + n_rows, + state_rows, + width); +} + +static int ds4_gpu_compressor_store_one_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos) { + if (!kv || !sc || !state_kv || !state_score || !model_map || + width == 0 || ratio == 0 || (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_dsv4_compressor_store_one_pipeline, + "kernel_dsv4_compressor_store_one"); + if (!pipeline) return 0; + + const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * row_bytes; + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + ds4_gpu_tensor_bytes(kv) < row_bytes || + ds4_gpu_tensor_bytes(sc) < row_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, + ape_offset, ape_bytes, + &ape_inner); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!apebuf || !kvbuf || !scbuf || !statekvbuf || !statescbuf) return 0; + + ds4_gpu_dsv4_compressor_store_one_args args = { + .width = width, + .ratio = ratio, + .pos = pos, + .ape_type = ape_type, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const NSUInteger nth = 256u; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; + [enc setBuffer:scbuf offset:ds4_gpu_tensor_offset(sc) atIndex:2]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:3]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:4]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)width + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return ds4_gpu_finish_command_buffer(cb, owned, "compressor one-row store"); +} + +int ds4_gpu_compressor_store_batch_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv || !sc || !state_kv || !state_score || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { + fprintf(stderr, "ds4: Metal compressor batch APE range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + if (!kvbuf || !scbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + fprintf(stderr, "ds4: Metal compressor batch store received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const uint64_t total_elems64 = (uint64_t)n_tokens * width; + if (total_elems64 > UINT32_MAX || state_rows > INT32_MAX) { + fprintf(stderr, "ds4: Metal compressor batch store received too many elements\n"); + return 0; + } + const uint32_t total_elems = (uint32_t)total_elems64; + const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, + &g_compressor_store_ape_bytes, + scratch_bytes, + "ds4_compressor_store_ape") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + scratch_bytes, + "ds4_compressor_store_score")) { + return 0; + } + + int32_t rows_stack[16]; + int32_t *rows = rows_stack; + if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { + rows = malloc((size_t)n_tokens * sizeof(*rows)); + if (!rows) { + fprintf(stderr, "ds4: failed to allocate compressor set_rows index list\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tokens; t++) { + const uint32_t pos_mod = (pos0 + t) % ratio; + rows[t] = (int32_t)(ratio == 4u ? ratio + pos_mod : pos_mod); + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { + if (rows != rows_stack) free(rows); + return 0; + } + + int ok = 1; + uint32_t copied_rows = 0; + uint32_t pos_mod = pos0 % ratio; + while (ok && copied_rows < n_tokens) { + uint32_t seg_rows = ratio - pos_mod; + if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; + const uint32_t seg_elems = seg_rows * width; + const NSUInteger src_off = (NSUInteger)ape_inner + + (NSUInteger)pos_mod * width * elem_ape; + const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); + if (ape_type == 1u) { + ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } else { + ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } + copied_rows += seg_rows; + pos_mod = 0; + } + + if (ok) { + ok = ds4_gpu_encode_add_f32_1d(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_ape_buffer, + 0, + g_compressor_store_score_buffer, + 0, + total_elems); + } + if (ok) { + ok = ds4_gpu_encode_set_rows_f32_i32(cb, + state_kv, + kvbuf, + ds4_gpu_tensor_offset(kv), + rows, + n_tokens, + state_rows, + width); + } + if (ok) { + ok = ds4_gpu_encode_set_rows_f32_i32(cb, + state_score, + g_compressor_store_score_buffer, + 0, + rows, + n_tokens, + state_rows, + width); + } + if (rows != rows_stack) free(rows); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor batch DS4 store")) return 0; + } + + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_contiguous_3d_args( + uint32_t cols, + uint32_t rows, + uint32_t planes) { + const uint64_t row_bytes = (uint64_t)cols * sizeof(float); + const uint64_t plane_bytes = (uint64_t)rows * row_bytes; + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)cols, + .ne01 = (int32_t)rows, + .ne02 = (int32_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = plane_bytes, + .nb03 = (uint64_t)planes * plane_bytes, + .ne10 = (int32_t)cols, + .ne11 = (int32_t)rows, + .ne12 = (int32_t)planes, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = row_bytes, + .nb12 = plane_bytes, + .nb13 = (uint64_t)planes * plane_bytes, + .ne0 = (int32_t)cols, + .ne1 = (int32_t)rows, + .ne2 = (int32_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = (uint64_t)planes * plane_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_softmax_f32_contiguous( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows, + uint32_t planes) { + if (!cb || !src || !dst || width == 0 || rows == 0 || planes == 0) return 0; + + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t plane_bytes = (uint64_t)rows * row_bytes; + ds4_gpu_softmax_args args = { + .ne00 = (int32_t)width, + .ne01 = (int32_t)rows, + .ne02 = (int32_t)planes, + .nb01 = row_bytes, + .nb02 = plane_bytes, + .nb03 = (uint64_t)planes * plane_bytes, + .ne11 = (int32_t)width, + .ne12 = (int32_t)rows, + .ne13 = (int32_t)planes, + .nb11 = row_bytes, + .nb12 = plane_bytes, + .nb13 = (uint64_t)planes * plane_bytes, + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = (uint64_t)planes * plane_bytes, + .scale = 1.0f, + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 1, + }; + + id pipeline = + (width % 4u) == 0 ? g_soft_max_f32_4_pipeline : g_soft_max_f32_pipeline; + if (!pipeline) return 0; + + NSUInteger nth = 32u; + if ((width % 4u) == 0) { + while (nth < (NSUInteger)(width / 4u) && + nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { + nth *= 2u; + } + } else { + while (nth < (NSUInteger)width && + nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { + nth *= 2u; + } + } + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:src offset:src_off atIndex:2]; + [enc setBuffer:src offset:src_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + uint64_t kv_nb0, + uint64_t kv_nb1, + uint64_t kv_nb2, + id scorebuf, + NSUInteger score_offset, + uint64_t score_nb0, + uint64_t score_nb1, + uint64_t score_nb2, + uint32_t n_rows, + uint32_t head_dim) { + id outbuf = ds4_gpu_tensor_buffer(out); + if (!cb || !outbuf || !kvbuf || !scorebuf || n_rows == 0 || head_dim == 0 || + ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float)) { + return 0; + } + + const NSUInteger pack_bytes = (NSUInteger)n_rows * head_dim * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer, + &g_compressor_pool_product_bytes, + pack_bytes, + "ds4_compressor_pool_product") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer, + &g_compressor_pool_score_cont_bytes, + pack_bytes, + "ds4_compressor_pool_score_cont") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer, + &g_compressor_pool_softmax_bytes, + pack_bytes, + "ds4_compressor_pool_softmax")) { + return 0; + } + + const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float); + const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride; + + /* + * Keep the n_comp == 1 compressor path as the unfused graph sequence: + * + * score = soft_max(contiguous(score)) + * pooled = sum_rows(contiguous(kv) * score) + * + * The fused DS4 pool kernel is mathematically equivalent, but it reduces in + * a different order. That is enough to create ~1e-6 compressor differences + * and later FP8/routing flips, so this path intentionally keeps the same + * operation boundary and memory layout as the graph. + */ + ds4_gpu_bin_args mul_args = + ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1); + + return + ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + kvbuf, + kv_offset, + g_compressor_pool_product_buffer, + 0, + n_rows, + head_dim, + 1, + kv_nb0, + kv_nb1, + kv_nb2, + cont_row_stride, + cont_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + scorebuf, + score_offset, + g_compressor_pool_score_cont_buffer, + 0, + n_rows, + head_dim, + 1, + score_nb0, + score_nb1, + score_nb2, + cont_row_stride, + cont_plane_stride) && + ds4_gpu_encode_softmax_f32_contiguous(cb, + g_compressor_pool_score_cont_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + n_rows, + head_dim, + 1) && + ds4_gpu_encode_bin_f32_rows(cb, + g_mul_pipeline, + &mul_args, + g_compressor_pool_product_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + g_compressor_pool_product_buffer, + 0) && + ds4_gpu_encode_sum_rows_f32(cb, + g_compressor_pool_product_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out), + n_rows, + head_dim); +} + +static int ds4_gpu_encode_dsv4_softmax_pool( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + uint64_t kv_nb0, + uint64_t kv_nb1, + uint64_t kv_nb2, + id scorebuf, + NSUInteger score_offset, + uint64_t score_nb0, + uint64_t score_nb1, + uint64_t score_nb2, + uint32_t n_rows, + uint32_t head_dim, + uint32_t n_comp) { + id outbuf = ds4_gpu_tensor_buffer(out); + if (!cb || !outbuf || !kvbuf || !scorebuf || + n_rows == 0 || head_dim == 0 || n_comp == 0 || + ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * n_comp * sizeof(float)) { + return 0; + } + + if (n_comp == 1) { + return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(cb, + out, + kvbuf, + kv_offset, + kv_nb0, + kv_nb1, + kv_nb2, + scorebuf, + score_offset, + score_nb0, + score_nb1, + score_nb2, + n_rows, + head_dim); + } + + ds4_gpu_dsv4_softmax_pool_args args = { + .ne00 = (int64_t)n_rows, + .ne01 = (int64_t)head_dim, + .ne02 = (int64_t)n_comp, + .nb00 = kv_nb0, + .nb01 = kv_nb1, + .nb02 = kv_nb2, + .nb10 = score_nb0, + .nb11 = score_nb1, + .nb12 = score_nb2, + .ne0 = (int64_t)head_dim, + .ne1 = (int64_t)n_comp, + .nb0 = sizeof(float), + .nb1 = (uint64_t)head_dim * sizeof(float), + }; + const uint64_t n = (uint64_t)head_dim * n_comp; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_softmax_pool_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_concat_f32_dim1( + id cb, + id src0, + NSUInteger src0_offset, + uint32_t src0_rows, + uint64_t src0_row_stride, + id src1, + NSUInteger src1_offset, + uint32_t src1_rows, + uint64_t src1_row_stride, + id dst, + NSUInteger dst_offset, + uint32_t cols, + uint64_t dst_row_stride) { + if (!cb || !src0 || !src1 || !dst || cols == 0 || src0_rows == 0 || src1_rows == 0) { + return 0; + } + + const uint32_t rows = src0_rows + src1_rows; + const uint64_t src0_plane = (uint64_t)src0_rows * src0_row_stride; + const uint64_t src1_plane = (uint64_t)src1_rows * src1_row_stride; + const uint64_t dst_plane = (uint64_t)rows * dst_row_stride; + ds4_gpu_concat_args args = { + .ne00 = (int32_t)cols, + .ne01 = (int32_t)src0_rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src0_row_stride, + .nb02 = src0_plane, + .nb03 = src0_plane, + .ne10 = (int32_t)cols, + .ne11 = (int32_t)src1_rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_row_stride, + .nb12 = src1_plane, + .nb13 = src1_plane, + .ne0 = (int32_t)cols, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane, + .nb3 = dst_plane, + .dim = 1, + }; + + NSUInteger nth = cols < 1024u ? (NSUInteger)cols : 1024u; + const NSUInteger max_threads = g_concat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_concat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src0 offset:src0_offset atIndex:1]; + [enc setBuffer:src1 offset:src1_offset atIndex:2]; + [enc setBuffer:dst offset:dst_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) { + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_compressor_pack_ratio4_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static bool ds4_gpu_buffer_ranges_overlap( + id a, + NSUInteger a_offset, + uint64_t a_bytes, + id b, + NSUInteger b_offset, + uint64_t b_bytes) { + if (a != b || a_bytes == 0u || b_bytes == 0u) return false; + if (a_offset <= b_offset) { + return a_bytes > (uint64_t)(b_offset - a_offset); + } + return b_bytes > (uint64_t)(a_offset - b_offset); +} + +static int ds4_gpu_compressor_ratio4_direct_pool_mode( + uint32_t head_dim, + uint32_t n_comp) { + // One compressed row intentionally uses the legacy GGML reduction graph. + if (n_comp <= 1u) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static int ds4_gpu_encode_compressor_ratio4_direct_pool( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t n = (uint64_t)head_dim * n_comp; + if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + head_dim == 0u || n_comp <= 1u || n > UINT32_MAX || + ds4_gpu_tensor_bytes(out) < n * sizeof(float)) { + return 0; + } + + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_softmax_pool_ratio4_direct_pipeline, + "kernel_dsv4_softmax_pool_ratio4_direct"); + if (!pipeline) return 0; + + ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = { + .n_rows = 8, + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .pad = 0u, + }; + + NSUInteger nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0u) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_compressor_pack_ratio4( + id cb, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + !g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer || + head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) { + return 0; + } + + const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim; + if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0; + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_pack_ratio4_pipeline, + "kernel_dsv4_compressor_pack_ratio4"); + if (!pipeline) return 0; + + NSUInteger nth = head_dim; + if (nth > 256u) nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0) return 0; + + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .n_threads = (uint32_t)nth, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5]; + [enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_compressor_pool( + id cb, + ds4_gpu_tensor *out, + const ds4_gpu_tensor *state_kv, + const ds4_gpu_tensor *state_score, + uint32_t head_dim, + uint32_t ratio) { + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0; + + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t rows = coff * ratio; + const uint64_t state_bytes = (uint64_t)width * rows * sizeof(float); + if (ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return 0; + } + + if (ratio != 4u) { + const uint64_t row_stride = (uint64_t)width * sizeof(float); + return ds4_gpu_encode_dsv4_softmax_pool(cb, + out, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + row_stride, + sizeof(float), + (uint64_t)rows * row_stride, + statescbuf, + ds4_gpu_tensor_offset(state_score), + row_stride, + sizeof(float), + (uint64_t)rows * row_stride, + ratio, + head_dim, + 1); + } + + const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + packed_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + packed_bytes, + "ds4_compressor_pool_score")) { + return 0; + } + + const uint64_t state_row_stride = (uint64_t)width * sizeof(float); + const uint64_t pool_row_stride = (uint64_t)head_dim * sizeof(float); + const NSUInteger curr_offset = (NSUInteger)4u * state_row_stride + + (NSUInteger)head_dim * sizeof(float); + if (!ds4_gpu_encode_concat_f32_dim1(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + 4, + state_row_stride, + statekvbuf, + ds4_gpu_tensor_offset(state_kv) + curr_offset, + 4, + state_row_stride, + g_compressor_pool_kv_buffer, + 0, + head_dim, + pool_row_stride) || + !ds4_gpu_encode_concat_f32_dim1(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + 4, + state_row_stride, + statescbuf, + ds4_gpu_tensor_offset(state_score) + curr_offset, + 4, + state_row_stride, + g_compressor_pool_score_buffer, + 0, + head_dim, + pool_row_stride)) { + return 0; + } + + return ds4_gpu_encode_dsv4_softmax_pool(cb, + out, + g_compressor_pool_kv_buffer, + 0, + pool_row_stride, + sizeof(float), + packed_bytes, + g_compressor_pool_score_buffer, + 0, + pool_row_stride, + sizeof(float), + packed_bytes, + 8, + head_dim, + 1); +} + +static int ds4_gpu_encode_compressor_shift_ratio4( + id cb, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + uint32_t width) { + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!cb || !statekvbuf || !statescbuf || !g_dsv4_ratio4_shift_pipeline || width == 0) return 0; + + ds4_gpu_dsv4_ratio4_shift_args args = { .width = width }; + const uint32_t n = 4u * width; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_ratio4_shift_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:1]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_compressor_prefill_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t n_comp = n_tokens / ratio; + const uint32_t cutoff = n_comp * ratio; + const uint32_t rem = n_tokens - cutoff; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor prefill tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { + fprintf(stderr, "ds4: Metal compressor prefill received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + width, + state_rows, + ds4_gpu_negative_infinity()); + } + + if (ok && ratio == 4u) { + int32_t rows_prev[4] = { 0, 1, 2, 3 }; + const int have_prev = cutoff >= ratio ? 1 : 0; + const uint32_t prev_start = rem == 0 ? cutoff - ratio : cutoff - ratio; + if (have_prev) { + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)prev_start * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)prev_start * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + prev_start, + rows_prev, + 4, + state_rows); + } + if (ok && rem != 0) { + int32_t rows_cur[4]; + for (uint32_t i = 0; i < rem; i++) rows_cur[i] = (int32_t)(ratio + i); + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)cutoff * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)cutoff * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + cutoff, + rows_cur, + rem, + state_rows); + } + } else if (ok && rem != 0) { + int32_t rows[128]; + if (rem > (uint32_t)(sizeof(rows) / sizeof(rows[0]))) { + fprintf(stderr, "ds4: Metal compressor prefill remainder exceeds local row list\n"); + ok = 0; + } else { + for (uint32_t i = 0; i < rem; i++) rows[i] = (int32_t)i; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)cutoff * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)cutoff * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + cutoff, + rows, + rem, + state_rows); + } + } + + if (ok && n_comp != 0) { + const NSUInteger score_bytes = (NSUInteger)cutoff * width * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_bytes, + "ds4_compressor_store_score")) { + ok = 0; + } + if (ok) { + ok = ds4_gpu_encode_compressor_score_with_ape(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_score_buffer, + 0, + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + cutoff); + } + + if (ok && ratio == 4u) { + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_comp * 4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; + const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); + if (ok && !use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))) { + ok = 0; + } + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_kv_buffer, + 0, + head_dim, + 8u * n_comp, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_score_buffer, + 0, + head_dim, + 8u * n_comp, + ds4_gpu_negative_infinity()); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_kv_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_score_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_pool_kv_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + 0, + g_compressor_pool_score_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool) { + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + g_compressor_pool_kv_buffer, + 0, + (uint64_t)head_dim * sizeof(float), + sizeof(float), + 8ull * head_dim * sizeof(float), + g_compressor_pool_score_buffer, + 0, + (uint64_t)head_dim * sizeof(float), + sizeof(float), + 8ull * head_dim * sizeof(float), + 8, + head_dim, + n_comp); + } + } else if (ok) { + const uint64_t row_stride = (uint64_t)width * sizeof(float); + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + row_stride, + sizeof(float), + (uint64_t)ratio * row_stride, + g_compressor_store_score_buffer, + 0, + row_stride, + sizeof(float), + (uint64_t)ratio * row_stride, + ratio, + head_dim, + n_comp); + } + } + + if (ok && n_comp != 0) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, + comp_cache, + model_map, + model_size, + norm_offset, + head_dim, + n_comp, + rms_eps) != 0; + } + if (ok && n_comp != 0 && n_rot != 0) { + ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( + n_comp, 1, head_dim, n_rot, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + cb = ds4_gpu_command_buffer(&owned); + ok = cb && !owned && + ds4_gpu_encode_rope_tail_inplace(cb, + compbuf, + ds4_gpu_tensor_offset(comp_cache), + &rope_args, + n_comp, + 1, + head_dim, + pos0, + ratio); + } + if (ok && n_comp != 0 && quantize_fp8) { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint32_t n_comp = n_tokens / ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor replay tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + ds4_gpu_tensor_bytes(comp_cache) < comp_bytes) { + fprintf(stderr, "ds4: Metal compressor replay received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float); + const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_tokens * width * sizeof(float); + const uint64_t direct_state_bytes = + (uint64_t)4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + (ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + direct_state_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statescbuf, + ds4_gpu_tensor_offset(state_score), + direct_state_bytes)); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; + if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_bytes, + "ds4_compressor_store_score") || + (!use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))))) { + ok = 0; + } + + if (ok) { + ok = ds4_gpu_encode_compressor_score_with_ape(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_score_buffer, + 0, + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + n_tokens); + } + + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv); + const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score); + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_kv_buffer, + 0, + head_dim, + 8u * n_comp, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_score_buffer, + 0, + head_dim, + 8u * n_comp, + ds4_gpu_negative_infinity()); + } + + if (ok && !use_direct_pool && !use_pack_fusion) { + /* + * The aligned nonzero ratio-4 path replays the current ubatch + * compressor, but seeds the first compressed row with the previous + * compressor state. Rows 0..3 are the previous half, rows 4..7 are + * the current half. + */ + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + statekvbuf, + state_off, + g_compressor_pool_kv_buffer, + 0, + head_dim, + ratio, + 1, + src_row_stride, + (uint64_t)ratio * src_row_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + statescbuf, + state_score_off, + g_compressor_pool_score_buffer, + 0, + head_dim, + ratio, + 1, + src_row_stride, + (uint64_t)ratio * src_row_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_kv_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_score_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_pool_kv_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + 0, + g_compressor_pool_score_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool) { + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + g_compressor_pool_kv_buffer, + 0, + dst_row_stride, + sizeof(float), + dst_plane_stride, + g_compressor_pool_score_buffer, + 0, + dst_row_stride, + sizeof(float), + dst_plane_stride, + 8, + head_dim, + n_comp); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, + comp_cache, + model_map, + model_size, + norm_offset, + head_dim, + n_comp, + rms_eps) != 0; + } + if (ok && n_rot != 0) { + ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( + n_comp, 1, head_dim, n_rot, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + cb = ds4_gpu_command_buffer(&owned); + ok = cb && !owned && + ds4_gpu_encode_rope_tail_inplace(cb, + compbuf, + ds4_gpu_tensor_offset(comp_cache), + &rope_args, + n_comp, + 1, + head_dim, + pos0, + ratio); + } + if (ok && quantize_fp8) { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; + } + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + state_off, + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + state_score_off, + width, + state_rows, + ds4_gpu_negative_infinity()); + } + if (ok) { + int32_t rows_prev[4] = { 0, 1, 2, 3 }; + const uint32_t prev_start = n_tokens - ratio; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)prev_start * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)prev_start * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + prev_start, + rows_prev, + ratio, + state_rows); + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_prefill_state_ratio4_tensor( + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv_tail, + const ds4_gpu_tensor *sc_tail, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t pos0) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map || + head_dim == 0 || (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + @autoreleasepool { + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape; + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { + fprintf(stderr, "ds4: Metal compressor prefill-state APE range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv_tail); + id scbuf = ds4_gpu_tensor_buffer(sc_tail); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv_tail) < tail_bytes || + ds4_gpu_tensor_bytes(sc_tail) < tail_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + fprintf(stderr, "ds4: Metal compressor prefill-state received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + width, + state_rows, + ds4_gpu_negative_infinity()); + } + if (ok) { + int32_t rows[4] = { 0, 1, 2, 3 }; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv_tail), + scbuf, + ds4_gpu_tensor_offset(sc_tail), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + rows, + ratio, + state_rows); + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_update_tensor( + const ds4_gpu_tensor *kv_cur, + const ds4_gpu_tensor *sc_cur, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + ds4_gpu_tensor *comp_cache, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos, + uint32_t comp_row, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps, + bool state_already_stored) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || + !model_map || head_dim == 0 || ratio == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv_cur); + id scbuf = ds4_gpu_tensor_buffer(sc_cur); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + if (!kvbuf || !scbuf || !compbuf || + ds4_gpu_tensor_bytes(kv_cur) < kv_bytes || + ds4_gpu_tensor_bytes(sc_cur) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + (emit && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { + fprintf(stderr, "ds4: Metal compressor update received undersized buffers\n"); + return 0; + } + + if (!state_already_stored) { + const bool use_store_one = + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL; + const int store_ok = use_store_one + ? ds4_gpu_compressor_store_one_tensor(kv_cur, + sc_cur, + state_kv, + state_score, + model_map, + model_size, + ape_offset, + ape_type, + width, + ratio, + pos) + : ds4_gpu_compressor_store_batch_tensor(kv_cur, + sc_cur, + state_kv, + state_score, + model_map, + model_size, + ape_offset, + ape_type, + head_dim, + ratio, + pos, + 1); + if (!store_ok) { + return 0; + } + } + if (!emit) return 1; + + ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( + comp_cache, + (uint64_t)comp_row * head_dim * sizeof(float), + (uint64_t)head_dim * sizeof(float)); + if (!comp_row_view) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + int ok = cb && + ds4_gpu_encode_compressor_pool(cb, + comp_row_view, + state_kv, + state_score, + head_dim, + ratio); + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool"); + if (ok) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, + comp_row_view, + model_map, + model_size, + norm_offset, + head_dim, + 1, + rms_eps) != 0; + } + if (ok) { + const uint32_t comp_pos = pos + 1u - ratio; + ok = ds4_gpu_rope_tail_tensor(comp_row_view, + 1, + 1, + head_dim, + n_rot, + comp_pos, + n_ctx_orig, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow) != 0; + } + if (ok && ratio == 4u) { + cb = ds4_gpu_command_buffer(&owned); + ok = cb && + ds4_gpu_encode_compressor_shift_ratio4(cb, + state_kv, + state_score, + width); + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor ratio4 state shift"); + } + ds4_gpu_tensor_free(comp_row_view); + if (!ok) return 0; + } + + return 1; +} + +static int ds4_gpu_encode_fill_f32_rows( + id cb, + id buf, + NSUInteger offset, + uint32_t width, + uint32_t rows, + float value) { + if (!cb || !buf || width == 0 || rows == 0 || (width & 3u) != 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, 1, 0.0f, 0.0f); + args.val = value; + + NSUInteger nth_max = g_unary_fill_pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = (NSUInteger)args.ne00; + if (nth > nth_max) nth = nth_max; + if (nth == 0) nth = 1u; + const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_unary_fill_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:buf offset:offset atIndex:1]; + [enc setBuffer:buf offset:offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, + (NSUInteger)args.ne02, + (NSUInteger)args.ne03) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_attention_output_q8_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + 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, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 32u) != 0 || (low_dim % 32u) != 0 || low_dim > UINT32_MAX) { + fprintf(stderr, "ds4: Metal attention output batch received invalid q8 dimensions\n"); + return 0; + } + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t row_b_bytes = (low_dim / 32u) * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_bytes = out_dim * row_b_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || + out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { + fprintf(stderr, "ds4: Metal attention output batch weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal attention output batch received undersized buffers\n"); + return 0; + } + (void)group_tmp; + (void)low_tmp; + + const bool use_direct_low = + n_tokens < 32u && getenv("DS4_METAL_DISABLE_ATTN_OUT_LOW_DIRECT") == NULL; + /* The exported TensorOps attention-output kernel is a 64-token tile. + * Keep this on full tiles only; smaller multiples of 32 use the legacy + * path instead of relying on cooperative tensor partial RHS bounds. */ + const bool use_mpp_low = + n_tokens >= 32u && + (n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 && + ds4_gpu_use_mpp_attn_out_low_matmul(); + const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); + id group_ids_buffer = nil; + if (!use_direct_low && !use_mpp_low) { + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); + if (!group_ids_buffer) { + return 0; + } + } else { + if (!ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + return 0; + } + group_ids_buffer = g_attn_out_group_ids_buffer; + } + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + const bool attn_out_profile = + getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") != NULL && g_batch_cb != nil; + if (ok && attn_out_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + ok = false; + } else { + cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = false; + } + } + double attn_out_t0 = attn_out_profile ? ds4_gpu_now_ms() : 0.0; +#define DS4_METAL_PROFILE_ATTN_OUT_STAGE(name) do { \ + if (ok && attn_out_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = false; \ + } else { \ + const double now_ms = ds4_gpu_now_ms(); \ + fprintf(stderr, \ + "ds4: Metal attention output stage tokens=%u %s=%.3f ms\n", \ + n_tokens, (name), now_ms - attn_out_t0); \ + attn_out_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = false; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb || owned) ok = false; \ + } \ + } \ + } \ + } while (0) + + if (ok) { + /* + * Batched attention-output projections switch from the vector + * kernel to the SIMD matrix kernel once the batch has at least 32 + * tokens. This preserves the single-token generation path while + * keeping prefill accumulation stable. + */ + if (use_mpp_low) { + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + /* + * Direct RHS lets MPP read the dense low-rank activation tile + * directly from device memory instead of staging a second + * threadgroup tile. The retained attention-output path is the + * 64-token direct-RHS kernel; the older staged-RHS and 32-token + * variants were not kept as alternate runtime modes. + */ + const char *attn_out_pipeline_name = + "kernel_attn_out_low_q8_0_mpp_direct_rhs_n64"; + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false); + ok = ds4_gpu_encode_attn_out_low_q8_mpp(cb, + mm_pipeline, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low)) != 0; + if (!ok) { + ds4_gpu_warn_mpp_fallback(); + if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); + } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + group_ids_buffer = g_attn_out_group_ids_buffer; + } + if (group_ids_buffer) { + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id fallback_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + fallback_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } + } + } + } else if (n_tokens >= 32u && ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + mm_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } else if (use_direct_low) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = (int32_t)n_tokens, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true) != 0; + } else { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = (int32_t)n_tokens, + .nbi1 = (uint64_t)n_groups * sizeof(int32_t), + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); + ok = ds4_gpu_encode_mul_mv_id(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0, + 32u * 2u * sizeof(float), + 4, + true) != 0; + } + } + DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj"); + + if (ok) { + ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, + out_b_offset, + low_dim, out_dim, low, n_tokens) != 0; + } + DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj"); + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } +#undef DS4_METAL_PROFILE_ATTN_OUT_STAGE + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_q4_K_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint32_t out_b_type, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + if (n_tokens < 32u) return 0; + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 256u) != 0 || (low_dim % 256u) != 0 || low_dim > UINT32_MAX) { + return 0; + } + + uint64_t row_a_bytes = 0; + uint64_t row_b_bytes = 0; + if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, (uint32_t)group_dim, &row_a_bytes) || + !ds4_gpu_quant_row_bytes(out_b_type, (uint32_t)low_dim, &row_b_bytes)) { + return 0; + } + + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_bytes = out_dim * row_b_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || + out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { + fprintf(stderr, "ds4: Metal Q4 attention output batch weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q4 attention output batch received undersized buffers\n"); + return 0; + } + (void)group_tmp; + (void)low_tmp; + + const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); + id group_ids_buffer = nil; + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids"); + } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + group_ids_buffer = g_attn_out_group_ids_buffer; + } + if (!group_ids_buffer) return 0; + + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = false; + + if (ok) { + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + mm_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } + + if (ok) { + ok = ds4_gpu_matmul_quant_tensor(out, + model_map, + model_size, + out_b_offset, + out_b_type, + low_dim, + out_dim, + low, + n_tokens) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_q8_batch_f16_tensor( + ds4_gpu_tensor *out_h, + 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, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + (void)out_h; (void)low; (void)model_map; (void)model_size; + (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank; + (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; + return 0; +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((full_in_dim & 31u) != 0 || (k_off & 31u) != 0 || (k_cnt & 31u) != 0 || + k_cnt == 0 || k_off + k_cnt > full_in_dim || + full_in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal Q8_0 kslice matmul received undersized buffers\n"); + return 0; + } + const uint64_t row_bytes = (full_in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 kslice weights are outside the mapped model\n"); + return 0; + } + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, &inner); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + /* Same matvec kernel as the full projection: ne00 bounds the k loop + * while nb01/nb02 keep the full-row stride, so each row reads only + * the owned k window. */ + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(full_in_dim, out_dim); + mv_args.ne00 = (int32_t)k_cnt; + mv_args.ne10 = (int32_t)k_cnt; + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)(inner + (k_off / 32u) * 34u) atIndex:1]; + [enc setBuffer:xbuf + offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 kslice matvec")) { + return 0; + } + return 1; + } +} + +int ds4_gpu_matmul_quant_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + 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) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_matmul_q8_0_kslice_tensor(out, + model_map, + model_size, + weight_offset, + full_in_dim, + k_off, + k_cnt, + out_dim, + x, + x_elem_off); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + full_in_dim == 0 || k_cnt == 0 || out_dim == 0 || + k_off + k_cnt > full_in_dim || + full_in_dim > UINT32_MAX || k_cnt > UINT32_MAX || + out_dim > UINT32_MAX) { + return 0; + } + + uint64_t block_elems = 0; + uint64_t block_bytes = 0; + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + block_elems = 256u; + block_bytes = 144u; + } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { + block_elems = 32u; + block_bytes = 18u; + } else { + fprintf(stderr, "ds4: Metal quant kslice received unsupported type %u\n", weight_type); + return 0; + } + if ((full_in_dim % block_elems) != 0 || + (k_off % block_elems) != 0 || + (k_cnt % block_elems) != 0) { + fprintf(stderr, "ds4: Metal quant kslice dimensions are not block aligned\n"); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal quant kslice matmul received undersized buffers\n"); + return 0; + } + + const uint64_t row_bytes = (full_in_dim / block_elems) * block_bytes; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant kslice weights are outside the mapped model\n"); + return 0; + } + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, &inner); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (weight_type == DS4_METAL_TENSOR_Q4_K && + (k_cnt % 256u) == 0 && + getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { + const int16_t nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); + if (pipeline) { + ds4_gpu_q8_0_matvec_args args = { + .ne00 = (int32_t)k_cnt, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)k_cnt, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = k_cnt * sizeof(float), + .nb12 = k_cnt * sizeof(float), + .nb13 = k_cnt * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; + const uint64_t rows_ptg = (uint64_t)nsg * 2u; + const uint64_t w_skip = (k_off / block_elems) * block_bytes; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)(inner + w_skip) atIndex:1]; + [enc setBuffer:xbuf + offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:32 atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K kslice matvec")) return 0; + return 1; + } + } + + fprintf(stderr, "ds4: Metal quant kslice has no kernel for type %u\n", + weight_type); + return 0; + } +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || group_cnt == 0 || + group0 + group_cnt > n_groups_total || + (group_dim % 32u) != 0 || ((rank * group_cnt) % 32u) != 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim_total = (uint64_t)n_groups_total * rank; + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t a_group_bytes = rank * row_a_bytes; + const uint64_t out_a_bytes = (uint64_t)n_groups_total * a_group_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { + fprintf(stderr, "ds4: Metal TP attention output weights are outside the mapped model\n"); + return 0; + } + /* The heads buffer holds only the owned groups, compact at its + * base (the head slice keeps q/attention output halves packed). */ + if (ds4_gpu_tensor_bytes(heads) < (uint64_t)group_cnt * group_dim * sizeof(float) || + ds4_gpu_tensor_bytes(low) < (uint64_t)group_cnt * rank * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal TP attention output received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, &out_a_inner); + if (!out_a_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + + /* Low projection for the owned groups only: identical dispatch to + * the single-node direct path with the weight base, heads input and + * group count shifted to the slice. The owned low half lands + * compactly at low[0 .. group_cnt*rank). */ + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)group_cnt, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)group_cnt, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = a_group_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)group_cnt, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)group_cnt, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + int ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)(out_a_inner + (uint64_t)group0 * a_group_bytes), + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true); + if (!ok) return 0; + + /* Expand projection over the owned k window only; the result is this + * rank's partial attention block output. */ + return ds4_gpu_matmul_q8_0_kslice_tensor(out, model_map, model_size, + out_b_offset, + low_dim_total, + (uint64_t)group0 * rank, + (uint64_t)group_cnt * rank, + out_dim, low, 0); + } +} + +int ds4_gpu_attention_output_low_q8_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || + n_groups == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 32u) != 0 || low_dim > UINT32_MAX) { + fprintf(stderr, "ds4: Metal attention output low received invalid q8 dimensions\n"); + return 0; + } + + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { + fprintf(stderr, "ds4: Metal attention output low weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = low_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes) { + fprintf(stderr, "ds4: Metal attention output low received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + + if (ok) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_low_q4_K_slice_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || + group_cnt == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + if ((group_dim % 256u) != 0) { + fprintf(stderr, "ds4: Metal attention output low received invalid Q4_K dimensions\n"); + return 0; + } + + uint64_t row_a_bytes = 0; + if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, + (uint32_t)group_dim, + &row_a_bytes)) { + return 0; + } + if (rank > UINT64_MAX / row_a_bytes) return 0; + const uint64_t group_weight_bytes = rank * row_a_bytes; + if (group0 > UINT64_MAX / group_weight_bytes || + group_cnt > UINT64_MAX / group_weight_bytes) { + return 0; + } + const uint64_t group_skip = (uint64_t)group0 * group_weight_bytes; + const uint64_t out_a_bytes = (uint64_t)group_cnt * group_weight_bytes; + if (out_a_offset > UINT64_MAX - group_skip || + out_a_offset + group_skip > model_size || + out_a_bytes > model_size - (out_a_offset + group_skip)) { + fprintf(stderr, "ds4: Metal Q4 attention output low weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)group_cnt * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)group_cnt * rank * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes) { + fprintf(stderr, "ds4: Metal Q4 attention output low received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset + group_skip, + out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + + if (ok) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)group_cnt, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)group_cnt, + .nb00 = 1, + .nb01 = row_a_bytes, + .nb02 = group_weight_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)group_cnt, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)group_cnt, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + const NSUInteger nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q4_K_f32", (int16_t)nsg); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u, + nsg, + false) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +static NSUInteger ds4_gpu_align_up_ns(NSUInteger value, NSUInteger align) { + return (value + align - 1u) & ~(align - 1u); +} + +static int ds4_gpu_encode_cpy_f32_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(float)); + const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f32_pipeline); + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f32_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_col_stride, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = src_col_stride, + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f32_f16_pipeline + : g_cpy_f32_f16_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_2d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint64_t src_row_stride, + uint64_t dst_row_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = (uint64_t)rows * src_row_stride, + .nb03 = (uint64_t)rows * src_row_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = (uint64_t)rows * dst_row_stride, + .nb3 = (uint64_t)rows * dst_row_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f16_f16_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(uint16_t), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f16_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f16_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f16_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f16_f32_pipeline + : g_cpy_f16_f32_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_copy_to_f16_1d( + id cb, + id src, + NSUInteger src_off, + bool src_is_f16, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst) return 0; + if (n == 0) return 1; + if (!src_is_f16) { + return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n); + } + + if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) { + const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u; + const NSUInteger nth = ds4_gpu_cpy_threads( + (uint32_t)work_items, + g_cpy_contig_f16_f16_pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline]; + [enc setBytes:&n length:sizeof(n) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder(); + id blit = [cb blitCommandEncoder]; + if (!blit) return 0; + [blit copyFromBuffer:src + sourceOffset:src_off + toBuffer:dst + destinationOffset:dst_off + size:(NSUInteger)n * sizeof(uint16_t)]; + [blit endEncoding]; + return 1; +} + +static int ds4_gpu_encode_copy_raw_ring_to_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + uint32_t head_dim, + id dst, + NSUInteger dst_offset) { + if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || head_dim == 0) { + return 0; + } + + const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t); + const uint32_t tail_rows = raw_cap - raw_start < n_raw + ? raw_cap - raw_start + : n_raw; + const uint32_t head_rows = n_raw - tail_rows; + const uint64_t tail_count = (uint64_t)tail_rows * head_dim; + const uint64_t head_count = (uint64_t)head_rows * head_dim; + const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes; + const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes; + if (tail_count > UINT32_MAX || head_count > UINT32_MAX || + raw_inner > NSUIntegerMax - raw_offset || + dst_inner > NSUIntegerMax - dst_offset) { + return 0; + } + + if (tail_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset + (NSUInteger)raw_inner, + dst, + dst_offset, + (uint32_t)tail_count)) { + return 0; + } + if (head_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset, + dst, + dst_offset + (NSUInteger)dst_inner, + (uint32_t)head_count)) { + return 0; + } + return 1; +} + +static int ds4_gpu_encode_flash_kv_stage_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + id comp, + NSUInteger comp_offset, + bool comp_is_f16, + uint32_t n_comp, + uint32_t head_dim, + id dst, + NSUInteger dst_offset, + id mask, + NSUInteger mask_offset, + id pad, + NSUInteger pad_offset, + bool fuse_pad, + bool shared_pad, + bool *did_fuse_pad) { + if (did_fuse_pad) *did_fuse_pad = false; + if (!cb || !raw || !comp || !dst || raw_cap == 0 || + raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap || + n_comp == 0 || head_dim == 0) { + return 0; + } + + const bool force = + getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL; + const bool supported_shape = + comp_is_f16 && head_dim == 512u && raw_cap <= UINT32_MAX / 128u; + const uint64_t row_vecs64 = 128u; + const uint64_t comp_count64 = (uint64_t)n_comp * head_dim; + const uint64_t dst_comp_inner64 = + (uint64_t)n_raw * head_dim * sizeof(uint16_t); + if (comp_count64 > UINT32_MAX || + dst_comp_inner64 > NSUIntegerMax - dst_offset) { + return 0; + } + const uint64_t total_vecs64 = + ((uint64_t)n_raw + n_comp) * row_vecs64; + const bool valid_grid = + total_vecs64 != 0 && total_vecs64 <= UINT32_MAX; + const bool eligible = + supported_shape && valid_grid && !g_quality_mode && !disabled && + g_flash_kv_stage_f16_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || force); + const bool component_disabled = eligible && + (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 || + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0); + const bool use_fusion = eligible && !component_disabled; + const bool use_pad_fusion = + use_fusion && fuse_pad && mask != nil && pad != nil && + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL; + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal gathered KV staging kernel was not selected\n"); + return 0; + } + + if (use_fusion) { + ds4_gpu_flash_kv_stage_f16_args args = { + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_raw = n_raw, + .n_comp = n_comp, + .pad_rows = use_pad_fusion ? 32u : 0u, + .shared_pad = use_pad_fusion && shared_pad ? 1u : 0u, + }; + const NSUInteger total_vecs = (NSUInteger)total_vecs64 + + (use_pad_fusion ? 32u * 128u + 32u : 0u); + NSUInteger nth = + g_flash_kv_stage_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > total_vecs) nth = total_vecs; + if (nth == 0) return 0; + const NSUInteger groups = (total_vecs + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_flash_kv_stage_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:raw offset:raw_offset atIndex:1]; + [enc setBuffer:comp offset:comp_offset atIndex:2]; + [enc setBuffer:dst offset:dst_offset atIndex:3]; + [enc setBuffer:(use_pad_fusion ? mask : dst) + offset:(use_pad_fusion ? mask_offset : dst_offset) + atIndex:4]; + [enc setBuffer:(use_pad_fusion ? pad : dst) + offset:(use_pad_fusion ? pad_offset : dst_offset) + atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (did_fuse_pad) *did_fuse_pad = use_pad_fusion; + return 1; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + raw, + raw_offset, + raw_cap, + raw_start, + n_raw, + head_dim, + dst, + dst_offset)) { + return 0; + } + return ds4_gpu_encode_copy_to_f16_1d( + cb, + comp, + comp_offset, + comp_is_f16, + dst, + dst_offset + (NSUInteger)dst_comp_inner64, + (uint32_t)comp_count64); +} + +int ds4_gpu_flash_kv_stage_f16_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + const ds4_gpu_tensor *comp, + uint32_t comp_is_f16, + uint32_t n_comp, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || n_comp == 0 || + comp_is_f16 == 0 || head_dim != 512u) { + return 0; + } + + @autoreleasepool { + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * + (comp_is_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t dst_bytes = + ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t); + id rawbuf = ds4_gpu_tensor_buffer(raw); + id compbuf = ds4_gpu_tensor_buffer(comp); + id dstbuf = ds4_gpu_tensor_buffer(dst); + if (!rawbuf || !compbuf || !dstbuf || + ds4_gpu_tensor_bytes(raw) < raw_bytes || + ds4_gpu_tensor_bytes(comp) < comp_bytes || + ds4_gpu_tensor_bytes(dst) < dst_bytes) { + fprintf(stderr, + "ds4: Metal gathered KV staging received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp), + comp_is_f16 != 0, + n_comp, + head_dim, + dstbuf, + ds4_gpu_tensor_offset(dst), + nil, + 0, + nil, + 0, + false, + false, + NULL)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer( + cb, owned, "gathered KV staging")) { + return 0; + } + } + return 1; +} + +static int ds4_gpu_encode_fill_f16_1d( + id cb, + id buf, + NSUInteger offset, + uint32_t n, + float value) { + if (!cb || !buf || n == 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(n, 1, 0, 0.0f, 0.0f); + args.val = value; + + NSUInteger nth = (NSUInteger)n; + const NSUInteger max_threads = g_unary_fill_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth > 256u) nth = 256u; + if (nth == 0) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_unary_fill_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:buf offset:offset atIndex:1]; + [enc setBuffer:buf offset:offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_raw_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (head_dim != 512 || n_head == 0 || n_raw == 0 || raw_cap < n_raw) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t heads_bytes = q_bytes; + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < heads_bytes) { + fprintf(stderr, "ds4: Metal DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_raw, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_raw * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + memset([mask_buffer contents], 0, mask_bytes); + + id pad_pipeline = nil; + if ((n_raw % ncpsg) != 0) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, (n_raw % ncpsg) != 0, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { + return 0; + } + + if ((n_raw % ncpsg) != 0) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = 1, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_head * row_bytes, + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = 1, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +/* Rectangular causal/window mask for raw prefill: q covers rows + * [q_row0, q_row0 + n_q) of an n_kv-token block whose keys all live in + * rows [0, n_kv). The square prefill case is q_row0 == 0, n_q == n_kv. */ +static void ds4_gpu_fill_raw_prefill_mask( + uint16_t *mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window) { + const uint16_t neg_inf_half = 0xfc00u; + for (uint32_t q = 0; q < n_q; q++) { + const uint32_t qpos = q_row0 + q; + uint16_t *row = mask + (uint64_t)q * n_kv; + for (uint32_t k = 0; k < n_kv; k++) { + const bool causal = k <= qpos; + const bool in_window = window == 0 || qpos - k < window; + row[k] = causal && in_window ? 0u : neg_inf_half; + } + } +} + +static void ds4_gpu_fill_glm_prefill_mask( + uint16_t *mask, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len) { + const uint16_t neg_inf_half = 0xfc00u; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * cache_len; + for (uint32_t k = 0; k < cache_len; k++) { + row[k] = k <= qpos ? 0u : neg_inf_half; + } + } +} + +static id ds4_gpu_glm_prefill_mask_buffer( + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + NSUInteger mask_bytes) { + const int same_shape = + g_glm_flash_attn_mask_valid && + g_glm_flash_attn_mask_buffer && + g_glm_flash_attn_mask_bytes >= mask_bytes && + g_glm_flash_attn_mask_pos0 == pos0 && + g_glm_flash_attn_mask_tokens == n_tokens && + g_glm_flash_attn_mask_cache_len == cache_len; + if (same_shape) return g_glm_flash_attn_mask_buffer; + + if (g_glm_flash_attn_mask_buffer) { + [g_transient_buffers addObject:g_glm_flash_attn_mask_buffer]; + g_glm_flash_attn_mask_buffer = nil; + } + g_glm_flash_attn_mask_bytes = 0; + g_glm_flash_attn_mask_valid = 0; + if (!ds4_gpu_ensure_scratch_buffer(&g_glm_flash_attn_mask_buffer, + &g_glm_flash_attn_mask_bytes, + mask_bytes, + "ds4_glm_flash_attn_mask")) { + return nil; + } + + ds4_gpu_fill_glm_prefill_mask((uint16_t *)[g_glm_flash_attn_mask_buffer contents], + pos0, + n_tokens, + cache_len); + g_glm_flash_attn_mask_pos0 = pos0; + g_glm_flash_attn_mask_tokens = n_tokens; + g_glm_flash_attn_mask_cache_len = cache_len; + g_glm_flash_attn_mask_valid = 1; + return g_glm_flash_attn_mask_buffer; +} + +static void ds4_gpu_fill_raw_decode_batch_mask( + uint16_t *mask, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t pos0, + uint32_t window) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t last_pos = pos0 + n_tokens - 1u; + /* The caller has already copied the SWA ring into logical order when it + * wraps, so key row k represents first_raw_pos + k. */ + const uint32_t first_raw_pos = last_pos + 1u - n_raw; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * n_raw; + for (uint32_t k = 0; k < n_raw; k++) { + const uint32_t kpos = first_raw_pos + k; + const bool causal = kpos <= qpos; + const bool in_window = causal && (window == 0 || qpos - kpos < window); + row[k] = causal && in_window ? 0u : neg_inf_half; + } + } +} + +static void ds4_gpu_fill_mixed_decode_batch_mask( + uint16_t *mask, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t n_comp, + uint32_t pos0, + uint32_t window, + uint32_t ratio) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t n_keys = n_raw + n_comp; + const uint32_t last_pos = pos0 + n_tokens - 1u; + /* Raw keys are laid out by logical position; compressed keys follow them. */ + const uint32_t first_raw_pos = last_pos + 1u - n_raw; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * n_keys; + for (uint32_t k = 0; k < n_raw; k++) { + const uint32_t kpos = first_raw_pos + k; + const bool causal = kpos <= qpos; + const bool in_window = causal && (window == 0 || qpos - kpos < window); + row[k] = causal && in_window ? 0u : neg_inf_half; + } + const uint32_t n_visible = (qpos + 1u) / ratio; + for (uint32_t c = 0; c < n_comp; c++) { + row[n_raw + c] = c < n_visible ? 0u : neg_inf_half; + } + } +} + +/* Rectangular causal/window + compressed-key visibility mask: q covers rows + * [q_row0, q_row0 + n_q) of an n_tokens-token chunk whose raw keys all stay + * resident, followed by n_comp compressed keys. The square prefill case is + * q_row0 == 0, n_q == n_tokens. */ +static void ds4_gpu_fill_static_mixed_prefill_mask( + uint16_t *mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t n_keys = n_tokens + n_comp; + for (uint32_t q = 0; q < n_q; q++) { + const uint32_t qpos = q_row0 + q; + uint16_t *row = mask + (uint64_t)q * n_keys; + for (uint32_t k = 0; k < n_tokens; k++) { + const bool causal = k <= qpos; + const bool in_window = window == 0 || qpos - k < window; + row[k] = causal && in_window ? 0u : neg_inf_half; + } + + const uint32_t n_visible = (qpos + 1u) / ratio; + for (uint32_t c = 0; c < n_comp; c++) { + row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half; + } + } +} + +/* Static-mixed prefill FlashAttention over a rectangular problem: q holds + * n_q query rows for token positions [q_row0, q_row0 + n_q) of the chunk, + * while the keys stay full (all n_tokens raw rows plus n_comp compressed + * rows). The classic square prefill is q_row0 == 0, n_q == n_tokens. */ +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_tokens == 0 || ratio == 0) { + return 0; + } + + const uint32_t n_keys = n_tokens + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask + ? (uint64_t)n_comp * (q_row0 + n_q) * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill static mixed DS4 non-vector FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool bc_mask = (n_q % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_q * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + const uint32_t mask_cache_kind = ratio == 4u + ? DS4_GPU_PREFILL_MASK_CACHE_RATIO4 + : (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u); + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + use_comp_mask == 0u && mask_cache_kind != 0u && + q_row0 == 0u && n_q == n_tokens + ? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind, + n_tokens, + n_comp, + n_keys, + window, + ratio, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created) + : NULL; + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + 0u, + n_tokens, + n_tokens, + n_comp, + window, + ratio); + mask_cache->valid = true; + } + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { + return 0; + } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "static_mixed_nonvec", (name), n_q, n_comp, n_keys, \ + n_head, head_dim, window, ratio, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + if (n_comp && + !ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_tokens * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + if (n_comp) { + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); + } + + if (!mask_cache) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio); + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + if (use_comp_mask && n_comp != 0) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask) + + (NSUInteger)((uint64_t)q_row0 * n_comp * sizeof(float)), + mask_buffer, + (NSUInteger)n_tokens * sizeof(uint16_t), + n_comp, + n_q, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_q, + .ne30 = (int32_t)n_keys, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = nil; + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_q, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_q * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_q, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || ratio == 0) { + return 0; + } + + const uint32_t n_keys = n_tokens + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill static mixed DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nrows = (NSUInteger)n_tokens * n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "static_mixed_vec", (name), n_tokens, n_comp, n_keys, \ + n_head, head_dim, window, ratio, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + if (n_comp) { + if (!ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_tokens * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); + } + + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + if (use_comp_mask && n_comp != 0) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + mask_buffer, + (NSUInteger)n_tokens * sizeof(uint16_t), + n_comp, + n_tokens, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); + } + + id pad_pipeline = nil; + id enc = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, has_kvpad, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + /* The vector sibling below only handles the small square case; every + * rectangular (TP row-split) problem goes through the long path. */ + if (n_tokens >= 20 || q_row0 != 0 || n_q != n_tokens) { + return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); + } + return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); +} + +/* Raw prefill FlashAttention over a rectangular problem: q holds n_q query + * rows that correspond to token positions [q_row0, q_row0 + n_q) of the + * chunk, raw_kv holds all n_kv key rows, and heads receives one output row + * per query row. The classic square prefill is q_row0 == 0, n_q == n_kv. */ +static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill raw DS4 non-vector FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_kv % ncpsg) != 0; + const bool bc_mask = (n_q % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_q * (NSUInteger)n_kv * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_kv * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_kv + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + q_row0 == 0u && n_q == n_kv + ? ds4_gpu_get_zero_prefix_prefill_mask_cache( + DS4_GPU_PREFILL_MASK_CACHE_RAW, + n_kv, + 0u, + n_kv, + window, + 0u, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created) + : NULL; + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], + 0u, n_kv, n_kv, window); + mask_cache->valid = true; + } + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { + return 0; + } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "raw_nonvec", (name), n_q, 0, n_kv, \ + n_head, head_dim, window, 0, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!mask_cache) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], + q_row0, n_q, n_kv, window); + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_kv * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_kv, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_kv * row_bytes_f16, + .nb13 = (uint64_t)n_kv * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_kv * row_bytes_f16, + .nb23 = (uint64_t)n_kv * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_q, + .ne30 = (int32_t)n_kv, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = nil; + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_q, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_q * n_head * row_bytes, + .ne11 = (int32_t)n_kv, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_kv * row_bytes_f16, + .nb13 = (uint64_t)n_kv * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_kv * row_bytes_f16, + .nb23 = (uint64_t)n_kv * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_q, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_raw_heads( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { + return 0; + } + /* The vector sibling below only handles the small square case; every + * rectangular (TP row-split) problem goes through the non-vector path. */ + if (n_kv >= 20 || q_row0 != 0 || n_q != n_kv) { + return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + q_row0, + n_q, + n_kv, + window, + n_head, + head_dim); + } + const uint32_t n_tokens = n_q; + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill raw DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_tokens, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_f16_offset = 0; + const NSUInteger kv_f16_bytes = (NSUInteger)n_tokens * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_tokens * n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "raw_vec", (name), n_tokens, 0, n_tokens, \ + n_head, head_dim, window, 0, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], 0, n_tokens, n_tokens, window); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + + id pad_pipeline = nil; + if ((n_tokens % ncpsg) != 0) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, true, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + kv_f16_offset, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + + if ((n_tokens % ncpsg) != 0) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_tokens, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_tokens * row_bytes_f16, + .nb13 = (uint64_t)n_tokens * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_tokens * row_bytes_f16, + .nb23 = (uint64_t)n_tokens * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_tokens, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_tokens * row_bytes_f16, + .nb13 = (uint64_t)n_tokens * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_tokens * row_bytes_f16, + .nb23 = (uint64_t)n_tokens * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_gathered_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t n_keys = n_raw + n_comp; + if (head_dim != 512 || n_head == 0 || n_raw == 0 || n_keys == 0 || + raw_cap < n_raw || n_keys < n_raw) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : nil; + id headsbuf = ds4_gpu_tensor_buffer(heads); + id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : nil; + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + (n_comp && !compbuf) || + (use_mask && !maskbuf) || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes || + (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes)) { + fprintf(stderr, "ds4: Metal gathered DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + const bool use_persistent_zero_mask = + use_mask == 0u && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) && + getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL; + + if (!(use_persistent_zero_mask + ? ds4_gpu_ensure_zero_attention_mask(mask_bytes) + : ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer, + &g_flash_attn_mask_bytes, + mask_bytes, + "ds4_flash_attn_mask")) || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + id flash_mask_buffer = use_persistent_zero_mask + ? g_flash_attn_zero_mask_buffer + : g_flash_attn_mask_buffer; + + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool use_shared_kvpad = + has_kvpad && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) && + getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL; + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, has_kvpad, + use_shared_kvpad, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!use_persistent_zero_mask && + !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) { + return 0; + } + if (use_mask && n_comp && + !ds4_gpu_encode_cpy_f32_f16_1d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + flash_mask_buffer, + (NSUInteger)n_raw * sizeof(uint16_t), + n_comp)) { + return 0; + } + + bool pad_fused = false; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + n_comp, + head_dim, + g_flash_attn_kv_buffer, + 0, + flash_mask_buffer, + 0, + g_flash_attn_pad_buffer, + 0, + has_kvpad, + use_shared_kvpad, + &pad_fused)) { + return 0; + } + + id pad_pipeline = nil; + if (has_kvpad && !pad_fused) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + + if (has_kvpad && !pad_fused) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = 1, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = 1, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim, + bool noncausal) { + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal decode raw batch FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_raw % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_raw * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_raw + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk")) { + return 0; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { + return 0; + } + + if (noncausal) { + memset([mask_buffer contents], 0, mask_bytes); + } else { + ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_raw, + pos0, + window); + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)n_raw, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (n_comp == 0) { + return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + window, + n_head, + head_dim, + false); + } + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + ratio == 0 || !comp_kv || (use_comp_mask && !comp_mask)) { + return 0; + } + + const uint32_t n_keys = n_raw + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = ds4_gpu_tensor_buffer(comp_kv); + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal decode mixed batch FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk")) { + return 0; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0) || + !ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_raw * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + + ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_raw, + n_comp, + pos0, + window, + ratio); + if (use_comp_mask) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + mask_buffer, + (NSUInteger)n_raw * sizeof(uint16_t), + n_comp, + n_tokens, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)n_keys, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +int ds4_gpu_attention_prefill_raw_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_kv == 0) return 0; + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_raw_heads(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + q_row0, + n_q, + n_kv, + window, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill raw attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_raw_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + return ds4_gpu_attention_prefill_raw_heads_range_tensor(heads, + model_map, + model_size, + sinks_offset, + q, + raw_kv, + 0, + n_tokens, + n_tokens, + window, + n_head, + head_dim); +} + +int ds4_gpu_attention_decode_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + window, + n_head, + head_dim, + false)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode raw batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || n_head == 0 || head_dim == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); + if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal noncausal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + sink_bytes, + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_tokens, + 0, + n_raw, + raw_cap, + raw_start, + 0, + n_head, + head_dim, + true)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph noncausal raw batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_decode_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + ratio == 0 || (n_comp != 0 && !comp_kv) || + (use_comp_mask != 0 && !comp_mask)) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode mixed batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *topk, + uint32_t n_tokens, + uint32_t pos0, + 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 n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !model_map || !q || !raw_kv || !comp_kv || !topk || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + n_comp == 0 || top_k == 0 || top_k > n_comp || (top_k & (top_k - 1u)) != 0 || + ratio == 0 || n_head == 0 || head_dim != 512) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal indexed attention sinks range is outside the mapped model\n"); + return 0; + } + + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t row_bytes_f16 = (uint64_t)head_dim * sizeof(uint16_t); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * row_bytes; + const uint64_t raw_bytes = (uint64_t)raw_cap * row_bytes; + const uint64_t comp_bytes = (uint64_t)n_comp * (comp_kv_f16 ? row_bytes_f16 : row_bytes); + const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = ds4_gpu_tensor_buffer(comp_kv); + id topkbuf = ds4_gpu_tensor_buffer(topk); + id headsbuf = ds4_gpu_tensor_buffer(heads); + if (!qbuf || !rawbuf || !compbuf || !topkbuf || !headsbuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || + ds4_gpu_tensor_bytes(topk) < topk_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal indexed mixed attention received undersized buffers\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + id sort_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, + "kernel_dsv4_sort_i32_rows_asc"); + const bool decode_one_token = n_tokens == 1u; + id attn_pipeline = + decode_one_token ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8_rb16") : + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8"); + if (!sort_pipeline || !attn_pipeline) return 0; + if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) { + fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n"); + return 0; + } + /* + * Fast decode attends to the same full top-k compressed rows but keeps + * them in score order, avoiding a chronological sort dispatch. + * --quality restores the sorted order for stricter reproducibility. + */ + const bool skip_decode_sort = !g_quality_mode && decode_one_token; + if (!skip_decode_sort && + !ds4_gpu_ensure_scratch_buffer(&g_indexed_topk_buffer, + &g_indexed_topk_bytes, + (NSUInteger)topk_bytes, + "ds4_indexed_topk_sorted")) { + return 0; + } + + ds4_gpu_dsv4_topk_mask_args sort_args = { + .ne00 = (int64_t)top_k, + .ne01 = (int64_t)n_tokens, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)top_k * sizeof(int32_t), + .ne0 = (int64_t)top_k, + .ne1 = (int64_t)n_tokens, + .nb0 = sizeof(int32_t), + .nb1 = (uint64_t)top_k * sizeof(int32_t), + }; + ds4_gpu_dsv4_indexed_attention_args attn_args = { + .n_tokens = n_tokens, + .n_head = n_head, + .n_raw = n_raw, + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_comp = n_comp, + .top_k = top_k, + .pos0 = pos0, + .window = window, + .ratio = ratio, + .comp_kv_f16 = comp_kv_f16 ? 1u : 0u, + .pad0 = 0, + .q_token_stride = (uint64_t)n_head * row_bytes, + .q_head_stride = row_bytes, + .raw_row_stride = row_bytes, + .comp_row_stride = comp_kv_f16 ? row_bytes_f16 : row_bytes, + .topk_token_stride = (uint64_t)top_k * sizeof(int32_t), + .dst_token_stride = (uint64_t)n_head * row_bytes, + .dst_head_stride = row_bytes, + .scale = 1.0f / sqrtf((float)head_dim), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = nil; + if (!skip_decode_sort) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:sort_pipeline]; + [enc setBytes:&sort_args length:sizeof(sort_args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:g_indexed_topk_buffer offset:0 atIndex:2]; + [enc setThreadgroupMemoryLength:(NSUInteger)top_k * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(top_k, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; + [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer + offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 + atIndex:4]; + [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; + [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * + 128u * 4u * sizeof(uint16_t) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, ((NSUInteger)n_head + 7u) / 8u, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_tokens == 0 || + ratio == 0 || (n_comp != 0 && !comp_kv)) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + NULL, + 0, + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill static mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_static_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + return ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(heads, + model_map, + model_size, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); +} + +int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !comp_kv || !comp_mask || !model_map || + n_tokens == 0 || n_comp == 0 || ratio == 0) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + 1, + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill masked mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_decode_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !model_map || !q || !raw_kv || + n_raw == 0 || n_head == 0 || head_dim == 0 || + raw_cap < n_raw || raw_start >= raw_cap || + n_raw > UINT32_MAX - n_comp || n_raw + n_comp > 8192u || + (n_comp != 0 && !comp_kv) || + (use_mask != 0 && !comp_mask)) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); + if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal graph attention heads sink range is outside the mapped model\n"); + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal graph attention heads received undersized buffers\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner); + if (!sinks_buf) return 0; + + if (n_comp == 0) { + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_raw_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_raw, + raw_cap, + raw_start, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph raw attention heads")) return 0; + return 1; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_gathered_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_raw, + raw_cap, + raw_start, + comp_kv, + comp_kv_f16, + n_comp, + comp_mask, + use_mask, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph attention heads")) return 0; + } + + return 1; +} diff --git a/models/deepseek/metal/host/hc.inc b/models/deepseek/metal/host/hc.inc new file mode 100644 index 0000000000..e8d013794d --- /dev/null +++ b/models/deepseek/metal/host/hc.inc @@ -0,0 +1,1457 @@ +int ds4_gpu_hc_split_sinkhorn_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *mix, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_hc == 0 || n_hc > 16) return 0; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t mix_tensor_bytes = ds4_gpu_tensor_bytes(mix); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (!mixbuf || !outbuf || + mix_tensor_bytes < mix_bytes || + out_tensor_bytes < mix_bytes) { + fprintf(stderr, "ds4: Metal HC split received undersized activation buffers\n"); + return 0; + } + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset) { + fprintf(stderr, "ds4: Metal HC split parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + if (!scalebuf || !basebuf) return 0; + + uint64_t n_rows64 = mix_tensor_bytes / mix_bytes; + const uint64_t out_rows64 = out_tensor_bytes / mix_bytes; + if (out_rows64 < n_rows64) n_rows64 = out_rows64; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC split row count is outside supported range\n"); + return 0; + } + + ds4_gpu_hc_split_args args = { + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb01 = mix_bytes, + .nb1 = mix_bytes, + .eps = eps, + }; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_rows64)); + const NSUInteger n_tg = ((NSUInteger)n_rows64 + nth - 1u) / nth; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_split_sinkhorn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sinkhorn")) return 0; + } + + return 1; +} + +static int ds4_gpu_hc_weighted_sum_strided( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + uint64_t weight_offset, + uint64_t weight_row_stride, + uint32_t n_embd, + uint32_t n_hc, + const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0 || + weight_row_stride < (uint64_t)n_hc * sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id wbuf = ds4_gpu_tensor_buffer(weights); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC weighted sum output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / out_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC weighted sum token count is outside supported range\n"); + return 0; + } + + const uint64_t x_row_values = (uint64_t)n_hc * n_embd; + if (x_row_values == 0 || + x_row_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (x_row_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / ((uint64_t)n_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC weighted sum activation size overflow\n"); + return 0; + } + + const uint64_t x_bytes = n_tokens64 * x_row_values * sizeof(float); + const uint64_t w_last = weight_offset + + (n_tokens64 - 1u) * weight_row_stride + + (uint64_t)n_hc * sizeof(float); + if (!xbuf || !wbuf || !outbuf || + ds4_gpu_tensor_bytes(residual_hc) < x_bytes || + ds4_gpu_tensor_bytes(weights) < w_last) { + fprintf(stderr, "ds4: Metal HC weighted sum received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_weighted_sum_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_w0 = sizeof(float), + .nb_w1 = weight_row_stride, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + }; + const uint64_t n_elem = (uint64_t)n_embd * n_tokens64; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_weighted_sum_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) + (NSUInteger)weight_offset atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, label)) return 0; + } + + return 1; +} + +int ds4_gpu_hc_weighted_sum_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_hc) { + return ds4_gpu_hc_weighted_sum_strided(out, + residual_hc, + weights, + 0, + (uint64_t)n_hc * sizeof(float), + n_embd, + n_hc, + "HC weighted sum"); +} + +int ds4_gpu_hc_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + const void *model_map, + uint64_t model_size, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + float norm_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool supported_shape = + n_hc == 4u && (n_embd == 4096u || n_embd == 7168u); + const bool auto_shape = + n_embd == 4096u && ds4_gpu_device_name_contains("M3"); + const bool use_fusion = + supported_shape && !g_quality_mode && !disabled && + g_hc_weighted_sum_norm_pipeline != nil && + (auto_shape || force); + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n"); + return 0; + } + if (!use_fusion) return 0; + + @autoreleasepool { + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes; + const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id wbuf = ds4_gpu_tensor_buffer(weights); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + if (!xbuf || !wbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(out) != out_row_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n"); + return 0; + } + if (norm_weight_offset > model_size || + out_row_bytes > model_size - norm_weight_offset) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n"); + return 0; + } + + uint64_t norm_inner = 0; + id normwbuf = ds4_gpu_wrap_model_range( + model_map, model_size, norm_weight_offset, + out_row_bytes, &norm_inner); + if (!normwbuf) return 0; + + ds4_gpu_hc_weighted_sum_norm_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int64_t)n_hc, + .n_tokens = 1, + .nb_x0 = sizeof(float), + .nb_x1 = out_row_bytes, + .nb_x2 = residual_bytes, + .nb_w0 = sizeof(float), + .nb_w1 = weight_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .nb_norm1 = out_row_bytes, + .norm_eps = norm_eps, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); + const NSUInteger shared_bytes = + ((NSUInteger)n_embd + 32u) * sizeof(float); + if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) { + if (require) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n"); + } + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC sum/RMSNorm fused")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_hc_weighted_sum_split_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + return ds4_gpu_hc_weighted_sum_strided(out, + residual_hc, + split, + 0, + mix_hc * sizeof(float), + n_embd, + n_hc, + "HC weighted sum split"); +} + +/* Release decode fused HC pre-sublayer operation. The graph driver owns the + * optional reference fallback so this function stays a direct fused dispatch. */ +int ds4_gpu_hc_split_weighted_sum_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !split || !mix || !residual_hc || !model_map || + n_embd == 0 || n_hc == 0) { + return 0; + } + if (n_hc != 4) { + fprintf(stderr, "ds4: Metal fused HC split/sum is specialized for HC=4\n"); + return 0; + } + + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id splitbuf = ds4_gpu_tensor_buffer(split); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || + out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal fused HC split/sum output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX || + n_rows64 > UINT64_MAX / mix_bytes || + n_rows64 > UINT64_MAX / residual_row_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum row count is outside supported range\n"); + return 0; + } + + const uint64_t mix_total_bytes = n_rows64 * mix_bytes; + const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; + if (!mixbuf || !splitbuf || !xbuf || !outbuf || + ds4_gpu_tensor_bytes(mix) < mix_total_bytes || + ds4_gpu_tensor_bytes(split) < mix_total_bytes || + ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum received undersized activation buffers\n"); + return 0; + } + + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset) { + fprintf(stderr, "ds4: Metal fused HC split/sum parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + if (!scalebuf || !basebuf) return 0; + + ds4_gpu_hc_split_weighted_sum_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb_mix1 = mix_bytes, + .nb_split1 = mix_bytes, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = residual_row_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .eps = eps, + }; + + NSUInteger nth = g_hc_split_weighted_sum_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n_embd) nth = (NSUInteger)n_embd; + if (nth == 0) nth = 1u; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_split_weighted_sum_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; + [enc setThreadgroupMemoryLength:(NSUInteger)n_hc * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum fused")) return 0; + } + + return 1; +} + +/* HC-pre plus the immediately following weighted RMSNorm, specialized for + * DS4's HC=4 shape. Both decode and batched prefill use this implementation; + * the kernel preserves their established single-row and batched scale formulas. */ +int ds4_gpu_hc_split_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float norm_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !norm_out || !split || !mix || !residual_hc || !model_map || + n_hc != 4 || (n_embd & 3u) != 0) { + return 0; + } + + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id splitbuf = ds4_gpu_tensor_buffer(split); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || + out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX || + n_rows64 > UINT64_MAX / mix_bytes || + n_rows64 > UINT64_MAX / residual_row_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm row count is outside supported range\n"); + return 0; + } + + const uint64_t mix_total_bytes = n_rows64 * mix_bytes; + const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; + const uint64_t out_total_bytes = n_rows64 * out_row_bytes; + if (!mixbuf || !splitbuf || !xbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(mix) < mix_total_bytes || + ds4_gpu_tensor_bytes(split) < mix_total_bytes || + ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_total_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm received undersized activation buffers\n"); + return 0; + } + + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset || + norm_weight_offset > model_size || out_row_bytes > model_size - norm_weight_offset) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + uint64_t norm_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, out_row_bytes, &norm_inner); + if (!scalebuf || !basebuf || !normwbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_hc_split_weighted_sum_norm_pipeline, + "kernel_dsv4_hc_split_weighted_sum_norm4"); + if (!pipeline) return 0; + + ds4_gpu_hc_split_weighted_sum_norm_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb_mix1 = mix_bytes, + .nb_split1 = mix_bytes, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = residual_row_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .nb_norm1 = out_row_bytes, + .eps = eps, + .norm_eps = norm_eps, + }; + + NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu threads but pipeline supports %lu\n", + (unsigned long)nth, + (unsigned long)pipeline.maxTotalThreadsPerThreadgroup); + return 0; + } + + const NSUInteger shared_bytes = ((NSUInteger)n_embd + 4u + 32u) * sizeof(float); + const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength]; + if (max_shared != 0 && shared_bytes > max_shared) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu bytes of threadgroup memory but device supports %lu\n", + (unsigned long)shared_bytes, + (unsigned long)max_shared); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:7]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum/norm fused")) return 0; + } + + return 1; +} + +int ds4_gpu_output_hc_weights_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *pre, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !pre || !model_map || n_hc == 0) return 0; + + @autoreleasepool { + if ((n_hc % 4u) != 0) { + fprintf(stderr, "ds4: Metal output HC weights requires a multiple-of-4 HC width\n"); + return 0; + } + + id prebuf = ds4_gpu_tensor_buffer(pre); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (row_bytes == 0 || out_tensor_bytes < row_bytes || out_tensor_bytes % row_bytes != 0) { + fprintf(stderr, "ds4: Metal output HC weights size is not a whole token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX || + n_tokens64 > UINT64_MAX / row_bytes) { + fprintf(stderr, "ds4: Metal output HC weights token count is outside supported range\n"); + return 0; + } + + const uint64_t bytes = n_tokens64 * row_bytes; + if (!prebuf || !outbuf || + ds4_gpu_tensor_bytes(pre) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal output HC weights received undersized buffers\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, + scale_offset, sizeof(float), + &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, + base_offset, row_bytes, + &base_inner); + if (!scalebuf || !basebuf) return 0; + + const bool force_weights4 = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool disable_weights4 = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL; + const bool require_weights4 = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u; + const bool use_weights4 = + weights4_shape && !g_quality_mode && !disable_weights4 && + g_output_hc_weights4_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || force_weights4) && + g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u; + if (require_weights4 && weights4_shape && !use_weights4) { + fprintf(stderr, + "ds4: required Metal output HC weights4 kernel was not selected\n"); + return 0; + } + + if (use_weights4) { + ds4_gpu_output_hc_weights4_args args = { + .post_scale = 1.0f, + .eps = eps, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_output_hc_weights4_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(2, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC weights4")) { + return 0; + } + return 1; + } + + const uint32_t n_tokens = (uint32_t)n_tokens64; + ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1); + ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc); + ds4_gpu_unary_args sigmoid_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 0.0f, 0.0f); + ds4_gpu_unary_args scale_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 1.0f, eps); + + NSUInteger mul_nth_max = g_bin_mul_scalar_pipeline.maxTotalThreadsPerThreadgroup; + if (mul_nth_max > 256u) mul_nth_max = 256u; + NSUInteger mul_nth = 1u; + while (2u * mul_nth < (NSUInteger)mul_args.ne0 && mul_nth < mul_nth_max) { + mul_nth *= 2u; + } + + NSUInteger add_nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup; + if (add_nth_max > 256u) add_nth_max = 256u; + NSUInteger add_nth = 1u; + while (2u * add_nth < (NSUInteger)add_args.ne0 && add_nth < add_nth_max) { + add_nth *= 2u; + } + + NSUInteger unary_nth_max = g_unary_sigmoid_pipeline.maxTotalThreadsPerThreadgroup; + if (unary_nth_max > 256u) unary_nth_max = 256u; + NSUInteger unary_nth = (NSUInteger)sigmoid_args.ne00; + if (unary_nth > unary_nth_max) unary_nth = unary_nth_max; + if (unary_nth == 0) unary_nth = 1u; + const NSUInteger unary_nk0 = ((NSUInteger)sigmoid_args.ne00 + unary_nth - 1u) / unary_nth; + const NSUInteger out_offset = ds4_gpu_tensor_offset(out); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + + [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; + [enc setBytes:&mul_args length:sizeof(mul_args) atIndex:0]; + [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:outbuf offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)mul_args.ne01, + (NSUInteger)mul_args.ne02, + (NSUInteger)mul_args.ne03) + threadsPerThreadgroup:MTLSizeMake(mul_nth, 1, 1)]; + + [enc setComputePipelineState:g_add_pipeline]; + [enc setBytes:&add_args length:sizeof(add_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:2]; + [enc setBuffer:outbuf offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)add_args.ne01, + (NSUInteger)add_args.ne02, + (NSUInteger)add_args.ne03) + threadsPerThreadgroup:MTLSizeMake(add_nth, 1, 1)]; + + [enc setComputePipelineState:g_unary_sigmoid_pipeline]; + [enc setBytes:&sigmoid_args length:sizeof(sigmoid_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:outbuf offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)sigmoid_args.ne01, + (NSUInteger)sigmoid_args.ne02, + (NSUInteger)sigmoid_args.ne03) + threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; + + [enc setComputePipelineState:g_unary_scale_pipeline]; + [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:outbuf offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)scale_args.ne01, + (NSUInteger)scale_args.ne02, + (NSUInteger)scale_args.ne03) + threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "output HC weights")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id postbuf = ds4_gpu_tensor_buffer(post); + id combbuf = ds4_gpu_tensor_buffer(comb); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t comb_values = (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + comb_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); + const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); + if (!blockbuf || !resbuf || !postbuf || !combbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(post) < post_bytes || + ds4_gpu_tensor_bytes(comb) < comb_bytes) { + fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = (uint64_t)n_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; + [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand")) return 0; + } + + return 1; +} + +/* Expand of the SUM of two block vectors — the TP attention combine folded + * into the HC expand (the kernel's has_add path adds them element-wise + * before the post/comb mixing, canonical rank order preserved by argument + * position). */ +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_embd == 0 || n_hc != 4) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id addbuf = ds4_gpu_tensor_buffer(block_add); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id postbuf = ds4_gpu_tensor_buffer(post); + id combbuf = ds4_gpu_tensor_buffer(comb); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t comb_values = (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + comb_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); + const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); + if (!blockbuf || !addbuf || !resbuf || !postbuf || !combbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(block_add) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(post) < post_bytes || + ds4_gpu_tensor_bytes(comb) < comb_bytes) { + fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = (uint64_t)n_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; + [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; + [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand split output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand split token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + mix_hc > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand split activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); + if (!blockbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes) { + fprintf(stderr, "ds4: Metal HC expand split received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand split")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_split_half_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split; + (void)n_embd; (void)n_hc; + return 0; +} + +int ds4_gpu_hc_expand_add_split_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 *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id addbuf = ds4_gpu_tensor_buffer(block_add); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand add split output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand add split token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + mix_hc > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand add split activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); + if (!blockbuf || !addbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(block_add) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes) { + fprintf(stderr, "ds4: Metal HC expand add split received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; + [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add split")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_add_split_half_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + (void)out_hc; (void)block_out; (void)block_add_h; (void)residual_hc; + (void)split; (void)n_embd; (void)n_hc; + return 0; +} + +int ds4_gpu_shared_down_hc_expand_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_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 *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !shared_out || !model_map || !shared_mid || !routed_out || + !residual_hc || !split || n_embd == 0 || n_hc == 0 || + n_hc != 4 || out_dim != n_embd || (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id midbuf = ds4_gpu_tensor_buffer(shared_mid); + id sharedbuf = ds4_gpu_tensor_buffer(shared_out); + id routedbuf = ds4_gpu_tensor_buffer(routed_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + const uint64_t shared_mid_bytes = in_dim * sizeof(float); + const uint64_t embd_bytes = out_dim * sizeof(float); + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t split_bytes = mix_hc * sizeof(float); + + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal shared-down HC fusion weight range is outside the mapped model\n"); + return 0; + } + if (!midbuf || !sharedbuf || !routedbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(shared_mid) < shared_mid_bytes || + ds4_gpu_tensor_bytes(shared_out) < embd_bytes || + ds4_gpu_tensor_bytes(routed_out) < embd_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes || + ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { + fprintf(stderr, "ds4: Metal shared-down HC fusion received undersized buffers\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + mv_args.nr0 = mv_dispatch.nr0; + + ds4_gpu_hc_expand_args hc_args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = 1, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_hc_expand4_q8_0", + mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(shared_mid) atIndex:3]; + [enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:4]; + [enc setBuffer:routedbuf offset:ds4_gpu_tensor_offset(routed_out) atIndex:5]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:7]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:8]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:9]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "shared-down HC expand fused")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_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, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !model_map || !x || !residual_hc || !split || + n_embd == 0 || n_hc == 0 || n_hc != 4 || out_dim != n_embd || + (in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t embd_bytes = out_dim * sizeof(float); + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t split_bytes = mix_hc * sizeof(float); + + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8 HC fusion weight range is outside the mapped model\n"); + return 0; + } + if (!xbuf || !blockbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(block_out) < embd_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes || + ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { + fprintf(stderr, "ds4: Metal Q8 HC fusion received undersized buffers\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + mv_args.nr0 = mv_dispatch.nr0; + + ds4_gpu_hc_expand_args hc_args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = 1, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_q8_hc_expand4_q8_0", + mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:6]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8 HC expand fused")) return 0; + } + + return 1; +} diff --git a/models/deepseek/metal/host/indexer.inc b/models/deepseek/metal/host/indexer.inc new file mode 100644 index 0000000000..a4cf9e4f56 --- /dev/null +++ b/models/deepseek/metal/host/indexer.inc @@ -0,0 +1,756 @@ +int ds4_gpu_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_head, + uint32_t head_dim, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !index_comp || + n_comp == 0 || n_head == 0 || head_dim == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)n_head * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t score_bytes = (uint64_t)n_comp * sizeof(float); + id qbuf = ds4_gpu_tensor_buffer(q); + id wbuf = ds4_gpu_tensor_buffer(weights); + id compbuf = ds4_gpu_tensor_buffer(index_comp); + id scorebuf = ds4_gpu_tensor_buffer(scores); + if (!qbuf || !wbuf || !compbuf || !scorebuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(index_comp) < comp_bytes || + ds4_gpu_tensor_bytes(scores) < score_bytes) { + fprintf(stderr, "ds4: Metal graph indexer score received undersized buffers\n"); + return 0; + } + + if (n_head == 64 && head_dim == 128) { + id direct_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline, + "kernel_dsv4_indexer_score_one_direct"); + if (!direct_pipeline) return 0; + + ds4_gpu_dsv4_indexer_scores_fused_args args = { + .n_comp = n_comp, + .n_tokens = 1, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = 0, + .ratio = 4, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .index_row_stride = (uint64_t)head_dim * sizeof(float), + .score_token_stride = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:direct_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer direct score")) return 0; + return 1; + } + + const uint64_t head_score_bytes = (uint64_t)n_comp * n_head * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_head_scores_buffer, + &g_indexer_head_scores_bytes, + (NSUInteger)head_score_bytes, + "ds4_indexer_head_scores")) { + return 0; + } + + ds4_gpu_q8_0_matvec_args dot_args = + ds4_gpu_make_f32_mv_args(head_dim, n_comp, n_head); + ds4_gpu_mv_dispatch dot_dispatch = + ds4_gpu_make_plain_mv_dispatch(head_dim, 1); + dot_args.nr0 = dot_dispatch.nr0; + id dot_pipeline = + ds4_gpu_get_mul_mv_pipeline(dot_dispatch.function_name, dot_dispatch.nsg); + if (!dot_pipeline) return 0; + ds4_gpu_dsv4_indexer_weighted_sum_args sum_args = { + .ne00 = (int64_t)n_comp, + .ne01 = 1, + .ne02 = (int64_t)n_head, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * sizeof(float), + .ne10 = (int64_t)n_head, + .ne11 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)n_head * sizeof(float), + .ne0 = (int64_t)n_comp, + .ne1 = 1, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:dot_pipeline]; + [enc setBytes:&dot_args length:sizeof(dot_args) atIndex:0]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:3]; + if (dot_dispatch.smem) { + [enc setThreadgroupMemoryLength:dot_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + (NSUInteger)dot_dispatch.nr0 - 1u) / (NSUInteger)dot_dispatch.nr0, + n_head, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dot_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_indexer_weighted_sum_pipeline]; + [enc setBytes:&sum_args length:sizeof(sum_args) atIndex:0]; + [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer score")) return 0; + } + + return 1; +} + +static int ds4_gpu_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !index_comp || + n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || ratio == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)n_tokens * n_head * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + id qbuf = ds4_gpu_tensor_buffer(q); + id wbuf = ds4_gpu_tensor_buffer(weights); + id compbuf = ds4_gpu_tensor_buffer(index_comp); + id scorebuf = ds4_gpu_tensor_buffer(scores); + if (!qbuf || !wbuf || !compbuf || !scorebuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(index_comp) < comp_bytes || + ds4_gpu_tensor_bytes(scores) < score_bytes) { + fprintf(stderr, "ds4: Metal graph indexer prefill scores received undersized buffers\n"); + return 0; + } + if (head_dim != 128) { + fprintf(stderr, "ds4: Metal fused DS4 indexer scores expect 128-wide rows\n"); + return 0; + } + /* + * The NAX/TensorOps score builder is a prefill-only win. At small + * batches and in one-token decode the setup cost is not amortized, so + * those paths keep the older direct/tiled score kernels. + */ + const bool use_nax = ds4_gpu_mpp_available() && n_tokens >= 16u; + id pipeline = ds4_gpu_get_pipeline( + use_nax ? "kernel_dsv4_indexer_scores_nax" : + (g_quality_mode ? "kernel_dsv4_indexer_scores_tiled_f32" + : "kernel_dsv4_indexer_scores_tiled")); + if (!pipeline) return 0; + + ds4_gpu_dsv4_indexer_scores_fused_args args = { + .n_comp = n_comp, + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = pos0, + .ratio = ratio, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .index_row_stride = (uint64_t)head_dim * sizeof(float), + .score_token_stride = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + if (use_nax) { + const NSUInteger q_shared = 2u * 32u * 32u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 32u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 15u) / 16u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } else if (g_quality_mode) { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer prefill scores")) return 0; + } + + return 1; +} + +int ds4_gpu_indexer_scores_prefill_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return ds4_gpu_indexer_scores_batch_tensor(scores, + q, + weights, + index_comp, + n_comp, + n_tokens, + 0, + n_head, + head_dim, + ratio, + scale); +} + +int ds4_gpu_indexer_scores_decode_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return ds4_gpu_indexer_scores_batch_tensor(scores, + q, + weights, + index_comp, + n_comp, + n_tokens, + pos0, + n_head, + head_dim, + ratio, + scale); +} + +int ds4_gpu_indexer_topk_tensor( + ds4_gpu_tensor *selected, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || top_k > n_comp) return 0; + + @autoreleasepool { + const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + const uint64_t selected_bytes = (uint64_t)top_k * n_tokens * sizeof(uint32_t); + id scorebuf = ds4_gpu_tensor_buffer(scores); + id selbuf = ds4_gpu_tensor_buffer(selected); + if (!scorebuf || !selbuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n"); + return 0; + } + NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; + if (max_threads == 0) max_threads = 256; + int32_t nth = 1; + while ((uint32_t)nth < n_comp && (uint64_t)2u * (uint64_t)nth <= (uint64_t)max_threads) { + nth *= 2; + } + const int32_t npr = (int32_t)((n_comp + (uint32_t)nth - 1u) / (uint32_t)nth); + const int32_t block_top_k = (int32_t)(top_k < (uint32_t)nth ? top_k : (uint32_t)nth); + int32_t work_width = (int32_t)top_k; + if (npr > 1) { + const int32_t last_block = (int32_t)n_comp - (npr - 1) * nth; + work_width = (npr - 1) * block_top_k + (last_block < block_top_k ? last_block : block_top_k); + } + const uint64_t scratch_row_bytes = (uint64_t)work_width * sizeof(uint32_t); + const bool one_pass = npr <= 1; + const uint64_t scratch_bytes = one_pass ? scratch_row_bytes * n_tokens : + 2u * scratch_row_bytes * n_tokens; + if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_topk_buffer, + &g_indexer_topk_bytes, + (NSUInteger)scratch_bytes, + "ds4_indexer_topk")) { + return 0; + } + + ds4_gpu_kargs_argsort args = { + .ne00 = (int32_t)n_comp, + .ne01 = (int32_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), + .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), + .ne0 = work_width, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .top_k = block_top_k, + }; + // kernel_argsort_f32_i32_desc stages the block's scores behind the + // index array: nth int32 indices + nth float scores. + const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u; + + NSUInteger cur_off = 0; + NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_argsort_f32_i32_desc_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; + [enc setBuffer:one_pass ? selbuf : g_indexer_topk_buffer + offset:one_pass ? ds4_gpu_tensor_offset(selected) : cur_off + atIndex:2]; + [enc setThreadgroupMemoryLength:smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)npr * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + int32_t len = block_top_k; + while (len < work_width) { + const int32_t nm = (work_width + 2 * len - 1) / (2 * len); + const bool final_merge = nm == 1; + NSUInteger merge_threads = g_argsort_merge_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; + if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u; + if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len; + if (merge_threads == 0) merge_threads = 1; + + ds4_gpu_kargs_argsort_merge merge_args = { + .ne00 = (int64_t)n_comp, + .ne01 = (int64_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), + .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), + .ne0 = work_width, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .top_k = nm == 1 ? (int32_t)top_k : work_width, + .len = len, + }; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_argsort_merge_f32_i32_desc_pipeline]; + [enc setBytes:&merge_args length:sizeof(merge_args) atIndex:0]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; + [enc setBuffer:g_indexer_topk_buffer offset:cur_off atIndex:2]; + [enc setBuffer:final_merge ? selbuf : g_indexer_topk_buffer + offset:final_merge ? ds4_gpu_tensor_offset(selected) : next_off + atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nm * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(merge_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + const NSUInteger tmp = cur_off; + cur_off = next_off; + next_off = tmp; + len <<= 1; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0; + } + + return 1; +} + +int ds4_gpu_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits, + uint32_t n_vocab) { + if (!out_idx || !logits || n_vocab == 0) return 0; + if (ds4_gpu_tensor_bytes(out_idx) < sizeof(int32_t) || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_vocab * sizeof(float)) { + fprintf(stderr, "ds4: Metal graph argmax received undersized buffers\n"); + return 0; + } + + return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1, 1); +} + +int ds4_gpu_dsv4_topk_mask_tensor( + ds4_gpu_tensor *mask, + const ds4_gpu_tensor *topk, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0) return 0; + + @autoreleasepool { + const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); + const uint64_t mask_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + id topkbuf = ds4_gpu_tensor_buffer(topk); + id maskbuf = ds4_gpu_tensor_buffer(mask); + if (!topkbuf || !maskbuf || + ds4_gpu_tensor_bytes(topk) < topk_bytes || + ds4_gpu_tensor_bytes(mask) < mask_bytes) { + fprintf(stderr, "ds4: Metal dsv4 top-k mask received undersized buffers\n"); + return 0; + } + + ds4_gpu_dsv4_topk_mask_args args = { + .ne00 = (int64_t)top_k, + .ne01 = (int64_t)n_tokens, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)top_k * sizeof(int32_t), + .ne0 = (int64_t)n_comp, + .ne1 = (int64_t)n_tokens, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_comp * sizeof(float), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_topk_mask_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)n_comp * n_tokens) + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_topk_mask_scatter_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)top_k * n_tokens) + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "dsv4 top-k mask")) return 0; + } + + return 1; +} + +static int ds4_gpu_matmul_q8_0_legacy_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_tok, + bool prefer_decode_mpp, + bool force_model_view) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = force_model_view ? + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset) : + ds4_gpu_wrap_q8_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + if (ds4_gpu_mpp_available() && + (prefer_decode_mpp || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) && + getenv("DS4_METAL_DISABLE_Q8_DECODE_MPP") == NULL && + (in_dim % 64u) == 0) { + const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs"; + id mpp_pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (mpp_pipeline) { + ds4_gpu_mul_mm_args args = + ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mpp_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1u, + ((NSUInteger)out_dim + 63u) / 64u, + 1u) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode MPP matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + const bool force_output_nr4 = + getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL; + const bool output_shape = + in_dim == 4096u && out_dim == 129280u; + const bool use_output_nr4 = + !g_quality_mode && (out_dim & 3u) == 0u && + (force_output_nr4 || + (output_shape && ds4_gpu_device_name_contains("M3"))) && + getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL; + if (use_output_nr4) { + mv_dispatch.function_name = + "kernel_mul_mv_q8_0_f32_nr4"; + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matvec")) { + return 0; + } + return 1; + } + + const uint64_t mv_ext_max_tokens = + ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); + if (n_tok <= mv_ext_max_tokens && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_name(1, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor mul_mv_ext")) { + return 0; + } + return 1; + } + + /* + * Dense Q8_0 prefill is the cleanest DS4 TensorOps shape: M/N/K are + * aligned and the RHS activation matrix is already dense. The retained + * kernel dequantizes each 64x32 weight tile to half in threadgroup + * memory, then uses direct-RHS MPP for the activation tile. This avoids + * staging RHS into threadgroup memory and was the direct replacement for + * the slower generic MPP prototype. + */ + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 64u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = nax_tile_n == 128u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128" + : (nax_tile_n == 64u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64" + : "kernel_mul_mm_q8_0_f32_nax_direct_rhs"); + id pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_q8_0_f32", bc_inp, bc_out); + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matmul")) { + return 0; + } + } + + return 1; +} diff --git a/models/deepseek/metal/host/moe.inc b/models/deepseek/metal/host/moe.inc new file mode 100644 index 0000000000..e5c794442d --- /dev/null +++ b/models/deepseek/metal/host/moe.inc @@ -0,0 +1,3756 @@ +int ds4_gpu_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t token, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_expert == 0 || n_expert_used == 0) return 0; + if (hash_mode && token >= hash_rows) return 0; + if (n_expert_groups > 1u || n_group_used > 0u) { + fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal router select received undersized buffers\n"); + return 0; + } + + uint64_t bias_inner = 0; + uint64_t hash_inner = 0; + id biasbuf = nil; + id hashbuf = nil; + NSUInteger bias_set_offset = 0; + NSUInteger hash_set_offset = 0; + if (has_bias && !hash_mode) { + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); + if (!biasbuf) return 0; + bias_set_offset = (NSUInteger)bias_inner; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); + hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); + if (!hashbuf) return 0; + hash_set_offset = (NSUInteger)hash_inner; + } + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + const int32_t token_i32 = (int32_t)token; + int ok = cb && + ds4_gpu_encode_router_select(cb, + selected, + weights, + probs, + logitsbuf, + ds4_gpu_tensor_offset(logits), + biasbuf, + bias_set_offset, + hashbuf, + hash_set_offset, + nil, + 0, + &token_i32, + hash_rows, + 1, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode); + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + if (!ok) return 0; + } + + return 1; +} + +int ds4_gpu_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits, + const ds4_gpu_tensor *tokens, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !tokens || !model_map || + n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; + if (n_expert_groups > 1u || n_group_used > 0u) { + fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + id tokensbuf = ds4_gpu_tensor_buffer(tokens); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal router batch select received undersized buffers\n"); + return 0; + } + + uint64_t bias_inner = 0; + uint64_t hash_inner = 0; + id biasbuf = nil; + id hashbuf = nil; + NSUInteger bias_set_offset = 0; + NSUInteger hash_set_offset = 0; + if (has_bias && !hash_mode) { + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); + if (!biasbuf) return 0; + bias_set_offset = (NSUInteger)bias_inner; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); + hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); + if (!hashbuf) return 0; + hash_set_offset = (NSUInteger)hash_inner; + } + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + int ok = cb && + ds4_gpu_encode_router_select(cb, + selected, + weights, + probs, + logitsbuf, + ds4_gpu_tensor_offset(logits), + biasbuf, + bias_set_offset, + hashbuf, + hash_set_offset, + tokensbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + hash_rows, + n_tokens, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode); + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + if (!ok) return 0; + } + + return 1; +} + +int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) { + if (n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED || + (!selected && n_selected != 0)) return 0; + for (uint32_t i = 0; i < n_selected; i++) { + g_routed_moe_selected_override[i] = selected[i]; + } + g_routed_moe_selected_override_n = n_selected; + return 1; +} + +int ds4_gpu_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *add_in, + uint32_t layer_index, + bool force_resident) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding: only the owned contiguous expert range is mapped, + * so bind from the owned base, validate only its bytes, and tell the + * kernels the first expert id present at that base. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * gate_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + + if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || + n_total_expert == 0 || n_expert == 0 || + n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_row_bytes == 0 || down_row_bytes == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32044); + return 0; + } + if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32047); return 0; } + ds4_gpu_stream_expert_cache_note_token(layer_index); + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id expertsbuf = ds4_gpu_tensor_buffer(experts); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id selected_exec_buf = selectedbuf; + NSUInteger selected_exec_off = ds4_gpu_tensor_offset(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < mid_bytes || + ds4_gpu_tensor_bytes(up) < mid_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal routed tensor MoE received undersized activation buffers\n"); + return 0; + } + if (n_expert > 1 && + (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < (uint64_t)n_expert * out_dim * sizeof(float))) { + fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n"); + return 0; + } + + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE tensor byte size overflow\n"); + return 0; + } + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + __unsafe_unretained id gate_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + __unsafe_unretained id up_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + __unsafe_unretained id down_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + ds4_gpu_stream_expert_cache_entry *stream_slot_entries[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { NULL }; + NSUInteger gate_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + NSUInteger up_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + NSUInteger down_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_gate_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_up_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_down_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + bool use_stream_expert_addr_table = false; + bool use_stream_expert_masked_addr_table = false; + bool use_stream_compact_addr_table = false; + bool use_stream_expert_cache = false; + bool use_stream_expert_split_candidate = false; + bool use_stream_expert_split_deferred = false; + bool stream_expert_split_completed = false; + uint32_t stream_expert_resident_mask = 0; + uint32_t stream_expert_missing_mask = 0; + __unsafe_unretained id gate_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id up_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id down_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + NSUInteger gate_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + NSUInteger up_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + NSUInteger down_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + __unsafe_unretained id gate_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id up_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id down_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + NSUInteger gate_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + NSUInteger up_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + NSUInteger down_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + __unsafe_unretained id gate_group24_bufs[24] = { nil }; + __unsafe_unretained id up_group24_bufs[24] = { nil }; + __unsafe_unretained id down_group24_bufs[24] = { nil }; + NSUInteger gate_group24_offsets[24] = { 0 }; + NSUInteger up_group24_offsets[24] = { 0 }; + NSUInteger down_group24_offsets[24] = { 0 }; + DS4MetalQ4ExpertTable *gate_table = nil; + DS4MetalQ4ExpertTable *up_table = nil; + DS4MetalQ4ExpertTable *down_table = nil; + id q4_table_layer_residency = nil; + int32_t selected_ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + + const uint32_t n_tokens = 1; + const uint32_t pair_rows = n_tokens * n_expert; + const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); + if ((n_expert > 1 && !expertsbuf && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch"))) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32148); + return 0; + } + + const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); + const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); + id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); + id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); + if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { + fprintf(stderr, "ds4: unsupported Metal routed MoE quant types gate=%u down=%u\n", + gate_type, down_type); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32158); + return 0; + } + + ds4_gpu_mul_mv_id_args gate_args = + ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens, gate_nr0); + ds4_gpu_mul_mv_id_args down_args = + ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, down_nr0); + /* Tensor-parallel expert ownership; non-TP calls keep tp_world at 1. */ + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = tp_expert_base_host; + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_addend = add_in != NULL; + down_args.tp_expert_base = tp_expert_base_host; + + const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); + const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); + const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); + const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); + const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); + const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); + int ok = 1; + const bool write_clamped_moe = + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL; + id pair_swiglu_pipeline = nil; + if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { + pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { + pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } + const bool fuse_pair_swiglu = + !g_quality_mode && + !write_clamped_moe && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + pair_swiglu_pipeline != nil; + id down_sum6_pipeline = nil; + if (down_type == DS4_METAL_TENSOR_Q2_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_Q4_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_IQ2_XXS && + g_tp_split_world == 2) { + /* IQ2 down-sum exists for the GLM TP resident split only; the + * streaming paths must keep their addr/masked chain (expert + * bytes are not at their model-map offsets when streamed). */ + down_sum6_pipeline = g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; + } + const bool direct_down_sum = + !g_quality_mode && + (n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2)) && + n_tokens == 1 && + down_sum6_pipeline != nil; + /* The expert-ownership split lives only in the fused id pair+sum6 + * kernels; every other routed variant would silently compute full + * sums on both ranks and double the combine. Fail fast instead. */ + if ((g_tp_split_world > 1 || add_in) && !(fuse_pair_swiglu && direct_down_sum)) { + fprintf(stderr, "ds4: tensor-parallel routed MoE requires the fused pair+sum6 decode path\n"); + return 0; + } + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + /* + * The grouped Q4 experiment keeps selected IDs on GPU, but it also walks + * every expert window in the layer. On PRO Q4 this measured far slower + * than the active selected-slot path, so keep it opt-in for profiling. + */ + const bool enable_q4_grouped_experts = + getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL; + const bool use_q4_grouped_experts = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert >= 128 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group_q4_k_sum6_pipeline != nil && + enable_q4_grouped_experts; + const uint32_t q4_expert_group_size = + use_q4_grouped_experts ? ds4_gpu_q4_expert_group_size(n_total_expert) : 0; + const bool q4_grouped_boundary = + use_q4_grouped_experts && + g_batch_cb != nil && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_BOUNDARY") == NULL; + const bool q4_grouped_cache_views = + getenv("DS4_METAL_Q4_GROUPED_CACHE_VIEWS") != NULL; + const uint32_t q4_group6_expert_group_size = 64; + const bool use_q4_group6_experts = + !use_q4_grouped_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group6_expert_group_size * 6u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL; + const uint32_t q4_group8_expert_group_size = 48; + const bool use_q4_group8_experts = + !use_q4_grouped_experts && + !use_q4_group6_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group8_expert_group_size * 8u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL; + const uint32_t q4_group24_expert_group_size = 16; + const bool use_q4_group24_experts = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group24_expert_group_size * 24u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + direct_down_sum && + g_moe_mul_mv_group24_q4_k_id_pipeline != nil && + g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL; + const bool q4_group24_exact_views = + use_q4_group24_experts && + getenv("DS4_METAL_Q4_GROUP24_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_Q4_GROUP24_BASE_VIEWS") == NULL; + const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; + const bool can_wrap_q4_exact_tensors = + max_buffer_len != 0 && + gate_tensor_bytes <= max_buffer_len && + down_tensor_bytes <= max_buffer_len; + /* + * The full-tensor Q4 ID path is the closest arithmetic analogue to IQ2, + * but PRO Q4 routed tensors are multi-GiB. Keep it opt-in: even with + * per-layer command boundaries it is much slower than binding only the + * six active experts on current M3 Ultra Metal. + */ + const bool enable_q4_exact_tensor_id = + getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; + const bool use_q4_exact_tensor_id = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + can_wrap_q4_exact_tensors && + fuse_pair_swiglu && + direct_down_sum && + enable_q4_exact_tensor_id; + const bool q4_exact_boundary = + use_q4_exact_tensor_id && + g_batch_cb != nil && + getenv("DS4_METAL_DISABLE_Q4_EXACT_BOUNDARY") == NULL; + const bool q4_expert_table_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_expert_address_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_table_queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(q4_expert_table_auto || + q4_expert_address_auto); + const bool use_q4_expert_address_table = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && + (getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || + q4_expert_address_auto) && + (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled() || + getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL; + const bool enable_q4_expert_table = + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + q4_expert_table_auto; + const bool use_q4_expert_table = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil && + enable_q4_expert_table && + (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled() || + getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL; + const bool q4_table_boundary = + use_q4_expert_table && + g_batch_cb != nil && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + !q4_table_queue_residency && + getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; + const bool enable_q4_gather_slots = + getenv("DS4_METAL_ENABLE_Q4_GATHER_SLOTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GATHER_SLOTS") == NULL; + const bool use_q4_gather_slots = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + !use_q4_expert_table && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group6_expert_group_size * 6u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_q4_gather_slots6_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && + enable_q4_gather_slots; + const bool use_q4_selected_slots = + !force_resident && + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + !use_q4_expert_table && + !use_q4_gather_slots && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert >= 128 && + (g_ssd_streaming_mode || + (gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes)) && + fuse_pair_swiglu && + direct_down_sum && + ds4_gpu_q4_selected_paths_allowed() && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL; + const bool use_iq2_selected_slots = + !force_resident && + g_ssd_streaming_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + n_expert == 6 && + n_tokens == 1 && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; + const bool use_iq2_stream_addr_table = + !force_resident && + g_ssd_streaming_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_IQ2_XXS && + n_expert <= DS4_METAL_MAX_ROUTED_EXPERT_USED && + n_tokens == 1 && + fuse_pair_swiglu && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_iq2_xxs_pipeline != nil && + getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL; + const bool use_selected_slots = + use_q4_selected_slots || use_iq2_selected_slots || use_iq2_stream_addr_table; + id slots_pair_swiglu_pipeline = + use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline : + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; + id slots_sum6_pipeline = + use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline : + g_moe_mul_mv_slots6_q4_k_sum6_pipeline; + const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE"); + if (!selected_profile_env) { + selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE"); + } + const char *selected_profile_layer_env = getenv("DS4_METAL_SELECTED_PROFILE_LAYER"); + if (!selected_profile_layer_env) { + selected_profile_layer_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE_LAYER"); + } + bool selected_profile_layer_match = true; + if (selected_profile_layer_env && selected_profile_layer_env[0]) { + char *end = NULL; + long layer = strtol(selected_profile_layer_env, &end, 10); + selected_profile_layer_match = + end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index; + } + const bool selected_profile = + use_selected_slots && + selected_profile_env != NULL && + selected_profile_layer_match; + const bool q4_selected_shared_event = + use_q4_selected_slots && + getenv("DS4_METAL_Q4_SELECTED_SHARED_EVENT") != NULL; + const bool q4_selected_base_views = + use_q4_selected_slots && + getenv("DS4_METAL_Q4_SELECTED_USE_BASE_VIEWS") != NULL && + getenv("DS4_METAL_Q4_SELECTED_EXACT_VIEWS") == NULL; + const bool q4_selected_transient_views = + use_q4_selected_slots && + !q4_selected_base_views && + getenv("DS4_METAL_Q4_SELECTED_TRANSIENT_VIEWS") != NULL; + const char *q4_selected_view_mode = + q4_selected_base_views ? "base" : + (q4_selected_transient_views ? "transient" : "cached"); + if (!use_selected_slots) { + g_routed_moe_selected_override_n = 0; + } + if (use_q4_expert_address_table) { + gate_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert); + up_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert); + down_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert); + if (!gate_table || !up_table || !down_table) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32546); + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_expert_address_auto); + } else if (use_q4_expert_table) { + gate_table = ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + up_table = ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + down_table = ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32573); + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_expert_table_auto); + } else if (use_q4_exact_tensor_id) { + gate_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_inner); + up_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_inner); + down_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32601); return 0; } + } else if (use_q4_group6_experts || use_q4_gather_slots) { + if ((uint64_t)q4_group6_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group6_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group6 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group6_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group6_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 6; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group6 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group6_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group6_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group6_offsets[i] = (NSUInteger)group_inner; + if (!gate_group6_bufs[i] || !up_group6_bufs[i] || !down_group6_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32643); + return 0; + } + } + if (use_q4_gather_slots) { + if (gate_expert_bytes > NSUIntegerMax || + down_expert_bytes > NSUIntegerMax || + gate_expert_bytes > UINT64_MAX / 6u || + down_expert_bytes > UINT64_MAX / 6u || + 6ull * gate_expert_bytes > NSUIntegerMax || + 6ull * down_expert_bytes > NSUIntegerMax) { + fprintf(stderr, "ds4: Metal routed MoE Q4 gather scratch byte size overflow\n"); + return 0; + } + const NSUInteger gate_slots_bytes = (NSUInteger)(6ull * gate_expert_bytes); + const NSUInteger down_slots_bytes = (NSUInteger)(6ull * down_expert_bytes); + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_gate_slots_buffer, + &g_moe_q4_gate_slots_bytes, + gate_slots_bytes, + "ds4_moe_q4_gate_slots") || + !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_up_slots_buffer, + &g_moe_q4_up_slots_bytes, + gate_slots_bytes, + "ds4_moe_q4_up_slots") || + !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_down_slots_buffer, + &g_moe_q4_down_slots_bytes, + down_slots_bytes, + "ds4_moe_q4_down_slots")) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32670); + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + gate_slot_bufs[i] = g_moe_q4_gate_slots_buffer; + up_slot_bufs[i] = g_moe_q4_up_slots_buffer; + down_slot_bufs[i] = g_moe_q4_down_slots_buffer; + gate_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); + up_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); + down_slot_offsets[i] = (NSUInteger)((uint64_t)i * down_expert_bytes); + } + } + } else if (use_q4_group8_experts) { + if ((uint64_t)q4_group8_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group8_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group8 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group8_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group8_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 8; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group8 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group8_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group8_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group8_offsets[i] = (NSUInteger)group_inner; + if (!gate_group8_bufs[i] || !up_group8_bufs[i] || !down_group8_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32727); + return 0; + } + } + } else if (use_q4_group24_experts) { + if ((uint64_t)q4_group24_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group24_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group24 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group24_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group24_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 24; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group24 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group24_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group24_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group24_offsets[i] = (NSUInteger)group_inner; + if (!gate_group24_bufs[i] || !up_group24_bufs[i] || !down_group24_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32794); + return 0; + } + } + } else if (use_selected_slots) { + const bool selected_timing = + selected_profile || + ds4_gpu_stream_expert_timing_summary_enabled(); + double selected_t0 = selected_timing ? ds4_gpu_now_ms() : 0.0; + double selected_read_ms = 0.0; + double selected_sync_ms = 0.0; + double selected_copy_ms = 0.0; + double selected_wrap_ms = 0.0; + uint64_t selected_cache_hits0 = g_stream_expert_cache_hits; + uint64_t selected_cache_misses0 = g_stream_expert_cache_misses; + uint64_t selected_cache_wraps0 = g_stream_expert_cache_wraps; + uint64_t selected_cache_evictions0 = g_stream_expert_cache_evictions; + const char *selected_id_source = "readback"; + bool selected_ids_available = true; + bool selected_exec_ids_from_host = false; + const int stream_expert_cache_size_known = + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes); + const bool use_iq2_full_expert_addr_table = + use_iq2_selected_slots && + ds4_gpu_stream_full_expert_addr_table_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + use_stream_expert_cache = + !use_iq2_full_expert_addr_table && + (use_iq2_selected_slots || use_iq2_stream_addr_table || use_q4_selected_slots) && + stream_expert_cache_size_known && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + if (use_iq2_stream_addr_table && !use_stream_expert_cache) { + fprintf(stderr, + "ds4: Metal IQ2/IQ2 streaming decode requires a non-empty expert cache\n"); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32831); + return 0; + } + const bool stream_split_ready = + use_stream_expert_cache && + ds4_gpu_stream_expert_split_ready(); + const bool use_stream_compact_addr = + use_stream_expert_cache && + use_iq2_selected_slots && + ds4_gpu_stream_compact_addr_requested() && + !stream_split_ready && + !ds4_gpu_stream_expert_masked_addr_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + use_stream_expert_split_candidate = + use_stream_expert_cache && + use_iq2_selected_slots && + !use_stream_compact_addr && + stream_split_ready && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; + const bool use_stream_hit_validator = + use_stream_expert_cache && + use_iq2_selected_slots && + ds4_gpu_stream_expert_hit_validator_requested() && + g_moe_stream_expert_cache_validate_pipeline != nil && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil && + ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (use_iq2_full_expert_addr_table) { + selected_id_source = "gpu-full-addr"; + selected_ids_available = false; + g_routed_moe_selected_override_n = 0; + ds4_gpu_stream_expert_cache_entry *full_entry = NULL; + if (!ds4_gpu_stream_full_expert_addr_table_prepare(model_map, + model_size, + layer_index, + n_total_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + &full_entry)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32880); + return 0; + } + for (uint32_t i = 0; i < n_expert; i++) { + stream_slot_entries[i] = full_entry; + } + use_stream_expert_addr_table = true; + } else { + const int replayed_selected_ids = + ds4_gpu_moe_selected_trace_replay(selected_ids, n_expert); + if (replayed_selected_ids < 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32890); + return 0; + } else if (replayed_selected_ids > 0) { + selected_id_source = "replay"; + selected_exec_ids_from_host = true; + g_routed_moe_selected_override_n = 0; + } else if (g_routed_moe_selected_override_n == n_expert) { + memcpy(selected_ids, + g_routed_moe_selected_override, + (size_t)n_expert * sizeof(selected_ids[0])); + selected_id_source = "override"; + selected_exec_ids_from_host = true; + g_routed_moe_selected_override_n = 0; + } else if (use_stream_hit_validator) { + g_routed_moe_selected_override_n = 0; + uint32_t validator_all_cached = 0; + uint32_t validator_miss_mask = 0; + uint32_t validator_invalid_mask = 0; + if (!ds4_gpu_stream_expert_cache_validate_selected(selected, + stream_gate_addr_buf, + stream_up_addr_buf, + stream_down_addr_buf, + n_total_expert, + n_expert, + selected_ids, + &validator_all_cached, + &validator_miss_mask, + &validator_invalid_mask)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32917); + return 0; + } + selected_id_source = + validator_invalid_mask != 0 ? "validator-invalid" : + (validator_miss_mask != 0 ? "validator-miss" : + (validator_all_cached != 0 ? "validator-hit" : "validator-miss")); + } else { + g_routed_moe_selected_override_n = 0; + if (g_batch_cb != nil) { + double selected_boundary_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (q4_selected_shared_event) { + if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32942); return 0; } + } else if (ds4_gpu_end_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32931); + return 0; + } + if (selected_timing) { + selected_sync_ms += + ds4_gpu_now_ms() - selected_boundary_t0; + } + double selected_copy_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32943); + return 0; + } + if (selected_timing) { + selected_copy_ms += + ds4_gpu_now_ms() - selected_copy_t0; + } + if (!q4_selected_shared_event) { + selected_boundary_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_begin_commands() == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32967); return 0; } + if (selected_timing) { + selected_sync_ms += + ds4_gpu_now_ms() - selected_boundary_t0; + } + } + } else { + double selected_copy_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32965); + return 0; + } + if (selected_timing) { + selected_copy_ms += + ds4_gpu_now_ms() - selected_copy_t0; + } + } + } + } + if (selected_timing) { + selected_read_ms = selected_sync_ms + selected_copy_ms; + selected_t0 = ds4_gpu_now_ms(); + } + + if (selected_ids_available) { + for (uint32_t i = 0; i < n_expert; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal routed MoE selected expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32986); + return 0; + } + } + ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, + selected_ids, + n_expert); + if (!ds4_gpu_moe_selected_trace_record(selected_ids, n_expert) || + !ds4_gpu_moe_selected_hotlist_record(layer_index, + selected_ids, + n_expert, + n_total_expert)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32997); + return 0; + } + + for (uint32_t i = 0; i < n_expert; i++) { + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); + return 0; + } + stream_gate_abs_offsets[i] = gate_offset + gate_rel; + stream_up_abs_offsets[i] = up_offset + gate_rel; + stream_down_abs_offsets[i] = down_offset + down_rel; + + if (use_stream_expert_cache) { + ds4_gpu_stream_expert_cache_entry *entry = NULL; + entry = ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer_index, + (uint32_t)selected_ids[i], + n_total_expert, + n_expert, + stream_gate_abs_offsets[i], + stream_up_abs_offsets[i], + stream_down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!entry) { + stream_expert_missing_mask |= 1u << i; + continue; + } + stream_expert_resident_mask |= 1u << i; + stream_slot_entries[i] = entry; + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + continue; + } + + uint64_t slot_inner = 0; + gate_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner)); + gate_slot_offsets[i] = (NSUInteger)slot_inner; + slot_inner = 0; + up_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner)); + up_slot_offsets[i] = (NSUInteger)slot_inner; + slot_inner = 0; + down_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner)); + down_slot_offsets[i] = (NSUInteger)slot_inner; + if (!gate_slot_bufs[i] || !up_slot_bufs[i] || !down_slot_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33105); + return 0; + } + } + } + if (use_stream_expert_cache) { + use_stream_expert_addr_table = + ((use_iq2_selected_slots && + ds4_gpu_stream_expert_addr_table_kernel_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil) || + use_iq2_stream_addr_table) && + ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (use_iq2_stream_addr_table && !use_stream_expert_addr_table) { + fprintf(stderr, + "ds4: Metal IQ2/IQ2 streaming decode could not prepare expert address buffers\n"); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33123); + return 0; + } + use_stream_expert_masked_addr_table = + use_stream_expert_addr_table && + use_iq2_selected_slots && + ds4_gpu_stream_expert_masked_addr_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; + use_stream_expert_split_deferred = + use_stream_expert_split_candidate && + use_stream_expert_masked_addr_table && + stream_expert_resident_mask != 0 && + stream_expert_missing_mask != 0 && + ds4_gpu_stream_expert_split_worthwhile(stream_expert_resident_mask, + stream_expert_missing_mask) && + g_batch_cb != nil && + getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL; + if (use_stream_expert_split_deferred) { + const ds4_gpu_stream_expert_table table = { + .model_map = model_map, + .model_size = model_size, + .layer = layer_index, + .n_total_expert = n_total_expert, + .gate_offset = gate_offset, + .up_offset = up_offset, + .down_offset = down_offset, + .gate_expert_bytes = gate_expert_bytes, + .down_expert_bytes = down_expert_bytes, + }; + if (!ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_ids, + n_expert)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33156); + return 0; + } + } + if (stream_expert_missing_mask != 0 && + !use_stream_expert_split_deferred) { + if (!ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_expert_missing_mask, + stream_slot_entries)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33175); + return 0; + } + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_expert_missing_mask & (1u << i)) == 0) continue; + ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; + if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); return 0; } + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + } + } + if (use_iq2_stream_addr_table && use_stream_expert_addr_table) { + for (uint32_t i = 0; i < n_expert; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; + if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33214); return 0; } + if (!ds4_gpu_stream_expert_cache_set_addr_slot_raw( + layer_index, + (uint32_t)selected_ids[i], + entry->gate_buffer, + entry->gate_inner, + entry->up_buffer, + entry->up_inner, + entry->down_buffer, + entry->down_inner)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); + return 0; + } + } + } + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + selected_ids, + n_expert); + if (use_stream_compact_addr) { + id compact_selected = nil; + if (!ds4_gpu_stream_compact_addr_prepare(layer_index, + stream_slot_entries, + n_expert, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + &compact_selected)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33223); + return 0; + } + selected_exec_buf = compact_selected; + selected_exec_off = 0; + use_stream_expert_addr_table = true; + use_stream_expert_masked_addr_table = false; + use_stream_compact_addr_table = true; + } + if (use_stream_expert_addr_table && + !use_stream_compact_addr_table && + selected_exec_ids_from_host) { + if (!ds4_gpu_stream_selected_ids_prepare(layer_index, + selected_ids, + n_expert, + &selected_exec_buf, + &selected_exec_off)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33239); + return 0; + } + } + } + if (selected_timing) { + selected_wrap_ms = ds4_gpu_now_ms() - selected_t0; + if (use_stream_expert_cache) { + ds4_gpu_stream_expert_timing_note_cache_class( + stream_expert_resident_mask, + stream_expert_missing_mask); + } + ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, + selected_copy_ms, + selected_wrap_ms); + } + if (selected_profile) { + const uint64_t selected_cache_hits = + g_stream_expert_cache_hits - selected_cache_hits0; + const uint64_t selected_cache_misses = + g_stream_expert_cache_misses - selected_cache_misses0; + const uint64_t selected_cache_wraps = + g_stream_expert_cache_wraps - selected_cache_wraps0; + const uint64_t selected_cache_evictions = + g_stream_expert_cache_evictions - selected_cache_evictions0; + const char *selected_path = + use_iq2_stream_addr_table ? "iq2/iq2" : + (use_iq2_selected_slots ? "iq2/q2" : "q4/q4"); + const char *selected_view_mode = + use_stream_expert_split_deferred ? "stream-split" : + use_stream_expert_masked_addr_table ? "stream-addr-mask" : + use_stream_compact_addr_table ? "stream-compact-addr" : + use_stream_expert_addr_table ? "stream-addr" : + (use_stream_expert_cache ? "stream-cache" : + (use_iq2_selected_slots ? "exact-cache" : q4_selected_view_mode)); + fprintf(stderr, + "ds4: Metal selected views layer=%u path=%s mode=%s ids=%s " + "experts=%d,%d,%d,%d,%d,%d expert_gate=%.2f MiB " + "expert_down=%.2f MiB read=%.3f ms bind=%.3f ms " + "cache_hits=%llu cache_misses=%llu cache_wraps=%llu cache_evictions=%llu\n", + layer_index, + selected_path, + selected_view_mode, + selected_id_source, + selected_ids_available ? selected_ids[0] : -1, + selected_ids_available ? selected_ids[1] : -1, + selected_ids_available ? selected_ids[2] : -1, + selected_ids_available ? selected_ids[3] : -1, + selected_ids_available ? selected_ids[4] : -1, + selected_ids_available ? selected_ids[5] : -1, + ds4_gpu_mib(gate_expert_bytes), + ds4_gpu_mib(down_expert_bytes), + selected_read_ms, + selected_wrap_ms, + (unsigned long long)selected_cache_hits, + (unsigned long long)selected_cache_misses, + (unsigned long long)selected_cache_wraps, + (unsigned long long)selected_cache_evictions); + } + } else if (!use_q4_grouped_experts) { + gate_buf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, gate_tensor_bytes, &gate_inner); + up_buf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, gate_tensor_bytes, &up_inner); + down_buf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, down_tensor_bytes, &down_inner); + if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33326); return 0; } + } + if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33305); + return 0; + } + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33337); return 0; } + if ((use_q4_expert_address_table || use_q4_expert_table) && + !ds4_gpu_use_model_residency_set(cb)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33314); + return 0; + } + if (!q4_table_queue_residency && + q4_table_layer_residency && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:q4_table_layer_residency]; + } + + const bool moe_one_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_ONE_STAGE_PROFILE", + "DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER", + layer_index); + const char *moe_one_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); + const char *moe_one_path = + use_q4_grouped_experts ? "q4_grouped_pair_swiglu" : + use_q4_group6_experts ? "q4_group6_pair_swiglu" : + use_q4_group8_experts ? "q4_group8_pair_swiglu" : + use_q4_group24_experts ? "q4_group24_split_gate_up" : + use_q4_exact_tensor_id ? "q4_exact_pair_swiglu" : + use_q4_expert_address_table ? "q4_addr_pair_swiglu" : + use_q4_expert_table ? "q4_table_pair_swiglu" : + use_q4_gather_slots ? "q4_gather_slots6_pair_swiglu" : + use_stream_expert_split_deferred ? "iq2_stream_split_pair_swiglu" : + use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" : + use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" : + use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" : + use_q4_selected_slots ? "q4_slots6_pair_swiglu" : + (fuse_pair_swiglu ? "pair_swiglu" : + ((!g_quality_mode && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || + (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline))) ? "pair" : "single")); + double moe_one_stage_t0 = moe_one_stage_profile ? ds4_gpu_now_ms() : 0.0; + if (moe_one_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33349); + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33380); return 0; } + moe_one_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_MOE_ONE_STAGE(name) do { \ + if (ok && moe_one_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !moe_one_stage_filter || !moe_one_stage_filter[0] || \ + strstr(stage_name, moe_one_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal routed MoE one stage layer=%u pairs=%u experts=%u " \ + "gate=%s down=%s path=%s %s=%.3f ms\n", \ + layer_index, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + moe_one_path, \ + stage_name, now_ms - moe_one_stage_t0); \ + } \ + moe_one_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + if (use_q4_gather_slots) { + ds4_gpu_q4_gather_slots6_args gate_gather_args = { + .expert_bytes = gate_expert_bytes, + .group_size = q4_group6_expert_group_size, + .n_slots = n_expert, + }; + ds4_gpu_q4_gather_slots6_args down_gather_args = { + .expert_bytes = down_expert_bytes, + .group_size = q4_group6_expert_group_size, + .n_slots = n_expert, + }; + ok = ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &gate_gather_args, + gate_group6_bufs, + gate_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_gate_slots_buffer, + 0) && + ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &gate_gather_args, + up_group6_bufs, + up_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_up_slots_buffer, + 0) && + ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &down_gather_args, + down_group6_bufs, + down_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_down_slots_buffer, + 0); + } + if (use_q4_gather_slots) { + DS4_METAL_PROFILE_MOE_ONE_STAGE("q4_gather"); + } + if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33455); return 0; } + if (use_q4_grouped_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { + const uint32_t expert_count = + q4_expert_group_size < n_total_expert - expert_base ? + q4_expert_group_size : n_total_expert - expert_base; + if ((uint64_t)expert_base > UINT64_MAX / gate_expert_bytes || + (uint64_t)expert_count > UINT64_MAX / gate_expert_bytes) { + ok = 0; + break; + } + const uint64_t group_rel = (uint64_t)expert_base * gate_expert_bytes; + const uint64_t group_bytes = (uint64_t)expert_count * gate_expert_bytes; + if (group_rel > UINT64_MAX - gate_offset || + group_rel > UINT64_MAX - up_offset) { + ok = 0; + break; + } + uint64_t gate_group_inner = 0; + uint64_t up_group_inner = 0; + id gate_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + group_rel, + group_bytes, + &gate_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + gate_offset + group_rel, + group_bytes, + &gate_group_inner); + id up_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + group_rel, + group_bytes, + &up_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + up_offset + group_rel, + group_bytes, + &up_group_inner); + if (!gate_group_buf || !up_group_buf) { + ok = 0; + break; + } + ds4_gpu_moe_expert_group_args group_args = { + .expert_base = expert_base, + .expert_count = expert_count, + .accumulate = 0, + .pad0 = 0, + }; + ok = ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(cb, + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + &group_args, + gate_group_buf, + (NSUInteger)gate_group_inner, + up_group_buf, + (NSUInteger)up_group_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else if (use_q4_expert_address_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(cb, + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_expert_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + q4_table_queue_residency); + } else if (use_q4_group6_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_group6_pair_swiglu(cb, + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_group6_bufs, + gate_group6_offsets, + up_group6_bufs, + up_group6_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_group8_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_group8_pair_swiglu(cb, + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_group8_bufs, + gate_group8_offsets, + up_group8_bufs, + up_group8_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_group24_experts) { + ok = ds4_gpu_encode_mul_mv_group24_id(cb, + g_moe_mul_mv_group24_q4_k_id_pipeline, + &gate_args, + gate_group24_bufs, + gate_group24_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false) && + ds4_gpu_encode_mul_mv_group24_id(cb, + g_moe_mul_mv_group24_q4_k_id_pipeline, + &gate_args, + up_group24_bufs, + up_group24_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else if (use_q4_gather_slots || use_selected_slots) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + if (use_stream_expert_addr_table) { + if (use_stream_expert_masked_addr_table) { + if (use_stream_expert_split_deferred) { + const bool stream_split_profile = + getenv("DS4_METAL_STREAMING_EXPERT_SPLIT_PROFILE") != NULL; + const bool stream_split_timing = + stream_split_profile || + ds4_gpu_stream_expert_timing_summary_enabled(); + double stream_split_t0 = + stream_split_timing ? ds4_gpu_now_ms() : 0.0; + ds4_gpu_stream_expert_split_args resident_pair_args = { + .active_mask = stream_expert_resident_mask, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &resident_pair_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + if (ok) { + ok = ds4_gpu_flush_commands(); + if (ok) { + cb = ds4_gpu_command_buffer(&owned); + if (!cb) ok = 0; + } + } + const double stream_split_resident_ms = + stream_split_timing ? ds4_gpu_now_ms() - stream_split_t0 : 0.0; + if (stream_split_timing) stream_split_t0 = ds4_gpu_now_ms(); + const double stream_split_missing_start_ms = stream_split_t0; + double stream_split_missing_load_ms = 0.0; + double stream_split_missing_slot_ms = 0.0; + double stream_split_missing_prune_ms = 0.0; + double stream_split_missing_addr_ms = 0.0; + double stream_split_missing_wait_ms = 0.0; + if (ok) { + ok = ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_expert_missing_mask, + stream_slot_entries); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_load_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + if (ok) { + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_expert_missing_mask & (1u << i)) == 0) continue; + ds4_gpu_stream_expert_cache_entry *entry = + stream_slot_entries[i]; + if (!entry) { + ok = 0; + break; + } + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + } + } + } + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_slot_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + if (ok) { + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + selected_ids, + n_expert); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_prune_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_addr_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + } + if (ok) { + /* + * The resident stage was submitted before the + * CPU read of missing experts so I/O can overlap + * with GPU work. The missing stage reuses the same + * gate/up/mid scratch buffers, so it must not + * execute until the resident command buffer has + * finished. The down/sum pass is issued once after + * all six mid slots exist; this keeps the final + * accumulation order stable regardless of the + * resident/missing split. + */ + ok = ds4_gpu_wait_pending_command_buffers( + "streaming expert split resident"); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_wait_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + } + const double stream_split_missing_ms = + stream_split_timing ? + ds4_gpu_now_ms() - stream_split_missing_start_ms : + 0.0; + ds4_gpu_stream_expert_split_args missing_pair_args = { + .active_mask = stream_expert_missing_mask, + .accumulate = 0u, + }; + ds4_gpu_stream_expert_split_args all_down_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + if (ok) { + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &missing_pair_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + if (ok) { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, + &down_args, + &all_down_args, + stream_slot_entries, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2); + } + stream_expert_split_completed = ok; + if (stream_split_timing) { + ds4_gpu_stream_expert_timing_note_split( + stream_expert_resident_mask, + stream_expert_missing_mask, + stream_split_resident_ms, + stream_split_missing_ms); + ds4_gpu_stream_expert_timing_note_split_missing_detail( + stream_split_missing_load_ms, + stream_split_missing_slot_ms, + stream_split_missing_prune_ms, + stream_split_missing_addr_ms, + stream_split_missing_wait_ms); + } + if (stream_split_profile) { + fprintf(stderr, + "ds4: Metal streaming expert split layer=%u " + "resident=0x%02x missing=0x%02x resident_submit=%.3f ms " + "missing_bind=%.3f ms\n", + layer_index, + stream_expert_resident_mask, + stream_expert_missing_mask, + stream_split_resident_ms, + stream_split_missing_ms); + } + } else { + ds4_gpu_stream_expert_split_args split_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &split_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else { + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, + &gate_args, + &act_args, + stream_slot_entries, + n_expert, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + nil, + nil); + } + } else { + ok = (!use_stream_expert_cache || + ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_slot_entries, + n_expert, + 0)) && + ds4_gpu_encode_mul_mv_slots6_pair_swiglu(cb, + slots_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_slot_bufs, + gate_slot_offsets, + up_slot_bufs, + up_slot_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else if (fuse_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, + pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (!g_quality_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + g_moe_mul_mv_id_iq2_xxs_pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else if (!g_quality_mode && + gate_type == DS4_METAL_TENSOR_Q4_K && + g_moe_mul_mv_id_q4_k_pair_pipeline) { + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + g_moe_mul_mv_id_q4_k_pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0) && + ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("gate_up"); + if (ok && (!fuse_pair_swiglu || use_q4_group24_experts)) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + clamp, + false); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : + (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); + if (ok && stream_expert_split_completed) { + /* The split path already wrote the resident partial output and + * accumulated the missing experts into out. */ + } else if (ok && use_q4_grouped_experts) { + bool first_group = true; + for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { + const uint32_t expert_count = + q4_expert_group_size < n_total_expert - expert_base ? + q4_expert_group_size : n_total_expert - expert_base; + if ((uint64_t)expert_base > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_count > UINT64_MAX / down_expert_bytes) { + ok = 0; + break; + } + const uint64_t group_rel = (uint64_t)expert_base * down_expert_bytes; + const uint64_t group_bytes = (uint64_t)expert_count * down_expert_bytes; + if (group_rel > UINT64_MAX - down_offset) { + ok = 0; + break; + } + uint64_t down_group_inner = 0; + id down_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + group_rel, + group_bytes, + &down_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + down_offset + group_rel, + group_bytes, + &down_group_inner); + if (!down_group_buf) { + ok = 0; + break; + } + ds4_gpu_moe_expert_group_args group_args = { + .expert_base = expert_base, + .expert_count = expert_count, + .accumulate = first_group ? 0u : 1u, + .pad0 = 0, + }; + ok = ds4_gpu_encode_mul_mv_group_q4_sum6(cb, + g_moe_mul_mv_group_q4_k_sum6_pipeline, + &down_args, + &group_args, + down_group_buf, + (NSUInteger)down_group_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + first_group = false; + } + } else if (ok && use_q4_expert_address_table) { + ok = ds4_gpu_encode_mul_mv_addr_q4_sum6(cb, + g_moe_mul_mv_addr_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_expert_table) { + ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, + g_moe_mul_mv_table_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + q4_table_queue_residency); + } else if (ok && use_q4_group6_experts) { + ok = ds4_gpu_encode_mul_mv_group6_sum6(cb, + g_moe_mul_mv_group6_q4_k_sum6_pipeline, + &down_args, + down_group6_bufs, + down_group6_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_group8_experts) { + ok = ds4_gpu_encode_mul_mv_group8_sum6(cb, + g_moe_mul_mv_group8_q4_k_sum6_pipeline, + &down_args, + down_group8_bufs, + down_group8_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_group24_experts) { + ok = ds4_gpu_encode_mul_mv_group24_sum6(cb, + g_moe_mul_mv_group24_q4_k_sum6_pipeline, + &down_args, + down_group24_bufs, + down_group24_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && (use_q4_gather_slots || use_selected_slots)) { + if (use_stream_expert_addr_table) { + if (down_type == DS4_METAL_TENSOR_IQ2_XXS) { + ok = ds4_gpu_encode_mul_mv_addr_iq2(cb, + g_moe_mul_mv_addr_iq2_xxs_pipeline, + &down_args, + stream_slot_entries, + n_expert, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selected_exec_buf, + selected_exec_off, + down_smem, + 2, + false); + } else if (use_stream_expert_masked_addr_table) { + ds4_gpu_stream_expert_split_args split_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, + &down_args, + &split_args, + stream_slot_entries, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2); + } else { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(cb, + g_moe_mul_mv_addr_q2_k_sum6_pipeline, + &down_args, + stream_slot_entries, + n_expert, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2, + nil); + } + } else { + ok = (!use_stream_expert_cache || + ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_slot_entries, + n_expert, + 0)) && + ds4_gpu_encode_mul_mv_slots6_sum6(cb, + slots_sum6_pipeline, + &down_args, + down_slot_bufs, + down_slot_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + down_smem, + 2); + } + } else if (ok && direct_down_sum) { + ok = ds4_gpu_encode_mul_mv_id_sum6(cb, + down_sum6_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + add_in ? ds4_gpu_tensor_buffer(add_in) : nil, + add_in ? ds4_gpu_tensor_offset(add_in) : 0, + down_smem, + 2); + } else if (ok) { + ok = ds4_gpu_encode_mul_mv_id(cb, + down_mv_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + down_nsg, + down_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("down"); + if (ok && n_expert > 1 && !direct_down_sum && !stream_expert_split_completed) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("sum"); + if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; } + if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34372); + return 0; + } + } +#undef DS4_METAL_PROFILE_MOE_ONE_STAGE + } + + return 1; +} + +int ds4_gpu_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + bool *mid_is_f16, + bool force_resident) { + (void)force_resident; + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding (see ds4_gpu_routed_moe_one_tensor): bind from the owned + * expert range and rebase ids in the kernels. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * gate_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || + n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || + n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED) { + return 0; + } + if (gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_row_bytes == 0 || down_row_bytes == 0) { + return 0; + } + if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) return 0; + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed batch MoE tensor byte size overflow\n"); + return 0; + } + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + /* + * PRO Q4 routed expert tensors are multi-GiB per layer. A one-token + * layer-slice prefill should use the one-token path so it can either bind + * exact tensor views with GPU-selected IDs or fall back to selected-expert + * views. Keep this guarded by tensor size so Flash and mixed-Flash Q4 keep + * their existing fast path. + */ + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + const bool can_single_token_q4_grouped = + getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL && + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_selected_slots = + ds4_gpu_q4_selected_paths_allowed() && + getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group6 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group8 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group24 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group24_q4_k_id_pipeline != nil && + g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil; + const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; + const bool can_single_token_q4_exact_tensor_id = + n_total_expert == 384 && + max_buffer_len != 0 && + gate_tensor_bytes <= max_buffer_len && + down_tensor_bytes <= max_buffer_len && + getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; + const bool enable_single_token_q4_expert_table = + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool can_single_token_q4_expert_table = + n_total_expert == 384 && + enable_single_token_q4_expert_table && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil; + const bool can_single_token_q4_expert_address_table = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil; + const bool use_single_token_q4_one_tensor = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_tokens == 1 && + n_expert == 6 && + n_total_expert >= 128 && + (g_ssd_streaming_mode || + (gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes)) && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + (can_single_token_q4_grouped || + can_single_token_q4_group6 || + can_single_token_q4_group8 || + can_single_token_q4_group24 || + can_single_token_q4_exact_tensor_id || + can_single_token_q4_expert_address_table || + can_single_token_q4_expert_table || + can_single_token_q4_selected_slots); + if (use_single_token_q4_one_tensor) { + if (mid_is_f16) *mid_is_f16 = false; + return ds4_gpu_routed_moe_one_tensor(out, + gate, + up, + mid, + experts, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + clamp, + x, + NULL, + layer_index, + false); + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id expertsbuf = ds4_gpu_tensor_buffer(experts); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)n_tokens * expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert * sizeof(int); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < mid_bytes || + ds4_gpu_tensor_bytes(up) < mid_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal routed batch MoE received undersized activation buffers\n"); + return 0; + } + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + DS4MetalQ4ExpertTable *gate_table = nil; + DS4MetalQ4ExpertTable *up_table = nil; + DS4MetalQ4ExpertTable *down_table = nil; + id q4_table_layer_residency = nil; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + id stream_overflow_gate = nil; + id stream_overflow_up = nil; + id stream_overflow_down = nil; + ds4_gpu_stream_expert_cache_entry + *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; + uint32_t stream_resource_count = 0; + uint32_t stream_unique = 0; + + const uint32_t pair_rows = n_tokens * n_expert; + const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); + + const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); + const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); + id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); + id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); + id gate_mm_pipeline = nil; + id up_mm_pipeline = nil; + id down_mm_pipeline = nil; + id pair_swiglu_mm_pipeline = nil; + if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { + fprintf(stderr, "ds4: unsupported Metal routed batch MoE quant types gate=%u down=%u\n", + gate_type, down_type); + return 0; + } + const bool use_iq2_batch_selected_addr = + ds4_gpu_stream_prefill_batch_selected_addr_enabled(n_tokens, + n_total_expert, + n_expert, + gate_type, + down_type) && + n_tokens > 1 && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + + ds4_gpu_mul_mv_id_args gate_args = + ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens, gate_nr0); + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = tp_expert_base_host; + ds4_gpu_mul_mv_id_args down_args = + ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, down_nr0); + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_expert_base = tp_expert_base_host; + const bool q4_batch_expert_table_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_batch_table_queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(q4_batch_expert_table_auto); + const bool enable_q4_batch_expert_table = + getenv("DS4_METAL_ENABLE_Q4_BATCH_EXPERT_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + q4_batch_expert_table_auto; + const bool use_q4_batch_expert_table = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens > 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + getenv("DS4_METAL_DISABLE_Q4_BATCH_EXPERT_TABLE") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil && + enable_q4_batch_expert_table && + (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_batch_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled()); + const bool use_mm_id = + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr && + n_tokens >= 32u && + ds4_gpu_mul_mm_id_map0_name(n_expert) != NULL; + /* + * MTP verification is neither normal decode nor large prefill: the + * target model must verify a tiny suffix (up to DSpark's 5-token + * block) in one layer-major pass. For that shape the prefill + * expert-major GEMM path + * is too large, but the decode pair kernels are exactly the right + * primitive: they read the same activation once and compute routed + * gate/up together for every selected expert row. Keep this limited to + * tiny batches so ordinary prefill keeps using the higher-throughput + * grouped matmul path. + */ + const bool use_tiny_pair_mv = + !g_quality_mode && + n_tokens <= 5u && + !use_q4_batch_expert_table && + !use_mm_id && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || + (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline)); + id tiny_pair_swiglu_pipeline = nil; + if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { + tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { + tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } + const bool use_tiny_pair_swiglu = + use_tiny_pair_mv && + tiny_pair_swiglu_pipeline != nil && + getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL; + ds4_gpu_mul_mm_id_map_args gate_map_args = { 0 }; + ds4_gpu_mul_mm_id_args gate_mm_args = { 0 }; + ds4_gpu_mul_mm_id_args down_mm_args = { 0 }; + id map_pipeline = nil; + /* + * The grouped routed-MoE matmul loads activation tiles as half before + * using SIMD-group MMA. Store the SwiGLU/route-weight intermediate in + * that same precision so the down projection avoids a large F32 mid + * write/read. --quality keeps the older F32 intermediate. + */ + const bool request_mid_f16 = + !g_quality_mode && + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr; + /* + * Fused gate+up grouped matmul with the SwiGLU epilogue. The IQ2 + * variant stays opt-in behind its env flag; the Q4_K variant is the + * default path — same MMA accumulation order and epilogue math as the + * separate GEMMs + swiglu pass, so the mid tensor is bit-identical. + */ + const bool use_mm_id_pair_swiglu = + use_mm_id && + g_tp_split_world != 2 && /* pair-swiglu mm kernel lacks expert ownership */ + request_mid_f16 && + n_expert == 6 && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) || + (gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K)) && + getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL; + if (use_mm_id) { + gate_map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + gate_mm_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + down_mm_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + request_mid_f16 ? sizeof(uint16_t) : sizeof(float)); + gate_mm_args.tp_rank = g_tp_split_rank; + gate_mm_args.tp_world = g_tp_split_world; + gate_mm_args.tp_expert_base = tp_expert_base_host; + down_mm_args.tp_rank = g_tp_split_rank; + down_mm_args.tp_world = g_tp_split_world; + down_mm_args.tp_expert_base = tp_expert_base_host; + + map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + down_mm_pipeline = request_mid_f16 ? + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) : + ds4_gpu_routed_mm_pipeline(down_type); + if (use_mm_id_pair_swiglu) { + pair_swiglu_mm_pipeline = + ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : + "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); + } + if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline || + (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) { + return 0; + } + } + + if (use_iq2_batch_selected_addr) { + const int had_batch = g_batch_cb != nil; + if (had_batch && ds4_gpu_end_commands() == 0) { + return 0; + } + g_stream_prefill_batch_selected_addr_building++; + if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( + model_map, + model_size, + layer_index, + selected, + n_tokens, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + stream_resources, + &stream_resource_count, + &stream_unique, + &stream_overflow_gate, + &stream_overflow_up, + &stream_overflow_down)) { + g_stream_prefill_batch_selected_addr_building--; + return 0; + } + g_stream_prefill_batch_selected_addr_building--; + if (stream_unique == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "produced no resident experts\n", + layer_index); + return 0; + } + for (uint32_t i = 0; i < stream_resource_count; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_resources[i]; + if (!entry || + !entry->valid || + !entry->gate_buffer || + !entry->up_buffer || + !entry->down_buffer) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "lost resident expert resource %u/%u during preparation\n", + layer_index, + i, + stream_resource_count); + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + if (had_batch && ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to reopen command batch after preparation\n", + layer_index); + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + + if (use_q4_batch_expert_table) { + gate_table = ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + up_table = ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + down_table = ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_batch_expert_table_auto); + if ((q4_batch_table_queue_residency || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL) && + !q4_table_layer_residency && + !ds4_gpu_q4_table_model_residency_enabled()) { + fprintf(stderr, "ds4: Metal Q4 batch expert table residency set is not available\n"); + return 0; + } + } else if (use_iq2_batch_selected_addr) { + if (n_expert > 1 && (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch")) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + } else { + if (n_expert > 1 && (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch")) { + return 0; + } + } + gate_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_inner); + up_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_inner); + down_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) return 0; + if (getenv("DS4_GLM_TP_DEBUG") && layer_index == 3) { + fprintf(stderr, + "ds4: batch mv binds l=%u rank=%d base=%d gate=%llu+%llu " + "up=%llu down=%llu inner=%llu/%llu/%llu ne02=%d nei0=%d nr0=%d\n", + layer_index, g_tp_split_rank, tp_expert_base_host, + (unsigned long long)gate_offset, + (unsigned long long)gate_tensor_bytes, + (unsigned long long)up_offset, + (unsigned long long)down_offset, + (unsigned long long)gate_inner, + (unsigned long long)up_inner, + (unsigned long long)down_inner, + gate_args.ne02, gate_args.nei0, gate_args.nr0); + } + } + + const bool q4_batch_table_boundary = + use_q4_batch_expert_table && + g_batch_cb != nil && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + !q4_batch_table_queue_residency && + getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; + if (q4_batch_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (use_q4_batch_expert_table && !ds4_gpu_use_model_residency_set(cb)) { + return 0; + } + if (!q4_batch_table_queue_residency && + q4_table_layer_residency && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:q4_table_layer_residency]; + } + const bool moe_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_STAGE_PROFILE", + "DS4_METAL_MOE_STAGE_PROFILE_LAYER", + layer_index); + const char *moe_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); + const char *moe_path = + use_q4_batch_expert_table ? "q4_table_pair_swiglu" : + use_iq2_batch_selected_addr ? "iq2_batch_stream_addr" : + use_mm_id_pair_swiglu ? "mm_id_pair_swiglu" : + use_mm_id ? "mm_id" : + use_tiny_pair_swiglu ? "tiny_pair_swiglu" : + (use_tiny_pair_mv ? "tiny_pair_mv" : "mv"); + double moe_stage_t0 = moe_stage_profile ? ds4_gpu_now_ms() : 0.0; + if (moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + moe_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_MOE_STAGE(name) do { \ + if (ok && moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !moe_stage_filter || !moe_stage_filter[0] || \ + strstr(stage_name, moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s path=%s mid=%s %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + moe_path, \ + request_mid_f16 ? "f16" : "f32", \ + stage_name, now_ms - moe_stage_t0); \ + } \ + moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); + const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); + const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); + const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); + const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); + const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); + id down_sum6_pipeline = nil; + if (down_type == DS4_METAL_TENSOR_Q2_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_Q4_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } + const bool direct_down_sum = + !g_quality_mode && + !use_q4_batch_expert_table && + !use_mm_id && + n_expert == 6 && + n_tokens <= 4u && + down_sum6_pipeline != nil; + int ok = 0; + if (use_iq2_batch_selected_addr) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( + cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, + &gate_args, + &act_args, + stream_resources, + stream_resource_count, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + stream_overflow_gate, + stream_overflow_up); + if (!ok) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to encode gate/up path tokens=%u unique=%u\n", + layer_index, + n_tokens, + stream_unique); + } + } else if (use_q4_batch_expert_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + q4_batch_table_queue_residency); + } else if (use_mm_id) { + /* + * The routed pair ids are the same for gate, up, and down. Build + * the expert-major work map once, then reuse it for all three + * batched expert matmuls. + */ + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &gate_map_args, + &gate_mm_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_MOE_STAGE("map"); + if (ok && use_mm_id_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(uint16_t), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb, + pair_swiglu_mm_pipeline, + &gate_mm_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights)); + DS4_METAL_PROFILE_MOE_STAGE("gate_up_fused"); + } else if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + gate_mm_pipeline, + &gate_mm_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + 8192u); + DS4_METAL_PROFILE_MOE_STAGE("gate"); + } + if (ok && !use_mm_id_pair_swiglu) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + up_mm_pipeline, + &gate_mm_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + 8192u); + DS4_METAL_PROFILE_MOE_STAGE("up"); + } + } else if (use_tiny_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, + tiny_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_tiny_pair_mv) { + id pair_pipeline = + gate_type == DS4_METAL_TENSOR_IQ2_XXS ? + g_moe_mul_mv_id_iq2_xxs_pair_pipeline : + g_moe_mul_mv_id_q4_k_pair_pipeline; + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0) && + ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_STAGE("gate_up"); + const bool use_fused_activation = !g_quality_mode && !use_q4_batch_expert_table; + const bool use_mid_f16 = + use_mm_id && + use_fused_activation && + request_mid_f16; + if (mid_is_f16) *mid_is_f16 = use_mid_f16; + if (ok && use_iq2_batch_selected_addr) { + /* The address-table pair kernel already wrote weighted SwiGLU rows into mid. */ + } else if (ok && use_q4_batch_expert_table) { + /* The table pair kernel already wrote weighted SwiGLU rows into mid. */ + } else if (ok && use_mm_id_pair_swiglu) { + /* The fused batch mm_id pair kernel already wrote weighted f16 SwiGLU rows into mid. */ + } else if (ok && use_tiny_pair_swiglu) { + /* The fused tiny pair kernel already wrote weighted F32 SwiGLU rows into mid. */ + } else if (ok && use_fused_activation) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + clamp, + use_mid_f16); + } else if (ok && clamp > 1.0e-6f) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + gatebuf, + ds4_gpu_tensor_offset(gate), + gatebuf, + ds4_gpu_tensor_offset(gate), + expert_mid_dim, + pair_rows, + 0, + -FLT_MAX, + clamp); + if (ok) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_silu_pipeline, + gatebuf, + ds4_gpu_tensor_offset(gate), + midbuf, + ds4_gpu_tensor_offset(mid), + expert_mid_dim, + pair_rows, + 1, + 0.0f, + 0.0f); + } + if (ok) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + upbuf, + ds4_gpu_tensor_offset(up), + upbuf, + ds4_gpu_tensor_offset(up), + expert_mid_dim, + pair_rows, + 0, + -clamp, + clamp); + } + if (ok) { + ds4_gpu_bin_args mul_args = + ds4_gpu_make_bin_same_rows_args(expert_mid_dim, pair_rows); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_mul_pipeline, + &mul_args, + midbuf, + ds4_gpu_tensor_offset(mid), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid)); + } + } else if (ok) { + ok = ds4_gpu_encode_swiglu_flat(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + (uint32_t)((uint64_t)pair_rows * expert_mid_dim)); + } + if (ok && !use_fused_activation && !use_q4_batch_expert_table) { + ds4_gpu_bin_args weight_args = + ds4_gpu_make_bin_rowwise_scalar_args(expert_mid_dim, pair_rows); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_bin_mul_scalar_pipeline, + &weight_args, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + midbuf, + ds4_gpu_tensor_offset(mid)); + } + DS4_METAL_PROFILE_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : + (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); + if (ok) { + if (use_iq2_batch_selected_addr) { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6( + cb, + g_moe_mul_mv_addr_q2_k_sum6_pipeline, + &down_args, + stream_resources, + stream_resource_count, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + stream_overflow_down); + if (!ok) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to encode down path tokens=%u unique=%u\n", + layer_index, + n_tokens, + stream_unique); + } + } else if (use_q4_batch_expert_table) { + ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, + g_moe_mul_mv_table_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + q4_batch_table_queue_residency); + } else if (direct_down_sum) { + ok = ds4_gpu_encode_mul_mv_id_sum6(cb, + down_sum6_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + nil, + 0, + down_smem, + 2); + } else if (use_mm_id) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + down_mm_pipeline, + &down_mm_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + 8192u); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + down_mv_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + down_nsg, + down_rows_per_group_is_nr0); + } + } + DS4_METAL_PROFILE_MOE_STAGE("down"); + if (ok && + n_expert > 1 && + !direct_down_sum && + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_MOE_STAGE("sum"); + if (!ok) { + fprintf(stderr, + "ds4: Metal routed batch MoE failed before submit layer=%u tokens=%u " + "gate=%s down=%s path=%s\n", + layer_index, + n_tokens, + ds4_gpu_metal_tensor_type_name(gate_type), + ds4_gpu_metal_tensor_type_name(down_type), + moe_path); + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "routed batch MoE")) { + if (use_iq2_batch_selected_addr) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + } + return 0; + } + if (use_iq2_batch_selected_addr) { + if (!owned) { + if (ds4_gpu_end_commands() == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to reopen command batch after execution\n", + layer_index); + return 0; + } + } + } + if (q4_batch_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + } +#undef DS4_METAL_PROFILE_MOE_STAGE + } + + return 1; +} diff --git a/models/deepseek/metal/shaders/attention_indexer.metal b/models/deepseek/metal/shaders/attention_indexer.metal new file mode 100644 index 0000000000..fa254bbfd3 --- /dev/null +++ b/models/deepseek/metal/shaders/attention_indexer.metal @@ -0,0 +1,1515 @@ +kernel void kernel_dsv4_router_weights_batch( + constant float &scale, + device const float *probs, + device const int32_t *selected, + device float *weights, + threadgroup volatile float *scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]]) { + if (tid >= 6) return; + + threadgroup volatile float *sum_scratch = scratch; + threadgroup volatile float *denom_scratch = scratch + 32; + threadgroup volatile float *div_scratch = scratch + 33; + const uint out_index = row * 6u + (uint)tid; + const int32_t expert = selected[out_index]; + const float p = probs[row * 256u + (uint)expert]; + + // Keep this sequence identical to kernel_sum_rows_f32_f32 for width 6. + if (sgitg == 0) { + sum_scratch[tiisg] = 0.0f; + } + float sumf = 0.0f; + sumf += p; + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_scratch[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = sum_scratch[tiisg]; + sumf = simd_sum(sumf); + + if (tid == 0) { + denom_scratch[0] = clamp(sumf, 6.103515625e-5f, INFINITY); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + div_scratch[tid] = p / denom_scratch[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + weights[out_index] = div_scratch[tid] * scale; +} + +// Decode router selection for one token after the existing +// sqrt(softplus(logit)) probability kernel has run. Bias affects only top-k +// selection. Route-weight normalization deliberately stays in the old one-token +// kernel: even tiny denominator-order changes here are amplified by 43 MoE +// layers, so this kernel only replaces the selection work. +kernel void kernel_dsv4_router_finalize_one( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256) return; + + threadgroup float *sel_scores = scratch; + threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256); + const float p = probs[tid]; + sel_scores[tid] = args.has_bias ? p + bias[tid] : p; + idx[tid] = (int32_t)tid; + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (args.hash_mode) { + if (tid == 0) { + const uint token = args.use_token_buffer ? (uint)tokens[0] : args.token; + const uint row = min(token, args.hash_rows - 1u); + device const int32_t *src = hash + row * 6u; + for (uint i = 0; i < 6; i++) { + selected[i] = src[i]; + } + } + } else { + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + const uint other = tid ^ j; + if (other > tid) { + if ((tid & k) == 0) { + if (sel_scores[(uint)idx[tid]] < sel_scores[(uint)idx[other]]) { + const int32_t tmp = idx[tid]; + idx[tid] = idx[other]; + idx[other] = tmp; + } + } else { + if (sel_scores[(uint)idx[tid]] > sel_scores[(uint)idx[other]]) { + const int32_t tmp = idx[tid]; + idx[tid] = idx[other]; + idx[other] = tmp; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + if (tid < 6) { + selected[tid] = idx[tid]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// M3 decode specialization for the non-hash one-token router. Scores and ids +// stay in registers. Intra-SIMD bitonic stages use shuffle-xor; the six stages +// that cross 32-lane SIMD groups exchange through alternating threadgroup +// banks. The next bank's publish barrier proves every prior-bank read finished; +// by the time a bank is reused two cross stages later, no reader can remain. +kernel void kernel_dsv4_router_finalize_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } +} + +// M3 decode specialization that extends the register/TG SIMD selection above +// through the existing six-value serial weight normalization. The selected ids +// cross the same device-memory boundary as the standalone weight kernel; +// volatile TG stores pin its left-fold and scaled-reciprocal rounding points. +kernel void kernel_dsv4_router_finalize_weights_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + device float *weights, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } + threadgroup_barrier(mem_flags::mem_device); + + threadgroup volatile float *norm_scratch = + (threadgroup volatile float *)scratch; + if (tid == 0) { + device const int32_t *s = selected; + norm_scratch[0] = 0.0f; + for (uint i = 0; i < 6; i++) { + norm_scratch[0] = norm_scratch[0] + probs[s[i]]; + } + norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); + norm_scratch[1] = 1.5f / norm_scratch[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 6) { + device const int32_t *s = selected; + weights[tid] = probs[s[tid]] * norm_scratch[1]; + } +} + +// Fills the dense compressed-attention mask with -inf. The selected top-k rows +// are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch. +kernel void kernel_dsv4_topk_mask( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * topk, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ic = gid % args.ne0; + const int64_t it = gid / args.ne0; + + (void)topk; + *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = -INFINITY; +} + +// Enables the selected compressed rows in the dense mask. This replaces the +// old O(n_comp * n_tokens * top_k) membership test with O(top_k * n_tokens) +// writes while preserving exactly the same 0/-inf mask consumed by attention. +kernel void kernel_dsv4_topk_mask_scatter( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * topk, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne00 * args.ne01; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ik = gid % args.ne00; + const int64_t it = gid / args.ne00; + const int32_t idx = *((device const int32_t *) (topk + ik*args.nb00 + it*args.nb01)); + if (idx >= 0 && (int64_t)idx < args.ne0) { + *((device float *) (dst + (int64_t)idx*args.nb0 + it*args.nb1)) = 0.0f; + } +} + +// Sorts each token's selected compressed rows by row id. The indexer selects by +// score, but attention scans compressed K/V in cache order in the dense graph. +// Sorting preserves that order while still letting the indexed attention kernel +// touch only the selected rows. +kernel void kernel_dsv4_sort_i32_rows_asc( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * src, + device char * dst, + threadgroup int32_t * row_tmp [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint n_threads [[threads_per_threadgroup]]) { + const uint top_k = (uint)args.ne00; + if (row >= (uint)args.ne01 || tid >= n_threads) { + return; + } + + for (uint i = tid; i < top_k; i += n_threads) { + row_tmp[i] = *((device const int32_t *) (src + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k = 2; k <= top_k; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + for (uint i = tid; i < top_k; i += n_threads) { + const uint other = i ^ j; + if (other > i && other < top_k) { + const int32_t a = row_tmp[i]; + const int32_t b = row_tmp[other]; + const bool up = (i & k) == 0; + if ((up && a > b) || (!up && a < b)) { + row_tmp[i] = b; + row_tmp[other] = a; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + for (uint i = tid; i < top_k; i += n_threads) { + *((device int32_t *) (dst + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)) = row_tmp[i]; + } +} + +static inline void dsv4_attend_f32_row_as_f16( + device const char *kv, + uint64_t row_stride, + uint row, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + device const float4 *kv4 = (device const float4 *)(kv + (uint64_t)row * row_stride); + const half4 k0 = (half4)kv4[lane + 0]; + const half4 k1 = (half4)kv4[lane + 32]; + const half4 k2 = (half4)kv4[lane + 64]; + const half4 k3 = (half4)kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_f32_row_as_f16( + threadgroup const float4 *kv4, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const half4 k0 = (half4)kv4[lane + 0]; + const half4 k1 = (half4)kv4[lane + 32]; + const half4 k2 = (half4)kv4[lane + 64]; + const half4 k3 = (half4)kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_f32_row_as_f16_at( + threadgroup const float4 *kv4, + uint row_in_tg, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + dsv4_attend_shared_f32_row_as_f16(kv4 + row_in_tg * 128u, + q0, q1, q2, q3, + scale, + lane, + M, S, + o0, o1, o2, o3); +} + +static inline void dsv4_attend_shared_h4_row( + threadgroup const half4 *kv4, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const half4 k0 = kv4[lane + 0]; + const half4 k1 = kv4[lane + 32]; + const half4 k2 = kv4[lane + 64]; + const half4 k3 = kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_h4_row_at( + threadgroup const half4 *kv4, + uint row_in_tg, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + dsv4_attend_shared_h4_row(kv4 + row_in_tg * 128u, + q0, q1, q2, q3, + scale, + lane, + M, S, + o0, o1, o2, o3); +} + +static inline half4 dsv4_load_cache_h4( + device const char *kv, + uint64_t row_stride, + uint row, + uint col, + bool f16_rows) { + device const char *base = kv + (uint64_t)row * row_stride; + if (f16_rows) { + return ((device const half4 *)base)[col]; + } + return (half4)((device const float4 *)base)[col]; +} + +static inline void dsv4_attend_sink( + float score, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + M = new_m; +} + +// DS4 ratio-4 indexed mixed attention. It replaces the dense top-k mask path: +// the threadgroup covers one token and eight heads. Top-k rows and local raw +// rows are the same for all heads of a token, so K/V is staged once in +// threadgroup memory and reused by the eight simdgroups. It keeps the DS4 F16 +// attention rounding by casting Q/K/V to half before the dot/value update. +kernel void kernel_dsv4_indexed_mixed_attention_heads8( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.x; + const uint head = tgpig.y * 8u + (uint)sg; + if (token >= args.n_tokens || head >= args.n_head) { + return; + } + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + uint first = max(first_raw_pos, window_first); + uint last = min(qpos, raw_last_pos); + + if (first <= last) { + for (uint pos = first; pos <= last; pos++) { + const uint logical = pos - first_raw_pos; + const uint row = (args.raw_start + logical) % args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row * args.raw_row_stride); + if (tid < 128) kv_shared[tid] = (half4)src[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint visible = (qpos + 1u) / args.ratio; + visible = min(visible, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token * args.topk_token_stride); + for (uint i = 0; i < args.top_k; i++) { + const int32_t idx = row_topk[i]; + if (idx < 0) { + continue; + } + if ((uint)idx >= visible) { + break; + } + if (tid < 128) { + kv_shared[tid] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + (uint)idx, + tid, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); + + const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; + device float4 *dst4 = (device float4 *)(dst + + (uint64_t)token * args.dst_token_stride + + (uint64_t)head * args.dst_head_stride); + dst4[lane + 0] = o0 * inv_s; + dst4[lane + 32] = o1 * inv_s; + dst4[lane + 64] = o2 * inv_s; + dst4[lane + 96] = o3 * inv_s; +} + +// Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8. +// Generation attends one token at a time, so the ratio-4 indexed path spends a +// visible amount of time repeatedly staging the same K/V row for the eight +// heads in a group. This variant stages sixteen selected rows at once and then +// consumes them sequentially, preserving the row order and online softmax math +// while cutting threadgroup barriers in the long top-k scan. +kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.x; + const uint head = tgpig.y * 8u + (uint)sg; + if (token >= args.n_tokens || head >= args.n_head) { + return; + } + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + uint first = max(first_raw_pos, window_first); + uint last = min(qpos, raw_last_pos); + + if (first <= last) { + for (uint pos0 = first; pos0 <= last; pos0 += 16u) { + const uint n_rows = min(16u, last - pos0 + 1u); + for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { + const uint r = off >> 7; + const uint c = off & 127u; + const uint logical = pos0 + r - first_raw_pos; + const uint row = (args.raw_start + logical) % args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row * args.raw_row_stride); + kv_shared[off] = (half4)src[c]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, + r, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint visible = (qpos + 1u) / args.ratio; + visible = min(visible, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token * args.topk_token_stride); + bool stop = false; + for (uint i = 0; i < args.top_k && !stop; i += 16u) { + uint rows[16]; + uint n_rows = 0; + for (uint j = 0; j < 16u && i + j < args.top_k; j++) { + const int32_t idx = row_topk[i + j]; + if (idx < 0) { + continue; + } + if ((uint)idx >= visible) { + stop = true; + break; + } + rows[n_rows++] = (uint)idx; + } + if (n_rows == 0) { + continue; + } + for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { + const uint r = off >> 7; + const uint c = off & 127u; + kv_shared[off] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + rows[r], + c, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, + r, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); + + const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; + device float4 *dst4 = (device float4 *)(dst + + (uint64_t)token * args.dst_token_stride + + (uint64_t)head * args.dst_head_stride); + dst4[lane + 0] = o0 * inv_s; + dst4[lane + 32] = o1 * inv_s; + dst4[lane + 64] = o2 * inv_s; + dst4[lane + 96] = o3 * inv_s; +} + +static inline float dsv4_indexer_dot128_shared_q( + float4 c0, + float4 c1, + float4 c2, + float4 c3, + threadgroup const float4 *q4, + ushort lane) { + float sum = 0.0f; + if (lane < 8) { + const ushort ib = lane >> 1; + const ushort il = lane & 1; + const ushort base = ib*8 + il*4; + sum += dot(c0, q4[base + 0]); + sum += dot(c1, q4[base + 1]); + sum += dot(c2, q4[base + 2]); + sum += dot(c3, q4[base + 3]); + } + return simd_sum(sum); +} + +// Tiled prefill score builder for the sparse-compressed attention indexer. +// +// The kernel covers an 8-token by 32-compressed-row rectangle: K is copied into +// threadgroup memory once, then reused for all 64 indexer heads, while simdgroup +// matrix multiply computes each 8x8 score subtile. +// +// It still writes the exact score matrix consumed by top-k: +// +// score[t,c] = sum_h relu(dot(Q[t,h], K[c])) * W[t,h] * scale +// +// Causal masking is applied on store so invisible compressed rows become -inf. +kernel void kernel_dsv4_indexer_scores_tiled_f32( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + threadgroup float *qtg = shared; // [8][128] + threadgroup float *ktg = qtg + TM*D; // [32][128] + threadgroup float *dot = ktg + TN*D; // [8][32] + + const uint last_token = min(t0 + TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + for (uint i = tid; i < TN*D; i += 128) { + const uint cc = i / D; + const uint d = i - cc*D; + const uint comp = c0 + cc; + float v = 0.0f; + if (comp < args.n_comp) { + device const float *row = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + v = row[d]; + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint row0 = cell0 >> 3; + const uint row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = t0 + row0; + const uint token1 = t0 + row1; + const uint comp0 = c0 + col0; + const uint comp1 = c0 + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + for (uint i = tid; i < TM*D; i += 128) { + const uint r = i / D; + const uint d = i - r*D; + const uint token = t0 + r; + float v = 0.0f; + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = qrow[d]; + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_float8x8 mq; + simdgroup_float8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[row0*TN + col0]; + acc0 += max(s, 0.0f) * (w[head] * args.scale); + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[row1*TN + col1]; + acc1 += max(s, 0.0f) * (w[head] * args.scale); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + comp0; + *dst = comp0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + comp1; + *dst = comp1 < visible ? acc1 : -INFINITY; + } +} + +kernel void kernel_dsv4_indexer_scores_tiled( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + // Q/K are staged as half but the dot accumulator and final score remain + // float. This is the one intentional precision tradeoff in the indexer: + // the indexer only ranks compressed rows for top-k selection, and long + // context profiling shows this score matrix dominates the prefill slope. + threadgroup half *qtg = (threadgroup half *)shared; // [8][128] + threadgroup half *ktg = qtg + TM*D; // [32][128] + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [8][32] + + const uint last_token = min(t0 + TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + // Stage compressed index rows once. Edge columns are zeroed so the matrix + // loads below can stay regular; guarded stores discard them. + for (uint i = tid; i < TN*D; i += 128) { + const uint cc = i / D; + const uint d = i - cc*D; + const uint comp = c0 + cc; + half v = half(0.0f); + if (comp < args.n_comp) { + device const float *row = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + v = half(row[d]); + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint row0 = cell0 >> 3; + const uint row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = t0 + row0; + const uint token1 = t0 + row1; + const uint comp0 = c0 + col0; + const uint comp1 = c0 + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + // Stage Q for the eight-token tile. Each 8x8 matrix load below reads a + // contiguous depth block from this layout. + for (uint i = tid; i < TM*D; i += 128) { + const uint r = i / D; + const uint d = i - r*D; + const uint token = t0 + r; + half v = half(0.0f); + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = half(qrow[d]); + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_half8x8 mq; + simdgroup_half8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[row0*TN + col0]; + acc0 += max(s, 0.0f) * (w[head] * args.scale); + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[row1*TN + col1]; + acc1 += max(s, 0.0f) * (w[head] * args.scale); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + comp0; + *dst = comp0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + comp1; + *dst = comp1 < visible ? acc1 : -INFINITY; + } +} + +#ifdef DS4_METAL_HAS_TENSOR +// Retained full-512 prefill indexer score path. This is the part of sparse +// compressed attention that maps cleanly to TensorOps: a regular token by +// compressed-row dot tile. The kernel intentionally leaves top-k selection and +// indexed attention semantics unchanged; all 512 selected rows remain available +// to the later attention kernel. +// +// Each matmul processes a pair of heads (TQ = 2 x TM q rows): the per-element +// dot is still a 128-deep reduction in 32-wide k-steps, so scores are +// bit-identical to single-head tiles while the run count halves. The q tile +// is double-buffered, so the next k-step's stage overlaps the current +// cooperative matmul and each pair needs 5 barriers instead of 10. q and k +// staging use one float4/half4 per lane (each thread covers one row of 8/32 +// consecutive elements), which is the same half(float) conversion per element +// as the scalar form. +kernel void kernel_dsv4_indexer_scores_nax( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup half *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + constexpr int TQ = 32; + constexpr int TN = 32; + constexpr int NK = 32; + constexpr int D = 128; + constexpr int NUM_THREADS = 128; + + // The 16-token x 32-row tile was the winning NAX shape in local sweeps. A + // wider 64-row compressed tile increased setup/cache pressure and was + // slower despite doing more work per dispatch. + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + threadgroup half *qtg = shared; // 2 x [TQ][NK] + threadgroup half *ktg = qtg + 2*TQ*NK; // [32][128] + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [TQ][TN], column-major + + const uint last_token = min(t0 + (uint)TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += NUM_THREADS) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + { + // One compressed row per 4 threads, 32 consecutive floats per thread. + const uint cc = tid / 4; + const uint comp = c0 + cc; + device const float *krow = nullptr; + if (comp < args.n_comp) { + krow = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + } + const uint d0 = (tid % 4) * 32; + FOR_UNROLL (uint j = 0; j < 8; j++) { + const float4 kv = krow ? *(device const float4 *)(krow + d0 + 4*j) + : float4(0.0f); + *(threadgroup half4 *)(ktg + cc*D + d0 + 4*j) = half4(kv); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc[4]; + #pragma unroll + for (uint j = 0; j < 4; j++) { + acc[j] = 0.0f; + } + + auto tq0 = tensor(qtg, dextents(NK, TQ)); + auto tq1 = tensor(qtg + TQ*NK, dextents(NK, TQ)); + auto tk = tensor(ktg, dextents(D, TN)); + auto td = tensor(dot, dextents(TQ, TN), array({1, TQ})); + + matmul2d< + matmul2d_descriptor(TN, TQ, NK, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + // One q row per 4 threads, 8 consecutive floats per thread. Row r covers + // head (r / TM) of the pair and token row (r % TM). + const uint q_r = tid / 4; + const uint q_k4 = (tid % 4) * 8; + const uint q_hl = q_r / TM; + const uint q_tr = q_r % TM; + const uint q_token = t0 + q_tr; + device const char *q_row_base = nullptr; + if (q_token < args.n_tokens) { + q_row_base = q + (uint64_t)q_token * args.q_token_stride; + } + + auto stage_q = [&](const uint head0, const uint loop_k, threadgroup half *buf) { + const uint head = head0 + q_hl; + half4 v0 = half4(0.0f); + half4 v1 = half4(0.0f); + if (q_row_base && head < args.n_head) { + device const float4 *src4 = (device const float4 *) + (q_row_base + (uint64_t)head * args.q_head_stride + + (uint64_t)(loop_k + q_k4) * sizeof(float)); + v0 = half4(src4[0]); + v1 = half4(src4[1]); + } + *(threadgroup half4 *)(buf + q_r*NK + q_k4) = v0; + *(threadgroup half4 *)(buf + q_r*NK + q_k4 + 4) = v1; + }; + + for (uint head0 = 0; head0 < args.n_head; head0 += 2) { + auto ct = mm.template get_destination_cooperative_tensor(); + #pragma unroll + for (uint16_t i = 0; i < ct.get_capacity(); i++) { + if (ct.is_valid_element(i)) { + ct[i] = 0.0f; + } + } + + stage_q(head0, 0, qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint qsel = 0; + FOR_UNROLL (uint i = 0; i < 4; i++) { + auto mk = tk.slice(i*NK, 0); + auto mq = (qsel ? tq1 : tq0).slice(0, 0); + mm.run(mk, mq, ct); + if (i < 3) { + qsel ^= 1u; + stage_q(head0, (i + 1)*NK, qsel ? qtg + TQ*NK : qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + ct.store(td); + threadgroup_barrier(mem_flags::mem_threadgroup); + + #pragma unroll + for (uint j = 0; j < 4; j++) { + const uint linear = (uint)tid + j*NUM_THREADS; + if (linear < TM*TN) { + const uint r = linear / TN; + const uint cc = linear - r*TN; + const uint token = t0 + r; + if (token < args.n_tokens) { + device const float *w = (device const float *)(weights + + (uint64_t)token * args.weights_token_stride); + acc[j] += max(dot[cc*TQ + r], 0.0f) * (w[head0] * args.scale); + if (head0 + 1 < args.n_head) { + acc[j] += max(dot[cc*TQ + TM + r], 0.0f) * (w[head0 + 1] * args.scale); + } + } + } + } + // No barrier here: the next pair's q stage and these dot reads touch + // different buffers, and the next q-stage barrier separates the next + // ct.store from these reads. + } + + #pragma unroll + for (uint j = 0; j < 4; j++) { + const uint linear = (uint)tid + j*NUM_THREADS; + if (linear >= TM*TN) { + continue; + } + const uint r = linear / TN; + const uint cc = linear - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + const uint visible = min((args.pos0 + token + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = comp < visible ? acc[j] : -INFINITY; + } + } +} +#endif + +// Collapses per-head indexer scores into one score per compressed row using the +// learned head weights. Negative head scores are clipped exactly as DS4 expects. +kernel void kernel_dsv4_indexer_weighted_sum( + constant ds4_metal_args_dsv4_indexer_weighted_sum & args, + device const char * scores, + device const char * weights, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ic = gid % args.ne0; + const int64_t it = gid / args.ne0; + + float acc = 0.0f; + for (int64_t ih = 0; ih < args.ne02; ++ih) { + const float s = *((device const float *) (scores + ic*args.nb00 + it*args.nb01 + ih*args.nb02)); + const float w = *((device const float *) (weights + ih*args.nb10 + it*args.nb11)); + acc += max(s, 0.0f) * (w * args.scale); + } + + *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = acc; +} + +// Adds the periodic compressor APE directly to projected scores. The legacy +// path materializes one repeated APE segment per period and then performs this +// same single F32 add; these kernels remove only that intermediate copy graph. +kernel void kernel_dsv4_compressor_score_ape_f32( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const float *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + ape[(uint64_t)ape_row*args.width + col]; +} + +kernel void kernel_dsv4_compressor_score_ape_f16( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const half *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + float(ape[(uint64_t)ape_row*args.width + col]); +} + +// Fused softmax-weighted pooling of compressed KV rows. It is used when several +// compressor rows are present; the one-row case deliberately follows the +// unfused softmax/mul/sum graph in Objective-C to keep identical reductions. +kernel void kernel_dsv4_softmax_pool( + constant ds4_metal_args_dsv4_softmax_pool & args, + device const char * kv, + device const char * score, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t id = gid % args.ne0; + const int64_t ic = gid / args.ne0; + + float max_s = -INFINITY; + for (int64_t ir = 0; ir < args.ne00; ++ir) { + const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); + max_s = max(max_s, s); + } + + float sum = 0.0f; + float acc = 0.0f; + for (int64_t ir = 0; ir < args.ne00; ++ir) { + const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); + const float w = exp(s - max_s); + const float v = *((device const float *) (kv + ir*args.nb00 + id*args.nb01 + ic*args.nb02)); + sum += w; + acc += v*w; + } + + *((device float *) (dst + id*args.nb0 + ic*args.nb1)) = acc/sum; +} + + + +// Tensor-parallel keep-alive: a few threadgroups of FMAs dispatched +// back-to-back on a side queue while TP decode runs. The per-layer gate +// stalls make the real workload look idle to the GPU power manager, which +// otherwise halves the clocks within a second (~2x decode regression); +// this holds them up for negligible bandwidth and a few watts. +kernel void kernel_dsv4_tp_keepalive( + device float * out, + constant uint & iters, + uint tid [[thread_position_in_grid]]) { + float a = out[tid]; + const float b = 1.000001f; + for (uint i = 0; i < iters; i++) { + a = fma(a, b, 0.000001f); + a = fma(a, b, -0.000001f); + } + out[tid] = a; +} + +// Tensor-parallel gate flag: publishes a sequence number to a slab slot the +// CPU service thread spin-reads, replacing the much slower shared-event +// signal for the GPU->CPU direction. Ordering against the partial-output +// kernels comes from the buffer hazard on the shared slab. +kernel void kernel_dsv4_tp_flag_set( + device atomic_uint & flag, + constant uint & value, + uint tid [[thread_position_in_grid]]) { + if (tid == 0) { + atomic_store_explicit(&flag, value, memory_order_relaxed); + } +} + +// Ratio-4 compressor pooling without materializing the [n_comp, 8, head_dim] +// KV and score packs. The row mapping and both reduction loops deliberately +// match kernel_dsv4_softmax_pool so the arithmetic order is unchanged. +kernel void kernel_dsv4_softmax_pool_ratio4_direct( + constant ds4_metal_args_dsv4_softmax_pool_ratio4_direct & args, + device const float * kv, + device const float * score, + device const float * state_kv, + device const float * state_score, + device float * dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t n = (uint64_t)args.head_dim * args.n_comp; + if ((uint64_t)gid >= n || args.head_dim == 0u) { + return; + } + + const uint64_t id = gid % args.head_dim; + const uint64_t ic = gid / args.head_dim; + const uint64_t input_row_stride = 2ull * args.head_dim; + + float max_s = -INFINITY; + float sum = 0.0f; + float acc = 0.0f; + if (ic != 0u) { + const int64_t token_base = (int64_t)ic * 4 - 4; + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + const float w = exp(s - max_s); + const float v = kv[src]; + sum += w; + acc += v*w; + } + } else { + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + } else if (args.replay != 0u) { + s = state_score[(uint64_t)ir * input_row_stride + id]; + } else { + s = -INFINITY; + } + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + float v; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + v = kv[src]; + } else if (args.replay != 0u) { + const uint64_t src = (uint64_t)ir * input_row_stride + id; + s = state_score[src]; + v = state_kv[src]; + } else { + s = -INFINITY; + v = 0.0f; + } + const float w = exp(s - max_s); + sum += w; + acc += v*w; + } + } + + dst[ic * args.head_dim + id] = acc/sum; +} diff --git a/models/deepseek/metal/shaders/control.metal b/models/deepseek/metal/shaders/control.metal new file mode 100644 index 0000000000..4d8337391c --- /dev/null +++ b/models/deepseek/metal/shaders/control.metal @@ -0,0 +1,114 @@ +kernel void kernel_dsv4_directional_steering_project_f32( + constant ds4_metal_args_dsv4_directional_steering_project & args, + device float *x, + device const float *directions, + threadgroup float *scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + if (row >= args.rows || args.width == 0) return; + + device float *xr = x + (uint64_t)row * args.width; + device const float *dir = directions + (uint64_t)args.layer * args.width; + const uint nth = args.n_threads; + + float sum = 0.0f; + for (uint i = tid; i < args.width; i += nth) { + sum += xr[i] * dir[i]; + } + scratch[tid] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float coeff = args.scale * scratch[0]; + for (uint i = tid; i < args.width; i += nth) { + xr[i] -= coeff * dir[i]; + } +} + +// Decode-only DS4 ratio-4 indexer score builder. One threadgroup owns one +// compressed row for the current token, stages that 128-wide row once, then +// walks the 64 indexer heads in four-head groups. This avoids materializing the +// intermediate [compressed rows x heads] score matrix used by the generic +// matvec + weighted-sum path. +kernel void kernel_dsv4_indexer_score_one_direct( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + if (row >= args.n_comp || args.n_head != 64u || args.head_dim != 128u) { + return; + } + + threadgroup float *ktg = shared; // [128] + threadgroup float *psum = ktg + 128u; // [4] + + if (tid < 128u) { + device const float *krow = (device const float *)(index_comp + + (uint64_t)row * args.index_row_stride); + ktg[tid] = krow[tid]; + } + + float acc = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head0 = 0; head0 < 64u; head0 += 4u) { + const uint head = head0 + (uint)sg; + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)head * args.q_head_stride); + threadgroup const float4 *k4 = (threadgroup const float4 *)ktg; + + float s = dot(q4[lane], k4[lane]); + s = simd_sum(s); + if (lane == 0) { + device const float *w = (device const float *)weights; + psum[sg] = max(s, 0.0f) * (w[head] * args.scale); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + acc += psum[0]; + acc += psum[1]; + acc += psum[2]; + acc += psum[3]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + device float *dst = (device float *)scores; + dst[row] = acc; + } +} + +// Decode router post-processing for one token. The selected expert ids are +// already known; this gathers their probabilities, normalizes by the selected +// sum, clamps the denominator like the reference path, and applies DS4's 1.5 +// expert-weight scale in one tiny dispatch. +kernel void kernel_dsv4_router_weights_one( + device const char *probs, + device const char *selected, + device char *weights, + uint tid [[thread_position_in_grid]]) { + if (tid >= 6) return; + + device const float *p = (device const float *)probs; + device const int *s = (device const int *)selected; + + float sum = 0.0f; + for (uint i = 0; i < 6; i++) { + sum += p[s[i]]; + } + sum = max(sum, 6.103515625e-5f); + + device float *w = (device float *)weights; + w[tid] = p[s[tid]] / sum * 1.5f; +} diff --git a/metal/dsv4_hc.metal b/models/deepseek/metal/shaders/hc.metal similarity index 100% rename from metal/dsv4_hc.metal rename to models/deepseek/metal/shaders/hc.metal diff --git a/metal/dsv4_kv.metal b/models/deepseek/metal/shaders/kv.metal similarity index 100% rename from metal/dsv4_kv.metal rename to models/deepseek/metal/shaders/kv.metal diff --git a/metal/dsv4_rope.metal b/models/deepseek/metal/shaders/rope.metal similarity index 100% rename from metal/dsv4_rope.metal rename to models/deepseek/metal/shaders/rope.metal diff --git a/models/deepseek/provider.c b/models/deepseek/provider.c new file mode 100644 index 0000000000..cf91574007 --- /dev/null +++ b/models/deepseek/provider.c @@ -0,0 +1,35 @@ +#include "provider.h" + +#include "../../ds4_model_provider_builtin.h" + +static const ds4_model_provider_v1 DS4_DEEPSEEK_PROVIDER = { + .abi_version = DS4_MODEL_PROVIDER_ABI_VERSION, + .struct_size = sizeof(ds4_model_provider_v1), + .id = "deepseek-v4", + .session_create = ds4_deepseek_session_create, + .session_destroy = ds4_deepseek_session_destroy, + .session_sync = ds4_deepseek_session_sync, + .session_eval = ds4_deepseek_session_eval, + .sessions_eval_batch = ds4_builtin_sessions_eval_batch, + .sessions_eval_batch_with_prefill = + ds4_builtin_sessions_eval_batch_with_prefill, + .session_eval_speculative = ds4_deepseek_session_eval_speculative, + .session_invalidate = ds4_deepseek_session_invalidate, + .session_rewind = ds4_deepseek_session_rewind, + .session_layer_slice_reset = ds4_deepseek_session_layer_slice_reset, + .session_eval_output_head = ds4_deepseek_session_eval_output_head, + .session_eval_layer_slice = ds4_deepseek_session_eval_layer_slice, + .session_payload_bytes = ds4_deepseek_session_payload_bytes, + .session_save_payload = ds4_deepseek_session_save_payload, + .session_load_payload = ds4_deepseek_session_load_payload, + .session_layer_payload_bytes = + ds4_deepseek_session_layer_payload_bytes, + .session_save_layer_payload = + ds4_deepseek_session_save_layer_payload, + .session_load_layer_payload = + ds4_deepseek_session_load_layer_payload, +}; + +const ds4_model_provider_v1 *ds4_deepseek_model_provider(void) { + return &DS4_DEEPSEEK_PROVIDER; +} diff --git a/models/deepseek/provider.h b/models/deepseek/provider.h new file mode 100644 index 0000000000..f8739d425a --- /dev/null +++ b/models/deepseek/provider.h @@ -0,0 +1,87 @@ +#ifndef DS4_DEEPSEEK_MODEL_PROVIDER_H +#define DS4_DEEPSEEK_MODEL_PROVIDER_H + +#include "../../ds4_model_provider.h" + +const ds4_model_provider_v1 *ds4_deepseek_model_provider(void); + +int ds4_deepseek_session_create(ds4_session **out, + ds4_engine *engine, + int context_size); +void ds4_deepseek_session_destroy(ds4_session *session); +int ds4_deepseek_session_sync(ds4_session *session, + const ds4_tokens *prompt, + char *err, + size_t errlen); +int ds4_deepseek_session_eval(ds4_session *session, + int token, + bool probe_support_model, + char *err, + size_t errlen); +int ds4_deepseek_session_eval_speculative( + ds4_session *session, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen); +void ds4_deepseek_session_invalidate(ds4_session *session); +void ds4_deepseek_session_rewind(ds4_session *session, int position); +int ds4_deepseek_session_layer_slice_reset(ds4_session *session, + char *err, + size_t errlen); +int ds4_deepseek_session_eval_output_head( + ds4_session *session, + const float *hidden_state, + uint32_t token_count, + float *logits, + char *err, + size_t errlen); +int ds4_deepseek_session_eval_layer_slice( + ds4_session *session, + const int *tokens, + uint32_t token_count, + uint32_t position, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hidden_state, + float *output_hidden_state, + bool output_logits, + float *logits, + char *err, + size_t errlen); +uint64_t ds4_deepseek_session_payload_bytes(ds4_session *session); +int ds4_deepseek_session_save_payload(ds4_session *session, + FILE *file, + char *err, + size_t errlen); +int ds4_deepseek_session_load_payload(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + char *err, + size_t errlen); +uint64_t ds4_deepseek_session_layer_payload_bytes( + ds4_session *session, + uint32_t layer_start, + uint32_t layer_end); +int ds4_deepseek_session_save_layer_payload( + ds4_session *session, + FILE *file, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); +int ds4_deepseek_session_load_layer_payload( + ds4_session *session, + FILE *file, + uint64_t payload_bytes, + const int *tokens, + uint32_t token_count, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); + +#endif diff --git a/rocm/ds4_rocm_attention.cuh b/models/deepseek/rocm/attention.cuh similarity index 100% rename from rocm/ds4_rocm_attention.cuh rename to models/deepseek/rocm/attention.cuh diff --git a/rocm/ds4_rocm_attention_launch.cuh b/models/deepseek/rocm/attention_launch.cuh similarity index 100% rename from rocm/ds4_rocm_attention_launch.cuh rename to models/deepseek/rocm/attention_launch.cuh diff --git a/rocm/ds4_rocm_compressor.cuh b/models/deepseek/rocm/compressor.cuh similarity index 100% rename from rocm/ds4_rocm_compressor.cuh rename to models/deepseek/rocm/compressor.cuh diff --git a/rocm/ds4_rocm_fp8_kv.cuh b/models/deepseek/rocm/fp8_kv.cuh similarity index 100% rename from rocm/ds4_rocm_fp8_kv.cuh rename to models/deepseek/rocm/fp8_kv.cuh diff --git a/rocm/ds4_rocm_fp8_kv_launch.cuh b/models/deepseek/rocm/fp8_kv_launch.cuh similarity index 100% rename from rocm/ds4_rocm_fp8_kv_launch.cuh rename to models/deepseek/rocm/fp8_kv_launch.cuh diff --git a/rocm/ds4_rocm_hc.cuh b/models/deepseek/rocm/hc.cuh similarity index 100% rename from rocm/ds4_rocm_hc.cuh rename to models/deepseek/rocm/hc.cuh diff --git a/rocm/ds4_rocm_hc_output_launch.cuh b/models/deepseek/rocm/hc_output_launch.cuh similarity index 100% rename from rocm/ds4_rocm_hc_output_launch.cuh rename to models/deepseek/rocm/hc_output_launch.cuh diff --git a/rocm/ds4_rocm_indexer.cuh b/models/deepseek/rocm/indexer.cuh similarity index 100% rename from rocm/ds4_rocm_indexer.cuh rename to models/deepseek/rocm/indexer.cuh diff --git a/rocm/ds4_rocm_output.cuh b/models/deepseek/rocm/output.cuh similarity index 100% rename from rocm/ds4_rocm_output.cuh rename to models/deepseek/rocm/output.cuh diff --git a/rocm/ds4_rocm_router.cuh b/models/deepseek/rocm/router.cuh similarity index 100% rename from rocm/ds4_rocm_router.cuh rename to models/deepseek/rocm/router.cuh diff --git a/models/glm/README.md b/models/glm/README.md new file mode 100644 index 0000000000..2df5794f6e --- /dev/null +++ b/models/glm/README.md @@ -0,0 +1,14 @@ +# GLM DSA integration + +This directory owns the GLM DSA model provider and its tailored inference +implementation: + +- `provider.c` exposes the whole-model lifecycle to the engine core. +- `cpu.inc` contains CPU reference kernels used by correctness diagnostics. +- `graph.inc` owns GPU graph state, allocation, prefill, decode, MTP, and + checkpoint orchestration. +- `cuda/`, `metal/`, and `rocm/` contain GLM-specific host and device + implementations. + +The provider calls these concrete paths directly. They do not implement a +generic kernel interface. diff --git a/models/glm/cpu.inc b/models/glm/cpu.inc new file mode 100644 index 0000000000..36751a633a --- /dev/null +++ b/models/glm/cpu.inc @@ -0,0 +1,816 @@ +/* + * GLM CPU reference and diagnostic inference kernels. + * + * Included exactly once by ds4.c. The production GLM path remains graph-only; + * these routines provide exact reference behavior for validation. + */ + +static void layer_glm_first_token_attention_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); + float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); + float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); + const uint64_t kv_blocks = (DS4_N_KV_LORA + 31) / 32; + int8_t *kvq = xmalloc((size_t)kv_blocks * 32); + float *kvscale = xmalloc((size_t)kv_blocks * sizeof(kvscale[0])); + + if (layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || + layer->attn_v_b->dim[0] != DS4_N_KV_LORA || + layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || + layer->attn_v_b->dim[2] != DS4_N_HEAD || + layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || + layer->attn_output->dim[1] != DS4_N_EMBD) { + ds4_die("GLM attention tensors have an unexpected layout"); + } + + rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0(kv_raw, model, layer->attn_kv_a_mqa, norm); + rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), + DS4_N_KV_LORA, DS4_RMS_EPS); + quantize_q8_0_activation(kv_norm, kvq, kvscale, DS4_N_KV_LORA); + + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + matvec_q8_0_3d_slice_prequant(heads + (uint64_t)h * DS4_N_VALUE_MLA, + model, + layer->attn_v_b, + kvq, + kvscale, + h); + } + matvec_q8_0(out, model, layer->attn_output, heads); + + free(kvscale); + free(kvq); + free(heads); + free(kv_norm); + free(kv_raw); + free(norm); +} + +static void layer_glm_first_token_attention_one_f32_ref( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); + float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); + float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); + + if (layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || + layer->attn_v_b->type != DS4_TENSOR_Q8_0 || + layer->attn_output->type != DS4_TENSOR_Q8_0 || + layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || + layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || + layer->attn_v_b->dim[0] != DS4_N_KV_LORA || + layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || + layer->attn_v_b->dim[2] != DS4_N_HEAD || + layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || + layer->attn_output->dim[1] != DS4_N_EMBD) { + ds4_die("GLM F32 attention reference found unexpected tensor layout"); + } + + rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0_f32_ref(kv_raw, model, layer->attn_kv_a_mqa, norm); + rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), + DS4_N_KV_LORA, DS4_RMS_EPS); + matvec_q8_0_f32_ref(heads, model, layer->attn_v_b, kv_norm); + matvec_q8_0_f32_ref(out, model, layer->attn_output, heads); + + free(heads); + free(kv_norm); + free(kv_raw); + free(norm); +} + +static void glm_k_b_project_f32_ref( + float * out, + const ds4_model * model, + const ds4_tensor * w, + const float * kv_norm) { + const uint32_t q_nope = DS4_N_KEY_MLA - DS4_N_ROT; + if (w->type != DS4_TENSOR_Q8_0 || + w->ndim != 3 || + w->dim[0] != q_nope || + w->dim[1] != DS4_N_KV_LORA || + w->dim[2] != DS4_N_HEAD) { + ds4_die("GLM k_b reference found unexpected tensor layout"); + } + + const uint8_t *data = tensor_data(model, w); + const uint64_t blocks = (q_nope + 31u) / 32u; + const uint64_t row_bytes = blocks * 34u; + memset(out, 0, (size_t)DS4_N_HEAD * q_nope * sizeof(out[0])); + + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + float *dst = out + (uint64_t)h * q_nope; + for (uint32_t j = 0; j < DS4_N_KV_LORA; j++) { + const uint8_t *row = + data + ((uint64_t)h * DS4_N_KV_LORA + j) * row_bytes; + const float xj = kv_norm[j]; + for (uint64_t b = 0; b < blocks; b++) { + uint16_t scale_bits; + memcpy(&scale_bits, row + b * 34u, sizeof(scale_bits)); + const int8_t *qs = (const int8_t *)(row + b * 34u + 2u); + const float d = f16_to_f32(scale_bits) * xj; + const uint32_t i0 = (uint32_t)b * 32u; + const uint32_t n = q_nope - i0 < 32u ? q_nope - i0 : 32u; + for (uint32_t i = 0; i < n; i++) { + dst[i0 + i] += d * (float)qs[i]; + } + } + } + } +} + +static void layer_glm_attention_prefill_f32_ref( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t n_tok, + uint32_t pos0, + uint32_t il) { + if (n_tok == 0) return; + const uint32_t qk_dim = DS4_N_KEY_MLA; + const uint32_t q_nope = qk_dim - DS4_N_ROT; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * qk_dim; + const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; + const uint64_t kv_raw_dim = layer->attn_kv_a_mqa ? layer->attn_kv_a_mqa->dim[1] : 0; + + if (!layer->attn_norm || + !layer->attn_q_a || + !layer->attn_q_a_norm || + !layer->attn_q_b || + !layer->attn_kv_a_mqa || + !layer->attn_kv_a_norm || + !layer->attn_k_b || + !layer->attn_v_b || + !layer->attn_output || + layer->attn_q_a->type != DS4_TENSOR_Q8_0 || + layer->attn_q_b->type != DS4_TENSOR_Q8_0 || + layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || + layer->attn_k_b->type != DS4_TENSOR_Q8_0 || + layer->attn_v_b->type != DS4_TENSOR_Q8_0 || + layer->attn_output->type != DS4_TENSOR_Q8_0 || + layer->attn_norm->type != DS4_TENSOR_F32 || + layer->attn_q_a_norm->type != DS4_TENSOR_F32 || + layer->attn_kv_a_norm->type != DS4_TENSOR_F32 || + layer->attn_q_a->dim[0] != DS4_N_EMBD || + layer->attn_q_a->dim[1] != DS4_N_LORA_Q || + layer->attn_q_a_norm->dim[0] != DS4_N_LORA_Q || + layer->attn_q_b->dim[0] != DS4_N_LORA_Q || + layer->attn_q_b->dim[1] != q_dim || + layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || + kv_raw_dim < (uint64_t)DS4_N_KV_LORA + DS4_N_ROT || + layer->attn_kv_a_norm->dim[0] != DS4_N_KV_LORA || + layer->attn_k_b->dim[0] != q_nope || + layer->attn_k_b->dim[1] != DS4_N_KV_LORA || + layer->attn_k_b->dim[2] != DS4_N_HEAD || + layer->attn_v_b->dim[0] != DS4_N_KV_LORA || + layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || + layer->attn_v_b->dim[2] != DS4_N_HEAD || + layer->attn_output->dim[0] != heads_dim || + layer->attn_output->dim[1] != DS4_N_EMBD) { + ds4_die("GLM prefill attention reference found unexpected tensor layout"); + } + + float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); + float *q_rank = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank[0])); + float *q_rank_norm = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank_norm[0])); + float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); + float *kv_raw = xmalloc((size_t)n_tok * kv_raw_dim * sizeof(kv_raw[0])); + float *kv_norm = xmalloc((size_t)n_tok * DS4_N_KV_LORA * sizeof(kv_norm[0])); + float *k_nope = xmalloc((size_t)n_tok * DS4_N_HEAD * q_nope * sizeof(k_nope[0])); + float *key_cache = xmalloc((size_t)n_tok * DS4_N_HEAD * qk_dim * sizeof(key_cache[0])); + float *value_cache = xmalloc((size_t)n_tok * heads_dim * sizeof(value_cache[0])); + float *heads = xmalloc((size_t)n_tok * heads_dim * sizeof(heads[0])); + float *k_rot = xmalloc((size_t)DS4_N_ROT * sizeof(k_rot[0])); + float *scores = xmalloc((size_t)n_tok * sizeof(scores[0])); + + for (uint32_t t = 0; t < n_tok; t++) { + const float *xt = x + (uint64_t)t * DS4_N_EMBD; + float *norm_t = norm + (uint64_t)t * DS4_N_EMBD; + float *qr_t = q_rank + (uint64_t)t * DS4_N_LORA_Q; + float *qrn_t = q_rank_norm + (uint64_t)t * DS4_N_LORA_Q; + float *q_t = q + (uint64_t)t * q_dim; + float *raw_t = kv_raw + (uint64_t)t * kv_raw_dim; + float *kvn_t = kv_norm + (uint64_t)t * DS4_N_KV_LORA; + float *kn_t = k_nope + (uint64_t)t * DS4_N_HEAD * q_nope; + float *kc_t = key_cache + (uint64_t)t * DS4_N_HEAD * qk_dim; + float *vc_t = value_cache + (uint64_t)t * heads_dim; + + rms_norm_weight(norm_t, xt, tensor_data(model, layer->attn_norm), + DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0_f32_ref(qr_t, model, layer->attn_q_a, norm_t); + rms_norm_weight(qrn_t, qr_t, tensor_data(model, layer->attn_q_a_norm), + DS4_N_LORA_Q, DS4_RMS_EPS); + matvec_q8_0_f32_ref(q_t, model, layer->attn_q_b, qrn_t); + rope_tail_layer_inplace(q_t, DS4_N_HEAD, qk_dim, DS4_N_ROT, + pos0 + t, il, false); + + matvec_q8_0_f32_ref(raw_t, model, layer->attn_kv_a_mqa, norm_t); + rms_norm_weight(kvn_t, raw_t, tensor_data(model, layer->attn_kv_a_norm), + DS4_N_KV_LORA, DS4_RMS_EPS); + glm_k_b_project_f32_ref(kn_t, model, layer->attn_k_b, kvn_t); + matvec_q8_0_f32_ref(vc_t, model, layer->attn_v_b, kvn_t); + + memcpy(k_rot, raw_t + DS4_N_KV_LORA, (size_t)DS4_N_ROT * sizeof(k_rot[0])); + rope_tail_layer_inplace(k_rot, 1, DS4_N_ROT, DS4_N_ROT, + pos0 + t, il, false); + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + float *kd = kc_t + (uint64_t)h * qk_dim; + memcpy(kd, kn_t + (uint64_t)h * q_nope, + (size_t)q_nope * sizeof(kd[0])); + memcpy(kd + q_nope, k_rot, (size_t)DS4_N_ROT * sizeof(kd[0])); + } + } + + const float scale = 1.0f / sqrtf((float)qk_dim); + for (uint32_t t = 0; t < n_tok; t++) { + const uint32_t visible = t + 1u; + for (uint32_t h = 0; h < DS4_N_HEAD; h++) { + const float *q_h = q + ((uint64_t)t * DS4_N_HEAD + h) * qk_dim; + float max_score = -FLT_MAX; + for (uint32_t s = 0; s < visible; s++) { + const float *k_h = + key_cache + ((uint64_t)s * DS4_N_HEAD + h) * qk_dim; + float dot = 0.0f; + for (uint32_t i = 0; i < qk_dim; i++) dot += q_h[i] * k_h[i]; + scores[s] = dot * scale; + if (scores[s] > max_score) max_score = scores[s]; + } + + float denom = 0.0f; + for (uint32_t s = 0; s < visible; s++) { + scores[s] = expf(scores[s] - max_score); + denom += scores[s]; + } + if (denom < 1.0e-20f) denom = 1.0e-20f; + + float *head_out = + heads + ((uint64_t)t * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; + for (uint32_t d = 0; d < DS4_N_VALUE_MLA; d++) { + float acc = 0.0f; + for (uint32_t s = 0; s < visible; s++) { + const float *v_h = + value_cache + ((uint64_t)s * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; + acc += scores[s] * v_h[d]; + } + head_out[d] = acc / denom; + } + } + + matvec_q8_0_f32_ref(out + (uint64_t)t * DS4_N_EMBD, + model, + layer->attn_output, + heads + (uint64_t)t * heads_dim); + } + + free(scores); + free(k_rot); + free(heads); + free(value_cache); + free(key_cache); + free(k_nope); + free(kv_norm); + free(kv_raw); + free(q); + free(q_rank_norm); + free(q_rank); + free(norm); +} + +static void layer_glm_dense_ffn_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + const uint64_t hidden = layer->ffn_gate->dim[1]; + const uint64_t in_dim = layer->ffn_gate->dim[0]; + const uint64_t blocks = (in_dim + 31) / 32; + float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); + float *up = xmalloc((size_t)hidden * sizeof(up[0])); + float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); + int8_t *xq = xmalloc((size_t)blocks * 32); + float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); + + if (layer->ffn_gate->type != DS4_TENSOR_Q8_0 || + layer->ffn_up->type != DS4_TENSOR_Q8_0 || + layer->ffn_down->type != DS4_TENSOR_Q8_0 || + layer->ffn_up->dim[0] != in_dim || + layer->ffn_up->dim[1] != hidden || + layer->ffn_down->dim[0] != hidden || + layer->ffn_down->dim[1] != DS4_N_EMBD) { + ds4_die("GLM dense FFN tensors have an unexpected layout"); + } + + quantize_q8_0_activation(x, xq, xscale, in_dim); + matvec_q8_0_pair_prequant(gate, up, model, layer->ffn_gate, layer->ffn_up, xq, xscale); + swiglu(mid, gate, up, hidden, 0.0f); + matvec_q8_0(out, model, layer->ffn_down, mid); + + free(xscale); + free(xq); + free(mid); + free(up); + free(gate); +} + +static void layer_glm_dense_ffn_one_f32_ref( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x) { + const uint64_t hidden = layer->ffn_gate->dim[1]; + const uint64_t in_dim = layer->ffn_gate->dim[0]; + float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); + float *up = xmalloc((size_t)hidden * sizeof(up[0])); + float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); + + if (layer->ffn_gate->type != DS4_TENSOR_Q8_0 || + layer->ffn_up->type != DS4_TENSOR_Q8_0 || + layer->ffn_down->type != DS4_TENSOR_Q8_0 || + layer->ffn_up->dim[0] != in_dim || + layer->ffn_up->dim[1] != hidden || + layer->ffn_down->dim[0] != hidden || + layer->ffn_down->dim[1] != DS4_N_EMBD) { + ds4_die("GLM F32 dense FFN reference found unexpected tensor layout"); + } + + matvec_q8_0_f32_ref(gate, model, layer->ffn_gate, x); + matvec_q8_0_f32_ref(up, model, layer->ffn_up, x); + swiglu(mid, gate, up, hidden, 0.0f); + matvec_q8_0_f32_ref(out, model, layer->ffn_down, mid); + + free(mid); + free(up); + free(gate); +} + +static void layer_glm_router_selected_experts( + int selected[DS4_MAX_EXPERT_USED], + float expert_weight[DS4_MAX_EXPERT_USED], + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x) { + float logits[DS4_MAX_EXPERT]; + float probs[DS4_MAX_EXPERT]; + float selection[DS4_MAX_EXPERT]; + const float *bias = tensor_data(model, layer->ffn_exp_probs_b); + + matvec_any(logits, model, layer->ffn_gate_inp, x); + for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { + probs[i] = sigmoid_stable(logits[i]); + selection[i] = probs[i] + bias[i]; + } + + topk_desc(selection, (int)DS4_N_EXPERT, (int)DS4_N_EXPERT_USED, selected); + + float sum = 0.0f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) { + ds4_die("GLM selected expert is outside router range"); + } + expert_weight[i] = probs[selected[i]]; + sum += expert_weight[i]; + } + if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + expert_weight[i] = expert_weight[i] / sum * DS4_EXPERT_WEIGHT_SCALE; + } +} + +typedef struct { + float *mid; + const float *x; + const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; + const uint8_t *up_base[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t out_dim; + uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; + uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; + uint32_t gate_type; + uint32_t up_type; + uint32_t n_expert; +} glm_routed_moe_f32_mid_ctx; + +static bool glm_graph_gate_pair_type_supported(uint32_t gate_type, uint32_t up_type) { + return gate_type == up_type && + (gate_type == DS4_TENSOR_IQ2_XXS || + gate_type == DS4_TENSOR_Q2_K || + gate_type == DS4_TENSOR_Q4_K || + gate_type == DS4_TENSOR_Q5_K); +} + +static bool glm_graph_down_type_supported(uint32_t down_type) { + return down_type == DS4_TENSOR_IQ2_XXS || + down_type == DS4_TENSOR_Q2_K || + down_type == DS4_TENSOR_Q4_K || + down_type == DS4_TENSOR_Q5_K || + down_type == DS4_TENSOR_Q6_K; +} + +static float glm_routed_moe_dot_f32(uint32_t type, int n, const uint8_t *row, const float *x) { + if (type == DS4_TENSOR_IQ2_XXS) { + return ds4_vec_dot_iq2_xxs_f32(n, (const block_iq2_xxs *)row, x); + } + if (type == DS4_TENSOR_Q2_K) { + return ds4_vec_dot_q2_K_f32(n, (const block_q2_K *)row, x); + } + if (type == DS4_TENSOR_Q4_K) { + return ds4_vec_dot_q4_K_f32(n, (const block_q4_K *)row, x); + } + if (type == DS4_TENSOR_Q5_K || type == DS4_TENSOR_Q6_K) { + return ds4_vec_dot_q5_q6_K_f32(type, n, row, x); + } + ds4_die("GLM F32 routed-MoE reference encountered unsupported expert tensor type"); + return 0.0f; +} + +static void glm_routed_moe_f32_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { + glm_routed_moe_f32_mid_ctx *ctx = vctx; + + for (uint64_t idx = row0; idx < row1; idx++) { + const uint32_t slot = (uint32_t)(idx / ctx->out_dim); + const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; + const uint8_t *gate_row = ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]; + const uint8_t *up_row = ctx->up_base[slot] + row * ctx->up_row_bytes[slot]; + const float gate = glm_routed_moe_dot_f32(ctx->gate_type, (int)ctx->in_dim, gate_row, ctx->x); + const float up = glm_routed_moe_dot_f32(ctx->up_type, (int)ctx->in_dim, up_row, ctx->x); + ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; + } +} + +typedef struct { + float *out; + const float *mid; + const uint8_t *down_base[DS4_MAX_EXPERT_USED]; + uint64_t in_dim; + uint64_t out_dim; + uint64_t down_row_bytes[DS4_MAX_EXPERT_USED]; + uint32_t down_type; + uint32_t n_expert; +} glm_routed_moe_f32_down_ctx; + +static void glm_routed_moe_f32_down_worker(void *vctx, uint64_t row0, uint64_t row1) { + glm_routed_moe_f32_down_ctx *ctx = vctx; + + for (uint64_t row = row0; row < row1; row++) { + float acc = 0.0f; + for (uint32_t slot = 0; slot < ctx->n_expert; slot++) { + const uint8_t *down_row = ctx->down_base[slot] + row * ctx->down_row_bytes[slot]; + acc += glm_routed_moe_dot_f32(ctx->down_type, + (int)ctx->in_dim, + down_row, + ctx->mid + (uint64_t)slot * ctx->in_dim); + } + ctx->out[row] = acc; + } +} + +static void layer_glm_routed_moe_one_f32_ref( + float *out, + float *mid_all, + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x, + const int *selected, + const float *expert_weight) { + if (!layer->ffn_gate_exps || !layer->ffn_up_exps || !layer->ffn_down_exps || + !tensor_is_routed_expert_type(layer->ffn_gate_exps->type) || + !tensor_is_routed_expert_type(layer->ffn_up_exps->type) || + !glm_graph_gate_pair_type_supported(layer->ffn_gate_exps->type, + layer->ffn_up_exps->type) || + !glm_graph_down_type_supported(layer->ffn_down_exps->type)) { + ds4_die("GLM F32 routed-MoE reference expects supported matching routed gate/up tensors and down tensors"); + } + + glm_routed_moe_f32_mid_ctx mid_ctx = { + .mid = mid_all, + .x = x, + .gate_type = layer->ffn_gate_exps->type, + .up_type = layer->ffn_up_exps->type, + .n_expert = DS4_N_EXPERT_USED, + }; + glm_routed_moe_f32_down_ctx down_ctx = { + .out = out, + .mid = mid_all, + .down_type = layer->ffn_down_exps->type, + .n_expert = DS4_N_EXPERT_USED, + }; + + uint64_t gate_in0 = 0, gate_out0 = 0; + uint64_t down_in0 = 0, down_out0 = 0; + for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { + uint64_t gate_in, gate_out, up_in, up_out, down_in, down_out; + if (selected[slot] < 0 || (uint32_t)selected[slot] >= DS4_N_EXPERT) { + ds4_die("GLM F32 routed-MoE reference selected expert is outside range"); + } + mid_ctx.gate_base[slot] = + tensor_expert_bytes(model, layer->ffn_gate_exps, (uint32_t)selected[slot], + &gate_in, &gate_out, &mid_ctx.gate_row_bytes[slot]); + mid_ctx.up_base[slot] = + tensor_expert_bytes(model, layer->ffn_up_exps, (uint32_t)selected[slot], + &up_in, &up_out, &mid_ctx.up_row_bytes[slot]); + down_ctx.down_base[slot] = + tensor_expert_bytes(model, layer->ffn_down_exps, (uint32_t)selected[slot], + &down_in, &down_out, &down_ctx.down_row_bytes[slot]); + if (gate_in != up_in || gate_out != up_out || + down_in != gate_out || down_out != DS4_N_EMBD) { + ds4_die("GLM F32 routed-MoE reference found mismatched expert layouts"); + } + if (slot == 0) { + gate_in0 = gate_in; + gate_out0 = gate_out; + down_in0 = down_in; + down_out0 = down_out; + } else if (gate_in != gate_in0 || gate_out != gate_out0 || + down_in != down_in0 || down_out != down_out0) { + ds4_die("GLM F32 routed-MoE reference expert layouts are not uniform"); + } + mid_ctx.expert_weight[slot] = expert_weight[slot]; + } + + if (gate_in0 != DS4_N_EMBD || gate_in0 % QK_K != 0 || + down_in0 != DS4_N_FF_EXP || down_in0 % QK_K != 0 || + gate_out0 != DS4_N_FF_EXP || down_out0 != DS4_N_EMBD) { + ds4_die("GLM F32 routed-MoE reference found unexpected GLM expert dimensions"); + } + + mid_ctx.in_dim = gate_in0; + mid_ctx.out_dim = gate_out0; + down_ctx.in_dim = down_in0; + down_ctx.out_dim = down_out0; + ds4_parallel_for((uint64_t)DS4_N_EXPERT_USED * gate_out0, + glm_routed_moe_f32_mid_worker, + &mid_ctx); + ds4_parallel_for(down_out0, glm_routed_moe_f32_down_worker, &down_ctx); +} + +static void layer_glm_shared_ffn_one_f32_ref( + float *out, + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x) { + const uint64_t in_dim = layer->ffn_gate_shexp ? layer->ffn_gate_shexp->dim[0] : 0; + const uint64_t hidden = layer->ffn_gate_shexp ? layer->ffn_gate_shexp->dim[1] : 0; + if (!layer->ffn_gate_shexp || + !layer->ffn_up_shexp || + !layer->ffn_down_shexp || + layer->ffn_gate_shexp->type != DS4_TENSOR_Q8_0 || + layer->ffn_up_shexp->type != DS4_TENSOR_Q8_0 || + layer->ffn_down_shexp->type != DS4_TENSOR_Q8_0 || + layer->ffn_up_shexp->dim[0] != in_dim || + layer->ffn_up_shexp->dim[1] != hidden || + layer->ffn_down_shexp->dim[0] != hidden || + layer->ffn_down_shexp->dim[1] != DS4_N_EMBD || + in_dim != DS4_N_EMBD || + hidden != DS4_N_FF_EXP) { + ds4_die("GLM F32 shared expert reference found unexpected tensor layout"); + } + + float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); + float *up = xmalloc((size_t)hidden * sizeof(up[0])); + float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); + + matvec_q8_0_f32_ref(gate, model, layer->ffn_gate_shexp, x); + matvec_q8_0_f32_ref(up, model, layer->ffn_up_shexp, x); + swiglu(mid, gate, up, hidden, 0.0f); + matvec_q8_0_f32_ref(out, model, layer->ffn_down_shexp, mid); + + free(mid); + free(up); + free(gate); +} + +static void layer_glm_ffn_one_f32_ref( + float *out, + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x, + uint32_t il) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + + rms_norm_weight(norm, x, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); + if (il < DS4_N_LEADING_DENSE) { + layer_glm_dense_ffn_one_f32_ref(out, model, layer, norm); + } else { + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + float *mid = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid[0])); + float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); + float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); + + layer_glm_router_selected_experts(selected, expert_weight, model, layer, norm); + layer_glm_routed_moe_one_f32_ref(moe, mid, model, layer, norm, + selected, expert_weight); + layer_glm_shared_ffn_one_f32_ref(shared, model, layer, norm); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = moe[i] + shared[i]; + + free(shared); + free(moe); + free(mid); + } + + free(norm); +} + +static void layer_glm_first_token_one_f32_ref( + float *out, + const ds4_model *model, + const ds4_layer_weights *layer, + const float *x, + uint32_t il) { + float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); + float *after_attn = xmalloc((size_t)DS4_N_EMBD * sizeof(after_attn[0])); + float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); + + layer_glm_first_token_attention_one_f32_ref(attn_out, model, layer, x); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) after_attn[i] = x[i] + attn_out[i]; + + layer_glm_ffn_one_f32_ref(ffn_out, model, layer, after_attn, il); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = after_attn[i] + ffn_out[i]; + + free(ffn_out); + free(after_attn); + free(attn_out); +} + +static void forward_glm_first_token_cpu_f32_ref( + float *out_hidden, + const ds4_model *model, + const ds4_weights *weights, + int token) { + float *cur = xmalloc((size_t)DS4_N_EMBD * sizeof(cur[0])); + float *next = xmalloc((size_t)DS4_N_EMBD * sizeof(next[0])); + const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + + embed_token_any(model, weights, token, cur); + for (uint32_t il = 0; il < normal_layers; il++) { + layer_glm_first_token_one_f32_ref(next, model, &weights->layer[il], cur, il); + float *tmp = cur; + cur = next; + next = tmp; + } + + memcpy(out_hidden, cur, (size_t)DS4_N_EMBD * sizeof(out_hidden[0])); + + free(next); + free(cur); +} + +static void output_logits_glm_one_f32_ref( + float *logits, + const ds4_model *model, + const ds4_weights *weights, + const float *hidden) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + + rms_norm_weight(norm, hidden, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0_f32_ref(logits, model, weights->output, norm); + + free(norm); +} + +static void layer_glm_routed_moe_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il) { + int selected[DS4_MAX_EXPERT_USED]; + float expert_weight[DS4_MAX_EXPERT_USED]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + float *mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid_all[0])); + block_q8_K *xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0])); + block_q8_K *midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0])); + + if (expert_in_dim != DS4_N_EMBD || expert_in_dim % QK_K != 0 || + down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { + ds4_die("GLM routed expert tensors have an unexpected layout"); + } + + memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); + ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); + layer_glm_router_selected_experts(selected, expert_weight, model, layer, x); + + matvec_experts_mid_prequant(mid_all, model, + layer->ffn_gate_exps, + layer->ffn_up_exps, + xq, + selected, + expert_weight, + DS4_N_EXPERT_USED, + 0.0f); + for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { + ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, + midq + (uint64_t)i * (down_in_dim / QK_K), + (int64_t)down_in_dim); + } + matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, + selected, DS4_N_EXPERT_USED); + + free(midq); + free(xq); + free(mid_all); + (void)il; +} + +static void layer_glm_sparse_ffn_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il) { + float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); + float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); + + layer_glm_routed_moe_one(moe, model, layer, x, il); + layer_shared_ffn_one(shared, model, layer, x); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = moe[i] + shared[i]; + + free(shared); + free(moe); +} + +static void layer_glm_ffn_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + + rms_norm_weight(norm, x, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); + if (il < DS4_N_LEADING_DENSE) { + layer_glm_dense_ffn_one(out, model, layer, norm); + } else { + layer_glm_sparse_ffn_one(out, model, layer, norm, il); + } + + free(norm); +} + +static void layer_glm_first_token_one( + float * out, + const ds4_model * model, + const ds4_layer_weights * layer, + const float * x, + uint32_t il) { + float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); + float *after_attn = xmalloc((size_t)DS4_N_EMBD * sizeof(after_attn[0])); + float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); + + layer_glm_first_token_attention_one(attn_out, model, layer, x); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) after_attn[i] = x[i] + attn_out[i]; + + layer_glm_ffn_one(ffn_out, model, layer, after_attn, il); + for (uint32_t i = 0; i < DS4_N_EMBD; i++) out[i] = after_attn[i] + ffn_out[i]; + + free(ffn_out); + free(after_attn); + free(attn_out); +} + +static void forward_glm_first_token_cpu( + float * out_hidden, + const ds4_model * model, + const ds4_weights * weights, + int token) { + float *cur = xmalloc((size_t)DS4_N_EMBD * sizeof(cur[0])); + float *next = xmalloc((size_t)DS4_N_EMBD * sizeof(next[0])); + const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + + embed_token_any(model, weights, token, cur); + for (uint32_t il = 0; il < normal_layers; il++) { + layer_glm_first_token_one(next, model, &weights->layer[il], cur, il); + float *tmp = cur; + cur = next; + next = tmp; + } + + memcpy(out_hidden, cur, (size_t)DS4_N_EMBD * sizeof(out_hidden[0])); + + free(next); + free(cur); +} + +static void output_logits_glm_one( + float * logits, + const ds4_model * model, + const ds4_weights * weights, + const float * hidden) { + float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); + + rms_norm_weight(norm, hidden, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); + matvec_q8_0(logits, model, weights->output, norm); + + free(norm); +} diff --git a/models/glm/cuda/kernels.inc b/models/glm/cuda/kernels.inc new file mode 100644 index 0000000000..94e89d5970 --- /dev/null +++ b/models/glm/cuda/kernels.inc @@ -0,0 +1,4578 @@ +__device__ __forceinline__ static float glm_rope_yarn_corr_factor_dev( + int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * logf(n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) / + (2.0f * logf(base)); +} +__device__ __forceinline__ static float glm_rope_yarn_ramp_dev( + float low, float high, int i0) { + const float y = (i0 / 2 - low) / fmaxf(0.001f, high - low); + return 1.0f - fminf(1.0f, fmaxf(0.0f, y)); +} +__device__ __forceinline__ static void glm_rope_yarn_dev( + float theta_extrap, float freq_scale, const float corr_dims[2], + int i0, float ext_factor, float mscale, + float *cos_theta, float *sin_theta) { + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + float ramp_mix = glm_rope_yarn_ramp_dev(corr_dims[0], corr_dims[1], i0) * + ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + *cos_theta = cosf(theta) * mscale; + *sin_theta = sinf(theta) * mscale; +} + +static int cuda_current_tier(void) { + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) return 0; + return dev; +} + +/* ===== GLM 5.2 stubs (to be implemented; fail loudly) ===== */ +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-parameter" +__global__ static void add3_kernel(float *out, const float *a, + const float *b, const float *c, + uint32_t n) { + uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = a[i] + b[i] + c[i]; +} + +extern "C" int ds4_gpu_add3_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + uint32_t n) { + if (!out || !a || !b || !c || n == 0 || + out->bytes < (uint64_t)n * sizeof(float) || + a->bytes < (uint64_t)n * sizeof(float) || + b->bytes < (uint64_t)n * sizeof(float) || + c->bytes < (uint64_t)n * sizeof(float)) { + return 0; + } + add3_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, (const float *)a->ptr, + (const float *)b->ptr, (const float *)c->ptr, n); + return cuda_ok(cudaGetLastError(), "add3 launch"); +} + +/* Fused decode residual: sum_out = a + b; norm_out = rmsnorm(sum) * w. + * Single row, one block (two-pass over n with a shared reduction). */ +__global__ static void glm_add_rms_norm_weight_kernel( + float *norm_out, + float *sum_out, + const float *a, + const float *b, + const float *w, + uint32_t n, + float eps) { + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + __shared__ float sh[32]; + float sumsq = 0.0f; + for (uint32_t i = tid; i < n; i += nth) { + const float v = a[i] + b[i]; + sum_out[i] = v; + sumsq += v * v; + } + for (int off = 16; off > 0; off >>= 1) { + sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off); + } + if ((tid & 31u) == 0u) sh[tid >> 5] = sumsq; + __syncthreads(); + if (tid < 32u) { + sumsq = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f; + for (int off = 16; off > 0; off >>= 1) { + sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off); + } + if (tid == 0u) sh[0] = sumsq; + } + __syncthreads(); + const float scale = rsqrtf(sh[0] / (float)n + eps); + for (uint32_t i = tid; i < n; i += nth) { + norm_out[i] = (sum_out[i] * scale) * w[i]; + } +} + +extern "C" int ds4_gpu_add_rms_norm_weight_tensor( + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *sum_out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps) { + if (!norm_out || !sum_out || !a || !b || !model_map || n == 0 || + norm_out->bytes < (uint64_t)n * sizeof(float) || + sum_out->bytes < (uint64_t)n * sizeof(float) || + a->bytes < (uint64_t)n * sizeof(float) || + b->bytes < (uint64_t)n * sizeof(float) || + weight_offset > model_size || + (uint64_t)n * sizeof(float) > model_size - weight_offset) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const float *w = (const float *)cuda_resolve_weight_ptr( + model_map, weight_offset, (uint64_t)n * sizeof(float), + logical_tier, "rms_weight"); + if (!w) return 0; + glm_add_rms_norm_weight_kernel<<<1, 1024>>>( + (float *)norm_out->ptr, (float *)sum_out->ptr, + (const float *)a->ptr, (const float *)b->ptr, w, n, eps); + return cuda_ok(cudaGetLastError(), "add rms norm weight"); +} + +extern "C" bool ds4_gpu_commands_active(void) { + return false; +} + +__global__ static void glm_embed_token_q8_0_kernel( + float *out, + const unsigned char *w, + uint32_t token, + uint32_t n_embd) { + uint32_t d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= n_embd) return; + const uint64_t row_blocks = n_embd / 32u; + const unsigned char *blk = + w + ((uint64_t)token * row_blocks + (d >> 5)) * 34u; + const float scale = __half2float(*(const __half *)blk); + out[d] = scale * (float)((const int8_t *)(blk + 2))[d & 31u]; +} + +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 (!out || !model_map || n_embd == 0 || (n_embd & 31u) != 0u || + token >= n_vocab) { + return 0; + } + if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ + fprintf(stderr, "ds4: embed_token_quant: unsupported type %u\n", + weight_type); + return 0; + } + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + if (weight_offset > model_size || + (uint64_t)n_vocab * row_bytes > model_size - weight_offset || + out->bytes < (uint64_t)n_embd * sizeof(float)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr( + model_map, weight_offset, (uint64_t)n_vocab * row_bytes, + logical_tier, "glm_token_embd"); + if (!w) return 0; + glm_embed_token_q8_0_kernel<<<(n_embd + 255) / 256, 256>>>( + (float *)out->ptr, w, token, n_embd); + return cuda_ok(cudaGetLastError(), "glm embed token launch"); +} + +__global__ static void glm_embed_tokens_q8_0_kernel( + float *out, + const int32_t *tokens, + const unsigned char *w, + uint32_t n_tokens, + uint32_t n_embd) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_embd; + if (gid >= n) return; + uint32_t t = gid / n_embd; + uint32_t d = gid - (uint64_t)t * n_embd; + int32_t tok = tokens[t]; + const uint64_t row_blocks = n_embd / 32u; + const unsigned char *blk = w + ((uint64_t)tok * row_blocks + (d >> 5)) * 34u; + const float scale = __half2float(*(const __half *)blk); + out[gid] = scale * (float)((const int8_t *)(blk + 2))[d & 31u]; +} + +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 (!out || !tokens || !model_map || n_tokens == 0 || n_embd == 0 || + (n_embd & 31u) != 0u) { + return 0; + } + if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ + fprintf(stderr, "ds4: embed_tokens_quant: unsupported type %u\n", + weight_type); + return 0; + } + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + if (weight_offset > model_size || + (uint64_t)n_vocab * row_bytes > model_size - weight_offset || + out->bytes < (uint64_t)n_tokens * n_embd * sizeof(float) || + tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr( + model_map, weight_offset, (uint64_t)n_vocab * row_bytes, + logical_tier, "glm_token_embd"); + if (!w) return 0; + uint64_t n = (uint64_t)n_tokens * n_embd; + glm_embed_tokens_q8_0_kernel<<<(n + 255) / 256, 256>>>( + (float *)out->ptr, + (const int32_t *)tokens->ptr, + w, n_tokens, n_embd); + return cuda_ok(cudaGetLastError(), "glm embed tokens launch"); +} + +extern "C" int ds4_gpu_flush_encoder(void) { + /* Metal encoder flush: CUDA kernels are already queued in stream + * order, nothing to split. */ + return 1; +} + +extern "C" int ds4_gpu_glm_attention_flash_staged_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_flash_staged_tensor\n"); + return 0; +} + +extern "C" int ds4_gpu_glm_attention_flash_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_flash_tensor\n"); + return 0; +} + +extern "C" int ds4_gpu_glm_attention_full_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_full_tensor\n"); + return 0; +} + +template +__device__ __forceinline__ static float2 glm_cache_rope_pair_f16_dev( + const CT *rope_cache, uint64_t rope_base, uint32_t r, + uint32_t row, uint32_t qk_rope, float freq_base, float freq_scale, + float ext_factor, float attn_factor, const float corr_dims[2]) { + const float theta_base = (float)row; + const float inv_ndims = -1.0f / (float)qk_rope; + const float theta = theta_base * powf(freq_base, inv_ndims * (float)r); + float ct, st; + glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)r, + ext_factor, attn_factor, &ct, &st); + const float x0 = (float)rope_cache[rope_base + r]; + const float x1 = (float)rope_cache[rope_base + r + 1u]; + return make_float2(x0 * ct - x1 * st, x0 * st + x1 * ct); +} + +/* Scalar-correct MLA attention: one warp per head, grid + * (ceil(n_head/8), n_tokens). The row loop handles either a contiguous + * causal range or an explicit selected-row list and mirrors the Metal online + * softmax so numerics stay comparable. */ +template +__global__ static void glm_attention_lora_causal_kernel( + float *lora_out, + const float *q, + const float *qk_low, + const CT *kv_lora_cache, + const CT *k_rope_cache, + const uint32_t *selected, + uint32_t cache_cap, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float scale) { + const uint32_t token = blockIdx.y; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t head = blockIdx.x * 8u + warp; + if (token >= n_tokens || head >= n_head || kv_lora_dim != 512u || + qk_rope != 64u) { + return; + } + const uint32_t visible = selected_rows + ? n_selected : min(n_selected, pos0 + token + 1u); + if (visible == 0u) return; + + const uint32_t qk_dim = qk_nope + qk_rope; + const float *qh = q + (uint64_t)token * n_head * qk_dim + + (uint64_t)head * qk_dim; + const float4 *low4 = (const float4 *)(qk_low + + (uint64_t)token * n_head * kv_lora_dim + + (uint64_t)head * kv_lora_dim); + + float4 low0 = low4[lane]; + float4 low1 = low4[lane + 32u]; + float4 low2 = low4[lane + 64u]; + float4 low3 = low4[lane + 96u]; + float4 qrope = make_float4(0.f, 0.f, 0.f, 0.f); + const uint32_t rope_vecs = qk_rope >> 2; /* 16 */ + if (lane < rope_vecs) { + qrope = *((const float4 *)(qh + qk_nope + lane * 4u)); + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (ext_factor != 0.0f) { + corr_dims[0] = fmaxf(0.0f, + floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig, + beta_fast, freq_base))); + corr_dims[1] = fminf((float)qk_rope - 1.0f, + ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig, + beta_slow, freq_base))); + } + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float4 o0 = make_float4(0.f,0.f,0.f,0.f); + float4 o1 = o0, o2 = o0, o3 = o0; + + for (uint32_t ri = 0u; ri < visible; ri++) { + const uint32_t row = selected_rows + ? selected[(uint64_t)token * n_selected + ri] : ri; + if (row >= cache_cap) continue; + const CT *kvrow = kv_lora_cache + (uint64_t)row * kv_lora_dim; + float partial = 0.0f; + { + const float4 k0 = make_float4( + (float)(kvrow[lane*4u+0u]), (float)(kvrow[lane*4u+1u]), + (float)(kvrow[lane*4u+2u]), (float)(kvrow[lane*4u+3u])); + const float4 k1 = make_float4( + (float)(kvrow[(lane+32u)*4u+0u]), (float)(kvrow[(lane+32u)*4u+1u]), + (float)(kvrow[(lane+32u)*4u+2u]), (float)(kvrow[(lane+32u)*4u+3u])); + const float4 k2 = make_float4( + (float)(kvrow[(lane+64u)*4u+0u]), (float)(kvrow[(lane+64u)*4u+1u]), + (float)(kvrow[(lane+64u)*4u+2u]), (float)(kvrow[(lane+64u)*4u+3u])); + const float4 k3 = make_float4( + (float)(kvrow[(lane+96u)*4u+0u]), (float)(kvrow[(lane+96u)*4u+1u]), + (float)(kvrow[(lane+96u)*4u+2u]), (float)(kvrow[(lane+96u)*4u+3u])); + partial += low0.x*k0.x + low0.y*k0.y + low0.z*k0.z + low0.w*k0.w; + partial += low1.x*k1.x + low1.y*k1.y + low1.z*k1.z + low1.w*k1.w; + partial += low2.x*k2.x + low2.y*k2.y + low2.z*k2.z + low2.w*k2.w; + partial += low3.x*k3.x + low3.y*k3.y + low3.z*k3.z + low3.w*k3.w; + if (lane < rope_vecs) { + const uint64_t rope_base = (uint64_t)row * qk_rope; + const uint32_t r = lane * 4u; + const float2 y0 = glm_cache_rope_pair_f16_dev( + k_rope_cache, rope_base, r, row, qk_rope, freq_base, + freq_scale, ext_factor, attn_factor, corr_dims); + const float2 y1 = glm_cache_rope_pair_f16_dev( + k_rope_cache, rope_base, r + 2u, row, qk_rope, + freq_base, freq_scale, ext_factor, attn_factor, + corr_dims); + partial += qrope.x*y0.x + qrope.y*y0.y + + qrope.z*y1.x + qrope.w*y1.y; + } + for (uint32_t off = 16u; off > 0u; off >>= 1u) { + partial += __shfl_xor_sync(0xffffffffu, partial, off); + } + const float score = partial * scale; + const float new_m = fmaxf(M, score); + const float old_scale = expf(M - new_m); + const float row_scale = expf(score - new_m); + o0.x = o0.x*old_scale + k0.x*row_scale; o0.y = o0.y*old_scale + k0.y*row_scale; + o0.z = o0.z*old_scale + k0.z*row_scale; o0.w = o0.w*old_scale + k0.w*row_scale; + o1.x = o1.x*old_scale + k1.x*row_scale; o1.y = o1.y*old_scale + k1.y*row_scale; + o1.z = o1.z*old_scale + k1.z*row_scale; o1.w = o1.w*old_scale + k1.w*row_scale; + o2.x = o2.x*old_scale + k2.x*row_scale; o2.y = o2.y*old_scale + k2.y*row_scale; + o2.z = o2.z*old_scale + k2.z*row_scale; o2.w = o2.w*old_scale + k2.w*row_scale; + o3.x = o3.x*old_scale + k3.x*row_scale; o3.y = o3.y*old_scale + k3.y*row_scale; + o3.z = o3.z*old_scale + k3.z*row_scale; o3.w = o3.w*old_scale + k3.w*row_scale; + S = S*old_scale + row_scale; + M = new_m; + } + } + + const float inv_s = S > 0.0f ? 1.0f / S : 0.0f; + float4 *out4 = (float4 *)(lora_out + + ((uint64_t)token * n_head + head) * kv_lora_dim); + o0.x*=inv_s; o0.y*=inv_s; o0.z*=inv_s; o0.w*=inv_s; + o1.x*=inv_s; o1.y*=inv_s; o1.z*=inv_s; o1.w*=inv_s; + o2.x*=inv_s; o2.y*=inv_s; o2.z*=inv_s; o2.w*=inv_s; + o3.x*=inv_s; o3.y*=inv_s; o3.z*=inv_s; o3.w*=inv_s; + out4[lane] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; +} + +extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + ds4_gpu_tensor *lora_out, + 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, + uint32_t n_tokens, + uint32_t pos0, + 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + n_tokens == 0 || n_head == 0 || kv_lora_dim != 512u || + qk_rope != 64u) { + fprintf(stderr, "ds4: glm attn lora causal: unsupported config " + "(n_tok=%u head=%u lora=%u rope=%u f16=%d)\n", + n_tokens, n_head, kv_lora_dim, qk_rope, (int)cache_f16); + return 0; + } + const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); + dim3 grid((n_head + 7u) / 8u, n_tokens, 1); + if (cache_f16) { + glm_attention_lora_causal_kernel<__half, false><<>>( + (float *)lora_out->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + NULL, cache_cap, + n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope, + qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, scale); + } else { + glm_attention_lora_causal_kernel<<>>( + (float *)lora_out->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + NULL, cache_cap, + n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope, + qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, scale); + } + return cuda_ok(cudaGetLastError(), "glm attn lora causal launch"); +} + +extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + ds4_gpu_tensor *lora_out, + 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 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float); + const uint64_t qk_dim = (uint64_t)qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !selected || n_tokens == 0u || n_selected == 0u || n_head == 0u || + kv_lora_dim != 512u || qk_rope != 64u || cache_cap == 0u || + selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t) || + q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || + qk_low->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || + kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem || + lora_out->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) { + fprintf(stderr, "ds4: glm attn lora selected: unsupported config " + "(n_tok=%u selected=%u head=%u lora=%u rope=%u f16=%d)\n", + n_tokens, n_selected, n_head, kv_lora_dim, qk_rope, + (int)cache_f16); + return 0; + } + const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); + dim3 grid((n_head + 7u) / 8u, n_tokens, 1); + if (cache_f16) { + glm_attention_lora_causal_kernel<__half, true><<>>( + (float *)lora_out->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, cache_cap, + n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope, + qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, scale); + } else { + glm_attention_lora_causal_kernel<<>>( + (float *)lora_out->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, cache_cap, + n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope, + qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, scale); + } + return cuda_ok(cudaGetLastError(), "glm attn lora selected launch"); +} + +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) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_batch_typed_tensor\n"); + return 0; +} + +extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + 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, + bool selected_rows_valid, + 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, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor\n"); + return 0; +} + +__device__ __forceinline__ static float glm_q8_0_dot_row_dev( + const char *row, const float *x, uint32_t n_cols) { + float acc = 0.0f; + const uint32_t nb = n_cols >> 5; + for (uint32_t b = 0; b < nb; b++) { + const char *blk = row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + const int8_t *q = (const int8_t *)(blk + 2); + float s = 0.0f; + #pragma unroll 8 + for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * x[b * 32u + k]; + acc += d * s; + } + return acc; +} + +template +__device__ __forceinline__ static float2 glm_cache_value_pair_dev( + const CT *p) { + return make_float2((float)p[0], (float)p[1]); +} + +template <> +__device__ __forceinline__ float2 glm_cache_value_pair_dev<__half>( + const __half *p) { + return __half22float2(*(const __half2 *)p); +} + +template <> +__device__ __forceinline__ float2 glm_cache_value_pair_dev( + const float *p) { + return *(const float2 *)p; +} + +/* Exact staged decode attention. The original fused kernel owns one block per + * head, which leaves more than half of an L40S idle. These stages preserve the + * fused kernel's arithmetic order for every score, softmax lane, lora output, + * and value-projection row while exposing independent rows/dimensions as + * separate blocks. */ +template +__global__ static void glm_attention_decode_weights_staged_kernel( + float *weights, + float *denom, + const float *q, + const float *qk_low, + const CT *kv_lora_cache, + const CT *k_rope_cache, + const uint32_t *selected, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + float scale, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool score_vec2) { + const uint32_t head = blockIdx.x; + const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u; + const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); + const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u); + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + if (head >= n_head || row_count == 0u) return; + const uint32_t qk_dim = qk_nope + qk_rope; + extern __shared__ float glm_dec_stage_sh[]; + float *red = glm_dec_stage_sh; + float *scores = glm_dec_stage_sh + 256u; + const float *qh = q + + ((uint64_t)token * n_head + head) * qk_dim; + const float *low = qk_low + + ((uint64_t)token * n_head + head) * kv_lora_dim; + + float corr_dims[2] = {0.0f, 0.0f}; + if (ext_factor != 0.0f) { + corr_dims[0] = fmaxf(0.0f, + floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, + (int)n_ctx_orig, beta_fast, freq_base))); + corr_dims[1] = fminf((float)qk_rope - 1.0f, + ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, + (int)n_ctx_orig, beta_slow, freq_base))); + } + + float local_max = -FLT_MAX; + for (uint32_t s = tid; s < row_count; s += nth) { + const uint32_t row = RANGE_TOK2 ? s : selected[s]; + float score = -FLT_MAX; + if (row < cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * kv_lora_dim; + if (score_vec2) { + for (uint32_t j = 0; j < kv_lora_dim; j += 2u) { + const float2 x = *(const float2 *)(low + j); + const float2 y = glm_cache_value_pair_dev( + kv_lora_cache + lora_base + j); + dotv += x.x * y.x; + dotv += x.y * y.y; + } + } else { + for (uint32_t j = 0; j < kv_lora_dim; j++) { + dotv += low[j] * (float)kv_lora_cache[lora_base + j]; + } + } + const uint64_t rope_base = (uint64_t)row * qk_rope; + for (uint32_t r = 0; r < qk_rope; r += 2u) { + const float2 y = glm_cache_rope_pair_f16_dev( + k_rope_cache, rope_base, r, row, qk_rope, + freq_base, freq_scale, ext_factor, attn_factor, + corr_dims); + dotv += qh[qk_nope + r] * y.x + + qh[qk_nope + r + 1u] * y.y; + } + score = dotv * scale; + } + scores[s] = score; + local_max = fmaxf(local_max, score); + } + red[tid] = local_max; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]); + __syncthreads(); + } + const float max_score = red[0]; + __syncthreads(); + + float local_sum = 0.0f; + for (uint32_t s = tid; s < row_count; s += nth) { + const float w = expf(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + __syncthreads(); + } + const uint64_t head_index = (uint64_t)token * n_head + head; + if (tid == 0u) denom[head_index] = fmaxf(red[0], 1.0e-20f); + float *head_weights = weights + head_index * score_stride; + for (uint32_t s = tid; s < row_count; s += nth) { + head_weights[s] = scores[s]; + } +} + +template +__global__ static void glm_attention_decode_lora_staged_kernel( + float *lora_sum, + const float *scores, + const float *denom, + const CT *kv_lora_cache, + const uint32_t *selected, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_lora_dim) { + const uint32_t head = blockIdx.y; + const uint32_t token = RANGE_TOK2 ? blockIdx.z : 0u; + const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); + const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u); + const uint32_t pair = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t j = pair * 2u; + if (head >= n_head || j >= kv_lora_dim) return; + const uint64_t head_index = (uint64_t)token * n_head + head; + const float *head_scores = scores + head_index * score_stride; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t s = 0; s < row_count; s++) { + const uint32_t row = RANGE_TOK2 ? s : selected[s]; + if (row < cache_cap) { + const float2 v = glm_cache_value_pair_dev( + kv_lora_cache + (uint64_t)row * kv_lora_dim + j); + const float w = head_scores[s]; + acc0 += w * v.x; + acc1 += w * v.y; + } + } + float *out = lora_sum + head_index * kv_lora_dim + j; + out[0] = acc0 / denom[head_index]; + out[1] = acc1 / denom[head_index]; +} + +template +__global__ static void glm_attention_decode_value_staged_kernel( + float *heads, + const float *lora_sum, + const char *value_weight, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim, + uint32_t value_row_bytes) { + const uint32_t head = blockIdx.y; + const uint32_t token = TOK2 ? blockIdx.z : 0u; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t nwarps = blockDim.x >> 5; + const uint32_t out_warp = blockIdx.x * nwarps + warp; + const uint32_t total_warps = gridDim.x * nwarps; + if (head >= n_head) return; + const float *low = lora_sum + + ((uint64_t)token * n_head + head) * kv_lora_dim; + float *out = heads + + ((uint64_t)token * n_head + head) * value_dim; + const uint32_t nblk = kv_lora_dim >> 5; + for (uint32_t d = out_warp; d < value_dim; d += total_warps) { + const char *row = value_weight + + ((uint64_t)head * value_dim + d) * value_row_bytes; + float acc = 0.0f; + for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) { + const char *b = row + (uint64_t)blk * 34u; + const float dscale = __half2float(*(const __half *)b); + const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u; + const float *xs = low + blk * 32u + (lane & 1u) * 16u; + float s = 0.0f; + #pragma unroll + for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k]; + acc += dscale * s; + } + for (int off = 16; off > 0; off >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, off); + } + if (lane == 0u) out[d] = acc; + } +} + +/* Single-token indexed MLA decode attention, one block per head. + * Fuses score (qk_low . kv_lora + q_rope . rope(k_rope@row)), softmax over + * the indexer-selected rows, the weighted kv_lora sum, and the per-head + * value projection (q8_0). Dynamic shared: red[256] + scores[n_selected] + + * lora_sum[kv_lora_dim]. */ +template +__global__ static void glm_attention_indexed_decode_kernel( + float *heads, + const float *q, + const float *qk_low, + const CT *kv_lora_cache, + const CT *k_rope_cache, + const char *value_weight, + const uint32_t *selected, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t value_row_bytes, + bool lora_vec2, + bool score_vec2, + float scale, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint32_t head = blockIdx.x; + const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u; + const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u); + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + if (head >= n_head || row_count == 0u) return; + const uint32_t qk_dim = qk_nope + qk_rope; + + extern __shared__ float glm_dec_sh[]; + float *red = glm_dec_sh; + float *scores = glm_dec_sh + 256u; + float *lora_sum = scores + row_count; + + const float *qh = q + + ((uint64_t)token * n_head + head) * qk_dim; + const float *low = qk_low + + ((uint64_t)token * n_head + head) * kv_lora_dim; + + float corr_dims[2] = {0.0f, 0.0f}; + if (ext_factor != 0.0f) { + corr_dims[0] = fmaxf(0.0f, + floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, + (int)n_ctx_orig, beta_fast, freq_base))); + corr_dims[1] = fminf((float)qk_rope - 1.0f, + ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, + (int)n_ctx_orig, beta_slow, freq_base))); + } + + float local_max = -FLT_MAX; + for (uint32_t s = tid; s < row_count; s += nth) { + const uint32_t row = RANGE_TOK2 ? s : selected[s]; + float score = -FLT_MAX; + if (row < cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * kv_lora_dim; + if (score_vec2) { + for (uint32_t j = 0; j < kv_lora_dim; j += 2u) { + const float2 x = *(const float2 *)(low + j); + const float2 y = glm_cache_value_pair_dev( + kv_lora_cache + lora_base + j); + dotv += x.x * y.x; + dotv += x.y * y.y; + } + } else { + for (uint32_t j = 0; j < kv_lora_dim; j++) { + dotv += low[j] * (float)kv_lora_cache[lora_base + j]; + } + } + const uint64_t rope_base = (uint64_t)row * qk_rope; + for (uint32_t r = 0; r < qk_rope; r += 2u) { + const float2 y = glm_cache_rope_pair_f16_dev( + k_rope_cache, rope_base, r, row, qk_rope, + freq_base, freq_scale, ext_factor, attn_factor, + corr_dims); + dotv += qh[qk_nope + r] * y.x + qh[qk_nope + r + 1u] * y.y; + } + score = dotv * scale; + } + scores[s] = score; + local_max = fmaxf(local_max, score); + } + red[tid] = local_max; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]); + __syncthreads(); + } + const float max_score = red[0]; + __syncthreads(); + + float local_sum = 0.0f; + for (uint32_t s = tid; s < row_count; s += nth) { + const float w = expf(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + __syncthreads(); + } + const float denom = fmaxf(red[0], 1.0e-20f); + __syncthreads(); + + if (lora_vec2) { + for (uint32_t j = tid * 2u; j < kv_lora_dim; j += nth * 2u) { + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t s = 0; s < row_count; s++) { + const uint32_t row = RANGE_TOK2 ? s : selected[s]; + if (row < cache_cap) { + const float2 v = glm_cache_value_pair_dev( + kv_lora_cache + (uint64_t)row * kv_lora_dim + j); + const float w = scores[s]; + acc0 += w * v.x; + acc1 += w * v.y; + } + } + lora_sum[j] = acc0 / denom; + lora_sum[j + 1u] = acc1 / denom; + } + } else { + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint32_t s = 0; s < row_count; s++) { + const uint32_t row = RANGE_TOK2 ? s : selected[s]; + if (row < cache_cap) { + acc += scores[s] * + (float)kv_lora_cache[(uint64_t)row * kv_lora_dim + j]; + } + } + lora_sum[j] = acc / denom; + } + } + __syncthreads(); + + float *out = heads + + ((uint64_t)token * n_head + head) * value_dim; + /* Warp-cooperative value projection: one warp per output dim, two + * lanes per q8_0 block (16 cols each). */ + const uint32_t nwarps = nth >> 5; + const uint32_t warp = tid >> 5; + const uint32_t lane = tid & 31u; + const uint32_t nblk = kv_lora_dim >> 5; + for (uint32_t d = warp; d < value_dim; d += nwarps) { + const char *row = value_weight + + ((uint64_t)head * value_dim + d) * value_row_bytes; + float acc = 0.0f; + for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) { + const char *b = row + (uint64_t)blk * 34u; + const float dscale = __half2float(*(const __half *)b); + const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u; + const float *xs = lora_sum + blk * 32u + (lane & 1u) * 16u; + float s = 0.0f; + #pragma unroll + for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k]; + acc += dscale * s; + } + for (int off = 16; off > 0; off >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, off); + } + if (lane == 0u) out[d] = acc; + } +} + +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) { + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0u || + value_dim == 0) { + return 0; + } + if (value_weight_type != 8u) { /* DS4_TENSOR_Q8_0 */ + fprintf(stderr, + "ds4: glm indexed decode attention: unsupported value type %u\n", + value_weight_type); + return 0; + } + const uint64_t value_row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u; + const uint64_t value_weight_bytes = + (uint64_t)n_head * value_dim * value_row_bytes; + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + return 0; + } + const uint64_t cache_elem = cache_f16 ? 2u : 4u; + if (heads->bytes < (uint64_t)n_head * value_dim * sizeof(float) || + q->bytes < (uint64_t)n_head * qk_dim * sizeof(float) || + qk_low->bytes < (uint64_t)n_head * kv_lora_dim * sizeof(float) || + kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem || + selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const char *vw = cuda_resolve_weight_ptr(model_map, value_weight_offset, + value_weight_bytes, logical_tier, "glm_v_b_decode"); + if (!vw) return 0; + const float scale = 1.0f / sqrtf((float)qk_dim); + const bool score_vec2 = getenv("DS4_GLM_ATTN_NO_SCORE_VEC2") == NULL; + const bool range_tok2 = + g_glm_mtp_verify_mode && + getenv("DS4_GLM_MTP_NO_ATTN_TOK2") == NULL && + n_selected < cache_cap && + heads->bytes >= 2u * (uint64_t)n_head * value_dim * sizeof(float) && + q->bytes >= 2u * (uint64_t)n_head * qk_dim * sizeof(float) && + qk_low->bytes >= + 2u * (uint64_t)n_head * kv_lora_dim * sizeof(float); + if (range_tok2 && n_selected < 512u) { + const bool lora_vec2 = + getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL; + const uint32_t shmem = + (256u + n_selected + 1u + kv_lora_dim) * + (uint32_t)sizeof(float); + const dim3 grid(n_head, 2u, 1u); + if (cache_f16) { + glm_attention_indexed_decode_kernel<__half, true> + <<>>( + (float *)heads->ptr, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + vw, (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, value_dim, + (uint32_t)value_row_bytes, + lora_vec2, score_vec2, scale, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + } else { + glm_attention_indexed_decode_kernel + <<>>( + (float *)heads->ptr, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + vw, (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, value_dim, + (uint32_t)value_row_bytes, + lora_vec2, score_vec2, scale, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + } + return cuda_ok(cudaGetLastError(), + "glm indexed decode attention tok2 range"); + } + if (n_selected >= 512u && + getenv("DS4_GLM_ATTN_NO_STAGED_DECODE") == NULL) { + const uint32_t token_count = range_tok2 ? 2u : 1u; + const uint32_t score_stride = n_selected + (range_tok2 ? 1u : 0u); + const uint64_t head_count = (uint64_t)token_count * n_head; + if (head_count > UINT64_MAX / score_stride || + head_count > UINT64_MAX / kv_lora_dim) { + return 0; + } + const uint64_t score_count = head_count * score_stride; + const uint64_t lora_count = head_count * kv_lora_dim; + if (score_count > UINT64_MAX - head_count - lora_count || + score_count + head_count + lora_count > + UINT64_MAX / sizeof(float)) { + return 0; + } + const uint64_t scratch_bytes = + (score_count + head_count + lora_count) * sizeof(float); + float *scratch = (float *)cuda_tmp_alloc_on( + ds4_tensor_device_idx(heads), scratch_bytes, + "glm staged decode attention"); + if (!scratch) return 0; + float *softmax_denom = scratch + score_count; + float *lora_sum = softmax_denom + head_count; + const uint32_t weight_shmem = + (256u + score_stride) * (uint32_t)sizeof(float); + const dim3 weight_grid(n_head, token_count, 1u); + if (cache_f16 && range_tok2) { + glm_attention_decode_weights_staged_kernel<__half, true> + <<>>( + scratch, softmax_denom, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, scale, n_ctx_orig, freq_base, + freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow, score_vec2); + } else if (cache_f16) { + glm_attention_decode_weights_staged_kernel<__half> + <<>>( + scratch, softmax_denom, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, scale, n_ctx_orig, freq_base, + freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow, score_vec2); + } else if (range_tok2) { + glm_attention_decode_weights_staged_kernel + <<>>( + scratch, softmax_denom, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, scale, n_ctx_orig, freq_base, + freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow, score_vec2); + } else { + glm_attention_decode_weights_staged_kernel + <<>>( + scratch, softmax_denom, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, scale, n_ctx_orig, freq_base, + freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow, score_vec2); + } + if (!cuda_ok(cudaGetLastError(), + "glm staged decode weights launch")) { + return 0; + } + dim3 lora_grid((kv_lora_dim / 2u + 63u) / 64u, + n_head, token_count); + if (cache_f16 && range_tok2) { + glm_attention_decode_lora_staged_kernel<__half, true> + <<>>( + lora_sum, scratch, softmax_denom, + (const __half *)kv_lora_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim); + } else if (cache_f16) { + glm_attention_decode_lora_staged_kernel<__half> + <<>>( + lora_sum, scratch, softmax_denom, + (const __half *)kv_lora_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim); + } else if (range_tok2) { + glm_attention_decode_lora_staged_kernel + <<>>( + lora_sum, scratch, softmax_denom, + (const float *)kv_lora_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim); + } else { + glm_attention_decode_lora_staged_kernel + <<>>( + lora_sum, scratch, softmax_denom, + (const float *)kv_lora_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim); + } + if (!cuda_ok(cudaGetLastError(), + "glm staged decode lora launch")) { + return 0; + } + dim3 value_grid((value_dim + 127u) / 128u, + n_head, token_count); + if (range_tok2) { + glm_attention_decode_value_staged_kernel + <<>>( + (float *)heads->ptr, lora_sum, vw, + n_head, kv_lora_dim, value_dim, + (uint32_t)value_row_bytes); + } else { + glm_attention_decode_value_staged_kernel + <<>>( + (float *)heads->ptr, lora_sum, vw, + n_head, kv_lora_dim, value_dim, + (uint32_t)value_row_bytes); + } + return cuda_ok(cudaGetLastError(), + "glm staged decode value launch"); + } + const bool lora_vec2 = getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL; + const uint32_t shmem = + (256u + n_selected + kv_lora_dim) * (uint32_t)sizeof(float); + if (cache_f16) { + glm_attention_indexed_decode_kernel<__half><<>>( + (float *)heads->ptr, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const __half *)kv_lora_cache->ptr, + (const __half *)k_rope_cache->ptr, + vw, (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes, + lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + } else { + glm_attention_indexed_decode_kernel<<>>( + (float *)heads->ptr, (const float *)q->ptr, + (const float *)qk_low->ptr, + (const float *)kv_lora_cache->ptr, + (const float *)k_rope_cache->ptr, + vw, (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, kv_lora_dim, + qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes, + lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + } + return cuda_ok(cudaGetLastError(), "glm indexed decode attention"); +} + +extern "C" int ds4_gpu_glm_build_kv_cache_flash_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + 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, + bool cache_f16) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_build_kv_cache_flash_tensor\n"); + return 0; +} + +extern "C" int ds4_gpu_glm_build_kv_cache_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + 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, + bool cache_f16) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_build_kv_cache_tensor\n"); + return 0; +} + +__global__ static void glm_fill_selected_range_batch_kernel( + uint32_t *selected, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t pad_row) { + uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t total = n_tokens * n_selected; + if (gid >= total || n_selected == 0u) return; + const uint32_t token = gid / n_selected; + const uint32_t slot = gid - token * n_selected; + const uint32_t visible = pos0 + token + 1u; + selected[gid] = slot < visible ? slot : pad_row; +} + +extern "C" int ds4_gpu_glm_fill_selected_range_batch_tensor( + ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t pad_row) { + if (!selected || n_tokens == 0 || n_selected == 0 || + selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t)) { + return 0; + } + const uint64_t total = (uint64_t)n_tokens * n_selected; + glm_fill_selected_range_batch_kernel<<<(unsigned)((total + 255) / 256), 256>>>( + (uint32_t *)selected->ptr, n_tokens, pos0, n_selected, pad_row); + return cuda_ok(cudaGetLastError(), "glm fill selected batch launch"); +} + +__global__ static void glm_fill_selected_range_kernel( + uint32_t *selected, uint32_t n_selected) { + uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + if (gid < n_selected) selected[gid] = gid; +} + +extern "C" int ds4_gpu_glm_fill_selected_range_tensor( + ds4_gpu_tensor *selected, + uint32_t n_selected) { + if (!selected || n_selected == 0 || + selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { + return 0; + } + glm_fill_selected_range_kernel<<<(n_selected + 255) / 256, 256>>>( + (uint32_t *)selected->ptr, n_selected); + return cuda_ok(cudaGetLastError(), "glm fill selected launch"); +} + +static int glm_rope_tail_offset_launch( + ds4_gpu_tensor *x, + uint32_t n_tokens, uint32_t n_head, uint32_t head_dim, + uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0, + uint32_t n_ctx_orig, float freq_base, float freq_scale, + float ext_factor, float attn_factor, + float beta_fast, float beta_slow, const char *what); + +extern "C" int ds4_gpu_glm_indexer_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim, + rot_dim, 0, pos0, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow, "glm indexer rope tail"); +} + +template +__global__ static void glm_indexer_scores_f32_kernel( + float *scores, + const float *q, + const float *weights, + const CT *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool causal) { + const uint32_t row = blockIdx.x; + const uint32_t token = blockIdx.y; + const uint32_t tid = threadIdx.x; + if (row >= n_rows || token >= n_tokens || tid >= 128u) return; + if (causal && row >= min(n_rows, pos0 + token + 1u)) { + if (tid == 0u) scores[(uint64_t)token * n_rows + row] = -INFINITY; + return; + } + + __shared__ float partial[128]; + float total = 0.0f; + const CT *krow = indexer_key_cache + (uint64_t)row * head_dim; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + + ((uint64_t)token * n_head + h) * head_dim; + float dot = tid < head_dim ? qh[tid] * (float)krow[tid] : 0.0f; + partial[tid] = dot; + __syncthreads(); + for (uint32_t stride = 64u; stride > 0u; stride >>= 1u) { + if (tid < stride) partial[tid] += partial[tid + stride]; + __syncthreads(); + } + if (tid == 0u) { + total += fmaxf(partial[0], 0.0f) * + weights[(uint64_t)token * n_head + h]; + } + __syncthreads(); + } + if (tid == 0u) { + scores[(uint64_t)token * n_rows + row] = total * scale; + } +} + +/* 16-token x 128-row indexer tile. Q and cached K are staged as fp16, + * matching the model's compact-cache precision; each head's MMA result and + * the weighted head reduction remain fp32. */ +template +__global__ static void glm_indexer_scores_wmma128_kernel( + float *scores, + const float *q, + const float *weights, + const CT *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool causal) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + const uint32_t row0 = blockIdx.x * 128u; + const uint32_t token0 = blockIdx.y * 16u; + const uint32_t tid = threadIdx.x; + const uint32_t warp = tid >> 5u; + if (tid >= 256u || head_dim != 128u) return; + + if (causal) { + const uint32_t last_token = min(token0 + 16u, n_tokens); + const uint32_t max_visible = last_token > token0 + ? min(pos0 + last_token, n_rows) : 0u; + if (row0 >= max_visible) { + for (uint32_t i = tid; i < 16u * 128u; i += 256u) { + const uint32_t token = token0 + (i >> 7u); + const uint32_t row = row0 + (i & 127u); + if (token < n_tokens && row < n_rows) { + scores[(uint64_t)token * n_rows + row] = -INFINITY; + } + } + return; + } + } + + __shared__ __half q_sh[16 * 128]; + __shared__ __half k_sh[128 * 128]; + __shared__ float dot_sh[8 * 16 * 16]; + float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint32_t i = tid; i < 128u * 128u; i += 256u) { + const uint32_t r = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t row = row0 + r; + const float v = row < n_rows + ? (float)indexer_key_cache[(uint64_t)row * head_dim + d] + : 0.0f; + k_sh[d + r * 128u] = __float2half(v); + } + __syncthreads(); + + for (uint32_t h = 0; h < n_head; h++) { + for (uint32_t i = tid; i < 16u * 128u; i += 256u) { + const uint32_t tr = i >> 7u; + const uint32_t d = i & 127u; + const uint32_t token = token0 + tr; + const float v = token < n_tokens + ? q[((uint64_t)token * n_head + h) * head_dim + d] + : 0.0f; + q_sh[i] = __float2half(v); + } + __syncthreads(); + + wmma::fragment q_frag; + wmma::fragment k_frag; + wmma::fragment dot_frag; + wmma::fill_fragment(dot_frag, 0.0f); + const uint32_t col0 = warp * 16u; + for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) { + wmma::load_matrix_sync(q_frag, q_sh + k0, 128); + wmma::load_matrix_sync(k_frag, + k_sh + col0 * 128u + k0, 128); + wmma::mma_sync(dot_frag, q_frag, k_frag, dot_frag); + } + wmma::store_matrix_sync(dot_sh + warp * 16u * 16u, + dot_frag, 16, wmma::mem_row_major); + __syncthreads(); + + const uint32_t local0 = tid & 255u; + const uint32_t token = token0 + (local0 >> 4u); + const float w = token < n_tokens + ? weights[(uint64_t)token * n_head + h] : 0.0f; + uint32_t slot = 0; + for (uint32_t i = tid; i < 8u * 16u * 16u; + i += 256u, slot++) { + const uint32_t row = row0 + (i >> 8u) * 16u + (i & 15u); + if (token < n_tokens && row < n_rows) { + acc[slot] += fmaxf(dot_sh[i], 0.0f) * w; + } + } + __syncthreads(); + } + + uint32_t slot = 0; + for (uint32_t i = tid; i < 8u * 16u * 16u; + i += 256u, slot++) { + const uint32_t local = i & 255u; + const uint32_t token = token0 + (local >> 4u); + const uint32_t row = row0 + (i >> 8u) * 16u + (local & 15u); + if (token < n_tokens && row < n_rows) { + float out = acc[slot] * scale; + if (causal && row >= pos0 + token + 1u) out = -INFINITY; + scores[(uint64_t)token * n_rows + row] = out; + } + } +#endif +} + +static int glm_indexer_scores_launch( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16, + bool causal) { + const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float); + if (!scores || !q || !weights || !indexer_key_cache || n_rows == 0u || + n_tokens == 0u || n_head == 0u || head_dim != 128u || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) || + indexer_key_cache->bytes < (uint64_t)n_rows * head_dim * cache_elem || + scores->bytes < (uint64_t)n_tokens * n_rows * sizeof(float)) { + return 0; + } + if (!g_quality_mode) { + dim3 grid((n_rows + 127u) / 128u, + (n_tokens + 15u) / 16u, 1); + if (cache_f16) { + glm_indexer_scores_wmma128_kernel<__half><<>>( + (float *)scores->ptr, (const float *)q->ptr, + (const float *)weights->ptr, + (const __half *)indexer_key_cache->ptr, + n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); + } else { + glm_indexer_scores_wmma128_kernel<<>>( + (float *)scores->ptr, (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)indexer_key_cache->ptr, + n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); + } + return cuda_ok(cudaGetLastError(), "glm indexer scores wmma launch"); + } + + dim3 grid(n_rows, n_tokens, 1); + if (cache_f16) { + glm_indexer_scores_f32_kernel<__half><<>>( + (float *)scores->ptr, (const float *)q->ptr, + (const float *)weights->ptr, + (const __half *)indexer_key_cache->ptr, + n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); + } else { + glm_indexer_scores_f32_kernel<<>>( + (float *)scores->ptr, (const float *)q->ptr, + (const float *)weights->ptr, + (const float *)indexer_key_cache->ptr, + n_rows, n_tokens, pos0, n_head, head_dim, scale, causal); + } + return cuda_ok(cudaGetLastError(), "glm indexer scores f32 launch"); +} + +extern "C" int ds4_gpu_glm_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache, + n_rows, 1u, 0u, n_head, head_dim, + scale, cache_f16, false); +} + +extern "C" int ds4_gpu_glm_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache, + n_rows, n_tokens, pos0, n_head, head_dim, + scale, cache_f16, true); +} + +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) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_k_b_project_typed_tensor\n"); + return 0; +} + +__global__ static void glm_kv_lora_rms_norm_kernel( + float *dst, + const float *src, + const float *w, + uint32_t n_tokens, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + float eps) { + const uint32_t row = blockIdx.x; + if (row >= n_tokens) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + const float *x = src + (uint64_t)row * kv_raw_dim; + float *out = dst + (uint64_t)row * kv_lora_dim; + __shared__ float scratch[256]; + float ss = 0.0f; + for (uint32_t i = tid; i < kv_lora_dim; i += nth) { + const float v = x[i]; + ss += v * v; + } + scratch[tid] = ss; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + __syncthreads(); + } + const float inv = rsqrtf(scratch[0] / (float)kv_lora_dim + eps); + for (uint32_t i = tid; i < kv_lora_dim; i += nth) { + out[i] = x[i] * inv * w[i]; + } +} + +extern "C" int ds4_gpu_glm_kv_lora_rms_norm_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + float eps) { + if (!dst || !src || !model_map || n_tokens == 0 || + kv_lora_dim == 0 || kv_lora_dim > kv_raw_dim) { + return 0; + } + const uint64_t wb = (uint64_t)kv_lora_dim * sizeof(float); + if (weight_offset > model_size || wb > model_size - weight_offset || + src->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) || + dst->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const float *w = (const float *)cuda_resolve_weight_ptr( + model_map, weight_offset, wb, logical_tier, "glm_kv_lora_norm"); + if (!w) return 0; + glm_kv_lora_rms_norm_kernel<<>>( + (float *)dst->ptr, (const float *)src->ptr, w, + n_tokens, kv_raw_dim, kv_lora_dim, eps); + return cuda_ok(cudaGetLastError(), "glm kv lora rms norm launch"); +} + + + +__global__ static void glm_qk_lowrank_q8_0_batch_kernel( + float *qk_low, + const char *weight, + const float *q, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim, + uint64_t row_bytes) { + const uint32_t head = blockIdx.x; + const uint32_t token = blockIdx.y; + if (head >= n_head || token >= n_tokens) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + const float *qh = q + (uint64_t)token * n_head * qk_dim + + (uint64_t)head * qk_dim; + float *out = qk_low + (uint64_t)token * n_head * kv_lora_dim + + (uint64_t)head * kv_lora_dim; + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + const char *row = weight + + ((uint64_t)head * kv_lora_dim + j) * row_bytes; + out[j] = glm_q8_0_dot_row_dev(row, qh, qk_nope); + } +} + +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 (!qk_low || !q || !model_map || n_tokens == 0 || n_head == 0 || + kv_lora_dim == 0 || qk_nope == 0 || (qk_nope & 31u) != 0u) { + return 0; + } + if (weight_type != 8u) { + fprintf(stderr, "ds4: glm qk_lowrank: unsupported type %u\n", + weight_type); + return 0; + } + const uint64_t row_bytes = ((uint64_t)qk_nope / 32u) * 34u; + const uint64_t wbytes = (uint64_t)n_head * kv_lora_dim * row_bytes; + if (weight_offset > model_size || wbytes > model_size - weight_offset || + q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || + qk_low->bytes < + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const char *w = (const char *)cuda_resolve_weight_ptr( + model_map, weight_offset, wbytes, logical_tier, "glm_k_b_qk"); + if (!w) return 0; + if (g_q8_dequant_gemm_enabled && g_cublas_ready && n_tokens >= 128u) { + /* Per-head strided-batched GEMM over a dequantized k_b: the + * per-(token,head) warp kernel was ~65ms/layer at 820 tokens. + * Scratch (executing device): [w_f16][q_f16][out_f32]. */ + const uint64_t wh_bytes = + (uint64_t)n_head * kv_lora_dim * qk_nope * sizeof(__half); + const uint64_t xh_off = (wh_bytes + 255u) & ~255ull; + const uint64_t xh_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(__half); + const uint64_t oo_off = (xh_off + xh_bytes + 255u) & ~255ull; + const uint64_t oo_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + void *tmp = cuda_tmp_alloc_on(logical_tier, oo_off + oo_bytes, + "glm qk_low gemm"); + if (tmp) { + __half *wh = (__half *)tmp; + __half *xh = (__half *)((char *)tmp + xh_off); + float *oo = (float *)((char *)tmp + oo_off); + const uint64_t total_blocks = + (uint64_t)n_head * kv_lora_dim * (qk_nope / 32u); + q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256>>>( + wh, (const unsigned char *)w, total_blocks, + qk_nope / 32u, qk_nope); + const uint64_t xn = (uint64_t)n_tokens * n_head * qk_dim; + f32_to_f16_kernel<<<(xn + 255u) / 256u, 256>>>( + xh, (const float *)q->ptr, xn); + if (cuda_ok(cudaGetLastError(), "glm qk_low gemm staging")) { + const float alpha = 1.0f; + const float beta = 0.0f; + cublasStatus_t st = cublasGemmStridedBatchedEx( + cuda_cublas_for_tier(logical_tier), + CUBLAS_OP_T, CUBLAS_OP_N, + (int)kv_lora_dim, (int)n_tokens, (int)qk_nope, + &alpha, + wh, CUDA_R_16F, (int)qk_nope, + (long long)((uint64_t)kv_lora_dim * qk_nope), + xh, CUDA_R_16F, (int)(n_head * qk_dim), + (long long)qk_dim, + &beta, + oo, CUDA_R_32F, (int)(n_head * kv_lora_dim), + (long long)kv_lora_dim, + (int)n_head, + CUDA_R_32F, CUBLAS_GEMM_DEFAULT); + if (st == CUBLAS_STATUS_SUCCESS && + cuda_ok(cudaMemcpyAsync(qk_low->ptr, oo, oo_bytes, + cudaMemcpyDeviceToDevice, 0), + "glm qk_low gemm out copy")) { + return 1; + } + fprintf(stderr, + "ds4: glm qk_low gemm failed (status %d); native path\n", + (int)st); + } + } + } + dim3 grid(n_head, n_tokens, 1); + glm_qk_lowrank_q8_0_batch_kernel<<>>( + (float *)qk_low->ptr, w, (const float *)q->ptr, + n_tokens, n_head, kv_lora_dim, qk_nope, qk_dim, row_bytes); + return cuda_ok(cudaGetLastError(), "glm qk lowrank batch launch"); +} + +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) { + return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, q, model_map, + model_size, + weight_offset, + weight_type, 1u, n_head, + kv_lora_dim, qk_nope, + qk_dim); +} + +/* Fused decode-path QKV norm + compact-KV store, one block per + * (token, part): part 0 rms-norms q into q_out, part 1 rms-norms + * kv_raw[:kv_lora_dim] into the kv_lora ring, part 2 copies the + * UNROTATED rope tail into the k_rope ring (roped at attention read). */ +__global__ static void glm_qkv_norm_store_compact_kv_kernel( + float *q_dst, + const float *q_src, + const float *q_w, + uint32_t q_n, + const float *kv_raw, + const float *kv_w, + char *kv_lora_cache, + char *k_rope_cache, + uint32_t pos0, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + int cache_f16, + float eps) { + const uint32_t token = blockIdx.x; + const uint32_t part = blockIdx.y; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + const uint32_t pos = pos0 + token; + + if (part == 2u) { + if (pos >= cache_cap) return; + const float *src = kv_raw + (uint64_t)token * kv_raw_dim + kv_lora_dim; + if (cache_f16) { + __half *dst = (__half *)k_rope_cache + (uint64_t)pos * qk_rope; + for (uint32_t i = tid; i < qk_rope; i += nth) { + dst[i] = __float2half(src[i]); + } + } else { + float *dst = (float *)k_rope_cache + (uint64_t)pos * qk_rope; + for (uint32_t i = tid; i < qk_rope; i += nth) { + dst[i] = src[i]; + } + } + return; + } + + const bool kv_task = part != 0u; + const uint32_t n = kv_task ? kv_lora_dim : q_n; + const float *x = kv_task ? kv_raw + (uint64_t)token * kv_raw_dim + : q_src + (uint64_t)token * q_n; + const float *w = kv_task ? kv_w : q_w; + + __shared__ float sh[32]; + float sumf = 0.0f; + for (uint32_t i = tid; i < n; i += nth) { + const float v = x[i]; + sumf += v * v; + } + for (int off = 16; off > 0; off >>= 1) { + sumf += __shfl_xor_sync(0xffffffffu, sumf, off); + } + if ((tid & 31u) == 0u) sh[tid >> 5] = sumf; + __syncthreads(); + if (tid < 32u) { + sumf = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f; + for (int off = 16; off > 0; off >>= 1) { + sumf += __shfl_xor_sync(0xffffffffu, sumf, off); + } + if (tid == 0u) sh[0] = sumf; + } + __syncthreads(); + const float scale = rsqrtf(sh[0] / (float)n + eps); + + if (!kv_task) { + float *y = q_dst + (uint64_t)token * q_n; + for (uint32_t i = tid; i < n; i += nth) { + y[i] = (x[i] * scale) * w[i]; + } + return; + } + + if (pos >= cache_cap) return; + if (cache_f16) { + __half *dst = (__half *)kv_lora_cache + (uint64_t)pos * kv_lora_dim; + for (uint32_t i = tid; i < kv_lora_dim; i += nth) { + dst[i] = __float2half((x[i] * scale) * w[i]); + } + } else { + float *dst = (float *)kv_lora_cache + (uint64_t)pos * kv_lora_dim; + for (uint32_t i = tid; i < kv_lora_dim; i += nth) { + dst[i] = (x[i] * scale) * w[i]; + } + } +} + +extern "C" int ds4_gpu_glm_qkv_norm_store_compact_kv_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_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_raw, + uint64_t kv_weight_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16, + float eps) { + if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || + !model_map || n_tokens == 0 || q_n == 0 || kv_lora_dim == 0 || + qk_rope == 0 || kv_raw_dim < kv_lora_dim + qk_rope || + q->bytes < (uint64_t)n_tokens * q_n * sizeof(float) || + q_out->bytes < (uint64_t)n_tokens * q_n * sizeof(float) || + kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float)) { + return 0; + } + if (q_weight_offset > model_size || + (uint64_t)q_n * sizeof(float) > model_size - q_weight_offset || + kv_weight_offset > model_size || + (uint64_t)kv_lora_dim * sizeof(float) > model_size - kv_weight_offset) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const float *q_w = (const float *)cuda_resolve_weight_ptr( + model_map, q_weight_offset, (uint64_t)q_n * sizeof(float), + logical_tier, "glm_q_norm"); + const float *kv_w = (const float *)cuda_resolve_weight_ptr( + model_map, kv_weight_offset, + (uint64_t)kv_lora_dim * sizeof(float), + logical_tier, "glm_kv_norm"); + if (!q_w || !kv_w) return 0; + dim3 grid(n_tokens, 3, 1); + glm_qkv_norm_store_compact_kv_kernel<<>>( + (float *)q_out->ptr, + (const float *)q->ptr, + q_w, + q_n, + (const float *)kv_raw->ptr, + kv_w, + (char *)kv_lora_cache->ptr, + (char *)k_rope_cache->ptr, + pos0, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope, + cache_f16 ? 1 : 0, eps); + return cuda_ok(cudaGetLastError(), "glm qkv norm store compact kv"); +} + +/* In-place interleaved-pair yarn rope on a [n_tokens][n_head][head_dim] + * f32 tensor, rotating rot_dim dims starting at rot_offset. Shared by the + * attention q tail (offset = head_dim - rot_dim) and the DSA indexer + * (offset = 0). Grid (n_head, n_tokens). */ +__global__ static void glm_rope_tail_offset_kernel( + float *x, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t rot_offset, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint32_t head = blockIdx.x; + const uint32_t token = blockIdx.y; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + if (rot_dim == 0u || (rot_dim & 1u) != 0u || + rot_offset > head_dim || rot_dim > head_dim - rot_offset) return; + + const uint32_t pos = pos0 + token; + float *row = x + ((uint64_t)token * n_head + head) * head_dim + rot_offset; + + float corr_dims[2] = {0.0f, 0.0f}; + if (ext_factor != 0.0f) { + corr_dims[0] = fmaxf(0.0f, + floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim, + (int)n_ctx_orig, beta_fast, freq_base))); + corr_dims[1] = fminf((float)rot_dim - 1.0f, + ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim, + (int)n_ctx_orig, beta_slow, freq_base))); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)rot_dim; + for (uint32_t i = tid * 2u; i < rot_dim; i += nth * 2u) { + const float theta = + theta_base * powf(freq_base, inv_ndims * (float)i); + float ct, st; + glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i, + ext_factor, attn_factor, &ct, &st); + const float x0 = row[i]; + const float x1 = row[i + 1u]; + row[i] = x0 * ct - x1 * st; + row[i + 1u] = x0 * st + x1 * ct; + } +} + +static int glm_rope_tail_offset_launch( + ds4_gpu_tensor *x, + uint32_t n_tokens, uint32_t n_head, uint32_t head_dim, + uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0, + uint32_t n_ctx_orig, float freq_base, float freq_scale, + float ext_factor, float attn_factor, + float beta_fast, float beta_slow, const char *what) { + if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || + rot_dim == 0 || (rot_dim & 1u) != 0u || rot_offset > head_dim || + rot_dim > head_dim - rot_offset || + x->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float)) { + return 0; + } + dim3 grid(n_head, n_tokens, 1); + glm_rope_tail_offset_kernel<<>>( + (float *)x->ptr, n_head, head_dim, rot_dim, rot_offset, + pos0, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), what); +} + +extern "C" int ds4_gpu_glm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (rot_dim > head_dim) return 0; + return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim, + rot_dim, head_dim - rot_dim, pos0, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow, "glm rope tail"); +} + +extern "C" int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor\n"); + return 0; +} + +/* Scalar-correct GLM routed MoE (q2_K experts): per (token, slot) block + * quantizes nothing - dots q2_K rows against a q8_K-quantized activation + * staged in shared memory. Grid: (n_tokens, n_expert). Mid buffer holds + * silu(gate)*up per slot; out accumulates expert_weight-scaled down rows. + */ +__global__ static void glm_routed_moe_batch_q2K_gateup_kernel( + float *mid, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t n_tokens, + uint32_t mid_token_stride) { + const uint32_t tok = blockIdx.x; + const uint32_t slot = blockIdx.y; + if (tok >= n_tokens || slot >= n_expert) return; + const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; + if (expert < 0) return; + const cuda_block_q8_K *xrow = xq + (uint64_t)tok * xq_blocks; + float *mrow = mid + (uint64_t)tok * mid_token_stride + + (uint64_t)slot * expert_mid_dim; + for (uint32_t r = threadIdx.x; r < expert_mid_dim; r += blockDim.x) { + const cuda_block_q2_K *gr = (const cuda_block_q2_K *)(gate_base + + (uint64_t)expert * gate_expert_bytes + (uint64_t)r * gate_row_bytes); + const cuda_block_q2_K *ur = (const cuda_block_q2_K *)(up_base + + (uint64_t)expert * up_expert_bytes + (uint64_t)r * up_row_bytes); + float g = 0.0f, u = 0.0f; + for (uint32_t b = 0; b < xq_blocks; b++) { + g += dev_dot_q2_K_q8_K_block(gr + b, xrow + b); + u += dev_dot_q2_K_q8_K_block(ur + b, xrow + b); + } + mrow[r] = (g / (1.0f + expf(-g))) * u; /* silu(g)*u */ + } +} + +__global__ static void glm_routed_moe_batch_q2K_down_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + const float *weights, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.y; + if (tok >= n_tokens) return; + const uint32_t r = blockIdx.x * blockDim.x + threadIdx.x; + if (r >= out_dim) return; + float acc = 0.0f; + for (uint32_t slot = 0; slot < n_expert; slot++) { + const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; + if (expert < 0) continue; + const float w = weights[(uint64_t)tok * n_expert + slot]; + const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base + + (uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes); + const cuda_block_q8_K *mrow = midq + + ((uint64_t)tok * n_expert + slot) * midq_blocks; + float s = 0.0f; + for (uint32_t b = 0; b < midq_blocks; b++) { + s += dev_dot_q2_K_q8_K_block(dr + b, mrow + b); + } + acc += w * s; + } + out[(uint64_t)tok * out_dim + r] = acc; +} + +/* Warp-per-row routed MoE (q2_K x q8_K). Each block stages the token's + * q8_K activation row in shared memory; one warp produces one mid row + * (gate dot + up dot + silu*mul fused). Grid: + * (expert_mid_dim/warps, n_expert, n_tokens). */ +__global__ static void glm_routed_moe_gateup_warp_kernel( + float *mid, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.z; + const uint32_t slot = blockIdx.y; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + if (tok >= n_tokens || slot >= n_expert) return; + const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; + + extern __shared__ unsigned int glm_moe_sh_u32[]; + { + const unsigned int *src = + (const unsigned int *)(xq + (uint64_t)tok * xq_blocks); + const uint32_t words = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u; + for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) { + glm_moe_sh_u32[i] = src[i]; + } + } + __syncthreads(); + if (expert < 0) return; + const cuda_block_q8_K *xrow = (const cuda_block_q8_K *)glm_moe_sh_u32; + + const uint32_t r = blockIdx.x * warps + warp; + if (r >= expert_mid_dim) return; + const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + + (uint64_t)r * gate_row_bytes; + const char *ur = up_base + (uint64_t)expert * up_expert_bytes + + (uint64_t)r * up_row_bytes; + float g = 0.0f, u = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + g += dev_dot_q2_K_q8_K_block( + (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b); + u += dev_dot_q2_K_q8_K_block( + (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b); + } + for (int off = 16; off > 0; off >>= 1) { + g += __shfl_down_sync(0xffffffffu, g, off); + u += __shfl_down_sync(0xffffffffu, u, off); + } + if (lane == 0u) { + mid[((uint64_t)tok * n_expert + slot) * expert_mid_dim + r] = + (g / (1.0f + expf(-g))) * u; + } +} + +/* Exact two-token gate/up with adjacent-token expert reuse. Token 0 owns an + * expert present in both rows; token 1 only launches work for unmatched + * experts. Each token keeps the native lane assignment and warp reduction. */ +__global__ static void glm_routed_moe_gateup_tok2_reuse_kernel( + float *mid, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *selected, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert) { + const uint32_t owner = blockIdx.y; + const uint32_t tok = owner / n_expert; + const uint32_t slot = owner - tok * n_expert; + const uint32_t warps = blockDim.x >> 5u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t lane = threadIdx.x & 31u; + if (tok >= 2u || slot >= n_expert) return; + + const int32_t expert = selected[owner]; + if (expert < 0) return; + int32_t mate_slot = -1; + for (uint32_t s = 0; s < n_expert; s++) { + if (selected[(uint64_t)(1u - tok) * n_expert + s] == expert) { + mate_slot = (int32_t)s; + break; + } + } + if (tok == 1u && mate_slot >= 0) return; + + const uint32_t np = tok == 0u && mate_slot >= 0 ? 2u : 1u; + const uint32_t pair0 = owner; + const uint32_t pair1 = n_expert + (uint32_t)mate_slot; + extern __shared__ unsigned int glm_moe_tok2_sh_u32[]; + const uint32_t words_per_row = + xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u; + const unsigned int *src0 = (const unsigned int *)( + xq + (uint64_t)tok * xq_blocks); + for (uint32_t i = threadIdx.x; i < words_per_row; + i += blockDim.x) { + glm_moe_tok2_sh_u32[i] = src0[i]; + } + if (np == 2u) { + const unsigned int *src1 = + (const unsigned int *)(xq + xq_blocks); + for (uint32_t i = threadIdx.x; i < words_per_row; + i += blockDim.x) { + glm_moe_tok2_sh_u32[words_per_row + i] = src1[i]; + } + } + __syncthreads(); + + const uint32_t r = blockIdx.x * warps + warp; + if (r >= expert_mid_dim) return; + const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + + (uint64_t)r * gate_row_bytes; + const char *ur = up_base + (uint64_t)expert * up_expert_bytes + + (uint64_t)r * up_row_bytes; + const cuda_block_q8_K *x0 = + (const cuda_block_q8_K *)glm_moe_tok2_sh_u32; + const cuda_block_q8_K *x1 = np == 2u + ? x0 + xq_blocks + : NULL; + float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + dev_dot_q2_K_q8_K_block8( + (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), + x0 + b, np == 2u ? x1 + b : NULL, + NULL, NULL, NULL, NULL, NULL, NULL, np, g); + dev_dot_q2_K_q8_K_block8( + (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), + x0 + b, np == 2u ? x1 + b : NULL, + NULL, NULL, NULL, NULL, NULL, NULL, np, u); + } + for (uint32_t p = 0; p < np; p++) { + for (int off = 16; off > 0; off >>= 1) { + g[p] += __shfl_down_sync(0xffffffffu, g[p], off); + u[p] += __shfl_down_sync(0xffffffffu, u[p], off); + } + if (lane == 0u) { + const uint32_t pair = p == 0u ? pair0 : pair1; + mid[(uint64_t)pair * expert_mid_dim + r] = + (g[p] / (1.0f + expf(-g[p]))) * u[p]; + } + } +} + +/* Warp-per-output-row down projection: stages all n_expert quantized mid + * rows for the token in shared memory, each warp accumulates one out row + * across every selected expert. Grid: (out_dim/warps, n_tokens). */ +__global__ static void glm_routed_moe_down_warp_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + const float *weights, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.y; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + if (tok >= n_tokens) return; + + extern __shared__ unsigned int glm_moe_sh_u32[]; + { + const unsigned int *src = (const unsigned int *) + (midq + (uint64_t)tok * n_expert * midq_blocks); + const uint32_t words = n_expert * midq_blocks * + (uint32_t)sizeof(cuda_block_q8_K) / 4u; + for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) { + glm_moe_sh_u32[i] = src[i]; + } + } + __syncthreads(); + const cuda_block_q8_K *msh = (const cuda_block_q8_K *)glm_moe_sh_u32; + + const uint32_t r = blockIdx.x * warps + warp; + if (r >= out_dim) return; + const uint32_t units = n_expert * midq_blocks; + float acc = 0.0f; + for (uint32_t idx = lane; idx < units; idx += 32u) { + const uint32_t slot = idx / midq_blocks; + const uint32_t b = idx - slot * midq_blocks; + const int32_t expert = selected[(uint64_t)tok * n_expert + slot]; + if (expert < 0) continue; + const float w = weights[(uint64_t)tok * n_expert + slot]; + const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base + + (uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes); + acc += w * dev_dot_q2_K_q8_K_block(dr + b, msh + slot * midq_blocks + b); + } + for (int off = 16; off > 0; off >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, off); + } + if (lane == 0u) out[(uint64_t)tok * out_dim + r] = acc; +} + +/* Expert-major routed MoE for prefill: build per-expert token lists, + * then walk rows expert-by-expert so weights stream once per layer and + * activations hit L2. pair = tok * n_expert + slot indexes selected/ + * weights/mid rows directly. */ +__global__ static void glm_moe_expert_map_kernel( + int32_t *counts, + int32_t *lists, + const int32_t *selected, + uint32_t n_pairs, + uint32_t n_total_expert, + uint32_t cap, + uint32_t pair_base) { + const uint32_t p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= n_pairs) return; + const uint32_t pair = pair_base + p; + const int32_t e = selected[pair]; + if (e < 0 || (uint32_t)e >= n_total_expert) return; + const int32_t idx = atomicAdd(&counts[e], 1); + lists[(uint64_t)e * cap + idx] = (int32_t)pair; +} + +__global__ static void glm_moe_build_expert_tiles8_kernel( + uint32_t *tile_total, + uint32_t *tile_experts, + uint32_t *tile_starts, + const int32_t *counts, + uint32_t n_total_expert) { + if (blockIdx.x != 0u || threadIdx.x != 0u) return; + uint32_t total = 0; + for (uint32_t e = 0; e < n_total_expert; e++) { + const uint32_t count = counts[e] > 0 ? (uint32_t)counts[e] : 0u; + const uint32_t ntiles = (count + 7u) / 8u; + for (uint32_t t = 0; t < ntiles; t++) { + tile_experts[total] = e; + tile_starts[total] = t * 8u; + total++; + } + } + *tile_total = total; +} + +/* Expert-tiled Q2_K gate/up for GLM prefill. One warp keeps the same + * block-to-lane assignment and reduction tree as the token-major W32 + * kernel, but evaluates eight pairs against each loaded expert row. */ +__global__ static void glm_routed_moe_gateup_expert_tile8_kernel( + float *mid, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *counts, + const int32_t *lists, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t cap) { + const uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t warps = blockDim.x >> 5u; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t expert = tile_experts[tile]; + const uint32_t local_start = tile_starts[tile]; + const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u; + + __shared__ uint32_t pair[8]; + __shared__ uint32_t tok[8]; + __shared__ uint32_t np; + if (threadIdx.x == 0u) { + uint32_t n = count - local_start; + if (n > 8u) n = 8u; + np = n; + for (uint32_t p = 0; p < n; p++) { + const uint32_t pr = + (uint32_t)lists[(uint64_t)expert * cap + local_start + p]; + pair[p] = pr; + tok[p] = pr / n_expert; + } + } + __syncthreads(); + + const uint32_t r = blockIdx.x * warps + warp; + if (r >= expert_mid_dim) return; + const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes + + (uint64_t)r * gate_row_bytes; + const char *ur = up_base + (uint64_t)expert * up_expert_bytes + + (uint64_t)r * up_row_bytes; + float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + const cuda_block_q8_K *x0 = np > 0u ? xq + (uint64_t)tok[0] * xq_blocks + b : NULL; + const cuda_block_q8_K *x1 = np > 1u ? xq + (uint64_t)tok[1] * xq_blocks + b : NULL; + const cuda_block_q8_K *x2 = np > 2u ? xq + (uint64_t)tok[2] * xq_blocks + b : NULL; + const cuda_block_q8_K *x3 = np > 3u ? xq + (uint64_t)tok[3] * xq_blocks + b : NULL; + const cuda_block_q8_K *x4 = np > 4u ? xq + (uint64_t)tok[4] * xq_blocks + b : NULL; + const cuda_block_q8_K *x5 = np > 5u ? xq + (uint64_t)tok[5] * xq_blocks + b : NULL; + const cuda_block_q8_K *x6 = np > 6u ? xq + (uint64_t)tok[6] * xq_blocks + b : NULL; + const cuda_block_q8_K *x7 = np > 7u ? xq + (uint64_t)tok[7] * xq_blocks + b : NULL; + dev_dot_q2_K_q8_K_block8( + (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), + x0, x1, x2, x3, x4, x5, x6, x7, np, g); + dev_dot_q2_K_q8_K_block8( + (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), + x0, x1, x2, x3, x4, x5, x6, x7, np, u); + } + for (uint32_t p = 0; p < np; p++) { + for (int off = 16; off > 0; off >>= 1) { + g[p] += __shfl_down_sync(0xffffffffu, g[p], off); + u[p] += __shfl_down_sync(0xffffffffu, u[p], off); + } + if (lane == 0u) { + mid[(uint64_t)pair[p] * expert_mid_dim + r] = + (g[p] / (1.0f + expf(-g[p]))) * u[p]; + } + } +} + +__global__ static void glm_routed_moe_gateup_expert_kernel( + float *mid, + const char *gate_base, + const char *up_base, + const cuda_block_q8_K *xq, + const int32_t *counts, + const int32_t *lists, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t cap) { + const uint32_t e = blockIdx.y; + const int32_t nt = counts[e]; + if (nt == 0) return; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t r = blockIdx.x * warps + warp; + if (r >= expert_mid_dim) return; + const char *gr = gate_base + (uint64_t)e * gate_expert_bytes + + (uint64_t)r * gate_row_bytes; + const char *ur = up_base + (uint64_t)e * up_expert_bytes + + (uint64_t)r * up_row_bytes; + const int32_t *lst = lists + (uint64_t)e * cap; + for (int32_t i = 0; i < nt; i++) { + const uint32_t pair = (uint32_t)lst[i]; + const cuda_block_q8_K *xrow = xq + (uint64_t)(pair / n_expert) * xq_blocks; + float g = 0.0f, u = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 32u) { + g += dev_dot_q2_K_q8_K_block( + (const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b); + u += dev_dot_q2_K_q8_K_block( + (const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b); + } + for (int off = 16; off > 0; off >>= 1) { + g += __shfl_down_sync(0xffffffffu, g, off); + u += __shfl_down_sync(0xffffffffu, u, off); + } + if (lane == 0u) { + mid[(uint64_t)pair * expert_mid_dim + r] = + (g / (1.0f + expf(-g))) * u; + } + } +} + +__global__ static void glm_routed_moe_down_expert_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *counts, + const int32_t *lists, + const float *weights, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t cap) { + const uint32_t e = blockIdx.y; + const int32_t nt = counts[e]; + if (nt == 0) return; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t r = blockIdx.x * warps + warp; + if (r >= out_dim) return; + const char *dr = down_base + (uint64_t)e * down_expert_bytes + + (uint64_t)r * down_row_bytes; + const int32_t *lst = lists + (uint64_t)e * cap; + for (int32_t i = 0; i < nt; i++) { + const uint32_t pair = (uint32_t)lst[i]; + const cuda_block_q8_K *mrow = midq + (uint64_t)pair * midq_blocks; + float s = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 32u) { + s += dev_dot_q2_K_q8_K_block( + (const cuda_block_q2_K *)(dr + (uint64_t)b * 84u), mrow + b); + } + for (int off = 16; off > 0; off >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, off); + } + if (lane == 0u) { + atomicAdd(&out[(uint64_t)(pair / n_expert) * out_dim + r], + weights[pair] * s); + } + } +} + +/* Expert-tiled down projection with an exact token-major reduction. The + * first kernel reuses each Q2_K row across eight routed pairs, but materializes + * each block dot. The second kernel applies the router weight and consumes the + * dots with the same lane assignment and warp tree as the native kernel. */ +__global__ static void glm_routed_moe_down_expert_tile8_terms_kernel( + float *terms, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *counts, + const int32_t *lists, + const uint32_t *tile_total, + const uint32_t *tile_experts, + const uint32_t *tile_starts, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t cap, + uint32_t pair_base) { + const uint32_t tile = blockIdx.y; + if (tile >= *tile_total) return; + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t expert = tile_experts[tile]; + const uint32_t local_start = tile_starts[tile]; + const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u; + + __shared__ uint32_t pair[8]; + __shared__ uint32_t np; + __shared__ cuda_block_q8_K mq[8][8]; + if (threadIdx.x == 0u) { + uint32_t n = count - local_start; + if (n > 8u) n = 8u; + np = n; + for (uint32_t p = 0; p < n; p++) { + pair[p] = (uint32_t)lists[ + (uint64_t)expert * cap + local_start + p]; + } + } + __syncthreads(); + for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) { + const uint32_t p = i / midq_blocks; + const uint32_t b = i - p * midq_blocks; + mq[p][b] = midq[(uint64_t)pair[p] * midq_blocks + b]; + } + __syncthreads(); + if (row >= out_dim || lane >= midq_blocks) return; + + const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + + (uint64_t)expert * down_expert_bytes + + (uint64_t)row * down_row_bytes); + for (uint32_t p = 0; p < np; p++) { + const uint32_t pr = pair[p]; + const float dot = dev_dot_q2_K_q8_K_block(wr + lane, &mq[p][lane]); + terms[((uint64_t)(pr - pair_base) * out_dim + row) * + midq_blocks + lane] = + dot; + } +} + +__global__ static void glm_routed_moe_down_terms_reduce_kernel( + float *out, + const float *terms, + const int32_t *selected, + const float *weights, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.y; + const uint32_t warp = threadIdx.x >> 5u; + const uint32_t lane = threadIdx.x & 31u; + const uint32_t row = blockIdx.x * 8u + warp; + if (tok >= n_tokens || row >= out_dim) return; + const uint32_t units = n_expert * midq_blocks; + float acc = 0.0f; + for (uint32_t idx = lane; idx < units; idx += 32u) { + const uint32_t slot = idx / midq_blocks; + const uint32_t b = idx - slot * midq_blocks; + const uint32_t pr = tok * n_expert + slot; + if (selected[pr] >= 0) { + acc += weights[pr] * + terms[((uint64_t)pr * out_dim + row) * midq_blocks + b]; + } + } + for (int off = 16; off > 0; off >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, off); + } + if (lane == 0u) out[(uint64_t)tok * out_dim + row] = acc; +} + +static int glm_routed_moe_finish_batch( + ds4_gpu_tensor *out, + float *out_work, + uint64_t out_bytes, + const char *what) { + if (!cuda_ok(cudaGetLastError(), what)) return 0; + if (out_work == (float *)out->ptr) return 1; + return cuda_ok(cudaMemcpyAsync(out->ptr, out_work, out_bytes, + cudaMemcpyDeviceToDevice, 0), + "glm routed moe local output copy"); +} + +extern "C" int ds4_gpu_glm_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride) { + (void)layer_index; (void)n_total_expert; + if (!out || !mid || !x || !selected || !weights || !model_map || + n_tokens == 0 || n_expert == 0 || + (expert_in_dim & 255u) != 0u || (expert_mid_dim & 255u) != 0u) { + return 0; + } + if (gate_type != 10u || up_type != 10u || down_type != 10u) { + fprintf(stderr, "ds4: glm routed moe: unsupported types %u/%u/%u\n", + gate_type, up_type, down_type); + return 0; + } + if (mid_token_stride != n_expert * expert_mid_dim) { + fprintf(stderr, + "ds4: glm routed moe: mid stride %u != %u (packed rows expected)\n", + mid_token_stride, n_expert * expert_mid_dim); + return 0; + } + const int logical_tier = cuda_current_tier(); + const char *gw = (const char *)cuda_resolve_weight_ptr(model_map, + gate_offset, (uint64_t)256 * gate_expert_bytes, logical_tier, + "glm_gate_exps"); + const char *uw = (const char *)cuda_resolve_weight_ptr(model_map, + up_offset, (uint64_t)256 * up_expert_bytes, logical_tier, + "glm_up_exps"); + const char *dw = (const char *)cuda_resolve_weight_ptr(model_map, + down_offset, (uint64_t)256 * down_expert_bytes, logical_tier, + "glm_down_exps"); + if (!gw || !uw || !dw) return 0; + + /* Stage 1: quantize x rows to q8_K (existing kernel). */ + const uint32_t xq_blocks = expert_in_dim / 256u; + const uint32_t midq_blocks = expert_mid_dim / 256u; + static ds4_gpu_tensor *xq_scratch[DS4_MAX_GPUS] = {0}; + static ds4_gpu_tensor *midq_scratch[DS4_MAX_GPUS] = {0}; + int dev = logical_tier; + const int scratch_tier = getenv("DS4_GLM_MOE_SCRATCH_TIER0") ? 0 : dev; + const uint64_t xq_bytes = (uint64_t)n_tokens * xq_blocks * + sizeof(cuda_block_q8_K); + const uint64_t midq_bytes = (uint64_t)n_tokens * n_expert * midq_blocks * + sizeof(cuda_block_q8_K); + if (!xq_scratch[dev] || xq_scratch[dev]->bytes < xq_bytes) { + if (xq_scratch[dev]) ds4_gpu_tensor_free(xq_scratch[dev]); + xq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, xq_bytes); + } + if (!midq_scratch[dev] || midq_scratch[dev]->bytes < midq_bytes) { + if (midq_scratch[dev]) ds4_gpu_tensor_free(midq_scratch[dev]); + midq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, midq_bytes); + } + if (!xq_scratch[dev] || !midq_scratch[dev]) return 0; + + static ds4_gpu_tensor *mid_local[DS4_MAX_GPUS] = {0}; + static ds4_gpu_tensor *out_local[DS4_MAX_GPUS] = {0}; + const uint64_t mid_work_bytes = + (uint64_t)n_tokens * mid_token_stride * sizeof(float); + const uint64_t out_work_bytes = + (uint64_t)n_tokens * out_dim * sizeof(float); + float *mid_work = (float *)mid->ptr; + float *out_work = (float *)out->ptr; + const bool use_local_batch_io = + n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_LOCAL_BATCH_IO"); + if (use_local_batch_io && ds4_tensor_device_idx(mid) != dev) { + if (!mid_local[dev] || mid_local[dev]->bytes < mid_work_bytes) { + if (mid_local[dev]) ds4_gpu_tensor_free(mid_local[dev]); + mid_local[dev] = + ds4_gpu_tensor_alloc_ptr_on(dev, mid_work_bytes); + } + if (!mid_local[dev]) return 0; + mid_work = (float *)mid_local[dev]->ptr; + } + if (use_local_batch_io && ds4_tensor_device_idx(out) != dev) { + if (!out_local[dev] || out_local[dev]->bytes < out_work_bytes) { + if (out_local[dev]) ds4_gpu_tensor_free(out_local[dev]); + out_local[dev] = + ds4_gpu_tensor_alloc_ptr_on(dev, out_work_bytes); + } + if (!out_local[dev]) return 0; + out_work = (float *)out_local[dev]->ptr; + } + { + dim3 gq(xq_blocks, n_tokens, 1); + q8_K_quantize_kernel<<>>( + (cuda_block_q8_K *)xq_scratch[dev]->ptr, + (const float *)x->ptr, expert_in_dim, n_tokens); + } + + static ds4_gpu_tensor *map_scratch[DS4_MAX_GPUS] = {0}; + static ds4_gpu_tensor *down_terms_scratch[DS4_MAX_GPUS] = {0}; + const bool use_expert_tile8 = + n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_EXPERT_TILE8"); + const bool use_expert_major = + n_tokens >= 16u && getenv("DS4_GLM_MOE_EXPERT_MAJOR"); + if (use_expert_tile8 || use_expert_major) { + const uint32_t cap = n_tokens; + const uint32_t n_pairs = n_tokens * n_expert; + const uint64_t counts_bytes = 256u * sizeof(int32_t); + const uint64_t lists_off = (counts_bytes + 255u) & ~255ull; + const uint64_t lists_bytes = + (uint64_t)256u * cap * sizeof(int32_t); + const uint32_t tile_capacity = + (n_pairs + 7u) / 8u + 256u; + const uint64_t tile_total_off = + (lists_off + lists_bytes + 255u) & ~255ull; + const uint64_t tile_experts_off = + (tile_total_off + sizeof(uint32_t) + 255u) & ~255ull; + const uint64_t tile_starts_off = + tile_experts_off + (uint64_t)tile_capacity * sizeof(uint32_t); + const uint64_t map_bytes = use_expert_tile8 + ? tile_starts_off + (uint64_t)tile_capacity * sizeof(uint32_t) + : lists_off + lists_bytes; + if (!map_scratch[dev] || map_scratch[dev]->bytes < map_bytes) { + if (map_scratch[dev]) ds4_gpu_tensor_free(map_scratch[dev]); + map_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(dev, map_bytes); + } + if (map_scratch[dev]) { + int32_t *counts = (int32_t *)map_scratch[dev]->ptr; + int32_t *lists = (int32_t *)((char *)map_scratch[dev]->ptr + lists_off); + cudaMemsetAsync(counts, 0, counts_bytes); + glm_moe_expert_map_kernel<<<(n_pairs + 255u) / 256u, 256>>>( + counts, lists, (const int32_t *)selected->ptr, + n_pairs, 256u, cap, 0u); + if (use_expert_tile8) { + uint32_t *tile_total = (uint32_t *)( + (char *)map_scratch[dev]->ptr + tile_total_off); + uint32_t *tile_experts = (uint32_t *)( + (char *)map_scratch[dev]->ptr + tile_experts_off); + uint32_t *tile_starts = (uint32_t *)( + (char *)map_scratch[dev]->ptr + tile_starts_off); + glm_moe_build_expert_tiles8_kernel<<<1, 1>>>( + tile_total, tile_experts, tile_starts, + counts, 256u); + dim3 ge1((expert_mid_dim + 7u) / 8u, + tile_capacity, 1); + glm_routed_moe_gateup_expert_tile8_kernel<<>>( + mid_work, gw, uw, + (const cuda_block_q8_K *)xq_scratch[dev]->ptr, + counts, lists, + tile_total, tile_experts, tile_starts, + gate_expert_bytes, gate_row_bytes, + up_expert_bytes, up_row_bytes, + xq_blocks, expert_mid_dim, n_expert, cap); + q8_K_quantize_kernel<<< + dim3(midq_blocks, n_tokens * n_expert, 1), 256>>>( + (cuda_block_q8_K *)midq_scratch[dev]->ptr, + mid_work, + expert_mid_dim, n_tokens * n_expert); + if (getenv("DS4_GLM_MOE_NO_DOWN_TILE8_EXACT") == NULL) { + const uint32_t max_chunk_tokens = 512u; + const uint32_t scratch_tokens = + n_tokens < max_chunk_tokens ? n_tokens : max_chunk_tokens; + const uint64_t term_count = + (uint64_t)scratch_tokens * n_expert * + out_dim * midq_blocks; + const uint64_t term_bytes = term_count * sizeof(float); + if (!down_terms_scratch[dev] || + down_terms_scratch[dev]->bytes < term_bytes) { + if (down_terms_scratch[dev]) { + ds4_gpu_tensor_free(down_terms_scratch[dev]); + } + down_terms_scratch[dev] = + ds4_gpu_tensor_alloc_ptr_on(dev, term_bytes); + } + if (down_terms_scratch[dev]) { + for (uint32_t token0 = 0; token0 < n_tokens; + token0 += max_chunk_tokens) { + uint32_t chunk_tokens = n_tokens - token0; + if (chunk_tokens > max_chunk_tokens) { + chunk_tokens = max_chunk_tokens; + } + const uint32_t pair_base = token0 * n_expert; + const uint32_t chunk_pairs = + chunk_tokens * n_expert; + uint32_t chunk_tile_capacity = tile_capacity; + if (n_tokens > max_chunk_tokens) { + chunk_tile_capacity = + (chunk_pairs + 7u) / 8u + 256u; + cudaMemsetAsync(counts, 0, counts_bytes); + glm_moe_expert_map_kernel<<< + (chunk_pairs + 255u) / 256u, 256>>>( + counts, lists, + (const int32_t *)selected->ptr, + chunk_pairs, 256u, cap, pair_base); + glm_moe_build_expert_tiles8_kernel<<<1, 1>>>( + tile_total, tile_experts, tile_starts, + counts, 256u); + } + dim3 gd1((out_dim + 31u) / 32u, + chunk_tile_capacity, 1); + glm_routed_moe_down_expert_tile8_terms_kernel<<< + gd1, 256>>>( + (float *)down_terms_scratch[dev]->ptr, + dw, + (const cuda_block_q8_K *)midq_scratch[dev]->ptr, + counts, lists, + tile_total, tile_experts, tile_starts, + down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, cap, + pair_base); + dim3 gd2((out_dim + 7u) / 8u, + chunk_tokens, 1); + glm_routed_moe_down_terms_reduce_kernel<<< + gd2, 256>>>( + out_work + (uint64_t)token0 * out_dim, + (const float *)down_terms_scratch[dev]->ptr, + (const int32_t *)selected->ptr + pair_base, + (const float *)weights->ptr + pair_base, + midq_blocks, out_dim, n_expert, + chunk_tokens); + } + return glm_routed_moe_finish_batch( + out, out_work, out_work_bytes, + "glm routed moe exact down tile8"); + } + } + const uint32_t warps = 8u; + dim3 ge2((out_dim + warps - 1u) / warps, + n_tokens, 1); + const uint32_t sh2 = n_expert * midq_blocks * + (uint32_t)sizeof(cuda_block_q8_K); + glm_routed_moe_down_warp_kernel<<< + ge2, warps * 32u, sh2>>>( + out_work, dw, + (const cuda_block_q8_K *)midq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, n_tokens); + return glm_routed_moe_finish_batch( + out, out_work, out_work_bytes, + "glm routed moe expert tile8"); + } + dim3 ge1((expert_mid_dim + 7u) / 8u, 256u, 1); + glm_routed_moe_gateup_expert_kernel<<>>( + mid_work, gw, uw, + (const cuda_block_q8_K *)xq_scratch[dev]->ptr, + counts, lists, + gate_expert_bytes, gate_row_bytes, + up_expert_bytes, up_row_bytes, + xq_blocks, expert_mid_dim, n_expert, cap); + q8_K_quantize_kernel<<>>( + (cuda_block_q8_K *)midq_scratch[dev]->ptr, + mid_work, expert_mid_dim, n_tokens * n_expert); + cudaMemsetAsync(out_work, 0, + (uint64_t)n_tokens * out_dim * sizeof(float)); + dim3 ge2((out_dim + 7u) / 8u, 256u, 1); + glm_routed_moe_down_expert_kernel<<>>( + out_work, dw, + (const cuda_block_q8_K *)midq_scratch[dev]->ptr, + counts, lists, (const float *)weights->ptr, + down_expert_bytes, down_row_bytes, + midq_blocks, out_dim, n_expert, cap); + return glm_routed_moe_finish_batch( + out, out_work, out_work_bytes, + "glm routed moe expert-major"); + } + } + if (n_tokens == 2u && + g_glm_mtp_verify_mode && + getenv("DS4_GLM_MTP_NO_MOE_TOK2") == NULL) { + const uint32_t warps = 8u; + dim3 g1((expert_mid_dim + warps - 1u) / warps, + 2u * n_expert, 1u); + const uint32_t sh1 = 2u * xq_blocks * + (uint32_t)sizeof(cuda_block_q8_K); + glm_routed_moe_gateup_tok2_reuse_kernel<<< + g1, warps * 32u, sh1>>>( + mid_work, gw, uw, + (const cuda_block_q8_K *)xq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + gate_expert_bytes, gate_row_bytes, + up_expert_bytes, up_row_bytes, + xq_blocks, expert_mid_dim, n_expert); + } else if (getenv("DS4_GLM_MOE_SCALAR")) { + dim3 g1(n_tokens, n_expert, 1); + glm_routed_moe_batch_q2K_gateup_kernel<<>>( + mid_work, gw, uw, + (const cuda_block_q8_K *)xq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes, + xq_blocks, expert_mid_dim, n_expert, n_tokens, mid_token_stride); + } else { + const uint32_t warps = 8u; + dim3 g1((expert_mid_dim + warps - 1u) / warps, n_expert, n_tokens); + const uint32_t sh1 = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K); + glm_routed_moe_gateup_warp_kernel<<>>( + mid_work, gw, uw, + (const cuda_block_q8_K *)xq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes, + xq_blocks, expert_mid_dim, n_expert, n_tokens); + } + + { + dim3 gq(midq_blocks, n_tokens * n_expert, 1); + q8_K_quantize_kernel<<>>( + (cuda_block_q8_K *)midq_scratch[dev]->ptr, + mid_work, expert_mid_dim, n_tokens * n_expert); + } + + if (getenv("DS4_GLM_MOE_SCALAR")) { + dim3 g2((out_dim + 127u) / 128u, n_tokens, 1); + glm_routed_moe_batch_q2K_down_kernel<<>>( + out_work, dw, + (const cuda_block_q8_K *)midq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + down_expert_bytes, down_row_bytes, midq_blocks, out_dim, + n_expert, n_tokens); + } else { + const uint32_t warps = 8u; + dim3 g2((out_dim + warps - 1u) / warps, n_tokens, 1); + const uint32_t sh2 = n_expert * midq_blocks * + (uint32_t)sizeof(cuda_block_q8_K); + glm_routed_moe_down_warp_kernel<<>>( + out_work, dw, + (const cuda_block_q8_K *)midq_scratch[dev]->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + down_expert_bytes, down_row_bytes, midq_blocks, out_dim, + n_expert, n_tokens); + } + return glm_routed_moe_finish_batch( + out, out_work, out_work_bytes, + "glm routed moe batch launch"); +} + +extern "C" int ds4_gpu_glm_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + bool force_resident) { + (void)force_resident; + return ds4_gpu_glm_routed_moe_batch_tensor(out, mid, + model_map, model_size, + gate_offset, up_offset, down_offset, + gate_type, up_type, down_type, + gate_expert_bytes, gate_row_bytes, + up_expert_bytes, up_row_bytes, + down_expert_bytes, down_row_bytes, + expert_in_dim, expert_mid_dim, out_dim, + selected, weights, n_total_expert, n_expert, layer_index, + x, 1, n_expert * expert_mid_dim); +} + +/* Parallel router select: 256 threads compute sigmoid probs, then top-k + * via k rounds of shared-memory argmax over probs+bias (value desc, index + * asc tie-break — matches the CPU topk_desc). One block per token. */ +__global__ static void glm_router_select_parallel_kernel( + int32_t *selected, + float *weights_out, + float *probs_out, + const float *bias, + const float *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.x; + const uint32_t tid = threadIdx.x; + if (tok >= n_tokens) return; + const float *lg = logits + (uint64_t)tok * n_expert; + float *probs = probs_out + (uint64_t)tok * n_expert; + int32_t *sel = selected + (uint64_t)tok * n_expert_used; + float *w = weights_out + (uint64_t)tok * n_expert_used; + + __shared__ float sh_v[256]; + __shared__ int sh_i[256]; + __shared__ float sh_sel_v[256]; + __shared__ float sh_sum; + + float my_v = -1e30f; + if (tid < n_expert) { + const float p = 1.0f / (1.0f + expf(-lg[tid])); + probs[tid] = p; + my_v = p + bias[tid]; + } + if (tid == 0u) sh_sum = 0.0f; + sh_sel_v[tid] = my_v; + __syncthreads(); + + for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { + sh_v[tid] = sh_sel_v[tid]; + sh_i[tid] = (int)tid; + __syncthreads(); + for (uint32_t step = 128u; step > 0u; step >>= 1u) { + if (tid < step) { + const float ov = sh_v[tid + step]; + const int oi = sh_i[tid + step]; + if (ov > sh_v[tid] || (ov == sh_v[tid] && oi < sh_i[tid])) { + sh_v[tid] = ov; + sh_i[tid] = oi; + } + } + __syncthreads(); + } + if (tid == 0u) { + const int best = sh_i[0]; + sel[k2] = best; + const float p = probs[best]; + w[k2] = p; + sh_sum += p; + sh_sel_v[best] = -1e30f; + } + __syncthreads(); + } + if (tid == 0u) { + float sum = sh_sum; + if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; + for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { + w[k2] = w[k2] / sum * expert_weight_scale; + } + } +} + +__global__ static void glm_router_select_batch_kernel( + int32_t *selected, + float *weights_out, + float *probs_out, + const float *bias, + const float *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + const uint32_t tok = blockIdx.x; + if (tok >= n_tokens || threadIdx.x != 0u) return; + const float *lg = logits + (uint64_t)tok * n_expert; + float *probs = probs_out + (uint64_t)tok * n_expert; + int32_t *sel = selected + (uint64_t)tok * n_expert_used; + float *w = weights_out + (uint64_t)tok * n_expert_used; + + for (uint32_t i = 0; i < n_expert; i++) { + const float p = 1.0f / (1.0f + expf(-lg[i])); + probs[i] = p; + } + /* top-k over probs+bias, ties by smaller index (matches CPU topk_desc) */ + bool taken[384]; + for (uint32_t i = 0; i < n_expert; i++) taken[i] = false; + float sum = 0.0f; + for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { + int best = -1; float bv = -1e30f; + for (uint32_t i = 0; i < n_expert; i++) { + if (taken[i]) continue; + const float v = probs[i] + bias[i]; + if (v > bv) { bv = v; best = (int)i; } + } + taken[best] = true; + sel[k2] = best; + w[k2] = probs[best]; + sum += probs[best]; + } + if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; + for (uint32_t k2 = 0; k2 < n_expert_used; k2++) { + w[k2] = w[k2] / sum * expert_weight_scale; + } +} + +extern "C" int ds4_gpu_glm_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + if (!selected || !weights || !probs || !logits || !model_map || + n_expert == 0 || n_expert > 384u || n_expert_used == 0 || + n_tokens == 0) { + return 0; + } + const uint64_t bb = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bb > model_size - bias_offset || + logits->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) || + selected->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(int32_t) || + weights->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(float) || + probs->bytes < (uint64_t)n_tokens * n_expert * sizeof(float)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const float *bias = (const float *)cuda_resolve_weight_ptr( + model_map, bias_offset, bb, logical_tier, "glm_exp_probs_b"); + if (!bias) return 0; + if (n_expert <= 256u && !getenv("DS4_GLM_ROUTER_SCALAR")) { + glm_router_select_parallel_kernel<<>>( + (int32_t *)selected->ptr, + (float *)weights->ptr, + (float *)probs->ptr, + bias, + (const float *)logits->ptr, + n_expert, n_expert_used, expert_weight_scale, n_tokens); + } else glm_router_select_batch_kernel<<>>( + (int32_t *)selected->ptr, + (float *)weights->ptr, + (float *)probs->ptr, + bias, + (const float *)logits->ptr, + n_expert, n_expert_used, expert_weight_scale, n_tokens); + return cuda_ok(cudaGetLastError(), "glm router select batch launch"); +} + +extern "C" int ds4_gpu_glm_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale) { + return ds4_gpu_glm_router_select_batch_tensor(selected, weights, probs, + model_map, model_size, + bias_offset, logits, + n_expert, n_expert_used, + expert_weight_scale, 1u); +} + +__global__ static void glm_store_compact_kv_kernel( + char *kv_lora_cache, + char *k_rope_cache, + const float *kv_norm, + const float *kv_raw, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + uint32_t cache_f16) { + const uint32_t token = blockIdx.x; + const uint32_t part = blockIdx.y; + if (token >= n_tokens || part > 1u) return; + const uint32_t pos = pos0 + token; + if (pos >= cache_cap) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + if (part == 0u) { + const float *src = kv_norm + (uint64_t)token * kv_lora_dim; + if (cache_f16) { + __half *dst = (__half *)(kv_lora_cache + + (uint64_t)pos * kv_lora_dim * sizeof(__half)); + for (uint32_t i = tid; i < kv_lora_dim; i += nth) + dst[i] = __float2half(src[i]); + } else { + float *dst = (float *)(kv_lora_cache + + (uint64_t)pos * kv_lora_dim * sizeof(float)); + for (uint32_t i = tid; i < kv_lora_dim; i += nth) + dst[i] = src[i]; + } + } else { + const float *src = kv_raw + + (uint64_t)token * kv_raw_dim + kv_lora_dim; + if (cache_f16) { + __half *dst = (__half *)(k_rope_cache + + (uint64_t)pos * qk_rope * sizeof(__half)); + for (uint32_t i = tid; i < qk_rope; i += nth) + dst[i] = __float2half(src[i]); + } else { + float *dst = (float *)(k_rope_cache + + (uint64_t)pos * qk_rope * sizeof(float)); + for (uint32_t i = tid; i < qk_rope; i += nth) + dst[i] = src[i]; + } + } +} + +extern "C" int ds4_gpu_glm_store_compact_kv_tensor( + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_norm, + const ds4_gpu_tensor *kv_raw, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16) { + if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || + n_tokens == 0 || kv_lora_dim == 0 || qk_rope == 0 || + kv_lora_dim + qk_rope > kv_raw_dim + qk_rope) { + return 0; + } + const uint64_t es = cache_f16 ? sizeof(__half) : sizeof(float); + if (kv_norm->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float) || + kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) || + kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * es || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * es) { + return 0; + } + dim3 grid(n_tokens, 2, 1); + glm_store_compact_kv_kernel<<>>( + (char *)kv_lora_cache->ptr, + (char *)k_rope_cache->ptr, + (const float *)kv_norm->ptr, + (const float *)kv_raw->ptr, + pos0, n_tokens, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope, + cache_f16 ? 1u : 0u); + return cuda_ok(cudaGetLastError(), "glm store compact kv launch"); +} + + + + + + + +__global__ static void glm_store_indexer_k_kernel( + char *cache, + const float *raw_k, + const float *w, + const float *b, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + uint32_t cache_f16) { + const uint32_t token = blockIdx.x; + if (token >= n_tokens) return; + const uint32_t pos = pos0 + token; + if (pos >= cache_cap) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + const float *src = raw_k + (uint64_t)token * head_dim; + + __shared__ float scratch[256]; + float sum = 0.0f; + for (uint32_t i = tid; i < head_dim; i += nth) sum += src[i]; + scratch[tid] = sum; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + __syncthreads(); + } + const float mean = scratch[0] / (float)head_dim; + __syncthreads(); + float ss = 0.0f; + for (uint32_t i = tid; i < head_dim; i += nth) { + const float d = src[i] - mean; + ss += d * d; + } + scratch[tid] = ss; + __syncthreads(); + for (uint32_t step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + __syncthreads(); + } + const float inv = rsqrtf(scratch[0] / (float)head_dim + eps); + + float corr_dims[2] = {0.0f, 0.0f}; + if (ext_factor != 0.0f) { + corr_dims[0] = fmaxf(0.0f, + floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig, + beta_fast, freq_base))); + corr_dims[1] = fminf((float)rot_dim - 1.0f, + ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig, + beta_slow, freq_base))); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)rot_dim; + + for (uint32_t i = tid; i < head_dim; i += nth) { + float v0, v1; bool pair = false; + if (i < rot_dim) { + if ((i & 1u) != 0u) continue; + const float theta = theta_base * powf(freq_base, inv_ndims * (float)i); + float ct, st; + glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i, + ext_factor, attn_factor, &ct, &st); + const float x0 = (src[i] - mean) * inv * w[i] + b[i]; + const float x1 = (src[i + 1u] - mean) * inv * w[i + 1u] + b[i + 1u]; + v0 = x0 * ct - x1 * st; + v1 = x0 * st + x1 * ct; + pair = true; + } else { + v0 = (src[i] - mean) * inv * w[i] + b[i]; + } + if (cache_f16) { + __half *dst = (__half *)(cache + (uint64_t)pos * head_dim * sizeof(__half)); + dst[i] = __float2half(v0); + if (pair) dst[i + 1u] = __float2half(v1); + } else { + float *dst = (float *)(cache + (uint64_t)pos * head_dim * sizeof(float)); + dst[i] = v0; + if (pair) dst[i + 1u] = v1; + } + } +} + +extern "C" int ds4_gpu_glm_store_indexer_k_tensor( + ds4_gpu_tensor *indexer_key_cache, + const ds4_gpu_tensor *raw_k, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t bias_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16) { + if (!indexer_key_cache || !raw_k || !model_map || n_tokens == 0 || + head_dim == 0 || head_dim > 256u || (rot_dim & 1u) != 0u) { + return 0; + } + const uint64_t wb = (uint64_t)head_dim * sizeof(float); + if (weight_offset > model_size || wb > model_size - weight_offset || + bias_offset > model_size || wb > model_size - bias_offset || + raw_k->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) || + indexer_key_cache->bytes < + (uint64_t)cache_cap * head_dim * + (cache_f16 ? sizeof(__half) : sizeof(float))) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const float *w = (const float *)cuda_resolve_weight_ptr( + model_map, weight_offset, wb, logical_tier, "glm_indexer_k_norm"); + const float *b = (const float *)cuda_resolve_weight_ptr( + model_map, bias_offset, wb, logical_tier, "glm_indexer_k_norm_b"); + if (!w || !b) return 0; + glm_store_indexer_k_kernel<<>>( + (char *)indexer_key_cache->ptr, + (const float *)raw_k->ptr, + w, b, pos0, n_tokens, cache_cap, head_dim, rot_dim, n_ctx_orig, + eps, freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow, cache_f16 ? 1u : 0u); + return cuda_ok(cudaGetLastError(), "glm store indexer k launch"); +} + +extern "C" int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_selected) { + if (!g_ssd_streaming_mode) return 1; + if (!table || !selected || n_selected == 0 || + selected->bytes < (uint64_t)n_selected * sizeof(int32_t)) { + return 0; + } + std::vector ids; + try { + ids.resize(n_selected); + } catch (...) { + return 0; + } + if (!cuda_ok(cudaMemcpy(ids.data(), selected->ptr, + (size_t)n_selected * sizeof(int32_t), + cudaMemcpyDeviceToHost), + "GLM streaming selected-id read")) { + return 0; + } + return cuda_stream_selected_cache_begin_load(table, ids.data(), n_selected); +} + +__global__ static void glm_value_project_q8_0_batch_heads_kernel( + float *heads, + const char *weight, + const float *lora, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim, + uint64_t row_bytes) { + const uint32_t head = blockIdx.x; + const uint32_t token = blockIdx.y; + if (head >= n_head || token >= n_tokens) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + extern __shared__ float xsh[]; + const float *src = lora + (uint64_t)token * n_head * kv_lora_dim + + (uint64_t)head * kv_lora_dim; + float *out = heads + (uint64_t)token * n_head * value_dim + + (uint64_t)head * value_dim; + for (uint32_t j = tid; j < kv_lora_dim; j += nth) xsh[j] = src[j]; + __syncthreads(); + for (uint32_t d = tid; d < value_dim; d += nth) { + const char *row = weight + ((uint64_t)head * value_dim + d) * row_bytes; + out[d] = glm_q8_0_dot_row_dev(row, xsh, kv_lora_dim); + } +} + +/* Reuse each Q8 row across a token tile while retaining the scalar kernel's + * block/k accumulation order independently for every output token. */ +template +__global__ static void glm_value_project_q8_0_batch_heads_tiled_kernel( + float *heads, + const char *weight, + const float *lora, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim, + uint64_t row_bytes) { + const uint32_t head = blockIdx.x; + const uint32_t token0 = blockIdx.y * token_tile; + if (head >= n_head || token0 >= n_tokens) return; + const uint32_t tid = threadIdx.x; + const uint32_t nth = blockDim.x; + extern __shared__ float xsh[]; + +#pragma unroll + for (uint32_t t = 0; t < token_tile; t++) { + const uint32_t token = token0 + t; + if (token >= n_tokens) break; + const float *src = lora + (uint64_t)token * n_head * kv_lora_dim + + (uint64_t)head * kv_lora_dim; + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + xsh[(uint64_t)t * kv_lora_dim + j] = src[j]; + } + } + __syncthreads(); + + for (uint32_t od = tid; od < value_dim; od += nth) { + const char *row = weight + + ((uint64_t)head * value_dim + od) * row_bytes; + float acc[token_tile] = { 0.0f }; + const uint32_t nb = kv_lora_dim >> 5; + for (uint32_t b = 0; b < nb; b++) { + const char *blk = row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + const int8_t *q = (const int8_t *)(blk + 2); + float s[token_tile] = { 0.0f }; +#pragma unroll 8 + for (uint32_t k = 0; k < 32u; k++) { + const float w = (float)q[k]; +#pragma unroll + for (uint32_t t = 0; t < token_tile; t++) { + if (token0 + t < n_tokens) { + s[t] += w * xsh[(uint64_t)t * kv_lora_dim + + b * 32u + k]; + } + } + } +#pragma unroll + for (uint32_t t = 0; t < token_tile; t++) { + if (token0 + t < n_tokens) acc[t] += d * s[t]; + } + } +#pragma unroll + for (uint32_t t = 0; t < token_tile; t++) { + const uint32_t token = token0 + t; + if (token < n_tokens) { + heads[(uint64_t)token * n_head * value_dim + + (uint64_t)head * value_dim + od] = acc[t]; + } + } + } +} + +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 (!heads || !lora || !model_map || n_tokens == 0 || n_head == 0 || + kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u || value_dim == 0) { + return 0; + } + if (weight_type != 8u) { + fprintf(stderr, "ds4: glm value project: unsupported type %u\n", + weight_type); + return 0; + } + const uint64_t row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u; + const uint64_t wbytes = (uint64_t)n_head * value_dim * row_bytes; + if (weight_offset > model_size || wbytes > model_size - weight_offset || + lora->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || + heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float)) { + return 0; + } + const int logical_tier = cuda_current_tier(); + const char *w = (const char *)cuda_resolve_weight_ptr( + model_map, weight_offset, wbytes, logical_tier, "glm_v_b"); + if (!w) return 0; + if (n_tokens >= 16u && getenv("DS4_GLM_VALUE_NO_TILE16") == NULL) { + dim3 grid(n_head, (n_tokens + 15u) / 16u, 1); + const size_t shmem = 16ull * kv_lora_dim * sizeof(float); + glm_value_project_q8_0_batch_heads_tiled_kernel<16><<>>( + (float *)heads->ptr, w, (const float *)lora->ptr, + n_tokens, n_head, kv_lora_dim, value_dim, row_bytes); + return cuda_ok(cudaGetLastError(), "glm value project tile16 launch"); + } + dim3 grid(n_head, n_tokens, 1); + const size_t shmem = (size_t)kv_lora_dim * sizeof(float); + glm_value_project_q8_0_batch_heads_kernel<<>>( + (float *)heads->ptr, w, (const float *)lora->ptr, + n_tokens, n_head, kv_lora_dim, value_dim, row_bytes); + return cuda_ok(cudaGetLastError(), "glm value project launch"); +} + +/* Decode-time (n_tok small) quant matvec. The Metal "mpp/model-view" + * variant is a bandwidth-tuned matvec; on CUDA the generic quant matmul + * already dispatches per type, so delegate. Revisit in the perf pass. */ +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); + +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) { + 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_matmul_quant_rows_scalar_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) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_matmul_quant_rows_scalar_tensor\n"); + return 0; +} + +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) { + switch (weight_type) { + case 8u: /* Q8_0 */ + return ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok); + case 1u: /* F16 */ + return ds4_gpu_matmul_f16_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok); + default: + fprintf(stderr, "ds4: matmul_quant: unsupported type %u\n", + weight_type); + return 0; + } +} + +extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) { + /* GLM graph memory guard: on this backend the model weights are + * distributed across all devices by the multi-tier placement, so the + * relevant budget is the aggregate VRAM. */ + int n = 0; + if (cudaGetDeviceCount(&n) != cudaSuccess || n <= 0) return 0; + size_t free_b = 0, total_b = 0; + if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) return 0; + return (uint64_t)total_b * (uint64_t)n; +} + +extern "C" int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) { + (void)selected; + (void)n_selected; + return 1; +} + +extern "C" void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { + (void)enabled; /* SSD streaming is not used on the CUDA backend */ +} + +extern "C" void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { + g_glm_mtp_verify_mode = enabled; +} + +extern "C" int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes) { + (void)max_tensor_bytes; + if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) { + return 0; + } + for (uint32_t i = 0; i < count; i++) { + if (offsets[i] > model_size || sizes[i] == 0 || + sizes[i] > model_size - offsets[i]) { + return 0; + } + } + if (!ds4_gpu_set_model_map(model_map, model_size)) return 0; + if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL) { + for (uint32_t i = 0; i < count; i++) { + (void)cuda_model_prefetch_range(model_map, model_size, + offsets[i], sizes[i]); + } + } + return 1; +} + +extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, + uint64_t in_dim, uint64_t out_dim, + const ds4_gpu_tensor *x, uint64_t n_tok, float clamp); + +extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate, up, mid, + model_map, model_size, gate_offset, up_offset, + in_dim, out_dim, x, 1, clamp); +} + +extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + (void)gate; (void)up; (void)mid; (void)model_map; (void)model_size; + (void)gate_offset; (void)up_offset; (void)in_dim; (void)out_dim; + (void)x; (void)n_tok; (void)clamp; + return 0; +} + +/* Fused single-token shared-expert gate+up+swiglu: one warp per output + * row computes both q8_0 dots against a shared-staged f32 x and writes + * silu(gate)*up directly. Falls back to the split path for n_tok > 1. */ +__global__ static void glm_shared_gate_up_swiglu_one_kernel( + float *mid, + const char *gw, + const char *uw, + const float *x, + uint32_t in_dim, + uint32_t out_dim, + float clamp) { + extern __shared__ float glm_sgu_sh[]; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) { + glm_sgu_sh[i] = x[i]; + } + __syncthreads(); + const uint32_t r = blockIdx.x * warps + warp; + if (r >= out_dim) return; + const uint32_t nblk = in_dim >> 5; + const uint64_t row_bytes = (uint64_t)nblk * 34u; + const char *grow = gw + (uint64_t)r * row_bytes; + const char *urow = uw + (uint64_t)r * row_bytes; + float g = 0.0f, u = 0.0f; + for (uint32_t blk = lane; blk < nblk; blk += 32u) { + const char *gb = grow + (uint64_t)blk * 34u; + const char *ub = urow + (uint64_t)blk * 34u; + const float gd = __half2float(*(const __half *)gb); + const float ud = __half2float(*(const __half *)ub); + const int8_t *gq = (const int8_t *)(gb + 2); + const int8_t *uq = (const int8_t *)(ub + 2); + const float *xs = glm_sgu_sh + blk * 32u; + float gs = 0.0f, us = 0.0f; + #pragma unroll 8 + for (int k = 0; k < 32; k++) { + gs += (float)gq[k] * xs[k]; + us += (float)uq[k] * xs[k]; + } + g += gd * gs; + u += ud * us; + } + for (int off = 16; off > 0; off >>= 1) { + g += __shfl_down_sync(0xffffffffu, g, off); + u += __shfl_down_sync(0xffffffffu, u, off); + } + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (g > clamp) g = clamp; + if (u > clamp) u = clamp; + if (u < -clamp) u = -clamp; + } + mid[r] = (g / (1.0f + expf(-g))) * u; + } +} + +/* Two-token verifier variant of the decode kernel above. Each warp loads a + * gate/up weight row once, while each token keeps the decode kernel's block + * order and warp reduction tree independently. */ +__global__ static void glm_shared_gate_up_swiglu_tok2_exact_kernel( + float *mid, + const char *gw, + const char *uw, + const float *x, + uint32_t in_dim, + uint32_t out_dim, + float clamp) { + extern __shared__ float glm_sgu2_sh[]; + float *x0 = glm_sgu2_sh; + float *x1 = glm_sgu2_sh + in_dim; + const uint32_t warps = blockDim.x >> 5; + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31u; + for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) { + x0[i] = x[i]; + x1[i] = x[in_dim + i]; + } + __syncthreads(); + const uint32_t r = blockIdx.x * warps + warp; + if (r >= out_dim) return; + const uint32_t nblk = in_dim >> 5; + const uint64_t row_bytes = (uint64_t)nblk * 34u; + const char *grow = gw + (uint64_t)r * row_bytes; + const char *urow = uw + (uint64_t)r * row_bytes; + float g0 = 0.0f, u0 = 0.0f; + float g1 = 0.0f, u1 = 0.0f; + for (uint32_t blk = lane; blk < nblk; blk += 32u) { + const char *gb = grow + (uint64_t)blk * 34u; + const char *ub = urow + (uint64_t)blk * 34u; + const float gd = __half2float(*(const __half *)gb); + const float ud = __half2float(*(const __half *)ub); + const int8_t *gq = (const int8_t *)(gb + 2); + const int8_t *uq = (const int8_t *)(ub + 2); + const float *xs0 = x0 + blk * 32u; + const float *xs1 = x1 + blk * 32u; + float gs0 = 0.0f, us0 = 0.0f; + float gs1 = 0.0f, us1 = 0.0f; + #pragma unroll 8 + for (int k = 0; k < 32; k++) { + const float gk = (float)gq[k]; + const float uk = (float)uq[k]; + gs0 += gk * xs0[k]; + us0 += uk * xs0[k]; + gs1 += gk * xs1[k]; + us1 += uk * xs1[k]; + } + g0 += gd * gs0; + u0 += ud * us0; + g1 += gd * gs1; + u1 += ud * us1; + } + for (int off = 16; off > 0; off >>= 1) { + g0 += __shfl_down_sync(0xffffffffu, g0, off); + u0 += __shfl_down_sync(0xffffffffu, u0, off); + g1 += __shfl_down_sync(0xffffffffu, g1, off); + u1 += __shfl_down_sync(0xffffffffu, u1, off); + } + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (g0 > clamp) g0 = clamp; + if (u0 > clamp) u0 = clamp; + if (u0 < -clamp) u0 = -clamp; + if (g1 > clamp) g1 = clamp; + if (u1 > clamp) u1 = clamp; + if (u1 < -clamp) u1 = -clamp; + } + mid[r] = (g0 / (1.0f + expf(-g0))) * u0; + mid[out_dim + r] = (g1 / (1.0f + expf(-g1))) * u1; + } +} + +extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + if (!gate || !up || !mid || !x || n_tok == 0) return 0; + if (n_tok == 2 && (in_dim & 31u) == 0u && + g_glm_mtp_verify_mode && + getenv("DS4_GLM_MTP_NO_SHARED_TOK2") == NULL && + mid->bytes >= 2u * out_dim * sizeof(float) && + x->bytes >= 2u * in_dim * sizeof(float)) { + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t wb = out_dim * row_bytes; + if (gate_offset <= model_size && wb <= model_size - gate_offset && + up_offset <= model_size && wb <= model_size - up_offset) { + const int logical_tier = cuda_current_tier(); + const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset, + wb, logical_tier, "glm_shared_gate"); + const char *uw = cuda_resolve_weight_ptr(model_map, up_offset, + wb, logical_tier, "glm_shared_up"); + if (gw && uw) { + const uint32_t warps = 8u; + const uint32_t sh = 2u * (uint32_t)in_dim * sizeof(float); + glm_shared_gate_up_swiglu_tok2_exact_kernel + <<<(unsigned)((out_dim + warps - 1u) / warps), + warps * 32u, sh>>>( + (float *)mid->ptr, gw, uw, (const float *)x->ptr, + (uint32_t)in_dim, (uint32_t)out_dim, clamp); + return cuda_ok(cudaGetLastError(), + "glm shared swiglu tok2 exact"); + } + } + } + if (n_tok == 1 && (in_dim & 31u) == 0u && + !getenv("DS4_GLM_SHARED_SPLIT") && + mid->bytes >= out_dim * sizeof(float) && + x->bytes >= in_dim * sizeof(float)) { + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t wb = out_dim * row_bytes; + if (gate_offset <= model_size && wb <= model_size - gate_offset && + up_offset <= model_size && wb <= model_size - up_offset) { + const int logical_tier = cuda_current_tier(); + const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset, + wb, logical_tier, "glm_shared_gate"); + const char *uw = cuda_resolve_weight_ptr(model_map, up_offset, + wb, logical_tier, "glm_shared_up"); + if (gw && uw) { + const uint32_t warps = 8u; + const uint32_t sh = (uint32_t)in_dim * sizeof(float); + glm_shared_gate_up_swiglu_one_kernel + <<<(unsigned)((out_dim + warps - 1u) / warps), + warps * 32u, sh>>>( + (float *)mid->ptr, gw, uw, (const float *)x->ptr, + (uint32_t)in_dim, (uint32_t)out_dim, clamp); + return cuda_ok(cudaGetLastError(), "glm shared swiglu one"); + } + } + } + if (!ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size, gate_offset, + in_dim, out_dim, x, n_tok) || + !ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size, up_offset, + in_dim, out_dim, x, n_tok)) { + return 0; + } + return ds4_gpu_swiglu_tensor(mid, gate, up, + (uint32_t)(out_dim * n_tok), clamp, 1.0f); +} + +extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + static ds4_gpu_tensor *gu_scratch[DS4_MAX_GPUS][2] = {{0}}; + const int dev = cuda_current_tier(); + const uint64_t need = out_dim * sizeof(float); + for (int i = 0; i < 2; i++) { + if (!gu_scratch[dev][i] || gu_scratch[dev][i]->bytes < need) { + if (gu_scratch[dev][i]) ds4_gpu_tensor_free(gu_scratch[dev][i]); + gu_scratch[dev][i] = ds4_gpu_tensor_alloc(need); + } + if (!gu_scratch[dev][i]) return 0; + } + return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + gu_scratch[dev][0], gu_scratch[dev][1], mid, + model_map, model_size, gate_offset, up_offset, + in_dim, out_dim, x, 1, clamp); +} + +extern "C" int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { + if (event_value) *event_value = 1; + return cuda_ok(cudaDeviceSynchronize(), "selected readback signal"); +} + +extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + return cuda_stream_selected_cache_begin_load(table, selected_ids, + n_selected); +} + +extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + (void)gate_expert_bytes; + (void)down_expert_bytes; + return 0; +} + +extern "C" int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t count) { + if (!dst || !src) return 0; + if (count == 0) return 1; + if (count > UINT64_MAX / sizeof(float) || + count > UINT64_MAX / sizeof(__half)) { + return 0; + } + const uint64_t src_bytes = count * sizeof(float); + const uint64_t dst_bytes = count * sizeof(__half); + if (src_offset > src->bytes || src_bytes > src->bytes - src_offset || + dst_offset > dst->bytes || dst_bytes > dst->bytes - dst_offset || + ds4_tensor_device_idx(dst) != ds4_tensor_device_idx(src)) { + return 0; + } + const int tier = ds4_tensor_device_idx(dst); + if (ds4_gpu_set_current_device(tier) != 0) return 0; + const uint64_t blocks = (count + 255u) / 256u; + if (blocks > UINT32_MAX) return 0; + f32_to_f16_kernel<<<(unsigned)blocks, 256>>>( + (__half *)((char *)dst->ptr + dst_offset), + (const float *)((const char *)src->ptr + src_offset), + count); + return cuda_ok(cudaGetLastError(), "tensor f32-to-f16 copy launch"); +} + +extern "C" int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor, + uint64_t offset, + void *data, + uint64_t bytes, + uint64_t event_value, + const char *label) { + (void)event_value; + if (!tensor || !data || offset > tensor->bytes || + bytes > tensor->bytes - offset) { + return 0; + } + if (!cuda_ok(cudaDeviceSynchronize(), + label ? label : "selected readback wait")) { + return 0; + } + return cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset, + (size_t)bytes, cudaMemcpyDeviceToHost), + "selected tensor read"); +} + +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) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_tp_big_gate_encode\n"); + return 0; +} + +extern "C" int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) { + fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_tp_gate_encode\n"); + return 0; +} + +extern "C" void ds4_gpu_tp_set_attn_head_split(int enabled) { + (void)enabled; /* Mac network-TP head split: no-op on CUDA */ +} + +extern "C" int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) { + (void)event_value; + return cuda_ok(cudaDeviceSynchronize(), + label ? label : "selected readback wait"); +} + +/* Compatibility surface shared with the canonical Metal/ROCm graph. CUDA + * either delegates to its equivalent primitive or reports an unavailable + * optional fast path so the graph can use its established fallback. */ +extern "C" int ds4_gpu_commit_and_wait_selected_readback( + uint64_t event_value, const char *label) { + (void)event_value; + return cuda_ok(cudaDeviceSynchronize(), + label ? label : "selected readback wait"); +} + +extern "C" int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) { + const int ok = ds4_gpu_set_model_fd(fd); + if (ok) g_model_fd_host_base = model_map; + return ok; +} + +extern "C" int ds4_gpu_pro_q4_expert_table_auto_available(void) { + return 0; +} + +extern "C" int ds4_gpu_preload_q4_expert_tables( + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, + uint64_t gate_expert_bytes, uint64_t down_expert_bytes, + uint32_t n_total_expert) { + (void)model_map; (void)model_size; + (void)gate_offset; (void)up_offset; (void)down_offset; + (void)gate_expert_bytes; (void)down_expert_bytes; + (void)n_total_expert; + return 1; +} + +extern "C" void ds4_gpu_set_glm_model(bool enabled) { + (void)enabled; +} + +extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { + g_ssd_streaming_mode = enabled ? 1 : 0; + cuda_stream_selected_cache_invalidate(); + if (!g_ssd_streaming_mode) cuda_stream_selected_cache_release(); +} + +extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { + (void)experts; +} + +extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { + (void)bytes; +} + +extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { + return 0; +} + +extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) { + return g_stream_selected_cache.valid ? + g_stream_selected_cache.compact_count : 0; +} + +extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { +} + +extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) { + cuda_stream_selected_cache_release(); +} + +extern "C" int ds4_gpu_stream_expert_cache_seed_selected( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + (void)table; (void)selected_ids; (void)n_selected; + return 1; +} + +extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_tokens, + uint32_t n_selected) { + if (n_tokens == 0 || n_selected == 0 || + (uint64_t)n_tokens * n_selected > UINT32_MAX) { + return 0; + } + return cuda_stream_selected_cache_begin_load( + table, selected_ids, n_tokens * n_selected); +} + +extern "C" int ds4_gpu_stream_expert_cache_seed_experts( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts) { + (void)table; (void)expert_ids; (void)expert_priorities; (void)n_experts; + return 1; +} + +extern "C" int ds4_gpu_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits, + uint32_t n_vocab) { + return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1u, 1u); +} + +extern "C" int ds4_gpu_embed_token_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd) { + return ds4_gpu_embed_token_quant_tensor(out, model_map, model_size, + weight_offset, 8u, n_vocab, + token, n_embd); +} + +extern "C" int ds4_gpu_embed_tokens_q8_0_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + return ds4_gpu_embed_tokens_quant_tensor(out, tokens, model_map, + model_size, weight_offset, 8u, + n_vocab, n_tokens, n_embd); +} + +extern "C" int ds4_gpu_glm_k_b_project_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 n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head) { + return ds4_gpu_glm_k_b_project_typed_tensor( + out, kv_norm, model_map, model_size, weight_offset, 8u, + n_tokens, kv_lora_dim, qk_nope, n_head); +} + +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) { + if (!x || x_elem_off > x->bytes / sizeof(float) || + k_cnt > x->bytes / sizeof(float) - x_elem_off) { + return 0; + } + ds4_gpu_tensor x_slice = *x; + x_slice.ptr = (char *)x->ptr + x_elem_off * sizeof(float); + x_slice.bytes = k_cnt * sizeof(float); + x_slice.owner = 0; + return ds4_gpu_matmul_q8_0_kslice_rows_tensor( + out, model_map, model_size, weight_offset, + full_in_dim, out_dim, k_off, k_cnt, &x_slice, 1u); +} + +extern "C" int ds4_gpu_matmul_quant_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + 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) { + if (weight_type != 8u) return 0; + return ds4_gpu_matmul_q8_0_kslice_tensor( + out, model_map, model_size, weight_offset, + full_in_dim, k_off, k_cnt, out_dim, x, x_elem_off); +} + +extern "C" int ds4_gpu_matmul_q8_0_f16_out_tensor( + ds4_gpu_tensor *out_h, + 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_tok) { + (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; + (void)in_dim; (void)out_dim; (void)x; (void)n_tok; + return 0; +} + +extern "C" int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + ds4_gpu_tensor *out, ds4_gpu_tensor *q_half, + 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_tok, uint32_t n_head, uint32_t 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) { + (void)out; (void)q_half; (void)model_map; (void)model_size; + (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; + (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0; + (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale; + (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; + (void)eps; + return 0; +} + +extern "C" int ds4_gpu_attention_prefill_raw_heads_range_tensor( + ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, + uint64_t sinks_offset, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, uint32_t q_row0, uint32_t n_q, + uint32_t n_kv, uint32_t window, uint32_t n_head, + uint32_t head_dim) { + (void)heads; (void)model_map; (void)model_size; (void)sinks_offset; + (void)q; (void)raw_kv; (void)q_row0; (void)n_q; (void)n_kv; + (void)window; (void)n_head; (void)head_dim; + return 0; +} + +extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( + ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, + uint64_t sinks_offset, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, uint32_t q_row0, uint32_t n_q, + uint32_t n_tokens, uint32_t n_comp, uint32_t window, + uint32_t ratio, uint32_t n_head, uint32_t head_dim) { + (void)heads; (void)model_map; (void)model_size; (void)sinks_offset; + (void)q; (void)raw_kv; (void)comp_kv; (void)comp_kv_f16; + (void)q_row0; (void)n_q; (void)n_tokens; (void)n_comp; (void)window; + (void)ratio; (void)n_head; (void)head_dim; + return 0; +} + +extern "C" int ds4_gpu_attention_output_q8_batch_f16_tensor( + ds4_gpu_tensor *out_h, 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, + uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) { + (void)out_h; (void)low; (void)model_map; (void)model_size; + (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank; + (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; + return 0; +} + +extern "C" int ds4_gpu_attention_output_q4_K_batch_tensor( + ds4_gpu_tensor *out, ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, ds4_gpu_tensor *low_tmp, + const void *model_map, uint64_t model_size, + uint64_t out_a_offset, uint64_t out_b_offset, uint32_t out_b_type, + uint64_t group_dim, uint64_t rank, uint32_t n_groups, + uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) { + (void)out; (void)low; (void)group_tmp; (void)low_tmp; + (void)model_map; (void)model_size; (void)out_a_offset; + (void)out_b_offset; (void)out_b_type; (void)group_dim; (void)rank; + (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; + return 0; +} + +extern "C" int ds4_gpu_attention_output_low_q4_K_slice_tensor( + ds4_gpu_tensor *low, const void *model_map, uint64_t model_size, + uint64_t out_a_offset, uint64_t group_dim, uint64_t rank, + uint32_t group0, uint32_t group_cnt, + const ds4_gpu_tensor *heads) { + (void)low; (void)model_map; (void)model_size; (void)out_a_offset; + (void)group_dim; (void)rank; (void)group0; (void)group_cnt; + (void)heads; + return 0; +} + +extern "C" int ds4_gpu_hc_expand_split_half_tensor( + ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out_h, + const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, + uint32_t n_embd, uint32_t n_hc) { + (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split; + (void)n_embd; (void)n_hc; + return 0; +} + +extern "C" int ds4_gpu_hc_expand_add_split_half_add_tensor( + ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add_h, + const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, + uint32_t n_embd, uint32_t n_hc) { + (void)out_hc; (void)block_out; (void)block_add_h; + (void)residual_hc; (void)split; (void)n_embd; (void)n_hc; + return 0; +} + +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_model_residency_skip(int skip) { + (void)skip; +} + +extern "C" uint64_t ds4_gpu_tp_big_gate_kick( + 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_big_gate_wait(uint64_t seq) { + (void)seq; + return 0; +} + +extern "C" int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) { + (void)layer; (void)rows; + return 0; +} +#pragma GCC diagnostic pop diff --git a/models/glm/graph.inc b/models/glm/graph.inc new file mode 100644 index 0000000000..70f8669b59 --- /dev/null +++ b/models/glm/graph.inc @@ -0,0 +1,9513 @@ +/* + * GLM DSA graph state, allocation, decode, prefill, and diagnostics. + * + * Included exactly once inside ds4.c's graph-backend conditional. No generic + * operator layer is introduced between this model integration and its kernels. + */ + +typedef struct { + uint32_t ctx_size; + uint32_t ctx_cap; + uint32_t normal_layers; + uint32_t layer_start; + uint32_t layer_end; + uint32_t layer_count; + uint64_t q_dim; + uint64_t q_nope; + uint64_t heads_dim; + uint64_t kv_raw_dim; + uint64_t dense_hidden_max; + uint64_t ffn_mid_elems; + + ds4_gpu_tensor *cur; + ds4_gpu_tensor *next; + ds4_gpu_tensor *attn_norm; + ds4_gpu_tensor *q_rank; + ds4_gpu_tensor *q_rank_norm; + ds4_gpu_tensor *q; + ds4_gpu_tensor *kv_raw; + ds4_gpu_tensor *kv_norm; + ds4_gpu_tensor *k_nope; + ds4_gpu_tensor *value; + ds4_gpu_tensor *heads; + ds4_gpu_tensor *attn_out; + ds4_gpu_tensor *after_attn; + ds4_gpu_tensor *ffn_norm; + ds4_gpu_tensor *ffn_gate; + ds4_gpu_tensor *ffn_up; + ds4_gpu_tensor *ffn_mid; + ds4_gpu_tensor *routed_gate; + ds4_gpu_tensor *routed_up; + ds4_gpu_tensor *routed_down; + ds4_gpu_tensor *ffn_out; + ds4_gpu_tensor *ffn_sum; + ds4_gpu_tensor *router_logits; + ds4_gpu_tensor *router_probs; + ds4_gpu_tensor *router_selected; + ds4_gpu_tensor *router_weights; + ds4_gpu_tensor *output_norm; + ds4_gpu_tensor *logits; + ds4_gpu_tensor *batch_router_logits; + ds4_gpu_tensor *batch_router_probs; + ds4_gpu_tensor *batch_router_selected; + ds4_gpu_tensor *batch_router_weights; + ds4_gpu_tensor *prefill_seed_router_selected; + uint32_t prefill_seed_tokens; + bool prefill_seed_layer_captured[DS4_MAX_LAYER]; + + ds4_gpu_tensor *prefill_tokens; + ds4_gpu_tensor *batch_cur; + ds4_gpu_tensor *batch_next; + ds4_gpu_tensor *batch_attn_norm; + ds4_gpu_tensor *batch_q_rank; + ds4_gpu_tensor *batch_q_rank_norm; + ds4_gpu_tensor *batch_q; + ds4_gpu_tensor *batch_kv_raw; + ds4_gpu_tensor *batch_kv_norm; + ds4_gpu_tensor *batch_k_nope; + ds4_gpu_tensor *batch_value; + ds4_gpu_tensor *batch_heads; + ds4_gpu_tensor *batch_attn_out; + ds4_gpu_tensor *batch_after_attn; + ds4_gpu_tensor *batch_ffn_norm; + ds4_gpu_tensor *batch_ffn_gate; + ds4_gpu_tensor *batch_ffn_up; + ds4_gpu_tensor *batch_shared_mid; + ds4_gpu_tensor *batch_ffn_mid; + ds4_gpu_tensor *batch_routed_gate; + ds4_gpu_tensor *batch_routed_up; + ds4_gpu_tensor *batch_routed_down; + ds4_gpu_tensor *batch_ffn_out; + bool batch_routed_mid_is_f16; + + uint32_t compact_cache_cap; + uint32_t indexed_prefill_cap; + uint32_t indexed_prefill_score_cap; + uint32_t indexer_full_layers; + ds4_gpu_tensor *indexer_k; + ds4_gpu_tensor *indexer_q; + ds4_gpu_tensor *indexer_weights; + ds4_gpu_tensor *indexer_scores; + ds4_gpu_tensor *indexer_selected; + ds4_gpu_tensor *qk_low; + ds4_gpu_tensor *attn_partial_lora; + ds4_gpu_tensor *attn_partial_ms; + ds4_gpu_tensor *batch_indexer_k; + ds4_gpu_tensor *batch_indexer_q; + ds4_gpu_tensor *batch_indexer_weights; + ds4_gpu_tensor *batch_indexer_scores; + ds4_gpu_tensor *batch_indexer_selected; + ds4_gpu_tensor *batch_qk_low; + ds4_gpu_tensor *batch_attn_lora; + ds4_gpu_tensor *layer_kv_lora_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_k_rope_cache[DS4_MAX_LAYER]; + /* GLM MTP (nextn block) drafting: private compact caches for the nextn + * layer (slot = absolute position; only [mtp_min_pos..pos] is ever + * selected) plus small scratch. Allocated lazily on first draft. */ + ds4_gpu_tensor *mtp_kv_lora_cache; + ds4_gpu_tensor *mtp_k_rope_cache; + ds4_gpu_tensor *mtp_concat; + ds4_gpu_tensor *mtp_selected; + float *mtp_logits_host; + int mtp_ready; + ds4_gpu_tensor *layer_indexer_key_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_key_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_value_cache[DS4_MAX_LAYER]; + bool full_kv_cache; + bool has_token_embd; + bool has_output_head; + bool quality; + bool ssd_streaming; + bool ssd_streaming_cold; + bool generic_routed_moe; + bool streaming_static_decode_map_current; + /* Tensor parallelism (50/50 expert sharding): tp_world 2 means + * this rank computes only its contiguous half of the routed experts + * and exchanges the 24KB routed-FFN partial at one gate per sparse + * layer. Views alias the engine's TP slab slots [layer*2 + FFN]. */ + uint32_t tp_world; + uint32_t tp_rank; + ds4_gpu_tensor **tp_out; + ds4_gpu_tensor **tp_in; + /* Prefill batch gate bounce buffers (shared storage; grow on demand). */ + ds4_gpu_tensor *tp_bounce_out; + ds4_gpu_tensor *tp_bounce_in; + /* CUDA multi-tier placement and device-local decode scratch mirrors. */ + const int *placement; +#define DS4_GLM_WS_SLOTS 29 + ds4_gpu_tensor *ws_mirror[DS4_MAX_GPUS][DS4_GLM_WS_SLOTS]; + ds4_gpu_tensor *ws_orig[DS4_GLM_WS_SLOTS]; + int ws_ready; + int ws_tier; +#define DS4_GLM_VERIFY_WS_SLOTS 28 + ds4_gpu_tensor *verify_ws_mirror[DS4_MAX_GPUS][DS4_GLM_VERIFY_WS_SLOTS]; + ds4_gpu_tensor *verify_ws_orig[DS4_GLM_VERIFY_WS_SLOTS]; + int verify_ws_ready; + int verify_ws_tier; +} ds4_glm_gpu_graph; + +static uint32_t glm_graph_model_context_limit(void) { + if (DS4_ROPE_ORIG_CTX > UINT32_MAX) return UINT32_MAX; + return (uint32_t)DS4_ROPE_ORIG_CTX; +} + +static double glm_graph_bytes_to_gib(uint64_t bytes) { + return (double)bytes / (1024.0 * 1024.0 * 1024.0); +} + +static uint64_t glm_graph_saturating_add_u64(uint64_t a, uint64_t b) { + return a > UINT64_MAX - b ? UINT64_MAX : a + b; +} + +static bool glm_graph_env_disabled(const char *name) { + const char *env = getenv(name); + if (!env || !env[0]) return false; + return strcmp(env, "0") == 0 || + strcasecmp(env, "false") == 0 || + strcasecmp(env, "off") == 0 || + strcasecmp(env, "no") == 0; +} + +static double glm_graph_env_double( + const char *name, + double fallback, + double min_value, + double max_value) { + const char *env = getenv(name); + if (!env || !env[0]) return fallback; + char *end = NULL; + errno = 0; + const double v = strtod(env, &end); + if (end == env || errno != 0 || !isfinite(v)) return fallback; + if (v < min_value) return min_value; + if (v > max_value) return max_value; + return v; +} + +static uint64_t glm_graph_host_memory_bytes(void) { +#if defined(__APPLE__) + uint64_t mem = 0; + size_t len = sizeof(mem); + if (sysctlbyname("hw.memsize", &mem, &len, NULL, 0) != 0) return 0; + return mem; +#else + return 0; +#endif +} + +static uint64_t glm_graph_streaming_active_model_bytes( + const ds4_weights *weights) { + if (!weights) return 0; + + uint64_t max_bytes = 0; + ds4_model_map_span_vec spans; + + if (weights_layer_has_required(&weights->layer[0], 0) && + weights_model_map_token_spans(weights, &spans)) { + max_bytes = model_map_span_vec_total_bytes(&spans); + free(spans.v); + } + if (weights_have_output_head(weights) && + weights_model_map_output_spans(weights, &spans)) { + const uint64_t bytes = model_map_span_vec_total_bytes(&spans); + if (bytes > max_bytes) max_bytes = bytes; + free(spans.v); + } + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!weights_model_map_spans(weights, il, il, false, &spans)) { + continue; + } + const uint64_t bytes = model_map_span_vec_total_bytes(&spans); + if (bytes > max_bytes) max_bytes = bytes; + free(spans.v); + } + + return max_bytes; +} + +/* TP shard bytes: dense weights plus this rank's routed-expert range. + * Zero when not sharding. Set during engine open, before the GLM memory + * guard runs. */ +static uint64_t g_tp_shard_model_bytes; + +/* A user-raised iogpu.wired_limit_mb is an explicit GPU budget grant; + * prefer it over the fraction/reserve heuristics. */ +static uint64_t glm_graph_wired_limit_bytes(void) { +#ifdef __APPLE__ + int64_t mb = 0; + size_t len = sizeof(mb); + if (sysctlbyname("iogpu.wired_limit_mb", &mb, &len, NULL, 0) != 0) return 0; + if (mb <= 0) return 0; + return (uint64_t)mb * 1024ull * 1024ull; +#else + return 0; +#endif +} + +static uint64_t glm_graph_model_bytes_for_guard( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + bool load_slice, + uint32_t layer_start, + uint32_t layer_end, + bool include_token, + bool include_output) { + if (!model) return 0; + /* Under TP, the sharded map bytes are authoritative regardless of + * how the caller frames the request (TP excludes real layer slicing, + * so any slice request here is the session's full-range accounting). */ + if (!ssd_streaming && g_tp_shard_model_bytes != 0) { + return g_tp_shard_model_bytes; + } + if (load_slice && weights) { + ds4_model_map_span_vec spans; + bool ok = false; + if (ssd_streaming) { + ok = weights_model_map_decode_static_slice_spans(weights, + layer_start, + layer_end, + include_token, + include_output, + &spans); + } else { + ok = weights_model_map_spans(weights, + layer_start, + layer_end, + include_output, + &spans); + } + if (ok) { + const uint64_t bytes = model_map_span_vec_total_bytes(&spans); + free(spans.v); + if (bytes != 0) return bytes; + } + } + if (!ssd_streaming) { + if (g_tp_shard_model_bytes != 0) return g_tp_shard_model_bytes; + return model->size; + } + const uint64_t active_bytes = glm_graph_streaming_active_model_bytes(weights); + return active_bytes != 0 ? active_bytes : model->size; +} + +static double glm_graph_memory_guard_default_reserve_gib( + uint64_t budget_base, + uint64_t model_bytes) { + const double base_gib = glm_graph_bytes_to_gib(budget_base); + const double model_gib = glm_graph_bytes_to_gib(model_bytes); + if (base_gib >= 480.0 && + base_gib <= 640.0 && + model_gib >= base_gib * 0.80) { + return 24.0; + } + return 32.0; +} + +static bool glm_graph_memory_guard_for_compact_cap( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + bool load_slice, + uint32_t layer_start, + uint32_t layer_end, + bool include_token, + bool include_output, + uint32_t ctx_size, + uint32_t compact_cap, + uint64_t transient_extra_bytes, + const char *phase) { + if (!model || glm_graph_env_disabled("DS4_GLM_MEMORY_GUARD")) return true; + + const uint64_t host_bytes = glm_graph_host_memory_bytes(); + uint64_t budget_base = host_bytes; + if (budget_base == 0) { + budget_base = ds4_gpu_recommended_working_set_size(); + } + if (budget_base == 0) return true; + const uint64_t wired_limit = glm_graph_wired_limit_bytes(); + + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx_size, ssd_streaming); + const ds4_context_memory mem = load_slice ? + glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx_size, + work_ctx, + compact_cap, + ssd_streaming, + layer_start, + layer_end) : + glm_graph_context_memory_estimate_for_compact_cap( + ctx_size, + work_ctx, + compact_cap, + ssd_streaming); + const uint64_t graph_bytes = mem.total_bytes; + const uint64_t model_bytes = + glm_graph_model_bytes_for_guard(model, + weights, + ssd_streaming, + load_slice, + layer_start, + layer_end, + include_token, + include_output); + uint64_t required = glm_graph_saturating_add_u64(model_bytes, graph_bytes); + required = glm_graph_saturating_add_u64(required, transient_extra_bytes); + + const double fraction = + glm_graph_env_double("DS4_GLM_MEMORY_GUARD_FRACTION", 0.99, 0.50, 1.00); + double default_reserve_gib = + glm_graph_memory_guard_default_reserve_gib(budget_base, model_bytes); +#ifdef DS4_ROCM_BUILD + if (load_slice && !ssd_streaming) { + /* The original fixed reserve protects Metal's shared host/GPU heap. + * A resident ROCm layer slice already accounts its exact model spans + * and owned graph state above. Keep proportional backend headroom for + * driver and temporary allocations without rejecting viable UMA + * slices merely because the heap is smaller than a high-memory Mac. */ + double rocm_reserve_gib = glm_graph_bytes_to_gib(budget_base) / 16.0; + if (rocm_reserve_gib < 8.0) rocm_reserve_gib = 8.0; + if (rocm_reserve_gib < default_reserve_gib) { + default_reserve_gib = rocm_reserve_gib; + } + } +#endif + const double reserve_gib = + glm_graph_env_double("DS4_GLM_MEMORY_GUARD_RESERVE_GB", + default_reserve_gib, + 0.0, + 1024.0); + const uint64_t fraction_budget = (uint64_t)((double)budget_base * fraction); + const uint64_t reserve_bytes = + (uint64_t)(reserve_gib * 1024.0 * 1024.0 * 1024.0); + const uint64_t reserve_budget = + reserve_bytes >= budget_base ? 0 : budget_base - reserve_bytes; + uint64_t budget = fraction_budget; + if (reserve_bytes != 0 && reserve_budget < budget) budget = reserve_budget; + if (wired_limit != 0) { + /* An explicitly raised iogpu.wired_limit_mb is the user granting + * the GPU that much wired memory; it overrides the heuristics + * (keep a small margin for non-model GPU allocations). */ + const uint64_t margin = 2ull * 1024ull * 1024ull * 1024ull; + const uint64_t wired_budget = + wired_limit > margin ? wired_limit - margin : wired_limit; + if (wired_budget > budget) budget = wired_budget; + } + + if (required <= budget) { + const char *report = getenv("DS4_GLM_MEMORY_GUARD_REPORT"); + if (report && report[0]) { + fprintf(stderr, + "ds4: GLM memory guard ctx=%u compact_cap=%u required=%.2f GiB " + "budget=%.2f GiB (model %.2f GiB, graph %.2f GiB, transient %.2f GiB)\n", + ctx_size, + mem.comp_cap, + glm_graph_bytes_to_gib(required), + glm_graph_bytes_to_gib(budget), + glm_graph_bytes_to_gib(model_bytes), + glm_graph_bytes_to_gib(graph_bytes), + glm_graph_bytes_to_gib(transient_extra_bytes)); + if (ssd_streaming && model_bytes != model->size) { + fprintf(stderr, + "ds4: GLM streaming guard uses active model span %.2f GiB " + "(full GGUF %.2f GiB)\n", + glm_graph_bytes_to_gib(model_bytes), + glm_graph_bytes_to_gib(model->size)); + } else if (load_slice && model_bytes != model->size) { + fprintf(stderr, + "ds4: GLM memory guard uses sliced model span %.2f GiB " + "(full GGUF %.2f GiB)\n", + glm_graph_bytes_to_gib(model_bytes), + glm_graph_bytes_to_gib(model->size)); + } + } + return true; + } + + fprintf(stderr, + "ds4: GLM memory guard refused ctx=%u compact_cap=%u %s\n", + ctx_size, + mem.comp_cap, + phase ? phase : "before Metal graph allocation"); + if (ssd_streaming && model_bytes != model->size) { + fprintf(stderr, + "ds4: streamed active model map: %.2f GiB " + "(full GGUF %.2f GiB)\n", + glm_graph_bytes_to_gib(model_bytes), + glm_graph_bytes_to_gib(model->size)); + } else if (load_slice && model_bytes != model->size) { + fprintf(stderr, + "ds4: sliced model map: %.2f GiB " + "(full GGUF %.2f GiB)\n", + glm_graph_bytes_to_gib(model_bytes), + glm_graph_bytes_to_gib(model->size)); + } else { + fprintf(stderr, + "ds4: model map: %.2f GiB\n", + glm_graph_bytes_to_gib(model_bytes)); + } + fprintf(stderr, + "ds4: graph cache/scratch: %.2f GiB " + "(full KV %.2f GiB, compact DSA %.2f GiB, scratch %.2f GiB)\n", + glm_graph_bytes_to_gib(graph_bytes), + glm_graph_bytes_to_gib(mem.raw_bytes), + glm_graph_bytes_to_gib(mem.compressed_bytes), + glm_graph_bytes_to_gib(mem.scratch_bytes)); + fprintf(stderr, + "ds4: required model+graph: %.2f GiB; guard budget: %.2f GiB " + "(base %.2f GiB, fraction %.2f, reserve %.2f GiB, transient %.2f GiB)\n", + glm_graph_bytes_to_gib(required), + glm_graph_bytes_to_gib(budget), + glm_graph_bytes_to_gib(budget_base), + fraction, + reserve_gib, + glm_graph_bytes_to_gib(transient_extra_bytes)); + fprintf(stderr, + "ds4: set DS4_GLM_MEMORY_GUARD=0 to bypass, use a smaller --ctx, " + "or use SSD streaming\n"); + return false; +} + +static bool glm_graph_memory_guard( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + uint32_t ctx_size) { + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx_size, ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); + return glm_graph_memory_guard_for_compact_cap( + model, + weights, + ssd_streaming, + false, + 0, + 0, + true, + true, + ctx_size, + compact_cap, + 0, + "before GLM graph allocation"); +} + +static bool glm_graph_memory_guard_with_transient( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + uint32_t ctx_size, + uint64_t transient_extra_bytes, + const char *phase) { + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx_size, ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); + return glm_graph_memory_guard_for_compact_cap( + model, + weights, + ssd_streaming, + false, + 0, + 0, + true, + true, + ctx_size, + compact_cap, + transient_extra_bytes, + phase); +} + +static bool glm_graph_memory_guard_slice( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + uint32_t layer_start, + uint32_t layer_end, + bool include_token, + bool include_output, + uint32_t ctx_size) { + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx_size, ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); + return glm_graph_memory_guard_for_compact_cap( + model, + weights, + ssd_streaming, + true, + layer_start, + layer_end, + include_token, + include_output, + ctx_size, + compact_cap, + 0, + "before GLM graph allocation"); +} + +static bool glm_graph_memory_guard_slice_with_transient( + const ds4_model *model, + const ds4_weights *weights, + bool ssd_streaming, + uint32_t layer_start, + uint32_t layer_end, + bool include_token, + bool include_output, + uint32_t ctx_size, + uint64_t transient_extra_bytes, + const char *phase) { + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx_size, ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx_size, work_ctx); + return glm_graph_memory_guard_for_compact_cap( + model, + weights, + ssd_streaming, + true, + layer_start, + layer_end, + include_token, + include_output, + ctx_size, + compact_cap, + transient_extra_bytes, + phase); +} + +static uint32_t glm_graph_full_attention_cap(uint32_t ctx_size, + bool ssd_streaming) { + uint32_t cap = ssd_streaming ? + DS4_GLM_METAL_STREAMING_FULL_ATTN_CONTEXT : + DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT; + if (ctx_size >= DS4_GLM_METAL_LONG_CONTEXT_THRESHOLD && + cap > DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT) { + cap = DS4_GLM_METAL_LONG_CONTEXT_FULL_ATTN_CONTEXT; + } + if (ctx_size > 0 && cap > ctx_size) cap = ctx_size; + if (cap == 0) cap = 1; + return cap; +} + +static uint32_t glm_graph_full_prefill_layer_flush_interval( + uint32_t n_tokens, + uint32_t command_rows, + bool logits_requested) { + /* Tiny logits-bearing passes (MTP verify, short prefills) must NOT + * flush per layer: 76 command-buffer round-trips cost ~35ms while the + * whole pass is ~70ms of GPU work. Real prefill chunks keep the + * interactive per-layer flush. */ + return (n_tokens > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || + command_rows > DS4_GLM_METAL_FULL_ATTN_LAYER_FLUSH_CONTEXT || + (logits_requested && n_tokens > 8u)) ? 1u : 0u; +} + +static uint32_t glm_graph_prefill_progress_flush_interval( + uint32_t layer_flush_interval, + uint32_t n_tokens, + ds4_session_progress_fn display_progress, + uint32_t work_total) { + if (layer_flush_interval != 0) return layer_flush_interval; + (void)n_tokens; + (void)display_progress; + (void)work_total; + return 0; +} + +static void glm_graph_report_prefill_display_progress( + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t absolute_base, + uint32_t work_done_base, + uint32_t n_tokens, + uint32_t layer_done, + uint32_t normal_layers, + uint32_t work_total, + bool allow_complete) { + if (!display_progress || work_total == 0) return; + + uint64_t chunk_done = 0; + if (normal_layers == 0 || layer_done >= normal_layers) { + chunk_done = n_tokens; + } else { + chunk_done = (uint64_t)n_tokens * (uint64_t)layer_done / + (uint64_t)normal_layers; + } + + uint64_t done = (uint64_t)work_done_base + chunk_done; + if (done > (uint64_t)work_total) done = work_total; + if (!allow_complete && done >= (uint64_t)work_total) { + done = work_total > 0 ? (uint64_t)work_total - 1u : 0u; + } + display_progress(display_progress_ud, + "prefill_display", + (int)((uint64_t)absolute_base + done), + (int)((uint64_t)absolute_base + (uint64_t)work_total)); +} + +static bool glm_graph_small_prefill_stage_sync( + uint32_t n_tokens, + bool logits_requested) { + return logits_requested && + n_tokens > 0 && + n_tokens <= DS4_GLM_METAL_SMALL_PREFILL_STAGE_SYNC_TOKENS; +} + +static uint32_t glm_graph_indexed_decode_split_min_block_rows(void) { + return 32u; +} + +static uint32_t glm_graph_indexed_decode_split_blocks(void) { + const uint32_t block_rows = glm_graph_indexed_decode_split_min_block_rows(); + const uint32_t top_k = glm_graph_indexer_top_k_limit(); + return (top_k + block_rows - 1u) / block_rows; +} + +static uint32_t glm_graph_indexed_decode_split_block_rows_for(uint32_t n_selected) { + return n_selected <= 1024u ? 32u : 128u; +} + +static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) { + const uint32_t block_rows = glm_graph_indexed_decode_split_block_rows_for(n_selected); + const uint32_t needed_blocks = + block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + return n_selected > 512u && + block_rows > 0 && + needed_blocks > 0 && + needed_blocks <= glm_graph_indexed_decode_split_blocks() && + glm_graph_indexed_decode_split_blocks() <= 64u && + (DS4_N_HEAD % 8u) == 0 && + DS4_N_KV_LORA == 512u && + DS4_N_ROT == 64u && + glm_graph_compact_cache_is_f16(); +} + +static bool glm_graph_prefill_stage_sync_boundary(void) { + if (ds4_gpu_end_commands() == 0) return false; + return ds4_gpu_begin_commands() != 0; +} + +static bool glm_graph_indexed_prefill_attention_boundary(void) { +#ifdef DS4_ROCM_BUILD + /* + * ROCm launches in this path are ordered on the default stream. The Metal + * backend still needs the encoder flush, but on ROCm it is a full-device + * synchronize and stalls every indexed-prefill layer. + */ + return true; +#else + return ds4_gpu_flush_encoder() != 0; +#endif +} + +static DS4_MAYBE_UNUSED bool glm_graph_env_truthy(const char *env) { + return env && + env[0] && + strcmp(env, "0") != 0 && + strcasecmp(env, "false") != 0 && + strcasecmp(env, "off") != 0 && + strcasecmp(env, "no") != 0; +} + +static bool glm_graph_streaming_prefill_sync_each_layer( + bool full_layer_prefill) { +#ifdef DS4_ROCM_BUILD + /* + * ROCm command boundaries are full device synchronizes. Compact streaming + * prefill can keep queued default-stream work alive across layer mappings: + * streamed model-range eviction synchronizes before freeing ranges, while + * selected-expert cache reuse/eviction is protected by reuse events. The + * full-layer expert cache is only double-buffered, so keep its old boundary. + */ + if (full_layer_prefill) return true; + const char *env = glm_graph_env_value( + "DS4_ROCM_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER", + "DS4_METAL_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER"); + if (!env) env = getenv("DS4_GLM_STREAMING_PREFILL_SYNC_EACH_LAYER"); + return glm_graph_env_truthy(env); +#else + (void)full_layer_prefill; + return true; +#endif +} + +static bool glm_graph_indexed_prefill_batch_available( + const ds4_glm_gpu_graph *g) { + return g && + g->compact_cache_cap != 0 && + g->indexed_prefill_cap != 0 && + g->indexed_prefill_score_cap != 0 && + g->batch_indexer_q && + g->batch_indexer_weights && + g->batch_indexer_scores && + g->batch_indexer_selected && + g->batch_qk_low && + g->batch_attn_lora; +} + +static bool glm_graph_indexed_prefill_batch_ready( + const ds4_glm_gpu_graph *g, + uint32_t pos) { + return glm_graph_indexed_prefill_batch_available(g) && + (!g->full_kv_cache || pos >= g->ctx_cap); +} + +static uint32_t glm_graph_limit_indexed_prefill_chunk( + uint32_t pos, + uint32_t chunk) { + const uint32_t top_k = glm_graph_indexer_top_k_limit(); + if (pos < top_k) { + const uint32_t bridge = top_k - pos; + if (bridge != 0 && chunk > bridge) chunk = bridge; + } + return chunk; +} + +static uint32_t glm_graph_indexed_prefill_chunk_tokens( + uint32_t full_attention_cap, + uint32_t compact_cap) { + (void)full_attention_cap; + uint32_t chunk = DS4_GLM_METAL_INDEXED_PREFILL_CHUNK_TOKENS; + if (compact_cap > 0 && chunk > compact_cap) chunk = compact_cap; + if (chunk == 0) chunk = 1; + return chunk; +} + +static uint32_t glm_graph_indexed_prefill_score_tokens( + uint32_t indexed_prefill_cap, + uint32_t compact_cap) { + if (indexed_prefill_cap == 0 || compact_cap == 0) return 0; + const uint32_t scratch_mb = DS4_GLM_METAL_INDEXED_PREFILL_SCORE_SCRATCH_MB; + const uint64_t budget_bytes = (uint64_t)scratch_mb * 1024ull * 1024ull; + uint64_t budget_rows = budget_bytes / ((uint64_t)compact_cap * sizeof(float)); + if (budget_rows == 0) budget_rows = 1; + if (budget_rows > indexed_prefill_cap) budget_rows = indexed_prefill_cap; + if (budget_rows > UINT32_MAX) budget_rows = UINT32_MAX; + return (uint32_t)budget_rows; +} + +static bool glm_graph_context_request(int ctx_size, uint32_t *ctx_out) { + if (!ctx_out || ctx_size <= 0) return false; + const uint32_t model_ctx = glm_graph_model_context_limit(); + if ((uint64_t)(uint32_t)ctx_size > (uint64_t)model_ctx) { + fprintf(stderr, + "ds4: GLM context %d exceeds model context %u\n", + ctx_size, + model_ctx); + return false; + } + *ctx_out = (uint32_t)ctx_size; + return true; +} + +static bool glm_graph_span_fits_context( + const ds4_glm_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens) { + return g && n_tokens > 0 && pos0 < g->ctx_size && n_tokens <= g->ctx_size - pos0; +} + +static bool glm_graph_span_fits_full_attention( + const ds4_glm_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens) { + return g && n_tokens > 0 && pos0 < g->ctx_cap && n_tokens <= g->ctx_cap - pos0; +} + +static void glm_graph_log_full_attention_limit( + const ds4_glm_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens) { + const uint32_t end = pos0 + n_tokens; + fprintf(stderr, + "ds4: GLM Metal full-attention work cap is %u tokens; " + "requested span [%u,%u) in ctx %u needs compact indexed attention\n", + g ? g->ctx_cap : 0, + pos0, + end, + g ? g->ctx_size : 0); +} + +static bool glm_graph_tensor_layout( + const ds4_tensor *t, + uint32_t type, + uint32_t ndim, + uint64_t dim0, + uint64_t dim1, + uint64_t dim2) { + if (!t || t->type != type || t->ndim != ndim) return false; + if (ndim > 0 && t->dim[0] != dim0) return false; + if (ndim > 1 && t->dim[1] != dim1) return false; + if (ndim > 2 && t->dim[2] != dim2) return false; + return true; +} + +static bool glm_graph_dense_tensor_layout( + const ds4_tensor *t, + uint32_t ndim, + uint64_t dim0, + uint64_t dim1, + uint64_t dim2) { + if (!t || !tensor_type_is_glm_dense_quant(t->type) || t->ndim != ndim) return false; + if (ndim > 0 && t->dim[0] != dim0) return false; + if (ndim > 1 && t->dim[1] != dim1) return false; + if (ndim > 2 && t->dim[2] != dim2) return false; + return true; +} + +static bool glm_graph_layer_uses_generic_routed_moe( + const ds4_layer_weights *l) { + return l && + l->ffn_gate_exps && + l->ffn_up_exps && + l->ffn_down_exps && + l->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS; +} + +static bool glm_graph_stream_map_token( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + if (!g || !g->ssd_streaming) return true; + g->streaming_static_decode_map_current = false; + return metal_graph_stream_map_token(model, weights); +} + +static bool glm_graph_stream_layer_expert_cache_supported( + const ds4_weights *weights, + const ds4_layer_weights *l, + uint32_t il) { + if (!weights || !l) return false; + if (il < DS4_N_LEADING_DENSE) return true; + return glm_stream_decode_experts_are_streamed(weights, l, il); +} + +static bool glm_graph_stream_prefill_expert_addr_supported( + const ds4_weights *weights, + const ds4_layer_weights *l, + uint32_t il, + uint32_t n_tokens) { + if (il < DS4_N_LEADING_DENSE) return true; + if (n_tokens <= 1) return false; +#ifdef DS4_ROCM_BUILD + /* + * ROCm selected-address batch prefill has pointer kernels for the + * IQ2-gate/Q2-down generic path and the uniform Q2_K GLM path. Q4_K still + * maps the full layer until matching pointer kernels exist. + */ + if (glm_stream_selected_expert_cache_supported(l, il)) return true; + return l && + l->ffn_gate_exps && + l->ffn_up_exps && + l->ffn_down_exps && + l->ffn_gate_exps->type == DS4_TENSOR_Q2_K && + l->ffn_up_exps->type == DS4_TENSOR_Q2_K && + l->ffn_down_exps->type == DS4_TENSOR_Q2_K && + glm_stream_expert_cache_addr_layout_supported(weights, l, il); +#else + return glm_stream_expert_cache_addr_supported(weights, l, il); +#endif +} + +static bool rocm_graph_glm_stream_prefill_full_layer_enabled( + const ds4_glm_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens); + +static bool glm_graph_stream_map_decode_layer( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il) { + if (!g || !g->ssd_streaming) return true; + g->streaming_static_decode_map_current = false; + if (glm_graph_env_present("DS4_ROCM_GLM_STREAMING_DECODE_FULL_LAYER_MAP", + "DS4_METAL_GLM_STREAMING_DECODE_FULL_LAYER_MAP") || + getenv("DS4_GLM_STREAMING_DECODE_FULL_LAYER_MAP") != NULL) { + return metal_graph_stream_map_layer(model, weights, il); + } + if (weights && il < DS4_N_LAYER && + glm_stream_resident_decode_layer_enabled(&weights->layer[il], il)) { + return metal_graph_stream_map_layer(model, weights, il); + } + if (weights && il < DS4_N_LAYER && + glm_graph_stream_layer_expert_cache_supported(weights, + &weights->layer[il], + il)) { + return metal_graph_stream_map_layer_decode(model, weights, il); + } + return metal_graph_stream_map_layer(model, weights, il); +} + +static bool glm_graph_stream_map_prefill_layer( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il, + uint32_t n_tokens, + bool full_layer_prefill) { + if (!g || !g->ssd_streaming) return true; + g->streaming_static_decode_map_current = false; + if (full_layer_prefill) { + const char *trace = glm_graph_env_value("DS4_ROCM_STREAMING_MAP_TRACE", + "DS4_METAL_STREAMING_MAP_TRACE"); + if (trace && trace[0] && strcmp(trace, "0") != 0) { + fprintf(stderr, + "ds4: GLM SSD prefill map layer=%u tokens=%u mode=full-prefill\n", + il, + n_tokens); + } +#ifdef DS4_ROCM_BUILD + if (weights && + il < DS4_N_LAYER && + rocm_graph_glm_stream_prefill_full_layer_enabled(g, + &weights->layer[il], + il, + n_tokens)) { + return metal_graph_stream_map_layer_decode(model, weights, il); + } +#endif + return metal_graph_stream_map_layer(model, weights, il); + } + const bool addr_supported = + weights && il < DS4_N_LAYER && + glm_graph_stream_prefill_expert_addr_supported(weights, + &weights->layer[il], + il, + n_tokens); + const char *trace = glm_graph_env_value("DS4_ROCM_STREAMING_MAP_TRACE", + "DS4_METAL_STREAMING_MAP_TRACE"); + if (trace && trace[0] && strcmp(trace, "0") != 0) { + fprintf(stderr, + "ds4: GLM SSD prefill map layer=%u tokens=%u mode=%s\n", + il, + n_tokens, + addr_supported ? "decode-expert-cache" : "full-layer"); + } + if (addr_supported) { + return metal_graph_stream_map_layer_decode(model, weights, il); + } + return metal_graph_stream_map_layer(model, weights, il); +} + +#ifdef DS4_ROCM_BUILD +enum { DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 1024 }; +#else +enum { DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS = 64 }; +#endif + +static uint32_t glm_graph_stream_prefill_full_layer_min_tokens(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_GLM_STREAMING_PREFILL_FULL_LAYER_MIN_TOKENS", + "DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER_MIN_TOKENS"); + if (!env) return DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS; + char *end = NULL; + errno = 0; + unsigned long v = strtoul(env, &end, 10); + if (end == env || errno != 0 || v == 0 || v > UINT32_MAX) { + return DS4_GLM_STREAM_PREFILL_FULL_LAYER_MIN_TOKENS; + } + return (uint32_t)v; +} + +static bool glm_graph_stream_prefill_full_layer_enabled( + const ds4_glm_gpu_graph *g, + uint32_t n_tokens) { + if (!g || !g->ssd_streaming) return false; + if (glm_graph_env_present( + "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER", + "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER")) { + return false; + } + if (glm_graph_env_present("DS4_ROCM_GLM_STREAMING_PREFILL_FULL_LAYER", + "DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER")) { + return true; + } + return n_tokens >= glm_graph_stream_prefill_full_layer_min_tokens(); +} + +static bool glm_graph_stream_prefill_full_layer_prepare_enabled( + const ds4_glm_gpu_graph *g, + bool full_layer_prefill) { + return g && + g->ssd_streaming && + full_layer_prefill && + !glm_graph_env_present( + "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER_PREPARE", + "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_FULL_LAYER_PREPARE"); +} + +#ifdef DS4_ROCM_BUILD +static bool rocm_graph_glm_stream_prefill_full_layer_enabled( + const ds4_glm_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens) { + return glm_graph_stream_prefill_full_layer_enabled(g, n_tokens) && + layer && + glm_stream_resident_decode_layer_supported(layer, il); +} + +static bool rocm_graph_glm_stream_layer_expert_load_start_next( + rocm_graph_stream_layer_expert_load *job, + const ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t first_il, + uint32_t last_il, + uint32_t n_tokens) { + if (!job || !model || !weights || first_il > last_il) return true; + if (job->active) return true; + for (uint32_t il = first_il; il <= last_il && il < DS4_N_LAYER; il++) { + const ds4_layer_weights *layer = &weights->layer[il]; + if (!rocm_graph_glm_stream_prefill_full_layer_enabled(g, + layer, + il, + n_tokens)) { + continue; + } + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!rocm_graph_stream_layer_expert_bytes(layer, + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + return rocm_graph_stream_layer_expert_load_start(job, + model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + } + return true; +} + +static bool rocm_graph_glm_stream_layer_expert_load_ready( + rocm_graph_stream_layer_expert_load *job, + const ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t il, + uint32_t n_tokens) { + if (!model || !weights || il >= DS4_N_LAYER) return false; + const ds4_layer_weights *layer = &weights->layer[il]; + if (!rocm_graph_glm_stream_prefill_full_layer_enabled(g, + layer, + il, + n_tokens)) { + return true; + } + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!rocm_graph_stream_layer_expert_bytes(layer, + &gate_expert_bytes, + &down_expert_bytes)) { + return false; + } + if (job && job->active) { + if (job->il != il) { + fprintf(stderr, + "ds4: GLM ROCm streaming full-layer expert load expected " + "layer %u but pending job is layer %u\n", + il, + job->il); + return false; + } + return rocm_graph_stream_layer_expert_load_join(job); + } + return rocm_graph_stream_layer_expert_load_sync(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); +} +#else +static bool rocm_graph_glm_stream_prefill_full_layer_enabled( + const ds4_glm_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens) { + (void)g; + (void)layer; + (void)il; + (void)n_tokens; + return false; +} +#endif + +static bool glm_graph_stream_map_output( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + if (!g || !g->ssd_streaming) return true; + g->streaming_static_decode_map_current = false; + return metal_graph_stream_map_output(model, weights); +} + +static bool glm_graph_validate_expert_layout( + const ds4_model *model, + const ds4_tensor *gate, + const ds4_tensor *up, + const ds4_tensor *down, + uint64_t *gate_row_bytes, + uint64_t *up_row_bytes, + uint64_t *down_row_bytes) { + if (!gate || !up || !down) return false; + if (!glm_graph_gate_pair_type_supported(gate->type, up->type) || + !glm_graph_down_type_supported(down->type) || + gate->ndim != 3 || up->ndim != 3 || down->ndim != 3 || + gate->dim[0] != DS4_N_EMBD || + gate->dim[1] != DS4_N_FF_EXP || + gate->dim[2] != DS4_N_EXPERT || + up->dim[0] != DS4_N_EMBD || + up->dim[1] != DS4_N_FF_EXP || + up->dim[2] != DS4_N_EXPERT || + down->dim[0] != DS4_N_FF_EXP || + down->dim[1] != DS4_N_EMBD || + down->dim[2] != DS4_N_EXPERT) { + return false; + } + + uint64_t gate_in = 0, gate_out = 0; + uint64_t up_in = 0, up_out = 0; + uint64_t down_in = 0, down_out = 0; + (void)tensor_expert_bytes(model, gate, 0, &gate_in, &gate_out, gate_row_bytes); + (void)tensor_expert_bytes(model, up, 0, &up_in, &up_out, up_row_bytes); + (void)tensor_expert_bytes(model, down, 0, &down_in, &down_out, down_row_bytes); + return gate_in == DS4_N_EMBD && + up_in == DS4_N_EMBD && + down_in == DS4_N_FF_EXP && + gate_out == DS4_N_FF_EXP && + up_out == DS4_N_FF_EXP && + down_out == DS4_N_EMBD; +} + +static bool glm_graph_validate_layer_layout( + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint64_t q_dim, + uint64_t q_nope, + uint64_t heads_dim, + uint64_t *kv_raw_dim_out, + uint64_t *dense_hidden_max) { + if (!l) return false; + const uint64_t kv_raw_dim = l->attn_kv_a_mqa ? l->attn_kv_a_mqa->dim[1] : 0; + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + if (!glm_graph_tensor_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0) || + !glm_graph_dense_tensor_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0) || + !glm_graph_tensor_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0) || + !glm_graph_dense_tensor_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0) || + !l->attn_kv_a_mqa || + !tensor_type_is_glm_dense_quant(l->attn_kv_a_mqa->type) || + l->attn_kv_a_mqa->ndim != 2 || + l->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || + kv_raw_dim < (uint64_t)DS4_N_KV_LORA + DS4_N_ROT || + !glm_graph_tensor_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_KV_LORA, 0, 0) || + !glm_graph_dense_tensor_layout(l->attn_k_b, 3, q_nope, DS4_N_KV_LORA, DS4_N_HEAD) || + !glm_graph_dense_tensor_layout(l->attn_v_b, 3, DS4_N_KV_LORA, DS4_N_VALUE_MLA, DS4_N_HEAD) || + !glm_graph_dense_tensor_layout(l->attn_output, 2, heads_dim, DS4_N_EMBD, 0) || + !glm_graph_dense_tensor_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, indexer_q_dim, 0) || + !glm_graph_dense_tensor_layout(l->indexer_attn_k, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 0) || + !glm_graph_tensor_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0) || + !glm_graph_tensor_layout(l->indexer_k_norm_b, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0) || + !glm_graph_tensor_layout(l->indexer_proj, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0) || + !glm_graph_tensor_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0)) { + fprintf(stderr, "ds4: GLM Metal graph found unexpected attention layout in layer %u\n", il); + return false; + } + if (kv_raw_dim > *kv_raw_dim_out) *kv_raw_dim_out = kv_raw_dim; + + if (il < DS4_N_LEADING_DENSE) { + const uint64_t hidden = l->ffn_gate ? l->ffn_gate->dim[1] : 0; + if (!l->ffn_gate || + !glm_graph_dense_tensor_layout(l->ffn_gate, 2, DS4_N_EMBD, hidden, 0) || + !glm_graph_dense_tensor_layout(l->ffn_up, 2, DS4_N_EMBD, hidden, 0) || + !glm_graph_dense_tensor_layout(l->ffn_down, 2, hidden, DS4_N_EMBD, 0)) { + fprintf(stderr, "ds4: GLM Metal graph found unexpected dense FFN layout in layer %u\n", il); + return false; + } + if (hidden > *dense_hidden_max) *dense_hidden_max = hidden; + } else { + uint64_t gate_row_bytes = 0, up_row_bytes = 0, down_row_bytes = 0; + if (!glm_graph_tensor_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0) || + !glm_graph_tensor_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0) || + !glm_graph_validate_expert_layout(model, + l->ffn_gate_exps, + l->ffn_up_exps, + l->ffn_down_exps, + &gate_row_bytes, + &up_row_bytes, + &down_row_bytes) || + !glm_graph_dense_tensor_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0) || + !glm_graph_dense_tensor_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0) || + !glm_graph_dense_tensor_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0)) { + fprintf(stderr, "ds4: GLM Metal graph found unexpected sparse FFN layout in layer %u\n", il); + return false; + } + (void)gate_row_bytes; + (void)up_row_bytes; + (void)down_row_bytes; + } + return true; +} + +static bool glm_graph_validate_layout( + const ds4_model *model, + const ds4_weights *weights, + ds4_glm_gpu_graph *g, + uint32_t layer_start, + uint32_t layer_end, + bool require_token_embd, + bool require_output_head) { + if (!model || !weights || !g) return false; + const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0 || DS4_N_ROT >= DS4_N_KEY_MLA) { + fprintf(stderr, "ds4: GLM Metal graph found unsupported layer/key dimensions\n"); + return false; + } + if (layer_end == UINT32_MAX) layer_end = normal_layers - 1u; + if (layer_start > layer_end || layer_end >= normal_layers) { + fprintf(stderr, + "ds4: GLM Metal graph found invalid layer slice %u:%u for %u normal layers\n", + layer_start, + layer_end, + normal_layers); + return false; + } + + g->has_token_embd = weights->token_embd != NULL; + g->has_output_head = weights_have_output_head(weights); + if (require_token_embd && !g->has_token_embd) { + fprintf(stderr, "ds4: GLM Metal graph layer slice requires token embeddings\n"); + return false; + } + if (g->has_token_embd && + !glm_graph_dense_tensor_layout(weights->token_embd, 2, + DS4_N_EMBD, DS4_N_VOCAB, 0)) { + fprintf(stderr, "ds4: GLM Metal graph found unexpected token embedding layout\n"); + return false; + } + if (require_output_head && !g->has_output_head) { + fprintf(stderr, "ds4: GLM Metal graph layer slice requires the output head\n"); + return false; + } + if (weights_have_partial_output_head(weights) && !g->has_output_head) { + fprintf(stderr, "ds4: GLM Metal graph found partial output head\n"); + return false; + } + if (g->has_output_head && + (!glm_graph_tensor_layout(weights->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0) || + !glm_graph_dense_tensor_layout(weights->output, 2, + DS4_N_EMBD, DS4_N_VOCAB, 0))) { + fprintf(stderr, "ds4: GLM Metal graph found unexpected output head layout\n"); + return false; + } + + g->normal_layers = normal_layers; + g->layer_start = layer_start; + g->layer_end = layer_end; + g->layer_count = layer_end - layer_start + 1u; + g->q_dim = (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA; + g->q_nope = (uint64_t)DS4_N_KEY_MLA - DS4_N_ROT; + g->heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; + g->dense_hidden_max = DS4_N_FF_EXP; + g->kv_raw_dim = 0; + for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { + if (!glm_graph_validate_layer_layout(model, + &weights->layer[il], + il, + g->q_dim, + g->q_nope, + g->heads_dim, + &g->kv_raw_dim, + &g->dense_hidden_max)) { + return false; + } + if (glm_graph_layer_uses_generic_routed_moe(&weights->layer[il])) { + g->generic_routed_moe = true; + } + } + const uint64_t sparse_mid_elems = (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; + g->ffn_mid_elems = + g->dense_hidden_max > sparse_mid_elems ? g->dense_hidden_max : sparse_mid_elems; + return g->layer_count > 0 && + g->kv_raw_dim >= (uint64_t)DS4_N_KV_LORA + DS4_N_ROT; +} + +/* Per-tier decode scratch keeps layer kernels off peer-mapped work buffers. */ +#define DS4_GLM_WS_FIELDS(X) \ + X(cur) X(next) X(attn_norm) X(q_rank) X(q_rank_norm) X(q) X(kv_raw) \ + X(kv_norm) X(k_nope) X(value) X(heads) X(attn_out) X(after_attn) \ + X(ffn_norm) X(ffn_gate) X(ffn_up) X(ffn_mid) X(ffn_out) X(ffn_sum) \ + X(router_logits) X(router_probs) X(router_selected) X(router_weights) \ + X(indexer_k) X(indexer_q) X(indexer_weights) X(indexer_scores) \ + X(indexer_selected) X(qk_low) + +static ds4_gpu_tensor **glm_graph_ws_slot(ds4_glm_gpu_graph *g, int i) { + int n = 0; +#define DS4_GLM_WS_SLOT_CASE(field) if (n++ == i) return &g->field; + DS4_GLM_WS_FIELDS(DS4_GLM_WS_SLOT_CASE) +#undef DS4_GLM_WS_SLOT_CASE + return NULL; +} + +static void glm_graph_ws_free(ds4_glm_gpu_graph *g) { + if (!g || !g->ws_ready) return; + for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); + if (slot) *slot = g->ws_orig[i]; + } + for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { + for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { + ds4_gpu_tensor_free(g->ws_mirror[tier][i]); + g->ws_mirror[tier][i] = NULL; + } + } + g->ws_ready = 0; + g->ws_tier = -1; +} + +static void glm_graph_ws_init(ds4_glm_gpu_graph *g) { + g->ws_ready = 0; + g->ws_tier = -1; + if (!g->placement) return; + bool used[DS4_MAX_GPUS] = { false }; + if (g->placement[0] >= 0 && g->placement[0] < DS4_MAX_GPUS) { + used[g->placement[0]] = true; + } + for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { + const int tier = g->placement[il + 1u]; + if (tier >= 0 && tier < DS4_MAX_GPUS) used[tier] = true; + } + for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); + if (!slot) return; + g->ws_orig[i] = *slot; + } + for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { + if (!used[tier]) continue; + for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { + const ds4_gpu_tensor *orig = g->ws_orig[i]; + if (!orig) continue; + g->ws_mirror[tier][i] = + ds4_gpu_tensor_alloc_ptr_on(tier, ds4_gpu_tensor_bytes(orig)); + if (!g->ws_mirror[tier][i]) { + fprintf(stderr, + "ds4: GLM per-tier working-set alloc failed (tier %d); " + "falling back to base buffers\n", + tier); + g->ws_ready = 1; + glm_graph_ws_free(g); + return; + } + } + } + g->ws_ready = 1; +} + +static bool glm_graph_ws_switch(ds4_glm_gpu_graph *g, + int tier, + bool carry_hidden) { + if (!g->placement) return true; + if (tier < 0 || tier >= DS4_MAX_GPUS || + ds4_gpu_set_current_device_fenced(tier) != 0) { + return false; + } + if (!g->ws_ready) return true; + if (g->ws_tier == tier) return true; + ds4_gpu_tensor *old_cur = g->cur; + for (int i = 0; i < DS4_GLM_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_ws_slot(g, i); + if (slot && g->ws_mirror[tier][i]) { + *slot = g->ws_mirror[tier][i]; + } + } + if (carry_hidden && old_cur && g->cur && g->cur != old_cur) { + const uint64_t bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + if (ds4_gpu_tensor_copy_async(g->cur, old_cur, bytes) == 0) { + return false; + } + } + g->ws_tier = tier; + return true; +} + +#define DS4_GLM_VERIFY_WS_FIELDS(X) \ + X(batch_cur) X(batch_next) X(batch_attn_norm) X(batch_q_rank) \ + X(batch_q_rank_norm) X(batch_q) X(batch_indexer_k) X(batch_kv_raw) \ + X(batch_kv_norm) X(batch_qk_low) X(batch_attn_lora) \ + X(batch_indexer_selected) X(batch_heads) X(batch_attn_out) \ + X(batch_after_attn) X(batch_ffn_norm) X(batch_ffn_gate) X(batch_ffn_up) \ + X(batch_shared_mid) X(batch_ffn_mid) X(batch_routed_gate) \ + X(batch_routed_up) X(batch_routed_down) X(batch_ffn_out) \ + X(batch_router_logits) X(batch_router_probs) X(batch_router_selected) \ + X(batch_router_weights) + +static ds4_gpu_tensor **glm_graph_verify_ws_slot(ds4_glm_gpu_graph *g, + int i) { + int n = 0; +#define DS4_GLM_VERIFY_WS_SLOT_CASE(field) if (n++ == i) return &g->field; + DS4_GLM_VERIFY_WS_FIELDS(DS4_GLM_VERIFY_WS_SLOT_CASE) +#undef DS4_GLM_VERIFY_WS_SLOT_CASE + return NULL; +} + +static void glm_graph_verify_ws_restore(ds4_glm_gpu_graph *g) { + if (!g || !g->verify_ws_ready) return; + for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); + if (slot) *slot = g->verify_ws_orig[i]; + } + g->verify_ws_tier = -1; +} + +static void glm_graph_verify_ws_free(ds4_glm_gpu_graph *g) { + if (!g || !g->verify_ws_ready) return; + glm_graph_verify_ws_restore(g); + for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { + for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { + ds4_gpu_tensor_free(g->verify_ws_mirror[tier][i]); + g->verify_ws_mirror[tier][i] = NULL; + } + } + memset(g->verify_ws_orig, 0, sizeof(g->verify_ws_orig)); + g->verify_ws_ready = 0; +} + +static bool glm_graph_verify_ws_init(ds4_glm_gpu_graph *g) { + if (!g) return false; + if (g->verify_ws_ready) return true; + if (!g->placement) { + g->verify_ws_ready = 1; + g->verify_ws_tier = -1; + return true; + } + bool used[DS4_MAX_GPUS] = { false }; + for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { + const int tier = g->placement[il + 1u]; + if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + used[tier] = true; + } + for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); + if (!slot) return false; + g->verify_ws_orig[i] = *slot; + } + for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { + if (!used[tier]) continue; + for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { + const ds4_gpu_tensor *orig = g->verify_ws_orig[i]; + if (!orig) continue; + const uint32_t cap = (i >= 9 && i <= 11) ? + g->indexed_prefill_cap : g->ctx_cap; + const uint64_t orig_bytes = ds4_gpu_tensor_bytes(orig); + if (cap == 0 || orig_bytes % cap != 0 || + orig_bytes / cap > UINT64_MAX / 2u) { + g->verify_ws_ready = 1; + glm_graph_verify_ws_free(g); + return false; + } + const uint64_t bytes = (orig_bytes / cap) * 2u; + g->verify_ws_mirror[tier][i] = + ds4_gpu_tensor_alloc_ptr_on(tier, bytes); + if (!g->verify_ws_mirror[tier][i]) { + g->verify_ws_ready = 1; + glm_graph_verify_ws_free(g); + return false; + } + } + } + g->verify_ws_ready = 1; + g->verify_ws_tier = -1; + return true; +} + +static bool glm_graph_verify_ws_switch(ds4_glm_gpu_graph *g, + int tier, + bool carry_hidden, + uint32_t n_rows) { + if (!g || n_rows == 0 || n_rows > 2) return false; + if (!g->placement) return true; + if (!glm_graph_verify_ws_init(g) || + tier < 0 || tier >= DS4_MAX_GPUS || + ds4_gpu_set_current_device_fenced(tier) != 0) { + return false; + } + if (g->verify_ws_tier == tier) return true; + ds4_gpu_tensor *old_cur = g->batch_cur; + for (int i = 0; i < DS4_GLM_VERIFY_WS_SLOTS; i++) { + ds4_gpu_tensor **slot = glm_graph_verify_ws_slot(g, i); + if (slot && g->verify_ws_mirror[tier][i]) { + *slot = g->verify_ws_mirror[tier][i]; + } + } + if (carry_hidden && old_cur && g->batch_cur && old_cur != g->batch_cur) { + const uint64_t bytes = + (uint64_t)n_rows * DS4_N_EMBD * sizeof(float); + if (ds4_gpu_tensor_copy_async(g->batch_cur, old_cur, bytes) == 0) { + return false; + } + } + g->verify_ws_tier = tier; + return true; +} +#undef DS4_GLM_VERIFY_WS_FIELDS + +static void glm_graph_free(ds4_glm_gpu_graph *g) { + glm_graph_verify_ws_free(g); + glm_graph_ws_free(g); + if (!g) return; + ds4_gpu_tensor_free(g->mtp_kv_lora_cache); + ds4_gpu_tensor_free(g->mtp_k_rope_cache); + ds4_gpu_tensor_free(g->mtp_concat); + ds4_gpu_tensor_free(g->mtp_selected); + free(g->mtp_logits_host); + g->mtp_kv_lora_cache = NULL; + g->mtp_k_rope_cache = NULL; + g->mtp_concat = NULL; + g->mtp_selected = NULL; + g->mtp_logits_host = NULL; + g->mtp_ready = 0; + for (uint32_t il = 0; il < DS4_MAX_LAYER; il++) { + ds4_gpu_tensor_free(g->layer_indexer_key_cache[il]); + ds4_gpu_tensor_free(g->layer_k_rope_cache[il]); + ds4_gpu_tensor_free(g->layer_kv_lora_cache[il]); + ds4_gpu_tensor_free(g->layer_value_cache[il]); + ds4_gpu_tensor_free(g->layer_key_cache[il]); + } + ds4_gpu_tensor_free(g->logits); + ds4_gpu_tensor_free(g->batch_router_weights); + ds4_gpu_tensor_free(g->prefill_seed_router_selected); + ds4_gpu_tensor_free(g->batch_router_selected); + ds4_gpu_tensor_free(g->batch_router_probs); + ds4_gpu_tensor_free(g->batch_router_logits); + ds4_gpu_tensor_free(g->batch_routed_down); + ds4_gpu_tensor_free(g->batch_routed_up); + ds4_gpu_tensor_free(g->batch_routed_gate); + ds4_gpu_tensor_free(g->batch_ffn_out); + ds4_gpu_tensor_free(g->batch_ffn_mid); + ds4_gpu_tensor_free(g->batch_shared_mid); + ds4_gpu_tensor_free(g->batch_ffn_up); + ds4_gpu_tensor_free(g->batch_ffn_gate); + ds4_gpu_tensor_free(g->batch_ffn_norm); + ds4_gpu_tensor_free(g->batch_after_attn); + ds4_gpu_tensor_free(g->batch_attn_out); + ds4_gpu_tensor_free(g->batch_heads); + ds4_gpu_tensor_free(g->batch_value); + ds4_gpu_tensor_free(g->batch_k_nope); + ds4_gpu_tensor_free(g->batch_kv_norm); + ds4_gpu_tensor_free(g->batch_kv_raw); + ds4_gpu_tensor_free(g->batch_attn_lora); + ds4_gpu_tensor_free(g->batch_qk_low); + ds4_gpu_tensor_free(g->batch_indexer_selected); + ds4_gpu_tensor_free(g->batch_indexer_scores); + ds4_gpu_tensor_free(g->batch_indexer_weights); + ds4_gpu_tensor_free(g->batch_indexer_q); + ds4_gpu_tensor_free(g->batch_indexer_k); + ds4_gpu_tensor_free(g->batch_q); + ds4_gpu_tensor_free(g->batch_q_rank_norm); + ds4_gpu_tensor_free(g->batch_q_rank); + ds4_gpu_tensor_free(g->batch_attn_norm); + ds4_gpu_tensor_free(g->batch_next); + ds4_gpu_tensor_free(g->batch_cur); + ds4_gpu_tensor_free(g->prefill_tokens); + ds4_gpu_tensor_free(g->output_norm); + ds4_gpu_tensor_free(g->router_weights); + ds4_gpu_tensor_free(g->router_selected); + ds4_gpu_tensor_free(g->router_probs); + ds4_gpu_tensor_free(g->router_logits); + ds4_gpu_tensor_free(g->ffn_sum); + ds4_gpu_tensor_free(g->ffn_out); + ds4_gpu_tensor_free(g->routed_down); + ds4_gpu_tensor_free(g->routed_up); + ds4_gpu_tensor_free(g->tp_bounce_out); + ds4_gpu_tensor_free(g->tp_bounce_in); + ds4_gpu_tensor_free(g->routed_gate); + ds4_gpu_tensor_free(g->ffn_mid); + ds4_gpu_tensor_free(g->ffn_up); + ds4_gpu_tensor_free(g->ffn_gate); + ds4_gpu_tensor_free(g->ffn_norm); + ds4_gpu_tensor_free(g->after_attn); + ds4_gpu_tensor_free(g->attn_out); + ds4_gpu_tensor_free(g->heads); + ds4_gpu_tensor_free(g->value); + ds4_gpu_tensor_free(g->k_nope); + ds4_gpu_tensor_free(g->kv_norm); + ds4_gpu_tensor_free(g->kv_raw); + ds4_gpu_tensor_free(g->attn_partial_ms); + ds4_gpu_tensor_free(g->attn_partial_lora); + ds4_gpu_tensor_free(g->qk_low); + ds4_gpu_tensor_free(g->indexer_selected); + ds4_gpu_tensor_free(g->indexer_scores); + ds4_gpu_tensor_free(g->indexer_weights); + ds4_gpu_tensor_free(g->indexer_q); + ds4_gpu_tensor_free(g->indexer_k); + ds4_gpu_tensor_free(g->q); + ds4_gpu_tensor_free(g->q_rank_norm); + ds4_gpu_tensor_free(g->q_rank); + ds4_gpu_tensor_free(g->attn_norm); + ds4_gpu_tensor_free(g->next); + ds4_gpu_tensor_free(g->cur); + memset(g, 0, sizeof(*g)); +} + +static bool glm_graph_ensure_compact_cache( + const ds4_glm_gpu_graph *g, + uint32_t needed_rows) { + if (!g || needed_rows == 0) return false; + if (needed_rows <= g->ctx_cap && g->compact_cache_cap == 0) return true; + if (needed_rows <= g->compact_cache_cap) return true; + fprintf(stderr, + "ds4: GLM compact DSA cache capacity %u is smaller than required row %u " + "(ctx=%u, full_cap=%u)\n", + g->compact_cache_cap, + needed_rows, + g->ctx_size, + g->ctx_cap); + return false; +} + +static bool glm_graph_warm_compact_indexer_store( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t warm_pos) { + if (!g || !model || !weights) return false; + if (g->compact_cache_cap == 0 || g->indexer_full_layers == 0) return true; + if (!g->indexer_k) return false; + if (warm_pos >= g->compact_cache_cap) warm_pos = g->compact_cache_cap - 1u; + + if (ds4_gpu_tensor_fill_f32(g->indexer_k, + 0.0f, + DS4_N_INDEXER_HEAD_DIM) == 0) { + return false; + } + + const bool profile = false; + const double t0 = 0.0; + bool ok = ds4_gpu_begin_commands() != 0; + uint32_t warmed = 0; + for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { + if (!glm_graph_layer_uses_full_indexer(il)) continue; + const ds4_layer_weights *l = &weights->layer[il]; + if (!g->layer_indexer_key_cache[il] || + !l->indexer_k_norm || + !l->indexer_k_norm_b) { + ok = false; + break; + } + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + ok = ds4_gpu_glm_store_indexer_k_tensor( + g->layer_indexer_key_cache[il], + g->indexer_k, + model->map, + model->size, + l->indexer_k_norm->abs_offset, + l->indexer_k_norm_b->abs_offset, + warm_pos, + 1, + g->compact_cache_cap, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + 0, + 1.0e-6f, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + glm_graph_compact_cache_is_f16()) != 0; + if (ok) warmed++; + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + + if (profile) { + fprintf(stderr, + "ds4: GLM compact indexer warmup pos=%u layers=%u %.3f ms\n", + warm_pos, + warmed, + (now_sec() - t0) * 1000.0); + } + return ok; +} + +static bool glm_graph_alloc_slice( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int ctx_size, + bool ssd_streaming, + bool ssd_streaming_cold, + uint64_t streaming_transient_guard_bytes, + uint32_t layer_start, + uint32_t layer_end, + bool require_token_embd, + bool require_output_head) { + if (!g || !model || !weights || ctx_size <= 0) return false; + const int *placement = g->placement; + memset(g, 0, sizeof(*g)); + g->placement = placement; + g->ssd_streaming = ssd_streaming; + g->ssd_streaming_cold = ssd_streaming_cold; + + if (!glm_graph_context_request(ctx_size, &g->ctx_size)) return false; + if (!glm_graph_memory_guard_slice_with_transient( + model, + weights, + g->ssd_streaming, + layer_start, + layer_end, + require_token_embd, + require_output_head, + g->ctx_size, + streaming_transient_guard_bytes, + "before GLM graph allocation")) { + return false; + } + if (!glm_graph_validate_layout(model, + weights, + g, + layer_start, + layer_end, + require_token_embd, + require_output_head)) { + return false; + } + g->ctx_cap = glm_graph_full_attention_cap(g->ctx_size, + g->ssd_streaming); + g->full_kv_cache = glm_graph_expanded_kv_cache_enabled(g->ssd_streaming); + g->compact_cache_cap = + glm_graph_compact_cache_initial_cap(g->ctx_size, g->ctx_cap); + g->indexed_prefill_cap = + g->compact_cache_cap != 0 ? + glm_graph_indexed_prefill_chunk_tokens(g->ctx_cap, g->compact_cache_cap) : + 0; + g->indexed_prefill_score_cap = + glm_graph_indexed_prefill_score_tokens(g->indexed_prefill_cap, + g->compact_cache_cap); + g->indexer_full_layers = + glm_graph_full_indexer_layer_count_range(g->layer_start, + g->layer_end); + if (g->ctx_size > g->ctx_cap) { + fprintf(stderr, + "ds4: GLM Metal session ctx=%u (model max=%u); " + "full-attention prefill/work cap=%u; compact indexed decode is used beyond the cap\n", + g->ctx_size, + glm_graph_model_context_limit(), + g->ctx_cap); + } + + const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + const uint64_t q_rank_bytes = (uint64_t)DS4_N_LORA_Q * sizeof(float); + const uint64_t q_bytes = g->q_dim * sizeof(float); + const uint64_t kv_raw_bytes = g->kv_raw_dim * sizeof(float); + const uint64_t kv_norm_bytes = (uint64_t)DS4_N_KV_LORA * sizeof(float); + const uint64_t k_nope_bytes = (uint64_t)DS4_N_HEAD * g->q_nope * sizeof(float); + const uint64_t heads_bytes = g->heads_dim * sizeof(float); + const uint64_t indexer_k_bytes = (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float); + const uint64_t indexer_q_bytes = + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float); + const uint64_t indexer_weights_bytes = + (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); + const uint64_t indexer_work_cap = + g->compact_cache_cap != 0 ? g->compact_cache_cap : g->ctx_cap; + const uint64_t indexer_scores_bytes = indexer_work_cap * sizeof(float); + const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); + const uint64_t indexer_selected_bytes = + (uint64_t)indexer_top_k * sizeof(uint32_t); + const uint64_t qk_low_bytes = + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float); + const uint32_t split_attn_blocks = glm_graph_indexed_decode_split_blocks(); + const uint64_t attn_partial_lora_bytes = + (uint64_t)split_attn_blocks * qk_low_bytes; + const uint64_t attn_partial_ms_bytes = + (uint64_t)split_attn_blocks * DS4_N_HEAD * 2u * sizeof(float); + const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); + const uint64_t full_kv_elem_bytes = glm_graph_full_kv_cache_elem_bytes(); + const uint64_t key_cache_bytes = g->full_kv_cache ? + (uint64_t)g->ctx_cap * g->q_dim * full_kv_elem_bytes : 0; + const uint64_t value_cache_bytes = g->full_kv_cache ? + (uint64_t)g->ctx_cap * g->heads_dim * full_kv_elem_bytes : 0; + const uint64_t compact_kv_lora_bytes = + (uint64_t)g->compact_cache_cap * DS4_N_KV_LORA * + glm_graph_compact_cache_elem_bytes(); + const uint64_t compact_k_rope_bytes = + (uint64_t)g->compact_cache_cap * DS4_N_ROT * + glm_graph_compact_cache_elem_bytes(); + const uint64_t compact_indexer_key_bytes = + (uint64_t)g->compact_cache_cap * DS4_N_INDEXER_HEAD_DIM * + glm_graph_compact_cache_elem_bytes(); + const uint64_t batch_rows = + g->full_kv_cache || g->indexed_prefill_cap == 0 ? + g->ctx_cap : + g->indexed_prefill_cap; + const uint64_t indexed_batch_rows = g->indexed_prefill_cap; + const uint64_t indexed_score_rows = g->indexed_prefill_score_cap; + const uint64_t batch_indexer_q_bytes = indexed_batch_rows * indexer_q_bytes; + const uint64_t batch_indexer_weights_bytes = indexed_batch_rows * indexer_weights_bytes; + const uint64_t batch_indexer_scores_bytes = + indexed_score_rows * indexer_work_cap * sizeof(float); + const uint64_t batch_indexer_selected_bytes = + indexed_batch_rows * indexer_top_k * sizeof(uint32_t); + const uint64_t batch_qk_low_bytes = indexed_batch_rows * qk_low_bytes; + const uint64_t batch_attn_lora_bytes = + indexed_batch_rows * (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float); + const uint64_t routed_mid_bytes = + (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float); + const uint64_t routed_down_bytes = + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); + const double cache_gib = + (double)(g->layer_count * (key_cache_bytes + value_cache_bytes)) / + (1024.0 * 1024.0 * 1024.0); + if (g->full_kv_cache) { + fprintf(stderr, + "ds4: GLM graph allocating full-attention KV cache: work_ctx=%u layers=%u:%u (%u) %s %.2f GiB\n", + g->ctx_cap, + g->layer_start, + g->layer_end, + g->layer_count, + "f16", + cache_gib); + } else { + fprintf(stderr, + "ds4: GLM graph using compact DSA KV only; expanded full-attention KV cache is skipped\n"); + } +#ifdef DS4_ROCM_BUILD + if (glm_graph_env_truthy( + getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE"))) { + fprintf(stderr, + "ds4: ROCm GLM one-token layer slices use the optimized token graph\n"); + } +#endif + if (g->compact_cache_cap != 0) { + const uint64_t compact_kv_total = + (uint64_t)g->layer_count * (compact_kv_lora_bytes + compact_k_rope_bytes); + const uint64_t compact_indexer_total = + (uint64_t)g->indexer_full_layers * compact_indexer_key_bytes; + const double compact_gib = + (double)(compact_kv_total + compact_indexer_total) / + (1024.0 * 1024.0 * 1024.0); + fprintf(stderr, + "ds4: GLM graph allocating compact DSA cache: rows=%u logical_ctx=%u kv_layers=%u indexer_layers=%u %s %.2f GiB\n", + g->compact_cache_cap, + g->ctx_size, + g->layer_count, + g->indexer_full_layers, + glm_graph_compact_cache_is_f16() ? "f16" : "f32", + compact_gib); + fprintf(stderr, + "ds4: GLM compact indexed prefill chunk=%u score_rows=%u score_scratch=%.2f MiB\n", + g->indexed_prefill_cap, + g->indexed_prefill_score_cap, + (double)batch_indexer_scores_bytes / (1024.0 * 1024.0)); + } + + bool ok = true; +#define DS4_GLM_GRAPH_ALLOC_TENSOR(var, bytes_) \ + do { \ + (var) = ds4_gpu_tensor_alloc((bytes_)); \ + if (!(var)) { \ + fprintf(stderr, "ds4: GLM Metal graph could not allocate %s\n", #var); \ + ok = false; \ + } \ + } while (0) + + DS4_GLM_GRAPH_ALLOC_TENSOR(g->cur, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->next, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_norm, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->q_rank, q_rank_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->q_rank_norm, q_rank_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->q, q_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_k, indexer_k_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_q, indexer_q_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_weights, indexer_weights_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_scores, indexer_scores_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_selected, indexer_selected_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->qk_low, qk_low_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_lora, attn_partial_lora_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_partial_ms, attn_partial_ms_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_raw, kv_raw_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->kv_norm, kv_norm_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->k_nope, k_nope_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->value, heads_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->heads, heads_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->attn_out, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->after_attn, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_norm, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_gate, g->dense_hidden_max * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_up, g->dense_hidden_max * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_mid, g->ffn_mid_elems * sizeof(float)); + if (g->generic_routed_moe) { + DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_gate, routed_mid_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_up, routed_mid_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->routed_down, routed_down_bytes); + } + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_out, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->ffn_sum, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_logits, (uint64_t)DS4_N_EXPERT * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_probs, (uint64_t)DS4_N_EXPERT * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->output_norm, emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->logits, logits_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_logits, batch_rows * DS4_N_EXPERT * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_probs, batch_rows * DS4_N_EXPERT * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_selected, batch_rows * DS4_N_EXPERT_USED * sizeof(int32_t)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_router_weights, batch_rows * DS4_N_EXPERT_USED * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->prefill_seed_router_selected, + (uint64_t)DS4_N_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * + sizeof(int32_t)); + + DS4_GLM_GRAPH_ALLOC_TENSOR(g->prefill_tokens, batch_rows * sizeof(int32_t)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_cur, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_next, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_norm, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q_rank, batch_rows * q_rank_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q_rank_norm, batch_rows * q_rank_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_q, batch_rows * q_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_k, batch_rows * indexer_k_bytes); + if (g->compact_cache_cap != 0 && g->indexed_prefill_cap != 0) { + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_q, batch_indexer_q_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_weights, batch_indexer_weights_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_scores, batch_indexer_scores_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_selected, batch_indexer_selected_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_qk_low, batch_qk_low_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_lora, batch_attn_lora_bytes); + } + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_kv_raw, batch_rows * kv_raw_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_kv_norm, batch_rows * kv_norm_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_k_nope, batch_rows * k_nope_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_value, batch_rows * heads_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_heads, batch_rows * heads_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_attn_out, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_after_attn, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_norm, batch_rows * emb_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_gate, batch_rows * g->dense_hidden_max * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_up, batch_rows * g->dense_hidden_max * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_shared_mid, batch_rows * DS4_N_FF_EXP * sizeof(float)); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_mid, batch_rows * g->ffn_mid_elems * sizeof(float)); + if (g->generic_routed_moe) { + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_gate, batch_rows * routed_mid_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_up, batch_rows * routed_mid_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_routed_down, batch_rows * routed_down_bytes); + } + DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_ffn_out, batch_rows * emb_bytes); + + for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { + int cache_tier = 0; + if (g->placement && g->placement[il + 1u] >= 0 && + g->placement[il + 1u] < DS4_MAX_GPUS) { + cache_tier = g->placement[il + 1u]; + } +#define DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(var, bytes_) \ + do { \ + (var) = ds4_gpu_tensor_alloc_ptr_on(cache_tier, (bytes_)); \ + if (!(var)) { \ + fprintf(stderr, "ds4: GLM graph could not allocate %s on tier %d\n", \ + #var, cache_tier); \ + ok = false; \ + } \ + } while (0) + if (g->full_kv_cache) { + DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_key_cache[il], key_cache_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_value_cache[il], value_cache_bytes); + } + if (g->compact_cache_cap != 0) { + DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_kv_lora_cache[il], compact_kv_lora_bytes); + DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_k_rope_cache[il], compact_k_rope_bytes); + if (glm_graph_layer_uses_full_indexer(il)) { + DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_indexer_key_cache[il], + compact_indexer_key_bytes); + } + } +#undef DS4_GLM_GRAPH_ALLOC_TENSOR_TIER + } +#undef DS4_GLM_GRAPH_ALLOC_TENSOR + + if (!ok) { + glm_graph_free(g); + return false; + } + glm_graph_ws_init(g); + return true; +} + +static bool glm_graph_alloc( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int ctx_size, + bool ssd_streaming, + bool ssd_streaming_cold) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0) { + fprintf(stderr, "ds4: GLM Metal graph found no normal transformer layers\n"); + return false; + } + return glm_graph_alloc_slice(g, + model, + weights, + ctx_size, + ssd_streaming, + ssd_streaming_cold, + 0, + 0, + normal_layers - 1u, + true, + true); +} + +static uint32_t glm_graph_weight_type_for_offset( + const ds4_model *model, + uint64_t weight_offset); + +static int glm_graph_matmul_q8_0_decode_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + bool ssd_streaming) { + if (!model) return 0; + const uint32_t weight_type = glm_graph_weight_type_for_offset(model, weight_offset); + if (ssd_streaming) { + return ds4_gpu_matmul_quant_tensor(out, + model->map, + model->size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + 1); + } + return ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(out, + model->map, + model->size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + 1); +} + +static bool glm_graph_q8_decode_profile_enabled(uint32_t il, const char *label) { + (void)il; + (void)label; + return false; +} + +static uint32_t glm_graph_weight_type_for_offset( + const ds4_model *model, + uint64_t weight_offset) { + if (!model || !model->tensors) return DS4_TENSOR_Q8_0; + for (uint64_t i = 0; i < model->n_tensors; i++) { + const ds4_tensor *t = &model->tensors[i]; + if (t->abs_offset == weight_offset) return t->type; + } + return DS4_TENSOR_Q8_0; +} + +static bool glm_graph_weights_are_q8_0( + const ds4_model *model, + uint64_t offset_a, + uint64_t offset_b) { + return glm_graph_weight_type_for_offset(model, offset_a) == DS4_TENSOR_Q8_0 && + glm_graph_weight_type_for_offset(model, offset_b) == DS4_TENSOR_Q8_0; +} + +static int glm_graph_matmul_q8_0_decode_profiled_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t il, + uint32_t pos, + const char *label, + bool ssd_streaming) { + const bool profile = glm_graph_q8_decode_profile_enabled(il, label); + if (profile) { + if (ds4_gpu_end_commands() == 0) return 0; + if (ds4_gpu_begin_commands() == 0) return 0; + } + const double t0 = profile ? now_sec() : 0.0; + int ok = glm_graph_matmul_q8_0_decode_tensor(out, + model, + weight_offset, + in_dim, + out_dim, + x, + ssd_streaming); + if (profile) { + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM Q8 decode profile layer=%u pos=%u label=%s in=%llu out=%llu %.3f ms\n", + il, + pos, + label ? label : "?", + (unsigned long long)in_dim, + (unsigned long long)out_dim, + (now - t0) * 1000.0); + if (ok) ok = ds4_gpu_begin_commands() != 0; + } + return ok; +} + +static bool glm_graph_encode_output_head_from( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const ds4_gpu_tensor *hidden) { + bool ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, + hidden, + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->logits, + model, + weights->output->abs_offset, + DS4_N_EMBD, + DS4_N_VOCAB, + g->output_norm, + g->ssd_streaming) != 0; + return ok; +} + +static bool glm_graph_encode_output_head( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + return glm_graph_encode_output_head_from(g, model, weights, g->cur); +} + +static bool glm_graph_forward_output_head( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const ds4_gpu_tensor *hidden, + float *logits_out) { + if (!g || !model || !weights || !hidden || !logits_out) return false; + bool ok = ds4_gpu_begin_commands() != 0; + if (ok) ok = glm_graph_encode_output_head_from(g, model, weights, hidden); + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok && glm_debug_hidden_dump_layer() < 0) + glm_debug_dump_hidden_row(hidden, 0); + if (ok) { + ok = ds4_gpu_tensor_read(g->logits, + 0, + logits_out, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + return ok; +} + +static bool glm_graph_profile_stage( + bool enabled, + const char *part, + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + double *stage_t0) { + if (!enabled) return true; + if (!stage_t0) return false; + return metal_graph_layer_stage_profile_boundary(part, stage, il, pos0, n_tokens, stage_t0); +} + +static bool glm_graph_profile_router_selection( + ds4_glm_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos) { + if (!g_expert_profile.active) return true; + if (!g || !layer || !g->router_selected || !g->router_weights) return false; + + if (ds4_gpu_end_commands() == 0) { + fprintf(stderr, + "ds4: failed to end GLM Metal command batch for expert profile readback\n"); + return false; + } + + int32_t selected[DS4_MAX_EXPERT_USED] = {0}; + float weights[DS4_MAX_EXPERT_USED] = {0}; + const bool read_ok = + ds4_gpu_tensor_read(g->router_selected, + 0, + selected, + (uint64_t)DS4_N_EXPERT_USED * sizeof(selected[0])) != 0 && + ds4_gpu_tensor_read(g->router_weights, + 0, + weights, + (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) != 0; + + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: failed to resume GLM Metal command batch after expert profile readback\n"); + return false; + } + if (!read_ok) { + fprintf(stderr, "ds4: failed to read GLM Metal router tensors for expert profile\n"); + return false; + } + + ds4_expert_profile_record(il, + pos, + selected, + weights, + layer->ffn_gate_tid2eid != NULL); + return true; +} + +static bool glm_graph_profile_router_selection_batch( + ds4_glm_gpu_graph *g, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens) { + if (!g_expert_profile.active) return true; + if (!g || !layer || !g->batch_router_selected || + !g->batch_router_weights || n_tokens == 0) { + return false; + } + + const size_t selected_count = (size_t)n_tokens * DS4_N_EXPERT_USED; + if (n_tokens != 0 && selected_count / n_tokens != DS4_N_EXPERT_USED) { + return false; + } + if (selected_count > SIZE_MAX / sizeof(int32_t) || + selected_count > SIZE_MAX / sizeof(float)) { + return false; + } + + if (ds4_gpu_end_commands() == 0) { + fprintf(stderr, + "ds4: failed to end GLM Metal command batch for batch expert profile readback\n"); + return false; + } + + int32_t *selected = xmalloc(selected_count * sizeof(selected[0])); + float *weights = xmalloc(selected_count * sizeof(weights[0])); + const bool read_ok = + ds4_gpu_tensor_read(g->batch_router_selected, + 0, + selected, + (uint64_t)selected_count * sizeof(selected[0])) != 0 && + ds4_gpu_tensor_read(g->batch_router_weights, + 0, + weights, + (uint64_t)selected_count * sizeof(weights[0])) != 0; + + if (ds4_gpu_begin_commands() == 0) { + free(weights); + free(selected); + fprintf(stderr, + "ds4: failed to resume GLM Metal command batch after batch expert profile readback\n"); + return false; + } + if (!read_ok) { + free(weights); + free(selected); + fprintf(stderr, "ds4: failed to read GLM Metal batch router tensors for expert profile\n"); + return false; + } + + for (uint32_t t = 0; t < n_tokens; t++) { + const size_t off = (size_t)t * DS4_N_EXPERT_USED; + ds4_expert_profile_record(il, + pos0 + t, + selected + off, + weights + off, + layer->ffn_gate_tid2eid != NULL); + } + + free(weights); + free(selected); + return true; +} + +static bool glm_graph_prefill_stage_boundary( + bool stage_profile, + bool stage_sync, + const char *part, + const char *stage, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + double *stage_t0) { + if (stage_profile) { + return glm_graph_profile_stage(true, part, stage, il, pos0, n_tokens, stage_t0); + } + if (stage_sync) return glm_graph_prefill_stage_sync_boundary(); + return true; +} + +static int glm_graph_routed_moe_one_dispatch( + const ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *x, + bool force_resident) { + if (!g || !model || !l) return 0; + /* Under the TP expert split only the ownership-aware kernels may run: + * the generic mul_mv_id family and the GLM q2_K resident pair/down. + * Anything else would silently compute the full expert set. */ + if (g->tp_world == 2 && + !glm_graph_layer_uses_generic_routed_moe(l) && + l->ffn_gate_exps->type != DS4_TENSOR_Q2_K) { + fprintf(stderr, + "ds4: GLM TP split lacks ownership-aware kernels for expert type %u (layer %u)\n", + l->ffn_gate_exps->type, il); + return 0; + } + if (glm_graph_layer_uses_generic_routed_moe(l)) { + if (!g->routed_gate || !g->routed_up || !g->routed_down || + l->ffn_gate_exps->type != l->ffn_up_exps->type) { + if (getenv("DS4_GLM_TP_DEBUG")) { + fprintf(stderr, + "ds4: glm dispatch guard: gate=%p up=%p down=%p types=%u/%u\n", + (void *)g->routed_gate, (void *)g->routed_up, + (void *)g->routed_down, + l->ffn_gate_exps->type, l->ffn_up_exps->type); + } + return 0; + } + return ds4_gpu_routed_moe_one_tensor(out, + g->routed_gate, + g->routed_up, + mid, + g->routed_down, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + selected, + weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + 0.0f, + x, + NULL, + il, + force_resident); + } + + return ds4_gpu_glm_routed_moe_one_tensor(out, + mid, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_up_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + selected, + weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + il, + x, + force_resident); +} + +/* Post-compute visibility for GLM TP debugging: the combine stashes the + * selected-ids contents pointer; by exchange time the router kernels have + * completed, so the service thread sees final ids. */ +static const int32_t *g_glm_tp_debug_ids DS4_MAYBE_UNUSED; + +/* After the TP ownership-split batch routed MoE, exchange the + * per-token routed partial rows with the peer through shared bounce + * buffers (one gate per sparse layer per chunk) and rebuild the full + * routed output with a commutative add. */ +/* The routed batch dispatch writes its local partial DIRECTLY into the + * shared bounce buffer (graph scratch may be private/untracked on M5, so + * a blit from it is not reliable); ensure capacity before dispatching. */ +static bool glm_graph_tp_batch_bounce_ready(ds4_glm_gpu_graph *g, + uint32_t n_tokens) { + const uint64_t bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); + if (!g->tp_bounce_out || ds4_gpu_tensor_bytes(g->tp_bounce_out) < bytes) { + ds4_gpu_tensor_free(g->tp_bounce_out); + ds4_gpu_tensor_free(g->tp_bounce_in); + g->tp_bounce_out = ds4_gpu_tensor_alloc(bytes); + g->tp_bounce_in = ds4_gpu_tensor_alloc(bytes); + } + return g->tp_bounce_out && g->tp_bounce_in; +} + +static bool glm_graph_tp_batch_ffn_combine( + ds4_glm_gpu_graph *g, + uint32_t il, + ds4_gpu_tensor *ffn_out, + uint32_t n_tokens) { + const uint64_t bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); + if (!g->tp_bounce_out || !g->tp_bounce_in) return false; + if (getenv("DS4_GLM_ABLATE_COMBINE")) { + /* Timing probe: local half only, no exchange (garbage output; + * both ranks must set the env or the gates desync). */ + return ds4_gpu_add_tensor(ffn_out, + g->tp_bounce_out, + g->tp_bounce_out, + (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; + } + if (!ds4_gpu_tp_big_gate_encode(il, n_tokens, + g->tp_bounce_out, g->tp_bounce_in, + bytes)) { + return false; + } + return ds4_gpu_add_tensor(ffn_out, + g->tp_bounce_out, + g->tp_bounce_in, + (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; +} + +static int glm_graph_routed_moe_batch_dispatch( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool force_resident, + bool direct_scalar_q4) { + if (!g || !model || !l) return 0; + g->batch_routed_mid_is_f16 = false; + + if (glm_graph_layer_uses_generic_routed_moe(l)) { + if (!g->batch_routed_gate || !g->batch_routed_up || !g->batch_routed_down || + l->ffn_gate_exps->type != l->ffn_up_exps->type) { + return 0; + } + return ds4_gpu_routed_moe_batch_tensor(out, + g->batch_routed_gate, + g->batch_routed_up, + mid, + g->batch_routed_down, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + selected, + weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + 0.0f, + x, + il, + n_tokens, + &g->batch_routed_mid_is_f16, + force_resident); + } + + if (direct_scalar_q4) { + return ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( + out, + mid, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_up_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + selected, + weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + il, + x, + n_tokens, + mid_token_stride); + } + + return ds4_gpu_glm_routed_moe_batch_tensor( + out, + mid, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_up_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + selected, + weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + il, + x, + n_tokens, + mid_token_stride, + force_resident); +} + +static bool glm_graph_disable_add3_residual(void); + +static bool glm_graph_use_streaming_selected_async_load( + const ds4_glm_gpu_graph *g) { + if (!g || !g->ssd_streaming) return false; + if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_SELECTED_ASYNC_LOAD", + "DS4_METAL_DISABLE_STREAMING_SELECTED_ASYNC_LOAD") || + glm_graph_env_present("DS4_ROCM_DISABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD", + "DS4_METAL_DISABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD")) { + return false; + } +#ifdef DS4_ROCM_BUILD + return true; +#else + return getenv("DS4_METAL_ENABLE_GLM_STREAMING_SELECTED_ASYNC_LOAD") != NULL; +#endif +} + +typedef struct glm_graph_streaming_async_profile { + uint64_t async_calls; + uint64_t sync_calls; + double async_total_ms; + double sync_total_ms; + double async_signal_start_ms; + double async_flush_router_ms; + double async_shared_ms; + double async_flush_shared_ms; + double async_finish_ms; + double async_routed_ms; + double async_post_ms; + double sync_early_load_ms; + double sync_shared_ms; + double sync_routed_ms; + double sync_post_ms; +} glm_graph_streaming_async_profile; + +static glm_graph_streaming_async_profile g_glm_streaming_async_profile; +static bool g_glm_streaming_async_profile_registered; + +static bool glm_graph_streaming_async_profile_enabled(void) { + return glm_graph_env_present("DS4_ROCM_GLM_STREAMING_ASYNC_PROFILE", + "DS4_METAL_GLM_STREAMING_ASYNC_PROFILE"); +} + +static void glm_graph_streaming_async_profile_print(void) { + const glm_graph_streaming_async_profile *p = + &g_glm_streaming_async_profile; + if (p->async_calls == 0 && p->sync_calls == 0) return; + const double async_calls = p->async_calls ? (double)p->async_calls : 1.0; + const double sync_calls = p->sync_calls ? (double)p->sync_calls : 1.0; + fprintf(stderr, + "ds4: GLM streaming async profile async_calls=%llu " + "total=%.3f ms avg=%.3f ms signal_start=%.3f ms " + "flush_router=%.3f ms shared=%.3f ms flush_shared=%.3f ms " + "finish=%.3f ms routed=%.3f ms post=%.3f ms\n", + (unsigned long long)p->async_calls, + p->async_total_ms, + p->async_total_ms / async_calls, + p->async_signal_start_ms, + p->async_flush_router_ms, + p->async_shared_ms, + p->async_flush_shared_ms, + p->async_finish_ms, + p->async_routed_ms, + p->async_post_ms); + fprintf(stderr, + "ds4: GLM streaming sync profile calls=%llu " + "total=%.3f ms avg=%.3f ms early_load=%.3f ms " + "shared=%.3f ms routed=%.3f ms post=%.3f ms\n", + (unsigned long long)p->sync_calls, + p->sync_total_ms, + p->sync_total_ms / sync_calls, + p->sync_early_load_ms, + p->sync_shared_ms, + p->sync_routed_ms, + p->sync_post_ms); +} + +static void glm_graph_streaming_async_profile_register(void) { + if (g_glm_streaming_async_profile_registered) return; + if (!glm_graph_streaming_async_profile_enabled()) return; + atexit(glm_graph_streaming_async_profile_print); + g_glm_streaming_async_profile_registered = true; +} + +static double glm_graph_streaming_async_profile_ms(void) { + return now_sec() * 1000.0; +} + +/* Timing-only skip-ablation for the GLM decode layer (comma list in + * DS4_GLM_DECODE_ABLATE): the skipped stage's output buffer keeps stale + * contents, so the run produces garbage text but every remaining dispatch + * (and every TP gate) still executes. Whole-token time deltas against a + * baseline run are the only reliable per-stage cost measurement — the + * stage profiler's per-stage command-buffer splits inflate small stages. */ +#define DS4_GLM_ABLATE_ATTN_OUT (1u << 0) +#define DS4_GLM_ABLATE_ATTN_CORE (1u << 1) +#define DS4_GLM_ABLATE_QPATH (1u << 2) +#define DS4_GLM_ABLATE_INDEXER (1u << 3) +#define DS4_GLM_ABLATE_ROUTED (1u << 4) +#define DS4_GLM_ABLATE_SHARED (1u << 5) +#define DS4_GLM_ABLATE_QKLOW (1u << 6) + +static uint32_t glm_decode_ablate_mask(void) { + static int cached = -1; + if (cached < 0) { + uint32_t mask = 0; + const char *env = getenv("DS4_GLM_DECODE_ABLATE"); + if (env) { + if (strstr(env, "attn_out")) mask |= DS4_GLM_ABLATE_ATTN_OUT; + if (strstr(env, "attn_core")) mask |= DS4_GLM_ABLATE_ATTN_CORE; + if (strstr(env, "qpath")) mask |= DS4_GLM_ABLATE_QPATH; + if (strstr(env, "indexer")) mask |= DS4_GLM_ABLATE_INDEXER; + if (strstr(env, "routed")) mask |= DS4_GLM_ABLATE_ROUTED; + if (strstr(env, "shared")) mask |= DS4_GLM_ABLATE_SHARED; + if (strstr(env, "qklow")) mask |= DS4_GLM_ABLATE_QKLOW; + if (mask) { + fprintf(stderr, "ds4: GLM decode ablation active (mask 0x%x) — output is garbage, timing only\n", mask); + } + } + cached = (int)mask; + } + return (uint32_t)cached; +} + +static bool glm_graph_encode_shared_swiglu_one( + ds4_gpu_tensor *mid, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos, + const ds4_gpu_tensor *x, + bool ssd_streaming, + bool stage_profile, + double *stage_t0) { + if (!mid || !gate || !up || !model || !l || !x || + !l->ffn_gate_shexp || !l->ffn_up_shexp) { + return false; + } + + bool ok = true; + if (glm_graph_weights_are_q8_0(model, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset)) { + ok = ds4_gpu_shared_mid_swiglu_q8_0_tensor( + mid, + model->map, + model->size, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + x, + 0.0f) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "shared_gate_up_swiglu", + il, + pos, + 1, + stage_t0); + return ok; + } + + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(gate, + model, + l->ffn_gate_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + x, + il, + pos, + "shared_gate", + ssd_streaming) != 0; + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(up, + model, + l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + x, + il, + pos, + "shared_up", + ssd_streaming) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "shared_gate_up", + il, + pos, + 1, + stage_t0); + if (ok) ok = ds4_gpu_swiglu_tensor(mid, + gate, + up, + DS4_N_FF_EXP, + 0.0f, + 1.0f) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "shared_swiglu", + il, + pos, + 1, + stage_t0); + return ok; +} + +static bool glm_graph_encode_sparse_ffn_one( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos, + const ds4_gpu_tensor *ffn_norm, + const ds4_gpu_tensor *after_attn, + ds4_gpu_tensor *next, + ds4_gpu_tensor *ffn_gate, + ds4_gpu_tensor *ffn_up, + ds4_gpu_tensor *ffn_mid, + ds4_gpu_tensor *ffn_out, + ds4_gpu_tensor *ffn_sum, + ds4_gpu_tensor *tmp, + bool stage_profile, + double *stage_t0) { + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); + (void)gate_in; + (void)up_in; + (void)down_in; + + bool ok = ds4_gpu_matmul_f32_tensor(g->router_logits, + model->map, + model->size, + l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, + DS4_N_EXPERT, + ffn_norm, + 1) != 0; + if (ok) ok = ds4_gpu_glm_router_select_tensor(g->router_selected, + g->router_weights, + g->router_probs, + model->map, + model->size, + l->ffn_exp_probs_b->abs_offset, + g->router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "router", + il, + pos, + 1, + stage_t0); + if (ok) ok = glm_graph_profile_router_selection(g, l, il, pos); + const bool resident_decode_layer = + g->ssd_streaming && glm_stream_resident_decode_layer_enabled(l, il); + const bool generic_streaming_selected_cache = + g->ssd_streaming && + !resident_decode_layer && + glm_graph_layer_uses_generic_routed_moe(l); + const bool uniform_streaming_selected_cache = + g->ssd_streaming && + !resident_decode_layer && + l->ffn_gate_exps->type == l->ffn_up_exps->type && + l->ffn_gate_exps->type == l->ffn_down_exps->type && + (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || + l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); + const bool streaming_selected_cache = + generic_streaming_selected_cache || + uniform_streaming_selected_cache; + const bool shared_first = streaming_selected_cache; + metal_graph_selected_async_load async_load = {0}; + bool async_load_started = false; + const bool async_profile = + streaming_selected_cache && + glm_graph_streaming_async_profile_enabled(); + if (async_profile) glm_graph_streaming_async_profile_register(); + const double stream_total_t0 = + async_profile ? glm_graph_streaming_async_profile_ms() : 0.0; + double stream_t0 = stream_total_t0; + bool async_path_profiled = false; + if (ok && streaming_selected_cache) { + const ds4_gpu_stream_expert_table table = { + .model_map = model->map, + .model_size = model->size, + .layer = il, + .n_total_expert = DS4_N_EXPERT, + .gate_offset = l->ffn_gate_exps->abs_offset, + .up_offset = l->ffn_up_exps->abs_offset, + .down_offset = l->ffn_down_exps->abs_offset, + .gate_expert_bytes = gate_out * gate_row_bytes, + .down_expert_bytes = down_out * down_row_bytes, + }; + const bool async_selected_load = +#ifdef DS4_ROCM_BUILD + streaming_selected_cache && +#else + glm_graph_layer_uses_generic_routed_moe(l) && +#endif + glm_graph_use_streaming_selected_async_load(g); + async_path_profiled = false; + uint64_t selected_event = 0; + if (async_selected_load) { + if (ds4_gpu_signal_selected_readback_ready(&selected_event) != 0) { + async_load_started = metal_graph_selected_async_load_start_tensor( + &async_load, + g->router_selected, + model, + l, + il, + selected_event, + gate_out * gate_row_bytes, + down_out * down_row_bytes); + async_path_profiled = async_profile && async_load_started; + } + if (async_profile) { + g_glm_streaming_async_profile.async_signal_start_ms += + glm_graph_streaming_async_profile_ms() - stream_t0; + stream_t0 = glm_graph_streaming_async_profile_ms(); + } +#ifndef DS4_ROCM_BUILD + if (ok && async_load_started) { + ok = ds4_gpu_flush_commands() != 0; + } +#endif + if (async_profile) { + g_glm_streaming_async_profile.async_flush_router_ms += + glm_graph_streaming_async_profile_ms() - stream_t0; + stream_t0 = glm_graph_streaming_async_profile_ms(); + } + } + if (!async_load_started) { + if (async_selected_load && selected_event != 0) { + ok = ds4_gpu_wait_selected_readback_ready( + selected_event, + "selected-id sync expert load fallback") != 0; + } + if (ok) { + ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + &table, + g->router_selected, + DS4_N_EXPERT_USED) != 0; + } + if (async_profile) { + g_glm_streaming_async_profile.sync_early_load_ms += + glm_graph_streaming_async_profile_ms() - stream_t0; + stream_t0 = glm_graph_streaming_async_profile_ms(); + } + } + } + if (ok && shared_first) { + ok = glm_graph_encode_shared_swiglu_one(ffn_mid, + ffn_gate, + ffn_up, + model, + l, + il, + pos, + ffn_norm, + g->ssd_streaming, + stage_profile, + stage_t0); + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + ffn_mid, + il, + pos, + "shared_down", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "shared_down", + il, + pos, + 1, + stage_t0); + } + if (async_profile) { + const double now_ms = glm_graph_streaming_async_profile_ms(); + if (async_path_profiled) { + g_glm_streaming_async_profile.async_shared_ms += + now_ms - stream_t0; + } else { + g_glm_streaming_async_profile.sync_shared_ms += + now_ms - stream_t0; + } + stream_t0 = now_ms; + } + if (async_load_started) { + bool flush_ok = true; +#ifndef DS4_ROCM_BUILD + flush_ok = ds4_gpu_flush_commands() != 0; +#endif + if (async_profile) { + g_glm_streaming_async_profile.async_flush_shared_ms += + glm_graph_streaming_async_profile_ms() - stream_t0; + stream_t0 = glm_graph_streaming_async_profile_ms(); + } + const bool finish_ok = metal_graph_selected_async_load_finish(&async_load); + ok = ok && flush_ok && finish_ok; + if (async_profile) { + g_glm_streaming_async_profile.async_finish_ms += + glm_graph_streaming_async_profile_ms() - stream_t0; + stream_t0 = glm_graph_streaming_async_profile_ms(); + } + } + /* 50/50 TP: this rank's routed partial goes straight into the + * slab out slot, the gate exchanges it with the peer's half, and the + * commutative add rebuilds the full routed output on both ranks + * bit-identically. The shared expert and everything else stay + * replicated, so no other exchange is needed. */ + const bool tp_split_ffn = g->tp_world == 2 && g->tp_out && g->tp_in; + const uint32_t tp_ffn_slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; + ds4_gpu_tensor *routed_dst = tp_split_ffn ? g->tp_out[tp_ffn_slot] : ffn_out; + if (ok && tp_split_ffn && g->ssd_streaming) { + fprintf(stderr, "ds4: GLM tensor parallelism requires resident weights\n"); + ok = false; + } + if (!ok && tp_split_ffn && getenv("DS4_GLM_TP_DEBUG")) { + fprintf(stderr, "ds4: glm sparse ffn: failed before routed dispatch (layer %u)\n", il); + } + if (ok && !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { + ok = glm_graph_routed_moe_one_dispatch( + g, + model, + l, + il, + routed_dst, + ffn_mid, + gate_out * gate_row_bytes, + gate_row_bytes, + up_out * up_row_bytes, + up_row_bytes, + down_out * down_row_bytes, + down_row_bytes, + g->router_selected, + g->router_weights, + ffn_norm, + resident_decode_layer) != 0; + } + if (ok && tp_split_ffn) { + ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_FFN) != 0; + if (ok) ok = ds4_gpu_add_tensor(ffn_out, + g->tp_out[tp_ffn_slot], + g->tp_in[tp_ffn_slot], + DS4_N_EMBD) != 0; + if (!ok) fprintf(stderr, "ds4: GLM TP gate/combine failed (layer %u)\n", il); + } else if (!ok && tp_split_ffn) { + fprintf(stderr, "ds4: GLM TP routed dispatch failed before the gate (layer %u)\n", il); + } + if (async_profile) { + const double now_ms = glm_graph_streaming_async_profile_ms(); + if (async_path_profiled) { + g_glm_streaming_async_profile.async_routed_ms += + now_ms - stream_t0; + } else { + g_glm_streaming_async_profile.sync_routed_ms += + now_ms - stream_t0; + } + stream_t0 = now_ms; + } + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "routed_moe", + il, + pos, + 1, + stage_t0); + if (ok && !shared_first && + !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_SHARED)) { + ok = glm_graph_encode_shared_swiglu_one(ffn_mid, + ffn_gate, + ffn_up, + model, + l, + il, + pos, + ffn_norm, + g->ssd_streaming, + stage_profile, + stage_t0); + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_sum, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + ffn_mid, + il, + pos, + "shared_down", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "shared_down", + il, + pos, + 1, + stage_t0); + } + if (ok && !glm_graph_disable_add3_residual()) { + ok = ds4_gpu_add3_tensor(next, + after_attn, + ffn_out, + ffn_sum, + DS4_N_EMBD) != 0; + } else if (ok) { + ok = ds4_gpu_add_tensor(tmp, + ffn_out, + ffn_sum, + DS4_N_EMBD) != 0; + if (ok) ok = ds4_gpu_add_tensor(next, + after_attn, + tmp, + DS4_N_EMBD) != 0; + } + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "residual", + il, + pos, + 1, + stage_t0); + if (async_profile) { + const double now_ms = glm_graph_streaming_async_profile_ms(); + if (async_path_profiled) { + g_glm_streaming_async_profile.async_calls++; + g_glm_streaming_async_profile.async_post_ms += now_ms - stream_t0; + g_glm_streaming_async_profile.async_total_ms += + now_ms - stream_total_t0; + } else { + g_glm_streaming_async_profile.sync_calls++; + g_glm_streaming_async_profile.sync_post_ms += now_ms - stream_t0; + g_glm_streaming_async_profile.sync_total_ms += + now_ms - stream_total_t0; + } + } + return ok; +} + +static bool glm_graph_encode_ffn_one_normed_from( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos, + const ds4_gpu_tensor *ffn_norm, + const ds4_gpu_tensor *after_attn, + ds4_gpu_tensor *next, + ds4_gpu_tensor *ffn_gate, + ds4_gpu_tensor *ffn_up, + ds4_gpu_tensor *ffn_mid, + ds4_gpu_tensor *ffn_out, + ds4_gpu_tensor *ffn_sum, + ds4_gpu_tensor *tmp, + bool stage_profile, + double *stage_t0) { + if (!g || !model || !l || !ffn_norm || !after_attn || !next || + !ffn_gate || !ffn_up || !ffn_mid || !ffn_out || + !ffn_sum || !tmp) { + return false; + } + + if (il < DS4_N_LEADING_DENSE) { + const uint64_t hidden = l->ffn_gate->dim[1]; + const bool can_fuse_gate_up = + glm_graph_weights_are_q8_0(model, + l->ffn_gate->abs_offset, + l->ffn_up->abs_offset); + const bool fused_gate_up = can_fuse_gate_up && + ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + ffn_gate, + ffn_up, + ffn_mid, + model->map, + model->size, + l->ffn_gate->abs_offset, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + ffn_norm, + 0.0f) != 0; + bool ok = fused_gate_up; + if (fused_gate_up) { + ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "dense_gate_up_swiglu", + il, + pos, + 1, + stage_t0); + } else { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_gate, + model, + l->ffn_gate->abs_offset, + DS4_N_EMBD, + hidden, + ffn_norm, + il, + pos, + "dense_gate", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_up, + model, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + ffn_norm, + il, + pos, + "dense_up", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "dense_gate_up", + il, + pos, + 1, + stage_t0); + if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, + ffn_gate, + ffn_up, + (uint32_t)hidden, + 0.0f, + 1.0f) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "dense_swiglu", + il, + pos, + 1, + stage_t0); + } + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_out, + model, + l->ffn_down->abs_offset, + hidden, + DS4_N_EMBD, + ffn_mid, + il, + pos, + "dense_down", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "dense_down", + il, + pos, + 1, + stage_t0); + if (ok) ok = ds4_gpu_add_tensor(next, + after_attn, + ffn_out, + DS4_N_EMBD) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "residual", + il, + pos, + 1, + stage_t0); + return ok; + } + + return glm_graph_encode_sparse_ffn_one(g, + model, + l, + il, + pos, + ffn_norm, + after_attn, + next, + ffn_gate, + ffn_up, + ffn_mid, + ffn_out, + ffn_sum, + tmp, + stage_profile, + stage_t0); +} + +static bool glm_graph_encode_ffn_one_from( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos, + const ds4_gpu_tensor *after_attn, + ds4_gpu_tensor *next, + ds4_gpu_tensor *ffn_norm, + ds4_gpu_tensor *ffn_gate, + ds4_gpu_tensor *ffn_up, + ds4_gpu_tensor *ffn_mid, + ds4_gpu_tensor *ffn_out, + ds4_gpu_tensor *ffn_sum, + ds4_gpu_tensor *tmp, + bool stage_profile, + double *stage_t0) { + if (!g || !model || !l || !after_attn || !next || + !ffn_norm || !ffn_gate || !ffn_up || !ffn_mid || !ffn_out || + !ffn_sum || !tmp) { + return false; + } + + bool ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, + after_attn, + model->map, + model->size, + l->ffn_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_profile_stage(stage_profile, + "glm_decode_ffn", + "ffn_norm", + il, + pos, + 1, + stage_t0); + if (!ok) return false; + + return glm_graph_encode_ffn_one_normed_from(g, + model, + l, + il, + pos, + ffn_norm, + after_attn, + next, + ffn_gate, + ffn_up, + ffn_mid, + ffn_out, + ffn_sum, + tmp, + stage_profile, + stage_t0); +} + +static ds4_gpu_tensor *glm_graph_tensor_row_view_strided( + ds4_gpu_tensor *base, + uint32_t row, + uint64_t stride_values, + uint64_t row_values) { + return ds4_gpu_tensor_view(base, + (uint64_t)row * stride_values * sizeof(float), + row_values * sizeof(float)); +} + +static uint32_t glm_graph_q8_stripe_tokens(void) { + return 2048u; +} + +static bool glm_graph_flash_attention_prefill_enabled(void) { + return getenv("DS4_GLM_DISABLE_FLASH_PREFILL") == NULL; +} + +static uint32_t glm_graph_flash_attention_prefill_min_tokens(void) { + return 24u; +} + +static bool glm_graph_use_flash_attention_prefill(uint32_t n_tokens) { + return glm_graph_flash_attention_prefill_enabled() && + n_tokens >= glm_graph_flash_attention_prefill_min_tokens(); +} + +static bool glm_graph_use_flash_attention_staged_kv( + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len) { + return pos0 == 0 && + n_tokens == cache_len; +} + +static bool glm_graph_force_indexed_decode(void) { + return false; +} + +static bool glm_graph_disable_indexed_decode(void) { + return false; +} + +static bool glm_graph_decode_uses_indexed_attention(const ds4_glm_gpu_graph *g, + uint32_t pos, + const float *logits_out) { + return g && g->compact_cache_cap != 0 && + (!g->full_kv_cache || + pos >= g->ctx_cap || + glm_graph_force_indexed_decode() || + (logits_out != NULL && !glm_graph_disable_indexed_decode())); +} + +static bool glm_graph_decode_updates_dense_cache(const ds4_glm_gpu_graph *g, + uint32_t pos, + const float *logits_out) { + return g && pos < g->ctx_cap && + !glm_graph_decode_uses_indexed_attention(g, pos, logits_out); +} + +static bool glm_graph_indexed_prefill_scalar_kernels(void) { + return false; +} + +static bool glm_graph_indexed_prefill_scalar_indexer(void) { + return false; +} + +static bool glm_graph_indexed_prefill_batch_indexer(void) { + return true; +} + +static bool glm_graph_indexed_prefill_scalar_attn(void) { + return false; +} + +static bool glm_graph_indexed_prefill_batch_qk_low(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_attn_kernel(void) { + return true; +} + +static uint32_t glm_graph_indexed_prefill_batch_attn_slice_tokens(void) { + return 2048u; +} + +static bool glm_graph_indexer_qat(void) { + return false; +} + +static bool glm_graph_indexed_prefill_batch_ffn(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_ffn_norm(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_routed_moe(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_router_select(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_residual(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_f32_rows(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_q8_rows(void) { + return true; +} + +static bool glm_graph_indexed_prefill_batch_shared_expert(void) { + return true; +} + +static bool glm_graph_matmul_q8_0_tensor( + ds4_gpu_tensor *out, + const ds4_model *model, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!out || !model || !x || n_tokens == 0) return false; + + const uint32_t q8_stripe_tokens = glm_graph_q8_stripe_tokens(); + if (n_tokens <= q8_stripe_tokens) { + return ds4_gpu_matmul_quant_tensor(out, + model->map, + model->size, + weight_offset, + glm_graph_weight_type_for_offset(model, weight_offset), + in_dim, + out_dim, + x, + n_tokens) != 0; + } + if (in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return false; + } + + uint32_t done = 0; + while (done < n_tokens) { + uint32_t chunk = n_tokens - done; + if (chunk > q8_stripe_tokens) chunk = q8_stripe_tokens; + /* + * The Q8 prefill TensorOps path needs token counts divisible by 32. + * For a final chunk like 1736 rows, keep the aligned 1728 rows on that + * path and leave only the tiny tail to the small-batch kernel. Avoid + * splitting small batches where the extra launch would dominate. + */ + if (chunk >= 256u && chunk == n_tokens - done) { + const uint32_t tail = chunk & 31u; + if (tail != 0u && tail <= 16u) chunk -= tail; + } + + ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( + x, + (uint64_t)done * in_dim * sizeof(float), + (uint64_t)chunk * in_dim * sizeof(float)); + ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( + out, + (uint64_t)done * out_dim * sizeof(float), + (uint64_t)chunk * out_dim * sizeof(float)); + const bool ok = x_view && out_view && + ds4_gpu_matmul_quant_tensor(out_view, + model->map, + model->size, + weight_offset, + glm_graph_weight_type_for_offset(model, weight_offset), + in_dim, + out_dim, + x_view, + chunk) != 0; + ds4_gpu_tensor_free(out_view); + ds4_gpu_tensor_free(x_view); + if (!ok) return false; + done += chunk; + } + return true; +} + +static bool glm_graph_matmul_q8_0_rows_scalar( + ds4_gpu_tensor *out, + const ds4_model *model, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!out || !model || !x || n_tokens == 0) return false; + if (in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return false; + } + + if (glm_graph_indexed_prefill_batch_q8_rows() && + ds4_gpu_matmul_quant_rows_scalar_tensor(out, + model->map, + model->size, + weight_offset, + glm_graph_weight_type_for_offset(model, weight_offset), + in_dim, + out_dim, + x, + n_tokens) != 0) { + return true; + } + + for (uint32_t t = 0; t < n_tokens; t++) { + ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( + (ds4_gpu_tensor *)x, + (uint64_t)t * in_dim * sizeof(float), + in_dim * sizeof(float)); + ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( + out, + (uint64_t)t * out_dim * sizeof(float), + out_dim * sizeof(float)); + const bool ok = x_view && out_view && + ds4_gpu_matmul_quant_tensor(out_view, + model->map, + model->size, + weight_offset, + glm_graph_weight_type_for_offset(model, weight_offset), + in_dim, + out_dim, + x_view, + 1) != 0; + ds4_gpu_tensor_free(out_view); + ds4_gpu_tensor_free(x_view); + if (!ok) return false; + } + return true; +} + +static bool glm_graph_matmul_f32_rows_scalar( + ds4_gpu_tensor *out, + const ds4_model *model, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!out || !model || !x || n_tokens == 0) return false; + if (in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return false; + } + + if (glm_graph_indexed_prefill_batch_f32_rows()) { + return ds4_gpu_matmul_f32_tensor(out, + model->map, + model->size, + weight_offset, + in_dim, + out_dim, + x, + n_tokens) != 0; + } + + for (uint32_t t = 0; t < n_tokens; t++) { + ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( + (ds4_gpu_tensor *)x, + (uint64_t)t * in_dim * sizeof(float), + in_dim * sizeof(float)); + ds4_gpu_tensor *out_view = ds4_gpu_tensor_view( + out, + (uint64_t)t * out_dim * sizeof(float), + out_dim * sizeof(float)); + const bool ok = x_view && out_view && + ds4_gpu_matmul_f32_tensor(out_view, + model->map, + model->size, + weight_offset, + in_dim, + out_dim, + x_view, + 1) != 0; + ds4_gpu_tensor_free(out_view); + ds4_gpu_tensor_free(x_view); + if (!ok) return false; + } + return true; +} + +static bool glm_graph_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const ds4_model *model, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + float clamp) { + if (!gate || !up || !mid || !model || !x || n_tokens == 0) return false; + if (!glm_graph_weights_are_q8_0(model, gate_offset, up_offset)) return false; + + const uint32_t q8_stripe_tokens = glm_graph_q8_stripe_tokens(); + if (n_tokens == 1) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(gate, + up, + mid, + model->map, + model->size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp) != 0; + } + if (n_tokens <= q8_stripe_tokens) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate, + up, + mid, + model->map, + model->size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + n_tokens, + clamp) != 0; + } + if (in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return false; + } + + uint32_t done = 0; + while (done < n_tokens) { + uint32_t chunk = n_tokens - done; + if (chunk > q8_stripe_tokens) chunk = q8_stripe_tokens; + + ds4_gpu_tensor *x_view = ds4_gpu_tensor_view( + x, + (uint64_t)done * in_dim * sizeof(float), + (uint64_t)chunk * in_dim * sizeof(float)); + ds4_gpu_tensor *gate_view = ds4_gpu_tensor_view( + gate, + (uint64_t)done * out_dim * sizeof(float), + (uint64_t)chunk * out_dim * sizeof(float)); + ds4_gpu_tensor *up_view = ds4_gpu_tensor_view( + up, + (uint64_t)done * out_dim * sizeof(float), + (uint64_t)chunk * out_dim * sizeof(float)); + ds4_gpu_tensor *mid_view = ds4_gpu_tensor_view( + mid, + (uint64_t)done * out_dim * sizeof(float), + (uint64_t)chunk * out_dim * sizeof(float)); + const bool ok = x_view && gate_view && up_view && mid_view && + ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate_view, + up_view, + mid_view, + model->map, + model->size, + gate_offset, + up_offset, + in_dim, + out_dim, + x_view, + chunk, + clamp) != 0; + ds4_gpu_tensor_free(mid_view); + ds4_gpu_tensor_free(up_view); + ds4_gpu_tensor_free(gate_view); + ds4_gpu_tensor_free(x_view); + if (!ok) return false; + done += chunk; + } + return true; +} + +static bool glm_graph_indexed_prefill_grouped_moe_default( + const ds4_glm_gpu_graph *g) { + return g && + !g->quality; +} + +static uint32_t glm_graph_streaming_prefill_cache_seed_k( + const ds4_glm_gpu_graph *g) { + const bool enabled = + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_CACHE_SEED", + "DS4_METAL_ENABLE_STREAMING_PREFILL_CACHE_SEED"); + if (!g || + !g->ssd_streaming || + !enabled) { + return 0; + } + + uint32_t k = 1; + const char *env = glm_graph_env_value("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_K", + "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_K"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0') { + if (v == 0) return 0; + k = v > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS ? + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS : (uint32_t)v; + } + } + return k; +} + +static bool glm_graph_streaming_prefill_cache_seed_enabled( + const ds4_glm_gpu_graph *g) { + return glm_graph_streaming_prefill_cache_seed_k(g) != 0; +} + +static void glm_graph_reset_prefill_seed_capture(ds4_glm_gpu_graph *g) { + if (!g) return; + g->prefill_seed_tokens = 0; + memset(g->prefill_seed_layer_captured, + 0, + sizeof(g->prefill_seed_layer_captured)); +} + +static bool glm_graph_streaming_expert_cache_seed_layer_expected( + const ds4_glm_gpu_graph *g, + const ds4_weights *weights, + const ds4_layer_weights *layer, + uint32_t il) { + if (!g || + !g->ssd_streaming || + g->quality || + !weights || + !layer || + !layer->ffn_gate_exps || + !layer->ffn_up_exps || + !layer->ffn_down_exps) { + return false; + } + if (glm_stream_resident_decode_layer_enabled(layer, il)) return false; + return glm_stream_selected_expert_cache_supported(layer, il) || + glm_stream_expert_cache_addr_layout_supported(weights, layer, il); +} + +static bool glm_graph_capture_prefill_seed_router_selected( + ds4_glm_gpu_graph *g, + uint32_t il, + uint32_t n_tokens) { + uint32_t k = glm_graph_streaming_prefill_cache_seed_k(g); + if (k == 0) return true; + if (!g->prefill_seed_router_selected || + !g->batch_router_selected || + il >= DS4_N_LAYER || + il >= DS4_MAX_LAYER || + n_tokens == 0 || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return false; + } + if (k > n_tokens) k = n_tokens; + + const uint64_t bytes = (uint64_t)k * DS4_N_EXPERT_USED * sizeof(int32_t); + const uint64_t src_off = (uint64_t)(n_tokens - k) * + DS4_N_EXPERT_USED * sizeof(int32_t); + const uint64_t dst_off = (uint64_t)il * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * sizeof(int32_t); + if (ds4_gpu_tensor_copy(g->prefill_seed_router_selected, + dst_off, + g->batch_router_selected, + src_off, + bytes) == 0) { + return false; + } + g->prefill_seed_tokens = k; + g->prefill_seed_layer_captured[il] = true; + return true; +} + +static bool glm_graph_seed_streaming_expert_cache_from_prefill( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights) { + if (!glm_graph_streaming_prefill_cache_seed_enabled(g)) return true; + const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; + if (seed_tokens == 0) return true; + if (!g || + !model || + !weights || + !g->prefill_seed_router_selected || + seed_tokens > DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS || + DS4_N_LAYER > DS4_MAX_LAYER || + DS4_N_EXPERT_USED == 0 || + DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { + return false; + } + bool any_captured = false; + for (uint32_t il = g->layer_start; + il <= g->layer_end && il < DS4_N_LAYER && il < DS4_MAX_LAYER; + il++) { + if (g->prefill_seed_layer_captured[il]) { + any_captured = true; + break; + } + } + if (!any_captured) return true; + + int32_t selected[DS4_MAX_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_MAX_EXPERT_USED]; + const uint64_t bytes = (uint64_t)DS4_N_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * sizeof(selected[0]); + if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, + 0, + selected, + bytes) == 0) { + return false; + } + + const bool profile = + glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_PROFILE", + "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE"); + const double t0 = profile ? now_sec() : 0.0; + uint32_t seeded_layers = 0; + uint32_t seeded_rows = 0; + for (uint32_t il = g->layer_start; + il <= g->layer_end && il < DS4_N_LAYER; + il++) { + if (il >= DS4_MAX_LAYER || !g->prefill_seed_layer_captured[il]) { + continue; + } + const ds4_layer_weights *layer = &weights->layer[il]; + if (!glm_graph_streaming_expert_cache_seed_layer_expected(g, + weights, + layer, + il)) { + continue; + } + + uint64_t gate_expert_bytes = 0; + uint64_t down_expert_bytes = 0; + if (!streaming_layer_gate_down_expert_bytes(layer, + &gate_expert_bytes, + &down_expert_bytes)) { + fprintf(stderr, + "ds4: GLM prefill expert-cache seed byte size overflow at layer %u\n", + il); + return false; + } + + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + for (uint32_t row = 0; row < seed_tokens; row++) { + const size_t sel_off = ((size_t)il * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS + + row) * DS4_N_EXPERT_USED; + if (ds4_gpu_stream_expert_cache_seed_selected( + &table, + selected + sel_off, + DS4_N_EXPERT_USED) == 0) { + return false; + } + seeded_rows++; + } + seeded_layers++; + } + if (profile) { + fprintf(stderr, + "ds4: GLM streaming prefill expert-cache seed k=%u layers=%u rows=%u time=%.3f ms\n", + seed_tokens, + seeded_layers, + seeded_rows, + (now_sec() - t0) * 1000.0); + } + return true; +} + +static bool glm_graph_seed_streaming_expert_cache_from_full_layer( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t n_tokens, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + bool full_layer_prefill) { +#ifdef DS4_ROCM_BUILD + uint32_t seed_tokens = glm_graph_streaming_prefill_cache_seed_k(g); + if (seed_tokens == 0) return true; + if (!full_layer_prefill || + !model || + !weights || + !layer || + !g || + !g->batch_router_selected || + n_tokens == 0 || + il >= DS4_N_LAYER || + il >= DS4_MAX_LAYER || + gate_expert_bytes == 0 || + down_expert_bytes == 0 || + !g->prefill_seed_layer_captured[il] || + !glm_graph_streaming_expert_cache_seed_layer_expected(g, + weights, + layer, + il)) { + return true; + } + if (seed_tokens > n_tokens) seed_tokens = n_tokens; + const ds4_gpu_stream_expert_table table = + graph_stream_expert_table_make(model, + layer, + il, + gate_expert_bytes, + down_expert_bytes); + if (ds4_gpu_stream_expert_cache_seed_from_layer_selected( + &table, + g->batch_router_selected, + n_tokens, + seed_tokens, + DS4_N_EXPERT_USED) != 0) { + g->prefill_seed_layer_captured[il] = false; + return true; + } + + static bool warned = false; + if (!warned) { + fprintf(stderr, + "ds4: GLM ROCm full-layer prefill expert-cache seed skipped; " + "falling back to end-of-prefill selected seed\n"); + warned = true; + } + return true; +#else + (void)g; + (void)model; + (void)weights; + (void)layer; + (void)il; + (void)n_tokens; + (void)gate_expert_bytes; + (void)down_expert_bytes; + (void)full_layer_prefill; + return true; +#endif +} + +static bool glm_graph_disable_add3_residual(void); + +static bool glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos0, + const ds4_gpu_tensor *after_attn, + ds4_gpu_tensor *next, + uint32_t n_tokens, + bool stage_profile, + bool stage_sync, + double *stage_t0) { + if (!g || !model || !l || !after_attn || !next || + !g->batch_ffn_norm || + !g->batch_router_logits || + !g->batch_router_probs || + !g->batch_router_selected || + !g->batch_router_weights || + !g->batch_ffn_out || + !g->batch_ffn_mid || + n_tokens <= 1 || + il < DS4_N_LEADING_DENSE || + g->ffn_mid_elems > UINT32_MAX) { + return false; + } + + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); + (void)gate_in; + (void)up_in; + (void)down_in; + + bool ok = glm_graph_matmul_f32_rows_scalar(g->batch_router_logits, + model, + l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, + DS4_N_EXPERT, + g->batch_ffn_norm, + n_tokens); + const bool use_batch_router_select = + glm_graph_indexed_prefill_batch_router_select(); + if (ok && use_batch_router_select) { + ok = ds4_gpu_glm_router_select_batch_tensor(g->batch_router_selected, + g->batch_router_weights, + g->batch_router_probs, + model->map, + model->size, + l->ffn_exp_probs_b->abs_offset, + g->batch_router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE, + n_tokens) != 0; + } + for (uint32_t t = 0; ok && !use_batch_router_select && t < n_tokens; t++) { + ds4_gpu_tensor *logits_view = + glm_graph_tensor_row_view_strided(g->batch_router_logits, + t, + DS4_N_EXPERT, + DS4_N_EXPERT); + ds4_gpu_tensor *probs_view = + glm_graph_tensor_row_view_strided(g->batch_router_probs, + t, + DS4_N_EXPERT, + DS4_N_EXPERT); + ds4_gpu_tensor *selected_view = + ds4_gpu_tensor_view(g->batch_router_selected, + (uint64_t)t * DS4_N_EXPERT_USED * sizeof(int32_t), + (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + ds4_gpu_tensor *weights_view = + glm_graph_tensor_row_view_strided(g->batch_router_weights, + t, + DS4_N_EXPERT_USED, + DS4_N_EXPERT_USED); + ok = logits_view && probs_view && selected_view && weights_view; + if (ok) { + ok = ds4_gpu_glm_router_select_tensor(selected_view, + weights_view, + probs_view, + model->map, + model->size, + l->ffn_exp_probs_b->abs_offset, + logits_view, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE) != 0; + } + ds4_gpu_tensor_free(weights_view); + ds4_gpu_tensor_free(selected_view); + ds4_gpu_tensor_free(probs_view); + ds4_gpu_tensor_free(logits_view); + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_indexed_ffn", + "router", + il, + pos0, + n_tokens, + stage_t0); + if (ok) ok = glm_graph_profile_router_selection_batch(g, + l, + il, + pos0, + n_tokens); + if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, + il, + n_tokens); + metal_graph_debug_dump_tensor("glm_indexed_router_logits", + g->batch_router_logits, + (uint64_t)n_tokens * DS4_N_EXPERT, + il, + pos0); + metal_graph_debug_dump_i32_tensor("glm_indexed_router_selected", + g->batch_router_selected, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + metal_graph_debug_dump_tensor("glm_indexed_router_weights", + g->batch_router_weights, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + + const bool tp_batch_split_ffn2 = g->tp_world == 2; + if (ok && tp_batch_split_ffn2) { + ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); + } + if (ok) { + const bool use_grouped_moe = + glm_graph_indexed_prefill_grouped_moe_default(g); + ok = glm_graph_routed_moe_batch_dispatch( + g, + model, + l, + il, + tp_batch_split_ffn2 ? g->tp_bounce_out : g->batch_ffn_out, + g->batch_ffn_mid, + gate_out * gate_row_bytes, + gate_row_bytes, + up_out * up_row_bytes, + up_row_bytes, + down_out * down_row_bytes, + down_row_bytes, + g->batch_router_selected, + g->batch_router_weights, + g->batch_ffn_norm, + n_tokens, + (uint32_t)g->ffn_mid_elems, + false, + !use_grouped_moe) != 0; + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_indexed_ffn", + "routed_moe", + il, + pos0, + n_tokens, + stage_t0); + metal_graph_debug_dump_tensor("glm_indexed_routed_out", + g->batch_ffn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + + const bool use_batch_residual = + glm_graph_indexed_prefill_batch_residual(); + const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + if (use_batch_residual && residual_elems > UINT32_MAX) return false; + + bool shared_expert_done = false; + if (ok && + use_batch_residual && + glm_graph_indexed_prefill_batch_shared_expert() && + g->batch_ffn_gate && + g->batch_ffn_up && + g->batch_shared_mid && + glm_graph_weights_are_q8_0(model, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset) && + ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( + g->batch_ffn_gate, + g->batch_ffn_up, + g->batch_shared_mid, + model->map, + model->size, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + g->batch_ffn_norm, + n_tokens, + 0.0f) != 0) { + shared_expert_done = + glm_graph_matmul_q8_0_rows_scalar(g->batch_attn_out, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + g->batch_shared_mid, + n_tokens); + } + + for (uint32_t t = 0; ok && !shared_expert_done && t < n_tokens; t++) { + ds4_gpu_tensor *ffn_norm_view = + glm_graph_tensor_row_view_strided(g->batch_ffn_norm, + t, + DS4_N_EMBD, + DS4_N_EMBD); + ds4_gpu_tensor *shared_out_view = use_batch_residual ? + glm_graph_tensor_row_view_strided(g->batch_attn_out, + t, + DS4_N_EMBD, + DS4_N_EMBD) : + NULL; + ds4_gpu_tensor *after_attn_view = !use_batch_residual ? + glm_graph_tensor_row_view_strided((ds4_gpu_tensor *)after_attn, + t, + DS4_N_EMBD, + DS4_N_EMBD) : + NULL; + ds4_gpu_tensor *routed_out_view = !use_batch_residual ? + glm_graph_tensor_row_view_strided(g->batch_ffn_out, + t, + DS4_N_EMBD, + DS4_N_EMBD) : + NULL; + ds4_gpu_tensor *next_view = !use_batch_residual ? + glm_graph_tensor_row_view_strided(next, + t, + DS4_N_EMBD, + DS4_N_EMBD) : + NULL; + ok = ffn_norm_view && + (use_batch_residual ? (shared_out_view != NULL) : + (after_attn_view && routed_out_view && next_view)); + if (ok && glm_graph_weights_are_q8_0(model, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset)) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + g->ffn_gate, + g->ffn_up, + g->ffn_mid, + model->map, + model->size, + l->ffn_gate_shexp->abs_offset, + l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + ffn_norm_view, + 0.0f) != 0; + } else if (ok) { + ok = glm_graph_matmul_q8_0_tensor(g->ffn_gate, + model, + l->ffn_gate_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + ffn_norm_view, + 1); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->ffn_up, + model, + l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + ffn_norm_view, + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(g->ffn_mid, + g->ffn_gate, + g->ffn_up, + DS4_N_FF_EXP, + 0.0f, + 1.0f) != 0; + } + if (ok) ok = glm_graph_matmul_q8_0_tensor(use_batch_residual ? + shared_out_view : + g->ffn_sum, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + g->ffn_mid, + 1); + if (ok && !use_batch_residual) { + ok = ds4_gpu_add_tensor(g->attn_out, + routed_out_view, + g->ffn_sum, + DS4_N_EMBD) != 0; + } + if (ok && !use_batch_residual) { + ok = ds4_gpu_add_tensor(next_view, + after_attn_view, + g->attn_out, + DS4_N_EMBD) != 0; + } + + ds4_gpu_tensor_free(next_view); + ds4_gpu_tensor_free(routed_out_view); + ds4_gpu_tensor_free(shared_out_view); + ds4_gpu_tensor_free(ffn_norm_view); + ds4_gpu_tensor_free(after_attn_view); + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_indexed_ffn", + "shared_expert", + il, + pos0, + n_tokens, + stage_t0); + if (ok && use_batch_residual) { + if (!glm_graph_disable_add3_residual()) { + ok = ds4_gpu_add3_tensor(next, + after_attn, + g->batch_ffn_out, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + } else { + ok = ds4_gpu_add_tensor(g->batch_heads, + g->batch_ffn_out, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + if (ok) ok = ds4_gpu_add_tensor(next, + after_attn, + g->batch_heads, + (uint32_t)residual_elems) != 0; + } + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_indexed_ffn", + "residual", + il, + pos0, + n_tokens, + stage_t0); + metal_graph_debug_dump_tensor("glm_indexed_next", + next, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + + (void)pos0; + return ok; +} + +static bool glm_graph_upload_tokens( + ds4_gpu_tensor *out_tokens, + const int *tokens, + uint32_t n_tokens) { + if (!out_tokens || !tokens || n_tokens == 0) return false; + + int32_t *ids = xmalloc((size_t)n_tokens * sizeof(ids[0])); + for (uint32_t i = 0; i < n_tokens; i++) ids[i] = (int32_t)tokens[i]; + const bool ok = ds4_gpu_tensor_write(out_tokens, + 0, + ids, + (uint64_t)n_tokens * sizeof(ids[0])) != 0; + free(ids); + return ok; +} + +static bool glm_graph_disable_add3_residual(void) { + return false; +} + +static bool glm_graph_encode_ffn_batch( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const ds4_layer_weights *l, + uint32_t il, + uint32_t pos0, + ds4_gpu_tensor *after_attn, + ds4_gpu_tensor *next, + uint32_t n_tokens, + bool full_layer_prefill, + bool stage_profile, + bool stage_sync, + double *stage_t0) { + if (!g || !model || !weights || !l || !after_attn || !next || n_tokens == 0) return false; + + bool ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, + after_attn, + model->map, + model->size, + l->ffn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "ffn_norm", + il, + pos0, + n_tokens, + stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_norm", + g->batch_ffn_norm, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } + if (!ok) return false; + + if (il < DS4_N_LEADING_DENSE) { + const uint64_t hidden = l->ffn_gate->dim[1]; + if (hidden == 0 || hidden > UINT32_MAX / n_tokens) return false; + const uint32_t mid_elems = (uint32_t)(hidden * n_tokens); + const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + if (residual_elems > UINT32_MAX) return false; + + const bool fused_gate_up = glm_graph_shared_gate_up_swiglu_q8_0_tensor( + g->batch_ffn_gate, + g->batch_ffn_up, + g->batch_ffn_mid, + model, + l->ffn_gate->abs_offset, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + g->batch_ffn_norm, + n_tokens, + 0.0f); + if (fused_gate_up) { + ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "dense_gate_up_swiglu", + il, + pos0, + n_tokens, + stage_t0); + } else { + ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_gate, + model, + l->ffn_gate->abs_offset, + DS4_N_EMBD, + hidden, + g->batch_ffn_norm, + n_tokens); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, + model, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + g->batch_ffn_norm, + n_tokens); + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "dense_gate_up", + il, + pos0, + n_tokens, + stage_t0); + if (ok) ok = ds4_gpu_swiglu_tensor(g->batch_ffn_mid, + g->batch_ffn_gate, + g->batch_ffn_up, + mid_elems, + 0.0f, + 1.0f) != 0; + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "dense_swiglu", + il, + pos0, + n_tokens, + stage_t0); + } + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_out, + model, + l->ffn_down->abs_offset, + hidden, + DS4_N_EMBD, + g->batch_ffn_mid, + n_tokens); + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "dense_down", + il, + pos0, + n_tokens, + stage_t0); + if (ok) ok = ds4_gpu_add_tensor(next, + after_attn, + g->batch_ffn_out, + (uint32_t)residual_elems) != 0; + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "residual", + il, + pos0, + n_tokens, + stage_t0); + return ok; + } + + if (g->ffn_mid_elems > UINT32_MAX || + g->dense_hidden_max < DS4_N_FF_EXP || + (uint64_t)n_tokens > UINT32_MAX / DS4_N_FF_EXP) { + return false; + } + const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + if (residual_elems > UINT32_MAX) return false; + + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); + (void)gate_in; + (void)up_in; + (void)down_in; + + ok = ds4_gpu_matmul_f32_tensor(g->batch_router_logits, + model->map, + model->size, + l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, + DS4_N_EXPERT, + g->batch_ffn_norm, + n_tokens) != 0; + if (ok) ok = ds4_gpu_glm_router_select_batch_tensor(g->batch_router_selected, + g->batch_router_weights, + g->batch_router_probs, + model->map, + model->size, + l->ffn_exp_probs_b->abs_offset, + g->batch_router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE, + n_tokens) != 0; + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "router", + il, + pos0, + n_tokens, + stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_router_logits", + g->batch_router_logits, + (uint64_t)n_tokens * DS4_N_EXPERT, + il, + pos0); + metal_graph_debug_dump_tensor("glm_ffn_router_probs", + g->batch_router_probs, + (uint64_t)n_tokens * DS4_N_EXPERT, + il, + pos0); + metal_graph_debug_dump_i32_tensor("glm_ffn_router_selected", + g->batch_router_selected, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + metal_graph_debug_dump_tensor("glm_ffn_router_weights", + g->batch_router_weights, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + } + if (ok) ok = glm_graph_profile_router_selection_batch(g, + l, + il, + pos0, + n_tokens); + const bool tp_batch_split_ffn = g->tp_world == 2; + if (ok && tp_batch_split_ffn) { + ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); + } + if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, + il, + n_tokens); + if (ok) ok = glm_graph_seed_streaming_expert_cache_from_full_layer( + g, + model, + weights, + l, + il, + n_tokens, + gate_out * gate_row_bytes, + down_out * down_row_bytes, + full_layer_prefill); + bool shared_done = false; +#define DS4_GLM_ENCODE_FFN_BATCH_SHARED() do { \ + if (ok) { \ + const bool fused_shared = glm_graph_shared_gate_up_swiglu_q8_0_tensor( \ + g->batch_ffn_gate, \ + g->batch_ffn_up, \ + g->batch_shared_mid, \ + model, \ + l->ffn_gate_shexp->abs_offset, \ + l->ffn_up_shexp->abs_offset, \ + DS4_N_EMBD, \ + DS4_N_FF_EXP, \ + g->batch_ffn_norm, \ + n_tokens, \ + 0.0f); \ + if (fused_shared) { \ + ok = glm_graph_prefill_stage_boundary(stage_profile, \ + stage_sync, \ + "glm_ffn", \ + "shared_gate_up_swiglu", \ + il, \ + pos0, \ + n_tokens, \ + stage_t0); \ + } else { \ + ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_gate, \ + model, \ + l->ffn_gate_shexp->abs_offset, \ + DS4_N_EMBD, \ + DS4_N_FF_EXP, \ + g->batch_ffn_norm, \ + n_tokens); \ + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, \ + model, \ + l->ffn_up_shexp->abs_offset, \ + DS4_N_EMBD, \ + DS4_N_FF_EXP, \ + g->batch_ffn_norm, \ + n_tokens); \ + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ + stage_sync, \ + "glm_ffn", \ + "shared_gate_up", \ + il, \ + pos0, \ + n_tokens, \ + stage_t0); \ + if (ok) ok = ds4_gpu_swiglu_tensor(g->batch_shared_mid, \ + g->batch_ffn_gate, \ + g->batch_ffn_up, \ + n_tokens * DS4_N_FF_EXP, \ + 0.0f, \ + 1.0f) != 0; \ + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ + stage_sync, \ + "glm_ffn", \ + "shared_swiglu", \ + il, \ + pos0, \ + n_tokens, \ + stage_t0); \ + } \ + } \ + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, \ + model, \ + l->ffn_down_shexp->abs_offset, \ + DS4_N_FF_EXP, \ + DS4_N_EMBD, \ + g->batch_shared_mid, \ + n_tokens); \ + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, \ + stage_sync, \ + "glm_ffn", \ + "shared_down", \ + il, \ + pos0, \ + n_tokens, \ + stage_t0); \ + if (ok) shared_done = true; \ + } while (0) +#ifdef DS4_ROCM_BUILD + rocm_graph_batch_selected_async_load rocm_batch_selected_async = {0}; + bool rocm_batch_selected_async_started = false; + const bool rocm_batch_selected_shared_overlap = + ok && + g->ssd_streaming && + !g->quality && + n_tokens > 1 && + !full_layer_prefill && + !glm_graph_env_present( + "DS4_ROCM_DISABLE_GLM_STREAMING_PREFILL_SELECTED_ASYNC_LOAD", + "DS4_METAL_DISABLE_GLM_STREAMING_PREFILL_SELECTED_ASYNC_LOAD") && + !glm_graph_env_present( + "DS4_ROCM_DISABLE_STREAMING_PREFILL_SELECTED_ASYNC_LOAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_SELECTED_ASYNC_LOAD") && + glm_graph_stream_prefill_expert_addr_supported(weights, l, il, n_tokens); + if (rocm_batch_selected_shared_overlap) { + uint64_t selected_event = 0; + if (ds4_gpu_signal_selected_readback_ready(&selected_event) != 0) { + rocm_batch_selected_async_started = + rocm_graph_batch_selected_async_load_start( + &rocm_batch_selected_async, + g->batch_router_selected, + model, + l, + il, + n_tokens, + selected_event, + gate_out * gate_row_bytes, + down_out * down_row_bytes); + } + } + if (rocm_batch_selected_async_started) { + DS4_GLM_ENCODE_FFN_BATCH_SHARED(); + const bool finish_ok = + rocm_graph_batch_selected_async_load_finish(&rocm_batch_selected_async); + if (!finish_ok) rocm_batch_selected_async_started = false; + } +#endif + if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { /* ablate: keep the gate */ } else + if (ok) ok = glm_graph_routed_moe_batch_dispatch( + g, + model, + l, + il, + tp_batch_split_ffn ? g->tp_bounce_out : g->batch_ffn_out, + g->batch_ffn_mid, + gate_out * gate_row_bytes, + gate_row_bytes, + up_out * up_row_bytes, + up_row_bytes, + down_out * down_row_bytes, + down_row_bytes, + g->batch_router_selected, + g->batch_router_weights, + g->batch_ffn_norm, + n_tokens, + (uint32_t)g->ffn_mid_elems, + full_layer_prefill, + false) != 0; + if (ok && g->tp_world == 2) { + ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_ffn_out, n_tokens); + if (!ok) fprintf(stderr, "ds4: GLM TP batch gate failed (layer %u)\n", il); + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "routed_moe", + il, + pos0, + n_tokens, + stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_routed_out", + g->batch_ffn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } + if (ok && !shared_done) DS4_GLM_ENCODE_FFN_BATCH_SHARED(); +#undef DS4_GLM_ENCODE_FFN_BATCH_SHARED + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_shared_out", + g->batch_attn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } + if (ok && !glm_graph_disable_add3_residual()) { + ok = ds4_gpu_add3_tensor(next, + after_attn, + g->batch_ffn_out, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + } else if (ok) { + ok = ds4_gpu_add_tensor(g->batch_heads, + g->batch_ffn_out, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + if (ok) ok = ds4_gpu_add_tensor(next, + after_attn, + g->batch_heads, + (uint32_t)residual_elems) != 0; + } + if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, + stage_sync, + "glm_ffn", + "residual", + il, + pos0, + n_tokens, + stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_next", + next, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } + return ok; +} + + +static bool glm_graph_begin_commands_if_needed(void); + +/* ------------------------------------------------------------------------ + * GLM MTP (nextn block) drafting. + * + * blk.(N_LAYER-1) is GLM 5.2's multi-token-prediction block: a full + * attention+MoE layer fed with eh_proj(concat(enorm(embed(token[p+1])), + * hnorm(h[p]))), predicting token[p+2] through the shared output head. + * It keeps a private compact KV cache (slot = absolute position; only + * positions >= mtp_min_pos are ever selected, so the unwritten prompt + * range is never read). Under TP the routed experts are combined over + * the BIG-gate exchange, never the decode row gate, so the RDMA row-gate + * schedule stays intact. + * --------------------------------------------------------------------- */ +static bool glm_graph_mtp_ensure(ds4_glm_gpu_graph *g) { + if (g->mtp_ready) return true; + if (g->compact_cache_cap == 0) return false; + const uint64_t elem = glm_graph_compact_cache_elem_bytes(); + const uint64_t kv_bytes = (uint64_t)g->compact_cache_cap * DS4_N_KV_LORA * elem; + const uint64_t rope_bytes = (uint64_t)g->compact_cache_cap * DS4_N_ROT * elem; + g->mtp_kv_lora_cache = ds4_gpu_tensor_alloc(kv_bytes); + g->mtp_k_rope_cache = ds4_gpu_tensor_alloc(rope_bytes); + g->mtp_concat = ds4_gpu_tensor_alloc(2ull * DS4_N_EMBD * sizeof(float)); + g->mtp_selected = ds4_gpu_tensor_alloc((uint64_t)g->compact_cache_cap * sizeof(int32_t)); + g->mtp_logits_host = malloc((size_t)DS4_N_VOCAB * sizeof(float)); + if (!g->mtp_kv_lora_cache || !g->mtp_k_rope_cache || !g->mtp_concat || + !g->mtp_selected || !g->mtp_logits_host) { + return false; + } + g->mtp_ready = 1; + return true; +} + +/* One MTP step at (absolute) position pos: consumes the main model's last + * hidden h[pos] (expected in g->cur, pre output-norm) and next_token + * (= token[pos+1]), writes the nextn KV at slot pos, and returns the + * drafted token[pos+2] by greedy argmax. Clobbers the decode scratch. */ +static bool glm_graph_mtp_step( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int next_token, + uint32_t pos, + uint32_t min_pos, + int *draft_out) { + if (!g || !model || !weights || !draft_out) return false; + if (DS4_N_NEXTN_PREDICT == 0) return false; + if (pos >= g->compact_cache_cap || min_pos > pos) { + fprintf(stderr, "ds4: glm mtp: pos %u/min %u out of range (cap %u)\n", + pos, min_pos, g->compact_cache_cap); + return false; + } + const uint32_t il = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + const ds4_layer_weights *l = &weights->layer[il]; + if (!l->nextn_eh_proj || !l->nextn_enorm || !l->nextn_hnorm || + !l->nextn_shared_head_norm || !l->ffn_gate_exps) { + fprintf(stderr, "ds4: glm mtp: nextn weights missing at layer %u\n", il); + return false; + } + const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + const uint32_t n_selected = pos - min_pos + 1u; + bool input_ready = false; + + if (g->placement) { + const int embedding_tier = g->placement[0]; + const int mtp_tier = g->placement[il + 1u]; + bool handoff_ok = glm_graph_ws_switch(g, embedding_tier, true); + if (handoff_ok) handoff_ok = glm_graph_begin_commands_if_needed(); + if (handoff_ok) { + handoff_ok = ds4_gpu_embed_token_quant_tensor( + g->next, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + (uint32_t)next_token, + DS4_N_EMBD) != 0; + } + if (handoff_ok) handoff_ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + ds4_gpu_tensor *embedded_next = handoff_ok ? g->next : NULL; + if (handoff_ok) handoff_ok = glm_graph_ws_switch(g, mtp_tier, true); + if (handoff_ok) { + handoff_ok = ds4_gpu_tensor_copy_async( + g->next, embedded_next, + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } + if (!handoff_ok) { + fprintf(stderr, "ds4: glm mtp: multi-tier input handoff failed\n"); + return false; + } + input_ready = true; + } + if (!glm_graph_mtp_ensure(g)) { + fprintf(stderr, "ds4: glm mtp: ensure failed (cap %u)\n", g->compact_cache_cap); + return false; + } + + /* Draft attention window: absolute cache slots [min_pos..pos]. */ + { + int32_t *sel = malloc((size_t)n_selected * sizeof(int32_t)); + if (!sel) return false; + for (uint32_t i = 0; i < n_selected; i++) sel[i] = (int32_t)(min_pos + i); + const int wr = ds4_gpu_tensor_write(g->mtp_selected, 0, sel, + (uint64_t)n_selected * sizeof(int32_t)); + free(sel); + if (!wr) { + fprintf(stderr, "ds4: glm mtp: selected write failed (%u)\n", n_selected); + return false; + } + } + + ds4_gpu_tensor *enorm_view = + ds4_gpu_tensor_view(g->mtp_concat, 0, (uint64_t)DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *hnorm_view = + ds4_gpu_tensor_view(g->mtp_concat, + (uint64_t)DS4_N_EMBD * sizeof(float), + (uint64_t)DS4_N_EMBD * sizeof(float)); + if (!enorm_view || !hnorm_view) { + ds4_gpu_tensor_free(enorm_view); + ds4_gpu_tensor_free(hnorm_view); + fprintf(stderr, "ds4: glm mtp: concat views failed\n"); + return false; + } + + const char *mtp_stage = "begin"; +#define DS4_GLM_MTP_STAGE(name_) do { if (ok) mtp_stage = (name_); } while (0) + bool ok = glm_graph_begin_commands_if_needed(); + /* MTP input: concat(enorm(embed(next_token)), hnorm(h)) -> eh_proj. */ + if (ok && !input_ready) { + ok = ds4_gpu_embed_token_quant_tensor(g->next, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + (uint32_t)next_token, + DS4_N_EMBD) != 0; + } + DS4_GLM_MTP_STAGE("enorm"); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(enorm_view, + g->next, + model->map, + model->size, + l->nextn_enorm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("hnorm"); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(hnorm_view, + g->cur, + model->map, + model->size, + l->nextn_hnorm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("eh_proj"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->cur, + model, + l->nextn_eh_proj->abs_offset, + 2ull * DS4_N_EMBD, + DS4_N_EMBD, + g->mtp_concat, + false); + /* nextn attention (full causal over the MTP window, no indexer). */ + DS4_GLM_MTP_STAGE("attn_norm"); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->attn_norm, + g->cur, + model->map, + model->size, + l->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("q_a"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->attn_norm, + false); + DS4_GLM_MTP_STAGE("q_a_norm"); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->q_rank_norm, + g->q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("q_b"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->q_rank_norm, + false); + DS4_GLM_MTP_STAGE("rope"); + if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->q, + 1, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_ROT, + pos, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_GLM_MTP_STAGE("kv_a"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + false); + DS4_GLM_MTP_STAGE("kv_norm"); + if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, + g->kv_raw, + model->map, + model->size, + l->attn_kv_a_norm->abs_offset, + 1, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("kv_store"); + if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(g->mtp_kv_lora_cache, + g->mtp_k_rope_cache, + g->kv_norm, + g->kv_raw, + pos, + 1, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16()) != 0; + DS4_GLM_MTP_STAGE("qk_low"); + if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor(g->qk_low, + g->q, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_KEY_MLA) != 0; + DS4_GLM_MTP_STAGE("attention"); + if (ok) ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor(g->heads, + g->q, + g->qk_low, + g->mtp_kv_lora_cache, + g->mtp_k_rope_cache, + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + g->mtp_selected, + n_selected, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_GLM_MTP_STAGE("attn_out"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->heads, + false); + DS4_GLM_MTP_STAGE("ffn_norm"); + if (ok) ok = ds4_gpu_add_rms_norm_weight_tensor(g->ffn_norm, + g->after_attn, + g->cur, + g->attn_out, + model->map, + model->size, + l->ffn_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + /* nextn sparse FFN: router + split routed experts (BIG-gate combine + * under TP) + shared expert. */ + DS4_GLM_MTP_STAGE("router"); + if (ok) ok = ds4_gpu_matmul_f32_tensor(g->router_logits, + model->map, + model->size, + l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, + DS4_N_EXPERT, + g->ffn_norm, + 1) != 0; + DS4_GLM_MTP_STAGE("router_select"); + if (ok) ok = ds4_gpu_glm_router_select_tensor(g->router_selected, + g->router_weights, + g->router_probs, + model->map, + model->size, + l->ffn_exp_probs_b->abs_offset, + g->router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE) != 0; + DS4_GLM_MTP_STAGE("routed"); + if (ok) { + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); +#ifdef DS4_ROCM_BUILD + if (g->ssd_streaming) { + const ds4_gpu_stream_expert_table table = { + .model_map = model->map, + .model_size = model->size, + .layer = il, + .n_total_expert = DS4_N_EXPERT, + .gate_offset = l->ffn_gate_exps->abs_offset, + .up_offset = l->ffn_up_exps->abs_offset, + .down_offset = l->ffn_down_exps->abs_offset, + .gate_expert_bytes = gate_out * gate_row_bytes, + .down_expert_bytes = down_out * down_row_bytes, + }; + ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + &table, g->router_selected, DS4_N_EXPERT_USED) != 0; + } +#endif + const bool tp_split = g->tp_world == 2 && g->tp_out && g->tp_in; + ds4_gpu_tensor *routed_dst = g->ffn_out; + if (tp_split) { + ok = glm_graph_tp_batch_bounce_ready(g, 1); + routed_dst = g->tp_bounce_out; + } + if (ok) ok = glm_graph_routed_moe_one_dispatch(g, + model, + l, + il, + routed_dst, + g->ffn_mid, + gate_out * gate_row_bytes, + gate_row_bytes, + up_out * up_row_bytes, + up_row_bytes, + down_out * down_row_bytes, + down_row_bytes, + g->router_selected, + g->router_weights, + g->ffn_norm, + false) != 0; + if (ok && tp_split) { + ok = glm_graph_tp_batch_ffn_combine(g, il, g->ffn_out, 1); + } + } + DS4_GLM_MTP_STAGE("shared"); + if (ok) ok = glm_graph_encode_shared_swiglu_one(g->ffn_mid, + g->ffn_gate, + g->ffn_up, + model, + l, + il, + pos, + g->ffn_norm, + false, + false, + NULL); + DS4_GLM_MTP_STAGE("shared_down"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->ffn_sum, + model, + l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + g->ffn_mid, + false); + DS4_GLM_MTP_STAGE("residual"); + if (ok) ok = ds4_gpu_add3_tensor(g->next, + g->after_attn, + g->ffn_out, + g->ffn_sum, + DS4_N_EMBD) != 0; + /* Shared output head behind the nextn head norm. */ + DS4_GLM_MTP_STAGE("head_norm"); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, + g->next, + model->map, + model->size, + l->nextn_shared_head_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_MTP_STAGE("head"); + if (ok) ok = glm_graph_matmul_q8_0_decode_tensor(g->logits, + model, + weights->output->abs_offset, + DS4_N_EMBD, + DS4_N_VOCAB, + g->output_norm, + false); + DS4_GLM_MTP_STAGE("end"); + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok) { + ok = ds4_gpu_tensor_read(g->logits, + 0, + g->mtp_logits_host, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + ds4_gpu_tensor_free(enorm_view); + ds4_gpu_tensor_free(hnorm_view); + if (!ok) { + fprintf(stderr, "ds4: glm mtp step failed at stage '%s' (pos %u)\n", + mtp_stage, pos); + return false; + } +#undef DS4_GLM_MTP_STAGE + int best = 0; + float best_v = g->mtp_logits_host[0]; + for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { + if (g->mtp_logits_host[i] > best_v) { + best_v = g->mtp_logits_host[i]; + best = (int)i; + } + } + *draft_out = best; + return true; +} + +static bool glm_graph_forward_token( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + const float *input_hc, + uint32_t pos, + float *output_hc, + float *logits_out, + bool defer_completion); + + +/* Decode-style verify pass for tiny row counts (MTP): the indexed batch + * fn measures ~1.4ms/layer at n=2 (gate-profile: gpu-wait 1.22ms/layer) + * while decode does the same math in 0.79ms. This pass mirrors the + * decode encoders at n rows over the batch scratch, attends causally + * over the compact caches (valid while pos+n fits the indexer window), + * and reuses the batch FFN encoder (routed split + big-gate combine). + * KV/indexer caches are updated exactly like the indexed batch path. */ +static bool glm_graph_verify_rows( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + uint32_t pos, + uint32_t n, + float *output_hc, + float *logits_out) { + if (!g || !model || !weights || !tokens || n == 0 || + g->compact_cache_cap == 0 || + pos + n > g->compact_cache_cap || + !g->batch_cur || !g->batch_next || !g->prefill_tokens) { + return false; + } + const uint32_t executable = glm_graph_normal_layer_count(); + ds4_gpu_tensor *cur = g->batch_cur; + ds4_gpu_tensor *nxt = g->batch_next; + if (!ds4_gpu_tensor_write(g->prefill_tokens, 0, tokens, + (uint64_t)n * sizeof(int32_t))) { + return false; + } + if (g->placement && + !glm_graph_verify_ws_switch(g, g->placement[0], false, n)) { + glm_graph_verify_ws_restore(g); + return false; + } + cur = g->batch_cur; + nxt = g->batch_next; + bool ok = glm_graph_begin_commands_if_needed(); + if (ok) ok = ds4_gpu_embed_tokens_quant_tensor(cur, + g->prefill_tokens, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + n, + DS4_N_EMBD) != 0; + for (uint32_t il = 0; ok && il < executable; il++) { + if (g->placement) { + g->batch_cur = cur; + g->batch_next = nxt; + ok = glm_graph_verify_ws_switch(g, + g->placement[il + 1u], + il != 0u, + n); + if (ok) { + ok = glm_graph_ws_switch(g, + g->placement[il + 1u], + false); + } + cur = g->batch_cur; + nxt = g->batch_next; + } + const ds4_layer_weights *l = &weights->layer[il]; + const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, + cur, + model->map, + model->size, + l->attn_norm->abs_offset, + DS4_N_EMBD, + n, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->batch_attn_norm, + n); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, + g->batch_q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + n, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->batch_q_rank_norm, + n); + if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->batch_q, + n, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_ROT, + pos, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok && glm_graph_layer_uses_full_indexer(il)) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, + model, + l->indexer_attn_k->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD_DIM, + cur, + n); + if (ok) ok = ds4_gpu_glm_store_indexer_k_tensor( + g->layer_indexer_key_cache[il], + g->batch_indexer_k, + model->map, + model->size, + l->indexer_k_norm->abs_offset, + l->indexer_k_norm_b->abs_offset, + pos, + n, + g->compact_cache_cap, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + 0, + 1.0e-6f, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + glm_graph_compact_cache_is_f16()) != 0; + } + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->batch_attn_norm, + n); + if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, + g->batch_kv_raw, + model->map, + model->size, + l->attn_kv_a_norm->abs_offset, + n, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + g->batch_kv_norm, + g->batch_kv_raw, + pos, + n, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16()) != 0; + if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_batch_tensor(g->batch_qk_low, + g->batch_q, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + n, + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_KEY_MLA) != 0; + if (ok) ok = ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + g->batch_attn_lora, + g->batch_q, + g->batch_qk_low, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + n, + pos, + pos + n, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_glm_value_project_typed_batch_heads_tensor( + g->batch_heads, + g->batch_attn_lora, + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + n, + DS4_N_HEAD, + DS4_N_KV_LORA, + DS4_N_VALUE_MLA) != 0; + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->batch_heads, + n); + if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, + cur, + g->batch_attn_out, + (uint32_t)((uint64_t)n * DS4_N_EMBD)) != 0; + if (ok) ok = glm_graph_encode_ffn_batch(g, + model, + weights, + l, + il, + pos, + g->batch_after_attn, + nxt, + n, + false, + false, + false, + NULL); + if (ok) { + ds4_gpu_tensor *tmp = cur; + cur = nxt; + nxt = tmp; + } + } + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + if (ok && output_hc) { + ok = ds4_gpu_tensor_read(cur, + 0, + output_hc, + (uint64_t)n * DS4_N_EMBD * sizeof(float)) != 0; + } + if (ok && logits_out) { + ds4_gpu_tensor *last = glm_graph_tensor_row_view_strided(cur, + n - 1u, + DS4_N_EMBD, + DS4_N_EMBD); + ok = last != NULL && + glm_graph_forward_output_head(g, model, weights, last, logits_out); + ds4_gpu_tensor_free(last); + } + glm_graph_verify_ws_restore(g); + return ok; +} + +static bool glm_graph_forward_tokens( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + const float *input_hc, + uint32_t pos0, + uint32_t n_tokens, + float *output_hc, + float *logits_out, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t display_absolute_base, + uint32_t work_done_base, + uint32_t work_total) { + if (!g || !model || !weights || !tokens || + n_tokens == 0 || + g->layer_count == 0 || + !glm_graph_span_fits_context(g, pos0, n_tokens)) { + return false; + } + if (!glm_graph_span_fits_full_attention(g, pos0, n_tokens)) { + glm_graph_log_full_attention_limit(g, pos0, n_tokens); + return false; + } + if (!g->full_kv_cache) { + fprintf(stderr, + "ds4: GLM full-attention prefill was requested without an expanded KV cache\n"); + return false; + } + for (uint32_t i = 0; i < n_tokens; i++) { + if (tokens[i] < 0 || tokens[i] >= (int)DS4_N_VOCAB) return false; + } + if (!input_hc && !g->has_token_embd) return false; + if (logits_out && !g->has_output_head) return false; + glm_graph_reset_prefill_seed_capture(g); + const uint32_t n_rows = pos0 + n_tokens; + const bool trace = glm_graph_full_prefill_trace_enabled(); + const bool trace_all = trace && glm_graph_full_prefill_trace_all(); + const double trace_slow_ms = trace ? + (double)glm_graph_full_prefill_trace_slow_ms() : 0.0; + const double trace_chunk_t0 = trace ? now_sec() : 0.0; + if (trace) { + glm_graph_full_prefill_tracef( + "chunk begin pos=%u tokens=%u rows=%u compact_cap=%u work_base=%u work_total=%u", + pos0, + n_tokens, + n_rows, + g->compact_cache_cap, + work_done_base, + work_total); + } + if (g->compact_cache_cap != 0) { + const double trace_cache_t0 = trace ? now_sec() : 0.0; + if (!glm_graph_ensure_compact_cache(g, n_rows)) { + if (trace) { + glm_graph_full_prefill_tracef( + "ensure_cache failed pos=%u tokens=%u rows=%u compact_cap=%u", + pos0, + n_tokens, + n_rows, + g->compact_cache_cap); + } + return false; + } + if (trace) { + const double ms = (now_sec() - trace_cache_t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms) { + glm_graph_full_prefill_tracef( + "ensure_cache done pos=%u tokens=%u rows=%u compact_cap=%u %.3f ms", + pos0, + n_tokens, + n_rows, + g->compact_cache_cap, + ms); + } + } + } + + const double trace_upload_t0 = trace ? now_sec() : 0.0; + bool ok = glm_graph_upload_tokens(g->prefill_tokens, tokens, n_tokens); + if (trace) { + const double ms = (now_sec() - trace_upload_t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_full_prefill_tracef( + "upload_tokens %s pos=%u tokens=%u %.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + ds4_gpu_tensor *cur = g->batch_cur; + ds4_gpu_tensor *next = g->batch_next; + ds4_gpu_tensor *last_hidden = NULL; + + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + 0, + g->layer_count, + work_total, + true); + + const bool stage_sync = + glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); + const uint32_t layer_flush_interval = stage_sync ? 0u : + glm_graph_full_prefill_layer_flush_interval(n_tokens, + n_rows, + logits_out != NULL); + const uint32_t progress_flush_interval = + glm_graph_prefill_progress_flush_interval(layer_flush_interval, + n_tokens, + display_progress, + work_total); + const bool progress_requested = display_progress && work_total > 0; + const uint32_t drain_interval = + (progress_requested && progress_flush_interval != 0) ? + glm_graph_full_prefill_drain_interval() : 0u; + metal_graph_stream_prepare_slot layer_prepare_slots[DS4_STREAM_PREFILL_MAX_PREPARE_AHEAD]; + memset(layer_prepare_slots, 0, sizeof(layer_prepare_slots)); + const bool full_layer_prefill = + glm_graph_stream_prefill_full_layer_enabled(g, n_tokens); + ds4_gpu_set_glm_streaming_prefill_full_layer(full_layer_prefill); + const bool streaming_prefill_sync_each_layer = + !g->ssd_streaming || + glm_graph_streaming_prefill_sync_each_layer(full_layer_prefill); +#ifdef DS4_ROCM_BUILD + rocm_graph_stream_layer_expert_load rocm_full_layer_load; + memset(&rocm_full_layer_load, 0, sizeof(rocm_full_layer_load)); +#endif + const bool full_layer_prepare_base = + glm_graph_stream_prefill_full_layer_prepare_enabled(g, + full_layer_prefill); + const bool layer_pagein = + full_layer_prepare_base && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_PAGEIN") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PAGEIN"); + const bool layer_readahead = + full_layer_prepare_base && + !layer_pagein && + glm_graph_env_present("DS4_ROCM_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_ENABLE_STREAMING_PREFILL_LAYER_READAHEAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_READAHEAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); + const bool layer_pread = + full_layer_prepare_base && + !layer_pagein && + !layer_readahead && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREAD", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREAD") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE"); + const bool layer_madvise = + full_layer_prepare_base && + !layer_pagein && + !layer_pread && + !layer_readahead && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_PREPARE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_PREPARE") && + !glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_PREFILL_LAYER_MADVISE", + "DS4_METAL_DISABLE_STREAMING_PREFILL_LAYER_MADVISE"); + const bool layer_prepare = + layer_pagein || layer_pread || layer_readahead || layer_madvise; + const bool layer_prepare_overlap = + layer_prepare && + metal_graph_stream_prefill_layer_pagein_overlap_enabled(); + const bool full_layer_flush_intermediate = full_layer_prefill; + const uint32_t layer_prepare_ahead = + layer_prepare && layer_prepare_overlap ? + metal_graph_stream_prefill_layer_prepare_ahead() : 1u; + if (trace) { + glm_graph_full_prefill_tracef( + "mode pos=%u tokens=%u stage_sync=%u layer_flush_interval=%u progress_flush_interval=%u drain_interval=%u", + pos0, + n_tokens, + stage_sync ? 1u : 0u, + layer_flush_interval, + progress_flush_interval, + drain_interval); + } + if (ok && layer_prepare && g->layer_count > 0 && + !metal_graph_stream_prepare_start_if_needed(NULL, + model, + weights, + g->layer_start, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + full_layer_prefill && + rocm_graph_glm_stream_prefill_full_layer_enabled( + g, + &weights->layer[g->layer_start], + g->layer_start, + n_tokens), + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + } +#ifdef DS4_ROCM_BUILD + if (ok && + full_layer_prefill && + !rocm_graph_glm_stream_layer_expert_load_start_next( + &rocm_full_layer_load, + g, + model, + weights, + g->layer_start, + g->layer_end, + n_tokens)) { + ok = false; + } +#endif + if (ok) { + const double t0 = trace ? now_sec() : 0.0; + if (input_hc) { + ok = ds4_gpu_tensor_write(cur, + 0, + input_hc, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; + } else { + ok = glm_graph_stream_map_token(g, model, weights); + } + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (trace) { + const double ms = (now_sec() - t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_full_prefill_tracef( + "begin_commands%s %s pos=%u tokens=%u %.3f ms", + input_hc ? "_from_hidden" : "", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + } + if (ok && !input_hc) { + const double t0 = trace ? now_sec() : 0.0; + ok = ds4_gpu_embed_tokens_quant_tensor(cur, + g->prefill_tokens, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + n_tokens, + DS4_N_EMBD) != 0; + if (trace) { + const double ms = (now_sec() - t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_full_prefill_tracef( + "embed %s pos=%u tokens=%u %.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + } + if (ok && g->ssd_streaming && streaming_prefill_sync_each_layer) { + ok = ds4_gpu_end_commands() != 0; + } +#define DS4_GLM_PROFILE_PREFILL_STAGE(part_, name_) do { \ + if (ok && layer_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos0, n_tokens, &layer_stage_t0); \ + } else if (ok && stage_sync) { \ + ok = glm_graph_prefill_stage_sync_boundary(); \ + } \ + } while (0) + for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { + const uint32_t slice_layer_done = il - g->layer_start + 1u; + if (layer_prepare && + !metal_graph_stream_prepare_join_layer(NULL, + model, + weights, + il, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + full_layer_prefill && + rocm_graph_glm_stream_prefill_full_layer_enabled( + g, + &weights->layer[il], + il, + n_tokens), + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + break; + } +#ifdef DS4_ROCM_BUILD + if (full_layer_prefill && + !rocm_graph_glm_stream_layer_expert_load_ready( + &rocm_full_layer_load, + g, + model, + weights, + il, + n_tokens)) { + ok = false; + break; + } + if (full_layer_prefill && + !rocm_graph_glm_stream_layer_expert_load_start_next( + &rocm_full_layer_load, + g, + model, + weights, + il + 1u, + g->layer_end, + n_tokens)) { + ok = false; + break; + } +#endif + if (g->ssd_streaming) { + ok = glm_graph_stream_map_prefill_layer(g, + model, + weights, + il, + n_tokens, + full_layer_prefill); + if (ok && layer_prepare && layer_prepare_overlap) { + bool started_future = false; + for (uint32_t ahead = 1; ahead <= layer_prepare_ahead; ahead++) { + if (il + ahead > g->layer_end) break; + started_future = true; + if (!metal_graph_stream_prepare_start_if_needed(NULL, + model, + weights, + il + ahead, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + full_layer_prefill && + rocm_graph_glm_stream_prefill_full_layer_enabled( + g, + &weights->layer[il + ahead], + il + ahead, + n_tokens), + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + break; + } + } + if (ok && !started_future && logits_out) { + metal_graph_stream_readahead_output(model, weights); + } + } + if (ok && !ds4_gpu_commands_active()) { + ok = ds4_gpu_begin_commands() != 0; + } + } + const ds4_layer_weights *l = &weights->layer[il]; + const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + const uint32_t cache_len = pos0 + n_tokens; + const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); + double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; + const double trace_layer_t0 = trace ? now_sec() : 0.0; + bool trace_layer_flushed = false; + if (trace && trace_all) { + glm_graph_full_prefill_tracef( + "layer begin layer=%u pos=%u tokens=%u rows=%u", + il, + pos0, + n_tokens, + n_rows); + } + if (residual_elems > UINT32_MAX) { + ok = false; + break; + } + if (layer_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_attn", + NULL, + il, + pos0, + n_tokens, + &layer_stage_t0); + } + + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, + cur, + model->map, + model->size, + l->attn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attn_norm"); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->batch_attn_norm, + n_tokens); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, + g->batch_q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->batch_q_rank_norm, + n_tokens); + if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, + n_tokens, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_ROT, + pos0, + 0, + false, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "q_path"); + if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il)) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, + model, + l->indexer_attn_k->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD_DIM, + cur, + n_tokens); + if (ok) { + ok = ds4_gpu_glm_store_indexer_k_tensor( + g->layer_indexer_key_cache[il], + g->batch_indexer_k, + model->map, + model->size, + l->indexer_k_norm->abs_offset, + l->indexer_k_norm_b->abs_offset, + pos0, + n_tokens, + g->compact_cache_cap, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + 0, + 1.0e-6f, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + glm_graph_compact_cache_is_f16()) != 0; + } + } + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "indexer_k"); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->batch_attn_norm, + n_tokens); + if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, + g->batch_kv_raw, + model->map, + model->size, + l->attn_kv_a_norm->abs_offset, + n_tokens, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_RMS_EPS) != 0; + if (ok && g->compact_cache_cap != 0) { + ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + g->batch_kv_norm, + g->batch_kv_raw, + pos0, + n_tokens, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16()) != 0; + } + if (ok) ok = ds4_gpu_glm_k_b_project_typed_tensor(g->batch_k_nope, + g->batch_kv_norm, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + n_tokens, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_HEAD) != 0; + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_value, + model, + l->attn_v_b->abs_offset, + DS4_N_KV_LORA, + g->heads_dim, + g->batch_kv_norm, + n_tokens); + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_path"); + const bool flash_requested = glm_graph_use_flash_attention_prefill(n_tokens); + const bool use_staged_flash_kv = + flash_requested && glm_graph_use_flash_attention_staged_kv(pos0, n_tokens, cache_len); + const bool use_flash_attn = flash_requested; + if (ok) { + if (use_staged_flash_kv) { + ok = ds4_gpu_glm_build_kv_cache_flash_tensor(g->layer_key_cache[il], + g->layer_value_cache[il], + g->batch_kv_raw, + g->batch_k_nope, + g->batch_value, + pos0, + n_tokens, + g->ctx_cap, + DS4_N_HEAD, + kv_raw_dim, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + true) != 0; + } else { + ok = ds4_gpu_glm_build_kv_cache_tensor(g->layer_key_cache[il], + g->layer_value_cache[il], + g->batch_kv_raw, + g->batch_k_nope, + g->batch_value, + pos0, + n_tokens, + g->ctx_cap, + DS4_N_HEAD, + kv_raw_dim, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + true) != 0; + } + } + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_cache"); + if (ok) { + if (use_flash_attn) { + if (use_staged_flash_kv) { + ok = ds4_gpu_glm_attention_flash_staged_tensor(g->batch_heads, + g->batch_q, + g->layer_key_cache[il], + g->layer_value_cache[il], + pos0, + n_tokens, + cache_len, + g->ctx_cap, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_VALUE_MLA, + true) != 0; + } else { + ok = ds4_gpu_glm_attention_flash_tensor(g->batch_heads, + g->batch_q, + g->layer_key_cache[il], + g->layer_value_cache[il], + pos0, + n_tokens, + cache_len, + g->ctx_cap, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_VALUE_MLA, + true) != 0; + } + } else { + ok = ds4_gpu_glm_attention_full_tensor(g->batch_heads, + g->batch_q, + g->layer_key_cache[il], + g->layer_value_cache[il], + pos0, + n_tokens, + cache_len, + g->ctx_cap, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_VALUE_MLA, + true) != 0; + } + } + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attention"); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens); + if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, + cur, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attn_output"); + if (ok) ok = glm_graph_encode_ffn_batch(g, + model, + weights, + l, + il, + pos0, + g->batch_after_attn, + next, + n_tokens, + full_layer_prefill, + layer_stage_profile, + stage_sync, + layer_stage_profile ? &layer_stage_t0 : NULL); + if (ok) { + ds4_gpu_tensor *tmp = cur; + cur = next; + next = tmp; + } + if (ok && glm_debug_hidden_dump_layer_match(il)) { + ok = ds4_gpu_end_commands() != 0; + if (ok) { + for (uint32_t r = 0; r < n_tokens; r++) + glm_debug_dump_hidden_layer(cur, r, il, pos0 + r); + glm_debug_dump_raw_layer(g->batch_router_selected, "sel", + (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int32_t), + il, -1); + glm_debug_dump_raw_layer(g->batch_router_weights, "selw", + (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(float), + il, -1); + ok = ds4_gpu_begin_commands() != 0; + } + } + if (ok && + !g->ssd_streaming && + progress_flush_interval != 0 && + (il < g->layer_end || progress_requested) && + (slice_layer_done % progress_flush_interval) == 0) { + const uint32_t work_done = + work_done_base + (uint32_t)(((uint64_t)n_tokens * slice_layer_done) / g->layer_count); + const bool drain_now = + drain_interval != 0 && + il < g->layer_end && + (slice_layer_done % drain_interval) == 0; + const char *command_action = drain_now ? "drain" : "flush"; + const double trace_command_t0 = trace ? now_sec() : 0.0; + if (trace && (trace_all || drain_now)) { + glm_graph_full_prefill_tracef( + "layer %s begin layer=%u pos=%u tokens=%u work=%u/%u", + command_action, + il, + pos0, + n_tokens, + work_done, + work_total); + } + if (drain_now) { + ok = ds4_gpu_end_commands() != 0; + if (ok) ok = ds4_gpu_begin_commands() != 0; + } else { + ok = ds4_gpu_flush_commands() != 0; + } + if (trace) { + const double trace_command_done = now_sec(); + const double command_ms = (trace_command_done - trace_command_t0) * 1000.0; + const double layer_ms = (trace_command_done - trace_layer_t0) * 1000.0; + trace_layer_flushed = true; + if (trace_all || drain_now || + command_ms >= trace_slow_ms || layer_ms >= trace_slow_ms || !ok) { + glm_graph_full_prefill_tracef( + "layer %s %s layer=%u pos=%u tokens=%u command=%.3f ms layer_total=%.3f ms work=%u/%u", + command_action, + ok ? "done" : "failed", + il, + pos0, + n_tokens, + command_ms, + layer_ms, + work_done, + work_total); + } + } + if (ok) { + const bool progress_completed = + drain_interval == 0 || drain_now; + if (progress_completed) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + slice_layer_done, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + } + } + if (g->ssd_streaming) { + if (streaming_prefill_sync_each_layer) { + if (ok && full_layer_flush_intermediate && + il < g->layer_end) { + ok = ds4_gpu_flush_commands() != 0; + } else if (ok) { + ok = ds4_gpu_end_commands() != 0; + } + else (void)ds4_gpu_synchronize(); + } + else if (!ok) (void)ds4_gpu_synchronize(); + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + slice_layer_done, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + } + if (ok && + g->ssd_streaming && + layer_prepare && + !layer_prepare_overlap) { + if (il < g->layer_end) { + if (!metal_graph_stream_prepare_start_if_needed(NULL, + model, + weights, + il + 1u, + n_tokens, + layer_madvise, + layer_pread, + layer_readahead, + full_layer_prefill && + rocm_graph_glm_stream_prefill_full_layer_enabled( + g, + &weights->layer[il + 1u], + il + 1u, + n_tokens), + layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + } + } else if (logits_out) { + metal_graph_stream_readahead_output(model, weights); + } + } + if (trace && ok) { + const double layer_ms = (now_sec() - trace_layer_t0) * 1000.0; + if (trace_all || (!trace_layer_flushed && layer_ms >= trace_slow_ms)) { + glm_graph_full_prefill_tracef( + "layer end layer=%u pos=%u tokens=%u flushed=%u layer_total=%.3f ms", + il, + pos0, + n_tokens, + trace_layer_flushed ? 1u : 0u, + layer_ms); + } + } + } +#undef DS4_GLM_PROFILE_PREFILL_STAGE + if (ok && !g->ssd_streaming) { + const double trace_end_t0 = trace ? now_sec() : 0.0; + if (trace) { + glm_graph_full_prefill_tracef( + "chunk end_commands begin pos=%u tokens=%u", + pos0, + n_tokens); + } + ok = ds4_gpu_end_commands() != 0; + if (trace) { + const double end_ms = (now_sec() - trace_end_t0) * 1000.0; + const double chunk_ms = (now_sec() - trace_chunk_t0) * 1000.0; + glm_graph_full_prefill_tracef( + "chunk end_commands %s pos=%u tokens=%u end=%.3f ms chunk_total=%.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + end_ms, + chunk_ms); + } + } else if (!ok) { + if (trace) { + glm_graph_full_prefill_tracef( + "chunk failed before end pos=%u tokens=%u elapsed=%.3f ms", + pos0, + n_tokens, + (now_sec() - trace_chunk_t0) * 1000.0); + } +#ifdef DS4_ROCM_BUILD + (void)rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load); + if (full_layer_prefill) { + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); + } +#endif + if (layer_prepare) { + (void)metal_graph_stream_prepare_join_all(layer_prepare_slots, + layer_prepare_ahead); + } + (void)ds4_gpu_synchronize(); + } + if (ok && layer_prepare && + !metal_graph_stream_prepare_join_all(layer_prepare_slots, + layer_prepare_ahead)) { + ok = false; + } +#ifdef DS4_ROCM_BUILD + if (!rocm_graph_stream_layer_expert_load_join(&rocm_full_layer_load)) { + ok = false; + } + if (full_layer_prefill) { + (void)ds4_gpu_stream_expert_cache_release_layer_cache(); + } +#endif + if (ok && + g->ssd_streaming && + !streaming_prefill_sync_each_layer && + !output_hc && + !logits_out) { + ok = ds4_gpu_end_commands() != 0; + } + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + g->layer_count, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + if (ok && output_hc) { + ok = ds4_gpu_tensor_read(cur, + 0, + output_hc, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; + } + if (ok && logits_out) { + ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, + model, + weights); + } + if (ok && logits_out) { + last_hidden = glm_graph_tensor_row_view_strided(cur, + n_tokens - 1u, + DS4_N_EMBD, + DS4_N_EMBD); + ok = last_hidden != NULL; + if (ok && g->ssd_streaming) ok = glm_graph_stream_map_output(g, model, weights); + if (ok) ok = glm_graph_forward_output_head(g, model, weights, last_hidden, logits_out); + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + g->layer_count, + g->layer_count, + work_total, + true); + } + } + ds4_gpu_tensor_free(last_hidden); + ds4_gpu_set_glm_streaming_prefill_full_layer(false); + return ok; +} + +static uint32_t glm_graph_prefill_chunk_tokens(uint32_t full_attention_cap) { + return full_attention_cap ? full_attention_cap : DS4_GLM_METAL_FULL_ATTN_DEFAULT_CONTEXT; +} + +static bool glm_graph_forward_indexed_tokens( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + const float *input_hc, + uint32_t pos0, + uint32_t n_tokens, + float *output_hc, + float *logits_out, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t display_absolute_base, + uint32_t work_done_base, + uint32_t work_total) { + if (!g || !model || !weights || !tokens || + g->compact_cache_cap == 0 || + g->indexed_prefill_cap == 0 || + g->indexed_prefill_score_cap == 0 || + !g->batch_indexer_q || + !g->batch_indexer_weights || + !g->batch_indexer_scores || + !g->batch_indexer_selected || + !g->batch_qk_low || + n_tokens == 0 || + g->layer_count == 0 || + n_tokens > g->indexed_prefill_cap || + !glm_graph_span_fits_context(g, pos0, n_tokens)) { + return false; + } + const uint32_t n_rows = pos0 + n_tokens; + const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); + if (pos0 < indexer_top_k && n_rows > indexer_top_k) { + return false; + } + const uint32_t indexed_selected_count = + n_rows <= indexer_top_k ? n_rows : indexer_top_k; + const bool use_causal_range_select = n_rows <= indexer_top_k; + const bool trace = glm_graph_indexed_prefill_trace_enabled(); + const bool trace_all = trace && glm_graph_indexed_prefill_trace_all(); + const double trace_slow_ms = trace ? + (double)glm_graph_indexed_prefill_trace_slow_ms() : 0.0; + const double trace_chunk_t0 = trace ? now_sec() : 0.0; + if (trace) { + glm_graph_indexed_prefill_tracef( + "chunk begin pos=%u tokens=%u rows=%u selected=%u compact_cap=%u score_cap=%u work_base=%u work_total=%u", + pos0, + n_tokens, + n_rows, + indexed_selected_count, + g->compact_cache_cap, + g->indexed_prefill_score_cap, + work_done_base, + work_total); + } + const double trace_cache_t0 = trace ? now_sec() : 0.0; + if (!glm_graph_ensure_compact_cache(g, n_rows)) { + if (trace) { + glm_graph_indexed_prefill_tracef( + "ensure_cache failed pos=%u tokens=%u rows=%u compact_cap=%u", + pos0, + n_tokens, + n_rows, + g->compact_cache_cap); + } + return false; + } + if (trace) { + const double ms = (now_sec() - trace_cache_t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms) { + glm_graph_indexed_prefill_tracef( + "ensure_cache done pos=%u tokens=%u rows=%u compact_cap=%u %.3f ms", + pos0, + n_tokens, + n_rows, + g->compact_cache_cap, + ms); + } + } + for (uint32_t i = 0; i < n_tokens; i++) { + if (tokens[i] < 0 || tokens[i] >= (int)DS4_N_VOCAB) return false; + } + if (!input_hc && !g->has_token_embd) return false; + if (logits_out && !g->has_output_head) return false; + glm_graph_reset_prefill_seed_capture(g); + + const double trace_upload_t0 = trace ? now_sec() : 0.0; + bool ok = glm_graph_upload_tokens(g->prefill_tokens, tokens, n_tokens); + if (trace) { + const double ms = (now_sec() - trace_upload_t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_indexed_prefill_tracef( + "upload_tokens %s pos=%u tokens=%u %.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + ds4_gpu_tensor *cur = g->batch_cur; + ds4_gpu_tensor *next = g->batch_next; + ds4_gpu_tensor *last_hidden = NULL; + + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + 0, + g->layer_count, + work_total, + true); + + const bool use_all_scalar_kernels = + n_tokens == 1u && glm_graph_indexed_prefill_scalar_kernels(); + const bool use_scalar_indexer = + use_all_scalar_kernels || + glm_graph_indexed_prefill_scalar_indexer() || + !glm_graph_indexed_prefill_batch_indexer(); + const bool force_scalar_attn = + use_all_scalar_kernels || glm_graph_indexed_prefill_scalar_attn(); + const bool use_batch_qk_low = + !force_scalar_attn && glm_graph_indexed_prefill_batch_qk_low(); + const bool use_batch_attn_kernel = + !force_scalar_attn && glm_graph_indexed_prefill_batch_attn_kernel(); + const bool use_split_value_proj = + use_batch_attn_kernel && + g->batch_attn_lora; + /* Tensor-parallel attention head split: each rank computes half the + * heads in the qk-low / attention-lora / value-project kernels, the + * unowned half of batch_heads stays zero, and the full-width attn + * output projection yields partials combined over the big-gate + * exchange (same commutative add as the routed-FFN combine). Only the + * split-value-proj batch chain has head ownership. */ + const bool tp_attn_head_split = + g->tp_world == 2 && + use_batch_attn_kernel && + use_split_value_proj && + (DS4_N_HEAD % 16u) == 0u && + n_tokens >= glm_tp_head_split_min(); /* small batches replicate; + * the floor is env-tunable for correctness + * isolation (DS4_GLM_TP_HEAD_SPLIT_MIN). */ + const bool use_batch_q_rank_proj = true; + const bool use_batch_q_proj = true; + const bool use_batch_indexer_k_proj = true; + const bool use_batch_kv_proj = true; + const bool use_batch_indexer_q_proj = true; + const bool use_batch_indexer_weights_proj = true; + const bool use_batch_attn_out_proj = true; + const bool use_batch_ffn = glm_graph_indexed_prefill_batch_ffn(); + const bool stage_sync = + glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); + const uint32_t layer_flush_interval = stage_sync ? 0u : + glm_graph_full_prefill_layer_flush_interval(n_tokens, + n_tokens, + logits_out != NULL); + const uint32_t progress_flush_interval = + glm_graph_prefill_progress_flush_interval(layer_flush_interval, + n_tokens, + display_progress, + work_total); + const uint32_t drain_interval = + progress_flush_interval != 0 ? glm_graph_indexed_prefill_drain_interval() : 0u; + const bool progress_requested = display_progress && work_total > 0; + ds4_gpu_set_glm_streaming_prefill_full_layer(false); + const bool streaming_prefill_sync_each_layer = + !g->ssd_streaming || + glm_graph_streaming_prefill_sync_each_layer(false); + + if (trace) { + glm_graph_indexed_prefill_tracef( + "mode pos=%u tokens=%u scalar_indexer=%u batch_qk_low=%u batch_attn=%u split_value=%u batch_ffn=%u progress_flush_interval=%u drain_interval=%u", + pos0, + n_tokens, + use_scalar_indexer ? 1u : 0u, + use_batch_qk_low ? 1u : 0u, + use_batch_attn_kernel ? 1u : 0u, + use_split_value_proj ? 1u : 0u, + use_batch_ffn ? 1u : 0u, + progress_flush_interval, + drain_interval); + } + + if (ok) { + const double t0 = trace ? now_sec() : 0.0; + if (input_hc) { + ok = ds4_gpu_tensor_write(cur, + 0, + input_hc, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; + } else { + ok = glm_graph_stream_map_token(g, model, weights); + } + if (ok) ok = ds4_gpu_begin_commands() != 0; + if (trace) { + const double ms = (now_sec() - t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_indexed_prefill_tracef( + "begin_commands%s %s pos=%u tokens=%u %.3f ms", + input_hc ? "_from_hidden" : "", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + } + if (ok && !input_hc) { + const double t0 = trace ? now_sec() : 0.0; + ok = ds4_gpu_embed_tokens_quant_tensor(cur, + g->prefill_tokens, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + n_tokens, + DS4_N_EMBD) != 0; + if (trace) { + const double ms = (now_sec() - t0) * 1000.0; + if (trace_all || ms >= trace_slow_ms || !ok) { + glm_graph_indexed_prefill_tracef( + "embed %s pos=%u tokens=%u %.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + ms); + } + } + } + if (ok && g->ssd_streaming && streaming_prefill_sync_each_layer) { + ok = ds4_gpu_end_commands() != 0; + } + +#define DS4_GLM_PROFILE_INDEXED_STAGE(part_, name_) do { \ + if (ok && trace) { \ + const double _trace_stage_now = now_sec(); \ + const double _trace_stage_ms = (_trace_stage_now - trace_stage_t0) * 1000.0; \ + if (trace_all || _trace_stage_ms >= trace_slow_ms) { \ + glm_graph_indexed_prefill_tracef( \ + "stage layer=%u pos=%u tokens=%u %s.%s encode %.3f ms", \ + il, \ + pos0, \ + n_tokens, \ + (part_), \ + (name_), \ + _trace_stage_ms); \ + } \ + trace_stage_t0 = _trace_stage_now; \ + } \ + if (ok && layer_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos0, n_tokens, &layer_stage_t0); \ + } else if (ok && stage_sync) { \ + ok = glm_graph_prefill_stage_sync_boundary(); \ + } \ + } while (0) + ds4_gpu_tensor *last_indexer_selected = NULL; + uint32_t last_indexer_selected_count = 0; + if (ok && tp_attn_head_split) { + /* The unowned head range of batch_heads must be exactly zero so the + * full-width attn-output matmul produces partial sums. Owned heads + * are rewritten every layer, so one fill per chunk suffices. */ + ok = ds4_gpu_tensor_fill_f32(g->batch_heads, 0.0f, + (uint64_t)n_tokens * g->heads_dim) != 0; + } + ds4_gpu_tp_set_attn_head_split(tp_attn_head_split ? 1 : 0); + for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { + const uint32_t slice_layer_done = il - g->layer_start + 1u; + if (g->ssd_streaming) { + ok = glm_graph_stream_map_prefill_layer(g, + model, + weights, + il, + n_tokens, + false); + if (ok) ok = ds4_gpu_begin_commands() != 0; + } + const ds4_layer_weights *l = &weights->layer[il]; + const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); + double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; + double trace_stage_t0 = trace ? now_sec() : 0.0; + const double trace_layer_t0 = trace_stage_t0; + const bool trace_full_indexer = glm_graph_layer_uses_full_indexer(il); + bool trace_layer_flushed = false; + if (trace && (trace_all || trace_full_indexer)) { + glm_graph_indexed_prefill_tracef( + "layer begin layer=%u pos=%u tokens=%u rows=%u selected=%u full_indexer=%u", + il, + pos0, + n_tokens, + n_rows, + indexed_selected_count, + trace_full_indexer ? 1u : 0u); + } + if (residual_elems > UINT32_MAX) { + ok = false; + break; + } + if (layer_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", + NULL, + il, + pos0, + n_tokens, + &layer_stage_t0); + } + + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, + cur, + model->map, + model->size, + l->attn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attn_norm"); + if (ok) { + if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_QPATH)) { /* ablate */ } else + ok = (use_batch_q_rank_proj ? + glm_graph_matmul_q8_0_tensor(g->batch_q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->batch_attn_norm, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(g->batch_q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->batch_attn_norm, + n_tokens)); + } + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_q_rank_norm, + g->batch_q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) { + ok = (use_batch_q_proj ? + glm_graph_matmul_q8_0_tensor(g->batch_q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->batch_q_rank_norm, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(g->batch_q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->batch_q_rank_norm, + n_tokens)); + } + if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, + n_tokens, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_ROT, + pos0, + 0, + false, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "q_path"); + + if (ok && glm_graph_layer_uses_full_indexer(il)) { + ok = (use_batch_indexer_k_proj ? + glm_graph_matmul_q8_0_tensor(g->batch_indexer_k, + model, + l->indexer_attn_k->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD_DIM, + cur, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(g->batch_indexer_k, + model, + l->indexer_attn_k->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD_DIM, + cur, + n_tokens)); + if (ok) { + ok = ds4_gpu_glm_store_indexer_k_tensor( + g->layer_indexer_key_cache[il], + g->batch_indexer_k, + model->map, + model->size, + l->indexer_k_norm->abs_offset, + l->indexer_k_norm_b->abs_offset, + pos0, + n_tokens, + g->compact_cache_cap, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + 0, + 1.0e-6f, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + glm_graph_compact_cache_is_f16()) != 0; + } + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_k"); + + if (ok) { + ok = (use_batch_kv_proj ? + glm_graph_matmul_q8_0_tensor(g->batch_kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->batch_attn_norm, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(g->batch_kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->batch_attn_norm, + n_tokens)); + } + if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->batch_kv_norm, + g->batch_kv_raw, + model->map, + model->size, + l->attn_kv_a_norm->abs_offset, + n_tokens, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_RMS_EPS) != 0; + if (ok) { + ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + g->batch_kv_norm, + g->batch_kv_raw, + pos0, + n_tokens, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16()) != 0; + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "kv_path"); + + if (ok && glm_graph_layer_uses_full_indexer(il)) { + if (ok && !use_causal_range_select) { + ok = (use_batch_indexer_q_proj ? + glm_graph_matmul_q8_0_tensor(g->batch_indexer_q, + model, + l->indexer_attn_q_b->abs_offset, + DS4_N_LORA_Q, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, + g->batch_q_rank_norm, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(g->batch_indexer_q, + model, + l->indexer_attn_q_b->abs_offset, + DS4_N_LORA_Q, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, + g->batch_q_rank_norm, + n_tokens)); + if (ok) ok = ds4_gpu_glm_indexer_rope_tail_tensor(g->batch_indexer_q, + n_tokens, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos0, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok && glm_graph_indexer_qat()) { + ok = ds4_gpu_dsv4_indexer_qat_tensor(g->batch_indexer_q, + n_tokens * DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + if (ok) { + ok = (use_batch_indexer_weights_proj ? + ds4_gpu_matmul_f32_tensor(g->batch_indexer_weights, + model->map, + model->size, + l->indexer_proj->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD, + cur, + n_tokens) != 0 : + glm_graph_matmul_f32_rows_scalar(g->batch_indexer_weights, + model, + l->indexer_proj->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD, + cur, + n_tokens)); + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_q_weights"); + } + if (ok) { + if (use_causal_range_select) { + ok = ds4_gpu_glm_fill_selected_range_batch_tensor( + g->batch_indexer_selected, + n_tokens, + pos0, + indexed_selected_count, + g->compact_cache_cap) != 0; + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_range"); + } else if (use_scalar_indexer) { + const float indexer_scale = + 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + for (uint32_t t = 0; ok && t < n_tokens; t++) { + const uint32_t visible = pos0 + t + 1u; + ds4_gpu_tensor *indexer_q_view = + glm_graph_tensor_row_view_strided( + g->batch_indexer_q, + t, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); + ds4_gpu_tensor *indexer_weights_view = + glm_graph_tensor_row_view_strided(g->batch_indexer_weights, + t, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD); + ds4_gpu_tensor *scores_view = + ds4_gpu_tensor_view(g->batch_indexer_scores, + 0, + (uint64_t)visible * sizeof(float)); + ds4_gpu_tensor *selected_view = + ds4_gpu_tensor_view(g->batch_indexer_selected, + (uint64_t)t * indexed_selected_count * sizeof(uint32_t), + (uint64_t)indexed_selected_count * sizeof(uint32_t)); + ok = indexer_q_view && indexer_weights_view && scores_view && selected_view; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create indexer row views at layer %u token %u\n", il, t); + if (ok) { + int rc = ds4_gpu_glm_indexer_score_one_tensor( + scores_view, + indexer_q_view, + indexer_weights_view, + g->layer_indexer_key_cache[il], + visible, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + indexer_scale, + glm_graph_compact_cache_is_f16()); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill indexer scores failed at layer %u token %u\n", il, t); + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_score_scalar"); + if (ok) { + int rc = ds4_gpu_indexer_topk_tensor(selected_view, + scores_view, + visible, + 1, + indexed_selected_count); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill topk failed at layer %u token %u\n", il, t); + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_topk_scalar"); + ds4_gpu_tensor_free(selected_view); + ds4_gpu_tensor_free(scores_view); + ds4_gpu_tensor_free(indexer_weights_view); + ds4_gpu_tensor_free(indexer_q_view); + } + } else { + const float indexer_scale = + 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + const uint32_t score_cap = + g->indexed_prefill_score_cap != 0 ? + g->indexed_prefill_score_cap : + g->indexed_prefill_cap; + const uint64_t indexer_q_row_bytes = + (uint64_t)DS4_N_INDEXER_HEAD * + DS4_N_INDEXER_HEAD_DIM * + sizeof(float); + const uint64_t indexer_weights_row_bytes = + (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); + const uint64_t selected_row_bytes = + (uint64_t)indexed_selected_count * sizeof(uint32_t); + for (uint32_t t0 = 0; ok && t0 < n_tokens; ) { + uint32_t slice = n_tokens - t0; + if (slice > score_cap) slice = score_cap; + if (slice == 0) { + ok = false; + break; + } + + ds4_gpu_tensor *indexer_q_view = + ds4_gpu_tensor_view(g->batch_indexer_q, + (uint64_t)t0 * indexer_q_row_bytes, + (uint64_t)slice * indexer_q_row_bytes); + ds4_gpu_tensor *indexer_weights_view = + ds4_gpu_tensor_view(g->batch_indexer_weights, + (uint64_t)t0 * indexer_weights_row_bytes, + (uint64_t)slice * indexer_weights_row_bytes); + ds4_gpu_tensor *selected_view = + ds4_gpu_tensor_view(g->batch_indexer_selected, + (uint64_t)t0 * selected_row_bytes, + (uint64_t)slice * selected_row_bytes); + ok = indexer_q_view && indexer_weights_view && selected_view; + if (!ok) { + fprintf(stderr, + "ds4: GLM indexed prefill failed to create indexer score slice views at layer %u token %u\n", + il, + t0); + } + if (ok) { + ok = ds4_gpu_glm_indexer_scores_batch_tensor( + g->batch_indexer_scores, + indexer_q_view, + indexer_weights_view, + g->layer_indexer_key_cache[il], + n_rows, + slice, + pos0 + t0, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + indexer_scale, + glm_graph_compact_cache_is_f16()) != 0; + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_score"); + if (ok) { + ok = ds4_gpu_indexer_topk_tensor(selected_view, + g->batch_indexer_scores, + n_rows, + slice, + indexed_selected_count) != 0; + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_topk"); + ds4_gpu_tensor_free(selected_view); + ds4_gpu_tensor_free(indexer_weights_view); + ds4_gpu_tensor_free(indexer_q_view); + t0 += slice; + } + } + } + if (ok) { + last_indexer_selected = g->batch_indexer_selected; + last_indexer_selected_count = indexed_selected_count; + } + } else if (ok && (!last_indexer_selected || last_indexer_selected_count == 0)) { + ok = false; + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_select"); + metal_graph_debug_dump_tensor("glm_indexed_q", + g->batch_q, + (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KEY_MLA, + il, + pos0); + if (use_batch_qk_low) { + if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_CORE)) { /* ablate */ } else + if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_batch_tensor(g->batch_qk_low, + g->batch_q, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + n_tokens, + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_KEY_MLA) != 0; + } else { + for (uint32_t t = 0; ok && t < n_tokens; t++) { + ds4_gpu_tensor *q_view = + glm_graph_tensor_row_view_strided(g->batch_q, + t, + g->q_dim, + g->q_dim); + ds4_gpu_tensor *qk_low_view = + glm_graph_tensor_row_view_strided( + g->batch_qk_low, + t, + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); + ok = q_view && qk_low_view; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create qk-low row views at layer %u token %u\n", il, t); + if (ok) { + int rc = ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low_view, + q_view, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_KEY_MLA); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill qk-low failed at layer %u token %u\n", il, t); + } + ds4_gpu_tensor_free(qk_low_view); + ds4_gpu_tensor_free(q_view); + } + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "qk_low"); + metal_graph_debug_dump_tensor("glm_indexed_qk_low", + g->batch_qk_low, + (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KV_LORA, + il, + pos0); + if (ok && use_batch_attn_kernel) ok = glm_graph_indexed_prefill_attention_boundary(); + + if (use_batch_attn_kernel) { + const uint32_t attn_slice_cap = + glm_graph_indexed_prefill_batch_attn_slice_tokens(); + for (uint32_t t0 = 0; ok && t0 < n_tokens; ) { + uint32_t slice = n_tokens - t0; + if (slice > attn_slice_cap) slice = attn_slice_cap; + + ds4_gpu_tensor *q_view = + ds4_gpu_tensor_view(g->batch_q, + (uint64_t)t0 * g->q_dim * sizeof(float), + (uint64_t)slice * g->q_dim * sizeof(float)); + ds4_gpu_tensor *qk_low_view = + ds4_gpu_tensor_view(g->batch_qk_low, + (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), + (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)); + ds4_gpu_tensor *heads_view = + ds4_gpu_tensor_view(g->batch_heads, + (uint64_t)t0 * g->heads_dim * sizeof(float), + (uint64_t)slice * g->heads_dim * sizeof(float)); + ds4_gpu_tensor *attn_lora_view = use_split_value_proj ? + ds4_gpu_tensor_view(g->batch_attn_lora, + (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), + (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)) : + NULL; + ds4_gpu_tensor *selected_view = + ds4_gpu_tensor_view(last_indexer_selected, + (uint64_t)t0 * last_indexer_selected_count * sizeof(uint32_t), + (uint64_t)slice * last_indexer_selected_count * sizeof(uint32_t)); + ok = q_view && qk_low_view && heads_view && selected_view && + (!use_split_value_proj || attn_lora_view); + if (!ok) { + fprintf(stderr, "ds4: GLM sliced indexed prefill failed to create attention views at layer %u token %u\n", il, t0); + } + if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_CORE)) { /* ablate */ } else if (ok && use_split_value_proj) { + int rc = 0; + if (use_causal_range_select) { + rc = ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + attn_lora_view, + q_view, + qk_low_view, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + slice, + pos0 + t0, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW); + } else { + rc = ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + attn_lora_view, + q_view, + qk_low_view, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + selected_view, + slice, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW); + } + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill attention-lora failed at layer %u token %u\n", il, t0); + if (ok && layer_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", + "attention_lora", + il, + pos0 + t0, + slice, + &layer_stage_t0); + } + if (ok) { + rc = ds4_gpu_glm_value_project_typed_batch_heads_tensor( + heads_view, + attn_lora_view, + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + slice, + DS4_N_HEAD, + DS4_N_KV_LORA, + DS4_N_VALUE_MLA); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill value project failed at layer %u token %u\n", il, t0); + } + if (ok && layer_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", + "value_project", + il, + pos0 + t0, + slice, + &layer_stage_t0); + } + } else if (ok) { + int rc = ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads_view, + q_view, + qk_low_view, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + selected_view, + slice, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM sliced indexed prefill indexed attention failed at layer %u token %u\n", il, t0); + if (ok && layer_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_indexed_attn", + "attention_fused", + il, + pos0 + t0, + slice, + &layer_stage_t0); + } + } + ds4_gpu_tensor_free(selected_view); + ds4_gpu_tensor_free(attn_lora_view); + ds4_gpu_tensor_free(heads_view); + ds4_gpu_tensor_free(qk_low_view); + ds4_gpu_tensor_free(q_view); + t0 += slice; + } + if (ok) ok = glm_graph_indexed_prefill_attention_boundary(); + } else { + for (uint32_t t = 0; ok && t < n_tokens; t++) { + ds4_gpu_tensor *q_view = + glm_graph_tensor_row_view_strided(g->batch_q, + t, + g->q_dim, + g->q_dim); + ds4_gpu_tensor *qk_low_view = + glm_graph_tensor_row_view_strided( + g->batch_qk_low, + t, + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); + ds4_gpu_tensor *heads_view = + glm_graph_tensor_row_view_strided(g->batch_heads, + t, + g->heads_dim, + g->heads_dim); + ds4_gpu_tensor *selected_view = + ds4_gpu_tensor_view(last_indexer_selected, + (uint64_t)t * last_indexer_selected_count * sizeof(uint32_t), + (uint64_t)last_indexer_selected_count * sizeof(uint32_t)); + ok = q_view && qk_low_view && heads_view && selected_view; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create attention row views at layer %u token %u\n", il, t); + if (ok) { + int rc = ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads_view, + q_view, + qk_low_view, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + selected_view, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW); + ok = rc != 0; + if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill indexed attention failed at layer %u token %u\n", il, t); + } + ds4_gpu_tensor_free(selected_view); + ds4_gpu_tensor_free(heads_view); + ds4_gpu_tensor_free(qk_low_view); + ds4_gpu_tensor_free(q_view); + } + } + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attention"); + metal_graph_debug_dump_tensor("glm_indexed_heads", + g->batch_heads, + (uint64_t)n_tokens * g->heads_dim, + il, + pos0); + if (ok && tp_attn_head_split) { + ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); + } + if (ok) { + /* Under the head split the projection input has zeros in the + * unowned head columns, so the result is this rank's partial; + * it must land in the shared bounce tensor for the exchange. */ + ds4_gpu_tensor *attn_out_dst = + tp_attn_head_split ? g->tp_bounce_out : g->batch_attn_out; + if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_OUT)) { /* ablate */ } else + ok = (use_batch_attn_out_proj ? + glm_graph_matmul_q8_0_tensor(attn_out_dst, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens) : + glm_graph_matmul_q8_0_rows_scalar(attn_out_dst, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens)); + } + if (ok && tp_attn_head_split) { + ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_attn_out, n_tokens); + } + if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, + cur, + g->batch_attn_out, + (uint32_t)residual_elems) != 0; + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attn_output"); + metal_graph_debug_dump_tensor("glm_indexed_after_attn", + g->batch_after_attn, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + metal_graph_debug_dump_tensor("glm_indexed_attn_out", + g->batch_attn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + if (ok && use_batch_ffn) { + ok = glm_graph_encode_ffn_batch(g, + model, + weights, + l, + il, + pos0, + g->batch_after_attn, + next, + n_tokens, + false, + layer_stage_profile, + stage_sync, + layer_stage_profile ? &layer_stage_t0 : NULL); + } else if (ok) { + const bool use_batch_ffn_norm = + n_tokens > 1 && glm_graph_indexed_prefill_batch_ffn_norm(); + if (use_batch_ffn_norm) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, + g->batch_after_attn, + model->map, + model->size, + l->ffn_norm->abs_offset, + DS4_N_EMBD, + n_tokens, + DS4_RMS_EPS) != 0; + } + if (use_batch_ffn_norm) { + DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_ffn", "ffn_norm"); + } + if (ok && + use_batch_ffn_norm && + il >= DS4_N_LEADING_DENSE && + glm_graph_indexed_prefill_batch_routed_moe()) { + ok = glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( + g, + model, + l, + il, + pos0, + g->batch_after_attn, + next, + n_tokens, + layer_stage_profile, + stage_sync, + layer_stage_profile ? &layer_stage_t0 : NULL); + } else for (uint32_t t = 0; ok && t < n_tokens; t++) { + ds4_gpu_tensor *after_attn_view = + glm_graph_tensor_row_view_strided(g->batch_after_attn, + t, + DS4_N_EMBD, + DS4_N_EMBD); + ds4_gpu_tensor *ffn_norm_view = use_batch_ffn_norm ? + glm_graph_tensor_row_view_strided(g->batch_ffn_norm, + t, + DS4_N_EMBD, + DS4_N_EMBD) : + NULL; + ds4_gpu_tensor *next_view = + glm_graph_tensor_row_view_strided(next, + t, + DS4_N_EMBD, + DS4_N_EMBD); + ok = after_attn_view && next_view && + (!use_batch_ffn_norm || ffn_norm_view); + if (ok && use_batch_ffn_norm) { + ok = glm_graph_encode_ffn_one_normed_from(g, + model, + l, + il, + pos0 + t, + ffn_norm_view, + after_attn_view, + next_view, + g->ffn_gate, + g->ffn_up, + g->ffn_mid, + g->ffn_out, + g->ffn_sum, + g->attn_out, + false, + NULL); + } else if (ok) { + ok = glm_graph_encode_ffn_one_from(g, + model, + l, + il, + pos0 + t, + after_attn_view, + next_view, + g->ffn_norm, + g->ffn_gate, + g->ffn_up, + g->ffn_mid, + g->ffn_out, + g->ffn_sum, + g->attn_out, + false, + NULL); + } + ds4_gpu_tensor_free(next_view); + ds4_gpu_tensor_free(ffn_norm_view); + ds4_gpu_tensor_free(after_attn_view); + } + } + if (ok) { + ds4_gpu_tensor *tmp = cur; + cur = next; + next = tmp; + } + if (ok && glm_debug_hidden_dump_layer_match(il)) { + ok = ds4_gpu_end_commands() != 0; + if (ok) { + for (uint32_t r = 0; r < n_tokens; r++) + glm_debug_dump_hidden_layer(cur, r, il, pos0 + r); + glm_debug_dump_raw_layer(g->batch_router_selected, "sel", + (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int32_t), + il, -1); + glm_debug_dump_raw_layer(g->batch_router_weights, "selw", + (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(float), + il, -1); + ok = ds4_gpu_begin_commands() != 0; + } + } + if (ok && + !g->ssd_streaming && + progress_flush_interval != 0 && + (il < g->layer_end || progress_requested) && + (slice_layer_done % progress_flush_interval) == 0) { + const uint32_t work_done = + work_done_base + (uint32_t)(((uint64_t)n_tokens * slice_layer_done) / g->layer_count); + const bool drain_now = + drain_interval != 0 && + il < g->layer_end && + (slice_layer_done % drain_interval) == 0; + const char *command_action = drain_now ? "drain" : "flush"; + const double trace_command_t0 = trace ? now_sec() : 0.0; + if (trace && (trace_all || trace_full_indexer || drain_now)) { + glm_graph_indexed_prefill_tracef( + "layer %s begin layer=%u pos=%u tokens=%u work=%u/%u", + command_action, + il, + pos0, + n_tokens, + work_done, + work_total); + } + if (drain_now) { + ok = ds4_gpu_end_commands() != 0; + if (ok) ok = ds4_gpu_begin_commands() != 0; + } else { + ok = ds4_gpu_flush_commands() != 0; + } + if (trace) { + const double trace_command_done = now_sec(); + const double command_ms = (trace_command_done - trace_command_t0) * 1000.0; + const double layer_ms = (trace_command_done - trace_layer_t0) * 1000.0; + trace_layer_flushed = true; + if (trace_all || trace_full_indexer || drain_now || + command_ms >= trace_slow_ms || layer_ms >= trace_slow_ms || !ok) { + glm_graph_indexed_prefill_tracef( + "layer %s %s layer=%u pos=%u tokens=%u command=%.3f ms layer_total=%.3f ms work=%u/%u", + command_action, + ok ? "done" : "failed", + il, + pos0, + n_tokens, + command_ms, + layer_ms, + work_done, + work_total); + } + } + if (ok) { + const bool progress_completed = + drain_interval == 0 || drain_now; + if (progress_completed) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + slice_layer_done, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + } + } + if (g->ssd_streaming) { + if (streaming_prefill_sync_each_layer) { + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + } else if (!ok) { + (void)ds4_gpu_synchronize(); + } + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + slice_layer_done, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + } + if (trace && ok) { + const double layer_ms = (now_sec() - trace_layer_t0) * 1000.0; + if (trace_all || (!trace_layer_flushed && layer_ms >= trace_slow_ms)) { + glm_graph_indexed_prefill_tracef( + "layer end layer=%u pos=%u tokens=%u flushed=%u layer_total=%.3f ms", + il, + pos0, + n_tokens, + trace_layer_flushed ? 1u : 0u, + layer_ms); + } + } + } + ds4_gpu_tp_set_attn_head_split(0); +#undef DS4_GLM_PROFILE_INDEXED_STAGE + if (ok && !g->ssd_streaming) { + const double trace_end_t0 = trace ? now_sec() : 0.0; + if (trace) { + glm_graph_indexed_prefill_tracef( + "chunk end_commands begin pos=%u tokens=%u", + pos0, + n_tokens); + } + ok = ds4_gpu_end_commands() != 0; + if (trace) { + const double end_ms = (now_sec() - trace_end_t0) * 1000.0; + const double chunk_ms = (now_sec() - trace_chunk_t0) * 1000.0; + glm_graph_indexed_prefill_tracef( + "chunk end_commands %s pos=%u tokens=%u end=%.3f ms chunk_total=%.3f ms", + ok ? "done" : "failed", + pos0, + n_tokens, + end_ms, + chunk_ms); + } + } else if (!ok) { + if (trace) { + glm_graph_indexed_prefill_tracef( + "chunk failed before end pos=%u tokens=%u elapsed=%.3f ms", + pos0, + n_tokens, + (now_sec() - trace_chunk_t0) * 1000.0); + } + (void)ds4_gpu_synchronize(); + } + if (ok && + g->ssd_streaming && + !streaming_prefill_sync_each_layer && + !output_hc && + !logits_out) { + ok = ds4_gpu_end_commands() != 0; + } + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + g->layer_count, + g->layer_count, + work_total, + logits_out == NULL && output_hc == NULL); + } + if (ok && output_hc) { + ok = ds4_gpu_tensor_read(cur, + 0, + output_hc, + (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)) != 0; + } + if (ok && logits_out) { + ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, + model, + weights); + } + if (ok && logits_out) { + last_hidden = glm_graph_tensor_row_view_strided(cur, + n_tokens - 1u, + DS4_N_EMBD, + DS4_N_EMBD); + ok = last_hidden != NULL; + if (ok && g->ssd_streaming) ok = glm_graph_stream_map_output(g, model, weights); + if (ok) ok = glm_graph_forward_output_head(g, model, weights, last_hidden, logits_out); + if (ok) { + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base, + n_tokens, + g->layer_count, + g->layer_count, + work_total, + true); + } + } + ds4_gpu_tensor_free(last_hidden); + ds4_gpu_set_glm_streaming_prefill_full_layer(false); + return ok; +} + +static bool glm_graph_use_streaming_token_prefill( + const ds4_glm_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens); +static uint32_t glm_graph_streaming_token_prefill_max_tokens(void); +static bool glm_graph_prefill_token_major( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + uint32_t pos0, + uint32_t n_tokens, + float *logits_out, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t display_absolute_base, + uint32_t work_done_base, + uint32_t work_total); + +static bool glm_graph_prefill_range( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + uint32_t pos0, + uint32_t n_tokens, + float *logits_out, + ds4_session_progress_fn progress, + void *progress_ud, + uint32_t progress_total) { + if (n_tokens == 0) return true; + if (!glm_graph_span_fits_context(g, pos0, n_tokens)) return false; + const uint32_t chunk_max = glm_graph_prefill_chunk_tokens(g->ctx_cap); + uint32_t done = 0; + while (done < n_tokens) { + const uint32_t pos = pos0 + done; + if (!g->full_kv_cache) { + const uint32_t remaining = n_tokens - done; + uint32_t chunk = 1; + if (glm_graph_use_streaming_token_prefill(g, pos, remaining)) { + chunk = remaining; + const uint32_t token_prefill_max = + glm_graph_streaming_token_prefill_max_tokens(); + if (token_prefill_max != 0 && chunk > token_prefill_max) { + chunk = token_prefill_max; + } + float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; + if (!glm_graph_prefill_token_major(g, + model, + weights, + tokens + done, + pos, + chunk, + dst_logits, + progress, + progress_ud, + pos0, + done, + n_tokens)) { + return false; + } + } else if (glm_graph_indexed_prefill_batch_ready(g, pos)) { + chunk = remaining; + if (chunk > g->indexed_prefill_cap) chunk = g->indexed_prefill_cap; + chunk = glm_graph_limit_indexed_prefill_chunk(pos, chunk); + if (chunk == 0) chunk = 1; + float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; + if (!glm_graph_forward_indexed_tokens(g, + model, + weights, + tokens + done, + NULL, + pos, + chunk, + NULL, + dst_logits, + progress, + progress_ud, + pos0, + done, + n_tokens)) { + return false; + } + } else { + float *dst_logits = (done + 1u == n_tokens) ? logits_out : NULL; + if (!glm_graph_forward_token(g, + model, + weights, + tokens[done], + NULL, + pos, + NULL, + dst_logits, + false)) { + return false; + } + } + done += chunk; + if (progress) { + const uint32_t current = pos0 + done; + progress(progress_ud, + "prefill_chunk", + current, + progress_total ? progress_total : pos0 + n_tokens); + } + continue; + } + if (pos >= g->ctx_cap) { + if (g->compact_cache_cap == 0) { + glm_graph_log_full_attention_limit(g, pos, n_tokens - done); + return false; + } + while (done < n_tokens) { + const uint32_t cur_pos = pos0 + done; + const bool use_indexed_batch = + glm_graph_indexed_prefill_batch_ready(g, cur_pos); + if (use_indexed_batch) { + uint32_t chunk = n_tokens - done; + if (chunk > g->indexed_prefill_cap) chunk = g->indexed_prefill_cap; + chunk = glm_graph_limit_indexed_prefill_chunk(cur_pos, chunk); + float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; + if (!glm_graph_forward_indexed_tokens(g, + model, + weights, + tokens + done, + NULL, + cur_pos, + chunk, + NULL, + dst_logits, + progress, + progress_ud, + pos0, + done, + n_tokens)) { + return false; + } + done += chunk; + } else { + float *dst_logits = (done + 1u == n_tokens) ? logits_out : NULL; + if (!glm_graph_forward_token(g, + model, + weights, + tokens[done], + NULL, + cur_pos, + NULL, + dst_logits, + false)) { + return false; + } + done++; + } + if (progress) { + const uint32_t current = pos0 + done; + progress(progress_ud, + "prefill_chunk", + current, + progress_total ? progress_total : pos0 + n_tokens); + } + } + return true; + } + uint32_t chunk = n_tokens - done; + const uint32_t full_remaining = g->ctx_cap - pos; + if (chunk > full_remaining) chunk = full_remaining; + if (chunk > chunk_max) chunk = chunk_max; + float *dst_logits = (done + chunk == n_tokens) ? logits_out : NULL; + if (glm_graph_use_streaming_token_prefill(g, pos, chunk)) { + if (!glm_graph_prefill_token_major(g, + model, + weights, + tokens + done, + pos, + chunk, + dst_logits, + progress, + progress_ud, + pos0, + done, + n_tokens)) { + return false; + } + } else if (!glm_graph_forward_tokens(g, + model, + weights, + tokens + done, + NULL, + pos, + chunk, + NULL, + dst_logits, + progress, + progress_ud, + pos0, + done, + n_tokens)) { + return false; + } + done += chunk; + if (progress) { + const uint32_t current = pos0 + done; + progress(progress_ud, + "prefill_chunk", + current, + progress_total ? progress_total : pos0 + n_tokens); + } + } + return true; +} + +/* + * For very short GLM SSD-streaming prefills, Metal still benefits from the + * token-major path because it reuses the normal decode graph and warms the + * decode expert cache. On ROCm/Strix Halo the indexed batch prefill is faster + * now that streamed batch routing and expert cache seeding are implemented, so + * ROCm defaults to canonical batch prefill unless the env override below opts + * token-major prefill back in. + */ +enum { DS4_GLM_STREAM_PREFILL_TOKEN_MAJOR_MAX_TOKENS = 64 }; + +static uint32_t glm_graph_streaming_token_prefill_default_max_tokens(void) { +#ifdef DS4_ROCM_BUILD + return 0; +#else + return DS4_GLM_STREAM_PREFILL_TOKEN_MAJOR_MAX_TOKENS; +#endif +} + +static uint32_t glm_graph_streaming_token_prefill_max_tokens(void) { + const char *env = glm_graph_env_value( + "DS4_ROCM_GLM_STREAMING_TOKEN_PREFILL_MAX", + "DS4_METAL_GLM_STREAMING_TOKEN_PREFILL_MAX"); + if (!env || !env[0]) env = getenv("DS4_GLM_STREAMING_TOKEN_PREFILL_MAX"); + const uint32_t default_max = + glm_graph_streaming_token_prefill_default_max_tokens(); + if (!env || !env[0]) return default_max; + char *end = NULL; + errno = 0; + unsigned long v = strtoul(env, &end, 10); + if (end == env || errno != 0 || v > UINT32_MAX) { + return default_max; + } + return (uint32_t)v; +} + +static bool glm_graph_use_streaming_token_prefill( + const ds4_glm_gpu_graph *g, + uint32_t pos0, + uint32_t n_tokens) { + if (!g || !g->ssd_streaming || g->quality || n_tokens == 0) return false; + if (getenv("DS4_GLM_DISABLE_STREAMING_TOKEN_PREFILL") != NULL || + glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_TOKEN_PREFILL", + "DS4_METAL_GLM_DISABLE_STREAMING_TOKEN_PREFILL")) { + return false; + } + if (!glm_graph_span_fits_full_attention(g, pos0, n_tokens)) return false; + const uint32_t max_tokens = glm_graph_streaming_token_prefill_max_tokens(); + return max_tokens != 0 && n_tokens <= max_tokens; +} + +static bool glm_graph_prefill_token_major( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + uint32_t pos0, + uint32_t n_tokens, + float *logits_out, + ds4_session_progress_fn display_progress, + void *display_progress_ud, + uint32_t display_absolute_base, + uint32_t work_done_base, + uint32_t work_total) { + if (!g || !model || !weights || !tokens || n_tokens == 0) return false; + for (uint32_t i = 0; i < n_tokens; i++) { + const bool last = i + 1u == n_tokens; + float *dst_logits = (last && logits_out) ? logits_out : NULL; + if (!glm_graph_forward_token(g, + model, + weights, + tokens[i], + NULL, + pos0 + i, + NULL, + dst_logits, + false)) { + return false; + } + glm_graph_report_prefill_display_progress(display_progress, + display_progress_ud, + display_absolute_base, + work_done_base + i + 1u, + n_tokens, + 0, + 1, + work_total, + false); + } + return true; +} + +static bool glm_graph_maybe_warm_compact_indexer_after_prefill( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint32_t next_pos) { + if (!g || !model || !weights) return false; + if (g->compact_cache_cap == 0 || + g->indexer_full_layers == 0 || + next_pos < g->ctx_cap) { + return true; + } + if (next_pos >= g->ctx_size) return true; + return glm_graph_warm_compact_indexer_store(g, model, weights, next_pos); +} + +static bool glm_graph_begin_commands_if_needed(void) { + return ds4_gpu_commands_active() || ds4_gpu_begin_commands() != 0; +} + +static bool glm_graph_end_commands_if_active(void) { + return !ds4_gpu_commands_active() || ds4_gpu_end_commands() != 0; +} + +static bool glm_graph_streaming_decode_sync_each_layer(void) { +#ifdef DS4_ROCM_BUILD + const char *env = glm_graph_env_value( + "DS4_ROCM_GLM_STREAMING_DECODE_SYNC_EACH_LAYER", + "DS4_METAL_GLM_STREAMING_DECODE_SYNC_EACH_LAYER"); + if (!env) env = getenv("DS4_GLM_STREAMING_DECODE_SYNC_EACH_LAYER"); + return glm_graph_env_truthy(env); +#else + return true; +#endif +} + +static bool glm_graph_forward_token( + ds4_glm_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int token, + const float *input_hc, + uint32_t pos, + float *output_hc, + float *logits_out, + bool defer_completion) { +#define DS4_GLM_FT_FAIL(why) do { \ + if (getenv("DS4_GLM_TP_DEBUG")) \ + fprintf(stderr, "ds4: glm forward_token fail pos=%u: %s\n", pos, why); \ + } while (0) + if (!g || !model || !weights || + token < 0 || token >= (int)DS4_N_VOCAB || + g->layer_count == 0 || + pos >= g->ctx_size || + (defer_completion && + (g->ssd_streaming || !ds4_gpu_commands_active()))) { + DS4_GLM_FT_FAIL("arg guard"); + return false; + } + if (!input_hc && !g->has_token_embd) { DS4_GLM_FT_FAIL("no token embd"); return false; } + if (logits_out && !g->has_output_head) { DS4_GLM_FT_FAIL("no output head"); return false; } + const bool use_indexed_attention = + glm_graph_decode_uses_indexed_attention(g, pos, logits_out); + uint32_t decode_layer_flush_interval = 0; + if (logits_out != NULL) { + decode_layer_flush_interval = use_indexed_attention ? 4u : 32u; + const char *dfi = getenv("DS4_GLM_DECODE_FLUSH_INTERVAL"); + if (dfi && dfi[0]) { + int v = atoi(dfi); + decode_layer_flush_interval = v <= 0 ? 0u : (uint32_t)v; + } + if (decode_layer_flush_interval > g->layer_count) { + decode_layer_flush_interval = g->layer_count; + } + if (defer_completion) decode_layer_flush_interval = 0; + } + if (pos >= g->ctx_cap && !use_indexed_attention) { + glm_graph_log_full_attention_limit(g, pos, 1); + DS4_GLM_FT_FAIL("full attention limit"); + return false; + } + if (g->compact_cache_cap != 0 && + !glm_graph_ensure_compact_cache(g, pos + 1u)) { + DS4_GLM_FT_FAIL("compact cache ensure"); + return false; + } + + const bool decode_output_profile = false; + const bool merge_indexed_output = + logits_out != NULL && use_indexed_attention && !decode_output_profile; + double decode_output_stage_t0 = decode_output_profile ? now_sec() : 0.0; + const bool decode_flush_profile = false; + uint32_t decode_flush_layer0 = 0; + double decode_flush_stage_t0 = decode_flush_profile ? now_sec() : 0.0; + + const bool static_decode_map = + !input_hc && + g->has_token_embd && + g->ssd_streaming && + metal_graph_stream_decode_static_map_enabled(); + const bool static_map_state_cache = + static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); + const bool streaming_decode_sync_each_layer = + g->ssd_streaming && + !static_decode_map && + glm_graph_streaming_decode_sync_each_layer(); + bool ok = true; + if (input_hc) { + ok = ds4_gpu_tensor_write(g->cur, + 0, + input_hc, + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } else if (static_decode_map) { + if (!static_map_state_cache || !g->streaming_static_decode_map_current) { + ok = metal_graph_stream_map_decode_static_all(model, weights); + if (ok) g->streaming_static_decode_map_current = static_map_state_cache; + } + } else { + ok = glm_graph_stream_map_token(g, model, weights); + } + if (ok) ok = glm_graph_begin_commands_if_needed(); + if (ok && !input_hc) { + if (g->placement) { + ok = glm_graph_ws_switch(g, g->placement[0], false); + } + } + if (ok && !input_hc) { + ok = ds4_gpu_embed_token_quant_tensor(g->cur, + model->map, + model->size, + weights->token_embd->abs_offset, + weights->token_embd->type, + DS4_N_VOCAB, + (uint32_t)token, + DS4_N_EMBD) != 0; + } + if (ok && streaming_decode_sync_each_layer) { + ok = ds4_gpu_end_commands() != 0; + } + const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); + ds4_gpu_tensor *last_indexer_selected = NULL; + uint32_t last_indexer_selected_count = 0; +#define DS4_GLM_PROFILE_DECODE_STAGE(part_, name_) do { \ + if (ok && decode_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary((part_), (name_), il, pos, 1, &decode_stage_t0); \ + } \ + } while (0) + uint32_t glm_ft_fail_il = UINT32_MAX; + for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { + if (g->placement) { + ok = glm_graph_ws_switch(g, g->placement[il + 1u], true); + if (!ok) break; + } + glm_ft_fail_il = il; + const uint32_t slice_layer_done = il - g->layer_start + 1u; + if (g->ssd_streaming) { + if (!static_decode_map) { + ok = glm_graph_stream_map_decode_layer(g, model, weights, il); + } + if (ok) ok = glm_graph_begin_commands_if_needed(); + } + const ds4_layer_weights *l = &weights->layer[il]; + const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; + const float rope_base = layer_rope_freq_base(il); + const float rope_scale = layer_rope_freq_scale(il); + const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); + double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; + if (decode_stage_profile) { + ok = metal_graph_layer_stage_profile_boundary("glm_decode_attn", + NULL, + il, + pos, + 1, + &decode_stage_t0); + } + + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->attn_norm, + g->cur, + model->map, + model->size, + l->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_norm"); + const uint32_t decode_ablate = glm_decode_ablate_mask(); + if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q_rank, + model, + l->attn_q_a->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + g->attn_norm, + il, + pos, + "attn_q_a", + g->ssd_streaming) != 0; + } + const bool fuse_qkv_norm_store = use_indexed_attention && + !decode_stage_profile && + g->compact_cache_cap != 0; + const bool fuse_qkv_norm = !decode_stage_profile && !fuse_qkv_norm_store; + if (ok && fuse_qkv_norm_store) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_store", + g->ssd_streaming) != 0; + if (ok) { + ok = ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + g->q_rank_norm, + g->q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + g->kv_raw, + l->attn_kv_a_norm->abs_offset, + pos, + 1, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16(), + DS4_RMS_EPS) != 0; + } + } else if (ok && fuse_qkv_norm) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_norm", + g->ssd_streaming) != 0; + if (ok) { + ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->q_rank_norm, + g->q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + g->kv_norm, + g->kv_raw, + l->attn_kv_a_norm->abs_offset, + DS4_N_KV_LORA, + 1, + DS4_RMS_EPS) != 0; + } + } else if (ok) { + ok = ds4_gpu_rms_norm_weight_tensor(g->q_rank_norm, + g->q_rank, + model->map, + model->size, + l->attn_q_a_norm->abs_offset, + DS4_N_LORA_Q, + DS4_RMS_EPS) != 0; + } + if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->q_rank_norm, + il, + pos, + "attn_q_b", + g->ssd_streaming) != 0; + if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->q, + 1, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_ROT, + pos, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "q_path"); + if (ok) metal_graph_debug_dump_tensor("glm_decode_q", + g->q, + (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA, + il, + pos); + if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il) && + !(decode_ablate & DS4_GLM_ABLATE_INDEXER)) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->indexer_k, + model, + l->indexer_attn_k->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD_DIM, + g->cur, + il, + pos, + "indexer_k", + g->ssd_streaming) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k_proj"); + if (ok) { + ok = ds4_gpu_glm_store_indexer_k_tensor( + g->layer_indexer_key_cache[il], + g->indexer_k, + model->map, + model->size, + l->indexer_k_norm->abs_offset, + l->indexer_k_norm_b->abs_offset, + pos, + 1, + g->compact_cache_cap, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + 0, + 1.0e-6f, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + glm_graph_compact_cache_is_f16()) != 0; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k_store"); + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k"); + if (ok && !fuse_qkv_norm && !fuse_qkv_norm_store) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a", + g->ssd_streaming) != 0; + if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, + g->kv_raw, + model->map, + model->size, + l->attn_kv_a_norm->abs_offset, + 1, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_RMS_EPS) != 0; + } + if (ok && g->compact_cache_cap != 0 && !fuse_qkv_norm_store) { + ok = ds4_gpu_glm_store_compact_kv_tensor(g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + g->kv_norm, + g->kv_raw, + pos, + 1, + g->compact_cache_cap, + kv_raw_dim, + DS4_N_KV_LORA, + DS4_N_ROT, + glm_graph_compact_cache_is_f16()) != 0; + } + if (use_indexed_attention) { + if (ok && glm_graph_layer_uses_full_indexer(il)) { + const uint32_t visible = pos + 1u; + if (ok && visible <= indexer_top_k) { + ok = ds4_gpu_glm_fill_selected_range_tensor(g->indexer_selected, + visible) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_fill"); + last_indexer_selected_count = visible; + } else if (ok && (decode_ablate & DS4_GLM_ABLATE_INDEXER)) { + /* Ablation: valid selected ids without the score/topk + * chain, so downstream attention timing stays real. */ + ok = ds4_gpu_glm_fill_selected_range_tensor(g->indexer_selected, + indexer_top_k) != 0; + last_indexer_selected_count = indexer_top_k; + } else if (ok) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->indexer_q, + model, + l->indexer_attn_q_b->abs_offset, + DS4_N_LORA_Q, + (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, + g->q_rank_norm, + il, + pos, + "indexer_q", + g->ssd_streaming) != 0; + if (ok) ok = ds4_gpu_glm_indexer_rope_tail_tensor(g->indexer_q, + 1, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok && glm_graph_indexer_qat()) { + ok = ds4_gpu_dsv4_indexer_qat_tensor(g->indexer_q, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_q"); + if (ok) ok = ds4_gpu_matmul_f32_tensor(g->indexer_weights, + model->map, + model->size, + l->indexer_proj->abs_offset, + DS4_N_EMBD, + DS4_N_INDEXER_HEAD, + g->cur, + 1) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_weights"); + const float indexer_scale = + 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + ok = ds4_gpu_glm_indexer_score_one_tensor(g->indexer_scores, + g->indexer_q, + g->indexer_weights, + g->layer_indexer_key_cache[il], + visible, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + indexer_scale, + glm_graph_compact_cache_is_f16()) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_scores"); + if (ok) ok = ds4_gpu_indexer_topk_tensor(g->indexer_selected, + g->indexer_scores, + visible, + 1, + indexer_top_k) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_topk"); + last_indexer_selected_count = indexer_top_k; + } + if (ok) last_indexer_selected = g->indexer_selected; + } else if (ok && (!last_indexer_selected || last_indexer_selected_count == 0)) { + ok = false; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_select"); + if (ok && !(decode_ablate & (DS4_GLM_ABLATE_ATTN_CORE | DS4_GLM_ABLATE_QKLOW))) { + ok = ds4_gpu_glm_qk_lowrank_typed_tensor(g->qk_low, + g->q, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_KEY_MLA) != 0; + if (ok) metal_graph_debug_dump_tensor("glm_decode_qk_low", + g->qk_low, + (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, + il, + pos); + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_path"); + if (ok && (decode_ablate & DS4_GLM_ABLATE_ATTN_CORE)) { + /* Skip the indexed attention kernels; zero heads so the + * rest of the layer stays finite (timing-only). */ + ok = ds4_gpu_tensor_fill_f32(g->heads, 0.0f, + (uint64_t)g->heads_dim) != 0; + } else if (ok && glm_graph_indexed_decode_split_group8_available(last_indexer_selected_count)) { + const uint32_t split_block_rows = + glm_graph_indexed_decode_split_block_rows_for(last_indexer_selected_count); + const uint32_t split_blocks = + (last_indexer_selected_count + split_block_rows - 1u) / split_block_rows; + ok = ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(g->heads, + g->attn_partial_lora, + g->attn_partial_ms, + g->q, + g->qk_low, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + last_indexer_selected, + last_indexer_selected_count, + true, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + split_block_rows, + split_blocks, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } else if (ok) { + ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor(g->heads, + g->q, + g->qk_low, + g->layer_kv_lora_cache[il], + g->layer_k_rope_cache[il], + model->map, + model->size, + l->attn_v_b->abs_offset, + l->attn_v_b->type, + last_indexer_selected, + last_indexer_selected_count, + g->compact_cache_cap, + glm_graph_compact_cache_is_f16(), + DS4_N_HEAD, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + } else { + if (ok) ok = ds4_gpu_glm_k_b_project_typed_tensor(g->k_nope, + g->kv_norm, + model->map, + model->size, + l->attn_k_b->abs_offset, + l->attn_k_b->type, + 1, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_HEAD) != 0; + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->value, + model, + l->attn_v_b->abs_offset, + DS4_N_KV_LORA, + g->heads_dim, + g->kv_norm, + il, + pos, + "attn_v_b", + g->ssd_streaming) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_path"); + if (ok) ok = ds4_gpu_glm_build_kv_cache_tensor(g->layer_key_cache[il], + g->layer_value_cache[il], + g->kv_raw, + g->k_nope, + g->value, + pos, + 1, + g->ctx_cap, + DS4_N_HEAD, + kv_raw_dim, + DS4_N_KV_LORA, + (uint32_t)g->q_nope, + DS4_N_ROT, + DS4_N_VALUE_MLA, + 0, + rope_base, + rope_scale, + 0.0f, + 1.0f, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + true) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_cache"); + if (ok) ok = ds4_gpu_glm_attention_full_tensor(g->heads, + g->q, + g->layer_key_cache[il], + g->layer_value_cache[il], + pos, + 1, + pos + 1u, + g->ctx_cap, + DS4_N_HEAD, + DS4_N_KEY_MLA, + DS4_N_VALUE_MLA, + true) != 0; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attention"); + if (ok) metal_graph_debug_dump_tensor("glm_decode_heads", + g->heads, + g->heads_dim, + il, + pos); + if (ok && !(decode_ablate & DS4_GLM_ABLATE_ATTN_OUT)) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->heads, + il, + pos, + "attn_o", + g->ssd_streaming) != 0; + } + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_output"); + if (ok) ok = ds4_gpu_add_rms_norm_weight_tensor(g->ffn_norm, + g->after_attn, + g->cur, + g->attn_out, + model->map, + model->size, + l->ffn_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_ffn", "ffn_norm"); + if (ok) ok = glm_graph_encode_ffn_one_normed_from(g, + model, + l, + il, + pos, + g->ffn_norm, + g->after_attn, + g->next, + g->ffn_gate, + g->ffn_up, + g->ffn_mid, + g->ffn_out, + g->ffn_sum, + g->attn_out, + decode_stage_profile, + decode_stage_profile ? &decode_stage_t0 : NULL); + if (ok) { + ds4_gpu_tensor *tmp = g->cur; + g->cur = g->next; + g->next = tmp; + } + if (ok && glm_debug_hidden_dump_layer_match(il)) { + ok = ds4_gpu_end_commands() != 0; + if (ok) { + glm_debug_dump_hidden_layer(g->cur, 0, il, pos); + glm_debug_dump_raw_layer(g->router_selected, "sel", + (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t), + il, (int)pos); + glm_debug_dump_raw_layer(g->router_weights, "selw", + (uint64_t)DS4_N_EXPERT_USED * sizeof(float), + il, (int)pos); + ok = ds4_gpu_begin_commands() != 0; + } + } + if (ok && + !g->ssd_streaming && + decode_layer_flush_interval != 0 && + il < g->layer_end && + (slice_layer_done % decode_layer_flush_interval) == 0) { + if (decode_flush_profile) { + ok = ds4_gpu_flush_commands() != 0; + if (ok) ok = ds4_gpu_synchronize() != 0; + if (ok) { + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM decode layer flush pos=%u layers=%u..%u %.3f ms\n", + pos, + decode_flush_layer0, + il, + (now - decode_flush_stage_t0) * 1000.0); + decode_flush_layer0 = il + 1u; + decode_flush_stage_t0 = now; + ok = ds4_gpu_begin_commands() != 0; + } + } else { + ok = ds4_gpu_flush_commands() != 0; + } + } + if (streaming_decode_sync_each_layer) { + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + } + } +#undef DS4_GLM_PROFILE_DECODE_STAGE + if (ok && (merge_indexed_output || + (defer_completion && logits_out != NULL))) { + if (g->ssd_streaming) { + if (!static_decode_map) { + ok = glm_graph_stream_map_output(g, model, weights); + } + if (ok) ok = glm_graph_begin_commands_if_needed(); + } + ok = glm_graph_encode_output_head(g, model, weights); + if (g->ssd_streaming) { + if (ok) ok = glm_graph_end_commands_if_active(); + else (void)ds4_gpu_synchronize(); + } + } + if (!g->ssd_streaming && !defer_completion) { + if (ok) ok = ds4_gpu_end_commands() != 0; + else (void)ds4_gpu_synchronize(); + } else if (!ok) { + (void)ds4_gpu_synchronize(); + } + if (decode_output_profile) { + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM decode output profile pos=%u layers=%.3f ms\n", + pos, + (now - decode_output_stage_t0) * 1000.0); + decode_output_stage_t0 = now; + } + if (ok && output_hc && !defer_completion) { + ok = ds4_gpu_tensor_read(g->cur, + 0, + output_hc, + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } + if (ok && logits_out && !defer_completion) { + if (use_indexed_attention) { + if (!merge_indexed_output) { + if (g->ssd_streaming && !static_decode_map) { + ok = glm_graph_stream_map_output(g, model, weights); + } + if (ok) ok = glm_graph_begin_commands_if_needed(); + if (ok) ok = glm_graph_encode_output_head(g, model, weights); + if (ok) ok = glm_graph_end_commands_if_active(); + else (void)ds4_gpu_synchronize(); + if (decode_output_profile) { + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM decode output profile pos=%u output_head=%.3f ms\n", + pos, + (now - decode_output_stage_t0) * 1000.0); + decode_output_stage_t0 = now; + } + } + if (ok) { + if (glm_debug_hidden_dump_layer() < 0) + glm_debug_dump_hidden_row(g->cur, 0); + ok = ds4_gpu_tensor_read(g->logits, + 0, + logits_out, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + if (decode_output_profile) { + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM decode output profile pos=%u logits_read=%.3f ms\n", + pos, + (now - decode_output_stage_t0) * 1000.0); + } + } + } else { + if (g->ssd_streaming && !static_decode_map) { + ok = glm_graph_stream_map_output(g, model, weights); + } + if (ok && ds4_gpu_commands_active()) { + ok = ds4_gpu_end_commands() != 0; + } + if (ok) ok = glm_graph_forward_output_head(g, model, weights, g->cur, logits_out); + if (decode_output_profile) { + const double now = now_sec(); + fprintf(stderr, + "ds4: GLM decode output profile pos=%u fallback_output=%.3f ms\n", + pos, + (now - decode_output_stage_t0) * 1000.0); + } + } + } + if (ok && + !logits_out && + !output_hc && + g->ssd_streaming && + !streaming_decode_sync_each_layer) { + ok = ds4_gpu_end_commands() != 0; + } else if (ok && !logits_out && g->ssd_streaming) { + ok = glm_graph_end_commands_if_active(); + } else if (!ok) { + (void)ds4_gpu_synchronize(); + } + if (!ok && getenv("DS4_GLM_TP_DEBUG")) { + fprintf(stderr, + "ds4: glm forward_token fail pos=%u around layer %u\n", + pos, glm_ft_fail_il); + } + (void)glm_ft_fail_il; + return ok; +#undef DS4_GLM_FT_FAIL +} + +static int glm_metal_first_token_logits( + const ds4_model *model, + const ds4_weights *weights, + int token, + float *logits_out) { + if (!model || !weights || !logits_out) return 1; + if (token < 0 || token >= (int)DS4_N_VOCAB) { + fprintf(stderr, "ds4: GLM token %d is outside vocab\n", token); + return 1; + } + if (!weights->token_embd || weights->token_embd->type != DS4_TENSOR_Q8_0 || + !weights->output_norm || weights->output_norm->type != DS4_TENSOR_F32 || + !weights->output || weights->output->type != DS4_TENSOR_Q8_0 || + weights->output_norm->dim[0] != DS4_N_EMBD || + weights->output->dim[0] != DS4_N_EMBD || + weights->output->dim[1] != DS4_N_VOCAB) { + fprintf(stderr, "ds4: GLM Metal first-token path found unexpected embedding/output layout\n"); + return 1; + } + if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) { + fprintf(stderr, "ds4: GLM Metal first-token path has no normal transformer layers\n"); + return 1; + } + + const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; + uint64_t kv_raw_dim = 0; + uint64_t dense_hidden_max = DS4_N_FF_EXP; + bool generic_routed_moe = false; + for (uint32_t il = 0; il < normal_layers; il++) { + const ds4_layer_weights *l = &weights->layer[il]; + if (l->attn_kv_a_mqa && l->attn_kv_a_mqa->dim[1] > kv_raw_dim) { + kv_raw_dim = l->attn_kv_a_mqa->dim[1]; + } + if (il < DS4_N_LEADING_DENSE && l->ffn_gate && + l->ffn_gate->dim[1] > dense_hidden_max) { + dense_hidden_max = l->ffn_gate->dim[1]; + } + if (glm_graph_layer_uses_generic_routed_moe(l)) generic_routed_moe = true; + } + if (kv_raw_dim < DS4_N_KV_LORA) { + fprintf(stderr, "ds4: GLM Metal first-token path found no valid KV projection\n"); + return 1; + } + + const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + const uint64_t sparse_mid_elems = (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP; + const uint64_t ffn_mid_elems = + dense_hidden_max > sparse_mid_elems ? dense_hidden_max : sparse_mid_elems; + const uint64_t routed_mid_bytes = + (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float); + const uint64_t routed_down_bytes = + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); + const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); + + ds4_gpu_tensor *cur = NULL; + ds4_gpu_tensor *attn_norm = NULL; + ds4_gpu_tensor *kv_raw = NULL; + ds4_gpu_tensor *kv_norm = NULL; + ds4_gpu_tensor *heads = NULL; + ds4_gpu_tensor *attn_out = NULL; + ds4_gpu_tensor *after_attn = NULL; + ds4_gpu_tensor *ffn_norm = NULL; + ds4_gpu_tensor *ffn_gate = NULL; + ds4_gpu_tensor *ffn_up = NULL; + ds4_gpu_tensor *ffn_mid = NULL; + ds4_gpu_tensor *routed_gate = NULL; + ds4_gpu_tensor *routed_up = NULL; + ds4_gpu_tensor *routed_down = NULL; + ds4_gpu_tensor *ffn_out = NULL; + ds4_gpu_tensor *ffn_sum = NULL; + ds4_gpu_tensor *next = NULL; + ds4_gpu_tensor *router_logits = NULL; + ds4_gpu_tensor *router_probs = NULL; + ds4_gpu_tensor *router_selected = NULL; + ds4_gpu_tensor *router_weights = NULL; + ds4_gpu_tensor *logits = NULL; + + int ok = 1; +#define DS4_GLM_FIRST_ALLOC_TENSOR(var, bytes_) \ + do { \ + (var) = ds4_gpu_tensor_alloc((bytes_)); \ + if (!(var)) { \ + fprintf(stderr, "ds4: GLM Metal first-token path could not allocate %s\n", #var); \ + ok = 0; \ + } \ + } while (0) + + DS4_GLM_FIRST_ALLOC_TENSOR(cur, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(attn_norm, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(kv_raw, kv_raw_dim * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(kv_norm, (uint64_t)DS4_N_KV_LORA * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(heads, heads_dim * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(attn_out, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(after_attn, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_norm, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_gate, dense_hidden_max * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_up, dense_hidden_max * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_mid, ffn_mid_elems * sizeof(float)); + if (generic_routed_moe) { + DS4_GLM_FIRST_ALLOC_TENSOR(routed_gate, routed_mid_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(routed_up, routed_mid_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(routed_down, routed_down_bytes); + } + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_out, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(ffn_sum, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(next, emb_bytes); + DS4_GLM_FIRST_ALLOC_TENSOR(router_logits, (uint64_t)DS4_N_EXPERT * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(router_probs, (uint64_t)DS4_N_EXPERT * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + DS4_GLM_FIRST_ALLOC_TENSOR(router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); + DS4_GLM_FIRST_ALLOC_TENSOR(logits, logits_bytes); +#undef DS4_GLM_FIRST_ALLOC_TENSOR + + if (ok) { + ok = ds4_gpu_embed_token_q8_0_tensor(cur, + model->map, + model->size, + weights->token_embd->abs_offset, + DS4_N_VOCAB, + (uint32_t)token, + DS4_N_EMBD); + } + for (uint32_t il = 0; ok && il < normal_layers; il++) { + const ds4_layer_weights *gl = &weights->layer[il]; + const uint64_t gl_kv_raw_dim = gl->attn_kv_a_mqa ? gl->attn_kv_a_mqa->dim[1] : 0; + if (!gl->attn_norm || + !gl->attn_kv_a_mqa || + !gl->attn_kv_a_norm || + !gl->attn_v_b || + !gl->attn_output || + !gl->ffn_norm || + gl_kv_raw_dim < DS4_N_KV_LORA || + gl_kv_raw_dim > kv_raw_dim || + gl->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || + gl->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || + gl->attn_v_b->type != DS4_TENSOR_Q8_0 || + gl->attn_v_b->dim[0] != DS4_N_KV_LORA || + gl->attn_v_b->dim[1] != DS4_N_VALUE_MLA || + gl->attn_v_b->dim[2] != DS4_N_HEAD || + gl->attn_output->type != DS4_TENSOR_Q8_0 || + gl->attn_output->dim[0] != heads_dim || + gl->attn_output->dim[1] != DS4_N_EMBD) { + fprintf(stderr, + "ds4: GLM Metal first-token path found unexpected attention layout in layer %u\n", + il); + ok = 0; + break; + } + + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(attn_norm, cur, + model->map, model->size, + gl->attn_norm->abs_offset, + DS4_N_EMBD, DS4_RMS_EPS); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(kv_raw, + model->map, + model->size, + gl->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + gl_kv_raw_dim, + attn_norm, + 1); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(kv_norm, kv_raw, + model->map, model->size, + gl->attn_kv_a_norm->abs_offset, + DS4_N_KV_LORA, DS4_RMS_EPS); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(heads, + model->map, + model->size, + gl->attn_v_b->abs_offset, + DS4_N_KV_LORA, + heads_dim, + kv_norm, + 1); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(attn_out, + model->map, + model->size, + gl->attn_output->abs_offset, + heads_dim, + DS4_N_EMBD, + heads, + 1); + if (ok) ok = ds4_gpu_add_tensor(after_attn, cur, attn_out, DS4_N_EMBD); + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, after_attn, + model->map, model->size, + gl->ffn_norm->abs_offset, + DS4_N_EMBD, DS4_RMS_EPS); + if (il < DS4_N_LEADING_DENSE) { + const uint64_t gl_ffn_hidden = gl->ffn_gate ? gl->ffn_gate->dim[1] : 0; + if (!gl->ffn_gate || + !gl->ffn_up || + !gl->ffn_down || + gl->ffn_gate->type != DS4_TENSOR_Q8_0 || + gl->ffn_up->type != DS4_TENSOR_Q8_0 || + gl->ffn_down->type != DS4_TENSOR_Q8_0 || + gl->ffn_gate->dim[0] != DS4_N_EMBD || + gl->ffn_up->dim[0] != DS4_N_EMBD || + gl->ffn_up->dim[1] != gl_ffn_hidden || + gl->ffn_down->dim[0] != gl_ffn_hidden || + gl->ffn_down->dim[1] != DS4_N_EMBD || + gl_ffn_hidden > dense_hidden_max) { + fprintf(stderr, + "ds4: GLM Metal first-token path found unexpected dense FFN layout in layer %u\n", + il); + ok = 0; + break; + } + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_gate, + model->map, + model->size, + gl->ffn_gate->abs_offset, + DS4_N_EMBD, + gl_ffn_hidden, + ffn_norm, + 1); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_up, + model->map, + model->size, + gl->ffn_up->abs_offset, + DS4_N_EMBD, + gl_ffn_hidden, + ffn_norm, + 1); + if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, ffn_gate, ffn_up, + (uint32_t)gl_ffn_hidden, 0.0f, 1.0f); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_out, + model->map, + model->size, + gl->ffn_down->abs_offset, + gl_ffn_hidden, + DS4_N_EMBD, + ffn_mid, + 1); + if (ok) ok = ds4_gpu_add_tensor(next, after_attn, ffn_out, DS4_N_EMBD); + } else { + const uint32_t gl_gate_type = gl->ffn_gate_exps ? gl->ffn_gate_exps->type : 0; + const uint32_t gl_up_type = gl->ffn_up_exps ? gl->ffn_up_exps->type : 0; + const bool gl_gate_pair_supported = + glm_graph_gate_pair_type_supported(gl_gate_type, gl_up_type); + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + + if (!gl->ffn_gate_inp || + !gl->ffn_exp_probs_b || + !gl->ffn_gate_exps || + !gl->ffn_up_exps || + !gl->ffn_down_exps || + !gl->ffn_gate_shexp || + !gl->ffn_up_shexp || + !gl->ffn_down_shexp || + gl->ffn_gate_inp->type != DS4_TENSOR_F32 || + gl->ffn_gate_inp->dim[0] != DS4_N_EMBD || + gl->ffn_gate_inp->dim[1] != DS4_N_EXPERT || + gl->ffn_exp_probs_b->type != DS4_TENSOR_F32 || + gl->ffn_exp_probs_b->dim[0] != DS4_N_EXPERT || + !gl_gate_pair_supported || + !glm_graph_down_type_supported(gl->ffn_down_exps->type) || + gl->ffn_gate_shexp->type != DS4_TENSOR_Q8_0 || + gl->ffn_up_shexp->type != DS4_TENSOR_Q8_0 || + gl->ffn_down_shexp->type != DS4_TENSOR_Q8_0 || + gl->ffn_gate_shexp->dim[0] != DS4_N_EMBD || + gl->ffn_gate_shexp->dim[1] != DS4_N_FF_EXP || + gl->ffn_up_shexp->dim[0] != DS4_N_EMBD || + gl->ffn_up_shexp->dim[1] != DS4_N_FF_EXP || + gl->ffn_down_shexp->dim[0] != DS4_N_FF_EXP || + gl->ffn_down_shexp->dim[1] != DS4_N_EMBD || + sparse_mid_elems > ffn_mid_elems) { + fprintf(stderr, + "ds4: GLM Metal first-token path found unexpected sparse FFN layout in layer %u\n", + il); + ok = 0; + break; + } + + (void)tensor_expert_bytes(model, gl->ffn_gate_exps, 0, + &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, gl->ffn_up_exps, 0, + &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, gl->ffn_down_exps, 0, + &down_in, &down_out, &down_row_bytes); + if (gate_in != DS4_N_EMBD || + up_in != DS4_N_EMBD || + down_in != DS4_N_FF_EXP || + gate_out != DS4_N_FF_EXP || + up_out != DS4_N_FF_EXP || + down_out != DS4_N_EMBD) { + fprintf(stderr, + "ds4: GLM Metal first-token path found unexpected expert strides in layer %u\n", + il); + ok = 0; + break; + } + + if (ok) ok = ds4_gpu_matmul_f32_tensor(router_logits, + model->map, + model->size, + gl->ffn_gate_inp->abs_offset, + DS4_N_EMBD, + DS4_N_EXPERT, + ffn_norm, + 1); + if (ok) ok = ds4_gpu_glm_router_select_tensor(router_selected, + router_weights, + router_probs, + model->map, + model->size, + gl->ffn_exp_probs_b->abs_offset, + router_logits, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE); + if (ok) { + const ds4_gpu_stream_expert_table table = { + .model_map = model->map, + .model_size = model->size, + .layer = il, + .n_total_expert = DS4_N_EXPERT, + .gate_offset = gl->ffn_gate_exps->abs_offset, + .up_offset = gl->ffn_up_exps->abs_offset, + .down_offset = gl->ffn_down_exps->abs_offset, + .gate_expert_bytes = gate_out * gate_row_bytes, + .down_expert_bytes = down_out * down_row_bytes, + }; + ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + &table, + router_selected, + DS4_N_EXPERT_USED) != 0; + } + ds4_glm_gpu_graph route_g = { + .routed_gate = routed_gate, + .routed_up = routed_up, + .routed_down = routed_down, + .ssd_streaming = false, + }; + if (ok) ok = glm_graph_routed_moe_one_dispatch( + &route_g, + model, + gl, + il, + ffn_out, + ffn_mid, + gate_out * gate_row_bytes, + gate_row_bytes, + up_out * up_row_bytes, + up_row_bytes, + down_out * down_row_bytes, + down_row_bytes, + router_selected, + router_weights, + ffn_norm, + false); + if (ok) ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ffn_gate, + ffn_up, + ffn_mid, + model->map, + model->size, + gl->ffn_gate_shexp->abs_offset, + gl->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + DS4_N_FF_EXP, + ffn_norm, + 0.0f); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_sum, + model->map, + model->size, + gl->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, + DS4_N_EMBD, + ffn_mid, + 1); + if (ok) ok = ds4_gpu_add_tensor(attn_out, ffn_out, ffn_sum, DS4_N_EMBD); + if (ok) ok = ds4_gpu_add_tensor(next, after_attn, attn_out, DS4_N_EMBD); + } + + if (ok) { + ds4_gpu_tensor *tmp = cur; + cur = next; + next = tmp; + } + } + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(ffn_norm, cur, + model->map, model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, DS4_RMS_EPS); + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(logits, + model->map, + model->size, + weights->output->abs_offset, + DS4_N_EMBD, + DS4_N_VOCAB, + ffn_norm, + 1); + if (ok) ok = ds4_gpu_tensor_read(logits, 0, logits_out, logits_bytes) != 0; + + ds4_gpu_tensor_free(router_weights); + ds4_gpu_tensor_free(router_selected); + ds4_gpu_tensor_free(router_probs); + ds4_gpu_tensor_free(router_logits); + ds4_gpu_tensor_free(logits); + ds4_gpu_tensor_free(next); + ds4_gpu_tensor_free(ffn_sum); + ds4_gpu_tensor_free(ffn_out); + ds4_gpu_tensor_free(routed_down); + ds4_gpu_tensor_free(routed_up); + ds4_gpu_tensor_free(routed_gate); + ds4_gpu_tensor_free(ffn_mid); + ds4_gpu_tensor_free(ffn_up); + ds4_gpu_tensor_free(ffn_gate); + ds4_gpu_tensor_free(ffn_norm); + ds4_gpu_tensor_free(after_attn); + ds4_gpu_tensor_free(attn_out); + ds4_gpu_tensor_free(heads); + ds4_gpu_tensor_free(kv_norm); + ds4_gpu_tensor_free(kv_raw); + ds4_gpu_tensor_free(attn_norm); + ds4_gpu_tensor_free(cur); + return ok ? 0 : 1; +} + +static DS4_MAYBE_UNUSED int generate_glm_metal_first_token( + const ds4_model * model, + const ds4_vocab * vocab, + const ds4_weights * weights, + const token_vec * prompt, + int n_predict, + int ctx_size, + ds4_token_emit_fn emit, + ds4_generation_done_fn done, + void * emit_ud) { + fprintf(stderr, "ds4: using GLM Metal first-token generation path\n"); + + if (prompt->len != 1 || prompt->len > ctx_size) { + fprintf(stderr, + "ds4: GLM Metal generation currently supports exactly one prompt token; " + "multi-token prefill needs the GLM KV/DSA graph\n"); + return 1; + } + if (n_predict <= 0) { + if (done) done(emit_ud); + return 0; + } + if (n_predict > 1) { + fprintf(stderr, + "ds4: GLM Metal generation currently emits only the first generated token; " + "stopping after one token\n"); + } + + float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); + const double t0 = now_sec(); + const int rc = glm_metal_first_token_logits(model, weights, prompt->v[0], logits); + const double t1 = now_sec(); + if (rc != 0) { + free(logits); + return 1; + } + + if (getenv("DS4_TRACE_TOP") != NULL) { + print_top_logits(stderr, "GLM first-token", vocab, logits, DS4_N_VOCAB, 10); + } + const int token = sample_argmax(logits, DS4_N_VOCAB); + if (!vocab_token_is_generation_stop(vocab, token) && emit) emit(emit_ud, token); + if (done) done(emit_ud); + + const double eval_s = t1 - t0; + ds4_log(stderr, + DS4_LOG_TIMING, + "ds4: GLM first-token eval: %.2f t/s\n", + eval_s > 0.0 ? 1.0 / eval_s : 0.0); + + free(logits); + return 0; +} + +static int generate_glm_metal_argmax( + const ds4_model * model, + const ds4_vocab * vocab, + const ds4_weights * weights, + const token_vec * prompt, + int n_predict, + int ctx_size, + bool quality, + bool ssd_streaming, + bool ssd_streaming_cold, + uint32_t ssd_streaming_preload_experts, + uint64_t ssd_streaming_cache_bytes, + uint64_t ssd_streaming_prefill_headroom_bytes, + ds4_token_emit_fn emit, + ds4_generation_done_fn done, + void * emit_ud, + ds4_session_progress_fn progress, + void * progress_ud) { + fprintf(stderr, "ds4: using GLM full-attention argmax generation path\n"); + + if (!prompt || prompt->len <= 0 || prompt->len > ctx_size) { + fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); + return 1; + } + if (n_predict <= 0) { + if (done) done(emit_ud); + return 0; + } + + ds4_glm_gpu_graph g = {0}; + if (!glm_graph_alloc(&g, + model, + weights, + ctx_size, + ssd_streaming, + ssd_streaming_cold)) { + fprintf(stderr, "ds4: failed to allocate GLM graph runtime\n"); + return 1; + } + g.quality = quality; + if ((uint32_t)prompt->len >= g.ctx_size) { + fprintf(stderr, + "ds4: prompt length %d leaves no GLM context room (ctx %u)\n", + prompt->len, + g.ctx_size); + glm_graph_free(&g); + return 1; + } + const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; + if (memory_report) ds4_gpu_print_memory_report("after GLM graph alloc"); + + float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); + bool ok = true; + const bool seed_before_prefill = + ssd_streaming && + !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL", + "DS4_METAL_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL"); + const double t_prefill0 = now_sec(); + if (seed_before_prefill) { + ds4_gpu_graph seed_graph; + memset(&seed_graph, 0, sizeof(seed_graph)); + seed_graph.quality = quality; + seed_graph.ssd_streaming = ssd_streaming; + seed_graph.ssd_streaming_cold = ssd_streaming_cold; + seed_graph.streaming_preload_experts = ssd_streaming_preload_experts; + ok = metal_graph_seed_streaming_expert_cache_from_hotlist(&seed_graph, + model, + weights); + } + if (ok) { + ok = glm_graph_prefill_range(&g, + model, + weights, + prompt->v, + 0, + (uint32_t)prompt->len, + logits, + progress, + progress_ud, + (uint32_t)prompt->len); + } + const double t_prefill1 = now_sec(); + if (memory_report) ds4_gpu_print_memory_report("after GLM prefill"); + if (!ok) { + fprintf(stderr, "ds4: GLM prefill failed\n"); + free(logits); + glm_graph_free(&g); + return 1; + } +#ifdef DS4_ROCM_BUILD + /* + * Decode is SSD-read bound, so the prefill expert headroom is worth more + * as extra dynamic cache once prefill is done. Opt out with + * DS4_ROCM_GLM_STREAMING_GROW_CACHE_AFTER_PREFILL=0. + */ + const char *grow_cache_env = + getenv("DS4_ROCM_GLM_STREAMING_GROW_CACHE_AFTER_PREFILL"); + if (ssd_streaming && + ssd_streaming_cache_bytes != 0 && + ssd_streaming_prefill_headroom_bytes != 0 && + (grow_cache_env == NULL || glm_graph_env_truthy(grow_cache_env))) { + uint64_t budget_bytes = 0; + uint64_t per_expert_bytes = 0; + if (ssd_streaming_cache_bytes <= + UINT64_MAX - ssd_streaming_prefill_headroom_bytes) { + budget_bytes = + ssd_streaming_cache_bytes + ssd_streaming_prefill_headroom_bytes; + } + const uint32_t grown_budget = + ds4_streaming_cache_experts_for_byte_budget(weights, + budget_bytes, + &per_expert_bytes); + const uint32_t current_budget = + ds4_gpu_stream_expert_cache_configured_count(); + if (grown_budget > current_budget) { + ds4_gpu_set_streaming_expert_cache_budget(grown_budget); + fprintf(stderr, + "ds4: ROCm GLM streaming expert cache grew after prefill: " + "%u -> %u experts (%.2f GiB)\n", + current_budget, + grown_budget, + (double)((uint64_t)grown_budget * per_expert_bytes) / + 1073741824.0); + } + } +#else + (void)ssd_streaming_cache_bytes; + (void)ssd_streaming_prefill_headroom_bytes; +#endif + + int n_generated = 0; + int n_decode_eval = 0; + uint32_t pos = (uint32_t)prompt->len; + const bool token_timing = getenv("DS4_TOKEN_TIMING") != NULL; + const double t_decode0 = now_sec(); + for (int i = 0; i < n_predict && pos < g.ctx_size; i++) { + if (getenv("DS4_TRACE_TOP") != NULL) { + char label[64]; + snprintf(label, sizeof(label), "GLM step %d", i); + print_top_logits(stderr, label, vocab, logits, DS4_N_VOCAB, 10); + } + const int token = sample_argmax(logits, DS4_N_VOCAB); + if (vocab_token_is_generation_stop(vocab, token)) break; + if (emit) emit(emit_ud, token); + n_generated++; + + if (i == n_predict - 1 || pos + 1u >= g.ctx_size) { + pos++; + break; + } + + const double t_eval0 = token_timing ? now_sec() : 0.0; + ok = glm_graph_forward_token(&g, model, weights, token, NULL, pos, + NULL, logits, false); + if (!ok) { + fprintf(stderr, "ds4: GLM decode failed at position %u\n", pos); + free(logits); + glm_graph_free(&g); + return 1; + } + if (token_timing) { + const double t_eval1 = now_sec(); + fprintf(stderr, + "ds4: GLM decode eval %d took %.3f ms\n", + n_decode_eval + 1, + (t_eval1 - t_eval0) * 1000.0); + } + n_decode_eval++; + pos++; + } + const double t_decode1 = now_sec(); + if (done) done(emit_ud); + + const double prefill_s = t_prefill1 - t_prefill0; + const double decode_s = t_decode1 - t_decode0; + ds4_log(stderr, + DS4_LOG_TIMING, + "ds4: GLM prefill: %.2f t/s, generation: %.2f t/s\n", + prefill_s > 0.0 ? (double)prompt->len / prefill_s : 0.0, + decode_s > 0.0 ? (double)n_generated / decode_s : 0.0); + + if (memory_report) ds4_gpu_print_memory_report("before GLM graph free"); + free(logits); + glm_graph_free(&g); + return 0; +} diff --git a/models/glm/metal/host/kernels.inc b/models/glm/metal/host/kernels.inc new file mode 100644 index 0000000000..c95c1a7b17 --- /dev/null +++ b/models/glm/metal/host/kernels.inc @@ -0,0 +1,5340 @@ +int ds4_gpu_glm_kv_lora_rms_norm_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_raw, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !kv_raw || !model_map || + n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || + kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 || + !isfinite(eps) || eps < 0.0f) { + return 0; + } + + @autoreleasepool { + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + if (!rawbuf || !outbuf || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal GLM KV RMS norm received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM KV RMS norm weight range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + n_tokens == 1u && + weight_bytes <= (1ull << 20) && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t weight_inner = 0; + id weightbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_kv_lora_rms_norm_pipeline, + "kernel_glm_kv_lora_rms_norm"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_kv_lora_rms_norm_args args = { + .n_tokens = n_tokens, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .eps = eps, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(kv_lora_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV RMS norm")) return 0; + } + + return 1; +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !kv_norm || !model_map || + n_tokens == 0 || kv_lora_dim == 0 || + qk_nope == 0 || n_head == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id kvbuf = ds4_gpu_tensor_buffer(kv_norm); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM k_b projection received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + if (!outbuf || !kvbuf || + ds4_gpu_tensor_bytes(kv_norm) < kv_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal GLM k_b projection received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM k_b projection range is outside the mapped model\n"); + return 0; + } + const NSUInteger q_blocks = ((NSUInteger)qk_nope + 31u) / 32u; + if (q_blocks > 8u) { + fprintf(stderr, "ds4: Metal GLM k_b projection q width is too large for the tiled kernel\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_k_b_project_pipeline, + "kernel_glm_k_b_project_q8_0"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_k_b_project_args args = { + .n_tokens = n_tokens, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .n_head = n_head, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .pad1 = 0, + .pad2 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, q_blocks, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM k_b projection")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_k_b_project_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 n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head) { + return ds4_gpu_glm_k_b_project_typed_tensor(out, + kv_norm, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + kv_lora_dim, + qk_nope, + n_head); +} + +int ds4_gpu_glm_store_compact_kv_tensor( + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_norm, + const ds4_gpu_tensor *kv_raw, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || + n_tokens == 0 || cache_cap == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id kvnormbuf = ds4_gpu_tensor_buffer(kv_norm); + id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t kv_norm_bytes = + (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t kv_raw_bytes = + (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + if (!kvcachebuf || !ropecachebuf || !kvnormbuf || !kvrawbuf || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(kv_norm) < kv_norm_bytes || + ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { + fprintf(stderr, "ds4: Metal GLM compact KV store received undersized buffers\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_store_compact_kv_pipeline, + "kernel_glm_store_compact_kv"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_store_compact_kv_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_rope = qk_rope, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad1 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvnormbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:1]; + [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 2, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM compact KV store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qkv_norm_store_compact_kv_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_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_raw, + uint64_t kv_weight_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || + !model_map || n_tokens == 0 || cache_cap == 0 || + q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id qoutbuf = ds4_gpu_tensor_buffer(q_out); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t kv_raw_bytes = + (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + if (!qbuf || !qoutbuf || !kvcachebuf || !ropecachebuf || !kvrawbuf || + ds4_gpu_tensor_bytes(q) < q_row_bytes * n_tokens || + ds4_gpu_tensor_bytes(q_out) < q_row_bytes * n_tokens || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { + fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store received undersized buffers\n"); + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) { + fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store weight range is outside the mapped model\n"); + return 0; + } + + uint64_t q_weight_inner = 0; + id q_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, + q_weight_offset, q_row_bytes, + &q_weight_inner); + if (!q_weightbuf) return 0; + uint64_t kv_weight_inner = 0; + id kv_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, + kv_weight_offset, kv_weight_bytes, + &kv_weight_inner); + if (!kv_weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_qkv_norm_store_compact_kv_pipeline, + "kernel_glm_qkv_norm_store_compact_kv"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_qkv_norm_store_compact_kv_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .q_n = q_n, + .q_n4 = q_n / 4u, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .kv_lora_n4 = kv_lora_dim / 4u, + .qk_rope = qk_rope, + .cache_f16 = cache_f16 ? 1u : 0u, + .eps = eps, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:q_weightbuf offset:(NSUInteger)q_weight_inner atIndex:2]; + [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; + [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:4]; + [enc setBuffer:kv_weightbuf offset:(NSUInteger)kv_weight_inner atIndex:5]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:6]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:7]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 3, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM fused q/kv norm compact store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_store_indexer_k_tensor( + ds4_gpu_tensor *indexer_key_cache, + const ds4_gpu_tensor *raw_k, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t bias_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!indexer_key_cache || !raw_k || !model_map || + n_tokens == 0 || cache_cap == 0 || + head_dim == 0 || rot_dim == 0 || + rot_dim > head_dim || (rot_dim & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + id rawbuf = ds4_gpu_tensor_buffer(raw_k); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t cache_bytes = + (uint64_t)cache_cap * head_dim * cache_elem_bytes; + const uint64_t raw_bytes = + (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + if (!cachebuf || !rawbuf || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes || + ds4_gpu_tensor_bytes(raw_k) < raw_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer K store received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || norm_bytes > model_size - weight_offset || + bias_offset > model_size || norm_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM indexer K norm range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + uint64_t bias_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, norm_bytes, + &weight_inner); + id biasbuf = + ds4_gpu_wrap_model_range(model_map, model_size, + bias_offset, norm_bytes, + &bias_inner); + if (!weightbuf || !biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_store_indexer_k_pipeline, + "kernel_glm_store_indexer_k"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_store_indexer_k_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .head_dim = head_dim, + .rot_dim = rot_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0, + .eps = eps, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .pad1 = 0.0f, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:3]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:4]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer K store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_build_kv_cache_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + 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, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || + n_tokens == 0 || cache_cap == 0 || n_head == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + kv_lora_dim + qk_rope > kv_raw_dim || + qk_dim < qk_nope || (qk_rope & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id knbuf = ds4_gpu_tensor_buffer(k_nope); + id vbuf = ds4_gpu_tensor_buffer(value); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(k_nope) < kn_bytes || + ds4_gpu_tensor_bytes(value) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM KV cache builder received undersized buffers\n"); + return 0; + } + + const bool decode_group4 = + n_tokens == 1u && + n_head >= 4u && + ds4_gpu_env_bool("DS4_METAL_DISABLE_GLM_DECODE_KV_GROUP4") <= 0; + id pipeline = + decode_group4 ? + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_decode_group4_pipeline, + "kernel_glm_build_kv_cache_decode_group4") : + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_pipeline, + "kernel_glm_build_kv_cache"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_build_kv_cache_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .n_head = n_head, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0u, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; + [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; + const NSUInteger group_y = decode_group4 ? + ((NSUInteger)n_head + 3u) / 4u : + (NSUInteger)n_head; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, group_y, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV cache build")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_build_kv_cache_flash_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + 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, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || + pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + kv_lora_dim + qk_rope > kv_raw_dim || + qk_dim < qk_nope || (qk_rope & 1u) != 0 || + n_tokens > cache_cap || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id knbuf = ds4_gpu_tensor_buffer(k_nope); + id vbuf = ds4_gpu_tensor_buffer(value); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(k_nope) < kn_bytes || + ds4_gpu_tensor_bytes(value) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM staged KV cache builder received undersized buffers\n"); + return 0; + } + + const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); + const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); + const NSUInteger key_f16_offset = 0; + const NSUInteger key_f16_bytes = + (NSUInteger)n_tokens * (NSUInteger)n_head * q_row_bytes_f16; + const NSUInteger value_f16_offset = key_f16_bytes; + const NSUInteger value_f16_bytes = + (NSUInteger)n_tokens * (NSUInteger)n_head * v_row_bytes_f16; + const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_glm_flash_attn_kv_f16")) { + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_flash_pipeline, + "kernel_glm_build_kv_cache_flash"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_build_kv_cache_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .n_head = n_head, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0u, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; + [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:6]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM staged KV cache build")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_attention_flash_tensor_impl( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16, + int kv_pre_staged) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim != 256u || value_dim != 256u || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0 || + cache_len > ds4_gpu_glm_flash_attention_max_cache_len()) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + if (!headsbuf || !qbuf || !keybuf || !valbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM FlashAttention received undersized buffers\n"); + return 0; + } + const uint64_t key_elems = (uint64_t)cache_len * n_head * qk_dim; + const uint64_t value_elems = (uint64_t)cache_len * n_head * value_dim; + if (key_elems > UINT32_MAX || value_elems > UINT32_MAX) { + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = 4; + const bool has_kvpad = (cache_len % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger q_row_bytes = (NSUInteger)qk_dim * sizeof(float); + const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); + const NSUInteger v_row_bytes = (NSUInteger)value_dim * sizeof(float); + const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)cache_len * sizeof(uint16_t); + const NSUInteger key_f16_offset = 0; + const NSUInteger key_f16_bytes = + (NSUInteger)cache_len * (NSUInteger)n_head * q_row_bytes_f16; + const NSUInteger value_f16_offset = key_f16_bytes; + const NSUInteger value_f16_bytes = + (NSUInteger)cache_len * (NSUInteger)n_head * v_row_bytes_f16; + const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * ((NSUInteger)n_head * (q_row_bytes_f16 + v_row_bytes_f16) + + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)cache_len + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_glm_prefill_mask_buffer(pos0, n_tokens, cache_len, mask_bytes); + if (!mask_buffer) return 0; + if (kv_pre_staged) { + if (!g_flash_attn_kv_buffer || g_flash_attn_kv_bytes < kv_f16_bytes) { + fprintf(stderr, "ds4: GLM staged FlashAttention KV scratch is missing\n"); + return 0; + } + } else if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_glm_flash_attn_kv_f16")) { + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_glm_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_glm_flash_attn_blk")) { + return 0; + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk256_dv256", + true, false, false, false, has_kvpad, bc_mask, + (int32_t)qk_dim, + (int32_t)value_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!kv_pre_staged) { + const bool copied = cache_f16 ? + (ds4_gpu_encode_cpy_f16_f16_3d(cb, + keybuf, + ds4_gpu_tensor_offset(key_cache), + g_flash_attn_kv_buffer, + key_f16_offset, + qk_dim, + cache_len, + n_head, + (uint64_t)n_head * q_row_bytes_f16, + q_row_bytes_f16, + q_row_bytes_f16, + (uint64_t)cache_len * q_row_bytes_f16) && + ds4_gpu_encode_cpy_f16_f16_3d(cb, + valbuf, + ds4_gpu_tensor_offset(value_cache), + g_flash_attn_kv_buffer, + value_f16_offset, + value_dim, + cache_len, + n_head, + (uint64_t)n_head * v_row_bytes_f16, + v_row_bytes_f16, + v_row_bytes_f16, + (uint64_t)cache_len * v_row_bytes_f16)) : + (ds4_gpu_encode_cpy_f32_f16_3d(cb, + keybuf, + ds4_gpu_tensor_offset(key_cache), + g_flash_attn_kv_buffer, + key_f16_offset, + qk_dim, + cache_len, + n_head, + (uint64_t)n_head * q_row_bytes, + q_row_bytes, + q_row_bytes_f16, + (uint64_t)cache_len * q_row_bytes_f16) && + ds4_gpu_encode_cpy_f32_f16_3d(cb, + valbuf, + ds4_gpu_tensor_offset(value_cache), + g_flash_attn_kv_buffer, + value_f16_offset, + value_dim, + cache_len, + n_head, + (uint64_t)n_head * v_row_bytes, + v_row_bytes, + v_row_bytes_f16, + (uint64_t)cache_len * v_row_bytes_f16)); + if (!copied) { + return 0; + } + } + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)cache_len, + .ne_12_2 = (int32_t)n_head, + .ne_12_3 = 1, + .nb11 = q_row_bytes_f16, + .nb12 = (uint64_t)cache_len * q_row_bytes_f16, + .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, + .nb21 = v_row_bytes_f16, + .nb22 = (uint64_t)cache_len * v_row_bytes_f16, + .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)cache_len, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * q_row_bytes, + .nb02 = q_row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * q_row_bytes, + .ne11 = (int32_t)cache_len, + .ne_12_2 = (int32_t)n_head, + .ne_12_3 = 1, + .ns10 = (int32_t)qk_dim, + .nb11 = q_row_bytes_f16, + .nb12 = (uint64_t)cache_len * q_row_bytes_f16, + .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, + .ns20 = (int32_t)value_dim, + .nb21 = v_row_bytes_f16, + .nb22 = (uint64_t)cache_len * v_row_bytes_f16, + .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)qk_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(value_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)qk_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM FlashAttention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_flash_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + return ds4_gpu_glm_attention_flash_tensor_impl(heads, + q, + key_cache, + value_cache, + pos0, + n_tokens, + cache_len, + cache_cap, + n_head, + qk_dim, + value_dim, + cache_f16, + 0); +} + +int ds4_gpu_glm_attention_flash_staged_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + if (pos0 != 0 || n_tokens != cache_len) return 0; + return ds4_gpu_glm_attention_flash_tensor_impl(heads, + q, + key_cache, + value_cache, + pos0, + n_tokens, + cache_len, + cache_cap, + n_head, + qk_dim, + value_dim, + cache_f16, + 1); +} + +int ds4_gpu_glm_attention_full_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim == 0 || value_dim == 0 || + (qk_dim & 3u) != 0 || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0 || + cache_len > ds4_gpu_glm_full_attention_max_cache_len()) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + if (!headsbuf || !qbuf || !keybuf || !valbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM attention received undersized buffers\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_full_pipeline, + "kernel_glm_attention_full"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const uint32_t full_attn_mode = 2u; + ds4_gpu_glm_attention_full_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_len = cache_len, + .cache_cap = cache_cap, + .n_head = n_head, + .qk_dim = qk_dim, + .value_dim = value_dim, + .pad0 = full_attn_mode, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad1 = 0u, + .pad2 = 0u, + .scale = 1.0f / sqrtf((float)qk_dim), + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:2]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:3]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; + [enc setThreadgroupMemoryLength:(256u + (NSUInteger)cache_len) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM full attention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_fill_selected_range_tensor( + ds4_gpu_tensor *selected, + uint32_t n_selected) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || n_selected == 0) return 0; + + @autoreleasepool { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM selected range received undersized buffer\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_fill_selected_range_pipeline, + "kernel_glm_fill_selected_range"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_fill_selected_range_args args = { + .n_selected = n_selected, + }; + const NSUInteger nth = 256u; + const NSUInteger n_groups = ((NSUInteger)n_selected + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_fill_selected_range_batch_tensor( + ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t pad_row) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || n_tokens == 0 || n_selected == 0) return 0; + + @autoreleasepool { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t total = (uint64_t)n_tokens * n_selected; + if (n_tokens != 0 && total / n_tokens != n_selected) return 0; + if (total > UINT64_MAX / sizeof(uint32_t)) return 0; + const uint64_t selected_bytes = total * sizeof(uint32_t); + if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM selected range batch received undersized buffer\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_fill_selected_range_batch_pipeline, + "kernel_glm_fill_selected_range_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_fill_selected_range_batch_args args = { + .n_tokens = n_tokens, + .pos0 = pos0, + .n_selected = n_selected, + .pad_row = pad_row, + }; + const NSUInteger nth = 256u; + const NSUInteger n_groups = ((NSUInteger)total + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range batch")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_rope_tail_offset_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t rot_offset, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || + rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset || + (rot_dim & 1u) != 0 || + pos0 > UINT32_MAX - n_tokens || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal %s received undersized buffer\n", + label ? label : "GLM RoPE"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_rope_tail_pipeline, + "kernel_glm_indexer_rope_tail_f32"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_rope_tail_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .rot_dim = rot_dim, + .rot_offset = rot_offset, + .pos0 = pos0, + .n_ctx_orig = n_ctx_orig, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(rot_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, label ? label : "GLM RoPE")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (rot_dim > head_dim) return 0; + return ds4_gpu_glm_rope_tail_offset_tensor(x, + n_tokens, + n_head, + head_dim, + rot_dim, + head_dim - rot_dim, + pos0, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + "GLM RoPE"); +} + +int ds4_gpu_glm_indexer_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_rope_tail_offset_tensor(x, + n_tokens, + n_head, + head_dim, + rot_dim, + 0, + pos0, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + "GLM indexer RoPE"); +} + +int ds4_gpu_glm_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_head == 0 || head_dim == 0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + + @autoreleasepool { + id scoresbuf = ds4_gpu_tensor_buffer(scores); + id qbuf = ds4_gpu_tensor_buffer(q); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t score_bytes = (uint64_t)n_rows * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t weights_bytes = (uint64_t)n_head * sizeof(float); + const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; + if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer score received undersized buffers\n"); + return 0; + } + + ds4_gpu_glm_indexer_score_one_args args = { + .n_rows = n_rows, + .n_head = n_head, + .head_dim = head_dim, + .cache_f16 = cache_f16 ? 1u : 0u, + .scale = scale, + }; + + if (n_head == 32u && head_dim == 128u) { + id direct_pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_score_one_direct_pipeline, + "kernel_glm_indexer_score_one_direct"); + if (!direct_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:direct_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer direct score")) return 0; + return 1; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_score_one_pipeline, + "kernel_glm_indexer_score_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer score")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128 || + pos0 >= n_rows || n_tokens > n_rows - pos0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + + @autoreleasepool { + id scoresbuf = ds4_gpu_tensor_buffer(scores); + id qbuf = ds4_gpu_tensor_buffer(q); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t score_bytes = (uint64_t)n_rows * n_tokens * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_head * sizeof(float); + const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; + if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer batch scores received undersized buffers\n"); + return 0; + } + + const bool force_scalar = g_quality_mode; + const bool use_tiled_f32 = false; + const bool use_tiled = !force_scalar && n_tokens >= 8u && + n_head == 32u && head_dim == 128u; + id pipeline = + use_tiled + ? ds4_gpu_hot_pipeline(use_tiled_f32 ? g_glm_indexer_scores_tiled_f32_pipeline + : g_glm_indexer_scores_tiled_pipeline, + use_tiled_f32 ? "kernel_glm_indexer_scores_tiled_f32" + : "kernel_glm_indexer_scores_tiled") + : ds4_gpu_hot_pipeline(g_glm_indexer_scores_batch_pipeline, + "kernel_glm_indexer_scores_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_indexer_scores_batch_args args = { + .n_rows = n_rows, + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = pos0, + .cache_f16 = cache_f16 ? 1u : 0u, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .score_token_stride = (uint64_t)n_rows * sizeof(float), + .scale = scale, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + if (use_tiled) { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + if (use_tiled_f32) { + [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * + sizeof(float) atIndex:0]; + } else { + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else { + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer batch scores")) return 0; + } + + return 1; +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + if (!qk_low || !q || !model_map || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_nope > qk_dim) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(qk_low); + id qbuf = ds4_gpu_tensor_buffer(q); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM qk lowrank received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t out_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + if (!outbuf || !qbuf || + ds4_gpu_tensor_bytes(qk_low) < out_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes) { + fprintf(stderr, "ds4: Metal GLM qk lowrank received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM qk lowrank range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_glm52 = + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + row_bytes == 204u && + weight_type == DS4_METAL_TENSOR_Q8_0; + /* Coalesced simdgroup variant: lanes split the 192-dot so weight + * reads coalesce, and 2048 threadgroups replace 64. The thread- + * per-row kernels measured ~7.5x off the weight-bandwidth floor + * (7.2ms of the decode token by skip-ablation). Covers the GLM 5.2 + * shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8 + * fast path above never engaged there). */ + const int use_glm52_sg = + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + ((weight_type == DS4_METAL_TENSOR_Q8_0 && row_bytes == 204u) || + (weight_type == DS4_METAL_TENSOR_Q4_0 && row_bytes == 108u)) && + getenv("DS4_METAL_DISABLE_GLM_QKLOW_SG") == NULL && + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52_sg") != nil; + if (getenv("DS4_METAL_GLM_QKLOW_DEBUG")) { + static int printed = 0; + if (!printed) { + printed = 1; + fprintf(stderr, "ds4: qk_lowrank decode path: use_glm52=%d sg=%d n_head=%u kv=%u nope=%u dim=%u rb=%llu type=%u\n", + use_glm52, use_glm52_sg, n_head, kv_lora_dim, qk_nope, qk_dim, + (unsigned long long)row_bytes, weight_type); + } + } + id pipeline = + use_glm52_sg ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52_sg") : + use_glm52 ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52") : + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_pipeline, + "kernel_glm_qk_lowrank_q8_0"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_qk_lowrank_args args = { + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_dim = qk_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .pad1 = 0, + .pad2 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; + if (use_glm52_sg) { + /* 8 simdgroups x 2 rows per threadgroup: (64, 512/16) grid. */ + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, + (NSUInteger)(kv_lora_dim / 16u), + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + } else { + if (use_glm52) { + [enc setThreadgroupMemoryLength:192u * sizeof(float) atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM qk lowrank")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qk_lowrank_q8_0_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 n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low, + q, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_head, + kv_lora_dim, + qk_nope, + qk_dim); +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + if (!qk_low || !q || !model_map || + n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_nope > qk_dim) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(qk_low); + id qbuf = ds4_gpu_tensor_buffer(q); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t out_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + if (!outbuf || !qbuf || + ds4_gpu_tensor_bytes(qk_low) < out_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_glm52_t4 = + n_tokens >= 4u && + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + row_bytes == 204u && + weight_type == DS4_METAL_TENSOR_Q8_0; + id pipeline = + use_glm52_t4 ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") : + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint32_t head_base = 0; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); + ds4_gpu_glm_qk_lowrank_batch_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_dim = qk_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .head_base = head_base, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; + if (use_glm52_t4) { + [enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + ((NSUInteger)n_tokens + 3u) / 4u, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } else { + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch qk lowrank")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qk_lowrank_q8_0_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 n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, + q, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + n_head, + kv_lora_dim, + qk_nope, + qk_dim); +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !lora || !model_map || + n_tokens == 0 || n_head == 0 || + kv_lora_dim == 0 || value_dim == 0) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id lorabuf = ds4_gpu_tensor_buffer(lora); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, kv_lora_dim, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM batch value project received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * value_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t heads_bytes = + (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + if (!headsbuf || !lorabuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(lora) < lora_bytes) { + fprintf(stderr, "ds4: Metal GLM batch value project received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM batch value project range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_mma = + n_head == 64u && + kv_lora_dim == 512u && + value_dim == 256u && + row_bytes == 544u && + weight_type == DS4_METAL_TENSOR_Q8_0 && + n_tokens >= 32u; + id pipeline = + use_mma ? + ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_mma_pipeline, + "kernel_glm_value_project_q8_0_batch_heads_mma") : + ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_pipeline, + "kernel_glm_value_project_q8_0_batch_heads"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint32_t head_base = 0; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); + ds4_gpu_glm_qk_lowrank_batch_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = 0, + .qk_dim = value_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .head_base = head_base, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:2]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3]; + if (use_mma) { + [enc setThreadgroupMemoryLength:16u * 1024u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tokens + 31u) / 32u, + ((NSUInteger)value_dim + 63u) / 64u, + (NSUInteger)head_count) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } else { + [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch value project")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_value_project_q8_0_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 n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim) { + return ds4_gpu_glm_value_project_typed_batch_heads_tensor(heads, + lora, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + n_head, + kv_lora_dim, + value_dim); +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + value_dim == 0 || qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM indexed attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM indexed attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_pipeline, + "kernel_glm_attention_indexed_decode"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + ds4_gpu_glm_attention_indexed_decode_args args = { + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .value_type = value_weight_type, + }; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + const NSUInteger scratch_floats = + 256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim; + [enc setThreadgroupMemoryLength:scratch_floats * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed attention decode")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_decode_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, + 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) { + return ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + 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); +} + +int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + 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, + bool selected_rows_valid, + 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, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + const uint32_t needed_blocks = + block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + if (!heads || !partial_lora || !partial_ms || !q || !qk_low || + !kv_lora_cache || !k_rope_cache || !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || (n_head % 8u) != 0 || + kv_lora_dim != 512u || + qk_nope == 0 || qk_rope != 64u || + value_dim == 0 || qk_dim < qk_nope || + block_rows == 0u || needed_blocks == 0u || + n_blocks < needed_blocks || n_blocks > 64u || + !cache_f16 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id partial_lorabuf = ds4_gpu_tensor_buffer(partial_lora); + id partial_msbuf = ds4_gpu_tensor_buffer(partial_ms); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = sizeof(uint16_t); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); + const uint64_t partial_lora_bytes = + (uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float); + const uint64_t partial_ms_bytes = + (uint64_t)n_blocks * n_head * 2u * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!headsbuf || !partial_lorabuf || !partial_msbuf || + !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(partial_lora) < partial_lora_bytes || + ds4_gpu_tensor_bytes(partial_ms) < partial_ms_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + const bool use_valid_fullheads = + selected_rows_valid && (n_head % 8u) == 0u; + id partial_pipeline = + ds4_gpu_hot_pipeline(use_valid_fullheads ? + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline : + g_glm_attention_indexed_decode_split_group8_partial_pipeline, + use_valid_fullheads ? + "kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads" : + "kernel_glm_attention_indexed_decode_split_group8_partial"); + const bool use_reduce16 = + n_blocks == 16u && block_rows == 128u && + n_selected == 2048u && value_dim == 256u; + id reduce_pipeline = + ds4_gpu_hot_pipeline(use_reduce16 ? + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline : + g_glm_attention_indexed_decode_split_group8_reduce_pipeline, + use_reduce16 ? + "kernel_glm_attention_indexed_decode_split_group8_reduce16" : + "kernel_glm_attention_indexed_decode_split_group8_reduce"); + if (!partial_pipeline || !reduce_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_decode_split_args args = { + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = 1u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .block_rows = block_rows, + .n_blocks = n_blocks, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .value_type = value_weight_type, + }; + const NSUInteger stage_rows = 16u; + const NSUInteger kv_vecs = (NSUInteger)kv_lora_dim / 4u; + const NSUInteger rope_vecs = (NSUInteger)qk_rope / 4u; + const NSUInteger partial_scratch_bytes = + stage_rows * kv_vecs * sizeof(uint16_t) * 4u + + stage_rows * rope_vecs * sizeof(float) * 4u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:partial_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:6]; + [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:7]; + [enc setThreadgroupMemoryLength:partial_scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head / 8u, (NSUInteger)n_blocks, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + const NSUInteger reduce_scratch_floats = 256u + 64u + (NSUInteger)kv_lora_dim; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:1]; + [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:2]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:3]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; + [enc setThreadgroupMemoryLength:reduce_scratch_floats * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM split grouped indexed attention decode")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + 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, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + bool selected_rows_valid, + 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, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(heads, + partial_lora, + partial_ms, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + selected, + n_selected, + selected_rows_valid, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + value_dim, + n_ctx_orig, + block_rows, + n_blocks, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); +} + +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 (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + value_dim == 0 || qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = + (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = + (uint64_t)n_tokens * n_selected * sizeof(uint32_t); + if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + const uint64_t grouped_attn = (n_tokens >= 128u) ? 8u : 2u; + const NSUInteger q2_bit_words = ((NSUInteger)cache_cap + 31u) / 32u; + const NSUInteger q2_scratch_bytes = + q2_bit_words * sizeof(uint32_t) + + 4u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float); + const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; + const bool q2_fits = max_tg_mem == 0 || q2_scratch_bytes <= max_tg_mem; + const bool use_q2_group4 = grouped_attn == 4u && n_tokens >= 2u && q2_fits; + const bool use_group8 = + grouped_attn == 8u || + (grouped_attn == 4u && !use_q2_group4); + const bool use_group2 = grouped_attn == 1u || grouped_attn == 2u; + id pipeline = use_q2_group4 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_q2_group4_pipeline, + "kernel_glm_attention_indexed_batch_q2_group4") : + (use_group8 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, + "kernel_glm_attention_indexed_batch_group8") : + (use_group2 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group2_pipeline, + "kernel_glm_attention_indexed_batch_group2") : + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_pipeline, + "kernel_glm_attention_indexed_batch"))); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .value_type = value_weight_type, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + const NSUInteger scratch_bytes = use_q2_group4 ? + q2_scratch_bytes : + (use_group8 ? + (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float)) : + ((use_group2 ? + (512u + 2u * (NSUInteger)n_selected + 2u * (NSUInteger)kv_lora_dim) : + (256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim)) * sizeof(float))); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + if (use_q2_group4) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 3u) / 4u, + ((NSUInteger)n_tokens + 1u) / 2u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else if (use_group8) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + } else if (use_group2) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 1u) / 2u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } else { + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_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, + 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) { + return ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + 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); +} + +int ds4_gpu_sort_i32_rows_asc_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint32_t row_width, + uint32_t n_rows) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !src || row_width == 0 || n_rows == 0 || + (row_width & (row_width - 1u)) != 0) { + return 0; + } + + @autoreleasepool { + id pipeline = + ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, + "kernel_dsv4_sort_i32_rows_asc"); + if (!pipeline) return 0; + + const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t); + id srcbuf = ds4_gpu_tensor_buffer(src); + id dstbuf = ds4_gpu_tensor_buffer(dst); + if (!srcbuf || !dstbuf || + ds4_gpu_tensor_bytes(src) < bytes || + ds4_gpu_tensor_bytes(dst) < bytes) { + fprintf(stderr, "ds4: Metal row sort received undersized buffers\n"); + return 0; + } + + NSUInteger threads = (NSUInteger)row_width; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (max_threads != 0 && threads > max_threads) threads = max_threads; + if (threads == 0) return 0; + + const NSUInteger scratch_bytes = (NSUInteger)row_width * sizeof(int32_t); + const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; + if (max_tg_mem != 0 && scratch_bytes > max_tg_mem) { + fprintf(stderr, "ds4: Metal row sort scratch exceeds threadgroup memory limit\n"); + return 0; + } + + ds4_gpu_dsv4_topk_mask_args args = { + .ne00 = (int64_t)row_width, + .ne01 = (int64_t)n_rows, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)row_width * sizeof(int32_t), + .ne0 = (int64_t)row_width, + .ne1 = (int64_t)n_rows, + .nb0 = sizeof(int32_t), + .nb1 = (uint64_t)row_width * sizeof(int32_t), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1]; + [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:2]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "sort i32 rows asc")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( + ds4_gpu_tensor *lora_out, + 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 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool selected_rows_valid) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id lorabuf = ds4_gpu_tensor_buffer(lora_out); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = + (uint64_t)n_tokens * n_selected * sizeof(uint32_t); + if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(lora_out) < lora_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention-lora received undersized buffers\n"); + return 0; + } + + const bool use_vec_lora = + cache_f16 && kv_lora_dim == 512u && qk_rope == 64u; + const bool full_head_groups = (n_head % 8u) == 0u; + id pipeline = nil; + if (use_vec_lora && selected_rows_valid && full_head_groups) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); + } else if (use_vec_lora && selected_rows_valid) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); + } else if (use_vec_lora) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec"); + } else { + pipeline = ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, + "kernel_glm_attention_indexed_batch_group8"); + } + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = kv_lora_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = 0, + .value_type = 1u, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); + const NSUInteger scratch_bytes = use_vec_lora ? + (16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + + 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) : + (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float)); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + if (use_vec_lora) { + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:6]; + } else { + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7]; + } + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention-lora")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + ds4_gpu_tensor *lora_out, + 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, + uint32_t n_tokens, + uint32_t pos0, + 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + pos0 > n_selected || n_tokens > n_selected - pos0 || + n_head == 0 || kv_lora_dim != 512u || + qk_nope == 0 || qk_rope != 64u || + qk_dim < qk_nope || !cache_f16 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id lorabuf = ds4_gpu_tensor_buffer(lora_out); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + const uint64_t cache_elem_bytes = sizeof(uint16_t); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || + ds4_gpu_tensor_bytes(lora_out) < lora_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes) { + fprintf(stderr, "ds4: Metal GLM causal batch attention-lora received undersized buffers\n"); + return 0; + } + + const bool full_head_groups = (n_head % 8u) == 0u; + id pipeline = full_head_groups ? + ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads") : + ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = 1u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = kv_lora_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = 0, + .value_type = 1u, + .pos0 = pos0, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); + const NSUInteger scratch_bytes = + 16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + + 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM causal indexed batch attention-lora")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_lora_tensor( + ds4_gpu_tensor *lora_out, + 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 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + false); +} + +int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + ds4_gpu_tensor *lora_out, + 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 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 n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + true); +} + +int ds4_gpu_glm_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_expert == 0 || n_expert > 256u || + n_expert_used == 0 || n_expert_used > n_expert) { + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal GLM router received undersized buffers\n"); + return 0; + } + + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM router bias range is outside the mapped model\n"); + return 0; + } + const bool exact_bias_view = + bias_bytes <= (1ull << 20) && + getenv("DS4_METAL_DISABLE_DECODE_ROUTER_BIAS_EXACT_VIEWS") == NULL; + uint64_t bias_inner = 0; + id biasbuf = exact_bias_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + bias_offset, + bias_bytes, + &bias_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + bias_offset, + bias_bytes, + &bias_inner); + if (!biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, + "kernel_glm_router_select_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_router_select_one_args args = { + .n_expert = n_expert, + .n_expert_used = n_expert_used, + .expert_weight_scale = expert_weight_scale, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; + [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_tokens == 0 || + n_expert == 0 || n_expert > 256u || + n_expert_used == 0 || n_expert_used > n_expert) { + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + const uint64_t logits_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(int32_t); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(float); + const uint64_t probs_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < logits_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(probs) < probs_bytes) { + fprintf(stderr, "ds4: Metal GLM batch router received undersized buffers\n"); + return 0; + } + + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM batch router bias range is outside the mapped model\n"); + return 0; + } + uint64_t bias_inner = 0; + id biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, + bias_offset, bias_bytes, + &bias_inner); + if (!biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, + "kernel_glm_router_select_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_router_select_one_args args = { + .n_expert = n_expert, + .n_expert_used = n_expert_used, + .expert_weight_scale = expert_weight_scale, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; + [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0; + } + + return 1; +} + +static bool ds4_gpu_glm_gate_pair_type_supported( + uint32_t gate_type, + uint32_t up_type) { + return gate_type == up_type && + (gate_type == DS4_METAL_TENSOR_Q2_K || + gate_type == DS4_METAL_TENSOR_Q4_K || + gate_type == DS4_METAL_TENSOR_Q5_K); +} + +static bool ds4_gpu_glm_down_type_supported(uint32_t down_type) { + return down_type == DS4_METAL_TENSOR_Q2_K || + down_type == DS4_METAL_TENSOR_Q4_K || + down_type == DS4_METAL_TENSOR_Q5_K || + down_type == DS4_METAL_TENSOR_Q6_K; +} + +int ds4_gpu_glm_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + bool force_resident) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding: only the owned contiguous expert range is mapped, + * so bind from the owned base, validate only its bytes, and tell the + * kernels the first expert id present at that base. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + + if (!out || !mid || !model_map || !selected || !weights || !x || + n_total_expert == 0 || n_expert == 0 || n_expert > 256u || + n_expert > n_total_expert || + expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || + gate_expert_bytes == 0 || gate_row_bytes == 0 || + up_expert_bytes == 0 || up_row_bytes == 0 || + down_expert_bytes == 0 || down_row_bytes == 0 || + (expert_in_dim % 256u) != 0 || + (expert_mid_dim % 256u) != 0 || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type)) { + return 0; + } + + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || + (uint64_t)out_dim > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE tensor byte size overflow\n"); + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || + up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || + down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE received inconsistent expert strides\n"); + return 0; + } + if (layer_index == 3u) { + ds4_gpu_stream_expert_cache_note_decode_token(); + } + if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset || + up_offset > model_size || up_tensor_bytes > model_size - up_offset || + down_offset > model_size || down_tensor_bytes > model_size - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed MoE tensor range is outside the mapped model\n"); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int32_t) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal GLM routed MoE received undersized activation buffers\n"); + return 0; + } + + const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; + const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; + const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; + const BOOL down_simd_q4 = down_scalar_q4; + const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; + const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; + const BOOL stream_addr_q2 = + gate_pair_q2 && down_scalar_q2 && + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && + g_glm_q2_k_addr_down_f32_pipeline != nil; + const BOOL stream_addr_q4 = + !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && + g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && + g_glm_q4_k_addr_down_f32_pipeline != nil; + BOOL use_stream_expert_addr_table = + g_ssd_streaming_mode && + !force_resident && + getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL && + (stream_addr_q2 || stream_addr_q4) && + layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + n_expert <= 8u && + ds4_gpu_stream_expert_cache_configured_budget() >= n_expert && + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + int32_t stream_selected_ids[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + ds4_gpu_stream_expert_cache_entry *stream_entries[8] = { + NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL + }; + uint64_t stream_gate_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint64_t stream_up_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint64_t stream_down_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t stream_missing_mask = 0; + uint32_t stream_entry_count = 0; + uint32_t stream_resident_mask = 0; + BOOL use_stream_split_deferred = false; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + const int glm_stream_timing = + ds4_gpu_stream_expert_timing_summary_enabled(); + + if (use_stream_expert_addr_table) { + const int had_batch = g_batch_cb != nil; + int stream_ok = 1; + const int have_prefetched_selected = + ds4_gpu_glm_stream_selected_prefetch_take(model_map, + model_size, + layer_index, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + stream_selected_ids); + if (have_prefetched_selected) { + if (had_batch && g_batch_has_work && + g_stream_expert_pending_load.active && + ds4_gpu_flush_commands() == 0) { + stream_ok = 0; + } + } else { + if (had_batch && ds4_gpu_end_commands() == 0) { + stream_ok = 0; + } + if (stream_ok && + ds4_gpu_tensor_read(selected, + 0, + stream_selected_ids, + (uint64_t)n_expert * sizeof(stream_selected_ids[0])) == 0) { + stream_ok = 0; + } + } + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + if (stream_selected_ids[i] < 0 || + (uint32_t)stream_selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal GLM routed MoE selected expert id %d is outside 0..%u\n", + stream_selected_ids[i], + n_total_expert); + stream_ok = 0; + } + } + if (stream_ok) { + ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, + stream_selected_ids, + n_expert); + if (!ds4_gpu_moe_selected_hotlist_record(layer_index, + stream_selected_ids, + n_expert, + n_total_expert)) { + stream_ok = 0; + } + } + if (stream_ok) { + g_glm_stream_expert_addr_table_building++; + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + const uint64_t expert_id = (uint64_t)(uint32_t)stream_selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); + stream_ok = 0; + break; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); + stream_ok = 0; + break; + } + stream_gate_abs_offsets[i] = gate_offset + gate_rel; + stream_up_abs_offsets[i] = up_offset + gate_rel; + stream_down_abs_offsets[i] = down_offset + down_rel; + stream_entries[i] = + ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer_index, + (uint32_t)stream_selected_ids[i], + n_total_expert, + n_expert, + stream_gate_abs_offsets[i], + stream_up_abs_offsets[i], + stream_down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!stream_entries[i]) { + stream_missing_mask |= 1u << i; + } else { + stream_resident_mask |= 1u << i; + } + } + if (stream_ok && glm_stream_timing) { + ds4_gpu_stream_expert_timing_note_cache_class( + stream_resident_mask, + stream_missing_mask); + } + use_stream_split_deferred = + stream_ok && + getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_SPLIT") == NULL && + stream_missing_mask != 0 && + stream_resident_mask != 0 && + ds4_gpu_stream_expert_split_worthwhile(stream_resident_mask, + stream_missing_mask) && + ds4_gpu_stream_expert_split_ready() && + ((gate_pair_q2 && + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline != nil) || + (!gate_pair_q2 && !gate_pair_q5 && + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline != nil)); + if (use_stream_split_deferred) { + const ds4_gpu_stream_expert_table table = { + .model_map = model_map, + .model_size = model_size, + .layer = layer_index, + .n_total_expert = n_total_expert, + .gate_offset = gate_offset, + .up_offset = up_offset, + .down_offset = down_offset, + .gate_expert_bytes = gate_expert_bytes, + .down_expert_bytes = down_expert_bytes, + }; + if (!ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + stream_selected_ids, + n_expert)) { + stream_ok = 0; + } + } + if (stream_ok && stream_missing_mask != 0 && + !use_stream_split_deferred && + !ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + stream_selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_missing_mask, + stream_entries)) { + fprintf(stderr, + "ds4: Metal GLM streaming expert cache failed to load " + "layer=%u missing=0x%x budget=%u\n", + layer_index, + stream_missing_mask, + ds4_gpu_stream_expert_cache_configured_budget()); + stream_ok = 0; + } + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_entries[i]; + if (!entry) { + if (use_stream_split_deferred && + (stream_missing_mask & (1u << i)) != 0) { + continue; + } + stream_ok = 0; + break; + } + if (!ds4_gpu_stream_expert_cache_set_addr_slot( + layer_index, + (uint32_t)stream_selected_ids[i], + entry->gate_buffer, + entry->gate_inner, + entry->up_buffer, + entry->up_inner, + entry->down_buffer, + entry->down_inner)) { + stream_ok = 0; + break; + } + stream_entry_count++; + } + if (stream_ok && + !ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf)) { + stream_ok = 0; + } + g_glm_stream_expert_addr_table_building--; + } + if (!have_prefetched_selected && had_batch && + ds4_gpu_begin_commands() == 0) { + stream_ok = 0; + } + if (!stream_ok) return 0; + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + stream_selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + stream_selected_ids, + n_expert); + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = nil; + id upbuf = nil; + id downbuf = nil; + if (!use_stream_expert_addr_table) { + gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + } + + id pair_pipeline = + (use_stream_split_deferred ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32_masked") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32_masked")) : + use_stream_expert_addr_table ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32")) : + gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, + "kernel_glm_q2_K_pair_swiglu_f32") : + gate_pair_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, + "kernel_glm_q5_K_pair_swiglu_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu2_f32")); + id down_pipeline = + (use_stream_expert_addr_table ? + (down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, + "kernel_glm_q2_K_addr_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, + "kernel_glm_q4_K_addr_down_f32")) : + down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, + "kernel_glm_q2_K_down_f32") : + down_scalar_q4 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, + "kernel_glm_q4_K_down_f32") : + down_simd_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, + "kernel_glm_q5_K_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, + "kernel_glm_q6_K_down_f32")); + if (!pair_pipeline || !down_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer( + "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE", + "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE_LAYER", + layer_index); + const char *glm_moe_stage_filter = + getenv("DS4_METAL_GLM_MOE_STAGE_PROFILE_FILTER"); + const char *glm_pair_path = + use_stream_expert_addr_table ? + (gate_pair_q2 ? "q2_stream_addr_swiglu" : + "q4_stream_addr_swiglu") : + gate_pair_q2 ? "q2_scalar_swiglu" : + gate_pair_q5 ? "q5_pair_simd_swiglu" : "q4_pair2_simd_swiglu"; + const char *glm_down_path = + use_stream_expert_addr_table ? + (down_scalar_q2 ? "q2_stream_addr_down" : + "q4_stream_addr_down_simd") : + down_scalar_q2 ? "q2_down_simd" : + down_scalar_q4 ? "q4_down_simd" : + down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + int ok = 1; +#define DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM routed MoE one stage layer=%u tokens=1 experts=%u " \ + "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ + layer_index, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + glm_pair_path, glm_down_path, \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ds4_gpu_glm_routed_moe_args args = { + .tp_rank = g_tp_split_rank, + .tp_world = g_tp_split_world, + .tp_expert_base = tp_expert_base_host, + .in_dim = expert_in_dim, + .mid_dim = expert_mid_dim, + .out_dim = out_dim, + .n_total_expert = n_total_expert, + .n_expert_used = n_expert, + .n_tokens = 1, + .mid_token_stride = n_expert * expert_mid_dim, + .down_type = down_type, + .gate_expert_bytes = gate_expert_bytes, + .gate_row_bytes = gate_row_bytes, + .up_expert_bytes = up_expert_bytes, + .up_row_bytes = up_row_bytes, + .down_expert_bytes = down_expert_bytes, + .down_row_bytes = down_row_bytes, + }; + const NSUInteger pair_x_groups = + gate_pair_q2 ? (use_stream_expert_addr_table ? + (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u)) : + gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : + use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : + (NSUInteger)((expert_mid_dim + 1u) / 2u); + const NSUInteger pair_threadgroup_bytes = 0u; + const NSUInteger pair_threads = 64u; + const NSUInteger down_x_groups = + down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : + down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : + (NSUInteger)out_dim; + const NSUInteger down_threadgroup_bytes = + (down_scalar_q2 || down_simd) ? 0u : 256u * sizeof(float); + const NSUInteger down_threads = + (down_scalar_q2 || down_simd) ? 64u : 256u; + if (use_stream_expert_addr_table && + !ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_entries, + use_stream_split_deferred ? n_expert : stream_entry_count, + use_stream_split_deferred ? stream_resident_mask : 0)) { + return 0; + } + + const int glm_stream_split_timing = + use_stream_split_deferred && glm_stream_timing; + double glm_stream_split_t0 = + glm_stream_split_timing ? ds4_gpu_now_ms() : 0.0; + double glm_stream_split_resident_ms = 0.0; + double glm_stream_split_missing_load_ms = 0.0; + double glm_stream_split_missing_wait_ms = 0.0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + if (use_stream_split_deferred) { + [enc setBytes:&stream_resident_mask length:sizeof(stream_resident_mask) atIndex:1]; + [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; + [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; + } else { + [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner + atIndex:1]; + [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner + atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + } + if (use_stream_expert_addr_table) { + const uint32_t use_count = + use_stream_split_deferred ? n_expert : stream_entry_count; + const uint32_t use_mask = + use_stream_split_deferred ? stream_resident_mask : 0; + for (uint32_t i = 0; i < use_count; i++) { + if (use_mask != 0 && (use_mask & (1u << i)) == 0) continue; + [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + 1) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("pair"); + + if (ok && use_stream_split_deferred) { + id resident_cb = nil; + const int resident_owned = owned; + if (owned) { + resident_cb = cb; + [resident_cb commit]; + cb = nil; + } else { + ok = ds4_gpu_flush_commands(); + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_resident_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + } + if (ok) { + ok = ds4_gpu_stream_expert_pending_load_finish(stream_entries); + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_missing_load_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + } + if (ok) { + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + stream_selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + stream_selected_ids, + n_expert); + ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + } + if (ok) { + if (resident_owned) { + ok = ds4_gpu_wait_command_buffer( + resident_cb, + "GLM streaming split resident pair"); + ds4_gpu_stream_expert_cache_note_owned_completed(); + } else { + ok = ds4_gpu_wait_pending_command_buffers( + "GLM streaming split resident pair"); + } + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_missing_wait_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + ds4_gpu_stream_expert_timing_note_split( + stream_resident_mask, + stream_missing_mask, + glm_stream_split_resident_ms, + glm_stream_split_missing_load_ms + + glm_stream_split_missing_wait_ms); + ds4_gpu_stream_expert_timing_note_split_missing_detail( + glm_stream_split_missing_load_ms, + 0.0, + 0.0, + 0.0, + glm_stream_split_missing_wait_ms); + } + if (ok && + !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_entries, + n_expert, + stream_missing_mask)) { + ok = 0; + } + if (ok) { + cb = ds4_gpu_command_buffer(&owned); + if (!cb) ok = 0; + } + if (ok) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBytes:&stream_missing_mask length:sizeof(stream_missing_mask) atIndex:1]; + [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; + [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_missing_mask & (1u << i)) == 0) continue; + [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + 1) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + } + + if (!ok) return 0; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:down_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner + atIndex:1]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + if (use_stream_expert_addr_table) { + const uint32_t use_count = + use_stream_split_deferred ? n_expert : stream_entry_count; + for (uint32_t i = 0; i < use_count; i++) { + [enc useResource:stream_entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + } + if (down_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("down"); + + if (!ok) return 0; + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed MoE")) return 0; +#undef DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE + } + + return 1; +} + +static bool ds4_gpu_glm_grouped_moe_fast_default(void) { + return !g_quality_mode; +} + +static bool ds4_gpu_glm_routed_moe_batch_grouped_available( + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint32_t n_expert, + uint32_t n_tokens) { + if (n_tokens < 32u || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type) || + ds4_gpu_mul_mm_id_map0_name(n_expert) == NULL) { + return false; + } + const bool fast_default = ds4_gpu_glm_grouped_moe_fast_default(); + if (!fast_default) { + return false; + } + if (n_tokens < 96u) return false; + + return ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)) != nil && + ds4_gpu_routed_mm_pipeline(gate_type) != nil && + ds4_gpu_routed_mm_pipeline(up_type) != nil && + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) != nil; +} + +static bool ds4_gpu_glm_grouped_moe_layer_enabled(uint32_t layer_index) { + (void)layer_index; + return true; +} + +static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return 0; + } + if (n_expert > UINT32_MAX / n_tokens || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + return 0; + } + + const uint32_t pair_rows = n_tokens * n_expert; + if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || + (uint64_t)pair_rows > UINT64_MAX / out_dim || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + return 0; + } + + const bool mid_f16 = true; + const NSUInteger mm_id_threadgroup_bytes = 8192u; + const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; + const uint64_t down_values = (uint64_t)pair_rows * out_dim; + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + if (compact_mid_values > UINT64_MAX / sizeof(float) || + compact_mid_values > UINT64_MAX / (mid_f16 ? sizeof(uint16_t) : sizeof(float)) || + down_values > UINT64_MAX / sizeof(float) || + x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); + const uint64_t mid_bytes = compact_mid_values * (mid_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t down_scratch_bytes = down_values * sizeof(float); + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || + gate_scratch_bytes > NSUIntegerMax || + gate_scratch_bytes * 2ull > NSUIntegerMax || + down_scratch_bytes > NSUIntegerMax) { + return 0; + } + + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM grouped routed MoE received undersized activation buffers\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, + &g_moe_gate_scratch_bytes, + (NSUInteger)(gate_scratch_bytes * 2ull), + "ds4_glm_moe_gate_up_scratch")) { + return 0; + } + if (n_expert > 1 && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_glm_moe_down_scratch")) { + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + id upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + id downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + id gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + id up_pipeline = ds4_gpu_routed_mm_pipeline(up_type); + id down_pipeline = + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type); + if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { + return 0; + } + + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args gate_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args up_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + up_row_bytes, up_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args down_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + mid_f16 ? sizeof(uint16_t) : sizeof(float)); + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = (int32_t)first_expert; + up_args.tp_rank = g_tp_split_rank; + up_args.tp_world = g_tp_split_world; + up_args.tp_expert_base = (int32_t)first_expert; + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_expert_base = (int32_t)first_expert; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + + int ok = 1; +#define DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM grouped routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s mid=%s %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + mid_f16 ? "f16" : "f32", \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &map_args, + &gate_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("map"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + gate_pipeline, + &gate_args, + gatebuf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + 0, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("gate"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + up_pipeline, + &up_args, + upbuf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("up"); + if (ok) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + g_moe_gate_scratch_buffer, + 0, + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + 0.0f, + mid_f16); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + down_pipeline, + &down_args, + downbuf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("down"); + if (ok && n_expert > 1) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("sum"); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM grouped routed batch MoE")) { + return 0; + } +#undef DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE + } + + return 1; +} + +static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + id gate_addrs, + id up_addrs, + id down_addrs, + id overflow_gate, + id overflow_up, + id overflow_down, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count) { + if (n_expert > UINT32_MAX / n_tokens || + !gate_addrs || !up_addrs || !down_addrs || + !ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return 0; + } + + const uint32_t pair_rows = n_tokens * n_expert; + if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || + (uint64_t)pair_rows > UINT64_MAX / out_dim || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + return 0; + } + + const bool mid_f16 = true; + const NSUInteger mm_id_threadgroup_bytes = 8192u; + const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; + const uint64_t down_values = (uint64_t)pair_rows * out_dim; + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + if (compact_mid_values > UINT64_MAX / sizeof(float) || + compact_mid_values > UINT64_MAX / sizeof(uint16_t) || + down_values > UINT64_MAX / sizeof(float) || + x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); + const uint64_t mid_bytes = compact_mid_values * sizeof(uint16_t); + const uint64_t down_scratch_bytes = down_values * sizeof(float); + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || + gate_scratch_bytes > NSUIntegerMax || + gate_scratch_bytes * 2ull > NSUIntegerMax || + down_scratch_bytes > NSUIntegerMax) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM grouped-address routed MoE received undersized activation buffers\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, + &g_moe_gate_scratch_bytes, + (NSUInteger)(gate_scratch_bytes * 2ull), + "ds4_glm_moe_gate_up_scratch")) { + return 0; + } + if (n_expert > 1 && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_glm_moe_down_scratch")) { + return 0; + } + + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + id gate_pipeline = + ds4_gpu_routed_mm_addr_pipeline(gate_type); + id up_pipeline = + ds4_gpu_routed_mm_addr_pipeline(up_type); + id down_pipeline = + ds4_gpu_routed_mm_addr_f16_rhs_pipeline(down_type); + if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { + return 0; + } + + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args gate_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args up_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + up_row_bytes, up_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args down_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + sizeof(uint16_t)); + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(resources, + resource_count, + 0)) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + + int ok = 1; +#define DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM grouped-address routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s mid=f16 %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &map_args, + &gate_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("map"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + gate_pipeline, + &gate_args, + gate_addrs, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + 0, + mm_id_threadgroup_bytes, + resources, + resource_count, + 0, + overflow_gate); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("gate"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + up_pipeline, + &up_args, + up_addrs, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + mm_id_threadgroup_bytes, + resources, + resource_count, + 1, + overflow_up); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("up"); + if (ok) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + g_moe_gate_scratch_buffer, + 0, + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + 0.0f, + mid_f16); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + down_pipeline, + &down_args, + down_addrs, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + mm_id_threadgroup_bytes, + resources, + resource_count, + 2, + overflow_down); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("down"); + if (ok && n_expert > 1) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("sum"); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, + "GLM grouped-address routed batch MoE")) { + return 0; + } +#undef DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE + } + + return 1; +} + +static int ds4_gpu_glm_routed_moe_batch_tensor_impl( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool allow_grouped, + bool force_scalar_q4_pair) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !mid || !model_map || !selected || !weights || !x || + n_tokens == 0 || + n_total_expert == 0 || n_expert == 0 || n_expert > 256u || + n_expert > n_total_expert || + expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || + gate_expert_bytes == 0 || gate_row_bytes == 0 || + up_expert_bytes == 0 || up_row_bytes == 0 || + down_expert_bytes == 0 || down_row_bytes == 0 || + (expert_in_dim % 256u) != 0 || + (expert_mid_dim % 256u) != 0 || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type)) { + return 0; + } + + const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim; + if ((uint64_t)mid_token_stride < per_token_mid || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || + (uint64_t)out_dim > UINT64_MAX / down_row_bytes || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor byte size overflow\n"); + return 0; + } + + const uint64_t full_gate_tensor_bytes = + (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t full_up_tensor_bytes = + (uint64_t)n_total_expert * up_expert_bytes; + const uint64_t full_down_tensor_bytes = + (uint64_t)n_total_expert * down_expert_bytes; + if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || + up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || + down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE received inconsistent expert strides\n"); + return 0; + } + if (gate_offset > model_size || full_gate_tensor_bytes > model_size - gate_offset || + up_offset > model_size || full_up_tensor_bytes > model_size - up_offset || + down_offset > model_size || full_down_tensor_bytes > model_size - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor range is outside the mapped model\n"); + return 0; + } + + if (allow_grouped && + (!g_ssd_streaming_mode || + ds4_gpu_glm_streaming_prefill_full_layer_active()) && + ds4_gpu_glm_grouped_moe_layer_enabled(layer_index) && + ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return ds4_gpu_glm_routed_moe_batch_grouped_tensor(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens); + } + + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + const uint64_t mid_values = + (uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid; + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + if (x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float) || + mid_values > UINT64_MAX / sizeof(float) || + selected_values > UINT64_MAX / sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE activation byte size overflow\n"); + return 0; + } + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t mid_bytes = mid_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE received undersized activation buffers\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = nil; + id upbuf = nil; + id downbuf = nil; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + id stream_overflow_gate = nil; + id stream_overflow_up = nil; + id stream_overflow_down = nil; + ds4_gpu_stream_expert_cache_entry + *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; + uint32_t stream_resource_count = 0; + uint32_t stream_unique = 0; + + const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; + const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; + const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; + const BOOL down_simd_q4 = down_scalar_q4; + const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; + const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; + const BOOL stream_addr_q2 = + gate_pair_q2 && down_scalar_q2 && + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && + g_glm_q2_k_addr_down_f32_pipeline != nil; + const BOOL stream_addr_q4 = + !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && + g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && + g_glm_q4_k_addr_down_f32_pipeline != nil; + BOOL use_stream_expert_addr_table = + g_ssd_streaming_mode && + !ds4_gpu_glm_streaming_prefill_full_layer_active() && + n_tokens > 1 && + (stream_addr_q2 || stream_addr_q4) && + layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + n_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED && + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + const BOOL use_stream_grouped_addr_table = + use_stream_expert_addr_table && + allow_grouped && + getenv("DS4_METAL_GLM_DISABLE_STREAMING_GROUPED_ADDR_PREFILL") == NULL && + ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens); + const BOOL enable_q4_pair4 = true; + const BOOL q4_scalar_pair = false; + const BOOL q4_pair2 = + !use_stream_expert_addr_table && + !gate_pair_q5 && !q4_scalar_pair && + (force_scalar_q4_pair || !enable_q4_pair4); + id pair_pipeline = + use_stream_expert_addr_table ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32")) : + gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, + "kernel_glm_q2_K_pair_swiglu_f32") : + gate_pair_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, + "kernel_glm_q5_K_pair_swiglu_f32") : + (q4_scalar_pair ? + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu_f32") : + q4_pair2 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu4_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu4_f32")); + id down_pipeline = + use_stream_expert_addr_table ? + (down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, + "kernel_glm_q2_K_addr_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, + "kernel_glm_q4_K_addr_down_f32")) : + down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, + "kernel_glm_q2_K_down_f32") : + down_scalar_q4 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, + "kernel_glm_q4_K_down_f32") : + down_simd_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, + "kernel_glm_q5_K_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, + "kernel_glm_q6_K_down_f32"); + if (!pair_pipeline || !down_pipeline) return 0; + + if (use_stream_expert_addr_table) { + const int had_batch = g_batch_cb != nil; + if (had_batch && ds4_gpu_end_commands() == 0) { + return 0; + } + if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( + model_map, + model_size, + layer_index, + selected, + n_tokens, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + stream_resources, + &stream_resource_count, + &stream_unique, + &stream_overflow_gate, + &stream_overflow_up, + &stream_overflow_down)) { + return 0; + } + if (stream_unique == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (had_batch && ds4_gpu_begin_commands() == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (use_stream_grouped_addr_table) { + return ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( + out, + mid, + stream_gate_addr_buf, + stream_up_addr_buf, + stream_down_addr_buf, + stream_overflow_gate, + stream_overflow_up, + stream_overflow_down, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + stream_resources, + stream_resource_count); + } + } else { + gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + const char *glm_pair_path = use_stream_expert_addr_table ? + (gate_pair_q2 ? "q2_stream_addr_swiglu" : + "q4_stream_addr_swiglu") : + gate_pair_q2 ? "q2_scalar_swiglu" : + gate_pair_q5 ? "q5_pair_simd_swiglu" : + (q4_scalar_pair ? "q4_scalar_swiglu" : + (q4_pair2 ? "q4_pair2_simd_swiglu" : + "q4_pair4_simd_swiglu")); + const char *glm_down_path = + use_stream_expert_addr_table ? + (down_scalar_q2 ? "q2_stream_addr_down" : "q4_stream_addr_down_simd") : + down_scalar_q2 ? "q2_down_scalar" : + down_scalar_q4 ? "q4_down_simd" : + down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + int ok = 1; +#define DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM routed MoE batch stage layer=%u tokens=%u experts=%u " \ + "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ + layer_index, n_tokens, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + glm_pair_path, glm_down_path, \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ds4_gpu_glm_routed_moe_args args = { + .tp_rank = g_tp_split_rank, + .tp_world = g_tp_split_world, + .tp_expert_base = (int32_t)first_expert, + .in_dim = expert_in_dim, + .mid_dim = expert_mid_dim, + .out_dim = out_dim, + .n_total_expert = n_total_expert, + .n_expert_used = n_expert, + .n_tokens = n_tokens, + .mid_token_stride = mid_token_stride, + .down_type = down_type, + .gate_expert_bytes = gate_expert_bytes, + .gate_row_bytes = gate_row_bytes, + .up_expert_bytes = up_expert_bytes, + .up_row_bytes = up_row_bytes, + .down_expert_bytes = down_expert_bytes, + .down_row_bytes = down_row_bytes, + }; + const NSUInteger pair_x_groups = + gate_pair_q2 ? (use_stream_expert_addr_table ? + (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u)) : + gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : + use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : + q4_scalar_pair ? (NSUInteger)expert_mid_dim : + q4_pair2 ? (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u); + const NSUInteger pair_threadgroup_bytes = + q4_scalar_pair ? 512u * sizeof(float) : 0u; + const NSUInteger pair_threads = + q4_scalar_pair ? 256u : 64u; + const NSUInteger down_x_groups = + down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : + down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : + (NSUInteger)out_dim; + const NSUInteger down_threadgroup_bytes = + down_simd ? 0u : 256u * sizeof(float); + const NSUInteger down_threads = + down_simd ? 64u : 256u; + if (use_stream_expert_addr_table && + !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_resources, + stream_resource_count, + 0)) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner + atIndex:1]; + [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner + atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + if (use_stream_expert_addr_table) { + for (uint32_t i = 0; i < stream_resource_count; i++) { + [enc useResource:stream_resources[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_resources[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (stream_overflow_gate) [enc useResource:stream_overflow_gate usage:MTLResourceUsageRead]; + if (stream_overflow_up) [enc useResource:stream_overflow_up usage:MTLResourceUsageRead]; + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + (NSUInteger)n_tokens) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("pair"); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:down_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner + atIndex:1]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + if (use_stream_expert_addr_table) { + for (uint32_t i = 0; i < stream_resource_count; i++) { + [enc useResource:stream_resources[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (stream_overflow_down) [enc useResource:stream_overflow_down usage:MTLResourceUsageRead]; + } + if (down_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("down"); + + if (!ok) return 0; + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed batch MoE")) return 0; +#undef DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE + } + + return 1; +} + +int ds4_gpu_glm_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool force_resident) { + (void)force_resident; + return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + mid_token_stride, + true, + false); +} + +int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride) { + return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + mid_token_stride, + false, + false); +} diff --git a/metal/dsv4_misc.metal b/models/glm/metal/shaders/kernels.metal similarity index 71% rename from metal/dsv4_misc.metal rename to models/glm/metal/shaders/kernels.metal index b24167a19f..901f941dfe 100644 --- a/metal/dsv4_misc.metal +++ b/models/glm/metal/shaders/kernels.metal @@ -1,485 +1,3 @@ -struct ds4_metal_args_dsv4_topk_mask { - int64_t ne00; - int64_t ne01; - uint64_t nb00; - uint64_t nb01; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; -}; - -struct ds4_metal_args_dsv4_indexer_weighted_sum { - int64_t ne00; - int64_t ne01; - int64_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - int64_t ne10; - int64_t ne11; - uint64_t nb10; - uint64_t nb11; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; - float scale; -}; - -struct ds4_metal_args_dsv4_softmax_pool { - int64_t ne00; - int64_t ne01; - int64_t ne02; - uint64_t nb00; - uint64_t nb01; - uint64_t nb02; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - int64_t ne0; - int64_t ne1; - uint64_t nb0; - uint64_t nb1; -}; - -struct ds4_metal_args_dsv4_softmax_pool_ratio4_direct { - int64_t n_rows; - uint32_t head_dim; - uint32_t n_comp; - uint32_t replay; - uint32_t pad; -}; - -struct ds4_metal_args_dsv4_compressor_score_ape { - uint32_t width; - uint32_t ratio; - uint32_t pos0; - uint32_t n_tokens; -}; - -struct ds4_metal_args_dsv4_indexed_attention { - uint32_t n_tokens; - uint32_t n_head; - uint32_t n_raw; - uint32_t raw_cap; - uint32_t raw_start; - uint32_t n_comp; - uint32_t top_k; - uint32_t pos0; - uint32_t window; - uint32_t ratio; - uint32_t comp_kv_f16; - uint32_t pad0; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t raw_row_stride; - uint64_t comp_row_stride; - uint64_t topk_token_stride; - uint64_t dst_token_stride; - uint64_t dst_head_stride; - float scale; -}; - -struct ds4_metal_args_dsv4_indexer_scores_fused { - uint32_t n_comp; - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t pos0; - uint32_t ratio; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t weights_token_stride; - uint64_t index_row_stride; - uint64_t score_token_stride; - float scale; -}; - -struct ds4_metal_args_dsv4_router_select_one { - uint32_t has_bias; - uint32_t hash_mode; - uint32_t use_token_buffer; - uint32_t token; - uint32_t hash_rows; -}; - -struct ds4_metal_args_glm_router_select_one { - uint32_t n_expert; - uint32_t n_expert_used; - float expert_weight_scale; - uint32_t pad0; -}; - -struct ds4_metal_args_glm_kv_lora_rms_norm { - uint32_t n_tokens; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - float eps; -}; - -struct ds4_metal_args_glm_k_b_project { - uint32_t n_tokens; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t n_head; - uint32_t row_bytes; - uint32_t weight_type; - uint32_t pad1; - uint32_t pad2; -}; - -struct ds4_metal_args_glm_build_kv_cache { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t n_head; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_rope; - uint32_t value_dim; - uint32_t n_ctx_orig; - uint32_t cache_f16; - uint32_t pad0; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -}; - -struct ds4_metal_args_glm_store_compact_kv { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t qk_rope; - uint32_t cache_f16; - uint32_t pad1; -}; - -struct ds4_metal_args_glm_qkv_norm_store_compact_kv { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t q_n; - uint32_t q_n4; - uint32_t kv_raw_dim; - uint32_t kv_lora_dim; - uint32_t kv_lora_n4; - uint32_t qk_rope; - uint32_t cache_f16; - float eps; - uint32_t pad0; -}; - -struct ds4_metal_args_glm_store_indexer_k { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_cap; - uint32_t head_dim; - uint32_t rot_dim; - uint32_t n_ctx_orig; - uint32_t cache_f16; - uint32_t pad0; - float eps; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - float pad1; -}; - -struct ds4_metal_args_glm_attention_full { - uint32_t pos0; - uint32_t n_tokens; - uint32_t cache_len; - uint32_t cache_cap; - uint32_t n_head; - uint32_t qk_dim; - uint32_t value_dim; - uint32_t pad0; - uint32_t cache_f16; - uint32_t pad1; - uint32_t pad2; - float scale; -}; - -struct ds4_metal_args_glm_fill_selected_range { - uint32_t n_selected; -}; - -struct ds4_metal_args_glm_fill_selected_range_batch { - uint32_t n_tokens; - uint32_t pos0; - uint32_t n_selected; - uint32_t pad_row; -}; - -struct ds4_metal_args_glm_indexer_rope_tail { - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t rot_dim; - uint32_t rot_offset; - uint32_t pos0; - uint32_t n_ctx_orig; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; -}; - -struct ds4_metal_args_glm_indexer_score_one { - uint32_t n_rows; - uint32_t n_head; - uint32_t head_dim; - uint32_t cache_f16; - float scale; -}; - -struct ds4_metal_args_glm_indexer_scores_batch { - uint32_t n_rows; - uint32_t n_tokens; - uint32_t n_head; - uint32_t head_dim; - uint32_t pos0; - uint32_t cache_f16; - uint64_t q_token_stride; - uint64_t q_head_stride; - uint64_t weights_token_stride; - uint64_t score_token_stride; - float scale; -}; - -struct ds4_metal_args_glm_qk_lowrank { - uint32_t n_head; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_dim; - uint32_t row_bytes; - uint32_t weight_type; - uint32_t pad1; - uint32_t pad2; -}; - -struct ds4_metal_args_glm_qk_lowrank_batch { - uint32_t n_tokens; - uint32_t n_head; - uint32_t kv_lora_dim; - uint32_t qk_nope; - uint32_t qk_dim; - uint32_t row_bytes; - uint32_t weight_type; - /* First head this dispatch computes: under tensor-parallel head split - * each rank covers a contiguous half of the heads; buffers and weights - * keep full-model layout and are indexed by absolute head. */ - uint32_t head_base; -}; - -struct ds4_metal_args_glm_attention_indexed_decode { - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t value_type; -}; - -struct ds4_metal_args_glm_attention_indexed_decode_split { - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - uint32_t block_rows; - uint32_t n_blocks; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t value_type; -}; - -struct ds4_metal_args_glm_attention_indexed_batch { - uint32_t n_tokens; - uint32_t n_selected; - uint32_t cache_cap; - uint32_t 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; - uint32_t value_row_bytes; - uint32_t value_type; - uint32_t pos0; - float scale; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - uint32_t head_base; -}; - -struct ds4_metal_args_dsv4_directional_steering_project { - uint32_t width; - uint32_t rows; - uint32_t layer; - uint32_t n_threads; - float scale; -}; - -// Optional directional steering projection. -// -// Each threadgroup owns one 4096-wide token row, computes -// dot(row, direction[layer]), then subtracts scale * direction * dot in-place. -// Positive scales remove a concept direction; negative scales amplify it. The -// kernel is not used unless a steering file and nonzero scale are provided. -kernel void kernel_dsv4_directional_steering_project_f32( - constant ds4_metal_args_dsv4_directional_steering_project & args, - device float *x, - device const float *directions, - threadgroup float *scratch [[threadgroup(0)]], - uint row [[threadgroup_position_in_grid]], - uint tid [[thread_position_in_threadgroup]]) { - if (row >= args.rows || args.width == 0) return; - - device float *xr = x + (uint64_t)row * args.width; - device const float *dir = directions + (uint64_t)args.layer * args.width; - const uint nth = args.n_threads; - - float sum = 0.0f; - for (uint i = tid; i < args.width; i += nth) { - sum += xr[i] * dir[i]; - } - scratch[tid] = sum; - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint step = nth >> 1; step > 0; step >>= 1) { - if (tid < step) scratch[tid] += scratch[tid + step]; - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - const float coeff = args.scale * scratch[0]; - for (uint i = tid; i < args.width; i += nth) { - xr[i] -= coeff * dir[i]; - } -} - -// Decode-only DS4 ratio-4 indexer score builder. One threadgroup owns one -// compressed row for the current token, stages that 128-wide row once, then -// walks the 64 indexer heads in four-head groups. This avoids materializing the -// intermediate [compressed rows x heads] score matrix used by the generic -// matvec + weighted-sum path. -kernel void kernel_dsv4_indexer_score_one_direct( - constant ds4_metal_args_dsv4_indexer_scores_fused & args, - device const char *q, - device const char *weights, - device const char *index_comp, - device char *scores, - threadgroup float *shared [[threadgroup(0)]], - uint row [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]], - ushort sg [[simdgroup_index_in_threadgroup]]) { - if (row >= args.n_comp || args.n_head != 64u || args.head_dim != 128u) { - return; - } - - threadgroup float *ktg = shared; // [128] - threadgroup float *psum = ktg + 128u; // [4] - - if (tid < 128u) { - device const float *krow = (device const float *)(index_comp + - (uint64_t)row * args.index_row_stride); - ktg[tid] = krow[tid]; - } - - float acc = 0.0f; - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint head0 = 0; head0 < 64u; head0 += 4u) { - const uint head = head0 + (uint)sg; - device const float4 *q4 = (device const float4 *)(q + - (uint64_t)head * args.q_head_stride); - threadgroup const float4 *k4 = (threadgroup const float4 *)ktg; - - float s = dot(q4[lane], k4[lane]); - s = simd_sum(s); - if (lane == 0) { - device const float *w = (device const float *)weights; - psum[sg] = max(s, 0.0f) * (w[head] * args.scale); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid == 0) { - acc += psum[0]; - acc += psum[1]; - acc += psum[2]; - acc += psum[3]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - if (tid == 0) { - device float *dst = (device float *)scores; - dst[row] = acc; - } -} - -// Decode router post-processing for one token. The selected expert ids are -// already known; this gathers their probabilities, normalizes by the selected -// sum, clamps the denominator like the reference path, and applies DS4's 1.5 -// expert-weight scale in one tiny dispatch. -kernel void kernel_dsv4_router_weights_one( - device const char *probs, - device const char *selected, - device char *weights, - uint tid [[thread_position_in_grid]]) { - if (tid >= 6) return; - - device const float *p = (device const float *)probs; - device const int *s = (device const int *)selected; - - float sum = 0.0f; - for (uint i = 0; i < 6; i++) { - sum += p[s[i]]; - } - sum = max(sum, 6.103515625e-5f); - - device float *w = (device float *)weights; - w[tid] = p[s[tid]] / sum * 1.5f; -} static inline float ds4_glm_router_sigmoid(float x) { if (x >= 0.0f) { @@ -4640,1518 +4158,3 @@ kernel void kernel_glm_router_select_one( // Six active lanes deliberately match kernel_sum_rows_f32_f32's reduction // topology. The denominator and divided weights cross threadgroup storage // boundaries so division cannot be reassociated with the final scale. -kernel void kernel_dsv4_router_weights_batch( - constant float &scale, - device const float *probs, - device const int32_t *selected, - device float *weights, - threadgroup volatile float *scratch [[threadgroup(0)]], - uint row [[threadgroup_position_in_grid]], - ushort tid [[thread_position_in_threadgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]], - ushort tiisg [[thread_index_in_simdgroup]]) { - if (tid >= 6) return; - - threadgroup volatile float *sum_scratch = scratch; - threadgroup volatile float *denom_scratch = scratch + 32; - threadgroup volatile float *div_scratch = scratch + 33; - const uint out_index = row * 6u + (uint)tid; - const int32_t expert = selected[out_index]; - const float p = probs[row * 256u + (uint)expert]; - - // Keep this sequence identical to kernel_sum_rows_f32_f32 for width 6. - if (sgitg == 0) { - sum_scratch[tiisg] = 0.0f; - } - float sumf = 0.0f; - sumf += p; - sumf = simd_sum(sumf); - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tiisg == 0) { - sum_scratch[sgitg] = sumf; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - sumf = sum_scratch[tiisg]; - sumf = simd_sum(sumf); - - if (tid == 0) { - denom_scratch[0] = clamp(sumf, 6.103515625e-5f, INFINITY); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - div_scratch[tid] = p / denom_scratch[0]; - threadgroup_barrier(mem_flags::mem_threadgroup); - weights[out_index] = div_scratch[tid] * scale; -} - -// Decode router selection for one token after the existing -// sqrt(softplus(logit)) probability kernel has run. Bias affects only top-k -// selection. Route-weight normalization deliberately stays in the old one-token -// kernel: even tiny denominator-order changes here are amplified by 43 MoE -// layers, so this kernel only replaces the selection work. -kernel void kernel_dsv4_router_finalize_one( - constant ds4_metal_args_dsv4_router_select_one & args, - device const float *probs, - device const float *bias, - device const int32_t *hash, - device const int32_t *tokens, - device int32_t *selected, - threadgroup float *scratch [[threadgroup(0)]], - uint tid [[thread_position_in_threadgroup]]) { - if (tid >= 256) return; - - threadgroup float *sel_scores = scratch; - threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256); - const float p = probs[tid]; - sel_scores[tid] = args.has_bias ? p + bias[tid] : p; - idx[tid] = (int32_t)tid; - threadgroup_barrier(mem_flags::mem_threadgroup); - - if (args.hash_mode) { - if (tid == 0) { - const uint token = args.use_token_buffer ? (uint)tokens[0] : args.token; - const uint row = min(token, args.hash_rows - 1u); - device const int32_t *src = hash + row * 6u; - for (uint i = 0; i < 6; i++) { - selected[i] = src[i]; - } - } - } else { - for (uint k = 2; k <= 256; k <<= 1) { - for (uint j = k >> 1; j > 0; j >>= 1) { - const uint other = tid ^ j; - if (other > tid) { - if ((tid & k) == 0) { - if (sel_scores[(uint)idx[tid]] < sel_scores[(uint)idx[other]]) { - const int32_t tmp = idx[tid]; - idx[tid] = idx[other]; - idx[other] = tmp; - } - } else { - if (sel_scores[(uint)idx[tid]] > sel_scores[(uint)idx[other]]) { - const int32_t tmp = idx[tid]; - idx[tid] = idx[other]; - idx[other] = tmp; - } - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - } - if (tid < 6) { - selected[tid] = idx[tid]; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); -} - -// M3 decode specialization for the non-hash one-token router. Scores and ids -// stay in registers. Intra-SIMD bitonic stages use shuffle-xor; the six stages -// that cross 32-lane SIMD groups exchange through alternating threadgroup -// banks. The next bank's publish barrier proves every prior-bank read finished; -// by the time a bank is reused two cross stages later, no reader can remain. -kernel void kernel_dsv4_router_finalize_one_simd( - constant ds4_metal_args_dsv4_router_select_one & args, - device const float *probs, - device const float *bias, - device const int32_t *hash, - device const int32_t *tokens, - device int32_t *selected, - threadgroup float *scratch [[threadgroup(0)]], - uint tid [[thread_position_in_threadgroup]]) { - if (tid >= 256 || args.hash_mode) return; - - (void)hash; - (void)tokens; - threadgroup float *score0_tg = scratch; - threadgroup int32_t *idx0_tg = - (threadgroup int32_t *)(scratch + 256); - threadgroup float *score1_tg = scratch + 512; - threadgroup int32_t *idx1_tg = - (threadgroup int32_t *)(scratch + 768); - const float p = probs[tid]; - float score = args.has_bias ? p + bias[tid] : p; - int32_t idx = (int32_t)tid; - uint cross_stage = 0; - - for (uint k = 2; k <= 256; k <<= 1) { - for (uint j = k >> 1; j > 0; j >>= 1) { - float peer_score; - int32_t peer_idx; - bool take_peer; - const bool lower = (tid & j) == 0; - const bool descending = (tid & k) == 0; - - if (j < 32) { - peer_score = simd_shuffle_xor(score, (ushort)j); - peer_idx = simd_shuffle_xor(idx, (ushort)j); - take_peer = descending - ? (lower ? score < peer_score : score > peer_score) - : (lower ? score > peer_score : score < peer_score); - if (take_peer) { - score = peer_score; - idx = peer_idx; - } - } else { - threadgroup float *score_tg = - (cross_stage & 1u) != 0u ? score1_tg : score0_tg; - threadgroup int32_t *idx_tg = - (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; - score_tg[tid] = score; - idx_tg[tid] = idx; - threadgroup_barrier(mem_flags::mem_threadgroup); - - const uint other = tid ^ j; - peer_score = score_tg[other]; - peer_idx = idx_tg[other]; - take_peer = descending - ? (lower ? score < peer_score : score > peer_score) - : (lower ? score > peer_score : score < peer_score); - if (take_peer) { - score = peer_score; - idx = peer_idx; - } - cross_stage++; - } - } - } - - if (tid < 6) { - selected[tid] = idx; - } -} - -// M3 decode specialization that extends the register/TG SIMD selection above -// through the existing six-value serial weight normalization. The selected ids -// cross the same device-memory boundary as the standalone weight kernel; -// volatile TG stores pin its left-fold and scaled-reciprocal rounding points. -kernel void kernel_dsv4_router_finalize_weights_one_simd( - constant ds4_metal_args_dsv4_router_select_one & args, - device const float *probs, - device const float *bias, - device const int32_t *hash, - device const int32_t *tokens, - device int32_t *selected, - device float *weights, - threadgroup float *scratch [[threadgroup(0)]], - uint tid [[thread_position_in_threadgroup]]) { - if (tid >= 256 || args.hash_mode) return; - - (void)hash; - (void)tokens; - threadgroup float *score0_tg = scratch; - threadgroup int32_t *idx0_tg = - (threadgroup int32_t *)(scratch + 256); - threadgroup float *score1_tg = scratch + 512; - threadgroup int32_t *idx1_tg = - (threadgroup int32_t *)(scratch + 768); - const float p = probs[tid]; - float score = args.has_bias ? p + bias[tid] : p; - int32_t idx = (int32_t)tid; - uint cross_stage = 0; - - for (uint k = 2; k <= 256; k <<= 1) { - for (uint j = k >> 1; j > 0; j >>= 1) { - float peer_score; - int32_t peer_idx; - bool take_peer; - const bool lower = (tid & j) == 0; - const bool descending = (tid & k) == 0; - - if (j < 32) { - peer_score = simd_shuffle_xor(score, (ushort)j); - peer_idx = simd_shuffle_xor(idx, (ushort)j); - take_peer = descending - ? (lower ? score < peer_score : score > peer_score) - : (lower ? score > peer_score : score < peer_score); - if (take_peer) { - score = peer_score; - idx = peer_idx; - } - } else { - threadgroup float *score_tg = - (cross_stage & 1u) != 0u ? score1_tg : score0_tg; - threadgroup int32_t *idx_tg = - (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; - score_tg[tid] = score; - idx_tg[tid] = idx; - threadgroup_barrier(mem_flags::mem_threadgroup); - - const uint other = tid ^ j; - peer_score = score_tg[other]; - peer_idx = idx_tg[other]; - take_peer = descending - ? (lower ? score < peer_score : score > peer_score) - : (lower ? score > peer_score : score < peer_score); - if (take_peer) { - score = peer_score; - idx = peer_idx; - } - cross_stage++; - } - } - } - - if (tid < 6) { - selected[tid] = idx; - } - threadgroup_barrier(mem_flags::mem_device); - - threadgroup volatile float *norm_scratch = - (threadgroup volatile float *)scratch; - if (tid == 0) { - device const int32_t *s = selected; - norm_scratch[0] = 0.0f; - for (uint i = 0; i < 6; i++) { - norm_scratch[0] = norm_scratch[0] + probs[s[i]]; - } - norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); - norm_scratch[1] = 1.5f / norm_scratch[0]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (tid < 6) { - device const int32_t *s = selected; - weights[tid] = probs[s[tid]] * norm_scratch[1]; - } -} - -// Fills the dense compressed-attention mask with -inf. The selected top-k rows -// are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch. -kernel void kernel_dsv4_topk_mask( - constant ds4_metal_args_dsv4_topk_mask & args, - device const char * topk, - device char * dst, - uint gid [[thread_position_in_grid]]) { - const int64_t n = args.ne0 * args.ne1; - if ((int64_t) gid >= n) { - return; - } - - const int64_t ic = gid % args.ne0; - const int64_t it = gid / args.ne0; - - (void)topk; - *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = -INFINITY; -} - -// Enables the selected compressed rows in the dense mask. This replaces the -// old O(n_comp * n_tokens * top_k) membership test with O(top_k * n_tokens) -// writes while preserving exactly the same 0/-inf mask consumed by attention. -kernel void kernel_dsv4_topk_mask_scatter( - constant ds4_metal_args_dsv4_topk_mask & args, - device const char * topk, - device char * dst, - uint gid [[thread_position_in_grid]]) { - const int64_t n = args.ne00 * args.ne01; - if ((int64_t) gid >= n) { - return; - } - - const int64_t ik = gid % args.ne00; - const int64_t it = gid / args.ne00; - const int32_t idx = *((device const int32_t *) (topk + ik*args.nb00 + it*args.nb01)); - if (idx >= 0 && (int64_t)idx < args.ne0) { - *((device float *) (dst + (int64_t)idx*args.nb0 + it*args.nb1)) = 0.0f; - } -} - -// Sorts each token's selected compressed rows by row id. The indexer selects by -// score, but attention scans compressed K/V in cache order in the dense graph. -// Sorting preserves that order while still letting the indexed attention kernel -// touch only the selected rows. -kernel void kernel_dsv4_sort_i32_rows_asc( - constant ds4_metal_args_dsv4_topk_mask & args, - device const char * src, - device char * dst, - threadgroup int32_t * row_tmp [[threadgroup(0)]], - uint row [[threadgroup_position_in_grid]], - uint tid [[thread_position_in_threadgroup]], - uint n_threads [[threads_per_threadgroup]]) { - const uint top_k = (uint)args.ne00; - if (row >= (uint)args.ne01 || tid >= n_threads) { - return; - } - - for (uint i = tid; i < top_k; i += n_threads) { - row_tmp[i] = *((device const int32_t *) (src + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint k = 2; k <= top_k; k <<= 1) { - for (uint j = k >> 1; j > 0; j >>= 1) { - for (uint i = tid; i < top_k; i += n_threads) { - const uint other = i ^ j; - if (other > i && other < top_k) { - const int32_t a = row_tmp[i]; - const int32_t b = row_tmp[other]; - const bool up = (i & k) == 0; - if ((up && a > b) || (!up && a < b)) { - row_tmp[i] = b; - row_tmp[other] = a; - } - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - } - - for (uint i = tid; i < top_k; i += n_threads) { - *((device int32_t *) (dst + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)) = row_tmp[i]; - } -} - -static inline void dsv4_attend_f32_row_as_f16( - device const char *kv, - uint64_t row_stride, - uint row, - half4 q0, - half4 q1, - half4 q2, - half4 q3, - float scale, - ushort lane, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - device const float4 *kv4 = (device const float4 *)(kv + (uint64_t)row * row_stride); - const half4 k0 = (half4)kv4[lane + 0]; - const half4 k1 = (half4)kv4[lane + 32]; - const half4 k2 = (half4)kv4[lane + 64]; - const half4 k3 = (half4)kv4[lane + 96]; - - float score = dot((float4)q0, (float4)k0) + - dot((float4)q1, (float4)k1) + - dot((float4)q2, (float4)k2) + - dot((float4)q3, (float4)k3); - score = simd_sum(score) * scale; - - const float old_m = M; - const float new_m = max(M, score); - const float old_scale = exp(old_m - new_m); - const float row_scale = exp(score - new_m); - - S = S * old_scale + row_scale; - o0 *= old_scale; - o1 *= old_scale; - o2 *= old_scale; - o3 *= old_scale; - - o0 += (float4)k0 * row_scale; - o1 += (float4)k1 * row_scale; - o2 += (float4)k2 * row_scale; - o3 += (float4)k3 * row_scale; - M = new_m; -} - -static inline void dsv4_attend_shared_f32_row_as_f16( - threadgroup const float4 *kv4, - half4 q0, - half4 q1, - half4 q2, - half4 q3, - float scale, - ushort lane, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - const half4 k0 = (half4)kv4[lane + 0]; - const half4 k1 = (half4)kv4[lane + 32]; - const half4 k2 = (half4)kv4[lane + 64]; - const half4 k3 = (half4)kv4[lane + 96]; - - float score = dot((float4)q0, (float4)k0) + - dot((float4)q1, (float4)k1) + - dot((float4)q2, (float4)k2) + - dot((float4)q3, (float4)k3); - score = simd_sum(score) * scale; - - const float old_m = M; - const float new_m = max(M, score); - const float old_scale = exp(old_m - new_m); - const float row_scale = exp(score - new_m); - - S = S * old_scale + row_scale; - o0 *= old_scale; - o1 *= old_scale; - o2 *= old_scale; - o3 *= old_scale; - - o0 += (float4)k0 * row_scale; - o1 += (float4)k1 * row_scale; - o2 += (float4)k2 * row_scale; - o3 += (float4)k3 * row_scale; - M = new_m; -} - -static inline void dsv4_attend_shared_f32_row_as_f16_at( - threadgroup const float4 *kv4, - uint row_in_tg, - half4 q0, - half4 q1, - half4 q2, - half4 q3, - float scale, - ushort lane, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - dsv4_attend_shared_f32_row_as_f16(kv4 + row_in_tg * 128u, - q0, q1, q2, q3, - scale, - lane, - M, S, - o0, o1, o2, o3); -} - -static inline void dsv4_attend_shared_h4_row( - threadgroup const half4 *kv4, - half4 q0, - half4 q1, - half4 q2, - half4 q3, - float scale, - ushort lane, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - const half4 k0 = kv4[lane + 0]; - const half4 k1 = kv4[lane + 32]; - const half4 k2 = kv4[lane + 64]; - const half4 k3 = kv4[lane + 96]; - - float score = dot((float4)q0, (float4)k0) + - dot((float4)q1, (float4)k1) + - dot((float4)q2, (float4)k2) + - dot((float4)q3, (float4)k3); - score = simd_sum(score) * scale; - - const float old_m = M; - const float new_m = max(M, score); - const float old_scale = exp(old_m - new_m); - const float row_scale = exp(score - new_m); - - S = S * old_scale + row_scale; - o0 *= old_scale; - o1 *= old_scale; - o2 *= old_scale; - o3 *= old_scale; - - o0 += (float4)k0 * row_scale; - o1 += (float4)k1 * row_scale; - o2 += (float4)k2 * row_scale; - o3 += (float4)k3 * row_scale; - M = new_m; -} - -static inline void dsv4_attend_shared_h4_row_at( - threadgroup const half4 *kv4, - uint row_in_tg, - half4 q0, - half4 q1, - half4 q2, - half4 q3, - float scale, - ushort lane, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - dsv4_attend_shared_h4_row(kv4 + row_in_tg * 128u, - q0, q1, q2, q3, - scale, - lane, - M, S, - o0, o1, o2, o3); -} - -static inline half4 dsv4_load_cache_h4( - device const char *kv, - uint64_t row_stride, - uint row, - uint col, - bool f16_rows) { - device const char *base = kv + (uint64_t)row * row_stride; - if (f16_rows) { - return ((device const half4 *)base)[col]; - } - return (half4)((device const float4 *)base)[col]; -} - -static inline void dsv4_attend_sink( - float score, - thread float &M, - thread float &S, - thread float4 &o0, - thread float4 &o1, - thread float4 &o2, - thread float4 &o3) { - const float old_m = M; - const float new_m = max(M, score); - const float old_scale = exp(old_m - new_m); - const float row_scale = exp(score - new_m); - - S = S * old_scale + row_scale; - o0 *= old_scale; - o1 *= old_scale; - o2 *= old_scale; - o3 *= old_scale; - M = new_m; -} - -// DS4 ratio-4 indexed mixed attention. It replaces the dense top-k mask path: -// the threadgroup covers one token and eight heads. Top-k rows and local raw -// rows are the same for all heads of a token, so K/V is staged once in -// threadgroup memory and reused by the eight simdgroups. It keeps the DS4 F16 -// attention rounding by casting Q/K/V to half before the dot/value update. -kernel void kernel_dsv4_indexed_mixed_attention_heads8( - constant ds4_metal_args_dsv4_indexed_attention & args, - device const char *q, - device const char *raw_kv, - device const char *comp_kv, - device const char *topk, - device const char *sinks, - device char *dst, - threadgroup half4 *kv_shared [[threadgroup(0)]], - uint2 tgpig [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]], - ushort sg [[simdgroup_index_in_threadgroup]]) { - const uint token = tgpig.x; - const uint head = tgpig.y * 8u + (uint)sg; - if (token >= args.n_tokens || head >= args.n_head) { - return; - } - - device const float4 *q4 = (device const float4 *)(q + - (uint64_t)token * args.q_token_stride + - (uint64_t)head * args.q_head_stride); - const half4 q0 = (half4)q4[lane + 0]; - const half4 q1 = (half4)q4[lane + 32]; - const half4 q2 = (half4)q4[lane + 64]; - const half4 q3 = (half4)q4[lane + 96]; - - float M = -FLT_MAX/2.0f; - float S = 0.0f; - float4 o0 = 0.0f; - float4 o1 = 0.0f; - float4 o2 = 0.0f; - float4 o3 = 0.0f; - - const uint qpos = args.pos0 + token; - const uint last_pos = args.pos0 + args.n_tokens - 1u; - const uint first_raw_pos = last_pos + 1u - args.n_raw; - const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; - const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? - qpos + 1u - args.window : 0u; - uint first = max(first_raw_pos, window_first); - uint last = min(qpos, raw_last_pos); - - if (first <= last) { - for (uint pos = first; pos <= last; pos++) { - const uint logical = pos - first_raw_pos; - const uint row = (args.raw_start + logical) % args.raw_cap; - device const float4 *src = (device const float4 *)(raw_kv + - (uint64_t)row * args.raw_row_stride); - if (tid < 128) kv_shared[tid] = (half4)src[tid]; - threadgroup_barrier(mem_flags::mem_threadgroup); - dsv4_attend_shared_h4_row(kv_shared, - q0, q1, q2, q3, - args.scale, - lane, - M, S, - o0, o1, o2, o3); - threadgroup_barrier(mem_flags::mem_threadgroup); - } - } - - uint visible = (qpos + 1u) / args.ratio; - visible = min(visible, args.n_comp); - device const int32_t *row_topk = (device const int32_t *)(topk + - (uint64_t)token * args.topk_token_stride); - for (uint i = 0; i < args.top_k; i++) { - const int32_t idx = row_topk[i]; - if (idx < 0) { - continue; - } - if ((uint)idx >= visible) { - break; - } - if (tid < 128) { - kv_shared[tid] = dsv4_load_cache_h4(comp_kv, - args.comp_row_stride, - (uint)idx, - tid, - args.comp_kv_f16 != 0u); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - dsv4_attend_shared_h4_row(kv_shared, - q0, q1, q2, q3, - args.scale, - lane, - M, S, - o0, o1, o2, o3); - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); - - const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; - device float4 *dst4 = (device float4 *)(dst + - (uint64_t)token * args.dst_token_stride + - (uint64_t)head * args.dst_head_stride); - dst4[lane + 0] = o0 * inv_s; - dst4[lane + 32] = o1 * inv_s; - dst4[lane + 64] = o2 * inv_s; - dst4[lane + 96] = o3 * inv_s; -} - -// Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8. -// Generation attends one token at a time, so the ratio-4 indexed path spends a -// visible amount of time repeatedly staging the same K/V row for the eight -// heads in a group. This variant stages sixteen selected rows at once and then -// consumes them sequentially, preserving the row order and online softmax math -// while cutting threadgroup barriers in the long top-k scan. -kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16( - constant ds4_metal_args_dsv4_indexed_attention & args, - device const char *q, - device const char *raw_kv, - device const char *comp_kv, - device const char *topk, - device const char *sinks, - device char *dst, - threadgroup half4 *kv_shared [[threadgroup(0)]], - uint2 tgpig [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]], - ushort sg [[simdgroup_index_in_threadgroup]]) { - const uint token = tgpig.x; - const uint head = tgpig.y * 8u + (uint)sg; - if (token >= args.n_tokens || head >= args.n_head) { - return; - } - - device const float4 *q4 = (device const float4 *)(q + - (uint64_t)token * args.q_token_stride + - (uint64_t)head * args.q_head_stride); - const half4 q0 = (half4)q4[lane + 0]; - const half4 q1 = (half4)q4[lane + 32]; - const half4 q2 = (half4)q4[lane + 64]; - const half4 q3 = (half4)q4[lane + 96]; - - float M = -FLT_MAX/2.0f; - float S = 0.0f; - float4 o0 = 0.0f; - float4 o1 = 0.0f; - float4 o2 = 0.0f; - float4 o3 = 0.0f; - - const uint qpos = args.pos0 + token; - const uint last_pos = args.pos0 + args.n_tokens - 1u; - const uint first_raw_pos = last_pos + 1u - args.n_raw; - const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; - const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? - qpos + 1u - args.window : 0u; - uint first = max(first_raw_pos, window_first); - uint last = min(qpos, raw_last_pos); - - if (first <= last) { - for (uint pos0 = first; pos0 <= last; pos0 += 16u) { - const uint n_rows = min(16u, last - pos0 + 1u); - for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { - const uint r = off >> 7; - const uint c = off & 127u; - const uint logical = pos0 + r - first_raw_pos; - const uint row = (args.raw_start + logical) % args.raw_cap; - device const float4 *src = (device const float4 *)(raw_kv + - (uint64_t)row * args.raw_row_stride); - kv_shared[off] = (half4)src[c]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint r = 0; r < n_rows; r++) { - dsv4_attend_shared_h4_row_at(kv_shared, - r, - q0, q1, q2, q3, - args.scale, - lane, - M, S, - o0, o1, o2, o3); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - } - - uint visible = (qpos + 1u) / args.ratio; - visible = min(visible, args.n_comp); - device const int32_t *row_topk = (device const int32_t *)(topk + - (uint64_t)token * args.topk_token_stride); - bool stop = false; - for (uint i = 0; i < args.top_k && !stop; i += 16u) { - uint rows[16]; - uint n_rows = 0; - for (uint j = 0; j < 16u && i + j < args.top_k; j++) { - const int32_t idx = row_topk[i + j]; - if (idx < 0) { - continue; - } - if ((uint)idx >= visible) { - stop = true; - break; - } - rows[n_rows++] = (uint)idx; - } - if (n_rows == 0) { - continue; - } - for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { - const uint r = off >> 7; - const uint c = off & 127u; - kv_shared[off] = dsv4_load_cache_h4(comp_kv, - args.comp_row_stride, - rows[r], - c, - args.comp_kv_f16 != 0u); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint r = 0; r < n_rows; r++) { - dsv4_attend_shared_h4_row_at(kv_shared, - r, - q0, q1, q2, q3, - args.scale, - lane, - M, S, - o0, o1, o2, o3); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); - - const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; - device float4 *dst4 = (device float4 *)(dst + - (uint64_t)token * args.dst_token_stride + - (uint64_t)head * args.dst_head_stride); - dst4[lane + 0] = o0 * inv_s; - dst4[lane + 32] = o1 * inv_s; - dst4[lane + 64] = o2 * inv_s; - dst4[lane + 96] = o3 * inv_s; -} - -static inline float dsv4_indexer_dot128_shared_q( - float4 c0, - float4 c1, - float4 c2, - float4 c3, - threadgroup const float4 *q4, - ushort lane) { - float sum = 0.0f; - if (lane < 8) { - const ushort ib = lane >> 1; - const ushort il = lane & 1; - const ushort base = ib*8 + il*4; - sum += dot(c0, q4[base + 0]); - sum += dot(c1, q4[base + 1]); - sum += dot(c2, q4[base + 2]); - sum += dot(c3, q4[base + 3]); - } - return simd_sum(sum); -} - -// Tiled prefill score builder for the sparse-compressed attention indexer. -// -// The kernel covers an 8-token by 32-compressed-row rectangle: K is copied into -// threadgroup memory once, then reused for all 64 indexer heads, while simdgroup -// matrix multiply computes each 8x8 score subtile. -// -// It still writes the exact score matrix consumed by top-k: -// -// score[t,c] = sum_h relu(dot(Q[t,h], K[c])) * W[t,h] * scale -// -// Causal masking is applied on store so invisible compressed rows become -inf. -kernel void kernel_dsv4_indexer_scores_tiled_f32( - constant ds4_metal_args_dsv4_indexer_scores_fused & args, - device const char *q, - device const char *weights, - device const char *index_comp, - device char *scores, - threadgroup float *shared [[threadgroup(0)]], - uint2 tgpig [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]], - ushort sg [[simdgroup_index_in_threadgroup]]) { - constexpr uint TM = 8; - constexpr uint TN = 32; - constexpr uint TS = 8; - constexpr uint D = 128; - - const uint c0 = tgpig.x * TN; - const uint t0 = tgpig.y * TM; - - threadgroup float *qtg = shared; // [8][128] - threadgroup float *ktg = qtg + TM*D; // [32][128] - threadgroup float *dot = ktg + TN*D; // [8][32] - - const uint last_token = min(t0 + TM, args.n_tokens); - const uint max_visible = last_token > t0 ? - min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; - - if (c0 >= max_visible) { - for (uint i = tid; i < TM*TN; i += 128) { - const uint r = i / TN; - const uint cc = i - r*TN; - const uint token = t0 + r; - const uint comp = c0 + cc; - if (token < args.n_tokens && comp < args.n_comp) { - device float *dst = (device float *)(scores + - (uint64_t)token * args.score_token_stride) + comp; - *dst = -INFINITY; - } - } - return; - } - - for (uint i = tid; i < TN*D; i += 128) { - const uint cc = i / D; - const uint d = i - cc*D; - const uint comp = c0 + cc; - float v = 0.0f; - if (comp < args.n_comp) { - device const float *row = (device const float *)(index_comp + - (uint64_t)comp * args.index_row_stride); - v = row[d]; - } - ktg[i] = v; - } - - const uint cell0 = lane; - const uint cell1 = lane + 32u; - const uint row0 = cell0 >> 3; - const uint row1 = cell1 >> 3; - const uint sub0 = cell0 & 7u; - const uint sub1 = cell1 & 7u; - const uint col0 = (uint)sg * TS + sub0; - const uint col1 = (uint)sg * TS + sub1; - const uint token0 = t0 + row0; - const uint token1 = t0 + row1; - const uint comp0 = c0 + col0; - const uint comp1 = c0 + col1; - - float acc0 = 0.0f; - float acc1 = 0.0f; - - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint head = 0; head < args.n_head; head++) { - for (uint i = tid; i < TM*D; i += 128) { - const uint r = i / D; - const uint d = i - r*D; - const uint token = t0 + r; - float v = 0.0f; - if (token < args.n_tokens) { - device const float *qrow = (device const float *)(q + - (uint64_t)token * args.q_token_stride + - (uint64_t)head * args.q_head_stride); - v = qrow[d]; - } - qtg[i] = v; - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); - for (uint db = 0; db < D/TS; db++) { - simdgroup_float8x8 mq; - simdgroup_float8x8 mk; - simdgroup_load(mq, qtg + db*TS, D, 0, false); - simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); - simdgroup_multiply_accumulate(mdot, mq, mk, mdot); - } - - simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - if (token0 < args.n_tokens && comp0 < args.n_comp) { - device const float *w = (device const float *)(weights + - (uint64_t)token0 * args.weights_token_stride); - const float s = dot[row0*TN + col0]; - acc0 += max(s, 0.0f) * (w[head] * args.scale); - } - if (token1 < args.n_tokens && comp1 < args.n_comp) { - device const float *w = (device const float *)(weights + - (uint64_t)token1 * args.weights_token_stride); - const float s = dot[row1*TN + col1]; - acc1 += max(s, 0.0f) * (w[head] * args.scale); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - if (token0 < args.n_tokens && comp0 < args.n_comp) { - const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); - device float *dst = (device float *)(scores + - (uint64_t)token0 * args.score_token_stride) + comp0; - *dst = comp0 < visible ? acc0 : -INFINITY; - } - if (token1 < args.n_tokens && comp1 < args.n_comp) { - const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); - device float *dst = (device float *)(scores + - (uint64_t)token1 * args.score_token_stride) + comp1; - *dst = comp1 < visible ? acc1 : -INFINITY; - } -} - -kernel void kernel_dsv4_indexer_scores_tiled( - constant ds4_metal_args_dsv4_indexer_scores_fused & args, - device const char *q, - device const char *weights, - device const char *index_comp, - device char *scores, - threadgroup float *shared [[threadgroup(0)]], - uint2 tgpig [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]], - ushort lane [[thread_index_in_simdgroup]], - ushort sg [[simdgroup_index_in_threadgroup]]) { - constexpr uint TM = 8; - constexpr uint TN = 32; - constexpr uint TS = 8; - constexpr uint D = 128; - - const uint c0 = tgpig.x * TN; - const uint t0 = tgpig.y * TM; - - // Q/K are staged as half but the dot accumulator and final score remain - // float. This is the one intentional precision tradeoff in the indexer: - // the indexer only ranks compressed rows for top-k selection, and long - // context profiling shows this score matrix dominates the prefill slope. - threadgroup half *qtg = (threadgroup half *)shared; // [8][128] - threadgroup half *ktg = qtg + TM*D; // [32][128] - threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [8][32] - - const uint last_token = min(t0 + TM, args.n_tokens); - const uint max_visible = last_token > t0 ? - min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; - - if (c0 >= max_visible) { - for (uint i = tid; i < TM*TN; i += 128) { - const uint r = i / TN; - const uint cc = i - r*TN; - const uint token = t0 + r; - const uint comp = c0 + cc; - if (token < args.n_tokens && comp < args.n_comp) { - device float *dst = (device float *)(scores + - (uint64_t)token * args.score_token_stride) + comp; - *dst = -INFINITY; - } - } - return; - } - - // Stage compressed index rows once. Edge columns are zeroed so the matrix - // loads below can stay regular; guarded stores discard them. - for (uint i = tid; i < TN*D; i += 128) { - const uint cc = i / D; - const uint d = i - cc*D; - const uint comp = c0 + cc; - half v = half(0.0f); - if (comp < args.n_comp) { - device const float *row = (device const float *)(index_comp + - (uint64_t)comp * args.index_row_stride); - v = half(row[d]); - } - ktg[i] = v; - } - - const uint cell0 = lane; - const uint cell1 = lane + 32u; - const uint row0 = cell0 >> 3; - const uint row1 = cell1 >> 3; - const uint sub0 = cell0 & 7u; - const uint sub1 = cell1 & 7u; - const uint col0 = (uint)sg * TS + sub0; - const uint col1 = (uint)sg * TS + sub1; - const uint token0 = t0 + row0; - const uint token1 = t0 + row1; - const uint comp0 = c0 + col0; - const uint comp1 = c0 + col1; - - float acc0 = 0.0f; - float acc1 = 0.0f; - - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint head = 0; head < args.n_head; head++) { - // Stage Q for the eight-token tile. Each 8x8 matrix load below reads a - // contiguous depth block from this layout. - for (uint i = tid; i < TM*D; i += 128) { - const uint r = i / D; - const uint d = i - r*D; - const uint token = t0 + r; - half v = half(0.0f); - if (token < args.n_tokens) { - device const float *qrow = (device const float *)(q + - (uint64_t)token * args.q_token_stride + - (uint64_t)head * args.q_head_stride); - v = half(qrow[d]); - } - qtg[i] = v; - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); - for (uint db = 0; db < D/TS; db++) { - simdgroup_half8x8 mq; - simdgroup_half8x8 mk; - simdgroup_load(mq, qtg + db*TS, D, 0, false); - simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); - simdgroup_multiply_accumulate(mdot, mq, mk, mdot); - } - - simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - if (token0 < args.n_tokens && comp0 < args.n_comp) { - device const float *w = (device const float *)(weights + - (uint64_t)token0 * args.weights_token_stride); - const float s = dot[row0*TN + col0]; - acc0 += max(s, 0.0f) * (w[head] * args.scale); - } - if (token1 < args.n_tokens && comp1 < args.n_comp) { - device const float *w = (device const float *)(weights + - (uint64_t)token1 * args.weights_token_stride); - const float s = dot[row1*TN + col1]; - acc1 += max(s, 0.0f) * (w[head] * args.scale); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - if (token0 < args.n_tokens && comp0 < args.n_comp) { - const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); - device float *dst = (device float *)(scores + - (uint64_t)token0 * args.score_token_stride) + comp0; - *dst = comp0 < visible ? acc0 : -INFINITY; - } - if (token1 < args.n_tokens && comp1 < args.n_comp) { - const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); - device float *dst = (device float *)(scores + - (uint64_t)token1 * args.score_token_stride) + comp1; - *dst = comp1 < visible ? acc1 : -INFINITY; - } -} - -#ifdef DS4_METAL_HAS_TENSOR -// Retained full-512 prefill indexer score path. This is the part of sparse -// compressed attention that maps cleanly to TensorOps: a regular token by -// compressed-row dot tile. The kernel intentionally leaves top-k selection and -// indexed attention semantics unchanged; all 512 selected rows remain available -// to the later attention kernel. -// -// Each matmul processes a pair of heads (TQ = 2 x TM q rows): the per-element -// dot is still a 128-deep reduction in 32-wide k-steps, so scores are -// bit-identical to single-head tiles while the run count halves. The q tile -// is double-buffered, so the next k-step's stage overlaps the current -// cooperative matmul and each pair needs 5 barriers instead of 10. q and k -// staging use one float4/half4 per lane (each thread covers one row of 8/32 -// consecutive elements), which is the same half(float) conversion per element -// as the scalar form. -kernel void kernel_dsv4_indexer_scores_nax( - constant ds4_metal_args_dsv4_indexer_scores_fused & args, - device const char *q, - device const char *weights, - device const char *index_comp, - device char *scores, - threadgroup half *shared [[threadgroup(0)]], - uint2 tgpig [[threadgroup_position_in_grid]], - ushort tid [[thread_index_in_threadgroup]]) { - constexpr int TM = 16; - constexpr int TQ = 32; - constexpr int TN = 32; - constexpr int NK = 32; - constexpr int D = 128; - constexpr int NUM_THREADS = 128; - - // The 16-token x 32-row tile was the winning NAX shape in local sweeps. A - // wider 64-row compressed tile increased setup/cache pressure and was - // slower despite doing more work per dispatch. - const uint c0 = tgpig.x * TN; - const uint t0 = tgpig.y * TM; - - threadgroup half *qtg = shared; // 2 x [TQ][NK] - threadgroup half *ktg = qtg + 2*TQ*NK; // [32][128] - threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [TQ][TN], column-major - - const uint last_token = min(t0 + (uint)TM, args.n_tokens); - const uint max_visible = last_token > t0 ? - min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; - - if (c0 >= max_visible) { - for (uint i = tid; i < TM*TN; i += NUM_THREADS) { - const uint r = i / TN; - const uint cc = i - r*TN; - const uint token = t0 + r; - const uint comp = c0 + cc; - if (token < args.n_tokens && comp < args.n_comp) { - device float *dst = (device float *)(scores + - (uint64_t)token * args.score_token_stride) + comp; - *dst = -INFINITY; - } - } - return; - } - - { - // One compressed row per 4 threads, 32 consecutive floats per thread. - const uint cc = tid / 4; - const uint comp = c0 + cc; - device const float *krow = nullptr; - if (comp < args.n_comp) { - krow = (device const float *)(index_comp + - (uint64_t)comp * args.index_row_stride); - } - const uint d0 = (tid % 4) * 32; - FOR_UNROLL (uint j = 0; j < 8; j++) { - const float4 kv = krow ? *(device const float4 *)(krow + d0 + 4*j) - : float4(0.0f); - *(threadgroup half4 *)(ktg + cc*D + d0 + 4*j) = half4(kv); - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - float acc[4]; - #pragma unroll - for (uint j = 0; j < 4; j++) { - acc[j] = 0.0f; - } - - auto tq0 = tensor(qtg, dextents(NK, TQ)); - auto tq1 = tensor(qtg + TQ*NK, dextents(NK, TQ)); - auto tk = tensor(ktg, dextents(D, TN)); - auto td = tensor(dot, dextents(TQ, TN), array({1, TQ})); - - matmul2d< - matmul2d_descriptor(TN, TQ, NK, false, true, false, - matmul2d_descriptor::mode::multiply_accumulate), - execution_simdgroups<4>> mm; - - // One q row per 4 threads, 8 consecutive floats per thread. Row r covers - // head (r / TM) of the pair and token row (r % TM). - const uint q_r = tid / 4; - const uint q_k4 = (tid % 4) * 8; - const uint q_hl = q_r / TM; - const uint q_tr = q_r % TM; - const uint q_token = t0 + q_tr; - device const char *q_row_base = nullptr; - if (q_token < args.n_tokens) { - q_row_base = q + (uint64_t)q_token * args.q_token_stride; - } - - auto stage_q = [&](const uint head0, const uint loop_k, threadgroup half *buf) { - const uint head = head0 + q_hl; - half4 v0 = half4(0.0f); - half4 v1 = half4(0.0f); - if (q_row_base && head < args.n_head) { - device const float4 *src4 = (device const float4 *) - (q_row_base + (uint64_t)head * args.q_head_stride + - (uint64_t)(loop_k + q_k4) * sizeof(float)); - v0 = half4(src4[0]); - v1 = half4(src4[1]); - } - *(threadgroup half4 *)(buf + q_r*NK + q_k4) = v0; - *(threadgroup half4 *)(buf + q_r*NK + q_k4 + 4) = v1; - }; - - for (uint head0 = 0; head0 < args.n_head; head0 += 2) { - auto ct = mm.template get_destination_cooperative_tensor(); - #pragma unroll - for (uint16_t i = 0; i < ct.get_capacity(); i++) { - if (ct.is_valid_element(i)) { - ct[i] = 0.0f; - } - } - - stage_q(head0, 0, qtg); - threadgroup_barrier(mem_flags::mem_threadgroup); - - uint qsel = 0; - FOR_UNROLL (uint i = 0; i < 4; i++) { - auto mk = tk.slice(i*NK, 0); - auto mq = (qsel ? tq1 : tq0).slice(0, 0); - mm.run(mk, mq, ct); - if (i < 3) { - qsel ^= 1u; - stage_q(head0, (i + 1)*NK, qsel ? qtg + TQ*NK : qtg); - threadgroup_barrier(mem_flags::mem_threadgroup); - } - } - - ct.store(td); - threadgroup_barrier(mem_flags::mem_threadgroup); - - #pragma unroll - for (uint j = 0; j < 4; j++) { - const uint linear = (uint)tid + j*NUM_THREADS; - if (linear < TM*TN) { - const uint r = linear / TN; - const uint cc = linear - r*TN; - const uint token = t0 + r; - if (token < args.n_tokens) { - device const float *w = (device const float *)(weights + - (uint64_t)token * args.weights_token_stride); - acc[j] += max(dot[cc*TQ + r], 0.0f) * (w[head0] * args.scale); - if (head0 + 1 < args.n_head) { - acc[j] += max(dot[cc*TQ + TM + r], 0.0f) * (w[head0 + 1] * args.scale); - } - } - } - } - // No barrier here: the next pair's q stage and these dot reads touch - // different buffers, and the next q-stage barrier separates the next - // ct.store from these reads. - } - - #pragma unroll - for (uint j = 0; j < 4; j++) { - const uint linear = (uint)tid + j*NUM_THREADS; - if (linear >= TM*TN) { - continue; - } - const uint r = linear / TN; - const uint cc = linear - r*TN; - const uint token = t0 + r; - const uint comp = c0 + cc; - if (token < args.n_tokens && comp < args.n_comp) { - const uint visible = min((args.pos0 + token + 1u) / args.ratio, args.n_comp); - device float *dst = (device float *)(scores + - (uint64_t)token * args.score_token_stride) + comp; - *dst = comp < visible ? acc[j] : -INFINITY; - } - } -} -#endif - -// Collapses per-head indexer scores into one score per compressed row using the -// learned head weights. Negative head scores are clipped exactly as DS4 expects. -kernel void kernel_dsv4_indexer_weighted_sum( - constant ds4_metal_args_dsv4_indexer_weighted_sum & args, - device const char * scores, - device const char * weights, - device char * dst, - uint gid [[thread_position_in_grid]]) { - const int64_t n = args.ne0 * args.ne1; - if ((int64_t) gid >= n) { - return; - } - - const int64_t ic = gid % args.ne0; - const int64_t it = gid / args.ne0; - - float acc = 0.0f; - for (int64_t ih = 0; ih < args.ne02; ++ih) { - const float s = *((device const float *) (scores + ic*args.nb00 + it*args.nb01 + ih*args.nb02)); - const float w = *((device const float *) (weights + ih*args.nb10 + it*args.nb11)); - acc += max(s, 0.0f) * (w * args.scale); - } - - *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = acc; -} - -// Adds the periodic compressor APE directly to projected scores. The legacy -// path materializes one repeated APE segment per period and then performs this -// same single F32 add; these kernels remove only that intermediate copy graph. -kernel void kernel_dsv4_compressor_score_ape_f32( - constant ds4_metal_args_dsv4_compressor_score_ape & args, - device const float *score, - device const float *ape, - device float *dst, - uint gid [[thread_position_in_grid]]) { - const uint64_t total = (uint64_t)args.n_tokens * args.width; - if ((uint64_t)gid >= total) return; - - const uint token = gid / args.width; - const uint col = gid - token*args.width; - const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); - dst[gid] = score[gid] + ape[(uint64_t)ape_row*args.width + col]; -} - -kernel void kernel_dsv4_compressor_score_ape_f16( - constant ds4_metal_args_dsv4_compressor_score_ape & args, - device const float *score, - device const half *ape, - device float *dst, - uint gid [[thread_position_in_grid]]) { - const uint64_t total = (uint64_t)args.n_tokens * args.width; - if ((uint64_t)gid >= total) return; - - const uint token = gid / args.width; - const uint col = gid - token*args.width; - const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); - dst[gid] = score[gid] + float(ape[(uint64_t)ape_row*args.width + col]); -} - -// Fused softmax-weighted pooling of compressed KV rows. It is used when several -// compressor rows are present; the one-row case deliberately follows the -// unfused softmax/mul/sum graph in Objective-C to keep identical reductions. -kernel void kernel_dsv4_softmax_pool( - constant ds4_metal_args_dsv4_softmax_pool & args, - device const char * kv, - device const char * score, - device char * dst, - uint gid [[thread_position_in_grid]]) { - const int64_t n = args.ne0 * args.ne1; - if ((int64_t) gid >= n) { - return; - } - - const int64_t id = gid % args.ne0; - const int64_t ic = gid / args.ne0; - - float max_s = -INFINITY; - for (int64_t ir = 0; ir < args.ne00; ++ir) { - const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); - max_s = max(max_s, s); - } - - float sum = 0.0f; - float acc = 0.0f; - for (int64_t ir = 0; ir < args.ne00; ++ir) { - const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); - const float w = exp(s - max_s); - const float v = *((device const float *) (kv + ir*args.nb00 + id*args.nb01 + ic*args.nb02)); - sum += w; - acc += v*w; - } - - *((device float *) (dst + id*args.nb0 + ic*args.nb1)) = acc/sum; -} - - - -// Tensor-parallel keep-alive: a few threadgroups of FMAs dispatched -// back-to-back on a side queue while TP decode runs. The per-layer gate -// stalls make the real workload look idle to the GPU power manager, which -// otherwise halves the clocks within a second (~2x decode regression); -// this holds them up for negligible bandwidth and a few watts. -kernel void kernel_dsv4_tp_keepalive( - device float * out, - constant uint & iters, - uint tid [[thread_position_in_grid]]) { - float a = out[tid]; - const float b = 1.000001f; - for (uint i = 0; i < iters; i++) { - a = fma(a, b, 0.000001f); - a = fma(a, b, -0.000001f); - } - out[tid] = a; -} - -// Tensor-parallel gate flag: publishes a sequence number to a slab slot the -// CPU service thread spin-reads, replacing the much slower shared-event -// signal for the GPU->CPU direction. Ordering against the partial-output -// kernels comes from the buffer hazard on the shared slab. -kernel void kernel_dsv4_tp_flag_set( - device atomic_uint & flag, - constant uint & value, - uint tid [[thread_position_in_grid]]) { - if (tid == 0) { - atomic_store_explicit(&flag, value, memory_order_relaxed); - } -} - -// Ratio-4 compressor pooling without materializing the [n_comp, 8, head_dim] -// KV and score packs. The row mapping and both reduction loops deliberately -// match kernel_dsv4_softmax_pool so the arithmetic order is unchanged. -kernel void kernel_dsv4_softmax_pool_ratio4_direct( - constant ds4_metal_args_dsv4_softmax_pool_ratio4_direct & args, - device const float * kv, - device const float * score, - device const float * state_kv, - device const float * state_score, - device float * dst, - uint gid [[thread_position_in_grid]]) { - const uint64_t n = (uint64_t)args.head_dim * args.n_comp; - if ((uint64_t)gid >= n || args.head_dim == 0u) { - return; - } - - const uint64_t id = gid % args.head_dim; - const uint64_t ic = gid / args.head_dim; - const uint64_t input_row_stride = 2ull * args.head_dim; - - float max_s = -INFINITY; - float sum = 0.0f; - float acc = 0.0f; - if (ic != 0u) { - const int64_t token_base = (int64_t)ic * 4 - 4; - for (int64_t ir = 0; ir < args.n_rows; ++ir) { - const uint64_t token = (uint64_t)(token_base + ir); - const uint64_t src = token * input_row_stride + - ((uint64_t)ir >> 2u) * args.head_dim + id; - const float s = score[src]; - max_s = max(max_s, s); - } - - for (int64_t ir = 0; ir < args.n_rows; ++ir) { - const uint64_t token = (uint64_t)(token_base + ir); - const uint64_t src = token * input_row_stride + - ((uint64_t)ir >> 2u) * args.head_dim + id; - const float s = score[src]; - const float w = exp(s - max_s); - const float v = kv[src]; - sum += w; - acc += v*w; - } - } else { - for (int64_t ir = 0; ir < args.n_rows; ++ir) { - float s; - if (ir >= 4) { - const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + - args.head_dim + id; - s = score[src]; - } else if (args.replay != 0u) { - s = state_score[(uint64_t)ir * input_row_stride + id]; - } else { - s = -INFINITY; - } - max_s = max(max_s, s); - } - - for (int64_t ir = 0; ir < args.n_rows; ++ir) { - float s; - float v; - if (ir >= 4) { - const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + - args.head_dim + id; - s = score[src]; - v = kv[src]; - } else if (args.replay != 0u) { - const uint64_t src = (uint64_t)ir * input_row_stride + id; - s = state_score[src]; - v = state_kv[src]; - } else { - s = -INFINITY; - v = 0.0f; - } - const float w = exp(s - max_s); - sum += w; - acc += v*w; - } - } - - dst[ic * args.head_dim + id] = acc/sum; -} diff --git a/models/glm/provider.c b/models/glm/provider.c new file mode 100644 index 0000000000..5c7c629e26 --- /dev/null +++ b/models/glm/provider.c @@ -0,0 +1,32 @@ +#include "provider.h" + +#include "../../ds4_model_provider_builtin.h" + +static const ds4_model_provider_v1 DS4_GLM_PROVIDER = { + .abi_version = DS4_MODEL_PROVIDER_ABI_VERSION, + .struct_size = sizeof(ds4_model_provider_v1), + .id = "glm-dsa", + .session_create = ds4_glm_session_create, + .session_destroy = ds4_glm_session_destroy, + .session_sync = ds4_glm_session_sync, + .session_eval = ds4_glm_session_eval, + .sessions_eval_batch = ds4_builtin_sessions_eval_batch, + .sessions_eval_batch_with_prefill = + ds4_builtin_sessions_eval_batch_with_prefill, + .session_eval_speculative = ds4_glm_session_eval_speculative, + .session_invalidate = ds4_glm_session_invalidate, + .session_rewind = ds4_glm_session_rewind, + .session_layer_slice_reset = ds4_glm_session_layer_slice_reset, + .session_eval_output_head = ds4_glm_session_eval_output_head, + .session_eval_layer_slice = ds4_glm_session_eval_layer_slice, + .session_payload_bytes = ds4_glm_session_payload_bytes, + .session_save_payload = ds4_glm_session_save_payload, + .session_load_payload = ds4_glm_session_load_payload, + .session_layer_payload_bytes = ds4_glm_session_layer_payload_bytes, + .session_save_layer_payload = ds4_glm_session_save_layer_payload, + .session_load_layer_payload = ds4_glm_session_load_layer_payload, +}; + +const ds4_model_provider_v1 *ds4_glm_model_provider(void) { + return &DS4_GLM_PROVIDER; +} diff --git a/models/glm/provider.h b/models/glm/provider.h new file mode 100644 index 0000000000..4fb5481d94 --- /dev/null +++ b/models/glm/provider.h @@ -0,0 +1,87 @@ +#ifndef DS4_GLM_MODEL_PROVIDER_H +#define DS4_GLM_MODEL_PROVIDER_H + +#include "../../ds4_model_provider.h" + +const ds4_model_provider_v1 *ds4_glm_model_provider(void); + +int ds4_glm_session_create(ds4_session **out, + ds4_engine *engine, + int context_size); +void ds4_glm_session_destroy(ds4_session *session); +int ds4_glm_session_sync(ds4_session *session, + const ds4_tokens *prompt, + char *err, + size_t errlen); +int ds4_glm_session_eval(ds4_session *session, + int token, + bool probe_support_model, + char *err, + size_t errlen); +int ds4_glm_session_eval_speculative( + ds4_session *session, + int first_token, + int max_tokens, + int eos_token, + int *accepted, + int accepted_cap, + char *err, + size_t errlen); +void ds4_glm_session_invalidate(ds4_session *session); +void ds4_glm_session_rewind(ds4_session *session, int position); +int ds4_glm_session_layer_slice_reset(ds4_session *session, + char *err, + size_t errlen); +int ds4_glm_session_eval_output_head( + ds4_session *session, + const float *hidden_state, + uint32_t token_count, + float *logits, + char *err, + size_t errlen); +int ds4_glm_session_eval_layer_slice( + ds4_session *session, + const int *tokens, + uint32_t token_count, + uint32_t position, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hidden_state, + float *output_hidden_state, + bool output_logits, + float *logits, + char *err, + size_t errlen); +uint64_t ds4_glm_session_payload_bytes(ds4_session *session); +int ds4_glm_session_save_payload(ds4_session *session, + FILE *file, + char *err, + size_t errlen); +int ds4_glm_session_load_payload(ds4_session *session, + FILE *file, + uint64_t payload_bytes, + char *err, + size_t errlen); +uint64_t ds4_glm_session_layer_payload_bytes( + ds4_session *session, + uint32_t layer_start, + uint32_t layer_end); +int ds4_glm_session_save_layer_payload( + ds4_session *session, + FILE *file, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); +int ds4_glm_session_load_layer_payload( + ds4_session *session, + FILE *file, + uint64_t payload_bytes, + const int *tokens, + uint32_t token_count, + uint32_t layer_start, + uint32_t layer_end, + char *err, + size_t errlen); + +#endif diff --git a/rocm/ds4_rocm_glm.cuh b/models/glm/rocm/kernels.cuh similarity index 100% rename from rocm/ds4_rocm_glm.cuh rename to models/glm/rocm/kernels.cuh diff --git a/rocm/README.md b/rocm/README.md new file mode 100644 index 0000000000..dcb1f4a622 --- /dev/null +++ b/rocm/README.md @@ -0,0 +1,10 @@ +# ROCm implementation units + +This directory owns ROCm runtime code and low-level implementations reused by +model integrations: tensor storage, hipBLASLt matmul, quantization, embedding, +normalization/RoPE, shared-expert, and MoE launch paths. + +Model-specific ROCm implementations live under `models//rocm/`. +`ds4_rocm.cu` includes the shared and model-owned headers into one translation +unit, so the ownership split introduces no dispatch or wrapper layer between a +model and its custom kernels. diff --git a/tests/test_session_snapshot.c b/tests/test_session_snapshot.c new file mode 100644 index 0000000000..38105be8fe --- /dev/null +++ b/tests/test_session_snapshot.c @@ -0,0 +1,111 @@ +/* Model-backed provider snapshot round-trip. + * + * This is intentionally not part of `make test`: it loads a full supported + * model. Run with: + * + * DS4_TEST_MODEL=/path/to/model.gguf make test-session-snapshot + */ + +#include "ds4.h" + +#include +#include +#include + +#define TEST_CTX 512 + +static void fail(const char *what, const char *detail) { + fprintf(stderr, "FAIL: %s%s%s\n", + what, detail && detail[0] ? ": " : "", detail ? detail : ""); + exit(1); +} + +static void compare_logits(ds4_session *a, ds4_session *b, + float *a_logits, float *b_logits, int vocab, + const char *stage) { + if (ds4_session_copy_logits(a, a_logits, vocab) != vocab || + ds4_session_copy_logits(b, b_logits, vocab) != vocab) { + fail("copy logits", stage); + } + if (memcmp(a_logits, b_logits, (size_t)vocab * sizeof(float)) != 0) { + fail("logits changed across snapshot round-trip", stage); + } + if (ds4_session_argmax(a) != ds4_session_argmax(b)) { + fail("argmax changed across snapshot round-trip", stage); + } +} + +int main(void) { + const char *model = getenv("DS4_TEST_MODEL"); + if (!model || !model[0]) fail("DS4_TEST_MODEL is not set", NULL); + + ds4_engine_options opt = { + .model_path = model, +#if defined(__APPLE__) + .backend = DS4_BACKEND_METAL, +#else + .backend = DS4_BACKEND_CUDA, +#endif + .n_threads = 1, + .context_size = TEST_CTX, + }; + ds4_engine *engine = NULL; + if (ds4_engine_open(&engine, &opt) != 0) fail("engine open", NULL); + + ds4_tokens prompt = {0}; + ds4_encode_chat_prompt(engine, NULL, "Reply with exactly: OK", + DS4_THINK_NONE, &prompt); + + ds4_session *source = NULL; + ds4_session *restored = NULL; + char err[256] = {0}; + if (ds4_session_create(&source, engine, TEST_CTX) != 0 || + ds4_session_create(&restored, engine, TEST_CTX) != 0) { + fail("session create", NULL); + } + if (ds4_session_sync(source, &prompt, err, sizeof(err)) != 0) { + fail("source prefill", err); + } + + ds4_session_snapshot snapshot = {0}; + if (ds4_session_save_snapshot(source, &snapshot, err, sizeof(err)) != 0) { + fail("snapshot save", err); + } + if (ds4_session_load_snapshot(restored, &snapshot, err, sizeof(err)) != 0) { + fail("snapshot load", err); + } + if (ds4_session_pos(source) != ds4_session_pos(restored)) { + fail("checkpoint position changed across snapshot round-trip", NULL); + } + + const int vocab = ds4_engine_vocab_size(engine); + float *source_logits = malloc((size_t)vocab * sizeof(float)); + float *restored_logits = malloc((size_t)vocab * sizeof(float)); + if (!source_logits || !restored_logits) fail("logit allocation", NULL); + compare_logits(source, restored, source_logits, restored_logits, + vocab, "prefill"); + + const int token = ds4_session_argmax(source); + if (ds4_session_eval(source, token, err, sizeof(err)) != 0) { + fail("source decode", err); + } + if (ds4_session_eval(restored, token, err, sizeof(err)) != 0) { + fail("restored decode", err); + } + compare_logits(source, restored, source_logits, restored_logits, + vocab, "decode"); + + fprintf(stderr, + "test_session_snapshot: OK tokens=%d payload=%llu bytes\n", + ds4_session_pos(source), + (unsigned long long)snapshot.len); + + free(source_logits); + free(restored_logits); + ds4_session_snapshot_free(&snapshot); + ds4_session_free(source); + ds4_session_free(restored); + ds4_tokens_free(&prompt); + ds4_engine_close(engine); + return 0; +}