/* ========================================================================= * ds4.c - DeepSeek V4 inference engine. * ========================================================================= * * 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. * * Loading is mmap based. The loader parses only the GGUF header, metadata * table, and tensor directory. Tensor data stays in the kernel page cache * until inference touches it, or until Metal wraps slices of the mapping as * no-copy MTLBuffers. */ #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #if defined(__APPLE__) #include #endif #include #include #include #include "ds4.h" #include "ds4_distributed.h" #include "ds4_tp.h" /* Wave-2 multi-GPU types are needed in every build because the engine * struct embeds ds4_gpu_config and the placement table. ds4_layer_pack.h * is included unconditionally for the same reason (engine helpers call * the packer in multi-tier mode, but the headers are tiny and C-safe). */ #include "ds4_layer_pack.h" #include "ds4_gpu_mgpu.h" #define DS4_CUDA_TP_PEER_TMP_BYTES \ ((uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float) + 128u) static uint32_t metal_graph_cuda_tp_output_requested_ways(void) { const char *env = getenv("DS4_CUDA_TP_OUTPUT_WAYS"); if (!env || !env[0]) return 8; char *end = NULL; unsigned long v = strtoul(env, &end, 10); if (end != env && *end == '\0' && v >= 2 && v <= DS4_MAX_GPUS) { return (uint32_t)v; } fprintf(stderr, "ds4: invalid DS4_CUDA_TP_OUTPUT_WAYS=%s; expected 2..%d\n", env, DS4_MAX_GPUS); return 2; } static uint32_t metal_graph_cuda_tp_output_tiers_for_head( int head_tier, bool cuda_tp_output, int n_gpus, int tiers[DS4_MAX_GPUS]) { if (!tiers || !cuda_tp_output || n_gpus < 2 || (n_gpus & 1) != 0) return 0; if (head_tier < 0 || head_tier >= n_gpus) return 0; const int half = n_gpus / 2; if (head_tier >= half) return 0; const int partner = head_tier + half; if (partner >= n_gpus) return 0; uint32_t want = metal_graph_cuda_tp_output_requested_ways(); if (want > (uint32_t)n_gpus) want = (uint32_t)n_gpus; if (want < 2u) want = 2u; uint32_t n = 0; tiers[n++] = head_tier; if (partner != head_tier && n < want) tiers[n++] = partner; for (int pass = 0; pass < 2 && n < want; pass++) { const int start = pass == 0 ? half : 0; const int end = pass == 0 ? n_gpus : half; for (int t = start; t < end && n < want; t++) { bool seen = false; for (uint32_t i = 0; i < n; i++) { if (tiers[i] == t) { seen = true; break; } } if (!seen) tiers[n++] = t; } } return n; } #ifndef DS4_NO_GPU #include "ds4_gpu.h" #endif /* Non-CUDA builds (Mac/Metal, CPU-only) never link ds4_cuda.cu. Provide * stubs for the multi-GPU plumbing multi-GPU functions and globals declared in * ds4_gpu_mgpu.h. These keep the linker happy on Mac/Metal and on CPU * builds. None of these are reached at runtime in non-CUDA builds * because multi_tier == 1 requires graph_backend == true which on * non-CUDA paths means Metal, and the multi-tier branch only fires * when the caller supplied a non-NULL gpu_cfg — which only the new * ds4_engine_create_with_gpu_config can do, and no Metal caller does. * * We key off __APPLE__ + DS4_NO_GPU rather than the inverse of CUDA * because there's no positive "is CUDA build" macro and ds4_cuda.cu * is only compiled in the non-Apple, non-DS4_NO_GPU configuration. */ /* Apple (Metal) build: ds4_cuda.cu is not linked, but the engine compiles * code referencing some multi-GPU CLI symbols inside dead multi-tier branches. * Provide stubs so the linker is happy. CPU-only (DS4_NO_GPU) builds * never reach the multi-tier branches and never include ds4_gpu.h, so * we cannot reference ds4_tensor_range there — those stubs are guarded * by !DS4_NO_GPU below. */ #if defined(__APPLE__) && !defined(DS4_NO_GPU) int ds4_gpu_set_current_device(int logical_tier) { (void)logical_tier; return -1; } int ds4_gpu_set_current_device_fenced(int logical_tier) { (void)logical_tier; return -1; } void ds4_gpu_enable_q8_dequant_gemm(void) {} int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst, const ds4_gpu_tensor *src, uint64_t bytes) { (void)dst; (void)src; (void)bytes; return 0; } int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst, const ds4_gpu_tensor *src, uint64_t bytes) { return ds4_gpu_tensor_copy(dst, 0, src, 0, bytes); } int ds4_gpu_tensor_copy_xdev3_default_dst( ds4_gpu_tensor *dst0, const ds4_gpu_tensor *src0, uint64_t bytes0, ds4_gpu_tensor *dst1, const ds4_gpu_tensor *src1, uint64_t bytes1, ds4_gpu_tensor *dst2, const ds4_gpu_tensor *src2, uint64_t bytes2) { return (bytes0 == 0 || ds4_gpu_tensor_copy(dst0, 0, src0, 0, bytes0)) && (bytes1 == 0 || ds4_gpu_tensor_copy(dst1, 0, src1, 0, bytes1)) && (bytes2 == 0 || ds4_gpu_tensor_copy(dst2, 0, src2, 0, bytes2)); } int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size) { (void)model_map; (void)model_size; return 0; } int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias) { (void)map; (void)size; (void)bias; return 1; } int ds4_gpu_device_cache_support_tensors(int device_id, int exec_device_id, const ds4_tensor_range *ranges, int n_ranges, int main_model) { (void)device_id; (void)exec_device_id; (void)ranges; (void)n_ranges; (void)main_model; return 0; } uint64_t ds4_gpu_tier_free_vram(int logical_tier) { (void)logical_tier; return 0; } int ds4_gpu_device_cache_tensors(int device_id, const ds4_tensor_range *ranges, int n_ranges) { (void)device_id; (void)ranges; (void)n_ranges; return 1; } int ds4_gpu_init_multi(const ds4_gpu_config *cfg) { (void)cfg; return -1; } int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, uint64_t bytes) { (void)t; (void)device_id; (void)bytes; return -1; } ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes) { /* Metal / CPU build has no CUDA multi-tier — short-circuit tier 0 to * the legacy single-device allocator so the per-tier graph * allocation path (e.g. g->output_pre_by_tier[head_tier]) keeps * working byte-equivalent to pre-multi-tier Metal. Without this, * every multi-tier-aware allocation in metal_graph_alloc_raw_cap * returns NULL, the validation chain fails, and session_create * silently returns 1. */ if (tier == 0) return ds4_gpu_tensor_alloc(bytes); return NULL; } ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes) { if (tier == 0) return ds4_gpu_tensor_alloc_managed(bytes); return NULL; } void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t) { (void)t; } 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; } 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; } 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; } int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) { (void)src; (void)dst_tier; return -1; } 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; } int ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor( ds4_gpu_tensor *out0, ds4_gpu_tensor *out1, const void *model_map, uint64_t model_size, uint64_t weight0_offset, uint64_t weight1_offset, uint64_t in_dim, uint64_t out0_dim, uint64_t out1_dim, const ds4_gpu_tensor *x, uint32_t n_rows) { return ds4_gpu_matmul_q8_0_pair_tensor( out0, out1, model_map, model_size, weight0_offset, weight1_offset, in_dim, out0_dim, out1_dim, x, n_rows); } int ds4_gpu_matmul_f16_router_rows_exact_tensor( ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, const ds4_gpu_tensor *x, uint32_t n_rows) { return ds4_gpu_matmul_f16_tensor( out, model_map, model_size, weight_offset, 4096u, 256u, x, n_rows); } int ds4_gpu_q8_cache_suppressed(void) { return 0; } void ds4_gpu_set_q8_cache_suppressed(int suppressed) { (void)suppressed; } int ds4_gpu_set_decode_fast_attention(int enabled) { (void)enabled; return 0; } int ds4_gpu_set_decode_score_vec4(int enabled) { (void)enabled; return 0; } 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) { (void)selected; (void)values; (void)scores; (void)n_comp; (void)n_tokens; (void)index_offset; return 0; } 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) { (void)selected; (void)values; (void)model_map; (void)model_size; (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; (void)index_offset; return 0; } 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; } 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) { (void)n_rot; (void)pos0; (void)n_ctx_orig; (void)freq_base; (void)freq_scale; (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; if (fused_inv_rope) *fused_inv_rope = 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); } int ds4_gpu_tensor_device(const ds4_gpu_tensor *t) { (void)t; return -1; } ds4_gpu_ctx g_gpu[DS4_MAX_GPUS]; int g_n_gpus = 0; int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS]; #endif #if defined(DS4_NO_GPU) /* CPU-only build: even though no multi-tier code is reached, the engine * struct still embeds ds4_gpu_config and the global decls in * ds4_gpu_mgpu.h need matching definitions. We do not stub the * function symbols here because no caller references them in CPU * builds (every callsite is inside !DS4_NO_GPU). */ ds4_gpu_ctx g_gpu[DS4_MAX_GPUS]; int g_n_gpus = 0; int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS]; #endif #if defined(__ARM_NEON) #include #endif #ifndef M_PI #define M_PI 3.14159265358979323846 #endif #define DS4_NEG_INF (-1.0e30f) #define DS4_POS_INF ( 1.0e30f) #define DS4_DEFAULT_RMS_EPS ( 1.0e-6f) #define DS4_DEFAULT_HC_EPS ( 1.0e-6f) #define DS4_DEFAULT_SWIGLU_CLAMP_EXP (10.0f) #define DS4_DEFAULT_ROPE_FREQ_BASE (10000.0f) #define DS4_DEFAULT_ROPE_SCALE_FACTOR (16.0f) #define DS4_DEFAULT_ROPE_YARN_BETA_FAST (32.0f) #define DS4_DEFAULT_ROPE_YARN_BETA_SLOW (1.0f) #define DS4_DEFAULT_COMPRESS_ROPE_FREQ_BASE (160000.0f) #define DS4_DEFAULT_ROPE_ORIG_CTX UINT64_C(65536) static const char DS4_REASONING_EFFORT_MAX_PREFIX[] = "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"; /* DeepSeek recommends Think Max only with at least a 384K-token context window. * Below that size we keep ordinary thinking to avoid injecting a prompt that * asks for a reasoning budget the allocated context is not meant to hold. */ #define DS4_THINK_MAX_MIN_CONTEXT 393216u static bool ds4_backend_uses_graph(ds4_backend backend) { return backend == DS4_BACKEND_METAL || backend == DS4_BACKEND_CUDA; } static bool ds4_backend_supports_ssd_streaming(ds4_backend backend) { if (backend == DS4_BACKEND_METAL) return true; if (backend == DS4_BACKEND_CUDA) { #if defined(DS4_ROCM_BUILD) || (!defined(DS4_NO_GPU) && !defined(__APPLE__)) return true; #else return false; #endif } return false; } static bool ds4_backend_supports_streaming_auto_cache(ds4_backend backend) { if (backend == DS4_BACKEND_METAL) return true; #ifdef DS4_ROCM_BUILD if (backend == DS4_BACKEND_CUDA) return true; #else (void)backend; #endif return false; } static bool ds4_backend_supports_glm_streaming_full_layers(ds4_backend backend) { if (backend == DS4_BACKEND_METAL) return true; #ifdef DS4_ROCM_BUILD if (backend == DS4_BACKEND_CUDA) return true; #else (void)backend; #endif return false; } static bool glm_graph_env_present(const char *rocm_name, const char *metal_name) { #ifdef DS4_ROCM_BUILD if (rocm_name && getenv(rocm_name) != NULL) return true; #else (void)rocm_name; #endif return metal_name && getenv(metal_name) != NULL; } static const char *glm_graph_env_value(const char *rocm_name, const char *metal_name) { #ifdef DS4_ROCM_BUILD const char *rocm_env = rocm_name ? getenv(rocm_name) : NULL; if (rocm_env && rocm_env[0]) return rocm_env; #else (void)rocm_name; #endif const char *metal_env = metal_name ? getenv(metal_name) : NULL; return (metal_env && metal_env[0]) ? metal_env : NULL; } /* ========================================================================= * Model Shape Profiles. * ========================================================================= * * The weight binder and metadata validator select one of the known model * profiles below. Arrays reserve the maximum Pro dimensions; hot loops read * the active profile after GGUF validation. */ enum { DS4_MAX_LAYER = 79, DS4_MAX_EMBD = 7168, DS4_MAX_VOCAB = 154880, DS4_MAX_HEAD = 128, DS4_MAX_HEAD_KV = 1, DS4_MAX_HEAD_DIM = 576, DS4_MAX_VALUE_DIM = 512, DS4_MAX_ROT = 64, DS4_MAX_OUT_GROUP = 16, DS4_MAX_LORA_Q = 2048, DS4_MAX_LORA_O = 1024, DS4_MAX_EXPERT = 384, DS4_MAX_EXPERT_USED = 8, DS4_MAX_EXPERT_SHARED = 1, DS4_MAX_FF_EXP = 3072, DS4_MAX_HASH_LAYER = 3, DS4_MAX_SWA = 128, DS4_MAX_INDEXER_HEAD = 64, DS4_MAX_INDEXER_HEAD_DIM = 128, DS4_MAX_INDEXER_TOP_K = 2048, DS4_MAX_HC = 4, DS4_MAX_HC_SINKHORN_ITER = 20, }; typedef enum { DS4_MODEL_FAMILY_DEEPSEEK4 = 0, DS4_MODEL_FAMILY_GLM_DSA = 1, } ds4_model_family; typedef enum { DS4_VARIANT_FLASH = 0, DS4_VARIANT_PRO = 1, DS4_VARIANT_GLM52 = 2, } ds4_variant; typedef struct { const char *name; ds4_model_family family; ds4_variant variant; 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_out_group; uint32_t n_lora_q; uint32_t n_lora_o; uint32_t n_expert; uint32_t n_expert_used; uint32_t n_expert_shared; uint32_t n_ff_exp; uint32_t n_ff_dense; 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; uint32_t n_nextn_predict; uint32_t n_leading_dense; uint32_t n_kv_lora; uint32_t n_key_mla; uint32_t n_value_mla; float rms_eps; float hc_eps; float expert_weight_scale; float swiglu_clamp_exp; float rope_freq_base; float rope_scale_factor; float rope_yarn_beta_fast; float rope_yarn_beta_slow; float compress_rope_freq_base; uint64_t rope_orig_ctx; } ds4_shape; static const ds4_shape DS4_SHAPE_FLASH = { .name = "DeepSeek V4 Flash", .family = DS4_MODEL_FAMILY_DEEPSEEK4, .variant = DS4_VARIANT_FLASH, .n_layer = 43, .n_embd = 4096, .n_vocab = 129280, .n_head = 64, .n_head_kv = 1, .n_head_dim = 512, .n_value_dim = 512, .n_rot = 64, .n_out_group = 8, .n_lora_q = 1024, .n_lora_o = 1024, .n_expert = 256, .n_expert_used = 6, .n_expert_shared = 1, .n_ff_exp = 2048, .n_hash_layer = 3, .n_swa = 128, .n_indexer_head = 64, .n_indexer_head_dim = 128, .n_indexer_top_k = 512, .n_hc = 4, .n_hc_sinkhorn_iter = 20, .rms_eps = DS4_DEFAULT_RMS_EPS, .hc_eps = DS4_DEFAULT_HC_EPS, .expert_weight_scale = 1.5f, .swiglu_clamp_exp = DS4_DEFAULT_SWIGLU_CLAMP_EXP, .rope_freq_base = DS4_DEFAULT_ROPE_FREQ_BASE, .rope_scale_factor = DS4_DEFAULT_ROPE_SCALE_FACTOR, .rope_yarn_beta_fast = DS4_DEFAULT_ROPE_YARN_BETA_FAST, .rope_yarn_beta_slow = DS4_DEFAULT_ROPE_YARN_BETA_SLOW, .compress_rope_freq_base = DS4_DEFAULT_COMPRESS_ROPE_FREQ_BASE, .rope_orig_ctx = DS4_DEFAULT_ROPE_ORIG_CTX, }; static const ds4_shape DS4_SHAPE_PRO = { .name = "DeepSeek V4 Pro", .family = DS4_MODEL_FAMILY_DEEPSEEK4, .variant = DS4_VARIANT_PRO, .n_layer = 61, .n_embd = 7168, .n_vocab = 129280, .n_head = 128, .n_head_kv = 1, .n_head_dim = 512, .n_value_dim = 512, .n_rot = 64, .n_out_group = 16, .n_lora_q = 1536, .n_lora_o = 1024, .n_expert = 384, .n_expert_used = 6, .n_expert_shared = 1, .n_ff_exp = 3072, .n_hash_layer = 3, .n_swa = 128, .n_indexer_head = 64, .n_indexer_head_dim = 128, .n_indexer_top_k = 1024, .n_hc = 4, .n_hc_sinkhorn_iter = 20, .rms_eps = DS4_DEFAULT_RMS_EPS, .hc_eps = DS4_DEFAULT_HC_EPS, .expert_weight_scale = 2.5f, .swiglu_clamp_exp = DS4_DEFAULT_SWIGLU_CLAMP_EXP, .rope_freq_base = DS4_DEFAULT_ROPE_FREQ_BASE, .rope_scale_factor = DS4_DEFAULT_ROPE_SCALE_FACTOR, .rope_yarn_beta_fast = DS4_DEFAULT_ROPE_YARN_BETA_FAST, .rope_yarn_beta_slow = DS4_DEFAULT_ROPE_YARN_BETA_SLOW, .compress_rope_freq_base = DS4_DEFAULT_COMPRESS_ROPE_FREQ_BASE, .rope_orig_ctx = DS4_DEFAULT_ROPE_ORIG_CTX, }; static const ds4_shape DS4_SHAPE_GLM52 = { .name = "GLM 5.2", .family = DS4_MODEL_FAMILY_GLM_DSA, .variant = DS4_VARIANT_GLM52, .n_layer = 79, .n_embd = 6144, .n_vocab = 154880, .n_head = 64, .n_head_kv = 1, .n_head_dim = 576, .n_value_dim = 512, .n_rot = 64, .n_out_group = 0, .n_lora_q = 2048, .n_lora_o = 0, .n_expert = 256, .n_expert_used = 8, .n_expert_shared = 1, .n_ff_exp = 2048, .n_ff_dense = 12288, .n_hash_layer = 0, .n_swa = 0, .n_indexer_head = 32, .n_indexer_head_dim = 128, .n_indexer_top_k = 2048, .n_hc = 0, .n_hc_sinkhorn_iter = 0, .n_nextn_predict = 1, .n_leading_dense = 3, .n_kv_lora = 512, .n_key_mla = 256, .n_value_mla = 256, .rms_eps = 1.0e-5f, .hc_eps = 0.0f, .expert_weight_scale = 2.5f, .swiglu_clamp_exp = 0.0f, .rope_freq_base = 8000000.0f, .rope_scale_factor = 1.0f, .rope_yarn_beta_fast = 0.0f, .rope_yarn_beta_slow = 0.0f, .compress_rope_freq_base = 0.0f, .rope_orig_ctx = 1048576, }; static ds4_shape g_ds4_shape = { .name = "DeepSeek V4 Flash", .family = DS4_MODEL_FAMILY_DEEPSEEK4, .variant = DS4_VARIANT_FLASH, .n_layer = 43, .n_embd = 4096, .n_vocab = 129280, .n_head = 64, .n_head_kv = 1, .n_head_dim = 512, .n_value_dim = 512, .n_rot = 64, .n_out_group = 8, .n_lora_q = 1024, .n_lora_o = 1024, .n_expert = 256, .n_expert_used = 6, .n_expert_shared = 1, .n_ff_exp = 2048, .n_hash_layer = 3, .n_swa = 128, .n_indexer_head = 64, .n_indexer_head_dim = 128, .n_indexer_top_k = 512, .n_hc = 4, .n_hc_sinkhorn_iter = 20, .rms_eps = DS4_DEFAULT_RMS_EPS, .hc_eps = DS4_DEFAULT_HC_EPS, .expert_weight_scale = 1.5f, .swiglu_clamp_exp = DS4_DEFAULT_SWIGLU_CLAMP_EXP, .rope_freq_base = DS4_DEFAULT_ROPE_FREQ_BASE, .rope_scale_factor = DS4_DEFAULT_ROPE_SCALE_FACTOR, .rope_yarn_beta_fast = DS4_DEFAULT_ROPE_YARN_BETA_FAST, .rope_yarn_beta_slow = DS4_DEFAULT_ROPE_YARN_BETA_SLOW, .compress_rope_freq_base = DS4_DEFAULT_COMPRESS_ROPE_FREQ_BASE, .rope_orig_ctx = DS4_DEFAULT_ROPE_ORIG_CTX, }; static uint32_t g_ds4_compress_ratios[DS4_MAX_LAYER] = {0}; #define DS4_MODEL_SHAPE_NAME (g_ds4_shape.name) #define DS4_MODEL_FAMILY (g_ds4_shape.family) #define DS4_MODEL_VARIANT (g_ds4_shape.variant) #define DS4_N_LAYER (g_ds4_shape.n_layer) #define DS4_N_EMBD (g_ds4_shape.n_embd) #define DS4_N_VOCAB (g_ds4_shape.n_vocab) #define DS4_N_HEAD (g_ds4_shape.n_head) #define DS4_N_HEAD_KV (g_ds4_shape.n_head_kv) #define DS4_N_HEAD_DIM (g_ds4_shape.n_head_dim) #define DS4_N_VALUE_DIM (g_ds4_shape.n_value_dim) #define DS4_N_ROT (g_ds4_shape.n_rot) #define DS4_N_OUT_GROUP (g_ds4_shape.n_out_group) #define DS4_N_LORA_Q (g_ds4_shape.n_lora_q) #define DS4_N_LORA_O (g_ds4_shape.n_lora_o) #define DS4_N_EXPERT (g_ds4_shape.n_expert) #define DS4_N_EXPERT_USED (g_ds4_shape.n_expert_used) #define DS4_N_EXPERT_SHARED (g_ds4_shape.n_expert_shared) #define DS4_N_FF_EXP (g_ds4_shape.n_ff_exp) #define DS4_N_FF_DENSE (g_ds4_shape.n_ff_dense) #define DS4_N_HASH_LAYER (g_ds4_shape.n_hash_layer) #define DS4_N_SWA (g_ds4_shape.n_swa) #define DS4_N_INDEXER_HEAD (g_ds4_shape.n_indexer_head) #define DS4_N_INDEXER_HEAD_DIM (g_ds4_shape.n_indexer_head_dim) #define DS4_N_INDEXER_TOP_K (g_ds4_shape.n_indexer_top_k) #define DS4_N_HC (g_ds4_shape.n_hc) #define DS4_N_HC_SINKHORN_ITER (g_ds4_shape.n_hc_sinkhorn_iter) #define DS4_N_NEXTN_PREDICT (g_ds4_shape.n_nextn_predict) #define DS4_N_LEADING_DENSE (g_ds4_shape.n_leading_dense) #define DS4_N_KV_LORA (g_ds4_shape.n_kv_lora) #define DS4_N_KEY_MLA (g_ds4_shape.n_key_mla) #define DS4_N_VALUE_MLA (g_ds4_shape.n_value_mla) #define DS4_RMS_EPS (g_ds4_shape.rms_eps) #define DS4_HC_EPS (g_ds4_shape.hc_eps) #define DS4_EXPERT_WEIGHT_SCALE (g_ds4_shape.expert_weight_scale) #define DS4_SWIGLU_CLAMP_EXP (g_ds4_shape.swiglu_clamp_exp) #define DS4_ROPE_FREQ_BASE (g_ds4_shape.rope_freq_base) #define DS4_ROPE_SCALE_FACTOR (g_ds4_shape.rope_scale_factor) #define DS4_ROPE_YARN_BETA_FAST (g_ds4_shape.rope_yarn_beta_fast) #define DS4_ROPE_YARN_BETA_SLOW (g_ds4_shape.rope_yarn_beta_slow) #define DS4_COMPRESS_ROPE_FREQ_BASE (g_ds4_shape.compress_rope_freq_base) #define DS4_ROPE_ORIG_CTX (g_ds4_shape.rope_orig_ctx) static int g_ds4_lock_fd = -1; #if defined(__GNUC__) || defined(__clang__) #define DS4_MAYBE_UNUSED __attribute__((unused)) #else #define DS4_MAYBE_UNUSED #endif /* ========================================================================= * GGUF Quant Block Formats. * ========================================================================= * * These layouts and IQ2 tables match the GGUF quantized tensor format, * reduced to only the formats ds4.c currently reads or sizes: * - Q2_K routed down experts * - Q4_K routed experts in the high-memory variant * - Q5_K/Q6_K GLM routed experts * - IQ2_XXS routed gate/up experts * - Q8_K temporary activation blocks for dot products */ #define QK_K 256 typedef struct { uint8_t scales[QK_K / 16]; uint8_t qs[QK_K / 4]; uint16_t d; uint16_t dmin; } block_q2_K; typedef struct { uint16_t d; uint16_t dmin; uint8_t scales[12]; uint8_t qs[QK_K / 2]; } block_q4_K; typedef struct { uint16_t d; uint16_t dmin; uint8_t scales[12]; uint8_t qh[QK_K / 8]; uint8_t qs[QK_K / 2]; } block_q5_K; typedef struct { uint8_t ql[QK_K / 2]; uint8_t qh[QK_K / 4]; int8_t scales[QK_K / 16]; uint16_t d; } block_q6_K; typedef struct { float d; int8_t qs[QK_K]; int16_t bsums[QK_K / 16]; } block_q8_K; typedef struct { uint16_t d; uint16_t qs[QK_K / 8]; } block_iq2_xxs; #define DS4_STATIC_ASSERT(name, cond) typedef char name[(cond) ? 1 : -1] DS4_STATIC_ASSERT(ds4_block_q2_k_size, sizeof(block_q2_K) == 84); DS4_STATIC_ASSERT(ds4_block_q4_k_size, sizeof(block_q4_K) == 144); DS4_STATIC_ASSERT(ds4_block_q5_k_size, sizeof(block_q5_K) == 176); DS4_STATIC_ASSERT(ds4_block_q6_k_size, sizeof(block_q6_K) == 210); DS4_STATIC_ASSERT(ds4_block_q8_k_size, sizeof(block_q8_K) == 292); DS4_STATIC_ASSERT(ds4_block_iq2_xxs_size, sizeof(block_iq2_xxs) == 66); typedef struct { uint32_t ctx_size; uint32_t comp_cap; uint32_t attn_score_cap; uint32_t q8_cap; float *plain; float *cur; float *next; float *attn_cur; float *attn_norm; float *attn_residual; float *q; float *qr; float *qr_norm; float *kv_raw; float *kv; float *heads; float *attn_low; float *attn_out; float *after_attn_hc; float *attn_score; float *comp; float *index_comp; float *comp_kv_cur; float *comp_sc_cur; float *comp_pooled; bool *index_allowed; float *index_q; float *index_weights; float *index_scores; float *ffn_cur; float *ffn_norm; float *ffn_moe; float *ffn_shared; float *ffn_out; float *shared_gate; float *shared_up; float *shared_mid; float *routed_mid_all; block_q8_K *routed_xq; block_q8_K *routed_midq; int8_t *routed_q8_xq; float *routed_q8_xscale; int8_t *routed_q8_midq; float *routed_q8_midscale; int8_t *q8_xq; float *q8_xscale; float *hc_flat; float *output_flat; float *output_pre; float *output_weights; float *output_embd; float *output_norm; } ds4_cpu_decode_scratch; static const uint8_t kmask_iq2xs[8] = { 1, 2, 4, 8, 16, 32, 64, 128 }; static const uint8_t ksigns_iq2xs[128] = { 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, }; static const uint64_t iq2xxs_grid[256] = { 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, }; static int8_t iq2xxs_signed_grid[256][128][8]; static int8_t iq2xxs_signs[128][8]; static pthread_once_t iq2xxs_signed_grid_once = PTHREAD_ONCE_INIT; static void iq2xxs_signed_grid_init(void) { for (uint32_t s = 0; s < 128; s++) { const uint8_t signs = ksigns_iq2xs[s]; for (uint32_t j = 0; j < 8; j++) { iq2xxs_signs[s][j] = (int8_t)((signs & kmask_iq2xs[j]) ? -1 : 1); } } for (uint32_t g = 0; g < 256; g++) { const uint8_t *grid = (const uint8_t *)(iq2xxs_grid + g); for (uint32_t s = 0; s < 128; s++) { const uint8_t signs = ksigns_iq2xs[s]; for (uint32_t j = 0; j < 8; j++) { const int v = (int)grid[j]; iq2xxs_signed_grid[g][s][j] = (int8_t)((signs & kmask_iq2xs[j]) ? -v : v); } } } } static inline DS4_MAYBE_UNUSED int32_t dot_iq2_pair_16(const int8_t *grid0, const int8_t *grid1, const int8_t *q8) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) const int8x16_t gv = vcombine_s8(vld1_s8(grid0), vld1_s8(grid1)); const int32x4_t acc = vdotq_s32(vdupq_n_s32(0), gv, vld1q_s8(q8)); return vaddvq_s32(acc); #elif defined(__ARM_NEON) const int8x16_t gv = vcombine_s8(vld1_s8(grid0), vld1_s8(grid1)); const int8x16_t qv = vld1q_s8(q8); const int16x8_t p0 = vmull_s8(vget_low_s8(gv), vget_low_s8(qv)); const int16x8_t p1 = vmull_s8(vget_high_s8(gv), vget_high_s8(qv)); return vaddvq_s32(vaddq_s32(vpaddlq_s16(p0), vpaddlq_s16(p1))); #else int32_t sum = 0; for (uint32_t i = 0; i < 8; i++) sum += (int32_t)grid0[i] * (int32_t)q8[i]; for (uint32_t i = 0; i < 8; i++) sum += (int32_t)grid1[i] * (int32_t)q8[8 + i]; return sum; #endif } static inline DS4_MAYBE_UNUSED int32_t dot_q2_16(const uint8_t *q2, const int8_t *q8, int shift) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) const uint8x16_t packed = vld1q_u8(q2); uint8x16_t shifted; switch (shift) { case 0: shifted = packed; break; case 2: shifted = vshrq_n_u8(packed, 2); break; case 4: shifted = vshrq_n_u8(packed, 4); break; default: shifted = vshrq_n_u8(packed, 6); break; } const uint8x16_t vals_u = vandq_u8(shifted, vdupq_n_u8(3)); const int8x16_t vals = vreinterpretq_s8_u8(vals_u); const int8x16_t q8v = vld1q_s8(q8); const int32x4_t acc = vdotq_s32(vdupq_n_s32(0), q8v, vals); return vaddvq_s32(acc); #elif defined(__ARM_NEON) uint8_t vals_tmp[16]; for (uint32_t i = 0; i < 16; i++) vals_tmp[i] = (q2[i] >> shift) & 3; const int8x16_t vals = vreinterpretq_s8_u8(vld1q_u8(vals_tmp)); const int8x16_t q8v = vld1q_s8(q8); const int16x8_t p0 = vmull_s8(vget_low_s8(q8v), vget_low_s8(vals)); const int16x8_t p1 = vmull_s8(vget_high_s8(q8v), vget_high_s8(vals)); const int32x4_t s0 = vpaddlq_s16(p0); const int32x4_t s1 = vpaddlq_s16(p1); return vaddvq_s32(vaddq_s32(s0, s1)); #else int32_t sum = 0; for (uint32_t i = 0; i < 16; i++) sum += (int32_t)q8[i] * (int32_t)((q2[i] >> shift) & 3); return sum; #endif } /* ========================================================================= * Shared Helpers, Allocation Guards, Threads, and Cursor Reads. * ========================================================================= * * This section holds process-wide utilities used by all later stages: * fatal-error helpers, allocation wrappers, the persistent CPU worker pool, * and the small byte cursor used to parse GGUF metadata. */ #define DS4_GGUF_MAGIC 0x46554747u /* "GGUF", little endian. */ #define DS4_MAX_DIMS 8 typedef struct { const char *ptr; uint64_t len; } ds4_str; typedef ds4_tokens token_vec; typedef struct { const uint8_t *base; uint64_t size; uint64_t pos; char error[256]; } ds4_cursor; static void ds4_die(const char *msg) { fprintf(stderr, "ds4: %s\n", msg); exit(1); } /* Attention compression is read from GGUF metadata after validating that it * matches the exact layout expected for the loaded model shape. */ static uint32_t ds4_layer_compress_ratio(uint32_t il) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) return 0; if (il >= DS4_N_LAYER) ds4_die("DeepSeek4 layer index is outside the loaded model layout"); return g_ds4_compress_ratios[il]; } static uint32_t ds4_expected_layer_compress_ratio(uint32_t il) { if (il >= DS4_N_LAYER) ds4_die("DeepSeek4 layer index is outside the loaded model layout"); switch (DS4_MODEL_VARIANT) { case DS4_VARIANT_FLASH: if (il < 2) return 0; return (il & 1u) == 0 ? 4u : 128u; case DS4_VARIANT_PRO: if (il < 2) return 128u; return (il & 1u) == 0 ? 4u : 128u; default: ds4_die("unsupported DeepSeek4 model variant"); } return 0; } static void ds4_die_errno(const char *what, const char *path) { fprintf(stderr, "ds4: %s '%s': %s\n", what, path, strerror(errno)); exit(1); } static bool ds4_streq(ds4_str s, const char *z) { size_t n = strlen(z); return s.len == n && memcmp(s.ptr, z, n) == 0; } static bool ds4_str_starts_with(ds4_str s, const char *prefix) { size_t n = strlen(prefix); return s.len >= n && memcmp(s.ptr, prefix, n) == 0; } static bool ds4_str_contains(ds4_str s, const char *needle) { size_t n = strlen(needle); if (n == 0) return true; if (s.len < n) return false; for (uint64_t i = 0; i <= s.len - n; i++) { if (memcmp(s.ptr + i, needle, n) == 0) return true; } return false; } static bool ds4_str_eq(ds4_str a, ds4_str b) { return a.len == b.len && memcmp(a.ptr, b.ptr, a.len) == 0; } static uint64_t hash_bytes(const void *ptr, uint64_t len) { const uint8_t *p = ptr; uint64_t h = 1469598103934665603ull; for (uint64_t i = 0; i < len; i++) { h ^= p[i]; h *= 1099511628211ull; } return h; } static bool g_alloc_guard_enabled; static const char *g_alloc_guard_phase; static void ds4_alloc_guard_begin(const char *phase) { g_alloc_guard_phase = phase; g_alloc_guard_enabled = true; } static void ds4_alloc_guard_end(void) { g_alloc_guard_enabled = false; g_alloc_guard_phase = NULL; } static void ds4_alloc_guard_check(const char *op, size_t size) { if (!g_alloc_guard_enabled) return; fprintf(stderr, "ds4: internal allocation during %s: %s(%zu). " "CPU decode is expected to reuse preallocated scratch buffers.\n", g_alloc_guard_phase ? g_alloc_guard_phase : "guarded phase", op, size); exit(1); } static void *xcalloc(size_t n, size_t size) { ds4_alloc_guard_check("calloc", n * size); void *p = calloc(n, size); if (!p) ds4_die("out of memory"); return p; } static void *xmalloc(size_t size) { ds4_alloc_guard_check("malloc", size); void *p = malloc(size); if (!p) ds4_die("out of memory"); return p; } static char *ds4_strdup(const char *s) { size_t n = strlen(s); char *p = xmalloc(n + 1); memcpy(p, s, n + 1); return p; } static void *xrealloc(void *ptr, size_t size) { ds4_alloc_guard_check("realloc", size); void *p = realloc(ptr, size); if (!p) ds4_die("out of memory"); return p; } static void *xmalloc_zeroed(size_t n, size_t size) { if (size != 0 && n > SIZE_MAX / size) ds4_die("allocation size overflow"); const size_t total = n * size; void *p = xmalloc(total ? total : 1); /* * This is intentionally not calloc(). Large untouched calloc ranges may be * represented by the VM through shared zero-page bookkeeping. The CPU decode * KV cache grows one token at a time, so using calloc here can move thousands * of first-touch faults into generation. On Darwin we have observed this end * in a kernel cpt_mapcnt_inc overflow panic instead of a user-space error. * * Explicitly writing the zeroes while the cache is allocated keeps those VM * faults out of the token loop and gives the cache private resident pages. */ memset(p, 0, total); return p; } static double now_sec(void) { struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts); return (double)ts.tv_sec + (double)ts.tv_nsec * 1.0e-9; } /* ========================================================================= * Metal Routed Expert Locality Profiler. * ========================================================================= * * SSD streaming decode needs routed-expert locality data before a cache policy * is meaningful. This profiler is intentionally Metal-decode-only for now: * it reads the tiny selected-expert and route-weight tensors after router * selection, records histograms, and simulates per-layer latest-N unique expert * caches without changing normal inference when disabled. */ enum { DS4_EXPERT_PROFILE_MAX_CAPS = 10 }; static const uint32_t ds4_expert_profile_cap_candidates[DS4_EXPERT_PROFILE_MAX_CAPS] = { 1, 2, 4, 8, 16, 32, 64, 128, 256, 384 }; typedef struct { int expert_id; uint64_t count; double weight; } ds4_expert_profile_sort_entry; typedef struct { bool active; char *path; char *hotlist_path; char model_name[64]; uint32_t n_layer; uint32_t n_expert; uint32_t n_expert_used; uint32_t n_caps; uint32_t caps[DS4_EXPERT_PROFILE_MAX_CAPS]; uint64_t layer_records[DS4_MAX_LAYER]; uint64_t hist[DS4_MAX_LAYER][DS4_MAX_EXPERT]; double weight_hist[DS4_MAX_LAYER][DS4_MAX_EXPERT]; uint64_t cache_hits[DS4_EXPERT_PROFILE_MAX_CAPS][DS4_MAX_LAYER]; double cache_weight_hits[DS4_EXPERT_PROFILE_MAX_CAPS][DS4_MAX_LAYER]; int cache_entries[DS4_EXPERT_PROFILE_MAX_CAPS][DS4_MAX_LAYER][DS4_MAX_EXPERT]; uint32_t cache_count[DS4_EXPERT_PROFILE_MAX_CAPS][DS4_MAX_LAYER]; bool prev_valid[DS4_MAX_LAYER]; uint32_t prev_pos[DS4_MAX_LAYER]; int prev_selected[DS4_MAX_LAYER][DS4_MAX_EXPERT_USED]; uint64_t adjacent_pairs[DS4_MAX_LAYER]; double adjacent_overlap_sum[DS4_MAX_LAYER]; double adjacent_jaccard_sum[DS4_MAX_LAYER]; bool layer_is_hash[DS4_MAX_LAYER]; uint64_t total_records; uint64_t total_selections; double total_weight; } ds4_expert_profile; static ds4_expert_profile g_expert_profile; static int ds4_expert_profile_sort_cmp(const void *a, const void *b) { const ds4_expert_profile_sort_entry *ea = a; const ds4_expert_profile_sort_entry *eb = b; if (ea->count < eb->count) return 1; if (ea->count > eb->count) return -1; if (ea->weight < eb->weight) return 1; if (ea->weight > eb->weight) return -1; return ea->expert_id - eb->expert_id; } typedef struct { uint32_t layer; uint32_t expert; uint64_t count; double weight; } ds4_expert_hotlist_entry; static int ds4_expert_hotlist_sort_cmp(const void *a, const void *b) { const ds4_expert_hotlist_entry *ea = a; const ds4_expert_hotlist_entry *eb = b; if (ea->count < eb->count) return 1; if (ea->count > eb->count) return -1; if (ea->weight < eb->weight) return 1; if (ea->weight > eb->weight) 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; } #include "ds4_streaming_hotlist.inc" #include "ds4_streaming_hotlist_glm52.inc" static void ds4_json_write_string(FILE *fp, const char *s) { fputc('"', fp); if (s) { for (const unsigned char *p = (const unsigned char *)s; *p; p++) { switch (*p) { case '"': fputs("\\\"", fp); break; case '\\': fputs("\\\\", fp); break; case '\b': fputs("\\b", fp); break; case '\f': fputs("\\f", fp); break; case '\n': fputs("\\n", fp); break; case '\r': fputs("\\r", fp); break; case '\t': fputs("\\t", fp); break; default: if (*p < 0x20) fprintf(fp, "\\u%04x", (unsigned)*p); else fputc((char)*p, fp); break; } } } fputc('"', fp); } static void ds4_expert_profile_init(const char *path, const char *hotlist_path) { if ((!path || !path[0]) && (!hotlist_path || !hotlist_path[0])) return; if (g_expert_profile.active) { free(g_expert_profile.path); free(g_expert_profile.hotlist_path); memset(&g_expert_profile, 0, sizeof(g_expert_profile)); } g_expert_profile.active = true; if (path && path[0]) g_expert_profile.path = ds4_strdup(path); if (hotlist_path && hotlist_path[0]) { g_expert_profile.hotlist_path = ds4_strdup(hotlist_path); } snprintf(g_expert_profile.model_name, sizeof(g_expert_profile.model_name), "%s", DS4_MODEL_SHAPE_NAME); g_expert_profile.n_layer = DS4_N_LAYER; g_expert_profile.n_expert = DS4_N_EXPERT; g_expert_profile.n_expert_used = DS4_N_EXPERT_USED; for (uint32_t i = 0; i < DS4_EXPERT_PROFILE_MAX_CAPS; i++) { const uint32_t cap = ds4_expert_profile_cap_candidates[i]; if (cap == 0 || cap > DS4_N_EXPERT) continue; g_expert_profile.caps[g_expert_profile.n_caps++] = cap; } fprintf(stderr, "ds4: Metal expert locality profiler active (profile: %s, hotlist: %s)\n", g_expert_profile.path ? g_expert_profile.path : "disabled", g_expert_profile.hotlist_path ? g_expert_profile.hotlist_path : "disabled"); } static void ds4_expert_profile_cache_use( uint32_t ci, uint32_t il, int expert, float weight) { ds4_expert_profile *p = &g_expert_profile; const uint32_t cap = p->caps[ci]; int *entries = p->cache_entries[ci][il]; uint32_t *count = &p->cache_count[ci][il]; uint32_t found = UINT32_MAX; for (uint32_t i = 0; i < *count; i++) { if (entries[i] == expert) { found = i; break; } } if (found != UINT32_MAX) { p->cache_hits[ci][il]++; p->cache_weight_hits[ci][il] += weight; for (uint32_t i = found; i > 0; i--) entries[i] = entries[i - 1]; entries[0] = expert; return; } uint32_t n = *count; if (n < cap) { (*count)++; } else if (n > 0) { n = cap - 1; } for (uint32_t i = n; i > 0; i--) entries[i] = entries[i - 1]; entries[0] = expert; } static void ds4_expert_profile_record( uint32_t il, uint32_t pos, const int32_t selected[DS4_MAX_EXPERT_USED], const float weights[DS4_MAX_EXPERT_USED], bool is_hash) { ds4_expert_profile *p = &g_expert_profile; if (!p->active || il >= p->n_layer) return; p->layer_records[il]++; p->total_records++; if (is_hash) p->layer_is_hash[il] = true; if (p->prev_valid[il] && p->prev_pos[il] + 1u == pos) { uint32_t intersection = 0; for (uint32_t a = 0; a < p->n_expert_used; a++) { for (uint32_t b = 0; b < p->n_expert_used; b++) { if (p->prev_selected[il][a] == selected[b]) { intersection++; break; } } } const uint32_t union_size = 2u * p->n_expert_used - intersection; p->adjacent_pairs[il]++; p->adjacent_overlap_sum[il] += (double)intersection / (double)p->n_expert_used; if (union_size != 0) { p->adjacent_jaccard_sum[il] += (double)intersection / (double)union_size; } } p->prev_valid[il] = true; p->prev_pos[il] = pos; for (uint32_t slot = 0; slot < p->n_expert_used; slot++) { const int expert = selected[slot]; const float weight = weights[slot]; p->prev_selected[il][slot] = expert; if (expert < 0 || (uint32_t)expert >= p->n_expert) continue; p->hist[il][expert]++; p->weight_hist[il][expert] += weight; p->total_selections++; p->total_weight += weight; for (uint32_t ci = 0; ci < p->n_caps; ci++) { ds4_expert_profile_cache_use(ci, il, expert, weight); } } } static void ds4_expert_profile_write_cache_summary(FILE *fp) { ds4_expert_profile *p = &g_expert_profile; fputs(" \"cache_summary\": [\n", fp); for (uint32_t ci = 0; ci < p->n_caps; ci++) { uint64_t hits = 0; double weight_hits = 0.0; for (uint32_t il = 0; il < p->n_layer; il++) { hits += p->cache_hits[ci][il]; weight_hits += p->cache_weight_hits[ci][il]; } const double hit_rate = p->total_selections ? (double)hits / (double)p->total_selections : 0.0; const double weight_hit_rate = p->total_weight > 0.0 ? weight_hits / p->total_weight : 0.0; fprintf(fp, " {\"n\":%u,\"hits\":%" PRIu64 ",\"selections\":%" PRIu64 ",\"hit_rate\":%.6f,\"weighted_hit_rate\":%.6f}%s\n", p->caps[ci], hits, p->total_selections, hit_rate, weight_hit_rate, ci + 1u == p->n_caps ? "" : ","); } fputs(" ],\n", fp); } static void ds4_expert_profile_write_layer(FILE *fp, uint32_t il) { ds4_expert_profile *p = &g_expert_profile; ds4_expert_profile_sort_entry entries[DS4_MAX_EXPERT]; uint32_t unique = 0; uint64_t selections = 0; double weight_total = 0.0; for (uint32_t e = 0; e < p->n_expert; e++) { const uint64_t count = p->hist[il][e]; selections += count; weight_total += p->weight_hist[il][e]; if (count == 0) continue; entries[unique++] = (ds4_expert_profile_sort_entry) { .expert_id = (int)e, .count = count, .weight = p->weight_hist[il][e], }; } qsort(entries, unique, sizeof(entries[0]), ds4_expert_profile_sort_cmp); const double avg_overlap = p->adjacent_pairs[il] ? p->adjacent_overlap_sum[il] / (double)p->adjacent_pairs[il] : 0.0; const double avg_jaccard = p->adjacent_pairs[il] ? p->adjacent_jaccard_sum[il] / (double)p->adjacent_pairs[il] : 0.0; fprintf(fp, " {\"layer\":%u,\"hash_router\":%s,\"records\":%" PRIu64 ",\"selections\":%" PRIu64 ",\"unique_experts\":%u" ",\"avg_adjacent_overlap\":%.6f,\"avg_adjacent_jaccard\":%.6f,\n", il, p->layer_is_hash[il] ? "true" : "false", p->layer_records[il], selections, unique, avg_overlap, avg_jaccard); fputs(" \"top_experts\": [", fp); const uint32_t top_n = unique < 16u ? unique : 16u; for (uint32_t i = 0; i < top_n; i++) { const double pct = selections ? 100.0 * (double)entries[i].count / (double)selections : 0.0; const double weight_pct = weight_total > 0.0 ? 100.0 * entries[i].weight / weight_total : 0.0; if (i) fputc(',', fp); fprintf(fp, "{\"id\":%d,\"count\":%" PRIu64 ",\"pct\":%.4f" ",\"weight\":%.9g,\"weight_pct\":%.4f}", entries[i].expert_id, entries[i].count, pct, entries[i].weight, weight_pct); } fputs("],\n", fp); fputs(" \"cache\": [", fp); for (uint32_t ci = 0; ci < p->n_caps; ci++) { const uint64_t hits = p->cache_hits[ci][il]; const double weight_hits = p->cache_weight_hits[ci][il]; const double hit_rate = selections ? (double)hits / (double)selections : 0.0; const double weight_hit_rate = weight_total > 0.0 ? weight_hits / weight_total : 0.0; if (ci) fputc(',', fp); fprintf(fp, "{\"n\":%u,\"hits\":%" PRIu64 ",\"hit_rate\":%.6f" ",\"weighted_hit_rate\":%.6f}", p->caps[ci], hits, hit_rate, weight_hit_rate); } fputs("]}", fp); } static void ds4_expert_profile_write_hotlist_file(ds4_expert_profile *p) { if (!p || !p->hotlist_path || !p->hotlist_path[0]) return; const size_t cap = (size_t)p->n_layer * p->n_expert; ds4_expert_hotlist_entry *entries = xmalloc(cap * sizeof(entries[0])); size_t n = 0; for (uint32_t il = 0; il < p->n_layer; il++) { for (uint32_t expert = 0; expert < p->n_expert; expert++) { entries[n++] = (ds4_expert_hotlist_entry) { .layer = il, .expert = expert, .count = p->hist[il][expert], .weight = p->weight_hist[il][expert], }; } } qsort(entries, n, sizeof(entries[0]), ds4_expert_hotlist_sort_cmp); FILE *fp = fopen(p->hotlist_path, "wb"); if (!fp) { fprintf(stderr, "ds4: failed to open expert hotlist output %s: %s\n", p->hotlist_path, strerror(errno)); free(entries); return; } fprintf(fp, "# ds4 expert hotlist v1\n" "# model %s\n" "# layers %u\n" "# experts %u\n" "# layer_records %" PRIu64 "\n" "# selections %" PRIu64 "\n" "# columns: layer expert hits weight\n", p->model_name, p->n_layer, p->n_expert, p->total_records, p->total_selections); for (size_t i = 0; i < n; i++) { fprintf(fp, "%u %u %" PRIu64 " %.17g\n", entries[i].layer, entries[i].expert, entries[i].count, entries[i].weight); } free(entries); if (fclose(fp) != 0) { fprintf(stderr, "ds4: failed to close expert hotlist output %s: %s\n", p->hotlist_path, strerror(errno)); } else { fprintf(stderr, "ds4: wrote Metal expert hotlist to %s " "(%" PRIu64 " layer records, %" PRIu64 " selections)\n", p->hotlist_path, p->total_records, p->total_selections); } } static void ds4_expert_profile_close(void) { ds4_expert_profile *p = &g_expert_profile; if (!p->active) return; if (p->path) { FILE *fp = fopen(p->path, "wb"); if (!fp) { fprintf(stderr, "ds4: failed to open expert profile output %s: %s\n", p->path, strerror(errno)); } else { fputs("{\n", fp); fputs(" \"source\": \"ds4 Metal expert locality profile\",\n", fp); fputs(" \"model\": ", fp); ds4_json_write_string(fp, p->model_name); fputs(",\n", fp); fprintf(fp, " \"layers\": %u,\n" " \"experts\": %u,\n" " \"expert_used\": %u,\n" " \"layer_records\": %" PRIu64 ",\n" " \"selections\": %" PRIu64 ",\n", p->n_layer, p->n_expert, p->n_expert_used, p->total_records, p->total_selections); fputs(" \"cache_ns\": [", fp); for (uint32_t ci = 0; ci < p->n_caps; ci++) { if (ci) fputc(',', fp); fprintf(fp, "%u", p->caps[ci]); } fputs("],\n", fp); ds4_expert_profile_write_cache_summary(fp); fputs(" \"layers_detail\": [\n", fp); for (uint32_t il = 0; il < p->n_layer; il++) { ds4_expert_profile_write_layer(fp, il); fputs(il + 1u == p->n_layer ? "\n" : ",\n", fp); } fputs(" ]\n}\n", fp); if (fclose(fp) != 0) { fprintf(stderr, "ds4: failed to close expert profile output %s: %s\n", p->path, strerror(errno)); } else { fprintf(stderr, "ds4: wrote Metal expert locality profile to %s " "(%" PRIu64 " layer records, %" PRIu64 " selections)\n", p->path, p->total_records, p->total_selections); } } } ds4_expert_profile_write_hotlist_file(p); free(p->path); free(p->hotlist_path); memset(p, 0, sizeof(*p)); } static void sleep_sec(double sec) { if (sec <= 0.0 || !isfinite(sec)) return; struct timespec req; req.tv_sec = (time_t)sec; req.tv_nsec = (long)((sec - (double)req.tv_sec) * 1000000000.0); if (req.tv_nsec < 0) req.tv_nsec = 0; if (req.tv_nsec >= 1000000000L) { req.tv_sec++; req.tv_nsec -= 1000000000L; } /* Do not resume after EINTR: Ctrl+C should cut through throttling sleeps. */ (void)nanosleep(&req, &req); } static const char *ds4_log_color_code(ds4_log_type type) { switch (type) { case DS4_LOG_PREFILL: case DS4_LOG_TIMING: return "\x1b[36m"; case DS4_LOG_GENERATION: case DS4_LOG_OK: return "\x1b[32m"; case DS4_LOG_KVCACHE: return "\x1b[33m"; case DS4_LOG_TOOL: return "\x1b[90m"; case DS4_LOG_WARNING: return "\x1b[38;5;208m"; case DS4_LOG_ERROR: return "\x1b[31m"; default: return ""; } } bool ds4_log_is_tty(FILE *fp) { int fd = fileno(fp); return fd >= 0 && isatty(fd) != 0; } static void ds4_vlog(FILE *fp, ds4_log_type type, const char *fmt, va_list ap) { const bool colorize = type != DS4_LOG_DEFAULT && ds4_log_is_tty(fp); if (colorize) fputs(ds4_log_color_code(type), fp); vfprintf(fp, fmt, ap); if (colorize) fputs("\x1b[0m", fp); } void ds4_log(FILE *fp, ds4_log_type type, const char *fmt, ...) { va_list ap; va_start(ap, fmt); ds4_vlog(fp, type, fmt, ap); va_end(ap); } static bool write_f32_binary_file(const char *path, const float *data, uint64_t n) { FILE *fp = fopen(path, "wb"); if (!fp) { fprintf(stderr, "ds4: failed to open %s for writing: %s\n", path, strerror(errno)); return false; } const size_t nw = fwrite(data, sizeof(float), (size_t)n, fp); const bool ok = nw == (size_t)n && fclose(fp) == 0; if (!ok) { fprintf(stderr, "ds4: failed to write %s\n", path); return false; } return true; } static bool read_f32_binary_file(const char *path, float *data, uint64_t n) { struct stat st; if (stat(path, &st) != 0) { fprintf(stderr, "ds4: failed to stat %s: %s\n", path, strerror(errno)); return false; } if (st.st_size < 0 || (uint64_t)st.st_size != n * sizeof(float)) { fprintf(stderr, "ds4: %s has size %llu bytes, expected %llu bytes\n", path, (unsigned long long)st.st_size, (unsigned long long)(n * sizeof(float))); return false; } FILE *fp = fopen(path, "rb"); if (!fp) { fprintf(stderr, "ds4: failed to open %s for reading: %s\n", path, strerror(errno)); return false; } const size_t nr = fread(data, sizeof(float), (size_t)n, fp); const bool ok = nr == (size_t)n && fclose(fp) == 0; if (!ok) { fprintf(stderr, "ds4: failed to read %s\n", path); return false; } return true; } static bool cpu_directional_steering_enabled( const float *dirs, float scale); static void cpu_directional_steering_project_rows( float *x, const float *dirs, uint32_t il, uint32_t rows, float scale); typedef void (*ds4_parallel_fn)(void *ctx, uint64_t row0, uint64_t row1); #define DS4_MAX_THREADS 32 typedef struct { pthread_t threads[DS4_MAX_THREADS]; pthread_mutex_t mutex; pthread_cond_t work_cond; pthread_cond_t done_cond; uint32_t n_threads; uint32_t n_workers; uint32_t generation; uint32_t done; bool initialized; bool shutdown; ds4_parallel_fn fn; void *ctx; uint64_t n_rows; } ds4_thread_pool; static ds4_thread_pool g_pool; static __thread int g_parallel_depth; static uint32_t g_requested_threads; static void *ds4_worker_main(void *arg) { const uint32_t tid = (uint32_t)(uintptr_t)arg; uint32_t seen_generation = 0; for (;;) { pthread_mutex_lock(&g_pool.mutex); while (seen_generation == g_pool.generation && !g_pool.shutdown) { pthread_cond_wait(&g_pool.work_cond, &g_pool.mutex); } if (g_pool.shutdown) { pthread_mutex_unlock(&g_pool.mutex); return NULL; } seen_generation = g_pool.generation; ds4_parallel_fn fn = g_pool.fn; void *ctx = g_pool.ctx; const uint64_t n_rows = g_pool.n_rows; const uint32_t n_threads = g_pool.n_threads; pthread_mutex_unlock(&g_pool.mutex); const uint64_t rows_per_thread = (n_rows + n_threads - 1) / n_threads; const uint64_t row0 = (uint64_t)tid * rows_per_thread; uint64_t row1 = row0 + rows_per_thread; if (row1 > n_rows) row1 = n_rows; if (row0 < row1) { g_parallel_depth++; fn(ctx, row0, row1); g_parallel_depth--; } pthread_mutex_lock(&g_pool.mutex); g_pool.done++; if (g_pool.done == g_pool.n_workers) { pthread_cond_signal(&g_pool.done_cond); } pthread_mutex_unlock(&g_pool.mutex); } } /* Create the persistent CPU worker pool. Decode reuses these threads instead * of creating pthreads in the token loop. */ static void ds4_threads_init(void) { if (g_pool.initialized) return; pthread_once(&iq2xxs_signed_grid_once, iq2xxs_signed_grid_init); uint32_t n_threads = 12; const long online_cpus = sysconf(_SC_NPROCESSORS_ONLN); if (online_cpus > 0) { n_threads = online_cpus < 12 ? (uint32_t)online_cpus : 12; } const char *env = getenv("DS4_THREADS"); if (env && env[0]) { long v = strtol(env, NULL, 10); if (v > 0) n_threads = (uint32_t)v; } if (g_requested_threads > 0) n_threads = g_requested_threads; if (n_threads > DS4_MAX_THREADS) n_threads = DS4_MAX_THREADS; if (n_threads == 0) n_threads = 1; pthread_mutex_init(&g_pool.mutex, NULL); pthread_cond_init(&g_pool.work_cond, NULL); pthread_cond_init(&g_pool.done_cond, NULL); g_pool.n_threads = n_threads; g_pool.n_workers = n_threads > 0 ? n_threads - 1 : 0; g_pool.generation = 0; g_pool.done = 0; g_pool.shutdown = false; g_pool.initialized = true; for (uint32_t i = 1; i < n_threads; i++) { if (pthread_create(&g_pool.threads[i], NULL, ds4_worker_main, (void *)(uintptr_t)i) != 0) { ds4_die("failed to create worker thread"); } } } static void ds4_threads_shutdown(void) { if (!g_pool.initialized) return; pthread_mutex_lock(&g_pool.mutex); g_pool.shutdown = true; g_pool.generation++; pthread_cond_broadcast(&g_pool.work_cond); pthread_mutex_unlock(&g_pool.mutex); for (uint32_t i = 1; i < g_pool.n_threads; i++) { pthread_join(g_pool.threads[i], NULL); } pthread_cond_destroy(&g_pool.done_cond); pthread_cond_destroy(&g_pool.work_cond); pthread_mutex_destroy(&g_pool.mutex); memset(&g_pool, 0, sizeof(g_pool)); } /* Run a row-parallel CPU kernel, falling back to serial execution for small * jobs or nested calls where spawning more work would only add latency. */ static void ds4_parallel_for_min_rows(uint64_t n_rows, ds4_parallel_fn fn, void *ctx, uint64_t min_parallel_rows) { ds4_threads_init(); if (g_parallel_depth > 0 || g_pool.n_threads <= 1 || n_rows < min_parallel_rows) { fn(ctx, 0, n_rows); return; } pthread_mutex_lock(&g_pool.mutex); g_pool.fn = fn; g_pool.ctx = ctx; g_pool.n_rows = n_rows; g_pool.done = 0; g_pool.generation++; pthread_cond_broadcast(&g_pool.work_cond); const uint64_t rows_per_thread = (n_rows + g_pool.n_threads - 1) / g_pool.n_threads; uint64_t main_row1 = rows_per_thread; if (main_row1 > n_rows) main_row1 = n_rows; pthread_mutex_unlock(&g_pool.mutex); if (main_row1 > 0) { g_parallel_depth++; fn(ctx, 0, main_row1); g_parallel_depth--; } pthread_mutex_lock(&g_pool.mutex); while (g_pool.done < g_pool.n_workers) { pthread_cond_wait(&g_pool.done_cond, &g_pool.mutex); } pthread_mutex_unlock(&g_pool.mutex); } static void ds4_parallel_for(uint64_t n_rows, ds4_parallel_fn fn, void *ctx) { ds4_parallel_for_min_rows(n_rows, fn, ctx, 512); } static void cursor_error(ds4_cursor *c, const char *msg) { if (c->error[0] == '\0') { snprintf(c->error, sizeof(c->error), "%s at byte %" PRIu64, msg, c->pos); } } static bool cursor_has(ds4_cursor *c, uint64_t n) { if (n > c->size || c->pos > c->size - n) { cursor_error(c, "truncated GGUF file"); return false; } return true; } static bool cursor_read(ds4_cursor *c, void *dst, uint64_t n) { if (!cursor_has(c, n)) return false; memcpy(dst, c->base + c->pos, (size_t)n); c->pos += n; return true; } static bool cursor_skip(ds4_cursor *c, uint64_t n) { if (!cursor_has(c, n)) return false; c->pos += n; return true; } static bool cursor_u32(ds4_cursor *c, uint32_t *v) { return cursor_read(c, v, sizeof(*v)); } static bool cursor_u64(ds4_cursor *c, uint64_t *v) { return cursor_read(c, v, sizeof(*v)); } static bool cursor_string(ds4_cursor *c, ds4_str *s) { uint64_t len; if (!cursor_u64(c, &len)) return false; if (!cursor_has(c, len)) return false; s->ptr = (const char *)(c->base + c->pos); s->len = len; c->pos += len; return true; } static uint64_t align_up(uint64_t value, uint64_t alignment) { uint64_t rem = value % alignment; return rem == 0 ? value : value + alignment - rem; } /* ========================================================================= * GGUF Parsing and Model Mapping. * ========================================================================= * * The loader maps the model once, records metadata/tensor descriptors, and * leaves tensor bytes in place. Inference code accesses weights by adding * tensor offsets to the mapping instead of copying the GGUF into private * structures. */ enum { GGUF_VALUE_UINT8 = 0, GGUF_VALUE_INT8 = 1, GGUF_VALUE_UINT16 = 2, GGUF_VALUE_INT16 = 3, GGUF_VALUE_UINT32 = 4, GGUF_VALUE_INT32 = 5, GGUF_VALUE_FLOAT32 = 6, GGUF_VALUE_BOOL = 7, GGUF_VALUE_STRING = 8, GGUF_VALUE_ARRAY = 9, GGUF_VALUE_UINT64 = 10, GGUF_VALUE_INT64 = 11, GGUF_VALUE_FLOAT64 = 12, }; typedef struct { const char *name; uint32_t block_elems; uint32_t block_bytes; } gguf_type_info; static const gguf_type_info gguf_types[] = { [0] = {"f32", 1, 4}, [1] = {"f16", 1, 2}, [2] = {"q4_0", 32, 18}, [3] = {"q4_1", 32, 20}, [6] = {"q5_0", 32, 22}, [7] = {"q5_1", 32, 24}, [8] = {"q8_0", 32, 34}, [9] = {"q8_1", 32, 40}, [10] = {"q2_k", 256, 84}, [11] = {"q3_k", 256, 110}, [12] = {"q4_k", 256, 144}, [13] = {"q5_k", 256, 176}, [14] = {"q6_k", 256, 210}, [15] = {"q8_k", 256, 292}, [16] = {"iq2_xxs",256, 66}, [17] = {"iq2_xs", 256, 74}, [18] = {"iq3_xxs",256, 98}, [19] = {"iq1_s", 256, 110}, [20] = {"iq4_nl", 256, 50}, [21] = {"iq3_s", 256, 110}, [22] = {"iq2_s", 256, 82}, [23] = {"iq4_xs", 256, 136}, [24] = {"i8", 1, 1}, [25] = {"i16", 1, 2}, [26] = {"i32", 1, 4}, [27] = {"i64", 1, 8}, [28] = {"f64", 1, 8}, [29] = {"iq1_m", 256, 56}, [30] = {"bf16", 1, 2}, [39] = {"iq2_m", 256, 70}, }; enum { DS4_TENSOR_F32 = 0, DS4_TENSOR_F16 = 1, DS4_TENSOR_Q4_0 = 2, DS4_TENSOR_Q8_0 = 8, DS4_TENSOR_Q2_K = 10, DS4_TENSOR_Q4_K = 12, DS4_TENSOR_Q5_K = 13, DS4_TENSOR_Q6_K = 14, DS4_TENSOR_Q8_K = 15, DS4_TENSOR_IQ2_XXS = 16, DS4_TENSOR_IQ3_XXS = 18, DS4_TENSOR_IQ2_S = 22, DS4_TENSOR_I32 = 26, DS4_TENSOR_BF16 = 30, DS4_TENSOR_IQ2_M = 39, }; typedef struct { ds4_str key; uint32_t type; uint64_t value_pos; } ds4_kv; typedef struct { ds4_str name; uint32_t ndim; uint64_t dim[DS4_MAX_DIMS]; uint32_t type; uint64_t rel_offset; uint64_t abs_offset; uint64_t elements; uint64_t bytes; } ds4_tensor; typedef struct { int fd; const uint8_t *map; uint64_t size; uint32_t version; uint64_t n_kv; uint64_t n_tensors; uint64_t alignment; uint64_t tensor_data_pos; uint64_t max_tensor_bytes; ds4_kv *kv; ds4_tensor *tensors; } ds4_model; static uint64_t scalar_value_size(uint32_t type) { switch (type) { case GGUF_VALUE_UINT8: case GGUF_VALUE_INT8: case GGUF_VALUE_BOOL: return 1; case GGUF_VALUE_UINT16: case GGUF_VALUE_INT16: return 2; case GGUF_VALUE_UINT32: case GGUF_VALUE_INT32: case GGUF_VALUE_FLOAT32: return 4; case GGUF_VALUE_UINT64: case GGUF_VALUE_INT64: case GGUF_VALUE_FLOAT64: return 8; default: return 0; } } static bool skip_value(ds4_cursor *c, uint32_t type, int depth) { if (depth > 8) { cursor_error(c, "metadata array nesting is too deep"); return false; } uint64_t scalar = scalar_value_size(type); if (scalar != 0) return cursor_skip(c, scalar); if (type == GGUF_VALUE_STRING) { ds4_str ignored; return cursor_string(c, &ignored); } if (type == GGUF_VALUE_ARRAY) { uint32_t item_type; uint64_t len; if (!cursor_u32(c, &item_type)) return false; if (!cursor_u64(c, &len)) return false; uint64_t item_size = scalar_value_size(item_type); if (item_size != 0) { if (len > UINT64_MAX / item_size) { cursor_error(c, "metadata array is too large"); return false; } return cursor_skip(c, len * item_size); } for (uint64_t i = 0; i < len; i++) { if (!skip_value(c, item_type, depth + 1)) return false; } return true; } cursor_error(c, "unknown GGUF metadata type"); return false; } static const gguf_type_info *tensor_type(uint32_t type) { uint32_t n = sizeof(gguf_types) / sizeof(gguf_types[0]); if (type >= n || gguf_types[type].name == NULL) return NULL; return &gguf_types[type]; } static const char *tensor_type_name(uint32_t type) { const gguf_type_info *info = tensor_type(type); return info ? info->name : "unknown"; } static bool tensor_nbytes(uint32_t type, uint64_t elements, uint64_t *bytes) { const gguf_type_info *info = tensor_type(type); if (!info || info->block_elems == 0) return false; uint64_t blocks = (elements + info->block_elems - 1) / info->block_elems; if (blocks > UINT64_MAX / info->block_bytes) return false; *bytes = blocks * info->block_bytes; return true; } static ds4_cursor cursor_at(const ds4_model *m, uint64_t pos) { ds4_cursor c = { .base = m->map, .size = m->size, .pos = pos, .error = {0}, }; return c; } static ds4_kv *model_find_kv(const ds4_model *m, const char *key) { for (uint64_t i = 0; i < m->n_kv; i++) { if (ds4_streq(m->kv[i].key, key)) return &m->kv[i]; } return NULL; } static bool model_get_string(const ds4_model *m, const char *key, ds4_str *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv || kv->type != GGUF_VALUE_STRING) return false; ds4_cursor c = cursor_at(m, kv->value_pos); return cursor_string(&c, out); } static bool model_get_u32(const ds4_model *m, const char *key, uint32_t *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv || kv->type != GGUF_VALUE_UINT32) return false; ds4_cursor c = cursor_at(m, kv->value_pos); return cursor_u32(&c, out); } static bool model_get_token_id(const ds4_model *m, const char *key, int *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv) return false; ds4_cursor c = cursor_at(m, kv->value_pos); switch (kv->type) { case GGUF_VALUE_UINT32: { uint32_t v = 0; if (!cursor_u32(&c, &v) || v > (uint32_t)INT_MAX) return false; *out = (int)v; return true; } case GGUF_VALUE_INT32: { int32_t v = 0; if (!cursor_read(&c, &v, sizeof(v)) || v < 0) return false; *out = (int)v; return true; } case GGUF_VALUE_UINT64: { uint64_t v = 0; if (!cursor_u64(&c, &v) || v > (uint64_t)INT_MAX) return false; *out = (int)v; return true; } case GGUF_VALUE_INT64: { int64_t v = 0; if (!cursor_read(&c, &v, sizeof(v)) || v < 0 || v > (int64_t)INT_MAX) return false; *out = (int)v; return true; } default: return false; } } static bool model_get_u64_compat(const ds4_model *m, const char *key, uint64_t *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv) return false; ds4_cursor c = cursor_at(m, kv->value_pos); if (kv->type == GGUF_VALUE_UINT64) { return cursor_u64(&c, out); } if (kv->type == GGUF_VALUE_UINT32) { uint32_t v = 0; if (!cursor_u32(&c, &v)) return false; *out = v; return true; } return false; } static bool model_get_f32_compat(const ds4_model *m, const char *key, float *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv) return false; ds4_cursor c = cursor_at(m, kv->value_pos); if (kv->type == GGUF_VALUE_FLOAT32) { return cursor_read(&c, out, sizeof(*out)); } if (kv->type == GGUF_VALUE_FLOAT64) { double v = 0.0; if (!cursor_read(&c, &v, sizeof(v))) return false; *out = (float)v; return true; } if (kv->type == GGUF_VALUE_UINT32) { uint32_t v = 0; if (!cursor_u32(&c, &v)) return false; *out = (float)v; return true; } if (kv->type == GGUF_VALUE_INT32) { int32_t v = 0; if (!cursor_read(&c, &v, sizeof(v))) return false; *out = (float)v; return true; } return false; } static bool model_get_bool(const ds4_model *m, const char *key, bool *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv || kv->type != GGUF_VALUE_BOOL) return false; ds4_cursor c = cursor_at(m, kv->value_pos); uint8_t v = 0; if (!cursor_read(&c, &v, sizeof(v))) return false; *out = v != 0; return true; } typedef struct { uint32_t type; uint64_t len; uint64_t data_pos; } ds4_array_ref; static bool model_get_array(const ds4_model *m, const char *key, ds4_array_ref *out) { ds4_kv *kv = model_find_kv(m, key); if (!kv || kv->type != GGUF_VALUE_ARRAY) return false; ds4_cursor c = cursor_at(m, kv->value_pos); if (!cursor_u32(&c, &out->type)) return false; if (!cursor_u64(&c, &out->len)) return false; out->data_pos = c.pos; return true; } static void model_close(ds4_model *m) { if (!m) return; free(m->kv); free(m->tensors); if (m->map) munmap((void *)m->map, (size_t)m->size); if (m->fd >= 0) close(m->fd); memset(m, 0, sizeof(*m)); m->fd = -1; } static void model_prefetch_cpu_mapping(const ds4_model *m) { if (!m || !m->map || m->size == 0) return; /* * CPU generation touches expert weights according to router decisions, so a * long decode can fault in model pages that the prompt never touched. On * current Darwin kernels we have seen those late file-backed faults trigger * an OS-level VM panic in map-count accounting. This hint does not copy or * pin the GGUF; it just asks the kernel to start bringing the read-only * mapping into the page cache before token generation reaches it. */ #if defined(POSIX_MADV_WILLNEED) const int rc = posix_madvise((void *)m->map, (size_t)m->size, POSIX_MADV_WILLNEED); if (rc != 0) { ds4_log(stderr, DS4_LOG_WARNING, "ds4: warning: POSIX_MADV_WILLNEED failed for CPU model mapping: %s\n", strerror(rc)); } #else (void)m; #endif } /* Read the GGUF metadata table. Values stay in the mmap; we store offsets so * later validation can decode only the keys it needs. */ static void parse_metadata(ds4_model *m, ds4_cursor *c) { m->kv = calloc((size_t)m->n_kv, sizeof(m->kv[0])); if (!m->kv) ds4_die("out of memory while allocating metadata table"); m->alignment = 32; for (uint64_t i = 0; i < m->n_kv; i++) { ds4_kv *kv = &m->kv[i]; if (!cursor_string(c, &kv->key)) ds4_die(c->error); if (!cursor_u32(c, &kv->type)) ds4_die(c->error); kv->value_pos = c->pos; if (ds4_streq(kv->key, "general.alignment") && kv->type == GGUF_VALUE_UINT32) { ds4_cursor tmp = cursor_at(m, kv->value_pos); uint32_t alignment; if (cursor_u32(&tmp, &alignment) && alignment != 0) { m->alignment = alignment; } } if (!skip_value(c, kv->type, 0)) ds4_die(c->error); } } /* Read the tensor directory and convert relative GGUF offsets to absolute * mmap offsets. Tensor bytes are still never copied here. */ static void parse_tensors(ds4_model *m, ds4_cursor *c) { m->tensors = calloc((size_t)m->n_tensors, sizeof(m->tensors[0])); if (!m->tensors) ds4_die("out of memory while allocating tensor table"); for (uint64_t i = 0; i < m->n_tensors; i++) { ds4_tensor *t = &m->tensors[i]; if (!cursor_string(c, &t->name)) ds4_die(c->error); if (!cursor_u32(c, &t->ndim)) ds4_die(c->error); if (t->ndim == 0 || t->ndim > DS4_MAX_DIMS) { ds4_die("tensor has an unsupported number of dimensions"); } t->elements = 1; for (uint32_t d = 0; d < t->ndim; d++) { if (!cursor_u64(c, &t->dim[d])) ds4_die(c->error); if (t->dim[d] != 0 && t->elements > UINT64_MAX / t->dim[d]) { ds4_die("tensor element count overflow"); } t->elements *= t->dim[d]; } if (!cursor_u32(c, &t->type)) ds4_die(c->error); if (!cursor_u64(c, &t->rel_offset)) ds4_die(c->error); if (!tensor_nbytes(t->type, t->elements, &t->bytes)) { ds4_log(stderr, DS4_LOG_WARNING, "ds4: warning: tensor %.*s has unsupported GGUF type %u\n", (int)t->name.len, t->name.ptr, t->type); } } m->tensor_data_pos = align_up(c->pos, m->alignment); for (uint64_t i = 0; i < m->n_tensors; i++) { ds4_tensor *t = &m->tensors[i]; if (t->rel_offset > UINT64_MAX - m->tensor_data_pos) { ds4_die("tensor offset overflow"); } t->abs_offset = m->tensor_data_pos + t->rel_offset; if (t->bytes != 0 && (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset)) { ds4_die("tensor points outside GGUF file"); } if (t->bytes > m->max_tensor_bytes) { m->max_tensor_bytes = t->bytes; } } } /* Open and map the GGUF once. Metal needs a shared mapping for no-copy * MTLBuffers; CPU uses a private read-only mapping to avoid Darwin VM stress. * Tokenizer-only callers pass prefetch_cpu=false so inspecting tokens never * walks the huge tensor payload. */ static void model_open(ds4_model *m, const char *path, bool metal_mapping, bool prefetch_cpu) { memset(m, 0, sizeof(*m)); m->fd = -1; int fd = open(path, O_RDONLY); if (fd == -1) ds4_die_errno("cannot open model", path); struct stat st; if (fstat(fd, &st) == -1) ds4_die_errno("cannot stat model", path); if (st.st_size < 32) ds4_die("model file is too small to be GGUF"); /* * Metal wraps slices of this mapping as no-copy MTLBuffers, so the Metal * path keeps the file-backed shared mapping. The CPU path only reads the * weights through normal pointers and should not inherit Metal's VM policy: * use a private read-only mapping there. * * This is deliberately defensive against an OS-level Darwin VM bug observed * while the CPU backend streams the very large GGUF through a shared mmap: * the kernel can panic in VM map-count accounting instead of returning a * normal user-space failure. Keeping CPU inference off the shared mapping * avoids that VM accounting path while preserving normal file-backed reads. */ const int mmap_flags = metal_mapping ? MAP_SHARED : MAP_PRIVATE; void *map = mmap(NULL, (size_t)st.st_size, PROT_READ, mmap_flags, fd, 0); if (map == MAP_FAILED) ds4_die_errno("cannot mmap model", path); m->fd = fd; m->map = map; m->size = (uint64_t)st.st_size; ds4_cursor c = cursor_at(m, 0); uint32_t magic; if (!cursor_u32(&c, &magic)) ds4_die(c.error); if (magic != DS4_GGUF_MAGIC) ds4_die("model is not a GGUF file"); if (!cursor_u32(&c, &m->version)) ds4_die(c.error); if (!cursor_u64(&c, &m->n_tensors)) ds4_die(c.error); if (!cursor_u64(&c, &m->n_kv)) ds4_die(c.error); if (m->version != 3) ds4_die("only GGUF v3 is supported"); parse_metadata(m, &c); parse_tensors(m, &c); if (!metal_mapping && prefetch_cpu) model_prefetch_cpu_mapping(m); } static void print_size(uint64_t bytes) { const double gib = 1024.0 * 1024.0 * 1024.0; printf("%.2f GiB", (double)bytes / gib); } #define DS4_DSPARK_MAX_TARGET_LAYERS 8 #define DS4_DSPARK_MAX_STAGES 8 #define DS4_DSPARK_MAX_BLOCK_SIZE 16 typedef struct { uint32_t 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]; bool has_metadata; bool has_main_proj; bool has_main_norm; bool has_markov_head; bool has_confidence_head; bool has_final_head; bool has_block_size; bool has_markov_rank; bool has_noise_token_id; bool has_target_layers; } ds4_dspark_summary; typedef enum { DS4_SUPPORT_NONE = 0, DS4_SUPPORT_MTP_LEGACY, DS4_SUPPORT_DSPARK, } ds4_support_kind; static bool model_get_u32_any(const ds4_model *m, const char *const *keys, size_t nkeys, uint32_t *out) { for (size_t i = 0; i < nkeys; i++) { if (model_get_u32(m, keys[i], out)) return true; } return false; } static bool model_get_u32_array_any(const ds4_model *m, const char *const *keys, size_t nkeys, uint32_t *out, uint32_t cap, uint32_t *n_out) { for (size_t ik = 0; ik < nkeys; ik++) { ds4_array_ref arr = {0}; if (!model_get_array(m, keys[ik], &arr)) continue; if (arr.type != GGUF_VALUE_UINT32 && arr.type != GGUF_VALUE_INT32) continue; ds4_cursor c = cursor_at(m, arr.data_pos); uint32_t n = arr.len < cap ? (uint32_t)arr.len : cap; for (uint32_t i = 0; i < n; i++) { if (arr.type == GGUF_VALUE_UINT32) { if (!cursor_u32(&c, &out[i])) return false; } else { int32_t v = 0; if (!cursor_read(&c, &v, sizeof(v))) return false; if (v < 0) return false; out[i] = (uint32_t)v; } } *n_out = n; return true; } return false; } static bool ds4_tensor_mtp_stage(ds4_str name, uint32_t *stage) { if (!ds4_str_starts_with(name, "mtp.")) return false; uint64_t pos = 4; if (pos >= name.len || !isdigit((unsigned char)name.ptr[pos])) return false; uint32_t value = 0; while (pos < name.len && isdigit((unsigned char)name.ptr[pos])) { uint32_t digit = (uint32_t)(name.ptr[pos] - '0'); if (value > (UINT32_MAX - digit) / 10u) return false; value = value * 10u + digit; pos++; } if (pos >= name.len || name.ptr[pos] != '.') return false; *stage = value; return true; } static ds4_dspark_summary model_dspark_summary(const ds4_model *m) { static const char *const block_keys[] = { "deepseek4.dspark.block_size", "deepseek4.dspark_block_size", "dspark.block_size", }; static const char *const markov_keys[] = { "deepseek4.dspark.markov_rank", "deepseek4.dspark_markov_rank", "dspark.markov_rank", }; static const char *const noise_keys[] = { "deepseek4.dspark.noise_token_id", "deepseek4.dspark_noise_token_id", "dspark.noise_token_id", }; static const char *const target_keys[] = { "deepseek4.dspark.target_layer_ids", "deepseek4.dspark_target_layer_ids", "dspark.target_layer_ids", }; ds4_dspark_summary s = {0}; if (model_get_u32_any(m, block_keys, sizeof(block_keys) / sizeof(block_keys[0]), &s.block_size)) { s.has_metadata = true; s.has_block_size = true; } if (model_get_u32_any(m, markov_keys, sizeof(markov_keys) / sizeof(markov_keys[0]), &s.markov_rank)) { s.has_metadata = true; s.has_markov_rank = true; } if (model_get_u32_any(m, noise_keys, sizeof(noise_keys) / sizeof(noise_keys[0]), &s.noise_token_id)) { s.has_metadata = true; s.has_noise_token_id = true; } if (model_get_u32_array_any(m, target_keys, sizeof(target_keys) / sizeof(target_keys[0]), s.target_layers, DS4_DSPARK_MAX_TARGET_LAYERS, &s.target_layer_count)) { s.has_metadata = true; s.has_target_layers = true; } uint32_t max_stage = 0; bool have_stage = false; for (uint64_t i = 0; i < m->n_tensors; i++) { ds4_str name = m->tensors[i].name; uint32_t stage = 0; if (!ds4_tensor_mtp_stage(name, &stage)) continue; if (!have_stage || stage > max_stage) max_stage = stage; have_stage = true; if (ds4_str_contains(name, ".main_proj.")) s.has_main_proj = true; if (ds4_str_contains(name, ".main_norm.")) s.has_main_norm = true; if (ds4_str_contains(name, ".markov_head.")) s.has_markov_head = true; if (ds4_str_contains(name, ".confidence_head.")) s.has_confidence_head = true; if (ds4_str_contains(name, ".hc_head_") || ds4_str_contains(name, ".norm.weight")) { s.has_final_head = true; } } if (have_stage) s.stages = max_stage + 1u; return s; } static void model_print_dspark_summary(const ds4_model *m) { ds4_dspark_summary s = model_dspark_summary(m); if (!s.stages && !s.has_metadata) return; printf("mtp/dspark: stages=%u", s.stages); if (s.block_size) printf(" block=%u", s.block_size); if (s.markov_rank) printf(" markov_rank=%u", s.markov_rank); if (s.noise_token_id) printf(" noise_token=%u", s.noise_token_id); if (s.target_layer_count) { printf(" target_layers="); for (uint32_t i = 0; i < s.target_layer_count; i++) { printf("%s%u", i == 0 ? "" : ",", s.target_layers[i]); } } printf("\n"); if (s.has_main_proj || s.has_main_norm || s.has_markov_head || s.has_confidence_head || s.has_final_head) { printf("mtp/dspark tensors: main_proj=%s main_norm=%s markov=%s confidence=%s final_head=%s\n", s.has_main_proj ? "yes" : "no", s.has_main_norm ? "yes" : "no", s.has_markov_head ? "yes" : "no", s.has_confidence_head ? "yes" : "no", s.has_final_head ? "yes" : "no"); } } static void model_summary(const ds4_model *m) { ds4_str name = {0}; ds4_str arch = {0}; uint32_t layers = 0; uint64_t ctx_train = 0; uint32_t n_head = 0; uint32_t n_head_kv = 0; uint32_t head_dim = 0; uint32_t n_swa = 0; uint32_t indexer_heads = 0; uint32_t indexer_head_dim = 0; uint32_t indexer_top_k = 0; uint32_t n_expert = 0; uint32_t n_expert_used = 0; uint32_t n_expert_groups = 0; uint32_t n_group_used = 0; uint64_t tensor_bytes = 0; uint64_t params = 0; model_get_string(m, "general.name", &name); model_get_string(m, "general.architecture", &arch); if (!model_get_u32(m, "deepseek4.block_count", &layers)) { model_get_u32(m, "glm-dsa.block_count", &layers); } if (!model_get_u64_compat(m, "deepseek4.context_length", &ctx_train)) { model_get_u64_compat(m, "glm-dsa.context_length", &ctx_train); } if (!model_get_u32(m, "deepseek4.attention.head_count", &n_head)) { model_get_u32(m, "glm-dsa.attention.head_count", &n_head); } if (!model_get_u32(m, "deepseek4.attention.head_count_kv", &n_head_kv)) { model_get_u32(m, "glm-dsa.attention.head_count_kv", &n_head_kv); } if (!model_get_u32(m, "deepseek4.attention.key_length", &head_dim)) { model_get_u32(m, "glm-dsa.attention.key_length", &head_dim); } model_get_u32(m, "deepseek4.attention.sliding_window", &n_swa); if (!model_get_u32(m, "deepseek4.attention.indexer.head_count", &indexer_heads)) { model_get_u32(m, "glm-dsa.attention.indexer.head_count", &indexer_heads); } if (!model_get_u32(m, "deepseek4.attention.indexer.key_length", &indexer_head_dim)) { model_get_u32(m, "glm-dsa.attention.indexer.key_length", &indexer_head_dim); } if (!model_get_u32(m, "deepseek4.attention.indexer.top_k", &indexer_top_k)) { model_get_u32(m, "glm-dsa.attention.indexer.top_k", &indexer_top_k); } if (!model_get_u32(m, "deepseek4.expert_count", &n_expert)) { model_get_u32(m, "glm-dsa.expert_count", &n_expert); } if (!model_get_u32(m, "deepseek4.expert_used_count", &n_expert_used)) { model_get_u32(m, "glm-dsa.expert_used_count", &n_expert_used); } if (!model_get_u32(m, "deepseek4.expert_group_count", &n_expert_groups)) { model_get_u32(m, "glm-dsa.expert_group_count", &n_expert_groups); } if (!model_get_u32(m, "deepseek4.expert_group_used_count", &n_group_used)) { model_get_u32(m, "glm-dsa.expert_group_used_count", &n_group_used); } for (uint64_t i = 0; i < m->n_tensors; i++) { tensor_bytes += m->tensors[i].bytes; params += m->tensors[i].elements; } printf("model: %.*s\n", (int)name.len, name.ptr); printf("arch: %.*s\n", (int)arch.len, arch.ptr); printf("gguf: v%u, %" PRIu64 " metadata keys, %" PRIu64 " tensors\n", m->version, m->n_kv, m->n_tensors); if (layers) printf("layers: %u\n", layers); if (ctx_train) printf("train context: %" PRIu64 "\n", ctx_train); if (n_head || n_head_kv || head_dim || n_swa) { printf("attention: heads=%u kv_heads=%u head_dim=%u swa=%u\n", n_head, n_head_kv, head_dim, n_swa); } if (indexer_heads || indexer_head_dim || indexer_top_k) { printf("indexer: heads=%u head_dim=%u top_k=%u\n", indexer_heads, indexer_head_dim, indexer_top_k); } if (n_expert || n_expert_used || n_expert_groups || n_group_used) { printf("experts: count=%u used=%u groups=%u groups_used=%u\n", n_expert, n_expert_used, n_expert_groups, n_group_used); } model_print_dspark_summary(m); printf("file size: "); print_size(m->size); printf("\n"); printf("tensor bytes described by GGUF: "); print_size(tensor_bytes); printf("\n"); printf("logical parameters: %.2f B\n", (double)params / 1000000000.0); printf("tensor types:\n"); for (uint32_t type = 0; type < sizeof(gguf_types)/sizeof(gguf_types[0]); type++) { uint64_t count = 0; uint64_t bytes = 0; for (uint64_t i = 0; i < m->n_tensors; i++) { if (m->tensors[i].type == type) { count++; bytes += m->tensors[i].bytes; } } if (count != 0) { printf(" %-8s %5" PRIu64 " tensors, ", tensor_type_name(type), count); print_size(bytes); printf("\n"); } } } static ds4_tensor *model_find_tensor(const ds4_model *m, const char *name) { const size_t len = strlen(name); for (uint64_t i = 0; i < m->n_tensors; i++) { if (m->tensors[i].name.len == len && memcmp(m->tensors[i].name.ptr, name, len) == 0) { return &m->tensors[i]; } } return NULL; } static const char *support_kind_name(ds4_support_kind kind) { switch (kind) { case DS4_SUPPORT_MTP_LEGACY: return "legacy MTP"; case DS4_SUPPORT_DSPARK: return "DSpark"; case DS4_SUPPORT_NONE: return "none"; } return "unknown"; } static ds4_support_kind support_model_detect( const ds4_model *m, uint32_t *stages_out, ds4_dspark_summary *summary_out) { if (stages_out) *stages_out = 0; if (summary_out) memset(summary_out, 0, sizeof(*summary_out)); if (!m) return DS4_SUPPORT_NONE; ds4_dspark_summary s = model_dspark_summary(m); if (summary_out) *summary_out = s; if (stages_out) *stages_out = s.stages; if (s.stages >= 3 && s.has_main_proj && s.has_markov_head && s.has_confidence_head) { return DS4_SUPPORT_DSPARK; } if (model_find_tensor(m, "mtp.0.e_proj.weight") && model_find_tensor(m, "mtp.0.h_proj.weight") && model_find_tensor(m, "mtp.0.hc_head_base.weight")) { if (stages_out) *stages_out = s.stages ? s.stages : 1u; return DS4_SUPPORT_MTP_LEGACY; } return DS4_SUPPORT_NONE; } #ifndef DS4_NO_GPU #ifndef __APPLE__ typedef struct { uint64_t off; uint64_t end; } accelerator_tensor_span; static int accelerator_tensor_span_cmp(const void *a, const void *b) { const accelerator_tensor_span *sa = a; const accelerator_tensor_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 uint64_t accelerator_cuda_preload_span_bytes(void) { uint64_t mb = 1024; #ifndef DS4_ROCM_BUILD const char *env = getenv("DS4_CUDA_WEIGHT_PRELOAD_SPAN_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; } #endif if (mb < 64) mb = 64; if (mb > 4096) mb = 4096; return mb * 1048576ull; } static bool accelerator_span_filter_contains(uint64_t off, uint64_t bytes, const uint64_t *span_offsets, const uint64_t *span_sizes, uint32_t span_count) { if (span_count == 0) return true; if (bytes == 0) return true; const uint64_t end = off + bytes; if (end < off) return false; for (uint32_t i = 0; i < span_count; i++) { const uint64_t span_end = span_offsets[i] + span_sizes[i]; if (span_end < span_offsets[i]) return false; if (off >= span_offsets[i] && end <= span_end) return true; } return false; } static bool accelerator_prepare_model_tensor_spans(const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, uint32_t span_count, uint64_t *prepared_out) { uint64_t cap = m->n_tensors; if (cap == 0) { if (prepared_out) *prepared_out = 0; return true; } accelerator_tensor_span *spans = xmalloc((size_t)cap * sizeof(spans[0])); uint64_t nspan = 0; for (uint32_t i = 0; i < span_count; i++) { if (span_offsets[i] > m->size || span_sizes[i] == 0 || span_sizes[i] > m->size - span_offsets[i]) { free(spans); return false; } } for (uint64_t i = 0; i < m->n_tensors; i++) { const ds4_tensor *t = &m->tensors[i]; if (t->bytes == 0) continue; if (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset) { free(spans); return false; } if (!accelerator_span_filter_contains(t->abs_offset, t->bytes, span_offsets, span_sizes, span_count)) { continue; } spans[nspan++] = (accelerator_tensor_span){ .off = t->abs_offset, .end = t->abs_offset + t->bytes, }; } if (nspan == 0) { free(spans); if (prepared_out) *prepared_out = 0; return true; } qsort(spans, (size_t)nspan, sizeof(spans[0]), accelerator_tensor_span_cmp); const uint64_t max_span = accelerator_cuda_preload_span_bytes(); const int tty = ds4_log_is_tty(stderr); const uint64_t progress_step = (tty ? 2ull : 16ull) * 1073741824ull; uint64_t next_progress = progress_step; double last_progress = now_sec(); uint64_t prepared = 0; uint64_t merged = 0; #ifdef DS4_ROCM_BUILD const char *accelerator_name = "ROCm"; #else const char *accelerator_name = "CUDA"; #endif fprintf(stderr, "%sds4: %s preparing model tensor mappings%s", tty ? "\r\033[K" : "", accelerator_name, tty ? ": 0.00 GiB" : "\n"); fflush(stderr); for (uint64_t i = 0; i < nspan;) { uint64_t off = spans[i].off; uint64_t end = spans[i].end; i++; while (i < nspan && spans[i].off <= end + 65536u && spans[i].end - off <= max_span) { if (spans[i].end > end) end = spans[i].end; i++; } char label[96]; snprintf(label, sizeof(label), "tensor-span:%" PRIu64, merged); if (ds4_gpu_cache_model_range(m->map, m->size, off, end - off, label) == 0) { if (tty) fputc('\n', stderr); fprintf(stderr, "ds4: accelerator failed to prepare model tensor span %" PRIu64 " at offset %" PRIu64 "\n", merged, off); free(spans); return false; } prepared += end - off; merged++; const double now = now_sec(); if (prepared >= next_progress || now - last_progress >= (tty ? 2.0 : 10.0)) { if (tty) { fprintf(stderr, "\r\033[Kds4: %s preparing model tensor mappings: %.2f GiB", accelerator_name, (double)prepared / 1073741824.0); } else { fprintf(stderr, "ds4: %s prepared model tensor mappings %.2f GiB\n", accelerator_name, (double)prepared / 1073741824.0); } fflush(stderr); last_progress = now; while (next_progress <= prepared) next_progress += progress_step; } } if (tty) fputc('\n', stderr); free(spans); if (prepared_out) *prepared_out = prepared; return true; } static bool accelerator_cache_q8_tensors(const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, uint32_t span_count) { for (uint64_t i = 0; i < m->n_tensors; i++) { const ds4_tensor *t = &m->tensors[i]; if (t->bytes == 0) continue; if (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset) return false; if (!accelerator_span_filter_contains(t->abs_offset, t->bytes, span_offsets, span_sizes, span_count)) { continue; } char label[128]; snprintf(label, sizeof(label), "tensor:%.*s", (int)t->name.len, t->name.ptr); if (t->type == DS4_TENSOR_Q8_0 && t->ndim == 2 && ds4_gpu_cache_q8_f16_range(m->map, m->size, t->abs_offset, t->bytes, t->dim[0], t->dim[1], label) == 0) { fprintf(stderr, "ds4: accelerator failed to cache dequantized Q8 tensor %.*s\n", (int)t->name.len, t->name.ptr); return false; } } return true; } static bool accelerator_cache_model_tensors(ds4_backend backend, const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, uint32_t span_count) { if (backend != DS4_BACKEND_CUDA) return true; if (!m || !m->map || m->size == 0) return false; #ifndef DS4_ROCM_BUILD if (getenv("DS4_CUDA_DIRECT_MODEL") != NULL) { return true; } #endif const double t0 = now_sec(); uint64_t prepared = 0; if (!accelerator_prepare_model_tensor_spans(m, span_offsets, span_sizes, span_count, &prepared)) { return false; } if (!accelerator_cache_q8_tensors(m, span_offsets, span_sizes, span_count)) return false; const double t1 = now_sec(); #ifdef DS4_ROCM_BUILD const char *accelerator_name = "ROCm"; #else const char *accelerator_name = "CUDA"; #endif fprintf(stderr, "ds4: %s startup model preparation covered %.2f GiB of tensor spans in %.3fs\n", accelerator_name, (double)prepared / 1073741824.0, t1 - t0); return true; } #else static bool accelerator_cache_model_tensors(ds4_backend backend, const ds4_model *m, const uint64_t *span_offsets, const uint64_t *span_sizes, uint32_t span_count) { (void)backend; (void)m; (void)span_offsets; (void)span_sizes; (void)span_count; return true; } #endif #endif /* Return the in-place tensor payload inside the mapped GGUF. */ static const void *tensor_data(const ds4_model *m, const ds4_tensor *t) { return m->map + t->abs_offset; } /* Optional startup pass that touches tensor pages before timing generation. */ static void model_warm_weights(const ds4_model *m) { const uint64_t start = m->tensor_data_pos; const uint64_t end = m->size; if (start >= end) return; const uint64_t page = (uint64_t)sysconf(_SC_PAGESIZE); const uint8_t *p = m->map; volatile uint64_t checksum = 0; const double t0 = now_sec(); fprintf(stderr, "ds4: warming mapped tensor pages: %.2f GiB\n", (double)(end - start) / (1024.0 * 1024.0 * 1024.0)); #if defined(POSIX_MADV_WILLNEED) (void)posix_madvise((void *)(p + start), (size_t)(end - start), POSIX_MADV_WILLNEED); #endif for (uint64_t off = start; off < end; off += page) { checksum += p[off]; } checksum += p[end - 1]; const double t1 = now_sec(); fprintf(stderr, "ds4: warmed tensor pages in %.3fs (checksum=%llu)\n", t1 - t0, (unsigned long long)checksum); } /* ========================================================================= * 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; /* ========================================================================= * 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; } 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 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 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 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 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 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); } 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); } 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 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_Q2_K || type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_Q5_K || type == DS4_TENSOR_Q6_K || type == DS4_TENSOR_Q4_0; } 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); } tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } 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); } 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); } 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_F32 || 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); } 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_IQ3_XXS || type == DS4_TENSOR_IQ2_S || type == DS4_TENSOR_IQ2_M || type == DS4_TENSOR_Q2_K || type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_Q5_K || type == DS4_TENSOR_Q6_K; } 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_IQ3_XXS: return 30; /* 30 bytes per 256 block */ case DS4_TENSOR_IQ2_S: return 56; /* 56 bytes per 256 block */ case DS4_TENSOR_IQ2_M: return 56; /* 56 bytes per 256 block */ 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 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) { 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) { 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) { 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; 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; } *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 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; uint64_t max_bytes = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { uint64_t bytes = 0; if (streaming_layer_routed_expert_bytes(&weights->layer[il], &bytes)) { if (bytes > max_bytes) max_bytes = bytes; } } if (max_bytes == 0) return false; *per_expert_bytes_out = max_bytes; return true; } enum { DS4_STREAMING_PREFILL_HEADROOM_LAYERS = 2 }; static bool ds4_streaming_max_routed_layer_bytes( const ds4_weights *weights, uint64_t *layer_bytes_out) { if (layer_bytes_out) *layer_bytes_out = 0; if (!weights || !layer_bytes_out || DS4_N_EXPERT == 0) return false; uint64_t max_bytes = 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 > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { return false; } const uint64_t layer_bytes = per_expert_bytes * (uint64_t)DS4_N_EXPERT; if (layer_bytes > max_bytes) max_bytes = layer_bytes; } if (max_bytes == 0) return false; *layer_bytes_out = max_bytes; 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 layer_bytes = 0; if (!ds4_streaming_max_routed_layer_bytes(weights, &layer_bytes)) { return false; } if (layer_bytes > UINT64_MAX / (uint64_t)DS4_STREAMING_PREFILL_HEADROOM_LAYERS) { return false; } *bytes_out = layer_bytes * (uint64_t)DS4_STREAMING_PREFILL_HEADROOM_LAYERS; return true; } /* * 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 FIRST routed layer, so those * 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 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 (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 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; /* * 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 uint64_t ds4_add_sat_u64(uint64_t a, uint64_t b) { return a > UINT64_MAX - b ? UINT64_MAX : a + b; } 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 double ds4_bytes_to_gib(uint64_t bytes) { return (double)bytes / 1073741824.0; } 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)) { fprintf(stderr, "ds4: tensor %.*s has type %u (%s), expected a routed expert quant type\n", (int)t->name.len, t->name.ptr, t->type, tensor_type_name(t->type)); exit(1); } 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); } 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 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 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); } 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 (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 (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; } 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 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; } 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]; } return NULL; } /* 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); } 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 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); } 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 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; } 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 (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); if (w->token_embd) { if (w->token_embd->type != DS4_TENSOR_F16 && w->token_embd->type != DS4_TENSOR_Q8_0 && w->token_embd->type != DS4_TENSOR_Q4_K) { tensor_expect_layout(w->token_embd, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); } } 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); if (w->output_hc_fn->type != DS4_TENSOR_F16 && w->output_hc_fn->type != DS4_TENSOR_F32) { 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); } 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); } if (l->hc_attn_fn->type != DS4_TENSOR_F16 && l->hc_attn_fn->type != DS4_TENSOR_F32) { 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_f16_or_q8_0_layout(l->attn_compressor_ape, 2, comp_width, ratio, 0); tensor_expect_f16_or_q8_0_layout(l->attn_compressor_kv, 2, DS4_N_EMBD, comp_width, 0); tensor_expect_f16_or_q8_0_layout(l->attn_compressor_gate, 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_plain_layout(l->indexer_proj, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); tensor_expect_f16_or_q8_0_layout(l->indexer_compressor_ape, 2, index_width, ratio, 0); tensor_expect_f16_or_q8_0_layout(l->indexer_compressor_kv, 2, DS4_N_EMBD, index_width, 0); tensor_expect_f16_or_q8_0_layout(l->indexer_compressor_gate, 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); } if (l->hc_ffn_fn->type != DS4_TENSOR_F16 && l->hc_ffn_fn->type != DS4_TENSOR_F32) { 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); if (l->ffn_gate_inp->type != DS4_TENSOR_F16 && l->ffn_gate_inp->type != DS4_TENSOR_BF16) { 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: 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; } 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 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 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; } 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 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; 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"); uint32_t n_vocab = 0; if (!model_get_u32(m, "deepseek4.vocab_size", &n_vocab)) { if (!model_get_u32(m, "general.vocab_size", &n_vocab)) { n_vocab = 129280u; /* DeepSeek V4 Flash standard vocabulary 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) { 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"); } return; } 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"); } } 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) { 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; } 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); 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 (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_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 - DS4_N_LEADING_DENSE < 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; if (!weights_model_map_decode_static_spans(w, true, true, &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_q4_k(const ds4_model *m, const ds4_weights *w, int token, float *out) { ds4_tensor *te = w->token_embd; if (te->type != DS4_TENSOR_Q4_K || te->ndim != 2) { ds4_die("expected a 2D Q4_K 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 nb = (n + QK_K - 1) / QK_K; const block_q4_K *row = (const block_q4_K *)tensor_data(m, te) + (uint64_t)token * nb; for (uint64_t i = 0; i < nb; i++) { const float d = f16_to_f32(row[i].d); const float dmin = f16_to_f32(row[i].dmin); const uint8_t *qs = row[i].qs; const uint8_t *scales = row[i].scales; float *out_block = out + 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++) { if (j * 32 + l < (int)(n - i * QK_K)) { const int q = (qs[byte_off + l] >> shift) & 0x0F; out_block[j * 32 + l] = scale * (float)q - minv; } } } } } 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; case DS4_TENSOR_Q4_K: embed_token_q4_k(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_any(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_any(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_any(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); 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"); 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++) { 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++) { 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]; const uint64_t vocab_dim = weights->output->dim[1]; 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); 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++) { /* 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++) { 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) fprintf(stderr, "ds4: DBG layer %u router_selection failed\n", il); 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) { fprintf(stderr, "ds4: DBG layer %u entering cuda_stream_selected_load\n", il); ok = metal_graph_decode_cuda_selected_load(g, model, layer, il, gate_expert_bytes, down_expert_bytes); if (!ok) fprintf(stderr, "ds4: DBG layer %u cuda_selected_load failed\n", il); } if (selected_readahead_shared_delay) { fprintf(stderr, "ds4: DBG layer %u entering selected_readahead_shared_delay\n", il); if (ok) { ok = metal_graph_decode_selected_readahead_override(g, model, layer, il, gate_expert_bytes, down_expert_bytes); if (!ok) fprintf(stderr, "ds4: DBG layer %u readahead_override failed\n", il); } 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) { fprintf(stderr, "ds4: DBG layer %u entering overlap_selected_shared\n", il); uint64_t selected_event = 0; if (ok) ok = ds4_gpu_signal_selected_readback_ready(&selected_event) != 0; if (!ok) fprintf(stderr, "ds4: DBG layer %u signal_selected_readback failed\n", il); 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; fprintf(stderr, "ds4: DBG layer %u default routed_moe_one ok=%d gate_type=%d down_type=%d\n", il, (int)ok, (int)layer->ffn_gate_exps->type, (int)layer->ffn_down_exps->type); 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; if (!ok) fprintf(stderr, "ds4: DBG head stage rms_norm_plain failed\n"); 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; if (!ok) fprintf(stderr, "ds4: DBG head stage matmul_f16 output_hc_fn failed (type=%d)\n", (int)weights->output_hc_fn->type); 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; if (!ok) fprintf(stderr, "ds4: DBG head stage output_hc_weights failed\n"); 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; if (!ok) fprintf(stderr, "ds4: DBG head stage hc_weighted_sum failed\n"); } 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; if (!ok) fprintf(stderr, "ds4: DBG head stage rms_norm_weight failed\n"); } 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) { fprintf(stderr, "ds4: DBG dense quant output matmul failed (type=%d in=%llu out=%llu)\n", (int)weights->output->type, (unsigned long long)DS4_N_EMBD, (unsigned long long)vocab_dim); } } 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_BF16) { return ds4_gpu_matmul_bf16_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_any(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_any(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_any(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; } fprintf(stderr, "ds4: DBG eval streaming pos=%u token=%d logits=%d\n", pos, token, logits != NULL); 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) fprintf(stderr, "ds4: DBG eval begin_commands failed\n"); 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) fprintf(stderr, "ds4: DBG eval embed_token_hc failed (type=%d dim1=%lld token=%d)\n", (int)weights->token_embd->type, (long long)weights->token_embd->dim[1], token); } 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]); if (!ok) fprintf(stderr, "ds4: DBG streaming output head encode failed (batched)\n"); } const double t_encoded = (profile || throttle) ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; if (!ok) fprintf(stderr, "ds4: DBG batched end_commands failed\n"); 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; if (!ok) fprintf(stderr, "ds4: DBG batched logits tensor_read failed\n"); } 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; if (!ok) fprintf(stderr, "ds4: DBG eval embed end_commands failed\n"); 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; if (!ok) fprintf(stderr, "ds4: DBG eval layer %u begin_commands failed\n", il); 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); if (!ok) fprintf(stderr, "ds4: DBG eval layer %u encode_decode_layer failed\n", il); 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; if (!ok) fprintf(stderr, "ds4: DBG eval layer %u end_commands failed\n", il); 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) fprintf(stderr, "ds4: DBG head begin_commands failed\n"); 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; if (!ok && logits) fprintf(stderr, "ds4: DBG head end_commands failed\n"); 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; if (!ok) fprintf(stderr, "ds4: DBG head logits tensor_read failed\n"); } 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; fprintf(stderr, "ds4: DBG prefill loop i=%u pos=%u last=%d logits=%d\n", i, pos, (int)last, token_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++) { 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; fprintf(stderr, "ds4: DBG prefill_raw_swa n_tokens=%d streaming=%d\n", n_tokens, (int)metal_graph_use_streaming_decode_prefill_range(g, weights, 0, (uint32_t)n_tokens)); 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 uint32_t glm_graph_full_indexer_layer_count(uint32_t normal_layers) { return normal_layers ? glm_graph_full_indexer_layer_count_range(0, normal_layers - 1u) : 0; } 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( uint32_t ctx, uint32_t work_ctx, uint32_t compact_cap, bool ssd_streaming) { ds4_context_memory m = {0}; const uint32_t normal_layers = glm_graph_normal_layer_count(); 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)normal_layers * 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( normal_layers, glm_graph_full_indexer_layer_count(normal_layers), compact_cap); } m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; return m; } 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; const ds4_context_memory mem = ds4_context_memory_estimate_with_prefill_mode(e->backend, ctx_size, e->prefill_chunk, e->ssd_streaming); 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); 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 (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; /* 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; } } } 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; } 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 (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; } 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 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 (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 == 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]; 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]; 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]; } 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 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); } } /* 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); 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(); 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; } 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); } 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); } 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 + 1 >= ctx_size) { pos++; break; } 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); } 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: 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); free(logits); cpu_decode_scratch_free(&decode_scratch); kv_cache_free(&cache); return 0; } #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 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_model_map_token_spans(weights, &spans)) { max_bytes = model_map_span_vec_total_bytes(&spans); free(spans.v); } if (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 = 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); const double default_reserve_gib = glm_graph_memory_guard_default_reserve_gib(budget_base, model_bytes); 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"); } 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) 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) 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 && !shared_done) DS4_GLM_ENCODE_FFN_BATCH_SHARED(); #undef DS4_GLM_ENCODE_FFN_BATCH_SHARED 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); 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; } /* 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. */ static int generate_metal_graph_raw_swa( 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, int power_percent, uint32_t prefill_chunk, const char * directional_steering_file, 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 GPU graph generation with graph prefill\n"); if (prompt->len <= 0 || prompt->len > ctx_size) { fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); return 1; } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (power_percent > 0 && power_percent < 100) { fprintf(stderr, "ds4: --power is not supported by the GLM Metal path yet\n"); return 1; } if (prefill_chunk != 0) { fprintf(stderr, "ds4: --prefill-chunk is not supported by the GLM Metal path; " "GLM uses graph-selected prefill chunks\n"); return 1; } if ((directional_steering_file && directional_steering_file[0]) || directional_steering_attn != 0.0f || directional_steering_ffn != 0.0f) { fprintf(stderr, "ds4: directional steering is not supported by the GLM Metal path yet\n"); return 1; } return generate_glm_metal_argmax(model, vocab, weights, prompt, n_predict, ctx_size, quality, ssd_streaming, ssd_streaming_cold, ssd_streaming_preload_experts, ssd_streaming_cache_bytes, ssd_streaming_prefill_headroom_bytes, emit, done, emit_ud, progress, progress_ud); } const uint32_t prefill_cap = metal_graph_prefill_cap_for_prompt(prompt->len, prefill_chunk); const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, prefill_cap); if (prefill_cap < (uint32_t)prompt->len) { fprintf(stderr, "ds4: using chunked GPU prefill (%u-token chunks for %d prompt tokens)\n", prefill_cap, prompt->len); } 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, prefill_cap, false, NULL, false, NULL); if (!ok) { fprintf(stderr, "ds4: failed to allocate GPU graph runtime\n"); return 1; } g.quality = quality; g.ssd_streaming = ssd_streaming; g.ssd_streaming_cold = ssd_streaming_cold; g.streaming_preload_experts = ssd_streaming_preload_experts; g.power_percent = power_percent > 0 ? (uint32_t)power_percent : 100u; if (!metal_graph_load_directional_steering(&g, directional_steering_file, directional_steering_attn, directional_steering_ffn)) { metal_graph_free(&g); return 1; } const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; if (memory_report) ds4_gpu_print_memory_report("after graph alloc"); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); const bool trace_top = getenv("DS4_TRACE_TOP") != NULL; const bool token_timing = getenv("DS4_TOKEN_TIMING") != NULL; const double t_prefill0 = now_sec(); if (prefill_cap < (uint32_t)prompt->len) { ok = metal_graph_prefill_chunked(&g, model, weights, prompt, prompt->len, logits, false, progress, progress_ud, progress, progress_ud, NULL, NULL, NULL); } else { ok = metal_graph_prefill_raw_swa(&g, model, weights, prompt, prompt->len, logits, true, progress, progress_ud, NULL, NULL, NULL); } const double t_prefill1 = now_sec(); if (memory_report) ds4_gpu_print_memory_report("after prefill"); if (!ok) { free(logits); metal_graph_free(&g); return 1; } const char *dump_prefill_logits = getenv("DS4_METAL_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); metal_graph_free(&g); return 1; } fprintf(stderr, "ds4: wrote GPU prefill logits to %s\n", dump_prefill_logits); } int pos = prompt->len; int n_generated = 0; int n_decode_eval = 0; 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); } 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 + 1 >= ctx_size) { pos++; break; } const double t_eval0 = token_timing ? now_sec() : 0.0; ok = metal_graph_eval_token_raw_swa(&g, model, weights, (uint32_t)token, (uint32_t)pos, logits); if (!ok) break; if (token_timing) { const double t_eval1 = now_sec(); fprintf(stderr, "ds4: gpu 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: 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 graph free"); free(logits); metal_graph_free(&g); return ok ? 0 : 1; } #endif #ifdef DS4_NO_GPU ds4_context_memory ds4_context_memory_estimate_with_prefill_mode( ds4_backend backend, int ctx_size, uint32_t prefill_chunk, bool ssd_streaming) { (void)backend; (void)prefill_chunk; (void)ssd_streaming; ds4_context_memory m = {0}; uint32_t ctx = ctx_size > 0 ? (uint32_t)ctx_size : 1u; 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); } #endif /* Per-layer KV byte estimate at a given context size for the CUDA / * Metal graph backend. Mirrors the per-layer * accounting the graph backend uses at allocation time, so the multi-GPU * layer packer (engine_compute_entry_bytes) can price each layer * individually instead of dividing a global total uniformly across all * layers. Heterogeneous attention compression (raw / ratio-4 / * ratio-128) is then represented faithfully in the placement plan. * * Visible in both GPU and DS4_NO_GPU builds so the placement-test * harness (which builds ds4.c with -DDS4_NO_GPU) can exercise the same * packer code paths. * * Accounting (matches the graph branch of ds4_context_memory_estimate): * raw_bytes_per_layer = raw_cap * DS4_N_HEAD_DIM * sizeof(float) * applied to EVERY layer (the GPU graph * allocates a raw entry for every layer, * with the sliding-window cap applied at * execution time, not at planning time). * compressed_bytes_per_layer = ratio==0 -> 0 * ratio!=0 -> (ctx/ratio + 2) * DS4_N_HEAD_DIM * * sizeof(float) * (uses sizeof(float) rather * than the F16/F32 choice * made at GPU build time — * over-estimates on Apple * Metal (F16 cache) which is * the desired conservative * posture for the planner). * + (ratio==4 ? (ctx/ratio + 2) * * DS4_N_INDEXER_HEAD_DIM * * sizeof(float) * : 0) * * Per-tier scratch buffers (indexer_scores_by_tier, comp_mask_by_tier, * attn_comp_stage_by_tier, chunked-prefill batch_*_by_tier, head extras) * are NOT included here. They are reserved separately by * engine_per_tier_graph_overhead_bytes() and pre-subtracted from each * device's vram_bytes in engine_classify_multi_tier. Counting them per * layer here would double-count them once per layer. * * Invariant (by construction): sum over il of * engine_per_layer_kv_bytes_planner(il, ctx, prefill_chunk) == * DS4_N_LAYER * raw_cap * DS4_N_HEAD_DIM * sizeof(float) * + sum_over_ratio_layers(compressed_per_layer) */ /* Planner equivalents of metal_graph_prefill_cap_for_prompt and * metal_graph_raw_cap_for_context. The graph variants live inside * `#ifndef DS4_NO_GPU` blocks; these inline copies replicate the same numeric * math so the planner can be called from both build flavors. * * If the env-knob behavior in the graph helpers ever changes, update * these two helpers in lockstep. Tested via: * sum(engine_per_layer_kv_bytes_planner(il, ctx, prefill_chunk) * for il in 0..N) == * F32-sized ds4_context_memory_estimate(CUDA, ctx).total_bytes. */ static uint32_t engine_planner_prefill_cap(int prompt_len, uint32_t requested_chunk) { return ds4_prefill_cap_for_prompt(prompt_len, requested_chunk); } static uint32_t engine_planner_raw_cap(int ctx_size, uint32_t prefill_cap) { if (ctx_size <= 0) return 1; 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; /* Pad to 256-row multiple so the planner matches the graph layout. */ 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; /* align_up to 256 — inline since align_up is defined upstream. */ const uint64_t align = 256u; wanted = (wanted + align - 1u) & ~(align - 1u); if (wanted > 8192u) wanted = 8192u; uint32_t raw_cap = (uint32_t)wanted; if (raw_cap < raw_window) raw_cap = raw_window; /* Env override (matches metal_graph_raw_cap_for_context behavior). */ 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; } } return raw_cap; } /* Mirrors the GPU cache-storage choice while remaining visible to the * DS4_NO_GPU placement tests. */ #if defined(__APPLE__) #define DS4_PLANNER_ATTN_COMP_CACHE_F16 1 #else #define DS4_PLANNER_ATTN_COMP_CACHE_F16 0 #endif static bool engine_glm_layer_uses_full_indexer(uint32_t il) { if (il < DS4_N_LEADING_DENSE) return true; return il >= 6u && ((il - 6u) % 4u) == 0u; } /* GLM keeps one compact DSA row for every logical context position. These * caches are allocated on the same tier as their transformer layer, so they * must be priced per layer by the multi-GPU packer. */ static size_t engine_glm_per_layer_kv_bytes_planner(uint32_t il, int ctx_size) { if (ctx_size <= 0 || DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) return 0; const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; if (il >= normal_layers) return 0; const uint64_t elem_bytes = DS4_PLANNER_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float); uint64_t row_width = (uint64_t)DS4_N_KV_LORA + DS4_N_ROT; if (engine_glm_layer_uses_full_indexer(il)) { row_width += DS4_N_INDEXER_HEAD_DIM; } const uint64_t bytes = (uint64_t)(uint32_t)ctx_size * row_width * elem_bytes; return bytes > SIZE_MAX ? SIZE_MAX : (size_t)bytes; } static size_t engine_per_layer_kv_bytes_planner(uint32_t il, int ctx_size, uint32_t prefill_chunk) { if (ctx_size <= 0) return 0; if (il >= DS4_N_LAYER) return 0; const uint32_t ctx = (uint32_t)ctx_size; /* Raw KV cache: every layer gets a raw entry sized by raw_cap, the * same value the GPU graph requests per layer at * metal_graph_alloc_kv_cache_tensor_on(. , raw_cap * DS4_N_HEAD_DIM * * sizeof(float)). raw_cap factors in raw_window padding + prefill_cap * and clamps to [raw_window, 8192]. */ const uint32_t prefill_cap = engine_planner_prefill_cap((int)ctx, prefill_chunk); const uint32_t raw_cap = engine_planner_raw_cap((int)ctx, prefill_cap); size_t bytes = (size_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float); /* Compressed + indexer for ratio != 0 layers. Uses sizeof(float) * unconditionally (over-estimates on Apple Metal F16 cache, which * is the desired conservative posture per spec criterion 8). */ const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0) { const uint32_t layer_comp_cap = ctx / ratio + 2u; bytes += (size_t)layer_comp_cap * DS4_N_HEAD_DIM * sizeof(float); if (ratio == 4) { bytes += (size_t)layer_comp_cap * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } } /* Per-tier scratch (indexer_scores_by_tier, comp_mask_by_tier, * attn_comp_stage_by_tier, chunked-prefill batch_*_by_tier, head * extras) is accounted for separately by * engine_per_tier_graph_overhead_bytes() and pre-subtracted from * each device's vram_bytes in engine_classify_multi_tier. This * helper returns per-layer KV/index cache ONLY. Charging scratch * both here (per layer) and there (per tier) would double-count — * at large ctx the duplicate can falsely refuse layouts that * actually fit. */ return bytes; } /* 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). * * 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 * accounts for tensor weights + per-layer KV — it never reserved this * per-tier scratch, so a layout that fit by entry_bytes could still * late-OOM at session_create. This helper returns the exact byte total, * which is then pre-subtracted from EVERY device's vram_bytes inside * engine_classify_multi_tier (conservative: even tiers that end up unused * 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. * * Uses g_ds4_shape compile-time constants (DS4_N_*) and ctx-derived caps * from engine_planner_prefill_cap. The runtime computes some dims from * model weights (e.g. q_rank = layer->attn_q_a->dim[1]), but * tensor_expect_layout enforces those dims == DS4_N_LORA_Q etc., so the * compile-time aliases used here are byte-equivalent. * * Visible in both GPU and DS4_NO_GPU builds. */ static size_t engine_per_tier_graph_overhead_bytes(const ds4_engine *e) { #ifndef DS4_NO_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { const int est_ctx = e && e->placement_ctx_hint > 0 ? e->placement_ctx_hint : 4096; const uint32_t ctx = (uint32_t)est_ctx; const bool ssd_streaming = e && e->ssd_streaming; 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); const ds4_context_memory mem = glm_graph_context_memory_estimate_for_compact_cap( ctx, work_ctx, compact_cap, ssd_streaming); return mem.scratch_bytes > SIZE_MAX ? SIZE_MAX : (size_t)mem.scratch_bytes; } #endif /* Local dim aliases — same values the runtime per-tier scratch loop * reads. The runtime reads q_rank etc. from layer weights; the * config_validate_model path (called from ds4_engine_open_internal * before classify) enforces dim equivalence with DS4_N_*, so these * aliases are byte-equivalent. */ 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 = DS4_N_LORA_Q; 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 = DS4_N_OUT_GROUP ? (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP) : 0; const uint64_t shared_dim = DS4_N_FF_EXP; const uint64_t routed_mid_dim = DS4_N_FF_EXP; const uint64_t vocab_dim = DS4_N_VOCAB; uint64_t output_logits_elems = vocab_dim; const int planner_n_gpus = e ? e->gpu_cfg.n_gpus : 0; #if !defined(__APPLE__) const char *tp_output_env = getenv("DS4_CUDA_TP_OUTPUT"); const bool tp_decode_requested = e && e->cuda_tensor_parallel; const bool tp_output_requested = !tp_output_env || !tp_output_env[0] || strcmp(tp_output_env, "0") != 0; if (planner_n_gpus >= 2 && (planner_n_gpus & 1) == 0 && tp_decode_requested && tp_output_requested) { uint32_t output_ways = 8u; const char *ways_env = getenv("DS4_CUDA_TP_OUTPUT_WAYS"); if (ways_env && ways_env[0]) { char *end = NULL; const unsigned long parsed = strtoul(ways_env, &end, 10); output_ways = end != ways_env && *end == '\0' && parsed >= 2u && parsed <= DS4_MAX_GPUS ? (uint32_t)parsed : 2u; } if (output_ways > (uint32_t)planner_n_gpus) { output_ways = (uint32_t)planner_n_gpus; } const uint64_t max_shard_vocab = (vocab_dim + output_ways - 1u) / output_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; } } #else (void)planner_n_gpus; #endif 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; /* ctx-derived caps. The runtime uses prefill_cap derived from the * chunked-prefill batch size. The planner has no prompt yet, so it uses the * placement context and the engine's effective chunk setting. */ const int est_ctx = (e->placement_ctx_hint > 0) ? e->placement_ctx_hint : 4096; 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. * 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. */ 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); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } if (min_ratio == UINT32_MAX) { min_ratio = est_ctx > 0 ? (uint32_t)est_ctx : 1u; } uint32_t comp_cap = (uint32_t)est_ctx / min_ratio + 2u; if (comp_cap < 2u) comp_cap = 2u; uint32_t attn_comp_stage_cap = 0; if (DS4_PLANNER_ATTN_COMP_CACHE_F16) { attn_comp_stage_cap = prefill_cap / min_ratio + 2u; if (attn_comp_stage_cap < 2u) attn_comp_stage_cap = 2u; } const uint64_t pc = (uint64_t)prefill_cap; size_t total = 0; /* === Class P decode HC scratch (mirrors ds4.c:10664-10686). 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 */ total += hc_dim * sizeof(float); /* flat_hc_by_tier */ total += mix_hc * sizeof(float); /* hc_mix_by_tier */ total += mix_hc * sizeof(float); /* hc_split_by_tier */ /* hc_pre_by_tier, hc_post_by_tier, hc_comb_by_tier — VIEWS of hc_split. */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* attn_cur_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* attn_norm_by_tier */ total += q_rank * sizeof(float); /* qr_by_tier */ total += q_rank * sizeof(float); /* qr_norm_by_tier */ total += q_dim * sizeof(float); /* q_by_tier */ 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). === */ 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) { total += (uint64_t)attn_comp_stage_cap * DS4_N_HEAD_DIM * sizeof(float); /* attn_comp_stage_by_tier */ } total += indexer_q_dim * sizeof(float); /* indexer_q_by_tier */ total += (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); /* indexer_weights_by_tier */ total += (uint64_t)comp_cap * pc * sizeof(float); /* indexer_scores_by_tier */ total += (uint64_t)comp_cap * pc * sizeof(float); /* comp_mask_by_tier */ const uint64_t top_k = (uint64_t)(DS4_N_INDEXER_TOP_K ? DS4_N_INDEXER_TOP_K : 1u); total += top_k * pc * sizeof(uint32_t); /* comp_selected_by_tier */ total += q_dim * sizeof(float); /* heads_by_tier */ total += low_dim * sizeof(float); /* attn_low_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* attn_out_by_tier */ total += hc_dim * sizeof(float); /* after_attn_hc_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* ffn_cur_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* ffn_norm_by_tier */ total += shared_dim * sizeof(float); /* shared_gate_by_tier */ total += shared_dim * sizeof(float); /* shared_up_by_tier */ total += shared_dim * sizeof(float); /* shared_mid_by_tier */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* shared_out_by_tier */ total += (uint64_t)DS4_N_EXPERT * sizeof(float); /* router_logits_by_tier */ total += (uint64_t)DS4_N_EXPERT * sizeof(float); /* router_probs_by_tier */ total += (uint64_t)DS4_N_EXPERT_USED * sizeof(int); /* router_selected_by_tier */ total += (uint64_t)DS4_N_EXPERT_USED * sizeof(float); /* router_weights_by_tier */ total += (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* routed_gate */ total += (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* routed_up */ total += (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* routed_mid */ total += (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); /* routed_down */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* routed_out */ total += DS4_CUDA_TP_PEER_TMP_BYTES; /* tp_peer_tmp_by_tier */ total += hc_dim * sizeof(float); /* after_ffn_hc */ /* === Class P chunked-prefill batch scratch (mirrors allocation above). * These are the LARGEST per-tier allocations (e.g. batch_cur_hc = * pc * hc_dim * float = ~256 MiB at default prefill_cap=4096). Without * them the pre-subtract is meaningless for any non-trivial ctx. === */ total += pc * hc_dim * sizeof(float); /* batch_cur_hc_by_tier */ total += pc * hc_dim * sizeof(float); /* batch_next_hc_by_tier */ total += pc * hc_dim * sizeof(float); /* batch_flat_hc_by_tier */ total += pc * mix_hc * sizeof(float); /* batch_hc_mix_by_tier */ total += pc * mix_hc * sizeof(float); /* batch_hc_split_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_attn_cur_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_attn_norm_by_tier */ total += pc * q_rank * sizeof(float); /* batch_qr_by_tier */ total += pc * q_rank * sizeof(float); /* batch_qr_norm_by_tier */ total += pc * q_dim * sizeof(float); /* batch_q_by_tier */ total += pc * (uint64_t)DS4_N_HEAD_DIM * sizeof(float);/* batch_kv_raw_by_tier */ total += pc * (uint64_t)DS4_N_HEAD_DIM * sizeof(float);/* batch_kv_by_tier */ total += pc * comp_width_max * sizeof(float); /* batch_comp_kv_by_tier */ total += pc * comp_width_max * sizeof(float); /* batch_comp_sc_by_tier */ total += pc * indexer_q_dim * sizeof(float); /* batch_indexer_q_by_tier */ total += pc * (uint64_t)DS4_N_INDEXER_HEAD * sizeof(float); /* batch_indexer_weights */ total += pc * q_dim * sizeof(float); /* batch_heads_by_tier */ total += pc * low_dim * sizeof(float); /* batch_attn_low_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_attn_out_by_tier */ total += pc * group_dim * sizeof(float); /* batch_group_tmp_by_tier */ total += pc * (uint64_t)DS4_N_LORA_O * sizeof(float); /* batch_low_tmp_by_tier */ total += pc * hc_dim * sizeof(float); /* batch_after_attn_hc_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_ffn_cur_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_ffn_norm_by_tier */ total += pc * shared_dim * sizeof(float); /* batch_shared_gate_by_tier */ total += pc * shared_dim * sizeof(float); /* batch_shared_up_by_tier */ total += pc * shared_dim * sizeof(float); /* batch_shared_mid_by_tier */ total += pc * (uint64_t)DS4_N_EMBD * sizeof(float); /* batch_shared_out_by_tier */ total += pc * (uint64_t)DS4_N_EXPERT * sizeof(float); /* batch_router_logits_by_tier */ total += pc * (uint64_t)DS4_N_EXPERT * sizeof(float); /* batch_router_probs_by_tier */ total += pc * (uint64_t)DS4_N_EXPERT_USED * sizeof(int); /* batch_router_selected */ total += pc * (uint64_t)DS4_N_EXPERT_USED * sizeof(float); /* batch_router_weights */ total += pc * (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* batch_routed_gate */ total += pc * (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* batch_routed_up */ total += pc * (uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float); /* batch_routed_mid */ total += pc * (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float); /* batch_routed_down */ 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). * 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. === */ 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 */ total += (uint64_t)DS4_N_EMBD * sizeof(float); /* output_norm_by_tier */ /* Output-TP sessions may load DSpark after placement, so reserve its * block-verifier shard scratch conservatively even before support-model * detection. Plain decode allocates only the one-row subset. */ total += output_logits_elems * sizeof(float); /* logits_by_tier */ return total; } /* ========================================================================= * Engine API and Process Lock. * ========================================================================= * * The public entry points acquire the single instance lock, open the GGUF with * the backend-appropriate mmap policy, and expose tokenized prompt operations * to the CLI and server. */ const char *ds4_backend_name(ds4_backend backend) { switch (backend) { case DS4_BACKEND_METAL: return "metal"; case DS4_BACKEND_CUDA: #ifdef DS4_ROCM_BUILD return "rocm"; #else return "cuda"; #endif case DS4_BACKEND_CPU: return "cpu"; } return "unknown"; } static void ds4_linux_graph_backend_set_oom_score(ds4_backend backend) { #if defined(__linux__) && !defined(DS4_NO_GPU) static bool attempted = false; if (attempted) return; attempted = true; const int score = 1000; FILE *fp = fopen("/proc/self/oom_score_adj", "w"); if (!fp) { fprintf(stderr, "ds4: failed to set Linux %s backend oom_score_adj=%d: %s\n", ds4_backend_name(backend), score, strerror(errno)); return; } if (fprintf(fp, "%d\n", score) < 0) { const int err = errno; fclose(fp); fprintf(stderr, "ds4: failed to write Linux %s backend oom_score_adj=%d: %s\n", ds4_backend_name(backend), score, strerror(err)); return; } if (fclose(fp) != 0) { fprintf(stderr, "ds4: failed to close Linux %s backend oom_score_adj=%d: %s\n", ds4_backend_name(backend), score, strerror(errno)); return; } fprintf(stderr, "ds4: Linux %s backend set oom_score_adj=%d\n", ds4_backend_name(backend), score); #else (void)backend; #endif } bool ds4_think_mode_enabled(ds4_think_mode mode) { return mode == DS4_THINK_HIGH || mode == DS4_THINK_MAX; } const char *ds4_think_mode_name(ds4_think_mode mode) { switch (mode) { case DS4_THINK_NONE: return "none"; case DS4_THINK_HIGH: return "high"; case DS4_THINK_MAX: return "max"; } return "unknown"; } const char *ds4_think_max_prefix(void) { return DS4_REASONING_EFFORT_MAX_PREFIX; } uint32_t ds4_think_max_min_context(void) { return DS4_THINK_MAX_MIN_CONTEXT; } ds4_think_mode ds4_think_mode_for_context(ds4_think_mode mode, int ctx_size) { if (mode == DS4_THINK_MAX && (uint32_t)(ctx_size > 0 ? ctx_size : 0) < DS4_THINK_MAX_MIN_CONTEXT) { return DS4_THINK_HIGH; } return mode; } static void ds4_release_instance_lock(void) { if (g_ds4_lock_fd >= 0) { close(g_ds4_lock_fd); g_ds4_lock_fd = -1; } } /* Refuse to start a second ds4 process. The model can map tens of GiB, so a * stale accidental second run is more dangerous than a normal CLI error. */ static void ds4_acquire_instance_lock(void) { const char *path = getenv("DS4_LOCK_FILE"); if (!path || !path[0]) path = "/tmp/ds4.lock"; const int fd = open(path, O_RDWR | O_CREAT, 0600); if (fd < 0) { fprintf(stderr, "ds4: failed to open lock file %s: %s\n", path, strerror(errno)); exit(2); } (void)fcntl(fd, F_SETFD, FD_CLOEXEC); if (flock(fd, LOCK_EX | LOCK_NB) != 0) { if (errno == EWOULDBLOCK) { char buf[64]; const ssize_t n = pread(fd, buf, sizeof(buf) - 1, 0); long owner = -1; if (n > 0) { buf[n] = '\0'; char *end = NULL; owner = strtol(buf, &end, 10); } if (owner > 0) { fprintf(stderr, "ds4: another ds4 process is already running (pid %ld); refusing to start\n", owner); } else { fprintf(stderr, "ds4: another ds4 process is already running; refusing to start\n"); } close(fd); exit(2); } fprintf(stderr, "ds4: failed to lock %s: %s\n", path, strerror(errno)); close(fd); exit(2); } if (ftruncate(fd, 0) != 0) { fprintf(stderr, "ds4: failed to truncate lock file %s: %s\n", path, strerror(errno)); close(fd); exit(2); } dprintf(fd, "%ld\n", (long)getpid()); g_ds4_lock_fd = fd; atexit(ds4_release_instance_lock); } #ifndef DS4_NO_GPU typedef struct { uint32_t n_comp[DS4_MAX_LAYER]; uint32_t n_index_comp[DS4_MAX_LAYER]; uint32_t mtp_n_raw; uint32_t dspark_cache_start; uint32_t dspark_cache_token_start; uint32_t dspark_cache_len; bool light; } ds4_spec_frontier; typedef struct ds4_dspark_spec_stats { uint64_t cycles; uint64_t first_tokens; uint64_t proposed_tokens; uint64_t accepted_draft_tokens; uint64_t full_accepts; uint64_t partial_accepts; uint64_t first_misses; uint64_t no_draft; uint64_t no_room; uint64_t invalid_draft; uint64_t draft_len_hist[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; uint64_t accepted_len_hist[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; uint64_t scheduler_skips; uint64_t tail_skips; uint64_t verifier_unavailable; uint64_t verifier_errors; double target_ms; double saved_ms; double propose_ms; double propose_stage0_ms; double propose_setup_ms; double propose_cache_ms; double propose_chain_ms; double propose_hidden_ms; double propose_conf0_ms; double propose_logits_ms; double propose_markov_ms; double propose_confidence_ms; double snapshot_ms; double verify_ms; double verify_upload_ms; double verify_layer_ms; double verify_head_ms; double verify_read_ms; uint64_t verifier_fused_head; double replay_ms; double total_ms; } ds4_dspark_spec_stats; #endif struct ds4_session { ds4_engine *engine; ds4_dist_session *distributed; uint64_t tp_session_id; #ifndef DS4_NO_GPU ds4_gpu_graph graph; ds4_glm_gpu_graph glm_graph; bool glm_graph_ready; uint32_t glm_dense_cache_len; /* GLM MTP speculative state (--glm-mtp, greedy only). */ int glm_mtp_draft; int glm_mtp_have; int glm_spec_inside; uint32_t glm_mtp_min_pos; float *glm_mtp_hc; float *glm_mtp_logits0; ds4_spec_frontier greedy_splitkv_anchor; #endif ds4_kv_cache cpu_cache; ds4_cpu_decode_scratch cpu_scratch; token_vec checkpoint; token_vec greedy_splitkv_segment; float *logits; float *sample_probs; float *mtp_logits; int greedy_splitkv_anchor_len; #ifndef DS4_NO_GPU float *spec_row_logits; float *dspark_markov_bias; float *dspark_conf_features; size_t dspark_conf_features_cap; #endif int mtp_draft_token; #ifndef DS4_NO_GPU int dspark_draft_tokens[DS4_DSPARK_MAX_BLOCK_SIZE]; uint32_t dspark_draft_len; uint32_t dspark_sched_cycles; uint32_t dspark_sched_accepted; uint32_t dspark_sched_no_draft; uint32_t dspark_sched_skip; uint32_t dspark_sched_lifetime_accepted; double dspark_sched_life_extra_ms; double dspark_sched_life_saved_ms; double dspark_sched_extra_ms; double dspark_sched_saved_ms; double dspark_last_target_eval_ms; double dspark_last_propose_ms; float dspark_last_confidence0; bool dspark_draft_valid; bool dspark_sched_skipped_cycle; bool dspark_sched_long_accept_seen; bool dspark_last_confidence0_valid; ds4_dspark_spec_stats dspark_stats; #endif uint64_t mtp_probe_total; uint64_t mtp_probe_hit; 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; uint32_t prefill_cap; int ctx_size; bool checkpoint_valid; bool mtp_draft_valid; bool greedy_splitkv_anchor_valid; }; #ifndef DS4_NO_GPU static bool ds4_dspark_stats_enabled(void); static void ds4_dspark_stats_note_len( uint64_t hist[DS4_DSPARK_MAX_BLOCK_SIZE + 1u], uint32_t len) { if (len > DS4_DSPARK_MAX_BLOCK_SIZE) len = DS4_DSPARK_MAX_BLOCK_SIZE; hist[len]++; } static uint32_t ds4_dspark_env_u32(const char *name, uint32_t fallback) { const char *env = getenv(name); if (!env || !env[0]) return fallback; char *end = NULL; errno = 0; unsigned long v = strtoul(env, &end, 10); if (end == env || errno != 0 || v > UINT32_MAX) return fallback; return (uint32_t)v; } static bool ds4_dspark_scheduler_enabled(void) { const char *env = getenv("DS4_DSPARK_SCHEDULER"); return !env || !env[0] || strcmp(env, "0") != 0; } static uint32_t ds4_dspark_scheduler_window(void) { uint32_t v = ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_WINDOW", 4); return v ? v : 4; } static uint32_t ds4_dspark_scheduler_skip_cycles(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_SKIP", 2); } static uint32_t ds4_dspark_scheduler_slow_skip_cycles(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_SLOW_SKIP", 4); } static uint32_t ds4_dspark_scheduler_min_avg_milli(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_MIN_AVG_MILLI", 1500); } static uint32_t ds4_dspark_scheduler_max_ms_per_accept_milli(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_MAX_MS_PER_ACCEPT_MILLI", 28000); } static uint32_t ds4_dspark_scheduler_max_extra_saved_ratio_milli(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_MAX_EXTRA_SAVED_RATIO_MILLI", 1000); } static uint32_t ds4_dspark_scheduler_break_even_window(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_BREAK_EVEN_WINDOW", 0); } static uint32_t ds4_dspark_scheduler_no_draft_skip_cycles(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_NO_DRAFT_SKIP", 3); } static uint32_t ds4_dspark_scheduler_short_accept_no_draft_skip_cycles(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_SHORT_ACCEPT_NO_DRAFT_SKIP", 4); } static uint32_t ds4_dspark_scheduler_cold_low_confidence_skip_cycles(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_COLD_LOW_CONFIDENCE_SKIP", 7); } static uint32_t ds4_dspark_scheduler_tail_min_tokens(void) { return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_TAIL_MIN_TOKENS", 10); } static float ds4_dspark_scheduler_cold_low_confidence_threshold(void) { return (float)ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_COLD_LOW_CONFIDENCE_MILLI", 500) / 1000.0f; } static void ds4_session_dspark_scheduler_reset(ds4_session *s) { if (!s) return; s->dspark_sched_cycles = 0; s->dspark_sched_accepted = 0; s->dspark_sched_no_draft = 0; s->dspark_sched_extra_ms = 0.0; s->dspark_sched_saved_ms = 0.0; } static bool ds4_session_dspark_scheduler_should_skip(ds4_session *s) { if (!s || !ds4_dspark_scheduler_enabled()) return false; s->dspark_sched_skipped_cycle = false; if (s->dspark_sched_skip == 0) return false; s->dspark_sched_skip--; s->dspark_sched_skipped_cycle = true; s->dspark_stats.scheduler_skips++; if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler skip remaining=%u\n", s->dspark_sched_skip); } return true; } static void ds4_session_dspark_scheduler_note( ds4_session *s, uint32_t accepted_drafts, bool no_draft, double extra_ms) { if (!s || !ds4_dspark_scheduler_enabled()) return; if (s->dspark_sched_skipped_cycle) { s->dspark_sched_skipped_cycle = false; return; } s->dspark_sched_cycles++; s->dspark_sched_accepted += accepted_drafts; if (accepted_drafts != 0) { if (s->dspark_sched_lifetime_accepted <= UINT32_MAX - accepted_drafts) { s->dspark_sched_lifetime_accepted += accepted_drafts; } else { s->dspark_sched_lifetime_accepted = UINT32_MAX; } if (accepted_drafts > 2u) { s->dspark_sched_long_accept_seen = true; } } if (no_draft) s->dspark_sched_no_draft++; if (extra_ms > 0.0 && isfinite(extra_ms)) { s->dspark_sched_extra_ms += extra_ms; } if (accepted_drafts != 0 && s->dspark_last_target_eval_ms > 0.0 && isfinite(s->dspark_last_target_eval_ms)) { const double saved_ms = s->dspark_last_target_eval_ms * (double)accepted_drafts; s->dspark_sched_saved_ms += saved_ms; if (ds4_dspark_stats_enabled()) { s->dspark_stats.saved_ms += saved_ms; } } const uint32_t no_draft_skip = ds4_dspark_scheduler_no_draft_skip_cycles(); if (no_draft && no_draft_skip != 0) { uint32_t skip = no_draft_skip; if (s->dspark_sched_lifetime_accepted != 0 && !s->dspark_sched_long_accept_seen) { const uint32_t short_accept_skip = ds4_dspark_scheduler_short_accept_no_draft_skip_cycles(); if (skip < short_accept_skip) skip = short_accept_skip; } else if (s->dspark_sched_lifetime_accepted == 0 && s->dspark_last_confidence0_valid && s->dspark_last_confidence0 <= ds4_dspark_scheduler_cold_low_confidence_threshold()) { const uint32_t cold_low_conf_skip = ds4_dspark_scheduler_cold_low_confidence_skip_cycles(); if (skip < cold_low_conf_skip) skip = cold_low_conf_skip; } if (s->dspark_sched_skip < skip) { s->dspark_sched_skip = skip; } if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler no-draft pause skip=%u " "accepted_total=%u long_accept=%d confidence0=%s%.3f\n", s->dspark_sched_skip, s->dspark_sched_lifetime_accepted, s->dspark_sched_long_accept_seen ? 1 : 0, s->dspark_last_confidence0_valid ? "" : "n/a:", s->dspark_last_confidence0); } } const uint32_t window = ds4_dspark_scheduler_window(); const uint32_t break_even_window = ds4_dspark_scheduler_break_even_window(); const uint32_t max_extra_saved_ratio_milli = ds4_dspark_scheduler_max_extra_saved_ratio_milli(); const bool measured_unprofitable = max_extra_saved_ratio_milli != 0 && s->dspark_sched_accepted != 0 && s->dspark_sched_saved_ms > 0.0 && s->dspark_sched_extra_ms * 1000.0 > s->dspark_sched_saved_ms * (double)max_extra_saved_ratio_milli; if (break_even_window != 0 && s->dspark_sched_cycles >= break_even_window && measured_unprofitable) { s->dspark_sched_skip = ds4_dspark_scheduler_slow_skip_cycles(); if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler break-even pause cycles=%u " "accepted=%u saved=%.3fms extra=%.3fms skip=%u\n", s->dspark_sched_cycles, s->dspark_sched_accepted, s->dspark_sched_saved_ms, s->dspark_sched_extra_ms, s->dspark_sched_skip); } ds4_session_dspark_scheduler_reset(s); return; } if (s->dspark_sched_cycles < window) return; const uint64_t avg_milli = ((uint64_t)s->dspark_sched_accepted * 1000ull) / (uint64_t)s->dspark_sched_cycles; const uint32_t min_avg_milli = ds4_dspark_scheduler_min_avg_milli(); const bool low_accept = avg_milli < min_avg_milli; const bool many_no_draft = s->dspark_sched_no_draft * 2u >= s->dspark_sched_cycles; const uint32_t max_ms_per_accept_milli = ds4_dspark_scheduler_max_ms_per_accept_milli(); const double extra_per_accept_ms = s->dspark_sched_accepted != 0 ? s->dspark_sched_extra_ms / (double)s->dspark_sched_accepted : 0.0; const bool slow_accept = max_ms_per_accept_milli != 0 && s->dspark_sched_accepted != 0 && extra_per_accept_ms * 1000.0 > (double)max_ms_per_accept_milli; if (low_accept || many_no_draft || slow_accept || measured_unprofitable) { s->dspark_sched_skip = ds4_dspark_scheduler_skip_cycles(); if (many_no_draft || slow_accept || measured_unprofitable) { const uint32_t slow_skip = ds4_dspark_scheduler_slow_skip_cycles(); if (s->dspark_sched_skip < slow_skip) { s->dspark_sched_skip = slow_skip; } } if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler pause cycles=%u accepted=%u " "avg=%.3f no_draft=%u extra_per_accept=%.3fms " "saved=%.3fms extra=%.3fms skip=%u\n", s->dspark_sched_cycles, s->dspark_sched_accepted, (double)avg_milli / 1000.0, s->dspark_sched_no_draft, extra_per_accept_ms, s->dspark_sched_saved_ms, s->dspark_sched_extra_ms, s->dspark_sched_skip); } } ds4_session_dspark_scheduler_reset(s); } #endif /* ========================================================================= * Session Snapshot Payloads. * ========================================================================= * * The server disk cache stores a high-level file header, then delegates the * graph-specific payload below to the engine. This payload is intentionally * not mmaped: restoring a checkpoint copies bytes back into the already * allocated Metal tensors, preserving the same live graph buffers used by * normal prefill/decode. The raw SWA cache is serialized as the last logical * window only; suffix prefill writes its own raw rows before attention. The * compressed caches are serialized up to their live row counts because sparse * attention may select rows from the whole prefix. * * The payload is model-specific rather than self-describing. The fixed header * records enough shape information to reject a file written for a different * DS4 runtime, then the body writes: checkpoint tokens, last logits, per-layer * compressed row counts, raw SWA rows in logical order, compressed attention * rows, and the compressor/indexer frontiers. That is the minimum state needed * for the next token to match a session that had just prefetched the prefix. */ #define DS4_SESSION_IO_CHUNK (8u * 1024u * 1024u) static void payload_set_err(char *err, size_t errlen, const char *msg) { if (errlen != 0) snprintf(err, errlen, "%s", msg); } static void payload_put_u32(uint8_t out[4], uint32_t v) { out[0] = (uint8_t)(v); out[1] = (uint8_t)(v >> 8); out[2] = (uint8_t)(v >> 16); out[3] = (uint8_t)(v >> 24); } static uint32_t payload_get_u32(const uint8_t in[4]) { return (uint32_t)in[0] | ((uint32_t)in[1] << 8) | ((uint32_t)in[2] << 16) | ((uint32_t)in[3] << 24); } static int payload_write_bytes(FILE *fp, const void *ptr, uint64_t bytes, char *err, size_t errlen) { const uint8_t *p = ptr; while (bytes != 0) { const size_t n = bytes > (uint64_t)SIZE_MAX ? SIZE_MAX : (size_t)bytes; if (fwrite(p, 1, n, fp) != n) { payload_set_err(err, errlen, "failed to write session payload"); return 1; } p += n; bytes -= n; } return 0; } static DS4_MAYBE_UNUSED int payload_read_bytes(FILE *fp, void *ptr, uint64_t bytes, uint64_t *remaining, char *err, size_t errlen) { if (remaining && *remaining < bytes) { payload_set_err(err, errlen, "truncated session payload"); return 1; } const uint64_t original = bytes; uint8_t *p = ptr; while (bytes != 0) { const size_t n = bytes > (uint64_t)SIZE_MAX ? SIZE_MAX : (size_t)bytes; if (fread(p, 1, n, fp) != n) { payload_set_err(err, errlen, "failed to read session payload"); return 1; } p += n; bytes -= n; } if (remaining) *remaining -= original; return 0; } static DS4_MAYBE_UNUSED int payload_skip_bytes(FILE *fp, uint64_t bytes, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (remaining && *remaining < bytes) { payload_set_err(err, errlen, "truncated session payload"); return 1; } if (!buf || cap == 0) { payload_set_err(err, errlen, "session payload skip buffer is missing"); return 1; } const uint64_t original = bytes; while (bytes != 0) { const size_t n = bytes > (uint64_t)cap ? cap : (size_t)bytes; if (fread(buf, 1, n, fp) != n) { payload_set_err(err, errlen, "failed to skip session payload"); return 1; } bytes -= n; } if (remaining) *remaining -= original; return 0; } static DS4_MAYBE_UNUSED int payload_write_u32(FILE *fp, uint32_t v, char *err, size_t errlen) { uint8_t b[4]; payload_put_u32(b, v); return payload_write_bytes(fp, b, sizeof(b), err, errlen); } static DS4_MAYBE_UNUSED int payload_read_u32(FILE *fp, uint32_t *v, uint64_t *remaining, char *err, size_t errlen) { uint8_t b[4]; if (remaining && *remaining < sizeof(b)) { payload_set_err(err, errlen, "truncated session payload"); return 1; } if (fread(b, 1, sizeof(b), fp) != sizeof(b)) { payload_set_err(err, errlen, "failed to read session payload"); return 1; } if (remaining) *remaining -= sizeof(b); *v = payload_get_u32(b); return 0; } static int payload_copy_file_bytes(FILE *src, FILE *dst, uint64_t bytes, char *err, size_t errlen) { uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; while (bytes != 0) { const size_t n = bytes > DS4_SESSION_IO_CHUNK ? DS4_SESSION_IO_CHUNK : (size_t)bytes; if (fread(buf, 1, n, src) != n) { payload_set_err(err, errlen, "failed to read staged session payload"); rc = 1; break; } if (fwrite(buf, 1, n, dst) != n) { payload_set_err(err, errlen, "failed to write staged session payload"); rc = 1; break; } bytes -= n; } free(buf); return rc; } static DS4_MAYBE_UNUSED uint64_t layer_attn_state_bytes(uint32_t ratio) { const uint32_t coff = ratio == 4 ? 2u : 1u; return (uint64_t)coff * DS4_N_HEAD_DIM * coff * ratio * sizeof(float); } static DS4_MAYBE_UNUSED uint64_t layer_index_state_bytes(uint32_t ratio) { const uint32_t coff = ratio == 4 ? 2u : 1u; return (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM * coff * ratio * sizeof(float); } #ifndef DS4_NO_GPU /* Only the last logical sliding-window rows are needed from the raw cache. * The physical Metal tensor is a ring sized for ubatches, but after restore * the next suffix chunk will write its own raw rows before any attention read. * Compressed rows are different: sparse attention can select any row from the * prefix, so those are persisted up to their live row counts. */ static uint32_t session_raw_live_rows(const ds4_gpu_graph *g, uint32_t checkpoint_len) { uint32_t rows = g->raw_window ? g->raw_window : DS4_N_SWA; if (rows > g->raw_cap) rows = g->raw_cap; if (rows > checkpoint_len) rows = checkpoint_len; return rows; } /* Return the exact engine-owned payload size, excluding the server's KVC file * header and observability text. This is deliberately based on live row counts * rather than capacities so the disk cache scales with saved tokens, not with * the maximum context size used to allocate the graph. */ static uint64_t session_payload_live_tensor_bytes(const ds4_gpu_graph *g, uint32_t checkpoint_len) { uint64_t bytes = 0; const uint32_t raw_live = session_raw_live_rows(g, checkpoint_len); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { bytes += (uint64_t)raw_live * DS4_N_HEAD_DIM * sizeof(float); const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; bytes += (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM * sizeof(float); bytes += layer_attn_state_bytes(ratio); bytes += layer_attn_state_bytes(ratio); if (ratio == 4) { bytes += (uint64_t)g->layer_n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float); bytes += layer_index_state_bytes(ratio); bytes += layer_index_state_bytes(ratio); } } return bytes; } static uint32_t session_glm_full_live_rows(const ds4_glm_gpu_graph *g, uint32_t checkpoint_len, uint32_t dense_cache_len) { if (!g) return 0; if (!g->full_kv_cache) return 0; uint32_t rows = g->ctx_cap; if (rows > checkpoint_len) rows = checkpoint_len; if (rows > dense_cache_len) rows = dense_cache_len; return rows; } static uint32_t session_glm_compact_live_rows(const ds4_glm_gpu_graph *g, uint32_t checkpoint_len) { if (!g || g->compact_cache_cap == 0) return 0; uint32_t rows = g->compact_cache_cap; if (rows > checkpoint_len) rows = checkpoint_len; return rows; } static uint64_t session_glm_payload_live_tensor_bytes(const ds4_glm_gpu_graph *g, uint32_t full_live, uint32_t compact_live) { if (!g) return 0; uint64_t bytes = 0; for (uint32_t il = 0; il < g->normal_layers; il++) { bytes += (uint64_t)full_live * g->q_dim * sizeof(float); bytes += (uint64_t)full_live * g->heads_dim * sizeof(float); if (compact_live == 0) continue; bytes += (uint64_t)compact_live * DS4_N_KV_LORA * sizeof(float); bytes += (uint64_t)compact_live * DS4_N_ROT * sizeof(float); if (glm_graph_layer_uses_full_indexer(il)) { bytes += (uint64_t)compact_live * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } } return bytes; } static bool payload_u64_add(uint64_t *acc, uint64_t value) { if (!acc || *acc > UINT64_MAX - value) return false; *acc += value; return true; } static bool payload_u64_mul(uint64_t a, uint64_t b, uint64_t *out) { if (!out) return false; if (a != 0 && b > UINT64_MAX / a) return false; *out = a * b; return true; } static bool payload_u64_add_tensor_bytes(uint64_t *acc, uint64_t rows, uint64_t cols) { uint64_t elems = 0; uint64_t bytes = 0; return payload_u64_mul(rows, cols, &elems) && payload_u64_mul(elems, sizeof(float), &bytes) && payload_u64_add(acc, bytes); } static bool glm_layer_payload_tensor_bytes(uint32_t layer, uint32_t full_live, uint32_t key_dim, uint32_t value_dim, uint32_t compact_live, uint32_t index_live, uint64_t *out) { if (!out || layer >= glm_graph_normal_layer_count() || key_dim != DS4_N_KEY_MLA || value_dim != DS4_N_VALUE_MLA) return false; const bool has_indexer = glm_graph_layer_uses_full_indexer(layer); const uint32_t expected_index_live = compact_live != 0 && has_indexer ? compact_live : 0; if (index_live != expected_index_live) return false; uint64_t bytes = 0; if (!payload_u64_add_tensor_bytes(&bytes, full_live, (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA) || !payload_u64_add_tensor_bytes(&bytes, full_live, (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA)) return false; if (compact_live != 0) { if (!payload_u64_add_tensor_bytes(&bytes, compact_live, DS4_N_KV_LORA) || !payload_u64_add_tensor_bytes(&bytes, compact_live, DS4_N_ROT)) return false; if (has_indexer && !payload_u64_add_tensor_bytes(&bytes, index_live, DS4_N_INDEXER_HEAD_DIM)) return false; } *out = bytes; return true; } /* Accelerator tensors are copied through a fixed-size CPU buffer. We do not mmap the * cache file and we do not allocate a second graph-sized blob just to serialize * it; both would be poor fits for this very large model. */ static int payload_write_tensor_span(FILE *fp, const ds4_gpu_tensor *tensor, uint64_t offset, uint64_t bytes, uint8_t *buf, size_t cap, char *err, size_t errlen) { if (!tensor || offset > ds4_gpu_tensor_bytes(tensor) || bytes > ds4_gpu_tensor_bytes(tensor) - offset) { payload_set_err(err, errlen, "session tensor is smaller than the payload"); return 1; } uint64_t done = 0; while (done < bytes) { const size_t n = bytes - done > (uint64_t)cap ? cap : (size_t)(bytes - done); if (ds4_gpu_tensor_read(tensor, offset + done, buf, n) == 0) { payload_set_err(err, errlen, "failed to read accelerator session tensor"); return 1; } if (payload_write_bytes(fp, buf, n, err, errlen) != 0) return 1; done += n; } return 0; } static int payload_read_tensor_span(FILE *fp, ds4_gpu_tensor *tensor, uint64_t offset, uint64_t bytes, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (!tensor || offset > ds4_gpu_tensor_bytes(tensor) || bytes > ds4_gpu_tensor_bytes(tensor) - offset) { payload_set_err(err, errlen, "session tensor is smaller than the payload"); return 1; } uint64_t done = 0; while (done < bytes) { const size_t n = bytes - done > (uint64_t)cap ? cap : (size_t)(bytes - done); if (payload_read_bytes(fp, buf, n, remaining, err, errlen) != 0) return 1; if (ds4_gpu_tensor_write(tensor, offset + done, buf, n) == 0) { payload_set_err(err, errlen, "failed to restore accelerator session tensor"); return 1; } done += n; } return 0; } static DS4_MAYBE_UNUSED int payload_write_tensor_span_f16_as_f32(FILE *fp, const ds4_gpu_tensor *tensor, uint64_t offset_f16, uint64_t count, uint8_t *buf, size_t cap, char *err, size_t errlen) { if (!tensor || count > (UINT64_MAX / sizeof(uint16_t)) || count > (UINT64_MAX / sizeof(float)) || offset_f16 > ds4_gpu_tensor_bytes(tensor) || count * sizeof(uint16_t) > ds4_gpu_tensor_bytes(tensor) - offset_f16) { payload_set_err(err, errlen, "session tensor is smaller than the F16 payload"); return 1; } size_t cap_elems = cap / (sizeof(uint16_t) + sizeof(float)); cap_elems &= ~(size_t)1u; if (cap_elems == 0) { payload_set_err(err, errlen, "session tensor conversion buffer is too small"); return 1; } uint16_t *h = (uint16_t *)buf; float *f = (float *)(void *)(buf + cap_elems * sizeof(uint16_t)); uint64_t done = 0; while (done < count) { const size_t n = count - done > (uint64_t)cap_elems ? cap_elems : (size_t)(count - done); if (ds4_gpu_tensor_read(tensor, offset_f16 + done * sizeof(uint16_t), h, n * sizeof(uint16_t)) == 0) { payload_set_err(err, errlen, "failed to read Metal F16 session tensor"); return 1; } for (size_t i = 0; i < n; i++) f[i] = f16_to_f32(h[i]); if (payload_write_bytes(fp, f, (uint64_t)n * sizeof(float), err, errlen) != 0) return 1; done += n; } return 0; } static DS4_MAYBE_UNUSED int payload_read_tensor_span_f32_as_f16(FILE *fp, ds4_gpu_tensor *tensor, uint64_t offset_f16, uint64_t count, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (!tensor || count > (UINT64_MAX / sizeof(uint16_t)) || count > (UINT64_MAX / sizeof(float)) || offset_f16 > ds4_gpu_tensor_bytes(tensor) || count * sizeof(uint16_t) > ds4_gpu_tensor_bytes(tensor) - offset_f16) { payload_set_err(err, errlen, "session tensor is smaller than the F16 payload"); return 1; } size_t cap_elems = cap / (sizeof(uint16_t) + sizeof(float)); cap_elems &= ~(size_t)1u; if (cap_elems == 0) { payload_set_err(err, errlen, "session tensor conversion buffer is too small"); return 1; } uint16_t *h = (uint16_t *)buf; float *f = (float *)(void *)(buf + cap_elems * sizeof(uint16_t)); uint64_t done = 0; while (done < count) { const size_t n = count - done > (uint64_t)cap_elems ? cap_elems : (size_t)(count - done); if (payload_read_bytes(fp, f, (uint64_t)n * sizeof(float), remaining, err, errlen) != 0) return 1; for (size_t i = 0; i < n; i++) h[i] = f32_to_f16(f[i]); if (ds4_gpu_tensor_write(tensor, offset_f16 + done * sizeof(uint16_t), h, n * sizeof(uint16_t)) == 0) { payload_set_err(err, errlen, "failed to restore Metal F16 session tensor"); return 1; } done += n; } return 0; } static int payload_write_glm_compact_span(FILE *fp, const ds4_gpu_tensor *tensor, uint64_t count, uint8_t *buf, size_t cap, char *err, size_t errlen) { if (glm_graph_compact_cache_is_f16()) { return payload_write_tensor_span_f16_as_f32(fp, tensor, 0, count, buf, cap, err, errlen); } return payload_write_tensor_span(fp, tensor, 0, count * sizeof(float), buf, cap, err, errlen); } static int payload_read_glm_compact_span(FILE *fp, ds4_gpu_tensor *tensor, uint64_t count, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (glm_graph_compact_cache_is_f16()) { return payload_read_tensor_span_f32_as_f16(fp, tensor, 0, count, buf, cap, remaining, err, errlen); } return payload_read_tensor_span(fp, tensor, 0, count * sizeof(float), buf, cap, remaining, err, errlen); } static int payload_write_glm_full_kv_span(FILE *fp, const ds4_gpu_tensor *tensor, uint64_t count, uint8_t *buf, size_t cap, char *err, size_t errlen) { if (count == 0) return 0; return payload_write_tensor_span_f16_as_f32(fp, tensor, 0, count, buf, cap, err, errlen); } static int payload_read_glm_full_kv_span(FILE *fp, ds4_gpu_tensor *tensor, uint64_t count, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (count == 0) return 0; return payload_read_tensor_span_f32_as_f16(fp, tensor, 0, count, buf, cap, remaining, err, errlen); } static int payload_read_or_skip_glm_full_kv_span(FILE *fp, ds4_gpu_tensor *tensor, uint64_t count, uint8_t *buf, size_t cap, uint64_t *remaining, char *err, size_t errlen) { if (count == 0) return 0; if (tensor) { return payload_read_glm_full_kv_span(fp, tensor, count, buf, cap, remaining, err, errlen); } if (count > UINT64_MAX / sizeof(float)) { payload_set_err(err, errlen, "GLM full KV payload is too large"); return 1; } return payload_skip_bytes(fp, count * sizeof(float), buf, cap, remaining, err, errlen); } #endif static bool ds4_session_is_cpu(const ds4_session *s) { return s && s->engine && s->engine->backend == DS4_BACKEND_CPU; } static void ds4_session_dspark_capture_invalidate(ds4_session *s) { #ifndef DS4_NO_GPU if (!s) return; s->dspark_draft_valid = false; s->dspark_draft_len = 0; if (ds4_session_is_cpu(s)) return; metal_graph_dspark_capture_invalidate(&s->graph); #else (void)s; #endif } static void ds4_session_dspark_capture_note_checkpoint(ds4_session *s) { #ifndef DS4_NO_GPU if (!s || ds4_session_is_cpu(s)) return; ds4_gpu_graph *g = &s->graph; if (!g->dspark_capture_enabled) return; if (!s->checkpoint_valid || s->checkpoint.len <= 0 || !g->dspark_capture_valid) { metal_graph_dspark_capture_invalidate(g); return; } g->dspark_capture_checkpoint_len = (uint32_t)s->checkpoint.len; #else (void)s; #endif } static bool ds4_session_is_glm(const ds4_session *s) { return s && s->engine && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; } #ifndef DS4_NO_GPU static void ds4_session_glm_reset_dense_cache(ds4_session *s) { if (s) s->glm_dense_cache_len = 0; } static void ds4_session_glm_cap_dense_cache(ds4_session *s) { if (!s) return; if (s->glm_graph_ready && !s->glm_graph.full_kv_cache) { s->glm_dense_cache_len = 0; return; } uint32_t cap = s->checkpoint.len > 0 ? (uint32_t)s->checkpoint.len : 0; if (s->glm_graph_ready && cap > s->glm_graph.ctx_cap) cap = s->glm_graph.ctx_cap; if (s->glm_dense_cache_len > cap) s->glm_dense_cache_len = cap; } static void ds4_session_glm_note_dense_cache(ds4_session *s, uint32_t pos0, uint32_t n_tokens) { if (!s || n_tokens == 0) return; if (s->glm_graph_ready && !s->glm_graph.full_kv_cache) return; if (pos0 > s->glm_dense_cache_len) return; uint32_t end = pos0 + n_tokens; if (end < pos0) end = UINT32_MAX; if (s->glm_graph_ready && end > s->glm_graph.ctx_cap) end = s->glm_graph.ctx_cap; s->glm_dense_cache_len = end; ds4_session_glm_cap_dense_cache(s); } #endif static uint32_t ds4_model_normal_layer_count(void) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) { return (uint32_t)DS4_N_LAYER; } 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 session_cpu_raw_live_rows(const ds4_session *s) { if (!s || !s->checkpoint_valid) return 0; uint32_t rows = ds4_default_raw_cap((uint32_t)s->ctx_size); if (rows > (uint32_t)s->checkpoint.len) rows = (uint32_t)s->checkpoint.len; return rows; } static uint32_t session_cpu_comp_cap(const ds4_session *s) { if (!s) return 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_cache *layer = &s->cpu_cache.layer[il]; if (layer->compress_ratio == 4) return layer->comp_cap; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_cache *layer = &s->cpu_cache.layer[il]; if (layer->compress_ratio != 0) return layer->comp_cap; } return (uint32_t)s->ctx_size; } static uint64_t session_cpu_payload_live_tensor_bytes(const ds4_session *s) { uint64_t bytes = 0; const uint32_t raw_live = session_cpu_raw_live_rows(s); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_cache *layer = &s->cpu_cache.layer[il]; bytes += (uint64_t)raw_live * DS4_N_HEAD_DIM * sizeof(float); const uint32_t ratio = layer->compress_ratio; if (ratio == 0) continue; bytes += (uint64_t)layer->n_comp * DS4_N_HEAD_DIM * sizeof(float); bytes += layer_attn_state_bytes(ratio); bytes += layer_attn_state_bytes(ratio); if (ratio == 4) { bytes += (uint64_t)layer->n_index_comp * DS4_N_INDEXER_HEAD_DIM * sizeof(float); bytes += layer_index_state_bytes(ratio); bytes += layer_index_state_bytes(ratio); } } return bytes; } static void session_cpu_reset_cache(ds4_session *s) { kv_cache_free(&s->cpu_cache); kv_cache_init(&s->cpu_cache, (uint32_t)s->ctx_size, 0); } static bool ds4_layer_payload_range_valid(uint32_t layer_start, uint32_t layer_end) { const uint32_t n_layers = ds4_model_normal_layer_count(); return n_layers != 0 && layer_start <= layer_end && layer_end < n_layers; } uint64_t ds4_session_layer_payload_bytes(ds4_session *s, uint32_t layer_start, uint32_t layer_end) { 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)) { #ifdef DS4_NO_GPU (void)layer_start; (void)layer_end; return 0; #else if (!s->glm_graph_ready) return 0; const ds4_glm_gpu_graph *g = &s->glm_graph; if (layer_start != g->layer_start || layer_end != g->layer_end) return 0; const uint32_t checkpoint_len = (uint32_t)s->checkpoint.len; const uint32_t full_live = session_glm_full_live_rows(g, checkpoint_len, s->glm_dense_cache_len); const uint32_t compact_live = session_glm_compact_live_rows(g, checkpoint_len); uint64_t bytes = (uint64_t)DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS * sizeof(uint32_t); const uint32_t n_layers = layer_end - layer_start + 1u; bytes += (uint64_t)n_layers * sizeof(uint32_t); bytes += (uint64_t)n_layers * sizeof(uint32_t); for (uint32_t il = layer_start; il <= layer_end; il++) { const uint32_t index_live = compact_live != 0 && glm_graph_layer_uses_full_indexer(il) ? compact_live : 0; uint64_t layer_bytes = 0; if (!glm_layer_payload_tensor_bytes(il, full_live, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, compact_live, index_live, &layer_bytes) || !payload_u64_add(&bytes, layer_bytes)) return 0; } return bytes; #endif } #ifdef DS4_NO_GPU (void)layer_start; (void)layer_end; return 0; #else const ds4_gpu_graph *g = &s->graph; const uint32_t raw_live = session_raw_live_rows(g, (uint32_t)s->checkpoint.len); uint64_t bytes = (uint64_t)DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS * sizeof(uint32_t); const uint32_t n_layers = layer_end - layer_start + 1u; bytes += (uint64_t)n_layers * sizeof(uint32_t); bytes += (uint64_t)n_layers * sizeof(uint32_t); for (uint32_t il = layer_start; il <= layer_end; il++) { bytes += (uint64_t)raw_live * DS4_N_HEAD_DIM * sizeof(float); const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; bytes += (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM * sizeof(float); bytes += layer_attn_state_bytes(ratio); bytes += layer_attn_state_bytes(ratio); if (ratio == 4) { bytes += (uint64_t)g->layer_n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float); bytes += layer_index_state_bytes(ratio); bytes += layer_index_state_bytes(ratio); } } return bytes; #endif } 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 || !fp || !s->checkpoint_valid || !ds4_layer_payload_range_valid(layer_start, layer_end)) { payload_set_err(err, errlen, "invalid session layer payload save"); return 1; } if (ds4_session_is_cpu(s)) { payload_set_err(err, errlen, "distributed layer payloads require the graph backend"); return 1; } if (ds4_session_is_glm(s)) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (!s->glm_graph_ready) { payload_set_err(err, errlen, "GLM graph is not ready for layer snapshot"); return 1; } ds4_glm_gpu_graph *g = &s->glm_graph; if (layer_start != g->layer_start || layer_end != g->layer_end) { payload_set_err(err, errlen, "requested GLM layer snapshot does not match loaded slice"); return 1; } if (ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator before GLM layer snapshot"); return 1; } const uint32_t checkpoint_len = (uint32_t)s->checkpoint.len; const uint32_t full_live = session_glm_full_live_rows(g, checkpoint_len, s->glm_dense_cache_len); const uint32_t compact_live = session_glm_compact_live_rows(g, checkpoint_len); uint32_t header[DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS] = { DS4_SESSION_LAYER_PAYLOAD_MAGIC, DS4_SESSION_LAYER_PAYLOAD_VERSION, (uint32_t)s->ctx_size, s->prefill_cap, g->ctx_cap, g->ctx_cap, g->compact_cache_cap, checkpoint_len, g->normal_layers, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, layer_start, layer_end, full_live, }; for (uint32_t i = 0; i < DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS; i++) { if (payload_write_u32(fp, header[i], err, errlen) != 0) return 1; } for (uint32_t il = layer_start; il <= layer_end; il++) { if (payload_write_u32(fp, compact_live, err, errlen) != 0) return 1; } for (uint32_t il = layer_start; il <= layer_end; il++) { const uint32_t index_rows = compact_live != 0 && glm_graph_layer_uses_full_indexer(il) ? compact_live : 0; if (payload_write_u32(fp, index_rows, err, errlen) != 0) return 1; } uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = layer_start; rc == 0 && il <= layer_end; il++) { rc = payload_write_glm_full_kv_span(fp, g->layer_key_cache[il], (uint64_t)full_live * g->q_dim, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_glm_full_kv_span(fp, g->layer_value_cache[il], (uint64_t)full_live * g->heads_dim, buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc != 0 || compact_live == 0) continue; rc = payload_write_glm_compact_span(fp, g->layer_kv_lora_cache[il], (uint64_t)compact_live * DS4_N_KV_LORA, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_glm_compact_span(fp, g->layer_k_rope_cache[il], (uint64_t)compact_live * DS4_N_ROT, buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc == 0 && glm_graph_layer_uses_full_indexer(il)) { rc = payload_write_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)compact_live * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); } } free(buf); return rc; #endif } #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator before layer snapshot"); return 1; } ds4_gpu_graph *g = &s->graph; const uint32_t raw_live = session_raw_live_rows(g, (uint32_t)s->checkpoint.len); uint32_t header[DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS] = { DS4_SESSION_LAYER_PAYLOAD_MAGIC, DS4_SESSION_LAYER_PAYLOAD_VERSION, (uint32_t)s->ctx_size, s->prefill_cap, g->raw_cap, g->raw_window, g->comp_cap, (uint32_t)s->checkpoint.len, DS4_N_LAYER, DS4_N_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM, layer_start, layer_end, raw_live, }; for (uint32_t i = 0; i < DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS; i++) { if (payload_write_u32(fp, header[i], err, errlen) != 0) return 1; } for (uint32_t il = layer_start; il <= layer_end; il++) { if (payload_write_u32(fp, g->layer_n_comp[il], err, errlen) != 0) return 1; } for (uint32_t il = layer_start; il <= layer_end; il++) { if (payload_write_u32(fp, g->layer_n_index_comp[il], err, errlen) != 0) return 1; } uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = layer_start; rc == 0 && il <= layer_end; il++) { const uint32_t raw_first = (uint32_t)s->checkpoint.len - raw_live; for (uint32_t r = 0; rc == 0 && r < raw_live; r++) { const uint32_t pos = raw_first + r; const uint32_t phys = pos % g->raw_cap; rc = payload_write_tensor_span(fp, g->layer_raw_cache[il], (uint64_t)phys * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); } const uint32_t ratio = ds4_layer_compress_ratio(il); if (rc != 0 || ratio == 0) continue; if (DS4_GPU_ATTN_COMP_CACHE_F16) { rc = payload_write_tensor_span_f16_as_f32(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); } else { rc = payload_write_tensor_span(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_attn_state_kv[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_attn_state_score[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0 && ratio == 4) { rc = payload_write_tensor_span(fp, g->layer_index_comp_cache[il], 0, (uint64_t)g->layer_n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_index_state_kv[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_index_state_score[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); } } free(buf); return rc; #endif } 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 || !fp || !tokens || !ds4_layer_payload_range_valid(layer_start, layer_end)) { payload_set_err(err, errlen, "invalid session layer payload load"); return 1; } if (ds4_session_is_cpu(s)) { payload_set_err(err, errlen, "distributed layer payloads require the graph backend"); return 1; } if (ds4_session_is_glm(s)) { #ifdef DS4_NO_GPU (void)payload_bytes; (void)n_tokens; payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (!s->glm_graph_ready) { payload_set_err(err, errlen, "GLM graph is not ready for KV shard restore"); return 1; } ds4_glm_gpu_graph *g = &s->glm_graph; if (layer_start != g->layer_start || layer_end != g->layer_end) { payload_set_err(err, errlen, "requested GLM KV shard does not match loaded slice"); return 1; } uint64_t remaining = payload_bytes; uint32_t h[DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS]; for (uint32_t i = 0; i < DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS; i++) { if (payload_read_u32(fp, &h[i], &remaining, err, errlen) != 0) return 1; } if (h[0] != DS4_SESSION_LAYER_PAYLOAD_MAGIC || h[1] != DS4_SESSION_LAYER_PAYLOAD_VERSION) { payload_set_err(err, errlen, "unsupported session layer payload version"); return 1; } const uint32_t saved_ctx = h[2]; const uint32_t saved_prefill_cap = h[3]; const uint32_t saved_full_cap = h[4]; const uint32_t saved_full_window = h[5]; const uint32_t saved_compact_cap = h[6]; const uint32_t saved_tokens = h[7]; const uint32_t saved_layer_start = h[11]; const uint32_t saved_layer_end = h[12]; const uint32_t saved_full_live = h[13]; (void)saved_prefill_cap; if (saved_layer_start != layer_start || saved_layer_end != layer_end) { payload_set_err(err, errlen, "KV shard layer range does not match requested worker"); return 1; } if (saved_ctx > (uint32_t)s->ctx_size || saved_tokens != n_tokens || saved_tokens >= (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV shard does not fit current context"); return 1; } for (uint32_t i = 0; i < n_tokens; i++) { if (tokens[i] < 0 || (uint32_t)tokens[i] >= DS4_N_VOCAB) { payload_set_err(err, errlen, "KV shard token is outside vocabulary"); return 1; } } if (h[8] != g->normal_layers || h[9] != DS4_N_KEY_MLA || h[10] != DS4_N_VALUE_MLA) { payload_set_err(err, errlen, "KV shard was written for a different GLM layout"); return 1; } if (saved_full_cap == 0 || saved_full_cap > g->ctx_cap || saved_full_window != saved_full_cap) { payload_set_err(err, errlen, "KV shard GLM full-cache layout does not match current runtime"); return 1; } const uint32_t max_full_live = saved_tokens < saved_full_cap ? saved_tokens : saved_full_cap; if (saved_full_live > max_full_live || saved_full_live > g->ctx_cap) { payload_set_err(err, errlen, "KV shard GLM full-cache row count is invalid"); return 1; } if (saved_compact_cap > g->ctx_size) { payload_set_err(err, errlen, "KV shard GLM compact cache is larger than current context"); return 1; } if (g->compact_cache_cap != 0 && saved_tokens != 0 && saved_compact_cap == 0) { payload_set_err(err, errlen, "KV shard lacks the GLM compact cache required by this context"); return 1; } const uint32_t expected_compact_live = saved_compact_cap != 0 ? (saved_tokens < saved_compact_cap ? saved_tokens : saved_compact_cap) : 0; if (expected_compact_live != 0 && (!s->engine || !glm_graph_ensure_compact_cache(g, expected_compact_live))) { payload_set_err(err, errlen, "KV shard GLM compact cache could not be allocated"); return 1; } if (expected_compact_live > g->compact_cache_cap) { payload_set_err(err, errlen, "KV shard GLM compact row count is invalid"); return 1; } const uint32_t n_layers = layer_end - layer_start + 1u; uint32_t *n_comp = xcalloc(n_layers, sizeof(n_comp[0])); uint32_t *n_index_comp = xcalloc(n_layers, sizeof(n_index_comp[0])); for (uint32_t i = 0; i < n_layers; i++) { if (payload_read_u32(fp, &n_comp[i], &remaining, err, errlen) != 0) { free(n_comp); free(n_index_comp); return 1; } if (n_comp[i] != expected_compact_live) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "KV shard GLM compact row count does not match token count"); return 1; } } for (uint32_t i = 0; i < n_layers; i++) { const uint32_t il = layer_start + i; if (payload_read_u32(fp, &n_index_comp[i], &remaining, err, errlen) != 0) { free(n_comp); free(n_index_comp); return 1; } const uint32_t expected_index_rows = expected_compact_live != 0 && glm_graph_layer_uses_full_indexer(il) ? expected_compact_live : 0; if (n_index_comp[i] != expected_index_rows) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "KV shard GLM indexer row count does not match token count"); return 1; } } if (ds4_gpu_synchronize() == 0) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "failed to synchronize accelerator before GLM KV shard restore"); return 1; } s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_reset_dense_cache(s); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t i = 0; rc == 0 && i < n_layers; i++) { const uint32_t il = layer_start + i; rc = payload_read_or_skip_glm_full_kv_span(fp, g->layer_key_cache[il], (uint64_t)saved_full_live * g->q_dim, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_or_skip_glm_full_kv_span(fp, g->layer_value_cache[il], (uint64_t)saved_full_live * g->heads_dim, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc != 0 || expected_compact_live == 0) continue; rc = payload_read_glm_compact_span(fp, g->layer_kv_lora_cache[il], (uint64_t)n_comp[i] * DS4_N_KV_LORA, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_glm_compact_span(fp, g->layer_k_rope_cache[il], (uint64_t)n_comp[i] * DS4_N_ROT, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc == 0 && glm_graph_layer_uses_full_indexer(il)) { rc = payload_read_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)n_index_comp[i] * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } } free(buf); if (rc == 0 && remaining != 0) { payload_set_err(err, errlen, "KV shard has trailing payload bytes"); rc = 1; } if (rc == 0 && ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator after GLM KV shard restore"); rc = 1; } if (rc == 0) { token_vec_free(&s->checkpoint); memset(&s->checkpoint, 0, sizeof(s->checkpoint)); for (uint32_t i = 0; i < n_tokens; i++) token_vec_push(&s->checkpoint, tokens[i]); s->checkpoint_valid = true; s->mtp_draft_valid = false; s->glm_dense_cache_len = g->full_kv_cache ? saved_full_live : 0; } free(n_comp); free(n_index_comp); return rc; #endif } #ifdef DS4_NO_GPU (void)payload_bytes; (void)n_tokens; payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else uint64_t remaining = payload_bytes; uint32_t h[DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS]; for (uint32_t i = 0; i < DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS; i++) { if (payload_read_u32(fp, &h[i], &remaining, err, errlen) != 0) return 1; } if (h[0] != DS4_SESSION_LAYER_PAYLOAD_MAGIC || h[1] != DS4_SESSION_LAYER_PAYLOAD_VERSION) { payload_set_err(err, errlen, "unsupported session layer payload version"); return 1; } ds4_gpu_graph *g = &s->graph; const uint32_t saved_ctx = h[2]; const uint32_t saved_prefill_cap = h[3]; const uint32_t saved_raw_cap = h[4]; const uint32_t saved_raw_window = h[5]; const uint32_t saved_comp_cap = h[6]; const uint32_t saved_tokens = h[7]; const uint32_t saved_layer_start = h[11]; const uint32_t saved_layer_end = h[12]; const uint32_t saved_raw_live = h[13]; (void)saved_prefill_cap; if (saved_layer_start != layer_start || saved_layer_end != layer_end) { payload_set_err(err, errlen, "KV shard layer range does not match requested worker"); return 1; } if (saved_ctx > (uint32_t)s->ctx_size || saved_tokens != n_tokens || saved_tokens >= (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV shard does not fit current context"); return 1; } if (h[8] != DS4_N_LAYER || h[9] != DS4_N_HEAD_DIM || h[10] != DS4_N_INDEXER_HEAD_DIM) { payload_set_err(err, errlen, "KV shard was written for a different DS4 layout"); return 1; } if (saved_raw_window != g->raw_window) { payload_set_err(err, errlen, "KV shard graph chunk layout does not match current runtime"); return 1; } const uint32_t expected_raw_live = saved_tokens < saved_raw_window ? saved_tokens : saved_raw_window; if (saved_raw_cap == 0 || saved_raw_live != expected_raw_live || saved_raw_live > saved_raw_cap || saved_raw_live > g->raw_cap) { payload_set_err(err, errlen, "KV shard raw ring layout does not match current context"); return 1; } if (saved_comp_cap > g->comp_cap) { payload_set_err(err, errlen, "KV shard compressed cache is larger than current context"); return 1; } const uint32_t n_layers = layer_end - layer_start + 1u; uint32_t *n_comp = xcalloc(n_layers, sizeof(n_comp[0])); uint32_t *n_index_comp = xcalloc(n_layers, sizeof(n_index_comp[0])); for (uint32_t i = 0; i < n_layers; i++) { const uint32_t il = layer_start + i; if (payload_read_u32(fp, &n_comp[i], &remaining, err, errlen) != 0) { free(n_comp); free(n_index_comp); return 1; } if (n_comp[i] > saved_comp_cap || n_comp[i] > g->layer_comp_cap[il]) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "KV shard has invalid compressed row count"); return 1; } } for (uint32_t i = 0; i < n_layers; i++) { const uint32_t il = layer_start + i; if (payload_read_u32(fp, &n_index_comp[i], &remaining, err, errlen) != 0) { free(n_comp); free(n_index_comp); return 1; } if (n_index_comp[i] > saved_comp_cap || n_index_comp[i] > g->layer_comp_cap[il]) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "KV shard has invalid indexer row count"); return 1; } } if (ds4_gpu_synchronize() == 0) { free(n_comp); free(n_index_comp); payload_set_err(err, errlen, "failed to synchronize accelerator before KV shard restore"); return 1; } s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t i = 0; rc == 0 && i < n_layers; i++) { const uint32_t il = layer_start + i; const uint32_t raw_first = saved_tokens - saved_raw_live; for (uint32_t r = 0; rc == 0 && r < saved_raw_live; r++) { const uint32_t pos = raw_first + r; const uint32_t phys = pos % g->raw_cap; rc = payload_read_tensor_span(fp, g->layer_raw_cache[il], (uint64_t)phys * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } const uint32_t ratio = ds4_layer_compress_ratio(il); if (rc != 0 || ratio == 0) continue; if (DS4_GPU_ATTN_COMP_CACHE_F16) { rc = payload_read_tensor_span_f32_as_f16(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)n_comp[i] * DS4_N_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } else { rc = payload_read_tensor_span(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)n_comp[i] * DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_attn_state_kv[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_attn_state_score[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0 && ratio == 4) { rc = payload_read_tensor_span(fp, g->layer_index_comp_cache[il], 0, (uint64_t)n_index_comp[i] * DS4_N_INDEXER_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_index_state_kv[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_index_state_score[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } } free(buf); if (rc == 0 && remaining != 0) { payload_set_err(err, errlen, "KV shard has trailing payload bytes"); rc = 1; } if (rc == 0 && ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator after KV shard restore"); rc = 1; } if (rc == 0) { token_vec_free(&s->checkpoint); memset(&s->checkpoint, 0, sizeof(s->checkpoint)); for (uint32_t i = 0; i < n_tokens; i++) token_vec_push(&s->checkpoint, tokens[i]); for (uint32_t i = 0; i < n_layers; i++) { const uint32_t il = layer_start + i; g->layer_n_comp[il] = n_comp[i]; g->layer_n_index_comp[il] = n_index_comp[i]; } s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); } free(n_comp); free(n_index_comp); return rc; #endif } int ds4_engine_routed_quant_bits(ds4_engine *e) { if (!e) return 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_tensor *gate = e->weights.layer[il].ffn_gate_exps; if (!gate) continue; return gate->type == DS4_TENSOR_Q4_K ? 4 : 2; } return 0; } bool ds4_engine_has_output_head(ds4_engine *e) { return e && weights_have_output_head(&e->weights); } #ifndef DS4_NO_GPU static bool ds4_engine_glm_mtp_spec_enabled(const ds4_engine *e) { return e && e->backend != DS4_BACKEND_CPU && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && DS4_N_NEXTN_PREDICT != 0 && e->glm_mtp; } #endif bool ds4_engine_has_mtp(ds4_engine *e) { return e && e->backend != DS4_BACKEND_CPU && e->distributed.role == DS4_DISTRIBUTED_NONE && e->mtp_ready; } int ds4_engine_mtp_draft_tokens(ds4_engine *e) { if (e && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { return e->glm_mtp && DS4_N_NEXTN_PREDICT != 0 ? 2 : 0; } if (ds4_engine_has_mtp(e)) return e->mtp_draft_tokens; #ifndef DS4_NO_GPU if (e && e->backend != DS4_BACKEND_CPU && e->distributed.role == DS4_DISTRIBUTED_NONE && e->support_kind == DS4_SUPPORT_DSPARK && e->dspark && e->dspark_weights.block_size > 1) { return (int)e->dspark_weights.block_size; } #endif return 0; } const ds4_tokens *ds4_session_tokens(ds4_session *s) { return s ? &s->checkpoint : NULL; } #ifndef DS4_NO_GPU static void spec_frontier_free(ds4_spec_frontier *f) { if (!f) return; memset(f, 0, sizeof(*f)); } static bool spec_frontier_snapshot(ds4_spec_frontier *f, ds4_session *s) { memset(f, 0, sizeof(*f)); ds4_gpu_graph *g = &s->graph; if (!metal_graph_dspark_cache_current_window_valid(g)) return false; f->mtp_n_raw = g->mtp_n_raw; f->dspark_cache_start = g->dspark_cache_start; f->dspark_cache_token_start = g->dspark_cache_token_start; f->dspark_cache_len = g->dspark_cache_len; bool ok = ds4_gpu_begin_commands() != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { f->n_comp[il] = g->layer_n_comp[il]; f->n_index_comp[il] = g->layer_n_index_comp[il]; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint64_t ab = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); ok = ds4_gpu_tensor_copy(g->spec_attn_state_kv[il], 0, g->layer_attn_state_kv[il], 0, ab) != 0 && ds4_gpu_tensor_copy(g->spec_attn_state_score[il], 0, g->layer_attn_state_score[il], 0, ab) != 0; if (ratio == 4) { const uint64_t ib = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); ok = ok && ds4_gpu_tensor_copy(g->spec_index_state_kv[il], 0, g->layer_index_state_kv[il], 0, ib) != 0 && ds4_gpu_tensor_copy(g->spec_index_state_score[il], 0, g->layer_index_state_score[il], 0, ib) != 0; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); if (ok) return true; spec_frontier_free(f); return false; } static bool spec_frontier_restore(ds4_spec_frontier *f, ds4_session *s) { ds4_gpu_graph *g = &s->graph; if (!metal_graph_dspark_cache_window_valid(g, f->dspark_cache_token_start, f->dspark_cache_start, f->dspark_cache_len)) { return false; } bool ok = ds4_gpu_begin_commands() != 0; g->mtp_n_raw = f->mtp_n_raw; if (ok) { ok = metal_graph_dspark_cache_set_window(g, f->dspark_cache_token_start, f->dspark_cache_len); } for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { g->layer_n_comp[il] = f->n_comp[il]; g->layer_n_index_comp[il] = f->n_index_comp[il]; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint64_t ab = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); ok = ds4_gpu_tensor_copy(g->layer_attn_state_kv[il], 0, g->spec_attn_state_kv[il], 0, ab) != 0 && ds4_gpu_tensor_copy(g->layer_attn_state_score[il], 0, g->spec_attn_state_score[il], 0, ab) != 0; if (ok && ratio == 4) { const uint64_t ib = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); ok = ds4_gpu_tensor_copy(g->layer_index_state_kv[il], 0, g->spec_index_state_kv[il], 0, ib) != 0 && ds4_gpu_tensor_copy(g->layer_index_state_score[il], 0, g->spec_index_state_score[il], 0, ib) != 0; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); return ok; } /* Commit the prefix-1 state captured by the N=2 speculative verifier. * * The verifier has already advanced every layer through both draft tokens. On * a one-token accept the append-only compressed caches can keep the second * speculative row as invisible garbage, but the compressor frontiers and row * counters must be rewound to the exact state after draft[0]. This is the * cheap partial-accept path: copy a few small per-layer frontiers instead of * restoring the whole prefix and replaying a one-token target decode. */ static bool spec_frontier_commit_prefix1(ds4_session *s) { ds4_gpu_graph *g = &s->graph; bool ok = ds4_gpu_begin_commands() != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; g->layer_n_comp[il] = g->spec_prefix1_n_comp[il]; const uint64_t ab = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); ok = ds4_gpu_tensor_copy(g->layer_attn_state_kv[il], 0, g->spec_prefix1_attn_state_kv[il], 0, ab) != 0 && ds4_gpu_tensor_copy(g->layer_attn_state_score[il], 0, g->spec_prefix1_attn_state_score[il], 0, ab) != 0; if (ok && ratio == 4) { g->layer_n_index_comp[il] = g->spec_prefix1_n_index_comp[il]; const uint64_t ib = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); ok = ds4_gpu_tensor_copy(g->layer_index_state_kv[il], 0, g->spec_prefix1_index_state_kv[il], 0, ib) != 0 && ds4_gpu_tensor_copy(g->layer_index_state_score[il], 0, g->spec_prefix1_index_state_score[il], 0, ib) != 0; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); return ok; } static void session_greedy_splitkv_reset(ds4_session *s) { if (!s) return; s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = false; s->greedy_splitkv_anchor_len = 0; spec_frontier_free(&s->greedy_splitkv_anchor); } #endif uint64_t ds4_session_payload_bytes(ds4_session *s) { if (!s || !s->checkpoint_valid) return 0; if (s->distributed) return 0; if (ds4_session_is_cpu(s)) { uint64_t bytes = (uint64_t)DS4_SESSION_PAYLOAD_U32_FIELDS * sizeof(uint32_t); bytes += (uint64_t)s->checkpoint.len * sizeof(uint32_t); bytes += (uint64_t)DS4_N_VOCAB * sizeof(float); bytes += (uint64_t)DS4_N_LAYER * sizeof(uint32_t); bytes += (uint64_t)DS4_N_LAYER * sizeof(uint32_t); bytes += session_cpu_payload_live_tensor_bytes(s); return bytes; } if (ds4_session_is_glm(s)) { #ifdef DS4_NO_GPU return 0; #else if (!s->glm_graph_ready) return 0; const ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t checkpoint_len = (uint32_t)s->checkpoint.len; const uint32_t full_live = session_glm_full_live_rows(g, checkpoint_len, s->glm_dense_cache_len); const uint32_t compact_live = session_glm_compact_live_rows(g, checkpoint_len); uint64_t bytes = (uint64_t)DS4_SESSION_PAYLOAD_U32_FIELDS * sizeof(uint32_t); bytes += (uint64_t)s->checkpoint.len * sizeof(uint32_t); bytes += (uint64_t)DS4_N_VOCAB * sizeof(float); bytes += (uint64_t)g->normal_layers * sizeof(uint32_t); bytes += (uint64_t)g->normal_layers * sizeof(uint32_t); bytes += session_glm_payload_live_tensor_bytes(g, full_live, compact_live); return bytes; #endif } #ifdef DS4_NO_GPU return 0; #else const ds4_gpu_graph *g = &s->graph; uint64_t bytes = (uint64_t)DS4_SESSION_PAYLOAD_U32_FIELDS * sizeof(uint32_t); bytes += (uint64_t)s->checkpoint.len * sizeof(uint32_t); bytes += (uint64_t)DS4_N_VOCAB * sizeof(float); bytes += (uint64_t)DS4_N_LAYER * sizeof(uint32_t); bytes += (uint64_t)DS4_N_LAYER * sizeof(uint32_t); bytes += session_payload_live_tensor_bytes(g, (uint32_t)s->checkpoint.len); return bytes; #endif } int ds4_session_write_staged_payload(const ds4_session_payload_file *payload, FILE *fp, char *err, size_t errlen) { if (!payload || !payload->path || !fp) { payload_set_err(err, errlen, "invalid staged session payload"); return 1; } FILE *src = fopen(payload->path, "rb"); if (!src) { payload_set_err(err, errlen, "failed to open staged session payload"); return 1; } int rc = payload_copy_file_bytes(src, fp, payload->bytes, err, errlen); if (fclose(src) != 0 && rc == 0) { payload_set_err(err, errlen, "failed to close staged session payload"); return 1; } return rc; } void ds4_session_payload_file_free(ds4_session_payload_file *payload) { if (!payload) return; if (payload->path) { unlink(payload->path); free(payload->path); } memset(payload, 0, sizeof(*payload)); } int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out, char *err, size_t errlen) { if (!out) { payload_set_err(err, errlen, "invalid session payload staging request"); return 1; } memset(out, 0, sizeof(*out)); if (!s || !s->checkpoint_valid) { payload_set_err(err, errlen, "session has no valid checkpoint to stage"); return 1; } char tmpl[] = "/tmp/ds4-session-payload.XXXXXX"; int fd = mkstemp(tmpl); if (fd < 0) { payload_set_err(err, errlen, "failed to create staged session payload"); return 1; } FILE *fp = fdopen(fd, "wb"); if (!fp) { int saved = errno; close(fd); unlink(tmpl); if (errlen) snprintf(err, errlen, "failed to open staged session payload: %s", strerror(saved)); return 1; } int rc = ds4_session_save_payload(s, fp, err, errlen); if (rc == 0 && fflush(fp) != 0) { payload_set_err(err, errlen, "failed to flush staged session payload"); rc = 1; } off_t pos = -1; if (rc == 0) { pos = ftello(fp); if (pos < 0) { payload_set_err(err, errlen, "failed to measure staged session payload"); rc = 1; } } if (fclose(fp) != 0 && rc == 0) { payload_set_err(err, errlen, "failed to close staged session payload"); rc = 1; } if (rc != 0) { unlink(tmpl); return 1; } out->path = ds4_strdup(tmpl); out->bytes = (uint64_t)pos; return 0; } int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) { if (!s || !fp || !s->checkpoint_valid) { payload_set_err(err, errlen, "session has no valid checkpoint to save"); return 1; } if (s->distributed) { return ds4_dist_session_save_payload(s->distributed, s, fp, err, errlen); } if (ds4_session_is_glm(s)) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (!s->glm_graph_ready) { payload_set_err(err, errlen, "GLM graph is not ready for snapshot"); return 1; } if (ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator before GLM snapshot"); return 1; } ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t checkpoint_len = (uint32_t)s->checkpoint.len; const uint32_t full_live = session_glm_full_live_rows(g, checkpoint_len, s->glm_dense_cache_len); const uint32_t compact_live = session_glm_compact_live_rows(g, checkpoint_len); uint32_t header[DS4_SESSION_PAYLOAD_U32_FIELDS] = { DS4_SESSION_PAYLOAD_MAGIC, DS4_SESSION_PAYLOAD_VERSION, (uint32_t)s->ctx_size, s->prefill_cap, g->ctx_cap, g->ctx_cap, g->compact_cache_cap, checkpoint_len, g->normal_layers, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, DS4_N_VOCAB, full_live, }; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) { if (payload_write_u32(fp, header[i], err, errlen) != 0) return 1; } for (int i = 0; i < s->checkpoint.len; i++) { if (payload_write_u32(fp, (uint32_t)s->checkpoint.v[i], err, errlen) != 0) return 1; } if (payload_write_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), err, errlen) != 0) return 1; for (uint32_t il = 0; il < g->normal_layers; il++) { if (payload_write_u32(fp, compact_live, err, errlen) != 0) return 1; } for (uint32_t il = 0; il < g->normal_layers; il++) { const uint32_t index_rows = compact_live != 0 && glm_graph_layer_uses_full_indexer(il) ? compact_live : 0; if (payload_write_u32(fp, index_rows, err, errlen) != 0) return 1; } uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = 0; rc == 0 && il < g->normal_layers; il++) { rc = payload_write_glm_full_kv_span(fp, g->layer_key_cache[il], (uint64_t)full_live * g->q_dim, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_glm_full_kv_span(fp, g->layer_value_cache[il], (uint64_t)full_live * g->heads_dim, buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc != 0 || compact_live == 0) continue; rc = payload_write_glm_compact_span(fp, g->layer_kv_lora_cache[il], (uint64_t)compact_live * DS4_N_KV_LORA, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_glm_compact_span(fp, g->layer_k_rope_cache[il], (uint64_t)compact_live * DS4_N_ROT, buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc == 0 && glm_graph_layer_uses_full_indexer(il)) { rc = payload_write_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)compact_live * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); } } free(buf); return rc; #endif } if (ds4_session_is_cpu(s)) { const uint32_t raw_live = session_cpu_raw_live_rows(s); const uint32_t raw_cap = ds4_default_raw_cap((uint32_t)s->ctx_size); const uint32_t comp_cap = session_cpu_comp_cap(s); uint32_t header[DS4_SESSION_PAYLOAD_U32_FIELDS] = { DS4_SESSION_PAYLOAD_MAGIC, DS4_SESSION_PAYLOAD_VERSION, (uint32_t)s->ctx_size, s->prefill_cap, raw_cap, raw_cap, comp_cap, (uint32_t)s->checkpoint.len, DS4_N_LAYER, DS4_N_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM, DS4_N_VOCAB, raw_live, }; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) { if (payload_write_u32(fp, header[i], err, errlen) != 0) return 1; } for (int i = 0; i < s->checkpoint.len; i++) { if (payload_write_u32(fp, (uint32_t)s->checkpoint.v[i], err, errlen) != 0) return 1; } if (payload_write_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), err, errlen) != 0) return 1; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_write_u32(fp, s->cpu_cache.layer[il].n_comp, err, errlen) != 0) return 1; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_write_u32(fp, s->cpu_cache.layer[il].n_index_comp, err, errlen) != 0) return 1; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_cache *layer = &s->cpu_cache.layer[il]; if (raw_live > layer->n_raw) { payload_set_err(err, errlen, "CPU session raw cache has fewer live rows than checkpoint"); return 1; } const uint32_t raw_start = layer->n_raw - raw_live; if (payload_write_bytes(fp, layer->raw_kv + (uint64_t)raw_start * DS4_N_HEAD_DIM, (uint64_t)raw_live * DS4_N_HEAD_DIM * sizeof(float), err, errlen) != 0) return 1; const uint32_t ratio = layer->compress_ratio; if (ratio == 0) continue; if (payload_write_bytes(fp, layer->attn_comp_kv, (uint64_t)layer->n_comp * DS4_N_HEAD_DIM * sizeof(float), err, errlen) != 0) return 1; if (payload_write_bytes(fp, layer->attn_state_kv, layer_attn_state_bytes(ratio), err, errlen) != 0) return 1; if (payload_write_bytes(fp, layer->attn_state_score, layer_attn_state_bytes(ratio), err, errlen) != 0) return 1; if (ratio == 4) { if (payload_write_bytes(fp, layer->index_comp_kv, (uint64_t)layer->n_index_comp * DS4_N_INDEXER_HEAD_DIM * sizeof(float), err, errlen) != 0) return 1; if (payload_write_bytes(fp, layer->index_state_kv, layer_index_state_bytes(ratio), err, errlen) != 0) return 1; if (payload_write_bytes(fp, layer->index_state_score, layer_index_state_bytes(ratio), err, errlen) != 0) return 1; } } return 0; } #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator before snapshot"); return 1; } ds4_gpu_graph *g = &s->graph; const uint32_t raw_live = session_raw_live_rows(g, (uint32_t)s->checkpoint.len); /* Header fields: * 0 magic, 1 version, 2 ctx, 3 prefill chunk, 4 raw cap, * 5 raw window, 6 compressed cap, 7 token count, * 8 layers, 9 raw head dim, 10 indexer head dim, 11 vocab, * 12 live raw rows serialized below. */ uint32_t header[DS4_SESSION_PAYLOAD_U32_FIELDS] = { DS4_SESSION_PAYLOAD_MAGIC, DS4_SESSION_PAYLOAD_VERSION, (uint32_t)s->ctx_size, s->prefill_cap, g->raw_cap, g->raw_window, g->comp_cap, (uint32_t)s->checkpoint.len, DS4_N_LAYER, DS4_N_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM, DS4_N_VOCAB, raw_live, }; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) { if (payload_write_u32(fp, header[i], err, errlen) != 0) return 1; } for (int i = 0; i < s->checkpoint.len; i++) { if (payload_write_u32(fp, (uint32_t)s->checkpoint.v[i], err, errlen) != 0) return 1; } if (payload_write_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), err, errlen) != 0) return 1; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_write_u32(fp, g->layer_n_comp[il], err, errlen) != 0) return 1; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_write_u32(fp, g->layer_n_index_comp[il], err, errlen) != 0) return 1; } uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) { /* Write the raw ring in logical position order. The file does not care * where the rows happened to live physically in the source graph. */ const uint32_t raw_first = (uint32_t)s->checkpoint.len - raw_live; for (uint32_t r = 0; rc == 0 && r < raw_live; r++) { const uint32_t pos = raw_first + r; const uint32_t phys = pos % g->raw_cap; rc = payload_write_tensor_span(fp, g->layer_raw_cache[il], (uint64_t)phys * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); } const uint32_t ratio = ds4_layer_compress_ratio(il); if (rc != 0 || ratio == 0) continue; /* Compressed rows are append-only from row zero, so the live prefix is * contiguous. The two compressor state tensors hold the partial window * that will become the next compressed row. */ if (DS4_GPU_ATTN_COMP_CACHE_F16) { rc = payload_write_tensor_span_f16_as_f32(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); } else { rc = payload_write_tensor_span(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)g->layer_n_comp[il] * DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); } if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_attn_state_kv[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_attn_state_score[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0 && ratio == 4) { rc = payload_write_tensor_span(fp, g->layer_index_comp_cache[il], 0, (uint64_t)g->layer_n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_index_state_kv[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_index_state_score[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, err, errlen); } } free(buf); return rc; #endif } int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, char *err, size_t errlen) { if (!s || !fp) { payload_set_err(err, errlen, "invalid session payload load"); return 1; } if (s->distributed) { return ds4_dist_session_load_payload(s->distributed, s, fp, payload_bytes, err, errlen); } uint64_t remaining = payload_bytes; uint32_t h[DS4_SESSION_PAYLOAD_U32_FIELDS]; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) { if (payload_read_u32(fp, &h[i], &remaining, err, errlen) != 0) return 1; } if (h[0] != DS4_SESSION_PAYLOAD_MAGIC || h[1] != DS4_SESSION_PAYLOAD_VERSION) { payload_set_err(err, errlen, "unsupported session payload version"); return 1; } if (ds4_session_is_glm(s)) { #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else if (!s->glm_graph_ready) { payload_set_err(err, errlen, "GLM graph is not ready for KV restore"); return 1; } ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t saved_ctx = h[2]; const uint32_t saved_prefill_cap = h[3]; const uint32_t saved_full_cap = h[4]; const uint32_t saved_full_window = h[5]; const uint32_t saved_compact_cap = h[6]; const uint32_t saved_tokens = h[7]; const uint32_t saved_full_live = h[12]; if (saved_ctx > (uint32_t)s->ctx_size || saved_tokens >= (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV checkpoint does not fit current context"); return 1; } if (h[8] != g->normal_layers || h[9] != DS4_N_KEY_MLA || h[10] != DS4_N_VALUE_MLA || h[11] != DS4_N_VOCAB) { payload_set_err(err, errlen, "KV checkpoint was written for a different GLM layout"); return 1; } (void)saved_prefill_cap; if (saved_full_cap == 0 || saved_full_cap > g->ctx_cap || saved_full_window != saved_full_cap) { payload_set_err(err, errlen, "KV checkpoint GLM full-cache layout does not match current runtime"); return 1; } const uint32_t max_full_live = saved_tokens < saved_full_cap ? saved_tokens : saved_full_cap; if (saved_full_live > max_full_live || saved_full_live > g->ctx_cap) { payload_set_err(err, errlen, "KV checkpoint GLM full-cache row count is invalid"); return 1; } if (saved_compact_cap > g->ctx_size) { payload_set_err(err, errlen, "KV checkpoint GLM compact cache is larger than current context"); return 1; } if (g->compact_cache_cap != 0 && saved_tokens != 0 && saved_compact_cap == 0) { payload_set_err(err, errlen, "KV checkpoint lacks the GLM compact cache required by this context"); return 1; } const uint32_t expected_compact_live = saved_compact_cap != 0 ? (saved_tokens < saved_compact_cap ? saved_tokens : saved_compact_cap) : 0; if (expected_compact_live != 0) { if (!s->engine || !glm_graph_ensure_compact_cache(g, expected_compact_live)) { payload_set_err(err, errlen, "KV checkpoint GLM compact cache could not be allocated"); return 1; } } if (expected_compact_live > g->compact_cache_cap) { payload_set_err(err, errlen, "KV checkpoint GLM compact row count is invalid"); return 1; } token_vec new_checkpoint = {0}; for (uint32_t i = 0; i < saved_tokens; i++) { uint32_t tok = 0; if (payload_read_u32(fp, &tok, &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } token_vec_push(&new_checkpoint, (int)tok); } if (payload_read_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } uint32_t n_comp[DS4_MAX_LAYER]; uint32_t n_index_comp[DS4_MAX_LAYER]; for (uint32_t il = 0; il < g->normal_layers; il++) { if (payload_read_u32(fp, &n_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (n_comp[il] != expected_compact_live) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint GLM compact row count does not match token count"); return 1; } } for (uint32_t il = 0; il < g->normal_layers; il++) { if (payload_read_u32(fp, &n_index_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } const uint32_t expected_index_rows = expected_compact_live != 0 && glm_graph_layer_uses_full_indexer(il) ? expected_compact_live : 0; if (n_index_comp[il] != expected_index_rows) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint GLM indexer row count does not match token count"); return 1; } } if (ds4_gpu_synchronize() == 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "failed to synchronize accelerator before GLM KV restore"); return 1; } s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_reset_dense_cache(s); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = 0; rc == 0 && il < g->normal_layers; il++) { rc = payload_read_or_skip_glm_full_kv_span(fp, g->layer_key_cache[il], (uint64_t)saved_full_live * g->q_dim, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_or_skip_glm_full_kv_span(fp, g->layer_value_cache[il], (uint64_t)saved_full_live * g->heads_dim, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc != 0 || expected_compact_live == 0) continue; rc = payload_read_glm_compact_span(fp, g->layer_kv_lora_cache[il], (uint64_t)expected_compact_live * DS4_N_KV_LORA, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_glm_compact_span(fp, g->layer_k_rope_cache[il], (uint64_t)expected_compact_live * DS4_N_ROT, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc == 0 && glm_graph_layer_uses_full_indexer(il)) { rc = payload_read_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)expected_compact_live * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } } free(buf); if (rc != 0) { token_vec_free(&new_checkpoint); return 1; } if (remaining != 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has trailing payload bytes"); return 1; } if (ds4_gpu_synchronize() == 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "failed to synchronize accelerator after GLM KV restore"); return 1; } token_vec_free(&s->checkpoint); s->checkpoint = new_checkpoint; s->checkpoint_valid = true; s->mtp_draft_valid = false; s->glm_dense_cache_len = g->full_kv_cache ? saved_full_live : 0; return 0; #endif } if (ds4_session_is_cpu(s)) { const uint32_t saved_ctx = h[2]; const uint32_t saved_prefill_cap = h[3]; const uint32_t saved_raw_cap = h[4]; const uint32_t saved_raw_window = h[5]; const uint32_t saved_comp_cap = h[6]; const uint32_t saved_tokens = h[7]; const uint32_t saved_raw_live = h[12]; const uint32_t cpu_raw_cap = ds4_default_raw_cap((uint32_t)s->ctx_size); const uint32_t cpu_comp_cap = session_cpu_comp_cap(s); if (saved_ctx > (uint32_t)s->ctx_size || saved_tokens >= (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV checkpoint does not fit current context"); return 1; } if (h[8] != DS4_N_LAYER || h[9] != DS4_N_HEAD_DIM || h[10] != DS4_N_INDEXER_HEAD_DIM || h[11] != DS4_N_VOCAB) { payload_set_err(err, errlen, "KV checkpoint was written for a different DS4 layout"); return 1; } /* prefill_cap is scratch scheduling capacity, not durable KV layout. * Old checkpoints remain valid as long as the raw KV window matches. */ (void)saved_prefill_cap; if (saved_raw_window != cpu_raw_cap) { payload_set_err(err, errlen, "KV checkpoint graph chunk layout does not match current runtime"); return 1; } const uint32_t expected_raw_live = saved_tokens < saved_raw_window ? saved_tokens : saved_raw_window; if (saved_raw_cap == 0 || saved_raw_live != expected_raw_live || saved_raw_live > saved_raw_cap || saved_raw_live > cpu_raw_cap) { payload_set_err(err, errlen, "KV checkpoint raw ring layout does not match current context"); return 1; } if (saved_comp_cap > cpu_comp_cap) { payload_set_err(err, errlen, "KV checkpoint compressed cache is larger than current context"); return 1; } token_vec new_checkpoint = {0}; for (uint32_t i = 0; i < saved_tokens; i++) { uint32_t tok = 0; if (payload_read_u32(fp, &tok, &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } token_vec_push(&new_checkpoint, (int)tok); } if (payload_read_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } uint32_t n_comp[DS4_MAX_LAYER]; uint32_t n_index_comp[DS4_MAX_LAYER]; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_read_u32(fp, &n_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (n_comp[il] > saved_comp_cap || n_comp[il] > cpu_comp_cap) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has invalid compressed row count"); return 1; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_read_u32(fp, &n_index_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (n_index_comp[il] > saved_comp_cap || n_index_comp[il] > cpu_comp_cap) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has invalid indexer row count"); return 1; } } s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); session_cpu_reset_cache(s); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { ds4_layer_cache *layer = &s->cpu_cache.layer[il]; if (payload_read_bytes(fp, layer->raw_kv, (uint64_t)saved_raw_live * DS4_N_HEAD_DIM * sizeof(float), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } layer->n_raw = saved_raw_live; const uint32_t ratio = layer->compress_ratio; if (ratio == 0) continue; layer->n_comp = n_comp[il]; layer->n_index_comp = n_index_comp[il]; if (payload_read_bytes(fp, layer->attn_comp_kv, (uint64_t)n_comp[il] * DS4_N_HEAD_DIM * sizeof(float), &remaining, err, errlen) != 0 || payload_read_bytes(fp, layer->attn_state_kv, layer_attn_state_bytes(ratio), &remaining, err, errlen) != 0 || payload_read_bytes(fp, layer->attn_state_score, layer_attn_state_bytes(ratio), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (ratio == 4) { if (payload_read_bytes(fp, layer->index_comp_kv, (uint64_t)n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float), &remaining, err, errlen) != 0 || payload_read_bytes(fp, layer->index_state_kv, layer_index_state_bytes(ratio), &remaining, err, errlen) != 0 || payload_read_bytes(fp, layer->index_state_score, layer_index_state_bytes(ratio), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } } } if (remaining != 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has trailing payload bytes"); return 1; } token_vec_free(&s->checkpoint); s->checkpoint = new_checkpoint; s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); return 0; } #ifdef DS4_NO_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else ds4_gpu_graph *g = &s->graph; const uint32_t saved_ctx = h[2]; const uint32_t saved_prefill_cap = h[3]; const uint32_t saved_raw_cap = h[4]; const uint32_t saved_raw_window = h[5]; const uint32_t saved_comp_cap = h[6]; const uint32_t saved_tokens = h[7]; const uint32_t saved_raw_live = h[12]; if (saved_ctx > (uint32_t)s->ctx_size || saved_tokens >= (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV checkpoint does not fit current context"); return 1; } if (h[8] != DS4_N_LAYER || h[9] != DS4_N_HEAD_DIM || h[10] != DS4_N_INDEXER_HEAD_DIM || h[11] != DS4_N_VOCAB) { payload_set_err(err, errlen, "KV checkpoint was written for a different DS4 layout"); return 1; } /* prefill_cap is scratch scheduling capacity, not durable KV layout. * Old checkpoints remain valid as long as the raw KV window matches. */ (void)saved_prefill_cap; if (saved_raw_window != g->raw_window) { payload_set_err(err, errlen, "KV checkpoint graph chunk layout does not match current runtime"); return 1; } /* The raw rows in the file are logical rows. We can restore them into any * current ring with enough capacity, but the saved live count must be exactly * the last window implied by the saved token count. */ const uint32_t expected_raw_live = saved_tokens < saved_raw_window ? saved_tokens : saved_raw_window; if (saved_raw_cap == 0 || saved_raw_live != expected_raw_live || saved_raw_live > saved_raw_cap || saved_raw_live > g->raw_cap) { payload_set_err(err, errlen, "KV checkpoint raw ring layout does not match current context"); return 1; } if (saved_comp_cap > g->comp_cap) { payload_set_err(err, errlen, "KV checkpoint compressed cache is larger than current context"); return 1; } token_vec new_checkpoint = {0}; for (uint32_t i = 0; i < saved_tokens; i++) { uint32_t tok = 0; if (payload_read_u32(fp, &tok, &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } token_vec_push(&new_checkpoint, (int)tok); } if (payload_read_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float), &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } uint32_t n_comp[DS4_MAX_LAYER]; uint32_t n_index_comp[DS4_MAX_LAYER]; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_read_u32(fp, &n_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (n_comp[il] > saved_comp_cap || n_comp[il] > g->layer_comp_cap[il]) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has invalid compressed row count"); return 1; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (payload_read_u32(fp, &n_index_comp[il], &remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (n_index_comp[il] > saved_comp_cap || n_index_comp[il] > g->layer_comp_cap[il]) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has invalid indexer row count"); return 1; } } if (ds4_gpu_synchronize() == 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "failed to synchronize accelerator before KV restore"); return 1; } s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) { /* Rebuild the physical raw ring expected by the current graph. This is * why the file stores rows in logical order instead of dumping bytes from * the old ring layout. */ const uint32_t raw_first = saved_tokens - saved_raw_live; for (uint32_t r = 0; rc == 0 && r < saved_raw_live; r++) { const uint32_t pos = raw_first + r; const uint32_t phys = pos % g->raw_cap; rc = payload_read_tensor_span(fp, g->layer_raw_cache[il], (uint64_t)phys * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } const uint32_t ratio = ds4_layer_compress_ratio(il); if (rc != 0 || ratio == 0) continue; if (DS4_GPU_ATTN_COMP_CACHE_F16) { rc = payload_read_tensor_span_f32_as_f16(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)n_comp[il] * DS4_N_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } else { rc = payload_read_tensor_span(fp, g->layer_attn_comp_cache[il], 0, (uint64_t)n_comp[il] * DS4_N_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_attn_state_kv[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_attn_state_score[il], 0, layer_attn_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0 && ratio == 4) { rc = payload_read_tensor_span(fp, g->layer_index_comp_cache[il], 0, (uint64_t)n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_index_state_kv[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_index_state_score[il], 0, layer_index_state_bytes(ratio), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } } free(buf); if (rc != 0) { token_vec_free(&new_checkpoint); return 1; } if (remaining != 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint has trailing payload bytes"); return 1; } if (ds4_gpu_synchronize() == 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "failed to synchronize accelerator after KV restore"); return 1; } token_vec_free(&s->checkpoint); s->checkpoint = new_checkpoint; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { g->layer_n_comp[il] = n_comp[il]; g->layer_n_index_comp[il] = n_index_comp[il]; } s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); return 0; #endif } 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"); return 1; } if (s->distributed) { payload_set_err(err, errlen, "distributed session snapshots are not supported yet"); return 1; } const uint64_t bytes = ds4_session_payload_bytes(s); if (bytes == 0) { payload_set_err(err, errlen, "session has no valid checkpoint to snapshot"); return 1; } if (bytes > (uint64_t)SIZE_MAX) { payload_set_err(err, errlen, "session snapshot is too large for this platform"); return 1; } if (snap->cap < bytes) { uint8_t *p = realloc(snap->ptr, (size_t)bytes); if (!p) { payload_set_err(err, errlen, "out of memory while allocating session snapshot"); return 1; } snap->ptr = p; snap->cap = bytes; } FILE *fp = fmemopen(snap->ptr, (size_t)bytes, "wb"); if (!fp) { payload_set_err(err, errlen, "failed to open memory stream for session snapshot"); return 1; } const int rc = ds4_session_save_payload(s, fp, err, errlen); if (fclose(fp) != 0 && rc == 0) { payload_set_err(err, errlen, "failed to finalize memory session snapshot"); return 1; } if (rc != 0) return 1; snap->len = bytes; return 0; } int ds4_session_load_snapshot(ds4_session *s, const ds4_session_snapshot *snap, char *err, size_t errlen) { if (!s || !snap || !snap->ptr || snap->len == 0) { payload_set_err(err, errlen, "invalid session snapshot load"); return 1; } if (s->distributed) { payload_set_err(err, errlen, "distributed session snapshots are not supported yet"); return 1; } if (snap->len > (uint64_t)SIZE_MAX) { payload_set_err(err, errlen, "session snapshot is too large for this platform"); return 1; } FILE *fp = fmemopen((void *)snap->ptr, (size_t)snap->len, "rb"); if (!fp) { payload_set_err(err, errlen, "failed to open memory stream for session snapshot restore"); return 1; } const int rc = ds4_session_load_payload(s, fp, snap->len, err, errlen); if (fclose(fp) != 0 && rc == 0) { payload_set_err(err, errlen, "failed to close memory session snapshot"); return 1; } return rc; } void ds4_session_snapshot_free(ds4_session_snapshot *snap) { if (!snap) return; free(snap->ptr); memset(snap, 0, sizeof(*snap)); } void ds4_engine_dump_tokens(ds4_engine *e, const ds4_tokens *tokens) { dump_tokens(&e->vocab, tokens); } int ds4_dump_text_tokenization(const char *model_path, const char *text, FILE *fp) { ds4_model model; ds4_vocab vocab; token_vec tokens = {0}; if (!fp) fp = stdout; model_open(&model, model_path, false, false); config_validate_model(&model); vocab_load(&vocab, &model); tokenize_rendered_chat_vocab(&vocab, text ? text : "", &tokens); dump_tokens_fp(fp, &vocab, &tokens); token_vec_free(&tokens); vocab_free(&vocab); model_close(&model); return 0; } #ifndef DS4_NO_GPU static bool imatrix_read_text_file(const char *path, char **out, size_t *len_out) { *out = NULL; *len_out = 0; struct stat st; if (stat(path, &st) != 0) { fprintf(stderr, "ds4: failed to stat imatrix dataset %s: %s\n", path, strerror(errno)); return false; } if (st.st_size < 0 || (uint64_t)st.st_size > SIZE_MAX - 1) { fprintf(stderr, "ds4: imatrix dataset is too large: %s\n", path); return false; } FILE *fp = fopen(path, "rb"); if (!fp) { fprintf(stderr, "ds4: failed to open imatrix dataset %s: %s\n", path, strerror(errno)); return false; } size_t n = (size_t)st.st_size; char *buf = xmalloc(n + 1); if (n != 0 && fread(buf, 1, n, fp) != n) { fprintf(stderr, "ds4: failed to read imatrix dataset %s\n", path); fclose(fp); free(buf); return false; } if (fclose(fp) != 0) { fprintf(stderr, "ds4: failed to close imatrix dataset %s: %s\n", path, strerror(errno)); free(buf); return false; } buf[n] = '\0'; *out = buf; *len_out = n; return true; } static char *imatrix_trim_block(char *p, char *end) { while (p < end && isspace((unsigned char)*p)) p++; while (end > p && isspace((unsigned char)end[-1])) end--; *end = '\0'; return p; } #endif int ds4_engine_collect_imatrix(ds4_engine *e, const char *dataset_path, const char *output_path, int ctx_size, int max_prompts, int max_tokens) { #ifdef DS4_NO_GPU (void)e; (void)dataset_path; (void)output_path; (void)ctx_size; (void)max_prompts; (void)max_tokens; fprintf(stderr, "ds4: imatrix collection requires a graph backend build\n"); return 1; #else if (!e || !dataset_path || !output_path) return 1; if (e->backend != DS4_BACKEND_METAL || !e->metal_ready) { fprintf(stderr, "ds4: imatrix collection currently requires --metal\n"); return 1; } if (ctx_size <= 0) ctx_size = 32768; char *dataset = NULL; size_t dataset_len = 0; if (!imatrix_read_text_file(dataset_path, &dataset, &dataset_len)) return 1; const ds4_model *model = &e->model; const ds4_weights *weights = &e->weights; const uint32_t prefill_cap = metal_graph_prefill_cap_for_prompt(ctx_size, e->prefill_chunk); const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, prefill_cap); 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, prefill_cap, false, NULL, false, NULL); if (!ok) { fprintf(stderr, "ds4: failed to allocate imatrix Metal graph runtime\n"); free(dataset); return 1; } g.quality = e->quality; g.ssd_streaming = e->ssd_streaming; g.ssd_streaming_cold = e->ssd_streaming_cold; g.streaming_preload_experts = e->ssd_streaming_preload_experts; g.power_percent = (uint32_t)e->power_percent; ds4_imatrix_collector collector; if (!imatrix_collector_init(&collector, prefill_cap, dataset_path)) { fprintf(stderr, "ds4: failed to allocate imatrix collector\n"); metal_graph_free(&g); free(dataset); return 1; } fprintf(stderr, "ds4: collecting routed-MoE imatrix from %s (model=%s, layers=%u, experts=%u, ctx=%d, chunk=%u)\n", dataset_path, DS4_MODEL_SHAPE_NAME, DS4_N_LAYER, DS4_N_EXPERT, ctx_size, prefill_cap); int prompts_done = 0; int tokens_done = 0; char *cursor = dataset; const char *marker_lit = "===== DS4_IMATRIX_PROMPT"; while (*cursor) { char *start = cursor; char *marker = strstr(cursor, marker_lit); if (marker) { char *nl = strchr(marker, '\n'); if (!nl) break; start = nl + 1; } else if (prompts_done != 0) { break; } char *next = strstr(start, marker_lit); char *end = next ? next : dataset + dataset_len; char saved = *end; char *prompt_text = imatrix_trim_block(start, end); if (prompt_text[0] != '\0') { token_vec prompt = {0}; ds4_tokenize_rendered_chat(e, prompt_text, &prompt); if (prompt.len > ctx_size) prompt.len = ctx_size; if (max_tokens > 0 && prompt.len > max_tokens - tokens_done) { prompt.len = max_tokens - tokens_done; } if (prompt.len > 0) { if (!metal_graph_reset_prefill_state(&g)) { fprintf(stderr, "ds4: failed to reset imatrix graph state\n"); ok = false; } else if ((uint32_t)prompt.len > prefill_cap) { ok = metal_graph_prefill_chunked_range(&g, model, weights, &prompt, 0, (uint32_t)prompt.len, NULL, false, NULL, NULL, NULL, NULL, &collector, NULL, NULL, NULL); } else { ok = metal_graph_prefill_layer_major(&g, model, weights, &prompt, 0, (uint32_t)prompt.len, NULL, false, &collector, NULL, NULL); } if (!ok) { fprintf(stderr, "ds4: imatrix prefill failed at prompt %d\n", prompts_done + 1); token_vec_free(&prompt); *end = saved; break; } prompts_done++; tokens_done += prompt.len; if (prompts_done % 10 == 0) { fprintf(stderr, "ds4: imatrix prompts=%d tokens=%d routes=%llu\r", prompts_done, tokens_done, (unsigned long long)collector.observed_routes); fflush(stderr); } } token_vec_free(&prompt); } *end = saved; if (!next) break; cursor = next; if (max_prompts > 0 && prompts_done >= max_prompts) break; if (max_tokens > 0 && tokens_done >= max_tokens) break; } fputc('\n', stderr); if (ok) { ok = imatrix_collector_save(&collector, weights, output_path); if (ok) { fprintf(stderr, "ds4: wrote imatrix %s from %d prompts, %d tokens, %llu routed expert observations\n", output_path, prompts_done, tokens_done, (unsigned long long)collector.observed_routes); } } imatrix_collector_free(&collector); metal_graph_free(&g); free(dataset); return ok ? 0 : 1; #endif } #ifndef DS4_NO_GPU static bool ds4_session_greedy_splitkv_replay_exact( ds4_session *s, int current_token, int *top, char *err, size_t errlen) { if (!s || !s->greedy_splitkv_anchor_valid) return false; ds4_engine *e = s->engine; const int start = s->greedy_splitkv_anchor_len; if (!spec_frontier_restore(&s->greedy_splitkv_anchor, s)) { snprintf(err, errlen, "%s split-kv frontier restore failed", ds4_backend_name(e->backend)); return false; } s->checkpoint.len = start; s->checkpoint_valid = true; int replay_top = -1; const bool trust_replay = metal_graph_cuda_greedy_splitkv_trust_replay_requested(); if (metal_graph_cuda_greedy_splitkv_pair_replay_requested()) { const int n_seq = s->greedy_splitkv_segment.len + 1; int *seq = xmalloc((size_t)n_seq * sizeof(seq[0])); for (int i = 0; i < s->greedy_splitkv_segment.len; i++) { seq[i] = s->greedy_splitkv_segment.v[i]; } seq[n_seq - 1] = current_token; for (int i = 0; i < n_seq; ) { if (i + 1 < n_seq) { int top0 = -1; int top1 = -1; const bool second_is_current = (i + 1 == n_seq - 1); const uint32_t pos = (uint32_t)s->checkpoint.len; bool ok = metal_graph_verify_decode2_exact(&s->graph, &e->model, &e->weights, seq[i], seq[i + 1], pos, &top0, &top1, NULL, NULL); if (!ok) { snprintf(err, errlen, "%s split-kv paired exact replay failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(seq); return false; } if (!trust_replay && top0 != seq[i + 1]) { snprintf(err, errlen, "%s split-kv paired exact replay mismatch at pos=%d", ds4_backend_name(e->backend), start + i); s->checkpoint_valid = false; free(seq); return false; } if (!trust_replay && !second_is_current && i + 2 < n_seq && top1 != seq[i + 2]) { snprintf(err, errlen, "%s split-kv paired exact replay mismatch at pos=%d", ds4_backend_name(e->backend), start + i + 1); s->checkpoint_valid = false; free(seq); return false; } token_vec_push(&s->checkpoint, seq[i]); if (!second_is_current) { token_vec_push(&s->checkpoint, seq[i + 1]); } else { *top = top1; } i += 2; continue; } if (!metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, seq[i], (uint32_t)s->checkpoint.len, top, NULL, false, NULL, false)) { snprintf(err, errlen, "%s split-kv paired exact replay current token failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(seq); return false; } i++; } free(seq); s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = false; s->greedy_splitkv_anchor_len = 0; return true; } for (int i = 0; i < s->greedy_splitkv_segment.len; i++) { const int tok = s->greedy_splitkv_segment.v[i]; bool ok = false; if (trust_replay) { /* Opt-in experiment: rebuild exact state for already-emitted * tokens without validating their intermediate greedy tops. */ ok = metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, tok, (uint32_t)s->checkpoint.len, NULL); } else { ok = metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, tok, (uint32_t)s->checkpoint.len, &replay_top, NULL, false, NULL, false); } if (!ok) { snprintf(err, errlen, "%s split-kv exact replay failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return false; } token_vec_push(&s->checkpoint, tok); if (!trust_replay && i + 1 < s->greedy_splitkv_segment.len && replay_top != s->greedy_splitkv_segment.v[i + 1]) { snprintf(err, errlen, "%s split-kv exact replay mismatch at pos=%d", ds4_backend_name(e->backend), start + i); s->checkpoint_valid = false; return false; } } if (!trust_replay && s->greedy_splitkv_segment.len > 0 && replay_top != current_token) { snprintf(err, errlen, "%s split-kv exact replay mismatch before current token", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return false; } if (!metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, current_token, (uint32_t)s->checkpoint.len, top, NULL, false, NULL, false)) { snprintf(err, errlen, "%s split-kv exact replay current token failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return false; } s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = false; s->greedy_splitkv_anchor_len = 0; return true; } #endif int ds4_session_eval_argmax(ds4_session *s, int token, char *err, size_t errlen) { if (!s) return -1; if (ds4_session_is_cpu(s) || ds4_session_is_glm(s)) { if (ds4_session_eval(s, token, err, errlen) != 0) return -1; return ds4_session_argmax(s); } #ifdef DS4_NO_GPU (void)token; snprintf(err, errlen, "GPU support is not compiled in"); return -1; #else ds4_engine *e = s->engine; int top = -1; const uint32_t pos = (uint32_t)s->checkpoint.len; const bool splitkv_may_engage = metal_graph_cuda_greedy_splitkv_may_engage(&s->graph, pos); const bool splitkv_fallback = splitkv_may_engage && metal_graph_cuda_greedy_splitkv_fallback_requested(); const bool splitkv_top2 = splitkv_may_engage && metal_graph_cuda_greedy_splitkv_top2_requested(); const bool greedy_vec4 = !splitkv_may_engage && metal_graph_cuda_greedy_vec4_requested(); const bool vec4_fallback = greedy_vec4 && metal_graph_cuda_greedy_vec4_fallback_requested(); const bool approx_fallback = splitkv_fallback || vec4_fallback; bool anchor_ready = false; if (approx_fallback) { if (!s->greedy_splitkv_anchor_valid) { s->greedy_splitkv_anchor_len = (int)pos; s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = spec_frontier_snapshot(&s->greedy_splitkv_anchor, s); if (!s->greedy_splitkv_anchor_valid) { s->greedy_splitkv_anchor_len = 0; } } anchor_ready = s->greedy_splitkv_anchor_valid; if (!anchor_ready && ((splitkv_fallback && getenv("DS4_CUDA_GREEDY_SPLITKV_FALLBACK_LOG")) || (vec4_fallback && getenv("DS4_CUDA_GREEDY_VEC4_FALLBACK_LOG")))) { fprintf(stderr, "ds4: greedy %s frontier snapshot failed at pos=%u; using exact attention\n", vec4_fallback ? "vec4" : "split-kv", pos); } } metal_graph_top2_result top2; memset(&top2, 0, sizeof(top2)); const bool allow_fast_attention = !splitkv_fallback || anchor_ready; const bool use_vec4_attention = greedy_vec4 && (!vec4_fallback || anchor_ready); const int old_score_vec4 = ds4_gpu_set_decode_score_vec4(use_vec4_attention ? 1 : 0); bool ok = metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, (uint32_t)token, pos, &top, NULL, allow_fast_attention, (anchor_ready || splitkv_top2) ? &top2 : NULL, false); (void)ds4_gpu_set_decode_score_vec4(old_score_vec4); bool replayed_exact = false; if (!ok && anchor_ready) { ok = ds4_session_greedy_splitkv_replay_exact(s, token, &top, err, errlen); replayed_exact = ok; } if (ok && anchor_ready && top2.fast_attention) { const float margin = top2.value0 - top2.value1; const float threshold = metal_graph_cuda_greedy_splitkv_margin_threshold(); if (getenv("DS4_CUDA_GREEDY_SPLITKV_TRACE")) { fprintf(stderr, "ds4: greedy split-kv margin pos=%u top=%d second=%d margin=%.6f threshold=%.6f\n", pos, top2.id0, top2.id1, (double)margin, (double)threshold); } if (!top2.valid || !isfinite(margin) || margin < threshold) { const int replay_segment_len = s->greedy_splitkv_segment.len; ok = ds4_session_greedy_splitkv_replay_exact(s, token, &top, err, errlen); replayed_exact = ok; if (ok && getenv("DS4_CUDA_GREEDY_SPLITKV_FALLBACK_LOG")) { fprintf(stderr, "ds4: greedy split-kv exact replay pos=%u segment=%d margin=%.6f threshold=%.6f approx_top=%d second=%d exact_top=%d\n", pos, replay_segment_len, (double)margin, (double)threshold, top2.id0, top2.id1, top); } } } if (ok && vec4_fallback && anchor_ready && use_vec4_attention) { const float margin = top2.value0 - top2.value1; const float threshold = metal_graph_cuda_greedy_vec4_margin_threshold(); if (getenv("DS4_CUDA_GREEDY_VEC4_TRACE")) { fprintf(stderr, "ds4: greedy vec4 margin pos=%u top=%d second=%d margin=%.6f threshold=%.6f\n", pos, top2.id0, top2.id1, (double)margin, (double)threshold); } if (!top2.valid || !isfinite(margin) || margin < threshold) { const int replay_segment_len = s->greedy_splitkv_segment.len; ok = ds4_session_greedy_splitkv_replay_exact(s, token, &top, err, errlen); replayed_exact = ok; if (ok && getenv("DS4_CUDA_GREEDY_VEC4_FALLBACK_LOG")) { fprintf(stderr, "ds4: greedy vec4 exact replay pos=%u segment=%d margin=%.6f threshold=%.6f approx_top=%d second=%d exact_top=%d\n", pos, replay_segment_len, (double)margin, (double)threshold, top2.id0, top2.id1, top); } } } if (ok && approx_fallback && anchor_ready && !replayed_exact) { const uint32_t max_segment = vec4_fallback ? metal_graph_cuda_greedy_vec4_max_segment() : metal_graph_cuda_greedy_splitkv_max_segment(); if (max_segment > 0 && (uint32_t)s->greedy_splitkv_segment.len >= max_segment) { const int replay_segment_len = s->greedy_splitkv_segment.len; ok = ds4_session_greedy_splitkv_replay_exact(s, token, &top, err, errlen); replayed_exact = ok; if (ok && ((vec4_fallback && getenv("DS4_CUDA_GREEDY_VEC4_FALLBACK_LOG")) || (splitkv_fallback && getenv("DS4_CUDA_GREEDY_SPLITKV_FALLBACK_LOG")))) { fprintf(stderr, "ds4: greedy %s exact replay pos=%u segment=%d reason=max_segment max=%u approx_top=%d second=%d exact_top=%d\n", vec4_fallback ? "vec4" : "split-kv", pos, replay_segment_len, max_segment, top2.id0, top2.id1, top); } } } if (!ok) { if (errlen != 0 && err[0] == '\0') { snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); } s->checkpoint_valid = false; return -1; } token_vec_push(&s->checkpoint, token); if (approx_fallback && anchor_ready && !replayed_exact) { token_vec_push(&s->greedy_splitkv_segment, token); } else if (!approx_fallback || !anchor_ready) { session_greedy_splitkv_reset(s); } s->checkpoint_valid = true; s->mtp_draft_valid = false; return top; #endif } #ifndef DS4_NO_GPU static int ds4_session_eval_splitkv_spec_after_first( ds4_session *s, int max_extra_tokens, int eos_token, int *accepted, int accepted_cap, char *err, size_t errlen) { if (!s || max_extra_tokens < 2 || accepted_cap < 2) return 0; if (!metal_graph_cuda_splitkv_spec_requested()) return 0; ds4_engine *e = s->engine; const uint32_t start = (uint32_t)s->checkpoint.len; if ((uint64_t)start + 2u > (uint64_t)s->ctx_size) { if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec skip pos=%u reason=context-room\n", start); } return 0; } if (!metal_graph_cuda_splitkv_score_may_engage(&s->graph, start)) { if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec skip pos=%u reason=score-gate\n", start); } return 0; } const int draft0 = sample_argmax(s->logits, DS4_N_VOCAB); if (draft0 < 0 || draft0 == eos_token) { if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec skip pos=%u reason=%s draft0=%d\n", start, draft0 < 0 ? "no-draft0" : "eos", draft0); } return 0; } ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); if (!spec_frontier_snapshot(&frontier, s)) { if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec frontier snapshot failed at pos=%u; using exact decode\n", start); } return 0; } const bool timing = getenv("DS4_CUDA_SPLITKV_SPEC_TIMING") != NULL; const double t0 = timing ? now_sec() : 0.0; int draft1 = -1; bool draft_ok = metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, draft0, start, &draft1, NULL, true, NULL, true); const double t_draft = timing ? now_sec() : 0.0; const bool restored = spec_frontier_restore(&frontier, s); const double t_restore = timing ? now_sec() : 0.0; if (!restored) { snprintf(err, errlen, "%s split-kv spec frontier restore failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; spec_frontier_free(&frontier); return -1; } if (!draft_ok || draft1 < 0) { if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec draft failed at pos=%u; using exact decode\n", start); } spec_frontier_free(&frontier); return 0; } if (metal_graph_cuda_splitkv_spec_batch_verify_requested()) { float *row_logits = xmalloc( (size_t)2 * DS4_N_VOCAB * sizeof(row_logits[0])); int row0_top = -1; s->checkpoint.len = (int)start; token_vec_push(&s->checkpoint, draft0); token_vec_push(&s->checkpoint, draft1); bool ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights, &s->checkpoint, start, 2, true, false, &row0_top, row_logits, NULL); const double t_verify = timing ? now_sec() : 0.0; if (!ok) { s->checkpoint.len = (int)start; (void)spec_frontier_restore(&frontier, s); snprintf(err, errlen, "%s split-kv spec batch verifier failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(row_logits); spec_frontier_free(&frontier); return -1; } int n_accept = 0; if (row0_top == draft1) { memcpy(s->logits, row_logits + DS4_N_VOCAB, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); accepted[n_accept++] = draft0; accepted[n_accept++] = draft1; s->checkpoint_valid = true; s->mtp_draft_valid = false; } else { s->checkpoint.len = (int)start; if (!spec_frontier_commit_prefix1(s)) { (void)spec_frontier_restore(&frontier, s); snprintf(err, errlen, "%s split-kv spec batch prefix commit failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(row_logits); spec_frontier_free(&frontier); return -1; } memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); token_vec_push(&s->checkpoint, draft0); accepted[n_accept++] = draft0; s->checkpoint_valid = true; s->mtp_draft_valid = false; } if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec batch pos=%u draft0=%d draft1=%d verify_next=%d accepted=%d\n", start, draft0, draft1, row0_top, n_accept); } if (timing) { const double done = now_sec(); fprintf(stderr, "ds4: split-kv spec batch timing accepted=%d draft=%.3f ms restore=%.3f ms verify=%.3f ms total=%.3f ms\n", n_accept, (t_draft - t0) * 1000.0, (t_restore - t_draft) * 1000.0, (t_verify - t_restore) * 1000.0, (done - t0) * 1000.0); } free(row_logits); spec_frontier_free(&frontier); return n_accept; } const bool toponly_row0 = metal_graph_cuda_splitkv_spec_toponly_row0_requested(); float *row0_logits = toponly_row0 ? NULL : xmalloc((size_t)DS4_N_VOCAB * sizeof(row0_logits[0])); float *row1_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row1_logits[0])); int row0_top = -1; bool ok = metal_graph_verify_decode2_exact(&s->graph, &e->model, &e->weights, draft0, draft1, start, &row0_top, NULL, row0_logits, row1_logits); const double t_verify = timing ? now_sec() : 0.0; if (!ok) { (void)spec_frontier_restore(&frontier, s); snprintf(err, errlen, "%s split-kv spec exact verifier failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(row1_logits); free(row0_logits); spec_frontier_free(&frontier); return -1; } int n_accept = 0; if (row0_top == draft1) { memcpy(s->logits, row1_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); token_vec_push(&s->checkpoint, draft0); token_vec_push(&s->checkpoint, draft1); accepted[n_accept++] = draft0; accepted[n_accept++] = draft1; s->checkpoint_valid = true; s->mtp_draft_valid = false; } else { s->checkpoint.len = (int)start; if (toponly_row0) { if (!spec_frontier_restore(&frontier, s) || !metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, draft0, start, s->logits)) { snprintf(err, errlen, "%s split-kv spec row0 exact replay failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(row1_logits); spec_frontier_free(&frontier); return -1; } } else { if (!spec_frontier_commit_prefix1(s)) { (void)spec_frontier_restore(&frontier, s); snprintf(err, errlen, "%s split-kv spec prefix commit failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; free(row1_logits); free(row0_logits); spec_frontier_free(&frontier); return -1; } memcpy(s->logits, row0_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); } token_vec_push(&s->checkpoint, draft0); accepted[n_accept++] = draft0; s->checkpoint_valid = true; s->mtp_draft_valid = false; } if (getenv("DS4_CUDA_SPLITKV_SPEC_LOG")) { fprintf(stderr, "ds4: split-kv spec pos=%u draft0=%d draft1=%d exact_next=%d accepted=%d\n", start, draft0, draft1, row0_top, n_accept); } if (timing) { const double done = now_sec(); fprintf(stderr, "ds4: split-kv spec timing accepted=%d draft=%.3f ms restore=%.3f ms verify=%.3f ms total=%.3f ms\n", n_accept, (t_draft - t0) * 1000.0, (t_restore - t_draft) * 1000.0, (t_verify - t_restore) * 1000.0, (done - t0) * 1000.0); } free(row1_logits); free(row0_logits); spec_frontier_free(&frontier); return n_accept; } #endif /* Speculative decode state machine: * 1. commit the normal target token and use its logits to validate draft[0]; * 2. let MTP recursively draft a tiny suffix from its own raw-cache frontier; * 3. verify the suffix with the target graph, committing only the accepted * prefix and rolling back speculative Metal state on miss; * 4. fall back to ordinary one-token decode if the fast verifier cannot prove * the target stream. */ int ds4_engine_generate_argmax( ds4_engine *e, const ds4_tokens *prompt, int n_predict, int ctx_size, ds4_token_emit_fn emit, ds4_generation_done_fn done, void *emit_ud, ds4_session_progress_fn progress, void *progress_ud) { const ds4_model *model = &e->model; const ds4_vocab *vocab = &e->vocab; const ds4_weights *weights = &e->weights; if (ds4_backend_uses_graph(e->backend)) { #ifndef DS4_NO_GPU if (!e->metal_ready) { fprintf(stderr, "ds4: %s generation requested but the graph backend is unavailable\n", ds4_backend_name(e->backend)); return 1; } if (e->multi_tier) { ds4_session *s = NULL; char err[256] = {0}; const double t_prefill0 = now_sec(); if (ds4_session_create(&s, e, ctx_size) != 0) { fprintf(stderr, "ds4: failed to create multi-tier graph session\n"); return 1; } ds4_session_set_progress(s, progress, progress_ud); if (ds4_session_sync(s, prompt, err, sizeof(err)) != 0) { ds4_session_set_progress(s, NULL, NULL); fprintf(stderr, "ds4: multi-tier prefill failed: %s\n", err); ds4_session_free(s); return 1; } ds4_session_set_progress(s, NULL, NULL); const double t_prefill1 = now_sec(); int rc = 0; int n_generated = 0; const double t_decode0 = now_sec(); const bool greedy_top1 = metal_graph_tp_env_flag("DS4_CUDA_GREEDY_TOP1", true); int token = ds4_session_argmax(s); for (int i = 0; i < n_predict && ds4_session_pos(s) < ctx_size; i++) { if (token < 0) { fprintf(stderr, "ds4: multi-tier argmax failed\n"); rc = 1; break; } if (ds4_token_is_stop(e, token)) break; if (emit) emit(emit_ud, token); n_generated++; if (i == n_predict - 1 || ds4_session_pos(s) + 1 >= ctx_size) break; if (greedy_top1) { token = ds4_session_eval_argmax(s, token, err, sizeof(err)); } else if (ds4_session_eval(s, token, err, sizeof(err)) == 0) { token = ds4_session_argmax(s); } else { token = -1; } if (token < 0) { fprintf(stderr, "ds4: multi-tier decode failed: %s\n", err); rc = 1; break; } } const double t_decode1 = now_sec(); if (done) done(emit_ud); ds4_log(stderr, DS4_LOG_TIMING, "ds4: prefill: %.2f t/s, generation: %.2f t/s\n", (t_prefill1 - t_prefill0) > 0.0 ? (double)prompt->len / (t_prefill1 - t_prefill0) : 0.0, (t_decode1 - t_decode0) > 0.0 ? (double)n_generated / (t_decode1 - t_decode0) : 0.0); ds4_session_free(s); return rc; } return generate_metal_graph_raw_swa(model, vocab, weights, prompt, n_predict, ctx_size, e->quality, e->ssd_streaming, e->ssd_streaming_cold, e->ssd_streaming_preload_experts, e->ssd_streaming_cache_bytes, e->ssd_streaming_prefill_headroom_bytes, e->power_percent, e->prefill_chunk, e->directional_steering_file, e->directional_steering_attn_scale, e->directional_steering_ffn_scale, emit, done, emit_ud, progress, progress_ud); #else fprintf(stderr, "ds4: %s generation requested but this build has no graph backend support\n", ds4_backend_name(e->backend)); return 1; #endif } return generate_raw_swa_cpu(model, vocab, weights, prompt, n_predict, ctx_size, e->directional_steering_dirs, e->directional_steering_attn_scale, e->directional_steering_ffn_scale, emit, done, emit_ud, progress, progress_ud); } static int glm_metal_compare_f32( const char *name, const float *cpu, const float *gpu, uint32_t n, float tol) { float max_abs = 0.0f; uint32_t max_i = 0; double ss = 0.0; for (uint32_t i = 0; i < n; i++) { const float d = fabsf(gpu[i] - cpu[i]); if (d > max_abs) { max_abs = d; max_i = i; } ss += (double)d * (double)d; } const double rms = sqrt(ss / (double)n); printf(" %s: max_abs=%.9g at %u cpu=%.9g gpu=%.9g rms_abs=%.9g\n", name, max_abs, max_i, cpu[max_i], gpu[max_i], rms); if (max_abs > tol) { fprintf(stderr, "ds4: GLM Metal %s mismatch, max_abs %.9g > %.9g\n", name, max_abs, tol); return 0; } return 1; } static int glm_metal_compare_i32_list( const char *name, const int *cpu, const int32_t *gpu, uint32_t n) { int ok = 1; printf(" %s: cpu=[", name); for (uint32_t i = 0; i < n; i++) { printf("%s%d", i ? "," : "", cpu[i]); } printf("] gpu=["); for (uint32_t i = 0; i < n; i++) { printf("%s%d", i ? "," : "", (int)gpu[i]); if (cpu[i] != (int)gpu[i]) ok = 0; } printf("]\n"); if (!ok) { fprintf(stderr, "ds4: GLM Metal %s mismatch\n", name); return 0; } return 1; } #ifndef DS4_NO_GPU static void glm_metal_q8_diag_fill_input(float *x, uint64_t n_tok, uint64_t in_dim) { for (uint64_t t = 0; t < n_tok; t++) { for (uint64_t i = 0; i < in_dim; i++) { uint32_t s = (uint32_t)(0x9e3779b9u ^ (uint32_t)(t * 0x85ebca6bu) ^ (uint32_t)(i * 0xc2b2ae35u)); s ^= s >> 16; s *= 0x7feb352du; s ^= s >> 15; s *= 0x846ca68bu; s ^= s >> 16; const float centered = ((float)(int32_t)(s & 0xffffu) - 32768.0f) / 32768.0f; const float scale = 0.25f + 0.015625f * (float)((i + 3u * t) & 31u); x[t * in_dim + i] = centered * scale; } } } static void glm_metal_q8_diag_reference( float *out, const ds4_model *model, const ds4_tensor *w, const float *x, uint64_t n_tok) { const uint64_t out_dim = w->elements / w->dim[0]; for (uint64_t t = 0; t < n_tok; t++) { matvec_q8_0_f32_ref(out + t * out_dim, model, w, x + t * w->dim[0]); } } static int glm_metal_graph_test_q8_prefill_one( ds4_engine *e, const char *name, const ds4_tensor *w, int strict) { if (!w) return 1; if (w->type != DS4_TENSOR_Q8_0 || w->ndim < 2 || w->dim[0] == 0) { fprintf(stderr, "ds4: GLM Q8 prefill diagnostic found unexpected %s layout\n", name); return 0; } const ds4_model *model = &e->model; const uint64_t in_dim = w->dim[0]; const uint64_t out_dim = w->elements / in_dim; const uint32_t cases[] = { 16u, 17u, 31u, 32u }; int ok = 1; printf(" q8_prefill_diag: %s ndim=%u in=%llu out=%llu strict=%d\n", name, w->ndim, (unsigned long long)in_dim, (unsigned long long)out_dim, strict); for (uint32_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) { const uint32_t n_tok = cases[ci]; if (in_dim > UINT64_MAX / n_tok / sizeof(float) || out_dim > UINT64_MAX / n_tok / sizeof(float) || out_dim > UINT32_MAX / n_tok) { fprintf(stderr, "ds4: GLM Q8 prefill diagnostic size overflow in %s token case %u\n", name, n_tok); ok = 0; if (strict) break; continue; } const uint64_t x_bytes = (uint64_t)n_tok * in_dim * sizeof(float); const uint64_t out_bytes = (uint64_t)n_tok * out_dim * sizeof(float); if (x_bytes > SIZE_MAX || out_bytes > SIZE_MAX) { fprintf(stderr, "ds4: GLM Q8 prefill diagnostic host allocation is too large in %s token case %u\n", name, n_tok); ok = 0; if (strict) break; continue; } char label[128]; snprintf(label, sizeof(label), "q8_prefill_diag_%s_tok%u", name, n_tok); float *x_host = xmalloc((size_t)x_bytes); float *cpu_out = xmalloc((size_t)out_bytes); float *gpu_out = xmalloc((size_t)out_bytes); ds4_gpu_tensor *x_gpu = ds4_gpu_tensor_alloc(x_bytes); ds4_gpu_tensor *out_gpu = ds4_gpu_tensor_alloc(out_bytes); int case_ok = x_gpu && out_gpu; if (!case_ok) { fprintf(stderr, "ds4: GLM Q8 prefill diagnostic could not allocate %s token case %u\n", name, n_tok); } if (case_ok) { glm_metal_q8_diag_fill_input(x_host, n_tok, in_dim); glm_metal_q8_diag_reference(cpu_out, model, w, x_host, n_tok); case_ok = ds4_gpu_tensor_write(x_gpu, 0, x_host, x_bytes) != 0; } if (case_ok) { case_ok = ds4_gpu_matmul_q8_0_tensor(out_gpu, model->map, model->size, w->abs_offset, in_dim, out_dim, x_gpu, n_tok); } if (case_ok) { case_ok = ds4_gpu_tensor_read(out_gpu, 0, gpu_out, out_bytes) != 0; } if (case_ok) { case_ok = glm_metal_compare_f32(label, cpu_out, gpu_out, (uint32_t)((uint64_t)n_tok * out_dim), strict ? 5.0e-2f : 3.0e38f); } ds4_gpu_tensor_free(out_gpu); ds4_gpu_tensor_free(x_gpu); free(gpu_out); free(cpu_out); free(x_host); if (!case_ok) { ok = 0; if (strict) break; } } return ok || !strict; } static int glm_metal_graph_test_q8_prefill( ds4_engine *e, const ds4_layer_weights *layer) { return 1; if (!layer) { fprintf(stderr, "ds4: GLM Q8 prefill diagnostic requires layer weights\n"); return 0; } const int strict = 0; int ok = 1; const ds4_layer_weights *sparse = DS4_N_LEADING_DENSE < DS4_N_LAYER ? &e->weights.layer[DS4_N_LEADING_DENSE] : NULL; #define DS4_GLM_Q8_DIAG_ONE(label, tensor) \ do { \ if (!glm_metal_graph_test_q8_prefill_one(e, label, tensor, strict)) { \ ok = 0; \ if (strict) goto done; \ } \ } while (0) DS4_GLM_Q8_DIAG_ONE("layer0.attn_q_a", layer->attn_q_a); DS4_GLM_Q8_DIAG_ONE("layer0.attn_q_b", layer->attn_q_b); DS4_GLM_Q8_DIAG_ONE("layer0.attn_kv_a_mqa", layer->attn_kv_a_mqa); DS4_GLM_Q8_DIAG_ONE("layer0.attn_k_b", layer->attn_k_b); DS4_GLM_Q8_DIAG_ONE("layer0.attn_v_b", layer->attn_v_b); DS4_GLM_Q8_DIAG_ONE("layer0.attn_output", layer->attn_output); DS4_GLM_Q8_DIAG_ONE("layer0.ffn_gate", layer->ffn_gate); DS4_GLM_Q8_DIAG_ONE("layer0.ffn_up", layer->ffn_up); DS4_GLM_Q8_DIAG_ONE("layer0.ffn_down", layer->ffn_down); if (sparse) { DS4_GLM_Q8_DIAG_ONE("layer3.ffn_gate_shexp", sparse->ffn_gate_shexp); DS4_GLM_Q8_DIAG_ONE("layer3.ffn_up_shexp", sparse->ffn_up_shexp); DS4_GLM_Q8_DIAG_ONE("layer3.ffn_down_shexp", sparse->ffn_down_shexp); } done: #undef DS4_GLM_Q8_DIAG_ONE return strict ? ok : 1; } static int glm_metal_graph_test_multitok_attention( ds4_engine *e, const ds4_tokens *prompt, const ds4_layer_weights *layer) { if (!prompt || prompt->len < 2) { printf(" layer0_multitok_attention: skipped (prompt has fewer than 2 tokens)\n"); return 1; } const ds4_model *model = &e->model; const ds4_weights *weights = &e->weights; const uint32_t n_tok = prompt->len < 4 ? (uint32_t)prompt->len : 4u; 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 (!weights->token_embd || !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 || weights->token_embd->type != DS4_TENSOR_Q8_0 || 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) { fprintf(stderr, "ds4: GLM multi-token attention diagnostic found unexpected layer-0 layout\n"); return 0; } for (uint32_t t = 0; t < n_tok; t++) { if (prompt->v[t] < 0 || prompt->v[t] >= (int)DS4_N_VOCAB) { fprintf(stderr, "ds4: GLM multi-token attention token %d is outside vocab\n", prompt->v[t]); return 0; } } const uint64_t token_bytes = (uint64_t)n_tok * sizeof(int32_t); const uint64_t emb_bytes = (uint64_t)n_tok * DS4_N_EMBD * sizeof(float); const uint64_t q_rank_bytes = (uint64_t)n_tok * DS4_N_LORA_Q * sizeof(float); const uint64_t q_bytes = (uint64_t)n_tok * q_dim * sizeof(float); const uint64_t kv_raw_bytes = (uint64_t)n_tok * kv_raw_dim * sizeof(float); const uint64_t kv_norm_bytes = (uint64_t)n_tok * DS4_N_KV_LORA * sizeof(float); const uint64_t k_nope_bytes = (uint64_t)n_tok * DS4_N_HEAD * q_nope * sizeof(float); const uint64_t heads_bytes = (uint64_t)n_tok * heads_dim * sizeof(float); int32_t *tok_host = xmalloc((size_t)token_bytes); float *cpu_emb = xmalloc((size_t)emb_bytes); float *cpu_attn = xmalloc((size_t)emb_bytes); float *gpu_attn = xmalloc((size_t)emb_bytes); for (uint32_t t = 0; t < n_tok; t++) { tok_host[t] = (int32_t)prompt->v[t]; embed_token_any(model, weights, prompt->v[t], cpu_emb + (uint64_t)t * DS4_N_EMBD); } layer_glm_attention_prefill_f32_ref(cpu_attn, model, layer, cpu_emb, n_tok, 0, 0); ds4_gpu_tensor *tok = NULL; ds4_gpu_tensor *cur = NULL; ds4_gpu_tensor *attn_norm = NULL; ds4_gpu_tensor *q_rank = NULL; ds4_gpu_tensor *q_rank_norm = NULL; ds4_gpu_tensor *q = NULL; ds4_gpu_tensor *kv_raw = NULL; ds4_gpu_tensor *kv_norm = NULL; ds4_gpu_tensor *k_nope = NULL; ds4_gpu_tensor *value = NULL; ds4_gpu_tensor *key_cache = NULL; ds4_gpu_tensor *value_cache = NULL; ds4_gpu_tensor *heads = NULL; ds4_gpu_tensor *attn_out = NULL; int ok = 1; #define DS4_GLM_MT_ALLOC_TENSOR(var, bytes_) \ do { \ (var) = ds4_gpu_tensor_alloc((bytes_)); \ if (!(var)) { \ fprintf(stderr, "ds4: GLM multi-token attention diagnostic could not allocate %s\n", #var); \ ok = 0; \ } \ } while (0) DS4_GLM_MT_ALLOC_TENSOR(tok, token_bytes); DS4_GLM_MT_ALLOC_TENSOR(cur, emb_bytes); DS4_GLM_MT_ALLOC_TENSOR(attn_norm, emb_bytes); DS4_GLM_MT_ALLOC_TENSOR(q_rank, q_rank_bytes); DS4_GLM_MT_ALLOC_TENSOR(q_rank_norm, q_rank_bytes); DS4_GLM_MT_ALLOC_TENSOR(q, q_bytes); DS4_GLM_MT_ALLOC_TENSOR(kv_raw, kv_raw_bytes); DS4_GLM_MT_ALLOC_TENSOR(kv_norm, kv_norm_bytes); DS4_GLM_MT_ALLOC_TENSOR(k_nope, k_nope_bytes); DS4_GLM_MT_ALLOC_TENSOR(value, heads_bytes); DS4_GLM_MT_ALLOC_TENSOR(key_cache, q_bytes); DS4_GLM_MT_ALLOC_TENSOR(value_cache, heads_bytes); DS4_GLM_MT_ALLOC_TENSOR(heads, heads_bytes); DS4_GLM_MT_ALLOC_TENSOR(attn_out, emb_bytes); #undef DS4_GLM_MT_ALLOC_TENSOR const uint32_t rope_ctx = 0; const float rope_base = layer_rope_freq_base(0); const float rope_scale = layer_rope_freq_scale(0); const float rope_ext = 0.0f; const float rope_attn = 1.0f; if (ok) ok = ds4_gpu_tensor_write(tok, 0, tok_host, token_bytes) != 0; if (ok) ok = ds4_gpu_embed_tokens_q8_0_tensor(cur, tok, model->map, model->size, weights->token_embd->abs_offset, DS4_N_VOCAB, n_tok, DS4_N_EMBD); if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(attn_norm, cur, model->map, model->size, layer->attn_norm->abs_offset, DS4_N_EMBD, n_tok, DS4_RMS_EPS); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(q_rank, model->map, model->size, layer->attn_q_a->abs_offset, DS4_N_EMBD, DS4_N_LORA_Q, attn_norm, n_tok); if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(q_rank_norm, q_rank, model->map, model->size, layer->attn_q_a_norm->abs_offset, DS4_N_LORA_Q, n_tok, DS4_RMS_EPS); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(q, model->map, model->size, layer->attn_q_b->abs_offset, DS4_N_LORA_Q, q_dim, q_rank_norm, n_tok); if (ok) ok = ds4_gpu_rope_tail_tensor(q, n_tok, DS4_N_HEAD, qk_dim, DS4_N_ROT, 0, rope_ctx, false, rope_base, rope_scale, rope_ext, rope_attn, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(kv_raw, model->map, model->size, layer->attn_kv_a_mqa->abs_offset, DS4_N_EMBD, kv_raw_dim, attn_norm, n_tok); if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(kv_norm, kv_raw, model->map, model->size, layer->attn_kv_a_norm->abs_offset, n_tok, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, DS4_RMS_EPS); if (ok) ok = ds4_gpu_glm_k_b_project_tensor(k_nope, kv_norm, model->map, model->size, layer->attn_k_b->abs_offset, n_tok, DS4_N_KV_LORA, q_nope, DS4_N_HEAD); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(value, model->map, model->size, layer->attn_v_b->abs_offset, DS4_N_KV_LORA, heads_dim, kv_norm, n_tok); if (ok) ok = ds4_gpu_glm_build_kv_cache_tensor(key_cache, value_cache, kv_raw, k_nope, value, 0, n_tok, n_tok, DS4_N_HEAD, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, q_nope, DS4_N_ROT, DS4_N_VALUE_MLA, rope_ctx, rope_base, rope_scale, rope_ext, rope_attn, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, false); if (ok) { if (glm_graph_flash_attention_prefill_enabled()) { ok = ds4_gpu_glm_attention_flash_tensor(heads, q, key_cache, value_cache, 0, n_tok, n_tok, n_tok, DS4_N_HEAD, qk_dim, DS4_N_VALUE_MLA, false); } else { ok = ds4_gpu_glm_attention_full_tensor(heads, q, key_cache, value_cache, 0, n_tok, n_tok, n_tok, DS4_N_HEAD, qk_dim, DS4_N_VALUE_MLA, false); } } if (ok) ok = ds4_gpu_matmul_q8_0_tensor(attn_out, model->map, model->size, layer->attn_output->abs_offset, heads_dim, DS4_N_EMBD, heads, n_tok); if (ok) ok = ds4_gpu_tensor_read(attn_out, 0, gpu_attn, emb_bytes) != 0; if (ok) { printf(" layer0_multitok_attention: tokens=%u\n", n_tok); ok = glm_metal_compare_f32("layer0_multitok_attention", cpu_attn, gpu_attn, (uint32_t)((uint64_t)n_tok * DS4_N_EMBD), 5.0e-1f); } ds4_gpu_tensor_free(attn_out); ds4_gpu_tensor_free(heads); ds4_gpu_tensor_free(value_cache); ds4_gpu_tensor_free(key_cache); ds4_gpu_tensor_free(value); ds4_gpu_tensor_free(k_nope); ds4_gpu_tensor_free(kv_norm); ds4_gpu_tensor_free(kv_raw); ds4_gpu_tensor_free(q); ds4_gpu_tensor_free(q_rank_norm); ds4_gpu_tensor_free(q_rank); ds4_gpu_tensor_free(attn_norm); ds4_gpu_tensor_free(cur); ds4_gpu_tensor_free(tok); free(gpu_attn); free(cpu_attn); free(cpu_emb); free(tok_host); return ok; } /* * Validate the single-token decode attention path at pos > 0: prior tokens are * stored into the compact DSA caches with the batch store used by prefill, * then the last token runs through the exact decode kernel sequence * (rope tail, compact store, selected range, qk-lowrank, indexed decode * attention). The pos-0-only diagnostics cannot see position-dependent * decode bugs; this one can. */ static int glm_metal_graph_test_decode_attention( ds4_engine *e, const ds4_tokens *prompt, const ds4_layer_weights *layer) { if (!prompt || prompt->len < 2) { printf(" layer0_decode_attention: skipped (prompt has fewer than 2 tokens)\n"); return 1; } const ds4_model *model = &e->model; const ds4_weights *weights = &e->weights; const uint32_t n_tok = prompt->len < 8 ? (uint32_t)prompt->len : 8u; const uint32_t n_prev = n_tok - 1u; const uint32_t pos = n_tok - 1u; 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; const uint64_t cache_elem = glm_graph_compact_cache_elem_bytes(); const uint32_t cache_f16 = glm_graph_compact_cache_is_f16(); const float rope_base = layer_rope_freq_base(0); const float rope_scale = layer_rope_freq_scale(0); for (uint32_t t = 0; t < n_tok; t++) { if (prompt->v[t] < 0 || prompt->v[t] >= (int)DS4_N_VOCAB) { fprintf(stderr, "ds4: GLM decode attention token %d is outside vocab\n", prompt->v[t]); return 0; } } const uint64_t emb_bytes = (uint64_t)n_tok * DS4_N_EMBD * sizeof(float); int32_t *tok_host = xmalloc((size_t)n_prev * sizeof(int32_t)); float *cpu_emb = xmalloc((size_t)emb_bytes); float *cpu_attn = xmalloc((size_t)emb_bytes); float *gpu_attn = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); for (uint32_t t = 0; t < n_tok; t++) { embed_token_any(model, weights, prompt->v[t], cpu_emb + (uint64_t)t * DS4_N_EMBD); } for (uint32_t t = 0; t < n_prev; t++) tok_host[t] = (int32_t)prompt->v[t]; layer_glm_attention_prefill_f32_ref(cpu_attn, model, layer, cpu_emb, n_tok, 0, 0); ds4_gpu_tensor *tok = NULL; ds4_gpu_tensor *cur_b = NULL; ds4_gpu_tensor *attn_norm_b = NULL; ds4_gpu_tensor *kv_raw_b = NULL; ds4_gpu_tensor *kv_norm_b = NULL; ds4_gpu_tensor *cur1 = NULL; ds4_gpu_tensor *attn_norm1 = NULL; ds4_gpu_tensor *q_rank = NULL; ds4_gpu_tensor *q_rank_norm = NULL; ds4_gpu_tensor *q = NULL; ds4_gpu_tensor *kv_raw1 = NULL; ds4_gpu_tensor *kv_norm1 = NULL; ds4_gpu_tensor *kv_lora_cache = NULL; ds4_gpu_tensor *k_rope_cache = NULL; ds4_gpu_tensor *selected = NULL; ds4_gpu_tensor *qk_low = NULL; ds4_gpu_tensor *heads = NULL; ds4_gpu_tensor *attn_out = NULL; int ok = 1; #define DS4_GLM_DA_ALLOC_TENSOR(var, bytes_) \ do { \ (var) = ds4_gpu_tensor_alloc((bytes_)); \ if (!(var)) { \ fprintf(stderr, "ds4: GLM decode attention diagnostic could not allocate %s\n", #var); \ ok = 0; \ } \ } while (0) DS4_GLM_DA_ALLOC_TENSOR(tok, (uint64_t)n_prev * sizeof(int32_t)); DS4_GLM_DA_ALLOC_TENSOR(cur_b, (uint64_t)n_prev * DS4_N_EMBD * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(attn_norm_b, (uint64_t)n_prev * DS4_N_EMBD * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(kv_raw_b, (uint64_t)n_prev * kv_raw_dim * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(kv_norm_b, (uint64_t)n_prev * DS4_N_KV_LORA * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(cur1, (uint64_t)DS4_N_EMBD * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(attn_norm1, (uint64_t)DS4_N_EMBD * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(q_rank, (uint64_t)DS4_N_LORA_Q * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(q_rank_norm, (uint64_t)DS4_N_LORA_Q * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(q, q_dim * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(kv_raw1, kv_raw_dim * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(kv_norm1, (uint64_t)DS4_N_KV_LORA * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(kv_lora_cache, (uint64_t)n_tok * DS4_N_KV_LORA * cache_elem); DS4_GLM_DA_ALLOC_TENSOR(k_rope_cache, (uint64_t)n_tok * DS4_N_ROT * cache_elem); DS4_GLM_DA_ALLOC_TENSOR(selected, (uint64_t)n_tok * sizeof(int32_t)); DS4_GLM_DA_ALLOC_TENSOR(qk_low, (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(heads, heads_dim * sizeof(float)); DS4_GLM_DA_ALLOC_TENSOR(attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)); #undef DS4_GLM_DA_ALLOC_TENSOR /* Store rows 0..n_prev-1 into the compact caches with the batch store. */ if (ok) ok = ds4_gpu_tensor_write(tok, 0, tok_host, (uint64_t)n_prev * sizeof(int32_t)) != 0; if (ok) ok = ds4_gpu_embed_tokens_q8_0_tensor(cur_b, tok, model->map, model->size, weights->token_embd->abs_offset, DS4_N_VOCAB, n_prev, DS4_N_EMBD); if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(attn_norm_b, cur_b, model->map, model->size, layer->attn_norm->abs_offset, DS4_N_EMBD, n_prev, DS4_RMS_EPS); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(kv_raw_b, model->map, model->size, layer->attn_kv_a_mqa->abs_offset, DS4_N_EMBD, kv_raw_dim, attn_norm_b, n_prev); if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(kv_norm_b, kv_raw_b, model->map, model->size, layer->attn_kv_a_norm->abs_offset, n_prev, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, DS4_RMS_EPS); if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(kv_lora_cache, k_rope_cache, kv_norm_b, kv_raw_b, 0, n_prev, n_tok, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, DS4_N_ROT, cache_f16) != 0; /* Run the last token through the decode kernel sequence at pos > 0. */ if (ok) ok = ds4_gpu_embed_token_q8_0_tensor(cur1, model->map, model->size, weights->token_embd->abs_offset, DS4_N_VOCAB, (uint32_t)prompt->v[pos], DS4_N_EMBD) != 0; if (ok) ok = ds4_gpu_rms_norm_weight_tensor(attn_norm1, cur1, model->map, model->size, layer->attn_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; if (ok) ok = ds4_gpu_matmul_q8_0_tensor(q_rank, model->map, model->size, layer->attn_q_a->abs_offset, DS4_N_EMBD, DS4_N_LORA_Q, attn_norm1, 1); if (ok) ok = ds4_gpu_rms_norm_weight_tensor(q_rank_norm, q_rank, model->map, model->size, layer->attn_q_a_norm->abs_offset, DS4_N_LORA_Q, DS4_RMS_EPS) != 0; if (ok) ok = ds4_gpu_matmul_q8_0_tensor(q, model->map, model->size, layer->attn_q_b->abs_offset, DS4_N_LORA_Q, q_dim, q_rank_norm, 1); if (ok) ok = ds4_gpu_glm_rope_tail_tensor(q, 1, DS4_N_HEAD, qk_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) ok = ds4_gpu_matmul_q8_0_tensor(kv_raw1, model->map, model->size, layer->attn_kv_a_mqa->abs_offset, DS4_N_EMBD, kv_raw_dim, attn_norm1, 1); if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(kv_norm1, kv_raw1, model->map, model->size, layer->attn_kv_a_norm->abs_offset, 1, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, DS4_RMS_EPS); if (ok) ok = ds4_gpu_glm_store_compact_kv_tensor(kv_lora_cache, k_rope_cache, kv_norm1, kv_raw1, pos, 1, n_tok, (uint32_t)kv_raw_dim, DS4_N_KV_LORA, DS4_N_ROT, cache_f16) != 0; if (ok) ok = ds4_gpu_glm_fill_selected_range_tensor(selected, n_tok) != 0; if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor( qk_low, q, model->map, model->size, layer->attn_k_b->abs_offset, layer->attn_k_b->type, DS4_N_HEAD, DS4_N_KV_LORA, q_nope, qk_dim) != 0; if (ok) ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( heads, q, qk_low, kv_lora_cache, k_rope_cache, model->map, model->size, layer->attn_v_b->abs_offset, layer->attn_v_b->type, selected, n_tok, n_tok, cache_f16, DS4_N_HEAD, DS4_N_KV_LORA, 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; if (ok) ok = ds4_gpu_matmul_q8_0_tensor(attn_out, model->map, model->size, layer->attn_output->abs_offset, heads_dim, DS4_N_EMBD, heads, 1); if (ok) ok = ds4_gpu_tensor_read(attn_out, 0, gpu_attn, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; if (ok) { printf(" layer0_decode_attention: tokens=%u pos=%u\n", n_tok, pos); ok = glm_metal_compare_f32("layer0_decode_attention", cpu_attn + (uint64_t)pos * DS4_N_EMBD, gpu_attn, DS4_N_EMBD, 5.0e-1f); } ds4_gpu_tensor_free(attn_out); ds4_gpu_tensor_free(heads); ds4_gpu_tensor_free(qk_low); ds4_gpu_tensor_free(selected); ds4_gpu_tensor_free(k_rope_cache); ds4_gpu_tensor_free(kv_lora_cache); ds4_gpu_tensor_free(kv_norm1); ds4_gpu_tensor_free(kv_raw1); ds4_gpu_tensor_free(q); ds4_gpu_tensor_free(q_rank_norm); ds4_gpu_tensor_free(q_rank); ds4_gpu_tensor_free(attn_norm1); ds4_gpu_tensor_free(cur1); ds4_gpu_tensor_free(kv_norm_b); ds4_gpu_tensor_free(kv_raw_b); ds4_gpu_tensor_free(attn_norm_b); ds4_gpu_tensor_free(cur_b); ds4_gpu_tensor_free(tok); free(gpu_attn); free(cpu_attn); free(cpu_emb); free(tok_host); return ok; } static int glm_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt) { if (!prompt || prompt->len <= 0) { fprintf(stderr, "ds4: GLM Metal graph test requires a non-empty prompt\n"); return 1; } const int token = prompt->v[0]; if (token < 0 || token >= (int)DS4_N_VOCAB) { fprintf(stderr, "ds4: GLM Metal graph test token %d is outside vocab\n", token); return 1; } if (!e->weights.token_embd || e->weights.token_embd->type != DS4_TENSOR_Q8_0) { fprintf(stderr, "ds4: GLM Metal graph test requires Q8_0 token embeddings\n"); return 1; } if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) { fprintf(stderr, "ds4: GLM Metal graph test has no normal transformer layers\n"); return 1; } const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; const uint32_t sparse_il = normal_layers > 8u ? 8u : DS4_N_LEADING_DENSE; if (sparse_il >= normal_layers) { fprintf(stderr, "ds4: GLM Metal graph test found no sparse GLM layer\n"); return 1; } const ds4_model *model = &e->model; const ds4_weights *weights = &e->weights; const ds4_layer_weights *layer = &weights->layer[0]; const ds4_layer_weights *sparse_layer = &weights->layer[sparse_il]; 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; const uint64_t ffn_hidden = layer->ffn_gate ? layer->ffn_gate->dim[1] : 0; const uint64_t sparse_mid_dim = sparse_layer->ffn_gate_exps ? sparse_layer->ffn_gate_exps->dim[1] : 0; const uint32_t sparse_gate_type = sparse_layer->ffn_gate_exps ? sparse_layer->ffn_gate_exps->type : 0; const uint32_t sparse_up_type = sparse_layer->ffn_up_exps ? sparse_layer->ffn_up_exps->type : 0; const bool sparse_gate_pair_supported = glm_graph_gate_pair_type_supported(sparse_gate_type, sparse_up_type); if (!layer->attn_norm || !layer->attn_kv_a_mqa || !layer->attn_kv_a_norm || !layer->attn_v_b || !layer->attn_output || !layer->ffn_norm || !layer->ffn_gate || !layer->ffn_up || !layer->ffn_down || kv_raw_dim < 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] != heads_dim || layer->attn_output->dim[1] != DS4_N_EMBD || layer->ffn_gate->dim[0] != DS4_N_EMBD || layer->ffn_up->dim[0] != DS4_N_EMBD || layer->ffn_up->dim[1] != ffn_hidden || layer->ffn_down->dim[0] != ffn_hidden || layer->ffn_down->dim[1] != DS4_N_EMBD) { fprintf(stderr, "ds4: GLM Metal graph test found unexpected layer-0 tensor layout\n"); return 1; } if (!sparse_layer->ffn_gate_inp || !sparse_layer->ffn_exp_probs_b || !sparse_layer->ffn_norm || sparse_layer->ffn_gate_inp->type != DS4_TENSOR_F32 || sparse_layer->ffn_gate_inp->dim[0] != DS4_N_EMBD || sparse_layer->ffn_gate_inp->dim[1] != DS4_N_EXPERT || sparse_layer->ffn_exp_probs_b->type != DS4_TENSOR_F32 || sparse_layer->ffn_exp_probs_b->dim[0] != DS4_N_EXPERT || sparse_layer->ffn_norm->type != DS4_TENSOR_F32 || sparse_layer->ffn_norm->dim[0] != DS4_N_EMBD || !sparse_layer->ffn_gate_exps || !sparse_layer->ffn_up_exps || !sparse_layer->ffn_down_exps || !sparse_layer->ffn_gate_shexp || !sparse_layer->ffn_up_shexp || !sparse_layer->ffn_down_shexp || !sparse_gate_pair_supported || !glm_graph_down_type_supported(sparse_layer->ffn_down_exps->type) || sparse_layer->ffn_gate_shexp->type != DS4_TENSOR_Q8_0 || sparse_layer->ffn_up_shexp->type != DS4_TENSOR_Q8_0 || sparse_layer->ffn_down_shexp->type != DS4_TENSOR_Q8_0 || sparse_layer->ffn_gate_exps->dim[0] != DS4_N_EMBD || sparse_layer->ffn_gate_exps->dim[1] != DS4_N_FF_EXP || sparse_layer->ffn_gate_exps->dim[2] != DS4_N_EXPERT || sparse_layer->ffn_up_exps->dim[0] != DS4_N_EMBD || sparse_layer->ffn_up_exps->dim[1] != DS4_N_FF_EXP || sparse_layer->ffn_up_exps->dim[2] != DS4_N_EXPERT || sparse_layer->ffn_down_exps->dim[0] != DS4_N_FF_EXP || sparse_layer->ffn_down_exps->dim[1] != DS4_N_EMBD || sparse_layer->ffn_down_exps->dim[2] != DS4_N_EXPERT || sparse_layer->ffn_gate_shexp->dim[0] != DS4_N_EMBD || sparse_layer->ffn_gate_shexp->dim[1] != DS4_N_FF_EXP || sparse_layer->ffn_up_shexp->dim[0] != DS4_N_EMBD || sparse_layer->ffn_up_shexp->dim[1] != DS4_N_FF_EXP || sparse_layer->ffn_down_shexp->dim[0] != DS4_N_FF_EXP || sparse_layer->ffn_down_shexp->dim[1] != DS4_N_EMBD || sparse_mid_dim != DS4_N_FF_EXP) { fprintf(stderr, "ds4: GLM Metal graph test found unexpected layer-%u router layout\n", sparse_il); 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 * sparse_mid_dim; const uint64_t ffn_mid_elems = ffn_hidden > sparse_mid_elems ? ffn_hidden : 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 max_read_elems = ffn_hidden > DS4_N_EMBD ? ffn_hidden : (uint64_t)DS4_N_EMBD; const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); float *cpu_emb = xmalloc(emb_bytes); float *cpu_layer = xmalloc(emb_bytes); float *cpu_cur = xmalloc(emb_bytes); float *cpu_next = xmalloc(emb_bytes); float *cpu_sparse_attn = xmalloc(emb_bytes); float *cpu_sparse_after_attn = xmalloc(emb_bytes); float *cpu_sparse_norm = xmalloc(emb_bytes); float *cpu_sparse_mid = xmalloc((size_t)sparse_mid_elems * sizeof(cpu_sparse_mid[0])); float *cpu_sparse_moe = xmalloc(emb_bytes); float *cpu_sparse_shared = xmalloc(emb_bytes); float *cpu_sparse_ffn = xmalloc(emb_bytes); float *cpu_full_hidden = xmalloc(emb_bytes); float *gpu_full_hidden = xmalloc(emb_bytes); float *cpu_logits = xmalloc((size_t)logits_bytes); float *gpu_logits = xmalloc((size_t)logits_bytes); float *cpu_router_logits = xmalloc((size_t)DS4_N_EXPERT * sizeof(cpu_router_logits[0])); float *gpu_router_logits = xmalloc((size_t)DS4_N_EXPERT * sizeof(gpu_router_logits[0])); float *gpu_read = xmalloc((size_t)max_read_elems * sizeof(gpu_read[0])); int cpu_router_selected[DS4_MAX_EXPERT_USED] = {0}; int32_t gpu_router_selected[DS4_MAX_EXPERT_USED] = {0}; float cpu_router_weights[DS4_MAX_EXPERT_USED] = {0}; float gpu_router_weights[DS4_MAX_EXPERT_USED] = {0}; char label_router_logits[64]; char label_router_selected[64]; char label_router_weights[64]; char label_routed_moe[64]; char label_shared_expert[64]; char label_sparse_ffn[64]; snprintf(label_router_logits, sizeof(label_router_logits), "layer%u_router_logits", sparse_il); snprintf(label_router_selected, sizeof(label_router_selected), "layer%u_router_selected", sparse_il); snprintf(label_router_weights, sizeof(label_router_weights), "layer%u_router_weights", sparse_il); snprintf(label_routed_moe, sizeof(label_routed_moe), "layer%u_routed_moe", sparse_il); snprintf(label_shared_expert, sizeof(label_shared_expert), "layer%u_shared_expert", sparse_il); snprintf(label_sparse_ffn, sizeof(label_sparse_ffn), "layer%u_sparse_ffn", sparse_il); 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_ALLOC_TENSOR(var, bytes_) \ do { \ (var) = ds4_gpu_tensor_alloc((bytes_)); \ if (!(var)) { \ fprintf(stderr, "ds4: GLM Metal graph test could not allocate %s\n", #var); \ ok = 0; \ } \ } while (0) DS4_GLM_ALLOC_TENSOR(cur, emb_bytes); DS4_GLM_ALLOC_TENSOR(attn_norm, emb_bytes); DS4_GLM_ALLOC_TENSOR(kv_raw, kv_raw_dim * sizeof(float)); DS4_GLM_ALLOC_TENSOR(kv_norm, (uint64_t)DS4_N_KV_LORA * sizeof(float)); DS4_GLM_ALLOC_TENSOR(heads, heads_dim * sizeof(float)); DS4_GLM_ALLOC_TENSOR(attn_out, emb_bytes); DS4_GLM_ALLOC_TENSOR(after_attn, emb_bytes); DS4_GLM_ALLOC_TENSOR(ffn_norm, emb_bytes); DS4_GLM_ALLOC_TENSOR(ffn_gate, ffn_hidden * sizeof(float)); DS4_GLM_ALLOC_TENSOR(ffn_up, ffn_hidden * sizeof(float)); DS4_GLM_ALLOC_TENSOR(ffn_mid, ffn_mid_elems * sizeof(float)); if (glm_graph_layer_uses_generic_routed_moe(sparse_layer)) { DS4_GLM_ALLOC_TENSOR(routed_gate, routed_mid_bytes); DS4_GLM_ALLOC_TENSOR(routed_up, routed_mid_bytes); DS4_GLM_ALLOC_TENSOR(routed_down, routed_down_bytes); } DS4_GLM_ALLOC_TENSOR(ffn_out, emb_bytes); DS4_GLM_ALLOC_TENSOR(ffn_sum, emb_bytes); DS4_GLM_ALLOC_TENSOR(next, emb_bytes); DS4_GLM_ALLOC_TENSOR(router_logits, (uint64_t)DS4_N_EXPERT * sizeof(float)); DS4_GLM_ALLOC_TENSOR(router_probs, (uint64_t)DS4_N_EXPERT * sizeof(float)); DS4_GLM_ALLOC_TENSOR(router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); DS4_GLM_ALLOC_TENSOR(router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); DS4_GLM_ALLOC_TENSOR(logits, logits_bytes); #undef DS4_GLM_ALLOC_TENSOR if (ok) { embed_token_any(model, weights, token, cpu_emb); layer_glm_first_token_one(cpu_layer, model, layer, cpu_emb, 0); 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); } if (ok) ok = ds4_gpu_tensor_read(cur, 0, gpu_read, emb_bytes); if (ok) { printf("GLM Metal graph test token=%d", token); if (e->vocab.token && token < e->vocab.n_vocab) { const ds4_str s = e->vocab.token[token]; printf(" text=%.*s", (int)s.len, s.ptr); } printf("\n"); ok = glm_metal_compare_f32("embedding_q8_0", cpu_emb, gpu_read, DS4_N_EMBD, 1.0e-3f); } if (ok) ok = ds4_gpu_rms_norm_weight_tensor(attn_norm, cur, model->map, model->size, layer->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, layer->attn_kv_a_mqa->abs_offset, DS4_N_EMBD, kv_raw_dim, attn_norm, 1); if (ok) ok = ds4_gpu_rms_norm_weight_tensor(kv_norm, kv_raw, model->map, model->size, layer->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, layer->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, layer->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, layer->ffn_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_gate, model->map, model->size, layer->ffn_gate->abs_offset, DS4_N_EMBD, ffn_hidden, ffn_norm, 1); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_up, model->map, model->size, layer->ffn_up->abs_offset, DS4_N_EMBD, ffn_hidden, ffn_norm, 1); if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, ffn_gate, ffn_up, (uint32_t)ffn_hidden, 0.0f, 1.0f); if (ok) ok = ds4_gpu_matmul_q8_0_tensor(ffn_out, model->map, model->size, layer->ffn_down->abs_offset, ffn_hidden, DS4_N_EMBD, ffn_mid, 1); if (ok) ok = ds4_gpu_add_tensor(next, after_attn, ffn_out, DS4_N_EMBD); if (ok) ok = ds4_gpu_tensor_read(next, 0, gpu_read, emb_bytes); if (ok) { ok = glm_metal_compare_f32("layer0_hidden", cpu_layer, gpu_read, DS4_N_EMBD, 5.0e-2f); } if (ok) ok = glm_metal_graph_test_multitok_attention(e, prompt, layer); if (ok) ok = glm_metal_graph_test_decode_attention(e, prompt, layer); if (ok) ok = glm_metal_graph_test_q8_prefill(e, layer); if (ok) { memcpy(cpu_cur, cpu_emb, emb_bytes); for (uint32_t il = 0; il < sparse_il; il++) { layer_glm_first_token_one(cpu_next, model, &weights->layer[il], cpu_cur, il); float *tmp = cpu_cur; cpu_cur = cpu_next; cpu_next = tmp; } layer_glm_first_token_attention_one(cpu_sparse_attn, model, sparse_layer, cpu_cur); for (uint32_t i = 0; i < DS4_N_EMBD; i++) { cpu_sparse_after_attn[i] = cpu_cur[i] + cpu_sparse_attn[i]; } rms_norm_weight(cpu_sparse_norm, cpu_sparse_after_attn, tensor_data(model, sparse_layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); matvec_any(cpu_router_logits, model, sparse_layer->ffn_gate_inp, cpu_sparse_norm); layer_glm_router_selected_experts(cpu_router_selected, cpu_router_weights, model, sparse_layer, cpu_sparse_norm); ok = ds4_gpu_tensor_write(ffn_norm, 0, cpu_sparse_norm, emb_bytes) != 0; } if (ok) ok = ds4_gpu_matmul_f32_tensor(router_logits, model->map, model->size, sparse_layer->ffn_gate_inp->abs_offset, DS4_N_EMBD, DS4_N_EXPERT, ffn_norm, 1); if (ok) ok = ds4_gpu_tensor_read(router_logits, 0, gpu_router_logits, (uint64_t)DS4_N_EXPERT * sizeof(gpu_router_logits[0])) != 0; if (ok) { ok = glm_metal_compare_f32(label_router_logits, cpu_router_logits, gpu_router_logits, DS4_N_EXPERT, 2.0e-3f); } if (ok) ok = ds4_gpu_glm_router_select_tensor(router_selected, router_weights, router_probs, model->map, model->size, sparse_layer->ffn_exp_probs_b->abs_offset, router_logits, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE); if (ok) ok = ds4_gpu_tensor_read(router_selected, 0, gpu_router_selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(gpu_router_selected[0])) != 0; if (ok) ok = ds4_gpu_tensor_read(router_weights, 0, gpu_router_weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(gpu_router_weights[0])) != 0; if (ok) { ok = glm_metal_compare_i32_list(label_router_selected, cpu_router_selected, gpu_router_selected, DS4_N_EXPERT_USED); } if (ok) { ok = glm_metal_compare_f32(label_router_weights, cpu_router_weights, gpu_router_weights, DS4_N_EXPERT_USED, 1.0e-4f); } 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, sparse_layer->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); (void)tensor_expert_bytes(model, sparse_layer->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); (void)tensor_expert_bytes(model, sparse_layer->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 graph test found unexpected layer-%u expert strides\n", sparse_il); ok = 0; } else { const ds4_gpu_stream_expert_table table = { .model_map = model->map, .model_size = model->size, .layer = sparse_il, .n_total_expert = DS4_N_EXPERT, .gate_offset = sparse_layer->ffn_gate_exps->abs_offset, .up_offset = sparse_layer->ffn_up_exps->abs_offset, .down_offset = sparse_layer->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; if (ok) { layer_glm_routed_moe_one_f32_ref(cpu_sparse_moe, cpu_sparse_mid, model, sparse_layer, cpu_sparse_norm, cpu_router_selected, cpu_router_weights); } ds4_glm_gpu_graph route_g = { .routed_gate = routed_gate, .routed_up = routed_up, .routed_down = routed_down, .ssd_streaming = e->ssd_streaming, }; if (ok) ok = glm_graph_routed_moe_one_dispatch( &route_g, model, sparse_layer, sparse_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_tensor_read(ffn_out, 0, gpu_read, emb_bytes) != 0; if (ok) { ok = glm_metal_compare_f32(label_routed_moe, cpu_sparse_moe, gpu_read, DS4_N_EMBD, 1.0e-1f); } if (ok) { layer_glm_shared_ffn_one_f32_ref(cpu_sparse_shared, model, sparse_layer, cpu_sparse_norm); for (uint32_t i = 0; i < DS4_N_EMBD; i++) { cpu_sparse_ffn[i] = cpu_sparse_moe[i] + cpu_sparse_shared[i]; } ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( ffn_gate, ffn_up, ffn_mid, model->map, model->size, sparse_layer->ffn_gate_shexp->abs_offset, sparse_layer->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(next, model->map, model->size, sparse_layer->ffn_down_shexp->abs_offset, DS4_N_FF_EXP, DS4_N_EMBD, ffn_mid, 1); if (ok) ok = ds4_gpu_tensor_read(next, 0, gpu_read, emb_bytes) != 0; if (ok) { ok = glm_metal_compare_f32(label_shared_expert, cpu_sparse_shared, gpu_read, DS4_N_EMBD, 5.0e-2f); } if (ok) ok = ds4_gpu_add_tensor(after_attn, ffn_out, next, DS4_N_EMBD); if (ok) ok = ds4_gpu_tensor_read(after_attn, 0, gpu_read, emb_bytes) != 0; if (ok) { ok = glm_metal_compare_f32(label_sparse_ffn, cpu_sparse_ffn, gpu_read, DS4_N_EMBD, 5.0e-2f); } if (ok) { printf(" all_layers_reference: cpu_f32 normal_layers=%u\n", normal_layers); fflush(stdout); forward_glm_first_token_cpu_f32_ref(cpu_full_hidden, model, weights, token); output_logits_glm_one_f32_ref(cpu_logits, model, weights, cpu_full_hidden); 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->dim[0] != DS4_N_EMBD || 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->dim[0] != heads_dim || gl->attn_output->dim[1] != DS4_N_EMBD) { fprintf(stderr, "ds4: GLM Metal all-layer test 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 > ffn_hidden) { fprintf(stderr, "ds4: GLM Metal all-layer test 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 || (uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP > ffn_mid_elems) { fprintf(stderr, "ds4: GLM Metal all-layer test 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 all-layer test 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 = e->ssd_streaming, }; 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_tensor_read(cur, 0, gpu_full_hidden, emb_bytes) != 0; if (ok) { ok = glm_metal_compare_f32("all_layers_hidden", cpu_full_hidden, gpu_full_hidden, DS4_N_EMBD, 5.0f); } 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, gpu_logits, logits_bytes) != 0; if (ok) { ok = glm_metal_compare_f32("first_token_logits", cpu_logits, gpu_logits, DS4_N_VOCAB, 10.0f); } 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); free(gpu_read); free(gpu_router_logits); free(cpu_router_logits); free(gpu_logits); free(cpu_logits); free(gpu_full_hidden); free(cpu_full_hidden); free(cpu_sparse_ffn); free(cpu_sparse_shared); free(cpu_sparse_moe); free(cpu_sparse_mid); free(cpu_sparse_norm); free(cpu_sparse_after_attn); free(cpu_sparse_attn); free(cpu_next); free(cpu_cur); free(cpu_layer); free(cpu_emb); return ok ? 0 : 1; } #endif int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt) { #ifndef DS4_NO_GPU if (!e->metal_ready) { fprintf(stderr, "ds4: %s graph test requested but backend is unavailable\n", ds4_backend_name(e->backend)); return 1; } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { return glm_metal_graph_test(e, prompt); } return metal_graph_decode_test(&e->model, &e->weights, prompt, e->quality); #else (void)e; (void)prompt; fprintf(stderr, "ds4: graph test requested but this build has no graph backend support\n"); return 1; #endif } int ds4_engine_metal_graph_full_test(ds4_engine *e, const ds4_tokens *prompt) { #ifndef DS4_NO_GPU if (!e->metal_ready) { fprintf(stderr, "ds4: %s full graph test requested but backend is unavailable\n", ds4_backend_name(e->backend)); return 1; } return metal_graph_first_token_full_test(&e->model, &e->weights, prompt, e->quality); #else (void)e; (void)prompt; fprintf(stderr, "ds4: full graph test requested but this build has no graph backend support\n"); return 1; #endif } int ds4_engine_metal_graph_prompt_test(ds4_engine *e, const ds4_tokens *prompt, int ctx_size) { #ifndef DS4_NO_GPU if (!e->metal_ready) { fprintf(stderr, "ds4: %s prompt graph test requested but backend is unavailable\n", ds4_backend_name(e->backend)); return 1; } return metal_graph_prompt_logits_test(&e->model, &e->weights, prompt, ctx_size); #else (void)e; (void)prompt; (void)ctx_size; fprintf(stderr, "ds4: prompt graph test requested but this build has no graph backend support\n"); return 1; #endif } int ds4_engine_head_test(ds4_engine *e, const ds4_tokens *prompt) { if (!prompt || prompt->len <= 0) { fprintf(stderr, "ds4: head test requires a non-empty prompt\n"); return 1; } const ds4_model *model = &e->model; const ds4_vocab *vocab = &e->vocab; const ds4_weights *weights = &e->weights; const ds4_layer_weights *layer0 = &weights->layer[0]; float *prompt_embd = xmalloc((size_t)prompt->len * DS4_N_EMBD * sizeof(prompt_embd[0])); embed_prompt(model, weights, prompt, DS4_N_EMBD, prompt_embd); const uint32_t n_hc = DS4_N_HC; float *hc0 = xmalloc((size_t)DS4_N_EMBD * sizeof(hc0[0])); float *residual_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(residual_hc[0])); float hc_post[4]; float hc_comb[16]; layer_attn_pre_one(model, layer0, prompt_embd + (uint64_t)(prompt->len - 1) * DS4_N_EMBD, hc0, residual_hc, hc_post, hc_comb); print_vec_stats("blk.0 attn_pre", hc0, DS4_N_EMBD); float *attn_norm0 = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm0[0])); layer_attn_norm_one(attn_norm0, model, layer0, hc0); const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; float *q0 = xmalloc((size_t)q_dim * sizeof(q0[0])); layer_q_projection_normed_one(model, layer0, attn_norm0, q0); print_vec_stats("blk.0 q", q0, q_dim); float *kv0 = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv0[0])); layer_kv_projection_normed_one(model, layer0, attn_norm0, kv0); print_vec_stats("blk.0 kv", kv0, DS4_N_HEAD_DIM); rope_tail_layer_inplace(q0, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, (uint32_t)(prompt->len - 1), 0, false); rope_tail_layer_inplace(kv0, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, (uint32_t)(prompt->len - 1), 0, false); dsv4_fp8_kv_quantize_row_inplace_cpu(kv0, DS4_N_HEAD_DIM, DS4_N_ROT); f16_round_inplace_cpu(kv0, DS4_N_HEAD_DIM); float *attn_heads = xmalloc((size_t)q_dim * sizeof(attn_heads[0])); layer_attention_one(attn_heads, model, layer0, q0, kv0); print_vec_stats("blk.0 attn_heads", attn_heads, q_dim); rope_tail_layer_inplace(attn_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, (uint32_t)(prompt->len - 1), 0, true); float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); layer_grouped_out_one(attn_out, model, layer0, attn_heads); print_vec_stats("blk.0 attn_out", attn_out, DS4_N_EMBD); float *after_attn_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(after_attn_hc[0])); hc_post_one(after_attn_hc, attn_out, residual_hc, hc_post, hc_comb, DS4_N_EMBD, n_hc); print_vec_stats("blk.0 after_attn_hc", after_attn_hc, (uint64_t)n_hc * DS4_N_EMBD); float *after_ffn_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(after_ffn_hc[0])); layer_ffn_one(after_ffn_hc, model, layer0, after_attn_hc, 0, prompt->v[prompt->len - 1], NULL, 0.0f, true); print_vec_stats("blk.0 after_ffn_hc", after_ffn_hc, (uint64_t)n_hc * DS4_N_EMBD); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); output_logits_one(logits, model, weights, after_ffn_hc); print_vec_stats("logits", logits, DS4_N_VOCAB); int best[8]; for (int i = 0; i < 8; i++) best[i] = -1; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { for (int j = 0; j < 8; j++) { if (best[j] < 0 || logits[i] > logits[best[j]]) { for (int k = 7; k > j; k--) best[k] = best[k - 1]; best[j] = (int)i; break; } } } printf("top logits after native blk.0 slice:\n"); for (int i = 0; i < 8; i++) { printf(" %6d %9.4f %.*s\n", best[i], logits[best[i]], (int)vocab->token[best[i]].len, vocab->token[best[i]].ptr); } free(logits); free(after_ffn_hc); free(after_attn_hc); free(attn_out); free(attn_heads); free(kv0); free(q0); free(attn_norm0); free(residual_hc); free(hc0); free(prompt_embd); return 0; } int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt) { if (!prompt || prompt->len <= 0) { fprintf(stderr, "ds4: first-token test requires a non-empty prompt\n"); return 1; } const ds4_model *model = &e->model; const ds4_vocab *vocab = &e->vocab; const ds4_weights *weights = &e->weights; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { float *hidden = xmalloc((size_t)DS4_N_EMBD * sizeof(hidden[0])); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); forward_glm_first_token_cpu(hidden, model, weights, prompt->v[0]); print_vec_stats("first-token final_hidden", hidden, DS4_N_EMBD); output_logits_glm_one(logits, model, weights, hidden); print_vec_stats("first-token logits", logits, DS4_N_VOCAB); int best[8]; for (int i = 0; i < 8; i++) best[i] = -1; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { for (int j = 0; j < 8; j++) { if (best[j] < 0 || logits[i] > logits[best[j]]) { for (int k = 7; k > j; k--) best[k] = best[k - 1]; best[j] = (int)i; break; } } } printf("top logits after GLM first-token CPU pass:\n"); for (int i = 0; i < 8; i++) { printf(" %6d %9.4f %.*s\n", best[i], logits[best[i]], (int)vocab->token[best[i]].len, vocab->token[best[i]].ptr); } free(logits); free(hidden); return 0; } float *hc = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(hc[0])); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); forward_first_token_cpu(hc, model, weights, prompt->v[0]); print_vec_stats("first-token final_hc", hc, (uint64_t)DS4_N_HC * DS4_N_EMBD); output_logits_one(logits, model, weights, hc); print_vec_stats("first-token logits", logits, DS4_N_VOCAB); int best[8]; for (int i = 0; i < 8; i++) best[i] = -1; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { for (int j = 0; j < 8; j++) { if (best[j] < 0 || logits[i] > logits[best[j]]) { for (int k = 7; k > j; k--) best[k] = best[k - 1]; best[j] = (int)i; break; } } } printf("top logits after first-token whole-model CPU pass:\n"); for (int i = 0; i < 8; i++) { printf(" %6d %9.4f %.*s\n", best[i], logits[best[i]], (int)vocab->token[best[i]].len, vocab->token[best[i]].ptr); } free(logits); free(hc); return 0; } static bool ds4_engine_configure_streaming_auto_cache(ds4_engine *e) { #ifdef DS4_NO_GPU (void)e; return true; #else if (!e || !e->ssd_streaming || !ds4_backend_supports_ssd_streaming(e->backend) || e->ssd_streaming_cache_experts != 0 || e->ssd_streaming_cache_bytes != 0) { return true; } if (!ds4_backend_supports_streaming_auto_cache(e->backend)) { return true; } const uint64_t recommended = ds4_gpu_recommended_working_set_size(); if (recommended == 0) { fprintf(stderr, "ds4: SSD streaming auto cache: recommended working set unavailable; " "set --ssd-streaming-cache-experts N or NGB explicitly\n"); return false; } uint64_t non_routed_bytes = 0; if (!weights_streaming_non_routed_bytes(&e->weights, &non_routed_bytes)) { fprintf(stderr, "ds4: SSD streaming auto cache could not measure non-routed model weights\n"); return false; } uint64_t per_expert_bytes = 0; if (!ds4_streaming_routed_expert_bytes(&e->weights, &per_expert_bytes)) { fprintf(stderr, "ds4: SSD streaming auto cache could not measure routed expert size\n"); return false; } const uint64_t max_model_experts = (uint64_t)DS4_N_LAYER * (uint64_t)DS4_N_EXPERT; ds4_ssd_cache_plan plan; if (!ds4_ssd_auto_cache_plan(recommended, non_routed_bytes, per_expert_bytes, max_model_experts, &plan)) { fprintf(stderr, "ds4: SSD streaming auto cache could not compute a valid cache budget\n"); return false; } uint32_t cache_experts = plan.cache_experts; uint64_t effective_cache_bytes = plan.effective_cache_bytes; const bool glm_full_layer_reserve = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && ds4_backend_supports_glm_streaming_full_layers(e->backend); /* * The 12 GiB GLM graph cap is a Metal-era VM pressure guard. ROCm/Strix * reports its own working-set recommendation and is hurt badly by this cap: * it leaves less than one routed token's experts in the dynamic cache. */ const bool glm_auto_cap = glm_full_layer_reserve && e->backend == DS4_BACKEND_METAL; const uint64_t glm_auto_cap_bytes = 12ull * 1024ull * 1024ull * 1024ull; if (glm_auto_cap && effective_cache_bytes > glm_auto_cap_bytes) { uint64_t capped_experts = glm_auto_cap_bytes / per_expert_bytes; if (capped_experts == 0) capped_experts = 1; if (capped_experts > max_model_experts) capped_experts = max_model_experts; cache_experts = capped_experts > UINT32_MAX ? UINT32_MAX : (uint32_t)capped_experts; effective_cache_bytes = (uint64_t)cache_experts * per_expert_bytes; } #ifdef DS4_ROCM_BUILD uint64_t glm_rocm_guard_cap_bytes = 0; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && e->backend == DS4_BACKEND_CUDA) { const int requested_ctx = e->placement_ctx_hint > 0 ? e->placement_ctx_hint : 4096; uint32_t guard_ctx = 0; if (!glm_graph_context_request(requested_ctx, &guard_ctx)) return false; const uint32_t work_ctx = glm_graph_full_attention_cap(guard_ctx, true); const uint32_t compact_cap = glm_graph_compact_cache_initial_cap(guard_ctx, work_ctx); const ds4_context_memory graph_mem = glm_graph_context_memory_estimate_for_compact_cap( guard_ctx, work_ctx, compact_cap, true); uint64_t active_model_bytes = glm_graph_streaming_active_model_bytes(&e->weights); if (non_routed_bytes > active_model_bytes) { active_model_bytes = non_routed_bytes; } const double fraction = glm_graph_env_double( "DS4_GLM_MEMORY_GUARD_FRACTION", 0.99, 0.50, 1.00); const double reserve_gib = glm_graph_env_double( "DS4_GLM_MEMORY_GUARD_RESERVE_GB", glm_graph_memory_guard_default_reserve_gib( recommended, active_model_bytes), 0.0, 1024.0); const uint64_t fraction_budget = (uint64_t)((double)recommended * fraction); const uint64_t reserve_bytes = (uint64_t)(reserve_gib * 1024.0 * 1024.0 * 1024.0); const uint64_t reserve_budget = reserve_bytes >= recommended ? 0 : recommended - reserve_bytes; uint64_t guard_budget = fraction_budget; if (reserve_bytes != 0 && reserve_budget < guard_budget) { guard_budget = reserve_budget; } const uint64_t fixed_bytes = glm_graph_saturating_add_u64( active_model_bytes, graph_mem.total_bytes); if (guard_budget > fixed_bytes) { glm_rocm_guard_cap_bytes = guard_budget - fixed_bytes; } if (glm_rocm_guard_cap_bytes < per_expert_bytes) { fprintf(stderr, "ds4: GLM ROCm auto cache has no room after model, graph, " "and memory-guard reserves\n"); return false; } if (effective_cache_bytes > glm_rocm_guard_cap_bytes) { uint64_t capped_experts = glm_rocm_guard_cap_bytes / per_expert_bytes; if (capped_experts > max_model_experts) { capped_experts = max_model_experts; } cache_experts = capped_experts > UINT32_MAX ? UINT32_MAX : (uint32_t)capped_experts; effective_cache_bytes = (uint64_t)cache_experts * per_expert_bytes; } } #endif e->ssd_streaming_cache_experts = cache_experts; e->ssd_streaming_cache_bytes = effective_cache_bytes; fprintf(stderr, "ds4: SSD streaming auto cache budget\n"); fprintf(stderr, "ds4: %s recommends %.2f GiB working set\n", ds4_backend_name(e->backend), (double)recommended / 1073741824.0); fprintf(stderr, "ds4: using %.0f%% total for model + cached experts: %.2f GiB\n", recommended != 0 ? 100.0 * (double)plan.model_target_bytes / (double)recommended : 0.0, (double)plan.model_target_bytes / 1073741824.0); fprintf(stderr, "ds4: non-routed weights: %.2f GiB\n", (double)non_routed_bytes / 1073741824.0); fprintf(stderr, "ds4: routed expert size: %.2f MiB\n", (double)per_expert_bytes / 1048576.0); if (glm_full_layer_reserve) { fprintf(stderr, "ds4: expert budget before prefill/full-layer reserve: %u " "(%.2f GiB)\n", e->ssd_streaming_cache_experts, (double)effective_cache_bytes / 1073741824.0); } else { fprintf(stderr, "ds4: expert budget before prefill reserve: %u (%.2f GiB)\n", e->ssd_streaming_cache_experts, (double)effective_cache_bytes / 1073741824.0); } if (glm_auto_cap && plan.effective_cache_bytes != effective_cache_bytes) { fprintf(stderr, "ds4: GLM graph auto cache capped at %.2f GiB; pass " "--ssd-streaming-cache-experts NGB to override\n", (double)glm_auto_cap_bytes / 1073741824.0); } #ifdef DS4_ROCM_BUILD if (glm_rocm_guard_cap_bytes != 0 && plan.effective_cache_bytes != effective_cache_bytes) { fprintf(stderr, "ds4: GLM ROCm cache capped to %.2f GiB by the memory " "guard for ctx=%d\n", (double)effective_cache_bytes / 1073741824.0, e->placement_ctx_hint > 0 ? e->placement_ctx_hint : 4096); } #endif if (plan.model_target_bytes <= non_routed_bytes) { fprintf(stderr, "ds4: note: non-routed weights already fill the 80%% target; keeping a one-expert cache\n"); } return true; #endif } static uint32_t ds4_glm_streaming_normal_layer_count(void) { if (DS4_N_NEXTN_PREDICT != 0 && DS4_N_LAYER > DS4_N_NEXTN_PREDICT) { return DS4_N_LAYER - DS4_N_NEXTN_PREDICT; } return DS4_N_LAYER; } static uint32_t ds4_glm_streaming_supported_resident_prefix_layers( const ds4_weights *weights) { if (!weights) return 0; const uint32_t normal_layers = ds4_glm_streaming_normal_layer_count(); if (normal_layers <= DS4_N_LEADING_DENSE) return 0; uint32_t n = 0; for (uint32_t il = DS4_N_LEADING_DENSE; il < normal_layers; il++) { if (!glm_stream_resident_decode_layer_supported(&weights->layer[il], il)) { break; } n++; } return n; } static bool ds4_glm_streaming_resident_prefix_bytes( const ds4_weights *weights, uint32_t layers, uint64_t *bytes_out) { if (bytes_out) *bytes_out = 0; if (!weights || !bytes_out || DS4_N_EXPERT == 0) return false; uint64_t total = 0; for (uint32_t i = 0; i < layers; i++) { const uint32_t il = DS4_N_LEADING_DENSE + i; if (!glm_stream_resident_decode_layer_supported(&weights->layer[il], il)) { return false; } uint64_t per_expert_bytes = 0; if (!streaming_layer_routed_expert_bytes(&weights->layer[il], &per_expert_bytes) || 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 (total > UINT64_MAX - layer_bytes) return false; total += layer_bytes; } *bytes_out = total; return true; } static uint32_t ds4_glm_streaming_auto_full_layers( const ds4_weights *weights, uint32_t supported_layers, uint64_t total_budget_bytes) { if (!weights || supported_layers == 0 || total_budget_bytes == 0) { return 0; } const uint64_t max_auto_bytes = 10ull * 1024ull * 1024ull * 1024ull; uint64_t target_bytes = total_budget_bytes / 7ull; if (target_bytes > max_auto_bytes) target_bytes = max_auto_bytes; if (target_bytes == 0) return 0; uint32_t best = 0; for (uint32_t n = 1; n <= supported_layers; n++) { uint64_t bytes = 0; if (!ds4_glm_streaming_resident_prefix_bytes(weights, n, &bytes)) { break; } if (bytes > target_bytes) break; best = n; } return best; } static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { g_glm_streaming_full_resident_layers = 0; #ifdef DS4_NO_GPU (void)e; return true; #else if (!e || !e->ssd_streaming) return true; const bool glm_full_layer_streaming = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && ds4_backend_supports_glm_streaming_full_layers(e->backend); if (e->ssd_streaming_full_layers_set && e->ssd_streaming_full_layers != 0 && !glm_full_layer_streaming) { fprintf(stderr, "ds4: --ssd-streaming-full-layers is currently supported " "only for GLM graph SSD streaming; ignoring\n"); e->ssd_streaming_full_layers = 0; } uint64_t per_expert_bytes = 0; const bool need_expert_bytes = e->ssd_streaming_cache_bytes != 0 || (glm_full_layer_streaming && !e->ssd_streaming_full_layers_set) || e->ssd_streaming_full_layers != 0; if (need_expert_bytes && !ds4_streaming_routed_expert_bytes(&e->weights, &per_expert_bytes)) { fprintf(stderr, "ds4: SSD streaming could not measure routed expert size\n"); return false; } uint32_t full_layers = 0; uint64_t full_layers_bytes = 0; bool full_layers_auto = false; uint32_t supported = 0; uint64_t total_cache_bytes = e->ssd_streaming_cache_bytes; uint64_t prefill_headroom_bytes = 0; uint64_t budget_after_prefill_headroom = total_cache_bytes; if (total_cache_bytes != 0) { if (!ds4_streaming_prefill_headroom_bytes(&e->weights, &prefill_headroom_bytes)) { fprintf(stderr, "ds4: SSD streaming prefill headroom byte accounting failed\n"); return false; } if (prefill_headroom_bytes >= total_cache_bytes) { fprintf(stderr, "ds4: --ssd-streaming-cache-experts byte budget %.2f GiB " "is too small: two routed prefill layers need %.2f GiB\n", (double)total_cache_bytes / 1073741824.0, (double)prefill_headroom_bytes / 1073741824.0); return false; } budget_after_prefill_headroom = total_cache_bytes - prefill_headroom_bytes; } e->ssd_streaming_prefill_headroom_bytes = prefill_headroom_bytes; if (glm_full_layer_streaming) { supported = ds4_glm_streaming_supported_resident_prefix_layers(&e->weights); if (!e->ssd_streaming_full_layers_set && e->ssd_streaming_cache_bytes != 0) { /* * On ROCm/Strix, the dynamic selected-expert cache is a better use * of memory than pinning full routed layers by default. Manual * --ssd-streaming-full-layers remains available for experiments. */ #ifdef DS4_ROCM_BUILD if (e->backend == DS4_BACKEND_CUDA) { e->ssd_streaming_full_layers = 0; } else #endif { e->ssd_streaming_full_layers = ds4_glm_streaming_auto_full_layers( &e->weights, supported, budget_after_prefill_headroom); } full_layers_auto = true; } } if (e->ssd_streaming_full_layers != 0) { const uint32_t requested = e->ssd_streaming_full_layers; const uint32_t supported = ds4_glm_streaming_supported_resident_prefix_layers(&e->weights); full_layers = requested < supported ? requested : supported; if (full_layers != requested) { fprintf(stderr, "ds4: GLM SSD streaming full resident layers capped " "from %u to %u supported routed prefix layers\n", requested, full_layers); } if (total_cache_bytes != 0 && per_expert_bytes != 0) { /* * The dynamic cache must be large enough for batch prefill too: * selected-address prefill can see many unique experts for a * layer, not just the top-k used by one decode token. */ uint64_t min_dynamic_experts = DS4_N_EXPERT != 0 ? (uint64_t)DS4_N_EXPERT : 1u; if (min_dynamic_experts > UINT64_MAX / per_expert_bytes) { fprintf(stderr, "ds4: SSD streaming full-layer budget overflow\n"); return false; } const uint64_t min_dynamic_bytes = min_dynamic_experts * per_expert_bytes; const uint32_t before_budget_cap = full_layers; while (full_layers != 0) { uint64_t bytes = 0; if (!ds4_glm_streaming_resident_prefix_bytes(&e->weights, full_layers, &bytes)) { fprintf(stderr, "ds4: SSD streaming full-layer byte accounting failed\n"); return false; } if (budget_after_prefill_headroom > min_dynamic_bytes && bytes <= budget_after_prefill_headroom - min_dynamic_bytes) { full_layers_bytes = bytes; break; } full_layers--; } if (full_layers == 0) full_layers_bytes = 0; if (full_layers != before_budget_cap) { fprintf(stderr, "ds4: GLM SSD streaming full resident layers capped " "from %u to %u by %.2f GiB total expert budget\n", before_budget_cap, full_layers, (double)total_cache_bytes / 1073741824.0); } } else if (full_layers != 0 && !ds4_glm_streaming_resident_prefix_bytes(&e->weights, full_layers, &full_layers_bytes)) { fprintf(stderr, "ds4: SSD streaming full-layer byte accounting failed\n"); return false; } } if (total_cache_bytes != 0) { uint64_t dynamic_cache_bytes = budget_after_prefill_headroom; if (full_layers_bytes != 0) { if (full_layers_bytes >= dynamic_cache_bytes) { fprintf(stderr, "ds4: SSD streaming full-layer budget leaves no dynamic expert cache\n"); return false; } dynamic_cache_bytes -= full_layers_bytes; } uint64_t budget_expert_bytes = 0; const uint32_t budget = ds4_streaming_cache_experts_for_byte_budget( &e->weights, dynamic_cache_bytes, &budget_expert_bytes); if (budget == 0 || budget_expert_bytes == 0) { fprintf(stderr, "ds4: --ssd-streaming-cache-experts byte budget is too small or invalid for this model\n"); return false; } e->ssd_streaming_cache_experts = budget; e->ssd_streaming_cache_bytes = (uint64_t)budget * budget_expert_bytes; if (full_layers != 0) { fprintf(stderr, "ds4: GLM SSD streaming full resident layers: %u " "(%.2f GiB, %s)\n", full_layers, (double)full_layers_bytes / 1073741824.0, full_layers_auto ? "auto" : "explicit"); fprintf(stderr, "ds4: %s SSD streaming total expert budget %.2f GiB = " "%.2f GiB prefill headroom + %.2f GiB full layers + " "%.2f GiB dynamic cache (%u experts, %.2f MiB each)\n", ds4_backend_name(e->backend), (double)total_cache_bytes / 1073741824.0, (double)prefill_headroom_bytes / 1073741824.0, (double)full_layers_bytes / 1073741824.0, (double)e->ssd_streaming_cache_bytes / 1073741824.0, budget, (double)budget_expert_bytes / 1048576.0); } else { if (glm_full_layer_streaming) { const char *reason = e->ssd_streaming_full_layers_set ? "disabled by --ssd-streaming-full-layers 0" : "auto selected 0 layers"; fprintf(stderr, "ds4: GLM SSD streaming full resident layers: 0 (%s)\n", reason); } fprintf(stderr, "ds4: %s SSD streaming total expert budget %.2f GiB = " "%.2f GiB prefill headroom + %.2f GiB dynamic cache " "(%u experts, %.2f MiB each)\n", ds4_backend_name(e->backend), (double)total_cache_bytes / 1073741824.0, (double)prefill_headroom_bytes / 1073741824.0, (double)e->ssd_streaming_cache_bytes / 1073741824.0, budget, (double)budget_expert_bytes / 1048576.0); } } else if (full_layers != 0) { fprintf(stderr, "ds4: GLM SSD streaming full resident layers: %u " "(%.2f GiB, %s), dynamic cache budget: %u experts\n", full_layers, (double)full_layers_bytes / 1073741824.0, full_layers_auto ? "auto" : "explicit", e->ssd_streaming_cache_experts); } else if (glm_full_layer_streaming) { const char *reason = e->ssd_streaming_full_layers_set ? "disabled by --ssd-streaming-full-layers 0" : "not auto-tuned for expert-count cache budgets"; fprintf(stderr, "ds4: GLM SSD streaming full resident layers: 0 (%s)\n", reason); } e->ssd_streaming_full_layers = full_layers; e->ssd_streaming_full_layer_bytes = full_layers_bytes; g_glm_streaming_full_resident_layers = full_layers; return true; #endif } static bool ds4_engine_glm_streaming_memory_guard( const ds4_engine *e, bool load_slice, uint32_t load_layer_start, uint32_t load_layer_end, bool load_output, int ctx_size, const char *phase) { #ifdef DS4_NO_GPU (void)e; (void)load_slice; (void)load_layer_start; (void)load_layer_end; (void)load_output; (void)ctx_size; (void)phase; return true; #else if (!e || DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || ctx_size <= 0) { return true; } const uint64_t transient_extra_bytes = ds4_engine_streaming_transient_guard_bytes(e); if (transient_extra_bytes == 0) return true; uint32_t guard_ctx = 0; if (!glm_graph_context_request(ctx_size, &guard_ctx)) return false; if (load_slice) { return glm_graph_memory_guard_slice_with_transient( &e->model, &e->weights, e->ssd_streaming, load_layer_start, load_layer_end, load_layer_start == 0, load_output, guard_ctx, transient_extra_bytes, phase); } return glm_graph_memory_guard_with_transient( &e->model, &e->weights, e->ssd_streaming, guard_ctx, transient_extra_bytes, phase); #endif } static bool ds4_engine_preload_pro_q4_expert_tables( ds4_engine *e, bool load_slice, uint32_t load_layer_start, uint32_t load_layer_end) { #ifdef DS4_NO_GPU (void)e; (void)load_slice; (void)load_layer_start; (void)load_layer_end; return true; #else if (!e || e->backend != DS4_BACKEND_METAL || e->ssd_streaming || DS4_MODEL_VARIANT != DS4_VARIANT_PRO || getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_PRELOAD") != NULL) { return true; } if (!metal_graph_q4_non_streaming_opt_in_enabled()) { return true; } if (ds4_gpu_pro_q4_expert_table_auto_available() == 0 && getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL && getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") == NULL && getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL) { return true; } uint32_t start = load_slice ? load_layer_start : 0; uint32_t end = load_slice ? load_layer_end : DS4_N_LAYER - 1u; if (start >= DS4_N_LAYER) return true; if (end == UINT32_MAX || end >= DS4_N_LAYER) end = DS4_N_LAYER - 1u; if (end < start) return true; bool any = false; const double t0 = now_sec(); for (uint32_t il = start; il <= end; il++) { const ds4_layer_weights *layer = &e->weights.layer[il]; if (!layer->ffn_gate_exps || !layer->ffn_up_exps || !layer->ffn_down_exps) { continue; } if (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) { 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: PRO Q4 expert table preload 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 (!ds4_gpu_preload_q4_expert_tables(e->model.map, e->model.size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, layer->ffn_down_exps->abs_offset, gate_expert_bytes, down_expert_bytes, DS4_N_EXPERT)) { fprintf(stderr, "ds4: Metal failed to preload PRO Q4 expert tables for layer %u\n", il); return false; } any = true; } if (any) { fprintf(stderr, "ds4: Metal preloaded PRO Q4 expert tables for layers %u:%u in %.2fs\n", start, end, now_sec() - t0); } return true; #endif } /* TP sharding: touch the dense weights and only this rank's contiguous range * of every routed-expert blob, so the other range is never faulted in. * Replaces the whole-file residency request; the pages * fault through the same view buffers, streaming-style, but everything a * rank will ever read is pre-faulted here. */ static void model_warm_weights_sharded(const ds4_model *m, const ds4_weights *w, int rank) { typedef struct { uint64_t off, len; } skip_range; if (rank != 0 && rank != 1) return; skip_range *skips = xmalloc((size_t)DS4_N_LAYER * 3 * sizeof(*skips)); uint32_t n_skips = 0; uint64_t skip_bytes = 0; for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { const ds4_layer_weights *l = &w->layer[il]; 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 total_bytes = x->dim[2] * expert_bytes; const uint64_t low_bytes = (x->dim[2] / 2) * expert_bytes; /* Unowned range: rank 0 owns the low ids; rank 1 owns the high * ids and takes any odd-count remainder. */ skips[n_skips].off = x->abs_offset + (rank == 0 ? low_bytes : 0); skips[n_skips].len = rank == 0 ? total_bytes - low_bytes : low_bytes; skip_bytes += skips[n_skips].len; n_skips++; } } /* File order should already ascend, but do not rely on it. */ for (uint32_t i = 1; i < n_skips; i++) { skip_range key = skips[i]; uint32_t j = i; while (j > 0 && skips[j - 1].off > key.off) { skips[j] = skips[j - 1]; j--; } skips[j] = key; } const uint64_t page = (uint64_t)sysconf(_SC_PAGESIZE); const uint8_t *p = m->map; const uint64_t start = m->tensor_data_pos, end = m->size; fprintf(stderr, "ds4: warming sharded tensor pages (rank %d): %.2f of %.2f GiB\n", rank, (double)(end - start - skip_bytes) / 1073741824.0, (double)(end - start) / 1073741824.0); const double t0 = now_sec(); volatile uint64_t checksum = 0; uint64_t off = start; uint32_t si = 0; /* skips are naturally sorted: layers ascend and tensors within a layer * are laid out in file order. */ while (off < end) { uint64_t stop = end; while (si < n_skips && skips[si].off + skips[si].len <= off) si++; if (si < n_skips && skips[si].off <= off) { off = skips[si].off + skips[si].len; si++; continue; } if (si < n_skips && skips[si].off < stop) stop = skips[si].off; #if defined(POSIX_MADV_WILLNEED) (void)posix_madvise((void *)(p + off), (size_t)(stop - off), POSIX_MADV_WILLNEED); #endif for (; off < stop; off += page) checksum += p[off]; off = stop; } free(skips); fprintf(stderr, "ds4: sharded warm done in %.1fs (checksum=%llu)\n", now_sec() - t0, (unsigned long long)checksum); } /* ========================================================================= * Wave-2 multi-GPU placement scaffolding: engine placement helpers. * ========================================================================= * * These helpers compute and install the multi-GPU layer placement table. * They are reached only when ds4_engine_create_with_gpu_config is called * with a non-NULL ds4_gpu_config. When the caller passes NULL (every * existing caller — ds4_engine_open shim, ds4_test, ds4_cli, ds4_server, * ds4_bench, ds4_eval, ds4_agent), these helpers are not invoked and the * engine state is byte-equivalent to the pre-multi-GPU CLI main branch. */ /* Classify each model tensor by its placement entry. * * Tensor names live in ds4_str slices (ptr+len), NOT NUL-terminated. * We bound every comparison by name.len to avoid out-of-buffer reads. */ static int tensor_to_entry(const ds4_tensor *t, int n_layer) { const char *p = t->name.ptr; int n = (int)t->name.len; if (n <= 0 || !p) return 0; /* "blk.." prefix -> entry il + 1. */ if (n >= 5 && memcmp(p, "blk.", 4) == 0) { int i = 4; int il = 0; int digits = 0; while (i < n && p[i] >= '0' && p[i] <= '9') { if (digits >= 4) return 0; il = il * 10 + (p[i] - '0'); digits++; i++; } if (digits > 0 && i < n && p[i] == '.' && il >= 0 && il < n_layer) { return il + 1; } return 0; } /* Output-head tensors -> head bucket. weights_bind() requires six * top-level tensors: token_embd (embedding), and five head-tier * tensors output, output_norm, plus three output_hc_* (the * heavily-compressed head's base / fn / scale). All five live on the * output-head tier; missing any of them here would mis-size the * embedding bucket and the head bucket for the packer. */ static const char k_output[] = "output.weight"; static const char k_output_norm[] = "output_norm.weight"; static const char k_output_hc[] = "output_hc_"; if (n == (int)sizeof(k_output) - 1 && memcmp(p, k_output, n) == 0) { return n_layer + 1; } if (n == (int)sizeof(k_output_norm) - 1 && memcmp(p, k_output_norm, n) == 0) { return n_layer + 1; } if (n >= (int)sizeof(k_output_hc) - 1 && memcmp(p, k_output_hc, sizeof(k_output_hc) - 1) == 0) { return n_layer + 1; } /* "mtp." prefix -> head bucket (no current code path loads MTP into * e->model, but harmless). */ if (n >= 4 && memcmp(p, "mtp.", 4) == 0) return n_layer + 1; /* token_embd.weight and everything else falls into entry 0. */ return 0; } static bool engine_deepseek_routed_expert_tensor( const ds4_engine *e, const ds4_tensor *t, int entry, uint64_t *expert_bytes) { if (!e || !t || entry < 1 || entry > (int)DS4_N_LAYER || g_ds4_shape.family != DS4_MODEL_FAMILY_DEEPSEEK4 || DS4_N_EXPERT == 0u || (DS4_N_EXPERT & 1u) != 0u) { return false; } const ds4_layer_weights *layer = &e->weights.layer[entry - 1]; if (t != layer->ffn_gate_exps && t != layer->ffn_up_exps && t != layer->ffn_down_exps) { return false; } if (t->bytes == 0u || t->bytes % DS4_N_EXPERT != 0u) { fprintf(stderr, "ds4: routed expert tensor %.*s is not evenly divisible " "across %u experts\n", (int)t->name.len, t->name.ptr, DS4_N_EXPERT); return false; } if (expert_bytes) *expert_bytes = t->bytes / DS4_N_EXPERT; return true; } static bool engine_cuda_tp_decode_requested(const ds4_engine *e); static bool engine_cuda_tp_ep_requested(const ds4_engine *e); static bool engine_cuda_tp_output_env_requested(void); /* Compute per-entry byte footprint estimates. Walks the tensor table once * and adds each tensor's bytes to its entry bucket. Also adds a per-layer * KV/scratch estimate so the packer's budget math reflects runtime * requirements (not just weight bytes). Returns 0 on success. */ static int engine_compute_entry_bytes(const ds4_engine *e, size_t *out) { const int n_entries = DS4_N_LAYER + 2; const bool cuda_tp_ep = engine_cuda_tp_ep_requested(e); for (int i = 0; i < n_entries; i++) out[i] = 0; for (uint64_t i = 0; i < e->model.n_tensors; i++) { const ds4_tensor *t = &e->model.tensors[i]; if (t->bytes == 0) continue; int entry = tensor_to_entry(t, DS4_N_LAYER); if (entry < 0 || entry >= n_entries) entry = 0; if (cuda_tp_ep && engine_cuda_tp_output_env_requested() && t == e->weights.output) { /* Output TP stores one vocabulary-row slice per participating * tier. Those bytes are reserved per tier after the head tier is * known, rather than charging a full output matrix here. */ continue; } uint64_t expert_bytes = 0; if (cuda_tp_ep && engine_deepseek_routed_expert_tensor(e, t, entry, &expert_bytes)) { out[entry] += expert_bytes * (DS4_N_EXPERT / 2u); } else { out[entry] += t->bytes; } } /* Per-layer KV estimate. The CTX hint is plumbed from the CLI * (--ctx, --ctx-max, etc.) so the packer accounts for the actual * session size the user intends to allocate. Zero/unset falls back to * the legacy 4096 value for back-compat with callers that don't * populate the option (single-tier paths, tests). Per-tier scratch * is accounted separately (see engine_per_tier_graph_overhead_bytes * and its pre-subtract in engine_classify_multi_tier). */ const int est_ctx = (e->placement_ctx_hint > 0) ? e->placement_ctx_hint : 4096; ds4_context_memory mem = ds4_context_memory_estimate_with_prefill(DS4_BACKEND_CUDA, est_ctx, e->prefill_chunk); if (mem.total_bytes > 0) { /* mem.total_bytes is used only as a sentinel; cache values are * re-derived per layer so placement follows the active model's * allocation geometry. */ for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { out[il + 1] += DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA ? engine_glm_per_layer_kv_bytes_planner(il, est_ctx) : engine_per_layer_kv_bytes_planner(il, est_ctx, e->prefill_chunk); } } else { /* Fallback: 128 MiB per layer as a static estimate. */ const size_t fallback_per_layer = (size_t)128ull * 1024ull * 1024ull; for (uint32_t i = 1; i <= DS4_N_LAYER; i++) out[i] += fallback_per_layer; } return 0; } static bool engine_cuda_tp_decode_requested(const ds4_engine *e) { #if defined(__APPLE__) && !defined(DS4_TEST_HOOKS) (void)e; return false; #else return e && e->cuda_tensor_parallel && e->gpu_cfg.n_gpus >= 2 && (e->gpu_cfg.n_gpus & 1) == 0; #endif } static bool engine_cuda_tp_ep_requested(const ds4_engine *e) { #if defined(__APPLE__) && !defined(DS4_TEST_HOOKS) (void)e; return false; #else return engine_cuda_tp_decode_requested(e) && g_ds4_shape.family == DS4_MODEL_FAMILY_DEEPSEEK4 && DS4_N_EXPERT != 0u && (DS4_N_EXPERT & 1u) == 0u; #endif } static bool engine_cuda_tp_output_env_requested(void) { #if defined(__APPLE__) && !defined(DS4_TEST_HOOKS) return false; #else const char *env = getenv("DS4_CUDA_TP_OUTPUT"); return !env || !env[0] || strcmp(env, "0") != 0; #endif } /* Return one output-weight shard for logical_tier. Zero means the tier does * not participate, one means a valid span was returned, and -1 is an invalid * output tensor/configuration. This matches the row partition used by every * output-TP execution path. */ static int engine_cuda_tp_output_shard_span( const ds4_engine *e, int head_tier, int logical_tier, uint64_t *source_offset, uint64_t *bytes) { if (!e || !e->weights.output || !source_offset || !bytes) return -1; const ds4_tensor *output = e->weights.output; if (output->type != DS4_TENSOR_Q8_0 || output->ndim != 2 || output->dim[0] != DS4_N_EMBD || output->dim[1] < 2u || output->bytes == 0u || output->bytes % output->dim[1] != 0u) { return -1; } int tiers[DS4_MAX_GPUS] = {0}; const uint32_t ways = metal_graph_cuda_tp_output_tiers_for_head( head_tier, true, e->gpu_cfg.n_gpus, tiers); if (ways < 2u) return -1; for (uint32_t i = 0; i < ways; i++) { if (tiers[i] != logical_tier) continue; const uint64_t vocab = output->dim[1]; const uint64_t row_bytes = output->bytes / vocab; const uint64_t row0 = (vocab * i) / ways; const uint64_t row1 = (vocab * (i + 1u)) / ways; *source_offset = output->abs_offset + row0 * row_bytes; *bytes = (row1 - row0) * row_bytes; return *bytes != 0u ? 1 : -1; } *source_offset = 0u; *bytes = 0u; return 0; } static int engine_reserve_cuda_ep_output_shards( const ds4_engine *e, ds4_layer_pack_config *pcfg) { if (!e || !pcfg || !engine_cuda_tp_output_env_requested()) return 0; const int n_stages = pcfg->n_gpus / 2; const int head_tier = n_stages - 1; for (int tier = 0; tier < pcfg->n_gpus; tier++) { uint64_t source_offset = 0; uint64_t shard_bytes = 0; const int shard = engine_cuda_tp_output_shard_span( e, head_tier, tier, &source_offset, &shard_bytes); (void)source_offset; if (shard < 0 || (shard > 0 && shard_bytes > pcfg->gpu_budget_bytes[tier])) { fprintf(stderr, "ds4: CUDA EP output shard cannot fit tier %d " "(need %.2f GiB, budget %.2f GiB)\n", tier, (double)shard_bytes / 1073741824.0, (double)pcfg->gpu_budget_bytes[tier] / 1073741824.0); return -1; } if (shard > 0) pcfg->gpu_budget_bytes[tier] -= (size_t)shard_bytes; } return 0; } static void engine_adjust_output_head_for_cuda_tp(ds4_engine *e, const size_t *entry_bytes) { if (!e || !entry_bytes || e->n_placement_entries < (int)DS4_N_LAYER + 2) return; if (!engine_cuda_tp_decode_requested(e) || !engine_cuda_tp_output_env_requested()) { return; } const int half = e->gpu_cfg.n_gpus / 2; const int head_entry = (int)DS4_N_LAYER + 1; const int head_tier = e->placement[head_entry]; if (head_tier >= 0 && head_tier < half) return; size_t used[DS4_LAYER_PACK_MAX_GPUS] = {0}; for (int i = 0; i < e->n_placement_entries; i++) { if (i == head_entry) continue; const int dev = e->placement[i]; if (dev >= 0 && dev < e->gpu_cfg.n_gpus) used[dev] += entry_bytes[i]; } const int layer_tier = e->placement[DS4_N_LAYER]; int preferred = layer_tier; if (preferred >= half) preferred -= half; if (preferred < 0 || preferred >= half) preferred = -1; const size_t head_bytes = entry_bytes[head_entry]; int target = -1; for (int pass = 0; pass < 2 && target < 0; pass++) { const int begin = pass == 0 && preferred >= 0 ? preferred : 0; const int end = pass == 0 && preferred >= 0 ? preferred + 1 : half; for (int tier = begin; tier < end; tier++) { if (pass != 0 && tier == preferred) continue; const size_t budget = e->gpu_cfg.vram_bytes[tier]; if (head_bytes <= budget && used[tier] <= budget - head_bytes) { target = tier; break; } } } if (target >= 0) { fprintf(stderr, "ds4: CUDA output TP moved output head from tier %d to " "lower-half tier %d so the head can use partner tier %d\n", head_tier, target, target + half); e->placement[head_entry] = target; } } static int engine_compute_cuda_ep_placement( const size_t *entry_bytes, int n_entries, const ds4_layer_pack_config *pcfg, int *placement) { if (!entry_bytes || !pcfg || !placement || n_entries != (int)DS4_N_LAYER + 2 || pcfg->n_gpus < 2 || (pcfg->n_gpus & 1) != 0) { return -1; } const int n_stages = pcfg->n_gpus / 2; if (n_stages > (int)DS4_N_LAYER) return -1; uint64_t remaining_layers = 0; for (int entry = 1; entry <= (int)DS4_N_LAYER; entry++) { if (remaining_layers > UINT64_MAX - entry_bytes[entry]) return -1; remaining_layers += entry_bytes[entry]; } placement[0] = 0; int next_entry = 1; for (int stage = 0; stage < n_stages; stage++) { const int stages_left = n_stages - stage; const int layers_left = (int)DS4_N_LAYER - next_entry + 1; const uint64_t fixed_here = (stage == 0 ? entry_bytes[0] : 0u) + (stage == n_stages - 1 ? entry_bytes[n_entries - 1] : 0u); const uint64_t fixed_left = (stage == 0 ? entry_bytes[0] : 0u) + entry_bytes[n_entries - 1]; const uint64_t target_total = (remaining_layers + fixed_left + (uint64_t)stages_left - 1u) / (uint64_t)stages_left; const uint64_t target_layers = target_total > fixed_here ? target_total - fixed_here : 0u; int take = 0; uint64_t layer_bytes = 0; const int max_take = layers_left - (stages_left - 1); const uint64_t home_budget = pcfg->gpu_budget_bytes[stage]; const uint64_t partner_budget = pcfg->gpu_budget_bytes[stage + n_stages]; if (fixed_here > home_budget) { fprintf(stderr, "ds4: CUDA EP fixed weights do not fit stage %d home " "(need %.2f GiB, budget %.2f GiB)\n", stage, (double)fixed_here / 1073741824.0, (double)home_budget / 1073741824.0); return -1; } const uint64_t home_layer_budget = home_budget - fixed_here; const uint64_t layer_budget = home_layer_budget < partner_budget ? home_layer_budget : partner_budget; while (take < max_take) { const uint64_t next = entry_bytes[next_entry + take]; if (layer_bytes > UINT64_MAX - next) return -1; const uint64_t candidate = layer_bytes + next; if (candidate > layer_budget) break; if (take > 0) { const uint64_t old_diff = layer_bytes > target_layers ? layer_bytes - target_layers : target_layers - layer_bytes; const uint64_t new_diff = candidate > target_layers ? candidate - target_layers : target_layers - candidate; if (new_diff > old_diff) break; } layer_bytes = candidate; take++; } if (take == 0 || (stage == n_stages - 1 && take != max_take)) { fprintf(stderr, "ds4: CUDA EP cannot fit balanced stage %d in pair " "budgets %.2f/%.2f GiB (%d layers remain)\n", stage, (double)home_budget / 1073741824.0, (double)partner_budget / 1073741824.0, layers_left); return -1; } for (int i = 0; i < take; i++) placement[next_entry + i] = stage; next_entry += take; remaining_layers -= layer_bytes; } if (next_entry != (int)DS4_N_LAYER + 1) return -1; placement[n_entries - 1] = n_stages - 1; return 0; } /* Phase A: classify multi-tier on a freshly-opened engine (model loaded, * weights bound). Pure CPU — no GPU init required. Sets e->multi_tier, * e->n_placement_entries, e->placement[], and e->gpu_cfg. Returns 0 on * success. NULL config is a no-op. */ static int engine_classify_multi_tier(ds4_engine *e, const ds4_gpu_config *cfg) { if (!cfg) { e->multi_tier = 0; e->n_placement_entries = 0; return 0; } if (cfg->n_gpus <= 0 || cfg->n_gpus > DS4_LAYER_PACK_MAX_GPUS) return -1; /* Caller-bug guard: ds4_gpu_config has no auto-detect; a zero-init * struct with only n_gpus and device_indices populated would * classify every entry as CPU spill and the engine would refuse. * That outcome is never useful, so reject it loudly. Auto-detection * belongs in the CLI layer (CLI flag wiring maps --gpu-vram auto to * cudaMemGetInfo before constructing the config). */ size_t total_budget = 0; for (int d = 0; d < cfg->n_gpus; d++) total_budget += cfg->vram_bytes[d]; if (total_budget == 0) { fprintf(stderr, "ds4: ds4_gpu_config has n_gpus=%d but every vram_bytes[d]==0; " "caller must populate explicit budgets (auto-detect is the CLI's " "job, not the engine's)\n", cfg->n_gpus); return -1; } e->gpu_cfg = *cfg; /* Pre-subtract per-tier Class-P graph scratch from EVERY device * budget BEFORE the packer reads vram_bytes. * Conservative: tiers that end up unused still reserve the overhead, so * the packer cannot accept a layout that would later OOM at * session_create when metal_graph_alloc_raw_cap's per-tier scratch loop * runs. The reservation flows into the packer via e->gpu_cfg (NOT * the caller's cfg) — see the budget loop below. */ const size_t per_tier_overhead = engine_per_tier_graph_overhead_bytes(e); for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { if (e->gpu_cfg.vram_bytes[d] <= per_tier_overhead) { fprintf(stderr, "ds4: GPU%d budget %.2f GiB <= per-tier graph overhead " "%.2f GiB; tier cannot hold its own runtime scratch — " "refusing upfront.\n", e->gpu_cfg.device_indices[d], (double)e->gpu_cfg.vram_bytes[d] / (1024.0 * 1024.0 * 1024.0), (double)per_tier_overhead / (1024.0 * 1024.0 * 1024.0)); return -1; } e->gpu_cfg.vram_bytes[d] -= per_tier_overhead; } size_t entry_bytes[DS4_MAX_LAYER + 2]; if (engine_compute_entry_bytes(e, entry_bytes) != 0) return -1; ds4_layer_pack_config pcfg; memset(&pcfg, 0, sizeof(pcfg)); pcfg.n_gpus = e->gpu_cfg.n_gpus; const size_t cublas_workspace_overhead = (size_t)64ull * 1024ull * 1024ull; for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { /* Read post-subtract budgets from e->gpu_cfg (NOT the caller's * cfg) so the per-tier overhead pre-subtract actually flows into * the packer. */ size_t budget = e->gpu_cfg.vram_bytes[d]; size_t reserve = e->gpu_cfg.safety_margin_bytes + cublas_workspace_overhead; pcfg.gpu_budget_bytes[d] = budget > reserve ? budget - reserve : 0; } const bool cuda_tp_ep = engine_cuda_tp_ep_requested(e); if (cuda_tp_ep && engine_reserve_cuda_ep_output_shards(e, &pcfg) != 0) { return -1; } const int placement_rc = cuda_tp_ep ? engine_compute_cuda_ep_placement(entry_bytes, DS4_N_LAYER + 2, &pcfg, e->placement) : ds4_compute_layer_placement(entry_bytes, DS4_N_LAYER + 2, &pcfg, e->placement); if (placement_rc != 0) { return -1; } e->n_placement_entries = DS4_N_LAYER + 2; engine_adjust_output_head_for_cuda_tp(e, entry_bytes); int first_tier = e->placement[0]; int multi_tier = 0; for (int i = 1; i < e->n_placement_entries; i++) { if (e->placement[i] != first_tier) { multi_tier = 1; break; } } if (!multi_tier) { for (int i = 0; i < e->n_placement_entries; i++) { if (e->placement[i] == DS4_LAYER_PACK_CPU) { multi_tier = 1; break; } } } e->multi_tier = multi_tier; return 0; } #ifndef DS4_NO_GPU static int engine_append_device_cache_span( ds4_tensor_range **per_dev_ranges, int *per_dev_n, int *per_dev_cap, int logical_tier, int physical_device, uint64_t source_offset, uint64_t bytes) { if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS || bytes == 0u) return -1; if (per_dev_n[logical_tier] >= per_dev_cap[logical_tier]) { int new_cap = per_dev_cap[logical_tier] == 0 ? 64 : per_dev_cap[logical_tier] * 2; ds4_tensor_range *grow = realloc(per_dev_ranges[logical_tier], (size_t)new_cap * sizeof(*grow)); if (!grow) { fprintf(stderr, "ds4: out of memory growing per-device range list\n"); return -1; } per_dev_ranges[logical_tier] = grow; per_dev_cap[logical_tier] = new_cap; } per_dev_ranges[logical_tier][per_dev_n[logical_tier]].source_offset = source_offset; per_dev_ranges[logical_tier][per_dev_n[logical_tier]].bytes = bytes; per_dev_ranges[logical_tier][per_dev_n[logical_tier]].target_device = physical_device; per_dev_n[logical_tier]++; return 0; } static int engine_append_device_cache_range( ds4_tensor_range **per_dev_ranges, int *per_dev_n, int *per_dev_cap, int logical_tier, int physical_device, const ds4_tensor *t) { if (!t) return -1; return engine_append_device_cache_span(per_dev_ranges, per_dev_n, per_dev_cap, logical_tier, physical_device, t->abs_offset, t->bytes); } /* Install per-device selective caches for GPU-placed tensors. Skips any * tensor whose entry is placed on CPU — that case must be rejected at a * higher level for this PR (execution wiring not yet shipped). */ static int engine_install_per_device_caches(ds4_engine *e) { /* Prereq: register the host model map so ds4_gpu_device_cache_tensors * can resolve g_model_host_base. We use the no-copy variant so * DS4_CUDA_COPY_MODEL cannot reintroduce a full-model copy that would * defeat the per-device selective cache. */ if (!ds4_gpu_register_model_map_no_copy(e->model.map, e->model.size)) return -1; /* Per-logical-tier dynamic range lists. */ ds4_tensor_range *per_dev_ranges[DS4_MAX_GPUS] = {0}; int per_dev_n[DS4_MAX_GPUS] = {0}; int per_dev_cap[DS4_MAX_GPUS] = {0}; int rc = -1; const bool cuda_tp_decode = engine_cuda_tp_decode_requested(e); const bool cuda_tp_ep = engine_cuda_tp_ep_requested(e); const int tp_half = cuda_tp_decode ? e->gpu_cfg.n_gpus / 2 : 0; const bool cuda_tp_output = cuda_tp_decode && metal_graph_cuda_tp_output_requested(); int output_tp_tiers[DS4_MAX_GPUS] = {0}; const int output_head_entry = (int)DS4_N_LAYER + 1; const uint32_t output_tp_ways = cuda_tp_output ? metal_graph_cuda_tp_output_tiers_for_head(e->placement[output_head_entry], true, e->gpu_cfg.n_gpus, output_tp_tiers) : 0; if (cuda_tp_ep && !metal_graph_cuda_tp_moe_requested()) { fprintf(stderr, "ds4: CUDA tensor parallelism requires routed MoE TP, but it is disabled\n"); goto cleanup; } if (cuda_tp_ep && !metal_graph_cuda_tp_prefill_ffn_requested()) { fprintf(stderr, "ds4: CUDA tensor parallelism requires owned routed-MoE prefill, but it is disabled\n"); goto cleanup; } if (cuda_tp_ep) { fprintf(stderr, "ds4: CUDA decode TP half-resident routed experts enabled " "([0,%u) home, [%u,%u) partner); other layer tensors remain replicated\n", DS4_N_EXPERT / 2u, DS4_N_EXPERT / 2u, DS4_N_EXPERT); } else if (cuda_tp_decode) { fprintf(stderr, "ds4: CUDA decode TP cache duplication enabled for layer tensors " "(tier N -> tier N+%d)\n", tp_half); } if (cuda_tp_output) { if (output_tp_ways < 2u) { fprintf(stderr, "ds4: DS4_CUDA_TP_OUTPUT requires output-head home in " "lower-half tiers; output head is on tier %d\n", e->placement[output_head_entry]); goto cleanup; } fprintf(stderr, cuda_tp_ep ? "ds4: CUDA output TP stores one vocabulary shard per tier (%u ways)\n" : "ds4: CUDA output TP cache duplication enabled for output head (%u ways)\n", output_tp_ways); } for (uint64_t i = 0; i < e->model.n_tensors; i++) { const ds4_tensor *t = &e->model.tensors[i]; if (t->bytes == 0) continue; int entry = tensor_to_entry(t, DS4_N_LAYER); if (entry < 0 || entry >= e->n_placement_entries) entry = 0; int logical_tier = e->placement[entry]; if (logical_tier == DS4_LAYER_PACK_CPU) continue; /* CPU spill: skip here. */ if (logical_tier < 0 || logical_tier >= e->gpu_cfg.n_gpus) { fprintf(stderr, "ds4: placement tier %d out of range for tensor %.*s\n", logical_tier, (int)t->name.len, t->name.ptr ? t->name.ptr : ""); goto cleanup; } const int physical_device = g_gpu[logical_tier].device_id; if (cuda_tp_ep && cuda_tp_output && entry == output_head_entry && t == e->weights.output) { for (int tier = 0; tier < e->gpu_cfg.n_gpus; tier++) { uint64_t shard_offset = 0; uint64_t shard_bytes = 0; const int shard = engine_cuda_tp_output_shard_span( e, logical_tier, tier, &shard_offset, &shard_bytes); if (shard < 0) { fprintf(stderr, "ds4: invalid CUDA output TP shard for tier %d\n", tier); goto cleanup; } if (shard == 0) continue; if (engine_append_device_cache_span( per_dev_ranges, per_dev_n, per_dev_cap, tier, g_gpu[tier].device_id, shard_offset, shard_bytes) != 0) { goto cleanup; } } continue; } uint64_t expert_bytes = 0; const bool shard_experts = cuda_tp_ep && engine_deepseek_routed_expert_tensor( e, t, entry, &expert_bytes); const uint64_t home_offset = t->abs_offset; const uint64_t home_bytes = shard_experts ? expert_bytes * (DS4_N_EXPERT / 2u) : t->bytes; if (engine_append_device_cache_span(per_dev_ranges, per_dev_n, per_dev_cap, logical_tier, physical_device, home_offset, home_bytes) != 0) { goto cleanup; } if (cuda_tp_decode && entry >= 1 && entry <= (int)DS4_N_LAYER) { if (logical_tier >= tp_half) { fprintf(stderr, "ds4: CUDA tensor parallelism only supports layer homes in " "lower-half tiers; tensor %.*s is on tier %d\n", (int)t->name.len, t->name.ptr ? t->name.ptr : "", logical_tier); goto cleanup; } const int partner_tier = logical_tier + tp_half; const int partner_device = g_gpu[partner_tier].device_id; const uint64_t partner_offset = shard_experts ? t->abs_offset + home_bytes : t->abs_offset; const uint64_t partner_bytes = shard_experts ? t->bytes - home_bytes : t->bytes; if (engine_append_device_cache_span(per_dev_ranges, per_dev_n, per_dev_cap, partner_tier, partner_device, partner_offset, partner_bytes) != 0) { goto cleanup; } } if (cuda_tp_output && !cuda_tp_ep && entry == (int)DS4_N_LAYER + 1) { for (uint32_t j = 1; j < output_tp_ways; j++) { const int peer_tier = output_tp_tiers[j]; if (peer_tier == logical_tier) continue; if (peer_tier < 0 || peer_tier >= e->gpu_cfg.n_gpus) { fprintf(stderr, "ds4: DS4_CUDA_TP_OUTPUT selected invalid tier %d " "for tensor %.*s\n", peer_tier, (int)t->name.len, t->name.ptr ? t->name.ptr : ""); goto cleanup; } const int peer_device = g_gpu[peer_tier].device_id; if (engine_append_device_cache_range(per_dev_ranges, per_dev_n, per_dev_cap, peer_tier, peer_device, t) != 0) { goto cleanup; } } } } for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { if (per_dev_n[d] == 0) continue; const int physical_device = g_gpu[d].device_id; uint64_t cache_bytes = 0; for (int i = 0; i < per_dev_n[d]; i++) { if (cache_bytes > UINT64_MAX - per_dev_ranges[d][i].bytes) { fprintf(stderr, "ds4: CUDA cache byte accounting overflow\n"); goto cleanup; } cache_bytes += per_dev_ranges[d][i].bytes; } fprintf(stderr, "ds4: CUDA tier %d (device %d) selective weights: %.2f GiB in %d ranges\n", d, physical_device, (double)cache_bytes / 1073741824.0, per_dev_n[d]); int cache_rc = ds4_gpu_device_cache_tensors(physical_device, per_dev_ranges[d], per_dev_n[d]); if (cache_rc != 0) { fprintf(stderr, "ds4: ds4_gpu_device_cache_tensors failed for tier %d " "(physical device %d, %d ranges) rc=%d\n", d, physical_device, per_dev_n[d], cache_rc); goto cleanup; } } rc = 0; cleanup: for (int d = 0; d < DS4_MAX_GPUS; d++) free(per_dev_ranges[d]); return rc; } /* Pretty-print the layout via the multi-GPU plumbing helper, then emit a peer-access * summary by walking g_gpu_peer_ok[][]. Always called in multi-tier * mode so the operator sees what the packer decided. */ static void engine_print_layout(const ds4_engine *e) { size_t entry_bytes[DS4_MAX_LAYER + 2]; (void)engine_compute_entry_bytes(e, entry_bytes); size_t used[DS4_LAYER_PACK_MAX_GPUS] = {0}; size_t budget[DS4_LAYER_PACK_MAX_GPUS] = {0}; for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { budget[d] = e->gpu_cfg.vram_bytes[d]; } for (int i = 0; i < e->n_placement_entries; i++) { int dev = e->placement[i]; if (dev >= 0 && dev < e->gpu_cfg.n_gpus) used[dev] += entry_bytes[i]; } ds4_layer_pack_print(stderr, e->placement, e->n_placement_entries, DS4_N_LAYER, entry_bytes, used, budget, e->gpu_cfg.n_gpus); /* Show the per-tier graph scratch reservation so operators can * correlate "47 GiB free" with the smaller post-subtract budget the * packer actually had to spend. */ const size_t per_tier_overhead = engine_per_tier_graph_overhead_bytes(e); fprintf(stderr, "ds4: per-tier graph scratch reserved: %.2f GiB " "(pre-subtracted from each GPU budget)\n", (double)per_tier_overhead / (1024.0 * 1024.0 * 1024.0)); fprintf(stderr, "ds4: peer access matrix (validated):"); int peer_any = 0; for (int i = 0; i < g_n_gpus; i++) { for (int j = 0; j < g_n_gpus; j++) { if (i == j) continue; peer_any = 1; const char *mode = g_gpu_peer_ok[i][j] ? "DIRECT" : "BOUNCE"; fprintf(stderr, " %d->%d %s", i, j, mode); } } if (!peer_any) fprintf(stderr, " (single-device)"); fputc('\n', stderr); } /* Install GPU-side placement state. * * Ordering: * 1. Print layout FIRST so it is always visible to the operator, * even if subsequent steps fail. * 2. Detect any CPU-spill placement; this PR refuses to open * multi-tier engines with CPU spill because the execution-side * wiring (CPU↔GPU boundary materialization) lands in the * follow-up multi-GPU execution task. * 3. Register host model map and install per-device selective * caches. * * Returns 0 on success. */ /* Install the DSpark support model's tensors on one executor tier: the * TP partner of the output-head tier when decode TP is active (it has the * output-head copy from output TP and plenty of free VRAM), otherwise the * head tier itself. The strict multi-tier weight cache is offset-keyed, so * support tensors are registered with a disjoint offset bias. */ static int engine_install_dspark_support_cache(ds4_engine *e) { if (!e->multi_tier || e->support_kind != DS4_SUPPORT_DSPARK) return 0; if (!e->dspark) return 0; if (!e->mtp_model.map || e->mtp_model.n_tensors == 0) return 0; int exec_tier = e->placement[DS4_N_LAYER + 1]; if (exec_tier < 0 || exec_tier >= e->gpu_cfg.n_gpus) exec_tier = 0; const bool tp_decode = e->cuda_tensor_parallel; const int tp_half = e->gpu_cfg.n_gpus / 2; if (tp_decode && e->gpu_cfg.n_gpus >= 2 && exec_tier < tp_half) { exec_tier += tp_half; } if (tp_decode && e->gpu_cfg.n_gpus >= 2) { /* Prefer the partner tier with the most free VRAM so the support * weights stay local to the executor. */ uint64_t best_free = ds4_gpu_tier_free_vram(exec_tier); for (int t = tp_half; t < e->gpu_cfg.n_gpus; t++) { const uint64_t f = ds4_gpu_tier_free_vram(t); if (f > best_free) { best_free = f; exec_tier = t; } } } const char *tier_env = getenv("DS4_DSPARK_EXEC_TIER"); if (tier_env && tier_env[0]) { const int v = atoi(tier_env); if (v >= 0 && v < e->gpu_cfg.n_gpus) exec_tier = v; } e->dspark_exec_tier = exec_tier; const uint64_t bias = (e->model.size + 4095ull) & ~4095ull; if (!ds4_gpu_register_support_map(e->mtp_model.map, e->mtp_model.size, bias)) { fprintf(stderr, "ds4: failed to register DSpark support model map\n"); return -1; } /* Greedy pack: exec tier first, then the other TP-partner tiers. * Entries always claim the executor device; spilled tensors are read * through peer access. Per-tier budget leaves room for the graph * scratch and allocator slack that session_create allocates later. */ int order[DS4_MAX_GPUS]; int n_order = 0; order[n_order++] = exec_tier; for (int t = e->gpu_cfg.n_gpus - 1; t >= 0; t--) { if (t == exec_tier) continue; if (tp_decode && t < tp_half) continue; /* home tiers are packed full */ order[n_order++] = t; } uint64_t reserve = 4ull * 1024ull * 1024ull * 1024ull + (1ull << 29); { const char *renv = getenv("DS4_DSPARK_CACHE_RESERVE_GB"); if (renv && renv[0]) { const int gv = atoi(renv); if (gv >= 1 && gv <= 32) reserve = (uint64_t)gv << 30; } } const uint64_t range_cap = e->mtp_model.n_tensors > e->model.n_tensors ? e->mtp_model.n_tensors : e->model.n_tensors; ds4_tensor_range *ranges = xmalloc((size_t)range_cap * sizeof(ranges[0])); bool *placed = xmalloc((size_t)e->mtp_model.n_tensors * sizeof(placed[0])); memset(placed, 0, (size_t)e->mtp_model.n_tensors * sizeof(placed[0])); uint64_t remaining = 0; for (uint64_t ti = 0; ti < e->mtp_model.n_tensors; ti++) { if (e->mtp_model.tensors[ti].bytes != 0) remaining++; } for (int oi = 0; oi < n_order && remaining != 0; oi++) { const int tier = order[oi]; const uint64_t free_b = ds4_gpu_tier_free_vram(tier); /* The executor tier later hosts the per-tier graph scratch, the * dspark capture buffers, and cuBLAS workspace; spill tiers only * host the graph scratch. */ const uint64_t tier_reserve = tier == exec_tier ? reserve : (reserve > (1ull << 30) ? reserve - (1ull << 30) : reserve); uint64_t budget = free_b > tier_reserve ? free_b - tier_reserve : 0; fprintf(stderr, "ds4: DSpark support pack tier=%d free=%.2f GiB budget=%.2f GiB\n", tier, (double)free_b / 1073741824.0, (double)budget / 1073741824.0); if (budget == 0) continue; int n = 0; for (uint64_t ti = 0; ti < e->mtp_model.n_tensors; ti++) { const ds4_tensor *t = &e->mtp_model.tensors[ti]; if (placed[ti] || t->bytes == 0) continue; if (t->bytes > budget) continue; ranges[n].source_offset = t->abs_offset; ranges[n].bytes = t->bytes; ranges[n].target_device = g_gpu[tier].device_id; n++; placed[ti] = true; budget -= t->bytes; remaining--; } if (n == 0) continue; const int rc = ds4_gpu_device_cache_support_tensors( g_gpu[tier].device_id, g_gpu[exec_tier].device_id, ranges, n, 0); if (rc != 0) { fprintf(stderr, "ds4: DSpark support cache install failed on tier %d (rc=%d)\n", tier, rc); free(placed); free(ranges); return -1; } fprintf(stderr, "ds4: DSpark support tensors cached on tier %d (%d tensors%s)\n", tier, n, tier == exec_tier ? ", executor tier" : ", peer spill"); } /* The DSpark draft block embeds tokens through the BASE model's * embedding tensors, which normally live only on the embedding tier. * Install the embedding bucket (placement entry 0) on the executor * tier as well, through the normal main-model cache API. */ { int n = 0; for (uint64_t ti = 0; ti < e->model.n_tensors; ti++) { const ds4_tensor *t = &e->model.tensors[ti]; if (t->bytes == 0) continue; if (tensor_to_entry(t, DS4_N_LAYER) != 0) continue; ranges[n].source_offset = t->abs_offset; ranges[n].bytes = t->bytes; ranges[n].target_device = g_gpu[exec_tier].device_id; n++; } if (n != 0) { const int rc = ds4_gpu_device_cache_support_tensors( g_gpu[exec_tier].device_id, g_gpu[exec_tier].device_id, ranges, n, 1); if (rc != 0) { fprintf(stderr, "ds4: DSpark base embedding cache install failed on " "tier %d (rc=%d)\n", exec_tier, rc); free(placed); free(ranges); return -1; } fprintf(stderr, "ds4: DSpark base embedding bucket cached on tier %d " "(%d tensors)\n", exec_tier, n); } } free(placed); free(ranges); if (remaining != 0) { fprintf(stderr, "ds4: DSpark support cache could not place %llu tensors " "(insufficient free VRAM); disabling DSpark runtime\n", (unsigned long long)remaining); return -1; } return 0; } static int engine_install_gpu_placement(ds4_engine *e) { if (!e->multi_tier) return 0; engine_print_layout(e); int has_cpu_spill = 0; for (int i = 0; i < e->n_placement_entries; i++) { if (e->placement[i] == DS4_LAYER_PACK_CPU) { has_cpu_spill = 1; break; } } if (has_cpu_spill) { fprintf(stderr, "ds4: CPU-spill placement detected; CPU-tier execution wiring lands in\n" "ds4: cpu-spill execution (follow-up) (CPU-spill execution). Aborting engine creation.\n"); return -1; } if (engine_install_per_device_caches(e) != 0) return -1; return 0; } #endif /* !DS4_NO_GPU */ #ifdef DS4_TEST_HOOKS typedef struct { const char *name; uint64_t bytes; } ds4_test_fake_tensor; int ds4_test_tensor_to_entry(const char *name, int name_len) { ds4_tensor fake; memset(&fake, 0, sizeof(fake)); fake.name.ptr = name; fake.name.len = (uint64_t)(name_len > 0 ? name_len : 0); return tensor_to_entry(&fake, DS4_N_LAYER); } static int ds4_test_make_engine( ds4_engine *eng, const ds4_test_fake_tensor *tensors, int n_tensors, int placement_ctx_hint) { memset(eng, 0, sizeof(*eng)); eng->model.fd = -1; eng->placement_ctx_hint = placement_ctx_hint; eng->model.n_tensors = (uint64_t)(n_tensors > 0 ? n_tensors : 0); if (n_tensors <= 0) return 0; eng->model.tensors = calloc((size_t)n_tensors, sizeof(*eng->model.tensors)); if (!eng->model.tensors) return -1; for (int i = 0; i < n_tensors; i++) { eng->model.tensors[i].name.ptr = tensors[i].name; eng->model.tensors[i].name.len = tensors[i].name ? (uint64_t)strlen(tensors[i].name) : 0; eng->model.tensors[i].bytes = tensors[i].bytes; if (tensors[i].name && !strcmp(tensors[i].name, "output.weight")) { eng->model.tensors[i].type = DS4_TENSOR_Q8_0; eng->model.tensors[i].ndim = 2; eng->model.tensors[i].dim[0] = DS4_N_EMBD; eng->model.tensors[i].dim[1] = DS4_N_VOCAB; eng->weights.output = &eng->model.tensors[i]; } else if (tensors[i].name) { int il = -1; char suffix = '\0'; if (sscanf(tensors[i].name, "blk.%d.ffn_gate_exps.weight%c", &il, &suffix) == 1 && il >= 0 && il < (int)DS4_N_LAYER) { eng->weights.layer[il].ffn_gate_exps = &eng->model.tensors[i]; } } } return 0; } static int ds4_test_classify_multi_tier_impl( const ds4_test_fake_tensor *tensors, int n_tensors, const ds4_gpu_config *cfg, int placement_ctx_hint, bool cuda_tensor_parallel, int placement_out[DS4_MAX_LAYER + 2], int *out_multi_tier, int *out_n_entries) { ds4_engine eng; if (ds4_test_make_engine(&eng, tensors, n_tensors, placement_ctx_hint) != 0) { return -1; } eng.cuda_tensor_parallel = cuda_tensor_parallel; eng.prefill_chunk = ds4_effective_prefill_chunk(cuda_tensor_parallel, 0); const int rc = engine_classify_multi_tier(&eng, cfg); if (rc == 0) { if (out_multi_tier) *out_multi_tier = eng.multi_tier; if (out_n_entries) *out_n_entries = eng.n_placement_entries; if (placement_out) { for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER + 2u; i++) { placement_out[i] = eng.placement[i]; } } } free(eng.model.tensors); return rc; } int ds4_test_classify_multi_tier( const ds4_test_fake_tensor *tensors, int n_tensors, const ds4_gpu_config *cfg, int placement_out[DS4_MAX_LAYER + 2], int *out_multi_tier, int *out_n_entries) { return ds4_test_classify_multi_tier_impl( tensors, n_tensors, cfg, 0, false, placement_out, out_multi_tier, out_n_entries); } int ds4_test_classify_multi_tier_with_ctx( const ds4_test_fake_tensor *tensors, int n_tensors, const ds4_gpu_config *cfg, int placement_ctx_hint, int placement_out[DS4_MAX_LAYER + 2], int *out_multi_tier, int *out_n_entries) { return ds4_test_classify_multi_tier_impl( tensors, n_tensors, cfg, placement_ctx_hint, false, placement_out, out_multi_tier, out_n_entries); } int ds4_test_classify_multi_tier_with_ctx_cuda_tp( const ds4_test_fake_tensor *tensors, int n_tensors, const ds4_gpu_config *cfg, int placement_ctx_hint, int placement_out[DS4_MAX_LAYER + 2], int *out_multi_tier, int *out_n_entries) { return ds4_test_classify_multi_tier_impl( tensors, n_tensors, cfg, placement_ctx_hint, true, placement_out, out_multi_tier, out_n_entries); } void ds4_test_seed_compress_ratios(void) { for (uint32_t il = 0; il < DS4_N_LAYER; il++) { g_ds4_compress_ratios[il] = ds4_expected_layer_compress_ratio(il); } } void ds4_test_clear_compress_ratios(void) { memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); } uint32_t ds4_test_effective_prefill_chunk(bool cuda_tensor_parallel, uint32_t requested_chunk) { return ds4_effective_prefill_chunk(cuda_tensor_parallel, requested_chunk); } uint32_t ds4_test_planner_prefill_cap(int prompt_len, uint32_t prefill_chunk) { return engine_planner_prefill_cap(prompt_len, prefill_chunk); } uint32_t ds4_test_planner_raw_cap(int ctx_size, uint32_t prefill_cap) { return engine_planner_raw_cap(ctx_size, prefill_cap); } size_t ds4_test_per_tier_graph_overhead_bytes_with_prefill( int placement_ctx_hint, uint32_t prefill_chunk) { ds4_engine eng; memset(&eng, 0, sizeof(eng)); eng.model.fd = -1; eng.placement_ctx_hint = placement_ctx_hint; eng.prefill_chunk = prefill_chunk; return engine_per_tier_graph_overhead_bytes(&eng); } size_t ds4_test_per_tier_graph_overhead_bytes(int placement_ctx_hint) { return ds4_test_per_tier_graph_overhead_bytes_with_prefill( placement_ctx_hint, 0); } size_t ds4_test_compute_entry_bytes_sum_with_prefill( const ds4_test_fake_tensor *tensors, int n_tensors, int placement_ctx_hint, uint32_t prefill_chunk) { ds4_engine eng; if (ds4_test_make_engine(&eng, tensors, n_tensors, placement_ctx_hint) != 0) { return 0; } eng.prefill_chunk = prefill_chunk; size_t entry_bytes[DS4_MAX_LAYER + 2]; size_t sum = 0; if (engine_compute_entry_bytes(&eng, entry_bytes) == 0) { for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER + 2u; i++) { sum += entry_bytes[i]; } } free(eng.model.tensors); return sum; } size_t ds4_test_compute_entry_bytes_sum( const ds4_test_fake_tensor *tensors, int n_tensors, int placement_ctx_hint) { return ds4_test_compute_entry_bytes_sum_with_prefill( tensors, n_tensors, placement_ctx_hint, 0); } size_t ds4_test_glm_per_layer_kv_bytes(uint32_t layer, int ctx_size) { const ds4_shape saved_shape = g_ds4_shape; g_ds4_shape = DS4_SHAPE_GLM52; const size_t bytes = engine_glm_per_layer_kv_bytes_planner(layer, ctx_size); g_ds4_shape = saved_shape; return bytes; } int ds4_test_session_read_logits(ds4_session *s, float *out, uint64_t out_bytes) { if (!s || !out || out_bytes < (uint64_t)DS4_N_VOCAB * sizeof(float)) { return 1; } return ds4_session_copy_logits(s, out, (int)DS4_N_VOCAB) == (int)DS4_N_VOCAB ? 0 : 1; } const int *ds4_test_engine_placement(const ds4_engine *e) { return e ? e->placement : NULL; } #endif /* DS4_TEST_HOOKS */ static int engine_install_dspark_support_cache(ds4_engine *e); static int engine_install_gpu_placement(ds4_engine *e); static int ds4_engine_open_internal(ds4_engine **out, const ds4_engine_options *opt, const ds4_gpu_config *gpu_cfg); int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { return ds4_engine_open_internal(out, opt, NULL); } int ds4_engine_create_with_gpu_config(ds4_engine **out, const ds4_engine_options *opt, const struct ds4_gpu_config *gpu_cfg) { return ds4_engine_open_internal(out, opt, gpu_cfg); } static int ds4_engine_open_internal(ds4_engine **out, const ds4_engine_options *opt, const ds4_gpu_config *gpu_cfg) { ds4_engine *e = xcalloc(1, sizeof(*e)); e->model.fd = -1; e->mtp_model.fd = -1; e->backend = opt->backend; e->quality = opt->quality; e->glm_mtp = opt->glm_mtp; e->glm_mtp_timing = opt->glm_mtp_timing; e->dspark = opt->dspark; e->dspark_strict = opt->dspark_strict; e->cuda_tensor_parallel = opt->cuda_tensor_parallel; e->glm_tp_token_prefill = opt->tp.glm_token_prefill; e->ssd_streaming = opt->ssd_streaming; e->ssd_streaming_cold = opt->ssd_streaming_cold; e->ssd_streaming_full_layers_set = opt->ssd_streaming_full_layers_set; e->distributed = opt->distributed; e->power_percent = opt->power_percent > 0 ? opt->power_percent : 100; e->prefill_chunk = ds4_effective_prefill_chunk(opt->cuda_tensor_parallel, opt->prefill_chunk); e->ssd_streaming_cache_experts = opt->ssd_streaming_cache_experts; e->ssd_streaming_cache_bytes = opt->ssd_streaming_cache_bytes; e->ssd_streaming_full_layers = opt->ssd_streaming_full_layers; e->ssd_streaming_preload_experts = opt->ssd_streaming_preload_experts; if (e->power_percent > 100) e->power_percent = 100; e->mtp_draft_tokens = opt->mtp_draft_tokens > 0 ? opt->mtp_draft_tokens : 1; if (e->mtp_draft_tokens > 16) e->mtp_draft_tokens = 16; e->mtp_margin = opt->mtp_margin >= 0.0f ? opt->mtp_margin : 3.0f; if (opt->dspark_confidence_threshold_set) { e->dspark_confidence_threshold = opt->dspark_confidence_threshold; } else { e->dspark_confidence_threshold = 0.9f; } if (opt->cuda_tensor_parallel && (opt->backend != DS4_BACKEND_CUDA || !gpu_cfg || gpu_cfg->n_gpus < 2 || (gpu_cfg->n_gpus & 1) != 0)) { fprintf(stderr, "ds4: --cuda-tensor-parallel requires an even multi-GPU CUDA placement\n"); free(e); *out = NULL; return 1; } if (opt->dspark && (!opt->mtp_path || !opt->mtp_path[0])) { fprintf(stderr, "ds4: --dspark requires --mtp FILE\n"); free(e); *out = NULL; return 1; } if ((opt->directional_steering_attn != 0.0f || opt->directional_steering_ffn != 0.0f) && (!opt->directional_steering_file || !opt->directional_steering_file[0])) { fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); free(e); *out = NULL; return 1; } if (opt->directional_steering_file && opt->directional_steering_file[0]) { e->directional_steering_file = ds4_strdup(opt->directional_steering_file); e->directional_steering_attn_scale = opt->directional_steering_attn; e->directional_steering_ffn_scale = opt->directional_steering_ffn; } if (opt->n_threads > 0) g_requested_threads = (uint32_t)opt->n_threads; e->placement_ctx_hint = opt->placement_ctx_hint; e->share_session_prefill_workspace = opt->share_session_prefill_workspace; ds4_acquire_instance_lock(); if (opt->simulate_used_memory_bytes != 0 && !ds4_ssd_memory_lock_acquire(&e->simulated_memory, opt->simulate_used_memory_bytes)) { ds4_engine_close(e); *out = NULL; return 1; } bool load_slice = opt->load_slice; uint32_t load_layer_start = opt->load_layer_start; uint32_t load_layer_end = opt->load_layer_end; bool load_output = opt->load_output; if (opt->distributed.role != DS4_DISTRIBUTED_NONE && opt->distributed.layers.set) { load_slice = true; load_layer_start = opt->distributed.layers.start; load_layer_end = (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA && opt->distributed.layers.has_output) ? UINT32_MAX : opt->distributed.layers.end; load_output = opt->distributed.layers.has_output; } const bool graph_backend = ds4_backend_uses_graph(opt->backend); if (graph_backend) ds4_linux_graph_backend_set_oom_score(opt->backend); 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); if (e->cuda_tensor_parallel && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { fprintf(stderr, "ds4: --cuda-tensor-parallel is currently supported only for DeepSeek models\n"); ds4_engine_close(e); *out = NULL; return 1; } if (e->ssd_streaming && !ds4_backend_supports_ssd_streaming(e->backend)) { fprintf(stderr, "ds4: --ssd-streaming is currently supported only with --metal/--cuda/--rocm\n"); ds4_engine_close(e); *out = NULL; return 1; } const char *expert_profile_path = opt->expert_profile_path; if (!expert_profile_path || !expert_profile_path[0]) { expert_profile_path = getenv("DS4_EXPERT_PROFILE"); } const char *expert_hotlist_path = getenv("DS4_EXPERT_HOTLIST"); if ((expert_profile_path && expert_profile_path[0]) || (expert_hotlist_path && expert_hotlist_path[0])) { if (e->backend == DS4_BACKEND_METAL) { ds4_expert_profile_init(expert_profile_path, expert_hotlist_path); } else { fprintf(stderr, "ds4: expert profile/hotlist is Metal-only for now; ignoring for %s backend\n", ds4_backend_name(e->backend)); } } weights_bind(&e->weights, &e->model, load_slice, load_layer_start, load_layer_end, load_output); /* TP always maps one contiguous routed-expert half per rank. Decide * immediately after binding so memory guards account only the bytes this * rank owns (replicated dense weights plus its expert shard). */ #ifndef DS4_NO_GPU const bool tp_shard = opt->tp.role != DS4_TP_NONE && !e->ssd_streaming; const int tp_shard_rank = opt->tp.role == DS4_TP_WORKER ? 1 : 0; g_tp_shard_model_bytes = 0; if (tp_shard) { ds4_model_map_span_vec shard_spans; if (weights_model_map_sharded_spans(&e->weights, &e->model, tp_shard_rank, &shard_spans)) { g_tp_shard_model_bytes = model_map_span_vec_total_bytes(&shard_spans); free(shard_spans.v); } } #endif if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (opt->inspect_only) { *out = e; return 0; } if ((opt->directional_steering_file && opt->directional_steering_file[0]) || opt->directional_steering_attn != 0.0f || opt->directional_steering_ffn != 0.0f) { fprintf(stderr, "ds4: directional steering is not supported for GLM 5.2 yet\n"); ds4_engine_close(e); *out = NULL; return 1; } if (e->power_percent < 100) { fprintf(stderr, "ds4: --power is not supported for GLM 5.2 yet\n"); ds4_engine_close(e); *out = NULL; return 1; } if (opt->prefill_chunk != 0) { fprintf(stderr, "ds4: --prefill-chunk is not supported for GLM 5.2; " "GLM uses graph-selected prefill chunks\n"); ds4_engine_close(e); *out = NULL; return 1; } if (opt->mtp_path && opt->mtp_path[0]) { fprintf(stderr, "ds4: --mtp is not supported for GLM 5.2 yet\n"); ds4_engine_close(e); *out = NULL; return 1; } if (opt->first_token_test) { if (e->backend != DS4_BACKEND_CPU) { fprintf(stderr, "ds4: GLM first-token test is CPU-only; pass --cpu\n"); ds4_engine_close(e); *out = NULL; return 1; } vocab_load(&e->vocab, &e->model); *out = e; return 0; } bool glm_backend_supported = ds4_backend_uses_graph(e->backend); #ifdef DS4_ROCM_BUILD if (e->backend == DS4_BACKEND_CUDA && !e->ssd_streaming) { glm_backend_supported = false; } #endif if (!glm_backend_supported) { #ifdef DS4_ROCM_BUILD fprintf(stderr, "ds4: GLM 5.2 ROCm inference requires --ssd-streaming; " "use --inspect or --cpu --first-token-test for CPU diagnostics\n"); #else fprintf(stderr, "ds4: GLM 5.2 inference requires the Metal or CUDA graph " "backend; use --inspect or --cpu --first-token-test for CPU diagnostics\n"); #endif ds4_engine_close(e); *out = NULL; return 1; } #ifndef DS4_NO_GPU if (opt->context_size > 0) { uint32_t guard_ctx = 0; if (!glm_graph_context_request(opt->context_size, &guard_ctx) || !(load_slice ? glm_graph_memory_guard_slice(&e->model, &e->weights, e->ssd_streaming, load_layer_start, load_layer_end, load_layer_start == 0, load_output, guard_ctx) : glm_graph_memory_guard(&e->model, &e->weights, e->ssd_streaming, guard_ctx))) { ds4_engine_close(e); *out = NULL; return 1; } } #endif vocab_load(&e->vocab, &e->model); } else if (!opt->inspect_only) { vocab_load(&e->vocab, &e->model); } if (e->ssd_streaming && e->ssd_streaming_cache_bytes != 0) { const uint64_t requested_cache_bytes = e->ssd_streaming_cache_bytes; const uint64_t safe_cache_bytes = ds4_streaming_manual_cache_safe_bytes(e->backend, opt->context_size, e->prefill_chunk, e->ssd_streaming); if (safe_cache_bytes != 0 && e->ssd_streaming_cache_bytes > safe_cache_bytes) { e->ssd_streaming_cache_bytes = safe_cache_bytes; fprintf(stderr, "ds4: %s SSD streaming cache budget %.2f GiB capped to %.2f GiB " "to stay below the graph working-set pressure budget\n", ds4_backend_name(e->backend), (double)requested_cache_bytes / 1073741824.0, (double)e->ssd_streaming_cache_bytes / 1073741824.0); } } if (opt->inspect_only) { if (opt->mtp_path && opt->mtp_path[0] && opt->distributed.role == DS4_DISTRIBUTED_NONE) { model_open(&e->mtp_model, opt->mtp_path, false, false); ds4_dspark_summary dspark = {0}; e->support_kind = support_model_detect(&e->mtp_model, &e->support_stages, &dspark); if (e->support_kind == DS4_SUPPORT_DSPARK) { dspark_weights_bind_optional(&e->dspark_weights, &e->mtp_model, &dspark); } } *out = e; return 0; } if (e->backend == DS4_BACKEND_CPU && !cpu_load_directional_steering(e)) { ds4_engine_close(e); *out = NULL; return 1; } if (engine_classify_multi_tier(e, gpu_cfg) != 0) { fprintf(stderr, "ds4: failed to classify multi-tier placement\n"); ds4_engine_close(e); *out = NULL; return 1; } if (e->ssd_streaming && e->multi_tier) { fprintf(stderr, "ds4: --ssd-streaming is not compatible with multi-GPU placement\n"); ds4_engine_close(e); *out = NULL; return 1; } if (gpu_cfg && e->n_placement_entries > 0) { int spilled = 0; size_t spilled_bytes = 0; size_t entry_bytes_buf[DS4_MAX_LAYER + 2]; size_t used_bytes[DS4_LAYER_PACK_MAX_GPUS] = {0}; size_t budget_bytes[DS4_LAYER_PACK_MAX_GPUS] = {0}; const int have_entry_bytes = engine_compute_entry_bytes(e, entry_bytes_buf) == 0; if (have_entry_bytes) { for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { budget_bytes[d] = e->gpu_cfg.vram_bytes[d]; } for (int i = 0; i < e->n_placement_entries; i++) { const int dev = e->placement[i]; if (dev == DS4_LAYER_PACK_CPU) { spilled++; spilled_bytes += entry_bytes_buf[i]; } else if (dev >= 0 && dev < e->gpu_cfg.n_gpus) { used_bytes[dev] += entry_bytes_buf[i]; } } } if (spilled > 0) { size_t total_budget = 0; for (int d = 0; d < e->gpu_cfg.n_gpus; d++) { total_budget += e->gpu_cfg.vram_bytes[d]; } if (have_entry_bytes) { ds4_layer_pack_print(stderr, e->placement, e->n_placement_entries, DS4_N_LAYER, entry_bytes_buf, used_bytes, budget_bytes, e->gpu_cfg.n_gpus); } fprintf(stderr, "ds4: per-tier graph scratch reserved on each GPU: " "%.2f GiB (pre-subtracted from each --gpu-vram budget " "before packing).\n", (double)engine_per_tier_graph_overhead_bytes(e) / (1024.0 * 1024.0 * 1024.0)); fprintf(stderr, "ds4: CPU-spill placement detected; CPU-tier execution wiring " "is the wave-3b mgpu-graph-session-cpu-spill follow-up.\n"); fprintf(stderr, "ds4: --gpu-vram placement does not fit at the requested " "context (ctx hint = %d):\n" "ds4: %d placement entries spilled to CPU " "(%.2f GiB unaccommodated of %.2f GiB total per-device budget).\n" "ds4: Lower --ctx / --ctx-max, raise --gpu-vram budgets, or use " "--gpu-vram auto on a host with more free VRAM.\n" "ds4: Refusing upfront to avoid silent OOM at session_create.\n", e->placement_ctx_hint, spilled, (double)spilled_bytes / (1024.0 * 1024.0 * 1024.0), (double)total_budget / (1024.0 * 1024.0 * 1024.0)); ds4_engine_close(e); *out = NULL; return 1; } } if (opt->mtp_path && opt->mtp_path[0] && opt->distributed.role == DS4_DISTRIBUTED_NONE) { if (e->ssd_streaming) { fprintf(stderr, "ds4: --ssd-streaming is not compatible with --mtp yet\n"); ds4_engine_close(e); *out = NULL; return 1; } model_open(&e->mtp_model, opt->mtp_path, graph_backend, true); ds4_dspark_summary dspark = {0}; e->support_kind = support_model_detect(&e->mtp_model, &e->support_stages, &dspark); if (e->support_kind == DS4_SUPPORT_MTP_LEGACY) { if (opt->tp.role != DS4_TP_NONE) { fprintf(stderr, "ds4: legacy MTP support is ignored under tensor parallelism; " "using the normal TP decode path\n"); model_close(&e->mtp_model); e->support_kind = DS4_SUPPORT_NONE; e->support_stages = 0; } else { mtp_weights_bind(&e->mtp_weights, &e->mtp_model); e->mtp_ready = true; fprintf(stderr, "ds4: MTP support model loaded: %s (draft=%d)\n", opt->mtp_path, e->mtp_draft_tokens); } } else if (e->support_kind == DS4_SUPPORT_DSPARK) { dspark_weights_bind_optional(&e->dspark_weights, &e->mtp_model, &dspark); fprintf(stderr, "ds4: DSpark support model detected: %s " "(stages=%u block=%u markov_rank=%u tensors=%u missing=%u " "invalid=%u metadata_errors=%u); " "use --dspark to enable experimental runtime decode\n", opt->mtp_path, e->support_stages, dspark.block_size, dspark.markov_rank, e->dspark_weights.present_tensors, e->dspark_weights.missing_tensors, e->dspark_weights.invalid_tensors, e->dspark_weights.metadata_errors); } else { fprintf(stderr, "ds4: unsupported --mtp support model %s (detected=%s); " "expected legacy MTP or DSpark tensors\n", opt->mtp_path, support_kind_name(e->support_kind)); ds4_engine_close(e); *out = NULL; return 1; } } #ifndef DS4_NO_GPU if (e->backend == DS4_BACKEND_CUDA) { #ifdef __APPLE__ fprintf(stderr, "ds4: CUDA backend requested but this build is linked with Metal, not CUDA\n"); ds4_engine_close(e); *out = NULL; return 1; #endif } if (e->backend == DS4_BACKEND_METAL) { #ifndef __APPLE__ fprintf(stderr, "ds4: Metal backend requested but this build is linked with CUDA, not Metal\n"); ds4_engine_close(e); *out = NULL; return 1; #endif } /* With a raised wired limit the sharded span views (~97 GiB) fit the * GPU budget, so let the residency set pin them — that is what makes * the shard actually resident. Without the sysctl, fall back to lazy * faulting (slow but functional). */ if (graph_backend && tp_shard && glm_graph_wired_limit_bytes() == 0) { fprintf(stderr, "ds4: iogpu.wired_limit_mb is 0 -- TP expert shard will page " "lazily; raise it (e.g. sudo sysctl iogpu.wired_limit_mb=120000) " "for full residency\n"); ds4_gpu_model_residency_skip(1); } if (graph_backend) { if (e->multi_tier) { /* Wave-2 multi-tier branch. * * 1. Initialize all configured CUDA devices via ds4_gpu_init_multi * (the multi-GPU plumbing multi-device init that populates g_gpu[]). * 2. Skip the legacy ds4_gpu_set_model_map_range / accelerator * cache calls — those are single-tier behavior. * 3. Install per-device selective caches per the placement. * 4. Refuse engine creation with the documented notice: this PR * ships scaffolding only; execution wiring lands in * multi-GPU execution. */ /* ds4_gpu_init_multi returns 1 on success and 0 on failure * (matches ds4_gpu_init). */ if (ds4_gpu_init_multi(gpu_cfg) == 0) { fprintf(stderr, "ds4: ds4_gpu_init_multi failed; aborting multi-tier startup\n"); ds4_engine_close(e); *out = NULL; return 1; } e->metal_ready = true; ds4_gpu_set_quality(e->quality); (void)ds4_gpu_set_model_fd(e->model.fd); if (engine_install_gpu_placement(e) != 0) { ds4_engine_close(e); *out = NULL; return 1; } if (engine_install_dspark_support_cache(e) != 0) { ds4_engine_close(e); *out = NULL; return 1; } /* GPU-only multi-tier execution is now wired up * (B2-B6: per-tier graph allocation, dispatch loops, boundary * copies). CPU-spill placements were rejected by * engine_install_gpu_placement above with stderr naming the * CPU-spill execution follow-up. Skip the single-tier ds4_gpu_init / * set_model_map_range path below — those calls are explicitly * single-tier behavior (per the pre-B7 comment block). */ *out = e; return 0; } /* Single-tier path (every existing caller). Body is byte-equivalent * to pre-multi-GPU CLI main. */ e->metal_ready = ds4_gpu_init() != 0; if (!e->metal_ready) { fprintf(stderr, "ds4: %s backend unavailable; aborting startup\n", ds4_backend_name(e->backend)); ds4_engine_close(e); *out = NULL; return 1; } ds4_gpu_set_quality(e->quality); ds4_gpu_set_glm_model(DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA); ds4_gpu_set_ssd_streaming(e->ssd_streaming); if (!ds4_engine_configure_streaming_auto_cache(e)) { ds4_engine_close(e); *out = NULL; return 1; } if (!ds4_engine_configure_streaming_cache_budget(e)) { ds4_engine_close(e); *out = NULL; return 1; } if (!ds4_engine_glm_streaming_memory_guard( e, load_slice, load_layer_start, load_layer_end, load_output, opt->context_size, "after GLM streaming cache budget")) { ds4_engine_close(e); *out = NULL; return 1; } ds4_gpu_set_streaming_expert_cache_budget(e->ssd_streaming_cache_experts); if (e->ssd_streaming) { /* * Pin the expert cache's slab size class to the model's uniform * per-expert bytes, and count mixed-precision (boosted) layers: * those are served through mapped model views instead of the * cache (see weights_streaming_layer_experts_uniform). */ uint64_t slab_expert_bytes = 0; if (ds4_streaming_routed_expert_bytes(&e->weights, &slab_expert_bytes)) { ds4_gpu_set_streaming_expert_cache_expert_bytes(slab_expert_bytes); uint32_t routed = 0, boosted = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &e->weights.layer[il]; if (!l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps) continue; routed++; if (!weights_streaming_layer_experts_uniform(&e->weights, il)) boosted++; } if (boosted > 0) { fprintf(stderr, "ds4: SSD streaming mixed-precision model: %u/%u routed layers " "off the slab size class will bypass the expert cache and read " "experts via mapped model views\n", boosted, routed); } if (boosted * 2 > routed) { fprintf(stderr, "ds4: WARNING: the majority of routed layers (%u/%u) are off the " "slab size class (is the FIRST routed layer itself boosted?); " "expert-cache hit rate will be catastrophic\n", boosted, routed); } /* * Below one token's routed working set (uniform routed layers * x experts used) every token evicts entries it is about to * reuse, and prefill serves layer overflow through mapped * model views. Output stays byte-identical at any budget * (the addr-table kernels read the same bytes either way); * only throughput collapses, so warn instead of refusing. */ const uint64_t min_experts = (uint64_t)(routed - boosted) * DS4_N_EXPERT_USED; if (min_experts != 0 && e->ssd_streaming_cache_experts != 0 && e->ssd_streaming_cache_experts < 2u * min_experts) { fprintf(stderr, "ds4: WARNING: SSD streaming expert cache (%u experts) is " "under twice the per-token routed working set (%u layers " "x %u experts = %llu); expect heavy thrashing below " "%.2f GiB\n", e->ssd_streaming_cache_experts, routed - boosted, DS4_N_EXPERT_USED, (unsigned long long)min_experts, (double)(2u * min_experts * slab_expert_bytes) / 1073741824.0); } } } (void)ds4_gpu_set_model_fd(e->model.fd); int model_map_ok = 0; uint64_t *load_offsets = NULL; uint64_t *load_sizes = NULL; uint32_t load_span_count = 0; if (e->ssd_streaming) { const bool map_output = load_slice && load_output; ds4_model_map_span_vec spans; bool spans_ok = false; if (load_slice) { spans_ok = weights_model_map_decode_static_slice_spans( &e->weights, load_layer_start, load_layer_end, true, map_output, &spans); } else { spans_ok = weights_model_map_token_spans(&e->weights, &spans); } if (!spans_ok) { fprintf(stderr, "ds4: invalid SSD streaming initial token embedding map\n"); ds4_engine_close(e); *out = NULL; return 1; } uint64_t *offsets = xmalloc((size_t)spans.len * sizeof(offsets[0])); uint64_t *sizes = xmalloc((size_t)spans.len * sizeof(sizes[0])); uint64_t span_bytes = 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; span_bytes += sizes[i]; } load_offsets = offsets; load_sizes = sizes; load_span_count = spans.len; e->startup_model_span_bytes = span_bytes; if (load_slice) { char load_end[32]; if (map_output && load_layer_end == UINT32_MAX) { snprintf(load_end, sizeof(load_end), "output"); } else if (map_output) { snprintf(load_end, sizeof(load_end), "%u+output", load_layer_end); } else { snprintf(load_end, sizeof(load_end), "%u", load_layer_end); } fprintf(stderr, "ds4: SSD streaming initial %s model map restricted to token + non-routed layers %u:%s (%u spans, %.2f GiB tensor span)\n", ds4_backend_name(e->backend), load_layer_start, load_end, spans.len, (double)span_bytes / 1073741824.0); } else { fprintf(stderr, "ds4: SSD streaming initial %s model map restricted to token embedding (%u spans, %.2f GiB tensor span)\n", ds4_backend_name(e->backend), spans.len, (double)span_bytes / 1073741824.0); } model_map_ok = ds4_gpu_set_model_map_spans(e->model.map, e->model.size, load_offsets, load_sizes, load_span_count, spans.max_tensor_bytes); free(spans.v); } else if (load_slice) { const bool map_output = load_output; char load_end[32]; if (map_output && load_layer_end == UINT32_MAX) { snprintf(load_end, sizeof(load_end), "output"); } else if (map_output) { snprintf(load_end, sizeof(load_end), "%u+output", load_layer_end); } else { snprintf(load_end, sizeof(load_end), "%u", load_layer_end); } ds4_model_map_span_vec spans; if (!weights_model_map_spans(&e->weights, load_layer_start, load_layer_end, map_output, &spans)) { fprintf(stderr, "ds4: invalid model load layer slice %u:%s\n", load_layer_start, load_end); ds4_engine_close(e); *out = NULL; return 1; } uint64_t *offsets = xmalloc((size_t)spans.len * sizeof(offsets[0])); uint64_t *sizes = xmalloc((size_t)spans.len * sizeof(sizes[0])); uint64_t span_bytes = 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; span_bytes += sizes[i]; } load_offsets = offsets; load_sizes = sizes; load_span_count = spans.len; e->startup_model_span_bytes = span_bytes; fprintf(stderr, "ds4: restricting %s model map to layers %u:%s (%u spans, %.2f GiB tensor span)\n", ds4_backend_name(e->backend), load_layer_start, load_end, spans.len, (double)span_bytes / 1073741824.0); model_map_ok = ds4_gpu_set_model_map_spans(e->model.map, e->model.size, load_offsets, load_sizes, load_span_count, spans.max_tensor_bytes); free(spans.v); } else if (tp_shard) { ds4_model_map_span_vec spans; if (!weights_model_map_sharded_spans(&e->weights, &e->model, tp_shard_rank, &spans)) { fprintf(stderr, "ds4: sharded model span build failed\n"); ds4_engine_close(e); *out = NULL; return 1; } uint64_t *offsets = xmalloc((size_t)spans.len * sizeof(offsets[0])); uint64_t *sizes = xmalloc((size_t)spans.len * sizeof(sizes[0])); uint64_t span_bytes = 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; span_bytes += sizes[i]; } load_offsets = offsets; load_sizes = sizes; load_span_count = spans.len; e->startup_model_span_bytes = span_bytes; fprintf(stderr, "ds4: TP expert shard (rank %d): mapping %u spans, %.2f GiB of %.2f GiB\n", tp_shard_rank, spans.len, (double)span_bytes / 1073741824.0, (double)(e->model.size - e->model.tensor_data_pos) / 1073741824.0); model_map_ok = ds4_gpu_set_model_map_spans(e->model.map, e->model.size, load_offsets, load_sizes, load_span_count, spans.max_tensor_bytes); free(spans.v); } else { e->startup_model_span_bytes = e->model.size > e->model.tensor_data_pos ? e->model.size - e->model.tensor_data_pos : 0; model_map_ok = ds4_gpu_set_model_map_range(e->model.map, e->model.size, e->model.tensor_data_pos, e->model.size - e->model.tensor_data_pos, e->model.max_tensor_bytes); } if (!model_map_ok) { fprintf(stderr, "ds4: %s failed to map model views; aborting startup. " "This is commonly caused by insufficient memory or accelerator VM budget.\n", ds4_backend_name(e->backend)); free(load_offsets); free(load_sizes); ds4_engine_close(e); *out = NULL; return 1; } if (tp_shard) { model_warm_weights_sharded(&e->model, &e->weights, tp_shard_rank); } const bool support_model_runtime_ready = e->mtp_ready || (e->support_kind == DS4_SUPPORT_DSPARK && e->dspark); if (support_model_runtime_ready && !ds4_gpu_set_model_map_range(e->mtp_model.map, e->mtp_model.size, e->mtp_model.tensor_data_pos, e->mtp_model.size - e->mtp_model.tensor_data_pos, e->mtp_model.max_tensor_bytes)) { fprintf(stderr, "ds4: %s failed to map support model views; aborting startup. " "This is commonly caused by insufficient memory or accelerator VM budget.\n", ds4_backend_name(e->backend)); free(load_offsets); free(load_sizes); ds4_engine_close(e); *out = NULL; return 1; } if (!ds4_engine_preload_pro_q4_expert_tables(e, load_slice, load_layer_start, load_layer_end)) { free(load_offsets); free(load_sizes); ds4_engine_close(e); *out = NULL; return 1; } (void)ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map); if (!accelerator_cache_model_tensors(e->backend, &e->model, load_offsets, load_sizes, load_span_count)) { fprintf(stderr, "ds4: %s failed to prepare optional model cache\n", ds4_backend_name(e->backend)); free(load_offsets); free(load_sizes); ds4_engine_close(e); *out = NULL; return 1; } free(load_offsets); free(load_sizes); /* Also apply explicit optional Q8 preload settings to the runtime * support model when loaded. */ if (support_model_runtime_ready) { (void)ds4_gpu_set_model_fd_for_map(e->mtp_model.fd, e->mtp_model.map); if (!accelerator_cache_model_tensors(e->backend, &e->mtp_model, NULL, NULL, 0)) { fprintf(stderr, "ds4: %s failed to prepare optional support model cache\n", ds4_backend_name(e->backend)); ds4_engine_close(e); *out = NULL; return 1; } (void)ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map); } fprintf(stderr, "ds4: %s backend initialized for graph diagnostics\n", ds4_backend_name(e->backend)); } #else if (graph_backend) { fprintf(stderr, "ds4: %s backend requested but this build has no graph backend support; aborting startup\n", ds4_backend_name(e->backend)); ds4_engine_close(e); *out = NULL; return 1; } #endif if (!opt->inspect_only) { ds4_engine_print_startup_memory(e, opt->context_size); } *out = e; return 0; } void ds4_engine_summary(ds4_engine *e) { model_summary(&e->model); if (e->mtp_model.map) { printf("\nsupport model"); if (e->support_kind != DS4_SUPPORT_NONE) { printf(" (%s", support_kind_name(e->support_kind)); if (e->support_stages) printf(", stages=%u", e->support_stages); printf(")"); } printf(":\n"); model_summary(&e->mtp_model); if (e->support_kind == DS4_SUPPORT_DSPARK && e->dspark_weights.n_stages != 0) { printf("support binding: tensors=%u missing=%u invalid=%u metadata_errors=%u\n", e->dspark_weights.present_tensors, e->dspark_weights.missing_tensors, e->dspark_weights.invalid_tensors, e->dspark_weights.metadata_errors); } } } int ds4_engine_vocab_size(ds4_engine *e) { return e ? e->vocab.n_vocab : 0; } uint32_t ds4_engine_prefill_chunk(ds4_engine *e) { return e ? e->prefill_chunk : 0; } int ds4_engine_power(ds4_engine *e) { return e ? e->power_percent : 100; } int ds4_engine_set_power(ds4_engine *e, int power_percent) { if (!e || power_percent < 1 || power_percent > 100) return 1; e->power_percent = power_percent; return 0; } const char *ds4_engine_model_name(ds4_engine *e) { (void)e; return DS4_MODEL_SHAPE_NAME; } int ds4_engine_layer_count(ds4_engine *e) { (void)e; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) return 0; return (int)(DS4_N_LAYER - DS4_N_NEXTN_PREDICT); } return (int)DS4_N_LAYER; } uint32_t ds4_engine_layer_compress_ratio(ds4_engine *e, uint32_t layer) { (void)e; if (layer >= DS4_N_LAYER) return 0; return ds4_layer_compress_ratio(layer); } uint64_t ds4_engine_hidden_f32_values(ds4_engine *e) { (void)e; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) return DS4_N_EMBD; return (uint64_t)DS4_N_HC * DS4_N_EMBD; } bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e, uint32_t layer, uint32_t full_live, uint32_t key_dim, uint32_t value_dim, uint32_t compact_live, uint32_t index_live, uint64_t *out) { #ifdef DS4_NO_GPU (void)e; (void)layer; (void)full_live; (void)key_dim; (void)value_dim; (void)compact_live; (void)index_live; (void)out; return false; #else (void)e; if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) return false; return glm_layer_payload_tensor_bytes(layer, full_live, key_dim, value_dim, compact_live, index_live, out); #endif } int ds4_engine_model_id(ds4_engine *e) { (void)e; return (int)DS4_MODEL_VARIANT; } /* Decode gate firing schedule for the TP transport (see ds4_tp_identity): * DS4 fires ATTN+FFN gates on every layer (identity mapping); GLM fires a * single FFN gate per sparse layer, skipping the leading dense blocks. */ void ds4_engine_tp_gate_schedule(ds4_engine *e, uint32_t *start, uint32_t *step, uint32_t *per_token) { (void)e; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { /* One FFN gate per sparse layer of the NORMAL pass: the leading * dense blocks fire nothing and the trailing nextn/MTP block is * not part of the decode pass at all. */ *start = DS4_N_LEADING_DENSE * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; *step = DS4_TP_GATES_PER_LAYER; *per_token = DS4_N_LAYER - DS4_N_NEXTN_PREDICT - DS4_N_LEADING_DENSE; } else { *start = 0; *step = 1; *per_token = DS4_N_LAYER * DS4_TP_GATES_PER_LAYER; } } int ds4_engine_embd_dim(ds4_engine *e) { (void)e; return (int)DS4_N_EMBD; } uint64_t ds4_engine_model_bytes(ds4_engine *e) { return e->model.size; } int ds4_engine_tp_vocab_split(ds4_engine *e) { return e && e->tp.active && e->tp.vocab_split; } #if !defined(DS4_NO_GPU) && defined(__APPLE__) static int ds4_engine_tp_exchange(void *ud, uint32_t layer, uint32_t gate, uint64_t seq) { ds4_tp *tp = ud; const int ok = ds4_tp_gate_exchange(tp, layer, gate, seq); if (!ok) ds4_tp_mark_failed(tp); return ok; } static int ds4_engine_tp_batch_exchange(void *ud, uint32_t layer, uint32_t rows, uint64_t seq) { ds4_tp *tp = ud; const int ok = ds4_tp_batch_gate_exchange(tp, layer, rows, seq); if (!ok) ds4_tp_mark_failed(tp); return ok; } static int ds4_engine_tp_big_exchange(void *ud, uint32_t layer, uint64_t seq, const void *out, void *in, uint64_t bytes) { if (g_glm_tp_debug_ids && getenv("DS4_GLM_TP_DEBUG")) { fprintf(stderr, "ds4-tp: big gate l=%u ids[0..7]=%d %d %d %d %d %d %d %d\n", layer, g_glm_tp_debug_ids[0], g_glm_tp_debug_ids[1], g_glm_tp_debug_ids[2], g_glm_tp_debug_ids[3], g_glm_tp_debug_ids[4], g_glm_tp_debug_ids[5], g_glm_tp_debug_ids[6], g_glm_tp_debug_ids[7]); } ds4_tp *tp = ud; const int ok = ds4_tp_big_gate_exchange(tp, layer, seq, out, in, bytes); if (!ok) ds4_tp_mark_failed(tp); return ok; } #endif int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen) { #if defined(DS4_NO_GPU) || !defined(__APPLE__) (void)e; (void)tp; snprintf(err, errlen, "tensor parallelism requires the Metal backend"); return 0; #else if (e->backend != DS4_BACKEND_METAL) { snprintf(err, errlen, "tensor parallelism requires the Metal backend"); return 0; } if (e->tp.active) { snprintf(err, errlen, "tensor parallelism already bound"); return 0; } const uint32_t slots = (uint32_t)DS4_N_LAYER * DS4_TP_GATES_PER_LAYER; const uint64_t vec_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); const uint64_t slab_bytes = ds4_tp_slab_bytes((uint32_t)DS4_N_LAYER, (uint32_t)DS4_N_EMBD); e->tp.slab = ds4_gpu_tensor_alloc(slab_bytes); e->tp.zero_vec = ds4_gpu_tensor_alloc(vec_bytes); e->tp.out_views = calloc(slots, sizeof(*e->tp.out_views)); e->tp.in_views = calloc(slots, sizeof(*e->tp.in_views)); e->tp.batch_out_views = calloc((size_t)DS4_N_LAYER, sizeof(*e->tp.batch_out_views)); e->tp.batch_in_views = calloc((size_t)DS4_N_LAYER, sizeof(*e->tp.batch_in_views)); if (!e->tp.batch_out_views || !e->tp.batch_in_views) { snprintf(err, errlen, "tp: batch view table allocation failed"); return 0; } if (!e->tp.slab || !e->tp.zero_vec || !e->tp.out_views || !e->tp.in_views) { snprintf(err, errlen, "tp: slab allocation failed (%llu bytes)", (unsigned long long)slab_bytes); return 0; } memset(ds4_gpu_tensor_contents(e->tp.slab), 0, slab_bytes); memset(ds4_gpu_tensor_contents(e->tp.zero_vec), 0, vec_bytes); if (!ds4_tp_attach_slab(tp, ds4_gpu_tensor_contents(e->tp.slab), err, errlen)) return 0; for (uint32_t l = 0; l < (uint32_t)DS4_N_LAYER; l++) { for (uint32_t gate = 0; gate < DS4_TP_GATES_PER_LAYER; gate++) { const uint32_t slot = l * DS4_TP_GATES_PER_LAYER + gate; e->tp.out_views[slot] = ds4_gpu_tensor_view( e->tp.slab, ds4_tp_slab_out_offset(tp, l, gate), vec_bytes); e->tp.in_views[slot] = ds4_gpu_tensor_view( e->tp.slab, ds4_tp_slab_in_offset(tp, l, gate), vec_bytes); if (!e->tp.out_views[slot] || !e->tp.in_views[slot]) { snprintf(err, errlen, "tp: slab view creation failed"); return 0; } } e->tp.batch_out_views[l] = ds4_gpu_tensor_view( e->tp.slab, ds4_tp_slab_batch_out_offset(tp, l), (uint64_t)DS4_TP_BATCH_MAX_ROWS * vec_bytes); e->tp.batch_in_views[l] = ds4_gpu_tensor_view( e->tp.slab, ds4_tp_slab_batch_in_offset(tp, l), (uint64_t)DS4_TP_BATCH_MAX_ROWS * vec_bytes); if (!e->tp.batch_out_views[l] || !e->tp.batch_in_views[l]) { snprintf(err, errlen, "tp: batch slab view creation failed"); return 0; } } if (!ds4_gpu_tp_init((uint32_t)ds4_tp_rank(tp), e->tp.slab, ds4_tp_slab_gpu_flags_offset(tp), ds4_engine_tp_exchange, tp)) { snprintf(err, errlen, "tp: gate service init failed"); return 0; } ds4_gpu_tp_set_batch_exchange(ds4_engine_tp_batch_exchange); ds4_gpu_tp_set_big_exchange(ds4_engine_tp_big_exchange); /* GLM keeps its replicated output head unsplit in v0: the * leader computes full logits and nothing crosses the wire. */ e->tp.vocab_split = DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA; e->tp.ctx = tp; e->tp.rank = ds4_tp_rank(tp); e->tp.eval_seq = 0; e->tp.active = true; ds4_log(stderr, DS4_LOG_OK, "tensor parallelism bound: rank %d, 50/50 expert split, %s transport", e->tp.rank, ds4_tp_is_rdma(tp) ? "rdma" : "tcp"); return 1; #endif } bool ds4_engine_is_glm_dsa(ds4_engine *e) { (void)e; return DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; } void ds4_engine_close(ds4_engine *e) { if (!e) return; #if !defined(DS4_NO_GPU) && defined(__APPLE__) if (e->tp.active) { ds4_gpu_tp_shutdown(); const uint32_t slots = (uint32_t)DS4_N_LAYER * DS4_TP_GATES_PER_LAYER; for (uint32_t i = 0; i < slots; i++) { if (e->tp.out_views) ds4_gpu_tensor_free(e->tp.out_views[i]); if (e->tp.in_views) ds4_gpu_tensor_free(e->tp.in_views[i]); } for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER; i++) { if (e->tp.batch_out_views) ds4_gpu_tensor_free(e->tp.batch_out_views[i]); if (e->tp.batch_in_views) ds4_gpu_tensor_free(e->tp.batch_in_views[i]); } free(e->tp.batch_out_views); free(e->tp.batch_in_views); free(e->tp.out_views); free(e->tp.in_views); ds4_gpu_tensor_free(e->tp.zero_vec); ds4_gpu_tensor_free(e->tp.slab); memset(&e->tp, 0, sizeof(e->tp)); } #endif ds4_expert_profile_close(); weights_free(&e->weights); vocab_free(&e->vocab); ds4_threads_shutdown(); if (e->mtp_model.map) model_close(&e->mtp_model); model_close(&e->model); #ifndef DS4_NO_GPU ds4_gpu_cleanup(); #endif ds4_ssd_memory_lock_release(&e->simulated_memory); ds4_release_instance_lock(); free(e->directional_steering_dirs); free(e->directional_steering_file); free(e); } #ifndef DS4_NO_GPU static bool ds4_dspark_stats_enabled(void) { const char *env = getenv("DS4_DSPARK_STATS"); return env && env[0] && strcmp(env, "0") != 0; } static void ds4_format_len_hist( char *buf, size_t buflen, const uint64_t *hist) { if (!buf || buflen == 0 || !hist) return; size_t off = 0; buf[0] = '\0'; for (uint32_t i = 0; i <= DS4_DSPARK_MAX_BLOCK_SIZE; i++) { if (hist[i] == 0) continue; const int n = snprintf(buf + off, off < buflen ? buflen - off : 0, "%s%u:%llu", off ? "," : "", i, (unsigned long long)hist[i]); if (n < 0) break; if ((size_t)n >= (off < buflen ? buflen - off : 0)) { off = buflen - 1u; break; } off += (size_t)n; } if (off == 0) { snprintf(buf, buflen, "none"); } } static void ds4_session_print_dspark_stats(const ds4_session *s) { if (!s || !ds4_dspark_stats_enabled()) return; const ds4_dspark_spec_stats *st = &s->dspark_stats; if (st->cycles == 0 && st->propose_ms == 0.0) return; char draft_hist[192]; char accept_hist[192]; ds4_format_len_hist(draft_hist, sizeof(draft_hist), st->draft_len_hist); ds4_format_len_hist(accept_hist, sizeof(accept_hist), st->accepted_len_hist); const double accept_rate = st->proposed_tokens ? (100.0 * (double)st->accepted_draft_tokens / (double)st->proposed_tokens) : 0.0; const double avg_accept = st->cycles ? (double)st->accepted_draft_tokens / (double)st->cycles : 0.0; const double extra_ms = st->propose_ms + st->total_ms; const double net_saved_ms = st->saved_ms - extra_ms; fprintf(stderr, "ds4: DSpark stats cycles=%llu first_tokens=%llu proposed=%llu " "accepted_draft=%llu accept_rate=%.2f%% avg_accept=%.3f " "full=%llu partial=%llu miss_first=%llu no_draft=%llu " "no_room=%llu invalid=%llu scheduler_skips=%llu " "tail_skips=%llu verifier_unavailable=%llu errors=%llu time_ms propose=%.3f " "prop_stage0=%.3f prop_setup=%.3f prop_cache=%.3f " "prop_chain=%.3f prop_hidden=%.3f prop_conf0=%.3f " "prop_logits=%.3f prop_markov=%.3f prop_confidence=%.3f " "snapshot=%.3f verify=%.3f verify_upload=%.3f " "verify_layer=%.3f verify_head=%.3f verify_read=%.3f " "verify_fused_head=%llu replay=%.3f spec_total=%.3f " "target=%.3f saved=%.3f net_saved=%.3f " "draft_len_hist=%s accepted_len_hist=%s\n", (unsigned long long)st->cycles, (unsigned long long)st->first_tokens, (unsigned long long)st->proposed_tokens, (unsigned long long)st->accepted_draft_tokens, accept_rate, avg_accept, (unsigned long long)st->full_accepts, (unsigned long long)st->partial_accepts, (unsigned long long)st->first_misses, (unsigned long long)st->no_draft, (unsigned long long)st->no_room, (unsigned long long)st->invalid_draft, (unsigned long long)st->scheduler_skips, (unsigned long long)st->tail_skips, (unsigned long long)st->verifier_unavailable, (unsigned long long)st->verifier_errors, st->propose_ms, st->propose_stage0_ms, st->propose_setup_ms, st->propose_cache_ms, st->propose_chain_ms, st->propose_hidden_ms, st->propose_conf0_ms, st->propose_logits_ms, st->propose_markov_ms, st->propose_confidence_ms, st->snapshot_ms, st->verify_ms, st->verify_upload_ms, st->verify_layer_ms, st->verify_head_ms, st->verify_read_ms, (unsigned long long)st->verifier_fused_head, st->replay_ms, st->total_ms, st->target_ms, st->saved_ms, net_saved_ms, draft_hist, accept_hist); } #endif static bool ds4_session_tp_leader(const ds4_session *s) { return s && s->engine && s->engine->tp.active && s->engine->tp.rank == 0; } static int ds4_session_tp_register(ds4_session *s) { if (!ds4_session_tp_leader(s)) return 1; ds4_engine *e = s->engine; uint64_t id = ++e->tp.next_session_id; if (id == 0) id = ++e->tp.next_session_id; char err[256] = ""; if (!ds4_tp_send_session_create(e->tp.ctx, id, s->ctx_size) || !ds4_tp_wait_command_ack(e->tp.ctx, id, "session create", err, sizeof(err))) { fprintf(stderr, "ds4: %s\n", err[0] ? err : "tp: worker session create send failed"); return 0; } s->tp_session_id = id; return 1; } int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { if (!out || !e || ctx_size <= 0) return 1; if (e->backend == DS4_BACKEND_CPU) { if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { fprintf(stderr, "ds4: GLM sessions currently require a graph backend\n"); return 1; } if (e->distributed.role == DS4_DISTRIBUTED_COORDINATOR) { fprintf(stderr, "ds4: distributed coordinator sessions require the graph backend\n"); return 1; } ds4_session *s = xcalloc(1, sizeof(*s)); s->engine = e; s->ctx_size = ctx_size; s->prefill_cap = ds4_prefill_cap_for_prompt(ctx_size, e->prefill_chunk); kv_cache_init(&s->cpu_cache, (uint32_t)ctx_size, 0); cpu_decode_scratch_init(&s->cpu_scratch, (uint32_t)ctx_size); s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0])); s->sample_probs = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->sample_probs[0])); if (!ds4_session_tp_register(s)) { ds4_session_free(s); return 1; } *out = s; return 0; } #ifdef DS4_NO_GPU return 1; #else if (!ds4_backend_uses_graph(e->backend) || !e->metal_ready) return 1; ds4_session *s = xcalloc(1, sizeof(*s)); s->engine = e; s->ctx_size = ctx_size; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { 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; bool require_token_embd = true; bool require_output = true; if (e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set) { layer_start = e->distributed.layers.start; layer_end = e->distributed.layers.has_output ? (normal_layers ? normal_layers - 1u : 0u) : e->distributed.layers.end; require_token_embd = layer_start == 0; require_output = e->distributed.layers.has_output; } if (!normal_layers || layer_start > layer_end || layer_end >= normal_layers) { fprintf(stderr, "ds4: invalid GLM layer slice %u:%u\n", layer_start, layer_end); free(s); return 1; } s->glm_graph.placement = e->multi_tier ? e->placement : NULL; if (!glm_graph_alloc_slice(&s->glm_graph, &e->model, &e->weights, ctx_size, e->ssd_streaming, e->ssd_streaming_cold, ds4_engine_streaming_transient_guard_bytes(e), layer_start, layer_end, require_token_embd, require_output)) { free(s); return 1; } s->prefill_cap = s->glm_graph.ctx_cap; ds4_gpu_enable_q8_dequant_gemm(); s->glm_graph_ready = true; s->glm_graph.quality = e->quality; s->glm_graph.ssd_streaming = e->ssd_streaming; s->glm_graph.ssd_streaming_cold = e->ssd_streaming_cold; #if !defined(DS4_NO_GPU) && defined(__APPLE__) if (e->tp.active) { s->glm_graph.tp_world = 2; s->glm_graph.tp_rank = (uint32_t)e->tp.rank; s->glm_graph.tp_out = e->tp.out_views; s->glm_graph.tp_in = e->tp.in_views; } #endif if (e->ssd_streaming && !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL", "DS4_METAL_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL")) { ds4_gpu_graph seed_graph; memset(&seed_graph, 0, sizeof(seed_graph)); seed_graph.quality = e->quality; seed_graph.ssd_streaming = e->ssd_streaming; seed_graph.ssd_streaming_cold = e->ssd_streaming_cold; seed_graph.streaming_preload_experts = e->ssd_streaming_preload_experts; if (!metal_graph_seed_streaming_expert_cache_from_hotlist( &seed_graph, &e->model, &e->weights)) { glm_graph_free(&s->glm_graph); free(s); return 1; } } s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0])); s->sample_probs = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->sample_probs[0])); if (e->distributed.role == DS4_DISTRIBUTED_COORDINATOR) { char err[256]; if (ds4_dist_session_create(&s->distributed, e, &e->distributed, s, ctx_size, err, sizeof(err)) != 0) { fprintf(stderr, "ds4: failed to create distributed coordinator session: %s\n", err[0] ? err : "unknown error"); glm_graph_free(&s->glm_graph); free(s->glm_mtp_hc); free(s->glm_mtp_logits0); free(s->logits); free(s->sample_probs); free(s); return 1; } } if (!ds4_session_tp_register(s)) { ds4_session_free(s); return 1; } *out = s; return 0; } s->prefill_cap = metal_graph_prefill_cap_for_prompt(ctx_size, e->prefill_chunk); const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, s->prefill_cap); const ds4_layer_weights *shape_layer = weights_first_bound_layer(&e->weights); if (!shape_layer) { fprintf(stderr, "ds4: no transformer layers are loaded\n"); free(s); return 1; } const bool need_spec_verifier = e->mtp_ready || (e->support_kind == DS4_SUPPORT_DSPARK && e->dspark) || e->tp.active; /* TP worker mirrors the leader's verify blocks */ const int *placement = e->multi_tier ? e->placement : NULL; const ds4_gpu_graph *shared_prefill_workspace = e->share_session_prefill_workspace && e->shared_prefill_workspace_ready ? &e->shared_prefill_workspace : NULL; s->graph.dspark_exec_tier = e->multi_tier ? e->dspark_exec_tier : 0; if (!metal_graph_alloc_raw_cap(&s->graph, &e->weights, shape_layer, raw_cap, (uint32_t)ctx_size, s->prefill_cap, need_spec_verifier, placement, e->cuda_tensor_parallel, shared_prefill_workspace)) { free(s); return 1; } if (e->share_session_prefill_workspace && !e->shared_prefill_workspace_ready) { bool workspace_ok = true; for (int t = 0; workspace_ok && t < DS4_MAX_GPUS; t++) { if (!s->graph.batch_cur_hc_by_tier[t]) continue; workspace_ok = metal_graph_ensure_batch_ffn_out_on(&s->graph, t); } if (!workspace_ok) { fprintf(stderr, "ds4: failed to complete shared prefill workspace allocation\n"); metal_graph_free(&s->graph); free(s); return 1; } const uint64_t workspace_bytes = metal_graph_prefill_workspace_bytes(&s->graph); metal_graph_transfer_prefill_workspace( &e->shared_prefill_workspace, &s->graph); e->shared_prefill_workspace_ready = true; fprintf(stderr, "ds4: shared session prefill workspace enabled: " "cap=%u, %.2f GiB total GPU memory; each additional " "session aliases this allocation\n", e->shared_prefill_workspace.prefill_cap, (double)workspace_bytes / 1073741824.0); } s->graph.quality = e->quality; s->graph.ssd_streaming = e->ssd_streaming; s->graph.ssd_streaming_cold = e->ssd_streaming_cold; s->graph.streaming_preload_experts = e->ssd_streaming_preload_experts; if (e->tp.active) { s->graph.tp_world = 2; s->graph.tp_rank = (uint32_t)e->tp.rank; s->graph.tp_out = e->tp.out_views; s->graph.tp_in = e->tp.in_views; s->graph.tp_batch_out = e->tp.batch_out_views; s->graph.tp_batch_in = e->tp.batch_in_views; s->graph.tp_zero = e->tp.zero_vec; const uint64_t half = (uint64_t)DS4_N_VOCAB / 2u; s->graph.tp_logits_half = ds4_gpu_tensor_view( metal_graph_logits(&s->graph), (uint64_t)e->tp.rank * half * sizeof(float), half * sizeof(float)); if (!s->graph.tp_logits_half) { metal_graph_free(&s->graph); free(s); return 1; } } s->graph.power_percent = (uint32_t)e->power_percent; if (!metal_graph_load_directional_steering(&s->graph, e->directional_steering_file, e->directional_steering_attn_scale, e->directional_steering_ffn_scale)) { metal_graph_free(&s->graph); free(s); return 1; } if (e->support_kind == DS4_SUPPORT_DSPARK) { if (!metal_graph_configure_dspark_capture(&s->graph, &e->dspark_weights)) { fprintf(stderr, "ds4: failed to configure DSpark target-hidden capture\n"); metal_graph_free(&s->graph); free(s); return 1; } if (s->graph.dspark_capture_enabled) { fprintf(stderr, "ds4: DSpark target-hidden capture enabled: layers="); for (uint32_t i = 0; i < s->graph.dspark_target_layer_count; i++) { fprintf(stderr, "%s%u", i == 0 ? "" : ",", s->graph.dspark_target_layers[i]); } fprintf(stderr, "\n"); } } s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0])); s->sample_probs = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->sample_probs[0])); if (need_spec_verifier) { s->spec_row_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->spec_row_logits[0])); } if (e->support_kind == DS4_SUPPORT_DSPARK && e->dspark) { const uint64_t dspark_feature_count = (uint64_t)DS4_N_EMBD + (uint64_t)e->dspark_weights.markov_rank; if (dspark_feature_count <= (uint64_t)SIZE_MAX / sizeof(float)) { s->dspark_markov_bias = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->dspark_markov_bias[0])); s->dspark_conf_features = xmalloc((size_t)dspark_feature_count * sizeof(s->dspark_conf_features[0])); s->dspark_conf_features_cap = (size_t)dspark_feature_count; } } if (e->mtp_ready) { s->mtp_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->mtp_logits[0])); s->mtp_draft_token = -1; } if (e->distributed.role == DS4_DISTRIBUTED_COORDINATOR) { char err[256]; if (ds4_dist_session_create(&s->distributed, e, &e->distributed, s, ctx_size, err, sizeof(err)) != 0) { fprintf(stderr, "ds4: failed to create distributed coordinator session: %s\n", err[0] ? err : "unknown error"); metal_graph_free(&s->graph); free(s->logits); free(s->sample_probs); free(s->mtp_logits); free(s->spec_row_logits); free(s->dspark_markov_bias); free(s->dspark_conf_features); free(s); return 1; } } if (!ds4_session_tp_register(s)) { ds4_session_free(s); return 1; } *out = s; return 0; #endif } void ds4_session_free(ds4_session *s) { if (!s) return; if (ds4_session_tp_leader(s) && s->tp_session_id != 0 && !ds4_tp_failed(s->engine->tp.ctx)) { char err[256] = ""; if (!ds4_tp_send_session_destroy(s->engine->tp.ctx, s->tp_session_id) || !ds4_tp_wait_command_ack(s->engine->tp.ctx, s->tp_session_id, "session destroy", err, sizeof(err))) { fprintf(stderr, "ds4: %s\n", err[0] ? err : "tp: worker session destroy send failed"); } s->tp_session_id = 0; } #ifndef DS4_NO_GPU ds4_session_print_dspark_stats(s); #endif ds4_dist_session_free(s->distributed); if (ds4_session_is_cpu(s)) { kv_cache_free(&s->cpu_cache); cpu_decode_scratch_free(&s->cpu_scratch); } #ifndef DS4_NO_GPU else { if (ds4_session_is_glm(s)) { glm_graph_free(&s->glm_graph); } else { metal_graph_free(&s->graph); } } #endif token_vec_free(&s->checkpoint); token_vec_free(&s->greedy_splitkv_segment); free(s->logits); free(s->sample_probs); #ifndef DS4_NO_GPU free(s->glm_mtp_hc); free(s->glm_mtp_logits0); #endif free(s->mtp_logits); #ifndef DS4_NO_GPU free(s->spec_row_logits); free(s->dspark_markov_bias); free(s->dspark_conf_features); #endif free(s); } int ds4_session_distributed_route_ready(ds4_session *s, char *err, size_t errlen) { if (!s || !s->distributed) { if (errlen) snprintf(err, errlen, "session is not a distributed coordinator"); return -1; } return ds4_dist_session_route_ready(s->distributed, err, errlen); } int ds4_session_power(ds4_session *s) { if (!s || !s->engine) return 100; return s->engine->power_percent; } bool ds4_session_is_distributed(ds4_session *s) { return s && s->distributed != NULL; } int ds4_session_set_power(ds4_session *s, int power_percent) { if (!s || !s->engine || power_percent < 1 || power_percent > 100) return 1; #ifndef DS4_NO_GPU if (ds4_session_is_glm(s) && power_percent != 100) { fprintf(stderr, "ds4: session power throttling is not supported for GLM 5.2 yet\n"); return 1; } #endif s->engine->power_percent = power_percent; #ifndef DS4_NO_GPU if (!ds4_session_is_cpu(s) && !ds4_session_is_glm(s)) s->graph.power_percent = (uint32_t)power_percent; #endif return 0; } void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud) { if (!s) return; s->progress = fn; s->progress_ud = ud; } void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud) { if (!s) return; s->display_progress = fn; s->display_progress_ud = ud; } void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud) { if (!s) return; s->cancel = fn; s->cancel_ud = ud; } static bool ds4_session_cancelled(ds4_session *s) { return s && s->cancel && s->cancel(s->cancel_ud); } static bool ds4_session_cancelled_cb(void *ud) { return ds4_session_cancelled(ud); } void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total) { if (!s || !s->progress || !event) return; s->progress(s->progress_ud, event, current, total); } int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen) { if (!s) { if (errlen) snprintf(err, errlen, "missing layer-slice session"); return 1; } ds4_session_invalidate(s); if (ds4_session_is_cpu(s)) { session_cpu_reset_cache(s); return 0; } #ifdef DS4_NO_GPU if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else if (ds4_session_is_glm(s)) { s->checkpoint.len = 0; s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_reset_dense_cache(s); return 0; } if (!metal_graph_reset_prefill_state(&s->graph)) { if (errlen) snprintf(err, errlen, "%s layer-slice state reset failed", ds4_backend_name(s->engine->backend)); return 1; } s->graph.mtp_n_raw = 0; return 0; #endif } 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 || !hidden_hc || n_tokens == 0 || !logits) { if (errlen) snprintf(err, errlen, "invalid output-head hidden-state input"); return 1; } ds4_engine *e = s->engine; if (!weights_have_output_head(&e->weights)) { if (errlen) snprintf(err, errlen, "output head is not loaded"); return 1; } const uint64_t hidden_dim = ds4_engine_hidden_f32_values(e); const float *last_hc = hidden_hc + (uint64_t)(n_tokens - 1u) * hidden_dim; if (ds4_session_is_cpu(s)) { output_logits_one(logits, &e->model, &e->weights, last_hc); return 0; } #ifdef DS4_NO_GPU (void)e; if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else if (ds4_session_is_glm(s)) { ds4_glm_gpu_graph *gg = &s->glm_graph; bool ok = ds4_gpu_tensor_write(gg->cur, 0, last_hc, hidden_dim * sizeof(float)) != 0; if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = glm_graph_encode_output_head(gg, &e->model, &e->weights); if (ok) ok = ds4_gpu_end_commands() != 0; if (ok) ok = ds4_gpu_tensor_read(gg->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: synchronize after GLM output-head hidden-state failure also failed\n"); } if (errlen) snprintf(err, errlen, "%s GLM output-head hidden-state evaluation failed", ds4_backend_name(e->backend)); return 1; } return 0; } ds4_gpu_graph *g = &s->graph; bool ok = ds4_gpu_tensor_write(metal_graph_cur_hc(g), 0, last_hc, hidden_dim * sizeof(float)) != 0; if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, &e->model, &e->weights, e->weights.output->dim[1]); if (ok) ok = ds4_gpu_end_commands() != 0; if (ok) ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: synchronize after output-head hidden-state failure also failed\n"); } if (errlen) snprintf(err, errlen, "%s output-head hidden-state evaluation failed", ds4_backend_name(e->backend)); return 1; } return 0; #endif } static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, char *err, size_t errlen); #ifndef DS4_NO_GPU static int glm_session_logits_argmax(const float *logits) { int best = 0; float bv = logits[0]; for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { if (logits[i] > bv) { bv = logits[i]; best = (int)i; } } return best; } /* One GLM MTP speculative cycle (greedy): evaluate first_token, and when a * pending draft exists verify [first_token, draft] in one 2-row batch pass * (big gates only under TP). Returns the number of committed tokens (1 or * 2) with s->logits left at the last committed position, or -1 on error. * Both TP ranks run this off the same mirrored EVAL frame and derive * identical drafts/decisions from their identical logits. */ static int ds4_session_glm_spec_cycle(ds4_session *s, int first_token, int *accepted, int accepted_cap, char *err, size_t errlen) { ds4_engine *e = s->engine; ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t executable = glm_graph_normal_layer_count(); const bool timing = s->engine->glm_mtp_timing; if (!s->glm_mtp_hc) { s->glm_mtp_hc = malloc(2ull * DS4_N_EMBD * sizeof(float)); s->glm_mtp_logits0 = malloc((size_t)DS4_N_VOCAB * sizeof(float)); if (!s->glm_mtp_hc || !s->glm_mtp_logits0) { if (errlen) snprintf(err, errlen, "glm mtp: out of memory"); return -1; } } if (!s->glm_mtp_have || accepted_cap < 2 || pos + 2 > g->ctx_size || (g->compact_cache_cap != 0 && pos + 1 >= g->compact_cache_cap)) { /* Plain step, then seed the first draft from g->cur. */ s->glm_spec_inside = 1; const int rc = ds4_session_eval_internal(s, first_token, false, err, errlen); s->glm_spec_inside = 0; if (rc != 0) return -1; const int n1 = glm_session_logits_argmax(s->logits); if (s->glm_mtp_min_pos == 0 || s->glm_mtp_min_pos > pos) { s->glm_mtp_min_pos = pos; } int d = -1; s->glm_mtp_have = 0; if (glm_graph_mtp_step(g, &e->model, &e->weights, n1, pos, s->glm_mtp_min_pos, &d)) { s->glm_mtp_draft = d; s->glm_mtp_have = 1; } accepted[0] = first_token; return 1; } const double t0 = timing ? now_sec() : 0.0; s->glm_mtp_have = 0; const int d = s->glm_mtp_draft; int toks[2] = { first_token, d }; /* Verify MUST run in the same compact/indexed attention world the * decode path writes: the full-KV batch branch attends over rows the * indexed decode never populates, corrupting every prediction. */ (void)executable; /* Fast decode-style verify while the causal window covers everything * the decode path would attend (pre-indexer-selection regime); the * indexed batch path remains the fallback beyond it. */ bool verified = false; if (pos + 2u <= glm_graph_indexer_top_k_limit()) { verified = glm_graph_verify_rows(g, &e->model, &e->weights, toks, pos, 2, s->glm_mtp_hc, s->logits); } if (!verified) { if (!glm_graph_indexed_prefill_batch_ready(g, pos) || glm_graph_limit_indexed_prefill_chunk(pos, 2u) < 2u || !glm_graph_forward_indexed_tokens(g, &e->model, &e->weights, toks, NULL, pos, 2, s->glm_mtp_hc, s->logits, NULL, NULL, pos, 0, 2)) { if (errlen) snprintf(err, errlen, "glm mtp: verify failed"); s->checkpoint_valid = false; return -1; } } const double t1 = timing ? now_sec() : 0.0; /* Row0 logits through the shared head. */ if (!glm_graph_mtp_ensure(g)) { if (errlen) snprintf(err, errlen, "glm mtp: scratch alloc failed"); return -1; } ds4_gpu_tensor *h0_view = ds4_gpu_tensor_view(g->mtp_concat, 0, (uint64_t)DS4_N_EMBD * sizeof(float)); if (!h0_view) { if (errlen) snprintf(err, errlen, "glm mtp: hidden view failed"); return -1; } bool head_ok = ds4_gpu_tensor_write(h0_view, 0, s->glm_mtp_hc, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && glm_graph_forward_output_head(g, &e->model, &e->weights, h0_view, s->glm_mtp_logits0); ds4_gpu_tensor_free(h0_view); if (!head_ok) { if (errlen) snprintf(err, errlen, "glm mtp: row0 head failed"); return -1; } const int n1 = glm_session_logits_argmax(s->glm_mtp_logits0); token_vec_push(&s->checkpoint, first_token); s->checkpoint_valid = true; const int accept = n1 == d; int n_committed = 1; if (accept) { token_vec_push(&s->checkpoint, d); ds4_session_glm_note_dense_cache(s, pos, 2); n_committed = 2; /* s->logits already holds row1 (position pos+1) logits. */ const int n2 = glm_session_logits_argmax(s->logits); int dummy = -1, nd = -1; const bool cu = ds4_gpu_tensor_write(g->cur, 0, s->glm_mtp_hc, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && glm_graph_mtp_step(g, &e->model, &e->weights, d, pos, s->glm_mtp_min_pos, &dummy) && ds4_gpu_tensor_write(g->cur, 0, s->glm_mtp_hc + DS4_N_EMBD, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && glm_graph_mtp_step(g, &e->model, &e->weights, n2, pos + 1u, s->glm_mtp_min_pos, &nd); if (cu) { s->glm_mtp_draft = nd; s->glm_mtp_have = 1; } accepted[0] = first_token; accepted[1] = d; } else { ds4_session_glm_note_dense_cache(s, pos, 1); memcpy(s->logits, s->glm_mtp_logits0, (size_t)DS4_N_VOCAB * sizeof(float)); int nd = -1; const bool cu = ds4_gpu_tensor_write(g->cur, 0, s->glm_mtp_hc, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && glm_graph_mtp_step(g, &e->model, &e->weights, n1, pos, s->glm_mtp_min_pos, &nd); if (cu) { s->glm_mtp_draft = nd; s->glm_mtp_have = 1; } accepted[0] = first_token; } if (timing) { const double t2 = now_sec(); char *dt = ds4_token_text(e, d, NULL); char *nt = ds4_token_text(e, n1, NULL); fprintf(stderr, "ds4: glm mtp cycle: verify2 %.1f ms, head+draft %.1f ms, %s " "(draft %d '%s' vs true %d '%s')\n", (t1 - t0) * 1000.0, (t2 - t1) * 1000.0, accept ? "ACCEPT" : "reject", d, dt ? dt : "?", n1, nt ? nt : "?"); free(dt); free(nt); } return n_committed; } #endif static int ds4_session_slice_check_timeline( ds4_session *s, const int *tokens, uint32_t n_tokens, uint32_t pos0, char *err, size_t errlen) { if (!s || !tokens || n_tokens == 0) { if (errlen) snprintf(err, errlen, "invalid layer-slice token span"); return 1; } const uint32_t ctx_size = (uint32_t)s->ctx_size; if (pos0 > (uint32_t)INT_MAX || n_tokens > (uint32_t)INT_MAX || pos0 > ctx_size || n_tokens > ctx_size - pos0) { if (errlen) snprintf(err, errlen, "layer-slice token span exceeds context"); return 1; } if (!s->checkpoint_valid) { if (pos0 != 0) { if (errlen) snprintf(err, errlen, "layer-slice session needs reset before pos %u", pos0); return 1; } return 0; } if ((uint32_t)s->checkpoint.len != pos0) { if (errlen) snprintf(err, errlen, "layer-slice KV position mismatch: have %d want %u", s->checkpoint.len, pos0); return 1; } return 0; } static DS4_MAYBE_UNUSED void ds4_session_slice_commit_timeline(ds4_session *s, const int *tokens, uint32_t n_tokens) { for (uint32_t i = 0; i < n_tokens; i++) token_vec_push(&s->checkpoint, tokens[i]); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); } int ds4_session_eval_layer_slice(ds4_session *s, const int *tokens, uint32_t n_tokens, uint32_t pos0, uint32_t layer_start, uint32_t layer_end, const float *input_hc, float *output_hc, bool output_logits, float *logits, char *err, size_t errlen) { if (!s || !s->engine) { if (errlen) snprintf(err, errlen, "missing layer-slice session"); return 1; } const uint32_t executable_layers = ds4_model_normal_layer_count(); if (executable_layers == 0 || layer_start > layer_end || layer_end >= executable_layers) { if (errlen) snprintf(err, errlen, "invalid layer-slice layer range %u:%u", layer_start, layer_end); return 1; } if (layer_start != 0 && !input_hc) { if (errlen) snprintf(err, errlen, "layer-slice layer %u requires input hidden-state", layer_start); return 1; } if (output_logits && layer_end + 1u != executable_layers) { if (errlen) snprintf(err, errlen, "layer-slice logits require final transformer layer"); return 1; } if (output_logits && !logits) { if (errlen) snprintf(err, errlen, "layer-slice logits output is missing"); return 1; } if (!weights_layers_bound(&s->engine->weights, layer_start, layer_end)) { if (errlen) snprintf(err, errlen, "requested layer slice %u:%u is not loaded", layer_start, layer_end); return 1; } if (!input_hc && !s->engine->weights.token_embd) { if (errlen) snprintf(err, errlen, "token embedding is not loaded"); return 1; } if (output_logits && !weights_have_output_head(&s->engine->weights)) { if (errlen) snprintf(err, errlen, "output head is not loaded"); return 1; } /* A distributed prefill pipeline may need only the KV side effect for * non-final chunks. In that case both output_hc and logits are NULL. */ if (ds4_session_slice_check_timeline(s, tokens, n_tokens, pos0, err, errlen) != 0) { return 1; } if (ds4_session_is_cpu(s)) { if (errlen) snprintf(err, errlen, "layer slices require the graph backend"); s->checkpoint_valid = false; return 1; } #ifdef DS4_NO_GPU (void)output_hc; if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); s->checkpoint_valid = false; return 1; #else if (ds4_session_is_glm(s)) { ds4_engine *e = s->engine; ds4_glm_gpu_graph *g = &s->glm_graph; if (!s->glm_graph_ready) { if (errlen) snprintf(err, errlen, "%s GLM graph is not initialized", ds4_backend_name(e->backend)); return 1; } if (layer_start != g->layer_start || layer_end != g->layer_end) { if (errlen) snprintf(err, errlen, "requested GLM layer slice %u:%u does not match loaded slice %u:%u", layer_start, layer_end, g->layer_start, g->layer_end); s->checkpoint_valid = false; return 1; } if (n_tokens > s->prefill_cap) { if (errlen) snprintf(err, errlen, "GLM layer-slice chunk %u exceeds prefill cap %u", n_tokens, s->prefill_cap); return 1; } const uint64_t hidden_dim = DS4_N_EMBD; uint32_t done = 0; while (done < n_tokens) { const uint32_t pos = pos0 + done; const uint32_t remaining = n_tokens - done; const float *chunk_input = input_hc ? input_hc + (uint64_t)done * hidden_dim : NULL; float *chunk_output = output_hc ? output_hc + (uint64_t)done * hidden_dim : NULL; uint32_t chunk = 1; bool ok = false; /* * Raw decode starts from a token embedding. Continuation slices * receive a hidden vector from the previous node, so use the * batch-equivalent path even for a single token. */ if (remaining == 1 && pos > 0 && !input_hc && !output_hc) { float *chunk_logits = output_logits ? logits : NULL; ok = glm_graph_forward_token(g, &e->model, &e->weights, tokens[done], chunk_input, pos, chunk_output, chunk_logits, false); if (ok && glm_graph_decode_updates_dense_cache(g, pos, chunk_logits)) { ds4_session_glm_note_dense_cache(s, pos, 1); } } else if (g->full_kv_cache && pos < g->ctx_cap) { chunk = remaining; const uint32_t dense_remaining = g->ctx_cap - pos; if (chunk > dense_remaining) chunk = dense_remaining; float *chunk_logits = (output_logits && done + chunk == n_tokens) ? logits : NULL; ok = glm_graph_forward_tokens(g, &e->model, &e->weights, tokens + done, chunk_input, pos, chunk, chunk_output, chunk_logits, NULL, NULL, pos0, done, n_tokens); if (ok) ds4_session_glm_note_dense_cache(s, pos, chunk); } 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 *chunk_logits = (output_logits && done + chunk == n_tokens) ? logits : NULL; ok = glm_graph_forward_indexed_tokens(g, &e->model, &e->weights, tokens + done, chunk_input, pos, chunk, chunk_output, chunk_logits, NULL, NULL, pos0, done, n_tokens); } else { float *chunk_logits = (output_logits && done + 1u == n_tokens) ? logits : NULL; ok = glm_graph_forward_token(g, &e->model, &e->weights, tokens[done], chunk_input, pos, chunk_output, chunk_logits, false); chunk = 1; } if (!ok) { if (errlen) snprintf(err, errlen, "%s GLM layer-slice evaluation failed at pos %u", ds4_backend_name(e->backend), pos); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } done += chunk; } if (!output_hc && !output_logits && ds4_gpu_synchronize() == 0) { if (errlen) snprintf(err, errlen, "%s GLM layer-slice synchronization failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } ds4_session_slice_commit_timeline(s, tokens, n_tokens); return 0; } if (n_tokens > s->prefill_cap) { if (errlen) snprintf(err, errlen, "layer-slice chunk %u exceeds prefill cap %u", n_tokens, s->prefill_cap); return 1; } ds4_engine *e = s->engine; ds4_gpu_graph *g = &s->graph; if (!input_hc && !output_hc && output_logits && layer_start == 0 && layer_end + 1u == (uint32_t)DS4_N_LAYER) { bool ok = false; ds4_tokens span = {0}; if (pos0 == 0) { span.v = (int *)tokens; span.len = (int)n_tokens; span.cap = (int)n_tokens; ok = metal_graph_prefill_layer_major(g, &e->model, &e->weights, &span, 0, n_tokens, logits, false, NULL, NULL, NULL); } else if (n_tokens == 1) { ok = metal_graph_eval_token_raw_swa(g, &e->model, &e->weights, tokens[0], pos0, logits); } else { if (pos0 > (uint32_t)INT_MAX - n_tokens) { if (errlen) snprintf(err, errlen, "layer-slice full span is too large"); s->checkpoint_valid = false; return 1; } span.len = (int)(pos0 + n_tokens); span.cap = span.len; span.v = calloc((size_t)span.len, sizeof(span.v[0])); if (span.v) { for (uint32_t i = 0; i < n_tokens; i++) span.v[pos0 + i] = tokens[i]; ok = metal_graph_prefill_layer_major(g, &e->model, &e->weights, &span, pos0, n_tokens, logits, false, NULL, NULL, NULL); } free(span.v); } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: synchronize after layer-slice full failure also failed\n"); } if (errlen) snprintf(err, errlen, "%s layer-slice full evaluation failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return 1; } ds4_session_slice_commit_timeline(s, tokens, n_tokens); return 0; } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_bytes = (uint64_t)n_tokens * hc_dim * sizeof(float); if (n_tokens == 1 && pos0 > 0) { if (g->raw_cap == 0) { if (errlen) snprintf(err, errlen, "%s layer-slice decode has no raw KV cache", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return 1; } bool ok = true; if (g->ssd_streaming && !input_hc) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_token(&e->model, &e->weights); } if (input_hc) { ok = ds4_gpu_tensor_write(metal_graph_cur_hc(g), 0, input_hc, hc_dim * sizeof(float)) != 0; } if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok && !input_hc) { ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), e->model.map, e->model.size, e->weights.token_embd->abs_offset, (uint32_t)e->weights.token_embd->dim[1], (uint32_t)tokens[0], DS4_N_EMBD, DS4_N_HC) != 0; } const uint32_t raw_row = pos0 % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos0, 1); const uint32_t split_after_layers = metal_graph_token_split_after_layers(); uint32_t encoded_layers = 0; if (g->ssd_streaming) { if (ok) ok = ds4_gpu_end_commands() != 0; for (uint32_t il = layer_start; ok && il <= layer_end; il++) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_layer_decode(&e->model, &e->weights, il); if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = metal_graph_encode_decode_layer(g, &e->model, &e->weights.layer[il], il, pos0, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, tokens[0]); 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; } if (ok && output_logits) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_output(&e->model, &e->weights); if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, &e->model, &e->weights, e->weights.output->dim[1]); if (ok) ok = ds4_gpu_end_commands() != 0; } } else { for (uint32_t il = layer_start; ok && il <= layer_end; il++) { ok = metal_graph_encode_decode_layer(g, &e->model, &e->weights.layer[il], il, pos0, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, tokens[0]); 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; encoded_layers++; if (ok && split_after_layers != 0 && encoded_layers == split_after_layers && il < layer_end) { ok = ds4_gpu_flush_commands() != 0; } } if (ok && output_logits) { ok = metal_graph_encode_output_head(g, &e->model, &e->weights, e->weights.output->dim[1]); } if (ok) ok = ds4_gpu_end_commands() != 0; } if (ok && !output_hc && !output_logits) ok = ds4_gpu_synchronize() != 0; if (ok && output_hc) { ok = ds4_gpu_tensor_read(metal_graph_cur_hc(g), 0, output_hc, hc_dim * sizeof(float)) != 0; } if (ok && output_logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: synchronize after layer-slice decode failure also failed\n"); } if (errlen) snprintf(err, errlen, "%s layer-slice decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return 1; } ds4_session_slice_commit_timeline(s, tokens, n_tokens); return 0; } ds4_tokens span = { .v = (int *)tokens, .len = (int)n_tokens, .cap = (int)n_tokens, }; bool ok = true; if (g->ssd_streaming && !input_hc) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_token(&e->model, &e->weights); } if (ok) ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), &span, 0, n_tokens); if (ok && input_hc) { ok = ds4_gpu_tensor_write(metal_graph_batch_cur_hc(g), 0, input_hc, hc_bytes) != 0; } else if (ok) { ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), &e->model, &e->weights, &span, 0, n_tokens); } ds4_gpu_tensor *last_hc = NULL; ds4_gpu_tensor *saved_cur = NULL; const int src_tier = g->active_tier; const bool batch_selected_addr = g->ssd_streaming && layer_start == 0 && (metal_graph_stream_prefill_batch_selected_addr_enabled(g, &e->weights, n_tokens) || metal_graph_cuda_stream_prefill_batch_selected_addr_enabled(g, &e->weights, n_tokens)); if (g->ssd_streaming) { for (uint32_t il = layer_start; ok && il <= layer_end; il++) { g->streaming_static_decode_map_current = false; ok = batch_selected_addr ? metal_graph_stream_map_layer_decode(&e->model, &e->weights, il) : metal_graph_stream_map_layer(&e->model, &e->weights, il); if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = metal_graph_encode_layer_batch(g, &e->model, &e->weights.layer[il], il, pos0, n_tokens); } if (ok) ok = ds4_gpu_end_commands() != 0; } } else { if (ok) ok = ds4_gpu_begin_commands() != 0; for (uint32_t il = layer_start; ok && il <= layer_end; il++) { ok = metal_graph_encode_layer_batch(g, &e->model, &e->weights.layer[il], il, pos0, n_tokens); } } if (ok && output_logits) { saved_cur = g->cur_hc_by_tier[src_tier]; last_hc = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), n_tokens - 1u, hc_dim); ok = last_hc != NULL; if (ok && g->ssd_streaming) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_output(&e->model, &e->weights); } if (ok) { g->cur_hc_by_tier[src_tier] = last_hc; if (g->ssd_streaming) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, &e->model, &e->weights, e->weights.output->dim[1]); if (ok && g->ssd_streaming) ok = ds4_gpu_end_commands() != 0; g->cur_hc_by_tier[src_tier] = saved_cur; } } if (ok && !g->ssd_streaming) ok = ds4_gpu_end_commands() != 0; if (saved_cur) g->cur_hc_by_tier[src_tier] = saved_cur; if (last_hc) ds4_gpu_tensor_free(last_hc); if (ok && !output_hc && !output_logits) ok = ds4_gpu_synchronize() != 0; if (ok && output_hc) { ok = ds4_gpu_tensor_read(metal_graph_batch_cur_hc(g), 0, output_hc, hc_bytes) != 0; } if (ok && output_logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: synchronize after layer-slice failure also failed\n"); } if (errlen) snprintf(err, errlen, "%s layer-slice failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return 1; } ds4_session_slice_commit_timeline(s, tokens, n_tokens); return 0; #endif } #ifndef DS4_NO_GPU typedef struct { ds4_session *session; const ds4_tokens *prompt; ds4_session_progress_fn user; void *user_ud; } ds4_sync_progress; static void ds4_session_note_prefill_progress(void *ud, const char *event, int current, int total) { ds4_sync_progress *p = ud; if (!p || !p->session || !p->prompt) return; if (!strcmp(event, "prefill_chunk") && current > 0 && current <= p->prompt->len) { p->session->checkpoint.len = 0; for (int i = 0; i < current; i++) token_vec_push(&p->session->checkpoint, p->prompt->v[i]); p->session->checkpoint_valid = true; p->session->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(p->session); } if (p->user) p->user(p->user_ud, event, current, total); } #endif /* Bring the live backend state to exactly the supplied token prefix. * * ds4-server and the REPL are stateless at the text/API layer but stateful here: * they resend or rebuild the full transcript, and this function decides whether * the live checkpoint is a prefix. A matching prefix is extended in one of two * ways: * * - long suffix: batched layer-major prefill, aligned to absolute chunk * boundaries so compressor/indexer rows finalize in the same order as a * cold prompt; * - short suffix: ordinary one-token decode, which is faster below the * measured crossover and preserves exact autoregressive semantics. * * A non-matching prompt discards the checkpoint and prefills from token zero. */ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen); /* Under tensor parallelism the leader mirrors every public sync/eval to the * worker before doing the work itself, so both engines execute the same * graph sequence and the per-layer gates pair up. The worker acks a sync * once its matching prefill completes, surfacing worker-side failures * here instead of as a gate timeout mid-decode. */ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen) { const bool mirror = ds4_session_tp_leader(s); if (mirror && prompt && prompt->len > 0) { if (!ds4_tp_send_sync(s->engine->tp.ctx, s->tp_session_id, prompt->v, (uint32_t)prompt->len)) { snprintf(err, errlen, "tp: worker sync send failed"); return 1; } } int rc = ds4_session_sync_internal(s, prompt, err, errlen); #ifndef DS4_NO_GPU if (rc == 0) glm_debug_dump_prefill_logits(s->logits); if (rc == 0) { const char *kvp = getenv("DS4_GLM_KV_DUMP"); if (kvp && kvp[0] && s->glm_graph.layer_kv_lora_cache[0]) { const ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t rows = prompt ? (uint32_t)prompt->len : 0; const uint64_t eb = glm_graph_compact_cache_elem_bytes(); const struct { const char *sfx; ds4_gpu_tensor *t; uint64_t rb; } kd[2] = { { "lora0", g->layer_kv_lora_cache[0], DS4_N_KV_LORA * eb }, { "rope0", g->layer_k_rope_cache[0], DS4_N_ROT * eb }, }; for (int i = 0; i < 2 && rows; i++) { char fp[1024]; snprintf(fp, sizeof(fp), "%s.%s", kvp, kd[i].sfx); void *buf = malloc(rows * kd[i].rb); if (buf && ds4_gpu_tensor_read(kd[i].t, 0, buf, rows * kd[i].rb)) { FILE *f = fopen(fp, "wb"); if (f) { fwrite(buf, 1, rows * kd[i].rb, f); fclose(f); } } free(buf); } } } #endif if (mirror) { const bool worker_ok = ds4_tp_wait_command_ack( s->engine->tp.ctx, s->tp_session_id, "prefill sync", err, errlen); bool logits_ok = true; /* A successful worker sends its split logits even if the leader's * local prefill failed. Drain them to keep the control stream framed * before invalidating the mirrored session. */ if (worker_ok && s->engine->tp.vocab_split) { const uint32_t vhalf = (uint32_t)DS4_N_VOCAB / 2u; if (!ds4_tp_recv_logits_half(s->engine->tp.ctx, s->logits + vhalf, vhalf)) { snprintf(err, errlen, "tp: worker sync logits half missing"); logits_ok = false; } } if (rc != 0 || !worker_ok || !logits_ok) { ds4_session_invalidate(s); return rc != 0 ? rc : 1; } } return rc; } static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen) { if (!s || !prompt) { snprintf(err, errlen, "missing session or prompt"); return 1; } if (prompt->len <= 0) { snprintf(err, errlen, "empty prompt"); return 1; } if (prompt->len >= s->ctx_size) { snprintf(err, errlen, "prompt length %d exceeds context %d (one token of generation room is required)", prompt->len, s->ctx_size); return 1; } if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); return DS4_SESSION_SYNC_INTERRUPTED; } if (s->distributed) { const ds4_tokens *checkpoint = s->checkpoint_valid ? &s->checkpoint : NULL; return ds4_dist_session_sync(s->distributed, s, checkpoint, prompt, s->logits, err, errlen); } if (ds4_session_is_cpu(s)) { ds4_engine *e = s->engine; if (s->checkpoint_valid && prompt->len >= s->checkpoint.len && ds4_tokens_starts_with(prompt, &s->checkpoint)) { s->mtp_draft_valid = false; for (int i = s->checkpoint.len; i < prompt->len; i++) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = true; s->mtp_draft_valid = false; return DS4_SESSION_SYNC_INTERRUPTED; } forward_token_raw_swa_cpu_decode_scratch(s->logits, &e->model, &e->weights, &s->cpu_cache, prompt->v[i], (uint32_t)s->checkpoint.len, e->directional_steering_dirs, e->directional_steering_attn_scale, e->directional_steering_ffn_scale, &s->cpu_scratch); token_vec_push(&s->checkpoint, prompt->v[i]); if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i + 1, prompt->len); } s->checkpoint_valid = true; s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = false; return 0; } session_cpu_reset_cache(s); prefill_layer_major_cpu(s->logits, &e->model, &e->weights, &s->cpu_cache, prompt, e->directional_steering_dirs, e->directional_steering_attn_scale, e->directional_steering_ffn_scale); ds4_tokens_copy(&s->checkpoint, prompt); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); if (s->progress) s->progress(s->progress_ud, "prefill_chunk", prompt->len, prompt->len); return 0; } #ifdef DS4_NO_GPU (void)s; (void)prompt; snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else { 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)) { /* 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; { const char *tr = getenv("DS4_GLM_PREFILL_TRUNC"); if (tr && tr[0]) { const int tn = atoi(tr); if (tn > 0 && tn < prompt->len) { glm_trunc_prompt = *prompt; glm_trunc_prompt.len = tn; prompt = &glm_trunc_prompt; } } } if (!s->glm_graph_ready) { snprintf(err, errlen, "%s GLM graph is not initialized", backend_name); return 1; } if ((uint32_t)prompt->len >= s->glm_graph.ctx_size) { snprintf(err, errlen, "prompt length %d leaves no GLM Metal context room (ctx %u)", prompt->len, s->glm_graph.ctx_size); s->checkpoint_valid = false; ds4_session_glm_reset_dense_cache(s); return 1; } int start = 0; bool resumed_checkpoint = false; if (s->checkpoint_valid && prompt->len >= s->checkpoint.len && ds4_tokens_starts_with(prompt, &s->checkpoint)) { start = s->checkpoint.len; resumed_checkpoint = true; s->mtp_draft_valid = false; } else { s->checkpoint.len = 0; s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_reset_dense_cache(s); } uint32_t glm_exact_prefill_max = 64; { const char *em = getenv("DS4_GLM_TP_EXACT_PREFILL_MAX"); if (em && em[0]) glm_exact_prefill_max = (uint32_t)atoi(em); } if (s->engine->glm_tp_token_prefill || s->glm_graph.placement != NULL || (s->glm_graph.tp_world == 2 && (uint32_t)(prompt->len - start) <= glm_exact_prefill_max)) { /* Multi-tier GLM has per-tier decode workspaces and KV caches; * its large batch-prefill workspace is not mirrored across * devices. Small TP prompts also use this path to retain exact * single-node arithmetic. */ for (int i = start; i < prompt->len; i++) { const uint32_t pos = (uint32_t)i; const bool updates_dense = glm_graph_decode_updates_dense_cache(&s->glm_graph, pos, s->logits); const bool last = i + 1 == prompt->len; /* Always request logits: it forces the per-token commit * cadence of normal decode, so the per-layer gates execute * token by token on both ranks instead of piling up in one * giant open batch that times out at the late flush. */ if (!glm_graph_forward_token(&s->glm_graph, &s->engine->model, &s->engine->weights, prompt->v[i], NULL, pos, NULL, s->logits, false)) { snprintf(err, errlen, "%s GLM TP prefill failed at token %d", backend_name, i); s->checkpoint_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } token_vec_push(&s->checkpoint, prompt->v[i]); if (updates_dense) ds4_session_glm_note_dense_cache(s, pos, 1); if (s->progress && ((i - start) % 8 == 0 || last)) { s->progress(s->progress_ud, "prefill_chunk", i + 1 - start, prompt->len - start); } } s->checkpoint_valid = true; s->mtp_draft_valid = false; (void)resumed_checkpoint; glm_debug_dump_prefill_logits(s->logits); return 0; } int checkpoint_len_with_logits = s->checkpoint_valid ? s->checkpoint.len : 0; int suffix = prompt->len - start; const uint32_t resume_min = glm_graph_resume_prefill_min_tokens(); bool dense_cache_gap = (uint32_t)start > s->glm_dense_cache_len; const bool prompt_fits_dense_attention = s->glm_graph.full_kv_cache && (uint32_t)prompt->len <= s->glm_graph.ctx_cap; const bool sync_trace = getenv("DS4_GLM_SYNC_TRACE") != NULL; const bool indexed_batch_available = glm_graph_indexed_prefill_batch_available(&s->glm_graph); const bool indexed_resume_keeps_sparse_state = resumed_checkpoint && suffix > 0 && dense_cache_gap && glm_graph_decode_uses_indexed_attention(&s->glm_graph, (uint32_t)start, s->logits); if (sync_trace) { fprintf(stderr, "ds4: GLM sync start=%d prompt=%d suffix=%d checkpoint=%d dense_len=%u ctx_cap=%u dense_fit=%d resume_min=%u dense_gap=%d indexed_keep=%d indexed_batch=%d batch_ffn=%d\n", start, prompt->len, suffix, s->checkpoint_valid ? s->checkpoint.len : 0, s->glm_dense_cache_len, s->glm_graph.ctx_cap, prompt_fits_dense_attention ? 1 : 0, resume_min, dense_cache_gap ? 1 : 0, indexed_resume_keeps_sparse_state ? 1 : 0, indexed_batch_available ? 1 : 0, glm_graph_indexed_prefill_batch_ffn() ? 1 : 0); } /* If the live checkpoint advanced with indexed attention, the compact * cache already represents the live prefix but dense KV may lag behind. * Continue from compact/indexed state instead of replaying the old * assistant text just to repair dense KV before the new suffix. */ if (resumed_checkpoint && suffix > 0 && dense_cache_gap && !indexed_resume_keeps_sparse_state && s->glm_graph.full_kv_cache && s->glm_dense_cache_len < s->glm_graph.ctx_cap) { const uint32_t dense_end = (uint32_t)prompt->len < s->glm_graph.ctx_cap ? (uint32_t)prompt->len : s->glm_graph.ctx_cap; const uint32_t dense_progress_start = (uint32_t)start; const uint32_t dense_progress_total = (uint32_t)suffix; const uint32_t dense_chunk_max = glm_graph_prefill_chunk_tokens(s->glm_graph.ctx_cap); for (uint32_t dense_pos = s->glm_dense_cache_len; dense_pos < dense_end; ) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); if (!s->checkpoint_valid) { s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; } s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return DS4_SESSION_SYNC_INTERRUPTED; } uint32_t chunk = dense_end - dense_pos; if (chunk > dense_chunk_max) chunk = dense_chunk_max; float *chunk_logits = (dense_pos + chunk >= (uint32_t)prompt->len) ? s->logits : NULL; if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=dense_bridge pos=%u chunk=%u logits=%d\n", dense_pos, chunk, chunk_logits ? 1 : 0); } if (!glm_graph_forward_tokens(&s->glm_graph, &e->model, &e->weights, prompt->v + dense_pos, NULL, dense_pos, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, dense_progress_start, dense_pos > dense_progress_start ? dense_pos - dense_progress_start : 0, dense_progress_total)) { snprintf(err, errlen, "%s GLM dense prefill failed while rebuilding checkpoint cache", backend_name); if (!s->checkpoint_valid) { s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; } s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } const uint32_t dense_next = dense_pos + chunk; ds4_session_glm_note_dense_cache(s, dense_pos, chunk); if (dense_next > (uint32_t)start) { const uint32_t push_from = dense_pos > (uint32_t)start ? dense_pos : (uint32_t)start; if ((uint32_t)s->checkpoint.len != push_from) { snprintf(err, errlen, "%s GLM checkpoint/cache bridge is inconsistent", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } for (uint32_t j = push_from; j < dense_next; j++) { token_vec_push(&s->checkpoint, prompt->v[j]); } start = s->checkpoint.len; if (chunk_logits) { checkpoint_len_with_logits = s->checkpoint.len; s->checkpoint_valid = true; } else { s->checkpoint_valid = false; } s->mtp_draft_valid = false; } dense_pos = dense_next; if (s->progress) { s->progress(s->progress_ud, "prefill_chunk", (int)dense_pos, prompt->len); } } suffix = prompt->len - start; dense_cache_gap = (uint32_t)start > s->glm_dense_cache_len; } if (resumed_checkpoint && suffix > 0 && dense_cache_gap && (uint32_t)suffix >= resume_min && indexed_batch_available) { for (int i = start; i < prompt->len; ) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return DS4_SESSION_SYNC_INTERRUPTED; } const uint32_t pos = (uint32_t)s->checkpoint.len; uint32_t chunk = (uint32_t)(prompt->len - i); if (chunk > s->glm_graph.indexed_prefill_cap) { chunk = s->glm_graph.indexed_prefill_cap; } chunk = glm_graph_limit_indexed_prefill_chunk(pos, chunk); if (chunk == 0) chunk = 1; float *chunk_logits = (i + (int)chunk >= prompt->len) ? s->logits : NULL; if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=indexed_resume pos=%u chunk=%u logits=%d\n", pos, chunk, chunk_logits ? 1 : 0); } if (!glm_graph_forward_indexed_tokens(&s->glm_graph, &e->model, &e->weights, prompt->v + i, NULL, pos, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, (uint32_t)start, (uint32_t)(i - start), (uint32_t)suffix)) { snprintf(err, errlen, "%s GLM indexed prefill failed while extending checkpoint", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } for (uint32_t j = 0; j < chunk; j++) { token_vec_push(&s->checkpoint, prompt->v[i + (int)j]); } i += (int)chunk; if (chunk_logits) { checkpoint_len_with_logits = s->checkpoint.len; s->checkpoint_valid = true; } else { s->checkpoint_valid = false; } s->mtp_draft_valid = false; if (s->progress) { s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); } } if (!glm_graph_maybe_warm_compact_indexer_after_prefill( &s->glm_graph, &e->model, &e->weights, (uint32_t)s->checkpoint.len)) { snprintf(err, errlen, "%s GLM compact indexer warmup failed", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } glm_debug_dump_prefill_logits(s->logits); return 0; } if (resumed_checkpoint && suffix > 0 && (dense_cache_gap || (uint32_t)suffix < resume_min)) { for (int i = start; i < prompt->len; i++) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return DS4_SESSION_SYNC_INTERRUPTED; } const uint32_t pos = (uint32_t)s->checkpoint.len; const bool updates_dense = glm_graph_decode_updates_dense_cache(&s->glm_graph, pos, s->logits); if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=decode_resume pos=%u token_index=%d updates_dense=%d\n", pos, i, updates_dense ? 1 : 0); } if (!glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, prompt->v[i], NULL, pos, NULL, s->logits, false)) { snprintf(err, errlen, "%s GLM decode failed while extending checkpoint", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } token_vec_push(&s->checkpoint, prompt->v[i]); s->checkpoint_valid = true; s->mtp_draft_valid = false; if (updates_dense) ds4_session_glm_note_dense_cache(s, pos, 1); if (s->progress) { s->progress(s->progress_ud, "prefill_chunk", i + 1, prompt->len); } } if (!glm_graph_maybe_warm_compact_indexer_after_prefill( &s->glm_graph, &e->model, &e->weights, (uint32_t)s->checkpoint.len)) { snprintf(err, errlen, "%s GLM compact indexer warmup failed", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } glm_debug_dump_prefill_logits(s->logits); return 0; } const uint32_t chunk_max = glm_graph_prefill_chunk_tokens(s->glm_graph.ctx_cap); const uint32_t full_work_start = (uint32_t)start; const uint32_t full_work_total = (uint32_t)prompt->len > full_work_start ? (uint32_t)prompt->len - full_work_start : 0; for (int i = start; i < prompt->len; ) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); if (!s->checkpoint_valid) { s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; } s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return DS4_SESSION_SYNC_INTERRUPTED; } if (!s->glm_graph.full_kv_cache || (uint32_t)i >= s->glm_graph.ctx_cap) { const uint32_t remaining = (uint32_t)(prompt->len - i); const bool use_token_major_prefill = glm_graph_use_streaming_token_prefill(&s->glm_graph, (uint32_t)i, remaining); const bool use_indexed_batch = !use_token_major_prefill && glm_graph_indexed_prefill_batch_ready(&s->glm_graph, (uint32_t)i); uint32_t chunk = 1; bool prefill_ok = false; if (use_token_major_prefill) { 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 *chunk_logits = (i + (int)chunk >= prompt->len) ? s->logits : NULL; if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=full_compact_token_major pos=%d chunk=%u logits=%d\n", i, chunk, chunk_logits ? 1 : 0); } prefill_ok = glm_graph_prefill_token_major(&s->glm_graph, &e->model, &e->weights, prompt->v + i, (uint32_t)i, chunk, chunk_logits, s->display_progress, s->display_progress_ud, full_work_start, (uint32_t)i - full_work_start, full_work_total); } else if (use_indexed_batch) { chunk = (uint32_t)(prompt->len - i); if (chunk > s->glm_graph.indexed_prefill_cap) { chunk = s->glm_graph.indexed_prefill_cap; } chunk = glm_graph_limit_indexed_prefill_chunk((uint32_t)i, chunk); float *chunk_logits = (i + (int)chunk >= prompt->len) ? s->logits : NULL; if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=full_indexed pos=%d chunk=%u logits=%d\n", i, chunk, chunk_logits ? 1 : 0); } prefill_ok = glm_graph_forward_indexed_tokens(&s->glm_graph, &e->model, &e->weights, prompt->v + i, NULL, (uint32_t)i, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, full_work_start, (uint32_t)i - full_work_start, full_work_total); } else { float *chunk_logits = (i + 1 >= prompt->len) ? s->logits : NULL; if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=full_decode pos=%d logits=%d\n", i, chunk_logits ? 1 : 0); } prefill_ok = glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, prompt->v[i], NULL, (uint32_t)i, NULL, chunk_logits, false); } if (!prefill_ok) { snprintf(err, errlen, "%s GLM prefill failed", backend_name); if (!s->checkpoint_valid) { s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; } s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } for (uint32_t j = 0; j < chunk; j++) { token_vec_push(&s->checkpoint, prompt->v[i + (int)j]); } i += (int)chunk; if (i >= prompt->len) { checkpoint_len_with_logits = s->checkpoint.len; s->checkpoint_valid = true; } else { s->checkpoint_valid = false; } s->mtp_draft_valid = false; if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); continue; } uint32_t chunk = (uint32_t)(prompt->len - i); const uint32_t full_remaining = s->glm_graph.ctx_cap - (uint32_t)i; if (chunk > full_remaining) chunk = full_remaining; if (chunk > chunk_max) chunk = chunk_max; float *chunk_logits = (i + (int)chunk >= prompt->len) ? s->logits : NULL; const bool use_token_major_prefill = glm_graph_use_streaming_token_prefill(&s->glm_graph, (uint32_t)i, chunk); if (sync_trace) { fprintf(stderr, "ds4: GLM sync branch=%s pos=%d chunk=%u logits=%d\n", use_token_major_prefill ? "full_token_major" : "full_dense", i, chunk, chunk_logits ? 1 : 0); } bool prefill_ok = false; if (use_token_major_prefill) { prefill_ok = glm_graph_prefill_token_major(&s->glm_graph, &e->model, &e->weights, prompt->v + i, (uint32_t)i, chunk, chunk_logits, s->display_progress, s->display_progress_ud, full_work_start, (uint32_t)i - full_work_start, full_work_total); } else { prefill_ok = glm_graph_forward_tokens(&s->glm_graph, &e->model, &e->weights, prompt->v + i, NULL, (uint32_t)i, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, full_work_start, (uint32_t)i - full_work_start, full_work_total); } if (!prefill_ok) { snprintf(err, errlen, "%s GLM prefill failed", backend_name); if (!s->checkpoint_valid) { s->checkpoint.len = checkpoint_len_with_logits; s->checkpoint_valid = checkpoint_len_with_logits > 0; } s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } for (uint32_t j = 0; j < chunk; j++) { token_vec_push(&s->checkpoint, prompt->v[i + (int)j]); } ds4_session_glm_note_dense_cache(s, (uint32_t)i, chunk); i += (int)chunk; if (i >= prompt->len) { checkpoint_len_with_logits = s->checkpoint.len; s->checkpoint_valid = true; } else { s->checkpoint_valid = false; } s->mtp_draft_valid = false; if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); } if (!glm_graph_maybe_warm_compact_indexer_after_prefill( &s->glm_graph, &e->model, &e->weights, (uint32_t)s->checkpoint.len)) { snprintf(err, errlen, "%s GLM compact indexer warmup failed", backend_name); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } glm_debug_dump_prefill_logits(s->logits); return 0; } if (s->checkpoint_valid && prompt->len >= s->checkpoint.len && ds4_tokens_starts_with(prompt, &s->checkpoint)) { s->mtp_draft_valid = false; const int suffix = prompt->len - s->checkpoint.len; const uint32_t resume_min = metal_graph_resume_prefill_min_tokens(); if (suffix > 0 && (uint32_t)suffix >= resume_min) { bool cancelled = false; ds4_sync_progress progress = { .session = s, .prompt = prompt, .user = s->progress, .user_ud = s->progress_ud, }; bool ok = metal_graph_prefill_chunked_range(&s->graph, &e->model, &e->weights, prompt, (uint32_t)s->checkpoint.len, (uint32_t)suffix, s->logits, false, ds4_session_note_prefill_progress, &progress, s->display_progress, s->display_progress_ud, NULL, ds4_session_cancelled_cb, s, &cancelled); if (cancelled) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = true; s->mtp_draft_valid = false; return DS4_SESSION_SYNC_INTERRUPTED; } if (!ok) { snprintf(err, errlen, "%s resumed prefill failed while extending checkpoint", backend_name); s->checkpoint_valid = false; return 1; } ds4_tokens_copy(&s->checkpoint, prompt); s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); return 0; } for (int i = s->checkpoint.len; i < prompt->len; i++) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = true; s->mtp_draft_valid = false; return DS4_SESSION_SYNC_INTERRUPTED; } if (!metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, (uint32_t)prompt->v[i], (uint32_t)s->checkpoint.len, s->logits)) { snprintf(err, errlen, "%s decode failed while extending checkpoint", backend_name); s->checkpoint_valid = false; return 1; } token_vec_push(&s->checkpoint, prompt->v[i]); s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); } session_greedy_splitkv_reset(s); return 0; } bool ok; s->checkpoint_valid = false; s->checkpoint.len = 0; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); if (!metal_graph_reset_prefill_state(&s->graph)) { snprintf(err, errlen, "%s prefill state reset failed", backend_name); return 1; } if (s->prefill_cap < (uint32_t)prompt->len) { bool cancelled = false; ds4_sync_progress progress = { .session = s, .prompt = prompt, .user = s->progress, .user_ud = s->progress_ud, }; ok = metal_graph_prefill_chunked(&s->graph, &e->model, &e->weights, prompt, prompt->len, s->logits, false, ds4_session_note_prefill_progress, &progress, s->display_progress, s->display_progress_ud, ds4_session_cancelled_cb, s, &cancelled); if (cancelled) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = s->checkpoint.len > 0; s->mtp_draft_valid = false; return DS4_SESSION_SYNC_INTERRUPTED; } } else { bool cancelled = false; ok = metal_graph_prefill_raw_swa(&s->graph, &e->model, &e->weights, prompt, prompt->len, s->logits, false, s->display_progress, s->display_progress_ud, ds4_session_cancelled_cb, s, &cancelled); if (cancelled) { snprintf(err, errlen, "interrupted"); return DS4_SESSION_SYNC_INTERRUPTED; } } if (!ok) { snprintf(err, errlen, "%s prefill failed", backend_name); s->checkpoint_valid = false; return 1; } ds4_tokens_copy(&s->checkpoint, prompt); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); s->graph.mtp_n_raw = 0; session_greedy_splitkv_reset(s); return 0; } #endif } /* Return true when canonicalization would replace already-sampled tokens. * * A DS4 session checkpoint is more than a token vector: the backend state also * contains raw SWA rows, compressed KV rows, indexer rows, and compressor * frontiers. Replacing any part of the live tail requires restoring that whole * frontier first. Extending exactly at the live end is safe; rewriting behind * it is not an in-place operation. */ bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common) { if (live_len < 0 || canonical_len < 0 || common < 0) return true; if (common > live_len || common > canonical_len) return true; return common < live_len; } /* Replace the live suffix after a shared prefix. * * This is used after parsing a generated tool call. The model may have emitted * DSML in an order that is semantically valid but not byte-for-byte equal to the * canonical prompt we will see on the next request. Rewriting only the token * checkpoint is not enough: the backend still contains raw and compressed rows * for the old suffix. Until we have a real frontier snapshot at the * rewrite point, any replacement behind the live end reports that a rebuild is * needed without mutating the session. The server may still find an older disk KV * checkpoint before falling back to a full replay. */ ds4_session_rewrite_result ds4_session_rewrite_from_common( ds4_session *s, const ds4_tokens *prompt, int common, char *err, size_t errlen) { if (!s || !prompt) { snprintf(err, errlen, "missing session or prompt"); return DS4_SESSION_REWRITE_ERROR; } if (prompt->len <= 0) { snprintf(err, errlen, "empty prompt"); return DS4_SESSION_REWRITE_ERROR; } if (prompt->len >= s->ctx_size) { snprintf(err, errlen, "prompt length %d exceeds context %d (one token of generation room is required)", prompt->len, s->ctx_size); return DS4_SESSION_REWRITE_ERROR; } if (!s->checkpoint_valid) { snprintf(err, errlen, "session has no valid checkpoint"); return DS4_SESSION_REWRITE_ERROR; } if (common < 0 || common > s->checkpoint.len || common > prompt->len) { snprintf(err, errlen, "invalid rewrite prefix"); return DS4_SESSION_REWRITE_ERROR; } for (int i = 0; i < common; i++) { if (s->checkpoint.v[i] != prompt->v[i]) { snprintf(err, errlen, "rewrite prefix does not match live checkpoint"); return DS4_SESSION_REWRITE_ERROR; } } if (common == s->checkpoint.len) { return ds4_session_sync(s, prompt, err, errlen) == 0 ? DS4_SESSION_REWRITE_OK : DS4_SESSION_REWRITE_ERROR; } if (ds4_session_rewrite_requires_rebuild(s->checkpoint.len, prompt->len, common)) { snprintf(err, errlen, "rewrite needs rebuild: common=%d live=%d canonical=%d", common, s->checkpoint.len, prompt->len); return DS4_SESSION_REWRITE_REBUILD_NEEDED; } snprintf(err, errlen, "unexpected canonical rewrite state"); return DS4_SESSION_REWRITE_ERROR; } int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt) { if (!s->checkpoint_valid) return 0; int n = s->checkpoint.len < prompt->len ? s->checkpoint.len : prompt->len; int i = 0; while (i < n && s->checkpoint.v[i] == prompt->v[i]) i++; return i; } int ds4_session_argmax(ds4_session *s) { return sample_argmax(s->logits, DS4_N_VOCAB); } int ds4_session_argmax_excluding(ds4_session *s, int excluded_id) { if (!s || !s->logits) return -1; int best = -1; float best_logit = DS4_NEG_INF; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { if ((int)i == excluded_id) continue; const float v = s->logits[i]; if (best < 0 || v > best_logit) { best = (int)i; best_logit = v; } } return best; } int ds4_sample_logits(const float *logits, int n_vocab, float temperature, int top_k, float top_p, float min_p, uint64_t *rng) { if (!logits || n_vocab <= 0) return 0; float *scratch = xmalloc((size_t)n_vocab * sizeof(scratch[0])); const int token = sample_top_p_min_p(logits, (uint32_t)n_vocab, temperature, top_k, top_p, min_p, rng, scratch); free(scratch); return token; } int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng) { return sample_top_p_min_p(s->logits, DS4_N_VOCAB, temperature, top_k, top_p, min_p, rng, s->sample_probs); } int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k) { if (!s || !out || k <= 0) return 0; if (k > (int)DS4_N_VOCAB) k = (int)DS4_N_VOCAB; for (int i = 0; i < k; i++) { out[i].id = -1; out[i].logit = DS4_NEG_INF; out[i].logprob = DS4_NEG_INF; } float max_logit = DS4_NEG_INF; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { const float v = s->logits[i]; if (!isfinite(v)) continue; if (v > max_logit) max_logit = v; for (int j = 0; j < k; j++) { if (out[j].id < 0 || v > out[j].logit) { for (int l = k - 1; l > j; l--) out[l] = out[l - 1]; out[j].id = (int)i; out[j].logit = v; break; } } } if (!isfinite(max_logit)) return 0; double sum = 0.0; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { const float v = s->logits[i]; if (isfinite(v)) sum += exp((double)v - (double)max_logit); } const double logsum = (double)max_logit + log(sum); for (int i = 0; i < k && out[i].id >= 0; i++) { out[i].logprob = isfinite(out[i].logit) ? (float)((double)out[i].logit - logsum) : DS4_NEG_INF; } return k; } int ds4_session_token_logprob(ds4_session *s, int token, ds4_token_score *out) { if (!s || !out || token < 0 || token >= (int)DS4_N_VOCAB) return 0; float max_logit = DS4_NEG_INF; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { const float v = s->logits[i]; if (isfinite(v) && v > max_logit) max_logit = v; } if (!isfinite(max_logit)) return 0; double sum = 0.0; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { const float v = s->logits[i]; if (isfinite(v)) sum += exp((double)v - (double)max_logit); } const double logsum = (double)max_logit + log(sum); out->id = token; out->logit = s->logits[token]; out->logprob = isfinite(out->logit) ? (float)((double)out->logit - logsum) : DS4_NEG_INF; return 1; } int ds4_session_copy_logits(ds4_session *s, float *out, int cap) { if (!s || !out || cap < (int)DS4_N_VOCAB) return 0; memcpy(out, s->logits, (size_t)DS4_N_VOCAB * sizeof(out[0])); return (int)DS4_N_VOCAB; } int ds4_session_set_logits(ds4_session *s, const float *logits, int n) { if (!s || !logits || n != (int)DS4_N_VOCAB) return 1; memcpy(s->logits, logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); return 0; } /* Pay the one-time first-submission GPU cost (pipeline ramp plus model-heap * residency for the batched prefill kernels) outside any measured window. * The TP worker calls this right after creating its session: it otherwise * encodes no main-queue GPU work until the first mirrored sync arrives, so * the cost lands inside the leader-timed prefill (measured ~1.1 s per run * on the M5 Max pair -- the whole first-run TP deficit vs single * node, which pays the same cost before its timing window starts). */ void ds4_session_gpu_warmup(ds4_session *s) { #ifndef DS4_NO_GPU if (!s || ds4_session_is_cpu(s)) return; if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) return; if (!metal_graph_batch_hc_mix(&s->graph) || !metal_graph_batch_flat_hc(&s->graph)) return; if (!s->engine->weights.layer[0].hc_attn_fn) return; (void)metal_graph_warmup_prefill_kernels(&s->graph, &s->engine->model, &s->engine->weights, 32); #else (void)s; #endif } #ifndef DS4_NO_GPU static bool ds4_session_dspark_capture_current(const ds4_session *s) { if (!s || ds4_session_is_cpu(s) || !s->checkpoint_valid) return false; const ds4_gpu_graph *g = &s->graph; return g->dspark_capture_enabled && g->dspark_capture_valid && g->dspark_capture_checkpoint_len == (uint32_t)s->checkpoint.len; } static bool ds4_session_dspark_capture_batch_current(const ds4_session *s) { if (!s || ds4_session_is_cpu(s) || !s->checkpoint_valid) return false; const ds4_gpu_graph *g = &s->graph; if (!g->dspark_capture_enabled || !g->dspark_capture_batch_valid) { return false; } if (g->dspark_capture_batch_tokens == 0 || g->dspark_capture_batch_start > (uint32_t)s->checkpoint.len) { return false; } const uint32_t batch_end = g->dspark_capture_batch_start + g->dspark_capture_batch_tokens; return batch_end >= g->dspark_capture_batch_start && batch_end <= (uint32_t)s->checkpoint.len; } static void ds4_session_note_legacy_mtp_probe(ds4_session *s, int token, bool mtp_probe_log) { if (!s || !s->mtp_draft_valid) return; if (mtp_probe_log) { s->mtp_probe_total++; if (s->mtp_draft_token == token) s->mtp_probe_hit++; fprintf(stderr, "ds4: mtp probe token=%d draft=%d hit=%llu/%llu\n", token, s->mtp_draft_token, (unsigned long long)s->mtp_probe_hit, (unsigned long long)s->mtp_probe_total); } s->mtp_draft_valid = false; } static bool ds4_session_prepare_legacy_mtp_draft(ds4_session *s, int token, uint32_t pos, bool mtp_probe_log) { ds4_engine *e = s->engine; if (!e->mtp_ready || !s->mtp_logits || (e->mtp_draft_tokens <= 1 && !mtp_probe_log)) { return false; } int mtp_top = -1; if (metal_graph_eval_mtp_draft(&s->graph, &e->model, &e->weights, &e->mtp_model, &e->mtp_weights, token, pos, getenv("DS4_MTP_FULL_LOGITS") ? s->mtp_logits : NULL, &mtp_top)) { s->mtp_draft_token = mtp_top >= 0 ? mtp_top : sample_argmax(s->mtp_logits, DS4_N_VOCAB); s->mtp_draft_valid = true; return true; } if (mtp_probe_log) { fprintf(stderr, "ds4: mtp probe draft failed\n"); } return false; } static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s, int token, uint32_t pos) { const char *probe = getenv("DS4_DSPARK_PROBE"); const bool probe_log = probe && probe[0]; const bool enabled = s->engine->dspark; const char *fake_argmax = getenv("DS4_DSPARK_FAKE_ARGMAX_PROPOSAL"); const bool fake_argmax_enabled = enabled && fake_argmax && fake_argmax[0] && strcmp(fake_argmax, "0") != 0; const float confidence_threshold = s->engine->dspark_confidence_threshold; const bool stats_enabled = ds4_dspark_stats_enabled(); const bool scheduler_enabled = ds4_dspark_scheduler_enabled(); const bool time_enabled = stats_enabled || scheduler_enabled; const double stats_t0 = time_enabled ? now_sec() : 0.0; #define DS4_DSPARK_PROP_T0() (stats_enabled ? now_sec() : 0.0) #define DS4_DSPARK_PROP_ADD(field_, t0_) do { \ if (stats_enabled) { \ s->dspark_stats.field_ += (now_sec() - (t0_)) * 1000.0; \ } \ } while (0) s->dspark_draft_valid = false; s->dspark_draft_len = 0; s->dspark_last_confidence0 = 0.0f; s->dspark_last_confidence0_valid = false; if (scheduler_enabled) s->dspark_last_propose_ms = 0.0; if (enabled && !fake_argmax_enabled && ds4_session_dspark_scheduler_should_skip(s)) { (void)metal_graph_dspark_ring_maintain(&s->graph, &s->engine->mtp_model, &s->engine->dspark_weights, pos); const double propose_ms = time_enabled ? (now_sec() - stats_t0) * 1000.0 : 0.0; if (scheduler_enabled) s->dspark_last_propose_ms = propose_ms; if (stats_enabled) { s->dspark_stats.propose_ms += propose_ms; } return false; } if (probe_log || enabled) { const bool capture_ok = ds4_session_dspark_capture_current(s); const bool batch_capture_ok = ds4_session_dspark_capture_batch_current(s); const ds4_dspark_weights *dw = &s->engine->dspark_weights; const bool stage0_ready = dspark_stage0_weights_ready(&s->graph, dw); const bool runtime_fused_stage0_setup = enabled && !fake_argmax_enabled && !probe_log; const double stage0_t0 = DS4_DSPARK_PROP_T0(); bool stage0_ok = false; if (capture_ok && stage0_ready && !runtime_fused_stage0_setup) { stage0_ok = metal_graph_eval_dspark_stage0(&s->graph, &s->engine->mtp_model, dw); } else if (capture_ok && stage0_ready && runtime_fused_stage0_setup) { stage0_ok = true; } DS4_DSPARK_PROP_ADD(propose_stage0_ms, stage0_t0); const double setup_t0 = DS4_DSPARK_PROP_T0(); const bool draft_block_ready = dspark_draft_block_ready(&s->graph, &s->engine->weights, dw, token); const bool stage_input_ready = dspark_stage_input_ready(&s->graph, dw); bool stage_input_ok = false; if (stage0_ok && draft_block_ready && stage_input_ready) { if (runtime_fused_stage0_setup) { stage_input_ok = metal_graph_prepare_dspark_stage0_setup_block( &s->graph, &s->engine->model, &s->engine->weights, &s->engine->mtp_model, dw, token, pos); stage0_ok = stage_input_ok; } else { stage_input_ok = metal_graph_prepare_dspark_setup_block(&s->graph, &s->engine->model, &s->engine->weights, dw, token, pos); } } DS4_DSPARK_PROP_ADD(propose_setup_ms, setup_t0); const bool draft_cache_ready = dspark_stage_cache_ready(&s->graph, dw); uint32_t initial_cache_rows = 0; const uint64_t captured_batch_end = (uint64_t)s->graph.dspark_capture_batch_start + s->graph.dspark_capture_batch_tokens; const bool captured_batch_end_ok = s->graph.dspark_capture_batch_tokens != 0 && captured_batch_end <= UINT32_MAX; const bool initial_cache_ready = batch_capture_ok && draft_cache_ready && captured_batch_end_ok && (uint32_t)captured_batch_end == pos; const double cache_t0 = DS4_DSPARK_PROP_T0(); bool initial_cache_ok = false; if (initial_cache_ready) { initial_cache_ok = metal_graph_seed_dspark_initial_cache_from_prefill( &s->graph, &s->engine->mtp_model, dw, s->graph.dspark_capture_batch_start, s->graph.dspark_capture_batch_tokens, &initial_cache_rows); } bool cache_window_ok = draft_cache_ready; if (draft_cache_ready) { if (initial_cache_ready) { if (!initial_cache_ok) metal_graph_dspark_cache_reset(&s->graph); cache_window_ok = initial_cache_ok; } else { cache_window_ok = metal_graph_dspark_cache_crop_to_prefix(&s->graph, pos) && metal_graph_dspark_cache_ends_at(&s->graph, pos); } } DS4_DSPARK_PROP_ADD(propose_cache_ms, cache_t0); const bool noncausal_attn_ready = probe_log && stage_input_ok && draft_cache_ready && dspark_noncausal_attention_probe_ready(&s->graph, dw); const bool noncausal_attn_ok = noncausal_attn_ready && metal_graph_probe_dspark_noncausal_attention(&s->graph, &s->engine->mtp_model, dw); const bool stage_block_ready = stage_input_ok && draft_cache_ready && dspark_stage_block_ready(&s->graph, dw, 0); uint32_t stage_chain_done = 0; const bool stage_chain_ready = stage_input_ok && draft_cache_ready && cache_window_ok && dw->n_stages != 0 && dw->n_stages <= DS4_DSPARK_MAX_STAGES; uint32_t stage_cache_start = 0; uint32_t stage_cache_rows = 0; const double chain_t0 = DS4_DSPARK_PROP_T0(); bool stage_chain_ok = false; if (stage_chain_ready) { stage_chain_ok = metal_graph_eval_dspark_stage_chain(&s->graph, &s->engine->mtp_model, dw, pos, &stage_chain_done, &stage_cache_start, &stage_cache_rows); } DS4_DSPARK_PROP_ADD(propose_chain_ms, chain_t0); const bool stage_block_ok = stage_chain_done >= 1u; const bool base_logits_ready = stage_chain_ok && dspark_final_head_ready(&s->graph, &s->engine->weights, dw); int32_t markov_proposal[DS4_DSPARK_MAX_BLOCK_SIZE]; for (uint32_t i = 0; i < DS4_DSPARK_MAX_BLOCK_SIZE; i++) { markov_proposal[i] = -1; } uint32_t markov_proposal_len = 0; bool markov_ok = false; uint32_t confidence_len = 0; uint32_t confidence_prefix_len = 0; float confidence0 = 0.0f; bool confidence_ok = false; bool reuse_confidence0_markov = false; const bool runtime_confidence_precheck = base_logits_ready && !probe_log && confidence_threshold > 0.0f && dspark_confidence_probe_ready(dw); bool base_logits_ok = false; if (runtime_confidence_precheck) { const double hidden_t0 = DS4_DSPARK_PROP_T0(); const bool hidden_ok = metal_graph_eval_dspark_final_hidden(&s->graph, &s->engine->mtp_model, dw); DS4_DSPARK_PROP_ADD(propose_hidden_ms, hidden_t0); bool conf0_ok = false; if (hidden_ok) { const double conf0_t0 = DS4_DSPARK_PROP_T0(); conf0_ok = dspark_eval_confidence0_runtime(&s->graph, &s->engine->mtp_model, dw, token, s->dspark_conf_features, s->dspark_conf_features_cap, &confidence0); DS4_DSPARK_PROP_ADD(propose_conf0_ms, conf0_t0); } if (conf0_ok) { confidence_ok = true; confidence_len = 1; if (sigmoid_stable(confidence0) >= confidence_threshold) { confidence_prefix_len = 1; const double logits_t0 = DS4_DSPARK_PROP_T0(); base_logits_ok = metal_graph_eval_dspark_base_logits_from_hidden( &s->graph, &s->engine->model, &s->engine->weights, dw); DS4_DSPARK_PROP_ADD(propose_logits_ms, logits_t0); reuse_confidence0_markov = base_logits_ok && !dspark_disable_reuse_confidence0_markov(); } } } else if (base_logits_ready) { const double logits_t0 = DS4_DSPARK_PROP_T0(); base_logits_ok = metal_graph_eval_dspark_base_logits(&s->graph, &s->engine->model, &s->engine->weights, &s->engine->mtp_model, dw); DS4_DSPARK_PROP_ADD(propose_logits_ms, logits_t0); } const bool markov_ready = base_logits_ok && dspark_markov_probe_ready(dw); const bool lazy_runtime_confidence = markov_ready && !probe_log && confidence_threshold > 0.0f; if (lazy_runtime_confidence) { const double markov_t0 = DS4_DSPARK_PROP_T0(); markov_ok = dspark_apply_markov_confidence_lazy_runtime( &s->graph, &s->engine->mtp_model, dw, token, confidence_threshold, s->spec_row_logits, s->dspark_markov_bias, s->dspark_conf_features, s->dspark_conf_features_cap, markov_proposal, &markov_proposal_len, &confidence_len, &confidence_prefix_len, reuse_confidence0_markov, &confidence0); DS4_DSPARK_PROP_ADD(propose_markov_ms, markov_t0); confidence_ok = markov_ok; } else if (markov_ready) { const uint64_t logits_count = (uint64_t)dw->block_size * (uint64_t)DS4_N_VOCAB; if (logits_count != 0 && logits_count <= (uint64_t)SIZE_MAX / sizeof(float)) { const double markov_t0 = DS4_DSPARK_PROP_T0(); const uint64_t logits_bytes = logits_count * sizeof(float); float *logits = xmalloc((size_t)logits_bytes); float *markov_state = xmalloc((size_t)dw->markov_rank * sizeof(markov_state[0])); float *markov_bias = xmalloc((size_t)DS4_N_VOCAB * sizeof(markov_bias[0])); markov_ok = ds4_gpu_tensor_read(s->graph.spec_logits, 0, logits, logits_bytes) != 0 && dspark_apply_markov_greedy_probe(logits, &s->engine->mtp_model, dw, token, markov_state, markov_bias, markov_proposal, &markov_proposal_len); if (markov_ok && probe_log) { /* Runtime only needs the greedy proposal. Keep the GPU * writeback for probe mode, where spec_logits may be * inspected after Markov correction. */ markov_ok = ds4_gpu_tensor_write(s->graph.spec_logits, 0, logits, logits_bytes) != 0; } free(markov_bias); free(markov_state); free(logits); DS4_DSPARK_PROP_ADD(propose_markov_ms, markov_t0); } } const bool confidence_ready = !lazy_runtime_confidence && markov_ok && dspark_confidence_probe_ready(dw); if (confidence_ready) { const uint64_t hidden_count = (uint64_t)dw->block_size * (uint64_t)DS4_N_EMBD; const uint64_t feature_count = (uint64_t)DS4_N_EMBD + (uint64_t)dw->markov_rank; if (hidden_count != 0 && hidden_count <= (uint64_t)SIZE_MAX / sizeof(float) && feature_count <= (uint64_t)SIZE_MAX / sizeof(float)) { const double confidence_t0 = DS4_DSPARK_PROP_T0(); const uint64_t hidden_bytes = hidden_count * sizeof(float); float *hidden_rows = xmalloc((size_t)hidden_bytes); float *markov_state = xmalloc((size_t)dw->markov_rank * sizeof(markov_state[0])); float *features = xmalloc((size_t)feature_count * sizeof(features[0])); float *confidence_logits = xmalloc((size_t)dw->block_size * sizeof(confidence_logits[0])); confidence_ok = ds4_gpu_tensor_read(metal_graph_batch_ffn_norm(&s->graph), 0, hidden_rows, hidden_bytes) != 0 && dspark_eval_confidence_probe(confidence_logits, hidden_rows, &s->engine->mtp_model, dw, token, markov_proposal, markov_state, features, &confidence_len); if (confidence_ok && confidence_len != 0) { confidence0 = confidence_logits[0]; confidence_prefix_len = dspark_confident_prefix_len(confidence_logits, confidence_len, confidence_threshold); } free(confidence_logits); free(features); free(markov_state); free(hidden_rows); DS4_DSPARK_PROP_ADD(propose_confidence_ms, confidence_t0); } } if (markov_ok && markov_proposal_len != 0) { uint32_t proposal_len = markov_proposal_len; if (confidence_threshold > 0.0f) { proposal_len = confidence_ok ? confidence_prefix_len : 0; } s->dspark_draft_len = proposal_len; if (s->dspark_draft_len > DS4_DSPARK_MAX_BLOCK_SIZE) { s->dspark_draft_len = DS4_DSPARK_MAX_BLOCK_SIZE; } for (uint32_t i = 0; i < s->dspark_draft_len; i++) { s->dspark_draft_tokens[i] = markov_proposal[i]; } s->dspark_draft_valid = s->dspark_draft_len != 0; } if (confidence_ok && confidence_len != 0) { s->dspark_last_confidence0 = confidence0; s->dspark_last_confidence0_valid = true; } bool fake_argmax_ok = false; if (!s->dspark_draft_valid && fake_argmax_enabled) { s->dspark_draft_tokens[0] = sample_argmax(s->logits, DS4_N_VOCAB); s->dspark_draft_len = 1; s->dspark_draft_valid = true; fake_argmax_ok = true; } if (probe_log) { const char *stage0_status = stage0_ok ? "ok" : !capture_ok ? "capture-invalid" : !stage0_ready ? "unavailable" : "failed"; const char *draft_block_status = stage_input_ok ? "ok" : !stage0_ok ? "stage0-not-ready" : !draft_block_ready ? "unavailable" : !stage_input_ready ? "stage-input-not-ready" : "failed"; const char *stage_input_status = stage_input_ok ? "ok" : !stage0_ok ? "stage0-not-ready" : !draft_block_ready ? "draft-block-not-ready" : !stage_input_ready ? "unavailable" : "failed"; const char *draft_cache_status = draft_cache_ready ? "ok" : "unavailable"; const char *cache_window_status = !draft_cache_ready ? "draft-cache-not-ready" : cache_window_ok ? "ok" : initial_cache_ready ? "seed-failed" : "invalid"; const char *initial_cache_status = initial_cache_ok ? "ok" : !batch_capture_ok ? "batch-capture-invalid" : !draft_cache_ready ? "draft-cache-not-ready" : !initial_cache_ready ? "unsupported-window" : "failed"; const char *noncausal_attn_status = noncausal_attn_ok ? "ok" : !stage_input_ok ? "stage-input-not-ready" : !draft_cache_ready ? "draft-cache-not-ready" : !noncausal_attn_ready ? "unavailable" : "failed"; const char *stage_block_status = stage_block_ok ? "ok" : !stage_input_ok ? "stage-input-not-ready" : !draft_cache_ready ? "draft-cache-not-ready" : !stage_block_ready ? "unavailable" : "failed"; const char *stage_chain_status = stage_chain_ok ? "ok" : !stage_input_ok ? "stage-input-not-ready" : !draft_cache_ready ? "draft-cache-not-ready" : !cache_window_ok ? "cache-window-not-ready" : !stage_chain_ready ? "unavailable" : stage_chain_done != 0 ? "partial" : "failed"; const char *base_logits_status = base_logits_ok ? "ok" : !stage_chain_ok ? "stage-chain-not-ready" : !base_logits_ready ? "unavailable" : "failed"; const char *markov_status = markov_ok ? "ok" : !base_logits_ok ? "base-logits-not-ready" : !markov_ready ? "unavailable" : "failed"; const char *confidence_status = confidence_ok ? "ok" : !markov_ok ? "markov-not-ready" : !confidence_ready ? "unavailable" : "failed"; const char *fake_argmax_status = fake_argmax_ok ? "ok" : !fake_argmax_enabled ? "disabled" : s->dspark_draft_valid ? "real-proposal-present" : "skipped"; fprintf(stderr, "ds4: DSpark proposer pending token=%d pos=%u capture=%s " "batch_capture=%s batch_start=%u batch_tokens=%u " "stage0=%s draft_block=%s stage_input=%s draft_cache=%s " "initial_cache=%s initial_cache_rows=%u " "cache_window=%s cache_token_start=%u " "cache_raw_start=%u cache_len=%u " "noncausal_attn=%s stage_block=%s stage_chain=%s " "stage_chain_done=%u/%u stage_raw_start=%u cache_rows=%u " "base_logits=%s markov=%s " "proposal_len=%u proposal0=%d fake_argmax=%s confidence=%s " "confidence_len=%u confidence_prefix=%u " "confidence_threshold=%.3f confidence0=%.3f block=%u " "tensors=%u missing=%u invalid=%u metadata_errors=%u\n", token, pos, capture_ok ? "current" : "invalid", batch_capture_ok ? "current" : "invalid", s->graph.dspark_capture_batch_start, s->graph.dspark_capture_batch_tokens, stage0_status, draft_block_status, stage_input_status, draft_cache_status, initial_cache_status, initial_cache_rows, cache_window_status, s->graph.dspark_cache_token_start, s->graph.dspark_cache_start, s->graph.dspark_cache_len, noncausal_attn_status, stage_block_status, stage_chain_status, stage_chain_done, dw->n_stages, stage_cache_start, stage_cache_rows, base_logits_status, markov_status, s->dspark_draft_len, s->dspark_draft_len ? s->dspark_draft_tokens[0] : -1, fake_argmax_status, confidence_status, confidence_len, confidence_prefix_len, confidence_threshold, confidence0, dw->block_size, dw->present_tensors, dw->missing_tensors, dw->invalid_tensors, dw->metadata_errors); } } if (time_enabled) { const double propose_ms = (now_sec() - stats_t0) * 1000.0; if (scheduler_enabled) s->dspark_last_propose_ms = propose_ms; if (stats_enabled) s->dspark_stats.propose_ms += propose_ms; } #undef DS4_DSPARK_PROP_ADD #undef DS4_DSPARK_PROP_T0 return s->dspark_draft_valid; } /* Support weights and DSpark scratch live on the configured executor tier. */ static bool ds4_session_prepare_dspark_draft(ds4_session *s, int token, uint32_t pos) { ds4_gpu_graph *g = &s->graph; const int exec_tier = g->dspark_exec_tier >= 0 && g->dspark_exec_tier < DS4_MAX_GPUS ? g->dspark_exec_tier : 0; const int saved_tier = g->active_tier; if (exec_tier != saved_tier) { if (ds4_gpu_set_current_device(exec_tier) != 0) return false; g->active_tier = exec_tier; } const bool ok = ds4_session_prepare_dspark_draft_impl(s, token, pos); if (exec_tier != saved_tier) { g->active_tier = saved_tier; if (ds4_gpu_set_current_device(saved_tier) != 0) return false; } return ok; } static void ds4_session_prepare_support_draft(ds4_session *s, int token, uint32_t pos, bool probe_mtp, bool mtp_probe_log) { if (!s || !probe_mtp || !s->engine) return; switch (s->engine->support_kind) { case DS4_SUPPORT_MTP_LEGACY: (void)ds4_session_prepare_legacy_mtp_draft(s, token, pos, mtp_probe_log); break; case DS4_SUPPORT_DSPARK: (void)ds4_session_prepare_dspark_draft(s, token, pos); break; case DS4_SUPPORT_NONE: default: break; } } #endif static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, char *err, size_t errlen) { if (!s) return 1; if (s->distributed) { if (!s->checkpoint_valid) { if (errlen) snprintf(err, errlen, "distributed decode requires a valid checkpoint"); return 1; } (void)probe_mtp; return ds4_dist_session_eval(s->distributed, s, &s->checkpoint, token, s->logits, err, errlen); } if (ds4_session_is_cpu(s)) { ds4_engine *e = s->engine; forward_token_raw_swa_cpu_decode_scratch(s->logits, &e->model, &e->weights, &s->cpu_cache, token, (uint32_t)s->checkpoint.len, e->directional_steering_dirs, e->directional_steering_attn_scale, e->directional_steering_ffn_scale, &s->cpu_scratch); token_vec_push(&s->checkpoint, token); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); (void)probe_mtp; return 0; } #ifdef DS4_NO_GPU (void)s; (void)token; (void)probe_mtp; snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else ds4_engine *e = s->engine; if (ds4_session_is_glm(s)) { /* 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. */ if (!s->glm_spec_inside && s->glm_graph_ready && e->glm_mtp && DS4_N_NEXTN_PREDICT != 0 && e->tp.active && e->tp.rank != 0) { int acc[2]; const int rc = ds4_session_glm_spec_cycle(s, token, acc, 2, err, errlen); (void)probe_mtp; return rc < 0 ? 1 : 0; } if (!s->glm_graph_ready) { if (errlen) snprintf(err, errlen, "%s GLM graph is not initialized", ds4_backend_name(e->backend)); return 1; } if ((uint32_t)s->checkpoint.len >= s->glm_graph.ctx_size) { if (errlen) snprintf(err, errlen, "GLM Metal context reached (%u)", s->glm_graph.ctx_size); return 1; } const uint32_t pos = (uint32_t)s->checkpoint.len; const bool updates_dense = glm_graph_decode_updates_dense_cache(&s->glm_graph, pos, s->logits); if (!glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, token, NULL, pos, NULL, s->logits, false)) { if (errlen) snprintf(err, errlen, "%s GLM decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return 1; } token_vec_push(&s->checkpoint, token); s->checkpoint_valid = true; s->mtp_draft_valid = false; if (updates_dense) ds4_session_glm_note_dense_cache(s, pos, 1); /* MTP acceptance probe (timing/quality only, output untouched): * after each committed token, draft the token two positions ahead * from the nextn block and score it against the next greedy pick. * Both TP ranks must set the env (the draft's routed experts ride * a big-gate exchange). */ if (getenv("DS4_GLM_MTP_PROBE") && DS4_N_NEXTN_PREDICT != 0) { static int probe_hits, probe_total, probe_have, probe_draft; static uint32_t probe_min_pos = UINT32_MAX; int nmax = 0; float nbest = s->logits[0]; for (uint32_t i = 1; i < DS4_N_VOCAB; i++) { if (s->logits[i] > nbest) { nbest = s->logits[i]; nmax = (int)i; } } if (probe_have) { probe_total++; probe_hits += (probe_draft == nmax); if ((probe_total % 16) == 0) { fprintf(stderr, "ds4: glm mtp probe: %d/%d hits (%.1f%%)\n", probe_hits, probe_total, 100.0 * probe_hits / probe_total); } } if (probe_min_pos == UINT32_MAX) probe_min_pos = pos; int draft = -1; if (glm_graph_mtp_step(&s->glm_graph, &e->model, &e->weights, nmax, pos, probe_min_pos, &draft)) { probe_draft = draft; probe_have = 1; } else { probe_have = 0; fprintf(stderr, "ds4: glm mtp probe: draft step failed at pos %u\n", pos); } } (void)probe_mtp; return 0; } const bool mtp_probe_log = getenv("DS4_MTP_PROBE") != NULL; if (probe_mtp && e->support_kind == DS4_SUPPORT_MTP_LEGACY) { ds4_session_note_legacy_mtp_probe(s, token, mtp_probe_log); } const bool dspark_target_timing = e->support_kind == DS4_SUPPORT_DSPARK && (ds4_dspark_stats_enabled() || ds4_dspark_scheduler_enabled()); const double target_t0 = dspark_target_timing ? now_sec() : 0.0; if (!metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, (uint32_t)token, (uint32_t)s->checkpoint.len, s->logits)) { snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return 1; } if (dspark_target_timing) { const double target_ms = (now_sec() - target_t0) * 1000.0; s->dspark_last_target_eval_ms = target_ms; if (ds4_dspark_stats_enabled()) { s->dspark_stats.target_ms += target_ms; } } token_vec_push(&s->checkpoint, token); s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); ds4_session_prepare_support_draft(s, token, (uint32_t)(s->checkpoint.len - 1), probe_mtp, mtp_probe_log); return 0; #endif } /* 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 * the leader/worker lockstep survives every eval entry point. */ static int ds4_session_eval_probe_tp(ds4_session *s, int token, bool probe_mtp, char *err, size_t errlen) { if (ds4_session_tp_leader(s)) { ds4_engine *e = s->engine; if (!ds4_tp_send_eval(e->tp.ctx, s->tp_session_id, ++e->tp.eval_seq, token)) { snprintf(err, errlen, "tp: worker eval send failed"); return 1; } } int rc = ds4_session_eval_internal(s, token, probe_mtp, err, errlen); if (rc != 0 && ds4_session_tp_leader(s)) { ds4_session_invalidate(s); return rc; } #if !defined(DS4_NO_GPU) && defined(__APPLE__) if (rc == 0 && s->engine && s->engine->tp.active && ds4_gpu_tp_failed()) { snprintf(err, errlen, "tp: gate transport failed"); if (ds4_session_tp_leader(s)) ds4_session_invalidate(s); return 1; } #endif /* Vocab-split head: merge the halves after every eval (DS4 only). */ if (rc == 0 && s->engine && s->engine->tp.active && s->engine->tp.vocab_split) { const uint32_t vhalf = (uint32_t)DS4_N_VOCAB / 2u; if (s->engine->tp.rank == 0) { if (!ds4_tp_recv_logits_half(s->engine->tp.ctx, s->logits + vhalf, vhalf)) { snprintf(err, errlen, "tp: worker logits half missing"); ds4_session_invalidate(s); return 1; } } else { if (!ds4_tp_send_logits_half(s->engine->tp.ctx, s->logits + vhalf, vhalf)) { snprintf(err, errlen, "tp: logits half send failed"); return 1; } } } return rc; } int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen) { bool probe_mtp = true; #ifndef DS4_NO_GPU if (s && s->engine && s->engine->support_kind == DS4_SUPPORT_DSPARK) { probe_mtp = false; } #endif return ds4_session_eval_probe_tp(s, token, probe_mtp, err, errlen); } #ifndef DS4_NO_GPU static bool ds4_sessions_eval_batch_metal_supported( ds4_decode_item *items, int count, ds4_engine *e) { const char *tp_batch = getenv("DS4_METAL_TP_SESSION_BATCH"); if (!items || count < 2 || !e || e->backend != DS4_BACKEND_METAL || e->support_kind != DS4_SUPPORT_NONE || (e->tp.active && tp_batch && strcmp(tp_batch, "0") == 0) || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL || getenv("DS4_METAL_DECODE_STAGE_PROFILE") != NULL) { return false; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e || s->distributed || ds4_session_is_cpu(s) || !s->checkpoint_valid) { return false; } if (ds4_session_is_glm(s)) { if (!s->glm_graph_ready || s->glm_graph.ssd_streaming || glm_debug_hidden_dump_layer() >= 0 || e->glm_mtp || getenv("DS4_GLM_MTP_PROBE") != NULL) { return false; } } else if (s->graph.ssd_streaming) { return false; } } return true; } static bool metal_graph_native_session_batch_shared_supported( ds4_decode_item *items, int count, const ds4_engine *e) { const char *enabled = getenv("DS4_METAL_SESSION_BATCH_SHARED"); if ((enabled && enabled[0] && strcmp(enabled, "0") == 0) || !items || count < 2 || !e || e->tp.active || e->support_kind != DS4_SUPPORT_NONE || metal_graph_use_reference_shared_down_hc() || metal_graph_use_q4_selected_shared_overlap() || metal_graph_use_pro_q4_cpu_router() || getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (first->placement || first->ssd_streaming || first->quality || !first->shared_gate_up_swiglu_fuse || (uint32_t)count > first->prefill_cap || !metal_graph_batch_ffn_norm(first) || !metal_graph_batch_shared_gate(first) || !metal_graph_batch_shared_up(first) || !metal_graph_batch_shared_mid(first)) { return false; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &e->weights.layer[il]; 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) { return false; } } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || ds4_session_is_glm(s) || s->graph.placement || s->graph.ssd_streaming || s->graph.quality || s->graph.tp_world >= 2 || !s->graph.shared_gate_up_swiglu_fuse || s->graph.materialize_ffn_out || metal_graph_directional_steering_attn_enabled(&s->graph) || metal_graph_directional_steering_ffn_enabled(&s->graph)) { return false; } } return true; } static bool metal_graph_native_session_batch_qkv_supported( ds4_decode_item *items, int count, const ds4_engine *e) { const char *enabled = getenv("DS4_METAL_SESSION_BATCH_QKV"); if ((enabled && enabled[0] && strcmp(enabled, "0") == 0) || !items || count < 2 || !e || e->tp.active || metal_graph_use_reference_qkv_norm() || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (!metal_graph_batch_attn_norm(first) || !metal_graph_batch_qr(first) || !metal_graph_batch_kv_raw(first)) { return false; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &e->weights.layer[il]; if (!layer->attn_q_a || !layer->attn_kv || layer->attn_q_a->type != DS4_TENSOR_Q8_0 || layer->attn_kv->type != DS4_TENSOR_Q8_0 || layer->attn_q_a->dim[0] != DS4_N_EMBD || layer->attn_kv->dim[0] != DS4_N_EMBD || layer->attn_kv->dim[1] != DS4_N_HEAD_DIM) { return false; } } return true; } static bool metal_graph_encode_native_session_batch_shared( ds4_decode_item *items, int count, const ds4_model *model, const ds4_weights *weights, bool batch_qkv) { if (!items || count < 2 || !model || !weights) return false; ds4_gpu_graph *batch = &items[0].session->graph; const uint64_t norm_row_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); const char *stage = "embed"; uint32_t failed_layer = UINT32_MAX; bool ok = true; for (int i = 0; ok && i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; metal_graph_dspark_capture_begin(g); 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)items[i].token, DS4_N_EMBD, DS4_N_HC) != 0; } for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; const uint64_t shared_row_bytes = shared_dim * sizeof(float); failed_layer = il; if (batch_qkv) { stage = "decode-to-qkv"; for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); ok = metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, METAL_DECODE_LAYER_TO_QKV); } stage = "gather-qkv"; for (int i = 0; ok && i < count; i++) { ok = ds4_gpu_tensor_copy( metal_graph_batch_attn_norm(batch), (uint64_t)i * norm_row_bytes, metal_graph_attn_norm(&items[i].session->graph), 0, norm_row_bytes) != 0; } const uint64_t q_rank = layer->attn_q_a->dim[1]; const uint64_t qr_row_bytes = q_rank * sizeof(float); const uint64_t kv_row_bytes = (uint64_t)DS4_N_HEAD_DIM * sizeof(float); ds4_gpu_tensor *qr = NULL; ds4_gpu_tensor *kv_raw = NULL; if (ok) { qr = ds4_gpu_tensor_view( metal_graph_batch_qr(batch), 0, (uint64_t)count * qr_row_bytes); kv_raw = ds4_gpu_tensor_view( metal_graph_batch_kv_raw(batch), 0, (uint64_t)count * kv_row_bytes); ok = qr && kv_raw; } stage = "qkv"; if (ok) { ok = ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( qr, model->map, model->size, layer->attn_q_a->abs_offset, DS4_N_EMBD, q_rank, metal_graph_batch_attn_norm(batch), (uint32_t)count) != 0 && ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( kv_raw, model->map, model->size, layer->attn_kv->abs_offset, DS4_N_EMBD, DS4_N_HEAD_DIM, metal_graph_batch_attn_norm(batch), (uint32_t)count) != 0; } stage = "decode-from-qkv"; for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const int tier = g->active_tier; ds4_gpu_tensor *qr_row = ds4_gpu_tensor_view( qr, (uint64_t)i * qr_row_bytes, qr_row_bytes); ds4_gpu_tensor *kv_row = ds4_gpu_tensor_view( kv_raw, (uint64_t)i * kv_row_bytes, kv_row_bytes); ds4_gpu_tensor *saved_qr = g->qr_by_tier[tier]; ds4_gpu_tensor *saved_kv_raw = g->kv_raw_by_tier[tier]; ok = qr_row && kv_row; if (ok) { g->qr_by_tier[tier] = qr_row; g->kv_raw_by_tier[tier] = kv_row; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); ok = metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, METAL_DECODE_LAYER_FROM_QA_KV_RAW_TO_SHARED_MID); g->qr_by_tier[tier] = saved_qr; g->kv_raw_by_tier[tier] = saved_kv_raw; } ds4_gpu_tensor_free(kv_row); ds4_gpu_tensor_free(qr_row); } ds4_gpu_tensor_free(kv_raw); ds4_gpu_tensor_free(qr); } else { stage = "decode-to-shared"; for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); ok = metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, METAL_DECODE_LAYER_TO_SHARED_MID); } } stage = "gather-norm"; for (int i = 0; ok && i < count; i++) { ok = ds4_gpu_tensor_copy( metal_graph_batch_ffn_norm(batch), (uint64_t)i * norm_row_bytes, metal_graph_ffn_norm(&items[i].session->graph), 0, norm_row_bytes) != 0; } ds4_gpu_tensor *gate = NULL; ds4_gpu_tensor *up = NULL; ds4_gpu_tensor *mid = NULL; if (ok) { gate = ds4_gpu_tensor_view( metal_graph_batch_shared_gate(batch), 0, (uint64_t)count * shared_row_bytes); up = ds4_gpu_tensor_view( metal_graph_batch_shared_up(batch), 0, (uint64_t)count * shared_row_bytes); mid = ds4_gpu_tensor_view( metal_graph_batch_shared_mid(batch), 0, (uint64_t)count * shared_row_bytes); ok = gate && up && mid; } stage = "shared-gate-up"; if (ok) { ok = ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( gate, up, mid, model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_batch_ffn_norm(batch), (uint32_t)count, DS4_SWIGLU_CLAMP_EXP) != 0; } stage = "shared-down"; for (int i = 0; ok && i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; ds4_gpu_tensor *mid_row = ds4_gpu_tensor_view( mid, (uint64_t)i * shared_row_bytes, shared_row_bytes); ok = mid_row && 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, mid_row, metal_graph_routed_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; ds4_gpu_tensor_free(mid_row); } for (int i = 0; ok && i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; 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); } ds4_gpu_tensor_free(mid); ds4_gpu_tensor_free(up); ds4_gpu_tensor_free(gate); } stage = "output"; for (int i = 0; ok && i < count; i++) { ok = metal_graph_encode_output_head( &items[i].session->graph, model, weights, weights->output->dim[1]); } if (!ok) { fprintf(stderr, "ds4: Metal native shared session batch failed " "stage=%s layer=%u rows=%d\n", stage, failed_layer, count); } return ok; } static ds4_tp_batch_item *ds4_sessions_tp_batch_items( const ds4_decode_item *items, int count) { if (!items || count <= 0) return NULL; ds4_tp_batch_item *wire = calloc((size_t)count, sizeof(*wire)); if (!wire) return NULL; for (int i = 0; i < count; i++) { if (!items[i].session || items[i].session->tp_session_id == 0) { free(wire); return NULL; } wire[i].session_id = items[i].session->tp_session_id; wire[i].token = items[i].token; } return wire; } static bool ds4_sessions_tp_recv_logits( ds4_engine *e, ds4_session *prefill, ds4_decode_item *items, int count, char *err, size_t errlen) { if (!e || !e->tp.active || e->tp.rank != 0 || !e->tp.vocab_split) { return true; } const uint32_t vhalf = (uint32_t)DS4_N_VOCAB / 2u; if (prefill && !ds4_tp_recv_logits_half(e->tp.ctx, prefill->logits + vhalf, vhalf)) { if (err && errlen) snprintf(err, errlen, "tp: worker mixed-prefill logits missing"); return false; } for (int i = 0; i < count; i++) { if (!ds4_tp_recv_logits_half(e->tp.ctx, items[i].session->logits + vhalf, vhalf)) { if (err && errlen) { snprintf(err, errlen, "tp: worker batch logits missing for item %d", i); } return false; } } return true; } static int ds4_sessions_eval_batch_metal( ds4_decode_item *items, int count, ds4_engine *e, char *err, size_t errlen) { const bool mirror = e->tp.active && e->tp.rank == 0; if (mirror) { ds4_tp_batch_item *wire = ds4_sessions_tp_batch_items(items, count); const bool sent = wire && ds4_tp_send_eval_batch(e->tp.ctx, wire, (uint32_t)count); free(wire); if (!sent) { if (err && errlen) snprintf(err, errlen, "tp: worker decode batch send failed"); for (int i = 0; i < count; i++) { ds4_session_invalidate(items[i].session); } return 1; } } #if defined(__APPLE__) if (e->tp.active) ds4_gpu_tp_set_session_batch_mode(1); #endif bool ok = ds4_gpu_begin_commands() != 0; const bool native_shared = ok && metal_graph_native_session_batch_shared_supported(items, count, e); const bool native_qkv = native_shared && metal_graph_native_session_batch_qkv_supported(items, count, e); if (native_shared) { ok = metal_graph_encode_native_session_batch_shared( items, count, &e->model, &e->weights, native_qkv); } else { for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; const uint32_t pos = (uint32_t)s->checkpoint.len; if (ds4_session_is_glm(s)) { ok = glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, items[i].token, NULL, pos, NULL, s->logits, true); } else { ok = metal_graph_encode_token_raw_swa(&s->graph, &e->model, &e->weights, items[i].token, pos, true, false); } } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); #if defined(__APPLE__) if (e->tp.active) ds4_gpu_tp_set_session_batch_mode(0); if (ok && e->tp.active && ds4_gpu_tp_failed()) { if (err && errlen) snprintf(err, errlen, "tp: batch gate transport failed"); ok = false; } #endif for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_tensor *logits = ds4_session_is_glm(s) ? s->glm_graph.logits : metal_graph_logits(&s->graph); ok = ds4_gpu_tensor_read(logits, 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (ok && mirror) { ok = ds4_tp_wait_command_ack(e->tp.ctx, 0, "decode batch", err, errlen) && ds4_sessions_tp_recv_logits(e, NULL, items, count, err, errlen); } if (!ok) { for (int i = 0; i < count; i++) { ds4_session_invalidate(items[i].session); } if (err && errlen && !err[0]) { snprintf(err, errlen, "Metal batched decode failed"); } return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; const uint32_t pos = (uint32_t)s->checkpoint.len; const bool updates_dense = ds4_session_is_glm(s) && glm_graph_decode_updates_dense_cache(&s->glm_graph, pos, s->logits); token_vec_push(&s->checkpoint, items[i].token); s->checkpoint_valid = true; s->mtp_draft_valid = false; if (updates_dense) ds4_session_glm_note_dense_cache(s, pos, 1); if (!ds4_session_is_glm(s)) { ds4_session_dspark_capture_note_checkpoint(s); } } if (getenv("DS4_METAL_SESSION_BATCH_LOG") != NULL) { fprintf(stderr, "ds4: Metal session batch rows=%d family=%s " "native_shared=%d native_qkv=%d\n", count, ds4_session_is_glm(items[0].session) ? "glm" : "deepseek", native_shared ? 1 : 0, native_qkv ? 1 : 0); } return 0; } static bool ds4_sessions_eval_batch_with_prefill_metal_supported( ds4_decode_item *items, int count, ds4_session *prefill_session, const ds4_tokens *prefill_prompt) { if (!items || count <= 0 || !prefill_session || !prefill_prompt || !prefill_session->engine) { return false; } ds4_engine *e = prefill_session->engine; const char *tp_batch = getenv("DS4_METAL_TP_SESSION_BATCH"); if (e->backend != DS4_BACKEND_METAL || e->support_kind != DS4_SUPPORT_NONE || (e->tp.active && tp_batch && strcmp(tp_batch, "0") == 0) || ds4_session_is_cpu(prefill_session) || ds4_session_is_glm(prefill_session) || prefill_session->distributed || prefill_session->graph.ssd_streaming || ds4_session_cancelled(prefill_session) || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL || getenv("DS4_METAL_GRAPH_PREFILL_SPLIT_PROFILE") != NULL || getenv("DS4_METAL_LAYER_STAGE_PROFILE") != NULL) { return false; } const uint32_t start = (uint32_t)prefill_session->checkpoint.len; const uint32_t rows = (uint32_t)prefill_prompt->len - start; ds4_gpu_graph *g = &prefill_session->graph; if (rows < metal_graph_resume_prefill_min_tokens() || rows > g->prefill_cap || rows > g->raw_cap || (start % g->prefill_cap) + rows > g->prefill_cap) { return false; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e || s == prefill_session || s->distributed || ds4_session_is_cpu(s) || ds4_session_is_glm(s) || !s->checkpoint_valid || s->graph.ssd_streaming || ds4_session_cancelled(s)) { return false; } } return true; } /* Interleave one resumed Metal prefill chunk and independent decode rows by * layer in a single command epoch. Arithmetic stays on the established * per-sequence paths, so every session retains private cache and checkpoint * state while the scheduler avoids serial prefill/decode submissions. */ static int ds4_sessions_eval_batch_with_prefill_metal( ds4_decode_item *items, int count, ds4_session *prefill_session, const ds4_tokens *prefill_prompt, char *err, size_t errlen) { ds4_engine *e = prefill_session->engine; ds4_gpu_graph *pg = &prefill_session->graph; const uint32_t start = (uint32_t)prefill_session->checkpoint.len; const uint32_t rows = (uint32_t)prefill_prompt->len - start; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const bool mirror = e->tp.active && e->tp.rank == 0; if (mirror) { ds4_tp_batch_item *wire = ds4_sessions_tp_batch_items(items, count); const bool sent = wire && prefill_session->tp_session_id != 0 && ds4_tp_send_mixed_batch( e->tp.ctx, prefill_session->tp_session_id, prefill_prompt->v, (uint32_t)prefill_prompt->len, wire, (uint32_t)count); free(wire); if (!sent) { if (err && errlen) snprintf(err, errlen, "tp: worker mixed batch send failed"); ds4_session_invalidate(prefill_session); for (int i = 0; i < count; i++) { ds4_session_invalidate(items[i].session); } return 1; } } bool ok = metal_graph_set_active_tier_batch(pg, pg->emb_tier, rows) && metal_graph_upload_prompt_tokens( metal_graph_prefill_tokens(pg), prefill_prompt, start, rows) && metal_graph_warmup_prefill_kernels( pg, &e->model, &e->weights, rows) && metal_graph_set_active_tier_batch(pg, pg->emb_tier, rows) && metal_graph_upload_prompt_embeddings_hc( metal_graph_batch_cur_hc(pg), metal_graph_prefill_tokens(pg), &e->model, &e->weights, prefill_prompt, start, rows); metal_graph_dspark_capture_begin_prefill(pg); for (int i = 0; ok && i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; metal_graph_dspark_capture_begin(g); ok = metal_graph_set_active_tier_decode(g, g->emb_tier) && ds4_gpu_embed_token_hc_tensor( metal_graph_cur_hc(g), e->model.map, e->model.size, e->weights.token_embd->abs_offset, (uint32_t)e->weights.token_embd->dim[1], (uint32_t)items[i].token, DS4_N_EMBD, DS4_N_HC) != 0; } #if defined(__APPLE__) if (e->tp.active) ds4_gpu_tp_set_session_batch_mode(1); #endif 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( pg, &e->model, &e->weights.layer[il], il, start, rows); if (ok) { ok = metal_graph_dspark_capture_prefill_layer(pg, il, start, rows); } } /* Decode graphs retain their established session-major ordering. Some * decode kernels use the engine-shared prefill scratch as transient * reduction space, so alternating prefill and decode at layer granularity * would let graph-local ping-pong orientations alias the same allocation. */ for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { ok = metal_graph_encode_decode_layer( g, &e->model, &e->weights.layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].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); } } } const int prefill_src_tier = pg->active_tier; ds4_gpu_tensor *saved_prefill_cur = NULL; ds4_gpu_tensor *last_prefill_hc = NULL; if (ok) { saved_prefill_cur = pg->cur_hc_by_tier[prefill_src_tier]; last_prefill_hc = metal_graph_tensor_row_view( metal_graph_batch_cur_hc(pg), rows - 1u, hc_dim); ok = last_prefill_hc != NULL; } if (ok) { pg->cur_hc_by_tier[prefill_src_tier] = last_prefill_hc; ok = metal_graph_encode_output_head( pg, &e->model, &e->weights, e->weights.output->dim[1]); pg->cur_hc_by_tier[prefill_src_tier] = saved_prefill_cur; } for (int i = 0; ok && i < count; i++) { ok = metal_graph_encode_output_head( &items[i].session->graph, &e->model, &e->weights, e->weights.output->dim[1]); } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); #if defined(__APPLE__) if (e->tp.active) ds4_gpu_tp_set_session_batch_mode(0); if (ok && e->tp.active && ds4_gpu_tp_failed()) { if (err && errlen) snprintf(err, errlen, "tp: mixed batch gate transport failed"); ok = false; } #endif if (saved_prefill_cur) { pg->cur_hc_by_tier[prefill_src_tier] = saved_prefill_cur; } ds4_gpu_tensor_free(last_prefill_hc); if (ok) { ok = ds4_gpu_tensor_read( metal_graph_logits(pg), 0, prefill_session->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ok = ds4_gpu_tensor_read( metal_graph_logits(&s->graph), 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (ok && mirror) { ok = ds4_tp_wait_command_ack( e->tp.ctx, prefill_session->tp_session_id, "mixed prefill/decode batch", err, errlen) && ds4_sessions_tp_recv_logits( e, prefill_session, items, count, err, errlen); } if (!ok) { ds4_session_invalidate(prefill_session); for (int i = 0; i < count; i++) { ds4_session_invalidate(items[i].session); } if (err && errlen && !err[0]) { snprintf(err, errlen, "Metal mixed model batch failed"); } return 1; } ds4_tokens_copy(&prefill_session->checkpoint, prefill_prompt); prefill_session->checkpoint_valid = true; prefill_session->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(prefill_session); session_greedy_splitkv_reset(prefill_session); if (prefill_session->progress) { prefill_session->progress( prefill_session->progress_ud, "prefill_chunk", prefill_prompt->len, prefill_prompt->len); } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; token_vec_push(&s->checkpoint, items[i].token); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); } if (getenv("DS4_METAL_SESSION_BATCH_LOG") != NULL) { fprintf(stderr, "ds4: Metal mixed batch prefill_rows=%u decode_rows=%d\n", rows, count); } return 0; } #endif static int ds4_sessions_eval_batch_cuda(ds4_decode_item *items, int count, char *err, size_t errlen); static int ds4_sessions_eval_batch_with_prefill_cuda( ds4_decode_item *items, int count, ds4_session *prefill_session, const ds4_tokens *prefill_prompt, char *err, size_t errlen); int ds4_sessions_eval_batch(ds4_decode_item *items, int count, char *err, size_t errlen) { if (!items || count <= 0) { if (err && errlen) snprintf(err, errlen, "empty decode batch"); return 1; } if (count == 1) { return ds4_session_eval(items[0].session, items[0].token, err, errlen); } ds4_session *first = items[0].session; if (!first || !first->engine) { if (err && errlen) snprintf(err, errlen, "decode batch has no session"); return 1; } ds4_engine *e = first->engine; for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e) { if (err && errlen) { snprintf(err, errlen, "decode batch item %d belongs to a different engine", i); } return 1; } if (items[i].token < 0 || items[i].token >= (int)DS4_N_VOCAB) { if (err && errlen) { snprintf(err, errlen, "decode batch item %d has an invalid token", i); } return 1; } for (int j = 0; j < i; j++) { if (items[j].session == s) { if (err && errlen) { snprintf(err, errlen, "decode batch repeats session at items %d and %d", j, i); } return 1; } } if (s->checkpoint.len >= s->ctx_size) { if (err && errlen) { snprintf(err, errlen, "decode batch item %d reached its context limit", i); } return 1; } } #ifndef DS4_NO_GPU if (e->backend == DS4_BACKEND_CUDA) { return ds4_sessions_eval_batch_cuda(items, count, err, errlen); } if (ds4_sessions_eval_batch_metal_supported(items, count, e)) { return ds4_sessions_eval_batch_metal(items, count, e, err, errlen); } #endif /* Preserve logical all-or-nothing behavior even on the serialized path. * A failure can leave earlier members advanced, so force every member to * rebuild before it is used again. */ for (int i = 0; i < count; i++) { if (ds4_session_eval(items[i].session, items[i].token, err, errlen) != 0) { for (int j = 0; j < count; j++) { ds4_session_invalidate(items[j].session); } return 1; } } return 0; } int ds4_sessions_eval_batch_with_prefill( ds4_decode_item *items, int count, ds4_session *prefill_session, const ds4_tokens *prefill_prompt, char *err, size_t errlen) { if (!items || count <= 0 || !prefill_session || !prefill_prompt || !prefill_session->engine) { if (err && errlen) snprintf(err, errlen, "invalid mixed model batch"); return 1; } if (!prefill_session->checkpoint_valid || prefill_prompt->len <= prefill_session->checkpoint.len || prefill_prompt->len > prefill_session->ctx_size || !ds4_tokens_starts_with(prefill_prompt, &prefill_session->checkpoint)) { if (err && errlen) { snprintf(err, errlen, "mixed prefill must extend a valid session checkpoint"); } return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s == prefill_session || s->engine != prefill_session->engine || !s->checkpoint_valid || s->checkpoint.len >= s->ctx_size || items[i].token < 0 || items[i].token >= (int)DS4_N_VOCAB) { if (err && errlen) { snprintf(err, errlen, "invalid mixed decode item %d", i); } return 1; } for (int j = 0; j < i; j++) { if (items[j].session == s) { if (err && errlen) { snprintf(err, errlen, "mixed decode repeats session at items %d and %d", j, i); } return 1; } } } #ifndef DS4_NO_GPU if (prefill_session->engine->backend == DS4_BACKEND_CUDA) { return ds4_sessions_eval_batch_with_prefill_cuda( items, count, prefill_session, prefill_prompt, err, errlen); } if (ds4_sessions_eval_batch_with_prefill_metal_supported( items, count, prefill_session, prefill_prompt)) { return ds4_sessions_eval_batch_with_prefill_metal( items, count, prefill_session, prefill_prompt, err, errlen); } #endif int rc = ds4_session_sync(prefill_session, prefill_prompt, err, errlen); if (rc != 0) return rc; rc = ds4_sessions_eval_batch(items, count, err, errlen); if (rc != 0) ds4_session_invalidate(prefill_session); return rc; } #ifndef DS4_NO_GPU static int ds4_session_eval_dspark_speculative_argmax( ds4_session *s, int n_accept, int max_tokens, int eos_token, int *accepted, int accepted_cap, char *err, size_t errlen) { const bool spec_log = getenv("DS4_DSPARK_SPEC_LOG") != NULL; const bool stats_enabled = s && ds4_dspark_stats_enabled(); const bool scheduler_enabled = s && ds4_dspark_scheduler_enabled(); const double stats_t0 = (stats_enabled || scheduler_enabled) ? now_sec() : 0.0; #define DS4_DSPARK_STATS_FINISH() do { \ if (stats_enabled) { \ s->dspark_stats.total_ms += (now_sec() - stats_t0) * 1000.0; \ } \ } while (0) #define DS4_DSPARK_SCHED_EXTRA_MS() \ ((scheduler_enabled && stats_t0 != 0.0) ? \ s->dspark_last_propose_ms + (now_sec() - stats_t0) * 1000.0 : 0.0) if (stats_enabled) { s->dspark_stats.cycles++; if (n_accept > 0) s->dspark_stats.first_tokens++; } if (spec_log) { fprintf(stderr, "ds4: DSpark spec enter accepted=%d max=%d valid=%d len=%u pos=%d\n", n_accept, max_tokens, s ? (s->dspark_draft_valid ? 1 : 0) : 0, s ? s->dspark_draft_len : 0, s ? s->checkpoint.len : -1); } if (!s || !s->dspark_draft_valid || s->dspark_draft_len == 0) { if (stats_enabled) { s->dspark_stats.no_draft++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } if (s) { ds4_session_dspark_scheduler_note( s, 0, true, DS4_DSPARK_SCHED_EXTRA_MS()); } if (spec_log) { fprintf(stderr, "ds4: DSpark spec skip no-draft\n"); } DS4_DSPARK_STATS_FINISH(); return n_accept; } int draft_n = (int)s->dspark_draft_len; if (draft_n > max_tokens - n_accept) draft_n = max_tokens - n_accept; if (draft_n > accepted_cap - n_accept) draft_n = accepted_cap - n_accept; int room = s->ctx_size - s->checkpoint.len; if (draft_n > room - 1) draft_n = room - 1; if (draft_n <= 0) { s->dspark_draft_valid = false; s->dspark_draft_len = 0; if (stats_enabled) { s->dspark_stats.no_room++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } if (spec_log) { fprintf(stderr, "ds4: DSpark spec skip no-room\n"); } DS4_DSPARK_STATS_FINISH(); return n_accept; } int drafts[DS4_DSPARK_MAX_BLOCK_SIZE]; for (int i = 0; i < draft_n; i++) { drafts[i] = s->dspark_draft_tokens[i]; if (drafts[i] < 0 || drafts[i] >= (int)DS4_N_VOCAB) { s->dspark_draft_valid = false; s->dspark_draft_len = 0; if (stats_enabled) { s->dspark_stats.invalid_draft++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } if (spec_log) { fprintf(stderr, "ds4: DSpark spec skip invalid-draft index=%d token=%d\n", i, drafts[i]); } DS4_DSPARK_STATS_FINISH(); return n_accept; } } if (stats_enabled) { s->dspark_stats.proposed_tokens += (uint64_t)draft_n; ds4_dspark_stats_note_len(s->dspark_stats.draft_len_hist, (uint32_t)draft_n); } s->dspark_draft_valid = false; s->dspark_draft_len = 0; const int target_top = sample_argmax(s->logits, DS4_N_VOCAB); if (target_top != drafts[0]) { if (stats_enabled) { s->dspark_stats.first_misses++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } ds4_session_dspark_scheduler_note( s, 0, false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec miss first draft=%d base=%d\n", drafts[0], target_top); } DS4_DSPARK_STATS_FINISH(); return n_accept; } if (drafts[0] == eos_token) draft_n = 1; ds4_engine *e = s->engine; ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); int row_tops_buf[DS4_DSPARK_MAX_BLOCK_SIZE]; int *row_tops = draft_n > 1 ? row_tops_buf : NULL; float *row_logits = s->spec_row_logits; const int start = s->checkpoint.len; const double snapshot_t0 = stats_enabled ? now_sec() : 0.0; bool have_frontier = spec_frontier_snapshot(&frontier, s); if (stats_enabled) { s->dspark_stats.snapshot_ms += (now_sec() - snapshot_t0) * 1000.0; } bool ok = have_frontier && row_logits && (draft_n <= 1 || row_tops); bool verifier_may_have_mutated = false; bool tp_verify_sent = false; if (ok && ds4_session_tp_leader(s)) { /* Announce the block before mutating anything: the worker runs its * half of the verify and then waits for our commit decision. */ if (!ds4_tp_send_verify(e->tp.ctx, s->tp_session_id, drafts, (uint32_t)draft_n)) { snprintf(err, errlen, "tp: verify send failed"); spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } tp_verify_sent = true; } if (ok) { for (int i = 0; i < draft_n; i++) token_vec_push(&s->checkpoint, drafts[i]); verifier_may_have_mutated = true; ds4_verify_suffix_timing verify_timing; const double verify_t0 = stats_enabled ? now_sec() : 0.0; ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights, &s->checkpoint, (uint32_t)start, (uint32_t)draft_n, false, true, row_tops, NULL, stats_enabled ? &verify_timing : NULL); if (stats_enabled) { s->dspark_stats.verify_ms += (now_sec() - verify_t0) * 1000.0; s->dspark_stats.verify_upload_ms += verify_timing.upload_ms; s->dspark_stats.verify_layer_ms += verify_timing.layer_ms; s->dspark_stats.verify_head_ms += verify_timing.head_ms; s->dspark_stats.verify_read_ms += verify_timing.read_ms; if (verify_timing.fused_head) { s->dspark_stats.verifier_fused_head++; } } } int commit_drafts = 0; if (ok) { commit_drafts = 1; for (int i = 1; i < draft_n; i++) { if (row_tops[i - 1] != drafts[i]) break; commit_drafts++; } } bool final_logits_ok = false; if (ok && commit_drafts == draft_n) { const double read_t0 = stats_enabled ? now_sec() : 0.0; final_logits_ok = metal_graph_read_spec_logits_row(&s->graph, (uint32_t)(draft_n - 1), row_logits); if (stats_enabled) { s->dspark_stats.verify_read_ms += (now_sec() - read_t0) * 1000.0; } } if (ok && commit_drafts == draft_n && final_logits_ok) { if (tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, 1, 0)) { snprintf(err, errlen, "tp: verify commit send failed"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); int emitted_drafts = 0; for (int i = 0; i < draft_n && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; emitted_drafts++; if (drafts[i] == eos_token) break; } s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); if (stats_enabled) { s->dspark_stats.full_accepts++; s->dspark_stats.accepted_draft_tokens += (uint64_t)emitted_drafts; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)emitted_drafts); } ds4_session_dspark_scheduler_note( s, (uint32_t)emitted_drafts, false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec accept drafted=%d accepted=%d\n", draft_n, n_accept); } spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return n_accept; } if (verifier_may_have_mutated) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); if (!have_frontier || !spec_frontier_restore(&frontier, s)) { if (tp_verify_sent) (void)ds4_tp_send_verify_commit(e->tp.ctx, 0, 0); snprintf(err, errlen, "DSpark verifier rollback failed"); s->checkpoint_valid = false; if (stats_enabled) { s->dspark_stats.verifier_errors++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } } if (!ok) { if (tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, 0, 0)) { snprintf(err, errlen, "tp: verify commit send failed"); spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } if (stats_enabled) { s->dspark_stats.verifier_unavailable++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } if (spec_log) { fprintf(stderr, "ds4: DSpark spec verifier unavailable frontier=%d row_logits=%d row_tops=%d mutated=%d\n", have_frontier ? 1 : 0, row_logits ? 1 : 0, row_tops || draft_n <= 1 ? 1 : 0, verifier_may_have_mutated ? 1 : 0); } spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return n_accept; } /* Precompute the exact replay count (cap + eos cuts) so the worker can * run the same gated replay evals in lockstep. */ int replay_budget = commit_drafts; if (replay_budget > accepted_cap - n_accept) replay_budget = accepted_cap - n_accept; if (replay_budget < 0) replay_budget = 0; for (int i = 0; i < replay_budget; i++) { if (drafts[i] == eos_token) { replay_budget = i + 1; break; } } if (tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, 0, replay_budget)) { snprintf(err, errlen, "tp: verify commit send failed"); spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } const double replay_t0 = stats_enabled ? now_sec() : 0.0; int replayed_drafts = 0; for (int i = 0; i < replay_budget; i++) { ok = metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, drafts[i], (uint32_t)s->checkpoint.len, row_logits); if (!ok) { snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; if (stats_enabled) { s->dspark_stats.verifier_errors++; s->dspark_stats.replay_ms += (now_sec() - replay_t0) * 1000.0; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } token_vec_push(&s->checkpoint, drafts[i]); accepted[n_accept++] = drafts[i]; replayed_drafts++; if (drafts[i] == eos_token) break; } if (stats_enabled) { s->dspark_stats.replay_ms += (now_sec() - replay_t0) * 1000.0; } /* Vocab-split head: the last replay eval produced only our logits half; * merge the worker's before installing them as the session logits. */ if (replayed_drafts > 0 && tp_verify_sent && e->tp.vocab_split) { const uint32_t vhalf = (uint32_t)DS4_N_VOCAB / 2u; if (!ds4_tp_recv_logits_half(e->tp.ctx, row_logits + vhalf, vhalf)) { snprintf(err, errlen, "tp: replay logits half missing"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } } if (replayed_drafts > 0) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); if (stats_enabled) { if (replayed_drafts == draft_n) s->dspark_stats.full_accepts++; else s->dspark_stats.partial_accepts++; s->dspark_stats.accepted_draft_tokens += (uint64_t)replayed_drafts; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)replayed_drafts); } } else if (stats_enabled) { ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } ds4_session_dspark_scheduler_note( s, (uint32_t)replayed_drafts, false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec partial drafted=%d verified=%d accepted=%d\n", draft_n, commit_drafts, n_accept); } spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); #undef DS4_DSPARK_SCHED_EXTRA_MS #undef DS4_DSPARK_STATS_FINISH return n_accept; } #endif /* TP worker side of a mirrored speculative-verify block. Runs its half of the * same batch verify as the leader, including per-layer combine gates, purely * for KV, compressor, and indexer side effects; then it obeys the commit * frame: keep the pushed rows, or roll back and replay the accepted prefix * through the gated single-token decode in lockstep with the leader. */ int ds4_session_tp_spec_cycle(ds4_session *s, const int *drafts, int draft_n, char *err, size_t errlen) { #ifdef DS4_NO_GPU (void)s; (void)drafts; (void)draft_n; snprintf(err, errlen, "GPU support is not compiled in"); return 1; #else ds4_engine *e = s ? s->engine : NULL; if (!e || !e->tp.active || e->tp.rank != 1) { snprintf(err, errlen, "tp: spec cycle outside worker mode"); return 1; } if (draft_n <= 0 || draft_n > DS4_DSPARK_MAX_BLOCK_SIZE) { snprintf(err, errlen, "tp: bad verify block size %d", draft_n); return 1; } for (int i = 0; i < draft_n; i++) { if (drafts[i] < 0 || drafts[i] >= (int)DS4_N_VOCAB) { snprintf(err, errlen, "tp: bad draft token %d", drafts[i]); return 1; } } if (draft_n > s->ctx_size - s->checkpoint.len) { snprintf(err, errlen, "tp: verify block beyond context"); return 1; } ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); const int start = s->checkpoint.len; if (!spec_frontier_snapshot(&frontier, s)) { snprintf(err, errlen, "tp: frontier snapshot failed"); return 1; } for (int i = 0; i < draft_n; i++) token_vec_push(&s->checkpoint, drafts[i]); int row_tops[DS4_DSPARK_MAX_BLOCK_SIZE]; bool ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights, &s->checkpoint, (uint32_t)start, (uint32_t)draft_n, false, false, draft_n > 1 ? row_tops : NULL, NULL, NULL); int32_t full_accept = 0, replay_n = 0; if (!ds4_tp_recv_verify_commit(e->tp.ctx, &full_accept, &replay_n)) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: verify commit frame missing"); return 1; } if (full_accept) { spec_frontier_free(&frontier); if (!ok) { /* The leader committed a block our verify failed to apply: * the states have diverged and lockstep cannot continue. */ snprintf(err, errlen, "tp: worker verify failed on accepted block"); s->checkpoint_valid = false; return 1; } s->checkpoint_valid = true; return 0; } s->checkpoint.len = start; if (!spec_frontier_restore(&frontier, s)) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: verify rollback failed"); s->checkpoint_valid = false; return 1; } spec_frontier_free(&frontier); if (replay_n > draft_n) replay_n = draft_n; float *logits = s->spec_row_logits; float *scratch = NULL; if (!logits) { scratch = malloc((size_t)DS4_N_VOCAB * sizeof(float)); if (!scratch) { snprintf(err, errlen, "tp: replay logits alloc failed"); return 1; } logits = scratch; } for (int i = 0; i < replay_n; i++) { if (!metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, drafts[i], (uint32_t)s->checkpoint.len, logits)) { free(scratch); snprintf(err, errlen, "tp: replay decode failed"); s->checkpoint_valid = false; return 1; } token_vec_push(&s->checkpoint, drafts[i]); } if (replay_n > 0) { s->checkpoint_valid = true; if (e->tp.vocab_split) { const uint32_t vhalf = (uint32_t)DS4_N_VOCAB / 2u; if (!ds4_tp_send_logits_half(e->tp.ctx, logits + vhalf, vhalf)) { free(scratch); snprintf(err, errlen, "tp: replay logits half send failed"); return 1; } } } free(scratch); return 0; #endif } #ifndef DS4_NO_GPU static bool metal_graph_session_batch_moe_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { if (!items || count < 3 || !weights || DS4_N_EXPERT_USED != 6u || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement || !first->cuda_tp_decode || !first->cuda_tp_ep || first->quality || (uint32_t)count > first->prefill_cap) { 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_down_exps || layer->ffn_gate_exps->type != 12u || layer->ffn_down_exps->type != 12u) { return false; } const int home = first->placement[il + 1u]; const int partner = metal_graph_cuda_tp_partner_tier(home); if (partner < 0 || !g_gpu_peer_ok[partner][home] || !first->batch_ffn_norm_by_tier[home] || !first->batch_router_logits_by_tier[home] || !first->batch_router_selected_by_tier[home] || !first->batch_router_weights_by_tier[home] || !first->batch_routed_gate_by_tier[home] || !first->batch_routed_up_by_tier[home] || !first->batch_routed_mid_by_tier[home] || !first->batch_routed_down_by_tier[home] || !first->batch_routed_out_by_tier[home] || !first->batch_ffn_norm_by_tier[partner] || !first->batch_router_selected_by_tier[partner] || !first->batch_router_weights_by_tier[partner] || !first->batch_routed_gate_by_tier[partner] || !first->batch_routed_up_by_tier[partner] || !first->batch_routed_mid_by_tier[partner] || !first->batch_routed_down_by_tier[partner] || !first->batch_routed_out_by_tier[partner]) { return false; } for (int i = 1; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || !g->cuda_tp_decode || !g->cuda_tp_ep || g->quality || (uint32_t)count > g->prefill_cap || g->placement[il + 1u] != home) { return false; } } } return true; } static bool metal_graph_session_batch_shared_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { if (!items || count < 3 || !weights || metal_graph_use_reference_shared_down_hc()) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (first->quality || first->cuda_tp_shared || !first->shared_gate_up_swiglu_fuse || (uint32_t)count > first->prefill_cap) { return false; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; 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) { return false; } const int home = first->placement[il + 1u]; if (!first->batch_ffn_norm_by_tier[home] || !first->batch_routed_out_by_tier[home] || !first->batch_shared_gate_by_tier[home] || !first->batch_shared_up_by_tier[home] || !first->batch_shared_mid_by_tier[home] || !first->batch_shared_out_by_tier[home] || !first->batch_after_attn_hc_by_tier[home] || !first->batch_hc_split_by_tier[home] || !first->batch_next_hc_by_tier[home]) { return false; } for (int i = 0; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; const uint32_t pos = (uint32_t)items[i].session->checkpoint.len; if (g->quality || g->cuda_tp_shared || !g->shared_gate_up_swiglu_fuse || g->placement[il + 1u] != home || metal_graph_needs_ffn_out(g, il, pos) || metal_graph_directional_steering_ffn_enabled(g)) { return false; } } } return true; } static bool metal_graph_session_batch_attn_pre_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { if (!items || count < 3 || !weights || DS4_MODEL_VARIANT != DS4_VARIANT_FLASH || metal_graph_use_reference_hc_decode() || metal_graph_use_reference_hc_norm_decode()) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if ((uint32_t)count > first->prefill_cap) return false; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; if (!layer->hc_attn_fn || !layer->hc_attn_scale || !layer->hc_attn_base || !layer->attn_norm || layer->hc_attn_fn->type != DS4_TENSOR_F16) { return false; } const int home = first->placement[il + 1u]; if (!first->batch_cur_hc_by_tier[home] || !first->batch_flat_hc_by_tier[home] || !first->batch_hc_mix_by_tier[home] || !first->batch_hc_split_by_tier[home] || !first->batch_attn_cur_by_tier[home] || !first->batch_attn_norm_by_tier[home]) { return false; } for (int i = 1; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || g->placement[il + 1u] != home) return false; } } return true; } static bool metal_graph_session_batch_attn_core_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { #if defined(__APPLE__) (void)items; (void)count; (void)weights; return false; #else if (!items || count < 3 || !weights || count > (int)DS4_GPU_ATTENTION_DECODE_BATCH_MAX || DS4_N_HEAD_DIM != 512u || metal_graph_attn_comp_cache_is_f16() || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL || getenv("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE") != NULL || getenv("DS4_CUDA_SPLITKV_DECODE") != NULL || getenv("DS4_CUDA_DECODE_HEADS8_ONLINE") != NULL || getenv("DS4_CUDA_DECODE_SCORE4") != NULL || getenv("DS4_CUDA_DECODE_SCORE8") != NULL || getenv("DS4_CUDA_NO_DECODE_VALUE512") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_GRAPH") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_LDG") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4_PLAIN") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_DIM2") != NULL || getenv("DS4_CUDA_EXACT_SCORE_SPLIT_FUSE_INV_ROPE") != NULL || 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 false; } ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement || first->quality || first->cuda_tp_attn_heads || first->cuda_tp_q || first->decode_stage_profile || first->decode_index_stage_profile || (uint32_t)count > first->prefill_cap) { return false; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; if (!g->placement || g->quality || g->cuda_tp_attn_heads || g->cuda_tp_q || g->decode_stage_profile || g->decode_index_stage_profile || metal_graph_directional_steering_attn_enabled(g) || s->checkpoint.len < 1 || (uint32_t)count > g->prefill_cap) { return false; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; if (!layer->attn_q_a || !layer->attn_sinks || layer->attn_q_a->dim[1] == 0u) { return false; } const int home = first->placement[il + 1u]; if (home < 0 || !first->batch_qr_by_tier[home] || !first->batch_qr_norm_by_tier[home] || !first->batch_q_by_tier[home] || !first->batch_kv_raw_by_tier[home] || !first->batch_kv_by_tier[home] || !first->batch_comp_kv_by_tier[home] || !first->batch_comp_sc_by_tier[home] || !first->batch_indexer_q_by_tier[home] || !first->batch_indexer_weights_by_tier[home] || !first->batch_heads_by_tier[home] || !first->batch_attn_low_by_tier[home] || !first->batch_attn_out_by_tier[home] || !first->batch_after_attn_hc_by_tier[home]) { return false; } for (int i = 1; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (g->placement[il + 1u] != home || g->cuda_tp_attn != first->cuda_tp_attn) { return false; } } } return true; #endif } static bool metal_graph_session_batch_attn_post_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { #if defined(__APPLE__) (void)items; (void)count; (void)weights; return false; #else if (!items || count < 3 || !weights || (DS4_N_OUT_GROUP & 1u) != 0u || metal_graph_use_reference_attn_out_hc()) { return false; } const uint32_t tp_groups = DS4_N_OUT_GROUP / 2u; const uint64_t tp_width = (uint64_t)tp_groups * DS4_N_LORA_O; if (tp_width == 0u || (tp_width % 32u) != 0u) return false; ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement || first->quality || !first->cuda_tp_attn || first->cuda_tp_attn_heads || !first->cuda_tp_attn_peer_read || first->decode_stage_profile || (uint32_t)count > first->prefill_cap) { return false; } for (int i = 0; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || g->quality || !g->cuda_tp_attn || g->cuda_tp_attn_heads || !g->cuda_tp_attn_peer_read || g->decode_stage_profile || metal_graph_directional_steering_attn_enabled(g)) { return false; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; const int home = first->placement[il + 1u]; const int partner = metal_graph_cuda_tp_partner_tier(home); if (home < 0 || partner < 0 || !g_gpu_peer_ok[home][partner] || !g_gpu_peer_ok[partner][home] || !layer->attn_output_a || !layer->attn_output_b || layer->attn_output_a->type != DS4_TENSOR_Q8_0 || layer->attn_output_b->type != DS4_TENSOR_Q8_0 || !first->batch_heads_by_tier[home] || !first->batch_attn_low_by_tier[home] || !first->batch_attn_out_by_tier[home] || !first->batch_attn_low_by_tier[partner] || !first->batch_attn_out_by_tier[partner] || !first->batch_shared_out_by_tier[home] || !first->batch_cur_hc_by_tier[home] || !first->batch_hc_split_by_tier[home] || !first->batch_after_attn_hc_by_tier[home]) { return false; } for (int i = 1; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (g->placement[il + 1u] != home) return false; } } return true; #endif } static bool metal_graph_session_batch_qkv_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { if (!items || count < 3 || !weights || count > (int)DS4_GPU_ATTENTION_DECODE_BATCH_MAX || DS4_N_HEAD_KV != 1u || metal_graph_use_reference_qkv_norm() || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement || first->quality || first->cuda_tp_q || !first->cuda_qkv_pair || !first->cuda_qkv_kv_rope_fuse || !first->cuda_q_norm_rope_fuse || first->decode_stage_profile || (uint32_t)count > first->prefill_cap) { return false; } for (int i = 0; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || g->quality || g->cuda_tp_q || !g->cuda_qkv_pair || !g->cuda_qkv_kv_rope_fuse || !g->cuda_q_norm_rope_fuse || g->decode_stage_profile) { return false; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; const int home = first->placement[il + 1u]; if (home < 0 || !layer->attn_q_a || !layer->attn_q_b || !layer->attn_kv || !layer->attn_q_a_norm || !layer->attn_kv_a_norm || layer->attn_q_a->type != DS4_TENSOR_Q8_0 || layer->attn_q_b->type != DS4_TENSOR_Q8_0 || layer->attn_kv->type != DS4_TENSOR_Q8_0 || !first->batch_attn_norm_by_tier[home] || !first->batch_qr_by_tier[home] || !first->batch_qr_norm_by_tier[home] || !first->batch_q_by_tier[home] || !first->batch_kv_raw_by_tier[home] || !first->batch_kv_by_tier[home]) { return false; } for (int i = 1; i < count; i++) { if (items[i].session->graph.placement[il + 1u] != home) { return false; } } } return true; } static bool metal_graph_session_batch_kv_store_supported( ds4_decode_item *items, int count) { if (!items || count < 2 || count > (int)DS4_GPU_ATTENTION_DECODE_BATCH_MAX || metal_graph_use_reference_kv_decode()) { return false; } ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement || first->cuda_tp_attn_cache_dup) return false; for (int i = 0; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || g->raw_cap == 0u || g->cuda_tp_attn_cache_dup) { return false; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const int home = first->placement[il + 1u]; if (g->placement[il + 1u] != home || !g->layer_raw_cache[il] || ds4_gpu_tensor_device(g->layer_raw_cache[il]) != home) { return false; } } } return true; } /* Batch the cache-independent HC preparation at the start of an attention * block. The narrow F16 HC projection stays rowwise for exact arithmetic; * both HC transforms around it execute over all session rows together. */ static bool metal_graph_encode_attn_pre_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t row_base, bool scatter_rows) { if (!items || count < 2 || !model || !layer || il >= DS4_N_LAYER) { return false; } ds4_gpu_graph *g = &items[0].session->graph; const int home_tier = g->placement[il + 1u]; 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 hc_bytes = (uint64_t)count * hc_dim * sizeof(float); const uint64_t mix_bytes = (uint64_t)count * mix_hc * sizeof(float); const uint64_t norm_bytes = (uint64_t)count * DS4_N_EMBD * sizeof(float); bool ok = home_tier >= 0 && (uint64_t)row_base + (uint32_t)count <= g->prefill_cap; for (int i = 0; ok && i < count; i++) { ok = metal_graph_set_active_tier_decode( &items[i].session->graph, home_tier); } ds4_gpu_tensor cur_hc, flat_hc, hc_mix, hc_split; ds4_gpu_tensor attn_cur, attn_norm; if (ok) { ok = metal_graph_borrow_tensor_view(&cur_hc, metal_graph_batch_cur_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes) && metal_graph_borrow_tensor_view(&flat_hc, metal_graph_batch_flat_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes) && metal_graph_borrow_tensor_view(&hc_mix, metal_graph_batch_hc_mix(g), (uint64_t)row_base * mix_hc * sizeof(float), mix_bytes) && metal_graph_borrow_tensor_view(&hc_split, metal_graph_batch_hc_split(g), (uint64_t)row_base * mix_hc * sizeof(float), mix_bytes) && metal_graph_borrow_tensor_view(&attn_cur, metal_graph_batch_attn_cur(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&attn_norm, metal_graph_batch_attn_norm(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes); } for (int i = 0; ok && i < count; i++) { ds4_gpu_graph *srcg = &items[i].session->graph; ds4_gpu_tensor dst_row; ok = srcg->active_tier == home_tier && metal_graph_borrow_tensor_view( &dst_row, &cur_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)); if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &dst_row, metal_graph_cur_hc(srcg), dst_row.bytes) != 0; } } if (ok) { ok = ds4_gpu_rms_norm_plain_rows_tensor( &flat_hc, &cur_hc, (uint32_t)hc_dim, (uint32_t)count, DS4_RMS_EPS) != 0; } for (int i = 0; ok && i < count; i++) { ds4_gpu_tensor out_row, in_row; ok = metal_graph_borrow_tensor_view( &out_row, &hc_mix, (uint64_t)i * mix_hc * sizeof(float), mix_hc * sizeof(float)) && metal_graph_borrow_tensor_view( &in_row, &flat_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)); if (ok) { ok = metal_graph_matmul_plain_tensor( &out_row, model, layer->hc_attn_fn, hc_dim, mix_hc, &in_row, 1); } } if (ok) { ok = ds4_gpu_hc_split_weighted_sum_norm_tensor( &attn_cur, &attn_norm, &hc_split, &hc_mix, &cur_hc, 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; } for (int i = 0; scatter_rows && ok && i < count; i++) { ds4_gpu_graph *dstg = &items[i].session->graph; ds4_gpu_tensor norm_row, split_row; ok = metal_graph_borrow_tensor_view( &norm_row, &attn_norm, (uint64_t)i * DS4_N_EMBD * sizeof(float), (uint64_t)DS4_N_EMBD * sizeof(float)) && metal_graph_borrow_tensor_view( &split_row, &hc_split, (uint64_t)i * mix_hc * sizeof(float), mix_hc * sizeof(float)); if (ok) { ok = ds4_gpu_tensor_copy_xdev3_default_dst( metal_graph_attn_norm(dstg), &norm_row, norm_row.bytes, metal_graph_hc_split(dstg), &split_row, split_row.bytes, NULL, NULL, 0) != 0; } } return ok; } /* Batch Q/KV work while preserving each session's position. The optional KV * finalizer writes contiguous rows directly to the private raw caches. */ static bool metal_graph_encode_qkv_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t row_base, bool store_private_kv) { if (!items || count < 2 || !model || !layer || il >= DS4_N_LAYER || count > (int)DS4_GPU_ATTENTION_DECODE_BATCH_MAX) { return false; } ds4_gpu_graph *g = &items[0].session->graph; const int home = g->placement[il + 1u]; const uint32_t rows = (uint32_t)count; 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; if (home < 0 || (uint64_t)row_base + rows > g->prefill_cap) { return false; } for (int i = 0; i < count; i++) { if (!metal_graph_set_active_tier_decode( &items[i].session->graph, home)) { return false; } } ds4_gpu_tensor attn_norm, qr, qr_norm, q, kv_raw, kv; bool ok = metal_graph_borrow_tensor_view( &attn_norm, metal_graph_batch_attn_norm(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), (uint64_t)rows * DS4_N_EMBD * sizeof(float)) && metal_graph_borrow_tensor_view( &qr, metal_graph_batch_qr(g), (uint64_t)row_base * q_rank * sizeof(float), (uint64_t)rows * q_rank * sizeof(float)) && metal_graph_borrow_tensor_view( &qr_norm, metal_graph_batch_qr_norm(g), (uint64_t)row_base * q_rank * sizeof(float), (uint64_t)rows * q_rank * sizeof(float)) && metal_graph_borrow_tensor_view( &q, metal_graph_batch_q(g), (uint64_t)row_base * q_dim * sizeof(float), (uint64_t)rows * q_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &kv_raw, metal_graph_batch_kv_raw(g), (uint64_t)row_base * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)rows * DS4_N_HEAD_DIM * sizeof(float)) && metal_graph_borrow_tensor_view( &kv, metal_graph_batch_kv(g), (uint64_t)row_base * DS4_N_HEAD_DIM * sizeof(float), (uint64_t)rows * DS4_N_HEAD_DIM * sizeof(float)); ds4_gpu_attention_decode_row positions[DS4_GPU_ATTENTION_DECODE_BATCH_MAX]; memset(positions, 0, sizeof(positions)); for (uint32_t i = 0; i < rows; i++) { positions[i].pos = (uint32_t)items[i].session->checkpoint.len; } const bool compressed = ds4_layer_compress_ratio(il) != 0u; 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); } if (ok) { ok = ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor( &qr, &kv_raw, model->map, model->size, layer->attn_q_a->abs_offset, layer->attn_kv->abs_offset, DS4_N_EMBD, q_rank, DS4_N_HEAD_DIM, &attn_norm, rows) != 0; } if (ok) { ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( &qr_norm, &qr, model->map, model->size, layer->attn_q_a_norm->abs_offset, (uint32_t)q_rank, &kv, &kv_raw, layer->attn_kv_a_norm->abs_offset, DS4_N_HEAD_DIM, rows, DS4_RMS_EPS) != 0; } if (ok) { ok = ds4_gpu_rope_tail_decode_rows_tensor( &kv, positions, rows, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0u, 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_matmul_q8_0_decode_rows_exact_tensor( &q, model->map, model->size, layer->attn_q_b->abs_offset, q_rank, q_dim, &qr_norm, rows) != 0; } if (ok) { ok = ds4_gpu_head_rms_norm_tensor( &q, rows, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; } if (ok) { ok = ds4_gpu_rope_tail_decode_rows_tensor( &q, positions, rows, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0u, false, freq_base, freq_scale, ext_factor, attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; } if (ok && store_private_kv) { ds4_gpu_tensor *raw_caches[DS4_GPU_ATTENTION_DECODE_BATCH_MAX]; uint32_t raw_caps[DS4_GPU_ATTENTION_DECODE_BATCH_MAX]; uint32_t raw_rows[DS4_GPU_ATTENTION_DECODE_BATCH_MAX]; for (uint32_t i = 0; i < rows; i++) { ds4_gpu_graph *row_graph = &items[i].session->graph; raw_caches[i] = row_graph->layer_raw_cache[il]; raw_caps[i] = row_graph->raw_cap; raw_rows[i] = positions[i].pos % raw_caps[i]; } ok = ds4_gpu_kv_fp8_store_raw_decode_rows_tensor( &kv, raw_caches, raw_caps, raw_rows, rows, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; } return ok; } typedef struct { int tier; bool active; bool full_attention_rows; ds4_gpu_tensor attn_norm; ds4_gpu_tensor hc_split; ds4_gpu_tensor hc_pre; ds4_gpu_tensor hc_post; ds4_gpu_tensor hc_comb; ds4_gpu_tensor qr; ds4_gpu_tensor qr_norm; ds4_gpu_tensor q; ds4_gpu_tensor kv_raw; ds4_gpu_tensor kv; ds4_gpu_tensor comp_kv_cur; ds4_gpu_tensor comp_sc_cur; ds4_gpu_tensor indexer_q; ds4_gpu_tensor indexer_weights; ds4_gpu_tensor heads; ds4_gpu_tensor attn_low; ds4_gpu_tensor attn_out; ds4_gpu_tensor after_attn_hc; ds4_gpu_tensor *saved_attn_norm; ds4_gpu_tensor *saved_hc_split; ds4_gpu_tensor *saved_hc_pre; ds4_gpu_tensor *saved_hc_post; ds4_gpu_tensor *saved_hc_comb; ds4_gpu_tensor *saved_qr; ds4_gpu_tensor *saved_qr_norm; ds4_gpu_tensor *saved_q; ds4_gpu_tensor *saved_kv_raw; ds4_gpu_tensor *saved_kv; ds4_gpu_tensor *saved_comp_kv_cur; ds4_gpu_tensor *saved_comp_sc_cur; ds4_gpu_tensor *saved_indexer_q; ds4_gpu_tensor *saved_indexer_weights; ds4_gpu_tensor *saved_heads; ds4_gpu_tensor *saved_attn_low; ds4_gpu_tensor *saved_attn_out; ds4_gpu_tensor *saved_after_attn_hc; } metal_graph_attn_pre_alias; static bool metal_graph_bind_attn_pre_batch_row( ds4_gpu_graph *g, ds4_gpu_graph *batch_graph, const ds4_layer_weights *layer, uint32_t row, bool full_attention_rows, metal_graph_attn_pre_alias *alias) { if (!g || !batch_graph || !layer || !alias || g->active_tier < 0 || g->active_tier != batch_graph->active_tier) { return false; } memset(alias, 0, sizeof(*alias)); alias->tier = g->active_tier; const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; if (!metal_graph_borrow_tensor_view( &alias->attn_norm, metal_graph_batch_attn_norm(batch_graph), (uint64_t)row * DS4_N_EMBD * sizeof(float), (uint64_t)DS4_N_EMBD * sizeof(float)) || !metal_graph_borrow_tensor_view( &alias->hc_split, metal_graph_batch_hc_split(batch_graph), (uint64_t)row * mix_hc * sizeof(float), mix_hc * sizeof(float)) || !metal_graph_borrow_tensor_view( &alias->hc_pre, &alias->hc_split, 0, (uint64_t)DS4_N_HC * sizeof(float)) || !metal_graph_borrow_tensor_view( &alias->hc_post, &alias->hc_split, (uint64_t)DS4_N_HC * sizeof(float), (uint64_t)DS4_N_HC * sizeof(float)) || !metal_graph_borrow_tensor_view( &alias->hc_comb, &alias->hc_split, 2ull * DS4_N_HC * sizeof(float), (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float))) { return false; } if (full_attention_rows) { 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 comp_width_max = 2ull * DS4_N_HEAD_DIM; const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; #define DS4_BIND_ATTN_BATCH_ROW(member, batch_tensor, stride) \ metal_graph_borrow_tensor_view( \ &alias->member, (batch_tensor), \ (uint64_t)row * (stride) * sizeof(float), \ (stride) * sizeof(float)) if (!DS4_BIND_ATTN_BATCH_ROW(qr, metal_graph_batch_qr(batch_graph), q_rank) || !DS4_BIND_ATTN_BATCH_ROW(qr_norm, metal_graph_batch_qr_norm(batch_graph), q_rank) || !DS4_BIND_ATTN_BATCH_ROW(q, metal_graph_batch_q(batch_graph), q_dim) || !DS4_BIND_ATTN_BATCH_ROW(kv_raw, metal_graph_batch_kv_raw(batch_graph), DS4_N_HEAD_DIM) || !DS4_BIND_ATTN_BATCH_ROW(kv, metal_graph_batch_kv(batch_graph), DS4_N_HEAD_DIM) || !DS4_BIND_ATTN_BATCH_ROW(comp_kv_cur, metal_graph_batch_comp_kv(batch_graph), comp_width_max) || !DS4_BIND_ATTN_BATCH_ROW(comp_sc_cur, metal_graph_batch_comp_sc(batch_graph), comp_width_max) || !DS4_BIND_ATTN_BATCH_ROW(indexer_q, metal_graph_batch_indexer_q(batch_graph), indexer_q_dim) || !DS4_BIND_ATTN_BATCH_ROW(indexer_weights, metal_graph_batch_indexer_weights(batch_graph), DS4_N_INDEXER_HEAD) || !DS4_BIND_ATTN_BATCH_ROW(heads, metal_graph_batch_heads(batch_graph), q_dim) || !DS4_BIND_ATTN_BATCH_ROW(attn_low, metal_graph_batch_attn_low(batch_graph), low_dim) || !DS4_BIND_ATTN_BATCH_ROW(attn_out, metal_graph_batch_attn_out(batch_graph), DS4_N_EMBD) || !DS4_BIND_ATTN_BATCH_ROW(after_attn_hc, metal_graph_batch_after_attn_hc(batch_graph), hc_dim)) { #undef DS4_BIND_ATTN_BATCH_ROW return false; } #undef DS4_BIND_ATTN_BATCH_ROW } const int t = alias->tier; alias->saved_attn_norm = g->attn_norm_by_tier[t]; alias->saved_hc_split = g->hc_split_by_tier[t]; alias->saved_hc_pre = g->hc_pre_by_tier[t]; alias->saved_hc_post = g->hc_post_by_tier[t]; alias->saved_hc_comb = g->hc_comb_by_tier[t]; g->attn_norm_by_tier[t] = &alias->attn_norm; g->hc_split_by_tier[t] = &alias->hc_split; g->hc_pre_by_tier[t] = &alias->hc_pre; g->hc_post_by_tier[t] = &alias->hc_post; g->hc_comb_by_tier[t] = &alias->hc_comb; if (full_attention_rows) { alias->saved_qr = g->qr_by_tier[t]; alias->saved_qr_norm = g->qr_norm_by_tier[t]; alias->saved_q = g->q_by_tier[t]; alias->saved_kv_raw = g->kv_raw_by_tier[t]; alias->saved_kv = g->kv_by_tier[t]; alias->saved_comp_kv_cur = g->comp_kv_cur_by_tier[t]; alias->saved_comp_sc_cur = g->comp_sc_cur_by_tier[t]; alias->saved_indexer_q = g->indexer_q_by_tier[t]; alias->saved_indexer_weights = g->indexer_weights_by_tier[t]; alias->saved_heads = g->heads_by_tier[t]; alias->saved_attn_low = g->attn_low_by_tier[t]; alias->saved_attn_out = g->attn_out_by_tier[t]; alias->saved_after_attn_hc = g->after_attn_hc_by_tier[t]; g->qr_by_tier[t] = &alias->qr; g->qr_norm_by_tier[t] = &alias->qr_norm; g->q_by_tier[t] = &alias->q; g->kv_raw_by_tier[t] = &alias->kv_raw; g->kv_by_tier[t] = &alias->kv; g->comp_kv_cur_by_tier[t] = &alias->comp_kv_cur; g->comp_sc_cur_by_tier[t] = &alias->comp_sc_cur; g->indexer_q_by_tier[t] = &alias->indexer_q; g->indexer_weights_by_tier[t] = &alias->indexer_weights; g->heads_by_tier[t] = &alias->heads; g->attn_low_by_tier[t] = &alias->attn_low; g->attn_out_by_tier[t] = &alias->attn_out; g->after_attn_hc_by_tier[t] = &alias->after_attn_hc; alias->full_attention_rows = true; } alias->active = true; return true; } static void metal_graph_unbind_attn_pre_batch_row( ds4_gpu_graph *g, metal_graph_attn_pre_alias *alias) { if (!g || !alias || !alias->active) return; const int t = alias->tier; g->attn_norm_by_tier[t] = alias->saved_attn_norm; g->hc_split_by_tier[t] = alias->saved_hc_split; g->hc_pre_by_tier[t] = alias->saved_hc_pre; g->hc_post_by_tier[t] = alias->saved_hc_post; g->hc_comb_by_tier[t] = alias->saved_hc_comb; if (alias->full_attention_rows) { g->qr_by_tier[t] = alias->saved_qr; g->qr_norm_by_tier[t] = alias->saved_qr_norm; g->q_by_tier[t] = alias->saved_q; g->kv_raw_by_tier[t] = alias->saved_kv_raw; g->kv_by_tier[t] = alias->saved_kv; g->comp_kv_cur_by_tier[t] = alias->saved_comp_kv_cur; g->comp_sc_cur_by_tier[t] = alias->saved_comp_sc_cur; g->indexer_q_by_tier[t] = alias->saved_indexer_q; g->indexer_weights_by_tier[t] = alias->saved_indexer_weights; g->heads_by_tier[t] = alias->saved_heads; g->attn_low_by_tier[t] = alias->saved_attn_low; g->attn_out_by_tier[t] = alias->saved_attn_out; g->after_attn_hc_by_tier[t] = alias->saved_after_attn_hc; } alias->active = false; } static bool metal_graph_encode_attention_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t row_base) { #if defined(__APPLE__) (void)items; (void)count; (void)model; (void)layer; (void)il; (void)row_base; return false; #else if (!items || count < 2 || !model || !layer || il >= DS4_N_LAYER || count > (int)DS4_GPU_ATTENTION_DECODE_BATCH_MAX) { return false; } ds4_gpu_graph *batch_graph = &items[0].session->graph; const int home = batch_graph->active_tier; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; ds4_gpu_tensor q_rows, head_rows; if (home < 0 || (uint64_t)row_base + (uint32_t)count > batch_graph->prefill_cap || !metal_graph_borrow_tensor_view( &q_rows, metal_graph_batch_q(batch_graph), (uint64_t)row_base * q_dim * sizeof(float), (uint64_t)count * q_dim * sizeof(float)) || !metal_graph_borrow_tensor_view( &head_rows, metal_graph_batch_heads(batch_graph), (uint64_t)row_base * q_dim * sizeof(float), (uint64_t)count * q_dim * sizeof(float))) { return false; } ds4_gpu_attention_decode_row rows[DS4_GPU_ATTENTION_DECODE_BATCH_MAX]; memset(rows, 0, sizeof(rows)); for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); const uint32_t ratio = ds4_layer_compress_ratio(il); const uint32_t n_comp = g->layer_n_comp[il]; const bool indexed = ratio == 4u && n_comp > metal_graph_decode_indexer_sparse_threshold(g) && g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K; const uint32_t n_selected = indexed ? (DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] ? DS4_N_INDEXER_TOP_K : g->layer_n_index_comp[il]) : 0u; ds4_gpu_tensor *selected = indexed ? metal_graph_comp_selected(g) : NULL; if (g->active_tier != home || !g->layer_raw_cache[il] || (n_comp != 0u && !g->layer_attn_comp_cache[il]) || (indexed && (!selected || n_selected == 0u))) { return false; } rows[i].raw_kv = (uint64_t)(uintptr_t)g->layer_raw_cache[il]->ptr; rows[i].comp_kv = (uint64_t)(uintptr_t)( n_comp ? g->layer_attn_comp_cache[il]->ptr : g->layer_raw_cache[il]->ptr); rows[i].topk = selected ? (uint64_t)(uintptr_t)selected->ptr : 0u; rows[i].pos = pos; rows[i].n_raw = n_raw; rows[i].raw_cap = g->raw_cap; rows[i].raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); rows[i].n_comp = n_comp; rows[i].top_k = n_selected; rows[i].window = indexed ? g->raw_window : 0u; rows[i].ratio = indexed ? ratio : 0u; rows[i].indexed = indexed ? 1u : 0u; } const bool compressed = ds4_layer_compress_ratio(il) != 0u; 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); } return ds4_gpu_attention_decode_rows_rope_tensor( &head_rows, model->map, model->size, layer->attn_sinks->abs_offset, &q_rows, rows, (uint32_t)count, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0u, freq_base, freq_scale, ext_factor, attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; #endif } /* Keep grouped decode rows contiguous through the two TP output halves and * HC expansion. Each projection uses the batch-one warp reduction per row. */ static bool metal_graph_encode_attn_post_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t row_base) { #if defined(__APPLE__) (void)items; (void)count; (void)model; (void)layer; (void)il; (void)row_base; return false; #else if (!items || count < 2 || !model || !layer || il >= DS4_N_LAYER || (DS4_N_OUT_GROUP & 1u) != 0u) { return false; } ds4_gpu_graph *g = &items[0].session->graph; const int home = g->placement[il + 1u]; const int partner = metal_graph_cuda_tp_partner_tier(home); const uint32_t rows = (uint32_t)count; const uint32_t tp_groups = DS4_N_OUT_GROUP / 2u; const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); 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 tp_low_dim = (uint64_t)tp_groups * DS4_N_LORA_O; 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; if (home < 0 || partner < 0 || (uint64_t)row_base + rows > g->prefill_cap) { return false; } for (int i = 0; i < count; i++) { if (!metal_graph_set_active_tier_decode( &items[i].session->graph, home)) { return false; } } ds4_gpu_tensor heads, home_low, home_out; ds4_gpu_tensor peer_low, peer_out, peer_return; ds4_gpu_tensor residual_hc, split, after_attn_hc; const uint64_t out_bytes = (uint64_t)rows * DS4_N_EMBD * sizeof(float); bool ok = metal_graph_borrow_tensor_view( &heads, metal_graph_batch_heads(g), (uint64_t)row_base * q_dim * sizeof(float), (uint64_t)rows * q_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &home_low, metal_graph_batch_attn_low(g), (uint64_t)row_base * low_dim * sizeof(float), (uint64_t)rows * tp_low_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &home_out, metal_graph_batch_attn_out(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), out_bytes) && metal_graph_borrow_tensor_view( &peer_low, g->batch_attn_low_by_tier[partner], (uint64_t)row_base * low_dim * sizeof(float), (uint64_t)rows * tp_low_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &peer_out, g->batch_attn_out_by_tier[partner], (uint64_t)row_base * DS4_N_EMBD * sizeof(float), out_bytes) && metal_graph_borrow_tensor_view( &peer_return, metal_graph_batch_shared_out(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), out_bytes) && metal_graph_borrow_tensor_view( &residual_hc, metal_graph_batch_cur_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), (uint64_t)rows * hc_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &split, metal_graph_batch_hc_split(g), (uint64_t)row_base * mix_hc * sizeof(float), (uint64_t)rows * mix_hc * sizeof(float)) && metal_graph_borrow_tensor_view( &after_attn_hc, metal_graph_batch_after_attn_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), (uint64_t)rows * hc_dim * sizeof(float)); if (ok) ok = ds4_gpu_tensor_wait_xdev_default(&heads, partner) != 0; if (ok) ok = ds4_gpu_set_current_device(partner) == 0; if (ok) { ok = ds4_gpu_attention_output_low_q8_rows_exact_tensor( &peer_low, model->map, model->size, layer->attn_output_a->abs_offset, group_dim, DS4_N_LORA_O, DS4_N_OUT_GROUP, tp_groups, tp_groups, &heads, rows) != 0; } if (ok) { ok = ds4_gpu_matmul_q8_0_kslice_rows_tensor( &peer_out, model->map, model->size, layer->attn_output_b->abs_offset, low_dim, DS4_N_EMBD, tp_low_dim, tp_low_dim, &peer_low, rows) != 0; } if (ds4_gpu_set_current_device(home) != 0) ok = false; if (ok) { ok = ds4_gpu_attention_output_low_q8_rows_exact_tensor( &home_low, model->map, model->size, layer->attn_output_a->abs_offset, group_dim, DS4_N_LORA_O, DS4_N_OUT_GROUP, 0u, tp_groups, &heads, rows) != 0; } if (ok) { ok = ds4_gpu_matmul_q8_0_kslice_rows_tensor( &home_out, model->map, model->size, layer->attn_output_b->abs_offset, low_dim, DS4_N_EMBD, 0u, tp_low_dim, &home_low, rows) != 0; } if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &peer_return, &peer_out, out_bytes) != 0; } if (ok) { ok = ds4_gpu_hc_expand_add_split_tensor( &after_attn_hc, &home_out, &peer_return, &residual_hc, &split, DS4_N_EMBD, DS4_N_HC) != 0; } return ok; #endif } static bool metal_graph_session_batch_ffn_pre_supported( ds4_decode_item *items, int count, const ds4_weights *weights) { if (!items || count < 3 || !weights) return false; ds4_gpu_graph *first = &items[0].session->graph; if ((uint32_t)count > first->prefill_cap || !metal_graph_prefill_tokens(first) || ds4_gpu_tensor_bytes(metal_graph_prefill_tokens(first)) < (uint64_t)count * sizeof(int32_t)) { return false; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; if (!layer->hc_ffn_fn || !layer->ffn_gate_inp || layer->hc_ffn_fn->type != DS4_TENSOR_F16 || layer->ffn_gate_inp->type != DS4_TENSOR_F16 || layer->ffn_gate_inp->dim[0] != 4096u || layer->ffn_gate_inp->dim[1] != 256u) { return false; } const int home = first->placement[il + 1u]; if (!first->batch_after_attn_hc_by_tier[home] || !first->batch_flat_hc_by_tier[home] || !first->batch_hc_mix_by_tier[home] || !first->batch_hc_split_by_tier[home] || !first->batch_ffn_cur_by_tier[home] || !first->batch_ffn_norm_by_tier[home] || !first->batch_router_logits_by_tier[home] || !first->batch_router_probs_by_tier[home] || !first->batch_router_selected_by_tier[home] || !first->batch_router_weights_by_tier[home]) { return false; } for (int i = 1; i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; if (!g->placement || g->placement[il + 1u] != home) return false; } } return true; } /* Gather the cache-independent FFN input rows after attention, then execute * HC mixing, normalization, router projection, and selection as row batches. * Routed and shared experts consume these batch tensors directly. */ static bool metal_graph_encode_ffn_pre_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t row_base, bool gather_after_attn_rows) { if (!items || count < 2 || !model || !layer) return false; ds4_gpu_graph *g = &items[0].session->graph; const int home_tier = g->active_tier; 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 hc_bytes = (uint64_t)count * hc_dim * sizeof(float); const uint64_t mix_bytes = (uint64_t)count * mix_hc * sizeof(float); const uint64_t norm_bytes = (uint64_t)count * DS4_N_EMBD * sizeof(float); const uint64_t router_bytes = (uint64_t)count * DS4_N_EXPERT * sizeof(float); const uint64_t selected_bytes = (uint64_t)count * DS4_N_EXPERT_USED * sizeof(int32_t); const uint64_t weights_bytes = (uint64_t)count * DS4_N_EXPERT_USED * sizeof(float); ds4_gpu_tensor after_attn_hc, flat_hc, hc_mix, hc_split; ds4_gpu_tensor ffn_cur, ffn_norm, router_logits, router_probs; ds4_gpu_tensor router_selected, router_weights, router_tokens; bool ok = home_tier >= 0 && (uint64_t)row_base + (uint32_t)count <= g->prefill_cap && metal_graph_borrow_tensor_view(&after_attn_hc, metal_graph_batch_after_attn_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes) && metal_graph_borrow_tensor_view(&flat_hc, metal_graph_batch_flat_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes) && metal_graph_borrow_tensor_view(&hc_mix, metal_graph_batch_hc_mix(g), (uint64_t)row_base * mix_hc * sizeof(float), mix_bytes) && metal_graph_borrow_tensor_view(&hc_split, metal_graph_batch_hc_split(g), (uint64_t)row_base * mix_hc * sizeof(float), mix_bytes) && metal_graph_borrow_tensor_view(&ffn_cur, metal_graph_batch_ffn_cur(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&ffn_norm, metal_graph_batch_ffn_norm(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&router_logits, metal_graph_batch_router_logits(g), (uint64_t)row_base * DS4_N_EXPERT * sizeof(float), router_bytes) && metal_graph_borrow_tensor_view(&router_probs, metal_graph_batch_router_probs(g), (uint64_t)row_base * DS4_N_EXPERT * sizeof(float), router_bytes) && metal_graph_borrow_tensor_view(&router_selected, metal_graph_batch_router_selected(g), (uint64_t)row_base * DS4_N_EXPERT_USED * sizeof(int32_t), selected_bytes) && metal_graph_borrow_tensor_view(&router_weights, metal_graph_batch_router_weights(g), (uint64_t)row_base * DS4_N_EXPERT_USED * sizeof(float), weights_bytes) && metal_graph_borrow_tensor_view(&router_tokens, metal_graph_prefill_tokens(g), (uint64_t)row_base * sizeof(int32_t), (uint64_t)count * sizeof(int32_t)); for (int i = 0; gather_after_attn_rows && ok && i < count; i++) { ds4_gpu_graph *srcg = &items[i].session->graph; ds4_gpu_tensor dst_row; ok = srcg->active_tier == home_tier && metal_graph_borrow_tensor_view( &dst_row, &after_attn_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)); if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &dst_row, metal_graph_after_attn_hc(srcg), dst_row.bytes) != 0; } } if (ok) { ok = ds4_gpu_rms_norm_plain_rows_tensor( &flat_hc, &after_attn_hc, (uint32_t)hc_dim, (uint32_t)count, DS4_RMS_EPS) != 0; } for (int i = 0; ok && i < count; i++) { ds4_gpu_tensor out_row, in_row; ok = metal_graph_borrow_tensor_view( &out_row, &hc_mix, (uint64_t)i * mix_hc * sizeof(float), mix_hc * sizeof(float)) && metal_graph_borrow_tensor_view( &in_row, &flat_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)); if (ok) { ok = metal_graph_matmul_plain_tensor( &out_row, model, layer->hc_ffn_fn, hc_dim, mix_hc, &in_row, 1); } } if (ok) { ok = ds4_gpu_hc_split_weighted_sum_norm_tensor( &ffn_cur, &ffn_norm, &hc_split, &hc_mix, &after_attn_hc, 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 = ds4_gpu_matmul_f16_router_rows_exact_tensor( &router_logits, model->map, model->size, layer->ffn_gate_inp->abs_offset, &ffn_norm, (uint32_t)count) != 0; } if (ok) { ok = ds4_gpu_router_select_batch_tensor( &router_selected, &router_weights, &router_probs, 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, &router_logits, &router_tokens, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, (uint32_t)count) != 0; } return ok; } /* Batch the routed part of one decode layer across independent sessions. The * router and KV-bearing attention path have already run in each session. We * gather only normalized rows and routing metadata, execute the existing * expert-sorted multi-row kernels on both EP owners, then reduce each row in * the exact decode slot grouping before scattering it back. */ static bool metal_graph_encode_routed_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t row_base, bool gather_rows, bool scatter_rows) { if (!items || count < 2 || !model || !layer || DS4_N_EXPERT_USED != 6u) { return false; } ds4_gpu_graph *g = &items[0].session->graph; const int home_tier = g->active_tier; const int partner_tier = g->cuda_tp_decode ? metal_graph_cuda_tp_partner_tier(home_tier) : -1; if (!g->cuda_tp_ep || partner_tier < 0 || (uint64_t)row_base + (uint32_t)count > g->prefill_cap || !g_gpu_peer_ok[partner_tier][home_tier]) { return false; } 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 uint64_t norm_bytes = (uint64_t)count * DS4_N_EMBD * sizeof(float); const uint64_t selected_bytes = (uint64_t)count * DS4_N_EXPERT_USED * sizeof(int32_t); const uint64_t weights_bytes = (uint64_t)count * DS4_N_EXPERT_USED * sizeof(float); const uint64_t slots_bytes = (uint64_t)count * DS4_N_EXPERT_USED * routed_out_dim * sizeof(float); const uint64_t output_bytes = (uint64_t)count * routed_out_dim * sizeof(float); const uint64_t norm_offset = (uint64_t)row_base * DS4_N_EMBD * sizeof(float); const uint64_t selected_offset = (uint64_t)row_base * DS4_N_EXPERT_USED * sizeof(int32_t); const uint64_t weights_offset = (uint64_t)row_base * DS4_N_EXPERT_USED * sizeof(float); const uint64_t slots_offset = (uint64_t)row_base * DS4_N_EXPERT_USED * routed_out_dim * sizeof(float); const uint64_t output_offset = (uint64_t)row_base * routed_out_dim * sizeof(float); ds4_gpu_tensor local_norm, local_selected, local_weights, original_selected; ds4_gpu_tensor local_out, local_down; ds4_gpu_tensor peer_norm, peer_selected, peer_weights, peer_out, peer_down; bool ok = metal_graph_borrow_tensor_view(&local_norm, metal_graph_batch_ffn_norm(g), norm_offset, norm_bytes) && metal_graph_borrow_tensor_view(&local_selected, metal_graph_batch_router_selected(g), selected_offset, selected_bytes) && metal_graph_borrow_tensor_view(&local_weights, metal_graph_batch_router_weights(g), weights_offset, weights_bytes) && metal_graph_borrow_tensor_view(&original_selected, metal_graph_batch_router_logits(g), selected_offset, selected_bytes) && metal_graph_borrow_tensor_view(&local_out, metal_graph_batch_routed_out(g), output_offset, output_bytes) && metal_graph_borrow_tensor_view(&local_down, metal_graph_batch_routed_down(g), slots_offset, slots_bytes) && metal_graph_borrow_tensor_view(&peer_norm, g->batch_ffn_norm_by_tier[partner_tier], norm_offset, norm_bytes) && metal_graph_borrow_tensor_view(&peer_selected, g->batch_router_selected_by_tier[partner_tier], selected_offset, selected_bytes) && metal_graph_borrow_tensor_view(&peer_weights, g->batch_router_weights_by_tier[partner_tier], weights_offset, weights_bytes) && metal_graph_borrow_tensor_view(&peer_out, g->batch_routed_out_by_tier[partner_tier], output_offset, output_bytes) && metal_graph_borrow_tensor_view(&peer_down, g->batch_routed_down_by_tier[partner_tier], slots_offset, slots_bytes); for (int i = 0; gather_rows && ok && i < count; i++) { ds4_gpu_graph *srcg = &items[i].session->graph; if (srcg->active_tier != home_tier) { ok = false; break; } ds4_gpu_tensor norm_row, weights_row, original_row; ok = metal_graph_borrow_tensor_view( &norm_row, &local_norm, (uint64_t)i * DS4_N_EMBD * sizeof(float), (uint64_t)DS4_N_EMBD * sizeof(float)) && metal_graph_borrow_tensor_view( &weights_row, &local_weights, (uint64_t)i * DS4_N_EXPERT_USED * sizeof(float), (uint64_t)DS4_N_EXPERT_USED * sizeof(float)) && metal_graph_borrow_tensor_view( &original_row, &original_selected, (uint64_t)i * DS4_N_EXPERT_USED * sizeof(int32_t), (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); if (ok) { ok = ds4_gpu_tensor_copy_xdev3_default_dst( &norm_row, metal_graph_ffn_norm(srcg), norm_row.bytes, &original_row, metal_graph_router_selected(srcg), original_row.bytes, &weights_row, metal_graph_router_weights(srcg), weights_row.bytes) != 0; } } if (ok && gather_rows) { ok = ds4_gpu_tensor_copy_xdev_default( &local_selected, &original_selected, selected_bytes) != 0; } else if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &original_selected, &local_selected, selected_bytes) != 0; } if (ok) { ok = ds4_gpu_tensor_copy_xdev3_default_dst( &peer_norm, &local_norm, norm_bytes, &peer_selected, &local_selected, selected_bytes, &peer_weights, &local_weights, weights_bytes) != 0; } bool peer_mid_is_f16 = false; if (ok) ok = ds4_gpu_set_current_device(partner_tier) == 0; if (ok) { ok = ds4_gpu_routed_moe_batch_owned_tensor( &peer_out, g->batch_routed_gate_by_tier[partner_tier], g->batch_routed_up_by_tier[partner_tier], g->batch_routed_mid_by_tier[partner_tier], &peer_down, 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, DS4_N_EXPERT_USED, DS4_N_EXPERT / 2u, DS4_N_EXPERT - DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &peer_norm, il, (uint32_t)count, &peer_mid_is_f16) != 0; } if (ds4_gpu_set_current_device(home_tier) != 0) ok = false; if (ok) { ok = ds4_gpu_routed_moe_batch_owned_tensor( &local_out, metal_graph_batch_routed_gate(g), metal_graph_batch_routed_up(g), metal_graph_batch_routed_mid(g), &local_down, 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, DS4_N_EXPERT_USED, 0, DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &local_norm, il, (uint32_t)count, &g->batch_routed_mid_is_f16) != 0; } if (ok) { ok = ds4_gpu_tensor_wait_xdev_default(&peer_down, home_tier) != 0; } const bool combine_rows = metal_graph_tp_env_flag( "DS4_CUDA_SESSION_BATCH_MOE_COMBINE_ROWS", true); if (ok && combine_rows) { ok = ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( &local_out, &local_down, &peer_down, &original_selected, (uint32_t)routed_out_dim, DS4_N_EXPERT / 2u, (uint32_t)count) != 0; } for (int i = 0; ok && (!combine_rows || scatter_rows) && i < count; i++) { const uint64_t slot_stride = (uint64_t)DS4_N_EXPERT_USED * routed_out_dim * sizeof(float); const uint64_t out_stride = routed_out_dim * sizeof(float); ds4_gpu_tensor out_row, home_slots, peer_slots, selected_row; ok = metal_graph_borrow_tensor_view( &out_row, &local_out, (uint64_t)i * out_stride, out_stride) && metal_graph_borrow_tensor_view( &home_slots, &local_down, (uint64_t)i * slot_stride, slot_stride) && metal_graph_borrow_tensor_view( &peer_slots, &peer_down, (uint64_t)i * slot_stride, slot_stride) && metal_graph_borrow_tensor_view( &selected_row, &original_selected, (uint64_t)i * DS4_N_EXPERT_USED * sizeof(int32_t), (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); if (ok && !combine_rows) { ok = ds4_gpu_routed_moe_owned_slots_combine_tensor( &out_row, &home_slots, &peer_slots, &selected_row, (uint32_t)routed_out_dim, DS4_N_EXPERT / 2u) != 0; } if (ok && scatter_rows) { ok = ds4_gpu_tensor_copy_xdev_default( metal_graph_routed_out(&items[i].session->graph), &out_row, out_stride) != 0; } } return ok; } /* Run one owned-expert dispatch over prefill rows followed by exact decode * rows. The expert kernels see the larger combined batch, but reduction is * path-specific: prefill keeps the existing two-owner partial sum and decode * keeps its fixed six-slot order in the shared row workspace. */ 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) { if (!g || decode_count < 0 || (decode_count > 0 && !decode_items) || !model || !layer || prefill_rows == 0 || DS4_N_EXPERT_USED != 6u) { return false; } const uint32_t decode_rows = (uint32_t)decode_count; const uint32_t total_rows = prefill_rows + decode_rows; const int home_tier = g->active_tier; const int partner_tier = g->cuda_tp_decode ? metal_graph_cuda_tp_partner_tier(home_tier) : -1; if (total_rows < prefill_rows || total_rows > g->prefill_cap || !g->cuda_tp_ep || partner_tier < 0 || !g_gpu_peer_ok[partner_tier][home_tier] || !metal_graph_ensure_batch_ffn_out_on(g, home_tier)) { return false; } for (int i = 0; i < decode_count; i++) { if (!decode_items[i].session || decode_items[i].session->graph.active_tier != home_tier) { return false; } } 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 uint64_t norm_bytes = (uint64_t)total_rows * DS4_N_EMBD * sizeof(float); const uint64_t selected_bytes = (uint64_t)total_rows * DS4_N_EXPERT_USED * sizeof(int32_t); const uint64_t weights_bytes = (uint64_t)total_rows * DS4_N_EXPERT_USED * sizeof(float); const uint64_t slots_bytes = (uint64_t)total_rows * DS4_N_EXPERT_USED * routed_out_dim * sizeof(float); const uint64_t output_bytes = (uint64_t)total_rows * routed_out_dim * sizeof(float); const uint64_t prefill_output_bytes = (uint64_t)prefill_rows * routed_out_dim * sizeof(float); ds4_gpu_tensor local_norm, local_selected, local_weights, original_selected; ds4_gpu_tensor local_out, local_down, prefill_peer_out; ds4_gpu_tensor peer_norm, peer_selected, peer_weights, peer_out, peer_down; ds4_gpu_tensor peer_prefill_out; bool ok = metal_graph_borrow_tensor_view(&local_norm, metal_graph_batch_ffn_norm(g), 0, norm_bytes) && metal_graph_borrow_tensor_view(&local_selected, metal_graph_batch_router_selected(g), 0, selected_bytes) && metal_graph_borrow_tensor_view(&local_weights, metal_graph_batch_router_weights(g), 0, weights_bytes) && metal_graph_borrow_tensor_view(&original_selected, metal_graph_batch_router_logits(g), 0, selected_bytes) && metal_graph_borrow_tensor_view(&local_out, metal_graph_batch_routed_out(g), 0, output_bytes) && metal_graph_borrow_tensor_view(&local_down, metal_graph_batch_routed_down(g), 0, slots_bytes) && metal_graph_borrow_tensor_view(&prefill_peer_out, g->batch_ffn_out_by_tier[home_tier], 0, prefill_output_bytes) && metal_graph_borrow_tensor_view(&peer_norm, g->batch_ffn_norm_by_tier[partner_tier], 0, norm_bytes) && metal_graph_borrow_tensor_view(&peer_selected, g->batch_router_selected_by_tier[partner_tier], 0, selected_bytes) && metal_graph_borrow_tensor_view(&peer_weights, g->batch_router_weights_by_tier[partner_tier], 0, weights_bytes) && metal_graph_borrow_tensor_view(&peer_out, g->batch_routed_out_by_tier[partner_tier], 0, output_bytes) && metal_graph_borrow_tensor_view(&peer_down, g->batch_routed_down_by_tier[partner_tier], 0, slots_bytes) && metal_graph_borrow_tensor_view(&peer_prefill_out, &peer_out, 0, prefill_output_bytes); if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &original_selected, &local_selected, selected_bytes) != 0 && ds4_gpu_tensor_copy_xdev3_default_dst( &peer_norm, &local_norm, norm_bytes, &peer_selected, &local_selected, selected_bytes, &peer_weights, &local_weights, weights_bytes) != 0; } bool peer_mid_is_f16 = false; if (ok) ok = ds4_gpu_set_current_device(partner_tier) == 0; if (ok) { ok = ds4_gpu_routed_moe_batch_owned_tensor( &peer_out, g->batch_routed_gate_by_tier[partner_tier], g->batch_routed_up_by_tier[partner_tier], g->batch_routed_mid_by_tier[partner_tier], &peer_down, 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, DS4_N_EXPERT_USED, DS4_N_EXPERT / 2u, DS4_N_EXPERT - DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &peer_norm, il, total_rows, &peer_mid_is_f16) != 0; } if (ds4_gpu_set_current_device(home_tier) != 0) ok = false; if (ok) { ok = ds4_gpu_routed_moe_batch_owned_tensor( &local_out, metal_graph_batch_routed_gate(g), metal_graph_batch_routed_up(g), metal_graph_batch_routed_mid(g), &local_down, 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, DS4_N_EXPERT_USED, 0, DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &local_norm, il, total_rows, &g->batch_routed_mid_is_f16) != 0; } if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( &prefill_peer_out, &peer_prefill_out, prefill_output_bytes) != 0; } if (ok) { ds4_gpu_tensor prefill_home_out; ok = metal_graph_borrow_tensor_view( &prefill_home_out, &local_out, 0, prefill_output_bytes) && ds4_gpu_add_tensor( &prefill_home_out, &prefill_home_out, &prefill_peer_out, (uint32_t)((uint64_t)prefill_rows * routed_out_dim)) != 0; } if (ok && decode_rows > 0) { ok = ds4_gpu_tensor_wait_xdev_default(&peer_down, home_tier) != 0; } const uint64_t slot_stride = (uint64_t)DS4_N_EXPERT_USED * routed_out_dim * sizeof(float); const uint64_t out_stride = routed_out_dim * sizeof(float); const bool combine_rows = metal_graph_tp_env_flag( "DS4_CUDA_SESSION_BATCH_MOE_COMBINE_ROWS", true); const bool scatter_rows = metal_graph_tp_env_flag( "DS4_CUDA_MIXED_ROUTED_SCATTER", false); if (ok && combine_rows && decode_rows > 0) { ds4_gpu_tensor decode_out, home_slots, peer_slots, selected; ok = metal_graph_borrow_tensor_view( &decode_out, &local_out, prefill_output_bytes, (uint64_t)decode_rows * out_stride) && metal_graph_borrow_tensor_view( &home_slots, &local_down, (uint64_t)prefill_rows * slot_stride, (uint64_t)decode_rows * slot_stride) && metal_graph_borrow_tensor_view( &peer_slots, &peer_down, (uint64_t)prefill_rows * slot_stride, (uint64_t)decode_rows * slot_stride) && metal_graph_borrow_tensor_view( &selected, &original_selected, (uint64_t)prefill_rows * DS4_N_EXPERT_USED * sizeof(int32_t), (uint64_t)decode_rows * DS4_N_EXPERT_USED * sizeof(int32_t)); if (ok) { ok = ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( &decode_out, &home_slots, &peer_slots, &selected, (uint32_t)routed_out_dim, DS4_N_EXPERT / 2u, decode_rows) != 0; } } for (int i = 0; ok && (!combine_rows || scatter_rows) && i < decode_count; i++) { const uint32_t row = prefill_rows + (uint32_t)i; ds4_gpu_tensor out_row, home_slots, peer_slots, selected_row; ok = metal_graph_borrow_tensor_view( &out_row, &local_out, (uint64_t)row * out_stride, out_stride) && metal_graph_borrow_tensor_view( &home_slots, &local_down, (uint64_t)row * slot_stride, slot_stride) && metal_graph_borrow_tensor_view( &peer_slots, &peer_down, (uint64_t)row * slot_stride, slot_stride) && metal_graph_borrow_tensor_view( &selected_row, &original_selected, (uint64_t)row * DS4_N_EXPERT_USED * sizeof(int32_t), (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); if (ok && !combine_rows) { ok = ds4_gpu_routed_moe_owned_slots_combine_tensor( &out_row, &home_slots, &peer_slots, &selected_row, (uint32_t)routed_out_dim, DS4_N_EXPERT / 2u) != 0; } if (ok && scatter_rows) { ok = ds4_gpu_tensor_copy_xdev_default( metal_graph_routed_out(&decode_items[i].session->graph), &out_row, out_stride) != 0; } } return ok; } static bool metal_graph_encode_shared_rows_exact( ds4_gpu_tensor *shared_gate, ds4_gpu_tensor *shared_up, ds4_gpu_tensor *shared_mid, ds4_gpu_tensor *shared_out, ds4_gpu_tensor *norm, const ds4_model *model, const ds4_layer_weights *layer, uint32_t rows) { if (!shared_gate || !shared_up || !shared_mid || !shared_out || !norm || !model || !layer || rows == 0) { return false; } const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; bool ok = ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor( shared_gate, shared_up, model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, shared_dim, norm, rows) != 0; if (ok) { ok = ds4_gpu_swiglu_tensor( shared_mid, shared_gate, shared_up, (uint32_t)((uint64_t)rows * shared_dim), DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } if (ok) { ok = ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, shared_mid, rows) != 0; } return ok; } /* Finish the cache-independent half of one FFN for all session rows together. * The routed outputs and normalized inputs already occupy the first graph's * shared prefill workspace. Gather only the per-session HC residual/mix, * execute the shared expert as exact row batches, and scatter final HC rows. */ static bool metal_graph_encode_shared_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_layer_weights *layer, uint32_t row_base, bool gather_hc_rows) { if (!items || count < 2 || !model || !layer) return false; ds4_gpu_graph *g = &items[0].session->graph; const int home_tier = g->active_tier; 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 norm_bytes = (uint64_t)count * DS4_N_EMBD * sizeof(float); const uint64_t shared_bytes = (uint64_t)count * shared_dim * sizeof(float); const uint64_t hc_bytes = (uint64_t)count * hc_dim * sizeof(float); const uint64_t split_bytes = (uint64_t)count * mix_hc * sizeof(float); ds4_gpu_tensor norm, routed, shared_gate, shared_up, shared_mid, shared_out; ds4_gpu_tensor residual_hc, split, next_hc; bool ok = home_tier >= 0 && (uint64_t)row_base + (uint32_t)count <= g->prefill_cap && metal_graph_borrow_tensor_view(&norm, metal_graph_batch_ffn_norm(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&routed, metal_graph_batch_routed_out(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&shared_gate, metal_graph_batch_shared_gate(g), (uint64_t)row_base * shared_dim * sizeof(float), shared_bytes) && metal_graph_borrow_tensor_view(&shared_up, metal_graph_batch_shared_up(g), (uint64_t)row_base * shared_dim * sizeof(float), shared_bytes) && metal_graph_borrow_tensor_view(&shared_mid, metal_graph_batch_shared_mid(g), (uint64_t)row_base * shared_dim * sizeof(float), shared_bytes) && metal_graph_borrow_tensor_view(&shared_out, metal_graph_batch_shared_out(g), (uint64_t)row_base * DS4_N_EMBD * sizeof(float), norm_bytes) && metal_graph_borrow_tensor_view(&residual_hc, metal_graph_batch_after_attn_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes) && metal_graph_borrow_tensor_view(&split, metal_graph_batch_hc_split(g), (uint64_t)row_base * mix_hc * sizeof(float), split_bytes) && metal_graph_borrow_tensor_view(&next_hc, metal_graph_batch_next_hc(g), (uint64_t)row_base * hc_dim * sizeof(float), hc_bytes); for (int i = 0; gather_hc_rows && ok && i < count; i++) { ds4_gpu_graph *srcg = &items[i].session->graph; ds4_gpu_tensor residual_row, split_row; ok = srcg->active_tier == home_tier && metal_graph_borrow_tensor_view( &residual_row, &residual_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)) && metal_graph_borrow_tensor_view( &split_row, &split, (uint64_t)i * mix_hc * sizeof(float), mix_hc * sizeof(float)); if (ok) { ok = ds4_gpu_tensor_copy_xdev3_default_dst( &residual_row, metal_graph_after_attn_hc(srcg), residual_row.bytes, &split_row, metal_graph_hc_split(srcg), split_row.bytes, NULL, NULL, 0) != 0; } } if (ok) { ok = metal_graph_encode_shared_rows_exact( &shared_gate, &shared_up, &shared_mid, &shared_out, &norm, model, layer, (uint32_t)count); } if (ok) { ok = ds4_gpu_hc_expand_add_split_tensor( &next_hc, &routed, &shared_out, &residual_hc, &split, DS4_N_EMBD, DS4_N_HC) != 0; } for (int i = 0; ok && i < count; i++) { ds4_gpu_tensor next_row; ok = metal_graph_borrow_tensor_view( &next_row, &next_hc, (uint64_t)i * hc_dim * sizeof(float), hc_dim * sizeof(float)); if (ok) { ok = ds4_gpu_tensor_copy_xdev_default( metal_graph_after_ffn_hc(&items[i].session->graph), &next_row, next_row.bytes) != 0; } } return ok; } /* Submit independent decode graphs stage by stage. A whole-graph ordering * puts session A's all-device output head ahead of session B's first stage, * which serializes the nominal pipeline. Stage-major submission leaves each * session's kernel arithmetic untouched while allowing B's previous stage to * overlap A's next stage. */ static bool metal_graph_encode_session_pipeline_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_weights *weights) { if (!items || count <= 0 || !model || !weights) return false; ds4_gpu_graph *first = &items[0].session->graph; if (!first->placement) return false; const bool group_routed_moe = metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_MOE", true) && metal_graph_session_batch_moe_supported(items, count, weights); const bool group_shared_ffn = group_routed_moe && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_SHARED", true) && metal_graph_session_batch_shared_supported(items, count, weights); const bool group_ffn_pre = group_shared_ffn && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_FFN_PRE", true) && metal_graph_session_batch_ffn_pre_supported(items, count, weights); const bool group_attn_pre = group_ffn_pre && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_PRE", true) && metal_graph_session_batch_attn_pre_supported(items, count, weights); const bool alias_attn_pre = group_attn_pre && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_ALIAS", true); const bool group_attn_core = alias_attn_pre && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_CORE", true) && metal_graph_session_batch_attn_core_supported(items, count, weights); const bool group_qkv = group_attn_core && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_QKV", true) && metal_graph_session_batch_qkv_supported(items, count, weights); const bool group_kv_store = group_qkv && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_KV_STORE", true) && metal_graph_session_batch_kv_store_supported(items, count); const bool group_attn_post = group_attn_core && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_POST", true) && metal_graph_session_batch_attn_post_supported(items, count, weights); bool ok = true; if (group_ffn_pre) { int32_t *router_tokens = malloc((size_t)count * sizeof(*router_tokens)); if (!router_tokens) return false; for (int i = 0; i < count; i++) router_tokens[i] = items[i].token; ok = ds4_gpu_tensor_write( metal_graph_prefill_tokens(first), 0, router_tokens, (uint64_t)count * sizeof(*router_tokens)) != 0; free(router_tokens); } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; if (!g->placement || g->raw_cap == 0) return false; metal_graph_dspark_capture_begin(g); if (!metal_graph_set_active_tier_decode(g, g->emb_tier)) return false; 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)items[i].token, DS4_N_EMBD, DS4_N_HC) != 0; } uint32_t stage_begin = 0; while (ok && stage_begin < DS4_N_LAYER) { const int stage_tier = first->placement[stage_begin + 1u]; uint32_t stage_end = stage_begin + 1u; while (stage_end < DS4_N_LAYER && first->placement[stage_end + 1u] == stage_tier) { stage_end++; } if (group_routed_moe) { for (uint32_t il = stage_begin; ok && il < stage_end; il++) { if (group_attn_pre) { ok = metal_graph_encode_attn_pre_session_batch( items, count, model, &weights->layer[il], il, 0, !alias_attn_pre); } if (ok && group_qkv) { ok = metal_graph_encode_qkv_session_batch( items, count, model, &weights->layer[il], il, 0, group_kv_store); } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); if (g->placement[il + 1u] != stage_tier) { ok = false; break; } metal_graph_attn_pre_alias attn_alias = {0}; if (alias_attn_pre) { ok = metal_graph_bind_attn_pre_batch_row( g, first, &weights->layer[il], (uint32_t)i, group_attn_core, &attn_alias); } if (ok) { ok = metal_graph_encode_decode_layer_phase( g, model, &weights->layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, group_kv_store ? METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN : group_qkv ? METAL_DECODE_LAYER_FROM_QKV_TO_ATTN : group_attn_core ? METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN : group_attn_pre ? METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN : group_ffn_pre ? METAL_DECODE_LAYER_TO_FFN : METAL_DECODE_LAYER_TO_ROUTER); } metal_graph_unbind_attn_pre_batch_row(g, &attn_alias); } if (ok && group_attn_core) { ok = metal_graph_encode_attention_session_batch( items, count, model, &weights->layer[il], il, 0); } if (ok && group_attn_post) { ok = metal_graph_encode_attn_post_session_batch( items, count, model, &weights->layer[il], il, 0); } for (int i = 0; ok && group_attn_core && !group_attn_post && i < count; i++) { ds4_gpu_graph *g = &items[i].session->graph; const uint32_t pos = (uint32_t)items[i].session->checkpoint.len; 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_attn_pre_alias attn_alias = {0}; ok = metal_graph_bind_attn_pre_batch_row( g, first, &weights->layer[il], (uint32_t)i, true, &attn_alias); if (ok) { ok = metal_graph_encode_decode_layer_phase( g, model, &weights->layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, METAL_DECODE_LAYER_FROM_ATTN_TO_FFN); } metal_graph_unbind_attn_pre_batch_row(g, &attn_alias); } if (ok && group_ffn_pre) { ok = metal_graph_encode_ffn_pre_session_batch( items, count, model, &weights->layer[il], 0, !group_attn_core); } if (ok) { ok = metal_graph_encode_routed_session_batch( items, count, model, &weights->layer[il], il, 0, !group_ffn_pre, !group_shared_ffn); } if (ok && group_shared_ffn) { ok = metal_graph_encode_shared_session_batch( items, count, model, &weights->layer[il], 0, !group_ffn_pre); } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; if (!group_shared_ffn) { const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); ok = metal_graph_encode_decode_layer_phase( g, model, &weights->layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, items[i].token, METAL_DECODE_LAYER_FROM_ROUTER); } 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); } } } } else { for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_gpu_graph *g = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); if (g->placement[stage_begin + 1u] != stage_tier) { ok = false; break; } for (uint32_t il = stage_begin; ok && il < stage_end; il++) { if (g->placement[il + 1u] != stage_tier) { ok = false; break; } 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, items[i].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); } } } } stage_begin = stage_end; } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ok = metal_graph_encode_output_head(&s->graph, model, weights, weights->output->dim[1]); } return ok; } static DS4_MAYBE_UNUSED bool metal_graph_mixed_workspace_compatible( const ds4_gpu_graph *owner, const ds4_gpu_graph *member) { if (!owner || !member || metal_graph_prefill_tokens(owner) != metal_graph_prefill_tokens(member)) { return false; } for (int t = 0; t < DS4_MAX_GPUS; t++) { ds4_gpu_tensor *owner_cur = owner->batch_cur_hc_by_tier[t]; ds4_gpu_tensor *owner_next = owner->batch_next_hc_by_tier[t]; ds4_gpu_tensor *member_cur = member->batch_cur_hc_by_tier[t]; ds4_gpu_tensor *member_next = member->batch_next_hc_by_tier[t]; const bool same_orientation = owner_cur == member_cur && owner_next == member_next; const bool opposite_orientation = owner_cur == member_next && owner_next == member_cur; if ((!same_orientation && !opposite_orientation) || owner->batch_after_attn_hc_by_tier[t] != member->batch_after_attn_hc_by_tier[t] || owner->batch_ffn_norm_by_tier[t] != member->batch_ffn_norm_by_tier[t] || owner->batch_router_selected_by_tier[t] != member->batch_router_selected_by_tier[t] || owner->batch_routed_down_by_tier[t] != member->batch_routed_down_by_tier[t] || owner->batch_shared_out_by_tier[t] != member->batch_shared_out_by_tier[t]) { return false; } } return true; } static void metal_graph_align_mixed_workspace( const ds4_gpu_graph *owner, ds4_gpu_graph *member) { for (int t = 0; t < DS4_MAX_GPUS; t++) { member->batch_cur_hc_by_tier[t] = owner->batch_cur_hc_by_tier[t]; member->batch_next_hc_by_tier[t] = owner->batch_next_hc_by_tier[t]; } } /* The Q4 model-backed oracle is exact through 512 prefill rows on the * layer-interleaved path. Larger waves currently change prefill logits and * must use the serialized public-API fallback. */ static DS4_MAYBE_UNUSED uint32_t metal_graph_mixed_routed_max_prefill_rows(void) { const char *value = getenv("DS4_CUDA_MIXED_ROUTED_MAX_PREFILL"); if (!value || !value[0]) return 512u; char *end = NULL; unsigned long parsed = strtoul(value, &end, 10); if (end == value || *end != '\0' || parsed > UINT32_MAX) return 512u; return (uint32_t)parsed; } static bool metal_graph_mixed_prefill_decode_supported( ds4_session *prefill_session, const token_vec *prompt, uint32_t start, uint32_t prefill_rows, ds4_decode_item *decode_items, int decode_count, const ds4_weights *weights) { #if defined(__APPLE__) (void)prefill_session; (void)prompt; (void)start; (void)prefill_rows; (void)decode_items; (void)decode_count; (void)weights; return false; #else if (!prefill_session || !prompt || !decode_items || decode_count < 3 || !weights || prefill_rows < metal_graph_resume_prefill_min_tokens() || prefill_rows > metal_graph_mixed_routed_max_prefill_rows() || start > (uint32_t)prompt->len || prefill_rows > (uint32_t)prompt->len - start) { return false; } ds4_gpu_graph *g = &prefill_session->graph; ds4_gpu_graph *batch_graph = &decode_items[0].session->graph; if (!g->placement || !batch_graph->placement || !g->cuda_tp_ep || !g->cuda_tp_prefill_ffn || (uint64_t)prefill_rows + (uint32_t)decode_count > g->prefill_cap || prefill_rows > g->raw_cap || (start % g->prefill_cap) + prefill_rows > g->prefill_cap || metal_graph_directional_steering_ffn_enabled(g) || getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL || getenv("DS4_METAL_LAYER_STAGE_PROFILE") != NULL) { return false; } if (!metal_graph_mixed_workspace_compatible(g, batch_graph)) return false; if (!metal_graph_session_batch_moe_supported( decode_items, decode_count, weights) || !metal_graph_session_batch_shared_supported( decode_items, decode_count, weights) || !metal_graph_session_batch_ffn_pre_supported( decode_items, decode_count, weights) || !metal_graph_session_batch_attn_pre_supported( decode_items, decode_count, weights)) { return false; } for (int i = 0; i < decode_count; i++) { ds4_session *s = decode_items[i].session; if (!s || s == prefill_session || s->engine != prefill_session->engine || !s->checkpoint_valid || s->graph.prefill_cap != g->prefill_cap || !metal_graph_mixed_workspace_compatible(g, &s->graph)) { return false; } } return true; #endif } /* Execute one resumed prefill quantum together with one exact decode batch. * Attention and non-routed FFN arithmetic remain on their original paths. */ static bool metal_graph_eval_mixed_prefill_decode( ds4_session *prefill_session, const token_vec *prompt, uint32_t start, uint32_t prefill_rows, ds4_decode_item *decode_items, int decode_count, const ds4_model *model, const ds4_weights *weights) { if (!metal_graph_mixed_prefill_decode_supported( prefill_session, prompt, start, prefill_rows, decode_items, decode_count, weights)) { return false; } const bool group_attn_core = metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_CORE", true) && metal_graph_session_batch_attn_core_supported( decode_items, decode_count, weights); const bool group_attn_post = group_attn_core && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_ATTN_POST", true) && metal_graph_session_batch_attn_post_supported( decode_items, decode_count, weights); const bool group_qkv = group_attn_core && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_QKV", true) && metal_graph_session_batch_qkv_supported( decode_items, decode_count, weights); const bool group_kv_store = group_qkv && metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_KV_STORE", true) && metal_graph_session_batch_kv_store_supported( decode_items, decode_count); ds4_gpu_graph *g = &prefill_session->graph; for (int i = 0; i < decode_count; i++) { metal_graph_align_mixed_workspace( g, &decode_items[i].session->graph); } g->batch_token_offset = 0; if (!metal_graph_set_active_tier_batch(g, g->emb_tier, prefill_rows)) { return false; } bool ok = metal_graph_upload_prompt_tokens( metal_graph_prefill_tokens(g), prompt, start, prefill_rows); if (ok) { int32_t *tokens = malloc((size_t)decode_count * sizeof(*tokens)); if (!tokens) return false; for (int i = 0; i < decode_count; i++) tokens[i] = decode_items[i].token; ok = ds4_gpu_tensor_write( metal_graph_prefill_tokens(g), (uint64_t)prefill_rows * sizeof(int32_t), tokens, (uint64_t)decode_count * sizeof(*tokens)) != 0; free(tokens); } if (ok) { ok = metal_graph_warmup_prefill_kernels( g, model, weights, prefill_rows); } if (ok) { ok = metal_graph_set_active_tier_batch(g, g->emb_tier, prefill_rows); } if (ok) { ok = metal_graph_upload_prompt_embeddings_hc( metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, weights, prompt, start, prefill_rows); } metal_graph_dspark_capture_begin_prefill(g); for (int i = 0; ok && i < decode_count; i++) { ds4_gpu_graph *dg = &decode_items[i].session->graph; metal_graph_dspark_capture_begin(dg); ok = metal_graph_set_active_tier_decode(dg, dg->emb_tier); if (ok) { ok = ds4_gpu_embed_token_hc_tensor( metal_graph_cur_hc(dg), model->map, model->size, weights->token_embd->abs_offset, (uint32_t)weights->token_embd->dim[1], (uint32_t)decode_items[i].token, DS4_N_EMBD, DS4_N_HC) != 0; } } if (ok) ok = ds4_gpu_begin_commands() != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; const int home_tier = g->placement[il + 1u]; ok = metal_graph_set_active_tier_batch(g, home_tier, prefill_rows) && metal_graph_layer_stage_profile_start(il); if (ok) { ok = metal_graph_encode_layer_attention_batch( g, model, layer, il, start, prefill_rows); } if (ok) { ok = metal_graph_encode_attn_pre_session_batch( decode_items, decode_count, model, layer, il, prefill_rows, false); } if (ok && group_qkv) { ok = metal_graph_encode_qkv_session_batch( decode_items, decode_count, model, layer, il, prefill_rows, group_kv_store); } for (int i = 0; ok && i < decode_count; i++) { ds4_session *s = decode_items[i].session; ds4_gpu_graph *dg = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % dg->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(dg, pos, 1); metal_graph_attn_pre_alias alias = {0}; ok = dg->placement[il + 1u] == home_tier && metal_graph_bind_attn_pre_batch_row( dg, g, layer, prefill_rows + (uint32_t)i, group_attn_core, &alias); if (ok) { ok = metal_graph_encode_decode_layer_phase( dg, model, layer, il, pos, dg->layer_raw_cache[il], dg->raw_cap, raw_row, n_raw, decode_items[i].token, !group_attn_core ? METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN : group_kv_store ? METAL_DECODE_LAYER_FROM_KV_STORE_TO_ATTN : group_qkv ? METAL_DECODE_LAYER_FROM_QKV_TO_ATTN : METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_ATTN); } metal_graph_unbind_attn_pre_batch_row(dg, &alias); } if (ok && group_attn_core) { ok = metal_graph_encode_attention_session_batch( decode_items, decode_count, model, layer, il, prefill_rows); } if (ok && group_attn_post) { ok = metal_graph_encode_attn_post_session_batch( decode_items, decode_count, model, layer, il, prefill_rows); } for (int i = 0; ok && group_attn_core && !group_attn_post && i < decode_count; i++) { ds4_session *s = decode_items[i].session; ds4_gpu_graph *dg = &s->graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const uint32_t raw_row = pos % dg->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(dg, pos, 1); metal_graph_attn_pre_alias alias = {0}; ok = metal_graph_bind_attn_pre_batch_row( dg, g, layer, prefill_rows + (uint32_t)i, true, &alias); if (ok) { ok = metal_graph_encode_decode_layer_phase( dg, model, layer, il, pos, dg->layer_raw_cache[il], dg->raw_cap, raw_row, n_raw, decode_items[i].token, METAL_DECODE_LAYER_FROM_ATTN_TO_FFN); } metal_graph_unbind_attn_pre_batch_row(dg, &alias); } if (ok) { ok = metal_graph_encode_ffn_pre_session_batch( decode_items, decode_count, model, layer, prefill_rows, !group_attn_core); } if (ok) { ok = metal_graph_encode_layer_ffn_batch( g, model, layer, il, start, prefill_rows, NULL, 0); } if (ok) { ok = metal_graph_encode_routed_session_batch( decode_items, decode_count, model, layer, il, prefill_rows, false, false); } if (ok) { ok = metal_graph_encode_shared_session_batch( decode_items, decode_count, model, layer, prefill_rows, false); } 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; ok = metal_graph_dspark_capture_prefill_layer( g, il, start, prefill_rows); } for (int i = 0; ok && i < decode_count; i++) { ds4_gpu_graph *dg = &decode_items[i].session->graph; ds4_gpu_tensor *tmp = metal_graph_cur_hc(dg); dg->cur_hc_by_tier[dg->active_tier] = metal_graph_after_ffn_hc(dg); dg->after_ffn_hc_by_tier[dg->active_tier] = tmp; ok = metal_graph_dspark_capture_decode_layer(dg, il); } } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const int prefill_src_tier = g->active_tier; ds4_gpu_tensor *saved_prefill_cur = NULL; ds4_gpu_tensor *last_prefill_hc = NULL; if (ok) { saved_prefill_cur = g->cur_hc_by_tier[prefill_src_tier]; last_prefill_hc = metal_graph_tensor_row_view( metal_graph_batch_cur_hc(g), prefill_rows - 1u, hc_dim); ok = last_prefill_hc != NULL; } if (ok) { g->cur_hc_by_tier[prefill_src_tier] = last_prefill_hc; ok = metal_graph_encode_output_head( g, model, weights, weights->output->dim[1]); g->cur_hc_by_tier[prefill_src_tier] = saved_prefill_cur; } for (int i = 0; ok && i < decode_count; i++) { ds4_session *s = decode_items[i].session; ok = metal_graph_encode_output_head( &s->graph, model, weights, weights->output->dim[1]); } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); g->cur_hc_by_tier[prefill_src_tier] = saved_prefill_cur; ds4_gpu_tensor_free(last_prefill_hc); g->batch_token_offset = 0; if (ok && g->active_tier != prefill_src_tier) { ok = metal_graph_set_active_tier_no_copy(g, prefill_src_tier); } if (ok) { ok = ds4_gpu_tensor_read( metal_graph_logits(g), 0, prefill_session->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } for (int i = 0; ok && i < decode_count; i++) { ds4_session *s = decode_items[i].session; ok = ds4_gpu_tensor_read( metal_graph_logits(&s->graph), 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (ok) ok = metal_graph_cuda_tp_attn_cache_sync_all(g); return ok; } #endif static int ds4_sessions_eval_batch_cuda(ds4_decode_item *items, int count, char *err, size_t errlen) { if (!items || count <= 0) { if (err && errlen) snprintf(err, errlen, "empty decode batch"); return 1; } if (count == 1) { return ds4_session_eval(items[0].session, items[0].token, err, errlen); } ds4_session *first = items[0].session; if (!first || !first->engine) { if (err && errlen) snprintf(err, errlen, "decode batch has no session"); return 1; } ds4_engine *e = first->engine; for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e) { if (err && errlen) { snprintf(err, errlen, "decode batch item %d belongs to a different engine", i); } return 1; } for (int j = 0; j < i; j++) { if (items[j].session == s) { if (err && errlen) { snprintf(err, errlen, "decode batch repeats session at items %d and %d", j, i); } return 1; } } if (ds4_session_pos(s) >= ds4_session_ctx(s)) { if (err && errlen) { snprintf(err, errlen, "decode batch item %d reached its context limit", i); } return 1; } } #ifndef DS4_NO_GPU /* The DeepSeek CUDA graph uses one default stream per device plus ordered * cross-device events. Encoding several independent graphs before the * final synchronization lets their packets fill different pipeline stages * while preserving the exact one-token kernels and per-session KV order. */ if (e->backend == DS4_BACKEND_CUDA && !ds4_session_is_glm(first) && e->support_kind == DS4_SUPPORT_NONE) { bool ok = ds4_gpu_begin_commands() != 0; for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; if (ds4_session_is_cpu(s) || ds4_session_is_glm(s) || !s->checkpoint_valid || s->engine->support_kind != DS4_SUPPORT_NONE) { ok = false; } } const char *interleave = getenv("DS4_CUDA_SESSION_BATCH_INTERLEAVE"); const bool use_pipeline = !interleave || !interleave[0] || strcmp(interleave, "0") != 0; if (ok && use_pipeline && first->graph.placement) { ok = metal_graph_encode_session_pipeline_batch( items, count, &e->model, &e->weights); } else { for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ok = metal_graph_encode_token_raw_swa( &s->graph, &e->model, &e->weights, items[i].token, (uint32_t)s->checkpoint.len, true, false); } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ok = ds4_gpu_tensor_read( metal_graph_logits(&s->graph), 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { for (int i = 0; i < count; i++) { items[i].session->checkpoint_valid = false; items[i].session->mtp_draft_valid = false; } if (err && errlen) { snprintf(err, errlen, "%s batched decode failed", ds4_backend_name(e->backend)); } return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; token_vec_push(&s->checkpoint, items[i].token); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); } return 0; } #endif /* Correctness fallback. If a later item fails, invalidate every member: * earlier sessions may already have advanced and must be rebuilt rather * than exposing a partially committed logical batch. */ for (int i = 0; i < count; i++) { if (ds4_session_eval(items[i].session, items[i].token, err, errlen) != 0) { for (int j = 0; j < count; j++) ds4_session_invalidate(items[j].session); return 1; } } return 0; } #ifdef DS4_TEST_HOOKS static uint64_t ds4_test_mixed_native_evals = 0; uint64_t ds4_test_mixed_native_count(void) { return ds4_test_mixed_native_evals; } #endif static int ds4_sessions_eval_batch_with_prefill_cuda( ds4_decode_item *items, int count, ds4_session *prefill_session, const ds4_tokens *prefill_prompt, char *err, size_t errlen) { if (!items || count <= 0 || !prefill_session || !prefill_prompt || !prefill_session->engine) { if (err && errlen) snprintf(err, errlen, "invalid mixed model batch"); return 1; } ds4_engine *e = prefill_session->engine; if (!prefill_session->checkpoint_valid || prefill_prompt->len <= prefill_session->checkpoint.len || prefill_prompt->len >= prefill_session->ctx_size || !ds4_tokens_starts_with(prefill_prompt, &prefill_session->checkpoint)) { if (err && errlen) { snprintf(err, errlen, "mixed prefill must extend a valid session checkpoint"); } return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s == prefill_session || s->engine != e || !s->checkpoint_valid || s->checkpoint.len >= s->ctx_size) { if (err && errlen) { snprintf(err, errlen, "invalid mixed decode item %d", i); } return 1; } for (int j = 0; j < i; j++) { if (items[j].session == s) { if (err && errlen) { snprintf(err, errlen, "mixed decode repeats session at items %d and %d", j, i); } return 1; } } } #ifndef DS4_NO_GPU const uint32_t start = (uint32_t)prefill_session->checkpoint.len; const uint32_t prefill_rows = (uint32_t)prefill_prompt->len - start; const bool native_requested = metal_graph_tp_env_flag("DS4_CUDA_MIXED_PREFILL_DECODE", true); const bool native_supported = native_requested && e->backend == DS4_BACKEND_CUDA && !ds4_session_is_glm(prefill_session) && e->support_kind == DS4_SUPPORT_NONE && metal_graph_mixed_prefill_decode_supported( prefill_session, prefill_prompt, start, prefill_rows, items, count, &e->weights); if (native_supported) { bool ok = metal_graph_eval_mixed_prefill_decode( prefill_session, prefill_prompt, start, prefill_rows, items, count, &e->model, &e->weights); if (!ok) { prefill_session->checkpoint_valid = false; prefill_session->mtp_draft_valid = false; for (int i = 0; i < count; i++) { items[i].session->checkpoint_valid = false; items[i].session->mtp_draft_valid = false; } if (err && errlen) snprintf(err, errlen, "CUDA mixed model batch failed"); return 1; } ds4_tokens_copy(&prefill_session->checkpoint, prefill_prompt); prefill_session->checkpoint_valid = true; prefill_session->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(prefill_session); session_greedy_splitkv_reset(prefill_session); if (prefill_session->progress) { prefill_session->progress( prefill_session->progress_ud, "prefill_chunk", prefill_prompt->len, prefill_prompt->len); } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; token_vec_push(&s->checkpoint, items[i].token); s->checkpoint_valid = true; s->mtp_draft_valid = false; ds4_session_dspark_capture_note_checkpoint(s); } #ifdef DS4_TEST_HOOKS ds4_test_mixed_native_evals++; #endif return 0; } #endif int rc = ds4_session_sync( prefill_session, prefill_prompt, err, errlen); if (rc != 0) return rc; rc = ds4_sessions_eval_batch(items, count, err, errlen); if (rc != 0) ds4_session_invalidate(prefill_session); return rc; } int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, int max_tokens, int eos_token, int *accepted, int accepted_cap, char *err, size_t errlen) { 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; accepted[0] = first_token; return 1; } if (ds4_session_is_cpu(s)) { (void)max_tokens; (void)eos_token; if (!accepted || accepted_cap <= 0) return 0; if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } if (ds4_session_is_glm(s)) { (void)max_tokens; (void)eos_token; if (!accepted || accepted_cap <= 0) return 0; #ifndef DS4_NO_GPU if (s->engine->glm_mtp && DS4_N_NEXTN_PREDICT != 0 && s->glm_graph_ready) { if (ds4_session_tp_leader(s)) { ds4_engine *ge = s->engine; if (!ds4_tp_send_eval(ge->tp.ctx, s->tp_session_id, ++ge->tp.eval_seq, first_token)) { snprintf(err, errlen, "tp: worker eval send failed"); return -1; } } int rc = ds4_session_glm_spec_cycle(s, first_token, accepted, accepted_cap, err, errlen); #if defined(__APPLE__) if (rc >= 0 && s->engine && s->engine->tp.active && ds4_gpu_tp_failed()) { snprintf(err, errlen, "tp: gate transport failed"); return -1; } #endif return rc; } #endif if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } #ifdef DS4_NO_GPU (void)s; (void)first_token; (void)max_tokens; (void)eos_token; (void)accepted; (void)accepted_cap; snprintf(err, errlen, "GPU support is not compiled in"); return -1; #else ds4_engine *e = s->engine; if (ds4_session_is_glm(s) && 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, accepted, cycle_cap, err, errlen); } /* * MTP in DeepSeek V4 is a speculative drafter, not a replacement sampler. * The target model still defines the exact output stream. A cycle starts * by accepting one normal target token, then asks the MTP block to propose * a short suffix. The suffix is useful only if the target model can verify * several proposed positions together; running ordinary decode once per * draft token is correctness-safe but cannot be faster than baseline. */ const bool strict_dspark = e->support_kind == DS4_SUPPORT_DSPARK && (e->quality || e->dspark_strict); bool can_prepare_support_draft = !strict_dspark && first_token != eos_token && max_tokens > 1 && accepted_cap > 1; if (can_prepare_support_draft && e->tp.active && e->support_kind == DS4_SUPPORT_MTP_LEGACY) { can_prepare_support_draft = false; } bool dspark_tail_skip = false; if (can_prepare_support_draft && e->support_kind == DS4_SUPPORT_DSPARK && ds4_dspark_scheduler_enabled()) { const uint32_t tail_min = ds4_dspark_scheduler_tail_min_tokens(); if (tail_min != 0 && (uint32_t)max_tokens < tail_min) { can_prepare_support_draft = false; dspark_tail_skip = true; if (ds4_dspark_stats_enabled()) s->dspark_stats.tail_skips++; if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler tail skip max=%d min=%u\n", max_tokens, tail_min); } } } if (ds4_session_eval_probe_tp(s, first_token, can_prepare_support_draft, err, errlen) != 0) return -1; int n_accept = 0; accepted[n_accept++] = first_token; if (first_token == eos_token || max_tokens == 1 || n_accept >= accepted_cap) return n_accept; if (strict_dspark) return n_accept; if (dspark_tail_skip) return n_accept; if (e->support_kind == DS4_SUPPORT_DSPARK) { return ds4_session_eval_dspark_speculative_argmax(s, n_accept, max_tokens, eos_token, accepted, accepted_cap, err, errlen); } if (metal_graph_cuda_splitkv_spec_requested() && (!e->mtp_ready || e->mtp_draft_tokens <= 1)) { int extra = ds4_session_eval_splitkv_spec_after_first( s, max_tokens - n_accept, eos_token, accepted + n_accept, accepted_cap - n_accept, err, errlen); if (extra < 0) return -1; return n_accept + extra; } /* Legacy MTP verify is not TP-mirrored; fall back to per-token decode. */ if (e->tp.active) return n_accept; if (!e->mtp_ready || !s->mtp_draft_valid || e->mtp_draft_tokens <= 1) return n_accept; int draft_cap = e->mtp_draft_tokens; if (draft_cap > max_tokens - n_accept) draft_cap = max_tokens - n_accept; if (draft_cap > accepted_cap - n_accept) draft_cap = accepted_cap - n_accept; int room = s->ctx_size - s->checkpoint.len; if (draft_cap > room - 1) draft_cap = room - 1; if (draft_cap <= 0) return n_accept; int drafts[16]; int draft_n = 1; drafts[0] = s->mtp_draft_token; s->mtp_draft_valid = false; const bool strict_mtp = e->quality || getenv("DS4_MTP_STRICT") != NULL; float mtp_margin_threshold = e->mtp_margin; const char *mtp_margin_env = getenv("DS4_MTP_MIN_MARGIN"); if (mtp_margin_env && mtp_margin_env[0]) { char *end = NULL; float v = strtof(mtp_margin_env, &end); if (end != mtp_margin_env && v >= 0.0f) mtp_margin_threshold = v; } const bool mtp_timing = getenv("DS4_MTP_TIMING") != NULL; const bool mtp_conf_log = getenv("DS4_MTP_CONF_LOG") != NULL; const bool mtp_need_logits = mtp_conf_log || getenv("DS4_MTP_FULL_LOGITS") != NULL || (!strict_mtp && mtp_margin_threshold > 0.0f); const double mtp_t0 = mtp_timing ? now_sec() : 0.0; double mtp_t_after_draft = mtp_t0; float mtp_last_margin = 0.0f; int mtp_last_top0 = -1, mtp_last_top1 = -1; /* * The first proposed token is verified for free: ds4_session_eval() just * produced the base logits for the committed prefix. If MTP disagrees at * this point there is no suffix to verify, so the exact behavior is to emit * only first_token and skip all speculative work. */ if (sample_argmax(s->logits, DS4_N_VOCAB) != drafts[0]) { if (getenv("DS4_MTP_SPEC_LOG")) { fprintf(stderr, "ds4: mtp spec miss first draft=%d\n", drafts[0]); } return n_accept; } if (drafts[0] == eos_token) draft_cap = 1; const uint32_t mtp_base_raw = s->graph.mtp_n_raw; /* * MTP has its own raw SWA cache. Recursive drafting writes speculative * future rows into it; after verification, rows beyond the accepted prefix * must become invisible. We do not copy/rollback the cache body because the * next draft attempt will overwrite future slots. A counter is enough. */ #define DS4_MTP_KEEP_ACCEPTED(n_) do { \ uint32_t keep_ = mtp_base_raw + (uint32_t)(n_); \ if (keep_ > s->graph.raw_window) keep_ = s->graph.raw_window; \ s->graph.mtp_n_raw = keep_; \ } while (0) for (; draft_n < draft_cap; draft_n++) { ds4_gpu_tensor *prev_hc = (draft_n & 1) ? s->graph.mtp_state_hc : s->graph.mtp_next_hc; ds4_gpu_tensor *out_hc = (draft_n & 1) ? s->graph.mtp_next_hc : s->graph.mtp_state_hc; int mtp_top = -1; if (!metal_graph_eval_mtp_draft_from_hc(&s->graph, &e->model, &e->weights, &e->mtp_model, &e->mtp_weights, prev_hc, out_hc, drafts[draft_n - 1], (uint32_t)(s->checkpoint.len + draft_n - 1), mtp_need_logits ? s->mtp_logits : NULL, &mtp_top)) { return n_accept; } drafts[draft_n] = mtp_top >= 0 ? mtp_top : sample_argmax(s->mtp_logits, DS4_N_VOCAB); if (drafts[draft_n] == eos_token) { draft_n++; break; } } if (mtp_conf_log && draft_n > 1) { float v0 = 0.0f, v1 = 0.0f; logits_top2(s->mtp_logits, DS4_N_VOCAB, &mtp_last_top0, &v0, &mtp_last_top1, &v1); mtp_last_margin = v0 - v1; } if (mtp_timing) mtp_t_after_draft = now_sec(); if (!strict_mtp && draft_n == 2 && mtp_margin_threshold > 0.0f) { if (!mtp_conf_log) { float v0 = 0.0f, v1 = 0.0f; logits_top2(s->mtp_logits, DS4_N_VOCAB, &mtp_last_top0, &v0, &mtp_last_top1, &v1); mtp_last_margin = v0 - v1; } if (mtp_last_margin < mtp_margin_threshold) { float *row_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row_logits[0])); const int start = s->checkpoint.len; const double verify_t0 = mtp_timing ? now_sec() : 0.0; bool ok = metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, drafts[0], (uint32_t)start, row_logits); if (!ok) { free(row_logits); snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return -1; } memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); free(row_logits); token_vec_push(&s->checkpoint, drafts[0]); accepted[n_accept++] = drafts[0]; s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(1); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { const double done = now_sec(); fprintf(stderr, "ds4: mtp timing margin-skip drafted=2 committed=1 margin=%.3f threshold=%.3f draft=%.3f ms verify=%.3f ms total=%.3f ms\n", mtp_last_margin, mtp_margin_threshold, (mtp_t_after_draft - mtp_t0) * 1000.0, (done - verify_t0) * 1000.0, (done - mtp_t0) * 1000.0); } return n_accept; } } /* * The useful N=2 verifier is the tiny batch path: it verifies two target * positions in one layer-major pass and commits prefix-1 directly on a * partial accept. Like the rest of the non-quality Metal path, it may pick * a different greedy token when batched reductions perturb nearly-tied * logits. --quality / DS4_MTP_STRICT selects the exact decode verifier, * which preserves the one-token target stream but is not a speed win. */ const bool use_decode2_exact = draft_n == 2 && strict_mtp && getenv("DS4_MTP_BATCH_VERIFY") == NULL; if (use_decode2_exact) { ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); float *row_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row_logits[0])); float *row0_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row0_logits[0])); const int start = s->checkpoint.len; int row0_top = -1; const double snapshot_t0 = mtp_timing ? now_sec() : 0.0; bool have_frontier = spec_frontier_snapshot(&frontier, s); const double snapshot_done = mtp_timing ? now_sec() : 0.0; bool ok = have_frontier; if (ok) { ok = metal_graph_verify_decode2_exact(&s->graph, &e->model, &e->weights, drafts[0], drafts[1], (uint32_t)start, &row0_top, NULL, row0_logits, row_logits); } const double verify_done = mtp_timing ? now_sec() : 0.0; if (ok && row0_top == drafts[1]) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); token_vec_push(&s->checkpoint, drafts[0]); token_vec_push(&s->checkpoint, drafts[1]); accepted[n_accept++] = drafts[0]; if (n_accept < accepted_cap) accepted[n_accept++] = drafts[1]; s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(2); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { fprintf(stderr, "ds4: mtp timing decode2 drafted=2 committed=2 draft=%.3f ms snapshot=%.3f ms verify=%.3f ms total=%.3f ms\n", (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (verify_done - snapshot_done) * 1000.0, (now_sec() - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row0_logits); free(row_logits); return n_accept; } if (ok) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); ok = spec_frontier_commit_prefix1(s); } if (ok) memcpy(s->logits, row0_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); if (ok) { token_vec_push(&s->checkpoint, drafts[0]); accepted[n_accept++] = drafts[0]; s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(1); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { const double replay_done = now_sec(); fprintf(stderr, "ds4: mtp timing decode2 drafted=2 committed=1 draft=%.3f ms snapshot=%.3f ms verify=%.3f ms prefix=%.3f ms total=%.3f ms\n", (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (verify_done - snapshot_done) * 1000.0, (replay_done - verify_done) * 1000.0, (replay_done - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row0_logits); free(row_logits); return n_accept; } if (have_frontier) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); (void)spec_frontier_restore(&frontier, s); } spec_frontier_free(&frontier); free(row0_logits); free(row_logits); if (getenv("DS4_MTP_SPEC_LOG")) { fprintf(stderr, "ds4: mtp decode2 verifier failed, falling back to sequential\n"); } } if (!use_decode2_exact) { ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); int *row_tops = xmalloc((size_t)draft_n * sizeof(row_tops[0])); float *row_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row_logits[0])); const int start = s->checkpoint.len; /* * The production MTP depth is two. Prefix-1 capture makes partial * accepts cheap, but it copies per-layer compressor frontiers even when * both draft tokens are accepted. Full accepts are the path that makes * MTP worthwhile, so by default we snapshot before the verifier and * replay one token on partial accept. DS4_MTP_CAPTURE_PREFIX1 restores * the older no-replay partial path for measurement. */ const bool capture_prefix1 = draft_n == 2 && (!strict_mtp || getenv("DS4_MTP_CAPTURE_PREFIX1") != NULL); const bool exact_replay_debug = getenv("DS4_MTP_EXACT_REPLAY") != NULL; const bool snapshot_required = draft_n > 2 || (draft_n == 2 && (!capture_prefix1 || exact_replay_debug)) || getenv("DS4_MTP_FORCE_SNAPSHOT") != NULL; bool have_frontier = false; bool ok = true; bool verifier_may_have_mutated = false; const double snapshot_t0 = mtp_timing ? now_sec() : 0.0; if (snapshot_required) { have_frontier = spec_frontier_snapshot(&frontier, s); ok = have_frontier; } const double snapshot_done = mtp_timing ? now_sec() : 0.0; if (ok) { for (int i = 0; i < draft_n; i++) token_vec_push(&s->checkpoint, drafts[i]); verifier_may_have_mutated = true; ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights, &s->checkpoint, (uint32_t)start, (uint32_t)draft_n, capture_prefix1, false, row_tops, NULL, NULL); } const double micro_verify_done = mtp_timing ? now_sec() : 0.0; if (ok) { int commit_drafts = 1; for (int i = 1; i < draft_n; i++) { if (row_tops[i - 1] != drafts[i]) break; commit_drafts++; } if (mtp_conf_log) { fprintf(stderr, "ds4: mtp conf drafted=%d committed=%d mtp_top=%d runner=%d margin=%.6f target_next=%d draft_next=%d\n", draft_n, commit_drafts, mtp_last_top0, mtp_last_top1, mtp_last_margin, draft_n > 1 ? row_tops[0] : -1, draft_n > 1 ? drafts[1] : -1); } if (exact_replay_debug && have_frontier) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); ok = spec_frontier_restore(&frontier, s); if (ok) { int replayed = 0; for (; replayed < commit_drafts && ok; replayed++) { ok = metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, drafts[replayed], (uint32_t)(start + replayed), row_logits); if (ok) token_vec_push(&s->checkpoint, drafts[replayed]); } if (ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); for (int i = 0; i < replayed && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; if (drafts[i] == eos_token) break; } s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(replayed); ds4_session_dspark_capture_note_checkpoint(s); spec_frontier_free(&frontier); free(row_logits); free(row_tops); return n_accept; } } } if (commit_drafts == draft_n) { ok = metal_graph_read_spec_logits_row(&s->graph, (uint32_t)(draft_n - 1), row_logits); if (ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); for (int i = 0; i < draft_n && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; if (drafts[i] == eos_token) break; } s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(draft_n); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { fprintf(stderr, "ds4: mtp timing micro drafted=%d committed=%d draft=%.3f ms snapshot=%.3f ms verify=%.3f ms total=%.3f ms\n", draft_n, draft_n, (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (micro_verify_done - snapshot_done) * 1000.0, (now_sec() - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row_logits); free(row_tops); return n_accept; } } if (draft_n == 2 && commit_drafts == 1 && capture_prefix1) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); const double prefix_t0 = mtp_timing ? now_sec() : 0.0; ok = spec_frontier_commit_prefix1(s); const double prefix_done = mtp_timing ? now_sec() : 0.0; if (ok) ok = metal_graph_read_spec_logits_row(&s->graph, 0, row_logits); if (ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); accepted[n_accept++] = drafts[0]; s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(1); token_vec_push(&s->checkpoint, drafts[0]); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { fprintf(stderr, "ds4: mtp timing micro drafted=%d committed=%d draft=%.3f ms snapshot=%.3f ms verify=%.3f ms prefix=%.3f ms total=%.3f ms noreplay=1\n", draft_n, commit_drafts, (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (micro_verify_done - snapshot_done) * 1000.0, (prefix_done - prefix_t0) * 1000.0, (now_sec() - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row_logits); free(row_tops); return n_accept; } } else { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); ok = have_frontier && spec_frontier_restore(&frontier, s); } if (ok && draft_n == 2 && commit_drafts == 1) { ok = metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights, drafts[0], (uint32_t)start, row_logits); if (ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); accepted[n_accept++] = drafts[0]; s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(1); token_vec_push(&s->checkpoint, drafts[0]); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { const double replay_done = now_sec(); fprintf(stderr, "ds4: mtp timing micro drafted=%d committed=%d draft=%.3f ms snapshot=%.3f ms verify=%.3f ms exact_replay=%.3f ms total=%.3f ms\n", draft_n, commit_drafts, (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (micro_verify_done - snapshot_done) * 1000.0, (replay_done - micro_verify_done) * 1000.0, (replay_done - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row_logits); free(row_tops); return n_accept; } } if (ok) { for (int i = 0; i < commit_drafts; i++) token_vec_push(&s->checkpoint, drafts[i]); ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights, &s->checkpoint, (uint32_t)start, (uint32_t)commit_drafts, false, false, row_tops, NULL, NULL); if (ok) ok = metal_graph_read_spec_logits_row(&s->graph, (uint32_t)(commit_drafts - 1), row_logits); if (ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); for (int i = 0; i < commit_drafts && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; if (drafts[i] == eos_token) break; } s->checkpoint_valid = true; s->mtp_draft_valid = false; DS4_MTP_KEEP_ACCEPTED(commit_drafts); ds4_session_dspark_capture_note_checkpoint(s); if (mtp_timing) { const double replay_done = now_sec(); fprintf(stderr, "ds4: mtp timing micro drafted=%d committed=%d draft=%.3f ms snapshot=%.3f ms verify=%.3f ms replay=%.3f ms total=%.3f ms\n", draft_n, commit_drafts, (mtp_t_after_draft - mtp_t0) * 1000.0, (snapshot_done - snapshot_t0) * 1000.0, (micro_verify_done - snapshot_done) * 1000.0, (replay_done - micro_verify_done) * 1000.0, (replay_done - mtp_t0) * 1000.0); } spec_frontier_free(&frontier); free(row_logits); free(row_tops); return n_accept; } } } s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); if (have_frontier) { (void)spec_frontier_restore(&frontier, s); } else if (!verifier_may_have_mutated) { /* Snapshot setup failed before the verifier touched Metal state. * Fall through to the exact sequential verifier below. */ } else { snprintf(err, errlen, "MTP verifier failed"); s->checkpoint_valid = false; DS4_MTP_KEEP_ACCEPTED(0); spec_frontier_free(&frontier); free(row_logits); free(row_tops); return -1; } spec_frontier_free(&frontier); free(row_logits); free(row_tops); if (getenv("DS4_MTP_SPEC_LOG")) { fprintf(stderr, "ds4: mtp spec micro verifier failed, falling back to sequential\n"); } } /* * Safety fallback: if the production microbatch verifier fails, verify * drafts with the exact normal one-token decode path instead of returning * wrong state. This path is deliberately slow and should not be selected * during normal --mtp operation. */ int verified = 0; int target_top = sample_argmax(s->logits, DS4_N_VOCAB); bool logits_on_host = true; const double seq_t0 = mtp_timing ? now_sec() : 0.0; for (int i = 0; i < draft_n && n_accept < accepted_cap; i++) { if (target_top != drafts[i]) { if (getenv("DS4_MTP_SPEC_LOG")) { fprintf(stderr, "ds4: mtp spec seq miss at=%d draft=%d base=%d drafted=%d accepted=%d\n", i, drafts[i], target_top, draft_n, n_accept); } break; } if (!metal_graph_eval_token_raw_swa_top(&s->graph, &e->model, &e->weights, drafts[i], (uint32_t)s->checkpoint.len, &target_top, NULL, false, NULL, false)) { snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return -1; } token_vec_push(&s->checkpoint, drafts[i]); logits_on_host = false; accepted[n_accept++] = drafts[i]; verified++; if (drafts[i] == eos_token) break; } if (verified > 0 && !logits_on_host) { if (ds4_gpu_tensor_read(metal_graph_logits(&s->graph), 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(s->logits[0])) == 0) { snprintf(err, errlen, "%s logits readback failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; return -1; } logits_on_host = true; } (void)logits_on_host; DS4_MTP_KEEP_ACCEPTED(verified); if (verified > 0) ds4_session_dspark_capture_note_checkpoint(s); #undef DS4_MTP_KEEP_ACCEPTED if (mtp_timing) { fprintf(stderr, "ds4: mtp timing seq drafted=%d verified=%d draft=%.3f ms verify=%.3f ms total=%.3f ms\n", draft_n, verified, (mtp_t_after_draft - mtp_t0) * 1000.0, (now_sec() - seq_t0) * 1000.0, (now_sec() - mtp_t0) * 1000.0); } if (getenv("DS4_MTP_SPEC_LOG")) { if (verified == draft_n) { fprintf(stderr, "ds4: mtp spec seq accept drafted=%d accepted=%d\n", draft_n, n_accept); } else { fprintf(stderr, "ds4: mtp spec seq partial drafted=%d verified=%d accepted=%d\n", draft_n, verified, n_accept); } } return n_accept; #endif } void ds4_session_invalidate(ds4_session *s) { if (!s) return; if (ds4_session_tp_leader(s) && !ds4_tp_failed(s->engine->tp.ctx)) { (void)ds4_tp_send_invalidate(s->engine->tp.ctx, s->tp_session_id); } 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 } void ds4_session_rewind(ds4_session *s, int pos) { 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); } if (pos < 0) pos = 0; 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 } int ds4_session_pos(ds4_session *s) { return s->checkpoint.len; } int ds4_session_ctx(ds4_session *s) { return s->ctx_size; } int ds4_session_prefill_cap(ds4_session *s) { return s ? (int)s->prefill_cap : 0; }