/* ========================================================================= * 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 #include #endif #include #include #include #include "ds4.h" #include "ds4_tool_text.h" #include "ds4_distributed.h" #include "ds4_image.h" #include "ds4_engram.h" #include "ds4_tp.h" #if !defined(DS4_NO_GPU) && !defined(DS4_ROCM_BUILD) #define DS4_HAS_DEEPSEEK41_GPU 1 #endif #if !defined(DS4_NO_GPU) && !defined(DS4_ROCM_BUILD) #define DS4_HAS_QWEN4_GPU 1 #ifdef __APPLE__ #define DS4_HAS_QWEN4_METAL 1 #endif #endif #ifdef DS4_ROCM_BUILD #include "ds4_linux_memory.h" #endif #ifdef DS4_TEST_HOOKS static uint64_t ds4_test_ds41_native_evals; static uint64_t ds4_test_mixed_native_evals; uint64_t ds4_test_ds41_batch_count(void) { return ds4_test_ds41_native_evals; } uint64_t ds4_test_mixed_native_count(void) { return ds4_test_mixed_native_evals; } #endif /* TP context for the verify-block RDMA window (set with the gate callbacks). */ #if !defined(DS4_NO_GPU) && defined(__APPLE__) static ds4_tp *g_tp_block_ctx; #endif /* 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) /* Decode-island graph capture is CUDA-only; Metal decodes eagerly. */ int ds4_gpu_decode_graphs_supported(void) { return 0; } int ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key) { (void)key; return -1; } int ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key) { (void)key; return -1; } void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key) { (void)key; } void ds4_gpu_decode_graphs_invalidate(void) {} 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 = 248320, DS4_MAX_HEAD = 128, DS4_MAX_HEAD_KV = 2, 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 = 512, DS4_MAX_EXPERT_USED = 10, 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, DS4_MAX_KDA_HEAD = 64, DS4_MAX_KDA_HEAD_DIM = 128, DS4_MAX_KDA_CONV = 4, DS4_MAX_PLE_NGRAM = 8, DS4_MAX_PLE_HEADS = 64, }; typedef enum { DS4_MODEL_FAMILY_DEEPSEEK4 = 0, DS4_MODEL_FAMILY_GLM_DSA = 1, DS4_MODEL_FAMILY_DEEPSEEK41 = 2, DS4_MODEL_FAMILY_QWEN4_EXP = 3, } ds4_model_family; typedef enum { DS4_VARIANT_FLASH = 0, DS4_VARIANT_PRO = 1, DS4_VARIANT_GLM52 = 2, DS4_VARIANT_GLM53 = 3, DS4_VARIANT_FLASH41 = 4, DS4_VARIANT_QWEN4_EXP = 5, DS4_VARIANT_QWEN4_MINI = 6, } 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; uint32_t n_kda_head; uint32_t n_kda_head_dim; uint32_t n_kda_conv; uint32_t n_hc_lowrank; uint32_t n_lin_k_head; uint32_t n_lin_v_head; uint32_t n_lin_head_dim; uint32_t n_lin_conv; uint32_t n_full_attn_interval; uint32_t n_ple_layer; uint32_t n_ple_ngram; uint32_t n_ple_heads_per_ngram; uint32_t n_ple_conv; uint32_t n_ple_head_dim; int32_t ple_eos_id; 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; float kda_gate_lower_bound; 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_FLASH41 = { .name = "DeepSeek V4.1 Flash", .family = DS4_MODEL_FAMILY_DEEPSEEK41, .variant = DS4_VARIANT_FLASH41, .n_layer = 40, .n_embd = 5120, .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 = 1280, .n_lora_o = 1024, .n_expert = 384, .n_expert_used = 6, .n_expert_shared = 1, .n_ff_exp = 2304, .n_swa = 128, .n_indexer_head = 32, .n_indexer_head_dim = 128, .n_indexer_top_k = 512, .n_hc = 4, .n_hc_sinkhorn_iter = 20, .rms_eps = 1.0e-20f, .hc_eps = 1.0e-6f, .expert_weight_scale = 1.5f, .swiglu_clamp_exp = 10.0f, .rope_freq_base = 10000.0f, .rope_scale_factor = 16.0f, .rope_yarn_beta_fast = 32.0f, .rope_yarn_beta_slow = 1.0f, .compress_rope_freq_base = 160000.0f, .rope_orig_ctx = 65536, }; 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 const ds4_shape DS4_SHAPE_GLM53 = { .name = "GLM 5.3 Flash", .family = DS4_MODEL_FAMILY_GLM_DSA, .variant = DS4_VARIANT_GLM53, .n_layer = 46, .n_embd = 4096, .n_vocab = 154880, .n_head = 64, .n_head_kv = 1, /* Raw DSA cache row: 512 latent dimensions and no RoPE tail. */ .n_head_dim = 512, .n_value_dim = 256, .n_rot = 0, .n_lora_q = 1536, .n_expert = 288, .n_expert_used = 8, .n_expert_shared = 1, .n_ff_exp = 2048, .n_ff_dense = 12288, .n_indexer_head = 32, .n_indexer_head_dim = 128, .n_indexer_top_k = 2048, .n_hc = 4, .n_hc_sinkhorn_iter = 20, .n_nextn_predict = 1, .n_leading_dense = 3, .n_kv_lora = 512, .n_key_mla = 256, .n_value_mla = 256, .n_kda_head = 64, .n_kda_head_dim = 128, .n_kda_conv = 4, .rms_eps = 1.0e-5f, .hc_eps = 1.0e-6f, .expert_weight_scale = 2.5f, .swiglu_clamp_exp = 10.0f, .rope_orig_ctx = 1048576, .kda_gate_lower_bound = -5.0f, }; /* Qwen3.8-Flash-Next: 48 trunk layers (3 GDN + 1 gated GQA attention with * a QSA indexer, repeating) plus one MTP block, four hyper-connection * streams, and a hashed n-gram PLE layer at trunk index 1. n_layer counts * the MTP block like GLM; the validator trims it for files without one. */ static const ds4_shape DS4_SHAPE_QWEN4_EXP = { .name = "Qwen3.8 Flash Next", .family = DS4_MODEL_FAMILY_QWEN4_EXP, .variant = DS4_VARIANT_QWEN4_EXP, .n_layer = 49, .n_embd = 2560, .n_vocab = 248320, .n_head = 24, .n_head_kv = 2, .n_head_dim = 256, .n_value_dim = 256, .n_rot = 64, .n_expert = 512, .n_expert_used = 10, .n_expert_shared = 1, .n_ff_exp = 640, .n_indexer_head = 4, .n_indexer_head_dim = 128, .n_indexer_top_k = 2048, .n_hc = 4, .n_hc_lowrank = 320, .n_nextn_predict = 1, .n_lin_k_head = 16, .n_lin_v_head = 48, .n_lin_head_dim = 128, .n_lin_conv = 4, .n_full_attn_interval = 4, .n_ple_layer = 1, .n_ple_ngram = 3, .n_ple_heads_per_ngram = 8, .n_ple_conv = 4, .n_ple_head_dim = 160, .ple_eos_id = 248044, .rms_eps = 1.0e-6f, .rope_freq_base = 10000000.0f, .rope_orig_ctx = 262144, }; /* Synthetic test model: same structure at toy widths, real tokenizer. */ static const ds4_shape DS4_SHAPE_QWEN4_MINI = { .name = "Qwen3.8 Flash Next mini", .family = DS4_MODEL_FAMILY_QWEN4_EXP, .variant = DS4_VARIANT_QWEN4_MINI, .n_layer = 9, .n_embd = 64, .n_vocab = 248320, .n_head = 4, .n_head_kv = 2, .n_head_dim = 32, .n_value_dim = 32, .n_rot = 8, .n_expert = 32, .n_expert_used = 10, .n_expert_shared = 1, .n_ff_exp = 32, .n_indexer_head = 4, .n_indexer_head_dim = 32, .n_indexer_top_k = 8, .n_hc = 4, .n_hc_lowrank = 8, .n_nextn_predict = 1, .n_lin_k_head = 2, .n_lin_v_head = 6, .n_lin_head_dim = 32, .n_lin_conv = 4, .n_full_attn_interval = 4, .n_ple_layer = 1, .n_ple_ngram = 3, .n_ple_heads_per_ngram = 8, .n_ple_conv = 4, .n_ple_head_dim = 4, .ple_eos_id = 248044, .rms_eps = 1.0e-6f, .rope_freq_base = 10000000.0f, .rope_orig_ctx = 262144, }; 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 bool g_ds4_flash_vision_exp = false; 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_N_KDA_HEAD (g_ds4_shape.n_kda_head) #define DS4_N_KDA_HEAD_DIM (g_ds4_shape.n_kda_head_dim) #define DS4_N_KDA_CONV (g_ds4_shape.n_kda_conv) #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_KDA_GATE_LOWER_BOUND (g_ds4_shape.kda_gate_lower_bound) #define DS4_ROPE_ORIG_CTX (g_ds4_shape.rope_orig_ctx) #define DS4_N_HC_LOWRANK (g_ds4_shape.n_hc_lowrank) #define DS4_N_LIN_K_HEAD (g_ds4_shape.n_lin_k_head) #define DS4_N_LIN_V_HEAD (g_ds4_shape.n_lin_v_head) #define DS4_N_LIN_HEAD_DIM (g_ds4_shape.n_lin_head_dim) #define DS4_N_LIN_CONV (g_ds4_shape.n_lin_conv) #define DS4_N_FULL_ATTN_INTERVAL (g_ds4_shape.n_full_attn_interval) #define DS4_N_PLE_LAYER (g_ds4_shape.n_ple_layer) #define DS4_N_PLE_NGRAM (g_ds4_shape.n_ple_ngram) #define DS4_N_PLE_HEADS_PER_NGRAM (g_ds4_shape.n_ple_heads_per_ngram) #define DS4_N_PLE_CONV (g_ds4_shape.n_ple_conv) #define DS4_N_PLE_HEAD_DIM (g_ds4_shape.n_ple_head_dim) #define DS4_PLE_EOS_ID (g_ds4_shape.ple_eos_id) #define DS4_N_PLE_HEADS ((DS4_N_PLE_NGRAM - 1u) * DS4_N_PLE_HEADS_PER_NGRAM) #define DS4_N_LIN_CONV_DIM (2u * DS4_N_LIN_K_HEAD * DS4_N_LIN_HEAD_DIM + DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM) /* PLE n-gram hash constants come from the GGUF (uint64 multipliers, per-head * prime table sizes and row offsets); see qwen4_ple_step. */ typedef struct { uint32_t n_heads; uint64_t multiplier[DS4_MAX_PLE_NGRAM]; uint32_t head_offset[DS4_MAX_PLE_HEADS]; uint32_t head_vocab[DS4_MAX_PLE_HEADS]; uint64_t n_rows; } ds4_qwen4_ple_hash; static ds4_qwen4_ple_hash g_ds4_qwen4_ple; static bool ds4_model_is_glm53(void) { return DS4_MODEL_VARIANT == DS4_VARIANT_GLM53; } static uint32_t directional_steering_layer_count(void) { if (DS4_N_LAYER <= DS4_N_NEXTN_PREDICT) return 0; return DS4_N_LAYER - DS4_N_NEXTN_PREDICT; } static bool ds4_glm53_layer_is_kda(uint32_t il) { return ds4_model_is_glm53() && il + DS4_N_NEXTN_PREDICT < DS4_N_LAYER && il % 4u != 3u; } static bool ds4_model_is_qwen4(void) { return DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_QWEN4_EXP; } /* Trunk layers repeat 3 GDN + 1 full attention; the MTP block is attention. */ static bool ds4_qwen4_layer_is_linear(uint32_t il) { return ds4_model_is_qwen4() && il + DS4_N_NEXTN_PREDICT < DS4_N_LAYER && (il + 1u) % DS4_N_FULL_ATTN_INTERVAL != 0u; } static bool ds4_qwen4_layer_is_ple(uint32_t il) { return ds4_model_is_qwen4() && il == DS4_N_PLE_LAYER; } static bool ds4_qwen4_layer_is_nextn(uint32_t il) { return ds4_model_is_qwen4() && DS4_N_NEXTN_PREDICT != 0 && il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER; } 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 * - MXFP4 routed experts preserved from native checkpoints * - Q8_K temporary activation blocks for dot products */ #define QK_K 256 #define QK_MXFP4 32 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; typedef struct { uint8_t e; uint8_t qs[QK_MXFP4 / 2]; } block_mxfp4; #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); DS4_STATIC_ASSERT(ds4_block_mxfp4_size, sizeof(block_mxfp4) == 17); 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 && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) 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; case DS4_VARIANT_FLASH41: return il < 2 ? 0u : il < 20 ? 2u : 1u; default: ds4_die("unsupported DeepSeek4 model variant"); } return 0; } static bool ds41_kv_source(uint32_t il) { return il == 2 || il == 8 || il == 14 || il == 20; } static bool ds41_index_source(uint32_t il) { return ds41_kv_source(il) || il == 24 || il == 28 || il == 32 || il == 36; } static bool ds41_engram_layer(uint32_t il) { return il == 1 || il == 14; } 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] = {"mxfp4", 32, 17}, }; enum { DS4_TENSOR_F32 = 0, DS4_TENSOR_F16 = 1, DS4_TENSOR_Q4_0 = 2, DS4_TENSOR_Q4_1 = 3, 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_I8 = 24, DS4_TENSOR_I32 = 26, DS4_TENSOR_BF16 = 30, DS4_TENSOR_MXFP4 = 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 ds4_model { int fd; const uint8_t *map; uint64_t size; /* Addressable weights, excluding disk-only n-grams. */ uint64_t file_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; int ngram_fd; const ds4_tensor *ngram_tensor; } 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; if (m->ngram_tensor && m->ngram_fd >= 0) close(m->ngram_fd); 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; m->ngram_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) { /* n_kv comes from the header. Every entry consumes at least one byte in the * file, so a count larger than the bytes remaining cannot be real; reject it * before calloc so a tiny file can't request an enormous allocation. */ if (m->n_kv > c->size - c->pos) ds4_die("GGUF metadata count exceeds file size"); 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) { /* As in parse_metadata: each tensor directory entry needs at least one byte * in the file, so reject a count larger than the bytes remaining before the * allocation. */ if (m->n_tensors > c->size - c->pos) ds4_die("GGUF tensor count exceeds file size"); 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; } } } /* Engram is deliberately outside the weight mapping, not merely absent from * a residency list. Startup warming and any future weight-view code must not * turn its 189 GiB of random-access rows into a resident model allocation. */ static void model_unmap_engram(ds4_model *m) { ds4_str arch = {0}, encoding = {0}; if (!model_get_string(m, "general.architecture", &arch) || !ds4_streq(arch, "deepseek41")) return; if (!model_get_string(m, "deepseek41.engram.encoding", &encoding) || !ds4_streq(encoding, "e4m3_e8m0_32_row264")) ds4_die("unsupported V4.1 Engram row encoding"); const ds4_tensor *tables[2] = {NULL, NULL}; uint64_t resident_end = m->tensor_data_pos; m->max_tensor_bytes = 0; for (uint64_t i = 0; i < m->n_tensors; i++) { const ds4_tensor *t = &m->tensors[i]; int index = ds4_streq(t->name, "blk.1.engram_embd.weight") ? 0 : ds4_streq(t->name, "blk.14.engram_embd.weight") ? 1 : -1; if (index >= 0) { if (tables[index] || t->ndim != 2 || t->type != 24 || t->dim[0] != 264 || !t->dim[1] || t->dim[1] > UINT32_MAX) ds4_die("invalid or duplicate V4.1 Engram tensor"); tables[index] = t; } else { if (!t->bytes) ds4_die("unsupported V4.1 resident tensor type"); if (t->abs_offset + t->bytes > resident_end) resident_end = t->abs_offset + t->bytes; if (t->bytes > m->max_tensor_bytes) m->max_tensor_bytes = t->bytes; } } if (!tables[0] || !tables[1]) ds4_die("V4.1 GGUF is missing Engram tables"); const uint64_t start = tables[0]->abs_offset; const long page = sysconf(_SC_PAGESIZE); if (page <= 0 || start % (uint64_t)page || resident_end > start || tables[1]->abs_offset < start + tables[0]->bytes) ds4_die("V4.1 Engram tables must follow the page-aligned main weights"); if (munmap((void *)(m->map + start), (size_t)(m->size - start)) != 0) ds4_die_errno("cannot unmap disk-only region", "Engram"); m->size = start; } /* Like V4.1 Engram, the n-gram table must trail the resident weights. This * prevents warming or a future whole-model GPU view from faulting it in. */ static void model_unmap_qwen_ngrams(ds4_model *m, const char *path) { ds4_str arch = {0}; if (!model_get_string(m, "general.architecture", &arch) || !ds4_streq(arch, "qwen4exp")) return; const ds4_tensor *table = NULL; uint64_t resident_end = m->tensor_data_pos; m->max_tensor_bytes = 0; for (uint64_t i = 0; i < m->n_tensors; i++) { const ds4_tensor *t = &m->tensors[i]; if (ds4_streq(t->name, "per_layer_token_embd.weight")) { if (table || t->type != DS4_TENSOR_BF16 || t->ndim != 2 || !t->dim[0] || t->dim[0] > 160 || !t->dim[1] || t->dim[1] > UINT32_MAX) ds4_die("Qwen requires original BF16 n-grams; repack with gguf-tools/qwen4_native_ngrams.py"); table = t; } else { if (!t->bytes) ds4_die("unsupported Qwen resident tensor type"); if (t->abs_offset + t->bytes > resident_end) resident_end = t->abs_offset + t->bytes; if (t->bytes > m->max_tensor_bytes) m->max_tensor_bytes = t->bytes; } } /* Tokenizer inspection can still open old main-only files. Inference * rejects them when binding its required n-gram tensor. */ if (!table) return; const long page = sysconf(_SC_PAGESIZE); if (page <= 0 || table->abs_offset % (uint64_t)page || resident_end > table->abs_offset) ds4_die("Qwen n-grams must follow page-aligned weights; repack with gguf-tools/qwen4_native_ngrams.py"); int fd = open(path, O_RDONLY | O_CLOEXEC); struct stat main_st, table_st; if (fd < 0) ds4_die_errno("cannot open n-gram table", path); if (fstat(fd, &table_st) || fstat(m->fd, &main_st) || table_st.st_dev != main_st.st_dev || table_st.st_ino != main_st.st_ino || table_st.st_size < 0 || (uint64_t)table_st.st_size != m->file_size) { close(fd); ds4_die("Qwen GGUF changed while opening its n-grams"); } #ifdef __APPLE__ if (fcntl(fd, F_NOCACHE, 1) || fcntl(fd, F_RDAHEAD, 0)) ds4_die_errno("cannot configure n-gram disk reads", path); #elif defined(POSIX_FADV_RANDOM) (void)posix_fadvise(fd, (off_t)table->abs_offset, (off_t)table->bytes, POSIX_FADV_RANDOM); #ifdef POSIX_FADV_NOREUSE (void)posix_fadvise(fd, (off_t)table->abs_offset, (off_t)table->bytes, POSIX_FADV_NOREUSE); #endif #endif if (munmap((void *)(m->map + table->abs_offset), (size_t)(m->size - table->abs_offset))) ds4_die_errno("cannot unmap disk-only n-grams", path); m->size = table->abs_offset; m->ngram_fd = fd; m->ngram_tensor = table; } /* 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; m->ngram_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; m->file_size = m->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); model_unmap_engram(m); model_unmap_qwen_ngrams(m, path); 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 #if defined(__APPLE__) || (!defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU)) /* Seed plus five drafts needs five intermediate compressor frontiers. */ #define DS4_SPEC_PREFIX_SLOTS 5 #else #define DS4_SPEC_PREFIX_SLOTS 4 #endif static bool ds4_dspark_rocm_gfx1151_fast_path(void) { #if defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) const char *env = getenv("DS4_ROCM_DSPARK_FAST"); if (env && env[0]) return env[0] != '0'; return ds4_gpu_dspark_gfx1151_fast_path() != 0; #else return false; #endif } static bool ds4_dspark_rocm_gfx1151_reference_alignment(void) { #if defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) return ds4_gpu_dspark_gfx1151_fast_path() != 0; #else return false; #endif } 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); } const bool v41 = ds4_streq(arch, "deepseek41"); if (v41) { model_get_u32(m, "deepseek41.num_hidden_layers", &layers); model_get_u64_compat(m, "deepseek41.max_position_embeddings", &ctx_train); model_get_u32(m, "deepseek41.num_attention_heads", &n_head); model_get_u32(m, "deepseek41.num_key_value_heads", &n_head_kv); model_get_u32(m, "deepseek41.head_dim", &head_dim); model_get_u32(m, "deepseek41.sliding_window", &n_swa); model_get_u32(m, "deepseek41.index_n_heads", &indexer_heads); model_get_u32(m, "deepseek41.index_head_dim", &indexer_head_dim); model_get_u32(m, "deepseek41.index_topk", &indexer_top_k); model_get_u32(m, "deepseek41.n_routed_experts", &n_expert); model_get_u32(m, "deepseek41.num_experts_per_tok", &n_expert_used); } for (uint64_t i = 0; i < m->n_tensors; i++) { tensor_bytes += m->tensors[i].bytes; params += m->tensors[i].elements; /* Each validated disk-only Engram row contains 256 weights and eight * scale bytes. Scales count toward storage, not model parameters. */ if (v41 && m->tensors[i].abs_offset >= m->size) params -= m->tensors[i].dim[1] * (DS4_ENGRAM_ROW_BYTES - DS4_ENGRAM_DIM); } 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->file_size ? m->file_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; } static bool support_model_checkpoint_compatible(const ds4_model *m) { static const char vision_exp_revision[] = "e46e16bf6035c6f317eb2ac7458eb0362926d402"; ds4_str variant = {0}; const bool has_variant = model_get_string(m, "deepseek4.checkpoint_variant", &variant); if (g_ds4_flash_vision_exp) { ds4_str revision = {0}; if (!has_variant || !ds4_streq(variant, "vision-exp") || !model_get_string(m, "general.source.revision", &revision) || !ds4_streq(revision, vision_exp_revision)) { fprintf(stderr, "ds4: DSpark support model does not match the pinned " "DeepSeek V4 Flash Vision-Exp checkpoint\n"); return false; } } else if (has_variant && ds4_streq(variant, "vision-exp")) { fprintf(stderr, "ds4: Vision-Exp DSpark support model cannot be used with " "this DeepSeek checkpoint\n"); return false; } return true; } #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; } #ifndef DS4_ROCM_BUILD 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; } #endif 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 == m->ngram_tensor) continue; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && (ds4_streq(t->name, "blk.1.engram_embd.weight") || ds4_streq(t->name, "blk.14.engram_embd.weight"))) continue; if (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset) { free(spans); return false; } #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (ds4_gpu_model_range_replaced(m->map, t->abs_offset, t->bytes)) { continue; } #endif /* Expert sharding owns tensor slices, not necessarily whole tensors. */ for (uint32_t j = 0; j < (span_count ? span_count : 1u); j++) { uint64_t off = t->abs_offset, end = off + t->bytes; if (span_count) { if (off < span_offsets[j]) off = span_offsets[j]; const uint64_t limit = span_offsets[j] + span_sizes[j]; if (end > limit) end = limit; } if (end <= off) continue; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (ds4_gpu_model_range_replaced(m->map, off, end - off)) continue; #endif if (nspan == cap) { if (cap > SIZE_MAX / (2u * sizeof(*spans))) { free(spans); return false; } cap *= 2u; spans = xrealloc(spans, (size_t)cap * sizeof(*spans)); } spans[nspan++] = (accelerator_tensor_span){.off = off, .end = end}; } } 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; } #ifndef DS4_ROCM_BUILD 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 || t->type != DS4_TENSOR_Q8_0 || t->ndim != 2) 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 (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; } #endif 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; } #ifndef DS4_ROCM_BUILD if (!m->ngram_tensor && !accelerator_cache_q8_tensors(m, span_offsets, span_sizes, span_count)) return false; #endif /* ROCm expands Q8 weights lazily, after required session and support-model * allocations. An eager optional cache can otherwise starve those buffers * on unified-memory devices. */ 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) { if (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset) ds4_die("tensor is outside the mapped weights; Engram requires disk row reads"); 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 inline float ds4_e8m0_to_f32(uint8_t e) { const uint32_t bits = e == 0 ? 0x00400000u : (uint32_t)e << 23; float value; memcpy(&value, &bits, sizeof(value)); return value; } static const float ds4_mxfp4_values[16] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, }; static float ds4_vec_dot_mxfp4_f32(int n, const block_mxfp4 *x, const float *y) { const int nb = n / QK_MXFP4; float sumf = 0.0f; for (int ib = 0; ib < nb; ib++) { const float d = ds4_e8m0_to_f32(x[ib].e); const float *yb = y + (uint64_t)ib * QK_MXFP4; for (int j = 0; j < QK_MXFP4 / 2; j++) { const uint8_t q = x[ib].qs[j]; sumf += d * ds4_mxfp4_values[q & 0x0fu] * yb[j]; sumf += d * ds4_mxfp4_values[q >> 4] * yb[j + QK_MXFP4 / 2]; } } 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 *kda_q; ds4_tensor *kda_k; ds4_tensor *kda_v; ds4_tensor *kda_q_conv; ds4_tensor *kda_k_conv; ds4_tensor *kda_v_conv; ds4_tensor *kda_f_a; ds4_tensor *kda_f_b; ds4_tensor *kda_dt_bias; ds4_tensor *kda_a_log; ds4_tensor *kda_beta; ds4_tensor *kda_g_a; ds4_tensor *kda_g_b; ds4_tensor *kda_o_norm; ds4_tensor *kda_output; 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_exp_probs_vl; 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_tensor *engram_kv; ds4_tensor *engram_q_norm; ds4_tensor *engram_k_norm; /* Qwen3.8-Flash-Next */ ds4_tensor *hc_attn_norm; ds4_tensor *hc_attn_down; ds4_tensor *hc_attn_up; ds4_tensor *hc_attn_inject; ds4_tensor *hc_ffn_norm; ds4_tensor *hc_ffn_down; ds4_tensor *hc_ffn_up; ds4_tensor *hc_ffn_inject; ds4_tensor *attn_q; ds4_tensor *attn_k; ds4_tensor *attn_v; ds4_tensor *attn_q_norm; ds4_tensor *attn_k_norm; ds4_tensor *indexer_q_proj; ds4_tensor *indexer_k_proj; ds4_tensor *indexer_q_norm; ds4_tensor *lin_qkv; ds4_tensor *lin_gate; ds4_tensor *lin_conv; ds4_tensor *lin_dt_bias; ds4_tensor *lin_a; ds4_tensor *lin_beta; ds4_tensor *lin_alpha; ds4_tensor *lin_norm; ds4_tensor *lin_out; ds4_tensor *ple_key; ds4_tensor *ple_value; ds4_tensor *ple_norm_key; ds4_tensor *ple_norm_query; ds4_tensor *ple_norm_conv; ds4_tensor *ple_conv; ds4_tensor *ffn_gate_inp_shexp; ds4_tensor *nextn_hc_head_norm; ds4_tensor *nextn_hc_head_down; ds4_tensor *nextn_hc_head_up; } 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_tensor *output_hc_norm; ds4_tensor *output_hc_down; ds4_tensor *output_hc_up; ds4_tensor *ple_embd; 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 || type == DS4_TENSOR_BF16; } static bool tensor_type_is_dense_quant(uint32_t type) { return type == DS4_TENSOR_Q8_0 || type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_Q4_0; } static void 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 bf16, 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_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_Q2_K || type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_Q5_K || type == DS4_TENSOR_Q6_K || type == DS4_TENSOR_MXFP4; } static DS4_MAYBE_UNUSED uint64_t routed_expert_block_bytes(uint32_t type) { switch (type) { case DS4_TENSOR_Q8_0: return 34; case DS4_TENSOR_IQ2_XXS: return sizeof(block_iq2_xxs); case DS4_TENSOR_Q2_K: return sizeof(block_q2_K); case DS4_TENSOR_Q4_K: return sizeof(block_q4_K); case DS4_TENSOR_Q5_K: return sizeof(block_q5_K); case DS4_TENSOR_Q6_K: return sizeof(block_q6_K); case DS4_TENSOR_MXFP4: return sizeof(block_mxfp4); 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; /* Mixed-precision models can put an outlier quant at the first routed * layer owned by a distributed slice. Choosing that first layer as the * slab class makes every ordinary layer bypass the cache. Use the most * common local size class instead (ties retain the earliest class). */ uint64_t best_bytes = 0; uint32_t best_count = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { uint64_t candidate = 0; if (!streaming_layer_routed_expert_bytes(&weights->layer[il], &candidate)) { continue; } uint32_t count = 0; for (uint32_t jl = 0; jl < DS4_N_LAYER; jl++) { uint64_t bytes = 0; if (streaming_layer_routed_expert_bytes(&weights->layer[jl], &bytes) && bytes == candidate) { count++; } } if (count > best_count) { best_bytes = candidate; best_count = count; } } if (best_count == 0) return false; *per_expert_bytes_out = best_bytes; return true; } enum { DS4_STREAMING_PREFILL_HEADROOM_LAYERS = 2 }; static bool ds4_streaming_cacheable_expert_count( const ds4_weights *weights, uint64_t *experts_out, uint32_t *layers_out) { if (experts_out) *experts_out = 0; if (layers_out) *layers_out = 0; if (!weights || !experts_out || DS4_N_EXPERT == 0) return false; uint64_t slab_bytes = 0; if (!ds4_streaming_routed_expert_bytes(weights, &slab_bytes)) return false; uint32_t layers = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { uint64_t per_expert_bytes = 0; if (!streaming_layer_routed_expert_bytes(&weights->layer[il], &per_expert_bytes)) { continue; } if (per_expert_bytes == slab_bytes) layers++; } if (layers == 0 || (uint64_t)layers > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { return false; } *experts_out = (uint64_t)layers * (uint64_t)DS4_N_EXPERT; if (layers_out) *layers_out = layers; return true; } static bool ds4_streaming_prefill_headroom_bytes( const ds4_weights *weights, uint64_t *bytes_out) { if (bytes_out) *bytes_out = 0; if (!weights || !bytes_out) return false; uint64_t per_expert_bytes = 0; uint64_t cacheable_experts = 0; uint32_t cacheable_layers = 0; if (!ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes) || !ds4_streaming_cacheable_expert_count(weights, &cacheable_experts, &cacheable_layers)) { return false; } (void)cacheable_experts; const uint32_t reserve_layers = cacheable_layers < DS4_STREAMING_PREFILL_HEADROOM_LAYERS ? cacheable_layers : DS4_STREAMING_PREFILL_HEADROOM_LAYERS; if (per_expert_bytes > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { return false; } const uint64_t layer_bytes = per_expert_bytes * (uint64_t)DS4_N_EXPERT; if (reserve_layers != 0 && layer_bytes > UINT64_MAX / (uint64_t)reserve_layers) return false; *bytes_out = layer_bytes * (uint64_t)reserve_layers; return true; } /* * Mixed-precision ("boosted") GGUFs upcast a few layers' routed experts to a * bigger quant (e.g. Q4_K among IQ2 layers). The streaming expert cache is a * single-size-class slab allocator sized from the dominant local routed-layer * size class, so other layers can never be served from it: they must read * expert weights through the mapped-model views instead. A layer is "uniform" * iff its per-expert bytes match the slab class. */ static DS4_MAYBE_UNUSED bool weights_streaming_layer_experts_uniform( const ds4_weights *w, uint32_t il) { uint64_t base = 0; uint64_t bytes = 0; if (!w || il >= DS4_N_LAYER) return true; const ds4_layer_weights *l = &w->layer[il]; if (!streaming_layer_routed_expert_bytes(l, &bytes)) return true; if (!ds4_streaming_routed_expert_bytes(w, &base)) return true; return bytes == base; } static 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 || DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { return w && w->output_norm && w->output; } if (ds4_model_is_qwen4()) { return w && w->output_hc_norm && w->output_hc_down && w->output_hc_up && 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 || DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { return w && (w->output_norm || w->output); } if (ds4_model_is_qwen4()) { return w && (w->output_hc_norm || w->output_hc_down || w->output_hc_up || 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->ffn_norm) return false; if (ds4_model_is_glm53() && il + DS4_N_NEXTN_PREDICT < DS4_N_LAYER && (!l->hc_attn_fn || !l->hc_attn_scale || !l->hc_attn_base || !l->hc_ffn_fn || !l->hc_ffn_scale || !l->hc_ffn_base)) { return false; } if (ds4_glm53_layer_is_kda(il)) { if (!l->kda_q || !l->kda_k || !l->kda_v || !l->kda_q_conv || !l->kda_k_conv || !l->kda_v_conv || !l->kda_f_a || !l->kda_f_b || !l->kda_dt_bias || !l->kda_a_log || !l->kda_beta || !l->kda_g_a || !l->kda_g_b || !l->kda_o_norm || !l->kda_output) { return false; } } else { if (!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) { return false; } if (ds4_model_is_glm53() && (!l->indexer_compressor_ape || !l->indexer_compressor_gate)) { 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_qwen4_layer_has_required(const ds4_layer_weights *l, uint32_t il) { if (!l) return false; if (!l->hc_attn_norm || !l->hc_attn_down || !l->hc_attn_up || !l->hc_attn_inject || !l->hc_ffn_norm || !l->hc_ffn_down || !l->hc_ffn_up || !l->hc_ffn_inject) { return false; } if (ds4_qwen4_layer_is_linear(il)) { if (!l->lin_qkv || !l->lin_gate || !l->lin_conv || !l->lin_dt_bias || !l->lin_a || !l->lin_beta || !l->lin_alpha || !l->lin_norm || !l->lin_out) { return false; } } else { if (!l->attn_q || !l->attn_k || !l->attn_v || !l->attn_output || !l->attn_q_norm || !l->attn_k_norm || !l->indexer_q_proj || !l->indexer_k_proj || !l->indexer_q_norm || !l->indexer_k_norm) { return false; } } if (ds4_qwen4_layer_is_ple(il) && (!l->ple_key || !l->ple_value || !l->ple_norm_key || !l->ple_norm_query || !l->ple_norm_conv || !l->ple_conv)) { return false; } if (!l->ffn_gate_inp || !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || !l->ffn_gate_inp_shexp || !l->ffn_gate_shexp || !l->ffn_up_shexp || !l->ffn_down_shexp) { return false; } if (ds4_qwen4_layer_is_nextn(il) && (!l->nextn_eh_proj || !l->nextn_enorm || !l->nextn_hnorm || !l->nextn_hc_head_norm || !l->nextn_hc_head_down || !l->nextn_hc_head_up)) { return false; } return true; } /* Dense Qwen projections: Q8_0, F16 or F32, plus BF16 and Q4_0 from the upstream GGUF. */ static bool tensor_type_is_qwen4_dense(uint32_t type) { return type == DS4_TENSOR_Q8_0 || type == DS4_TENSOR_F16 || type == DS4_TENSOR_F32 || type == DS4_TENSOR_BF16 || type == DS4_TENSOR_Q4_0; } static void tensor_expect_qwen4_dense_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_qwen4_dense(t->type)) { fprintf(stderr, "ds4: tensor %.*s has type %u, expected Q8_0, Q4_0, F16, BF16 or F32\n", (int)t->name.len, t->name.ptr, t->type); exit(1); } tensor_expect_layout(t, t->type, ndim, d0, d1, d2); } static void tensor_expect_qwen4_expert_layout( const ds4_tensor *t, uint64_t d0, uint64_t d1, uint64_t d2) { if (!t) ds4_die("internal error: missing tensor while validating layout"); if (!tensor_is_routed_expert_type(t->type) && t->type != DS4_TENSOR_F16 && t->type != DS4_TENSOR_F32 && t->type != DS4_TENSOR_Q4_0) { fprintf(stderr, "ds4: routed expert tensor %.*s has unsupported type %u\n", (int)t->name.len, t->name.ptr, t->type); exit(1); } tensor_expect_layout(t, t->type, 3, d0, d1, d2); } static void weights_validate_qwen4_layout( const ds4_weights *w, uint32_t layer_start, uint32_t layer_end, bool require_token_embd, bool require_output) { const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM; const uint64_t lin_k_dim = (uint64_t)DS4_N_LIN_K_HEAD * DS4_N_LIN_HEAD_DIM; const uint64_t lin_v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; 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 Qwen layout"); if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in Qwen 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 Qwen weight layout validation"); } if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); if (w->token_embd) { tensor_expect_qwen4_dense_layout(w->token_embd, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); } if (require_token_embd && !w->ple_embd) ds4_die("Qwen GGUF lacks its n-grams; repack with gguf-tools/qwen4_native_ngrams.py"); if (w->ple_embd) { const uint32_t t = w->ple_embd->type; if (t != DS4_TENSOR_BF16) { fprintf(stderr, "ds4: n-gram embeddings must use original BF16, got type %u\n", t); exit(1); } if (w->ple_embd->ndim != 2 || w->ple_embd->dim[0] != DS4_N_PLE_HEAD_DIM || w->ple_embd->dim[1] < g_ds4_qwen4_ple.n_rows) { fprintf(stderr, "ds4: per_layer_token_embd layout [%" PRIu64 ", %" PRIu64 "] does not " "cover %u x %" PRIu64 " hash rows\n", w->ple_embd->dim[0], w->ple_embd->dim[1], DS4_N_PLE_HEAD_DIM, g_ds4_qwen4_ple.n_rows); exit(1); } } const bool have_output = weights_have_output_head(w); if (require_output && !have_output) ds4_die("required output head tensors are missing"); if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); if (have_output) { tensor_expect_layout(w->output_hc_norm, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_qwen4_dense_layout(w->output_hc_down, 2, hc_dim, DS4_N_HC_LOWRANK, 0); tensor_expect_qwen4_dense_layout(w->output_hc_up, 2, DS4_N_HC_LOWRANK, hc_dim, 0); tensor_expect_qwen4_dense_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_qwen4_layer_has_required(l, il)) { fprintf(stderr, "ds4: required Qwen tensors for layer %u are missing\n", il); exit(1); } tensor_expect_layout(l->hc_attn_norm, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_qwen4_dense_layout(l->hc_attn_down, 2, hc_dim, DS4_N_HC_LOWRANK, 0); tensor_expect_qwen4_dense_layout(l->hc_attn_up, 2, DS4_N_HC_LOWRANK, hc_dim, 0); tensor_expect_qwen4_dense_layout(l->hc_attn_inject, 2, hc_dim, DS4_N_HC, 0); tensor_expect_layout(l->hc_ffn_norm, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_qwen4_dense_layout(l->hc_ffn_down, 2, hc_dim, DS4_N_HC_LOWRANK, 0); tensor_expect_qwen4_dense_layout(l->hc_ffn_up, 2, DS4_N_HC_LOWRANK, hc_dim, 0); tensor_expect_qwen4_dense_layout(l->hc_ffn_inject, 2, hc_dim, DS4_N_HC, 0); if (ds4_qwen4_layer_is_linear(il)) { tensor_expect_qwen4_dense_layout(l->lin_qkv, 2, DS4_N_EMBD, 2u * lin_k_dim + lin_v_dim, 0); tensor_expect_qwen4_dense_layout(l->lin_gate, 2, DS4_N_EMBD, lin_v_dim, 0); tensor_expect_layout(l->lin_conv, DS4_TENSOR_F32, 2, DS4_N_LIN_CONV, 2u * lin_k_dim + lin_v_dim, 0); tensor_expect_layout(l->lin_dt_bias, DS4_TENSOR_F32, 1, DS4_N_LIN_V_HEAD, 0, 0); tensor_expect_layout(l->lin_a, DS4_TENSOR_F32, 1, DS4_N_LIN_V_HEAD, 0, 0); tensor_expect_qwen4_dense_layout(l->lin_beta, 2, DS4_N_EMBD, DS4_N_LIN_V_HEAD, 0); tensor_expect_qwen4_dense_layout(l->lin_alpha, 2, DS4_N_EMBD, DS4_N_LIN_V_HEAD, 0); tensor_expect_layout(l->lin_norm, DS4_TENSOR_F32, 1, DS4_N_LIN_HEAD_DIM, 0, 0); tensor_expect_qwen4_dense_layout(l->lin_out, 2, lin_v_dim, DS4_N_EMBD, 0); } else { tensor_expect_qwen4_dense_layout(l->attn_q, 2, DS4_N_EMBD, 2u * q_dim, 0); tensor_expect_qwen4_dense_layout(l->attn_k, 2, DS4_N_EMBD, kv_dim, 0); tensor_expect_qwen4_dense_layout(l->attn_v, 2, DS4_N_EMBD, kv_dim, 0); tensor_expect_qwen4_dense_layout(l->attn_output, 2, q_dim, DS4_N_EMBD, 0); tensor_expect_layout(l->attn_q_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); tensor_expect_layout(l->attn_k_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); tensor_expect_qwen4_dense_layout(l->indexer_q_proj, 2, DS4_N_EMBD, index_q_dim, 0); tensor_expect_qwen4_dense_layout(l->indexer_k_proj, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 0); tensor_expect_layout(l->indexer_q_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); tensor_expect_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); } if (ds4_qwen4_layer_is_ple(il)) { tensor_expect_qwen4_dense_layout(l->ple_key, 2, DS4_N_EMBD, hc_dim, 0); tensor_expect_qwen4_dense_layout(l->ple_value, 2, DS4_N_EMBD, DS4_N_EMBD, 0); tensor_expect_layout(l->ple_norm_key, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_layout(l->ple_norm_query, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_layout(l->ple_norm_conv, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_layout(l->ple_conv, l->ple_conv->type == DS4_TENSOR_F16 ? DS4_TENSOR_F16 : DS4_TENSOR_F32, 2, DS4_N_PLE_CONV, hc_dim, 0); } tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); tensor_expect_qwen4_expert_layout(l->ffn_gate_exps, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); tensor_expect_qwen4_expert_layout(l->ffn_up_exps, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); /* Q2_K down rows store 640 logical inputs in three 256-value blocks. * Activations remain 640 wide; only the weight row stride is padded. */ const uint32_t down_width = l->ffn_down_exps->type == DS4_TENSOR_Q2_K ? (DS4_N_FF_EXP + 255u) / 256u * 256u : DS4_N_FF_EXP; tensor_expect_qwen4_expert_layout(l->ffn_down_exps, down_width, 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); } /* llama.cpp writes the shared-expert gate as a 1-D [n_embd] tensor */ if (l->ffn_gate_inp_shexp->ndim == 1) { tensor_expect_layout(l->ffn_gate_inp_shexp, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); } else { tensor_expect_layout(l->ffn_gate_inp_shexp, DS4_TENSOR_F32, 2, DS4_N_EMBD, 1, 0); } tensor_expect_qwen4_dense_layout(l->ffn_gate_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); tensor_expect_qwen4_dense_layout(l->ffn_up_shexp, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); tensor_expect_qwen4_dense_layout(l->ffn_down_shexp, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); if (ds4_qwen4_layer_is_nextn(il)) { tensor_expect_qwen4_dense_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, hc_dim, 0, 0); tensor_expect_layout(l->nextn_hc_head_norm, DS4_TENSOR_F32, 1, hc_dim, 0, 0); tensor_expect_qwen4_dense_layout(l->nextn_hc_head_down, 2, hc_dim, DS4_N_HC_LOWRANK, 0); tensor_expect_qwen4_dense_layout(l->nextn_hc_head_up, 2, DS4_N_HC_LOWRANK, hc_dim, 0); } } } static bool weights_layer_has_required(const ds4_layer_weights *l, uint32_t il) { if (!l) return false; if (ds4_model_is_qwen4()) { return weights_qwen4_layer_has_required(l, il); } 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 (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { if (!l->ffn_exp_probs_b) return false; if (ds41_kv_source(il) && (!l->attn_compressor_kv || !l->attn_compressor_norm || !l->indexer_attn_k || !l->indexer_k_norm || (ratio == 2 && !l->attn_compressor_gate))) return false; if (ds41_index_source(il) && (!l->indexer_attn_q_b || !l->indexer_proj)) return false; if (ds41_engram_layer(il) && (!l->engram_kv || !l->engram_q_norm || !l->engram_k_norm)) return false; return true; } 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; 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 kda_projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_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); if (ds4_model_is_glm53() && il + DS4_N_NEXTN_PREDICT < DS4_N_LAYER) { tensor_expect_layout(l->hc_attn_fn, DS4_TENSOR_BF16, 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->hc_ffn_fn, DS4_TENSOR_BF16, 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); } if (ds4_glm53_layer_is_kda(il)) { tensor_expect_glm_dense_quant_layout(l->kda_q, 2, DS4_N_EMBD, kda_projection, 0); tensor_expect_glm_dense_quant_layout(l->kda_k, 2, DS4_N_EMBD, kda_projection, 0); tensor_expect_glm_dense_quant_layout(l->kda_v, 2, DS4_N_EMBD, kda_projection, 0); tensor_expect_layout(l->kda_q_conv, DS4_TENSOR_F32, 3, DS4_N_KDA_CONV, 1, kda_projection); tensor_expect_layout(l->kda_k_conv, DS4_TENSOR_F32, 3, DS4_N_KDA_CONV, 1, kda_projection); tensor_expect_layout(l->kda_v_conv, DS4_TENSOR_F32, 3, DS4_N_KDA_CONV, 1, kda_projection); tensor_expect_glm_dense_quant_layout(l->kda_f_a, 2, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, 0); tensor_expect_glm_dense_quant_layout(l->kda_f_b, 2, DS4_N_KDA_HEAD_DIM, kda_projection, 0); tensor_expect_layout(l->kda_dt_bias, DS4_TENSOR_F32, 1, kda_projection, 0, 0); tensor_expect_layout(l->kda_a_log, DS4_TENSOR_F32, 1, DS4_N_KDA_HEAD, 0, 0); tensor_expect_glm_dense_quant_layout(l->kda_beta, 2, DS4_N_EMBD, DS4_N_KDA_HEAD, 0); tensor_expect_glm_dense_quant_layout(l->kda_g_a, 2, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, 0); tensor_expect_glm_dense_quant_layout(l->kda_g_b, 2, DS4_N_KDA_HEAD_DIM, kda_projection, 0); tensor_expect_layout(l->kda_o_norm, DS4_TENSOR_F32, 1, DS4_N_KDA_HEAD_DIM, 0, 0); tensor_expect_glm_dense_quant_layout(l->kda_output, 2, kda_projection, DS4_N_EMBD, 0); } else { 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_KV_LORA + DS4_N_ROT, 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); if (ds4_model_is_glm53()) { tensor_expect_glm_dense_quant_layout( l->indexer_proj, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); } else { tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); } if (ds4_model_is_glm53()) { tensor_expect_layout(l->indexer_compressor_ape, DS4_TENSOR_BF16, 2, DS4_N_INDEXER_HEAD_DIM, 4, 0); tensor_expect_layout(l->indexer_compressor_gate, DS4_TENSOR_BF16, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, 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; } if (ds4_model_is_qwen4()) { weights_validate_qwen4_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) { 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) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) { tensor_expect_layout(w->output_hc_base, DS4_TENSOR_F32, 1, DS4_N_HC, 0, 0); tensor_expect_layout(w->output_hc_fn, DS4_TENSOR_F16, 2, hc_dim, DS4_N_HC, 0); tensor_expect_layout(w->output_hc_scale, DS4_TENSOR_F32, 1, 1, 0, 0); } tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); tensor_expect_dense_quant_layout(w->output, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); } for (uint32_t il = layer_start; il <= layer_end; il++) { const ds4_layer_weights *l = &w->layer[il]; const uint32_t ratio = ds4_layer_compress_ratio(il); if (!weights_layer_has_required(l, il)) { fprintf(stderr, "ds4: required tensors for layer %u are missing\n", il); exit(1); } tensor_expect_layout(l->hc_attn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); tensor_expect_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); tensor_expect_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); tensor_expect_dense_quant_layout(l->attn_q_a, 2, DS4_N_EMBD, DS4_N_LORA_Q, 0); tensor_expect_layout(l->attn_q_a_norm, DS4_TENSOR_F32, 1, DS4_N_LORA_Q, 0, 0); tensor_expect_dense_quant_layout(l->attn_q_b, 2, DS4_N_LORA_Q, q_dim, 0); tensor_expect_dense_quant_layout(l->attn_kv, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); tensor_expect_layout(l->attn_kv_a_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); tensor_expect_layout(l->attn_sinks, DS4_TENSOR_F32, 1, DS4_N_HEAD, 0, 0); tensor_expect_dense_quant_layout(l->attn_output_a, 2, DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP), out_low_dim, 0); tensor_expect_dense_quant_layout(l->attn_output_b, 2, out_low_dim, DS4_N_EMBD, 0); if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { /* The V4.1 grouped output kernels consume Q8 blocks directly. */ 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); if (ds41_kv_source(il)) { tensor_expect_layout(l->attn_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); tensor_expect_layout(l->attn_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); if (ratio == 2) tensor_expect_layout(l->attn_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_HEAD_DIM, 0); tensor_expect_layout(l->indexer_attn_k, DS4_TENSOR_F16, 2, DS4_N_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM, 0); tensor_expect_layout(l->indexer_k_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); } if (ds41_index_source(il)) { tensor_expect_layout(l->indexer_attn_q_b, DS4_TENSOR_F16, 2, DS4_N_LORA_Q, DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, 0); tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); } if (ds41_engram_layer(il)) { tensor_expect_layout(l->engram_kv, DS4_TENSOR_F16, 2, 24u * 256u, (DS4_N_HC + 1u) * DS4_N_EMBD, 0); tensor_expect_layout(l->engram_q_norm, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_HC, 0); tensor_expect_layout(l->engram_k_norm, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_HC, 0); } } else if (ratio != 0) { const uint32_t coff = ratio == 4 ? 2u : 1u; const uint64_t comp_width = (uint64_t)coff * DS4_N_HEAD_DIM; tensor_expect_layout(l->attn_compressor_ape, DS4_TENSOR_F16, 2, comp_width, ratio, 0); tensor_expect_layout(l->attn_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); tensor_expect_layout(l->attn_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, comp_width, 0); tensor_expect_layout(l->attn_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); } if (ratio == 4) { const uint64_t index_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; const uint64_t index_width = 2u * DS4_N_INDEXER_HEAD_DIM; tensor_expect_f16_or_q8_0_layout(l->indexer_attn_q_b, 2, DS4_N_LORA_Q, index_q_dim, 0); tensor_expect_layout(l->indexer_proj, DS4_TENSOR_F16, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 0); tensor_expect_layout(l->indexer_compressor_ape, DS4_TENSOR_F16, 2, index_width, ratio, 0); tensor_expect_layout(l->indexer_compressor_kv, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); tensor_expect_layout(l->indexer_compressor_gate, DS4_TENSOR_F16, 2, DS4_N_EMBD, index_width, 0); tensor_expect_layout(l->indexer_compressor_norm, DS4_TENSOR_F32, 1, DS4_N_INDEXER_HEAD_DIM, 0, 0); } tensor_expect_layout(l->hc_ffn_fn, DS4_TENSOR_F16, 2, hc_dim, hc_mix_dim, 0); tensor_expect_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0); tensor_expect_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix_dim, 0, 0); tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); tensor_expect_layout(l->ffn_gate_inp, DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 ? DS4_TENSOR_F32 : 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_epsilon(const char *name, float got, float expected) { const float tolerance = fmaxf(fabsf(expected) * 1.0e-5f, 1.0e-30f); if (fabsf(got - expected) <= tolerance) return; fprintf(stderr, "ds4: expected %s=%.9g for %s, got %.9g\n", name, (double)expected, DS4_MODEL_SHAPE_NAME, (double)got); exit(1); } static void config_expect_bool(const char *name, bool got, bool expected) { if (got == expected) return; fprintf(stderr, "ds4: expected %s=%s for %s, got %s\n", name, expected ? "true" : "false", DS4_MODEL_SHAPE_NAME, got ? "true" : "false"); exit(1); } static void config_validate_fixed_shape(uint32_t n_layer) { config_expect_u32("block_count", n_layer, DS4_N_LAYER); } /* Validate metadata values that affect semantics: attention shape, HC count, * expert routing, RoPE scaling, compression ratios, and SwiGLU clamp. */ static void config_validate_deepseek4_model(const ds4_model *m) { const uint32_t n_layer = required_u32(m, "deepseek4.block_count"); const uint32_t n_embd = required_u32(m, "deepseek4.embedding_length"); const uint32_t n_vocab = required_u32(m, "deepseek4.vocab_size"); const uint32_t n_head = required_u32(m, "deepseek4.attention.head_count"); const uint32_t n_head_kv = required_u32(m, "deepseek4.attention.head_count_kv"); const uint32_t n_head_dim = required_u32(m, "deepseek4.attention.key_length"); const uint32_t n_value_dim = required_u32(m, "deepseek4.attention.value_length"); const uint32_t n_rot = required_u32(m, "deepseek4.rope.dimension_count"); const uint32_t n_lora_q = required_u32(m, "deepseek4.attention.q_lora_rank"); const uint32_t n_lora_o = required_u32(m, "deepseek4.attention.output_lora_rank"); const uint32_t n_out_group = required_u32(m, "deepseek4.attention.output_group_count"); const uint32_t n_expert = required_u32(m, "deepseek4.expert_count"); const uint32_t n_expert_used = required_u32(m, "deepseek4.expert_used_count"); const uint32_t n_ff_exp = required_u32(m, "deepseek4.expert_feed_forward_length"); const uint32_t n_expert_shared = required_u32(m, "deepseek4.expert_shared_count"); const uint32_t n_hash_layer = required_u32(m, "deepseek4.hash_layer_count"); uint32_t n_expert_groups = 0; uint32_t n_group_used = 0; model_get_u32(m, "deepseek4.expert_group_count", &n_expert_groups); model_get_u32(m, "deepseek4.expert_group_used_count", &n_group_used); const uint32_t n_swa = required_u32(m, "deepseek4.attention.sliding_window"); const uint32_t n_indexer_head = required_u32(m, "deepseek4.attention.indexer.head_count"); const uint32_t n_indexer_head_dim = required_u32(m, "deepseek4.attention.indexer.key_length"); const uint32_t n_indexer_top_k = required_u32(m, "deepseek4.attention.indexer.top_k"); const uint32_t n_hc = required_u32(m, "deepseek4.hyper_connection.count"); const uint32_t n_hc_sinkhorn_iter = required_u32(m, "deepseek4.hyper_connection.sinkhorn_iterations"); ds4_select_shape_from_metadata(n_layer, n_embd, n_vocab, n_head, n_head_kv, n_head_dim, n_value_dim, n_rot, n_lora_q, n_lora_o, n_out_group, n_expert, n_expert_used, n_ff_exp, n_expert_shared, n_hash_layer, n_swa, n_indexer_head, n_indexer_head_dim, n_indexer_top_k, n_hc, n_hc_sinkhorn_iter); g_ds4_flash_vision_exp = false; ds4_str checkpoint_variant = {0}; if (model_get_string(m, "deepseek4.checkpoint_variant", &checkpoint_variant)) { if (!ds4_streq(checkpoint_variant, "vision-exp")) { ds4_die("unsupported deepseek4.checkpoint_variant"); } if (g_ds4_shape.variant != DS4_VARIANT_FLASH || !required_bool(m, "deepseek4.vision.sidecar_required")) { ds4_die("invalid DeepSeek Vision-Exp model metadata"); } ds4_str source_revision = {0}; if (!model_get_string(m, "general.source.revision", &source_revision) || !ds4_streq(source_revision, "e46e16bf6035c6f317eb2ac7458eb0362926d402")) { ds4_die("unexpected DeepSeek Vision-Exp source revision"); } g_ds4_shape.name = "DeepSeek V4 Flash Vision Experimental"; g_ds4_shape.rms_eps = 1.0e-20f; g_ds4_flash_vision_exp = true; } 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_epsilon("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_glm53_layer_types(const ds4_model *m) { const char *key = "glm5-next.layer_types"; ds4_array_ref arr; if (!model_get_array(m, key, &arr) || (arr.type != GGUF_VALUE_UINT32 && arr.type != GGUF_VALUE_INT32) || arr.len != DS4_N_LAYER) { ds4_die("glm5-next.layer_types must be an int32/uint32 array with one entry per layer"); } 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 value = 0; if (!cursor_read(&c, &value, sizeof(value))) ds4_die(c.error); if (value < 0) ds4_die("glm5-next.layer_types contains a negative value"); got = (uint32_t)value; } const uint32_t expected = ds4_glm53_layer_is_kda(il) ? 0u : 1u; if (got != expected) { fprintf(stderr, "ds4: unexpected GLM 5.3 attention type at layer %u: got %u, expected %u\n", il, got, expected); exit(1); } } } static void config_validate_glm53_model(const ds4_model *m) { g_ds4_shape = DS4_SHAPE_GLM53; memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); config_expect_u32("block_count", required_u32(m, "glm5-next.block_count"), DS4_N_LAYER); config_expect_u32("trunk_block_count", required_u32(m, "glm5-next.trunk_block_count"), DS4_N_LAYER - DS4_N_NEXTN_PREDICT); config_expect_u32("nextn_predict_layers", required_u32(m, "glm5-next.nextn_predict_layers"), DS4_N_NEXTN_PREDICT); config_expect_u64("context_length", required_u64_compat(m, "glm5-next.context_length"), DS4_ROPE_ORIG_CTX); config_expect_u32("embedding_length", required_u32(m, "glm5-next.embedding_length"), DS4_N_EMBD); config_expect_u32("vocab_size", required_u32(m, "glm5-next.vocab_size"), DS4_N_VOCAB); config_expect_u32("feed_forward_length", required_u32(m, "glm5-next.feed_forward_length"), DS4_N_FF_DENSE); config_expect_u32("expert_feed_forward_length", required_u32(m, "glm5-next.expert_feed_forward_length"), DS4_N_FF_EXP); config_expect_u32("expert_count", required_u32(m, "glm5-next.expert_count"), DS4_N_EXPERT); config_expect_u32("expert_used_count", required_u32(m, "glm5-next.expert_used_count"), DS4_N_EXPERT_USED); config_expect_u32("expert_shared_count", required_u32(m, "glm5-next.expert_shared_count"), DS4_N_EXPERT_SHARED); config_expect_u32("leading_dense_block_count", required_u32(m, "glm5-next.leading_dense_block_count"), DS4_N_LEADING_DENSE); config_expect_f32("expert_weights_scale", required_f32(m, "glm5-next.expert_weights_scale"), DS4_EXPERT_WEIGHT_SCALE); config_expect_bool("expert_weights_norm", required_bool(m, "glm5-next.expert_weights_norm"), true); config_expect_f32("swiglu_limit", required_f32(m, "glm5-next.swiglu_limit"), DS4_SWIGLU_CLAMP_EXP); config_expect_f32("attention.layer_norm_rms_epsilon", required_f32(m, "glm5-next.attention.layer_norm_rms_epsilon"), DS4_RMS_EPS); config_expect_u32("attention.head_count", required_u32(m, "glm5-next.attention.head_count"), DS4_N_HEAD); config_expect_u32("attention.key_length", required_u32(m, "glm5-next.attention.key_length"), DS4_N_KEY_MLA); config_expect_u32("attention.value_length", required_u32(m, "glm5-next.attention.value_length"), DS4_N_VALUE_MLA); config_expect_u32("attention.q_lora_rank", required_u32(m, "glm5-next.attention.q_lora_rank"), DS4_N_LORA_Q); config_expect_u32("attention.kv_lora_rank", required_u32(m, "glm5-next.attention.kv_lora_rank"), DS4_N_KV_LORA); config_expect_u32("attention.rope_dimension_count", required_u32(m, "glm5-next.attention.rope_dimension_count"), DS4_N_ROT); config_expect_u32("attention.indexer.head_count", required_u32(m, "glm5-next.attention.indexer.head_count"), DS4_N_INDEXER_HEAD); config_expect_u32("attention.indexer.key_length", required_u32(m, "glm5-next.attention.indexer.key_length"), DS4_N_INDEXER_HEAD_DIM); config_expect_u32("attention.indexer.top_k", required_u32(m, "glm5-next.attention.indexer.top_k"), DS4_N_INDEXER_TOP_K); config_expect_u32("attention.indexer.pool_size", required_u32(m, "glm5-next.attention.indexer.pool_size"), 4); config_expect_u32("linear_attention.head_count", required_u32(m, "glm5-next.linear_attention.head_count"), DS4_N_KDA_HEAD); config_expect_u32("linear_attention.head_dimension", required_u32(m, "glm5-next.linear_attention.head_dimension"), DS4_N_KDA_HEAD_DIM); config_expect_u32("linear_attention.conv_kernel", required_u32(m, "glm5-next.linear_attention.conv_kernel"), DS4_N_KDA_CONV); config_expect_f32("linear_attention.gate_lower_bound", required_f32(m, "glm5-next.linear_attention.gate_lower_bound"), DS4_KDA_GATE_LOWER_BOUND); config_expect_u32("hyper_connection.count", required_u32(m, "glm5-next.hyper_connection.count"), DS4_N_HC); config_expect_u32("hyper_connection.sinkhorn_iterations", required_u32(m, "glm5-next.hyper_connection.sinkhorn_iterations"), DS4_N_HC_SINKHORN_ITER); config_expect_f32("hyper_connection.epsilon", required_f32(m, "glm5-next.hyper_connection.epsilon"), DS4_HC_EPS); config_validate_glm53_layer_types(m); } static void config_expect_u32_array(const ds4_model *m, const char *key, const uint32_t *expected, uint32_t count) { ds4_array_ref arr; if (!model_get_array(m, key, &arr) || arr.type != GGUF_VALUE_UINT32 || arr.len != count) { fprintf(stderr, "ds4: %s must be a uint32 array of length %u\n", key, count); exit(1); } ds4_cursor c = cursor_at(m, arr.data_pos); for (uint32_t i = 0; i < count; i++) { uint32_t got; if (!cursor_u32(&c, &got)) ds4_die(c.error); config_expect_u32(key, got, expected[i]); } } static void config_validate_deepseek41_model(const ds4_model *m) { g_ds4_shape = DS4_SHAPE_FLASH41; memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); for (uint32_t il = 0; il < DS4_N_LAYER; il++) g_ds4_compress_ratios[il] = ds4_expected_layer_compress_ratio(il); const struct { const char *key; uint32_t value; } ints[] = { {"vocab_size", DS4_N_VOCAB}, {"hidden_size", DS4_N_EMBD}, {"moe_intermediate_size", DS4_N_FF_EXP}, {"num_hidden_layers", DS4_N_LAYER}, {"num_attention_heads", DS4_N_HEAD}, {"num_key_value_heads", DS4_N_HEAD_KV}, {"head_dim", DS4_N_HEAD_DIM}, {"qk_rope_head_dim", DS4_N_ROT}, {"q_lora_rank", DS4_N_LORA_Q}, {"o_lora_rank", DS4_N_LORA_O}, {"o_groups", DS4_N_OUT_GROUP}, {"n_routed_experts", DS4_N_EXPERT}, {"n_shared_experts", DS4_N_EXPERT_SHARED}, {"num_experts_per_tok", DS4_N_EXPERT_USED}, {"max_position_embeddings", 1048576}, {"sliding_window", DS4_N_SWA}, {"index_n_heads", DS4_N_INDEXER_HEAD}, {"index_head_dim", DS4_N_INDEXER_HEAD_DIM}, {"index_topk", DS4_N_INDEXER_TOP_K}, {"candidate_source_layer_id", 20}, {"candidate_topk_blocks", 2048}, {"candidate_block_size", 8}, {"hc_mult", DS4_N_HC}, {"hc_sinkhorn_iters", DS4_N_HC_SINKHORN_ITER}, {"rope_theta", 10000}, {"compress_rope_theta", 160000}, {"engram.compressed_vocab_size", 99092}, }; for (size_t i = 0; i < sizeof(ints) / sizeof(ints[0]); i++) { char key[128]; snprintf(key, sizeof(key), "deepseek41.%s", ints[i].key); config_expect_u32(key, required_u32(m, key), ints[i].value); } const struct { const char *key; float value; } floats[] = { {"rms_norm_eps", DS4_RMS_EPS}, {"hc_eps", DS4_HC_EPS}, {"swiglu_limit", DS4_SWIGLU_CLAMP_EXP}, {"routed_scaling_factor", DS4_EXPERT_WEIGHT_SCALE}, {"rope_scaling.factor", DS4_ROPE_SCALE_FACTOR}, {"rope_scaling.beta_fast", DS4_ROPE_YARN_BETA_FAST}, {"rope_scaling.beta_slow", DS4_ROPE_YARN_BETA_SLOW}, {"rope_scaling.original_max_position_embeddings", DS4_ROPE_ORIG_CTX}, }; for (size_t i = 0; i < sizeof(floats) / sizeof(floats[0]); i++) { char key[128]; snprintf(key, sizeof(key), "deepseek41.%s", floats[i].key); config_expect_epsilon(key, required_f32(m, key), floats[i].value); } const struct { const char *key, *value; } strings[] = { {"deepseek41.scoring_func", "sqrtsoftplus"}, {"deepseek41.hidden_act", "silu"}, {"deepseek41.topk_method", "noaux_tc"}, {"deepseek41.engram.encoding", "e4m3_e8m0_32_row264"}, }; for (size_t i = 0; i < sizeof(strings) / sizeof(strings[0]); i++) { ds4_str got = {0}; if (!model_get_string(m, strings[i].key, &got) || !ds4_streq(got, strings[i].value)) { fprintf(stderr, "ds4: unsupported %s\n", strings[i].key); exit(1); } } config_expect_bool("deepseek41.norm_topk_prob", required_bool(m, "deepseek41.norm_topk_prob"), true); const uint32_t kv[] = {2, 8, 14, 20}, index[] = {2, 8, 14, 20, 24, 28, 32, 36}; const uint32_t engram[] = {1, 14}, rows[] = {384006168, 384016682}; config_expect_u32_array(m, "deepseek41.kv_source_layer_ids", kv, 4); config_expect_u32_array(m, "deepseek41.index_source_layer_ids", index, 8); config_expect_u32_array(m, "deepseek41.engram.layer_ids", engram, 2); config_expect_u32_array(m, "deepseek41.engram.rows", rows, 2); config_expect_u32_array(m, "deepseek41.compress_ratios", g_ds4_compress_ratios, DS4_N_LAYER); for (uint32_t i = 0; i < 2; i++) { ds4_tensor *t = required_tensorf(m, "blk.%u.engram_embd.weight", engram[i]); tensor_expect_layout(t, DS4_TENSOR_I8, 2, 264, rows[i], 0); if (t->abs_offset < m->size) ds4_die("Engram table is still mapped"); } } static void config_read_qwen4_u64_array( const ds4_model *m, const char *key, uint64_t *out, uint32_t max, uint32_t *n_out) { ds4_array_ref arr; if (!model_get_array(m, key, &arr)) { fprintf(stderr, "ds4: required array metadata key is missing: %s\n", key); exit(1); } if (arr.len > max) { fprintf(stderr, "ds4: metadata array %s has %" PRIu64 " entries, at most %u supported\n", key, arr.len, max); exit(1); } ds4_cursor c = cursor_at(m, arr.data_pos); for (uint64_t i = 0; i < arr.len; i++) { uint64_t v = 0; if (arr.type == GGUF_VALUE_UINT64 || arr.type == GGUF_VALUE_INT64) { if (!cursor_read(&c, &v, sizeof(v))) ds4_die(c.error); } else if (arr.type == GGUF_VALUE_UINT32 || arr.type == GGUF_VALUE_INT32) { uint32_t v32 = 0; if (!cursor_u32(&c, &v32)) ds4_die(c.error); v = v32; } else { fprintf(stderr, "ds4: metadata array %s has an unsupported element type %u\n", key, arr.type); exit(1); } out[i] = v; } *n_out = (uint32_t)arr.len; } static float g_qwen4_rope_freq[32]; static float g_qwen4_rope_mscale = 1.0f; static uint32_t g_qwen4_native_ctx = 0; static bool g_qwen4_rope_yarn = false; /* Rotary inverse frequencies of the DS4_N_ROT/2 pairs. DS4_QWEN4_YARN_FACTOR=f * (f > 1) applies static YaRN over the native context (HF * _compute_yarn_parameters with beta_fast 32 / beta_slow 1), the model card's * recipe for prompts beyond 262k tokens; it costs some quality on short text, * so it stays off unless asked for. */ static void qwen4_rope_configure(uint32_t native_ctx) { const uint32_t n_rot = DS4_N_ROT, half = n_rot / 2u; const double base = DS4_ROPE_FREQ_BASE; const char *env = getenv("DS4_QWEN4_YARN_FACTOR"); const double factor = env ? atof(env) : 0.0; const bool yarn = factor > 1.0 && native_ctx > 0; double low = 0.0, high = 0.0; if (yarn) { low = fmax(0.0, floor(n_rot * log(native_ctx / (32.0 * 2.0 * M_PI)) / (2.0 * log(base)))); high = fmin((double)n_rot - 1.0, ceil(n_rot * log(native_ctx / (2.0 * M_PI)) / (2.0 * log(base)))); if (high == low) high += 0.001; } for (uint32_t i = 0; i < 32u; i++) { double f = i < half ? pow(base, -2.0 * (double)i / (double)n_rot) : 0.0; if (yarn && i < half) { const double extrap = 1.0 - fmin(1.0, fmax(0.0, ((double)i - low) / (high - low))); f = (f / factor) * (1.0 - extrap) + f * extrap; } g_qwen4_rope_freq[i] = (float)f; } g_qwen4_rope_mscale = yarn ? (float)(0.1 * log(factor) + 1.0) : 1.0f; g_qwen4_native_ctx = native_ctx; g_qwen4_rope_yarn = yarn; if (yarn) { fprintf(stderr, "ds4: Qwen3.8 YaRN factor %g over %u native tokens (pairs %g..%g blended, mscale %.4f)\n", factor, native_ctx, low, high, g_qwen4_rope_mscale); } #ifdef DS4_HAS_QWEN4_GPU ds4_gpu_qwen4_set_rope(g_qwen4_rope_freq, half, g_qwen4_rope_mscale); #endif } static void config_validate_qwen4_model(const ds4_model *m) { const uint32_t n_embd = required_u32(m, "qwen4exp.embedding_length"); if (n_embd == DS4_SHAPE_QWEN4_EXP.n_embd) { g_ds4_shape = DS4_SHAPE_QWEN4_EXP; } else if (n_embd == DS4_SHAPE_QWEN4_MINI.n_embd) { g_ds4_shape = DS4_SHAPE_QWEN4_MINI; } else { fprintf(stderr, "ds4: unsupported qwen4exp embedding_length %u\n", n_embd); exit(1); } memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); memset(&g_ds4_qwen4_ple, 0, sizeof(g_ds4_qwen4_ple)); /* upstream files carry no MTP block: trim the profile to the trunk */ uint32_t nextn = 0; (void)model_get_u32(m, "qwen4exp.nextn_predict_layers", &nextn); if (nextn > DS4_N_NEXTN_PREDICT) { fprintf(stderr, "ds4: qwen4exp.nextn_predict_layers %u exceeds %u\n", nextn, DS4_N_NEXTN_PREDICT); exit(1); } g_ds4_shape.n_layer -= DS4_N_NEXTN_PREDICT - nextn; g_ds4_shape.n_nextn_predict = nextn; config_expect_u32("block_count", required_u32(m, "qwen4exp.block_count"), DS4_N_LAYER); config_expect_u32("attention.head_count", required_u32(m, "qwen4exp.attention.head_count"), DS4_N_HEAD); config_expect_u32("attention.head_count_kv", required_u32(m, "qwen4exp.attention.head_count_kv"), DS4_N_HEAD_KV); config_expect_u32("attention.key_length", required_u32(m, "qwen4exp.attention.key_length"), DS4_N_HEAD_DIM); config_expect_u32("attention.value_length", required_u32(m, "qwen4exp.attention.value_length"), DS4_N_VALUE_DIM); config_expect_u32("rope.dimension_count", required_u32(m, "qwen4exp.rope.dimension_count"), DS4_N_ROT); config_expect_f32("rope.freq_base", required_f32(m, "qwen4exp.rope.freq_base"), DS4_ROPE_FREQ_BASE); qwen4_rope_configure(required_u32(m, "qwen4exp.context_length")); config_expect_epsilon("attention.layer_norm_rms_epsilon", required_f32(m, "qwen4exp.attention.layer_norm_rms_epsilon"), DS4_RMS_EPS); config_expect_u32("expert_count", required_u32(m, "qwen4exp.expert_count"), DS4_N_EXPERT); config_expect_u32("expert_used_count", required_u32(m, "qwen4exp.expert_used_count"), DS4_N_EXPERT_USED); config_expect_u32("expert_feed_forward_length", required_u32(m, "qwen4exp.expert_feed_forward_length"), DS4_N_FF_EXP); config_expect_u32("expert_shared_feed_forward_length", required_u32(m, "qwen4exp.expert_shared_feed_forward_length"), DS4_N_FF_EXP); config_expect_u32("ssm.conv_kernel", required_u32(m, "qwen4exp.ssm.conv_kernel"), DS4_N_LIN_CONV); config_expect_u32("ssm.state_size", required_u32(m, "qwen4exp.ssm.state_size"), DS4_N_LIN_HEAD_DIM); config_expect_u32("ssm.group_count", required_u32(m, "qwen4exp.ssm.group_count"), DS4_N_LIN_K_HEAD); config_expect_u32("ssm.time_step_rank", required_u32(m, "qwen4exp.ssm.time_step_rank"), DS4_N_LIN_V_HEAD); config_expect_u32("ssm.inner_size", required_u32(m, "qwen4exp.ssm.inner_size"), DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM); config_expect_u32("full_attention_interval", required_u32(m, "qwen4exp.full_attention_interval"), DS4_N_FULL_ATTN_INTERVAL); config_expect_u32("hyper_connection.count", required_u32(m, "qwen4exp.hyper_connection.count"), DS4_N_HC); config_expect_u32("hyper_connection.low_rank", required_u32(m, "qwen4exp.hyper_connection.low_rank"), DS4_N_HC_LOWRANK); config_expect_u32("attention.indexer.head_count", required_u32(m, "qwen4exp.attention.indexer.head_count"), DS4_N_INDEXER_HEAD); config_expect_u32("attention.indexer.key_length", required_u32(m, "qwen4exp.attention.indexer.key_length"), DS4_N_INDEXER_HEAD_DIM); config_expect_u32("attention.indexer.top_k", required_u32(m, "qwen4exp.attention.indexer.top_k"), DS4_N_INDEXER_TOP_K); /* 0 marks GDN layers; attention layers (and the MTP block) run QSA with ratio 4 */ { uint64_t ratios[DS4_MAX_LAYER]; uint32_t n = 0; config_read_qwen4_u64_array(m, "qwen4exp.attention.compress_ratios", ratios, DS4_MAX_LAYER, &n); if (n < DS4_N_LAYER) ds4_die("qwen4exp.attention.compress_ratios is shorter than the layer count"); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t expected = ds4_qwen4_layer_is_linear(il) ? 0u : 4u; if (ratios[il] != expected) { fprintf(stderr, "ds4: unexpected qwen4exp compress ratio at layer %u: got %" PRIu64 ", expected %u\n", il, ratios[il], expected); exit(1); } g_ds4_compress_ratios[il] = (uint32_t)ratios[il]; } } /* PLE n-gram hash: one layer, constants carried verbatim */ { uint64_t vals[DS4_MAX_PLE_HEADS]; uint32_t n = 0; config_read_qwen4_u64_array(m, "qwen4exp.ple.layers", vals, DS4_MAX_PLE_HEADS, &n); if (n != 1 || vals[0] != DS4_N_PLE_LAYER) { fprintf(stderr, "ds4: qwen4exp.ple.layers must be [%u]\n", DS4_N_PLE_LAYER); exit(1); } config_expect_u32("ple.ngram_size", required_u32(m, "qwen4exp.ple.ngram_size"), DS4_N_PLE_NGRAM); config_expect_u32("ple.heads_per_ngram", required_u32(m, "qwen4exp.ple.heads_per_ngram"), DS4_N_PLE_HEADS_PER_NGRAM); config_expect_u32("ple.conv_kernel", required_u32(m, "qwen4exp.ple.conv_kernel"), DS4_N_PLE_CONV); config_expect_u32("ple.eos_token_id", required_u32(m, "qwen4exp.ple.eos_token_id"), (uint32_t)DS4_PLE_EOS_ID); config_expect_u32("embedding_length_per_layer_input", required_u32(m, "qwen4exp.embedding_length_per_layer_input"), DS4_N_PLE_HEAD_DIM); const uint32_t n_heads = DS4_N_PLE_HEADS; if (n_heads > DS4_MAX_PLE_HEADS || DS4_N_PLE_NGRAM > DS4_MAX_PLE_NGRAM) { ds4_die("PLE head or n-gram count exceeds the compiled limit"); } config_read_qwen4_u64_array(m, "qwen4exp.ple.layer_multipliers", vals, DS4_MAX_PLE_HEADS, &n); if (n != DS4_N_PLE_NGRAM) ds4_die("qwen4exp.ple.layer_multipliers has the wrong length"); for (uint32_t i = 0; i < n; i++) g_ds4_qwen4_ple.multiplier[i] = vals[i]; config_read_qwen4_u64_array(m, "qwen4exp.ple.head_offsets", vals, DS4_MAX_PLE_HEADS, &n); if (n != n_heads) ds4_die("qwen4exp.ple.head_offsets has the wrong length"); for (uint32_t i = 0; i < n; i++) { if (vals[i] > UINT32_MAX) ds4_die("PLE head offset does not fit 32 bits"); g_ds4_qwen4_ple.head_offset[i] = (uint32_t)vals[i]; } config_read_qwen4_u64_array(m, "qwen4exp.ple.head_vocab_sizes", vals, DS4_MAX_PLE_HEADS, &n); if (n != n_heads) ds4_die("qwen4exp.ple.head_vocab_sizes has the wrong length"); for (uint32_t i = 0; i < n; i++) { if (vals[i] == 0 || vals[i] > UINT32_MAX) ds4_die("PLE head vocab size is out of range"); g_ds4_qwen4_ple.head_vocab[i] = (uint32_t)vals[i]; const uint64_t end = (uint64_t)g_ds4_qwen4_ple.head_offset[i] + vals[i]; if (end > g_ds4_qwen4_ple.n_rows) g_ds4_qwen4_ple.n_rows = end; } g_ds4_qwen4_ple.n_heads = n_heads; } } static void config_validate_model(const ds4_model *m) { g_ds4_flash_vision_exp = false; ds4_str arch = {0}; if (model_get_string(m, "general.architecture", &arch)) { if (ds4_streq(arch, "deepseek41")) { config_validate_deepseek41_model(m); return; } if (ds4_streq(arch, "glm-dsa")) { config_validate_glm_dsa_model(m); return; } if (ds4_streq(arch, "glm5-next")) { config_validate_glm53_model(m); return; } if (ds4_streq(arch, "qwen4exp")) { config_validate_qwen4_model(m); return; } } config_validate_deepseek4_model(m); } #ifndef DS4_NO_GPU static ds4_tensor *vision_required_tensor( const ds4_model *m, const char *name, uint32_t ndim, const uint64_t *dims) { ds4_tensor *t = required_tensor(m, name); if (t->type != DS4_TENSOR_BF16 || t->ndim != ndim) { fprintf(stderr, "ds4: vision tensor %s has type %s/rank %u, expected BF16/rank %u\n", name, tensor_type_name(t->type), t->ndim, ndim); exit(1); } for (uint32_t d = 0; d < ndim; d++) { if (t->dim[d] == dims[d]) continue; fprintf(stderr, "ds4: vision tensor %s has dim[%u]=%" PRIu64 ", expected %" PRIu64 "\n", name, d, t->dim[d], dims[d]); exit(1); } return t; } static uint64_t vision_required_offset( const ds4_model *m, const char *name, uint32_t ndim, const uint64_t *dims) { return vision_required_tensor(m, name, ndim, dims)->abs_offset; } static void vision_weights_bind( ds4_glm53_vision_weights *w, const ds4_model *m) { ds4_str arch = {0}; if (!model_get_string(m, "general.architecture", &arch) || !ds4_streq(arch, "glm5-next-vision")) { ds4_die("--vision file is not a GLM-5.3 vision encoder GGUF"); } if (m->n_tensors != 347u) { fprintf(stderr, "ds4: vision GGUF has %" PRIu64 " tensors, expected 347\n", m->n_tensors); exit(1); } config_expect_u32("vision block_count", required_u32(m, "glm5-next-vision.block_count"), 24u); config_expect_u32("vision embedding_length", required_u32(m, "glm5-next-vision.embedding_length"), 1024u); config_expect_u32("vision feed_forward_length", required_u32(m, "glm5-next-vision.feed_forward_length"), 4096u); config_expect_u32("vision head_count", required_u32(m, "glm5-next-vision.attention.head_count"), 16u); config_expect_u32("vision projection_length", required_u32(m, "glm5-next-vision.projection_length"), 4096u); config_expect_u32("vision projection feed_forward_length", required_u32(m, "glm5-next-vision.projection.feed_forward_length"), 10240u); config_expect_u32("vision patch_size", required_u32(m, "glm5-next-vision.patch_size"), 14u); config_expect_u32("vision temporal_patch_size", required_u32(m, "glm5-next-vision.temporal_patch_size"), 2u); config_expect_u32("vision spatial_merge_size", required_u32(m, "glm5-next-vision.spatial_merge_size"), 2u); static const uint64_t d1024[] = {1024u}; static const uint64_t d4096[] = {4096u}; static const uint64_t d64[] = {64u}; static const uint64_t d3072[] = {3072u}; static const uint64_t d1024_1024[] = {1024u, 1024u}; static const uint64_t d1024_3072[] = {1024u, 3072u}; static const uint64_t d1024_4096[] = {1024u, 4096u}; static const uint64_t d4096_1024[] = {4096u, 1024u}; static const uint64_t d4096_4096[] = {4096u, 4096u}; static const uint64_t d4096_10240[] = {4096u, 10240u}; static const uint64_t d10240_4096[] = {10240u, 4096u}; static const uint64_t patch_dims[] = {14u, 14u, 2u, 3u, 1024u}; static const uint64_t downsample_dims[] = {2u, 2u, 1024u, 4096u}; memset(w, 0, sizeof(*w)); w->patch_weight = vision_required_offset( m, "model.visual.patch_embed.proj.weight", 5, patch_dims); w->patch_bias = vision_required_offset( m, "model.visual.patch_embed.proj.bias", 1, d1024); w->post_norm = vision_required_offset( m, "model.visual.post_layernorm.weight", 1, d1024); w->downsample_weight = vision_required_offset( m, "model.visual.downsample.weight", 4, downsample_dims); w->downsample_bias = vision_required_offset( m, "model.visual.downsample.bias", 1, d4096); w->merger_proj = vision_required_offset( m, "model.visual.merger.proj.weight", 2, d4096_4096); w->merger_norm = vision_required_offset( m, "model.visual.merger.post_projection_norm.weight", 1, d4096); w->merger_norm_bias = vision_required_offset( m, "model.visual.merger.post_projection_norm.bias", 1, d4096); w->merger_gate = vision_required_offset( m, "model.visual.merger.gate_proj.weight", 2, d4096_10240); w->merger_up = vision_required_offset( m, "model.visual.merger.up_proj.weight", 2, d4096_10240); w->merger_down = vision_required_offset( m, "model.visual.merger.down_proj.weight", 2, d10240_4096); for (uint32_t il = 0; il < DS4_GLM53_VISION_LAYERS; il++) { char name[160]; #define VISION_LAYER_OFFSET(field_, suffix_, rank_, dims_) do { \ int n = snprintf(name, sizeof(name), \ "model.visual.blocks.%u.%s", il, suffix_); \ if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("vision tensor name overflow"); \ w->layer[il].field_ = vision_required_offset(m, name, rank_, dims_); \ } while (0) VISION_LAYER_OFFSET(norm1, "norm1.weight", 1, d1024); VISION_LAYER_OFFSET(qkv_weight, "attn.qkv.weight", 2, d1024_3072); VISION_LAYER_OFFSET(qkv_bias, "attn.qkv.bias", 1, d3072); VISION_LAYER_OFFSET(q_norm, "attn.q_norm.weight", 1, d64); VISION_LAYER_OFFSET(k_norm, "attn.k_norm.weight", 1, d64); VISION_LAYER_OFFSET(attn_proj_weight, "attn.proj.weight", 2, d1024_1024); VISION_LAYER_OFFSET(attn_proj_bias, "attn.proj.bias", 1, d1024); VISION_LAYER_OFFSET(norm2, "norm2.weight", 1, d1024); VISION_LAYER_OFFSET(gate_weight, "mlp.gate_proj.weight", 2, d1024_4096); VISION_LAYER_OFFSET(gate_bias, "mlp.gate_proj.bias", 1, d4096); VISION_LAYER_OFFSET(up_weight, "mlp.up_proj.weight", 2, d1024_4096); VISION_LAYER_OFFSET(up_bias, "mlp.up_proj.bias", 1, d4096); VISION_LAYER_OFFSET(down_weight, "mlp.down_proj.weight", 2, d4096_1024); VISION_LAYER_OFFSET(down_bias, "mlp.down_proj.bias", 1, d1024); #undef VISION_LAYER_OFFSET } } static ds4_tensor *deepseek4_vision_required_tensor( const ds4_model *m, const char *name, uint32_t type, uint32_t ndim, const uint64_t *dims) { ds4_tensor *t = required_tensor(m, name); if (t->type != type || t->ndim != ndim) { fprintf(stderr, "ds4: DeepSeek vision tensor %s has type %s/rank %u, " "expected %s/rank %u\n", name, tensor_type_name(t->type), t->ndim, tensor_type_name(type), ndim); exit(1); } for (uint32_t d = 0; d < ndim; d++) { if (t->dim[d] == dims[d]) continue; fprintf(stderr, "ds4: DeepSeek vision tensor %s has dim[%u]=%" PRIu64 ", expected %" PRIu64 "\n", name, d, t->dim[d], dims[d]); exit(1); } return t; } static uint64_t deepseek4_vision_required_offset( const ds4_model *m, const char *name, uint32_t type, uint32_t ndim, const uint64_t *dims) { return deepseek4_vision_required_tensor( m, name, type, ndim, dims)->abs_offset; } static void deepseek4_vision_weights_bind( ds4_deepseek4_vision_weights *w, const ds4_model *m) { const bool v41 = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41; const char *revision = v41 ? "df42c109f1defefcbfcedbe7d905718a12266e40" : "e46e16bf6035c6f317eb2ac7458eb0362926d402"; const uint32_t projection = v41 ? 5120u : 4096u; ds4_str arch = {0}, source_revision = {0}, variant = {0}; if (!model_get_string(m, "general.architecture", &arch) || !ds4_streq(arch, "deepseek4-vision")) { ds4_die("--vision file is not a DeepSeek V4 vision encoder GGUF"); } if (!model_get_string(m, "general.source.revision", &source_revision) || !ds4_streq(source_revision, revision) || !model_get_string(m, "deepseek4-vision.checkpoint_variant", &variant) || !ds4_streq(variant, v41 ? "v4.1-flash" : "vision-exp")) { ds4_die("DeepSeek vision sidecar does not match the language checkpoint"); } const uint32_t expected_tensors = v41 ? 266u : 316u; if (m->n_tensors != expected_tensors) { fprintf(stderr, "ds4: DeepSeek vision GGUF has %" PRIu64 " tensors, expected %u\n", m->n_tensors, expected_tensors); exit(1); } config_expect_u32("DeepSeek vision block_count", required_u32(m, "deepseek4-vision.block_count"), 32u); config_expect_u32("DeepSeek vision embedding_length", required_u32(m, "deepseek4-vision.embedding_length"), 1024u); config_expect_u32("DeepSeek vision feed_forward_length", required_u32(m, "deepseek4-vision.feed_forward_length"), 2816u); config_expect_u32("DeepSeek vision head_count", required_u32(m, "deepseek4-vision.attention.head_count"), 16u); config_expect_u32("DeepSeek vision projection_length", required_u32(m, "deepseek4-vision.projection_length"), projection); config_expect_u32("DeepSeek vision patch_size", required_u32(m, "deepseek4-vision.patch_size"), 14u); config_expect_u32("DeepSeek vision downsample_ratio", required_u32(m, "deepseek4-vision.downsample_ratio"), 3u); config_expect_u32("DeepSeek vision max_tokens", required_u32(m, "deepseek4-vision.image.max_tokens"), v41 ? 1024u : 384u); config_expect_u32("DeepSeek vision min_pixels", required_u32(m, "deepseek4-vision.image.min_pixels"), v41 ? 295936u : 147456u); config_expect_u32("DeepSeek vision max_width_height_ratio", required_u32(m, "deepseek4-vision.image.max_width_height_ratio"), v41 ? 0u : 8u); if (v41) config_expect_u32("DeepSeek V4.1 image token", required_u32(m, "deepseek4-vision.image.token_id"), 129264u); config_expect_u32("DeepSeek vision language layers", required_u32(m, "deepseek4-vision.language.block_count"), v41 ? 40u : 43u); config_expect_u32("DeepSeek vision language experts", required_u32(m, "deepseek4-vision.language.expert_count"), v41 ? 384u : 256u); config_expect_epsilon( "DeepSeek vision RMS epsilon", required_f32(m, "deepseek4-vision.attention.layer_norm_rms_epsilon"), 1.0e-6f); static const uint64_t d256[] = {256u}; static const uint64_t d1024[] = {1024u}; static const uint64_t d2816_1024[] = {2816u, 1024u}; static const uint64_t d588_1024[] = {588u, 1024u}; static const uint64_t d1024_1024[] = {1024u, 1024u}; static const uint64_t d1024_3072[] = {1024u, 3072u}; static const uint64_t d1024_5632[] = {1024u, 5632u}; static const uint64_t d3072[] = {3072u}; const uint64_t dproj[] = {projection}; const uint64_t dproj_proj[] = {projection, projection}; const uint64_t d9216_proj[] = {9216u, projection}; memset(w, 0, sizeof(*w)); w->projection_dim = projection; #define DEEPSEEK4_VISION_OFFSET(field_, name_, type_, rank_, dims_) \ w->field_ = deepseek4_vision_required_offset( \ m, name_, type_, rank_, dims_) DEEPSEEK4_VISION_OFFSET(patch_weight, "vision.patch_embed.proj.weight", DS4_TENSOR_BF16, 2, d588_1024); DEEPSEEK4_VISION_OFFSET(patch_bias, "vision.patch_embed.proj.bias", DS4_TENSOR_BF16, 1, d1024); DEEPSEEK4_VISION_OFFSET(post_norm, "vision.norm.weight", DS4_TENSOR_BF16, 1, d1024); DEEPSEEK4_VISION_OFFSET(aligner_w1, "aligner.w1.weight", DS4_TENSOR_BF16, 2, d9216_proj); DEEPSEEK4_VISION_OFFSET(aligner_w1_bias, "aligner.w1.bias", DS4_TENSOR_BF16, 1, dproj); DEEPSEEK4_VISION_OFFSET(aligner_w2, "aligner.w2.weight", DS4_TENSOR_BF16, 2, dproj_proj); DEEPSEEK4_VISION_OFFSET(aligner_w2_bias, "aligner.w2.bias", DS4_TENSOR_BF16, 1, dproj); DEEPSEEK4_VISION_OFFSET(image_start, "image_start", DS4_TENSOR_BF16, 1, dproj); if (!v41) DEEPSEEK4_VISION_OFFSET(image_pad, "image_pad", DS4_TENSOR_BF16, 1, dproj); DEEPSEEK4_VISION_OFFSET(image_newline, "image_newline", DS4_TENSOR_BF16, 1, dproj); DEEPSEEK4_VISION_OFFSET(image_end, "image_end", DS4_TENSOR_BF16, 1, dproj); #undef DEEPSEEK4_VISION_OFFSET for (uint32_t il = 0; il < DS4_DEEPSEEK4_VISION_LAYERS; il++) { char name[128]; #define DEEPSEEK4_VISION_LAYER_OFFSET(field_, suffix_, rank_, dims_) do { \ int n = snprintf(name, sizeof(name), \ "vision.blocks.%u.%s", il, suffix_); \ if (n < 0 || (size_t)n >= sizeof(name)) \ ds4_die("DeepSeek vision tensor name overflow"); \ w->layer[il].field_ = deepseek4_vision_required_offset( \ m, name, DS4_TENSOR_BF16, rank_, dims_); \ } while (0) DEEPSEEK4_VISION_LAYER_OFFSET(norm1, "norm1.weight", 1, d1024); DEEPSEEK4_VISION_LAYER_OFFSET(qkv_weight, "attn.wqkv.weight", 2, d1024_3072); DEEPSEEK4_VISION_LAYER_OFFSET(qkv_bias, "attn.wqkv.bias", 1, d3072); DEEPSEEK4_VISION_LAYER_OFFSET(attn_proj_weight, "attn.wo.weight", 2, d1024_1024); DEEPSEEK4_VISION_LAYER_OFFSET(attn_proj_bias, "attn.wo.bias", 1, d1024); DEEPSEEK4_VISION_LAYER_OFFSET(norm2, "norm2.weight", 1, d1024); DEEPSEEK4_VISION_LAYER_OFFSET(mlp_w1, "mlp.w1.weight", 2, d1024_5632); DEEPSEEK4_VISION_LAYER_OFFSET(mlp_w2, "mlp.w2.weight", 2, d2816_1024); #undef DEEPSEEK4_VISION_LAYER_OFFSET } /* V4.1 stores visual routing biases in the language GGUF. */ if (v41) return; for (uint32_t il = 0; il < DS4_DEEPSEEK4_LANGUAGE_LAYERS; il++) { char name[128]; int n = snprintf(name, sizeof(name), "layers.%u.ffn.gate.bias_vl", il); if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("DeepSeek visual router tensor name overflow"); w->visual_router_bias[il] = deepseek4_vision_required_offset( m, name, DS4_TENSOR_F32, 1, d256); } for (uint32_t stage = 0; stage < DS4_DEEPSEEK4_MTP_LAYERS; stage++) { char name[128]; int n = snprintf(name, sizeof(name), "mtp.%u.ffn.gate.bias_vl", stage); if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("DeepSeek MTP visual router tensor name overflow"); w->mtp_visual_router_bias[stage] = deepseek4_vision_required_offset( m, name, DS4_TENSOR_F32, 1, d256); } for (uint32_t il = 0; il < 3u; il++) { char name[128]; int n = snprintf(name, sizeof(name), "layers.%u.ffn.gate.bias", il); if (n < 0 || (size_t)n >= sizeof(name)) ds4_die("DeepSeek hash router bias tensor name overflow"); w->hash_router_bias[il] = deepseek4_vision_required_offset( m, name, DS4_TENSOR_F32, 1, d256); } } /* Qwen3.8 vision: llama.cpp's Qwen3-VL mmproj (clip / qwen3vl_merger). * Matmul weights may be f32, f16 or q8_0; everything else is f32. */ static uint64_t qwen4_vision_dense(const ds4_model *m, const char *name, uint64_t in_dim, uint64_t out_dim, uint32_t *type) { const ds4_tensor *t = required_tensor(m, name); uint64_t in = 1; for (uint32_t d = 0; d + 1 < t->ndim; d++) in *= t->dim[d]; if ((t->type != DS4_TENSOR_F32 && t->type != DS4_TENSOR_F16 && t->type != DS4_TENSOR_Q8_0) || t->ndim < 2 || in != in_dim || t->dim[t->ndim - 1] != out_dim) { fprintf(stderr, "ds4: vision tensor %s has type %s/shape [%" PRIu64 " x %" PRIu64 "], expected " "f32|f16|q8_0 [%" PRIu64 " x %" PRIu64 "]\n", name, tensor_type_name(t->type), in, t->ndim ? t->dim[t->ndim - 1] : 0, in_dim, out_dim); exit(1); } *type = (uint32_t)t->type; return t->abs_offset; } static void qwen4_vision_weights_bind(ds4_qwen4_vision_weights *w, const ds4_model *m) { ds4_str arch = {0}, proj = {0}; if (!model_get_string(m, "general.architecture", &arch) || !ds4_streq(arch, "clip") || !model_get_string(m, "clip.projector_type", &proj) || !ds4_streq(proj, "qwen3vl_merger")) { ds4_die("--vision file is not a Qwen3-VL mmproj GGUF (clip / qwen3vl_merger)"); } memset(w, 0, sizeof(*w)); w->n_embd = required_u32(m, "clip.vision.embedding_length"); w->n_ff = required_u32(m, "clip.vision.feed_forward_length"); w->n_head = required_u32(m, "clip.vision.attention.head_count"); w->n_patch = required_u32(m, "clip.vision.patch_size"); w->n_merge = required_u32(m, "clip.vision.spatial_merge_size"); w->n_out = required_u32(m, "clip.vision.projection_dim"); config_expect_u32("vision block_count", required_u32(m, "clip.vision.block_count"), DS4_QWEN4_VISION_LAYERS); config_expect_u32("vision embedding_length", w->n_embd, 1152u); config_expect_u32("vision head_count", w->n_head, 16u); config_expect_u32("vision patch_size", w->n_patch, 16u); config_expect_u32("vision spatial_merge_size", w->n_merge, 2u); w->eps = 1e-6f; (void)model_get_f32_compat(m, "clip.vision.attention.layer_norm_epsilon", &w->eps); const uint64_t E = w->n_embd, FF = w->n_ff, P = w->n_patch, ME = E * w->n_merge * w->n_merge, O = w->n_out; const uint64_t d_e[1] = { E }, d_ff[1] = { FF }, d_3e[1] = { 3u * E }; const uint64_t d_me[1] = { ME }, d_o[1] = { O }; const ds4_tensor *pos = required_tensor(m, "v.position_embd.weight"); if (pos->type != DS4_TENSOR_F32 || pos->ndim != 2 || pos->dim[0] != E) { ds4_die("vision tensor v.position_embd.weight must be f32 [embd, positions]"); } w->n_pos_side = (uint32_t)sqrt((double)pos->dim[1]); if ((uint64_t)w->n_pos_side * w->n_pos_side != pos->dim[1] || w->n_pos_side < 2) { ds4_die("vision position table is not square"); } w->pos_embd = pos->abs_offset; uint32_t patch_type1 = 0; w->patch_w0 = qwen4_vision_dense(m, "v.patch_embd.weight", 3u * P * P, E, &w->patch_type); w->patch_w1 = qwen4_vision_dense(m, "v.patch_embd.weight.1", 3u * P * P, E, &patch_type1); if (patch_type1 != w->patch_type || w->patch_type == DS4_TENSOR_Q8_0) ds4_die("vision patch weights must both be f32 or f16"); w->patch_b = deepseek4_vision_required_offset(m, "v.patch_embd.bias", DS4_TENSOR_F32, 1, d_e); w->post_ln_w = deepseek4_vision_required_offset(m, "v.post_ln.weight", DS4_TENSOR_F32, 1, d_e); w->post_ln_b = deepseek4_vision_required_offset(m, "v.post_ln.bias", DS4_TENSOR_F32, 1, d_e); w->mm0_w = qwen4_vision_dense(m, "mm.0.weight", ME, ME, &w->mm0_type); w->mm0_b = deepseek4_vision_required_offset(m, "mm.0.bias", DS4_TENSOR_F32, 1, d_me); w->mm2_w = qwen4_vision_dense(m, "mm.2.weight", ME, O, &w->mm2_type); w->mm2_b = deepseek4_vision_required_offset(m, "mm.2.bias", DS4_TENSOR_F32, 1, d_o); for (uint32_t il = 0; il < DS4_QWEN4_VISION_LAYERS; il++) { ds4_qwen4_vision_layer_weights *l = &w->layer[il]; char name[96]; #define VIS_NAME(suffix_) (snprintf(name, sizeof(name), "v.blk.%u.%s", il, suffix_), name) l->ln1_w = deepseek4_vision_required_offset(m, VIS_NAME("ln1.weight"), DS4_TENSOR_F32, 1, d_e); l->ln1_b = deepseek4_vision_required_offset(m, VIS_NAME("ln1.bias"), DS4_TENSOR_F32, 1, d_e); l->qkv_w = qwen4_vision_dense(m, VIS_NAME("attn_qkv.weight"), E, 3u * E, &l->qkv_type); l->qkv_b = deepseek4_vision_required_offset(m, VIS_NAME("attn_qkv.bias"), DS4_TENSOR_F32, 1, d_3e); l->out_w = qwen4_vision_dense(m, VIS_NAME("attn_out.weight"), E, E, &l->out_type); l->out_b = deepseek4_vision_required_offset(m, VIS_NAME("attn_out.bias"), DS4_TENSOR_F32, 1, d_e); l->ln2_w = deepseek4_vision_required_offset(m, VIS_NAME("ln2.weight"), DS4_TENSOR_F32, 1, d_e); l->ln2_b = deepseek4_vision_required_offset(m, VIS_NAME("ln2.bias"), DS4_TENSOR_F32, 1, d_e); l->up_w = qwen4_vision_dense(m, VIS_NAME("ffn_up.weight"), E, FF, &l->up_type); l->up_b = deepseek4_vision_required_offset(m, VIS_NAME("ffn_up.bias"), DS4_TENSOR_F32, 1, d_ff); l->down_w = qwen4_vision_dense(m, VIS_NAME("ffn_down.weight"), FF, E, &l->down_type); l->down_b = deepseek4_vision_required_offset(m, VIS_NAME("ffn_down.bias"), DS4_TENSOR_F32, 1, d_e); #undef VIS_NAME } } #ifdef DS4_HAS_QWEN4_GPU /* Learned 2D position table (side x side) resampled to the patch grid with * bilinear, align_corners interpolation; rows in 2x2 window order. */ static void qwen4_vision_pos_embed(const float *table, uint32_t side, uint32_t n_embd, uint32_t grid_h, uint32_t grid_w, float *out) { size_t row = 0; for (uint32_t by = 0; by < grid_h / 2; by++) { for (uint32_t bx = 0; bx < grid_w / 2; bx++) { for (uint32_t my = 0; my < 2; my++) { for (uint32_t mx = 0; mx < 2; mx++, row++) { const uint32_t y = by * 2 + my, x = bx * 2 + mx; const double sy = grid_h > 1 ? (double)y * (side - 1) / (grid_h - 1) : 0.0; const double sx = grid_w > 1 ? (double)x * (side - 1) / (grid_w - 1) : 0.0; const uint32_t y0 = (uint32_t)sy, x0 = (uint32_t)sx; const uint32_t y1 = y0 + 1 < side ? y0 + 1 : y0, x1 = x0 + 1 < side ? x0 + 1 : x0; const float wy = (float)(sy - y0), wx = (float)(sx - x0); const float *t00 = table + ((size_t)y0 * side + x0) * n_embd; const float *t01 = table + ((size_t)y0 * side + x1) * n_embd; const float *t10 = table + ((size_t)y1 * side + x0) * n_embd; const float *t11 = table + ((size_t)y1 * side + x1) * n_embd; float *dst = out + row * n_embd; for (uint32_t d = 0; d < n_embd; d++) { dst[d] = (1.0f - wy) * ((1.0f - wx) * t00[d] + wx * t01[d]) + wy * ((1.0f - wx) * t10[d] + wx * t11[d]); } } } } } } static uint32_t qwen4_vision_max_tokens(void) { const char *env = getenv("DS4_QWEN4_IMAGE_MAX_TOKENS"); const long v = env ? atol(env) : 0; return v >= 64 && v <= 16383 ? (uint32_t)v : 1024u; /* 4 patches per token, < 65536 patches */ } /* preprocess + encode one image with the bound mmproj; out gets malloc'd * [tokens][n_out] embeddings */ static int qwen4_vision_encode_image(const ds4_model *vm, const ds4_qwen4_vision_weights *w, const ds4_image *image, uint32_t min_tokens, uint32_t max_tokens, float **out, ds4_image_patches *patches, char *error, size_t error_cap) { *out = NULL; if (!ds4_image_preprocess_qwen4(patches, image, min_tokens, max_tokens, error, error_cap)) return 0; const uint64_t n = patches->patch_count; float *pos = malloc((size_t)n * w->n_embd * sizeof(float)); float *emb = malloc((size_t)patches->image_token_count * w->n_out * sizeof(float)); int ok = pos && emb; if (ok) { qwen4_vision_pos_embed((const float *)((const char *)vm->map + w->pos_embd), w->n_pos_side, w->n_embd, patches->grid_height, patches->grid_width, pos); ok = ds4_gpu_qwen4_vision_encode(emb, patches->patches, pos, patches->patch_count, patches->grid_width, vm->map, vm->size, w); } free(pos); if (!ok) { free(emb); if (error && error_cap) snprintf(error, error_cap, "Qwen3.8 vision inference failed"); return 0; } *out = emb; return 1; } #endif #endif static void weights_bind_output( ds4_weights *w, const ds4_model *m, bool required, bool optional) { if (ds4_model_is_qwen4()) { if (required) { w->output_hc_norm = required_tensor(m, "output_hc_norm.weight"); w->output_hc_down = required_tensor(m, "output_hc_down.weight"); w->output_hc_up = required_tensor(m, "output_hc_up.weight"); w->output = required_tensor(m, "output.weight"); } else if (optional) { w->output_hc_norm = model_find_tensor(m, "output_hc_norm.weight"); w->output_hc_down = model_find_tensor(m, "output_hc_down.weight"); w->output_hc_up = model_find_tensor(m, "output_hc_up.weight"); w->output = model_find_tensor(m, "output.weight"); } } else if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { if (required) { w->output_norm = required_tensor(m, "output_norm.weight"); w->output = required_tensor(m, "output.weight"); } else if (optional) { w->output_norm = model_find_tensor(m, "output_norm.weight"); w->output = model_find_tensor(m, "output.weight"); } } else if (required) { w->output_hc_base = required_tensor(m, "output_hc_base.weight"); w->output_hc_fn = required_tensor(m, "output_hc_fn.weight"); w->output_hc_scale = required_tensor(m, "output_hc_scale.weight"); w->output_norm = required_tensor(m, "output_norm.weight"); w->output = required_tensor(m, "output.weight"); } else if (optional) { w->output_hc_base = model_find_tensor(m, "output_hc_base.weight"); w->output_hc_fn = model_find_tensor(m, "output_hc_fn.weight"); w->output_hc_scale = model_find_tensor(m, "output_hc_scale.weight"); w->output_norm = model_find_tensor(m, "output_norm.weight"); w->output = model_find_tensor(m, "output.weight"); } if (optional && weights_have_partial_output_head(w) && !weights_have_output_head(w)) { ds4_die("partial output head in GGUF"); } } static void weights_bind_glm_dsa_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); if (ds4_model_is_glm53() && il + DS4_N_NEXTN_PREDICT < DS4_N_LAYER) { 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->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); } if (ds4_glm53_layer_is_kda(il)) { l->kda_q = required_tensorf(m, "blk.%u.kda_q.weight", il); l->kda_k = required_tensorf(m, "blk.%u.kda_k.weight", il); l->kda_v = required_tensorf(m, "blk.%u.kda_v.weight", il); l->kda_q_conv = required_tensorf(m, "blk.%u.kda_q_conv.weight", il); l->kda_k_conv = required_tensorf(m, "blk.%u.kda_k_conv.weight", il); l->kda_v_conv = required_tensorf(m, "blk.%u.kda_v_conv.weight", il); l->kda_f_a = required_tensorf(m, "blk.%u.kda_f_a.weight", il); l->kda_f_b = required_tensorf(m, "blk.%u.kda_f_b.weight", il); l->kda_dt_bias = required_tensorf(m, "blk.%u.kda_dt_bias.weight", il); l->kda_a_log = required_tensorf(m, "blk.%u.kda_a_log.weight", il); l->kda_beta = required_tensorf(m, "blk.%u.kda_beta.weight", il); l->kda_g_a = required_tensorf(m, "blk.%u.kda_g_a.weight", il); l->kda_g_b = required_tensorf(m, "blk.%u.kda_g_b.weight", il); l->kda_o_norm = required_tensorf(m, "blk.%u.kda_o_norm.weight", il); l->kda_output = required_tensorf(m, "blk.%u.kda_output.weight", il); } else { 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); if (ds4_model_is_glm53()) { l->indexer_compressor_ape = required_tensorf(m, "blk.%u.indexer.pool_ape.weight", il); l->indexer_compressor_gate = required_tensorf(m, "blk.%u.indexer.pool_gate.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_qwen4_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { l->hc_attn_norm = required_tensorf(m, "blk.%u.hc_attn_norm.weight", il); l->hc_attn_down = required_tensorf(m, "blk.%u.hc_attn_down.weight", il); l->hc_attn_up = required_tensorf(m, "blk.%u.hc_attn_up.weight", il); l->hc_attn_inject = required_tensorf(m, "blk.%u.hc_attn_inject.weight", il); l->hc_ffn_norm = required_tensorf(m, "blk.%u.hc_ffn_norm.weight", il); l->hc_ffn_down = required_tensorf(m, "blk.%u.hc_ffn_down.weight", il); l->hc_ffn_up = required_tensorf(m, "blk.%u.hc_ffn_up.weight", il); l->hc_ffn_inject = required_tensorf(m, "blk.%u.hc_ffn_inject.weight", il); if (ds4_qwen4_layer_is_linear(il)) { l->lin_qkv = required_tensorf(m, "blk.%u.attn_qkv.weight", il); l->lin_gate = required_tensorf(m, "blk.%u.attn_gate.weight", il); l->lin_conv = required_tensorf(m, "blk.%u.ssm_conv1d.weight", il); l->lin_dt_bias = required_tensorf(m, "blk.%u.ssm_dt.bias", il); l->lin_a = required_tensorf(m, "blk.%u.ssm_a", il); l->lin_beta = required_tensorf(m, "blk.%u.ssm_beta.weight", il); l->lin_alpha = required_tensorf(m, "blk.%u.ssm_alpha.weight", il); l->lin_norm = required_tensorf(m, "blk.%u.ssm_norm.weight", il); l->lin_out = required_tensorf(m, "blk.%u.ssm_out.weight", il); } else { l->attn_q = required_tensorf(m, "blk.%u.attn_q.weight", il); l->attn_k = required_tensorf(m, "blk.%u.attn_k.weight", il); l->attn_v = required_tensorf(m, "blk.%u.attn_v.weight", il); l->attn_output = required_tensorf(m, "blk.%u.attn_output.weight", il); l->attn_q_norm = required_tensorf(m, "blk.%u.attn_q_norm.weight", il); l->attn_k_norm = required_tensorf(m, "blk.%u.attn_k_norm.weight", il); l->indexer_q_proj = required_tensorf(m, "blk.%u.indexer.q_proj.weight", il); l->indexer_k_proj = required_tensorf(m, "blk.%u.indexer.k_proj.weight", il); l->indexer_q_norm = required_tensorf(m, "blk.%u.indexer.q_norm.weight", il); l->indexer_k_norm = required_tensorf(m, "blk.%u.indexer.k_norm.weight", il); } if (ds4_qwen4_layer_is_ple(il)) { l->ple_key = required_tensorf(m, "blk.%u.ple_key.weight", il); l->ple_value = required_tensorf(m, "blk.%u.ple_value.weight", il); l->ple_norm_key = required_tensorf(m, "blk.%u.ple_norm_key.weight", il); l->ple_norm_query = required_tensorf(m, "blk.%u.ple_norm_query.weight", il); l->ple_norm_conv = required_tensorf(m, "blk.%u.ple_norm_conv.weight", il); l->ple_conv = required_tensorf(m, "blk.%u.ple_conv1d.weight", il); } l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", 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_inp_shexp = required_tensorf(m, "blk.%u.ffn_gate_inp_shexp.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_qwen4_layer_is_nextn(il)) { 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_hc_head_norm = required_tensorf(m, "blk.%u.nextn.hc_head_norm.weight", il); l->nextn_hc_head_down = required_tensorf(m, "blk.%u.nextn.hc_head_down.weight", il); l->nextn_hc_head_up = required_tensorf(m, "blk.%u.nextn.hc_head_up.weight", il); } } static void weights_bind_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { if (ds4_model_is_qwen4()) { weights_bind_qwen4_layer(l, m, il); return; } 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 (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { l->ffn_exp_probs_vl = tensor_by_namef(m, "blk.%u.exp_probs_b_vl.bias", il); tensor_expect_optional(l->ffn_exp_probs_vl, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); if (ds41_kv_source(il)) { l->attn_compressor_kv = required_tensorf(m, "blk.%u.attn_compressor_kv.weight", il); l->attn_compressor_norm = required_tensorf(m, "blk.%u.attn_compressor_norm.weight", il); if (compress_ratio == 2) l->attn_compressor_gate = required_tensorf(m, "blk.%u.attn_compressor_gate.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); } if (ds41_index_source(il)) { 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); } if (ds41_engram_layer(il)) { l->engram_kv = required_tensorf(m, "blk.%u.engram_kv.weight", il); l->engram_q_norm = required_tensorf(m, "blk.%u.engram_q_norm.weight", il); l->engram_k_norm = required_tensorf(m, "blk.%u.engram_k_norm.weight", il); } } else if (compress_ratio != 0) { l->attn_compressor_ape = required_tensorf(m, "blk.%u.attn_compressor_ape.weight", il); l->attn_compressor_kv = required_tensorf(m, "blk.%u.attn_compressor_kv.weight", il); l->attn_compressor_gate = required_tensorf(m, "blk.%u.attn_compressor_gate.weight", il); l->attn_compressor_norm = required_tensorf(m, "blk.%u.attn_compressor_norm.weight", il); } if (compress_ratio == 4) { l->indexer_attn_q_b = required_tensorf(m, "blk.%u.indexer.attn_q_b.weight", il); l->indexer_proj = required_tensorf(m, "blk.%u.indexer.proj.weight", il); l->indexer_compressor_ape = required_tensorf(m, "blk.%u.indexer_compressor_ape.weight", il); l->indexer_compressor_kv = required_tensorf(m, "blk.%u.indexer_compressor_kv.weight", il); l->indexer_compressor_gate = required_tensorf(m, "blk.%u.indexer_compressor_gate.weight", il); l->indexer_compressor_norm = required_tensorf(m, "blk.%u.indexer_compressor_norm.weight", il); } l->hc_ffn_fn = required_tensorf(m, "blk.%u.hc_ffn_fn.weight", il); l->hc_ffn_scale = required_tensorf(m, "blk.%u.hc_ffn_scale.weight", il); l->hc_ffn_base = required_tensorf(m, "blk.%u.hc_ffn_base.weight", il); l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); l->ffn_exp_probs_b = tensor_by_namef(m, "blk.%u.exp_probs_b.bias", il); l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); if (il < DS4_N_HASH_LAYER) { l->ffn_gate_tid2eid = required_tensorf(m, "blk.%u.ffn_gate_tid2eid.weight", il); } } /* Bind tensor names once into the fixed DS4 layer layout. This is the point * where stringly GGUF metadata becomes direct model-specific pointers. */ static void weights_bind( ds4_weights *w, const ds4_model *m, bool load_slice, uint32_t load_layer_start, uint32_t load_layer_end, bool require_output, bool optional_output) { memset(w, 0, sizeof(*w)); uint32_t executable_layers = DS4_N_LAYER; if ((DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || ds4_model_is_qwen4()) && DS4_N_LAYER > DS4_N_NEXTN_PREDICT) { executable_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; } uint32_t start = 0; uint32_t end = executable_layers - 1u; bool require_token_embd = true; if (load_slice) { if (load_layer_start >= executable_layers) ds4_die("invalid model load layer slice"); start = load_layer_start; end = load_layer_end == UINT32_MAX ? executable_layers - 1u : load_layer_end; if (end >= executable_layers || end < start) ds4_die("invalid model load layer slice"); require_token_embd = start == 0; } else { require_output = true; optional_output = false; } if (require_token_embd) { w->token_embd = required_tensor(m, "token_embd.weight"); } else { w->token_embd = model_find_tensor(m, "token_embd.weight"); } if (ds4_model_is_qwen4()) { w->ple_embd = model_find_tensor(m, "per_layer_token_embd.weight"); } weights_bind_output(w, m, require_output, optional_output); for (uint32_t il = start; il <= end; il++) { weights_bind_layer(&w->layer[il], m, il); } /* GLM nextn/MTP block(s): excluded from the executable pass but bound * so the drafter can run them. Only when the full model is loaded. */ if (!load_slice && (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || ds4_model_is_qwen4()) && 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->type == DS4_TENSOR_MXFP4) && 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->kda_q); DS4_INCLUDE_TENSOR(l->kda_k); DS4_INCLUDE_TENSOR(l->kda_v); DS4_INCLUDE_TENSOR(l->kda_q_conv); DS4_INCLUDE_TENSOR(l->kda_k_conv); DS4_INCLUDE_TENSOR(l->kda_v_conv); DS4_INCLUDE_TENSOR(l->kda_f_a); DS4_INCLUDE_TENSOR(l->kda_f_b); DS4_INCLUDE_TENSOR(l->kda_dt_bias); DS4_INCLUDE_TENSOR(l->kda_a_log); DS4_INCLUDE_TENSOR(l->kda_beta); DS4_INCLUDE_TENSOR(l->kda_g_a); DS4_INCLUDE_TENSOR(l->kda_g_b); DS4_INCLUDE_TENSOR(l->kda_o_norm); DS4_INCLUDE_TENSOR(l->kda_output); 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_exp_probs_vl); 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); DS4_INCLUDE_TENSOR(l->engram_kv); DS4_INCLUDE_TENSOR(l->engram_q_norm); DS4_INCLUDE_TENSOR(l->engram_k_norm); DS4_INCLUDE_TENSOR(l->hc_attn_norm); DS4_INCLUDE_TENSOR(l->hc_attn_down); DS4_INCLUDE_TENSOR(l->hc_attn_up); DS4_INCLUDE_TENSOR(l->hc_attn_inject); DS4_INCLUDE_TENSOR(l->hc_ffn_norm); DS4_INCLUDE_TENSOR(l->hc_ffn_down); DS4_INCLUDE_TENSOR(l->hc_ffn_up); DS4_INCLUDE_TENSOR(l->hc_ffn_inject); DS4_INCLUDE_TENSOR(l->attn_q); DS4_INCLUDE_TENSOR(l->attn_k); DS4_INCLUDE_TENSOR(l->attn_v); DS4_INCLUDE_TENSOR(l->attn_q_norm); DS4_INCLUDE_TENSOR(l->attn_k_norm); DS4_INCLUDE_TENSOR(l->indexer_q_proj); DS4_INCLUDE_TENSOR(l->indexer_k_proj); DS4_INCLUDE_TENSOR(l->indexer_q_norm); DS4_INCLUDE_TENSOR(l->lin_qkv); DS4_INCLUDE_TENSOR(l->lin_gate); DS4_INCLUDE_TENSOR(l->lin_conv); DS4_INCLUDE_TENSOR(l->lin_dt_bias); DS4_INCLUDE_TENSOR(l->lin_a); DS4_INCLUDE_TENSOR(l->lin_beta); DS4_INCLUDE_TENSOR(l->lin_alpha); DS4_INCLUDE_TENSOR(l->lin_norm); DS4_INCLUDE_TENSOR(l->lin_out); DS4_INCLUDE_TENSOR(l->ple_key); DS4_INCLUDE_TENSOR(l->ple_value); DS4_INCLUDE_TENSOR(l->ple_norm_key); DS4_INCLUDE_TENSOR(l->ple_norm_query); DS4_INCLUDE_TENSOR(l->ple_norm_conv); DS4_INCLUDE_TENSOR(l->ple_conv); DS4_INCLUDE_TENSOR(l->ffn_gate_inp_shexp); DS4_INCLUDE_TENSOR(l->nextn_hc_head_norm); DS4_INCLUDE_TENSOR(l->nextn_hc_head_down); DS4_INCLUDE_TENSOR(l->nextn_hc_head_up); #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->kda_q); DS4_INCLUDE_TENSOR(l->kda_k); DS4_INCLUDE_TENSOR(l->kda_v); DS4_INCLUDE_TENSOR(l->kda_q_conv); DS4_INCLUDE_TENSOR(l->kda_k_conv); DS4_INCLUDE_TENSOR(l->kda_v_conv); DS4_INCLUDE_TENSOR(l->kda_f_a); DS4_INCLUDE_TENSOR(l->kda_f_b); DS4_INCLUDE_TENSOR(l->kda_dt_bias); DS4_INCLUDE_TENSOR(l->kda_a_log); DS4_INCLUDE_TENSOR(l->kda_beta); DS4_INCLUDE_TENSOR(l->kda_g_a); DS4_INCLUDE_TENSOR(l->kda_g_b); DS4_INCLUDE_TENSOR(l->kda_o_norm); DS4_INCLUDE_TENSOR(l->kda_output); 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_exp_probs_vl); 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); DS4_INCLUDE_TENSOR(l->engram_kv); DS4_INCLUDE_TENSOR(l->engram_q_norm); DS4_INCLUDE_TENSOR(l->engram_k_norm); DS4_INCLUDE_TENSOR(l->hc_attn_norm); DS4_INCLUDE_TENSOR(l->hc_attn_down); DS4_INCLUDE_TENSOR(l->hc_attn_up); DS4_INCLUDE_TENSOR(l->hc_attn_inject); DS4_INCLUDE_TENSOR(l->hc_ffn_norm); DS4_INCLUDE_TENSOR(l->hc_ffn_down); DS4_INCLUDE_TENSOR(l->hc_ffn_up); DS4_INCLUDE_TENSOR(l->hc_ffn_inject); DS4_INCLUDE_TENSOR(l->attn_q); DS4_INCLUDE_TENSOR(l->attn_k); DS4_INCLUDE_TENSOR(l->attn_v); DS4_INCLUDE_TENSOR(l->attn_q_norm); DS4_INCLUDE_TENSOR(l->attn_k_norm); DS4_INCLUDE_TENSOR(l->indexer_q_proj); DS4_INCLUDE_TENSOR(l->indexer_k_proj); DS4_INCLUDE_TENSOR(l->indexer_q_norm); DS4_INCLUDE_TENSOR(l->lin_qkv); DS4_INCLUDE_TENSOR(l->lin_gate); DS4_INCLUDE_TENSOR(l->lin_conv); DS4_INCLUDE_TENSOR(l->lin_dt_bias); DS4_INCLUDE_TENSOR(l->lin_a); DS4_INCLUDE_TENSOR(l->lin_beta); DS4_INCLUDE_TENSOR(l->lin_alpha); DS4_INCLUDE_TENSOR(l->lin_norm); DS4_INCLUDE_TENSOR(l->lin_out); DS4_INCLUDE_TENSOR(l->ple_key); DS4_INCLUDE_TENSOR(l->ple_value); DS4_INCLUDE_TENSOR(l->ple_norm_key); DS4_INCLUDE_TENSOR(l->ple_norm_query); DS4_INCLUDE_TENSOR(l->ple_norm_conv); DS4_INCLUDE_TENSOR(l->ple_conv); DS4_INCLUDE_TENSOR(l->ffn_gate_inp_shexp); DS4_INCLUDE_TENSOR(l->nextn_hc_head_norm); DS4_INCLUDE_TENSOR(l->nextn_hc_head_down); DS4_INCLUDE_TENSOR(l->nextn_hc_head_up); #undef DS4_INCLUDE_TENSOR } static bool glm_stream_resident_decode_layer_supported( const ds4_layer_weights *l, uint32_t il) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || !l || il < DS4_N_LEADING_DENSE || !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps) { return false; } if (l->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && l->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS || l->ffn_down_exps->type == DS4_TENSOR_Q2_K)) { return true; } return l->ffn_gate_exps->type == l->ffn_up_exps->type && l->ffn_gate_exps->type == l->ffn_down_exps->type && (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); } static uint32_t g_glm_streaming_full_resident_start; static uint32_t g_glm_streaming_full_resident_layers; static bool g_glm_streaming_runtime_static_map_enabled; static bool glm_stream_resident_decode_layer_enabled( const ds4_layer_weights *l, uint32_t il) { if (!glm_stream_resident_decode_layer_supported(l, il)) return false; return g_glm_streaming_full_resident_layers != 0 && il >= g_glm_streaming_full_resident_start && il - g_glm_streaming_full_resident_start < g_glm_streaming_full_resident_layers; } static bool glm_stream_expert_cache_addr_layout_supported( const ds4_weights *w, const ds4_layer_weights *l, uint32_t il) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || !w || !l || il >= DS4_N_LAYER || il < DS4_N_LEADING_DENSE || !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > 8 || DS4_N_EXPERT < 128 || glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { return false; } if (!weights_streaming_layer_experts_uniform(w, il)) return false; if (l->ffn_gate_exps->type != l->ffn_up_exps->type) return false; const bool q2_addr = l->ffn_gate_exps->type == DS4_TENSOR_Q2_K && l->ffn_down_exps->type == DS4_TENSOR_Q2_K; const bool q4_addr = l->ffn_gate_exps->type == DS4_TENSOR_Q4_K && l->ffn_down_exps->type == DS4_TENSOR_Q4_K; return q2_addr || q4_addr; } static DS4_MAYBE_UNUSED bool glm_stream_expert_cache_addr_supported( const ds4_weights *w, const ds4_layer_weights *l, uint32_t il) { if (!glm_stream_expert_cache_addr_layout_supported(w, l, il)) { return false; } #ifdef DS4_NO_GPU return false; #else uint64_t gate_expert_bytes = 0; uint64_t down_expert_bytes = 0; if (!streaming_layer_gate_down_expert_bytes(l, &gate_expert_bytes, &down_expert_bytes)) { return false; } return ds4_gpu_stream_expert_cache_budget_for_expert_size( gate_expert_bytes, down_expert_bytes) >= DS4_N_EXPERT_USED; #endif } static bool glm_stream_selected_expert_cache_supported( const ds4_layer_weights *l, uint32_t il) { if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || !l || il < DS4_N_LEADING_DENSE || !l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > 8 || DS4_N_EXPERT < 128 || glm_graph_env_present("DS4_ROCM_MOE_WRITE_CLAMPED_ACT", "DS4_METAL_MOE_WRITE_CLAMPED_ACT") || glm_graph_env_present("DS4_ROCM_DISABLE_ROUTED_PAIR_SWIGLU_FUSION", "DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") || glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_EXPERT_CACHE", "DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE")) { return false; } if (l->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || l->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS) { return false; } if (l->ffn_down_exps->type == DS4_TENSOR_Q2_K) { #ifdef __APPLE__ /* Metal's IQ2/Q2 selected-slot down kernel sums six experts. The * separate IQ2/IQ2 address-table path below supports up to eight. */ if (DS4_N_EXPERT_USED != 6) return false; #endif 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); } static bool glm_stream_decode_expert_cache_ready( const ds4_weights *w, const ds4_layer_weights *l, uint32_t il) { if (glm_stream_expert_cache_addr_supported(w, l, il)) return true; if (!glm_stream_selected_expert_cache_supported(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 } /* * 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 bool weights_model_map_decode_static_supported(const ds4_weights *w) { if (!w) return false; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { return g_glm_streaming_runtime_static_map_enabled; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &w->layer[il]; if (!l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps) continue; if (!weights_streaming_layer_experts_uniform(w, il)) { return false; } } return true; } 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); model_map_span_vec_include_one(spans, w->output_hc_norm); model_map_span_vec_include_one(spans, w->output_hc_down); model_map_span_vec_include_one(spans, w->output_hc_up); /* N-gram rows are read from disk, never included in GPU weight views. */ } 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_static(spans, &w->layer[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_static(spans, &w->layer[il]); } if (include_output) model_map_span_vec_include_output(spans, w); return model_map_span_vec_finish(spans); } /* Runtime static maps are built after the streaming cache and resident-layer * policy is configured. Unlike the guard-only static spans above, they must * include routed tensors for resident layers and cache-incompatible layouts. */ static DS4_MAYBE_UNUSED bool weights_model_map_decode_runtime_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_runtime_slice_spans( const ds4_weights *w, uint32_t layer_start, uint32_t layer_end, bool include_token, bool include_output, ds4_model_map_span_vec *spans) { if (!w || !spans) return false; if (layer_start >= DS4_N_LAYER) return false; if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; if (layer_end >= DS4_N_LAYER || layer_end < layer_start) return false; memset(spans, 0, sizeof(*spans)); if (include_token) model_map_span_vec_include_one(spans, w->token_embd); for (uint32_t il = layer_start; il <= layer_end; il++) { model_map_span_vec_include_layer_decode(spans, w, il); } if (include_output) model_map_span_vec_include_output(spans, w); return model_map_span_vec_finish(spans); } static DS4_MAYBE_UNUSED uint64_t model_map_span_vec_total_bytes( const ds4_model_map_span_vec *spans) { if (!spans) return 0; uint64_t total = 0; for (uint32_t i = 0; i < spans->len; i++) { const uint64_t bytes = spans->v[i].end - spans->v[i].off; if (total > UINT64_MAX - bytes) return UINT64_MAX; total += bytes; } return total; } static DS4_MAYBE_UNUSED bool weights_streaming_non_routed_bytes( const ds4_weights *w, uint64_t *bytes_out) { if (bytes_out) *bytes_out = 0; if (!w || !bytes_out) return false; ds4_model_map_span_vec spans; const bool include_token = weights_layer_has_required(&w->layer[0], 0); if (!weights_model_map_decode_static_spans(w, include_token, weights_have_output_head(w), &spans)) { return false; } *bytes_out = model_map_span_vec_total_bytes(&spans); free(spans.v); return true; } static DS4_MAYBE_UNUSED bool weights_model_map_token_spans( const ds4_weights *w, ds4_model_map_span_vec *spans) { if (!w || !spans) return false; memset(spans, 0, sizeof(*spans)); model_map_span_vec_include_one(spans, w->token_embd); return model_map_span_vec_finish(spans); } static DS4_MAYBE_UNUSED bool weights_model_map_output_spans( const ds4_weights *w, ds4_model_map_span_vec *spans) { if (!w || !spans) return false; memset(spans, 0, sizeof(*spans)); model_map_span_vec_include_output(spans, w); return model_map_span_vec_finish(spans); } static void mtp_weights_bind(ds4_mtp_weights *w, const ds4_model *m) { memset(w, 0, sizeof(*w)); w->hc_head_base = required_tensor(m, "mtp.0.hc_head_base.weight"); w->hc_head_fn = required_tensor(m, "mtp.0.hc_head_fn.weight"); w->hc_head_scale = required_tensor(m, "mtp.0.hc_head_scale.weight"); w->e_proj = required_tensor(m, "mtp.0.e_proj.weight"); w->h_proj = required_tensor(m, "mtp.0.h_proj.weight"); w->enorm = required_tensor(m, "mtp.0.enorm.weight"); w->hnorm = required_tensor(m, "mtp.0.hnorm.weight"); w->norm = required_tensor(m, "mtp.0.norm.weight"); ds4_layer_weights *l = &w->block; l->hc_attn_fn = required_tensor(m, "mtp.0.hc_attn_fn.weight"); l->hc_attn_scale = required_tensor(m, "mtp.0.hc_attn_scale.weight"); l->hc_attn_base = required_tensor(m, "mtp.0.hc_attn_base.weight"); l->attn_norm = required_tensor(m, "mtp.0.attn_norm.weight"); l->attn_q_a = required_tensor(m, "mtp.0.attn_q_a.weight"); l->attn_q_a_norm = required_tensor(m, "mtp.0.attn_q_a_norm.weight"); l->attn_q_b = required_tensor(m, "mtp.0.attn_q_b.weight"); l->attn_kv = required_tensor(m, "mtp.0.attn_kv.weight"); l->attn_kv_a_norm = required_tensor(m, "mtp.0.attn_kv_a_norm.weight"); l->attn_sinks = required_tensor(m, "mtp.0.attn_sinks.weight"); l->attn_output_a = required_tensor(m, "mtp.0.attn_output_a.weight"); l->attn_output_b = required_tensor(m, "mtp.0.attn_output_b.weight"); l->hc_ffn_fn = required_tensor(m, "mtp.0.hc_ffn_fn.weight"); l->hc_ffn_scale = required_tensor(m, "mtp.0.hc_ffn_scale.weight"); l->hc_ffn_base = required_tensor(m, "mtp.0.hc_ffn_base.weight"); l->ffn_norm = required_tensor(m, "mtp.0.ffn_norm.weight"); l->ffn_gate_inp = required_tensor(m, "mtp.0.ffn_gate_inp.weight"); l->ffn_exp_probs_b = required_tensor(m, "mtp.0.exp_probs_b.bias"); l->ffn_gate_exps = required_tensor(m, "mtp.0.ffn_gate_exps.weight"); l->ffn_up_exps = required_tensor(m, "mtp.0.ffn_up_exps.weight"); l->ffn_down_exps = required_tensor(m, "mtp.0.ffn_down_exps.weight"); l->ffn_gate_shexp = required_tensor(m, "mtp.0.ffn_gate_shexp.weight"); l->ffn_up_shexp = required_tensor(m, "mtp.0.ffn_up_shexp.weight"); l->ffn_down_shexp = required_tensor(m, "mtp.0.ffn_down_shexp.weight"); mtp_weights_validate_layout(w); } static ds4_tensor *dspark_bind_tensor( ds4_dspark_weights *dw, const ds4_model *m, uint32_t stage, const char *suffix, bool required) { ds4_tensor *t = tensor_by_mtp_stage_suffix(m, stage, suffix); if (t) { dw->present_tensors++; } else if (required) { dw->missing_tensors++; } return t; } static void dspark_bind_block( ds4_dspark_weights *dw, ds4_layer_weights *l, const ds4_model *m, uint32_t stage) { l->hc_attn_fn = dspark_bind_tensor(dw, m, stage, "hc_attn_fn.weight", true); l->hc_attn_scale = dspark_bind_tensor(dw, m, stage, "hc_attn_scale.weight", true); l->hc_attn_base = dspark_bind_tensor(dw, m, stage, "hc_attn_base.weight", true); l->attn_norm = dspark_bind_tensor(dw, m, stage, "attn_norm.weight", true); l->attn_q_a = dspark_bind_tensor(dw, m, stage, "attn_q_a.weight", true); l->attn_q_a_norm = dspark_bind_tensor(dw, m, stage, "attn_q_a_norm.weight", true); l->attn_q_b = dspark_bind_tensor(dw, m, stage, "attn_q_b.weight", true); l->attn_kv = dspark_bind_tensor(dw, m, stage, "attn_kv.weight", true); l->attn_kv_a_norm = dspark_bind_tensor(dw, m, stage, "attn_kv_a_norm.weight", true); l->attn_sinks = dspark_bind_tensor(dw, m, stage, "attn_sinks.weight", true); l->attn_output_a = dspark_bind_tensor(dw, m, stage, "attn_output_a.weight", true); l->attn_output_b = dspark_bind_tensor(dw, m, stage, "attn_output_b.weight", true); l->hc_ffn_fn = dspark_bind_tensor(dw, m, stage, "hc_ffn_fn.weight", true); l->hc_ffn_scale = dspark_bind_tensor(dw, m, stage, "hc_ffn_scale.weight", true); l->hc_ffn_base = dspark_bind_tensor(dw, m, stage, "hc_ffn_base.weight", true); l->ffn_norm = dspark_bind_tensor(dw, m, stage, "ffn_norm.weight", true); l->ffn_gate_inp = dspark_bind_tensor(dw, m, stage, "ffn_gate_inp.weight", true); l->ffn_exp_probs_b = dspark_bind_tensor(dw, m, stage, "exp_probs_b.bias", true); l->ffn_gate_exps = dspark_bind_tensor(dw, m, stage, "ffn_gate_exps.weight", true); l->ffn_up_exps = dspark_bind_tensor(dw, m, stage, "ffn_up_exps.weight", true); l->ffn_down_exps = dspark_bind_tensor(dw, m, stage, "ffn_down_exps.weight", true); l->ffn_gate_shexp = dspark_bind_tensor(dw, m, stage, "ffn_gate_shexp.weight", true); l->ffn_up_shexp = dspark_bind_tensor(dw, m, stage, "ffn_up_shexp.weight", true); l->ffn_down_shexp = dspark_bind_tensor(dw, m, stage, "ffn_down_shexp.weight", true); } static void dspark_weights_bind_optional( ds4_dspark_weights *dw, const ds4_model *m, const ds4_dspark_summary *summary) { memset(dw, 0, sizeof(*dw)); if (!m || !summary) return; dw->n_stages = summary->stages < DS4_DSPARK_MAX_STAGES ? summary->stages : DS4_DSPARK_MAX_STAGES; dw->block_size = summary->block_size; dw->markov_rank = summary->markov_rank; dw->noise_token_id = summary->noise_token_id; dw->target_layer_count = summary->target_layer_count; dw->has_block_size = summary->has_block_size; dw->has_markov_rank = summary->has_markov_rank; dw->has_noise_token_id = summary->has_noise_token_id; dw->has_target_layers = summary->has_target_layers; memcpy(dw->target_layers, summary->target_layers, (size_t)dw->target_layer_count * sizeof(dw->target_layers[0])); if (summary->stages > DS4_DSPARK_MAX_STAGES) dw->missing_tensors++; for (uint32_t stage = 0; stage < dw->n_stages; stage++) { ds4_dspark_stage_weights *sw = &dw->stage[stage]; dspark_bind_block(dw, &sw->block, m, stage); if (stage == 0) { sw->main_proj = dspark_bind_tensor(dw, m, stage, "main_proj.weight", true); sw->main_norm = dspark_bind_tensor(dw, m, stage, "main_norm.weight", true); } } if (dw->n_stages != 0) { const uint32_t final_stage = dw->n_stages - 1u; ds4_dspark_stage_weights *sw = &dw->stage[final_stage]; sw->norm = dspark_bind_tensor(dw, m, final_stage, "norm.weight", true); sw->hc_head_base = dspark_bind_tensor(dw, m, final_stage, "hc_head_base.weight", true); sw->hc_head_fn = dspark_bind_tensor(dw, m, final_stage, "hc_head_fn.weight", true); sw->hc_head_scale = dspark_bind_tensor(dw, m, final_stage, "hc_head_scale.weight", true); sw->markov_w1 = dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w1.weight", true); sw->markov_w2 = dspark_bind_tensor(dw, m, final_stage, "markov_head.markov_w2.weight", true); sw->confidence_proj = dspark_bind_tensor(dw, m, final_stage, "confidence_head.proj.weight", true); } dspark_weights_validate_layout(dw); } static void weights_free(ds4_weights *w) { memset(w, 0, sizeof(*w)); } /* Load one token embedding row and expand it to float activations. */ static void embed_token_f16(const ds4_model *m, const ds4_weights *w, int token, float *out) { ds4_tensor *te = w->token_embd; if (te->type != DS4_TENSOR_F16 || te->ndim != 2) { ds4_die("expected a 2D F16 token embedding tensor"); } if (token < 0 || (uint64_t)token >= te->dim[1]) { ds4_die("token id is outside the embedding table"); } const uint16_t *base = tensor_data(m, te); const uint64_t stride = te->dim[0]; const uint16_t *row = base + (uint64_t)token * stride; for (uint64_t i = 0; i < stride; i++) { out[i] = f16_to_f32(row[i]); } } static void embed_token_q8_0(const ds4_model *m, const ds4_weights *w, int token, float *out) { ds4_tensor *te = w->token_embd; if (te->type != DS4_TENSOR_Q8_0 || te->ndim != 2) { ds4_die("expected a 2D Q8_0 token embedding tensor"); } if (token < 0 || (uint64_t)token >= te->dim[1]) { ds4_die("token id is outside the embedding table"); } const uint64_t n = te->dim[0]; const uint64_t blocks = (n + 31) / 32; const uint8_t *row = (const uint8_t *)tensor_data(m, te) + (uint64_t)token * blocks * 34; for (uint64_t b = 0; b < blocks; b++) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); const float scale = f16_to_f32(scale_bits); const int8_t *qs = (const int8_t *)(row + b * 34 + 2); const uint64_t i0 = b * 32; const uint64_t bn = n - i0 < 32 ? n - i0 : 32; for (uint64_t i = 0; i < bn; i++) { out[i0 + i] = scale * (float)qs[i]; } } } static void embed_token_any(const ds4_model *m, const ds4_weights *w, int token, float *out) { if (!w->token_embd) ds4_die("token embedding tensor is missing"); switch (w->token_embd->type) { case DS4_TENSOR_F16: embed_token_f16(m, w, token, out); break; case DS4_TENSOR_Q8_0: embed_token_q8_0(m, w, token, out); break; default: ds4_die("unsupported token embedding tensor type"); } } /* RMSNorm without a learned scale, used by hyper-connection control vectors. */ static void rms_norm_no_weight(float *out, const float *x, uint64_t n, float eps) { double ss = 0.0; for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale; } /* Standard DS4 RMSNorm with learned per-channel scale. */ static void rms_norm_weight(float *out, const float *x, const float *weight, uint64_t n, float eps) { double ss = 0.0; for (uint64_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; const float scale = 1.0f / sqrtf((float)(ss / (double)n) + eps); for (uint64_t i = 0; i < n; i++) out[i] = x[i] * scale * weight[i]; } /* Normalize each attention head independently after Q projection. */ static void head_rms_norm_inplace(float *x, uint32_t n_head, uint32_t head_dim, float eps) { for (uint32_t h = 0; h < n_head; h++) { float *head = x + (uint64_t)h * head_dim; double ss = 0.0; for (uint32_t i = 0; i < head_dim; i++) ss += (double)head[i] * head[i]; const float scale = 1.0f / sqrtf((float)(ss / (double)head_dim) + eps); for (uint32_t i = 0; i < head_dim; i++) head[i] *= scale; } } typedef struct { float *out; const uint16_t *data; const float *x; uint64_t in_dim; } matvec_f16_ctx; static inline float dot_f16_row(const uint16_t *row, const float *x, uint64_t n) { #if defined(__ARM_NEON) uint64_t i = 0; float32x4_t acc0 = vdupq_n_f32(0.0f); float32x4_t acc1 = vdupq_n_f32(0.0f); for (; i + 8 <= n; i += 8) { const float16x8_t hv = vreinterpretq_f16_u16(vld1q_u16(row + i)); const float32x4_t h0 = vcvt_f32_f16(vget_low_f16(hv)); const float32x4_t h1 = vcvt_f32_f16(vget_high_f16(hv)); acc0 = vfmaq_f32(acc0, h0, vld1q_f32(x + i)); acc1 = vfmaq_f32(acc1, h1, vld1q_f32(x + i + 4)); } float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); for (; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; return acc; #else float acc = 0.0f; for (uint64_t i = 0; i < n; i++) acc += f16_to_f32(row[i]) * x[i]; return acc; #endif } static void matvec_f16_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_f16_ctx *ctx = vctx; for (uint64_t o = row0; o < row1; o++) { const uint16_t *row = ctx->data + o * ctx->in_dim; ctx->out[o] = dot_f16_row(row, ctx->x, ctx->in_dim); } } /* Dense F16 matvec for small control projections such as HC and router heads. */ static void matvec_f16(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t out_dim = w->dim[1]; matvec_f16_ctx ctx = { .out = out, .data = tensor_data(m, w), .x = x, .in_dim = in_dim, }; const uint64_t ops = in_dim * out_dim; const uint64_t min_rows = ops >= 262144 ? 1 : 512; ds4_parallel_for_min_rows(out_dim, matvec_f16_worker, &ctx, min_rows); } static void matvec_f16_serial(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { if (w->type != 1 || w->ndim != 2) ds4_die("expected a 2D F16 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t out_dim = w->dim[1]; const uint16_t *data = tensor_data(m, w); for (uint64_t o = 0; o < out_dim; o++) { out[o] = dot_f16_row(data + o * in_dim, x, in_dim); } } typedef struct { float *out; const uint8_t *data; const int8_t *xq; const float *xscale; uint64_t in_dim; uint64_t row0; uint64_t blocks; } matvec_q8_0_ctx; typedef struct { float *out0; float *out1; const uint8_t *data0; const uint8_t *data1; const int8_t *xq; const float *xscale; uint64_t in_dim; uint64_t blocks; } matvec_q8_0_pair_ctx; typedef struct { float *out; const uint8_t *data; const int8_t *xq; const float *xscale; uint64_t in_dim; uint64_t blocks; uint64_t rank; } matvec_q8_0_grouped_ctx; typedef struct { float *out; const uint8_t *data; const int8_t *xq; const float *xscale; uint64_t n_tok; uint64_t n_groups; uint64_t group_dim; uint64_t blocks; uint64_t rank; } matmul_q8_0_grouped_batch_ctx; typedef struct { float *out; const uint8_t *data; const int8_t *xq; const float *xscale; uint64_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t blocks; } matmul_q8_0_batch_ctx; typedef struct { float *out0; float *out1; const uint8_t *data0; const uint8_t *data1; const int8_t *xq; const float *xscale; uint64_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t blocks; } matmul_q8_0_pair_batch_ctx; typedef struct { const float *x; int8_t *xq; float *xscale; uint64_t in_dim; uint64_t blocks; } quantize_q8_0_batch_ctx; static inline int32_t dot_i8_32(const int8_t *a, const int8_t *b, uint64_t n) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) if (n == 32) { int32x4_t acc = vdupq_n_s32(0); acc = vdotq_s32(acc, vld1q_s8(a), vld1q_s8(b)); acc = vdotq_s32(acc, vld1q_s8(a + 16), vld1q_s8(b + 16)); return vaddvq_s32(acc); } #endif int32_t sum = 0; for (uint64_t i = 0; i < n; i++) sum += (int32_t)a[i] * (int32_t)b[i]; return sum; } static inline float dot_q8_0_row( const uint8_t *row, const int8_t *xq, const float *xscale, uint64_t in_dim, uint64_t blocks) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) if ((in_dim & 31u) == 0) { float32x4_t accv0 = vdupq_n_f32(0.0f); float32x4_t accv1 = vdupq_n_f32(0.0f); uint64_t b = 0; for (; b + 1 < blocks; b += 2) { uint16_t scale_bits0; uint16_t scale_bits1; memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); const int8_t *xq0 = xq + b * 32; const int8_t *xq1 = xq + (b + 1) * 32; int32x4_t dot0 = vdupq_n_s32(0); dot0 = vdotq_s32(dot0, vld1q_s8(qs0), vld1q_s8(xq0)); dot0 = vdotq_s32(dot0, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + 16)); int32x4_t dot1 = vdupq_n_s32(0); dot1 = vdotq_s32(dot1, vld1q_s8(qs1), vld1q_s8(xq1)); dot1 = vdotq_s32(dot1, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + 16)); accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot0), f16_to_f32(scale_bits0) * xscale[b]); accv1 = vfmaq_n_f32(accv1, vcvtq_f32_s32(dot1), f16_to_f32(scale_bits1) * xscale[b + 1]); } if (b < blocks) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); const int8_t *qs = (const int8_t *)(row + b * 34 + 2); const int8_t *xqb = xq + b * 32; int32x4_t dot = vdupq_n_s32(0); dot = vdotq_s32(dot, vld1q_s8(qs), vld1q_s8(xqb)); dot = vdotq_s32(dot, vld1q_s8(qs + 16), vld1q_s8(xqb + 16)); accv0 = vfmaq_n_f32(accv0, vcvtq_f32_s32(dot), f16_to_f32(scale_bits) * xscale[b]); } return vaddvq_f32(vaddq_f32(accv0, accv1)); } #endif float acc = 0.0f; for (uint64_t b = 0; b < blocks; b++) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); const int8_t *qs = (const int8_t *)(row + b * 34 + 2); const uint64_t i0 = b * 32; const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; acc += f16_to_f32(scale_bits) * xscale[b] * (float)dot_i8_32(qs, xq + i0, n); } return acc; } static inline void dot_q8_0_row_2( const uint8_t *row, const int8_t *xq0, const float *xscale0, const int8_t *xq1, const float *xscale1, uint64_t in_dim, uint64_t blocks, float *out0, float *out1) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) if ((in_dim & 31u) == 0) { float32x4_t acc00 = vdupq_n_f32(0.0f); float32x4_t acc01 = vdupq_n_f32(0.0f); float32x4_t acc10 = vdupq_n_f32(0.0f); float32x4_t acc11 = vdupq_n_f32(0.0f); uint64_t b = 0; for (; b + 1 < blocks; b += 2) { uint16_t scale_bits0; uint16_t scale_bits1; memcpy(&scale_bits0, row + b * 34, sizeof(scale_bits0)); memcpy(&scale_bits1, row + (b + 1) * 34, sizeof(scale_bits1)); const int8_t *qs0 = (const int8_t *)(row + b * 34 + 2); const int8_t *qs1 = (const int8_t *)(row + (b + 1) * 34 + 2); int32x4_t d00 = vdupq_n_s32(0); d00 = vdotq_s32(d00, vld1q_s8(qs0), vld1q_s8(xq0 + b * 32)); d00 = vdotq_s32(d00, vld1q_s8(qs0 + 16), vld1q_s8(xq0 + b * 32 + 16)); int32x4_t d01 = vdupq_n_s32(0); d01 = vdotq_s32(d01, vld1q_s8(qs1), vld1q_s8(xq0 + (b + 1) * 32)); d01 = vdotq_s32(d01, vld1q_s8(qs1 + 16), vld1q_s8(xq0 + (b + 1) * 32 + 16)); int32x4_t d10 = vdupq_n_s32(0); d10 = vdotq_s32(d10, vld1q_s8(qs0), vld1q_s8(xq1 + b * 32)); d10 = vdotq_s32(d10, vld1q_s8(qs0 + 16), vld1q_s8(xq1 + b * 32 + 16)); int32x4_t d11 = vdupq_n_s32(0); d11 = vdotq_s32(d11, vld1q_s8(qs1), vld1q_s8(xq1 + (b + 1) * 32)); d11 = vdotq_s32(d11, vld1q_s8(qs1 + 16), vld1q_s8(xq1 + (b + 1) * 32 + 16)); const float s0 = f16_to_f32(scale_bits0); const float s1 = f16_to_f32(scale_bits1); acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), s0 * xscale0[b]); acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), s1 * xscale0[b + 1]); acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), s0 * xscale1[b]); acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), s1 * xscale1[b + 1]); } if (b < blocks) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); const int8_t *qs = (const int8_t *)(row + b * 34 + 2); int32x4_t d0 = vdupq_n_s32(0); d0 = vdotq_s32(d0, vld1q_s8(qs), vld1q_s8(xq0 + b * 32)); d0 = vdotq_s32(d0, vld1q_s8(qs + 16), vld1q_s8(xq0 + b * 32 + 16)); int32x4_t d1 = vdupq_n_s32(0); d1 = vdotq_s32(d1, vld1q_s8(qs), vld1q_s8(xq1 + b * 32)); d1 = vdotq_s32(d1, vld1q_s8(qs + 16), vld1q_s8(xq1 + b * 32 + 16)); const float s0 = f16_to_f32(scale_bits); acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), s0 * xscale0[b]); acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), s0 * xscale1[b]); } *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); return; } #endif *out0 = dot_q8_0_row(row, xq0, xscale0, in_dim, blocks); *out1 = dot_q8_0_row(row, xq1, xscale1, in_dim, blocks); } static inline DS4_MAYBE_UNUSED void dot_q8_0_row_pair( const uint8_t *row0, const uint8_t *row1, const int8_t *xq, const float *xscale, uint64_t in_dim, uint64_t blocks, float *out0, float *out1) { #if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) if ((in_dim & 31u) == 0) { float32x4_t acc00 = vdupq_n_f32(0.0f); float32x4_t acc01 = vdupq_n_f32(0.0f); float32x4_t acc10 = vdupq_n_f32(0.0f); float32x4_t acc11 = vdupq_n_f32(0.0f); uint64_t b = 0; for (; b + 1 < blocks; b += 2) { uint16_t s00, s01, s10, s11; memcpy(&s00, row0 + b * 34, sizeof(s00)); memcpy(&s01, row0 + (b + 1) * 34, sizeof(s01)); memcpy(&s10, row1 + b * 34, sizeof(s10)); memcpy(&s11, row1 + (b + 1) * 34, sizeof(s11)); const int8_t *xq0 = xq + b * 32; const int8_t *xq1 = xq + (b + 1) * 32; const int8x16_t xv00 = vld1q_s8(xq0); const int8x16_t xv01 = vld1q_s8(xq0 + 16); const int8x16_t xv10 = vld1q_s8(xq1); const int8x16_t xv11 = vld1q_s8(xq1 + 16); const int8_t *q00 = (const int8_t *)(row0 + b * 34 + 2); const int8_t *q01 = (const int8_t *)(row0 + (b + 1) * 34 + 2); const int8_t *q10 = (const int8_t *)(row1 + b * 34 + 2); const int8_t *q11 = (const int8_t *)(row1 + (b + 1) * 34 + 2); int32x4_t d00 = vdupq_n_s32(0); d00 = vdotq_s32(d00, vld1q_s8(q00), xv00); d00 = vdotq_s32(d00, vld1q_s8(q00 + 16), xv01); int32x4_t d01 = vdupq_n_s32(0); d01 = vdotq_s32(d01, vld1q_s8(q01), xv10); d01 = vdotq_s32(d01, vld1q_s8(q01 + 16), xv11); int32x4_t d10 = vdupq_n_s32(0); d10 = vdotq_s32(d10, vld1q_s8(q10), xv00); d10 = vdotq_s32(d10, vld1q_s8(q10 + 16), xv01); int32x4_t d11 = vdupq_n_s32(0); d11 = vdotq_s32(d11, vld1q_s8(q11), xv10); d11 = vdotq_s32(d11, vld1q_s8(q11 + 16), xv11); acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d00), f16_to_f32(s00) * xscale[b]); acc01 = vfmaq_n_f32(acc01, vcvtq_f32_s32(d01), f16_to_f32(s01) * xscale[b + 1]); acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d10), f16_to_f32(s10) * xscale[b]); acc11 = vfmaq_n_f32(acc11, vcvtq_f32_s32(d11), f16_to_f32(s11) * xscale[b + 1]); } if (b < blocks) { uint16_t s0, s1; memcpy(&s0, row0 + b * 34, sizeof(s0)); memcpy(&s1, row1 + b * 34, sizeof(s1)); const int8_t *xqb = xq + b * 32; const int8x16_t xv0 = vld1q_s8(xqb); const int8x16_t xv1 = vld1q_s8(xqb + 16); const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); int32x4_t d0 = vdupq_n_s32(0); d0 = vdotq_s32(d0, vld1q_s8(q0), xv0); d0 = vdotq_s32(d0, vld1q_s8(q0 + 16), xv1); int32x4_t d1 = vdupq_n_s32(0); d1 = vdotq_s32(d1, vld1q_s8(q1), xv0); d1 = vdotq_s32(d1, vld1q_s8(q1 + 16), xv1); acc00 = vfmaq_n_f32(acc00, vcvtq_f32_s32(d0), f16_to_f32(s0) * xscale[b]); acc10 = vfmaq_n_f32(acc10, vcvtq_f32_s32(d1), f16_to_f32(s1) * xscale[b]); } *out0 = vaddvq_f32(vaddq_f32(acc00, acc01)); *out1 = vaddvq_f32(vaddq_f32(acc10, acc11)); return; } #endif float acc0 = 0.0f; float acc1 = 0.0f; for (uint64_t b = 0; b < blocks; b++) { uint16_t s0_bits; uint16_t s1_bits; memcpy(&s0_bits, row0 + b * 34, sizeof(s0_bits)); memcpy(&s1_bits, row1 + b * 34, sizeof(s1_bits)); const int8_t *q0 = (const int8_t *)(row0 + b * 34 + 2); const int8_t *q1 = (const int8_t *)(row1 + b * 34 + 2); const uint64_t i0 = b * 32; const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; acc0 += f16_to_f32(s0_bits) * xscale[b] * (float)dot_i8_32(q0, xq + i0, n); acc1 += f16_to_f32(s1_bits) * xscale[b] * (float)dot_i8_32(q1, xq + i0, n); } *out0 = acc0; *out1 = acc1; } static void quantize_q8_0_activation(const float *x, int8_t *xq, float *scale, uint64_t n) { const uint64_t blocks = (n + 31) / 32; for (uint64_t b = 0; b < blocks; b++) { const uint64_t i0 = b * 32; const uint64_t bn = n - i0 < 32 ? n - i0 : 32; float amax = 0.0f; for (uint64_t i = 0; i < bn; i++) { const float ax = fabsf(x[i0 + i]); if (ax > amax) amax = ax; } const float d = amax / 127.0f; const float id = d != 0.0f ? 1.0f / d : 0.0f; scale[b] = d; for (uint64_t i = 0; i < bn; i++) { int v = (int)lrintf(x[i0 + i] * id); if (v > 127) v = 127; if (v < -128) v = -128; xq[i0 + i] = (int8_t)v; } for (uint64_t i = bn; i < 32 && i0 + i < blocks * 32; i++) { xq[i0 + i] = 0; } } } static void quantize_q8_0_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { quantize_q8_0_batch_ctx *ctx = vctx; for (uint64_t t = t0; t < t1; t++) { quantize_q8_0_activation(ctx->x + t * ctx->in_dim, ctx->xq + t * ctx->blocks * 32, ctx->xscale + t * ctx->blocks, ctx->in_dim); } } static void quantize_q8_0_activation_batch( const float *x, int8_t *xq, float *xscale, uint64_t n_tok, uint64_t in_dim) { quantize_q8_0_batch_ctx ctx = { .x = x, .xq = xq, .xscale = xscale, .in_dim = in_dim, .blocks = (in_dim + 31) / 32, }; ds4_parallel_for(n_tok, quantize_q8_0_batch_worker, &ctx); } static void matvec_q8_0_worker(void *vctx, uint64_t r0, uint64_t r1) { matvec_q8_0_ctx *ctx = vctx; for (uint64_t r = r0; r < r1; r++) { const uint64_t o = ctx->row0 + r; const uint8_t *row = ctx->data + o * ctx->blocks * 34; ctx->out[r] = dot_q8_0_row(row, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks); } } static void matvec_q8_0_pair_worker(void *vctx, uint64_t r0, uint64_t r1) { matvec_q8_0_pair_ctx *ctx = vctx; for (uint64_t r = r0; r < r1; r++) { const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; dot_q8_0_row_pair(row0, row1, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks, ctx->out0 + r, ctx->out1 + r); } } static void matvec_q8_0_grouped_worker(void *vctx, uint64_t r0, uint64_t r1) { matvec_q8_0_grouped_ctx *ctx = vctx; for (uint64_t idx = r0; idx < r1; idx++) { const uint64_t group = idx / ctx->rank; const uint64_t row_in_group = idx - group * ctx->rank; const uint64_t tensor_row = group * ctx->rank + row_in_group; const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; const int8_t *xq = ctx->xq + group * ctx->blocks * 32; const float *xscale = ctx->xscale + group * ctx->blocks; ctx->out[idx] = dot_q8_0_row(row, xq, xscale, ctx->in_dim, ctx->blocks); } } static void matmul_q8_0_grouped_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { matmul_q8_0_grouped_batch_ctx *ctx = vctx; for (uint64_t idx = r0; idx < r1; idx++) { const uint64_t group = idx / ctx->rank; const uint64_t row_in_group = idx - group * ctx->rank; const uint64_t tensor_row = group * ctx->rank + row_in_group; const uint8_t *row = ctx->data + tensor_row * ctx->blocks * 34; uint64_t t = 0; for (; t + 1 < ctx->n_tok; t += 2) { const uint64_t xbase0 = (t * ctx->n_groups + group) * ctx->blocks; const uint64_t xbase1 = ((t + 1) * ctx->n_groups + group) * ctx->blocks; dot_q8_0_row_2(row, ctx->xq + xbase0 * 32, ctx->xscale + xbase0, ctx->xq + xbase1 * 32, ctx->xscale + xbase1, ctx->group_dim, ctx->blocks, ctx->out + t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group, ctx->out + (t + 1) * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group); } for (; t < ctx->n_tok; t++) { const uint64_t xbase = (t * ctx->n_groups + group) * ctx->blocks; ctx->out[t * ctx->n_groups * ctx->rank + group * ctx->rank + row_in_group] = dot_q8_0_row(row, ctx->xq + xbase * 32, ctx->xscale + xbase, ctx->group_dim, ctx->blocks); } } } static void matmul_q8_0_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { matmul_q8_0_batch_ctx *ctx = vctx; for (uint64_t r = r0; r < r1; r++) { const uint8_t *row = ctx->data + r * ctx->blocks * 34; uint64_t t = 0; for (; t + 1 < ctx->n_tok; t += 2) { dot_q8_0_row_2(row, ctx->xq + t * ctx->blocks * 32, ctx->xscale + t * ctx->blocks, ctx->xq + (t + 1) * ctx->blocks * 32, ctx->xscale + (t + 1) * ctx->blocks, ctx->in_dim, ctx->blocks, ctx->out + t * ctx->out_dim + r, ctx->out + (t + 1) * ctx->out_dim + r); } for (; t < ctx->n_tok; t++) { ctx->out[t * ctx->out_dim + r] = dot_q8_0_row(row, ctx->xq + t * ctx->blocks * 32, ctx->xscale + t * ctx->blocks, ctx->in_dim, ctx->blocks); } } } static void matmul_q8_0_pair_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { matmul_q8_0_pair_batch_ctx *ctx = vctx; for (uint64_t r = r0; r < r1; r++) { const uint8_t *row0 = ctx->data0 + r * ctx->blocks * 34; const uint8_t *row1 = ctx->data1 + r * ctx->blocks * 34; uint64_t t = 0; for (; t + 1 < ctx->n_tok; t += 2) { const int8_t *xq0 = ctx->xq + t * ctx->blocks * 32; const float *xscale0 = ctx->xscale + t * ctx->blocks; const int8_t *xq1 = ctx->xq + (t + 1) * ctx->blocks * 32; const float *xscale1 = ctx->xscale + (t + 1) * ctx->blocks; dot_q8_0_row_2(row0, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, ctx->out0 + t * ctx->out_dim + r, ctx->out0 + (t + 1) * ctx->out_dim + r); dot_q8_0_row_2(row1, xq0, xscale0, xq1, xscale1, ctx->in_dim, ctx->blocks, ctx->out1 + t * ctx->out_dim + r, ctx->out1 + (t + 1) * ctx->out_dim + r); } for (; t < ctx->n_tok; t++) { const int8_t *xq = ctx->xq + t * ctx->blocks * 32; const float *xscale = ctx->xscale + t * ctx->blocks; dot_q8_0_row_pair(row0, row1, xq, xscale, ctx->in_dim, ctx->blocks, ctx->out0 + t * ctx->out_dim + r, ctx->out1 + t * ctx->out_dim + r); } } } /* Multiply selected Q8_0 rows by an activation that has already been quantized * once. This avoids repeated activation quantization for paired projections. */ static void matvec_q8_0_rows_prequant( float * out, const ds4_model * m, const ds4_tensor * w, const int8_t * xq, const float * xscale, uint64_t row0, uint64_t n_rows) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t out_dim = w->dim[1]; if (row0 > out_dim || n_rows > out_dim - row0) ds4_die("Q8_0 row range is outside tensor"); const uint64_t ctx_blocks = (in_dim + 31) / 32; matvec_q8_0_ctx ctx = { .out = out, .data = tensor_data(m, w), .xq = xq, .xscale = xscale, .in_dim = in_dim, .row0 = row0, .blocks = ctx_blocks, }; ds4_parallel_for(n_rows, matvec_q8_0_worker, &ctx); } static DS4_MAYBE_UNUSED void matvec_q8_0_prequant( float * out, const ds4_model * m, const ds4_tensor * w, const int8_t * xq, const float * xscale) { matvec_q8_0_rows_prequant(out, m, w, xq, xscale, 0, w->dim[1]); } static void matvec_q8_0_3d_slice_prequant( float * out, const ds4_model * m, const ds4_tensor * w, const int8_t * xq, const float * xscale, uint64_t slice) { if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); if (slice >= w->dim[2]) ds4_die("Q8_0 slice is outside tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t out_dim = w->dim[1]; const uint64_t blocks = (in_dim + 31) / 32; const uint64_t slice_bytes = out_dim * blocks * 34; const uint8_t *data = (const uint8_t *)tensor_data(m, w) + slice * slice_bytes; matvec_q8_0_ctx ctx = { .out = out, .data = data, .xq = xq, .xscale = xscale, .in_dim = in_dim, .row0 = 0, .blocks = blocks, }; ds4_parallel_for(out_dim, matvec_q8_0_worker, &ctx); } static DS4_MAYBE_UNUSED void matvec_q8_0_3d_slice( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint64_t slice) { if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3) ds4_die("expected a 3D Q8_0 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)blocks * 32); float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); quantize_q8_0_activation(x, xq, xscale, in_dim); matvec_q8_0_3d_slice_prequant(out, m, w, xq, xscale, slice); free(xscale); free(xq); } /* Compute two Q8_0 projections from the same input, used by gate/up and * compressor kv/score pairs. */ static void matvec_q8_0_pair_prequant( float * out0, float * out1, const ds4_model * m, const ds4_tensor * w0, const ds4_tensor * w1, const int8_t * xq, const float * xscale) { if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { ds4_die("expected two 2D Q8_0 tensors"); } if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { ds4_die("paired Q8_0 tensors do not have the same shape"); } const uint64_t in_dim = w0->dim[0]; matvec_q8_0_pair_ctx ctx = { .out0 = out0, .out1 = out1, .data0 = tensor_data(m, w0), .data1 = tensor_data(m, w1), .xq = xq, .xscale = xscale, .in_dim = in_dim, .blocks = (in_dim + 31) / 32, }; ds4_parallel_for(w0->dim[1], matvec_q8_0_pair_worker, &ctx); } static void matmul_q8_0_batch_prequant( float * out, const ds4_model * m, const ds4_tensor * w, const int8_t * xq, const float * xscale, uint64_t n_tok) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); matmul_q8_0_batch_ctx ctx = { .out = out, .data = tensor_data(m, w), .xq = xq, .xscale = xscale, .n_tok = n_tok, .in_dim = w->dim[0], .out_dim = w->dim[1], .blocks = (w->dim[0] + 31) / 32, }; ds4_parallel_for(ctx.out_dim, matmul_q8_0_batch_worker, &ctx); } static void matmul_q8_0_pair_batch_prequant( float * out0, float * out1, const ds4_model * m, const ds4_tensor * w0, const ds4_tensor * w1, const int8_t * xq, const float * xscale, uint64_t n_tok) { if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { ds4_die("expected two 2D Q8_0 tensors"); } if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { ds4_die("paired Q8_0 tensors do not have the same shape"); } matmul_q8_0_pair_batch_ctx ctx = { .out0 = out0, .out1 = out1, .data0 = tensor_data(m, w0), .data1 = tensor_data(m, w1), .xq = xq, .xscale = xscale, .n_tok = n_tok, .in_dim = w0->dim[0], .out_dim = w0->dim[1], .blocks = (w0->dim[0] + 31) / 32, }; ds4_parallel_for(ctx.out_dim, matmul_q8_0_pair_batch_worker, &ctx); } /* Batched Q8_0 matmul for prefill: quantize all token activations, then scan * weight rows once per output channel. */ static void matmul_q8_0_batch( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint64_t n_tok) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); matmul_q8_0_batch_prequant(out, m, w, xq, xscale, n_tok); free(xscale); free(xq); } static void matmul_q8_0_pair_batch( float * out0, float * out1, const ds4_model * m, const ds4_tensor * w0, const ds4_tensor * w1, const float * x, uint64_t n_tok) { if (w0->type != 8 || w1->type != 8 || w0->ndim != 2 || w1->ndim != 2) { ds4_die("expected two 2D Q8_0 tensors"); } if (w0->dim[0] != w1->dim[0] || w0->dim[1] != w1->dim[1]) { ds4_die("paired Q8_0 tensors do not have the same shape"); } const uint64_t in_dim = w0->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)n_tok * blocks * 32); float *xscale = xmalloc((size_t)n_tok * blocks * sizeof(xscale[0])); quantize_q8_0_activation_batch(x, xq, xscale, n_tok, in_dim); matmul_q8_0_pair_batch_prequant(out0, out1, m, w0, w1, xq, xscale, n_tok); free(xscale); free(xq); } static void matvec_q8_0_rows( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint64_t row0, uint64_t n_rows) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); const uint64_t in_dim = w->dim[0]; const uint64_t ctx_blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)ctx_blocks * 32); float *xscale = xmalloc((size_t)ctx_blocks * sizeof(xscale[0])); quantize_q8_0_activation(x, xq, xscale, in_dim); matvec_q8_0_rows_prequant(out, m, w, xq, xscale, row0, n_rows); free(xscale); free(xq); } /* Single-token Q8_0 matvec, used heavily in decode. */ static void matvec_q8_0(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { matvec_q8_0_rows(out, m, w, x, 0, w->dim[1]); } static inline float dot_q8_0_row_f32_ref( const uint8_t *row, const float *x, uint64_t in_dim, uint64_t blocks) { float acc = 0.0f; for (uint64_t b = 0; b < blocks; b++) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34, sizeof(scale_bits)); const int8_t *qs = (const int8_t *)(row + b * 34 + 2); const float d = f16_to_f32(scale_bits); const uint64_t i0 = b * 32; const uint64_t n = in_dim - i0 < 32 ? in_dim - i0 : 32; for (uint64_t i = 0; i < n; i++) { acc += d * (float)qs[i] * x[i0 + i]; } } return acc; } typedef struct { float *out; const uint8_t *data; const float *x; uint64_t in_dim; uint64_t blocks; } matvec_q8_0_f32_ref_ctx; static void matvec_q8_0_f32_ref_worker(void *vctx, uint64_t r0, uint64_t r1) { matvec_q8_0_f32_ref_ctx *ctx = vctx; const uint64_t row_bytes = ctx->blocks * 34; for (uint64_t r = r0; r < r1; r++) { ctx->out[r] = dot_q8_0_row_f32_ref(ctx->data + r * row_bytes, ctx->x, ctx->in_dim, ctx->blocks); } } static void matvec_q8_0_f32_ref( float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { if (w->type != DS4_TENSOR_Q8_0 || w->ndim < 2 || w->dim[0] == 0) { ds4_die("expected a Q8_0 tensor with matrix rows"); } matvec_q8_0_f32_ref_ctx ctx = { .out = out, .data = tensor_data(m, w), .x = x, .in_dim = w->dim[0], .blocks = (w->dim[0] + 31) / 32, }; ds4_parallel_for(w->elements / w->dim[0], matvec_q8_0_f32_ref_worker, &ctx); } static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x); /* Decode scratch owns this temporary activation quantization so generation * can assert that the hot path performs no malloc. */ static void cpu_decode_quantize_q8_0( ds4_cpu_decode_scratch * scratch, const float * x, uint64_t in_dim) { if (in_dim > scratch->q8_cap) ds4_die("CPU decode Q8_0 scratch buffer is too small"); quantize_q8_0_activation(x, scratch->q8_xq, scratch->q8_xscale, in_dim); } static void matvec_q8_0_decode_scratch( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, ds4_cpu_decode_scratch * scratch) { cpu_decode_quantize_q8_0(scratch, x, w->dim[0]); matvec_q8_0_prequant(out, m, w, scratch->q8_xq, scratch->q8_xscale); } static void matvec_q8_0_pair_decode_scratch( float * out0, float * out1, const ds4_model * m, const ds4_tensor * w0, const ds4_tensor * w1, const float * x, ds4_cpu_decode_scratch * scratch) { cpu_decode_quantize_q8_0(scratch, x, w0->dim[0]); matvec_q8_0_pair_prequant(out0, out1, m, w0, w1, scratch->q8_xq, scratch->q8_xscale); } static void matvec_any_decode_scratch( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, ds4_cpu_decode_scratch * scratch) { if (w->type == 8) { matvec_q8_0_decode_scratch(out, m, w, x, scratch); } else { matvec_any(out, m, w, x); } } static void matvec_q8_0_grouped_rows( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint32_t n_groups, uint64_t group_dim, uint64_t rank) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { ds4_die("grouped Q8_0 tensor has an unexpected layout"); } const uint64_t blocks = (group_dim + 31) / 32; int8_t *xq = xmalloc((size_t)n_groups * blocks * 32); float *xscale = xmalloc((size_t)n_groups * blocks * sizeof(xscale[0])); for (uint32_t g = 0; g < n_groups; g++) { quantize_q8_0_activation(x + (uint64_t)g * group_dim, xq + (uint64_t)g * blocks * 32, xscale + (uint64_t)g * blocks, group_dim); } matvec_q8_0_grouped_ctx ctx = { .out = out, .data = tensor_data(m, w), .xq = xq, .xscale = xscale, .in_dim = group_dim, .blocks = blocks, .rank = rank, }; ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); free(xscale); free(xq); } static void matvec_q8_0_grouped_rows_decode_scratch( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint32_t n_groups, uint64_t group_dim, uint64_t rank, ds4_cpu_decode_scratch * scratch) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { ds4_die("grouped Q8_0 tensor has an unexpected layout"); } if ((uint64_t)n_groups * group_dim > scratch->q8_cap) { ds4_die("CPU decode grouped Q8_0 scratch buffer is too small"); } const uint64_t blocks = (group_dim + 31) / 32; for (uint32_t g = 0; g < n_groups; g++) { quantize_q8_0_activation(x + (uint64_t)g * group_dim, scratch->q8_xq + (uint64_t)g * blocks * 32, scratch->q8_xscale + (uint64_t)g * blocks, group_dim); } matvec_q8_0_grouped_ctx ctx = { .out = out, .data = tensor_data(m, w), .xq = scratch->q8_xq, .xscale = scratch->q8_xscale, .in_dim = group_dim, .blocks = blocks, .rank = rank, }; ds4_parallel_for((uint64_t)n_groups * rank, matvec_q8_0_grouped_worker, &ctx); } static void matmul_q8_0_grouped_batch( float * out, const ds4_model * m, const ds4_tensor * w, const float * x, uint64_t n_tok, uint32_t n_groups, uint64_t group_dim, uint64_t rank) { if (w->type != 8 || w->ndim != 2) ds4_die("expected a 2D Q8_0 tensor"); if (w->dim[0] != group_dim || w->dim[1] < (uint64_t)n_groups * rank) { ds4_die("grouped Q8_0 tensor has an unexpected layout"); } const uint64_t blocks = (group_dim + 31) / 32; int8_t *xq = xmalloc((size_t)n_tok * n_groups * blocks * 32); float *xscale = xmalloc((size_t)n_tok * n_groups * blocks * sizeof(xscale[0])); for (uint64_t t = 0; t < n_tok; t++) { for (uint32_t g = 0; g < n_groups; g++) { const uint64_t xbase = (t * n_groups + g) * blocks; quantize_q8_0_activation(x + t * n_groups * group_dim + (uint64_t)g * group_dim, xq + xbase * 32, xscale + xbase, group_dim); } } matmul_q8_0_grouped_batch_ctx ctx = { .out = out, .data = tensor_data(m, w), .xq = xq, .xscale = xscale, .n_tok = n_tok, .n_groups = n_groups, .group_dim = group_dim, .blocks = blocks, .rank = rank, }; ds4_parallel_for((uint64_t)n_groups * rank, matmul_q8_0_grouped_batch_worker, &ctx); free(xscale); free(xq); } typedef struct { float *out; const float *data; const float *x; uint64_t in_dim; } matvec_f32_ctx; static void matvec_f32_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_f32_ctx *ctx = vctx; for (uint64_t o = row0; o < row1; o++) { double acc = 0.0; const float *row = ctx->data + o * ctx->in_dim; for (uint64_t i = 0; i < ctx->in_dim; i++) { acc += (double)row[i] * ctx->x[i]; } ctx->out[o] = (float)acc; } } static void matvec_f32(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { if (w->type != 0 || w->ndim != 2) ds4_die("expected a 2D F32 tensor"); matvec_f32_ctx ctx = { .out = out, .data = tensor_data(m, w), .x = x, .in_dim = w->dim[0], }; ds4_parallel_for(w->dim[1], matvec_f32_worker, &ctx); } /* Dispatch for dense F32/F16/Q8_0 tensors used by auxiliary projections. */ static void matvec_any(float *out, const ds4_model *m, const ds4_tensor *w, const float *x) { switch (w->type) { case 0: matvec_f32(out, m, w, x); break; case 1: matvec_f16(out, m, w, x); break; case 8: matvec_q8_0(out, m, w, x); break; default: ds4_die("unsupported tensor type for dense matvec"); } } static float tensor_1d_value(const ds4_model *m, const ds4_tensor *t, uint64_t i) { if (i >= t->elements) ds4_die("tensor scalar index is out of bounds"); if (t->type == 0) { const float *p = tensor_data(m, t); return p[i]; } if (t->type == 1) { const uint16_t *p = tensor_data(m, t); return f16_to_f32(p[i]); } ds4_die("unsupported tensor scalar type"); return 0.0f; } static float tensor_2d_value(const ds4_model *m, const ds4_tensor *t, uint64_t x, uint64_t y) { if (t->ndim != 2 || x >= t->dim[0] || y >= t->dim[1]) { ds4_die("tensor 2D index is out of bounds"); } return tensor_1d_value(m, t, y * t->dim[0] + x); } /* Locate one expert's 2D matrix inside a 3D GGUF expert tensor. */ static const uint8_t *tensor_expert_bytes( const ds4_model *m, const ds4_tensor *w, uint32_t expert, uint64_t *in_dim, uint64_t *out_dim, uint64_t *row_bytes) { if (w->ndim != 3) ds4_die("expected a 3D expert tensor"); if (expert >= w->dim[2]) ds4_die("expert id is outside expert tensor"); *in_dim = w->dim[0]; *out_dim = w->dim[1]; const gguf_type_info *info = tensor_type(w->type); if (!info || info->block_elems == 0) ds4_die("unsupported expert tensor type"); const uint64_t blocks = (*in_dim + info->block_elems - 1) / info->block_elems; *row_bytes = blocks * info->block_bytes; const uint64_t expert_bytes = *out_dim * *row_bytes; return (const uint8_t *)tensor_data(m, w) + (uint64_t)expert * expert_bytes; } typedef struct { float *out0; float *out1; const uint8_t *base0; const uint8_t *base1; const block_q8_K *xq; uint64_t in_dim; uint64_t row_bytes0; uint64_t row_bytes1; } matvec_iq2_xxs_pair_ctx; static void matvec_iq2_xxs_pair_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_iq2_xxs_pair_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { const block_iq2_xxs *br0 = (const block_iq2_xxs *)(ctx->base0 + row * ctx->row_bytes0); const block_iq2_xxs *br1 = (const block_iq2_xxs *)(ctx->base1 + row * ctx->row_bytes1); ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &ctx->out0[row], &ctx->out1[row], br0, br1, ctx->xq); } } /* Project one routed expert's gate and up matrices. Both are IQ2_XXS and * share the same Q8_K activation. */ static void matvec_iq2_xxs_expert_pair_prequant( float *out0, float *out1, const ds4_model *m, const ds4_tensor *w0, const ds4_tensor *w1, const block_q8_K *xq, uint32_t expert) { if (w0->type != 16 || w1->type != 16) ds4_die("expected IQ2_XXS expert tensors"); uint64_t in_dim0, out_dim0, row_bytes0; uint64_t in_dim1, out_dim1, row_bytes1; const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &row_bytes0); const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &row_bytes1); if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired IQ2_XXS expert tensors do not match"); if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); matvec_iq2_xxs_pair_ctx ctx = { .out0 = out0, .out1 = out1, .base0 = base0, .base1 = base1, .xq = xq, .in_dim = in_dim0, .row_bytes0 = row_bytes0, .row_bytes1 = row_bytes1, }; ds4_parallel_for(out_dim0, matvec_iq2_xxs_pair_worker, &ctx); } static float silu(float x); typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_iq2_xxs_mid_ctx; typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const int8_t *xq; const float *xscale; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t blocks; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q8_0_mid_ctx; typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q8_k_mid_ctx; static void matvec_iq2_xxs_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_iq2_xxs_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } static void matvec_q8_0_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_0_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; const uint8_t *gate_row = ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]; const uint8_t *up_row = ctx->up_base[slot] + row * ctx->up_row_bytes[slot]; dot_q8_0_row_pair(gate_row, up_row, ctx->xq, ctx->xscale, ctx->in_dim, ctx->blocks, &gate, &up); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } static void matvec_q8_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_k_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, ctx->xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } /* Build all selected expert hidden vectors: IQ2_XXS gate/up, clamp, SwiGLU, * and router weight. The down projection runs later on the quantized mids. */ static void matvec_iq2_xxs_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type != 16 || up_w->type != 16) ds4_die("expected IQ2_XXS expert tensors"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_iq2_xxs_mid_ctx ctx = { .mid = mid, .xq = xq, .clamp = clamp, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired IQ2_XXS expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("IQ2_XXS expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_iq2_xxs_mid_worker, &ctx); } static DS4_MAYBE_UNUSED void matvec_q8_0_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const int8_t *xq, const float *xscale, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type != DS4_TENSOR_Q8_0 || up_w->type != DS4_TENSOR_Q8_0) { ds4_die("expected Q8_0 expert tensors"); } if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q8_0_mid_ctx ctx = { .mid = mid, .xq = xq, .xscale = xscale, .clamp = clamp, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired Q8_0 expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("Q8_0 expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ctx.blocks = in_dim0 / 32u; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_0_mid_worker, &ctx); } static DS4_MAYBE_UNUSED void matvec_q8_k_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type != DS4_TENSOR_Q8_K || up_w->type != DS4_TENSOR_Q8_K) { ds4_die("expected Q8_K expert tensors"); } if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q8_k_mid_ctx ctx = { .mid = mid, .xq = xq, .clamp = clamp, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired Q8_K expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("Q8_K expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q8_k_mid_worker, &ctx); } typedef struct { float *out; const uint8_t *base; const block_q8_K *xq; uint64_t in_dim; uint64_t row_bytes; } matvec_q2_k_ctx; static void matvec_q2_k_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q2_k_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { const block_q2_K *br = (const block_q2_K *)(ctx->base + row * ctx->row_bytes); ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &ctx->out[row], br, ctx->xq); } } /* Single expert Q2_K down projection, kept mostly for tracing and diagnostics. */ static void matvec_q2_k_expert( float *out, const ds4_model *m, const ds4_tensor *w, const float *x, uint32_t expert) { if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); uint64_t in_dim, out_dim, row_bytes; const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); if (in_dim % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); matvec_q2_k_ctx ctx = { .out = out, .base = base, .xq = xq, .in_dim = in_dim, .row_bytes = row_bytes, }; ds4_parallel_for(out_dim, matvec_q2_k_worker, &ctx); free(xq); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q2_k_accum_ctx; static void matvec_q2_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q2_k_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { float v = 0.0f; const block_q2_K *br = (const block_q2_K *)(ctx->base[i] + row * ctx->row_bytes[i]); ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); acc += v; } ctx->out[row] = acc; } } /* Accumulate all selected experts' Q2_K down projections directly into the * 4096-wide MoE output. */ static void matvec_q2_k_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type != 10) ds4_die("expected a Q2_K expert tensor"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; const uint8_t *base[DS4_MAX_EXPERT_USED]; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("Q2_K expert tensors do not share a layout"); } } if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); const uint64_t n_blocks = in_dim0 / QK_K; matvec_q2_k_accum_ctx ctx = { .out = out, .in_dim = in_dim0, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { ctx.base[i] = base[i]; ctx.row_bytes[i] = row_bytes[i]; ctx.xq[i] = xq + (uint64_t)i * n_blocks; } ds4_parallel_for(out_dim0, matvec_q2_k_accum_worker, &ctx); } typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q2_k_mid_ctx; static void matvec_q2_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q2_k_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } static void matvec_q2_k_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type != DS4_TENSOR_Q2_K || up_w->type != DS4_TENSOR_Q2_K) { ds4_die("expected Q2_K expert tensors"); } if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q2_k_mid_ctx ctx = { .mid = mid, .xq = xq, .clamp = clamp, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired Q2_K expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("Q2_K expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if (in_dim0 % QK_K != 0) ds4_die("Q2_K expert row is not QK_K aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q2_k_mid_worker, &ctx); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const int8_t *xq[DS4_MAX_EXPERT_USED]; const float *xscale[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; uint64_t blocks; int n_expert; } matvec_q8_0_accum_ctx; static void matvec_q8_0_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_0_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { const uint8_t *br = ctx->base[i] + row * ctx->row_bytes[i]; acc += dot_q8_0_row(br, ctx->xq[i], ctx->xscale[i], ctx->in_dim, ctx->blocks); } ctx->out[row] = acc; } } static void matvec_q8_0_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const int8_t *xq, const float *xscale, const int *selected, int n_expert) { if (w->type != DS4_TENSOR_Q8_0) ds4_die("expected a Q8_0 expert tensor"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q8_0_accum_ctx ctx = { .out = out, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &ctx.row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("Q8_0 expert tensors do not share a layout"); } } if ((in_dim0 % 32u) != 0) ds4_die("Q8_0 expert row is not QK8_0 aligned"); const uint64_t blocks0 = in_dim0 / 32u; ctx.in_dim = in_dim0; ctx.blocks = blocks0; for (int i = 0; i < n_expert; i++) { ctx.xq[i] = xq + (uint64_t)i * blocks0 * 32u; ctx.xscale[i] = xscale + (uint64_t)i * blocks0; } ds4_parallel_for(out_dim0, matvec_q8_0_accum_worker, &ctx); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q8_k_accum_ctx; static void matvec_q8_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_k_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { float v = 0.0f; const block_q8_K *br = (const block_q8_K *)(ctx->base[i] + row * ctx->row_bytes[i]); ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); acc += v; } ctx->out[row] = acc; } } static void matvec_q8_k_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type != DS4_TENSOR_Q8_K) ds4_die("expected a Q8_K expert tensor"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q8_k_accum_ctx ctx = { .out = out, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; ctx.base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &ctx.row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("Q8_K expert tensors do not share a layout"); } } if (in_dim0 % QK_K != 0) ds4_die("Q8_K expert row is not QK_K aligned"); const uint64_t n_blocks = in_dim0 / QK_K; ctx.in_dim = in_dim0; for (int i = 0; i < n_expert; i++) { ctx.xq[i] = xq + (uint64_t)i * n_blocks; } ds4_parallel_for(out_dim0, matvec_q8_k_accum_worker, &ctx); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_iq2_xxs_accum_ctx; static void matvec_iq2_xxs_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_iq2_xxs_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { float v = 0.0f; const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[i] + row * ctx->row_bytes[i]); ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); acc += v; } ctx->out[row] = acc; } } static void matvec_iq2_xxs_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type != DS4_TENSOR_IQ2_XXS) ds4_die("expected an IQ2_XXS expert tensor"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; const uint8_t *base[DS4_MAX_EXPERT_USED]; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("IQ2_XXS expert tensors do not share a layout"); } } if (in_dim0 % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); const uint64_t n_blocks = in_dim0 / QK_K; matvec_iq2_xxs_accum_ctx ctx = { .out = out, .in_dim = in_dim0, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { ctx.base[i] = base[i]; ctx.row_bytes[i] = row_bytes[i]; ctx.xq[i] = xq + (uint64_t)i * n_blocks; } ds4_parallel_for(out_dim0, matvec_iq2_xxs_accum_worker, &ctx); } typedef struct { uint32_t token; uint32_t slot; } ds4_expert_pair; typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT]; const uint8_t *up_base[DS4_MAX_EXPERT]; const block_q8_K *xq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; const float *pair_weight; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT]; uint64_t up_row_bytes[DS4_MAX_EXPERT]; uint64_t xq_blocks; } matvec_q2_k_batch_mid_ctx; static void matvec_q2_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_q2_k_batch_mid_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q2_K *gate_row = (const block_q2_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); const block_q2_K *up_row = (const block_q2_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; float gate = 0.0f; float up = 0.0f; ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &up, up_row, xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; } } } typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT]; const uint8_t *up_base[DS4_MAX_EXPERT]; const block_q8_K *xq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; const float *pair_weight; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT]; uint64_t up_row_bytes[DS4_MAX_EXPERT]; uint64_t xq_blocks; } matvec_iq2_xxs_batch_mid_ctx; typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT]; const uint8_t *up_base[DS4_MAX_EXPERT]; const block_q8_K *xq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; const float *pair_weight; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t xq_blocks; uint64_t gate_row_bytes[DS4_MAX_EXPERT]; uint64_t up_row_bytes[DS4_MAX_EXPERT]; } matvec_q8_k_batch_mid_ctx; static void matvec_iq2_xxs_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_iq2_xxs_batch_mid_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_iq2_xxs *gate_row = (const block_iq2_xxs *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); const block_iq2_xxs *up_row = (const block_iq2_xxs *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; float gate = 0.0f; float up = 0.0f; ds4_vec_dot_iq2_xxs_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; } } } static void matvec_q8_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_q8_k_batch_mid_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q8_K *gate_row = (const block_q8_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); const block_q8_K *up_row = (const block_q8_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; float gate = 0.0f; float up = 0.0f; ds4_vec_dot_q8_K_pair_q8_K((int)ctx->in_dim, &gate, &up, gate_row, up_row, xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; } } } typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT]; const uint8_t *up_base[DS4_MAX_EXPERT]; const int8_t *xq; const float *xscale; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; const float *pair_weight; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t blocks; uint64_t gate_row_bytes[DS4_MAX_EXPERT]; uint64_t up_row_bytes[DS4_MAX_EXPERT]; } matvec_q8_0_batch_mid_ctx; static void matvec_q8_0_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_q8_0_batch_mid_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const uint8_t *gate_row = ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]; const uint8_t *up_row = ctx->up_base[expert] + row * ctx->up_row_bytes[expert]; for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; float gate = 0.0f; float up = 0.0f; dot_q8_0_row_pair(gate_row, up_row, ctx->xq + (uint64_t)pair.token * ctx->blocks * 32u, ctx->xscale + (uint64_t)pair.token * ctx->blocks, ctx->in_dim, ctx->blocks, &gate, &up); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; } } } typedef struct { const float *mid; block_q8_K *midq; uint64_t down_in_dim; uint64_t down_blocks; } quantize_mid_pairs_ctx; static void quantize_mid_pairs_worker(void *vctx, uint64_t p0, uint64_t p1) { quantize_mid_pairs_ctx *ctx = vctx; for (uint64_t p = p0; p < p1; p++) { ds4_quantize_row_q8_K(ctx->mid + p * ctx->down_in_dim, ctx->midq + p * ctx->down_blocks, (int64_t)ctx->down_in_dim); } } typedef struct { float *down_pair; const uint8_t *base[DS4_MAX_EXPERT]; const block_q8_K *midq; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t midq_blocks; } matvec_q2_k_batch_down_ctx; static DS4_MAYBE_UNUSED void matvec_q2_k_batch_down_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_q2_k_batch_down_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, ctx->down_pair + (uint64_t)pair_id * ctx->out_dim + row, br, xq); } } } typedef struct { float *moe; const uint8_t *base[DS4_MAX_EXPERT]; const block_q8_K *midq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint32_t n_active; uint32_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t midq_blocks; } matvec_q2_k_batch_accum_rows_ctx; static void matvec_q2_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q2_k_batch_accum_rows_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { for (uint32_t t = 0; t < ctx->n_tok; t++) { ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; } for (uint32_t ai = 0; ai < ctx->n_active; ai++) { const uint32_t expert = ctx->active_expert[ai]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q2_K *br = (const block_q2_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; float v = 0.0f; ds4_vec_dot_q2_K_q8_K((int)ctx->in_dim, &v, br, xq); ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; } } } } /* ========================================================================= * Q4_K routed expert matrix-vector products. * ========================================================================= */ typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q4_k_mid_ctx; static void matvec_q4_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q4_k_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]); ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq); const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]); ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } static void matvec_q4_k_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type != DS4_TENSOR_Q4_K || up_w->type != DS4_TENSOR_Q4_K) ds4_die("expected Q4_K expert tensors"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q4_k_mid_ctx ctx = { .mid = mid, .xq = xq, .clamp = clamp, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired Q4_K expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("Q4_K expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q4_k_mid_worker, &ctx); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; int n_expert; } matvec_q4_k_accum_ctx; static void matvec_q4_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q4_k_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { float v = 0.0f; const block_q4_K *br = (const block_q4_K *)(ctx->base[i] + row * ctx->row_bytes[i]); ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]); acc += v; } ctx->out[row] = acc; } } static void matvec_q4_k_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type != DS4_TENSOR_Q4_K) ds4_die("expected a Q4_K expert tensor"); if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; const uint8_t *base[DS4_MAX_EXPERT_USED]; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("Q4_K expert tensors do not share a layout"); } } if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); const uint64_t n_blocks = in_dim0 / QK_K; matvec_q4_k_accum_ctx ctx = { .out = out, .in_dim = in_dim0, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { ctx.base[i] = base[i]; ctx.row_bytes[i] = row_bytes[i]; ctx.xq[i] = xq + (uint64_t)i * n_blocks; } ds4_parallel_for(out_dim0, matvec_q4_k_accum_worker, &ctx); } static inline void ds4_vec_dot_q5_q6_K_q8_K( uint32_t type, int n, float *s, const uint8_t *x, const block_q8_K *y) { if (type == DS4_TENSOR_Q5_K) { ds4_vec_dot_q5_K_q8_K(n, s, (const block_q5_K *)x, y); } else if (type == DS4_TENSOR_Q6_K) { ds4_vec_dot_q6_K_q8_K(n, s, (const block_q6_K *)x, y); } else { ds4_die("expected a Q5_K or Q6_K tensor"); } } typedef struct { float *out; const uint8_t *base; const block_q8_K *xq; uint64_t in_dim; uint64_t row_bytes; uint32_t type; } matvec_q5_q6_k_ctx; static void matvec_q5_q6_k_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q5_q6_k_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &ctx->out[row], ctx->base + row * ctx->row_bytes, ctx->xq); } } typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT_USED]; const uint8_t *up_base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq; float expert_weight[DS4_MAX_EXPERT_USED]; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED]; uint64_t up_row_bytes[DS4_MAX_EXPERT_USED]; uint32_t gate_type; uint32_t up_type; int n_expert; } matvec_q5_q6_k_mid_ctx; static void matvec_q5_q6_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q5_q6_k_mid_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const int slot = (int)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)slot * ctx->out_dim; float gate = 0.0f; float up = 0.0f; ds4_vec_dot_q5_q6_K_q8_K(ctx->gate_type, (int)ctx->in_dim, &gate, ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot], ctx->xq); ds4_vec_dot_q5_q6_K_q8_K(ctx->up_type, (int)ctx->in_dim, &up, ctx->up_base[slot] + row * ctx->up_row_bytes[slot], ctx->xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot]; } } static void matvec_q5_q6_k_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if ((gate_w->type != DS4_TENSOR_Q5_K && gate_w->type != DS4_TENSOR_Q6_K) || (up_w->type != DS4_TENSOR_Q5_K && up_w->type != DS4_TENSOR_Q6_K)) { ds4_die("expected Q5_K/Q6_K expert tensors"); } if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; matvec_q5_q6_k_mid_ctx ctx = { .mid = mid, .xq = xq, .clamp = clamp, .gate_type = gate_w->type, .up_type = up_w->type, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i], &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]); ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i], &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]); if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) { ds4_die("paired Q5_K/Q6_K expert tensors do not match"); } if (i == 0) { in_dim0 = gate_in_dim; out_dim0 = gate_out_dim; } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) { ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); } ctx.expert_weight[i] = expert_weight[i]; } if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); ctx.in_dim = in_dim0; ctx.out_dim = out_dim0; ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q5_q6_k_mid_worker, &ctx); } static void matvec_q5_q6_k_expert( float *out, const ds4_model *m, const ds4_tensor *w, const float *x, uint32_t expert) { if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { ds4_die("expected a Q5_K or Q6_K expert tensor"); } uint64_t in_dim, out_dim, row_bytes; const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); if (in_dim % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); matvec_q5_q6_k_ctx ctx = { .out = out, .base = base, .xq = xq, .in_dim = in_dim, .row_bytes = row_bytes, .type = w->type, }; ds4_parallel_for(out_dim, matvec_q5_q6_k_worker, &ctx); free(xq); } typedef struct { float *out; const uint8_t *base[DS4_MAX_EXPERT_USED]; const block_q8_K *xq[DS4_MAX_EXPERT_USED]; uint64_t in_dim; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; uint32_t type; int n_expert; } matvec_q5_q6_k_accum_ctx; static void matvec_q5_q6_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q5_q6_k_accum_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { float acc = 0.0f; for (int i = 0; i < ctx->n_expert; i++) { float v = 0.0f; ds4_vec_dot_q5_q6_K_q8_K(ctx->type, (int)ctx->in_dim, &v, ctx->base[i] + row * ctx->row_bytes[i], ctx->xq[i]); acc += v; } ctx->out[row] = acc; } } static void matvec_q5_q6_k_experts_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type != DS4_TENSOR_Q5_K && w->type != DS4_TENSOR_Q6_K) { ds4_die("expected a Q5_K or Q6_K expert tensor"); } if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count"); uint64_t in_dim0 = 0; uint64_t out_dim0 = 0; const uint8_t *base[DS4_MAX_EXPERT_USED]; uint64_t row_bytes[DS4_MAX_EXPERT_USED]; for (int i = 0; i < n_expert; i++) { uint64_t in_dim, out_dim; base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]); if (i == 0) { in_dim0 = in_dim; out_dim0 = out_dim; } else if (in_dim != in_dim0 || out_dim != out_dim0) { ds4_die("Q5_K/Q6_K expert tensors do not share a layout"); } } if (in_dim0 % QK_K != 0) ds4_die("Q5_K/Q6_K expert row is not QK_K aligned"); const uint64_t n_blocks = in_dim0 / QK_K; matvec_q5_q6_k_accum_ctx ctx = { .out = out, .in_dim = in_dim0, .type = w->type, .n_expert = n_expert, }; for (int i = 0; i < n_expert; i++) { ctx.base[i] = base[i]; ctx.row_bytes[i] = row_bytes[i]; ctx.xq[i] = xq + (uint64_t)i * n_blocks; } ds4_parallel_for(out_dim0, matvec_q5_q6_k_accum_worker, &ctx); } /* Q4_K batch mid worker: same structure as IQ2_XXS batch but uses Q4_K dot. */ typedef struct { float *mid; const uint8_t *gate_base[DS4_MAX_EXPERT]; const uint8_t *up_base[DS4_MAX_EXPERT]; const block_q8_K *xq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; const float *pair_weight; float clamp; uint64_t in_dim; uint64_t out_dim; uint64_t gate_row_bytes[DS4_MAX_EXPERT]; uint64_t up_row_bytes[DS4_MAX_EXPERT]; uint64_t xq_blocks; } matvec_q4_k_batch_mid_ctx; static void matvec_q4_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) { matvec_q4_k_batch_mid_ctx *ctx = vctx; for (uint64_t task = task0; task < task1; task++) { const uint32_t active_idx = (uint32_t)(task / ctx->out_dim); const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim; const uint32_t expert = ctx->active_expert[active_idx]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]); const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks; float gate = 0.0f; float up = 0.0f; ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq); ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, xq); if (ctx->clamp > 1.0e-6f) { if (gate > ctx->clamp) gate = ctx->clamp; if (up > ctx->clamp) up = ctx->clamp; if (up < -ctx->clamp) up = -ctx->clamp; } ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id]; } } } /* Q4_K batch down accum worker: same structure as Q2_K batch but uses Q4_K dot. */ typedef struct { float *moe; const uint8_t *base[DS4_MAX_EXPERT]; const block_q8_K *midq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint32_t n_active; uint32_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t midq_blocks; } matvec_q4_k_batch_accum_rows_ctx; static void matvec_q4_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q4_k_batch_accum_rows_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { for (uint32_t t = 0; t < ctx->n_tok; t++) { ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; } for (uint32_t ai = 0; ai < ctx->n_active; ai++) { const uint32_t expert = ctx->active_expert[ai]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q4_K *br = (const block_q4_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; float v = 0.0f; ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, xq); ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; } } } } typedef struct { float *moe; const uint8_t *base[DS4_MAX_EXPERT]; const block_q8_K *midq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint32_t n_active; uint32_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t midq_blocks; } matvec_iq2_xxs_batch_accum_rows_ctx; static void matvec_iq2_xxs_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_iq2_xxs_batch_accum_rows_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { for (uint32_t t = 0; t < ctx->n_tok; t++) { ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; } for (uint32_t ai = 0; ai < ctx->n_active; ai++) { const uint32_t expert = ctx->active_expert[ai]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_iq2_xxs *br = (const block_iq2_xxs *)(ctx->base[expert] + row * ctx->row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; float v = 0.0f; ds4_vec_dot_iq2_xxs_q8_K((int)ctx->in_dim, &v, br, xq); ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; } } } } /* Dispatch: call the right gate/up mid builder based on tensor type. */ static void matvec_experts_mid_prequant( float *mid, const ds4_model *m, const ds4_tensor *gate_w, const ds4_tensor *up_w, const block_q8_K *xq, const int *selected, const float *expert_weight, int n_expert, float clamp) { if (gate_w->type == DS4_TENSOR_IQ2_XXS) { matvec_iq2_xxs_experts_mid_prequant(mid, m, gate_w, up_w, xq, selected, expert_weight, n_expert, clamp); } else if (gate_w->type == DS4_TENSOR_Q2_K) { matvec_q2_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, selected, expert_weight, n_expert, clamp); } else if (gate_w->type == DS4_TENSOR_Q4_K) { matvec_q4_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, selected, expert_weight, n_expert, clamp); } else if (gate_w->type == DS4_TENSOR_Q5_K || gate_w->type == DS4_TENSOR_Q6_K) { matvec_q5_q6_k_experts_mid_prequant(mid, m, gate_w, up_w, xq, selected, expert_weight, n_expert, clamp); } else { ds4_die("unsupported gate/up expert tensor type"); } } /* Dispatch: call the right down-projection accumulator based on tensor type. */ static void matvec_experts_down_accum_prequant( float *out, const ds4_model *m, const ds4_tensor *w, const block_q8_K *xq, const int *selected, int n_expert) { if (w->type == DS4_TENSOR_IQ2_XXS) { matvec_iq2_xxs_experts_accum_prequant(out, m, w, xq, selected, n_expert); } else if (w->type == DS4_TENSOR_Q2_K) { matvec_q2_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); } else if (w->type == DS4_TENSOR_Q4_K) { matvec_q4_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { matvec_q5_q6_k_experts_accum_prequant(out, m, w, xq, selected, n_expert); } else { ds4_die("unsupported down expert tensor type"); } } /* Dispatch: single-expert gate/up pair for tracing. */ static void matvec_expert_pair_prequant( float *out0, float *out1, const ds4_model *m, const ds4_tensor *w0, const ds4_tensor *w1, const block_q8_K *xq, uint32_t expert) { if (w0->type == DS4_TENSOR_IQ2_XXS) { matvec_iq2_xxs_expert_pair_prequant(out0, out1, m, w0, w1, xq, expert); } else if (w0->type == DS4_TENSOR_Q2_K) { uint64_t in_dim0, out_dim0, rb0; uint64_t in_dim1, out_dim1, rb1; const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); if (w1->type != DS4_TENSOR_Q2_K || in_dim0 != in_dim1 || out_dim0 != out_dim1) { ds4_die("paired Q2_K expert tensors do not match"); } for (uint64_t row = 0; row < out_dim0; row++) { const block_q2_K *gr = (const block_q2_K *)(base0 + row * rb0); ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out0[row], gr, xq); const block_q2_K *ur = (const block_q2_K *)(base1 + row * rb1); ds4_vec_dot_q2_K_q8_K((int)in_dim0, &out1[row], ur, xq); } } else if (w0->type == DS4_TENSOR_Q4_K) { uint64_t in_dim0, out_dim0, rb0; uint64_t in_dim1, out_dim1, rb1; const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q4_K expert tensors do not match"); for (uint64_t row = 0; row < out_dim0; row++) { const block_q4_K *gr = (const block_q4_K *)(base0 + row * rb0); ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out0[row], gr, xq); const block_q4_K *ur = (const block_q4_K *)(base1 + row * rb1); ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out1[row], ur, xq); } } else if (w0->type == DS4_TENSOR_Q5_K || w0->type == DS4_TENSOR_Q6_K) { uint64_t in_dim0, out_dim0, rb0; uint64_t in_dim1, out_dim1, rb1; const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0); const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1); if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q5_K/Q6_K expert tensors do not match"); for (uint64_t row = 0; row < out_dim0; row++) { ds4_vec_dot_q5_q6_K_q8_K(w0->type, (int)in_dim0, &out0[row], base0 + row * rb0, xq); ds4_vec_dot_q5_q6_K_q8_K(w1->type, (int)in_dim0, &out1[row], base1 + row * rb1, xq); } } else { ds4_die("unsupported gate/up expert tensor type"); } } /* Dispatch: single-expert down projection for tracing. */ static void matvec_expert_down( float *out, const ds4_model *m, const ds4_tensor *w, const float *x, uint32_t expert) { if (w->type == DS4_TENSOR_IQ2_XXS) { uint64_t in_dim, out_dim, row_bytes; const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); if (in_dim % QK_K != 0) ds4_die("IQ2_XXS expert row is not QK_K aligned"); block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); for (uint64_t row = 0; row < out_dim; row++) { const block_iq2_xxs *br = (const block_iq2_xxs *)(base + row * row_bytes); ds4_vec_dot_iq2_xxs_q8_K((int)in_dim, &out[row], br, xq); } free(xq); } else if (w->type == DS4_TENSOR_Q2_K) { matvec_q2_k_expert(out, m, w, x, expert); } else if (w->type == DS4_TENSOR_Q4_K) { uint64_t in_dim, out_dim, row_bytes; const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes); if (in_dim % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned"); block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0])); ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim); for (uint64_t row = 0; row < out_dim; row++) { const block_q4_K *br = (const block_q4_K *)(base + row * row_bytes); ds4_vec_dot_q4_K_q8_K((int)in_dim, &out[row], br, xq); } free(xq); } else if (w->type == DS4_TENSOR_Q5_K || w->type == DS4_TENSOR_Q6_K) { matvec_q5_q6_k_expert(out, m, w, x, expert); } else { ds4_die("unsupported down expert tensor type"); } } typedef struct { float *moe; const uint8_t *base[DS4_MAX_EXPERT]; const block_q8_K *midq; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint32_t n_active; uint32_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t midq_blocks; } matvec_q8_k_batch_accum_rows_ctx; static void matvec_q8_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_k_batch_accum_rows_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { for (uint32_t t = 0; t < ctx->n_tok; t++) { ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; } for (uint32_t ai = 0; ai < ctx->n_active; ai++) { const uint32_t expert = ctx->active_expert[ai]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const block_q8_K *br = (const block_q8_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]); for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks; float v = 0.0f; ds4_vec_dot_q8_K_q8_K((int)ctx->in_dim, &v, br, xq); ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; } } } } typedef struct { float *moe; const uint8_t *base[DS4_MAX_EXPERT]; const int8_t *midq; const float *midscale; const ds4_expert_pair *pairs; const uint32_t *pair_ids; const uint32_t *expert_offset; const uint32_t *active_expert; uint32_t n_active; uint32_t n_tok; uint64_t in_dim; uint64_t out_dim; uint64_t row_bytes[DS4_MAX_EXPERT]; uint64_t blocks; } matvec_q8_0_batch_accum_rows_ctx; static void matvec_q8_0_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) { matvec_q8_0_batch_accum_rows_ctx *ctx = vctx; for (uint64_t row = row0; row < row1; row++) { for (uint32_t t = 0; t < ctx->n_tok; t++) { ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f; } for (uint32_t ai = 0; ai < ctx->n_active; ai++) { const uint32_t expert = ctx->active_expert[ai]; const uint32_t begin = ctx->expert_offset[expert]; const uint32_t end = ctx->expert_offset[expert + 1]; const uint8_t *br = ctx->base[expert] + row * ctx->row_bytes[expert]; for (uint32_t i = begin; i < end; i++) { const uint32_t pair_id = ctx->pair_ids[i]; const ds4_expert_pair pair = ctx->pairs[pair_id]; const int8_t *xq = ctx->midq + (uint64_t)pair_id * ctx->blocks * 32u; const float *xscale = ctx->midscale + (uint64_t)pair_id * ctx->blocks; const float v = dot_q8_0_row(br, xq, xscale, ctx->in_dim, ctx->blocks); ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v; } } } } typedef struct { float *moe; const float *down_pair; uint32_t n_tok; uint64_t out_dim; } sum_down_pairs_ctx; static DS4_MAYBE_UNUSED void sum_down_pairs_worker(void *vctx, uint64_t row0, uint64_t row1) { sum_down_pairs_ctx *ctx = vctx; for (uint64_t idx = row0; idx < row1; idx++) { const uint32_t token = (uint32_t)(idx / ctx->out_dim); const uint64_t row = idx - (uint64_t)token * ctx->out_dim; float acc = 0.0f; for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { const uint64_t pair_id = (uint64_t)token * DS4_N_EXPERT_USED + slot; acc += ctx->down_pair[pair_id * ctx->out_dim + row]; } ctx->moe[idx] = acc; } } /* ========================================================================= * Hyper-Connection Transforms. * ========================================================================= * * DeepSeek V4 Flash keeps four hyper-connection streams per token. Before * attention or FFN, a learned small projection chooses how to reduce the HC * state into the 4096-wide sublayer input. After the sublayer, the post and * combine weights expand the result back into the four-stream HC state. */ /* Decode the HC control projection. The output contains pre weights, post * gates, and a small doubly-normalized combine matrix. */ static void hc_split_sinkhorn_one( float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; const float post_scale = scale[1]; const float comb_scale = scale[2]; for (int i = 0; i < n_hc; i++) { const float z = mix[i] * pre_scale + base[i]; out[i] = 1.0f / (1.0f + expf(-z)) + eps; } for (int i = 0; i < n_hc; i++) { const int off = n_hc + i; const float z = mix[off] * post_scale + base[off]; out[off] = 2.0f / (1.0f + expf(-z)); } float c[16 * 16]; for (int dst = 0; dst < n_hc; dst++) { float row_max = DS4_NEG_INF; for (int src = 0; src < n_hc; src++) { const int idx = src + dst * n_hc; const int off = 2 * n_hc + idx; const float v = mix[off] * comb_scale + base[off]; c[idx] = v; if (v > row_max) row_max = v; } float row_sum = 0.0f; for (int src = 0; src < n_hc; src++) { const int idx = src + dst * n_hc; const float v = expf(c[idx] - row_max); c[idx] = v; row_sum += v; } const float inv = 1.0f / row_sum; for (int src = 0; src < n_hc; src++) { const int idx = src + dst * n_hc; c[idx] = c[idx] * inv + eps; } } for (int src = 0; src < n_hc; src++) { float sum = 0.0f; for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; const float inv = 1.0f / (sum + eps); for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; } for (int iter = 1; iter < iters; iter++) { for (int dst = 0; dst < n_hc; dst++) { float sum = 0.0f; for (int src = 0; src < n_hc; src++) sum += c[src + dst * n_hc]; const float inv = 1.0f / (sum + eps); for (int src = 0; src < n_hc; src++) c[src + dst * n_hc] *= inv; } for (int src = 0; src < n_hc; src++) { float sum = 0.0f; for (int dst = 0; dst < n_hc; dst++) sum += c[src + dst * n_hc]; const float inv = 1.0f / (sum + eps); for (int dst = 0; dst < n_hc; dst++) c[src + dst * n_hc] *= inv; } } for (int i = 0; i < n_hc * n_hc; i++) out[2 * n_hc + i] = c[i]; } /* Reduce the four HC streams into the plain embedding vector consumed by a * normal attention or FFN sublayer. */ static void hc_weighted_sum_one( float * out, const float * x, const float * weights, uint32_t n_embd, uint32_t n_hc) { for (uint32_t d = 0; d < n_embd; d++) { float acc = 0.0f; for (uint32_t h = 0; h < n_hc; h++) { acc += x[(uint64_t)h * n_embd + d] * weights[h]; } out[d] = acc; } } /* HC pre step for one token. It normalizes the HC state, projects the control * vector, runs the Sinkhorn split, and emits the sublayer input plus post data. */ static void hc_pre_from_state_one_scratch( const ds4_model * model, const ds4_tensor * fn, const ds4_tensor * scale_tensor, const ds4_tensor * base_tensor, const float * residual_hc, float * out, float * post, float * comb, float * flat, bool serial_fn) { const uint32_t n_hc = DS4_N_HC; const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; float mix[24]; float split[24]; rms_norm_no_weight(flat, residual_hc, hc_dim, DS4_RMS_EPS); if (serial_fn) { matvec_f16_serial(mix, model, fn, flat); } else { matvec_f16(mix, model, fn, flat); } const float *scale = tensor_data(model, scale_tensor); const float *base = tensor_data(model, base_tensor); hc_split_sinkhorn_one(split, mix, scale, base, (int)n_hc, DS4_N_HC_SINKHORN_ITER, 1.0e-6f); hc_weighted_sum_one(out, residual_hc, split, DS4_N_EMBD, n_hc); memcpy(post, split + n_hc, n_hc * sizeof(post[0])); memcpy(comb, split + 2 * n_hc, n_hc * n_hc * sizeof(comb[0])); } static void hc_pre_from_state_one( const ds4_model * model, const ds4_tensor * fn, const ds4_tensor * scale_tensor, const ds4_tensor * base_tensor, const float * residual_hc, float * out, float * post, float * comb) { const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); hc_pre_from_state_one_scratch(model, fn, scale_tensor, base_tensor, residual_hc, out, post, comb, flat, false); free(flat); } static void layer_attn_pre_one( const ds4_model * model, const ds4_layer_weights * layer, const float * token_embd, float * out, float * residual_hc, float * post, float * comb) { const uint32_t n_hc = DS4_N_HC; for (uint32_t h = 0; h < n_hc; h++) { memcpy(residual_hc + (uint64_t)h * DS4_N_EMBD, token_embd, (size_t)DS4_N_EMBD * sizeof(token_embd[0])); } hc_pre_from_state_one(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, residual_hc, out, post, comb); } /* The input embedding starts all HC streams with the same token vector. */ static void hc_from_plain_embedding(float *out_hc, const float *x, uint32_t n_embd, uint32_t n_hc) { for (uint32_t h = 0; h < n_hc; h++) { memcpy(out_hc + (uint64_t)h * n_embd, x, (size_t)n_embd * sizeof(x[0])); } } /* HC post step for one sublayer output. It injects the new block output and * mixes the previous HC streams through the learned combine matrix. */ static void hc_post_one( float * out_hc, const float * block_out, const float * residual_hc, const float * post, const float * comb, uint32_t n_embd, uint32_t n_hc) { for (uint32_t dst = 0; dst < n_hc; dst++) { for (uint32_t d = 0; d < n_embd; d++) { float acc = block_out[d] * post[dst]; for (uint32_t src = 0; src < n_hc; src++) { /* The HC combine matrix is addressed as [dst_hc, src_hc]. */ acc += comb[dst + src * n_hc] * residual_hc[(uint64_t)src * n_embd + d]; } out_hc[(uint64_t)dst * n_embd + d] = acc; } } } typedef struct { float *out_hc; const float *block_out; const float *residual_hc; const float *post; const float *comb; uint64_t hc_dim; uint32_t n_embd; uint32_t n_hc; } hc_post_batch_ctx; static void hc_post_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { hc_post_batch_ctx *ctx = vctx; for (uint64_t t = t0; t < t1; t++) { hc_post_one(ctx->out_hc + t * ctx->hc_dim, ctx->block_out + t * ctx->n_embd, ctx->residual_hc + t * ctx->hc_dim, ctx->post + t * ctx->n_hc, ctx->comb + t * ctx->n_hc * ctx->n_hc, ctx->n_embd, ctx->n_hc); } } static void hc_post_batch( float * out_hc, const float * block_out, const float * residual_hc, const float * post, const float * comb, uint32_t n_tok, uint32_t n_embd, uint32_t n_hc) { hc_post_batch_ctx ctx = { .out_hc = out_hc, .block_out = block_out, .residual_hc = residual_hc, .post = post, .comb = comb, .hc_dim = (uint64_t)n_hc * n_embd, .n_embd = n_embd, .n_hc = n_hc, }; ds4_parallel_for_min_rows(n_tok, hc_post_batch_worker, &ctx, 1); } typedef struct { float *out_hc; const float *moe; const float *shared; const float *residual_hc; const float *post; const float *comb; uint64_t hc_dim; uint32_t n_embd; uint32_t n_hc; } hc_post_sum_batch_ctx; static void hc_post_sum_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { hc_post_sum_batch_ctx *ctx = vctx; for (uint64_t t = t0; t < t1; t++) { const float *moe = ctx->moe + t * ctx->n_embd; const float *shared = ctx->shared + t * ctx->n_embd; const float *residual = ctx->residual_hc + t * ctx->hc_dim; const float *post = ctx->post + t * ctx->n_hc; const float *comb = ctx->comb + t * ctx->n_hc * ctx->n_hc; float *out = ctx->out_hc + t * ctx->hc_dim; for (uint32_t dst = 0; dst < ctx->n_hc; dst++) { for (uint32_t d = 0; d < ctx->n_embd; d++) { float acc = (moe[d] + shared[d]) * post[dst]; for (uint32_t src = 0; src < ctx->n_hc; src++) { acc += comb[dst + src * ctx->n_hc] * residual[(uint64_t)src * ctx->n_embd + d]; } out[(uint64_t)dst * ctx->n_embd + d] = acc; } } } } static void hc_post_sum_batch( float * out_hc, const float * moe, const float * shared, const float * residual_hc, const float * post, const float * comb, uint32_t n_tok, uint32_t n_embd, uint32_t n_hc) { hc_post_sum_batch_ctx ctx = { .out_hc = out_hc, .moe = moe, .shared = shared, .residual_hc = residual_hc, .post = post, .comb = comb, .hc_dim = (uint64_t)n_hc * n_embd, .n_embd = n_embd, .n_hc = n_hc, }; ds4_parallel_for_min_rows(n_tok, hc_post_sum_batch_worker, &ctx, 1); } typedef struct { const ds4_model *model; const ds4_tensor *fn; const ds4_tensor *scale; const ds4_tensor *base; const ds4_tensor *norm_w; const float *inp_hc; float *residual_hc; float *cur; float *norm; float *post; float *comb; uint64_t hc_dim; uint32_t n_hc; } hc_pre_norm_batch_ctx; static void hc_pre_norm_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { hc_pre_norm_batch_ctx *ctx = vctx; const float *norm_w = tensor_data(ctx->model, ctx->norm_w); float *flat = xmalloc((size_t)ctx->hc_dim * sizeof(flat[0])); for (uint64_t t = t0; t < t1; t++) { const float *residual = ctx->inp_hc + t * ctx->hc_dim; if (ctx->residual_hc) { float *dst = ctx->residual_hc + t * ctx->hc_dim; memcpy(dst, residual, (size_t)ctx->hc_dim * sizeof(dst[0])); residual = dst; } hc_pre_from_state_one_scratch(ctx->model, ctx->fn, ctx->scale, ctx->base, residual, ctx->cur + t * DS4_N_EMBD, ctx->post + t * ctx->n_hc, ctx->comb + t * ctx->n_hc * ctx->n_hc, flat, true); rms_norm_weight(ctx->norm + t * DS4_N_EMBD, ctx->cur + t * DS4_N_EMBD, norm_w, DS4_N_EMBD, DS4_RMS_EPS); } free(flat); } /* Batched HC pre plus RMSNorm. Prefill uses this to keep the layer-major * token batch in contiguous arrays. */ static void hc_pre_norm_batch( const ds4_model * model, const ds4_tensor * fn, const ds4_tensor * scale, const ds4_tensor * base, const ds4_tensor * norm_w, const float * inp_hc, float * residual_hc, float * cur, float * norm, float * post, float * comb, uint32_t n_tok) { hc_pre_norm_batch_ctx ctx = { .model = model, .fn = fn, .scale = scale, .base = base, .norm_w = norm_w, .inp_hc = inp_hc, .residual_hc = residual_hc, .cur = cur, .norm = norm, .post = post, .comb = comb, .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, .n_hc = DS4_N_HC, }; ds4_parallel_for_min_rows(n_tok, hc_pre_norm_batch_worker, &ctx, 1); } static void layer_attn_norm_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x) { const float *attn_norm = tensor_data(model, layer->attn_norm); rms_norm_weight(out, x, attn_norm, DS4_N_EMBD, DS4_RMS_EPS); } /* ========================================================================= * Attention Projections, RoPE, and Attention Output. * ========================================================================= * * This block performs the attention half of a transformer layer: HC pre, * attention RMSNorm, Q and KV projections, layer-specific RoPE, sink-aware * attention over raw and compressed KV rows, and the grouped LoRA output * projection back to embedding width. */ /* Q projection is low-rank: Q8_0 into the model-specific LoRA-Q rank, * RMSNorm, then Q8_0 back to all attention heads. */ static void layer_q_projection_normed_one( const ds4_model * model, const ds4_layer_weights * layer, const float * norm, float * q) { const uint32_t q_rank = DS4_N_LORA_Q; float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); float *qr_norm = xmalloc((size_t)q_rank * sizeof(qr_norm[0])); const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); matvec_q8_0(qr, model, layer->attn_q_a, norm); rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); matvec_q8_0(q, model, layer->attn_q_b, qr_norm); head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); free(qr_norm); free(qr); } static void layer_q_projection_with_lora_one( const ds4_model * model, const ds4_layer_weights * layer, const float * norm, float * q, float * qr_norm) { const uint32_t q_rank = DS4_N_LORA_Q; float *qr = xmalloc((size_t)q_rank * sizeof(qr[0])); const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); matvec_q8_0(qr, model, layer->attn_q_a, norm); rms_norm_weight(qr_norm, qr, q_a_norm, q_rank, DS4_RMS_EPS); matvec_q8_0(q, model, layer->attn_q_b, qr_norm); head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); free(qr); } /* KV projection has one KV head of width 512, followed by a learned RMSNorm. */ static void layer_kv_projection_normed_one( const ds4_model * model, const ds4_layer_weights * layer, const float * normed, float * kv) { float *raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(raw[0])); const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); matvec_q8_0(raw, model, layer->attn_kv, normed); rms_norm_weight(kv, raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); free(raw); } static void layer_q_projection_with_lora_one_decode_scratch( const ds4_model * model, const ds4_layer_weights * layer, const float * norm, float * q, float * qr_norm, ds4_cpu_decode_scratch * scratch) { const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); matvec_q8_0_decode_scratch(scratch->qr, model, layer->attn_q_a, norm, scratch); rms_norm_weight(qr_norm, scratch->qr, q_a_norm, DS4_N_LORA_Q, DS4_RMS_EPS); matvec_q8_0_decode_scratch(q, model, layer->attn_q_b, qr_norm, scratch); head_rms_norm_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); } static void layer_kv_projection_normed_one_decode_scratch( const ds4_model * model, const ds4_layer_weights * layer, const float * normed, float * kv, ds4_cpu_decode_scratch * scratch) { const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); matvec_q8_0_decode_scratch(scratch->kv_raw, model, layer->attn_kv, normed, scratch); rms_norm_weight(kv, scratch->kv_raw, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); } static float rope_yarn_ramp(float low, float high, int i0) { const float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low); return 1.0f - fminf(1.0f, fmaxf(0.0f, y)); } static float rope_yarn_corr_dim(int n_dims, uint64_t n_ctx_orig, float n_rot, float base) { return (float)n_dims * logf((float)n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) / (2.0f * logf(base)); } static void rope_yarn_corr_dims(int n_dims, uint64_t n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]) { const float start = floorf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); const float end = ceilf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); dims[0] = fmaxf(0.0f, start); dims[1] = fminf((float)(n_dims - 1), end); } /* Apply DS4 RoPE only to the tail of each head. Compressed layers use the * long-context frequency base and scale; inverse mode rotates attention output * back before the grouped output projection. */ static void rope_tail_ext_inplace( float * x, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint64_t n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, bool inverse) { const uint32_t n_nope = head_dim - n_rot; const float theta_scale = powf(freq_base, -2.0f / (float)n_rot); const float sin_sign = inverse ? -1.0f : 1.0f; float corr_dims[2] = { 0.0f, 0.0f }; if (ext_factor != 0.0f) { rope_yarn_corr_dims((int)n_rot, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); } for (uint32_t h = 0; h < n_head; h++) { float *tail = x + (uint64_t)h * head_dim + n_nope; float theta_extrap = (float)pos; for (uint32_t i = 0; i < n_rot; i += 2) { const float theta_interp = freq_scale * theta_extrap; float theta = theta_interp; float mscale = attn_factor; if (ext_factor != 0.0f) { const float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], (int)i) * ext_factor; theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); } const float c = cosf(theta) * mscale; const float s = sin_sign * sinf(theta) * mscale; const float x0 = tail[i + 0]; const float x1 = tail[i + 1]; tail[i + 0] = x0 * c - x1 * s; tail[i + 1] = x0 * s + x1 * c; theta_extrap *= theta_scale; } } } /* Dense layers and compressed layers use different RoPE bases. */ static float layer_rope_freq_base(uint32_t il) { return ds4_layer_compress_ratio(il) != 0 && DS4_COMPRESS_ROPE_FREQ_BASE > 0.0f ? DS4_COMPRESS_ROPE_FREQ_BASE : DS4_ROPE_FREQ_BASE; } static float layer_rope_freq_scale(uint32_t il) { if (ds4_layer_compress_ratio(il) == 0 || DS4_ROPE_SCALE_FACTOR <= 0.0f) { return 1.0f; } return 1.0f / DS4_ROPE_SCALE_FACTOR; } static void rope_tail_layer_inplace( float * x, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint32_t il, bool inverse) { const bool compressed = ds4_layer_compress_ratio(il) != 0; const float freq_base = layer_rope_freq_base(il); const float freq_scale = layer_rope_freq_scale(il); const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; float attn_factor = 1.0f; if (ext_factor != 0.0f && freq_scale > 0.0f) { /* * This YaRN helper applies magnitude scaling internally. DeepSeek V4 * reference RoPE uses interpolation without that magnitude change, so * pass the inverse factor here and let the helper cancel itself out. */ attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); } rope_tail_ext_inplace(x, n_head, head_dim, n_rot, pos, compressed ? DS4_ROPE_ORIG_CTX : 0, freq_base, freq_scale, ext_factor, attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, inverse); } typedef struct { float *x; uint64_t stride; uint32_t n_head; uint32_t head_dim; uint32_t n_rot; uint32_t pos0; uint32_t il; bool inverse; } rope_tail_batch_ctx; static void rope_tail_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { rope_tail_batch_ctx *ctx = vctx; for (uint64_t tt = t0; tt < t1; tt++) { rope_tail_layer_inplace(ctx->x + tt * ctx->stride, ctx->n_head, ctx->head_dim, ctx->n_rot, ctx->pos0 + (uint32_t)tt, ctx->il, ctx->inverse); } } static void rope_tail_layer_batch_inplace( float *x, uint64_t stride, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t il, bool inverse, uint32_t n_tok) { rope_tail_batch_ctx ctx = { .x = x, .stride = stride, .n_head = n_head, .head_dim = head_dim, .n_rot = n_rot, .pos0 = pos0, .il = il, .inverse = inverse, }; ds4_parallel_for_min_rows(n_tok, rope_tail_batch_worker, &ctx, 1); } static inline float dot_f32(const float *a, const float *b, uint32_t n) { #if defined(__ARM_NEON) uint32_t i = 0; float32x4_t acc0 = vdupq_n_f32(0.0f); float32x4_t acc1 = vdupq_n_f32(0.0f); for (; i + 8 <= n; i += 8) { acc0 = vfmaq_f32(acc0, vld1q_f32(a + i), vld1q_f32(b + i)); acc1 = vfmaq_f32(acc1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); } float acc = vaddvq_f32(vaddq_f32(acc0, acc1)); for (; i < n; i++) acc += a[i] * b[i]; return acc; #else float acc = 0.0f; for (uint32_t i = 0; i < n; i++) acc += a[i] * b[i]; return acc; #endif } static inline void axpy_f32(float *y, const float *x, float a, uint32_t n) { #if defined(__ARM_NEON) uint32_t i = 0; const float32x4_t av = vdupq_n_f32(a); for (; i + 8 <= n; i += 8) { vst1q_f32(y + i, vfmaq_f32(vld1q_f32(y + i), av, vld1q_f32(x + i))); vst1q_f32(y + i + 4, vfmaq_f32(vld1q_f32(y + i + 4), av, vld1q_f32(x + i + 4))); } for (; i < n; i++) y[i] += a * x[i]; #else for (uint32_t i = 0; i < n; i++) y[i] += a * x[i]; #endif } static inline void scale_f32(float *x, float a, uint32_t n) { #if defined(__ARM_NEON) uint32_t i = 0; const float32x4_t av = vdupq_n_f32(a); for (; i + 8 <= n; i += 8) { vst1q_f32(x + i, vmulq_f32(vld1q_f32(x + i), av)); vst1q_f32(x + i + 4, vmulq_f32(vld1q_f32(x + i + 4), av)); } for (; i < n; i++) x[i] *= a; #else for (uint32_t i = 0; i < n; i++) x[i] *= a; #endif } static float sigmoid_stable(float x) { if (x >= 0.0f) { const float e = expf(-x); return 1.0f / (1.0f + e); } else { const float e = expf(x); return e / (1.0f + e); } } /* Sink-aware attention over a set of KV rows. The learned sink logit is part * of the softmax denominator but contributes no value vector. */ static void layer_attention_rows_one( float * out_heads, const ds4_model * model, const ds4_layer_weights * layer, const float * q, const float * kv_rows, uint32_t n_kv) { const float *sinks = tensor_data(model, layer->attn_sinks); const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); float score_stack[512]; float *score = n_kv <= 512 ? score_stack : xmalloc((size_t)n_kv * sizeof(score[0])); for (uint32_t h = 0; h < DS4_N_HEAD; h++) { const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; float max_score = sinks[h]; for (uint32_t r = 0; r < n_kv; r++) { const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; score[r] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[r] > max_score) max_score = score[r]; } float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); float denom = expf(sinks[h] - max_score); for (uint32_t r = 0; r < n_kv; r++) { const float weight = expf(score[r] - max_score); const float *kv = kv_rows + (uint64_t)r * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } const float inv = 1.0f / denom; scale_f32(oh, inv, DS4_N_HEAD_DIM); } if (score != score_stack) free(score); } static void layer_attention_one( float * out_heads, const ds4_model * model, const ds4_layer_weights * layer, const float * q, const float * kv) { layer_attention_rows_one(out_heads, model, layer, q, kv, 1); } /* Attention output projection is grouped: each group first maps its heads to * a 1024-rank low vector, then all groups are projected back to 4096. */ static void layer_grouped_out_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * heads) { const uint32_t n_groups = 8; const uint32_t group_heads = DS4_N_HEAD / n_groups; const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; const uint32_t rank = 1024; float *low = xcalloc((size_t)n_groups * rank, sizeof(low[0])); matvec_q8_0_grouped_rows(low, model, layer->attn_output_a, heads, n_groups, group_dim, rank); matvec_q8_0(out, model, layer->attn_output_b, low); free(low); } static void layer_grouped_out_one_decode_scratch( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * heads, ds4_cpu_decode_scratch * scratch) { const uint32_t n_groups = 8; const uint32_t group_heads = DS4_N_HEAD / n_groups; const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; const uint32_t rank = 1024; memset(scratch->attn_low, 0, (size_t)n_groups * rank * sizeof(scratch->attn_low[0])); matvec_q8_0_grouped_rows_decode_scratch(scratch->attn_low, model, layer->attn_output_a, heads, n_groups, group_dim, rank, scratch); matvec_q8_0_decode_scratch(out, model, layer->attn_output_b, scratch->attn_low, scratch); } static void layer_grouped_out_batch( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * heads, uint32_t n_tok) { const uint32_t n_groups = 8; const uint32_t group_heads = DS4_N_HEAD / n_groups; const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; const uint32_t rank = 1024; float *low = xcalloc((size_t)n_tok * n_groups * rank, sizeof(low[0])); matmul_q8_0_grouped_batch(low, model, layer->attn_output_a, heads, n_tok, n_groups, group_dim, rank); matmul_q8_0_batch(out, model, layer->attn_output_b, low, n_tok); free(low); } /* ========================================================================= * Mixture-of-Experts FFN. * ========================================================================= * * This is the FFN half of each layer. It includes the shared expert, routed * expert selection, IQ2_XXS gate/up projections, SwiGLU, Q2_K down projection, * and the HC post step that returns the result to four-stream state. */ static float silu(float x) { return x * sigmoid_stable(x); } static float softplus_stable(float x) { if (x > 20.0f) return x; if (x < -20.0f) return expf(x); return log1pf(expf(x)); } static void swiglu(float *out, const float *gate, const float *up, uint64_t n, float clamp) { for (uint64_t i = 0; i < n; i++) { float g = gate[i]; float u = up[i]; if (clamp > 1.0e-6f) { if (g > clamp) g = clamp; if (u > clamp) u = clamp; if (u < -clamp) u = -clamp; } out[i] = silu(g) * u; } } /* The shared expert is a normal Q8_0 SwiGLU MLP that runs for every token. */ static void layer_shared_ffn_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x) { float *gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])); float *up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])); float *mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])); const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)blocks * 32); float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); if (layer->ffn_up_shexp->type != 8 || layer->ffn_gate_shexp->type != 8 || layer->ffn_up_shexp->dim[0] != in_dim) { ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); } quantize_q8_0_activation(x, xq, xscale, in_dim); matvec_q8_0_pair_prequant(gate, up, model, layer->ffn_gate_shexp, layer->ffn_up_shexp, xq, xscale); swiglu(mid, gate, up, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP); matvec_q8_0(out, model, layer->ffn_down_shexp, mid); free(xscale); free(xq); free(mid); free(up); free(gate); } static void layer_shared_ffn_one_decode_scratch( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x, ds4_cpu_decode_scratch * scratch) { const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; if (layer->ffn_up_shexp->type != 8 || layer->ffn_gate_shexp->type != 8 || layer->ffn_up_shexp->dim[0] != in_dim) { ds4_die("shared expert gate/up tensors do not share a Q8_0 input layout"); } matvec_q8_0_pair_decode_scratch(scratch->shared_gate, scratch->shared_up, model, layer->ffn_gate_shexp, layer->ffn_up_shexp, x, scratch); swiglu(scratch->shared_mid, scratch->shared_gate, scratch->shared_up, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP); matvec_q8_0_decode_scratch(out, model, layer->ffn_down_shexp, scratch->shared_mid, scratch); } typedef struct { float *mid; const float *gate; const float *up; uint64_t n; float clamp; } swiglu_batch_ctx; static void swiglu_batch_worker(void *vctx, uint64_t t0, uint64_t t1) { swiglu_batch_ctx *ctx = vctx; for (uint64_t t = t0; t < t1; t++) { swiglu(ctx->mid + t * ctx->n, ctx->gate + t * ctx->n, ctx->up + t * ctx->n, ctx->n, ctx->clamp); } } static void layer_shared_ffn_batch( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x, uint32_t n_tok) { const uint64_t in_dim = layer->ffn_gate_shexp->dim[0]; const uint64_t hidden = layer->ffn_gate_shexp->dim[1]; if (layer->ffn_up_shexp->type != 8 || layer->ffn_gate_shexp->type != 8 || layer->ffn_down_shexp->type != 8 || layer->ffn_up_shexp->dim[0] != in_dim || layer->ffn_up_shexp->dim[1] != hidden || layer->ffn_down_shexp->dim[0] != hidden) { ds4_die("shared expert tensors do not share the expected Q8_0 layout"); } float *gate = xmalloc((size_t)n_tok * hidden * sizeof(gate[0])); float *up = xmalloc((size_t)n_tok * hidden * sizeof(up[0])); float *mid = xmalloc((size_t)n_tok * hidden * sizeof(mid[0])); matmul_q8_0_pair_batch(gate, up, model, layer->ffn_gate_shexp, layer->ffn_up_shexp, x, n_tok); swiglu_batch_ctx swiglu_ctx = { .mid = mid, .gate = gate, .up = up, .n = hidden, .clamp = DS4_SWIGLU_CLAMP_EXP, }; ds4_parallel_for(n_tok, swiglu_batch_worker, &swiglu_ctx); matmul_q8_0_batch(out, model, layer->ffn_down_shexp, mid, n_tok); free(mid); free(up); free(gate); } /* Early DS4 layers use token-id hash routing instead of top-k routing. */ static void layer_hash_selected_experts( int selected[DS4_MAX_EXPERT_USED], const ds4_model *model, const ds4_layer_weights *layer, int token) { ds4_tensor *t = layer->ffn_gate_tid2eid; if (!t) ds4_die("hash routing table is missing for this layer"); if (t->type != 26 || t->ndim != 2 || t->dim[0] != DS4_N_EXPERT_USED) { ds4_die("ffn_gate_tid2eid.weight has an unexpected layout"); } if (token < 0 || (uint64_t)token >= t->dim[1]) { ds4_die("token id is outside the hash routing table"); } const int32_t *table = tensor_data(model, t); const int32_t *row = table + (uint64_t)token * DS4_N_EXPERT_USED; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) selected[i] = row[i]; } /* Router scores use sqrt(softplus(logit)); normalization happens only after * the six selected experts are known. */ static void layer_router_probs_one( float probs[DS4_MAX_EXPERT], const ds4_model * model, const ds4_layer_weights * layer, const float * x) { float logits[DS4_MAX_EXPERT]; matvec_any(logits, model, layer->ffn_gate_inp, x); for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { probs[i] = sqrtf(softplus_stable(logits[i])); } } static void layer_hash_router_weights_from_probs( float weights_out[DS4_MAX_EXPERT_USED], const float probs[DS4_MAX_EXPERT], const int selected[DS4_MAX_EXPERT_USED]) { float sum = 0.0f; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) ds4_die("hash-selected expert is outside router range"); weights_out[i] = probs[selected[i]]; sum += weights_out[i]; } if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { weights_out[i] = weights_out[i] / sum * DS4_EXPERT_WEIGHT_SCALE; } } static void layer_hash_router_weights_one( float weights_out[DS4_MAX_EXPERT_USED], const ds4_model * model, const ds4_layer_weights * layer, const float * x, const int selected[DS4_MAX_EXPERT_USED]) { float probs[DS4_MAX_EXPERT]; layer_router_probs_one(probs, model, layer, x); layer_hash_router_weights_from_probs(weights_out, probs, selected); } static void topk_desc(const float *score, int n, int k, int *idx) { for (int i = 0; i < k; i++) idx[i] = -1; for (int i = 0; i < n; i++) { for (int j = 0; j < k; j++) { if (idx[j] < 0 || score[i] > score[idx[j]]) { for (int m = k - 1; m > j; m--) idx[m] = idx[m - 1]; idx[j] = i; break; } } } } /* Later layers choose the six experts by biased top-k, but weight them using * the unbiased router probabilities. */ static void layer_topk_selected_experts_from_probs( int selected[DS4_MAX_EXPERT_USED], float expert_weight[DS4_MAX_EXPERT_USED], const ds4_model *model, const ds4_layer_weights *layer, const float probs[DS4_MAX_EXPERT]); static void layer_topk_selected_experts( int selected[DS4_MAX_EXPERT_USED], float expert_weight[DS4_MAX_EXPERT_USED], const ds4_model *model, const ds4_layer_weights *layer, const float *x) { float probs[DS4_MAX_EXPERT] = {0}; layer_router_probs_one(probs, model, layer, x); layer_topk_selected_experts_from_probs(selected, expert_weight, model, layer, probs); } static void layer_topk_selected_experts_from_probs( int selected[DS4_MAX_EXPERT_USED], float expert_weight[DS4_MAX_EXPERT_USED], const ds4_model *model, const ds4_layer_weights *layer, const float probs[DS4_MAX_EXPERT]) { float selection[DS4_MAX_EXPERT]; memcpy(selection, probs, sizeof(selection)); if (layer->ffn_exp_probs_b) { const float *bias = tensor_data(model, layer->ffn_exp_probs_b); for (uint32_t i = 0; i < DS4_N_EXPERT; i++) selection[i] += bias[i]; } topk_desc(selection, (int)DS4_N_EXPERT, (int)DS4_N_EXPERT_USED, selected); float sum = 0.0f; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { expert_weight[i] = probs[selected[i]]; sum += expert_weight[i]; } if (sum < 6.103515625e-5f) sum = 6.103515625e-5f; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { expert_weight[i] = expert_weight[i] / sum * DS4_EXPERT_WEIGHT_SCALE; } } static void print_vec_stats(const char *name, const float *x, uint64_t n); /* Single-token routed MoE. It selects six experts, runs IQ2_XXS gate/up, * applies SwiGLU and router weights, then accumulates Q2_K down projections. */ static void layer_routed_moe_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x, uint32_t il, int token, float clamp, bool trace) { int selected[DS4_MAX_EXPERT_USED]; float expert_weight[DS4_MAX_EXPERT_USED]; float *gate = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(gate[0])) : NULL; float *up = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(up[0])) : NULL; float *mid = trace ? xmalloc((size_t)DS4_N_FF_EXP * sizeof(mid[0])) : NULL; float *mid_all = trace ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(mid_all[0])); float *down = trace ? xmalloc((size_t)DS4_N_EMBD * sizeof(down[0])) : NULL; const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const bool routed_q8_k = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; if (routed_q8_0) { if (trace) ds4_die("Q8_0 routed trace mode is not supported"); if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { ds4_die("Q8_0 expert input has an unexpected layout"); } } else { if (routed_q8_k && trace) ds4_die("Q8_K routed trace mode is not supported"); if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { ds4_die("routed expert down input has an unexpected layout"); } } block_q8_K *xq = routed_q8_0 ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0])); block_q8_K *midq = (trace || routed_q8_0) ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0])); memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); if (!routed_q8_0) { ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); } if (layer->ffn_gate_tid2eid) { layer_hash_selected_experts(selected, model, layer, token); layer_hash_router_weights_one(expert_weight, model, layer, x, selected); } else { layer_topk_selected_experts(selected, expert_weight, model, layer, x); } if (routed_q8_0) { const uint64_t x_blocks = expert_in_dim / 32u; int8_t *xq8 = xmalloc((size_t)x_blocks * 32u); float *xscale8 = xmalloc((size_t)x_blocks * sizeof(float)); quantize_q8_0_activation(x, xq8, xscale8, expert_in_dim); matvec_q8_0_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq8, xscale8, selected, expert_weight, DS4_N_EXPERT_USED, clamp); const uint64_t mid_blocks = down_in_dim / 32u; int8_t *midq8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * 32u); float *midscale8 = xmalloc((size_t)DS4_N_EXPERT_USED * mid_blocks * sizeof(float)); for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, midq8 + (uint64_t)i * mid_blocks * 32u, midscale8 + (uint64_t)i * mid_blocks, down_in_dim); } matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, midq8, midscale8, selected, DS4_N_EXPERT_USED); free(midscale8); free(midq8); free(xscale8); free(xq8); } else if (routed_q8_k) { matvec_q8_k_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq, selected, expert_weight, DS4_N_EXPERT_USED, clamp); for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, midq + (uint64_t)i * (down_in_dim / QK_K), (int64_t)down_in_dim); } matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); } else if (!trace) { matvec_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq, selected, expert_weight, DS4_N_EXPERT_USED, clamp); for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, midq + (uint64_t)i * (down_in_dim / QK_K), (int64_t)down_in_dim); } matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); } else { for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { const uint32_t expert = (uint32_t)selected[i]; matvec_expert_pair_prequant(gate, up, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq, expert); char name[64]; snprintf(name, sizeof(name), "blk.%u expert %u gate", il, expert); print_vec_stats(name, gate, DS4_N_FF_EXP); snprintf(name, sizeof(name), "blk.%u expert %u up", il, expert); print_vec_stats(name, up, DS4_N_FF_EXP); /* * DeepSeek V4 clamps routed expert gate/up values before SwiGLU and * applies the router weight before the down projection. */ const float limit = clamp; for (uint32_t j = 0; j < DS4_N_FF_EXP; j++) { if (limit > 1.0e-6f) { if (gate[j] > limit) gate[j] = limit; if (up[j] > limit) up[j] = limit; if (up[j] < -limit) up[j] = -limit; } mid[j] = silu(gate[j]) * up[j] * expert_weight[i]; } snprintf(name, sizeof(name), "blk.%u expert %u mid", il, expert); print_vec_stats(name, mid, DS4_N_FF_EXP); matvec_expert_down(down, model, layer->ffn_down_exps, mid, expert); snprintf(name, sizeof(name), "blk.%u expert %u down", il, expert); print_vec_stats(name, down, DS4_N_EMBD); for (uint32_t j = 0; j < DS4_N_EMBD; j++) out[j] += down[j]; } } free(midq); free(xq); free(down); free(mid_all); free(mid); free(up); free(gate); } /* Decode version of routed MoE: same math as layer_routed_moe_one(), but all * large temporaries come from the persistent scratch arena. */ static void layer_routed_moe_one_prealloc( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x, uint32_t il, int token, float clamp, float * mid_all, block_q8_K * xq, block_q8_K * midq, int8_t * q8_xq, float * q8_xscale, int8_t * q8_midq, float * q8_midscale) { int selected[DS4_MAX_EXPERT_USED]; float expert_weight[DS4_MAX_EXPERT_USED]; const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const bool routed_q8_k = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; if (routed_q8_0) { if ((expert_in_dim % 32u) != 0) ds4_die("Q8_0 expert input is not QK8_0 aligned"); if (down_in_dim != DS4_N_FF_EXP || (down_in_dim % 32u) != 0) { ds4_die("Q8_0 expert input has an unexpected layout"); } } else { if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { ds4_die("routed expert down input has an unexpected layout"); } } memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0])); if (layer->ffn_gate_tid2eid) { layer_hash_selected_experts(selected, model, layer, token); layer_hash_router_weights_one(expert_weight, model, layer, x, selected); } else { layer_topk_selected_experts(selected, expert_weight, model, layer, x); } if (routed_q8_0) { if (!q8_xq || !q8_xscale || !q8_midq || !q8_midscale) { ds4_die("missing Q8_0 routed decode scratch"); } quantize_q8_0_activation(x, q8_xq, q8_xscale, expert_in_dim); matvec_q8_0_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, q8_xq, q8_xscale, selected, expert_weight, DS4_N_EXPERT_USED, clamp); const uint64_t mid_blocks = down_in_dim / 32u; for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { quantize_q8_0_activation(mid_all + (uint64_t)i * down_in_dim, q8_midq + (uint64_t)i * mid_blocks * 32u, q8_midscale + (uint64_t)i * mid_blocks, down_in_dim); } matvec_q8_0_experts_accum_prequant(out, model, layer->ffn_down_exps, q8_midq, q8_midscale, selected, DS4_N_EXPERT_USED); (void)il; return; } if (!mid_all || !xq || !midq) ds4_die("missing routed decode scratch"); ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim); if (routed_q8_k) { matvec_q8_k_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq, selected, expert_weight, DS4_N_EXPERT_USED, clamp); } else { matvec_experts_mid_prequant(mid_all, model, layer->ffn_gate_exps, layer->ffn_up_exps, xq, selected, expert_weight, DS4_N_EXPERT_USED, clamp); } for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) { ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim, midq + (uint64_t)i * (down_in_dim / QK_K), (int64_t)down_in_dim); } if (routed_q8_k) { matvec_q8_k_experts_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); } else { matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED); } (void)il; } /* Prefill MoE groups token/expert pairs by expert so each active expert's * rows are scanned once for the whole token batch. */ static void layer_routed_moe_batch( float * moe, const ds4_model * model, const ds4_layer_weights * layer, const float * norm, const int * token_ids, uint32_t n_tok, uint32_t il, float clamp) { const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t expert_out_dim = layer->ffn_gate_exps->dim[1]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const uint64_t down_out_dim = layer->ffn_down_exps->dim[1]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const bool routed_q8_k = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; if (routed_q8_0) { if ((expert_in_dim % 32u) != 0 || (down_in_dim % 32u) != 0) { ds4_die("Q8_0 routed expert input is not QK8_0 aligned"); } } else { if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned"); if (down_in_dim % QK_K != 0) ds4_die("routed expert down input is not QK_K aligned"); } if (expert_out_dim != down_in_dim || down_out_dim != DS4_N_EMBD) { ds4_die("routed expert tensor layout is unexpected"); } const uint32_t total_pairs = n_tok * DS4_N_EXPERT_USED; uint32_t counts[DS4_MAX_EXPERT + 1] = {0}; uint32_t cursor[DS4_MAX_EXPERT] = {0}; uint32_t active_expert[DS4_MAX_EXPERT]; uint32_t n_active = 0; int *selected = xmalloc((size_t)total_pairs * sizeof(selected[0])); float *pair_weight = xmalloc((size_t)total_pairs * sizeof(pair_weight[0])); ds4_expert_pair *pairs = xmalloc((size_t)total_pairs * sizeof(pairs[0])); for (uint32_t t = 0; t < n_tok; t++) { int sel[DS4_MAX_EXPERT_USED]; float weights[DS4_MAX_EXPERT_USED]; if (layer->ffn_gate_tid2eid) { layer_hash_selected_experts(sel, model, layer, token_ids[t]); layer_hash_router_weights_one(weights, model, layer, norm + (uint64_t)t * expert_in_dim, sel); } else { layer_topk_selected_experts(sel, weights, model, layer, norm + (uint64_t)t * expert_in_dim); } for (uint32_t slot = 0; slot < DS4_N_EXPERT_USED; slot++) { const uint32_t pair_id = t * DS4_N_EXPERT_USED + slot; selected[pair_id] = sel[slot]; pair_weight[pair_id] = weights[slot]; pairs[pair_id] = (ds4_expert_pair){ .token = t, .slot = slot }; if (sel[slot] < 0 || (uint32_t)sel[slot] >= DS4_N_EXPERT) ds4_die("selected expert is outside range"); counts[(uint32_t)sel[slot] + 1]++; } } for (uint32_t e = 0; e < DS4_N_EXPERT; e++) { counts[e + 1] += counts[e]; cursor[e] = counts[e]; if (counts[e + 1] != counts[e]) active_expert[n_active++] = e; } uint32_t *pair_ids = xmalloc((size_t)total_pairs * sizeof(pair_ids[0])); for (uint32_t p = 0; p < total_pairs; p++) { const uint32_t e = (uint32_t)selected[p]; pair_ids[cursor[e]++] = p; } if (routed_q8_0) { const uint64_t x_blocks = expert_in_dim / 32u; int8_t *xq8 = xmalloc((size_t)n_tok * x_blocks * 32u); float *xscale8 = xmalloc((size_t)n_tok * x_blocks * sizeof(float)); for (uint32_t t = 0; t < n_tok; t++) { quantize_q8_0_activation(norm + (uint64_t)t * expert_in_dim, xq8 + (uint64_t)t * x_blocks * 32u, xscale8 + (uint64_t)t * x_blocks, expert_in_dim); } float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); matvec_q8_0_batch_mid_ctx mid_ctx = { .mid = mid, .xq = xq8, .xscale = xscale8, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .pair_weight = pair_weight, .clamp = clamp, .in_dim = expert_in_dim, .out_dim = expert_out_dim, .blocks = x_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { ds4_die("Q8_0 batch expert tensor layout mismatch"); } } ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_0_batch_mid_worker, &mid_ctx); const uint64_t mid_blocks = down_in_dim / 32u; int8_t *midq8 = xmalloc((size_t)total_pairs * mid_blocks * 32u); float *midscale8 = xmalloc((size_t)total_pairs * mid_blocks * sizeof(float)); for (uint32_t p = 0; p < total_pairs; p++) { quantize_q8_0_activation(mid + (uint64_t)p * down_in_dim, midq8 + (uint64_t)p * mid_blocks * 32u, midscale8 + (uint64_t)p * mid_blocks, down_in_dim); } free(mid); matvec_q8_0_batch_accum_rows_ctx down_ctx = { .moe = moe, .midq = midq8, .midscale = midscale8, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .n_active = n_active, .n_tok = n_tok, .in_dim = down_in_dim, .out_dim = down_out_dim, .blocks = mid_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t in_dim, out_dim; down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, &in_dim, &out_dim, &down_ctx.row_bytes[e]); if (in_dim != down_in_dim || out_dim != down_out_dim) { ds4_die("Q8_0 batch down expert tensor layout mismatch"); } } ds4_parallel_for(down_out_dim, matvec_q8_0_batch_accum_rows_worker, &down_ctx); free(midscale8); free(midq8); free(xscale8); free(xq8); free(pair_ids); free(pairs); free(pair_weight); free(selected); (void)il; return; } if (routed_q8_k) { const uint64_t xq_blocks = expert_in_dim / QK_K; block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); for (uint32_t t = 0; t < n_tok; t++) { ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, xq + (uint64_t)t * xq_blocks, (int64_t)expert_in_dim); } float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); matvec_q8_k_batch_mid_ctx mid_ctx = { .mid = mid, .xq = xq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .pair_weight = pair_weight, .clamp = clamp, .in_dim = expert_in_dim, .out_dim = expert_out_dim, .xq_blocks = xq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { ds4_die("Q8_K batch expert tensor layout mismatch"); } } ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q8_k_batch_mid_worker, &mid_ctx); const uint64_t midq_blocks = down_in_dim / QK_K; block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); quantize_mid_pairs_ctx quant_ctx = { .mid = mid, .midq = midq, .down_in_dim = down_in_dim, .down_blocks = midq_blocks, }; ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); free(mid); matvec_q8_k_batch_accum_rows_ctx down_ctx = { .moe = moe, .midq = midq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .n_active = n_active, .n_tok = n_tok, .in_dim = down_in_dim, .out_dim = down_out_dim, .midq_blocks = midq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t in_dim, out_dim; down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, &in_dim, &out_dim, &down_ctx.row_bytes[e]); if (in_dim != down_in_dim || out_dim != down_out_dim) { ds4_die("Q8_K batch down expert tensor layout mismatch"); } } ds4_parallel_for(down_out_dim, matvec_q8_k_batch_accum_rows_worker, &down_ctx); free(midq); free(pair_ids); free(xq); free(pairs); free(pair_weight); free(selected); (void)il; return; } const uint64_t xq_blocks = expert_in_dim / QK_K; block_q8_K *xq = xmalloc((size_t)n_tok * xq_blocks * sizeof(xq[0])); for (uint32_t t = 0; t < n_tok; t++) { ds4_quantize_row_q8_K(norm + (uint64_t)t * expert_in_dim, xq + (uint64_t)t * xq_blocks, (int64_t)expert_in_dim); } float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0])); const uint32_t gate_type = layer->ffn_gate_exps->type; /* Build mid vectors: dispatch based on gate/up tensor type. */ if (gate_type == DS4_TENSOR_IQ2_XXS) { matvec_iq2_xxs_batch_mid_ctx mid_ctx = { .mid = mid, .xq = xq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .pair_weight = pair_weight, .clamp = clamp, .in_dim = expert_in_dim, .out_dim = expert_out_dim, .xq_blocks = xq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_iq2_xxs_batch_mid_worker, &mid_ctx); } else if (gate_type == DS4_TENSOR_Q2_K) { matvec_q2_k_batch_mid_ctx mid_ctx = { .mid = mid, .xq = xq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .pair_weight = pair_weight, .clamp = clamp, .in_dim = expert_in_dim, .out_dim = expert_out_dim, .xq_blocks = xq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q2_k_batch_mid_worker, &mid_ctx); } else if (gate_type == DS4_TENSOR_Q4_K) { matvec_q4_k_batch_mid_ctx mid_ctx = { .mid = mid, .xq = xq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .pair_weight = pair_weight, .clamp = clamp, .in_dim = expert_in_dim, .out_dim = expert_out_dim, .xq_blocks = xq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t gate_in_dim, gate_out_dim; uint64_t up_in_dim, up_out_dim; mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e, &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]); mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e, &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]); if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim || gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q4_k_batch_mid_worker, &mid_ctx); } else { ds4_die("unsupported gate/up expert tensor type for batch"); } const uint64_t midq_blocks = down_in_dim / QK_K; block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0])); quantize_mid_pairs_ctx quant_ctx = { .mid = mid, .midq = midq, .down_in_dim = down_in_dim, .down_blocks = midq_blocks, }; ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx); free(mid); /* Down projection: dispatch based on down tensor type. */ const uint32_t down_type = layer->ffn_down_exps->type; if (down_type == DS4_TENSOR_IQ2_XXS) { matvec_iq2_xxs_batch_accum_rows_ctx down_ctx = { .moe = moe, .midq = midq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .n_active = n_active, .n_tok = n_tok, .in_dim = down_in_dim, .out_dim = down_out_dim, .midq_blocks = midq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t in_dim, out_dim; down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, &in_dim, &out_dim, &down_ctx.row_bytes[e]); if (in_dim != down_in_dim || out_dim != down_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for(down_out_dim, matvec_iq2_xxs_batch_accum_rows_worker, &down_ctx); } else if (down_type == DS4_TENSOR_Q2_K) { matvec_q2_k_batch_accum_rows_ctx down_ctx = { .moe = moe, .midq = midq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .n_active = n_active, .n_tok = n_tok, .in_dim = down_in_dim, .out_dim = down_out_dim, .midq_blocks = midq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t in_dim, out_dim; down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, &in_dim, &out_dim, &down_ctx.row_bytes[e]); if (in_dim != down_in_dim || out_dim != down_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for(down_out_dim, matvec_q2_k_batch_accum_rows_worker, &down_ctx); } else if (down_type == DS4_TENSOR_Q4_K) { matvec_q4_k_batch_accum_rows_ctx down_ctx = { .moe = moe, .midq = midq, .pairs = pairs, .pair_ids = pair_ids, .expert_offset = counts, .active_expert = active_expert, .n_active = n_active, .n_tok = n_tok, .in_dim = down_in_dim, .out_dim = down_out_dim, .midq_blocks = midq_blocks, }; for (uint32_t ai = 0; ai < n_active; ai++) { const uint32_t e = active_expert[ai]; uint64_t in_dim, out_dim; down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e, &in_dim, &out_dim, &down_ctx.row_bytes[e]); if (in_dim != down_in_dim || out_dim != down_out_dim) { ds4_die("batch expert tensor layout mismatch"); } } ds4_parallel_for(down_out_dim, matvec_q4_k_batch_accum_rows_worker, &down_ctx); } else { ds4_die("unsupported down expert tensor type for batch"); } free(midq); free(pair_ids); free(xq); free(pairs); free(pair_weight); free(selected); (void)il; } static void print_vec_stats(const char *name, const float *x, uint64_t n); /* Full FFN sublayer for one token: HC pre, RMSNorm, routed MoE, shared expert, * sum, and HC post. */ static void layer_ffn_one( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, uint32_t il, int token, const float * steering_dirs, float steering_scale, bool trace) { const uint32_t n_hc = DS4_N_HC; const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; const double t_start = profile ? now_sec() : 0.0; double t_hc = 0.0; double t_norm = 0.0; double t_routed = 0.0; double t_shared = 0.0; double t_post = 0.0; float *ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_cur[0])); float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); float *moe = xmalloc((size_t)DS4_N_EMBD * sizeof(moe[0])); float *shared = xmalloc((size_t)DS4_N_EMBD * sizeof(shared[0])); float *ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(ffn_out[0])); float post[4]; float comb[16]; double t0 = profile ? now_sec() : 0.0; hc_pre_from_state_one(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, inp_hc, ffn_cur, post, comb); if (profile) t_hc = now_sec() - t0; if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u ffn_cur", il); print_vec_stats(name, ffn_cur, DS4_N_EMBD); } t0 = profile ? now_sec() : 0.0; const float *ffn_norm = tensor_data(model, layer->ffn_norm); rms_norm_weight(norm, ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); if (profile) t_norm = now_sec() - t0; if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u ffn_norm", il); print_vec_stats(name, norm, DS4_N_EMBD); } t0 = profile ? now_sec() : 0.0; layer_routed_moe_one(moe, model, layer, norm, il, token, DS4_SWIGLU_CLAMP_EXP, trace); if (profile) t_routed = now_sec() - t0; if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u routed_moe", il); print_vec_stats(name, moe, DS4_N_EMBD); } t0 = profile ? now_sec() : 0.0; layer_shared_ffn_one(shared, model, layer, norm); if (profile) t_shared = now_sec() - t0; if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u shared_ffn", il); print_vec_stats(name, shared, DS4_N_EMBD); } t0 = profile ? now_sec() : 0.0; for (uint32_t i = 0; i < DS4_N_EMBD; i++) { ffn_out[i] = moe[i] + shared[i]; } cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, 1, steering_scale); if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u ffn_out", il); print_vec_stats(name, ffn_out, DS4_N_EMBD); } hc_post_one(out_hc, ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); if (profile) t_post = now_sec() - t0; if (trace) { char name[64]; snprintf(name, sizeof(name), "blk.%u ffn_post_hc", il); print_vec_stats(name, out_hc, (uint64_t)n_hc * DS4_N_EMBD); } if (profile) { fprintf(stderr, "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", il, t_hc * 1000.0, t_norm * 1000.0, t_routed * 1000.0, t_shared * 1000.0, t_post * 1000.0, (now_sec() - t_start) * 1000.0); } free(ffn_out); free(shared); free(moe); free(norm); free(ffn_cur); } /* Allocation-free decode FFN using the persistent CPU scratch buffers. */ static void layer_ffn_one_decode_scratch( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, uint32_t il, int token, const float * steering_dirs, float steering_scale, ds4_cpu_decode_scratch * scratch) { const uint32_t n_hc = DS4_N_HC; const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; const double t_start = profile ? now_sec() : 0.0; double t_hc = 0.0; double t_norm = 0.0; double t_routed = 0.0; double t_shared = 0.0; double t_post = 0.0; float post[4]; float comb[16]; double t0 = profile ? now_sec() : 0.0; hc_pre_from_state_one_scratch(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, inp_hc, scratch->ffn_cur, post, comb, scratch->hc_flat, false); if (profile) t_hc = now_sec() - t0; t0 = profile ? now_sec() : 0.0; const float *ffn_norm = tensor_data(model, layer->ffn_norm); rms_norm_weight(scratch->ffn_norm, scratch->ffn_cur, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); if (profile) t_norm = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_routed_moe_one_prealloc(scratch->ffn_moe, model, layer, scratch->ffn_norm, il, token, DS4_SWIGLU_CLAMP_EXP, scratch->routed_mid_all, scratch->routed_xq, scratch->routed_midq, scratch->routed_q8_xq, scratch->routed_q8_xscale, scratch->routed_q8_midq, scratch->routed_q8_midscale); if (profile) t_routed = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_shared_ffn_one_decode_scratch(scratch->ffn_shared, model, layer, scratch->ffn_norm, scratch); if (profile) t_shared = now_sec() - t0; t0 = profile ? now_sec() : 0.0; for (uint32_t i = 0; i < DS4_N_EMBD; i++) { scratch->ffn_out[i] = scratch->ffn_moe[i] + scratch->ffn_shared[i]; } cpu_directional_steering_project_rows(scratch->ffn_out, steering_dirs, il, 1, steering_scale); hc_post_one(out_hc, scratch->ffn_out, inp_hc, post, comb, DS4_N_EMBD, n_hc); if (profile) t_post = now_sec() - t0; if (profile) { fprintf(stderr, "ds4: decode detail layer %u ffn hc=%.3f norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f ms\n", il, t_hc * 1000.0, t_norm * 1000.0, t_routed * 1000.0, t_shared * 1000.0, t_post * 1000.0, (now_sec() - t_start) * 1000.0); } } static void layer_ffn_batch( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, const int * token_ids, uint32_t n_tok, uint32_t il, const float * steering_dirs, float steering_scale) { if (n_tok == 0) return; const uint32_t n_hc = DS4_N_HC; const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); const float *ffn_norm = tensor_data(model, layer->ffn_norm); for (uint32_t t = 0; t < n_tok; t++) { hc_pre_from_state_one(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, inp_hc + (uint64_t)t * hc_dim, ffn_cur + (uint64_t)t * DS4_N_EMBD, post + (uint64_t)t * n_hc, comb + (uint64_t)t * n_hc * n_hc); rms_norm_weight(norm + (uint64_t)t * DS4_N_EMBD, ffn_cur + (uint64_t)t * DS4_N_EMBD, ffn_norm, DS4_N_EMBD, DS4_RMS_EPS); } layer_routed_moe_batch(moe, model, layer, norm, token_ids, n_tok, il, DS4_SWIGLU_CLAMP_EXP); layer_shared_ffn_batch(shared, model, layer, norm, n_tok); if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { ffn_out[i] = moe[i] + shared[i]; } cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); hc_post_batch(out_hc, ffn_out, inp_hc, post, comb, n_tok, DS4_N_EMBD, n_hc); free(ffn_out); } else { hc_post_sum_batch(out_hc, moe, shared, inp_hc, post, comb, n_tok, DS4_N_EMBD, n_hc); } free(comb); free(post); free(shared); free(moe); free(norm); free(ffn_cur); } typedef struct { float *moe; const ds4_model *model; const ds4_layer_weights *layer; const float *norm; const int *token_ids; uint64_t expert_in_dim; uint64_t down_in_dim; uint32_t il; bool routed_q8_0; } routed_moe_tokens_ctx; static void routed_moe_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { routed_moe_tokens_ctx *ctx = vctx; const uint64_t q8_x_blocks = ctx->expert_in_dim / 32u; const uint64_t q8_mid_blocks = ctx->down_in_dim / 32u; float *routed_mid = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); block_q8_K *routed_xq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)(ctx->expert_in_dim / QK_K) * sizeof(routed_xq[0])); block_q8_K *routed_midq = ctx->routed_q8_0 ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (ctx->down_in_dim / QK_K) * sizeof(routed_midq[0])); int8_t *routed_q8_xq = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * 32u) : NULL; float *routed_q8_xscale = ctx->routed_q8_0 ? xmalloc((size_t)q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; int8_t *routed_q8_midq = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * 32u) : NULL; float *routed_q8_midscale = ctx->routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; for (uint64_t t = t0; t < t1; t++) { layer_routed_moe_one_prealloc(ctx->moe + t * DS4_N_EMBD, ctx->model, ctx->layer, ctx->norm + t * DS4_N_EMBD, ctx->il, ctx->token_ids[t], DS4_SWIGLU_CLAMP_EXP, routed_mid, routed_xq, routed_midq, routed_q8_xq, routed_q8_xscale, routed_q8_midq, routed_q8_midscale); } free(routed_q8_midscale); free(routed_q8_midq); free(routed_q8_xscale); free(routed_q8_xq); free(routed_midq); free(routed_xq); free(routed_mid); } static void layer_routed_moe_tokens_parallel( float * moe, const ds4_model * model, const ds4_layer_weights * layer, const float * norm, const int * token_ids, uint32_t n_tok, uint32_t il) { const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const bool routed_q8_k = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; if (routed_q8_k) { if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { ds4_die("Q8_K expert input has an unexpected layout"); } } routed_moe_tokens_ctx ctx = { .moe = moe, .model = model, .layer = layer, .norm = norm, .token_ids = token_ids, .expert_in_dim = expert_in_dim, .down_in_dim = down_in_dim, .il = il, .routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0, }; ds4_parallel_for_min_rows(n_tok, routed_moe_tokens_worker, &ctx, 1); } /* Default prefill FFN path. HC and shared expert are batched, while routed * experts can run either token-parallel or expert-grouped depending on size. */ static void layer_ffn_shared_batch( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, const int * token_ids, uint32_t n_tok, uint32_t il, const float * steering_dirs, float steering_scale) { const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; const double t_start = profile ? now_sec() : 0.0; double t_hc_norm = 0.0; double t_routed = 0.0; double t_shared = 0.0; double t_post = 0.0; const uint32_t n_hc = DS4_N_HC; float *ffn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_cur[0])); float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); float *moe = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(moe[0])); float *shared = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(shared[0])); float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const bool routed_q8_k = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_K && layer->ffn_up_exps->type == DS4_TENSOR_Q8_K && layer->ffn_down_exps->type == DS4_TENSOR_Q8_K; if (routed_q8_k) { if (expert_in_dim % QK_K != 0) ds4_die("Q8_K expert input is not QK_K aligned"); if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) { ds4_die("Q8_K expert input has an unexpected layout"); } } const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; const bool routed_token_parallel = getenv("DS4_ROUTED_TOKEN_PARALLEL") != NULL || (getenv("DS4_NO_ROUTED_TOKEN_PARALLEL") == NULL && n_tok >= 64); float *routed_mid = routed_token_parallel ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(routed_mid[0])); block_q8_K *routed_xq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(routed_xq[0])); block_q8_K *routed_midq = (routed_token_parallel || routed_q8_0) ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(routed_midq[0])); int8_t *routed_q8_xq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; float *routed_q8_xscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; int8_t *routed_q8_midq = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; float *routed_q8_midscale = (!routed_token_parallel && routed_q8_0) ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; double t0 = profile ? now_sec() : 0.0; hc_pre_norm_batch(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, layer->ffn_norm, inp_hc, NULL, ffn_cur, norm, post, comb, n_tok); if (profile) t_hc_norm = now_sec() - t0; t0 = profile ? now_sec() : 0.0; if (routed_token_parallel) { layer_routed_moe_tokens_parallel(moe, model, layer, norm, token_ids, n_tok, il); } else { for (uint32_t t = 0; t < n_tok; t++) { layer_routed_moe_one_prealloc(moe + (uint64_t)t * DS4_N_EMBD, model, layer, norm + (uint64_t)t * DS4_N_EMBD, il, token_ids[t], DS4_SWIGLU_CLAMP_EXP, routed_mid, routed_xq, routed_midq, routed_q8_xq, routed_q8_xscale, routed_q8_midq, routed_q8_midscale); } } if (profile) t_routed = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_shared_ffn_batch(shared, model, layer, norm, n_tok); if (profile) t_shared = now_sec() - t0; t0 = profile ? now_sec() : 0.0; if (cpu_directional_steering_enabled(steering_dirs, steering_scale)) { float *ffn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(ffn_out[0])); for (uint64_t i = 0; i < (uint64_t)n_tok * DS4_N_EMBD; i++) { ffn_out[i] = moe[i] + shared[i]; } cpu_directional_steering_project_rows(ffn_out, steering_dirs, il, n_tok, steering_scale); hc_post_batch(out_hc, ffn_out, inp_hc, post, comb, n_tok, DS4_N_EMBD, n_hc); free(ffn_out); } else { hc_post_sum_batch(out_hc, moe, shared, inp_hc, post, comb, n_tok, DS4_N_EMBD, n_hc); } if (profile) t_post = now_sec() - t0; if (profile) { fprintf(stderr, "ds4: prefill detail layer %u ffn hc_norm=%.3f routed=%.3f shared=%.3f post=%.3f total=%.3f\n", il, t_hc_norm, t_routed, t_shared, t_post, now_sec() - t_start); } free(comb); free(post); free(routed_q8_midscale); free(routed_q8_midq); free(routed_q8_xscale); free(routed_q8_xq); free(routed_midq); free(routed_xq); free(routed_mid); free(shared); free(moe); free(norm); free(ffn_cur); } typedef struct { float *out_hc; const ds4_model *model; const ds4_layer_weights *layer; const float *inp_hc; const int *token_ids; const float *steering_dirs; float steering_scale; uint64_t hc_dim; uint32_t il; } layer_ffn_tokens_ctx; static void layer_ffn_tokens_worker(void *vctx, uint64_t t0, uint64_t t1) { layer_ffn_tokens_ctx *ctx = vctx; for (uint64_t t = t0; t < t1; t++) { layer_ffn_one(ctx->out_hc + t * ctx->hc_dim, ctx->model, ctx->layer, ctx->inp_hc + t * ctx->hc_dim, ctx->il, ctx->token_ids[t], ctx->steering_dirs, ctx->steering_scale, false); } } static void layer_ffn_tokens_parallel( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, const int * token_ids, uint32_t n_tok, uint32_t il, const float * steering_dirs, float steering_scale) { layer_ffn_tokens_ctx ctx = { .out_hc = out_hc, .model = model, .layer = layer, .inp_hc = inp_hc, .token_ids = token_ids, .steering_dirs = steering_dirs, .steering_scale = steering_scale, .hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD, .il = il, }; ds4_parallel_for(n_tok, layer_ffn_tokens_worker, &ctx); } static void output_logits_one( float * logits, const ds4_model * model, const ds4_weights * weights, const float * inp_hc); /* ========================================================================= * KV Cache, Compressors, and CPU Layer Execution. * ========================================================================= * * The CPU path is the correctness reference. It maintains raw SWA KV rows, * optional compressed KV rows, the indexer mask for ratio-4 layers, and a * reusable decode scratch arena so token generation does not allocate in the * hot loop. */ typedef struct { float *raw_kv; uint32_t n_raw; uint32_t cap_raw; uint32_t compress_ratio; uint32_t comp_cap; uint32_t n_comp; float *attn_comp_kv; float *attn_state_kv; float *attn_state_score; uint32_t n_index_comp; float *index_comp_kv; float *index_state_kv; float *index_state_score; } ds4_layer_cache; typedef struct { ds4_layer_cache layer[DS4_MAX_LAYER]; uint32_t head_dim; } ds4_kv_cache; static uint32_t ds4_default_raw_cap(uint32_t ctx_size) { uint32_t raw_cap = DS4_N_SWA; if (raw_cap > ctx_size) raw_cap = ctx_size; if (raw_cap == 0) raw_cap = 1; return raw_cap; } #define DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK 2048u static uint32_t ds4_effective_prefill_chunk(bool cuda_tensor_parallel, uint32_t requested_chunk) { if (requested_chunk != 0) return requested_chunk; return cuda_tensor_parallel ? DS4_CUDA_TP_DEFAULT_PREFILL_CHUNK : 0; } static uint32_t ds4_prefill_cap_for_prompt(int prompt_len, uint32_t requested_chunk) { if (prompt_len <= 0) return 1; uint32_t cap = (uint32_t)prompt_len; if (requested_chunk != 0) { cap = requested_chunk; } else { const char *env = getenv("DS4_METAL_PREFILL_CHUNK"); if (env && env[0]) { char *endp = NULL; const long v = strtol(env, &endp, 10); if (endp != env) { if (v <= 0) return cap; cap = (uint32_t)v; } } else if (prompt_len > 4096) { cap = DS4_MODEL_VARIANT == DS4_VARIANT_PRO ? 8192u : 4096u; } } if (cap == 0) cap = 1; if (cap > (uint32_t)prompt_len) cap = (uint32_t)prompt_len; return cap; } /* Allocate all CPU decode temporaries once. This keeps generation deterministic * from the VM's point of view and makes accidental hot-loop malloc visible. */ static void cpu_decode_scratch_init(ds4_cpu_decode_scratch *scratch, uint32_t ctx_size) { memset(scratch, 0, sizeof(*scratch)); if (ctx_size == 0) ctx_size = 1; const uint32_t raw_cap = ds4_default_raw_cap(ctx_size); const uint32_t comp_cap = ctx_size / 4 + 2; const uint32_t attn_score_cap = raw_cap + comp_cap; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t q8_cap = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t q8_blocks = (q8_cap + 31u) / 32u; if ((DS4_N_EMBD % 32u) != 0 || (DS4_N_FF_EXP % 32u) != 0) { ds4_die("Q8_0 routed decode scratch dimensions are not QK8_0 aligned"); } const uint64_t routed_q8_x_blocks = DS4_N_EMBD / 32u; const uint64_t routed_q8_mid_blocks = DS4_N_FF_EXP / 32u; /* * The CPU decode path used to malloc/free dozens of medium-sized buffers * for every layer of every generated token. On macOS this can drive the VM * system through repeated map/unmap bookkeeping while the huge model mmap is * also being streamed, and we have observed kernel panics in VM accounting. * Keep decode scratch resident for the whole generation instead. */ scratch->ctx_size = ctx_size; scratch->comp_cap = comp_cap; scratch->attn_score_cap = attn_score_cap; scratch->q8_cap = (uint32_t)q8_cap; scratch->plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->cur = xmalloc((size_t)hc_dim * sizeof(float)); scratch->next = xmalloc((size_t)hc_dim * sizeof(float)); scratch->attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->attn_residual = xmalloc((size_t)hc_dim * sizeof(float)); scratch->q = xmalloc((size_t)q_dim * sizeof(float)); scratch->qr = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); scratch->qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(float)); scratch->kv_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); scratch->kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); scratch->heads = xmalloc((size_t)q_dim * sizeof(float)); scratch->attn_low = xmalloc((size_t)DS4_N_OUT_GROUP * DS4_N_LORA_O * sizeof(float)); scratch->attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); scratch->attn_score = xmalloc((size_t)attn_score_cap * sizeof(float)); scratch->comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); scratch->index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); scratch->comp_kv_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); scratch->comp_sc_cur = xmalloc((size_t)2u * DS4_N_HEAD_DIM * sizeof(float)); scratch->comp_pooled = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); scratch->index_allowed = xmalloc((size_t)comp_cap * sizeof(bool)); scratch->index_q = xmalloc((size_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); scratch->index_weights = xmalloc((size_t)DS4_N_INDEXER_HEAD * sizeof(float)); scratch->index_scores = xmalloc((size_t)comp_cap * sizeof(float)); scratch->ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->ffn_moe = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->ffn_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->shared_gate = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); scratch->shared_up = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); scratch->shared_mid = xmalloc((size_t)DS4_N_FF_EXP * sizeof(float)); scratch->routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); scratch->routed_xq = xmalloc((size_t)(DS4_N_EMBD / QK_K) * sizeof(block_q8_K)); scratch->routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (DS4_N_FF_EXP / QK_K) * sizeof(block_q8_K)); scratch->routed_q8_xq = xmalloc((size_t)routed_q8_x_blocks * 32u); scratch->routed_q8_xscale = xmalloc((size_t)routed_q8_x_blocks * sizeof(float)); scratch->routed_q8_midq = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u); scratch->routed_q8_midscale = xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(float)); scratch->q8_xq = xmalloc((size_t)q8_blocks * 32u); scratch->q8_xscale = xmalloc((size_t)q8_blocks * sizeof(float)); scratch->hc_flat = xmalloc((size_t)hc_dim * sizeof(float)); scratch->output_flat = xmalloc((size_t)hc_dim * sizeof(float)); scratch->output_pre = xmalloc((size_t)DS4_N_HC * sizeof(float)); scratch->output_weights = xmalloc((size_t)DS4_N_HC * sizeof(float)); scratch->output_embd = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); scratch->output_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); } static void cpu_decode_scratch_free(ds4_cpu_decode_scratch *scratch) { if (!scratch) return; free(scratch->output_norm); free(scratch->output_embd); free(scratch->output_weights); free(scratch->output_pre); free(scratch->output_flat); free(scratch->hc_flat); free(scratch->q8_xscale); free(scratch->q8_xq); free(scratch->routed_q8_midscale); free(scratch->routed_q8_midq); free(scratch->routed_q8_xscale); free(scratch->routed_q8_xq); free(scratch->routed_midq); free(scratch->routed_xq); free(scratch->routed_mid_all); free(scratch->shared_mid); free(scratch->shared_up); free(scratch->shared_gate); free(scratch->ffn_out); free(scratch->ffn_shared); free(scratch->ffn_moe); free(scratch->ffn_norm); free(scratch->ffn_cur); free(scratch->index_scores); free(scratch->index_weights); free(scratch->index_q); free(scratch->index_allowed); free(scratch->comp_pooled); free(scratch->comp_sc_cur); free(scratch->comp_kv_cur); free(scratch->index_comp); free(scratch->comp); free(scratch->attn_score); free(scratch->after_attn_hc); free(scratch->attn_out); free(scratch->attn_low); free(scratch->heads); free(scratch->kv); free(scratch->kv_raw); free(scratch->qr_norm); free(scratch->qr); free(scratch->q); free(scratch->attn_residual); free(scratch->attn_norm); free(scratch->attn_cur); free(scratch->next); free(scratch->cur); free(scratch->plain); memset(scratch, 0, sizeof(*scratch)); } /* Allocate per-layer KV state: a raw sliding window for all layers, plus * compressed attention/indexer caches for layers whose ratio is nonzero. */ static void kv_cache_init(ds4_kv_cache *cache, uint32_t ctx_size, uint32_t raw_cap) { memset(cache, 0, sizeof(*cache)); if (raw_cap == 0) raw_cap = ds4_default_raw_cap(ctx_size); if (raw_cap > ctx_size) raw_cap = ctx_size; if (raw_cap == 0) raw_cap = 1; cache->head_dim = DS4_N_HEAD_DIM; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); cache->layer[il].cap_raw = raw_cap; cache->layer[il].raw_kv = xmalloc_zeroed((size_t)raw_cap * DS4_N_HEAD_DIM, sizeof(float)); cache->layer[il].compress_ratio = ratio; if (ratio != 0) { const uint32_t coff = ratio == 4 ? 2u : 1u; const uint32_t comp_cap = ctx_size / ratio + 2; const uint32_t attn_width = coff * DS4_N_HEAD_DIM; const uint32_t attn_rows = coff * ratio; cache->layer[il].comp_cap = comp_cap; cache->layer[il].attn_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_HEAD_DIM, sizeof(float)); cache->layer[il].attn_state_kv = xmalloc_zeroed((size_t)attn_width * attn_rows, sizeof(float)); cache->layer[il].attn_state_score = xmalloc((size_t)attn_width * attn_rows * sizeof(float)); for (uint64_t i = 0; i < (uint64_t)attn_width * attn_rows; i++) { cache->layer[il].attn_state_score[i] = DS4_NEG_INF; } if (ratio == 4) { const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; const uint32_t index_rows = coff * ratio; cache->layer[il].index_comp_kv = xmalloc_zeroed((size_t)comp_cap * DS4_N_INDEXER_HEAD_DIM, sizeof(float)); cache->layer[il].index_state_kv = xmalloc_zeroed((size_t)index_width * index_rows, sizeof(float)); cache->layer[il].index_state_score = xmalloc((size_t)index_width * index_rows * sizeof(float)); for (uint64_t i = 0; i < (uint64_t)index_width * index_rows; i++) { cache->layer[il].index_state_score[i] = DS4_NEG_INF; } } } } } static void kv_cache_free(ds4_kv_cache *cache) { if (!cache) return; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { free(cache->layer[il].raw_kv); free(cache->layer[il].attn_comp_kv); free(cache->layer[il].attn_state_kv); free(cache->layer[il].attn_state_score); free(cache->layer[il].index_comp_kv); free(cache->layer[il].index_state_kv); free(cache->layer[il].index_state_score); } memset(cache, 0, sizeof(*cache)); } /* Append to the raw SWA cache. Once full, it slides by one row. */ static void kv_cache_push_raw(ds4_layer_cache *cache, const float *kv) { if (cache->n_raw < cache->cap_raw) { float *dst = cache->raw_kv + (uint64_t)cache->n_raw * DS4_N_HEAD_DIM; for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); cache->n_raw++; return; } memmove(cache->raw_kv, cache->raw_kv + DS4_N_HEAD_DIM, (size_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM * sizeof(cache->raw_kv[0])); float *dst = cache->raw_kv + (uint64_t)(cache->cap_raw - 1) * DS4_N_HEAD_DIM; for (uint32_t i = 0; i < DS4_N_HEAD_DIM; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); } static void kv_cache_push_comp(float *rows, uint32_t *n_rows, uint32_t cap_rows, uint32_t row_dim, const float *kv) { if (*n_rows >= cap_rows) ds4_die("compressed KV cache capacity exceeded"); float *dst = rows + (uint64_t)(*n_rows) * row_dim; for (uint32_t i = 0; i < row_dim; i++) dst[i] = f16_to_f32(f32_to_f16(kv[i])); (*n_rows)++; } /* After prefill, clear unused compressor state rows so decode starts from the * same partial-window state the streaming path would have produced. */ static void compressor_finish_prefill_state_cpu( float * state_kv, float * state_score, uint32_t head_dim, uint32_t compress_ratio, uint32_t n_tokens) { if (!state_kv || !state_score || head_dim == 0 || compress_ratio == 0) return; const uint32_t coff = compress_ratio == 4 ? 2u : 1u; const uint32_t width = coff * head_dim; const uint32_t rem = n_tokens % compress_ratio; const uint32_t clear_start = compress_ratio == 4 ? compress_ratio + rem : rem; const uint32_t clear_end = compress_ratio == 4 ? 2u * compress_ratio : compress_ratio; for (uint32_t row = clear_start; row < clear_end; row++) { float *kv = state_kv + (uint64_t)row * width; float *score = state_score + (uint64_t)row * width; memset(kv, 0, (size_t)width * sizeof(kv[0])); for (uint32_t i = 0; i < width; i++) score[i] = DS4_NEG_INF; } } static void kv_cache_finish_prefill_states(ds4_kv_cache *cache, uint32_t n_tokens) { for (uint32_t il = 0; il < DS4_N_LAYER; il++) { ds4_layer_cache *layer = &cache->layer[il]; const uint32_t ratio = layer->compress_ratio; if (ratio == 0) continue; compressor_finish_prefill_state_cpu(layer->attn_state_kv, layer->attn_state_score, DS4_N_HEAD_DIM, ratio, n_tokens); if (ratio == 4) { compressor_finish_prefill_state_cpu(layer->index_state_kv, layer->index_state_score, DS4_N_INDEXER_HEAD_DIM, ratio, n_tokens); } } } /* Pool the current compression window with a softmax over per-dimension scores. * Ratio-4 layers keep two lanes: attention compression and indexer compression. */ static void compressor_pool_decode_state( float * out, float * state_kv, float * state_score, uint32_t head_dim, uint32_t compress_ratio) { const uint32_t coff = compress_ratio == 4 ? 2u : 1u; const uint32_t width = coff * head_dim; for (uint32_t j = 0; j < head_dim; j++) { float max_score = DS4_NEG_INF; if (compress_ratio == 4) { for (uint32_t r = 0; r < compress_ratio; r++) { const float sp = state_score[(uint64_t)r * width + j]; const float sc = state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j]; if (sp > max_score) max_score = sp; if (sc > max_score) max_score = sc; } } else { for (uint32_t r = 0; r < compress_ratio; r++) { const float s = state_score[(uint64_t)r * width + j]; if (s > max_score) max_score = s; } } if (max_score <= DS4_NEG_INF * 0.5f) { out[j] = 0.0f; continue; } float denom = 0.0f; float sum = 0.0f; if (compress_ratio == 4) { for (uint32_t r = 0; r < compress_ratio; r++) { const float wp = expf(state_score[(uint64_t)r * width + j] - max_score); const float wc = expf(state_score[(uint64_t)(compress_ratio + r) * width + head_dim + j] - max_score); denom += wp + wc; sum += wp * state_kv[(uint64_t)r * width + j]; sum += wc * state_kv[(uint64_t)(compress_ratio + r) * width + head_dim + j]; } } else { for (uint32_t r = 0; r < compress_ratio; r++) { const float w = expf(state_score[(uint64_t)r * width + j] - max_score); denom += w; sum += w * state_kv[(uint64_t)r * width + j]; } } out[j] = denom > 0.0f ? sum / denom : 0.0f; } } /* Streaming compressor update for one token. It projects kv/score rows, * updates the rolling state, and emits a compressed KV row on ratio boundaries. */ static bool compressor_decode_one( float * out_comp, const ds4_model * model, const ds4_tensor * wkv, const ds4_tensor * wgate, const ds4_tensor * ape, const ds4_tensor * norm, const float * x, float * state_kv, float * state_score, uint32_t head_dim, uint32_t compress_ratio, uint32_t il, uint32_t pos) { const uint32_t coff = compress_ratio == 4 ? 2u : 1u; const uint32_t width = coff * head_dim; const uint32_t pos_mod = pos % compress_ratio; const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; const bool should_compress = ((pos + 1) % compress_ratio) == 0; float *kv_cur = xmalloc((size_t)width * sizeof(kv_cur[0])); float *sc_cur = xmalloc((size_t)width * sizeof(sc_cur[0])); if (wkv->type == 8 && wgate->type == 8 && wkv->ndim == 2 && wgate->ndim == 2 && wkv->dim[0] == wgate->dim[0]) { const uint64_t in_dim = wkv->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; int8_t *xq = xmalloc((size_t)blocks * 32); float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); quantize_q8_0_activation(x, xq, xscale, in_dim); matvec_q8_0_pair_prequant(kv_cur, sc_cur, model, wkv, wgate, xq, xscale); free(xscale); free(xq); } else { matvec_any(kv_cur, model, wkv, x); matvec_any(sc_cur, model, wgate, x); } for (uint32_t j = 0; j < width; j++) { sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); } memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); free(sc_cur); free(kv_cur); if (!should_compress) { return false; } float *pooled = xmalloc((size_t)head_dim * sizeof(pooled[0])); compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); double ss = 0.0; for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); for (uint32_t i = 0; i < head_dim; i++) { out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); } const uint32_t comp_pos = pos + 1 - compress_ratio; rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); if (head_dim == DS4_N_HEAD_DIM) { dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); } if (compress_ratio == 4) { for (uint32_t r = 0; r < compress_ratio; r++) { memcpy(state_kv + (uint64_t)r * width, state_kv + (uint64_t)(compress_ratio + r) * width, (size_t)width * sizeof(state_kv[0])); memcpy(state_score + (uint64_t)r * width, state_score + (uint64_t)(compress_ratio + r) * width, (size_t)width * sizeof(state_score[0])); } for (uint32_t r = 0; r < compress_ratio; r++) { memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, state_kv + (uint64_t)r * width, (size_t)width * sizeof(state_kv[0])); memcpy(state_score + (uint64_t)(compress_ratio + r) * width, state_score + (uint64_t)r * width, (size_t)width * sizeof(state_score[0])); } } free(pooled); return true; } static bool compressor_decode_one_decode_scratch( float * out_comp, const ds4_model * model, const ds4_tensor * wkv, const ds4_tensor * wgate, const ds4_tensor * ape, const ds4_tensor * norm, const float * x, float * state_kv, float * state_score, uint32_t head_dim, uint32_t compress_ratio, uint32_t il, uint32_t pos, ds4_cpu_decode_scratch * scratch) { const uint32_t coff = compress_ratio == 4 ? 2u : 1u; const uint32_t width = coff * head_dim; const uint32_t pos_mod = pos % compress_ratio; const uint32_t row = compress_ratio == 4 ? compress_ratio + pos_mod : pos_mod; const bool should_compress = ((pos + 1) % compress_ratio) == 0; if (width > 2u * DS4_N_HEAD_DIM) ds4_die("compressor scratch width is outside the fixed model layout"); float *kv_cur = scratch->comp_kv_cur; float *sc_cur = scratch->comp_sc_cur; if (wkv->type == 8 && wgate->type == 8 && wkv->ndim == 2 && wgate->ndim == 2 && wkv->dim[0] == wgate->dim[0]) { matvec_q8_0_pair_decode_scratch(kv_cur, sc_cur, model, wkv, wgate, x, scratch); } else { matvec_any_decode_scratch(kv_cur, model, wkv, x, scratch); matvec_any_decode_scratch(sc_cur, model, wgate, x, scratch); } for (uint32_t j = 0; j < width; j++) { sc_cur[j] += tensor_2d_value(model, ape, j, pos_mod); } memcpy(state_kv + (uint64_t)row * width, kv_cur, (size_t)width * sizeof(kv_cur[0])); memcpy(state_score + (uint64_t)row * width, sc_cur, (size_t)width * sizeof(sc_cur[0])); if (!should_compress) { return false; } float *pooled = scratch->comp_pooled; compressor_pool_decode_state(pooled, state_kv, state_score, head_dim, compress_ratio); double ss = 0.0; for (uint32_t i = 0; i < head_dim; i++) ss += (double)pooled[i] * pooled[i]; const float rms = 1.0f / sqrtf((float)(ss / (double)head_dim) + DS4_RMS_EPS); for (uint32_t i = 0; i < head_dim; i++) { out_comp[i] = pooled[i] * rms * tensor_1d_value(model, norm, i); } const uint32_t comp_pos = pos + 1 - compress_ratio; rope_tail_layer_inplace(out_comp, 1, head_dim, DS4_N_ROT, comp_pos, il, false); if (head_dim == DS4_N_HEAD_DIM) { dsv4_fp8_kv_quantize_row_inplace_cpu(out_comp, head_dim, DS4_N_ROT); } else if (head_dim == DS4_N_INDEXER_HEAD_DIM) { dsv4_indexer_qat_row_inplace_cpu(out_comp, head_dim); } if (compress_ratio == 4) { for (uint32_t r = 0; r < compress_ratio; r++) { memcpy(state_kv + (uint64_t)r * width, state_kv + (uint64_t)(compress_ratio + r) * width, (size_t)width * sizeof(state_kv[0])); memcpy(state_score + (uint64_t)r * width, state_score + (uint64_t)(compress_ratio + r) * width, (size_t)width * sizeof(state_score[0])); } for (uint32_t r = 0; r < compress_ratio; r++) { memcpy(state_kv + (uint64_t)(compress_ratio + r) * width, state_kv + (uint64_t)r * width, (size_t)width * sizeof(state_kv[0])); memcpy(state_score + (uint64_t)(compress_ratio + r) * width, state_score + (uint64_t)r * width, (size_t)width * sizeof(state_score[0])); } } return true; } /* Attention over raw SWA rows plus optional compressed rows. Ratio-4 layers * pass an indexer mask to hide compressed rows not selected for this token. */ static void layer_attention_mixed_one( float * out_heads, const ds4_model * model, const ds4_layer_weights * layer, const float * q, const float * raw_kv, uint32_t n_raw, const float * comp_kv, uint32_t n_comp, const bool * comp_allowed) { const float *sinks = tensor_data(model, layer->attn_sinks); const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); const uint32_t n_total = n_raw + n_comp; float score_stack[512]; float *score = n_total <= 512 ? score_stack : xmalloc((size_t)n_total * sizeof(score[0])); for (uint32_t h = 0; h < DS4_N_HEAD; h++) { const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; float max_score = sinks[h]; uint32_t idx = 0; for (uint32_t r = 0; r < n_raw; r++, idx++) { const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[idx] > max_score) max_score = score[idx]; } for (uint32_t r = 0; r < n_comp; r++, idx++) { if (comp_allowed && !comp_allowed[r]) { score[idx] = DS4_NEG_INF; continue; } const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[idx] > max_score) max_score = score[idx]; } float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); float denom = expf(sinks[h] - max_score); idx = 0; for (uint32_t r = 0; r < n_raw; r++, idx++) { const float weight = expf(score[idx] - max_score); const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } for (uint32_t r = 0; r < n_comp; r++, idx++) { if (score[idx] <= DS4_NEG_INF * 0.5f) continue; const float weight = expf(score[idx] - max_score); const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } const float inv = 1.0f / denom; scale_f32(oh, inv, DS4_N_HEAD_DIM); } if (score != score_stack) free(score); } static void layer_attention_mixed_one_decode_scratch( float * out_heads, const ds4_model * model, const ds4_layer_weights * layer, const float * q, const float * raw_kv, uint32_t n_raw, const float * comp_kv, uint32_t n_comp, const bool * comp_allowed, ds4_cpu_decode_scratch * scratch) { const float *sinks = tensor_data(model, layer->attn_sinks); const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); const uint32_t n_total = n_raw + n_comp; if (n_total > scratch->attn_score_cap) ds4_die("CPU decode attention score scratch buffer is too small"); float *score = scratch->attn_score; for (uint32_t h = 0; h < DS4_N_HEAD; h++) { const float *qh = q + (uint64_t)h * DS4_N_HEAD_DIM; float max_score = sinks[h]; uint32_t idx = 0; for (uint32_t r = 0; r < n_raw; r++, idx++) { const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[idx] > max_score) max_score = score[idx]; } for (uint32_t r = 0; r < n_comp; r++, idx++) { if (comp_allowed && !comp_allowed[r]) { score[idx] = DS4_NEG_INF; continue; } const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; score[idx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[idx] > max_score) max_score = score[idx]; } float *oh = out_heads + (uint64_t)h * DS4_N_HEAD_DIM; memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); float denom = expf(sinks[h] - max_score); idx = 0; for (uint32_t r = 0; r < n_raw; r++, idx++) { const float weight = expf(score[idx] - max_score); const float *kv = raw_kv + (uint64_t)r * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } for (uint32_t r = 0; r < n_comp; r++, idx++) { if (score[idx] <= DS4_NEG_INF * 0.5f) continue; const float weight = expf(score[idx] - max_score); const float *kv = comp_kv + (uint64_t)r * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } const float inv = 1.0f / denom; scale_f32(oh, inv, DS4_N_HEAD_DIM); } } typedef struct { float * out_heads; const ds4_model * model; const ds4_layer_weights * layer; const float * q; const float * raw_kv; const float * comp_kv; const uint32_t * comp_counts; const uint8_t * allowed_mask; const uint8_t * allowed_bits; uint64_t allowed_stride; uint32_t n_tok; uint32_t raw_cap; } layer_attention_prefix_batch_ctx; static inline bool attention_prefix_comp_allowed( const layer_attention_prefix_batch_ctx *ctx, uint32_t t, uint32_t c) { if (!ctx->allowed_bits || !ctx->allowed_mask || !ctx->allowed_mask[t]) return true; const uint8_t *bits = ctx->allowed_bits + (uint64_t)t * ctx->allowed_stride; return (bits[c >> 3] & (uint8_t)(1u << (c & 7u))) != 0; } static void layer_attention_prefix_batch_worker(void *vctx, uint64_t r0, uint64_t r1) { layer_attention_prefix_batch_ctx *ctx = vctx; const float *sinks = tensor_data(ctx->model, ctx->layer->attn_sinks); const float kq_scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); const uint32_t max_comp = ctx->comp_counts ? ctx->comp_counts[ctx->n_tok - 1] : 0; const uint32_t max_total = ctx->raw_cap + max_comp; float score_stack[2048]; float *score = max_total <= 2048 ? score_stack : xmalloc((size_t)max_total * sizeof(score[0])); for (uint64_t idx = r0; idx < r1; idx++) { const uint32_t t = (uint32_t)(idx / DS4_N_HEAD); const uint32_t h = (uint32_t)(idx - (uint64_t)t * DS4_N_HEAD); const uint32_t raw_count = t + 1 < ctx->raw_cap ? t + 1 : ctx->raw_cap; const uint32_t raw_start = t + 1 - raw_count; const uint32_t comp_count = ctx->comp_counts ? ctx->comp_counts[t] : 0; const float *qh = ctx->q + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; float max_score = sinks[h]; uint32_t sidx = 0; for (uint32_t r = 0; r < raw_count; r++, sidx++) { const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[sidx] > max_score) max_score = score[sidx]; } for (uint32_t c = 0; c < comp_count; c++, sidx++) { if (!attention_prefix_comp_allowed(ctx, t, c)) { score[sidx] = DS4_NEG_INF; continue; } const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; score[sidx] = dot_f32(qh, kv, DS4_N_HEAD_DIM) * kq_scale; if (score[sidx] > max_score) max_score = score[sidx]; } float *oh = ctx->out_heads + (uint64_t)t * DS4_N_HEAD * DS4_N_HEAD_DIM + (uint64_t)h * DS4_N_HEAD_DIM; memset(oh, 0, (size_t)DS4_N_HEAD_DIM * sizeof(oh[0])); float denom = expf(sinks[h] - max_score); sidx = 0; for (uint32_t r = 0; r < raw_count; r++, sidx++) { const float weight = expf(score[sidx] - max_score); const float *kv = ctx->raw_kv + (uint64_t)(raw_start + r) * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } for (uint32_t c = 0; c < comp_count; c++, sidx++) { if (score[sidx] <= DS4_NEG_INF * 0.5f) continue; const float weight = expf(score[sidx] - max_score); const float *kv = ctx->comp_kv + (uint64_t)c * DS4_N_HEAD_DIM; denom += weight; axpy_f32(oh, kv, weight, DS4_N_HEAD_DIM); } scale_f32(oh, 1.0f / denom, DS4_N_HEAD_DIM); } if (score != score_stack) free(score); } /* Prefix prefill attention for a fresh prompt. It computes each token's view * of the raw window and compressed rows without running the decode loop. */ static void layer_attention_prefix_batch( float * out_heads, const ds4_model * model, const ds4_layer_weights * layer, const float * q, const float * raw_kv, const float * comp_kv, const uint32_t * comp_counts, const uint8_t * allowed_mask, const uint8_t * allowed_bits, uint64_t allowed_stride, uint32_t n_tok, uint32_t raw_cap) { layer_attention_prefix_batch_ctx ctx = { .out_heads = out_heads, .model = model, .layer = layer, .q = q, .raw_kv = raw_kv, .comp_kv = comp_kv, .comp_counts = comp_counts, .allowed_mask = allowed_mask, .allowed_bits = allowed_bits, .allowed_stride = allowed_stride, .n_tok = n_tok, .raw_cap = raw_cap, }; ds4_parallel_for_min_rows((uint64_t)n_tok * DS4_N_HEAD, layer_attention_prefix_batch_worker, &ctx, 1); } /* Ratio-4 layers use an auxiliary indexer to select which compressed rows are * visible to attention. This is the CPU allocation-owning helper. */ static bool *indexer_allowed_decode_one( const ds4_model * model, const ds4_layer_weights * layer, const float * cur, const float * qr_norm, const float * index_comp, uint32_t n_comp, uint32_t il, uint32_t pos) { if (n_comp == 0) return NULL; bool *allowed = xcalloc(n_comp, sizeof(allowed[0])); const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; if (top_k == n_comp) { for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; return allowed; } const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; const uint32_t n_head = DS4_N_INDEXER_HEAD; float *q = xmalloc((size_t)head_dim * n_head * sizeof(q[0])); float *weights = xmalloc((size_t)n_head * sizeof(weights[0])); float *scores = xmalloc((size_t)n_comp * sizeof(scores[0])); matvec_any(q, model, layer->indexer_attn_q_b, qr_norm); rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); matvec_any(weights, model, layer->indexer_proj, cur); const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; for (uint32_t c = 0; c < n_comp; c++) { const float *kv = index_comp + (uint64_t)c * head_dim; float s = 0.0f; for (uint32_t h = 0; h < n_head; h++) { const float *qh = q + (uint64_t)h * head_dim; float dot = dot_f32(kv, qh, head_dim); if (dot < 0.0f) dot = 0.0f; s += dot * weights[h]; } scores[c] = s; } for (uint32_t k = 0; k < top_k; k++) { uint32_t best = 0; float best_score = DS4_NEG_INF; for (uint32_t c = 0; c < n_comp; c++) { if (!allowed[c] && scores[c] > best_score) { best = c; best_score = scores[c]; } } allowed[best] = true; } free(scores); free(weights); free(q); return allowed; } /* Scratch-backed indexer selection for decode. */ static bool *indexer_allowed_decode_one_decode_scratch( const ds4_model * model, const ds4_layer_weights * layer, const float * cur, const float * qr_norm, const float * index_comp, uint32_t n_comp, uint32_t il, uint32_t pos, ds4_cpu_decode_scratch * scratch) { if (n_comp == 0) return NULL; if (n_comp > scratch->comp_cap) ds4_die("CPU decode indexer scratch buffer is too small"); bool *allowed = scratch->index_allowed; memset(allowed, 0, (size_t)n_comp * sizeof(allowed[0])); const uint32_t top_k = DS4_N_INDEXER_TOP_K < n_comp ? DS4_N_INDEXER_TOP_K : n_comp; if (top_k == n_comp) { for (uint32_t i = 0; i < n_comp; i++) allowed[i] = true; return allowed; } const uint32_t head_dim = DS4_N_INDEXER_HEAD_DIM; const uint32_t n_head = DS4_N_INDEXER_HEAD; float *q = scratch->index_q; float *weights = scratch->index_weights; float *scores = scratch->index_scores; matvec_any_decode_scratch(q, model, layer->indexer_attn_q_b, qr_norm, scratch); rope_tail_layer_inplace(q, n_head, head_dim, DS4_N_ROT, pos, il, false); dsv4_indexer_qat_rows_inplace_cpu(q, n_head, head_dim); matvec_any_decode_scratch(weights, model, layer->indexer_proj, cur, scratch); const float scale = 1.0f / sqrtf((float)(head_dim * n_head)); for (uint32_t h = 0; h < n_head; h++) weights[h] *= scale; for (uint32_t c = 0; c < n_comp; c++) { const float *kv = index_comp + (uint64_t)c * head_dim; float s = 0.0f; for (uint32_t h = 0; h < n_head; h++) { const float *qh = q + (uint64_t)h * head_dim; float dot = dot_f32(kv, qh, head_dim); if (dot < 0.0f) dot = 0.0f; s += dot * weights[h]; } scores[c] = s; } for (uint32_t k = 0; k < top_k; k++) { uint32_t best = 0; float best_score = DS4_NEG_INF; for (uint32_t c = 0; c < n_comp; c++) { if (!allowed[c] && scores[c] > best_score) { best = c; best_score = scores[c]; } } allowed[best] = true; } return allowed; } /* Single-token attention sublayer with raw SWA cache and DS4 compression. */ static void layer_attention_raw_swa_one( float * after_attn_hc, const ds4_model * model, const ds4_layer_weights * layer, ds4_layer_cache * cache, const float * inp_hc, uint32_t il, uint32_t pos, const float * steering_dirs, float steering_scale) { const uint32_t n_hc = DS4_N_HC; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); float *q = xmalloc((size_t)q_dim * sizeof(q[0])); float *qr_norm = xmalloc((size_t)DS4_N_LORA_Q * sizeof(qr_norm[0])); float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); bool *comp_allowed = NULL; float post[4]; float comb[16]; memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); hc_pre_from_state_one(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, attn_residual, attn_cur, post, comb); layer_attn_norm_one(attn_norm, model, layer, attn_cur); layer_q_projection_with_lora_one(model, layer, attn_norm, q, qr_norm); layer_kv_projection_normed_one(model, layer, attn_norm, kv); rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); kv_cache_push_raw(cache, kv); const uint32_t ratio = cache->compress_ratio; if (ratio != 0) { float *comp = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp[0])); if (compressor_decode_one(comp, model, layer->attn_compressor_kv, layer->attn_compressor_gate, layer->attn_compressor_ape, layer->attn_compressor_norm, attn_norm, cache->attn_state_kv, cache->attn_state_score, DS4_N_HEAD_DIM, ratio, il, pos)) { kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); } free(comp); if (ratio == 4) { float *index_comp = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp[0])); if (compressor_decode_one(index_comp, model, layer->indexer_compressor_kv, layer->indexer_compressor_gate, layer->indexer_compressor_ape, layer->indexer_compressor_norm, attn_norm, cache->index_state_kv, cache->index_state_score, DS4_N_INDEXER_HEAD_DIM, ratio, il, pos)) { kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); } free(index_comp); comp_allowed = indexer_allowed_decode_one(model, layer, attn_norm, qr_norm, cache->index_comp_kv, cache->n_index_comp, il, pos); } layer_attention_mixed_one(heads, model, layer, q, cache->raw_kv, cache->n_raw, cache->attn_comp_kv, cache->n_comp, comp_allowed); } else { layer_attention_rows_one(heads, model, layer, q, cache->raw_kv, cache->n_raw); } rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); layer_grouped_out_one(attn_out, model, layer, heads); cpu_directional_steering_project_rows(attn_out, steering_dirs, il, 1, steering_scale); hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); free(comp_allowed); free(attn_out); free(heads); free(kv); free(qr_norm); free(q); free(attn_residual); free(attn_norm); free(attn_cur); } /* Batched prefill attention. It projects Q/KV for all tokens, streams them * through the same raw/compressed cache updates, then runs prefix attention. */ static void layer_attention_raw_swa_batch( float * after_attn_hc, const ds4_model * model, const ds4_layer_weights * layer, ds4_layer_cache * cache, const float * inp_hc, uint32_t n_tok, uint32_t il, uint32_t pos0, const float * steering_dirs, float steering_scale) { const bool profile = getenv("DS4_PREFILL_PROFILE_DETAIL") != NULL; const double t_start = profile ? now_sec() : 0.0; double t_hc_norm = 0.0; double t_q = 0.0; double t_kv = 0.0; double t_token_loop = 0.0; double t_tl_rope_cache = 0.0; double t_tl_compress = 0.0; double t_tl_indexer = 0.0; double t_tl_attn_rows = 0.0; double t_tl_inv_rope = 0.0; double t_out = 0.0; const uint32_t n_hc = DS4_N_HC; const uint64_t hc_dim = (uint64_t)n_hc * DS4_N_EMBD; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; float *attn_cur = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_cur[0])); float *attn_norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_norm[0])); float *attn_residual = xmalloc((size_t)n_tok * hc_dim * sizeof(attn_residual[0])); const uint32_t q_rank = DS4_N_LORA_Q; float *qr = xmalloc((size_t)n_tok * q_rank * sizeof(qr[0])); float *qr_norm = xmalloc((size_t)n_tok * q_rank * sizeof(qr_norm[0])); float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); float *kv_raw = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv_raw[0])); float *kv = xmalloc((size_t)n_tok * DS4_N_HEAD_DIM * sizeof(kv[0])); float *heads = NULL; float *attn_out = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(attn_out[0])); float *post = xmalloc((size_t)n_tok * n_hc * sizeof(post[0])); float *comb = xmalloc((size_t)n_tok * n_hc * n_hc * sizeof(comb[0])); const float *q_a_norm = tensor_data(model, layer->attn_q_a_norm); const float *kv_norm = tensor_data(model, layer->attn_kv_a_norm); double t0 = profile ? now_sec() : 0.0; hc_pre_norm_batch(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, layer->attn_norm, inp_hc, attn_residual, attn_cur, attn_norm, post, comb, n_tok); if (profile) t_hc_norm = now_sec() - t0; t0 = profile ? now_sec() : 0.0; matmul_q8_0_batch(qr, model, layer->attn_q_a, attn_norm, n_tok); for (uint32_t t = 0; t < n_tok; t++) { rms_norm_weight(qr_norm + (uint64_t)t * q_rank, qr + (uint64_t)t * q_rank, q_a_norm, q_rank, DS4_RMS_EPS); } matmul_q8_0_batch(q, model, layer->attn_q_b, qr_norm, n_tok); for (uint32_t t = 0; t < n_tok; t++) { head_rms_norm_inplace(q + (uint64_t)t * q_dim, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS); } if (profile) t_q = now_sec() - t0; t0 = profile ? now_sec() : 0.0; matmul_q8_0_batch(kv_raw, model, layer->attn_kv, attn_norm, n_tok); for (uint32_t t = 0; t < n_tok; t++) { rms_norm_weight(kv + (uint64_t)t * DS4_N_HEAD_DIM, kv_raw + (uint64_t)t * DS4_N_HEAD_DIM, kv_norm, DS4_N_HEAD_DIM, DS4_RMS_EPS); } if (profile) t_kv = now_sec() - t0; t0 = profile ? now_sec() : 0.0; const uint32_t ratio = cache->compress_ratio; const bool prefer_parallel_attn = getenv("DS4_PARALLEL_ATTN_ROWS") != NULL; const bool prefix_batch_attn = prefer_parallel_attn && getenv("DS4_NO_PARALLEL_ATTN_ROWS") == NULL && cache->n_raw == 0 && pos0 == 0; if (!prefix_batch_attn) { heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); } uint32_t batch_rope_max = 4096; const char *batch_rope_max_env = getenv("DS4_BATCHED_ROPE_MAX"); if (batch_rope_max_env && batch_rope_max_env[0]) { long v = strtol(batch_rope_max_env, NULL, 10); if (v >= 0 && v <= 65536) batch_rope_max = (uint32_t)v; } const bool batch_prefix_rope = prefix_batch_attn && getenv("DS4_NO_BATCHED_ROPE") == NULL && n_tok <= batch_rope_max; uint32_t *comp_counts = prefix_batch_attn ? xcalloc((size_t)n_tok, sizeof(comp_counts[0])) : NULL; uint8_t *allowed_mask = prefix_batch_attn && ratio == 4 ? xcalloc((size_t)n_tok, sizeof(allowed_mask[0])) : NULL; uint8_t *allowed_bits = NULL; const uint64_t allowed_stride = ratio == 4 ? ((uint64_t)cache->comp_cap + 7u) / 8u : 0; float *comp_scratch = NULL; float *index_comp_scratch = NULL; if (ratio != 0) { comp_scratch = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(comp_scratch[0])); if (ratio == 4) { index_comp_scratch = xmalloc((size_t)DS4_N_INDEXER_HEAD_DIM * sizeof(index_comp_scratch[0])); } } if (batch_prefix_rope) { double tx = profile ? now_sec() : 0.0; rope_tail_layer_batch_inplace(q, q_dim, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos0, il, false, n_tok); rope_tail_layer_batch_inplace(kv, DS4_N_HEAD_DIM, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos0, il, false, n_tok); if (profile) t_tl_rope_cache += now_sec() - tx; } for (uint32_t t = 0; t < n_tok; t++) { const uint32_t pos = pos0 + t; float *q_t = q + (uint64_t)t * q_dim; float *kv_t = kv + (uint64_t)t * DS4_N_HEAD_DIM; bool *comp_allowed = NULL; double tx = profile ? now_sec() : 0.0; if (!batch_prefix_rope) { rope_tail_layer_inplace(q_t, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); rope_tail_layer_inplace(kv_t, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); } dsv4_fp8_kv_quantize_row_inplace_cpu(kv_t, DS4_N_HEAD_DIM, DS4_N_ROT); kv_cache_push_raw(cache, kv_t); if (profile) t_tl_rope_cache += now_sec() - tx; if (ratio != 0) { tx = profile ? now_sec() : 0.0; float *comp = comp_scratch; const bool have_comp = compressor_decode_one(comp, model, layer->attn_compressor_kv, layer->attn_compressor_gate, layer->attn_compressor_ape, layer->attn_compressor_norm, attn_norm + (uint64_t)t * DS4_N_EMBD, cache->attn_state_kv, cache->attn_state_score, DS4_N_HEAD_DIM, ratio, il, pos); if (have_comp) { kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, comp); } if (ratio == 4) { float *index_comp = index_comp_scratch; const bool have_index_comp = compressor_decode_one(index_comp, model, layer->indexer_compressor_kv, layer->indexer_compressor_gate, layer->indexer_compressor_ape, layer->indexer_compressor_norm, attn_norm + (uint64_t)t * DS4_N_EMBD, cache->index_state_kv, cache->index_state_score, DS4_N_INDEXER_HEAD_DIM, ratio, il, pos); if (have_index_comp) { kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, index_comp); } if (profile) t_tl_compress += now_sec() - tx; tx = profile ? now_sec() : 0.0; comp_allowed = indexer_allowed_decode_one(model, layer, attn_norm + (uint64_t)t * DS4_N_EMBD, qr_norm + (uint64_t)t * q_rank, cache->index_comp_kv, cache->n_index_comp, il, pos); if (profile) t_tl_indexer += now_sec() - tx; } else { if (profile) t_tl_compress += now_sec() - tx; } if (comp_counts) comp_counts[t] = cache->n_comp; if (prefix_batch_attn && comp_allowed) { if (!allowed_bits) { allowed_bits = xcalloc((size_t)n_tok * allowed_stride, sizeof(allowed_bits[0])); } allowed_mask[t] = 1; uint8_t *bits = allowed_bits + (uint64_t)t * allowed_stride; for (uint32_t c = 0; c < cache->n_comp; c++) { if (comp_allowed[c]) bits[c >> 3] |= (uint8_t)(1u << (c & 7u)); } } if (!prefix_batch_attn) { tx = profile ? now_sec() : 0.0; layer_attention_mixed_one(heads + (uint64_t)t * q_dim, model, layer, q_t, cache->raw_kv, cache->n_raw, cache->attn_comp_kv, cache->n_comp, comp_allowed); if (profile) t_tl_attn_rows += now_sec() - tx; } } else { if (!prefix_batch_attn) { tx = profile ? now_sec() : 0.0; layer_attention_rows_one(heads + (uint64_t)t * q_dim, model, layer, q_t, cache->raw_kv, cache->n_raw); if (profile) t_tl_attn_rows += now_sec() - tx; } } if (!prefix_batch_attn) { tx = profile ? now_sec() : 0.0; rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); if (profile) t_tl_inv_rope += now_sec() - tx; } free(comp_allowed); } if (prefix_batch_attn) { double tx = profile ? now_sec() : 0.0; const float *comp_kv_for_prefix = cache->attn_comp_kv ? cache->attn_comp_kv : kv; if (!heads) { heads = xmalloc((size_t)n_tok * q_dim * sizeof(heads[0])); } layer_attention_prefix_batch(heads, model, layer, q, kv, comp_kv_for_prefix, comp_counts, allowed_mask, allowed_bits, allowed_stride, n_tok, cache->cap_raw); if (profile) t_tl_attn_rows += now_sec() - tx; tx = profile ? now_sec() : 0.0; if (batch_prefix_rope) { rope_tail_layer_batch_inplace(heads, q_dim, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos0, il, true, n_tok); } else { for (uint32_t t = 0; t < n_tok; t++) { rope_tail_layer_inplace(heads + (uint64_t)t * q_dim, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos0 + t, il, true); } } if (profile) t_tl_inv_rope += now_sec() - tx; } if (profile) t_token_loop = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_grouped_out_batch(attn_out, model, layer, heads, n_tok); cpu_directional_steering_project_rows(attn_out, steering_dirs, il, n_tok, steering_scale); hc_post_batch(after_attn_hc, attn_out, attn_residual, post, comb, n_tok, DS4_N_EMBD, n_hc); if (profile) t_out = now_sec() - t0; if (profile) { fprintf(stderr, "ds4: prefill detail layer %u attn hc_norm=%.3f q=%.3f kv=%.3f token_loop=%.3f out=%.3f total=%.3f\n", il, t_hc_norm, t_q, t_kv, t_token_loop, t_out, now_sec() - t_start); if (getenv("DS4_PREFILL_PROFILE_TOKEN") != NULL) { fprintf(stderr, "ds4: prefill token detail layer %u rope_cache=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f\n", il, t_tl_rope_cache, t_tl_compress, t_tl_indexer, t_tl_attn_rows, t_tl_inv_rope); } } free(allowed_bits); free(allowed_mask); free(comp_counts); free(index_comp_scratch); free(comp_scratch); free(comb); free(post); free(attn_out); free(heads); free(kv); free(kv_raw); free(q); free(qr_norm); free(qr); free(attn_residual); free(attn_norm); free(attn_cur); } /* Full transformer layer for one decode token: attention sublayer followed by * FFN sublayer, both operating on the HC state. */ static void layer_forward_raw_swa_one( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, ds4_layer_cache * cache, const float * inp_hc, uint32_t il, uint32_t pos, int token, const float * steering_dirs, float steering_attn_scale, float steering_ffn_scale, ds4_cpu_decode_scratch * scratch) { const uint32_t n_hc = DS4_N_HC; const bool profile = getenv("DS4_DECODE_PROFILE_DETAIL") != NULL; const double t_start = profile ? now_sec() : 0.0; double t_hc = 0.0; double t_q = 0.0; double t_kv = 0.0; double t_rope_cache = 0.0; double t_compress = 0.0; double t_indexer = 0.0; double t_attn_rows = 0.0; double t_inv_rope = 0.0; double t_out = 0.0; double t_post = 0.0; double t_ffn = 0.0; bool *comp_allowed = NULL; float post[4]; float comb[16]; double t0 = profile ? now_sec() : 0.0; memcpy(scratch->attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); hc_pre_from_state_one_scratch(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, scratch->attn_residual, scratch->attn_cur, post, comb, scratch->hc_flat, false); if (profile) t_hc = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_attn_norm_one(scratch->attn_norm, model, layer, scratch->attn_cur); const uint32_t ratio = cache->compress_ratio; layer_q_projection_with_lora_one_decode_scratch(model, layer, scratch->attn_norm, scratch->q, scratch->qr_norm, scratch); if (profile) t_q = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_kv_projection_normed_one_decode_scratch(model, layer, scratch->attn_norm, scratch->kv, scratch); if (profile) t_kv = now_sec() - t0; t0 = profile ? now_sec() : 0.0; rope_tail_layer_inplace(scratch->q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); rope_tail_layer_inplace(scratch->kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); dsv4_fp8_kv_quantize_row_inplace_cpu(scratch->kv, DS4_N_HEAD_DIM, DS4_N_ROT); kv_cache_push_raw(cache, scratch->kv); if (profile) t_rope_cache = now_sec() - t0; if (ratio != 0) { t0 = profile ? now_sec() : 0.0; if (compressor_decode_one_decode_scratch(scratch->comp, model, layer->attn_compressor_kv, layer->attn_compressor_gate, layer->attn_compressor_ape, layer->attn_compressor_norm, scratch->attn_norm, cache->attn_state_kv, cache->attn_state_score, DS4_N_HEAD_DIM, ratio, il, pos, scratch)) { kv_cache_push_comp(cache->attn_comp_kv, &cache->n_comp, cache->comp_cap, DS4_N_HEAD_DIM, scratch->comp); } if (ratio == 4) { if (compressor_decode_one_decode_scratch(scratch->index_comp, model, layer->indexer_compressor_kv, layer->indexer_compressor_gate, layer->indexer_compressor_ape, layer->indexer_compressor_norm, scratch->attn_norm, cache->index_state_kv, cache->index_state_score, DS4_N_INDEXER_HEAD_DIM, ratio, il, pos, scratch)) { kv_cache_push_comp(cache->index_comp_kv, &cache->n_index_comp, cache->comp_cap, DS4_N_INDEXER_HEAD_DIM, scratch->index_comp); } if (profile) t_compress = now_sec() - t0; } else if (profile) { t_compress = now_sec() - t0; } } if (ratio == 4) { t0 = profile ? now_sec() : 0.0; comp_allowed = indexer_allowed_decode_one_decode_scratch(model, layer, scratch->attn_norm, scratch->qr_norm, cache->index_comp_kv, cache->n_index_comp, il, pos, scratch); if (profile) t_indexer = now_sec() - t0; } t0 = profile ? now_sec() : 0.0; if (ratio != 0) { layer_attention_mixed_one_decode_scratch(scratch->heads, model, layer, scratch->q, cache->raw_kv, cache->n_raw, cache->attn_comp_kv, cache->n_comp, comp_allowed, scratch); } else { layer_attention_rows_one(scratch->heads, model, layer, scratch->q, cache->raw_kv, cache->n_raw); } if (profile) t_attn_rows = now_sec() - t0; t0 = profile ? now_sec() : 0.0; rope_tail_layer_inplace(scratch->heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); if (profile) t_inv_rope = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_grouped_out_one_decode_scratch(scratch->attn_out, model, layer, scratch->heads, scratch); cpu_directional_steering_project_rows(scratch->attn_out, steering_dirs, il, 1, steering_attn_scale); if (profile) t_out = now_sec() - t0; t0 = profile ? now_sec() : 0.0; hc_post_one(scratch->after_attn_hc, scratch->attn_out, scratch->attn_residual, post, comb, DS4_N_EMBD, n_hc); if (profile) t_post = now_sec() - t0; t0 = profile ? now_sec() : 0.0; layer_ffn_one_decode_scratch(out_hc, model, layer, scratch->after_attn_hc, il, token, steering_dirs, steering_ffn_scale, scratch); if (profile) t_ffn = now_sec() - t0; if (profile) { fprintf(stderr, "ds4: decode detail layer %u attn hc=%.3f q=%.3f kv=%.3f rope=%.3f compress=%.3f indexer=%.3f attn_rows=%.3f inv_rope=%.3f out=%.3f post=%.3f ffn=%.3f total=%.3f ms\n", il, t_hc * 1000.0, t_q * 1000.0, t_kv * 1000.0, t_rope_cache * 1000.0, t_compress * 1000.0, t_indexer * 1000.0, t_attn_rows * 1000.0, t_inv_rope * 1000.0, t_out * 1000.0, t_post * 1000.0, t_ffn * 1000.0, (now_sec() - t_start) * 1000.0); } } static void output_logits_one_decode_scratch( float * logits, const ds4_model * model, const ds4_weights * weights, const float * inp_hc, ds4_cpu_decode_scratch * scratch); /* CPU decode for one token through all 43 layers. The caller owns scratch and * cache lifetimes so no per-token allocations are needed. */ static void forward_token_raw_swa_cpu_decode_scratch( float * logits, const ds4_model * model, const ds4_weights * weights, ds4_kv_cache * cache, int token, uint32_t pos, const float * steering_dirs, float steering_attn_scale, float steering_ffn_scale, ds4_cpu_decode_scratch * scratch) { float *cur = scratch->cur; float *next = scratch->next; embed_token_f16(model, weights, token, scratch->plain); hc_from_plain_embedding(cur, scratch->plain, DS4_N_EMBD, DS4_N_HC); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { layer_forward_raw_swa_one(next, model, &weights->layer[il], &cache->layer[il], cur, il, pos, token, steering_dirs, steering_attn_scale, steering_ffn_scale, scratch); float *tmp = cur; cur = next; next = tmp; } if (logits) { output_logits_one_decode_scratch(logits, model, weights, cur, scratch); } } #ifndef DS4_NO_GPU static void forward_token_raw_swa_cpu( float * logits, const ds4_model * model, const ds4_weights * weights, ds4_kv_cache * cache, int token, uint32_t pos) { ds4_cpu_decode_scratch scratch; uint32_t ctx_guess = pos + 1; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t ratio = cache->layer[il].compress_ratio; if (ratio != 0 && cache->layer[il].comp_cap > 2) { const uint32_t ctx_from_comp = (cache->layer[il].comp_cap - 2u) * ratio; if (ctx_guess < ctx_from_comp) ctx_guess = ctx_from_comp; } } cpu_decode_scratch_init(&scratch, ctx_guess); forward_token_raw_swa_cpu_decode_scratch(logits, model, weights, cache, token, pos, NULL, 0.0f, 0.0f, &scratch); cpu_decode_scratch_free(&scratch); } #endif /* CPU prefill in layer-major order. All prompt tokens pass through layer 0, * then layer 1, etc., which exposes batch matmul opportunities. */ static void prefill_layer_major_cpu( float * logits, const ds4_model * model, const ds4_weights * weights, ds4_kv_cache * cache, const token_vec * prompt, const float * steering_dirs, float steering_attn_scale, float steering_ffn_scale) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t n_tok = (uint64_t)prompt->len; float *cur = xmalloc((size_t)n_tok * hc_dim * sizeof(cur[0])); float *next = xmalloc((size_t)n_tok * hc_dim * sizeof(next[0])); float *attn = xmalloc((size_t)n_tok * hc_dim * sizeof(attn[0])); float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); uint32_t ffn_batch = 128; const bool batched_attn = getenv("DS4_NO_BATCHED_ATTN") == NULL; const bool batched_ffn = getenv("DS4_BATCHED_FFN") != NULL; const bool parallel_ffn = getenv("DS4_PARALLEL_FFN") != NULL; const bool shared_batch_ffn = getenv("DS4_NO_SHARED_BATCH_FFN") == NULL; const char *batch_env = getenv("DS4_PREFILL_BATCH"); ds4_cpu_decode_scratch decode_scratch; bool decode_scratch_ready = false; if (batch_env && batch_env[0]) { long v = strtol(batch_env, NULL, 10); if (v > 0 && v < 4096) ffn_batch = (uint32_t)v; } for (uint64_t t = 0; t < n_tok; t++) { embed_token_f16(model, weights, prompt->v[t], plain); hc_from_plain_embedding(cur + t * hc_dim, plain, DS4_N_EMBD, DS4_N_HC); } free(plain); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { fprintf(stderr, "ds4: prefill layer %u/%u\r", il + 1, (uint32_t)DS4_N_LAYER); fflush(stderr); if (batched_attn) { layer_attention_raw_swa_batch(attn, model, &weights->layer[il], &cache->layer[il], cur, (uint32_t)n_tok, il, 0, steering_dirs, steering_attn_scale); if (batched_ffn) { for (uint64_t t = 0; t < n_tok; t += ffn_batch) { uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); layer_ffn_batch(next + t * hc_dim, model, &weights->layer[il], attn + t * hc_dim, prompt->v + t, nb, il, steering_dirs, steering_ffn_scale); } } else if (shared_batch_ffn) { layer_ffn_shared_batch(next, model, &weights->layer[il], attn, prompt->v, (uint32_t)n_tok, il, steering_dirs, steering_ffn_scale); } else if (parallel_ffn) { layer_ffn_tokens_parallel(next, model, &weights->layer[il], attn, prompt->v, (uint32_t)n_tok, il, steering_dirs, steering_ffn_scale); } else { for (uint64_t t = 0; t < n_tok; t++) { layer_ffn_one(next + t * hc_dim, model, &weights->layer[il], attn + t * hc_dim, il, prompt->v[t], steering_dirs, steering_ffn_scale, false); } } } else if (batched_ffn) { for (uint64_t t = 0; t < n_tok; t++) { layer_attention_raw_swa_one(attn + t * hc_dim, model, &weights->layer[il], &cache->layer[il], cur + t * hc_dim, il, (uint32_t)t, steering_dirs, steering_attn_scale); } for (uint64_t t = 0; t < n_tok; t += ffn_batch) { uint32_t nb = (uint32_t)((n_tok - t) < ffn_batch ? (n_tok - t) : ffn_batch); layer_ffn_batch(next + t * hc_dim, model, &weights->layer[il], attn + t * hc_dim, prompt->v + t, nb, il, steering_dirs, steering_ffn_scale); } } else { if (!decode_scratch_ready) { cpu_decode_scratch_init(&decode_scratch, (uint32_t)n_tok); decode_scratch_ready = true; } for (uint64_t t = 0; t < n_tok; t++) { layer_forward_raw_swa_one(next + t * hc_dim, model, &weights->layer[il], &cache->layer[il], cur + t * hc_dim, il, (uint32_t)t, prompt->v[t], steering_dirs, steering_attn_scale, steering_ffn_scale, &decode_scratch); } } float *tmp = cur; cur = next; next = tmp; } kv_cache_finish_prefill_states(cache, (uint32_t)n_tok); if (logits) { output_logits_one(logits, model, weights, cur + (n_tok - 1) * hc_dim); } if (decode_scratch_ready) cpu_decode_scratch_free(&decode_scratch); free(next); free(cur); free(attn); } /* Diagnostic first-token layer without cache history: the token attends only * to itself, useful for checking a minimal end-to-end slice. */ static void layer_forward_self_one( float * out_hc, const ds4_model * model, const ds4_layer_weights * layer, const float * inp_hc, uint32_t il, uint32_t pos, int token) { const uint32_t n_hc = DS4_N_HC; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; float *attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_cur[0])); float *attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_norm[0])); float *attn_residual = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(attn_residual[0])); float *q = xmalloc((size_t)q_dim * sizeof(q[0])); float *kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(kv[0])); float *heads = xmalloc((size_t)q_dim * sizeof(heads[0])); float *attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(attn_out[0])); float *after_attn_hc = xmalloc((size_t)n_hc * DS4_N_EMBD * sizeof(after_attn_hc[0])); float post[4]; float comb[16]; memcpy(attn_residual, inp_hc, (size_t)n_hc * DS4_N_EMBD * sizeof(inp_hc[0])); hc_pre_from_state_one(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, attn_residual, attn_cur, post, comb); layer_attn_norm_one(attn_norm, model, layer, attn_cur); layer_q_projection_normed_one(model, layer, attn_norm, q); layer_kv_projection_normed_one(model, layer, attn_norm, kv); rope_tail_layer_inplace(q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); rope_tail_layer_inplace(kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, false); dsv4_fp8_kv_quantize_row_inplace_cpu(kv, DS4_N_HEAD_DIM, DS4_N_ROT); f16_round_inplace_cpu(kv, DS4_N_HEAD_DIM); layer_attention_one(heads, model, layer, q, kv); rope_tail_layer_inplace(heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, il, true); layer_grouped_out_one(attn_out, model, layer, heads); hc_post_one(after_attn_hc, attn_out, attn_residual, post, comb, DS4_N_EMBD, n_hc); layer_ffn_one(out_hc, model, layer, after_attn_hc, il, token, NULL, 0.0f, false); free(after_attn_hc); free(attn_out); free(heads); free(kv); free(q); free(attn_residual); free(attn_norm); free(attn_cur); } static void forward_first_token_cpu( float * out_hc, const ds4_model * model, const ds4_weights * weights, int token) { float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); float *cur = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(cur[0])); float *next = xmalloc((size_t)DS4_N_HC * DS4_N_EMBD * sizeof(next[0])); embed_token_f16(model, weights, token, plain); hc_from_plain_embedding(cur, plain, DS4_N_EMBD, DS4_N_HC); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { layer_forward_self_one(next, model, &weights->layer[il], cur, il, 0, token); float *tmp = cur; cur = next; next = tmp; } memcpy(out_hc, cur, (size_t)DS4_N_HC * DS4_N_EMBD * sizeof(out_hc[0])); free(next); free(cur); free(plain); } /* Collapse final HC streams into the ordinary embedding vector before the * output norm and vocabulary projection. */ static void output_hc_head_one( float * out, const ds4_model * model, const ds4_weights * weights, const float * inp_hc) { const uint32_t n_hc = DS4_N_HC; const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * n_hc; float *flat = xmalloc((size_t)hc_dim * sizeof(flat[0])); float *pre = xmalloc((size_t)n_hc * sizeof(pre[0])); float *w = xmalloc((size_t)n_hc * sizeof(w[0])); rms_norm_no_weight(flat, inp_hc, hc_dim, DS4_RMS_EPS); matvec_f16(pre, model, weights->output_hc_fn, flat); const float *scale = tensor_data(model, weights->output_hc_scale); const float *base = tensor_data(model, weights->output_hc_base); for (uint32_t i = 0; i < n_hc; i++) { w[i] = sigmoid_stable(pre[i] * scale[0] + base[i]) + DS4_HC_EPS; } hc_weighted_sum_one(out, inp_hc, w, DS4_N_EMBD, n_hc); free(w); free(pre); free(flat); } /* Final language-model head: HC collapse, RMSNorm, and Q8_0 vocab projection. */ static void output_logits_one( float * logits, const ds4_model * model, const ds4_weights * weights, const float * inp_hc) { float *embd = xmalloc((size_t)DS4_N_EMBD * sizeof(embd[0])); float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); output_hc_head_one(embd, model, weights, inp_hc); rms_norm_weight(norm, embd, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS); matvec_q8_0(logits, model, weights->output, norm); free(norm); free(embd); } static void layer_glm_first_token_attention_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x) { float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); const uint64_t kv_blocks = (DS4_N_KV_LORA + 31) / 32; int8_t *kvq = xmalloc((size_t)kv_blocks * 32); float *kvscale = xmalloc((size_t)kv_blocks * sizeof(kvscale[0])); if (layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || layer->attn_v_b->dim[0] != DS4_N_KV_LORA || layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || layer->attn_v_b->dim[2] != DS4_N_HEAD || layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || layer->attn_output->dim[1] != DS4_N_EMBD) { ds4_die("GLM attention tensors have an unexpected layout"); } rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); matvec_q8_0(kv_raw, model, layer->attn_kv_a_mqa, norm); rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), DS4_N_KV_LORA, DS4_RMS_EPS); quantize_q8_0_activation(kv_norm, kvq, kvscale, DS4_N_KV_LORA); for (uint32_t h = 0; h < DS4_N_HEAD; h++) { matvec_q8_0_3d_slice_prequant(heads + (uint64_t)h * DS4_N_VALUE_MLA, model, layer->attn_v_b, kvq, kvscale, h); } matvec_q8_0(out, model, layer->attn_output, heads); free(kvscale); free(kvq); free(heads); free(kv_norm); free(kv_raw); free(norm); } static void layer_glm_first_token_attention_one_f32_ref( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x) { float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0])); float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0])); float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0])); float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0])); if (layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || layer->attn_v_b->type != DS4_TENSOR_Q8_0 || layer->attn_output->type != DS4_TENSOR_Q8_0 || layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA || layer->attn_v_b->dim[0] != DS4_N_KV_LORA || layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || layer->attn_v_b->dim[2] != DS4_N_HEAD || layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA || layer->attn_output->dim[1] != DS4_N_EMBD) { ds4_die("GLM F32 attention reference found unexpected tensor layout"); } rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); matvec_q8_0_f32_ref(kv_raw, model, layer->attn_kv_a_mqa, norm); rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm), DS4_N_KV_LORA, DS4_RMS_EPS); matvec_q8_0_f32_ref(heads, model, layer->attn_v_b, kv_norm); matvec_q8_0_f32_ref(out, model, layer->attn_output, heads); free(heads); free(kv_norm); free(kv_raw); free(norm); } static void glm_k_b_project_f32_ref( float * out, const ds4_model * model, const ds4_tensor * w, const float * kv_norm) { const uint32_t q_nope = DS4_N_KEY_MLA - DS4_N_ROT; if (w->type != DS4_TENSOR_Q8_0 || w->ndim != 3 || w->dim[0] != q_nope || w->dim[1] != DS4_N_KV_LORA || w->dim[2] != DS4_N_HEAD) { ds4_die("GLM k_b reference found unexpected tensor layout"); } const uint8_t *data = tensor_data(model, w); const uint64_t blocks = (q_nope + 31u) / 32u; const uint64_t row_bytes = blocks * 34u; memset(out, 0, (size_t)DS4_N_HEAD * q_nope * sizeof(out[0])); for (uint32_t h = 0; h < DS4_N_HEAD; h++) { float *dst = out + (uint64_t)h * q_nope; for (uint32_t j = 0; j < DS4_N_KV_LORA; j++) { const uint8_t *row = data + ((uint64_t)h * DS4_N_KV_LORA + j) * row_bytes; const float xj = kv_norm[j]; for (uint64_t b = 0; b < blocks; b++) { uint16_t scale_bits; memcpy(&scale_bits, row + b * 34u, sizeof(scale_bits)); const int8_t *qs = (const int8_t *)(row + b * 34u + 2u); const float d = f16_to_f32(scale_bits) * xj; const uint32_t i0 = (uint32_t)b * 32u; const uint32_t n = q_nope - i0 < 32u ? q_nope - i0 : 32u; for (uint32_t i = 0; i < n; i++) { dst[i0 + i] += d * (float)qs[i]; } } } } } static void layer_glm_attention_prefill_f32_ref( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x, uint32_t n_tok, uint32_t pos0, uint32_t il) { if (n_tok == 0) return; const uint32_t qk_dim = DS4_N_KEY_MLA; const uint32_t q_nope = qk_dim - DS4_N_ROT; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * qk_dim; const uint64_t heads_dim = (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA; const uint64_t kv_raw_dim = layer->attn_kv_a_mqa ? layer->attn_kv_a_mqa->dim[1] : 0; if (!layer->attn_norm || !layer->attn_q_a || !layer->attn_q_a_norm || !layer->attn_q_b || !layer->attn_kv_a_mqa || !layer->attn_kv_a_norm || !layer->attn_k_b || !layer->attn_v_b || !layer->attn_output || layer->attn_q_a->type != DS4_TENSOR_Q8_0 || layer->attn_q_b->type != DS4_TENSOR_Q8_0 || layer->attn_kv_a_mqa->type != DS4_TENSOR_Q8_0 || layer->attn_k_b->type != DS4_TENSOR_Q8_0 || layer->attn_v_b->type != DS4_TENSOR_Q8_0 || layer->attn_output->type != DS4_TENSOR_Q8_0 || layer->attn_norm->type != DS4_TENSOR_F32 || layer->attn_q_a_norm->type != DS4_TENSOR_F32 || layer->attn_kv_a_norm->type != DS4_TENSOR_F32 || layer->attn_q_a->dim[0] != DS4_N_EMBD || layer->attn_q_a->dim[1] != DS4_N_LORA_Q || layer->attn_q_a_norm->dim[0] != DS4_N_LORA_Q || layer->attn_q_b->dim[0] != DS4_N_LORA_Q || layer->attn_q_b->dim[1] != q_dim || layer->attn_kv_a_mqa->dim[0] != DS4_N_EMBD || kv_raw_dim < (uint64_t)DS4_N_KV_LORA + DS4_N_ROT || layer->attn_kv_a_norm->dim[0] != DS4_N_KV_LORA || layer->attn_k_b->dim[0] != q_nope || layer->attn_k_b->dim[1] != DS4_N_KV_LORA || layer->attn_k_b->dim[2] != DS4_N_HEAD || layer->attn_v_b->dim[0] != DS4_N_KV_LORA || layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA || layer->attn_v_b->dim[2] != DS4_N_HEAD || layer->attn_output->dim[0] != heads_dim || layer->attn_output->dim[1] != DS4_N_EMBD) { ds4_die("GLM prefill attention reference found unexpected tensor layout"); } float *norm = xmalloc((size_t)n_tok * DS4_N_EMBD * sizeof(norm[0])); float *q_rank = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank[0])); float *q_rank_norm = xmalloc((size_t)n_tok * DS4_N_LORA_Q * sizeof(q_rank_norm[0])); float *q = xmalloc((size_t)n_tok * q_dim * sizeof(q[0])); float *kv_raw = xmalloc((size_t)n_tok * kv_raw_dim * sizeof(kv_raw[0])); float *kv_norm = xmalloc((size_t)n_tok * DS4_N_KV_LORA * sizeof(kv_norm[0])); float *k_nope = xmalloc((size_t)n_tok * DS4_N_HEAD * q_nope * sizeof(k_nope[0])); float *key_cache = xmalloc((size_t)n_tok * DS4_N_HEAD * qk_dim * sizeof(key_cache[0])); float *value_cache = xmalloc((size_t)n_tok * heads_dim * sizeof(value_cache[0])); float *heads = xmalloc((size_t)n_tok * heads_dim * sizeof(heads[0])); float *k_rot = xmalloc((size_t)DS4_N_ROT * sizeof(k_rot[0])); float *scores = xmalloc((size_t)n_tok * sizeof(scores[0])); for (uint32_t t = 0; t < n_tok; t++) { const float *xt = x + (uint64_t)t * DS4_N_EMBD; float *norm_t = norm + (uint64_t)t * DS4_N_EMBD; float *qr_t = q_rank + (uint64_t)t * DS4_N_LORA_Q; float *qrn_t = q_rank_norm + (uint64_t)t * DS4_N_LORA_Q; float *q_t = q + (uint64_t)t * q_dim; float *raw_t = kv_raw + (uint64_t)t * kv_raw_dim; float *kvn_t = kv_norm + (uint64_t)t * DS4_N_KV_LORA; float *kn_t = k_nope + (uint64_t)t * DS4_N_HEAD * q_nope; float *kc_t = key_cache + (uint64_t)t * DS4_N_HEAD * qk_dim; float *vc_t = value_cache + (uint64_t)t * heads_dim; rms_norm_weight(norm_t, xt, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS); matvec_q8_0_f32_ref(qr_t, model, layer->attn_q_a, norm_t); rms_norm_weight(qrn_t, qr_t, tensor_data(model, layer->attn_q_a_norm), DS4_N_LORA_Q, DS4_RMS_EPS); matvec_q8_0_f32_ref(q_t, model, layer->attn_q_b, qrn_t); rope_tail_layer_inplace(q_t, DS4_N_HEAD, qk_dim, DS4_N_ROT, pos0 + t, il, false); matvec_q8_0_f32_ref(raw_t, model, layer->attn_kv_a_mqa, norm_t); rms_norm_weight(kvn_t, raw_t, tensor_data(model, layer->attn_kv_a_norm), DS4_N_KV_LORA, DS4_RMS_EPS); glm_k_b_project_f32_ref(kn_t, model, layer->attn_k_b, kvn_t); matvec_q8_0_f32_ref(vc_t, model, layer->attn_v_b, kvn_t); memcpy(k_rot, raw_t + DS4_N_KV_LORA, (size_t)DS4_N_ROT * sizeof(k_rot[0])); rope_tail_layer_inplace(k_rot, 1, DS4_N_ROT, DS4_N_ROT, pos0 + t, il, false); for (uint32_t h = 0; h < DS4_N_HEAD; h++) { float *kd = kc_t + (uint64_t)h * qk_dim; memcpy(kd, kn_t + (uint64_t)h * q_nope, (size_t)q_nope * sizeof(kd[0])); memcpy(kd + q_nope, k_rot, (size_t)DS4_N_ROT * sizeof(kd[0])); } } const float scale = 1.0f / sqrtf((float)qk_dim); for (uint32_t t = 0; t < n_tok; t++) { const uint32_t visible = t + 1u; for (uint32_t h = 0; h < DS4_N_HEAD; h++) { const float *q_h = q + ((uint64_t)t * DS4_N_HEAD + h) * qk_dim; float max_score = -FLT_MAX; for (uint32_t s = 0; s < visible; s++) { const float *k_h = key_cache + ((uint64_t)s * DS4_N_HEAD + h) * qk_dim; float dot = 0.0f; for (uint32_t i = 0; i < qk_dim; i++) dot += q_h[i] * k_h[i]; scores[s] = dot * scale; if (scores[s] > max_score) max_score = scores[s]; } float denom = 0.0f; for (uint32_t s = 0; s < visible; s++) { scores[s] = expf(scores[s] - max_score); denom += scores[s]; } if (denom < 1.0e-20f) denom = 1.0e-20f; float *head_out = heads + ((uint64_t)t * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; for (uint32_t d = 0; d < DS4_N_VALUE_MLA; d++) { float acc = 0.0f; for (uint32_t s = 0; s < visible; s++) { const float *v_h = value_cache + ((uint64_t)s * DS4_N_HEAD + h) * DS4_N_VALUE_MLA; acc += scores[s] * v_h[d]; } head_out[d] = acc / denom; } } matvec_q8_0_f32_ref(out + (uint64_t)t * DS4_N_EMBD, model, layer->attn_output, heads + (uint64_t)t * heads_dim); } free(scores); free(k_rot); free(heads); free(value_cache); free(key_cache); free(k_nope); free(kv_norm); free(kv_raw); free(q); free(q_rank_norm); free(q_rank); free(norm); } static void layer_glm_dense_ffn_one( float * out, const ds4_model * model, const ds4_layer_weights * layer, const float * x) { const uint64_t hidden = layer->ffn_gate->dim[1]; const uint64_t in_dim = layer->ffn_gate->dim[0]; const uint64_t blocks = (in_dim + 31) / 32; float *gate = xmalloc((size_t)hidden * sizeof(gate[0])); float *up = xmalloc((size_t)hidden * sizeof(up[0])); float *mid = xmalloc((size_t)hidden * sizeof(mid[0])); int8_t *xq = xmalloc((size_t)blocks * 32); float *xscale = xmalloc((size_t)blocks * sizeof(xscale[0])); if (layer->ffn_gate->type != DS4_TENSOR_Q8_0 || layer->ffn_up->type != DS4_TENSOR_Q8_0 || layer->ffn_down->type != DS4_TENSOR_Q8_0 || layer->ffn_up->dim[0] != in_dim || layer->ffn_up->dim[1] != hidden || layer->ffn_down->dim[0] != hidden || layer->ffn_down->dim[1] != DS4_N_EMBD) { ds4_die("GLM dense FFN tensors have an unexpected layout"); } quantize_q8_0_activation(x, xq, xscale, in_dim); matvec_q8_0_pair_prequant(gate, up, model, layer->ffn_gate, layer->ffn_up, xq, xscale); swiglu(mid, gate, up, hidden, DS4_SWIGLU_CLAMP_EXP); 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, DS4_SWIGLU_CLAMP_EXP); 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; float clamp; } 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 || gate_type == DS4_TENSOR_MXFP4); } 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 || down_type == DS4_TENSOR_MXFP4; } 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); } if (type == DS4_TENSOR_MXFP4) { return ds4_vec_dot_mxfp4_f32(n, (const block_mxfp4 *)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]; float gate = glm_routed_moe_dot_f32(ctx->gate_type, (int)ctx->in_dim, gate_row, ctx->x); float up = glm_routed_moe_dot_f32(ctx->up_type, (int)ctx->in_dim, up_row, ctx->x); 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]; } } 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, .clamp = DS4_SWIGLU_CLAMP_EXP, }; 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, DS4_SWIGLU_CLAMP_EXP); 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, DS4_SWIGLU_CLAMP_EXP); 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); static bool sample_build_probabilities(const float *logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, float *probs); static int sample_probabilities(const float *probs, uint32_t n_vocab, uint64_t *rng); #ifdef DS4_TEST_HOOKS static int sample_residual_probabilities(float *target_probs, const float *draft_probs, uint32_t n_vocab, uint64_t *rng); #endif /* ========================================================================= * 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 /* * Metal stores the persistent attention-compressed KV cache in F16. Producers * retain F32 staging, while checkpoints and debug dumps expand back to F32 for * the stable external format. */ #if defined(__APPLE__) #define DS4_GPU_ATTN_COMP_CACHE_F16 1 #else #define DS4_GPU_ATTN_COMP_CACHE_F16 0 #endif /* All GPU backends store GLM's persistent compact cache in F16. Producers * retain F32 staging, and the kernels receive the storage format explicitly. */ #define DS4_GPU_GLM_COMPACT_CACHE_F16 1 /* ========================================================================= * 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. Prefix slots retain the intermediate frontiers of a tiny * verifier block so partial accepts do not need target-model replay. */ 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_prefix_n_comp[DS4_SPEC_PREFIX_SLOTS][DS4_MAX_LAYER]; uint32_t spec_prefix_n_index_comp[DS4_SPEC_PREFIX_SLOTS][DS4_MAX_LAYER]; bool spec_capture_prefixes; 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 *index_comp_kv_cur_by_tier[DS4_MAX_GPUS]; ds4_gpu_tensor *index_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; /* DeepSeek Vision-Exp prefill routing state. The engine owns both the * sidecar mapping and this immutable descriptor for the graph lifetime. */ const void *vision_model_map; uint64_t vision_model_size; const ds4_deepseek4_vision_weights *deepseek4_vision_weights; const int32_t *prefill_host_tokens; uint32_t prefill_host_pos0; uint32_t prefill_host_count; bool prefill_has_visual; const ds4_vision_span *prefill_vision_spans; size_t prefill_vision_span_count; /* Deferred HC post/expand after a TP combine: recorded here and folded * into the next compound HC producer dispatch (one pass instead of two). * Anything else that reads the expanded state flushes it first. */ bool hc_expand_pending; ds4_gpu_tensor *hc_expand_out; const ds4_gpu_tensor *hc_expand_block_out; const ds4_gpu_tensor *hc_expand_block_add; const ds4_gpu_tensor *hc_expand_residual; const ds4_gpu_tensor *hc_expand_post; const ds4_gpu_tensor *hc_expand_comb; } 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(index_comp_kv_cur) DS4_GPU_GRAPH_CLASS_P_ACCESSOR(index_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 /* CPU-only builds do not include ds4_gpu.h: mirror the decode-graph key * type and stub the capture API (always eager). */ typedef struct ds4_decode_graph_key { uint32_t il; uint32_t island; uint32_t variant; uint32_t _pad; void *cur_hc; void *after_attn_hc; void *after_ffn_hc; void *attn_norm; } ds4_decode_graph_key; static inline int ds4_gpu_decode_graphs_supported(void) { return 0; } static inline int ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key) { (void)key; return -1; } static inline int ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key) { (void)key; return -1; } static inline void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key) { (void)key; } static inline void ds4_gpu_decode_graphs_invalidate(void) {} 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) { /* Captured decode-island graphs bake this graph's buffer addresses * into their kernel nodes; retire them before the buffers go away. */ ds4_gpu_decode_graphs_invalidate(); /* 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->index_comp_kv_cur_by_tier[t]); ds4_gpu_tensor_free(g->index_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 hidden-width direction per normal * transformer layer. The model's MTP predictor is not part of the file. When * enabled, the GPU 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 uint32_t n_layers = directional_steering_layer_count(); if (n_layers == 0) return false; const uint64_t n = (uint64_t)n_layers * DS4_N_EMBD; float *dirs = xmalloc((size_t)n * sizeof(dirs[0])); bool ok = read_f32_binary_file(path, dirs, n); if (ok) { /* Replicate the directions buffer onto every Class P tier slot that * has any other Class P scratch allocated (used_tier marker is the * presence of g->cur_hc_by_tier[t]). Single-tier: only slot 0. */ bool any = false; for (int t = 0; ok && t < DS4_MAX_GPUS; t++) { if (!g->cur_hc_by_tier[t]) continue; g->directional_steering_dirs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, n * sizeof(dirs[0])); ok = g->directional_steering_dirs_by_tier[t] != NULL && ds4_gpu_tensor_write(g->directional_steering_dirs_by_tier[t], 0, dirs, n * sizeof(dirs[0])) != 0; if (ok) any = true; } if (ok && !any) { /* No used tiers — graph not allocated yet. This shouldn't happen * given the call site ordering, but bail cleanly. */ ok = false; } } free(dirs); if (!ok) { fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); return false; } g->directional_steering_attn_scale = attn_scale; g->directional_steering_ffn_scale = ffn_scale; fprintf(stderr, "ds4: directional steering enabled: %s attn=%g ffn=%g\n", path, (double)attn_scale, (double)ffn_scale); return true; } static bool metal_graph_directional_steering_attn_enabled(const ds4_gpu_graph *g) { return g && metal_graph_directional_steering_dirs(g) && g->directional_steering_attn_scale != 0.0f; } static bool metal_graph_directional_steering_ffn_enabled(const ds4_gpu_graph *g) { return g && metal_graph_directional_steering_dirs(g) && g->directional_steering_ffn_scale != 0.0f; } static bool metal_graph_apply_directional_steering( ds4_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows, float scale) { if (!g || !metal_graph_directional_steering_dirs(g) || scale == 0.0f) return true; return ds4_gpu_directional_steering_project_tensor(x, metal_graph_directional_steering_dirs(g), il, DS4_N_EMBD, rows, scale) != 0; } static bool metal_graph_apply_directional_steering_attn( ds4_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows) { return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_attn_scale : 0.0f); } static bool metal_graph_apply_directional_steering_ffn( ds4_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows) { return metal_graph_apply_directional_steering(g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f); } static bool metal_graph_configure_dspark_capture( ds4_gpu_graph *g, const ds4_dspark_weights *dw) { if (!g || !dw || dw->target_layer_count == 0) return true; if (dw->target_layer_count > DS4_DSPARK_MAX_TARGET_LAYERS || DS4_N_HC == 0 || DS4_N_HC > DS4_MAX_HC) { return false; } g->dspark_hc_mean_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); g->dspark_hc_mean_rows = ds4_gpu_tensor_alloc((uint64_t)g->prefill_cap * DS4_N_HC * sizeof(float)); g->dspark_target_hidden = ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * DS4_N_EMBD * sizeof(float)); g->dspark_target_hidden_batch = ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * g->prefill_cap * DS4_N_EMBD * sizeof(float)); if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { g->dspark_stage0_packed = ds4_gpu_tensor_alloc(((uint64_t)dw->block_size + 1u) * dw->target_layer_count * DS4_N_EMBD * sizeof(float)); } g->dspark_stage0_proj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); g->dspark_main_x = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); if (!g->dspark_hc_mean_weights || !g->dspark_hc_mean_rows || !g->dspark_target_hidden || !g->dspark_target_hidden_batch || !g->dspark_stage0_proj || !g->dspark_main_x) { return false; } if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; g->dspark_draft_tokens = ds4_gpu_tensor_alloc((uint64_t)dw->block_size * sizeof(int32_t)); g->dspark_draft_hc = ds4_gpu_tensor_alloc((uint64_t)dw->block_size * hc_dim * sizeof(float)); g->dspark_target_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->dspark_stage_input_hc = ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * hc_dim * sizeof(float)); g->dspark_stage_output_hc = ds4_gpu_tensor_alloc((uint64_t)dw->block_size * hc_dim * sizeof(float)); g->dspark_position_ids = ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * sizeof(int32_t)); if (!g->dspark_draft_tokens || !g->dspark_draft_hc || !g->dspark_target_hc || !g->dspark_stage_input_hc || !g->dspark_stage_output_hc || !g->dspark_position_ids) { return false; } if (dw->n_stages != 0 && g->raw_cap != 0) { for (uint32_t stage = 0; stage < dw->n_stages; stage++) { g->dspark_raw_cache[stage] = ds4_gpu_tensor_alloc((uint64_t)g->raw_cap * DS4_N_HEAD_DIM * sizeof(float)); if (!g->dspark_raw_cache[stage]) return false; } g->dspark_cache_cap = g->raw_cap; g->dspark_cache_start = 0; g->dspark_cache_token_start = 0; g->dspark_cache_len = 0; } g->dspark_block_size = dw->block_size; } float mean[DS4_MAX_HC] = {0}; const float inv_hc = 1.0f / (float)DS4_N_HC; for (uint32_t i = 0; i < DS4_N_HC; i++) mean[i] = inv_hc; if (ds4_gpu_tensor_write(g->dspark_hc_mean_weights, 0, mean, (uint64_t)DS4_N_HC * sizeof(mean[0])) == 0) { return false; } const uint64_t mean_rows_count = (uint64_t)g->prefill_cap * DS4_N_HC; if (mean_rows_count == 0 || mean_rows_count > (uint64_t)SIZE_MAX / sizeof(float)) { return false; } float *mean_rows = xmalloc((size_t)mean_rows_count * sizeof(mean_rows[0])); for (uint64_t i = 0; i < mean_rows_count; i++) mean_rows[i] = inv_hc; const bool mean_rows_ok = ds4_gpu_tensor_write(g->dspark_hc_mean_rows, 0, mean_rows, mean_rows_count * sizeof(mean_rows[0])) != 0; free(mean_rows); if (!mean_rows_ok) return false; g->dspark_target_layer_count = dw->target_layer_count; memcpy(g->dspark_target_layers, dw->target_layers, (size_t)dw->target_layer_count * sizeof(g->dspark_target_layers[0])); g->dspark_capture_mask = 0; g->dspark_capture_checkpoint_len = 0; g->dspark_capture_batch_mask = 0; g->dspark_capture_batch_start = 0; g->dspark_capture_batch_tokens = 0; g->dspark_capture_valid = false; g->dspark_capture_batch_valid = false; g->dspark_capture_enabled = true; return true; } static uint64_t metal_graph_kv_cache_bytes_for_context(uint32_t ctx_size, uint32_t raw_cap) { uint64_t bytes = (uint64_t)DS4_N_LAYER * raw_cap * DS4_N_HEAD_DIM * sizeof(float); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint64_t comp_cap = (uint64_t)(ctx_size / ratio + 2u); bytes += comp_cap * DS4_N_HEAD_DIM * (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)); if (ratio == 4) { bytes += comp_cap * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } } return bytes; } static uint64_t metal_graph_context_bytes_for_kv_policy( uint32_t ctx_size, uint32_t raw_cap, uint32_t prefill_cap, uint64_t *kv_cache_bytes_out) { uint32_t min_ratio = UINT32_MAX; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; uint64_t comp_cap = (uint64_t)(ctx_size / min_ratio + 2u); if (comp_cap < 2u) comp_cap = 2u; const uint64_t kv_cache_bytes = metal_graph_kv_cache_bytes_for_context(ctx_size, raw_cap); if (kv_cache_bytes_out) *kv_cache_bytes_out = kv_cache_bytes; uint64_t bytes = kv_cache_bytes + 2ull * comp_cap * prefill_cap * sizeof(float); if (DS4_GPU_ATTN_COMP_CACHE_F16) { uint64_t attn_stage_cap = (uint64_t)(prefill_cap / min_ratio + 2u); if (attn_stage_cap < 2u) attn_stage_cap = 2u; bytes += attn_stage_cap * DS4_N_HEAD_DIM * sizeof(float); } return bytes; } static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor_on( bool managed, int tier, uint64_t bytes) { if (g_n_gpus <= 1) { return managed ? ds4_gpu_tensor_alloc_managed(bytes) : ds4_gpu_tensor_alloc(bytes); } return managed ? ds4_gpu_tensor_alloc_managed_on(tier, bytes) : ds4_gpu_tensor_alloc_ptr_on(tier, bytes); } static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor(bool managed, uint64_t bytes) { return metal_graph_alloc_kv_cache_tensor_on(managed, 0, bytes); } /* ========================================================================= * Metal Diagnostic Dump Hooks. * ========================================================================= * * The release path calls these after important stages, but they are no-ops * unless DS4_METAL_GRAPH_DUMP_PREFIX or DS4_ROCM_GRAPH_DUMP_PREFIX is set. * Dumping synchronizes and restarts the command batch, so it is intentionally * isolated here. */ typedef struct { int init; const char *prefix; const char *name; int layer_set; uint32_t layer; int pos_set; uint32_t pos; } metal_graph_debug_config; static const metal_graph_debug_config *metal_graph_debug_get_config(void) { static metal_graph_debug_config cfg; if (!cfg.init) { cfg.init = 1; cfg.prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", "DS4_METAL_GRAPH_DUMP_PREFIX"); if (cfg.prefix && !cfg.prefix[0]) cfg.prefix = NULL; cfg.name = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_NAME", "DS4_METAL_GRAPH_DUMP_NAME"); if (cfg.name && !cfg.name[0]) cfg.name = NULL; const char *layer_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_LAYER", "DS4_METAL_GRAPH_DUMP_LAYER"); if (layer_env && layer_env[0] && strcmp(layer_env, "all") != 0) { cfg.layer_set = 1; cfg.layer = (uint32_t)strtoul(layer_env, NULL, 10); } const char *pos_env = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_POS", "DS4_METAL_GRAPH_DUMP_POS"); if (pos_env && pos_env[0]) { cfg.pos_set = 1; cfg.pos = (uint32_t)strtoul(pos_env, NULL, 10); } } return &cfg; } static const char *metal_graph_debug_prefix_for(const char *name, uint32_t il, uint32_t pos) { const metal_graph_debug_config *cfg = metal_graph_debug_get_config(); if (!cfg->prefix) return NULL; if (cfg->name && strstr(cfg->name, name) == NULL) return NULL; if (cfg->layer_set && cfg->layer != il) return NULL; if (cfg->pos_set && cfg->pos != pos) return NULL; return cfg->prefix; } static bool metal_graph_debug_wants(const char *name, uint32_t il, uint32_t pos) { return metal_graph_debug_prefix_for(name, il, pos) != NULL; } static void metal_graph_debug_dump_tensor( const char *name, ds4_gpu_tensor *t, uint64_t n_f32, uint32_t il, uint32_t pos) { const char *prefix = metal_graph_debug_prefix_for(name, il, pos); if (glm_graph_env_present("DS4_ROCM_GRAPH_DUMP_TRACE", "DS4_METAL_GRAPH_DUMP_TRACE")) fprintf(stderr, "ds4: dump? name=%s il=%u pos=%u t=%p n=%llu wants=%d\n", name, il, pos, (void *)t, (unsigned long long)n_f32, metal_graph_debug_wants(name, il, pos)); if (!t || n_f32 == 0 || !metal_graph_debug_wants(name, il, pos)) return; if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); return; } float *buf = xmalloc((size_t)n_f32 * sizeof(buf[0])); if (ds4_gpu_tensor_read(t, 0, buf, n_f32 * sizeof(buf[0])) != 0) { char path[1024]; snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); if (write_f32_binary_file(path, buf, n_f32)) { fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); } } free(buf); if (ds4_gpu_begin_commands() == 0) { fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); } } static void metal_graph_debug_dump_f16_tensor( const char *name, ds4_gpu_tensor *t, uint64_t n_f16, uint32_t il, uint32_t pos) { const char *prefix = glm_graph_env_value("DS4_ROCM_GRAPH_DUMP_PREFIX", "DS4_METAL_GRAPH_DUMP_PREFIX"); if (!t || n_f16 == 0 || !metal_graph_debug_wants(name, il, pos)) return; if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); return; } uint16_t *hbuf = xmalloc((size_t)n_f16 * sizeof(hbuf[0])); float *fbuf = xmalloc((size_t)n_f16 * sizeof(fbuf[0])); if (ds4_gpu_tensor_read(t, 0, hbuf, n_f16 * sizeof(hbuf[0])) != 0) { for (uint64_t i = 0; i < n_f16; i++) fbuf[i] = f16_to_f32(hbuf[i]); char path[1024]; snprintf(path, sizeof(path), "%s_%s-%u_pos%u.bin", prefix, name, il, pos); if (write_f32_binary_file(path, fbuf, n_f16)) { fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); } } free(fbuf); free(hbuf); if (ds4_gpu_begin_commands() == 0) { fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); } } static void metal_graph_debug_dump_i32_tensor( const char *name, ds4_gpu_tensor *t, uint64_t n_i32, uint32_t il, uint32_t pos) { if (!t || n_i32 == 0) return; const char *prefix = metal_graph_debug_prefix_for(name, il, pos); if (!prefix) return; if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); return; } int32_t *buf = xmalloc((size_t)n_i32 * sizeof(buf[0])); if (ds4_gpu_tensor_read(t, 0, buf, n_i32 * sizeof(buf[0])) != 0) { char path[1024]; snprintf(path, sizeof(path), "%s_%s-%u_pos%u.i32", prefix, name, il, pos); FILE *fp = fopen(path, "wb"); if (fp) { if (fwrite(buf, sizeof(buf[0]), (size_t)n_i32, fp) == (size_t)n_i32) { fprintf(stderr, "ds4: dumped %s layer %u pos %u to %s\n", name, il, pos, path); } fclose(fp); } } free(buf); if (ds4_gpu_begin_commands() == 0) { fprintf(stderr, "ds4: failed to resume Metal command batch after dumping %s layer %u pos %u\n", name, il, pos); } } static bool metal_graph_needs_ffn_out(const ds4_gpu_graph *g, uint32_t il, uint32_t pos) { return metal_graph_directional_steering_ffn_enabled(g) || g->materialize_ffn_out || metal_graph_debug_wants("ffn_out", il, pos); } /* tier-aware lazy allocator. The Class P ffn_out scratch is * created on demand the first time a layer that materializes ffn_out runs * on a tier; subsequent visits to the same tier reuse the existing slot. * Single-tier paths: active_tier == 0 always, behavior unchanged. */ static bool metal_graph_ensure_ffn_out(ds4_gpu_graph *g) { const int t = g->active_tier; if (!g->ffn_out_by_tier[t]) { g->ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( t, (uint64_t)DS4_N_EMBD * sizeof(float)); } return g->ffn_out_by_tier[t] != NULL; } static bool metal_graph_ensure_batch_ffn_out_on(ds4_gpu_graph *g, int t) { if (t < 0 || t >= DS4_MAX_GPUS) return false; if (!g->batch_ffn_out_by_tier[t]) { g->batch_ffn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on( t, (uint64_t)g->prefill_cap * DS4_N_EMBD * sizeof(float)); } return g->batch_ffn_out_by_tier[t] != NULL; } static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) { return metal_graph_ensure_batch_ffn_out_on(g, g->active_tier); } static bool metal_graph_tp_env_flag(const char *name, bool dflt) { const char *env = getenv(name); if (!env || !env[0]) return dflt; return strcmp(env, "0") != 0; } static bool metal_graph_cuda_tp_attn_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN", true); #endif } static bool metal_graph_cuda_tp_attn_peer_read_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_PEER_READ", true); #endif } static bool metal_graph_cuda_tp_attn_heads_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_HEADS", false); #endif } static bool metal_graph_cuda_tp_attn_cache_dup_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_ATTN_CACHE_DUP", false); #endif } static bool metal_graph_cuda_tp_moe_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE", true); #endif } static bool metal_graph_cuda_tp_ep_pack_exact_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_PACK_EXACT", true); #endif } static bool metal_graph_cuda_tp_ep_direct_return_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DIRECT_RETURN", true); #endif } static bool metal_graph_cuda_tp_ep_delay_reduce_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_DELAY_REDUCE", true); #endif } static bool metal_graph_cuda_tp_ep_fused_hc_reduce_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_HC_REDUCE", true); #endif } static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_fused_shared_mid_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_EP_FUSED_SHARED_MID", true); #endif } static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_balanced_shared_mid_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag( "DS4_CUDA_TP_EP_BALANCED_SHARED_MID", true); #endif } static DS4_MAYBE_UNUSED bool metal_graph_cuda_tp_ep_dual_prequant_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag( "DS4_CUDA_TP_EP_DUAL_PREQUANT", true); #endif } static bool metal_graph_cuda_tp_moe_delay_reduce_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_DELAY_REDUCE", true); #endif } static bool metal_graph_cuda_tp_moe_pack_handoff_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PACK", false); #endif } static bool metal_graph_cuda_tp_moe_copy3_handoff_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_COPY3", false); #endif } static bool metal_graph_cuda_tp_moe_peer_read_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_READ", false); #endif } static bool metal_graph_cuda_tp_moe_peer_router_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_MOE_PEER_ROUTER", false); #endif } static bool metal_graph_cuda_tp_shared_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED", false); #endif } static bool metal_graph_cuda_tp_shared_fold_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_SHARED_FOLD", true); #endif } static bool metal_graph_cuda_tp_q_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_Q", false); #endif } static bool metal_graph_cuda_tp_output_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_OUTPUT", true); #endif } static bool metal_graph_cuda_greedy_split_top1_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLIT_TOP1"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLIT_TOP1", false); #endif } static bool metal_graph_cuda_output_fused_top1_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_OUTPUT_FUSED_TOP1", false); #endif } static bool metal_graph_cuda_verify_decode2_split_top1_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_VERIFY_DECODE2_SPLIT_TOP1"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_VERIFY_DECODE2_SPLIT_TOP1", false); #endif } static bool metal_graph_cuda_greedy_splitkv_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV", false); #endif } static bool metal_graph_cuda_greedy_vec4_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_VEC4", false); #endif } static bool metal_graph_cuda_splitkv_spec_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC", false); #endif } static bool metal_graph_cuda_splitkv_spec_toponly_row0_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_TOPONLY_ROW0"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_TOPONLY_ROW0", false); #endif } static bool metal_graph_cuda_splitkv_spec_batch_verify_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_SPLITKV_SPEC_BATCH_VERIFY"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_SPEC_BATCH_VERIFY", false); #endif } static float metal_graph_cuda_greedy_vec4_margin_threshold(void) { #if defined(__APPLE__) return 0.0f; #else const char *env = getenv("DS4_CUDA_GREEDY_VEC4_MARGIN"); if (env && env[0]) { char *end = NULL; double v = strtod(env, &end); while (end && isspace((unsigned char)*end)) end++; if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { return (float)v; } fprintf(stderr, "ds4: invalid DS4_CUDA_GREEDY_VEC4_MARGIN=%s; using 0.25\n", env); } return 0.25f; #endif } static bool metal_graph_cuda_greedy_vec4_fallback_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_VEC4_FALLBACK"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_cuda_greedy_vec4_margin_threshold() > 0.0f; #endif } static float metal_graph_cuda_greedy_splitkv_margin_threshold(void) { #if defined(__APPLE__) return 0.0f; #else const char *env = getenv("DS4_CUDA_GREEDY_SPLITKV_MARGIN"); if (env && env[0]) { char *end = NULL; double v = strtod(env, &end); while (end && isspace((unsigned char)*end)) end++; if (end != env && end && *end == '\0' && isfinite(v) && v >= 0.0) { return (float)v; } fprintf(stderr, "ds4: invalid DS4_CUDA_GREEDY_SPLITKV_MARGIN=%s; using 0.25\n", env); } return 0.25f; #endif } static bool metal_graph_cuda_greedy_splitkv_fallback_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_FALLBACK"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_cuda_greedy_splitkv_margin_threshold() > 0.0f; #endif } static bool metal_graph_cuda_greedy_splitkv_top2_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_TOP2"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TOP2", true); #endif } static bool metal_graph_cuda_greedy_splitkv_trust_replay_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_TRUST_REPLAY", false); #endif } static bool metal_graph_cuda_greedy_splitkv_pair_replay_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_GREEDY_SPLITKV_PAIR_REPLAY"); if (no && no[0] && strcmp(no, "0") != 0) return false; return metal_graph_tp_env_flag("DS4_CUDA_GREEDY_SPLITKV_PAIR_REPLAY", false); #endif } static DS4_MAYBE_UNUSED uint32_t metal_graph_cuda_greedy_max_segment(const char *name) { const char *env = getenv(name); if (env && env[0]) { char *end = NULL; unsigned long v = strtoul(env, &end, 10); while (end && isspace((unsigned char)*end)) end++; if (end != env && end && *end == '\0' && v <= INT32_MAX) { return (uint32_t)v; } fprintf(stderr, "ds4: invalid %s=%s; expected 0..%d, using disabled\n", name, env, INT32_MAX); } return 0; } static uint32_t metal_graph_cuda_greedy_splitkv_max_segment(void) { #if defined(__APPLE__) return 0; #else return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_SPLITKV_MAX_SEGMENT"); #endif } static uint32_t metal_graph_cuda_greedy_vec4_max_segment(void) { #if defined(__APPLE__) return 0; #else return metal_graph_cuda_greedy_max_segment("DS4_CUDA_GREEDY_VEC4_MAX_SEGMENT"); #endif } static uint32_t metal_graph_cuda_greedy_splitkv_min_score(void) { const char *env = getenv("DS4_CUDA_SPLITKV_MIN_SCORE"); if (env && env[0]) { char *end = NULL; unsigned long v = strtoul(env, &end, 10); while (end && isspace((unsigned char)*end)) end++; if (end != env && end && *end == '\0' && v <= UINT32_MAX) { return (uint32_t)v; } } return metal_graph_tp_env_flag("DS4_CUDA_SPLITKV_DECODE", false) ? 0u : 512u; } static bool metal_graph_cuda_q_norm_rope_fuse_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_Q_NORM_ROPE_FUSE", true); #endif } static bool metal_graph_cuda_qkv_kv_rope_fuse_requested(void) { #if defined(__APPLE__) return false; #else const char *no = getenv("DS4_CUDA_NO_QKV_KV_ROPE_FUSE"); if (no && no[0] && strcmp(no, "0") != 0) return false; if (getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL) return false; return metal_graph_tp_env_flag("DS4_CUDA_QKV_KV_ROPE_FUSE", true); #endif } static bool metal_graph_cuda_tp_prefill_ffn_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_FFN", true); #endif } static bool metal_graph_cuda_tp_prefill_attn_output_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_TP_PREFILL_ATTN_OUTPUT", true); #endif } static bool metal_graph_cuda_prefill_pipeline_requested(const ds4_gpu_graph *g) { #if defined(__APPLE__) (void)g; return false; #else const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE"); if (env && env[0]) return strcmp(env, "0") != 0; return g && g->cuda_tp_decode; #endif } static bool metal_graph_cuda_prefill_pipeline_q8_cache_requested(void) { #if defined(__APPLE__) return false; #else return metal_graph_tp_env_flag("DS4_CUDA_PREFILL_PIPELINE_Q8_CACHE", false); #endif } static uint32_t metal_graph_cuda_prefill_pipeline_microbatch(void) { const char *env = getenv("DS4_CUDA_PREFILL_PIPELINE_MB"); if (env && env[0]) { char *end = NULL; unsigned long v = strtoul(env, &end, 10); if (end != env && v > 0 && v <= UINT32_MAX) return (uint32_t)v; } return 512; } static int metal_graph_cuda_tp_partner_tier(int tier) { if (g_n_gpus < 2 || (g_n_gpus & 1) != 0) return -1; const int half = g_n_gpus / 2; if (tier < 0 || tier >= half) return -1; return tier + half; } static uint32_t metal_graph_cuda_tp_output_tiers( const ds4_gpu_graph *g, int tiers[DS4_MAX_GPUS]) { if (!g) return 0; return metal_graph_cuda_tp_output_tiers_for_head(g->head_tier, g->cuda_tp_output, g_n_gpus, tiers); } static uint64_t metal_graph_q8_0_row_bytes(uint64_t in_dim) { return ((in_dim + 31u) / 32u) * 34u; } /* ========================================================================= * Metal Release Graph Allocation. * ========================================================================= */ /* Allocate the Metal graph state for a chosen raw-cache capacity. The model * weights are not copied here; tensors reference the mapped GGUF. * * tier-aware per-layer allocation. * placement: when non-NULL, an array of DS4_N_LAYER + 2 logical tiers * (embedding, per-layer..., head). The per-layer KV / state allocations * in this function use placement[il + 1] as the home tier for each * layer il. When NULL (single-tier callers, diagnostic paths), all * per-layer allocations land on tier 0 — byte-equivalent to legacy. * * Single-tier (g_n_gpus <= 1) is byte-equivalent regardless of placement, * because metal_graph_alloc_kv_cache_tensor_on short-circuits to the * legacy 1-arg helpers when g_n_gpus <= 1. */ static bool metal_graph_alloc_raw_cap( ds4_gpu_graph *g, const ds4_weights *weights, const ds4_layer_weights *layer, uint32_t raw_cap, uint32_t ctx_size, uint32_t prefill_cap, bool enable_mtp, const int *placement, bool cuda_tensor_parallel, const ds4_gpu_graph *shared_prefill_workspace) { const int saved_dspark_exec_tier = g->dspark_exec_tier; memset(g, 0, sizeof(*g)); g->dspark_exec_tier = saved_dspark_exec_tier; g->owns_prefill_workspace = shared_prefill_workspace == NULL; g->cpu_router_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])); g->active_tier = placement ? -1 : 0; /* cache placement on the graph so the dispatch loops can * walk it without threading the engine pointer through every * kernel-dispatch wrapper. NULL in single-tier callers (placement * was already NULL on entry). */ g->placement = placement; g->cuda_tp_decode = placement && cuda_tensor_parallel; g->cuda_tp_attn = g->cuda_tp_decode && metal_graph_cuda_tp_attn_requested(); g->cuda_tp_attn_peer_read = metal_graph_cuda_tp_attn_peer_read_requested(); g->cuda_tp_attn_heads = g->cuda_tp_decode && metal_graph_cuda_tp_attn_heads_requested(); g->cuda_tp_attn_cache_dup = g->cuda_tp_attn_heads && metal_graph_cuda_tp_attn_cache_dup_requested(); g->cuda_tp_moe = g->cuda_tp_decode && metal_graph_cuda_tp_moe_requested(); g->cuda_tp_ep = g->cuda_tp_moe && cuda_tensor_parallel; g->cuda_tp_ep_pack_exact = g->cuda_tp_ep && metal_graph_cuda_tp_ep_pack_exact_requested(); g->cuda_tp_moe_delay_reduce = metal_graph_cuda_tp_moe_delay_reduce_requested(); g->cuda_tp_moe_copy3_handoff = metal_graph_cuda_tp_moe_copy3_handoff_requested(); g->cuda_tp_moe_pack_handoff = metal_graph_cuda_tp_moe_pack_handoff_requested(); g->cuda_tp_moe_peer_read = metal_graph_cuda_tp_moe_peer_read_requested(); g->cuda_tp_moe_peer_router = metal_graph_cuda_tp_moe_peer_router_requested(); g->cuda_tp_shared = g->cuda_tp_decode && metal_graph_cuda_tp_shared_requested(); g->cuda_tp_shared_fold = metal_graph_cuda_tp_shared_fold_requested(); g->cuda_tp_q = g->cuda_tp_decode && metal_graph_cuda_tp_q_requested(); g->cuda_tp_output = g->cuda_tp_decode && metal_graph_cuda_tp_output_requested(); g->cuda_tp_prefill_ffn = g->cuda_tp_decode && metal_graph_cuda_tp_prefill_ffn_requested(); g->cuda_tp_prefill_attn_output = g->cuda_tp_decode && metal_graph_cuda_tp_prefill_attn_output_requested(); g->cuda_q_norm_rope_fuse = metal_graph_cuda_q_norm_rope_fuse_requested(); g->cuda_qkv_kv_rope_fuse = metal_graph_cuda_qkv_kv_rope_fuse_requested(); g->cuda_qkv_pair = getenv("DS4_CUDA_NO_QKV_PAIR") == NULL; g->cuda_tp_attn_out_hc_fuse = getenv("DS4_CUDA_TP_ATTN_OUT_HC_FUSE") != NULL && getenv("DS4_CUDA_NO_TP_ATTN_OUT_HC_FUSE") == NULL; g->shared_gate_up_swiglu_fuse = getenv("DS4_METAL_DISABLE_SHARED_GATE_UP_SWIGLU_FUSION") == NULL; g->decode_stage_profile = getenv("DS4_METAL_DECODE_STAGE_PROFILE") != NULL; g->decode_index_stage_profile = getenv("DS4_METAL_INDEXER_STAGE_PROFILE") != NULL; g->output_stage_profile = getenv("DS4_METAL_OUTPUT_STAGE_PROFILE") != NULL; const bool enable_splitkv_spec = metal_graph_cuda_splitkv_spec_requested(); const bool enable_splitkv_batch_verify = enable_splitkv_spec && metal_graph_cuda_splitkv_spec_batch_verify_requested(); const bool enable_spec_logits = enable_mtp || enable_splitkv_batch_verify; const bool enable_prefix1_snapshot = enable_mtp || enable_splitkv_spec; const bool enable_frontier_snapshot = enable_mtp || enable_splitkv_spec || (metal_graph_cuda_greedy_splitkv_requested() && metal_graph_cuda_greedy_splitkv_fallback_requested()) || (metal_graph_cuda_greedy_vec4_requested() && metal_graph_cuda_greedy_vec4_fallback_requested()); if (g->cuda_tp_decode && metal_graph_cuda_tp_partner_tier(0) < 0) { fprintf(stderr, "ds4: CUDA tensor parallelism requires an even multi-GPU placement; " "have %d GPU tiers\n", g_n_gpus); metal_graph_free(g); return false; } if (g->cuda_tp_ep && (g_ds4_shape.family != DS4_MODEL_FAMILY_DEEPSEEK4 || (DS4_N_EXPERT & 1u) != 0u)) { fprintf(stderr, "ds4: CUDA tensor parallelism requires an even-expert DeepSeek model\n"); metal_graph_free(g); return false; } if (g->cuda_tp_ep) { fprintf(stderr, "ds4: CUDA routed MoE expert ownership enabled " "(half-resident decode and prefill)\n"); } g->mtp_enabled = enable_mtp; if (raw_cap == 0) raw_cap = 1; if (ctx_size == 0) ctx_size = raw_cap; if (prefill_cap == 0) prefill_cap = 1; uint32_t raw_window = DS4_N_SWA; if (raw_window > ctx_size) raw_window = ctx_size; if (raw_window == 0) raw_window = 1; if (raw_cap < raw_window) raw_cap = raw_window; if (raw_cap > ctx_size) raw_cap = ctx_size; if (raw_cap == 0) raw_cap = 1; g->raw_cap = raw_cap; g->raw_window = raw_window; g->prefill_cap = prefill_cap; uint32_t min_ratio = UINT32_MAX; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!weights_layer_has_required(&weights->layer[il], il)) continue; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } if (min_ratio == UINT32_MAX) min_ratio = ctx_size ? ctx_size : 1u; g->comp_cap = ctx_size / min_ratio + 2u; if (g->comp_cap < 2u) g->comp_cap = 2u; if (DS4_GPU_ATTN_COMP_CACHE_F16) { g->attn_comp_stage_cap = prefill_cap / min_ratio + 2u; if (g->attn_comp_stage_cap < 2u) g->attn_comp_stage_cap = 2u; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!weights_layer_has_required(&weights->layer[il], il)) { g->layer_comp_cap[il] = 0; continue; } const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) { g->layer_comp_cap[il] = 0; } else { g->layer_comp_cap[il] = ctx_size / ratio + 2u; if (g->layer_comp_cap[il] < 2u) g->layer_comp_cap[il] = 2u; } } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; const uint64_t q_rank = layer->attn_q_a->dim[1]; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; /* Distributed coordinators do not normally own the output head. The * logits workspace still has a fixed model-vocabulary shape, while the * actual head is encoded only on a node that bound its tensors. */ const uint64_t vocab_dim = weights->output ? weights->output->dim[1] : DS4_N_VOCAB; const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM ? DS4_N_HEAD_DIM : DS4_N_INDEXER_HEAD_DIM); const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; const uint64_t pc = prefill_cap; uint64_t kv_cache_bytes = 0; const uint64_t context_bytes = metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes); const bool managed_kv_cache = ds4_gpu_should_use_managed_kv_cache(kv_cache_bytes, context_bytes) != 0; if (managed_kv_cache) { /* * CUDA device allocations are fastest, but a million-token KV cache is * large enough to starve DGX Spark's unified CPU/GPU memory once the * model cache and driver allocations are present. For this one * long-lived cache class, managed memory restores the old demand-paged * behavior. It can be slower, but it keeps oversized contexts from * turning memory pressure into a machine-wide lockup. */ fprintf(stderr, "ds4: CUDA using managed KV cache for ctx=%u " "(kv cache %.2f GiB, context buffers %.2f GiB); " "this may degrade performance but is needed for very large contexts\n", ctx_size, (double)kv_cache_bytes / 1073741824.0, (double)context_bytes / 1073741824.0); } /* Class P decode HC scratch — replicated across every tier * the placement uses (per-tier kernel-scratch). Single-tier path * (placement == NULL) collapses to tier 0 only; _ptr_on(0, ...) short- * circuits to legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte- * equivalent. The hc_pre/hc_post/hc_comb buffers are VIEWS of hc_split * and therefore allocated per tier alongside their parent. */ bool used_tier[DS4_MAX_GPUS] = {0}; used_tier[0] = true; /* single-tier baseline always uses tier 0 */ if (placement) { for (uint32_t i = 0; i < (uint32_t)DS4_N_LAYER + 2u; i++) { const int p = placement[i]; if (p >= 0 && p < DS4_MAX_GPUS) used_tier[p] = true; } } if (g->cuda_tp_decode) { const int half = g_n_gpus / 2; for (int t = half; t < g_n_gpus; t++) { if (used_tier[t]) { fprintf(stderr, "ds4: CUDA tensor parallelism expects layer homes in lower-half " "tiers; placement already uses tier %d\n", t); metal_graph_free(g); return false; } } for (int t = 0; t < half; t++) { if (used_tier[t]) used_tier[t + half] = true; } fprintf(stderr, "ds4: CUDA decode TP enabled: pairing lower-half tiers with " "upper-half tiers\n"); } for (int t = 0; t < DS4_MAX_GPUS; t++) { if (!used_tier[t]) continue; g->cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); g->flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, hc_dim * sizeof(float)); g->hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); g->hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, mix_hc * sizeof(float)); g->hc_pre_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], 0, (uint64_t)DS4_N_HC * sizeof(float)); g->hc_post_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], (uint64_t)DS4_N_HC * sizeof(float), (uint64_t)DS4_N_HC * sizeof(float)); g->hc_comb_by_tier[t] = ds4_gpu_tensor_view(g->hc_split_by_tier[t], 2ull * DS4_N_HC * sizeof(float), (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); g->attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); g->attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_EMBD * sizeof(float)); g->qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); g->qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_rank * sizeof(float)); g->q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, q_dim * sizeof(float)); g->kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); g->kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); } bool state_init_ok = true; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { /* A distributed process owns only its bound layer slice. Persistent * KV state must follow that ownership just like the model tensors; * allocating every model layer here defeats split-model residency. */ if (!weights_layer_has_required(&weights->layer[il], il)) continue; /* per-layer Class L allocations land on the layer's * home tier. placement is NULL on single-tier / diagnostic paths * (all-tier-0); non-NULL on the engine path that opted into * multi-tier. layer_tier == 0 in single-tier mode is the * byte-equivalent path through metal_graph_alloc_kv_cache_tensor_on * and ds4_gpu_tensor_alloc_ptr_on. */ const int layer_tier = placement ? placement[il + 1] : 0; g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor_on( managed_kv_cache, layer_tier, (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); const int layer_tp_partner = g->cuda_tp_attn_cache_dup ? metal_graph_cuda_tp_partner_tier(layer_tier) : -1; if (layer_tp_partner >= 0) { g->layer_raw_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( managed_kv_cache, layer_tp_partner, (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); } const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0) { const uint32_t coff = ratio == 4 ? 2u : 1u; const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; const uint64_t attn_rows = (uint64_t)coff * ratio; g->layer_attn_comp_cache[il] = metal_graph_alloc_kv_cache_tensor_on( managed_kv_cache, layer_tier, (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); if (layer_tp_partner >= 0) { g->layer_attn_comp_cache_tp[il] = metal_graph_alloc_kv_cache_tensor_on( managed_kv_cache, layer_tp_partner, (uint64_t)g->layer_comp_cap[il] * DS4_N_HEAD_DIM * (DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float))); } g->layer_attn_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); g->layer_attn_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); if (enable_frontier_snapshot) { g->spec_attn_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); g->spec_attn_state_score[il] = ds4_gpu_tensor_alloc_ptr_on(layer_tier, attn_width * attn_rows * sizeof(float)); if (enable_prefix1_snapshot) { g->spec_prefix1_attn_state_kv[il] = ds4_gpu_tensor_alloc_ptr_on( layer_tier, DS4_SPEC_PREFIX_SLOTS * attn_width * attn_rows * sizeof(float)); g->spec_prefix1_attn_state_score[il] = ds4_gpu_tensor_alloc_ptr_on( layer_tier, DS4_SPEC_PREFIX_SLOTS * 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, DS4_SPEC_PREFIX_SLOTS * index_width * index_rows * sizeof(float)); g->spec_prefix1_index_state_score[il] = ds4_gpu_tensor_alloc_ptr_on( layer_tier, DS4_SPEC_PREFIX_SLOTS * 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)); /* Decode-only scratch for the quad compressor projection: the indexer * pair outputs must not alias the attention compressor outputs inside * the single fused dispatch. */ g->index_comp_kv_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, 2ull * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); g->index_comp_sc_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, 2ull * DS4_N_INDEXER_HEAD_DIM * 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-model must allocate and execute exactly the same * buffers as the plain decoder: no support-model mapping, no draft logits, * and no MTP scratch hidden behind otherwise unused tensors. */ if (enable_mtp) { g->mtp_embed = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); g->mtp_enorm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); g->mtp_eproj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); g->mtp_eproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_hnorm_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_hproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor( managed_kv_cache, (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); g->mtp_n_raw = 0; } if (enable_spec_logits) { const int spec_tier = g->dspark_exec_tier > 0 && g->dspark_exec_tier < DS4_MAX_GPUS ? g->dspark_exec_tier : 0; g->spec_logits = spec_tier ? ds4_gpu_tensor_alloc_ptr_on(spec_tier, (uint64_t)16 * DS4_N_VOCAB * sizeof(float)) : ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); } /* Class E — emb_tier captured from placement[0] (or 0 in * single-tier / diagnostic paths). _ptr_on(0, ...) short-circuits to the * legacy ds4_gpu_tensor_alloc when g_n_gpus <= 1 — byte-equivalent. */ g->emb_tier = placement ? placement[0] : 0; /* Class P chunked-prefill batch scratch — replicated across * every used tier. The cur/next pair (batch_cur_hc / batch_next_hc) is * ping-ponged per layer step on each tier; tier transitions copy via * ds4_gpu_tensor_copy_xdev (handled in B6). batch_ffn_out is lazily * allocated by metal_graph_ensure_batch_ffn_out (per-tier on first touch) * and included in the CUDA scratch estimate because TP prefill can use it * as the combined routed+shared FFN buffer. */ if (shared_prefill_workspace) { if (shared_prefill_workspace->prefill_cap < prefill_cap || shared_prefill_workspace->emb_tier != g->emb_tier) { fprintf(stderr, "ds4: shared prefill workspace is incompatible " "(capacity %u/%u, embedding tier %d/%d)\n", shared_prefill_workspace->prefill_cap, prefill_cap, shared_prefill_workspace->emb_tier, g->emb_tier); } else { metal_graph_copy_prefill_workspace_pointers( g, shared_prefill_workspace); } } else { g->prefill_tokens_by_tier[g->emb_tier] = ds4_gpu_tensor_alloc_ptr_on(g->emb_tier, pc * sizeof(int32_t)); for (int t = 0; t < DS4_MAX_GPUS; t++) { if (!used_tier[t]) continue; g->batch_cur_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); g->batch_next_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); g->batch_flat_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); g->batch_hc_mix_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); g->batch_hc_split_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * mix_hc * sizeof(float)); g->batch_attn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_attn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_qr_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); g->batch_qr_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_rank * sizeof(float)); g->batch_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); g->batch_kv_raw_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); g->batch_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_HEAD_DIM * sizeof(float)); g->batch_comp_kv_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); g->batch_comp_sc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * comp_width_max * sizeof(float)); g->batch_indexer_q_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * indexer_q_dim * sizeof(float)); g->batch_indexer_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_INDEXER_HEAD * sizeof(float)); g->batch_heads_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * q_dim * sizeof(float)); g->batch_attn_low_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * low_dim * sizeof(float)); g->batch_attn_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_group_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * group_dim * sizeof(float)); g->batch_low_tmp_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_LORA_O * sizeof(float)); g->batch_after_attn_hc_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * hc_dim * sizeof(float)); g->batch_ffn_cur_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_ffn_norm_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_shared_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); g->batch_shared_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); g->batch_shared_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * shared_dim * sizeof(float)); g->batch_shared_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); g->batch_router_logits_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); g->batch_router_probs_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT * sizeof(float)); g->batch_router_selected_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(int)); g->batch_router_weights_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * sizeof(float)); g->batch_routed_gate_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); g->batch_routed_up_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); g->batch_routed_mid_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); g->batch_routed_down_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); g->batch_routed_out_by_tier[t] = ds4_gpu_tensor_alloc_ptr_on(t, pc * DS4_N_EMBD * sizeof(float)); } if (DS4_GPU_ATTN_COMP_CACHE_F16) { g->batch_q_half = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(uint16_t)); } g->prefill_seed_router_selected = ds4_gpu_tensor_alloc( (uint64_t)DS4_N_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_N_EXPERT_USED * sizeof(int32_t)); } bool layer_cache_ok = true; for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { if (!weights_layer_has_required(&weights->layer[il], il)) continue; layer_cache_ok = g->layer_raw_cache[il] != NULL; if (layer_cache_ok && g->cuda_tp_attn_cache_dup) { layer_cache_ok = g->layer_raw_cache_tp[il] != NULL; } const uint32_t ratio = ds4_layer_compress_ratio(il); if (layer_cache_ok && ratio != 0) { layer_cache_ok = g->layer_attn_comp_cache[il] != NULL && (!g->cuda_tp_attn_cache_dup || g->layer_attn_comp_cache_tp[il] != NULL) && g->layer_attn_state_kv[il] != NULL && g->layer_attn_state_score[il] != NULL && (!enable_frontier_snapshot || (g->spec_attn_state_kv[il] != NULL && g->spec_attn_state_score[il] != NULL)) && (!enable_prefix1_snapshot || (g->spec_prefix1_attn_state_kv[il] != NULL && g->spec_prefix1_attn_state_score[il] != NULL)); } if (layer_cache_ok && ratio == 4) { layer_cache_ok = g->layer_index_comp_cache[il] != NULL && g->layer_index_state_kv[il] != NULL && g->layer_index_state_score[il] != NULL && (!enable_frontier_snapshot || (g->spec_index_state_kv[il] != NULL && g->spec_index_state_score[il] != NULL)) && (!enable_prefix1_snapshot || (g->spec_prefix1_index_state_kv[il] != NULL && g->spec_prefix1_index_state_score[il] != NULL)); } } /* Class P validation — check every used tier's slot. */ bool class_p_ok = true; for (int t = 0; class_p_ok && t < DS4_MAX_GPUS; t++) { if (!used_tier[t]) continue; class_p_ok = g->cur_hc_by_tier[t] && g->flat_hc_by_tier[t] && g->hc_mix_by_tier[t] && g->hc_split_by_tier[t] && g->hc_pre_by_tier[t] && g->hc_post_by_tier[t] && g->hc_comb_by_tier[t] && g->attn_cur_by_tier[t] && g->attn_norm_by_tier[t] && g->qr_by_tier[t] && g->qr_norm_by_tier[t] && g->q_by_tier[t] && g->kv_raw_by_tier[t] && g->kv_by_tier[t] && g->comp_kv_cur_by_tier[t] && g->comp_sc_cur_by_tier[t] && g->index_comp_kv_cur_by_tier[t] && g->index_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 } static bool metal_graph_stream_prefill_batch_selected_addr_layer_supported( const ds4_weights *weights, uint32_t il) { if (!weights || il >= DS4_N_LAYER) return false; const ds4_layer_weights *layer = &weights->layer[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, 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, il); return selected_iq2 || selected_q2; #elif defined(__APPLE__) return 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; #elif !defined(DS4_NO_GPU) 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; return q4 || iq2; #else 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_runtime_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 an intermediate verifier prefix. * * Speculation is only profitable if partial accepts are cheap. For compressed * attention the mutable frontier is just the small compressor state plus append * counters, so retain each intermediate state while the tiny verifier already * steps 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_prefix_attn_state( ds4_gpu_graph *g, uint32_t il, uint32_t slot) { if (!g->spec_capture_prefixes) return true; /* Leading layers without an attention compressor (ratio 0) carry no * frontier; mirror spec_frontier_snapshot/commit which skip them. */ if (ds4_layer_compress_ratio(il) == 0) return true; if (slot >= DS4_SPEC_PREFIX_SLOTS || !g->spec_prefix1_attn_state_kv[il] || !g->spec_prefix1_attn_state_score[il]) { return false; } const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); const uint64_t offset = (uint64_t)slot * bytes; g->spec_prefix_n_comp[slot][il] = g->layer_n_comp[il]; return ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_kv[il], offset, g->layer_attn_state_kv[il], 0, bytes) != 0 && ds4_gpu_tensor_copy(g->spec_prefix1_attn_state_score[il], offset, g->layer_attn_state_score[il], 0, bytes) != 0; } static bool metal_graph_capture_prefix_index_state( ds4_gpu_graph *g, uint32_t il, uint32_t slot) { if (!g->spec_capture_prefixes) return true; /* Index-compressor frontiers exist only on ratio-4 layers; * mirror spec_frontier_snapshot/commit which skip the rest. */ if (ds4_layer_compress_ratio(il) != 4) return true; if (slot >= DS4_SPEC_PREFIX_SLOTS || !g->spec_prefix1_index_state_kv[il] || !g->spec_prefix1_index_state_score[il]) { return false; } const uint64_t bytes = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); const uint64_t offset = (uint64_t)slot * bytes; g->spec_prefix_n_index_comp[slot][il] = g->layer_n_index_comp[il]; return ds4_gpu_tensor_copy(g->spec_prefix1_index_state_kv[il], offset, g->layer_index_state_kv[il], 0, bytes) != 0 && ds4_gpu_tensor_copy(g->spec_prefix1_index_state_score[il], offset, 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; } extern int ds4_gpu_kv_rope_fp8_fuse_available(void); extern int ds4_gpu_decode_attn_rope_fuse_available(void); extern int ds4_gpu_decode_attn_rope_fuse_used(void); extern void ds4_gpu_set_decode_attn_rope_fuse( uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow); extern int ds4_gpu_kv_rope_fp8_store_raw_tensor( ds4_gpu_tensor *kv, ds4_gpu_tensor *raw_cache, uint32_t raw_cap, uint32_t row, uint32_t head_dim, uint32_t n_rot, 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); 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( const ds4_gpu_graph *g) { #ifdef DS4_ROCM_BUILD /* * ROCm SSD streaming maps only the selected routed experts, not the full * Q4 expert table. Stage that compact selection while the shared expert * runs so decode reaches the existing pointer-backed MoE path. * * GLM owns a separate routed-MoE graph and must retain its validated * selected-expert policy. */ if (g && g->ssd_streaming && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) { return true; } #else (void)g; #endif 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->ssd_streaming && !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_decode_mxfp4_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_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4 && 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") && getenv("DS4_METAL_DISABLE_MXFP4_SELECTED_EXPERT_VIEWS") == NULL; } 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) || metal_graph_decode_mxfp4_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; const bool mxfp4 = layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4 && getenv("DS4_METAL_DISABLE_MXFP4_SELECTED_EXPERT_VIEWS") == NULL; return q4 || iq2 || mxfp4; #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 uint32_t metal_graph_streaming_builtin_hotness(uint32_t remaining, uint32_t total) { #ifdef __APPLE__ /* List rank is a preload preference, not thousands of observed routes. */ return 1u + (uint32_t)(31ull * remaining / total); #else (void)total; return remaining; #endif } 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], metal_graph_streaming_builtin_hotness(max_entries - loaded, max_entries), 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"); if (g_ds4_shape.variant == DS4_VARIANT_GLM52 && (!env || !env[0])) { return 0; } /* 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 buffers and trip the system watchdog * before decode begins. GLM demand-fills by default; 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; } #ifdef __APPLE__ /* Layer-major prefill has already made this layer's weights GPU-visible. * Preserve the built-in hot seed with model-to-cache blits before the next * layer replaces those mapped views, rather than rereading it after prefill. */ static bool metal_graph_seed_streaming_expert_cache_layer_from_mapped_hotlist( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, uint32_t il) { const char *hotlist_path = glm_graph_env_value( "DS4_ROCM_STREAMING_EXPERT_HOTLIST", "DS4_METAL_STREAMING_EXPERT_HOTLIST"); if (!metal_graph_streaming_expert_hotlist_enabled(g) || !model || !weights || il >= DS4_N_LAYER || (hotlist_path && hotlist_path[0]) || (g_ds4_shape.variant != DS4_VARIANT_FLASH && g_ds4_shape.variant != DS4_VARIANT_PRO)) { return true; } const ds4_layer_weights *layer = &weights->layer[il]; if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { return true; } 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) { fprintf(stderr, "ds4: streaming resident 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 uint32_t cache_budget = ds4_gpu_stream_expert_cache_budget_for_expert_size(gate_expert_bytes, down_expert_bytes); const uint32_t preload_count = metal_graph_streaming_expert_preload_count(g, cache_budget); if (preload_count == 0) return true; const uint16_t (*hotlist)[2] = NULL; uint32_t hotlist_count = 0; if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { hotlist = ds4_default_streaming_hotlist_flash; hotlist_count = ds4_default_streaming_hotlist_flash_count; } else { hotlist = ds4_default_streaming_hotlist_pro; hotlist_count = ds4_default_streaming_hotlist_pro_count; } const uint32_t target_count = preload_count < hotlist_count ? preload_count : hotlist_count; if (ds4_gpu_stream_expert_cache_current_count() >= target_count) { return true; } int32_t experts[DS4_MAX_EXPERT]; uint32_t priorities[DS4_MAX_EXPERT]; uint32_t n = 0; uint32_t loaded = 0; for (uint32_t i = 0; i < hotlist_count && loaded < preload_count; i++) { const uint32_t hot_layer = hotlist[i][0]; const uint32_t hot_expert = hotlist[i][1]; if (hot_layer >= DS4_N_LAYER || hot_expert >= DS4_N_EXPERT) continue; const uint32_t priority = metal_graph_streaming_builtin_hotness( preload_count - loaded, preload_count); loaded++; if (hot_layer != il) continue; if (n >= DS4_MAX_EXPERT) return false; experts[n] = (int32_t)hot_expert; priorities[n] = priority; n++; } if (n == 0) return true; const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, il, gate_expert_bytes, down_expert_bytes); return ds4_gpu_stream_expert_cache_seed_experts_gpu_copy(&table, experts, priorities, n) != 0; } #endif 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); const bool mxfp4_selected = metal_graph_decode_mxfp4_selected_slots_expected(g, layer); if (!q4_selected && !iq2_selected && !mxfp4_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; /* M5 ports default to the existing standard M1-M4 Metal path. Rollbacks * dominate the benchmark force-enables on pre-M5 devices so the aggregate * switch is a reliable whole-bundle control. */ static bool metal_graph_ported_m5_decode_feature_enabled( const char *pre_m5_disable_env, const char *m5_disable_env) { #if defined(__APPLE__) if (m5_disable_env && getenv(m5_disable_env) != NULL) return false; const bool pre_m5 = ds4_gpu_device_is_pre_m5_apple_silicon(); if (pre_m5 && (getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != NULL || (pre_m5_disable_env && getenv(pre_m5_disable_env) != NULL))) { return false; } return pre_m5 || ds4_gpu_device_is_m5_apple_silicon(); #else (void)pre_m5_disable_env; (void)m5_disable_env; return false; #endif } /* Deferred HC expand support (see ds4_gpu_graph::hc_expand_pending). */ static bool metal_graph_flush_hc_expand(ds4_gpu_graph *g) { if (!g->hc_expand_pending) return true; g->hc_expand_pending = false; return ds4_gpu_hc_expand_add_tensor(g->hc_expand_out, g->hc_expand_block_out, g->hc_expand_block_add, g->hc_expand_residual, g->hc_expand_post, g->hc_expand_comb, DS4_N_EMBD, DS4_N_HC) != 0; } static void metal_graph_defer_hc_expand( ds4_gpu_graph *g, ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb) { g->hc_expand_pending = true; g->hc_expand_out = out_hc; g->hc_expand_block_out = block_out; g->hc_expand_block_add = block_add; g->hc_expand_residual = residual_hc; g->hc_expand_post = post; g->hc_expand_comb = comb; } static bool metal_graph_hc_expand_fusion_eligible(const ds4_gpu_graph *g, bool decode_stage_profile) { #if defined(__APPLE__) return g->tp_world == 2 && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && !decode_stage_profile && metal_graph_debug_get_config()->prefix == NULL && ds4_gpu_device_is_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_HC_EXPAND_PRODUCER_FUSE") == NULL; #else (void)g; (void)decode_stage_profile; return false; #endif } 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) { /* Defer the post-attention inverse RoPE tail into the FlashAttention reduce * kernel, which already owns each head's whole row. Removes one * 64-threadgroup dispatch per layer. */ const bool fuse_attn_inv_rope = getenv("DS4_METAL_DISABLE_PRE_M5_ATTN_INV_ROPE_FUSE") == NULL && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon()) && ds4_gpu_decode_attn_rope_fuse_available() != 0; /* The backend's consumed flag is process-global and remains true after a * gathered-attention layer. Track whether this layer actually armed the * fusion so a following indexed-attention layer cannot mistake that stale * flag for its own and skip the required standalone inverse RoPE. */ bool attn_inv_rope_fuse_armed = false; /* 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; 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); /* A concurrent compute encoder removes all of the generic routed-MoE * function's implicit dispatch ordering. Keep admission deliberately * narrower than the normal decode path: fixed resident IQ2 pair-SwiGLU, * direct Q2 top-6 sum, no readback/profiling/TP detours. */ #if defined(__APPLE__) const bool parallel_ffn_route_eligible = phase == METAL_DECODE_LAYER_FULL && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_PARALLEL_FULL_FFN", "DS4_METAL_DISABLE_M5_PARALLEL_FULL_FFN") && !g->quality && g->tp_world < 2 && !g->ssd_streaming && !g->ssd_streaming_cold && !g->cuda_tp_decode && !g->cuda_tp_moe && !g->cuda_tp_shared && !decode_stage_profile && !metal_graph_directional_steering_ffn_enabled(g) && metal_graph_debug_get_config()->prefix == NULL && getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL && getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == 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 == 256u && DS4_N_EXPERT_USED == 6u && expert_in_dim == 4096u && expert_mid_dim == 2048u && routed_out_dim == 4096u && gate_row_bytes == 1056u && gate_expert_bytes == 2162688u && down_row_bytes == 672u && down_expert_bytes == 2752512u; #else const bool parallel_ffn_route_eligible = false; #endif const bool parallel_full_ffn_eligible = parallel_ffn_route_eligible && getenv("DS4_METAL_Q8_MV_NSG") == NULL && fuse_shared_gate_up && layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && !keep_ffn_out && !metal_graph_use_reference_shared_down_hc() && layer->ffn_gate_inp->type == DS4_TENSOR_F16 && layer->ffn_gate_inp->dim[0] == DS4_N_EMBD && layer->ffn_gate_inp->dim[1] == DS4_N_EXPERT; #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"); /* Decode-island CUDA graph capture (design ported from the Entrpi/ds4 * batched-serving fork). Two position-independent islands of the * decode layer are captured and replayed: island 0 is the hc-pre * stage (exactly the METAL_DECODE_LAYER_TO_QKV prefix), island 1 is * the attention-output-to-layer-end tail (exactly the * METAL_DECODE_LAYER_FROM_ATTN_TO_FFN suffix). Everything a capture * would freeze incorrectly -- TP, multi-tier, SSD streaming, debug * dumps, stage profiling, expert profiling, reference kernels, the * hash router's token argument -- disqualifies the layer and it runs * eagerly, byte-identical to before. A replayed graph re-runs the * captured kernels byte-for-byte, so replay output is bit-identical * to the eager encode it recorded. */ const bool decode_graphs_common_ok = !g->placement && g->tp_world <= 1 && !g->cuda_tp_decode && !g->ssd_streaming && !g->materialize_ffn_out && !decode_stage_profile && !g_expert_profile.active && metal_graph_debug_get_config()->prefix == NULL && ds4_gpu_decode_graphs_supported() != 0; const bool island_a_ok = decode_graphs_common_ok && phase == METAL_DECODE_LAYER_FULL && !tp_ablate_hcpre && DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && !metal_graph_use_reference_hc_norm_decode() && !metal_graph_hc_norm_fusion_check_enabled(); if (island_a_ok && ok) { ds4_decode_graph_key isla_key; memset(&isla_key, 0, sizeof(isla_key)); isla_key.il = il; isla_key.island = 0u; isla_key.cur_hc = (void *)metal_graph_cur_hc(g); isla_key.after_attn_hc = (void *)metal_graph_after_attn_hc(g); isla_key.after_ffn_hc = (void *)metal_graph_after_ffn_hc(g); isla_key.attn_norm = (void *)metal_graph_attn_norm(g); for (;;) { const int isla_state = ds4_gpu_decode_graph_begin(&isla_key); if (isla_state == 1) { return metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, token, METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN); } if (isla_state == 0) { const bool isla_ok = metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, token, METAL_DECODE_LAYER_TO_QKV); if (!isla_ok) { /* Capture consumed no work; the entry is retired and * the retry runs the island eagerly. */ ds4_gpu_decode_graph_abort(&isla_key); continue; } if (ds4_gpu_decode_graph_end(&isla_key) == 0) { return metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, token, METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN); } continue; } break; /* eager: disabled, warm pass, or retired entry */ } } int router_shared_done = 0; int router_only_done = 0; bool router_project_select_fused = false; 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) { bool attn_hc_producer_pre_norm_fused = false; if (ok && !tp_ablate_hcpre) { /* Fused norm+mix removes one decode dispatch per layer; the kernel * reproduces both reduction trees bit-exactly (see dsv4_hc.metal). */ const bool fuse_norm_mix = hc_dim == 16384u && mix_hc == 24u && layer->hc_attn_fn->type == DS4_TENSOR_F16 && !metal_graph_use_reference_hc_decode() && getenv("DS4_METAL_DISABLE_PRE_M5_HC_NORM_MIX_FUSE") == NULL && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon()) && ds4_gpu_hc_rms_norm_mix_f16_available() != 0; #if defined(__APPLE__) const bool fuse_producer_pre_norm = fuse_norm_mix && fuse_hc_norm && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_HC_PRODUCER_PRE_NORM_FUSE", "DS4_METAL_DISABLE_M5_HC_PRODUCER_PRE_NORM_FUSE"); if (g->hc_expand_pending && !(fuse_producer_pre_norm && g->hc_expand_out == metal_graph_cur_hc(g))) { ok = metal_graph_flush_hc_expand(g); } if (ok && fuse_producer_pre_norm) { const bool fuse_expand = g->hc_expand_pending && g->hc_expand_out == metal_graph_cur_hc(g); const int fused = fuse_expand ? ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor( metal_graph_hc_mix(g), metal_graph_attn_cur(g), metal_graph_attn_norm(g), metal_graph_hc_split(g), metal_graph_cur_hc(g), g->hc_expand_block_out, g->hc_expand_block_add, g->hc_expand_residual, g->hc_expand_post, g->hc_expand_comb, model->map, model->size, layer->hc_attn_fn->abs_offset, layer->hc_attn_scale->abs_offset, layer->hc_attn_base->abs_offset, layer->attn_norm->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_RMS_EPS, DS4_HC_EPS, DS4_RMS_EPS) : ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( metal_graph_hc_mix(g), metal_graph_attn_cur(g), metal_graph_attn_norm(g), metal_graph_hc_split(g), metal_graph_cur_hc(g), model->map, model->size, layer->hc_attn_fn->abs_offset, layer->hc_attn_scale->abs_offset, layer->hc_attn_base->abs_offset, layer->attn_norm->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_RMS_EPS, DS4_HC_EPS, DS4_RMS_EPS); if (fuse_expand) { if (fused > 0) g->hc_expand_pending = false; else if (fused == 0) ok = metal_graph_flush_hc_expand(g); } if (fused < 0) { ok = false; } else { attn_hc_producer_pre_norm_fused = fused > 0; } } #endif if (ok && !attn_hc_producer_pre_norm_fused) { if (fuse_norm_mix) { ok = ds4_gpu_hc_rms_norm_mix_f16_tensor( metal_graph_hc_mix(g), metal_graph_cur_hc(g), model->map, model->size, layer->hc_attn_fn->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_RMS_EPS) != 0; } else { 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) { if (!attn_hc_producer_pre_norm_fused) { 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; } bool kv_norm_store_fused = false; if (!resume_after_attn) { /* Fuse the KV RoPE tail into the FP8/raw finalizer: both were single * 64-thread threadgroups on the same 2 KB row, so the pair paid two full * dispatch launches. Deferring the RoPE and running it inside the * finalizer removes one dispatch per layer. */ bool fuse_kv_rope_store = false; bool qkv_pair_quad_fused = false; if (!resume_after_qkv) { bool qkv_pair_projected = resume_after_qa_kv_raw; /* M1-M5 decode fusion: the q_a/kv Q8 pair and the four F16 compressor * projections all read the same normalized attention input and write * disjoint outputs, so one dispatch covers both stages with unchanged * per-row reduction trees (see the kernel comment). Restricted to the * resident FULL phase so split-phase / CUDA / SSD flows keep their * original ordering. */ if (!resume_after_qa_kv_raw && ok && qkv_rms_fused && compressed && phase == METAL_DECODE_LAYER_FULL && !g->ssd_streaming && !g->ssd_streaming_cold && ds4_layer_compress_ratio(il) == 4u && layer->attn_q_a->type == DS4_TENSOR_Q8_0 && layer->attn_kv->type == DS4_TENSOR_Q8_0 && g->cuda_qkv_pair && !metal_graph_use_reference_qkv_pair_proj() && !metal_graph_use_reference_compressor_pair_proj() && layer->attn_compressor_kv && layer->attn_compressor_gate && layer->attn_compressor_ape && 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] == 2u * DS4_N_HEAD_DIM && layer->attn_compressor_gate->dim[1] == 2u * DS4_N_HEAD_DIM && layer->indexer_compressor_kv && layer->indexer_compressor_gate && layer->indexer_compressor_ape && 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] == 2u * DS4_N_INDEXER_HEAD_DIM && layer->indexer_compressor_gate->dim[1] == 2u * DS4_N_INDEXER_HEAD_DIM && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_QUAD_FUSE", "DS4_METAL_DISABLE_M5_QKV_PAIR_QUAD_FUSE")) { const int fused = ds4_gpu_qkv_pair_quad_compressor_store_tensor( metal_graph_qr(g), metal_graph_kv_raw(g), metal_graph_comp_kv_cur(g), metal_graph_comp_sc_cur(g), metal_graph_index_comp_kv_cur(g), metal_graph_index_comp_sc_cur(g), g->layer_attn_state_kv[il], g->layer_attn_state_score[il], g->layer_index_state_kv[il], g->layer_index_state_score[il], model->map, model->size, layer->attn_q_a->abs_offset, layer->attn_kv->abs_offset, layer->attn_compressor_kv->abs_offset, layer->attn_compressor_gate->abs_offset, layer->indexer_compressor_kv->abs_offset, layer->indexer_compressor_gate->abs_offset, layer->attn_compressor_ape->abs_offset, layer->attn_compressor_ape->type, layer->indexer_compressor_ape->abs_offset, layer->indexer_compressor_ape->type, DS4_N_EMBD, q_rank, DS4_N_HEAD_DIM, 2u * DS4_N_HEAD_DIM, 2u * DS4_N_INDEXER_HEAD_DIM, metal_graph_attn_norm(g), 4u, pos); if (fused < 0) { ok = false; } else if (fused > 0) { qkv_pair_projected = true; qkv_pair_quad_fused = true; } } if (!resume_after_qa_kv_raw && ok && !qkv_pair_quad_fused && qkv_rms_fused && compressed && phase == METAL_DECODE_LAYER_FULL && !g->ssd_streaming && !g->ssd_streaming_cold && ds4_layer_compress_ratio(il) == 128u && layer->attn_q_a->type == DS4_TENSOR_Q8_0 && layer->attn_kv->type == DS4_TENSOR_Q8_0 && g->cuda_qkv_pair && !metal_graph_use_reference_qkv_pair_proj() && !metal_graph_use_reference_compressor_pair_proj() && layer->attn_compressor_kv && layer->attn_compressor_gate && layer->attn_compressor_ape && 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] == DS4_N_HEAD_DIM && layer->attn_compressor_gate->dim[1] == DS4_N_HEAD_DIM && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_QKV_PAIR_COMPRESSOR_FUSE", "DS4_METAL_DISABLE_M5_QKV_PAIR_COMPRESSOR_FUSE")) { const int fused = ds4_gpu_qkv_pair_quad_compressor_store_tensor( metal_graph_qr(g), metal_graph_kv_raw(g), metal_graph_comp_kv_cur(g), metal_graph_comp_sc_cur(g), metal_graph_comp_kv_cur(g), metal_graph_comp_sc_cur(g), g->layer_attn_state_kv[il], g->layer_attn_state_score[il], g->layer_attn_state_kv[il], g->layer_attn_state_score[il], model->map, model->size, layer->attn_q_a->abs_offset, layer->attn_kv->abs_offset, layer->attn_compressor_kv->abs_offset, layer->attn_compressor_gate->abs_offset, layer->attn_compressor_kv->abs_offset, layer->attn_compressor_gate->abs_offset, layer->attn_compressor_ape->abs_offset, layer->attn_compressor_ape->type, layer->attn_compressor_ape->abs_offset, layer->attn_compressor_ape->type, DS4_N_EMBD, q_rank, DS4_N_HEAD_DIM, DS4_N_HEAD_DIM, 0u, metal_graph_attn_norm(g), 128u, pos); if (fused < 0) { ok = false; } else if (fused > 0) { qkv_pair_projected = true; qkv_pair_quad_fused = true; } } if (!resume_after_qa_kv_raw && ok && !qkv_pair_quad_fused && qkv_rms_fused && layer->attn_q_a->type == DS4_TENSOR_Q8_0 && layer->attn_kv->type == DS4_TENSOR_Q8_0 && 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 = metal_graph_matmul_dense_quant_tensor(metal_graph_qr(g), model, layer->attn_q_a, DS4_N_EMBD, q_rank, metal_graph_attn_norm(g), 1); } 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 = 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 && 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) { /* Triple fusion: q/kv norm + KV RoPE tail + FP8/raw store in one * dispatch; verbatim arithmetic (see metal/norm.metal). */ if (!metal_graph_tp_ablate("kv") && !kvnorm_dump && !metal_graph_use_reference_kv_decode() && !resume_after_kv_store && DS4_N_HEAD_KV == 1u && DS4_N_HEAD_DIM == 512u && DS4_N_ROT == 64u && raw_cache != NULL && raw_row < raw_cap && phase == METAL_DECODE_LAYER_FULL && getenv("DS4_METAL_DISABLE_PRE_M5_QKV_NORM_KV_STORE_FUSE") == NULL && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon()) && ds4_gpu_kv_rope_fp8_fuse_available() != 0) { #if defined(__APPLE__) { /* Fold the kv task into this layer's KV staging kernel * (byte-exact); DS4_METAL_DISABLE_KV_NORM_DEFER=1 keeps * the standalone two-task dispatch. */ static int defer_kv = -1; if (defer_kv < 0) defer_kv = getenv("DS4_METAL_DISABLE_KV_NORM_DEFER") == NULL; if (defer_kv && g->tp_world == 2) ds4_gpu_dsv4_qkv_norm_defer_kv_next(); } #endif kv_norm_store_fused = ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_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, raw_cache, raw_cap, raw_row, DS4_N_ROT, pos, 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) != 0; if (!kv_norm_store_fused) ok = false; kv_rope_fused = kv_rope_fused || kv_norm_store_fused; } if (ok && !kv_norm_store_fused) { 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 defined(__APPLE__) /* The deferred kv norm/RoPE/FP8 task (one threadgroup) runs concurrently * with the q_b stream; the section closes before the RoPE tail. */ const int kv_concurrent = ok ? ds4_gpu_kv_norm_task_begin_concurrent() : 0; #endif 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 defined(__APPLE__) if (kv_concurrent) ds4_gpu_kv_norm_task_end_concurrent(); #endif 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"); fuse_kv_rope_store = !tp_ablate_kv && !kv_rope_fused && !metal_graph_use_reference_kv_decode() && !resume_after_kv_store && getenv("DS4_METAL_DISABLE_PRE_M5_KV_ROPE_FP8_FUSE") == NULL && ds4_gpu_device_is_pre_m5_apple_silicon() && ds4_gpu_kv_rope_fp8_fuse_available() != 0; if (ok && !tp_ablate_kv && !kv_rope_fused && !fuse_kv_rope_store) { 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 && !kv_norm_store_fused) { ok = fuse_kv_rope_store ? (ds4_gpu_kv_rope_fp8_store_raw_tensor( metal_graph_kv(g), raw_cache, raw_cap, raw_row, 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) : 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 = qkv_pair_quad_fused; /* Quad projection: the attention and indexer compressor pairs share * the normalized input and the F16 matvec shape, so a single dispatch * covers all four matrices with unchanged per-row reduction trees. * Removes one dispatch per decode layer. Either preceding fused * QKV path may already have performed the same work. */ int quad_store = comp_state_already_stored ? 1 : 0; if (ok && quad_store == 0 && ratio == 4u && !metal_graph_use_reference_compressor_pair_proj() && getenv("DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_QUAD_STORE") == NULL && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon()) && layer->indexer_compressor_kv && layer->indexer_compressor_gate && layer->indexer_compressor_ape && 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] == 2u * DS4_N_INDEXER_HEAD_DIM && layer->indexer_compressor_gate->dim[1] == 2u * DS4_N_INDEXER_HEAD_DIM) { quad_store = ds4_gpu_matmul_f16_quad_compressor_store_tensor( metal_graph_comp_kv_cur(g), metal_graph_comp_sc_cur(g), metal_graph_index_comp_kv_cur(g), metal_graph_index_comp_sc_cur(g), g->layer_attn_state_kv[il], g->layer_attn_state_score[il], g->layer_index_state_kv[il], g->layer_index_state_score[il], model->map, model->size, layer->attn_compressor_kv->abs_offset, layer->attn_compressor_gate->abs_offset, layer->indexer_compressor_kv->abs_offset, layer->indexer_compressor_gate->abs_offset, layer->attn_compressor_ape->abs_offset, layer->attn_compressor_ape->type, layer->indexer_compressor_ape->abs_offset, layer->indexer_compressor_ape->type, DS4_N_EMBD, comp_width, 2u * DS4_N_INDEXER_HEAD_DIM, metal_graph_attn_norm(g), ratio, pos); } if (quad_store < 0) { ok = false; } else if (quad_store > 0) { comp_state_already_stored = true; } if (ok && quad_store == 0 && !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 (quad_store == 0) { 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]; /* Eligible Metal decode emit-path fusion: the pooled attention and * indexer compressor rows each need norm + RoPE + quantize (+ F16 * commit for attention) — seven tiny single-row dispatches per layer * every ratio-th token. The fused kernel runs them as one * two-threadgroup dispatch with each standalone kernel's reduction * tree preserved bit-exactly (see kernel_dsv4_comp_row_finalize_f32). * The updates below then stop after the pool (and the state shift). */ const bool comp_finalize_fuse = ok && emit && ratio == 4u && DS4_GPU_ATTN_COMP_CACHE_F16 && phase == METAL_DECODE_LAYER_FULL && !g->ssd_streaming && !g->ssd_streaming_cold && DS4_N_HEAD_DIM == 512u && DS4_N_INDEXER_HEAD_DIM == 128u && layer->attn_compressor_norm && layer->attn_compressor_norm->type == DS4_TENSOR_F32 && layer->indexer_compressor_norm && layer->indexer_compressor_norm->type == DS4_TENSOR_F32 && ds4_gpu_kv_rope_fp8_fuse_available() != 0 && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_COMP_FINALIZE_FUSE", "DS4_METAL_DISABLE_M5_COMP_FINALIZE_FUSE"); 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, true, comp_finalize_fuse) != 0; DS4_METAL_PROFILE_DECODE_STAGE("compressor_update"); if (ok && emit && !comp_finalize_fuse) { 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 = quad_store > 0; if (ok && quad_store == 0 && !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 (quad_store == 0) { 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, true, comp_finalize_fuse) != 0; DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_update"); if (ok && emit && comp_finalize_fuse) { /* One dispatch finalizes both pooled compressor rows; the * attention row lives at offset 0 of the staging tensor and * commits to the F16 cache, the indexer row finalizes in * place. A missing pipeline fails loudly rather than * skipping the finalize the updates deferred. */ ok = ds4_gpu_dsv4_comp_row_finalize_tensor( metal_graph_attn_comp_stage(g), g->layer_attn_comp_cache[il], comp_row, layer->attn_compressor_norm->abs_offset, g->layer_index_comp_cache[il], index_row, layer->indexer_compressor_norm->abs_offset, g->layer_attn_state_kv[il], g->layer_attn_state_score[il], g->layer_index_state_kv[il], g->layer_index_state_score[il], model->map, model->size, pos + 1u - ratio, 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, DS4_RMS_EPS) == 1; } if (ok && emit && !comp_finalize_fuse) { #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 { /* Defer the inverse RoPE tail into the FlashAttention reduce: each * reduce threadgroup already owns one head's entire row, so the * rotation is an intra-threadgroup dependency and the separate * 64-threadgroup RoPE dispatch disappears. */ if (fuse_attn_inv_rope) { ds4_gpu_set_decode_attn_rope_fuse( 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); attn_inv_rope_fuse_armed = true; } 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 && !(attn_inv_rope_fuse_armed && ds4_gpu_decode_attn_rope_fuse_used())) { 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; /* Decode island 1: the attention-output-to-layer-end tail, which is * position-independent (the hash router's token argument disqualifies * layers with ffn_gate_tid2eid). Capture and replay both go through * the METAL_DECODE_LAYER_FROM_ATTN_TO_FFN phase, which encodes * exactly this suffix; the phase check keeps the capture recursion * from re-entering this block. */ const bool island_b_ok = decode_graphs_common_ok && ok && (phase == METAL_DECODE_LAYER_FULL || phase == METAL_DECODE_LAYER_FROM_ATTN_PRE_TO_FFN) && fuse_attn_out_hc && !cuda_tp_attn && !cuda_tp_attn_heads_active && layer->ffn_gate_tid2eid == NULL; if (island_b_ok) { ds4_decode_graph_key islb_key; memset(&islb_key, 0, sizeof(islb_key)); islb_key.il = il; islb_key.island = 1u; islb_key.cur_hc = (void *)metal_graph_cur_hc(g); islb_key.after_attn_hc = (void *)metal_graph_after_attn_hc(g); islb_key.after_ffn_hc = (void *)metal_graph_after_ffn_hc(g); islb_key.attn_norm = (void *)metal_graph_attn_norm(g); for (;;) { const int islb_state = ds4_gpu_decode_graph_begin(&islb_key); if (islb_state == 1) return ok; if (islb_state == 0) { const bool islb_ok = metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, token, METAL_DECODE_LAYER_FROM_ATTN_TO_FFN); if (!islb_ok) { ds4_gpu_decode_graph_abort(&islb_key); continue; } if (ds4_gpu_decode_graph_end(&islb_key) == 0) return true; continue; } break; } } 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; #if defined(__APPLE__) /* The K-slice matvec that writes the partial may publish the gate's * checked flag itself (see ds4_gpu_tp_flag_fold_request). */ ds4_gpu_tp_flag_fold_request(il, DS4_TP_GATE_ATTN); #endif 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) { if (metal_graph_hc_expand_fusion_eligible(g, decode_stage_profile)) { metal_graph_defer_hc_expand(g, 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)); } else { 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); } bool ffn_hc_producer_pre_norm_fused = false; if (ok && !tp_ablate_hcpre) { const bool fuse_norm_mix = hc_dim == 16384u && mix_hc == 24u && layer->hc_ffn_fn->type == DS4_TENSOR_F16 && !metal_graph_use_reference_hc_decode() && getenv("DS4_METAL_DISABLE_PRE_M5_HC_NORM_MIX_FUSE") == NULL && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon()) && ds4_gpu_hc_rms_norm_mix_f16_available() != 0; #if defined(__APPLE__) const bool fuse_producer_pre_norm = fuse_norm_mix && fuse_hc_norm && metal_graph_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_HC_PRODUCER_PRE_NORM_FUSE", "DS4_METAL_DISABLE_M5_HC_PRODUCER_PRE_NORM_FUSE"); if (g->hc_expand_pending && !(fuse_producer_pre_norm && g->hc_expand_out == metal_graph_after_attn_hc(g))) { ok = metal_graph_flush_hc_expand(g); } if (ok && fuse_producer_pre_norm) { const bool fuse_expand = g->hc_expand_pending && g->hc_expand_out == metal_graph_after_attn_hc(g); const int fused = fuse_expand ? ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor( metal_graph_hc_mix(g), metal_graph_ffn_cur(g), metal_graph_ffn_norm(g), metal_graph_hc_split(g), metal_graph_after_attn_hc(g), g->hc_expand_block_out, g->hc_expand_block_add, g->hc_expand_residual, g->hc_expand_post, g->hc_expand_comb, model->map, model->size, layer->hc_ffn_fn->abs_offset, layer->hc_ffn_scale->abs_offset, layer->hc_ffn_base->abs_offset, layer->ffn_norm->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_RMS_EPS, DS4_HC_EPS, DS4_RMS_EPS) : ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor( metal_graph_hc_mix(g), metal_graph_ffn_cur(g), metal_graph_ffn_norm(g), metal_graph_hc_split(g), metal_graph_after_attn_hc(g), model->map, model->size, layer->hc_ffn_fn->abs_offset, layer->hc_ffn_scale->abs_offset, layer->hc_ffn_base->abs_offset, layer->ffn_norm->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_RMS_EPS, DS4_HC_EPS, DS4_RMS_EPS); if (fuse_expand) { if (fused > 0) g->hc_expand_pending = false; else if (fused == 0) ok = metal_graph_flush_hc_expand(g); } if (fused < 0) { ok = false; } else { ffn_hc_producer_pre_norm_fused = fused > 0; } } #endif if (ok && !ffn_hc_producer_pre_norm_fused) { if (fuse_norm_mix) { ok = ds4_gpu_hc_rms_norm_mix_f16_tensor( metal_graph_hc_mix(g), metal_graph_after_attn_hc(g), model->map, model->size, layer->hc_ffn_fn->abs_offset, (uint32_t)hc_dim, (uint32_t)mix_hc, DS4_RMS_EPS) != 0; } else { 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) { if (!ffn_hc_producer_pre_norm_fused) { 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")) { /* Fused router + shared-expert gate/up: one dispatch instead of two * on the same normalized FFN input; per-row math is unchanged (see * metal/dense.metal). Restricted to the full single-pass decode * phase so split-phase/graph flows keep the original ordering. */ if (phase == METAL_DECODE_LAYER_FULL && !g->ssd_streaming && !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 && layer->ffn_gate_inp->type == DS4_TENSOR_F16 && layer->ffn_gate_inp->dim[0] == DS4_N_EMBD && layer->ffn_gate_inp->dim[1] == DS4_N_EXPERT && (!ds4_gpu_device_is_pre_m5_apple_silicon() || getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_SHARED_FUSE") == NULL) && (ds4_gpu_device_is_pre_m5_apple_silicon() || ds4_gpu_device_is_m5_apple_silicon())) { #if defined(__APPLE__) const bool fuse_router_project_select = parallel_full_ffn_eligible && layer->ffn_gate_tid2eid == NULL && getenv("DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE") == NULL && ds4_gpu_device_is_m5_apple_silicon(); if (fuse_router_project_select) { const int fused = ds4_gpu_router_project_select_fused_tensor( metal_graph_router_logits(g), metal_graph_router_probs(g), metal_graph_router_selected(g), metal_graph_router_weights(g), model->map, model->size, layer->ffn_gate_inp->abs_offset, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, layer->ffn_exp_probs_b != NULL, metal_graph_ffn_norm(g)); if (fused < 0) { ok = false; } else { router_project_select_fused = fused > 0; } } #endif if (ok && !router_project_select_fused) { router_shared_done = ds4_gpu_router_shared_gate_up_q8_0_tensor( metal_graph_router_logits(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), model->map, model->size, layer->ffn_gate_inp->abs_offset, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, DS4_N_EXPERT, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP, parallel_full_ffn_eligible); } else { router_only_done = 1; } if (router_shared_done < 0) { ok = false; router_shared_done = 0; } if (router_shared_done != 0 && parallel_full_ffn_eligible) { router_only_done = 1; router_shared_done = 0; } } #if defined(__APPLE__) /* Two-rank TP: router projection + exact select as one dispatch (the * last-arriving threadgroup selects); the shared expert runs in the * concurrent FFN encoder instead of the router+shared fusion above. * DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE opts out. */ if (ok && g->tp_world == 2 && phase == METAL_DECODE_LAYER_FULL && !g->ssd_streaming && !g->quality && router_shared_done == 0 && router_only_done == 0 && !router_project_select_fused && layer->ffn_gate_inp->type == DS4_TENSOR_F16 && layer->ffn_gate_inp->dim[0] == DS4_N_EMBD && layer->ffn_gate_inp->dim[1] == DS4_N_EXPERT && layer->ffn_gate_tid2eid == NULL && !metal_graph_tp_ablate("router") && getenv("DS4_METAL_DISABLE_M5_ROUTER_PROJECT_SELECT_FUSE") == NULL && ds4_gpu_device_is_m5_apple_silicon()) { const int fused = ds4_gpu_router_project_select_fused_tensor( metal_graph_router_logits(g), metal_graph_router_probs(g), metal_graph_router_selected(g), metal_graph_router_weights(g), model->map, model->size, layer->ffn_gate_inp->abs_offset, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, layer->ffn_exp_probs_b != NULL, metal_graph_ffn_norm(g)); if (fused < 0) { ok = false; } else if (fused > 0) { /* Logits and select done; router_only_done stays 0 so the * single-machine parallel FFN path is not enabled. */ router_project_select_fused = true; } } #endif if (ok && router_shared_done == 0 && router_only_done == 0 && !router_project_select_fused) 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 && !router_project_select_fused) ok = ds4_gpu_router_select_tensor(metal_graph_router_selected(g), metal_graph_router_weights(g), metal_graph_router_probs(g), model->map, model->size, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, (uint32_t)token, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, 0, 0, layer->ffn_exp_probs_b != NULL, layer->ffn_gate_tid2eid != NULL, metal_graph_router_logits(g)) != 0; } if (ok) ok = metal_graph_decode_set_hash_selected_override(model, layer, il, (uint32_t)token, layer->ffn_gate_exps->bytes, layer->ffn_down_exps->bytes, g); } DS4_METAL_PROFILE_DECODE_STAGE("router"); if (ok) ok = metal_graph_profile_router_selection(g, layer, il, pos); if (ok) { metal_graph_debug_dump_tensor("ffn_moe_logits", metal_graph_router_logits(g), DS4_N_EXPERT, il, pos); metal_graph_debug_dump_tensor("ffn_moe_probs", metal_graph_router_probs(g), DS4_N_EXPERT, il, pos); metal_graph_debug_dump_i32_tensor("ffn_moe_topk", metal_graph_router_selected(g), DS4_N_EXPERT_USED, il, pos); metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", metal_graph_router_weights(g), DS4_N_EXPERT_USED, il, pos); } if (phase == METAL_DECODE_LAYER_TO_ROUTER) return ok; } const bool external_routed = phase == METAL_DECODE_LAYER_FROM_ROUTER; const bool 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(g) && 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 mxfp4_selected_shared_overlap = metal_graph_use_iq2_selected_shared_overlap(g) && metal_graph_decode_mxfp4_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 || mxfp4_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)) || (mxfp4_selected_shared_overlap && metal_graph_use_iq2_selected_async_load(g)) || cuda_selected_shared_overlap); const bool selected_readahead_shared_delay = ok && g->tp_world < 2 && !overlap_selected_shared && !decode_stage_profile && metal_graph_use_iq2_selected_readahead_shared_delay(g) && metal_graph_decode_iq2_selected_slots_expected(g, layer) && !metal_graph_decode_cpu_router_applicable(g, layer) && layer->ffn_gate_tid2eid == NULL && getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; const bool cuda_stream_selected_load = ok && !overlap_selected_shared && !selected_readahead_shared_delay && g->ssd_streaming && metal_graph_decode_cuda_selected_slots_expected(g, layer) && layer->ffn_gate_tid2eid == NULL && getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL; if (cuda_stream_selected_load) { ok = metal_graph_decode_cuda_selected_load(g, model, layer, il, gate_expert_bytes, down_expert_bytes); } if (selected_readahead_shared_delay) { if (ok) { ok = metal_graph_decode_selected_readahead_override(g, model, layer, il, gate_expert_bytes, down_expert_bytes); } if (ok && fuse_shared_gate_up) { ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP) != 0; } else if (ok) { if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), model, layer->ffn_gate_shexp, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), 1); if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), model, layer->ffn_up_shexp, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), 1); if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), metal_graph_routed_gate(g), metal_graph_routed_up(g), metal_graph_routed_mid(g), metal_graph_routed_down(g), model->map, model->size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, layer->ffn_down_exps->abs_offset, layer->ffn_gate_exps->type, layer->ffn_down_exps->type, gate_expert_bytes, gate_row_bytes, down_expert_bytes, down_row_bytes, (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, metal_graph_router_selected(g), metal_graph_router_weights(g), DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), NULL, il, false) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); } if (ok && fuse_shared_down_hc) { ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), metal_graph_shared_out(g), model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, metal_graph_shared_mid(g), metal_graph_routed_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), model, layer->ffn_down_shexp, shared_dim, DS4_N_EMBD, metal_graph_shared_mid(g), 1); } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (ok) { metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); } if (ok && keep_ffn_out) { ok = metal_graph_ensure_ffn_out(g) && ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; } if (ok && keep_ffn_out) { metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), metal_graph_ffn_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_post(g), metal_graph_hc_comb(g), DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && !fuse_shared_down_hc) { ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), metal_graph_routed_out(g), metal_graph_shared_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); if (ok) { metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); } return ok; } if (overlap_selected_shared) { uint64_t selected_event = 0; if (ok) ok = ds4_gpu_signal_selected_readback_ready(&selected_event) != 0; metal_graph_selected_async_load async_load = {0}; bool async_load_started = false; const bool async_early_commit = async_selected_load && metal_graph_use_iq2_selected_async_early_commit(g); if (ok && async_selected_load) { ok = metal_graph_selected_async_load_start(&async_load, g, model, layer, il, selected_event, gate_expert_bytes, down_expert_bytes); async_load_started = ok; } if (ok && async_early_commit) { ok = ds4_gpu_flush_commands() != 0; } if (ok && fuse_shared_gate_up) { ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP) != 0; } else if (ok) { if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_gate(g), model, layer->ffn_gate_shexp, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), 1); if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_up(g), model, layer->ffn_up_shexp, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), 1); if (ok) ok = ds4_gpu_swiglu_tensor(metal_graph_shared_mid(g), metal_graph_shared_gate(g), metal_graph_shared_up(g), shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); if (ok && !fuse_shared_down_hc) { ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), model, layer->ffn_down_shexp, shared_dim, DS4_N_EMBD, metal_graph_shared_mid(g), 1); } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (async_load_started) { const bool flush_ok = ds4_gpu_flush_commands() != 0; bool finish_ok = metal_graph_selected_async_load_finish(&async_load); if (!finish_ok && async_load.ids_ok) { /* The worker read valid ids but could not stage the load * (it is not allowed to wait on in-flight cache entries). * This thread is, so retry the same load synchronously. */ const ds4_gpu_stream_expert_table retry_table = graph_stream_expert_table_make(model, layer, il, gate_expert_bytes, down_expert_bytes); finish_ok = ds4_gpu_stream_expert_cache_begin_selected_load( &retry_table, async_load.selected_ids, DS4_N_EXPERT_USED) != 0 && ds4_gpu_routed_moe_set_selected_override( async_load.selected_ids, DS4_N_EXPERT_USED) != 0; } ok = ok && flush_ok && finish_ok; } else if (ok) { ok = ds4_gpu_commit_and_wait_selected_readback(selected_event, "selected-id shared-overlap") != 0; } if (ok && !async_load_started) { int32_t selected_ids[DS4_MAX_EXPERT_USED]; ok = ds4_gpu_tensor_read(metal_graph_router_selected(g), 0, selected_ids, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) != 0 && ds4_gpu_routed_moe_set_selected_override(selected_ids, DS4_N_EXPERT_USED) != 0; if (ok) { const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, il, gate_expert_bytes, down_expert_bytes); ok = ds4_gpu_stream_expert_cache_begin_selected_load( &table, selected_ids, DS4_N_EXPERT_USED) != 0; } } if (ok) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), metal_graph_routed_gate(g), metal_graph_routed_up(g), metal_graph_routed_mid(g), metal_graph_routed_down(g), model->map, model->size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, layer->ffn_down_exps->abs_offset, layer->ffn_gate_exps->type, layer->ffn_down_exps->type, gate_expert_bytes, gate_row_bytes, down_expert_bytes, down_row_bytes, (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, metal_graph_router_selected(g), metal_graph_router_weights(g), DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), NULL, il, false) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_up_clamped", metal_graph_routed_up(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", metal_graph_routed_mid(g), (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); metal_graph_debug_dump_tensor("ffn_moe_down", metal_graph_routed_down(g), (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); metal_graph_debug_dump_tensor("ffn_moe_out", metal_graph_routed_out(g), DS4_N_EMBD, il, pos); } if (ok && fuse_shared_down_hc) { ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(metal_graph_after_ffn_hc(g), metal_graph_shared_out(g), model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, metal_graph_shared_mid(g), metal_graph_routed_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (ok) { metal_graph_debug_dump_tensor("ffn_shexp", metal_graph_shared_out(g), DS4_N_EMBD, il, pos); } if (ok && keep_ffn_out) { ok = metal_graph_ensure_ffn_out(g) && ds4_gpu_add_tensor(metal_graph_ffn_out(g), metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; } if (ok && keep_ffn_out) { metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_ffn_out(g), il, 1); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), metal_graph_ffn_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_post(g), metal_graph_hc_comb(g), DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && !fuse_shared_down_hc) { ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), metal_graph_routed_out(g), metal_graph_shared_out(g), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); if (ok) { metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); } return ok; } /* The routed and shared-expert branches depend only on ffn_norm. Run * their two levels in one concurrent encoder on M5 TP: each rank still * computes exactly its owned routed experts and shared lane slice, but * neither independent memory stream waits behind the other. */ const bool parallel_tp_ffn_eligible = #if defined(__APPLE__) ok && tp_split_shared && getenv("DS4_METAL_DISABLE_M5_TP_PARALLEL_FFN") == NULL && ds4_gpu_device_is_m5_apple_silicon() && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && !decode_stage_profile && !keep_ffn_out && !metal_graph_directional_steering_ffn_enabled(g) && metal_graph_debug_get_config()->prefix == NULL && layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4 && 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 && (shared_dim % 64u) == 0u; #else false; #endif bool parallel_full_ffn = ok && parallel_full_ffn_eligible && router_shared_done == 0 && router_only_done != 0 && fuse_shared_down_hc; const bool parallel_tp_ffn = parallel_tp_ffn_eligible; if (parallel_full_ffn) { #if defined(__APPLE__) parallel_full_ffn = ds4_gpu_parallel_ffn_start( metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), metal_graph_shared_out(g), model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, layer->ffn_down_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP) != 0; #else parallel_full_ffn = false; #endif } if (parallel_tp_ffn) { #if defined(__APPLE__) const uint32_t tp_half = shared_dim / 2u; /* Balance the two ranks per token: the routed experts are owned by * expert id, so the rank owning more of the selected experts gets * fewer shared-expert lanes. The split kernels decide on the GPU * from the selected ids; shift = routed expert bytes over twice the * shared expert bytes. DS4_TP_STATIC_SHARED_SPLIT=1 keeps the * fixed halves; DS4_TP_SHARED_SPLIT_SHIFT_Q16 overrides the shift * (0 = fixed halves through the split kernels, for bit-exact * plumbing checks). */ parallel_full_ffn = false; if (g->tp_world == 2 && (shared_dim % 64u) == 0u && getenv("DS4_TP_STATIC_SHARED_SPLIT") == NULL) { const uint64_t routed_expert_bytes = 2u * gate_expert_bytes + down_expert_bytes; const uint64_t shared_bytes = 2u * ((uint64_t)DS4_N_EMBD / 32u) * 34u * shared_dim + ((uint64_t)shared_dim / 32u) * 34u * (uint64_t)DS4_N_EMBD; uint64_t shift_q16 = ((routed_expert_bytes << 16) + shared_bytes) / (2u * shared_bytes); const char *shift_env = getenv("DS4_TP_SHARED_SPLIT_SHIFT_Q16"); if (shift_env) shift_q16 = strtoull(shift_env, NULL, 10); if (shift_q16 > 65536u) shift_q16 = 65536u; parallel_full_ffn = ds4_gpu_parallel_ffn_start_split( metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), metal_graph_shared_out(g), model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, layer->ffn_down_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP, metal_graph_router_selected(g), (uint32_t)g->tp_rank, (uint32_t)g->tp_world, (uint32_t)DS4_N_EXPERT, (uint32_t)DS4_N_EXPERT_USED, (uint32_t)shift_q16) != 0; } if (!parallel_full_ffn) parallel_full_ffn = ds4_gpu_parallel_ffn_start_sliced( metal_graph_shared_gate(g), metal_graph_shared_up(g), metal_graph_shared_mid(g), metal_graph_shared_out(g), model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, layer->ffn_down_shexp->abs_offset, DS4_N_EMBD, shared_dim, g->tp_rank * tp_half, tp_half, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP) != 0; #else parallel_full_ffn = false; #endif } /* The serial TP fallback runs the shared branch first so routed sum6 can * fold it into the slab slot. The concurrent path keeps both partials * separate until its encoder joins, then performs the same local sum. */ const bool tp_fold_ffn = tp_split_shared && !parallel_full_ffn && !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; if (!ok && parallel_full_ffn) { #if defined(__APPLE__) ds4_gpu_parallel_ffn_abort(); #endif } 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) { if (parallel_full_ffn) { #if defined(__APPLE__) ds4_gpu_parallel_ffn_abort(); #endif } return ok; } if (parallel_full_ffn) { /* Gate/up/SwiGLU was encoded as the first independent dispatch in * the concurrent FFN encoder. */ } else 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 && router_shared_done == 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, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, metal_graph_ffn_norm(g), DS4_SWIGLU_CLAMP_EXP) != 0; } else if (router_shared_done == 0) { 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 (parallel_full_ffn) { #if defined(__APPLE__) const bool parallel_joined = ds4_gpu_parallel_ffn_finish() != 0; ok = ok && parallel_joined; #else ok = false; #endif if (ok && !tp_split_shared) { 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; } } 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) { #if defined(__APPLE__) ok = ds4_gpu_add_tensor_tp_flag(g->tp_out[tp_slot], metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD, il, DS4_TP_GATE_FFN) != 0; #else ok = ds4_gpu_add_tensor(g->tp_out[tp_slot], metal_graph_shared_out(g), metal_graph_routed_out(g), DS4_N_EMBD) != 0; #endif } 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 && metal_graph_hc_expand_fusion_eligible(g, decode_stage_profile)) { metal_graph_defer_hc_expand(g, 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_zero, metal_graph_after_attn_hc(g), metal_graph_hc_post(g), metal_graph_hc_comb(g)); } else if (ok && !cuda_tp_shared_fold && !fuse_shared_down_hc) { ok = ds4_gpu_hc_expand_add_split_tensor(metal_graph_after_ffn_hc(g), tp_ffn_a ? tp_ffn_a : metal_graph_routed_out(g), tp_ffn_a ? tp_ffn_b : (g->tp_world == 2 ? g->tp_zero : metal_graph_shared_out(g)), metal_graph_after_attn_hc(g), metal_graph_hc_split(g), DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); #undef DS4_METAL_PROFILE_DECODE_STAGE if (ok) { metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_after_ffn_hc(g), hc_dim, il, pos); } return ok; } static bool metal_graph_encode_decode_layer( ds4_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t pos, ds4_gpu_tensor *raw_cache, uint32_t raw_cap, uint32_t raw_row, uint32_t n_raw, int token) { return metal_graph_encode_decode_layer_phase( g, model, layer, il, pos, raw_cache, raw_cap, raw_row, n_raw, token, METAL_DECODE_LAYER_FULL); } static bool metal_graph_output_logits_head_matmul( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, ds4_gpu_tensor *norm_full, ds4_gpu_tensor *dst_logits, uint32_t n_tokens, uint64_t vocab_dim); /* Encode the final HC collapse, output norm, and vocab projection on Metal. */ static bool metal_graph_encode_output_head( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, uint64_t vocab_dim) { /* switch to head_tier before the output-head pipeline. * Single-tier (placement == NULL): no-op (head_tier == 0 == active_tier). * Note: head_tier was captured in metal_graph_alloc_raw_cap; this * helper consults it directly (and also covers the case where the * preceding decode layer ran on a different tier — copy_xdev ferries * the active cur_hc across the boundary). */ if (g->placement) { if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const bool output_stage_profile = g->output_stage_profile; double output_stage_t0 = output_stage_profile ? now_sec() : 0.0; #define DS4_METAL_PROFILE_OUTPUT_STAGE(name) do { \ if (ok && output_stage_profile) { \ ok = metal_graph_layer_stage_profile_boundary("output", (name), DS4_N_LAYER, 0, 1, &output_stage_t0); \ } \ } while (0) bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; DS4_METAL_PROFILE_OUTPUT_STAGE("hc_flat_norm"); if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), model->map, model->size, weights->output_hc_fn->abs_offset, hc_dim, DS4_N_HC, metal_graph_flat_hc(g), 1) != 0; DS4_METAL_PROFILE_OUTPUT_STAGE("hc_pre"); if (ok) { metal_graph_debug_dump_tensor("result_hc_pre", metal_graph_output_pre(g), DS4_N_HC, DS4_N_LAYER, 0); } if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), metal_graph_output_pre(g), model->map, model->size, weights->output_hc_scale->abs_offset, weights->output_hc_base->abs_offset, DS4_N_HC, DS4_HC_EPS) != 0; DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weights"); if (ok) { metal_graph_debug_dump_tensor("result_hc_weights", metal_graph_output_weights(g), DS4_N_HC, DS4_N_LAYER, 0); } 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; } 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) { ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), metal_graph_output_embd(g), model->map, model->size, weights->output_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; } DS4_METAL_PROFILE_OUTPUT_STAGE("output_norm"); if (ok) { metal_graph_debug_dump_tensor("result_norm", metal_graph_output_norm(g), DS4_N_EMBD, DS4_N_LAYER, 0); } if (ok && g->tp_world == 2 && g->tp_logits_half) { /* Vocab-split: this rank computes its half of the head rows into * its logits view; the halves are bit-identical to the full head * (same kernel, same rows) and the worker ships its half to the * leader after the eval. */ const uint64_t tp_vhalf = vocab_dim / 2u; uint64_t head_row_bytes = 0; ok = metal_graph_dense_quant_row_bytes(weights->output, DS4_N_EMBD, &head_row_bytes); if (ok) ok = metal_graph_matmul_dense_quant_abs(g->tp_logits_half, model, weights->output, weights->output->abs_offset + (uint64_t)g->tp_rank * tp_vhalf * head_row_bytes, DS4_N_EMBD, tp_vhalf, metal_graph_output_norm(g), 1); } else if (ok && g->cuda_tp_ep && g->cuda_tp_output) { ok = metal_graph_output_logits_head_matmul( g, model, weights, metal_graph_output_norm(g), metal_graph_logits(g), 1, vocab_dim); } else if (ok) { ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), model, weights->output, DS4_N_EMBD, vocab_dim, metal_graph_output_norm(g), 1); } if (ok) { metal_graph_debug_dump_tensor("result_output", metal_graph_logits(g), vocab_dim, DS4_N_LAYER, 0); } #undef DS4_METAL_PROFILE_OUTPUT_STAGE return ok; } /* Greedy-only output head: compute one local top-1 candidate per output TP * split and leave the full split logits on their owning tiers. This avoids * gathering the whole vocabulary row back to the head tier when the caller only * needs the next argmax token. */ static bool metal_graph_encode_output_head_split_top1( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, uint64_t vocab_dim, int cuda_tp_output_tiers[DS4_MAX_GPUS], uint32_t *cuda_tp_output_ways_out) { if (!g || !model || !weights || !cuda_tp_output_tiers || !cuda_tp_output_ways_out || vocab_dim > UINT32_MAX) { return false; } *cuda_tp_output_ways_out = 0; if (g->placement) { if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), model->map, model->size, weights->output_hc_fn->abs_offset, hc_dim, DS4_N_HC, metal_graph_flat_hc(g), 1) != 0; if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), metal_graph_output_pre(g), model->map, model->size, weights->output_hc_scale->abs_offset, weights->output_hc_base->abs_offset, DS4_N_HC, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), metal_graph_cur_hc(g), metal_graph_output_weights(g), DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), metal_graph_output_embd(g), model->map, model->size, weights->output_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; if (!ok) return false; const uint32_t cuda_tp_output_ways = g->cuda_tp_output ? metal_graph_cuda_tp_output_tiers(g, cuda_tp_output_tiers) : 0; const bool cuda_tp_output = g->cuda_tp_output && cuda_tp_output_ways >= 2u && weights->output->type == DS4_TENSOR_Q8_0 && weights->output->ndim == 2 && weights->output->dim[0] == DS4_N_EMBD && weights->output->dim[1] == vocab_dim && vocab_dim >= 2; if (!cuda_tp_output) return false; for (uint32_t i = 0; i < cuda_tp_output_ways; i++) { const int t = cuda_tp_output_tiers[i]; if (t < 0 || t >= DS4_MAX_GPUS || !g->output_norm_by_tier[t] || !g->logits_by_tier[t] || !g->comp_selected_by_tier[t] || !g->comp_mask_by_tier[t]) { return false; } } const bool fused_top1 = metal_graph_cuda_output_fused_top1_requested(); const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); uint64_t split_start[DS4_MAX_GPUS] = {0}; uint64_t split_count[DS4_MAX_GPUS] = {0}; ds4_gpu_tensor split_logits[DS4_MAX_GPUS]; memset(split_logits, 0, sizeof(split_logits)); for (uint32_t i = 0; ok && i < cuda_tp_output_ways; i++) { const int t = cuda_tp_output_tiers[i]; split_start[i] = (vocab_dim * (uint64_t)i) / cuda_tp_output_ways; const uint64_t split_end = (vocab_dim * (uint64_t)(i + 1u)) / cuda_tp_output_ways; split_count[i] = split_end - split_start[i]; if (split_count[i] == 0 || split_count[i] > UINT32_MAX) { ok = false; break; } if (fused_top1) { if (t != g->head_tier) { ok = ds4_gpu_tensor_copy_xdev( g->output_norm_by_tier[t], metal_graph_output_norm(g), (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; } } else if (t == g->head_tier) { ok = metal_graph_borrow_tensor_view(&split_logits[i], metal_graph_logits(g), split_start[i] * sizeof(float), split_count[i] * sizeof(float)); } else { ok = metal_graph_borrow_tensor_view(&split_logits[i], g->logits_by_tier[t], 0, split_count[i] * sizeof(float)) && ds4_gpu_tensor_copy_xdev( g->output_norm_by_tier[t], metal_graph_output_norm(g), (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; } } for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { const int t = cuda_tp_output_tiers[i]; ok = ds4_gpu_set_current_device(t) == 0; if (ok && fused_top1) { ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[t], g->comp_mask_by_tier[t], model->map, model->size, weights->output->abs_offset + split_start[i] * row_bytes, DS4_N_EMBD, split_count[i], g->output_norm_by_tier[t], (uint32_t)split_start[i]) != 0; } else if (ok) { ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[i], model->map, model->size, weights->output->abs_offset + split_start[i] * row_bytes, DS4_N_EMBD, split_count[i], g->output_norm_by_tier[t], 1) != 0; } if (ok && !fused_top1) { ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[t], g->comp_mask_by_tier[t], &split_logits[i], (uint32_t)split_count[i], 1, (uint32_t)split_start[i]) != 0; } } if (ok) ok = ds4_gpu_set_current_device(g->head_tier) == 0; if (ok && fused_top1) { ok = ds4_gpu_matmul_q8_0_top1_tensor(g->comp_selected_by_tier[g->head_tier], g->comp_mask_by_tier[g->head_tier], model->map, model->size, weights->output->abs_offset, DS4_N_EMBD, split_count[0], metal_graph_output_norm(g), (uint32_t)split_start[0]) != 0; } else if (ok) { ok = ds4_gpu_matmul_q8_0_tensor(&split_logits[0], model->map, model->size, weights->output->abs_offset, DS4_N_EMBD, split_count[0], metal_graph_output_norm(g), 1) != 0; } if (ok && !fused_top1) { ok = ds4_gpu_indexer_top1_value_tensor(g->comp_selected_by_tier[g->head_tier], g->comp_mask_by_tier[g->head_tier], &split_logits[0], (uint32_t)split_count[0], 1, (uint32_t)split_start[0]) != 0; } for (uint32_t i = 1; ok && i < cuda_tp_output_ways; i++) { const int t = cuda_tp_output_tiers[i]; ds4_gpu_tensor head_id_dst; ds4_gpu_tensor head_value_dst; ok = metal_graph_borrow_tensor_view(&head_id_dst, g->comp_selected_by_tier[g->head_tier], (uint64_t)i * sizeof(uint32_t), sizeof(uint32_t)) && metal_graph_borrow_tensor_view(&head_value_dst, g->comp_mask_by_tier[g->head_tier], (uint64_t)i * sizeof(float), sizeof(float)) && ds4_gpu_tensor_copy_xdev3(&head_id_dst, g->comp_selected_by_tier[t], sizeof(uint32_t), &head_value_dst, g->comp_mask_by_tier[t], sizeof(float), NULL, NULL, 0) != 0; } if (ok) { ok = ds4_gpu_set_current_device(g->head_tier) == 0; *cuda_tp_output_ways_out = cuda_tp_output_ways; } return ok; } static bool metal_graph_read_output_split_top1( ds4_gpu_graph *g, uint32_t output_ways, int *top_id) { if (!g || !top_id || output_ways == 0 || output_ways > DS4_MAX_GPUS) { return false; } bool have_best = false; uint32_t best_id = 0; float best_value = 0.0f; uint32_t cand_ids[DS4_MAX_GPUS] = {0}; float cand_values[DS4_MAX_GPUS] = {0.0f}; bool ok = ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], 0, cand_ids, (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], 0, cand_values, (uint64_t)output_ways * sizeof(cand_values[0])) != 0; for (uint32_t i = 0; ok && i < output_ways; i++) { const uint32_t cand_id = cand_ids[i]; const float cand_value = cand_values[i]; if (!have_best || cand_value > best_value || (cand_value == best_value && cand_id < best_id)) { have_best = true; best_id = cand_id; best_value = cand_value; } } ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; if (ok) *top_id = (int)best_id; return ok; } /* Batched output head for speculative verification. * * A target verifier only needs top-1 ids for intermediate draft rows and full * logits for the last accepted row. Running the normal one-row output head in * a loop serializes the HC collapse, output norm, and Q8 vocab projection. For * tiny MTP suffixes we instead process all rows together and let the GPU reduce * each row to a top id; the CPU reads back just those ids plus the last row's * logits needed to continue the exact target stream. */ /* Shared vocab-head matmul: pads small batches to 8 rows for the exact-mma Q8 * kernel and shards the vocabulary across output-TP tiers. */ static bool metal_graph_output_logits_head_matmul( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, ds4_gpu_tensor *norm_full, ds4_gpu_tensor *dst_logits, uint32_t n_tokens, uint64_t vocab_dim) { if (!g || !model || !weights || !norm_full || n_tokens == 0 || !dst_logits || ds4_gpu_tensor_bytes(dst_logits) < (uint64_t)n_tokens * vocab_dim * sizeof(float)) { return false; } const uint32_t head_rows = (n_tokens > 1 && n_tokens < 8 && ds4_gpu_tensor_bytes(dst_logits) >= 8u * vocab_dim * sizeof(float) && ds4_gpu_tensor_bytes(norm_full) >= 8u * DS4_N_EMBD * sizeof(float)) ? 8u : n_tokens; ds4_gpu_tensor *output_norm = ds4_gpu_tensor_view(norm_full, 0, (uint64_t)head_rows * DS4_N_EMBD * sizeof(float)); ds4_gpu_tensor *logits = ds4_gpu_tensor_view(dst_logits, 0, (uint64_t)head_rows * vocab_dim * sizeof(float)); bool ok = output_norm && logits; if (ok && head_rows > n_tokens) { ds4_gpu_tensor *pad = ds4_gpu_tensor_view(norm_full, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float), (uint64_t)(head_rows - n_tokens) * DS4_N_EMBD * sizeof(float)); ok = pad && ds4_gpu_tensor_fill_f32(pad, 0.0f, (uint64_t)(head_rows - n_tokens) * DS4_N_EMBD) != 0; ds4_gpu_tensor_free(pad); } /* Output TP for the speculative batch, mirroring the decode head: each * device matmuls its VRAM-resident vocab shard. Shard outputs land * compactly in logits_by_tier[t] ([head_rows x split]) and are gathered * into spec_logits rows. */ int tp_tiers[DS4_MAX_GPUS] = {0}; const uint32_t tp_ways = (ok && g->cuda_tp_output) ? metal_graph_cuda_tp_output_tiers(g, tp_tiers) : 0; bool tp_ok = ok && tp_ways >= 2u && weights->output->type == DS4_TENSOR_Q8_0 && weights->output->ndim == 2 && weights->output->dim[0] == DS4_N_EMBD && weights->output->dim[1] == vocab_dim && head_rows <= DS4_DSPARK_MAX_BLOCK_SIZE && getenv("DS4_DSPARK_VERIFY_HEAD_NO_TP") == NULL; for (uint32_t i = 0; tp_ok && i < tp_ways; i++) { const int t = tp_tiers[i]; tp_ok = t >= 0 && t < DS4_MAX_GPUS && g->logits_by_tier[t] && ds4_gpu_tensor_bytes(g->logits_by_tier[t]) >= (uint64_t)head_rows * ((vocab_dim + tp_ways - 1u) / tp_ways) * sizeof(float) && (t == g->active_tier || (g->batch_ffn_norm_by_tier[t] && ds4_gpu_tensor_bytes(g->batch_ffn_norm_by_tier[t]) >= (uint64_t)head_rows * DS4_N_EMBD * sizeof(float))); } if (tp_ok) { const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); const int home_tier = g->active_tier; uint64_t split_start[DS4_MAX_GPUS] = {0}; uint64_t split_count[DS4_MAX_GPUS] = {0}; for (uint32_t i = 0; ok && i < tp_ways; i++) { const int t = tp_tiers[i]; split_start[i] = (vocab_dim * (uint64_t)i) / tp_ways; const uint64_t split_end = (vocab_dim * (uint64_t)(i + 1u)) / tp_ways; split_count[i] = split_end - split_start[i]; if (split_count[i] == 0) { ok = false; break; } if (t != home_tier) { ok = ds4_gpu_tensor_copy_xdev( g->batch_ffn_norm_by_tier[t], output_norm, (uint64_t)head_rows * DS4_N_EMBD * sizeof(float)) != 0; } } for (uint32_t i = 0; ok && i < tp_ways; i++) { const int t = tp_tiers[i]; ok = ds4_gpu_set_current_device(t) == 0; if (!ok) break; ds4_gpu_tensor *shard_out = ds4_gpu_tensor_view(g->logits_by_tier[t], 0, (uint64_t)head_rows * split_count[i] * sizeof(float)); ds4_gpu_tensor *shard_in = t == home_tier ? NULL : ds4_gpu_tensor_view(g->batch_ffn_norm_by_tier[t], 0, (uint64_t)head_rows * DS4_N_EMBD * sizeof(float)); ok = shard_out && (t == home_tier || shard_in) && ds4_gpu_matmul_q8_0_tensor(shard_out, model->map, model->size, weights->output->abs_offset + split_start[i] * row_bytes, DS4_N_EMBD, split_count[i], t == home_tier ? output_norm : shard_in, head_rows) != 0; ds4_gpu_tensor_free(shard_in); ds4_gpu_tensor_free(shard_out); } if (ok) ok = ds4_gpu_set_current_device(home_tier) == 0; for (uint32_t i = 0; ok && i < tp_ways; i++) { const int t = tp_tiers[i]; for (uint32_t r = 0; ok && r < n_tokens; r++) { ds4_gpu_tensor *dst = ds4_gpu_tensor_view(dst_logits, ((uint64_t)r * vocab_dim + split_start[i]) * sizeof(float), split_count[i] * sizeof(float)); ds4_gpu_tensor *src = ds4_gpu_tensor_view(g->logits_by_tier[t], (uint64_t)r * split_count[i] * sizeof(float), split_count[i] * sizeof(float)); ok = dst && src && ds4_gpu_tensor_copy_xdev(dst, src, split_count[i] * sizeof(float)) != 0; ds4_gpu_tensor_free(src); ds4_gpu_tensor_free(dst); } } } else if (ok && !(g->cuda_tp_ep && g->cuda_tp_output)) { ok = metal_graph_matmul_dense_quant_tensor(logits, model, weights->output, DS4_N_EMBD, vocab_dim, output_norm, head_rows); } else if (ok) { /* The expert-parallel cache stores only output vocabulary shards, so * a single-device full-head fallback would access uncached weights. */ ok = false; } ds4_gpu_tensor_free(logits); ds4_gpu_tensor_free(output_norm); return ok; } static bool metal_graph_encode_output_head_batch( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, uint32_t n_tokens, uint64_t vocab_dim) { if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; ds4_gpu_tensor *output_pre = NULL; ds4_gpu_tensor *output_weights = NULL; ds4_gpu_tensor *output_embd = NULL; ds4_gpu_tensor *output_norm = NULL; bool ok = true; output_pre = ds4_gpu_tensor_view(metal_graph_batch_hc_mix(g), 0, (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); output_weights = ds4_gpu_tensor_view(metal_graph_batch_hc_split(g), 0, (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); output_embd = ds4_gpu_tensor_view(metal_graph_batch_ffn_cur(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); output_norm = ds4_gpu_tensor_view(metal_graph_batch_ffn_norm(g), 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); ok = output_pre && output_weights && output_embd && output_norm; if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g), metal_graph_batch_cur_hc(g), (uint32_t)hc_dim, n_tokens, DS4_RMS_EPS) != 0; if (ok) ok = ds4_gpu_matmul_f16_tensor(output_pre, model->map, model->size, weights->output_hc_fn->abs_offset, hc_dim, DS4_N_HC, metal_graph_batch_flat_hc(g), n_tokens) != 0; if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, output_pre, model->map, model->size, weights->output_hc_scale->abs_offset, weights->output_hc_base->abs_offset, DS4_N_HC, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, metal_graph_batch_cur_hc(g), output_weights, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(output_norm, output_embd, model->map, model->size, weights->output_norm->abs_offset, DS4_N_EMBD, n_tokens, DS4_RMS_EPS) != 0; if (ok) ok = metal_graph_output_logits_head_matmul( g, model, weights, metal_graph_batch_ffn_norm(g), g->spec_logits, n_tokens, vocab_dim); ds4_gpu_tensor_free(output_norm); ds4_gpu_tensor_free(output_embd); ds4_gpu_tensor_free(output_weights); ds4_gpu_tensor_free(output_pre); return ok; } static bool metal_graph_matmul_plain_tensor( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *w, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { if (w->type == DS4_TENSOR_F16) { return ds4_gpu_matmul_f16_tensor(out, model->map, model->size, w->abs_offset, in_dim, out_dim, x, n_tok) != 0; } if (w->type == DS4_TENSOR_F32) { return ds4_gpu_matmul_f32_tensor(out, model->map, model->size, w->abs_offset, in_dim, out_dim, x, n_tok) != 0; } if (w->type == DS4_TENSOR_Q8_0) { return ds4_gpu_matmul_q8_0_tensor(out, model->map, model->size, w->abs_offset, in_dim, out_dim, x, n_tok) != 0; } if (tensor_type_is_dense_quant(w->type)) { return ds4_gpu_matmul_quant_tensor(out, model->map, model->size, w->abs_offset, w->type, in_dim, out_dim, x, n_tok) != 0; } fprintf(stderr, "ds4: Metal plain matmul does not support %s\n", tensor_type_name(w->type)); return false; } static bool metal_graph_dense_quant_row_bytes( const ds4_tensor *w, uint64_t in_dim, uint64_t *row_bytes) { if (row_bytes) *row_bytes = 0; if (!w || !row_bytes || !tensor_type_is_dense_quant(w->type)) return false; return tensor_nbytes(w->type, in_dim, row_bytes); } static bool metal_graph_matmul_dense_quant_abs( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *w, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { if (!w || !tensor_type_is_dense_quant(w->type)) return false; return ds4_gpu_matmul_quant_tensor(out, model->map, model->size, weight_offset, w->type, in_dim, out_dim, x, n_tok) != 0; } static bool metal_graph_matmul_dense_quant_tensor( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *w, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { if (!w) return false; return metal_graph_matmul_dense_quant_abs(out, model, w, w->abs_offset, in_dim, out_dim, x, n_tok); } static bool metal_graph_matmul_dense_quant_kslice( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *w, uint64_t full_in_dim, uint64_t k_off, uint64_t k_cnt, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t x_elem_off) { if (!w || !tensor_type_is_dense_quant(w->type)) return false; return ds4_gpu_matmul_quant_kslice_tensor(out, model->map, model->size, w->abs_offset, w->type, full_in_dim, k_off, k_cnt, out_dim, x, x_elem_off) != 0; } static bool metal_graph_attention_output_dense_quant_low( ds4_gpu_tensor *low, ds4_gpu_graph *g, const ds4_model *model, const ds4_tensor *out_a, uint64_t group_dim, uint64_t rank, uint32_t group0, uint32_t group_cnt, const ds4_gpu_tensor *heads) { (void)g; if (!low || !model || !out_a || !heads || group_dim == 0 || rank == 0 || group_cnt == 0) { return false; } if (out_a->type == DS4_TENSOR_Q8_0 && group0 == 0) { return ds4_gpu_attention_output_low_q8_tensor(low, model->map, model->size, out_a->abs_offset, group_dim, rank, group_cnt, heads) != 0; } if (out_a->type == DS4_TENSOR_Q4_K) { return ds4_gpu_attention_output_low_q4_K_slice_tensor(low, model->map, model->size, out_a->abs_offset, group_dim, rank, group0, group_cnt, heads) != 0; } uint64_t row_bytes = 0; if (!metal_graph_dense_quant_row_bytes(out_a, group_dim, &row_bytes)) return false; const uint64_t group_weight_bytes = rank * row_bytes; bool ok = true; for (uint32_t i = 0; ok && i < group_cnt; i++) { ds4_gpu_tensor *head_view = ds4_gpu_tensor_view( heads, (uint64_t)i * group_dim * sizeof(float), group_dim * sizeof(float)); ds4_gpu_tensor *low_view = ds4_gpu_tensor_view( low, (uint64_t)i * rank * sizeof(float), rank * sizeof(float)); ok = head_view && low_view && metal_graph_matmul_dense_quant_abs(low_view, model, out_a, out_a->abs_offset + (uint64_t)(group0 + i) * group_weight_bytes, group_dim, rank, head_view, 1); ds4_gpu_tensor_free(low_view); ds4_gpu_tensor_free(head_view); } return ok; } static bool metal_graph_attention_output_dense_quant_tp( ds4_gpu_tensor *out, ds4_gpu_tensor *low, ds4_gpu_graph *g, const ds4_model *model, const ds4_tensor *out_a, const ds4_tensor *out_b, uint64_t group_dim, uint64_t rank, uint32_t n_groups_total, uint32_t group0, uint32_t group_cnt, uint64_t out_dim, const ds4_gpu_tensor *heads) { if (!out || !low || !g || !model || !out_a || !out_b || !heads || group0 + group_cnt > n_groups_total) { return false; } if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { return ds4_gpu_attention_output_q8_tp_tensor(out, low, model->map, model->size, out_a->abs_offset, out_b->abs_offset, group_dim, rank, n_groups_total, group0, group_cnt, out_dim, heads) != 0; } if (!metal_graph_attention_output_dense_quant_low(low, g, model, out_a, group_dim, rank, group0, group_cnt, heads)) { return false; } return metal_graph_matmul_dense_quant_kslice(out, model, out_b, (uint64_t)n_groups_total * rank, (uint64_t)group0 * rank, (uint64_t)group_cnt * rank, out_dim, low, 0); } static bool metal_graph_attention_output_dense_quant_batch( ds4_gpu_tensor *out, ds4_gpu_tensor *low, ds4_gpu_graph *g, const ds4_model *model, const ds4_tensor *out_a, const ds4_tensor *out_b, uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) { if (!out || !low || !g || !model || !out_a || !out_b || !heads || n_groups == 0 || n_tokens == 0) { return false; } if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { return ds4_gpu_attention_output_q8_batch_tensor(out, low, metal_graph_batch_group_tmp(g), metal_graph_batch_low_tmp(g), model->map, model->size, out_a->abs_offset, out_b->abs_offset, group_dim, rank, n_groups, out_dim, heads, n_tokens) != 0; } if (out_a->type == DS4_TENSOR_Q4_K && n_tokens >= 32u) { if (ds4_gpu_attention_output_q4_K_batch_tensor(out, low, metal_graph_batch_group_tmp(g), metal_graph_batch_low_tmp(g), model->map, model->size, out_a->abs_offset, out_b->abs_offset, out_b->type, group_dim, rank, n_groups, out_dim, heads, n_tokens) != 0) { return true; } } const uint64_t heads_row_elems = (uint64_t)n_groups * group_dim; const uint64_t low_row_elems = (uint64_t)n_groups * rank; bool ok = true; for (uint32_t t = 0; ok && t < n_tokens; t++) { ds4_gpu_tensor *heads_row = ds4_gpu_tensor_view( heads, (uint64_t)t * heads_row_elems * sizeof(float), heads_row_elems * sizeof(float)); ds4_gpu_tensor *low_row = ds4_gpu_tensor_view( low, (uint64_t)t * low_row_elems * sizeof(float), low_row_elems * sizeof(float)); ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( out, (uint64_t)t * out_dim * sizeof(float), out_dim * sizeof(float)); ok = heads_row && low_row && out_row && metal_graph_attention_output_dense_quant_low(low_row, g, model, out_a, group_dim, rank, 0, n_groups, heads_row); if (ok) ok = metal_graph_matmul_dense_quant_tensor(out_row, model, out_b, low_row_elems, out_dim, low_row, 1); ds4_gpu_tensor_free(out_row); ds4_gpu_tensor_free(low_row); ds4_gpu_tensor_free(heads_row); } return ok; } static bool metal_graph_matmul_q8_0_named_tensor( const char *module, uint32_t il, uint32_t pos0, ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *w, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { (void)module; (void)il; (void)pos0; return metal_graph_matmul_dense_quant_tensor(out, model, w, in_dim, out_dim, x, n_tok); } static bool metal_graph_encode_output_head_mtp( ds4_gpu_graph *g, const ds4_model *base_model, const ds4_weights *base_weights, const ds4_model *mtp_model, const ds4_mtp_weights *mtp, uint64_t vocab_dim) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), metal_graph_cur_hc(g), (uint32_t)hc_dim, DS4_RMS_EPS) != 0; if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_output_pre(g), mtp_model, mtp->hc_head_fn, hc_dim, DS4_N_HC, metal_graph_flat_hc(g), 1); if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), metal_graph_output_pre(g), mtp_model->map, mtp_model->size, mtp->hc_head_scale->abs_offset, mtp->hc_head_base->abs_offset, DS4_N_HC, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), metal_graph_cur_hc(g), metal_graph_output_weights(g), DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g), metal_graph_output_embd(g), mtp_model->map, mtp_model->size, mtp->norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_logits(g), base_model, base_weights->output, DS4_N_EMBD, vocab_dim, metal_graph_output_norm(g), 1); return ok; } /* ========================================================================= * Metal Diagnostic Comparisons. * ========================================================================= * * These routines deliberately allocate CPU-side reference buffers and read * Metal tensors back. They are not part of generation; command-line tests use * them to localize drift against the C reference pipeline. */ static void metal_graph_trace_layer_stages( ds4_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *layer, const float *cpu_in_hc, uint32_t il, int token) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t q_rank = layer->attn_q_a->dim[1]; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t shared_in_dim = layer->ffn_gate_shexp->dim[0]; const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); float *cpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); float *cpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float post[4]; float comb[16]; float ffn_post[4]; float ffn_comb[16]; int selected[DS4_MAX_EXPERT_USED]; float expert_weight[DS4_MAX_EXPERT_USED]; const uint64_t shared_blocks = (shared_in_dim + 31) / 32; int8_t *shared_xq = xmalloc((size_t)shared_blocks * 32); float *shared_xscale = xmalloc((size_t)shared_blocks * sizeof(float)); float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; hc_pre_from_state_one(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, cpu_in_hc, cpu_attn_cur, post, comb); layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, false); dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); layer_attention_one(cpu_heads, model, layer, cpu_q, cpu_kv); rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, il, true); layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_in_hc, post, comb, DS4_N_EMBD, DS4_N_HC); hc_pre_from_state_one(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, cpu_after_attn_hc, cpu_ffn_cur, ffn_post, ffn_comb); rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); quantize_q8_0_activation(cpu_ffn_norm, shared_xq, shared_xscale, shared_in_dim); matvec_q8_0_pair_prequant(cpu_shared_gate, cpu_shared_up, model, layer->ffn_gate_shexp, layer->ffn_up_shexp, shared_xq, shared_xscale); swiglu(cpu_shared_mid, cpu_shared_gate, cpu_shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP); matvec_q8_0(cpu_shared, model, layer->ffn_down_shexp, cpu_shared_mid); layer_routed_moe_one_prealloc(cpu_routed, model, layer, cpu_ffn_norm, il, token, DS4_SWIGLU_CLAMP_EXP, routed_mid_all, routed_xq, routed_midq, routed_q8_xq, routed_q8_xscale, routed_q8_midq, routed_q8_midscale); if (layer->ffn_gate_tid2eid) { layer_hash_selected_experts(selected, model, layer, token); layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); } else { layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); } for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; hc_post_one(cpu_after_ffn_hc, cpu_ffn_out, cpu_after_attn_hc, ffn_post, ffn_comb, DS4_N_EMBD, DS4_N_HC); float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_shared_gate = xmalloc((size_t)shared_dim * sizeof(float)); float *gpu_shared_up = xmalloc((size_t)shared_dim * sizeof(float)); float *gpu_shared_mid = xmalloc((size_t)shared_dim * sizeof(float)); float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); int gpu_selected[DS4_MAX_EXPERT_USED]; float gpu_expert_weight[DS4_MAX_EXPERT_USED]; bool ok = ds4_gpu_tensor_read(metal_graph_attn_cur(g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_attn_norm(g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_q(g), 0, gpu_q, q_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_kv(g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_attn_out(g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_after_attn_hc(g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_cur(g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_norm(g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_shared_gate(g), 0, gpu_shared_gate, shared_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_shared_up(g), 0, gpu_shared_up, shared_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_shared_mid(g), 0, gpu_shared_mid, shared_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_shared_out(g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_router_selected(g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && ds4_gpu_tensor_read(metal_graph_router_weights(g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && ds4_gpu_tensor_read(metal_graph_routed_mid(g), 0, gpu_routed_mid_all, (uint64_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_routed_out(g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_out(g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_cur_hc(g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0; if (ok) { fprintf(stderr, "ds4: Metal stage layer %u attn_cur=%g/%g attn_norm=%g/%g q=%g/%g kv=%g/%g attn_out=%g/%g after_attn_hc=%g/%g ffn_cur=%g/%g ffn_norm=%g/%g shared=%g/%g router_w=%g routed=%g/%g ffn_out=%g/%g after_ffn_hc=%g/%g\n", il, max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), rms_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), rms_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), max_abs_diff(cpu_q, gpu_q, q_dim), rms_abs_diff(cpu_q, gpu_q, q_dim), max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), rms_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), rms_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), rms_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), rms_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), rms_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), rms_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), rms_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim)); fprintf(stderr, "ds4: Metal shared layer %u gate=%g/%g up=%g/%g mid=%g/%g down=%g/%g\n", il, max_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), rms_abs_diff(cpu_shared_gate, gpu_shared_gate, shared_dim), max_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), rms_abs_diff(cpu_shared_up, gpu_shared_up, shared_dim), max_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), rms_abs_diff(cpu_shared_mid, gpu_shared_mid, shared_dim), max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), rms_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD)); fprintf(stderr, "ds4: Metal routed layer %u mid=%g/%g out=%g/%g\n", il, max_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), rms_abs_diff(routed_mid_all, gpu_routed_mid_all, DS4_N_EXPERT_USED * down_in_dim), max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), rms_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD)); if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { fprintf(stderr, "ds4: Metal stage layer %u router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", il, selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); } } free(gpu_after_ffn_hc); free(gpu_ffn_out); free(gpu_routed); free(gpu_routed_mid_all); free(gpu_shared); free(gpu_shared_mid); free(gpu_shared_up); free(gpu_shared_gate); free(gpu_ffn_norm); free(gpu_ffn_cur); free(gpu_after_attn_hc); free(gpu_attn_out); free(gpu_kv); free(gpu_q); free(gpu_attn_norm); free(gpu_attn_cur); free(routed_q8_midscale); free(routed_q8_midq); free(routed_q8_xscale); free(routed_q8_xq); free(routed_midq); free(routed_xq); free(routed_mid_all); free(shared_xscale); free(shared_xq); free(cpu_after_ffn_hc); free(cpu_ffn_out); free(cpu_routed); free(cpu_shared); free(cpu_shared_mid); free(cpu_shared_up); free(cpu_shared_gate); free(cpu_ffn_norm); free(cpu_ffn_cur); free(cpu_after_attn_hc); free(cpu_attn_out); free(cpu_heads); free(cpu_kv); free(cpu_qr_norm); free(cpu_q); free(cpu_attn_norm); free(cpu_attn_cur); } static int metal_graph_decode_test( const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, bool quality) { if (prompt->len <= 0) { fprintf(stderr, "ds4: Metal graph test needs a non-empty prompt\n"); return 1; } const int token = prompt->v[0]; const ds4_layer_weights *layer = &weights->layer[0]; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t q_rank = layer->attn_q_a->dim[1]; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; const uint64_t vocab_dim = weights->output->dim[1]; const bool routed_q8_0 = layer->ffn_gate_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_up_exps->type == DS4_TENSOR_Q8_0 && layer->ffn_down_exps->type == DS4_TENSOR_Q8_0; const uint64_t routed_q8_x_blocks = expert_in_dim / 32u; const uint64_t routed_q8_mid_blocks = down_in_dim / 32u; float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); float *cpu_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); float *cpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_qr_norm = xmalloc((size_t)q_rank * sizeof(float)); float *cpu_q = xmalloc((size_t)q_dim * sizeof(float)); float *cpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); float *cpu_heads = xmalloc((size_t)q_dim * sizeof(float)); float *cpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_ffn_post = xmalloc((size_t)DS4_N_HC * sizeof(float)); float *cpu_ffn_comb = xmalloc((size_t)DS4_N_HC * DS4_N_HC * sizeof(float)); float *cpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *gpu_attn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_attn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_q = xmalloc((size_t)q_dim * sizeof(float)); float *gpu_kv = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); float *gpu_raw = xmalloc((size_t)DS4_N_HEAD_DIM * sizeof(float)); float *gpu_attn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_after_attn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *gpu_ffn_cur = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_ffn_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_shared = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_routed = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_ffn_out = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *gpu_after_ffn_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); int gpu_selected[DS4_MAX_EXPERT_USED]; float gpu_expert_weight[DS4_MAX_EXPERT_USED]; float *routed_mid_all = xmalloc((size_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)); block_q8_K *routed_xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(block_q8_K)); block_q8_K *routed_midq = xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(block_q8_K)); int8_t *routed_q8_xq = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * 32u) : NULL; float *routed_q8_xscale = routed_q8_0 ? xmalloc((size_t)routed_q8_x_blocks * sizeof(routed_q8_xscale[0])) : NULL; int8_t *routed_q8_midq = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * 32u) : NULL; float *routed_q8_midscale = routed_q8_0 ? xmalloc((size_t)DS4_N_EXPERT_USED * routed_q8_mid_blocks * sizeof(routed_q8_midscale[0])) : NULL; int selected[DS4_MAX_EXPERT_USED]; float expert_weight[DS4_MAX_EXPERT_USED]; embed_token_f16(model, weights, token, plain); hc_from_plain_embedding(cpu_hc, plain, DS4_N_EMBD, DS4_N_HC); hc_pre_from_state_one(model, layer->hc_attn_fn, layer->hc_attn_scale, layer->hc_attn_base, cpu_hc, cpu_attn_cur, cpu_post, cpu_comb); layer_attn_norm_one(cpu_attn_norm, model, layer, cpu_attn_cur); layer_q_projection_with_lora_one(model, layer, cpu_attn_norm, cpu_q, cpu_qr_norm); layer_kv_projection_normed_one(model, layer, cpu_attn_norm, cpu_kv); rope_tail_layer_inplace(cpu_q, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); rope_tail_layer_inplace(cpu_kv, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, false); dsv4_fp8_kv_quantize_row_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM, DS4_N_ROT); f16_round_inplace_cpu(cpu_kv, DS4_N_HEAD_DIM); layer_attention_rows_one(cpu_heads, model, layer, cpu_q, cpu_kv, 1); rope_tail_layer_inplace(cpu_heads, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, 0, 0, true); layer_grouped_out_one(cpu_attn_out, model, layer, cpu_heads); hc_post_one(cpu_after_attn_hc, cpu_attn_out, cpu_hc, cpu_post, cpu_comb, DS4_N_EMBD, DS4_N_HC); hc_pre_from_state_one(model, layer->hc_ffn_fn, layer->hc_ffn_scale, layer->hc_ffn_base, cpu_after_attn_hc, cpu_ffn_cur, cpu_ffn_post, cpu_ffn_comb); rms_norm_weight(cpu_ffn_norm, cpu_ffn_cur, tensor_data(model, layer->ffn_norm), DS4_N_EMBD, DS4_RMS_EPS); layer_shared_ffn_one(cpu_shared, model, layer, cpu_ffn_norm); layer_routed_moe_one_prealloc(cpu_routed, model, layer, cpu_ffn_norm, 0, token, DS4_SWIGLU_CLAMP_EXP, routed_mid_all, routed_xq, routed_midq, routed_q8_xq, routed_q8_xscale, routed_q8_midq, routed_q8_midscale); if (layer->ffn_gate_tid2eid) { layer_hash_selected_experts(selected, model, layer, token); layer_hash_router_weights_one(expert_weight, model, layer, cpu_ffn_norm, selected); } else { layer_topk_selected_experts(selected, expert_weight, model, layer, cpu_ffn_norm); } for (uint32_t i = 0; i < DS4_N_EMBD; i++) cpu_ffn_out[i] = cpu_shared[i] + cpu_routed[i]; hc_post_one(cpu_after_ffn_hc, cpu_ffn_out, cpu_after_attn_hc, cpu_ffn_post, cpu_ffn_comb, DS4_N_EMBD, DS4_N_HC); output_logits_one(cpu_logits, model, weights, cpu_after_ffn_hc); ds4_gpu_graph g; bool ok = metal_graph_alloc(&g, weights, layer); g.quality = quality; g.materialize_ffn_out = true; if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), model->map, model->size, weights->token_embd->abs_offset, (uint32_t)weights->token_embd->dim[1], (uint32_t)token, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = metal_graph_encode_decode_layer(&g, model, layer, 0, 0, g.layer_raw_cache[0], g.raw_cap, 0, 1, token); if (ok) { /* Single-tier diagnostic: swap the active-tier slots so the head * pipeline reads the embedded hidden state from cur_hc. */ ds4_gpu_tensor *embedded_hc = g.cur_hc_by_tier[g.active_tier]; g.cur_hc_by_tier[g.active_tier] = g.after_ffn_hc_by_tier[g.active_tier]; g.after_ffn_hc_by_tier[g.active_tier] = embedded_hc; } if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); if (ok) ok = ds4_gpu_end_commands() != 0; if (ok) { ok = ds4_gpu_tensor_read(metal_graph_after_ffn_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_attn_cur(&g), 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_attn_norm(&g), 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_q(&g), 0, gpu_q, q_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_kv(&g), 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && ds4_gpu_tensor_read(g.layer_raw_cache[0], 0, gpu_raw, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_attn_out(&g), 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_after_attn_hc(&g), 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_cur(&g), 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_norm(&g), 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_shared_out(&g), 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_router_selected(&g), 0, gpu_selected, sizeof(gpu_selected)) != 0 && ds4_gpu_tensor_read(metal_graph_router_weights(&g), 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && ds4_gpu_tensor_read(metal_graph_routed_out(&g), 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_ffn_out(&g), 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; } if (ok) { fprintf(stderr, "ds4: Metal graph test layer0 diffs: embed_hc=%g hc_pre=%g attn_norm=%g q_rope=%g kv_rope=%g raw_cache=%g attn_out=%g after_attn_hc=%g ffn_cur=%g ffn_norm=%g shared=%g router_w=%g routed=%g ffn_out=%g after_ffn_hc=%g logits=%g\n", max_abs_diff(cpu_hc, gpu_hc, hc_dim), max_abs_diff(cpu_attn_cur, gpu_attn_cur, DS4_N_EMBD), max_abs_diff(cpu_attn_norm, gpu_attn_norm, DS4_N_EMBD), max_abs_diff(cpu_q, gpu_q, q_dim), max_abs_diff(cpu_kv, gpu_kv, DS4_N_HEAD_DIM), max_abs_diff(cpu_kv, gpu_raw, DS4_N_HEAD_DIM), max_abs_diff(cpu_attn_out, gpu_attn_out, DS4_N_EMBD), max_abs_diff(cpu_after_attn_hc, gpu_after_attn_hc, hc_dim), max_abs_diff(cpu_ffn_cur, gpu_ffn_cur, DS4_N_EMBD), max_abs_diff(cpu_ffn_norm, gpu_ffn_norm, DS4_N_EMBD), max_abs_diff(cpu_shared, gpu_shared, DS4_N_EMBD), max_abs_diff(expert_weight, gpu_expert_weight, DS4_N_EXPERT_USED), max_abs_diff(cpu_routed, gpu_routed, DS4_N_EMBD), max_abs_diff(cpu_ffn_out, gpu_ffn_out, DS4_N_EMBD), max_abs_diff(cpu_after_ffn_hc, gpu_after_ffn_hc, hc_dim), max_abs_diff(cpu_logits, gpu_logits, vocab_dim)); if (memcmp(selected, gpu_selected, sizeof(selected)) != 0) { fprintf(stderr, "ds4: Metal graph router selected mismatch: cpu=[%d,%d,%d,%d,%d,%d] gpu=[%d,%d,%d,%d,%d,%d]\n", selected[0], selected[1], selected[2], selected[3], selected[4], selected[5], gpu_selected[0], gpu_selected[1], gpu_selected[2], gpu_selected[3], gpu_selected[4], gpu_selected[5]); } print_vec_stats("metal graph q", gpu_q, q_dim); print_vec_stats("metal graph kv", gpu_kv, DS4_N_HEAD_DIM); print_vec_stats("metal graph routed", gpu_routed, DS4_N_EMBD); } else { fprintf(stderr, "ds4: Metal graph test failed while encoding first decode stages\n"); if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after graph test failure also failed\n"); } } metal_graph_free(&g); free(routed_q8_midscale); free(routed_q8_midq); free(routed_q8_xscale); free(routed_q8_xq); free(routed_midq); free(routed_xq); free(routed_mid_all); free(gpu_logits); free(gpu_after_ffn_hc); free(gpu_ffn_out); free(gpu_routed); free(gpu_shared); free(gpu_ffn_norm); free(gpu_ffn_cur); free(gpu_after_attn_hc); free(gpu_attn_out); free(gpu_raw); free(gpu_kv); free(gpu_q); free(gpu_attn_norm); free(gpu_attn_cur); free(gpu_hc); free(cpu_kv); free(cpu_q); free(cpu_attn_out); free(cpu_heads); free(cpu_ffn_norm); free(cpu_routed); free(cpu_logits); free(cpu_after_ffn_hc); free(cpu_ffn_out); free(cpu_shared); free(cpu_ffn_comb); free(cpu_ffn_post); free(cpu_ffn_cur); free(cpu_after_attn_hc); free(cpu_qr_norm); free(cpu_attn_norm); free(cpu_comb); free(cpu_post); free(cpu_attn_cur); free(cpu_hc); free(plain); return ok ? 0 : 1; } static int metal_graph_first_token_full_test( const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, bool quality) { if (prompt->len <= 0) { fprintf(stderr, "ds4: full Metal graph test needs a non-empty prompt\n"); return 1; } const int token = prompt->v[0]; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t vocab_dim = weights->output->dim[1]; float *cpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *gpu_hc = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); float *gpu_logits = xmalloc((size_t)vocab_dim * sizeof(float)); forward_first_token_cpu(cpu_hc, model, weights, token); output_logits_one(cpu_logits, model, weights, cpu_hc); ds4_gpu_graph g; bool ok = metal_graph_alloc(&g, weights, &weights->layer[0]); g.quality = quality; const bool trace_layers = getenv("DS4_METAL_GRAPH_TRACE_LAYERS") != NULL; if (trace_layers && ok) { g.materialize_ffn_out = true; const bool teacher_force = getenv("DS4_METAL_GRAPH_TEACHER_FORCE") != NULL; const char *stage_layer_env = getenv("DS4_METAL_GRAPH_TRACE_STAGE_LAYER"); const long stage_layer = stage_layer_env ? strtol(stage_layer_env, NULL, 10) : -1; float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(float)); float *cpu_cur = xmalloc((size_t)hc_dim * sizeof(float)); float *cpu_next = xmalloc((size_t)hc_dim * sizeof(float)); embed_token_f16(model, weights, token, plain); hc_from_plain_embedding(cpu_cur, plain, DS4_N_EMBD, DS4_N_HC); ok = ds4_gpu_begin_commands() != 0; if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), model->map, model->size, weights->token_embd->abs_offset, (uint32_t)weights->token_embd->dim[1], (uint32_t)token, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = ds4_gpu_end_commands() != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { if (teacher_force) { ok = ds4_gpu_tensor_write(metal_graph_cur_hc(&g), 0, cpu_cur, hc_dim * sizeof(float)) != 0; } ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); g.after_ffn_hc_by_tier[g.active_tier] = tmp; if (ok) ok = ds4_gpu_end_commands() != 0; layer_forward_self_one(cpu_next, model, &weights->layer[il], cpu_cur, il, 0, token); if (ok) ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0; if (ok) { fprintf(stderr, "ds4: Metal full graph layer %u%s hc_max=%g hc_rms=%g\n", il, teacher_force ? " teacher" : "", max_abs_diff(cpu_next, gpu_hc, hc_dim), rms_abs_diff(cpu_next, gpu_hc, hc_dim)); if (stage_layer == (long)il) { metal_graph_trace_layer_stages(&g, model, &weights->layer[il], cpu_cur, il, token); } } float *ctmp = cpu_cur; cpu_cur = cpu_next; cpu_next = ctmp; } if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); if (ok) ok = ds4_gpu_end_commands() != 0; free(cpu_next); free(cpu_cur); free(plain); } else { if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(&g), model->map, model->size, weights->token_embd->abs_offset, (uint32_t)weights->token_embd->dim[1], (uint32_t)token, DS4_N_EMBD, DS4_N_HC) != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); ds4_gpu_tensor *tmp = metal_graph_cur_hc(&g); g.cur_hc_by_tier[g.active_tier] = metal_graph_after_ffn_hc(&g); g.after_ffn_hc_by_tier[g.active_tier] = tmp; } if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); if (ok) ok = ds4_gpu_end_commands() != 0; } if (ok) { ok = ds4_gpu_tensor_read(metal_graph_cur_hc(&g), 0, gpu_hc, hc_dim * sizeof(float)) != 0 && ds4_gpu_tensor_read(metal_graph_logits(&g), 0, gpu_logits, vocab_dim * sizeof(float)) != 0; } if (ok) { const uint64_t cpu_top = argmax_f32(cpu_logits, vocab_dim); const uint64_t gpu_top = argmax_f32(gpu_logits, vocab_dim); fprintf(stderr, "ds4: Metal full first-token graph diffs: final_hc_max=%g final_hc_rms=%g logits_max=%g logits_rms=%g cpu_top=%llu gpu_top=%llu cpu_top_logit=%g gpu_top_logit=%g\n", max_abs_diff(cpu_hc, gpu_hc, hc_dim), rms_abs_diff(cpu_hc, gpu_hc, hc_dim), max_abs_diff(cpu_logits, gpu_logits, vocab_dim), rms_abs_diff(cpu_logits, gpu_logits, vocab_dim), (unsigned long long)cpu_top, (unsigned long long)gpu_top, cpu_logits[cpu_top], gpu_logits[gpu_top]); } else { fprintf(stderr, "ds4: Metal full first-token graph test failed\n"); if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after full graph failure also failed\n"); } } metal_graph_free(&g); free(gpu_logits); free(cpu_logits); free(gpu_hc); free(cpu_hc); return ok ? 0 : 1; } /* ========================================================================= * Metal Release Decode and Prefill. * ========================================================================= * * Everything below is the user-facing Metal backend. It uses the same layer * encoder as diagnostics, but diagnostics are not required for normal command * flow and their CPU reads stay outside these generation entry points. */ static DS4_MAYBE_UNUSED bool metal_graph_pre_m5_q2_decode_schedule_eligible( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t pos, bool allow_split_flush) { #if defined(__APPLE__) if (!g || !weights || !allow_split_flush || DS4_N_LAYER <= 32u || !ds4_gpu_device_is_pre_m5_apple_silicon() || pos < 128u || pos >= 2816u || g->quality || g->ssd_streaming || g->ssd_streaming_cold || g->placement != NULL || g->tp_world > 1u || getenv("DS4_METAL_DISABLE_PRE_M5_Q2_DECODE_SPLIT2_32") != NULL || getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != NULL) { return false; } const ds4_layer_weights *layer = &weights->layer[4]; return layer->ffn_gate_exps && layer->ffn_up_exps && layer->ffn_down_exps && 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 == 256u && DS4_N_EXPERT_USED == 6u && layer->ffn_gate_exps->dim[0] == 4096u && layer->ffn_gate_exps->dim[1] == 2048u && layer->ffn_up_exps->dim[0] == 4096u && layer->ffn_up_exps->dim[1] == 2048u && layer->ffn_down_exps->dim[0] == 2048u && layer->ffn_down_exps->dim[1] == 4096u && routed_expert_row_bytes(layer->ffn_gate_exps) == 1056u && routed_expert_row_bytes(layer->ffn_up_exps) == 1056u && routed_expert_row_bytes(layer->ffn_down_exps) == 672u; #else (void)g; (void)weights; (void)pos; (void)allow_split_flush; return false; #endif } 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 uint32_t metal_graph_token_adaptive_split_after_layers( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t pos, uint32_t second_split_after_layers, bool allow_split_flush) { const uint32_t split_after_layers = metal_graph_token_split_after_layers(); #if defined(__APPLE__) /* In the validated 128..2815 Q2 window, routed decode exposes enough GPU * work after layer two to overlap encoding before the layer-32 flush. */ const char *split_env = getenv("DS4_METAL_GRAPH_TOKEN_SPLIT_LAYERS"); if ((!split_env || !split_env[0]) && metal_graph_pre_m5_q2_decode_schedule_eligible( g, weights, pos, allow_split_flush)) { return 2u; } /* Layer 4 is a routed layer in the fixed DS4 tape and remains available in * the resident full-model path this schedule targets. */ const ds4_layer_weights *layer = weights ? &weights->layer[4] : NULL; const bool mxfp4_routed = layer && layer->ffn_gate_exps && layer->ffn_up_exps && layer->ffn_down_exps && layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4; /* Once the raw SWA window is full, the short-context 4/12 schedule can * submit its first useful work one routed layer earlier. Five balanced * exact-output blocks (three eval-length and two 1,024-token blocks) * favored 3/12; keep the established 4/16 and adaptive 5-layer schedules * unchanged outside this early second-split window. */ if (split_after_layers == 4u && second_split_after_layers == 12u && allow_split_flush && pos >= 128u && pos < 2048u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && g->tp_world != 2u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_SPLIT3") == NULL) { return 3u; } /* At the 2K frontier the fifth layer adds enough work to overlap command * encoding. Past 2816 the larger first command buffer starts losing to * the established four-layer split, so retain that schedule there. */ if (split_after_layers == 4u && second_split_after_layers == 0u && allow_split_flush && pos >= 2048u && pos < 2816u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && g->tp_world != 2u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_SPLIT5") == NULL) { return 5u; } #else (void)g; (void)weights; (void)pos; (void)second_split_after_layers; (void)allow_split_flush; #endif return split_after_layers; } static DS4_MAYBE_UNUSED bool metal_graph_decode_pipeline_fast_lookup_eligible( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t pos, bool allow_split_flush) { #if defined(__APPLE__) if (!g || !weights || !allow_split_flush || g->quality || g->ssd_streaming || g->ssd_streaming_cold || g->placement != NULL || g->tp_world > 1u || g->mtp_enabled || DS4_N_LAYER <= 4u) { return false; } /* Eval prompts commonly decode below the 2K benchmark frontier. The * mirror key is the complete pipeline specialization and is independent * of position, so the same lookup is valid there; keep an early-only * rollback for attribution without disabling the established 2K path. */ if (pos < 2048u && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_PIPELINE_FAST_LOOKUP") != NULL) { return false; } const ds4_layer_weights *layer = &weights->layer[4]; const bool mxfp4_routed = layer->ffn_gate_exps && layer->ffn_up_exps && layer->ffn_down_exps && layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4; return mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PIPELINE_FAST_LOOKUP") == NULL; #else (void)g; (void)weights; (void)pos; (void)allow_split_flush; return false; #endif } static uint32_t metal_graph_token_second_split_after_layers(void) { uint32_t split_after_layers = 0; #if defined(__APPLE__) const char *split_env = getenv("DS4_METAL_GRAPH_TOKEN_SECOND_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; } /* Automatic second command-buffer splits for eligible resident pre-M5 decode. * They close the idle bubble left while the host encodes a long tail after the * first submission. Only command-buffer boundaries move, so kernel math, * ordering, and outputs are unchanged. An explicit * DS4_METAL_GRAPH_TOKEN_SECOND_SPLIT_LAYERS always wins. */ static uint32_t metal_graph_token_adaptive_second_split_after_layers( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t pos, uint32_t env_second_split_after_layers, bool allow_split_flush) { #if defined(__APPLE__) const char *second_split_env = getenv("DS4_METAL_GRAPH_TOKEN_SECOND_SPLIT_LAYERS"); if (second_split_env && second_split_env[0]) { return env_second_split_after_layers; } if (metal_graph_token_split_after_layers() != 0u && metal_graph_pre_m5_q2_decode_schedule_eligible( g, weights, pos, allow_split_flush)) { return 32u; } /* Mirror the R4 adaptive-first-split guards: layer 4 is a routed layer in * the fixed DS4 tape and stands in for the resident full-model path this * schedule targets. */ const ds4_layer_weights *layer = weights ? &weights->layer[4] : NULL; const bool mxfp4_routed = layer && layer->ffn_gate_exps && layer->ffn_up_exps && layer->ffn_down_exps && layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4; const bool eligible = allow_split_flush && pos < 3328u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && g->tp_world != 2u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon(); /* Once the raw SWA window is full, short eval prompts leave less GPU work * in the four-layer prefix while the host still has the same remaining * tape to encode. Flushing again after layer 12 starts that middle work * soon enough to cover the early decode bubble; 12 measured better than * 8, 16, and 20 below position 2048. Positions below 128 retain the * single-tail schedule, which is slightly faster before the window fills. * Keep a separate rollback so the established 2K schedule remains * independently attributable. */ if (eligible && pos >= 128u && pos < 2048u && DS4_N_LAYER > 12u && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_SECOND_SPLIT12") == NULL) { return 12u; } /* The encode/GPU overlap bubble the second split closes persists past * 2815: attention cost grows only slowly through the compressed-KV * window (raw SWA saturates at 128 rows) while encode cost per layer is * constant, so the same 4/16 schedule stays profitable through 3327. */ if (eligible && pos >= 2048u && DS4_N_LAYER > 16u && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_SECOND_SPLIT16") == NULL) { return 16u; } #else (void)g; (void)weights; (void)pos; (void)allow_split_flush; #endif return env_second_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; } /* The physical ring also holds temporary draft rows; only trusted target features belong to this logical window. */ static bool metal_graph_dspark_cache_merge_target_range(ds4_gpu_graph *g, uint32_t start, uint32_t len) { if (!g || len == 0 || len > g->dspark_cache_cap || start > UINT32_MAX - len || DS4_N_SWA == 0 || !metal_graph_dspark_cache_current_window_valid(g)) return false; const uint32_t end = start + len; uint32_t first = start; if (g->dspark_cache_len != 0) { const uint32_t old_start = g->dspark_cache_token_start; const uint32_t old_end = old_start + g->dspark_cache_len; /* Incoming capture is authoritative: retain only an overlapping/adjacent prefix, never its old future. */ if (old_start < start && start <= old_end) first = old_start; } uint32_t limit = DS4_N_SWA; if (limit > g->dspark_cache_cap) limit = g->dspark_cache_cap; if (end - first > limit) first = end - limit; return metal_graph_dspark_cache_set_window(g, first, end - first); } static bool metal_graph_dspark_cache_target_prefix(ds4_gpu_graph *g, uint32_t feature_pos) { if (!g || DS4_N_SWA == 0 || !metal_graph_dspark_cache_crop_to_prefix(g, feature_pos) || !metal_graph_dspark_cache_ends_at(g, feature_pos)) return false; const uint32_t limit = DS4_N_SWA - 1u; if (g->dspark_cache_len > limit) { return metal_graph_dspark_cache_set_window(g, feature_pos - limit, limit); } return true; } 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; } static bool metal_graph_dspark_capture_commit_prefix( ds4_gpu_graph *g, uint32_t prefix_len) { if (!g || !g->dspark_capture_enabled || !g->dspark_capture_batch_valid || !g->dspark_target_hidden || !g->dspark_target_hidden_batch || prefix_len == 0 || prefix_len >= g->dspark_capture_batch_tokens) { return false; } const uint64_t embd_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); bool ok = ds4_gpu_begin_commands() != 0; metal_graph_dspark_capture_row_invalidate(g); for (uint32_t slot = 0; ok && slot < g->dspark_target_layer_count; slot++) { ds4_gpu_tensor *src = ds4_gpu_tensor_view( g->dspark_target_hidden_batch, (((uint64_t)slot * g->prefill_cap + prefix_len) * DS4_N_EMBD) * sizeof(float), embd_bytes); ds4_gpu_tensor *dst = ds4_gpu_tensor_view( g->dspark_target_hidden, (uint64_t)slot * embd_bytes, embd_bytes); ok = src && dst && ds4_gpu_tensor_copy(dst, 0, src, 0, embd_bytes) != 0; ds4_gpu_tensor_free(dst); ds4_gpu_tensor_free(src); if (ok) metal_graph_dspark_capture_note_slot(g, slot); } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); if (!ok) { metal_graph_dspark_capture_invalidate(g); return false; } g->dspark_capture_batch_tokens = prefix_len + 1u; g->dspark_capture_batch_valid = g->dspark_capture_batch_mask == metal_graph_dspark_capture_complete_mask(g); return g->dspark_capture_valid && g->dspark_capture_batch_valid; } /* 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; } #if defined(__APPLE__) /* Metal command encoding is serialized by the backend. Preserve a caller * hint across this nested scope, but keep it inactive unless this exact * token satisfies the decode-only guards. */ const int previous_pipeline_fast_lookup = ds4_gpu_set_decode_pipeline_fast_lookup(0); const bool pipeline_fast_lookup_active = metal_graph_decode_pipeline_fast_lookup_eligible( g, weights, pos, allow_split_flush); if (pipeline_fast_lookup_active) { (void)ds4_gpu_set_decode_pipeline_fast_lookup(1); } #endif 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 second_split_after_layers = metal_graph_token_adaptive_second_split_after_layers( g, weights, pos, metal_graph_token_second_split_after_layers(), allow_split_flush); const uint32_t split_after_layers = metal_graph_token_adaptive_split_after_layers( g, weights, pos, second_split_after_layers, allow_split_flush); #if defined(__APPLE__) const bool tp_split_flush_safe = g->tp_world == 2 && ds4_gpu_tp_decode_split_flush_safe() != 0 && getenv("DS4_METAL_DISABLE_TP_DECODE_SPLIT_FLUSH") == NULL; #else const bool tp_split_flush_safe = false; #endif for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { #if defined(__APPLE__) if (g->tp_world == 2) { /* Gate-time prefetch plans: this layer's FFN-side weights while * the attention gate waits, the next layer's attention-side * weights while the FFN gate waits (rank slices where split). */ const ds4_layer_weights *lw = &weights->layer[il]; const uint64_t q8_row = (uint64_t)DS4_N_EMBD / 32u * 34u; const bool moe_layer = lw->ffn_gate_shexp && lw->ffn_up_shexp && lw->ffn_gate_inp; const uint64_t shared_half = moe_layer ? lw->ffn_gate_shexp->dim[1] / 2u : 0u; /* Priority order; the backend cuts at this rank's measured window. */ /* With the GPU-decided shared split (default) this rank may end up * with the other half of the shared gate/up as well; list it last * so a rank with a long measured wait streams it in, and a rank * with a short wait never reaches it. */ static int shared_split_static = -1; if (shared_split_static < 0) { shared_split_static = getenv("DS4_TP_STATIC_SHARED_SPLIT") != NULL; } const uint64_t other_rank = (uint64_t)(1 - g->tp_rank); const uint64_t attn_off[7] = { moe_layer ? lw->ffn_gate_shexp->abs_offset + (uint64_t)g->tp_rank * shared_half * q8_row : 0u, moe_layer ? lw->ffn_up_shexp->abs_offset + (uint64_t)g->tp_rank * shared_half * q8_row : 0u, moe_layer ? lw->ffn_gate_inp->abs_offset : 0u, lw->hc_ffn_fn ? lw->hc_ffn_fn->abs_offset : 0u, moe_layer ? lw->ffn_down_shexp->abs_offset : 0u, moe_layer ? lw->ffn_gate_shexp->abs_offset + other_rank * shared_half * q8_row : 0u, moe_layer ? lw->ffn_up_shexp->abs_offset + other_rank * shared_half * q8_row : 0u, }; const uint64_t attn_bytes[7] = { shared_half * q8_row, shared_half * q8_row, (uint64_t)DS4_N_EMBD * DS4_N_EXPERT * sizeof(uint16_t), lw->hc_ffn_fn ? (uint64_t)lw->hc_ffn_fn->dim[0] * lw->hc_ffn_fn->dim[1] * sizeof(uint16_t) : 0u, moe_layer ? (uint64_t)lw->ffn_down_shexp->dim[1] * (lw->ffn_down_shexp->dim[0] / 32u * 34u) : 0u, shared_half * q8_row, shared_half * q8_row, }; (void)ds4_gpu_tp_gate_prefetch_plan(DS4_TP_GATE_ATTN, model->map, model->size, attn_off, attn_bytes, moe_layer ? (shared_split_static ? 5u : 7u) : 0u); if (il + 1u < DS4_N_LAYER && weights->layer[il + 1u].attn_q_a && weights->layer[il + 1u].attn_kv && weights->layer[il + 1u].attn_compressor_kv && weights->layer[il + 1u].attn_compressor_gate) { const ds4_layer_weights *nw = &weights->layer[il + 1u]; const uint64_t qb_row = (uint64_t)nw->attn_q_b->dim[0] / 32u * 34u; const uint64_t qb_half_rows = (uint64_t)nw->attn_q_b->dim[1] / 2u; const uint64_t ffn_off[6] = { nw->attn_q_a->abs_offset, nw->attn_kv->abs_offset, nw->attn_compressor_kv->abs_offset, nw->attn_compressor_gate->abs_offset, nw->hc_attn_fn ? nw->hc_attn_fn->abs_offset : 0u, nw->attn_q_b->abs_offset + (uint64_t)g->tp_rank * qb_half_rows * qb_row, }; const uint64_t ffn_bytes[6] = { (uint64_t)nw->attn_q_a->dim[1] * q8_row, (uint64_t)nw->attn_kv->dim[1] * q8_row, (uint64_t)nw->attn_compressor_kv->dim[0] * nw->attn_compressor_kv->dim[1] * sizeof(uint16_t), (uint64_t)nw->attn_compressor_gate->dim[0] * nw->attn_compressor_gate->dim[1] * sizeof(uint16_t), nw->hc_attn_fn ? (uint64_t)nw->hc_attn_fn->dim[0] * nw->hc_attn_fn->dim[1] * sizeof(uint16_t) : 0u, qb_half_rows * qb_row, }; (void)ds4_gpu_tp_gate_prefetch_plan(DS4_TP_GATE_FFN, model->map, model->size, ffn_off, ffn_bytes, 6); } } #endif 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 && g->hc_expand_pending && metal_graph_dspark_target_slot(g, il) >= 0) { ok = metal_graph_flush_hc_expand(g); } if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); /* Shared-event TP arrival is monotonic: a later command buffer could * satisfy an earlier gate before that gate's payload is ready. The * default single-session flag path instead publishes an exact value * in a distinct layer/gate slot and is safe to submit in-order. */ if (ok && allow_split_flush && (g->tp_world != 2 || tp_split_flush_safe) && ((split_after_layers != 0 && il + 1u == split_after_layers) || (second_split_after_layers != 0 && il + 1u == second_split_after_layers))) { ok = ds4_gpu_flush_commands() != 0; } } if (ok && need_logits) { if (ok) ok = metal_graph_flush_hc_expand(g); ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); } #if defined(__APPLE__) (void)ds4_gpu_set_decode_pipeline_fast_lookup( previous_pipeline_fast_lookup); #endif return ok; } static ds4_gpu_tensor *metal_graph_tensor_row_view( ds4_gpu_tensor *base, uint32_t row, uint64_t row_values) { return ds4_gpu_tensor_view(base, (uint64_t)row * row_values * sizeof(float), row_values * sizeof(float)); } /* Upload prompt token ids for kernels that need token-aware hash routing. */ static bool metal_graph_upload_prompt_tokens( ds4_gpu_tensor *out_tokens, const token_vec *prompt, uint32_t pos0, uint32_t n_tokens) { if (!out_tokens || pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) { return false; } int32_t *tokens = xmalloc((size_t)n_tokens * sizeof(tokens[0])); for (uint32_t i = 0; i < n_tokens; i++) tokens[i] = prompt->v[pos0 + i]; const bool ok = ds4_gpu_tensor_write(out_tokens, 0, tokens, (uint64_t)n_tokens * sizeof(tokens[0])) != 0; free(tokens); return ok; } /* Rebuild ratio-4 compressor state after chunked prefill so a following decode * token sees the same rolling compression window. */ static bool metal_graph_refresh_ratio4_compressor_state( ds4_gpu_graph *g, const ds4_model *model, ds4_gpu_tensor *state_kv, ds4_gpu_tensor *state_score, const ds4_tensor *kv_weight, const ds4_tensor *score_weight, const ds4_tensor *ape, uint32_t head_dim, uint32_t width, uint32_t pos0, uint32_t n_tokens) { if (n_tokens < 4) { return true; } if (!g || !model || !state_kv || !state_score || !kv_weight || !score_weight || !ape || head_dim == 0 || width == 0) { return false; } /* * The recurrent ratio-4 state is intentionally rebuilt from the last * four tokens using the small-batch projection kernel. The full-chunk * projection is already available, but it uses the matrix-matrix path; * mixing those two accumulation orders changes a few FP8 rounding * decisions in later chunks. */ ds4_gpu_tensor *tail_hc = ds4_gpu_tensor_view( metal_graph_batch_attn_norm(g), (uint64_t)(n_tokens - 4u) * DS4_N_EMBD * sizeof(float), 4ull * DS4_N_EMBD * sizeof(float)); bool ok = tail_hc != NULL; if (!ok) { fprintf(stderr, "ds4: ratio-4 compressor tail view creation failed\n"); } if (ok) { #if defined(__APPLE__) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_kv(g), model->map, model->size, kv_weight->abs_offset, DS4_N_EMBD, width, tail_hc, 4) != 0; if (ok) { ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_comp_sc(g), model->map, model->size, score_weight->abs_offset, DS4_N_EMBD, width, tail_hc, 4) != 0; } #else ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_batch_comp_kv(g), metal_graph_batch_comp_sc(g), model->map, model->size, kv_weight->abs_offset, score_weight->abs_offset, DS4_N_EMBD, width, tail_hc, 4) != 0; #endif if (!ok) { fprintf(stderr, "ds4: ratio-4 compressor tail projection failed\n"); } } if (ok) { ok = ds4_gpu_compressor_prefill_state_ratio4_tensor(state_kv, state_score, metal_graph_batch_comp_kv(g), metal_graph_batch_comp_sc(g), model->map, model->size, ape->abs_offset, ape->type, head_dim, pos0 + n_tokens - 4u) != 0; if (!ok) { fprintf(stderr, "ds4: ratio-4 compressor state refresh failed\n"); } } ds4_gpu_tensor_free(tail_hc); return ok; } /* CPU fallback for seeding batched HC state from token embeddings. It is still * useful for tiny speculative verifier batches where a separate GPU embedding * command buffer costs more than the small host write. */ static bool metal_graph_upload_prompt_embeddings_hc_cpu( ds4_gpu_tensor *out_hc, const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, uint32_t pos0, uint32_t n_tokens) { if (pos0 > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - pos0) return false; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t total = (uint64_t)n_tokens * hc_dim; float *hc = xmalloc((size_t)total * sizeof(hc[0])); float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); for (uint32_t t = 0; t < n_tokens; t++) { embed_token_f16(model, weights, prompt->v[pos0 + t], plain); float *dst = hc + (uint64_t)t * hc_dim; for (uint32_t h = 0; h < DS4_N_HC; h++) { memcpy(dst + (uint64_t)h * DS4_N_EMBD, plain, (size_t)DS4_N_EMBD * sizeof(plain[0])); } } const bool ok = ds4_gpu_tensor_write(out_hc, 0, hc, total * sizeof(hc[0])) != 0; free(plain); free(hc); return ok; } static const float *metal_graph_visual_embedding_row( const ds4_gpu_graph *g, uint32_t token_pos) { if (!g || !g->prefill_vision_spans) return NULL; for (size_t i = 0; i < g->prefill_vision_span_count; i++) { const ds4_vision_span *span = &g->prefill_vision_spans[i]; const uint64_t begin = span->token_start; const uint64_t end = begin + span->embedding.token_count; if (token_pos >= begin && token_pos < end && span->embedding.data) { return span->embedding.data + ((uint64_t)token_pos - begin) * DS4_N_EMBD; } } return NULL; } /* DeepSeek image prompts use typed synthetic token ids for routing and * attention. Seed those rows from the vision encoder output instead of the * vocabulary table; regular rows retain the ordinary token embedding. */ static bool metal_graph_upload_visual_prompt_embeddings_hc( ds4_gpu_graph *g, 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 (!g || !out_hc || !model || !weights || !prompt || 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; if (total > SIZE_MAX / sizeof(float)) return false; float *hc = xmalloc((size_t)total * sizeof(hc[0])); float *plain = xmalloc((size_t)DS4_N_EMBD * sizeof(plain[0])); bool ok = true; for (uint32_t t = 0; t < n_tokens; t++) { const int token = prompt->v[pos0 + t]; const float *source = NULL; if (token >= 0 && token < (int)DS4_N_VOCAB) { embed_token_f16(model, weights, token, plain); source = plain; } else { source = metal_graph_visual_embedding_row(g, pos0 + t); if (!source) { ok = false; break; } } 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, source, (size_t)DS4_N_EMBD * sizeof(source[0])); } } if (ok) { 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 #if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) if (weight->type == DS4_TENSOR_F16 && ds4_gpu_matmul_f16_rms_fold_tensor( out, model->map, model->size, weight->abs_offset, in_dim, 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, x, n_tokens, DS4_RMS_EPS)) { return true; } #endif 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; } static const int32_t *metal_graph_visual_tokens_for_batch( const ds4_gpu_graph *g, uint32_t pos0, uint32_t n_tokens) { if (!g || !g->deepseek4_vision_weights || !g->prefill_has_visual || !g->prefill_host_tokens || pos0 < g->prefill_host_pos0) return NULL; const uint32_t off = pos0 - g->prefill_host_pos0; if (off > g->prefill_host_count || n_tokens > g->prefill_host_count - off) return NULL; const int32_t *tokens = g->prefill_host_tokens + off; for (uint32_t i = 0; i < n_tokens; i++) { if (tokens[i] >= (int32_t)DS4_N_VOCAB) return tokens; } return NULL; } 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; const int32_t *visual_tokens = metal_graph_visual_tokens_for_batch(g, pos0, n_tokens); const bool visual_attention = visual_tokens != NULL; /* 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) && !visual_attention && 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"); const bool q_norm_debug = metal_graph_debug_wants("Qnorm", il, pos0); bool q_norm_rope_fused = false; if (ok && !q_norm_debug) { q_norm_rope_fused = ds4_gpu_head_rms_norm_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, DS4_RMS_EPS) != 0; } if (ok && !q_norm_rope_fused) { 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 && !q_norm_rope_fused) { 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 && !q_norm_rope_fused) { 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 (visual_attention) { ok = ds4_gpu_attention_visual_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], NULL, 0, NULL, 0, visual_tokens, DS4_N_VOCAB, n_tokens, pos0, n_tokens, g->raw_cap, 0, 0, g->raw_window, 0, DS4_N_HEAD, DS4_N_HEAD_DIM) != 0; } else 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 && visual_attention) { ok = ds4_gpu_attention_visual_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], NULL, 0, NULL, 0, visual_tokens, DS4_N_VOCAB, n_tokens, pos0, n_raw, g->raw_cap, raw_start, 0, g->raw_window, 0, DS4_N_HEAD, DS4_N_HEAD_DIM) != 0; } else 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 && !g->spec_capture_prefixes && (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, false, 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 + 1u < n_tokens && t < DS4_SPEC_PREFIX_SLOTS) { ok = metal_graph_capture_prefix_attn_state(g, il, t); } 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 && !g->spec_capture_prefixes && (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, false, 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 + 1u < n_tokens && t < DS4_SPEC_PREFIX_SLOTS) { ok = metal_graph_capture_prefix_index_state(g, il, t); } ds4_gpu_tensor_free(sc_view); ds4_gpu_tensor_free(kv_view); } } } } if (ratio == 4) DS4_METAL_PROFILE_ATTN_STAGE("indexer_setup"); if (ok && visual_attention) { const uint32_t n_raw = zero_prefix ? n_tokens : metal_graph_raw_span_for_batch(g, pos0, n_tokens); const uint32_t raw_start = zero_prefix ? 0u : metal_graph_raw_start_for_span( g, pos0 + n_tokens - 1u, n_raw); bool use_indexed_comp = false; if (!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 && 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 (zero_prefix) { 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; } else { 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) { 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) { ok = ds4_gpu_dsv4_topk_mask_tensor( metal_graph_comp_mask(g), metal_graph_comp_selected(g), n_comp, n_tokens, DS4_N_INDEXER_TOP_K) != 0; } use_indexed_comp = ok; } if (ok) { ok = ds4_gpu_attention_visual_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], n_comp ? g->layer_attn_comp_cache[il] : NULL, metal_graph_attn_comp_cache_is_f16(), use_indexed_comp ? metal_graph_comp_mask(g) : NULL, use_indexed_comp ? 1u : 0u, visual_tokens, DS4_N_VOCAB, 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; } if (ok && !batch_attention_done && !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 && !batch_attention_done && 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 && !batch_attention_done && 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; } if (!attn_out_f16) { 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) { 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 && g->deepseek4_vision_weights && g->prefill_has_visual) { ok = ds4_gpu_router_select_batch_visual_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, layer->ffn_exp_probs_b != NULL, layer->ffn_gate_tid2eid != NULL, g->vision_model_map, g->vision_model_size, g->deepseek4_vision_weights->visual_router_bias[il], metal_graph_batch_router_logits(g), router_tokens, DS4_N_VOCAB, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, n_tokens) != 0; } else 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), router_tokens, 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 tp_batched_moe = tp_split_batch_moe && !g->quality && !g->ssd_streaming && n_tokens > 1 && n_tokens <= 6 && ds4_gpu_device_is_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_TP_BATCH_MOE") == NULL; 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 && !tp_batched_moe) { /* Reference schedule: evaluate each row's owned experts separately. */ 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); } } else if (ok) { /* Batched kernels apply the same expert ownership to every row and * write directly into the registered exchange slab under TP. */ ok = ds4_gpu_routed_moe_batch_tensor(tp_batched_moe ? g->tp_batch_out[il] : 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 && tp_split_batch_moe && !cuda_tp_owned_batch_moe) { 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; } } 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) { metal_graph_debug_dump_tensor("tp_ffn_recv", metal_graph_batch_ffn_out(g), (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor("tp_ffn_sum", metal_graph_batch_routed_out(g), (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (!ok) { fprintf(stderr, "ds4: TP prefill FFN all-reduce failed (layer %u)\n", il); } } if (ok && keep_ffn_out) { ok = metal_graph_ensure_batch_ffn_out(g) && ds4_gpu_add_tensor(metal_graph_batch_ffn_out(g), metal_graph_batch_shared_out(g), metal_graph_batch_routed_out(g), (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; } if (ok && keep_ffn_out) { metal_graph_debug_dump_tensor("ffn_out", metal_graph_batch_ffn_out(g), (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = metal_graph_apply_directional_steering_ffn(g, metal_graph_batch_ffn_out(g), il, n_tokens); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, metal_graph_batch_ffn_out(g), metal_graph_batch_after_attn_hc(g), hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && shared_down_f16) { ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, metal_graph_batch_routed_out(g), g->batch_q_half, metal_graph_batch_after_attn_hc(g), hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && tp_row_split_ffn) { /* Shared expert already folded into the exchanged routed rows. */ ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, metal_graph_batch_routed_out(g), metal_graph_batch_after_attn_hc(g), hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { ok = ds4_gpu_hc_expand_add_split_tensor(next_hc_view, metal_graph_batch_routed_out(g), metal_graph_batch_shared_out(g), metal_graph_batch_after_attn_hc(g), hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_FFN_STAGE("hc_post"); if (ok) { metal_graph_debug_dump_tensor("hc_ffn_post", metal_graph_batch_next_hc(g), (uint64_t)n_tokens * hc_dim, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("hc_post"); ds4_gpu_tensor_free(tp_ffn_x); ds4_gpu_tensor_free(next_hc_view); ds4_gpu_tensor_free(ffn_cur_view); ds4_gpu_tensor_free(hc_split_view); ds4_gpu_tensor_free(hc_mix_view); #undef DS4_METAL_PROFILE_FFN_STAGE return ok; } /* Encode one complete layer for prefill by chaining attention and FFN batches. */ static bool metal_graph_encode_layer_batch( ds4_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *layer, uint32_t il, uint32_t pos0, uint32_t n_tokens) { if (g->placement) { const int this_tier = g->placement[il + 1u]; if (!metal_graph_set_active_tier_batch(g, this_tier, n_tokens)) { return false; } } bool ok = metal_graph_layer_stage_profile_start(il); if (ok) { ok = metal_graph_encode_layer_attention_batch(g, model, layer, il, pos0, n_tokens); } if (!ok) { fprintf(stderr, "ds4: gpu layer %u attention batch encode failed\n", il); } if (ok) { ok = metal_graph_encode_layer_ffn_batch(g, model, layer, il, pos0, n_tokens, NULL, 0); if (!ok) { fprintf(stderr, "ds4: gpu layer %u ffn batch encode failed\n", il); } } if (ok) { ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); g->batch_next_hc_by_tier[g->active_tier] = tmp; } return ok; } static bool metal_graph_eval_token_raw_swa_streaming( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, int token, uint32_t pos, float *logits) { if (g->raw_cap == 0) { fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); return false; } const bool profile = glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", "DS4_METAL_GRAPH_TOKEN_PROFILE"); const bool throttle = graph_power_throttle_enabled(g); const double t0 = (profile || throttle) ? now_sec() : 0.0; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); metal_graph_dspark_capture_begin(g); const bool static_decode_map = metal_graph_stream_decode_static_map_enabled() && weights_model_map_decode_static_supported(weights); const bool static_map_state_cache = static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); const bool batch_static_decode = static_decode_map && metal_graph_stream_decode_layer_batch_enabled(g); bool ok = true; if (static_decode_map) { if (!static_map_state_cache || !g->streaming_static_decode_map_current) { ok = metal_graph_stream_map_decode_static_all(model, weights); if (ok) g->streaming_static_decode_map_current = static_map_state_cache; } } else { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_token(model, weights); } if (ok && !static_decode_map && DS4_N_LAYER > 0) { metal_graph_stream_readahead_layer_decode(model, weights, 0); } if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), model->map, model->size, weights->token_embd->abs_offset, (uint32_t)weights->token_embd->dim[1], (uint32_t)token, DS4_N_EMBD, DS4_N_HC) != 0; } if (batch_static_decode) { for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { ok = metal_graph_encode_decode_layer(g, model, &weights->layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, token); if (ok) { ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); g->after_ffn_hc_by_tier[g->active_tier] = tmp; ok = metal_graph_dspark_capture_decode_layer(g, il); } } if (ok && logits) { ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); } const double t_encoded = (profile || throttle) ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { fprintf(stderr, "ds4: metal SSD streaming batched token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", pos, (t_encoded - t0) * 1000.0, (t_done - t_encoded) * 1000.0, (t_read - t_done) * 1000.0, (t_read - t0) * 1000.0, logits != NULL); } if (ok && throttle) { graph_power_note_decode_token(g, t_read - t0); } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after batched SSD streaming graph eval failure also failed\n"); } } return ok; } if (ok) ok = ds4_gpu_end_commands() != 0; double encode_s = 0.0; double execute_s = 0.0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const double tl0 = profile ? now_sec() : 0.0; if (!static_decode_map && !metal_graph_stream_map_layer_decode(model, weights, il)) { ok = false; break; } if (!static_decode_map && il + 1 < DS4_N_LAYER) { metal_graph_stream_readahead_layer_decode(model, weights, il + 1); } else if (!static_decode_map && logits) { metal_graph_stream_readahead_output(model, weights); } if (ok) ok = ds4_gpu_begin_commands() != 0; bool encoded_layer = false; if (ok) { ok = metal_graph_encode_decode_layer(g, model, &weights->layer[il], il, pos, g->layer_raw_cache[il], g->raw_cap, raw_row, n_raw, token); encoded_layer = true; } if (encoded_layer) { ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); g->after_ffn_hc_by_tier[g->active_tier] = tmp; if (ok) ok = metal_graph_dspark_capture_decode_layer(g, il); } const double tl_encoded = profile ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double tl_done = profile ? now_sec() : 0.0; if (profile) { encode_s += tl_encoded - tl0; execute_s += tl_done - tl_encoded; } } if (ok && logits && !static_decode_map) ok = metal_graph_stream_map_output(model, weights); const double t_head0 = profile ? now_sec() : 0.0; if (ok && logits) ok = ds4_gpu_begin_commands() != 0; if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); const double t_head_encoded = profile ? now_sec() : 0.0; if (ok && logits) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { if (logits) { encode_s += t_head_encoded - t_head0; execute_s += t_done - t_head_encoded; } fprintf(stderr, "ds4: metal SSD streaming token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", pos, encode_s * 1000.0, execute_s * 1000.0, (t_read - t_done) * 1000.0, (t_read - t0) * 1000.0, logits != NULL); } if (ok) graph_power_note_decode_token(g, t_read - t0); if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after SSD streaming graph eval failure also failed\n"); } } return ok; } /* Execute one Metal decode token and read back logits. */ static bool metal_graph_eval_token_raw_swa( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, int token, uint32_t pos, float *logits) { if (g && g->ssd_streaming) { return metal_graph_eval_token_raw_swa_streaming(g, model, weights, token, pos, logits); } const bool profile = glm_graph_env_present("DS4_ROCM_GRAPH_TOKEN_PROFILE", "DS4_METAL_GRAPH_TOKEN_PROFILE"); const bool throttle = graph_power_throttle_enabled(g); const double t0 = (profile || throttle) ? now_sec() : 0.0; bool ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, logits != NULL, true); const double t_encoded = (profile || throttle) ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits && g->tp_world == 2 && g->tp_logits_half) { const uint64_t tp_vhalf = (uint64_t)DS4_N_VOCAB / 2u; const uint64_t off = (uint64_t)g->tp_rank * tp_vhalf * sizeof(float); ok = ds4_gpu_tensor_read(metal_graph_logits(g), off, logits + g->tp_rank * tp_vhalf, tp_vhalf * sizeof(float)) != 0; } else if (ok && logits && !(g->tp_world == 2 && g->tp_rank == 1)) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { fprintf(stderr, "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d\n", pos, (t_encoded - t0) * 1000.0, (t_done - t_encoded) * 1000.0, (t_read - t_done) * 1000.0, (t_read - t0) * 1000.0, logits != NULL); } if (ok) graph_power_note_decode_token(g, t_read - t0); if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after graph eval failure also failed\n"); } } return ok; } static bool metal_graph_streaming_decode_prefill_wide_default( const ds4_weights *weights) { if (DS4_MODEL_VARIANT != DS4_VARIANT_FLASH || !weights || DS4_N_LAYER == 0) { return false; } const uint32_t type = weights->layer[0].ffn_gate_exps->type; return (type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_MXFP4) && weights->layer[0].ffn_up_exps->type == type && weights->layer[0].ffn_down_exps->type == type; } static uint32_t metal_graph_streaming_decode_prefill_max_tokens( const ds4_gpu_graph *g, const ds4_weights *weights) { (void)g; if (glm_graph_env_present("DS4_ROCM_DISABLE_STREAMING_DECODE_PREFILL", "DS4_METAL_DISABLE_STREAMING_DECODE_PREFILL")) { return 0; } const char *env = glm_graph_env_value( "DS4_ROCM_STREAMING_DECODE_PREFILL_MAX", "DS4_METAL_STREAMING_DECODE_PREFILL_MAX"); if (env && env[0]) { char *end = NULL; const long v = strtol(env, &end, 10); if (end != env) { if (v <= 0) return 0; if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; return (uint32_t)v; } } if (DS4_MODEL_VARIANT != DS4_VARIANT_PRO && DS4_MODEL_VARIANT != DS4_VARIANT_FLASH) { return 0u; } return metal_graph_streaming_decode_prefill_wide_default(weights) ? 64u : 18u; } static bool metal_graph_use_streaming_decode_prefill( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t n_tokens) { const uint32_t max_tokens = metal_graph_streaming_decode_prefill_max_tokens(g, weights); return g && g->ssd_streaming && !g->quality && n_tokens != 0 && max_tokens != 0 && n_tokens <= max_tokens; } static bool metal_graph_use_streaming_decode_prefill_range( const ds4_gpu_graph *g, const ds4_weights *weights, uint32_t start, uint32_t n_tokens) { /* * Short streamed prefill is latency-sensitive. Use the decode-style path * by default for SSD streaming, while keeping a cold-only escape hatch for * strict-vector tests that need canonical layer-major prefill semantics. */ if (start == 0) { if (glm_graph_env_present( "DS4_ROCM_DISABLE_STREAMING_COLD_DECODE_PREFILL", "DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL")) { return false; } } return metal_graph_use_streaming_decode_prefill(g, weights, n_tokens); } static bool metal_graph_prefill_decode_streaming_range( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, uint32_t start, uint32_t n_tokens, float *logits, bool show_progress, ds4_session_progress_fn progress, void *progress_ud, ds4_session_progress_fn display_progress, void *display_progress_ud, ds4_session_cancel_fn cancel, void *cancel_ud, bool *cancelled) { if (!metal_graph_use_streaming_decode_prefill(g, weights, n_tokens)) return false; if (!prompt || start > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - start) return false; if (start == 0) { ds4_gpu_stream_expert_cache_reset_route_hotness(); } const bool profile = glm_graph_env_present("DS4_ROCM_GRAPH_PREFILL_PROFILE", "DS4_METAL_GRAPH_PREFILL_PROFILE"); const double t0 = profile ? now_sec() : 0.0; /* * `prefill_chunk` is not just UI progress: ds4_session_sync() wraps it to * advance the live checkpoint, and ds4-server may save that checkpoint. * Decode-style prefill only reads logits for the final token, so report one * cacheable chunk at the end. `prefill_display` remains per-token UI only. */ if (progress) progress(progress_ud, "prefill_chunk", (int)start, prompt->len); if (display_progress) { display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); } for (uint32_t i = 0; i < n_tokens; i++) { if (cancel && cancel(cancel_ud)) { if (cancelled) *cancelled = true; return true; } const uint32_t pos = start + i; const bool last = i + 1u == n_tokens; float *token_logits = (last && logits) ? logits : NULL; if (!metal_graph_eval_token_raw_swa(g, model, weights, prompt->v[pos], pos, token_logits)) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after decode-style streaming prefill failure also failed\n"); } return false; } if (last && progress && logits) { progress(progress_ud, "prefill_chunk", (int)(pos + 1u), prompt->len); } if (display_progress) { display_progress(display_progress_ud, "prefill_display", (int)(pos + 1u), prompt->len); } if (cancel && cancel(cancel_ud)) { if (cancelled) *cancelled = true; return true; } if (show_progress) { fprintf(stderr, "ds4: gpu streaming prefill token %u/%u\r", i + 1u, n_tokens); fflush(stderr); } } if (show_progress) fputc('\n', stderr); if (profile) { const double t1 = now_sec(); fprintf(stderr, "ds4: gpu decode-style streaming prefill start=%u tokens=%u total=%.3f ms\n", start, n_tokens, (t1 - t0) * 1000.0); } return true; } static bool metal_graph_capture_prefill_seed_router_selected( ds4_gpu_graph *g, uint32_t il, uint32_t n_tokens) { uint32_t k = metal_graph_streaming_prefill_cache_seed_k(g); if (k == 0) return true; if (k > n_tokens) k = n_tokens; g->prefill_seed_tokens = k; if (!g->prefill_seed_router_selected || !metal_graph_batch_router_selected(g) || il >= DS4_N_LAYER || n_tokens == 0 || sizeof(int) != sizeof(int32_t)) { return false; } const uint64_t bytes = (uint64_t)k * DS4_N_EXPERT_USED * sizeof(int32_t); const uint64_t src_off = (uint64_t)(n_tokens - k) * DS4_N_EXPERT_USED * sizeof(int); const uint64_t dst_off = (uint64_t)il * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_N_EXPERT_USED * sizeof(int32_t); return ds4_gpu_tensor_copy(g->prefill_seed_router_selected, dst_off, metal_graph_batch_router_selected(g), src_off, bytes) != 0; } static bool metal_graph_seed_streaming_expert_cache_from_prefill( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights) { const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; if (!metal_graph_streaming_prefill_cache_seed_enabled(g)) return true; if (!model || !weights || !g->prefill_seed_router_selected || seed_tokens == 0) { return false; } int32_t selected[DS4_MAX_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_MAX_EXPERT_USED]; const uint64_t bytes = (uint64_t)DS4_N_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_N_EXPERT_USED * sizeof(selected[0]); if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, 0, selected, bytes) == 0) { return false; } const bool profile = glm_graph_env_present("DS4_ROCM_STREAMING_PREFILL_CACHE_SEED_PROFILE", "DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE"); const double t0 = profile ? now_sec() : 0.0; uint32_t seeded_layers = 0; uint32_t seeded_rows = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { continue; } const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { fprintf(stderr, "ds4: Metal prefill expert-cache seed byte size overflow at layer %u\n", il); return false; } const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, il, gate_expert_bytes, down_expert_bytes); for (uint32_t row = 0; row < seed_tokens; row++) { const size_t sel_off = ((size_t)il * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS + row) * DS4_N_EXPERT_USED; if (ds4_gpu_stream_expert_cache_seed_selected( &table, selected + sel_off, DS4_N_EXPERT_USED) == 0) { return false; } seeded_rows++; } seeded_layers++; } if (profile) { fprintf(stderr, "ds4: Metal streaming prefill expert-cache seed k=%u layers=%u rows=%u time=%.3f ms\n", seed_tokens, seeded_layers, seeded_rows, (now_sec() - t0) * 1000.0); } return true; } static bool metal_graph_seed_streaming_expert_cache_from_hotlist( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights) { if (!metal_graph_streaming_expert_hotlist_enabled(g)) return true; if (!model || !weights) return false; uint32_t cache_budget = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *layer = &weights->layer[il]; if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { continue; } const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { fprintf(stderr, "ds4: streaming expert hotlist budget byte size overflow at layer %u\n", il); return false; } const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; cache_budget = ds4_gpu_stream_expert_cache_budget_for_expert_size( gate_expert_bytes, down_expert_bytes); break; } if (cache_budget == 0) return true; const uint32_t preload_count = metal_graph_streaming_expert_preload_count(g, cache_budget); if (preload_count == 0) return true; const uint32_t current_count = ds4_gpu_stream_expert_cache_current_count(); const char *path = glm_graph_env_value("DS4_ROCM_STREAMING_EXPERT_HOTLIST", "DS4_METAL_STREAMING_EXPERT_HOTLIST"); const bool from_file = path && path[0]; const bool refresh_builtin_glm = !from_file && g_ds4_shape.variant == DS4_VARIANT_GLM52; const bool profile = glm_graph_env_present("DS4_ROCM_STREAMING_EXPERT_HOTLIST_PROFILE", "DS4_METAL_STREAMING_EXPERT_HOTLIST_PROFILE"); if (!from_file && !refresh_builtin_glm && current_count >= preload_count) { if (profile) { fprintf(stderr, "ds4: streaming expert hotlist seed skipped preload=%u current=%u\n", preload_count, current_count); } return true; } int32_t experts[DS4_MAX_LAYER][DS4_MAX_EXPERT]; uint32_t priorities[DS4_MAX_LAYER][DS4_MAX_EXPERT]; uint32_t counts[DS4_MAX_LAYER]; bool seen[DS4_MAX_LAYER][DS4_MAX_EXPERT]; memset(experts, 0, sizeof(experts)); memset(priorities, 0, sizeof(priorities)); memset(counts, 0, sizeof(counts)); memset(seen, 0, sizeof(seen)); uint32_t loaded = 0; if (from_file) { if (!metal_graph_streaming_expert_hotlist_load_file(path, preload_count, experts, priorities, counts, seen, &loaded)) { return false; } } else if (!metal_graph_streaming_expert_hotlist_load_default(preload_count, experts, priorities, counts, seen, &loaded)) { return false; } if (loaded == 0) return true; const double t0 = profile ? now_sec() : 0.0; uint32_t seeded_layers = 0; uint32_t seeded_experts = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t n = counts[il]; if (n == 0) continue; const ds4_layer_weights *layer = &weights->layer[il]; if (!metal_graph_streaming_expert_cache_seed_layer_expected(g, layer)) { continue; } const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); if (layer->ffn_gate_exps->dim[1] > UINT64_MAX / gate_row_bytes || layer->ffn_down_exps->dim[1] > UINT64_MAX / down_row_bytes) { fprintf(stderr, "ds4: streaming expert hotlist seed byte size overflow at layer %u\n", il); return false; } const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, il, gate_expert_bytes, down_expert_bytes); if (ds4_gpu_stream_expert_cache_seed_experts( &table, experts[il], priorities[il], n) == 0) { return false; } seeded_layers++; seeded_experts += n; } if (profile) { const char *source_name = NULL; if (from_file) { source_name = path; } else if (g_ds4_shape.variant == DS4_VARIANT_GLM52) { source_name = "built-in-glm52"; } else if (g_ds4_shape.variant == DS4_VARIANT_FLASH) { source_name = "built-in-flash"; } else if (g_ds4_shape.variant == DS4_VARIANT_PRO) { source_name = "built-in-pro"; } else { source_name = "built-in"; } fprintf(stderr, "ds4: streaming expert hotlist seed source=%s preload=%u loaded=%u layers=%u experts=%u time=%.3f ms\n", source_name, preload_count, loaded, seeded_layers, seeded_experts, (now_sec() - t0) * 1000.0); } return true; } typedef struct { int id0; int id1; float value0; float value1; bool valid; bool fast_attention; } metal_graph_top2_result; /* Greedy verifier helper. Speculative decoding only needs the target model's * top token after most accepted draft rows; the full vocabulary row is needed * once, for the final committed state that normal sampling will continue from. * Keeping intermediate rows device-resident avoids turning verification into a * sequence of large CPU readbacks. */ static bool metal_graph_eval_token_raw_swa_top( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, int token, uint32_t pos, int *top_id, float *logits, bool allow_split_top1, metal_graph_top2_result *top2, bool force_fast_attention) { if (!top_id) return false; if (top2) memset(top2, 0, sizeof(*top2)); const bool fast_attention = allow_split_top1 && logits == NULL && (force_fast_attention || metal_graph_cuda_greedy_splitkv_requested()); if (top2) top2->fast_attention = fast_attention; const int old_fast_attention = ds4_gpu_set_decode_fast_attention(fast_attention ? 1 : 0); const bool profile = getenv("DS4_METAL_GRAPH_TOKEN_PROFILE") != NULL; const double t0 = profile ? now_sec() : 0.0; const bool split_top1 = allow_split_top1 && logits == NULL && top2 == NULL && g->cuda_tp_output && metal_graph_cuda_greedy_split_top1_requested(); if (split_top1) { int output_tiers[DS4_MAX_GPUS] = {0}; uint32_t output_ways = 0; bool ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, false, true); if (ok) ok = metal_graph_encode_output_head_split_top1(g, model, weights, weights->output->dim[1], output_tiers, &output_ways); const double t_encoded = profile ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = profile ? now_sec() : 0.0; if (ok) { bool have_best = false; uint32_t best_id = 0; float best_value = 0.0f; uint32_t cand_ids[DS4_MAX_GPUS] = {0}; float cand_values[DS4_MAX_GPUS] = {0.0f}; ok = output_ways <= DS4_MAX_GPUS && ds4_gpu_tensor_read(g->comp_selected_by_tier[g->head_tier], 0, cand_ids, (uint64_t)output_ways * sizeof(cand_ids[0])) != 0 && ds4_gpu_tensor_read(g->comp_mask_by_tier[g->head_tier], 0, cand_values, (uint64_t)output_ways * sizeof(cand_values[0])) != 0; for (uint32_t i = 0; ok && i < output_ways; i++) { const uint32_t cand_id = cand_ids[i]; const float cand_value = cand_values[i]; if (ok && (!have_best || cand_value > best_value || (cand_value == best_value && cand_id < best_id))) { have_best = true; best_id = cand_id; best_value = cand_value; } } ok = ok && have_best && best_id <= (uint32_t)INT32_MAX; if (ok) *top_id = (int)best_id; } const double t_read = profile ? now_sec() : 0.0; if (profile) { fprintf(stderr, "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=0 top=1 split_top1=1\n", pos, (t_encoded - t0) * 1000.0, (t_done - t_encoded) * 1000.0, (t_read - t_done) * 1000.0, (t_read - t0) * 1000.0); } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after split-top graph eval failure also failed\n"); } } (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); return ok; } bool ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, true, true); if (ok) { ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), metal_graph_logits(g), DS4_N_VOCAB) != 0; } const double t_encoded = profile ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = profile ? now_sec() : 0.0; if (ok && top2) { uint32_t ids[2] = {0, 0}; float values[2] = {0.0f, 0.0f}; ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, ids, sizeof(ids)) != 0 && ds4_gpu_tensor_read(metal_graph_comp_mask(g), 0, values, sizeof(values)) != 0; if (ok && ids[0] <= (uint32_t)INT32_MAX && ids[1] <= (uint32_t)INT32_MAX) { top2->id0 = (int)ids[0]; top2->id1 = (int)ids[1]; top2->value0 = values[0]; top2->value1 = values[1]; top2->valid = isfinite(values[0]) && isfinite(values[1]); *top_id = top2->id0; } else { ok = false; } } else if (ok) { ok = ds4_gpu_tensor_read(metal_graph_comp_selected(g), 0, top_id, sizeof(*top_id)) != 0; } if (ok && logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = profile ? now_sec() : 0.0; if (profile) { fprintf(stderr, "ds4: metal graph token pos=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms logits=%d top=1 split_top1=0\n", pos, (t_encoded - t0) * 1000.0, (t_done - t_encoded) * 1000.0, (t_read - t_done) * 1000.0, (t_read - t0) * 1000.0, logits != NULL); } if (!ok) { if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: Metal synchronize after top-only graph eval failure also failed\n"); } } (void)ds4_gpu_set_decode_fast_attention(old_fast_attention); return ok; } static bool dspark_stage0_weights_ready( const ds4_gpu_graph *g, const ds4_dspark_weights *dw) { if (!g || !dw || dw->n_stages == 0 || dw->target_layer_count == 0 || dw->target_layer_count != g->dspark_target_layer_count || !g->dspark_target_hidden || !g->dspark_stage0_proj || !g->dspark_main_x) { return false; } const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; const ds4_tensor *main_proj = stage0->main_proj; const ds4_tensor *main_norm = stage0->main_norm; const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; return main_proj && main_norm && dspark_tensor_type_matches(main_proj->type, DS4_DSPARK_LAYOUT_DENSE) && main_norm->type == DS4_TENSOR_F32 && main_proj->ndim == 2 && main_proj->dim[0] == in_dim && main_proj->dim[1] == DS4_N_EMBD && main_norm->ndim == 1 && main_norm->dim[0] == DS4_N_EMBD; } static bool metal_graph_eval_dspark_stage0( ds4_gpu_graph *g, const ds4_model *dspark_model, const ds4_dspark_weights *dw) { if (!g || !dspark_model || !dw || !dspark_stage0_weights_ready(g, dw)) { return false; } const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; bool ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = metal_graph_matmul_plain_tensor(g->dspark_stage0_proj, dspark_model, stage0->main_proj, in_dim, DS4_N_EMBD, g->dspark_target_hidden, 1); } if (ok) { ok = ds4_gpu_rms_norm_weight_tensor(g->dspark_main_x, g->dspark_stage0_proj, dspark_model->map, dspark_model->size, stage0->main_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; } if (ok) ok = ds4_gpu_end_commands() != 0; if (!ok) (void)ds4_gpu_synchronize(); return ok; } static bool dspark_stage0_batch_ready( const ds4_gpu_graph *g, const ds4_dspark_weights *dw, uint32_t n_tokens) { if (!dspark_stage0_weights_ready(g, dw) || n_tokens == 0 || n_tokens > g->prefill_cap || !g->dspark_target_hidden_batch || !metal_graph_batch_ffn_cur(g) || !metal_graph_batch_ffn_norm(g) || !metal_graph_batch_cur_hc(g)) { return false; } const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; return ds4_gpu_tensor_bytes(metal_graph_batch_ffn_cur(g)) >= (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) >= (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float) && ds4_gpu_tensor_bytes(metal_graph_batch_cur_hc(g)) >= (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float) && in_dim <= SIZE_MAX / sizeof(float); } static bool metal_graph_pack_dspark_target_hidden_batch( ds4_gpu_graph *g, const ds4_dspark_weights *dw, ds4_gpu_tensor *packed, uint32_t n_tokens) { if (!g || !dw || !packed || n_tokens == 0 || n_tokens > g->prefill_cap || dw->target_layer_count != g->dspark_target_layer_count) { return false; } const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; const uint64_t packed_count = (uint64_t)n_tokens * in_dim; if (packed_count == 0 || packed_count > (uint64_t)SIZE_MAX / sizeof(float)) { return false; } return ds4_gpu_pack_slot_rows_f32_tensor(packed, g->dspark_target_hidden_batch, n_tokens, DS4_N_EMBD, dw->target_layer_count, g->prefill_cap) != 0; } static bool metal_graph_eval_dspark_stage0_batch( ds4_gpu_graph *g, const ds4_model *dspark_model, const ds4_dspark_weights *dw, uint32_t n_tokens, bool commands_open) { if (!g || !dspark_model || !dw || !dspark_stage0_batch_ready(g, dw, n_tokens)) { return false; } const ds4_dspark_stage_weights *stage0 = &dw->stage[0]; const uint64_t in_dim = (uint64_t)dw->target_layer_count * DS4_N_EMBD; const uint64_t packed_bytes = (uint64_t)n_tokens * in_dim * sizeof(float); bool packed_owned = false; ds4_gpu_tensor *packed = NULL; if (g->dspark_stage0_packed && ds4_gpu_tensor_bytes(g->dspark_stage0_packed) >= packed_bytes) { packed = g->dspark_stage0_packed; } else { packed = ds4_gpu_tensor_alloc(packed_bytes); packed_owned = true; } if (!packed) return false; bool ok = metal_graph_pack_dspark_target_hidden_batch(g, dw, packed, n_tokens); if (ok && !commands_open) ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = metal_graph_matmul_plain_tensor(metal_graph_batch_ffn_cur(g), dspark_model, stage0->main_proj, in_dim, DS4_N_EMBD, packed, n_tokens); } if (ok) { ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_ffn_norm(g), metal_graph_batch_ffn_cur(g), dspark_model->map, dspark_model->size, stage0->main_norm->abs_offset, DS4_N_EMBD, n_tokens, DS4_RMS_EPS) != 0; } if (ok && !commands_open) ok = ds4_gpu_end_commands() != 0; if (!ok && !commands_open) (void)ds4_gpu_synchronize(); if (packed_owned) ds4_gpu_tensor_free(packed); return ok; } static bool dspark_draft_block_ready( const ds4_gpu_graph *g, const ds4_weights *base_weights, const ds4_dspark_weights *dw, int token) { if (!g || !base_weights || !dw || !base_weights->token_embd || !g->dspark_draft_tokens || !g->dspark_draft_hc || dw->block_size == 0 || dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || g->dspark_block_size != dw->block_size || !dw->has_noise_token_id) { return false; } const uint32_t n_vocab = (uint32_t)base_weights->token_embd->dim[1]; return token >= 0 && (uint32_t)token < n_vocab && dw->noise_token_id < n_vocab; } static bool dspark_stage_input_ready( const ds4_gpu_graph *g, const ds4_dspark_weights *dw) { if (!g || !dw || dw->block_size == 0 || dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || g->dspark_block_size != dw->block_size || !g->dspark_main_x || !g->dspark_draft_hc || !g->dspark_target_hc || !g->dspark_stage_input_hc || !g->dspark_position_ids) { return false; } if (dw->block_size == UINT32_MAX) return false; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t rows = (uint64_t)dw->block_size + 1u; return ds4_gpu_tensor_bytes(g->dspark_target_hc) >= hc_dim * sizeof(float) && ds4_gpu_tensor_bytes(g->dspark_stage_input_hc) >= rows * hc_dim * sizeof(float) && ds4_gpu_tensor_bytes(g->dspark_position_ids) >= rows * sizeof(int32_t); } static bool dspark_stage_cache_ready( const ds4_gpu_graph *g, const ds4_dspark_weights *dw) { if (!g || !dw || dw->n_stages == 0 || dw->n_stages > DS4_DSPARK_MAX_STAGES || g->dspark_cache_cap == 0 || !metal_graph_dspark_cache_current_window_valid(g)) { return false; } const uint64_t bytes = (uint64_t)g->dspark_cache_cap * DS4_N_HEAD_DIM * sizeof(float); for (uint32_t stage = 0; stage < dw->n_stages; stage++) { if (!g->dspark_raw_cache[stage] || ds4_gpu_tensor_bytes(g->dspark_raw_cache[stage]) < bytes) { return false; } } return true; } static bool dspark_noncausal_attention_probe_ready( const ds4_gpu_graph *g, const ds4_dspark_weights *dw) { if (!g || !dw || dw->n_stages == 0 || dw->block_size == 0 || dw->block_size > DS4_DSPARK_MAX_BLOCK_SIZE || g->prefill_cap < dw->block_size + 1u || !metal_graph_batch_q(g) || !metal_graph_batch_heads(g) || !g->dspark_raw_cache[0]) { return false; } const ds4_layer_weights *block = &dw->stage[0].block; if (!block->attn_sinks) return false; const uint64_t rows = (uint64_t)dw->block_size + 1u; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; return ds4_gpu_tensor_bytes(metal_graph_batch_q(g)) >= rows * q_dim * sizeof(float) && ds4_gpu_tensor_bytes(metal_graph_batch_heads(g)) >= rows * q_dim * sizeof(float) && ds4_gpu_tensor_bytes(g->dspark_raw_cache[0]) >= rows * DS4_N_HEAD_DIM * sizeof(float); } static bool metal_graph_probe_dspark_noncausal_attention( ds4_gpu_graph *g, const ds4_model *dspark_model, const ds4_dspark_weights *dw) { if (!g || !dspark_model || !dspark_noncausal_attention_probe_ready(g, dw)) { return false; } const uint32_t rows = dw->block_size + 1u; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const ds4_layer_weights *block = &dw->stage[0].block; bool ok = ds4_gpu_tensor_fill_f32(metal_graph_batch_q(g), 0.0f, (uint64_t)rows * q_dim) != 0 && ds4_gpu_tensor_fill_f32(g->dspark_raw_cache[0], 0.0f, (uint64_t)rows * DS4_N_HEAD_DIM) != 0; if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) { ok = ds4_gpu_attention_noncausal_raw_batch_heads_tensor( metal_graph_batch_heads(g), dspark_model->map, dspark_model->size, block->attn_sinks->abs_offset, metal_graph_batch_q(g), g->dspark_raw_cache[0], rows, rows, rows, 0, DS4_N_HEAD, DS4_N_HEAD_DIM) != 0; } if (ok) ok = ds4_gpu_end_commands() != 0; if (!ok) (void)ds4_gpu_synchronize(); return ok; } static bool metal_graph_prepare_dspark_setup_block( ds4_gpu_graph *g, const ds4_model *base_model, const ds4_weights *base_weights, const ds4_dspark_weights *dw, int token, uint32_t pos) { if (!g || !base_model || !dspark_draft_block_ready(g, base_weights, dw, token) || !dspark_stage_input_ready(g, dw)) { return false; } if (pos > (uint32_t)INT32_MAX || dw->block_size > (uint32_t)INT32_MAX || pos > (uint32_t)INT32_MAX - dw->block_size) { return false; } const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_bytes = hc_dim * sizeof(float); int32_t positions[DS4_DSPARK_MAX_BLOCK_SIZE + 1u]; if (ds4_dspark_rocm_gfx1151_reference_alignment() && pos == 0) return false; positions[0] = (int32_t)(ds4_dspark_rocm_gfx1151_reference_alignment() ? pos - 1u : 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]; if (ds4_dspark_rocm_gfx1151_reference_alignment() && pos == 0) return false; positions[0] = (int32_t)(ds4_dspark_rocm_gfx1151_reference_alignment() ? pos - 1u : 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 bool direct_target_kv = ds4_dspark_rocm_gfx1151_reference_alignment(); if (direct_target_kv && (!metal_graph_batch_ffn_norm(g) || ds4_gpu_tensor_bytes(metal_graph_batch_ffn_norm(g)) < (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float))) { 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 && !direct_target_kv) 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 && !direct_target_kv) 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 (!direct_target_kv && 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 (!direct_target_kv) { 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; } /* DSpark target features already have stage-0 main_proj/main_norm. The reference applies wkv directly; support HC/attn_norm is for drafts. */ 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, direct_target_kv ? metal_graph_batch_ffn_norm(g) : 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; } const bool window_ok = ds4_dspark_rocm_gfx1151_reference_alignment() ? metal_graph_dspark_cache_merge_target_range(g, batch_start, n_tokens) : metal_graph_dspark_cache_set_window(g, batch_start, n_tokens); if (!window_ok) 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 bool align_rocm = ds4_dspark_rocm_gfx1151_reference_alignment(); if (align_rocm && pos == 0) return false; const uint32_t feature_pos = align_rocm ? pos - 1u : pos; 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 : (feature_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 && ds4_dspark_rocm_gfx1151_reference_alignment()) { /* Keep every draft row's HC/attn_norm result; only target row zero uses the already-projected main_x shared by all support stages. */ ok = ds4_gpu_tensor_copy(metal_graph_batch_attn_norm(g), 0, g->dspark_main_x, 0, (uint64_t)DS4_N_EMBD * sizeof(float)) != 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), 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, feature_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_final_hidden( ds4_gpu_graph *g, const ds4_model *dspark_model, const ds4_dspark_weights *dw, bool commands_open); 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, bool encode_final_hidden, 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 bool align_rocm = ds4_dspark_rocm_gfx1151_reference_alignment(); if (align_rocm && pos == 0) return false; const uint32_t feature_pos = align_rocm ? pos - 1u : pos; 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, feature_pos)) { return false; } if (cache_start_out) { *cache_start_out = support_len ? raw_start : (feature_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; } /* Final hidden rows consume only the last stage's GPU output. Keep them in * this command buffer when the CPU confidence precheck will need them. */ if (ok && encode_final_hidden) { ok = metal_graph_eval_dspark_final_hidden(g, dspark_model, dw, true); } 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; } if (align_rocm) { /* Draft rows remain temporary; only the committed target feature joins the history. */ return metal_graph_dspark_cache_merge_target_range(g, feature_pos, 1u); } 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) { const bool align_rocm = ds4_dspark_rocm_gfx1151_reference_alignment(); if (!g || !dspark_model || !dw || !g->dspark_capture_valid || (!align_rocm && 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 && align_rocm) { ok = metal_graph_dspark_cache_merge_target_range(g, pos, 1u); } else 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, uint32_t draft_cap) { 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 = draft_cap != 0 && draft_cap < dw->block_size ? draft_cap : 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, bool commands_open) { 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 && !commands_open) 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 && !commands_open) ok = ds4_gpu_end_commands() != 0; if (!ok && !commands_open) (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; } #ifndef __APPLE__ static bool dspark_apply_markov_greedy_gpu_runtime( ds4_gpu_graph *g, const ds4_model *dspark_model, const ds4_dspark_weights *dw, int first_prev_token, uint32_t draft_cap, int32_t proposal[DS4_DSPARK_MAX_BLOCK_SIZE], uint32_t *proposal_len) { if (proposal_len) *proposal_len = 0; if (!g || !g->spec_logits || !g->dspark_draft_tokens || !dspark_model || !dw || !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]; if (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) { return false; } const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); int32_t prev_token = first_prev_token; const uint32_t draft_count = draft_cap != 0 && draft_cap < dw->block_size ? draft_cap : dw->block_size; for (uint32_t draft = 0; draft < draft_count; draft++) { ds4_gpu_tensor *row_view = ds4_gpu_tensor_view(g->spec_logits, (uint64_t)draft * logits_bytes, logits_bytes); uint64_t gpu_key = 0; const bool 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 (!ok || gpu_key == 0 || gpu_token >= DS4_N_VOCAB) { return false; } proposal[draft] = (int32_t)gpu_token; prev_token = (int32_t)gpu_token; } if (proposal_len) *proposal_len = draft_count; return true; } #endif static bool dspark_apply_markov_greedy_probe( float *logits, const ds4_model *dspark_model, const ds4_dspark_weights *dw, int first_prev_token, uint32_t draft_cap, 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; const uint32_t draft_count = draft_cap != 0 && draft_cap < dw->block_size ? draft_cap : dw->block_size; for (uint32_t draft = 0; draft < draft_count; 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 = draft_count; 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; struct ds4_glm_gpu_graph; static bool imatrix_collect_glm_one( ds4_imatrix_collector *c, struct ds4_glm_gpu_graph *g, uint32_t il); 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_tensor_batch( ds4_imatrix_collector *c, ds4_gpu_tensor *ffn_norm, ds4_gpu_tensor *routed_mid, ds4_gpu_tensor *router_selected, bool routed_mid_is_f16, uint32_t il, uint32_t n_tokens) { if (!c || n_tokens == 0) return true; if (!ffn_norm || !routed_mid || !router_selected || 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 * (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 = routed_mid_is_f16 ? (void *)c->routed_mid_f16_buf : (void *)c->routed_mid_buf; if (ds4_gpu_tensor_read(ffn_norm, 0, c->ffn_norm_buf, norm_bytes) == 0 || ds4_gpu_tensor_read(routed_mid, 0, mid_dst, mid_bytes) == 0 || ds4_gpu_tensor_read(router_selected, 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 (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 bool imatrix_collect_layer_batch( ds4_imatrix_collector *c, ds4_gpu_graph *g, uint32_t il, uint32_t n_tokens) { if (!g) return false; return imatrix_collect_tensor_batch(c, metal_graph_batch_ffn_norm(g), metal_graph_batch_routed_mid(g), metal_graph_batch_router_selected(g), g->batch_routed_mid_is_f16, il, n_tokens); } 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) { memset(tmp, 0, (size_t)n_col * sizeof(tmp[0])); } 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; } int32_t entries = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (weights->layer[il].ffn_gate_exps) entries += 3; } imatrix_write_i32(fp, entries); 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; 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 uint32_t imatrix_collector_min_samples( const ds4_imatrix_collector *c, const ds4_weights *weights, uint32_t layer_limit) { bool found = false; uint32_t min_samples = UINT32_MAX; if (layer_limit > DS4_N_LAYER) layer_limit = DS4_N_LAYER; for (uint32_t il = 0; il < layer_limit; il++) { if (!weights->layer[il].ffn_gate_exps) continue; for (uint32_t expert = 0; expert < DS4_N_EXPERT; expert++) { uint32_t samples = c->gate_up_count[il][expert]; if (c->down_count[il][expert] < samples) { samples = c->down_count[il][expert]; } if (samples < min_samples) min_samples = samples; found = true; } } return found ? min_samples : 0; } static void imatrix_collector_report_coverage( const ds4_imatrix_collector *c, const ds4_weights *weights, uint32_t layer_limit) { uint32_t layers = 0; uint32_t covered = 0; uint32_t possible = 0; uint32_t min_layer = DS4_N_EXPERT; uint32_t max_layer = 0; uint32_t min_samples = UINT32_MAX; uint32_t max_samples = 0; uint32_t under_four = 0; uint32_t under_eight = 0; if (layer_limit > DS4_N_LAYER) layer_limit = DS4_N_LAYER; for (uint32_t il = 0; il < layer_limit; il++) { if (!weights->layer[il].ffn_gate_exps) continue; uint32_t layer_covered = 0; for (uint32_t expert = 0; expert < DS4_N_EXPERT; expert++) { uint32_t samples = c->gate_up_count[il][expert]; if (c->down_count[il][expert] < samples) { samples = c->down_count[il][expert]; } if (samples != 0) { layer_covered++; } if (samples < min_samples) min_samples = samples; if (samples > max_samples) max_samples = samples; if (samples < 4) under_four++; if (samples < 8) under_eight++; } layers++; covered += layer_covered; possible += DS4_N_EXPERT; if (layer_covered < min_layer) min_layer = layer_covered; if (layer_covered > max_layer) max_layer = layer_covered; } if (layers == 0) { min_layer = 0; min_samples = 0; } fprintf(stderr, "ds4: imatrix expert coverage %u/%u (%.2f%%), " "per sparse layer min=%u max=%u, samples min=%u max=%u, " "under4=%u under8=%u\n", covered, possible, possible ? 100.0 * (double)covered / (double)possible : 0.0, min_layer, max_layer, min_samples, max_samples, under_four, under_eight); if (under_eight != 0) { uint32_t shown = 0; fprintf(stderr, "ds4: imatrix expert slots under 8 samples:"); for (uint32_t samples = 0; samples < 8 && shown < 256; samples++) { for (uint32_t il = 0; il < layer_limit && shown < 256; il++) { if (!weights->layer[il].ffn_gate_exps) continue; for (uint32_t expert = 0; expert < DS4_N_EXPERT && shown < 256; expert++) { uint32_t count = c->gate_up_count[il][expert]; if (c->down_count[il][expert] < count) { count = c->down_count[il][expert]; } if (count != samples) continue; fprintf(stderr, " %u:%u=%u", il, expert, count); shown++; } } } if (shown < under_eight) { fprintf(stderr, " ... (%u more)", under_eight - shown); } fputc('\n', stderr); } } static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { memset(g->layer_n_comp, 0, sizeof(g->layer_n_comp)); memset(g->layer_n_index_comp, 0, sizeof(g->layer_n_index_comp)); g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); metal_graph_dspark_capture_invalidate(g); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!g->layer_raw_cache[il]) continue; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint32_t coff = ratio == 4 ? 2u : 1u; const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; const uint64_t attn_rows = (uint64_t)coff * ratio; if (!metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows)) return false; if (!metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows)) return false; if (ratio == 4) { const uint64_t index_width = (uint64_t)coff * DS4_N_INDEXER_HEAD_DIM; const uint64_t index_rows = (uint64_t)coff * ratio; if (!metal_tensor_fill_f32(g->layer_index_state_kv[il], 0.0f, index_width * index_rows)) return false; if (!metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows)) return false; } } return true; } /* Execute graph-backend prefill in layer-major order so intermediate * activations stay on the GPU and cache state is built exactly once. */ static void gpu_graph_report_prefill_display_progress( ds4_session_progress_fn display_progress, void *display_progress_ud, uint32_t start, uint32_t n_tokens, uint32_t layer_done, int total) { if (!display_progress) return; if (layer_done > (uint32_t)DS4_N_LAYER) layer_done = (uint32_t)DS4_N_LAYER; uint64_t done = (uint64_t)n_tokens * layer_done / (uint32_t)DS4_N_LAYER; if (layer_done == (uint32_t)DS4_N_LAYER) done = n_tokens; display_progress(display_progress_ud, "prefill_display", (int)(start + (uint32_t)done), total); } typedef struct { int tier; uint32_t first_layer; uint32_t end_layer; } metal_graph_prefill_stage; static bool metal_graph_build_prefill_stages( const ds4_gpu_graph *g, metal_graph_prefill_stage *stages, uint32_t *n_stages) { if (!g || !g->placement || !stages || !n_stages) return false; uint32_t ns = 0; int prev_tier = -1; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const int tier = g->placement[il + 1]; if (tier < 0 || tier >= DS4_MAX_GPUS) return false; if (il == 0 || tier != prev_tier) { if (ns >= DS4_MAX_GPUS) return false; stages[ns].tier = tier; stages[ns].first_layer = il; stages[ns].end_layer = il + 1u; ns++; prev_tier = tier; } else { stages[ns - 1u].end_layer = il + 1u; } } *n_stages = ns; return ns != 0; } static bool metal_graph_encode_prefill_stage_batch( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, const metal_graph_prefill_stage *stage, uint32_t pos0, uint32_t n_tokens) { if (!g || !model || !weights || !stage || n_tokens == 0) return false; if (!metal_graph_set_active_tier_no_copy(g, stage->tier)) return false; for (uint32_t il = stage->first_layer; il < stage->end_layer; il++) { if (g->placement && g->placement[il + 1] != stage->tier) return false; if (!metal_graph_encode_layer_batch(g, model, &weights->layer[il], il, pos0, n_tokens)) { return false; } if (g->pipeline_capture_chunk_len != 0 && !metal_graph_dspark_capture_prefill_rows( g, il, g->pipeline_capture_chunk_start, g->pipeline_capture_chunk_len, pos0, n_tokens)) { return false; } } return true; } static bool metal_graph_prefill_pipeline_stage_major( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, uint32_t start, uint32_t n_tokens, float *logits, bool show_progress, ds4_session_progress_fn display_progress, void *display_progress_ud) { if (!g || !model || !weights || !prompt || !g->placement || n_tokens == 0 || n_tokens > g->prefill_cap || start > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - start) { return false; } metal_graph_prefill_stage stages[DS4_MAX_GPUS]; uint32_t n_stages = 0; g->pipeline_capture_chunk_start = start; g->pipeline_capture_chunk_len = g->dspark_capture_enabled ? n_tokens : 0; if (!metal_graph_build_prefill_stages(g, stages, &n_stages) || n_stages < 2) { return false; } if (stages[0].tier != g->emb_tier) { return false; } uint32_t mb_cap = metal_graph_cuda_prefill_pipeline_microbatch(); if (mb_cap == 0 || mb_cap >= n_tokens) return false; if (mb_cap > g->prefill_cap) mb_cap = g->prefill_cap; const uint32_t n_mb = (n_tokens + mb_cap - 1u) / mb_cap; if (n_mb < 2) return false; if (display_progress) display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; bool ok = true; const double t0 = getenv("DS4_METAL_GRAPH_PREFILL_PROFILE") ? now_sec() : 0.0; const bool sequential = getenv("DS4_CUDA_PREFILL_PIPELINE_SEQUENTIAL") != NULL; const bool suppress_q8_cache = !metal_graph_cuda_prefill_pipeline_q8_cache_requested(); const int saved_q8_cache_suppressed = ds4_gpu_q8_cache_suppressed(); if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(1); ok = ds4_gpu_begin_commands() != 0; if (sequential) { for (uint32_t mb_i = 0; ok && mb_i < n_mb; mb_i++) { const uint32_t mb_off = mb_i * mb_cap; uint32_t mb_len = n_tokens - mb_off; if (mb_len > mb_cap) mb_len = mb_cap; const uint32_t pos0 = start + mb_off; for (uint32_t stage_i = 0; ok && stage_i < n_stages; stage_i++) { g->batch_token_offset = mb_off; if (stage_i == 0) { ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); ds4_gpu_tensor *tokens_view = NULL; if (ok) { tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), (uint64_t)mb_off * sizeof(int32_t), (uint64_t)mb_len * sizeof(int32_t)); ok = tokens_view != NULL; } if (ok) { ok = metal_graph_upload_prompt_embeddings_hc( g->batch_cur_hc_by_tier[stages[0].tier], tokens_view, model, weights, prompt, pos0, mb_len); } ds4_gpu_tensor_free(tokens_view); } if (ok) { ok = metal_graph_encode_prefill_stage_batch(g, model, weights, &stages[stage_i], pos0, mb_len); } if (ok && stage_i + 1u < n_stages) { ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && ds4_gpu_synchronize() != 0; } ok = src && dst && ds4_gpu_tensor_copy_xdev_ordered(dst, src, (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); if (ok && display_progress) { uint32_t done = mb_off + mb_len; if (done > n_tokens) done = n_tokens; display_progress(display_progress_ud, "prefill_display", (int)(start + done), prompt->len); } if (show_progress) { fprintf(stderr, "ds4: gpu sequential pipeline prefill microbatch %u/%u\r", mb_i + 1u, n_mb); fflush(stderr); } if (ok && mb_i + 1u < n_mb) ok = ds4_gpu_begin_commands() != 0; } } else { for (uint32_t wave = 0; ok && wave < n_mb + n_stages - 1u; wave++) { uint32_t smax = wave < n_stages ? wave : n_stages - 1u; for (int si = (int)smax; ok && si >= 0; si--) { const uint32_t stage_i = (uint32_t)si; const uint32_t mb_i = wave - stage_i; if (mb_i >= n_mb) continue; const uint32_t mb_off = mb_i * mb_cap; uint32_t mb_len = n_tokens - mb_off; if (mb_len > mb_cap) mb_len = mb_cap; const uint32_t pos0 = start + mb_off; g->batch_token_offset = mb_off; if (stage_i == 0) { ok = metal_graph_set_active_tier_no_copy(g, stages[0].tier); ds4_gpu_tensor *tokens_view = NULL; if (ok) { tokens_view = ds4_gpu_tensor_view(metal_graph_prefill_tokens(g), (uint64_t)mb_off * sizeof(int32_t), (uint64_t)mb_len * sizeof(int32_t)); ok = tokens_view != NULL; } if (ok) { ok = metal_graph_upload_prompt_embeddings_hc( g->batch_cur_hc_by_tier[stages[0].tier], tokens_view, model, weights, prompt, pos0, mb_len); } ds4_gpu_tensor_free(tokens_view); } if (ok) { ok = metal_graph_encode_prefill_stage_batch(g, model, weights, &stages[stage_i], pos0, mb_len); } if (ok && stage_i + 1u < n_stages) { ds4_gpu_tensor *src = g->batch_cur_hc_by_tier[stages[stage_i].tier]; ds4_gpu_tensor *dst = g->batch_cur_hc_by_tier[stages[stage_i + 1u].tier]; if (ok && getenv("DS4_CUDA_PREFILL_PIPELINE_SYNC_BOUNDARY") != NULL) { ok = metal_graph_set_active_tier_no_copy(g, stages[stage_i].tier) && ds4_gpu_synchronize() != 0; } ok = src && dst && ds4_gpu_tensor_copy_xdev_ordered(dst, src, (uint64_t)mb_len * hc_dim * sizeof(float)) != 0; } if (ok && display_progress && stage_i + 1u == n_stages) { uint32_t done = mb_off + mb_len; if (done > n_tokens) done = n_tokens; display_progress(display_progress_ud, "prefill_display", (int)(start + done), prompt->len); } } if (show_progress) { fprintf(stderr, "ds4: gpu pipeline prefill wave %u/%u\r", wave + 1u, n_mb + n_stages - 1u); fflush(stderr); } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); } if (show_progress) fputc('\n', stderr); g->batch_token_offset = 0; if (!ok) { if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); return false; } const uint32_t final_len = n_tokens - (n_mb - 1u) * mb_cap; const int src_tier = stages[n_stages - 1u].tier; if (!metal_graph_set_active_tier_no_copy(g, src_tier)) { if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); return false; } ds4_gpu_tensor *saved_cur = g->cur_hc_by_tier[src_tier]; ds4_gpu_tensor *last_hc = NULL; if (logits) { last_hc = metal_graph_tensor_row_view(g->batch_cur_hc_by_tier[src_tier], final_len - 1u, hc_dim); ok = last_hc != NULL; } if (ok && logits) { g->cur_hc_by_tier[src_tier] = last_hc; ok = ds4_gpu_begin_commands() != 0; } if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); if (ok && logits) ok = ds4_gpu_end_commands() != 0; else if (!ok) (void)ds4_gpu_synchronize(); g->cur_hc_by_tier[src_tier] = saved_cur; ds4_gpu_tensor_free(last_hc); if (g->placement && g->active_tier != src_tier) { ok = metal_graph_set_active_tier_no_copy(g, src_tier); } if (ok && logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (ok && display_progress) display_progress(display_progress_ud, "prefill_display", (int)(start + n_tokens), prompt->len); if (ok && t0 != 0.0) { const double t1 = now_sec(); fprintf(stderr, "ds4: gpu pipeline prefill total tokens=%u stages=%u mb=%u total=%.3f ms\n", n_tokens, n_stages, mb_cap, (t1 - t0) * 1000.0); } if (suppress_q8_cache) ds4_gpu_set_q8_cache_suppressed(saved_q8_cache_suppressed); return ok; } static bool metal_graph_prefill_layer_major( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, uint32_t start, uint32_t n_tokens, float *logits, bool show_progress, ds4_imatrix_collector *imatrix, ds4_session_progress_fn display_progress, void *display_progress_ud) { if (n_tokens == 0 || n_tokens > g->prefill_cap) return false; if (start > (uint32_t)prompt->len) return false; if (n_tokens > (uint32_t)prompt->len - start) return false; g->prefill_host_tokens = (const int32_t *)(prompt->v + start); g->prefill_host_pos0 = start; g->prefill_host_count = n_tokens; g->prefill_has_visual = false; if (g->deepseek4_vision_weights) { for (uint32_t i = 0; i < n_tokens; i++) { if (prompt->v[start + i] >= (int)DS4_N_VOCAB) { g->prefill_has_visual = true; break; } } } 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 && !g->prefill_has_visual && 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 = g->prefill_has_visual ? metal_graph_upload_visual_prompt_embeddings_hc( g, metal_graph_batch_cur_hc(g), model, weights, prompt, start, n_tokens) : 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) { const bool layer_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported(weights, 0); if (layer_prepare) { if (!metal_graph_stream_prepare_start_if_needed(g, model, weights, 0, n_tokens, layer_madvise, layer_pread, layer_readahead, layer_selected_addr, layer_prepare_slots, layer_prepare_ahead)) { return false; } } else { if (layer_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 = g->prefill_has_visual ? metal_graph_upload_visual_prompt_embeddings_hc( g, metal_graph_batch_cur_hc(g), model, weights, prompt, start, n_tokens) : 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; const bool layer_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported(weights, il); if (layer_prepare && !metal_graph_stream_prepare_join_layer(g, model, weights, il, n_tokens, layer_madvise, layer_pread, layer_readahead, layer_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 = layer_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; const uint32_t future_il = il + ahead; const bool future_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported( weights, future_il); if (!metal_graph_stream_prepare_start_if_needed(g, model, weights, future_il, n_tokens, layer_madvise, layer_pread, layer_readahead, future_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) { const bool next_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported( weights, il + 1); if (next_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); #ifdef __APPLE__ if (ok && !layer_selected_addr) { ok = metal_graph_seed_streaming_expert_cache_layer_from_mapped_hotlist( g, model, weights, il); } #endif 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); #ifdef __APPLE__ if (ok && !layer_selected_addr) { ok = metal_graph_seed_streaming_expert_cache_layer_from_mapped_hotlist( g, model, weights, il); } #endif 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) { const bool next_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported( weights, il + 1); if (!metal_graph_stream_prepare_start_if_needed(g, model, weights, il + 1, n_tokens, layer_madvise, layer_pread, layer_readahead, next_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 /* * The final layer can use its fully mapped prefill table after the * asynchronous selected-expert loader has already queued a resident-cache * batch. Drain that completed batch before hotlist seeding starts another * read-pool job set, otherwise both callers wait forever for ownership of * the single ROCm streaming read pool. */ if (g->ssd_streaming && !ds4_gpu_stream_expert_cache_finish_pending_batch()) { return false; } (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() && weights_model_map_decode_static_supported(weights); const bool static_map_state_cache = static_decode_map && metal_graph_stream_decode_static_map_state_cache_enabled(); g->streaming_static_decode_map_current = false; if (static_map_state_cache) { ok = metal_graph_stream_map_decode_static_all(model, weights); if (ok) g->streaming_static_decode_map_current = true; } else { ok = metal_graph_stream_map_output(model, weights); } } if (ok && logits) { g->cur_hc_by_tier[g->active_tier] = last_hc; ok = ds4_gpu_begin_commands() != 0; } if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); const double t_head_encoded = profile ? now_sec() : 0.0; if (ok && logits) ok = ds4_gpu_end_commands() != 0; const double t_head_done = profile ? now_sec() : 0.0; g->cur_hc_by_tier[g->active_tier] = saved_cur; if (last_hc) ds4_gpu_tensor_free(last_hc); if (!ok) return false; const double t_before_read = profile ? now_sec() : 0.0; if (logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (profile) { const double t_read = now_sec(); encode_s += t_head_encoded - t_head0; execute_s += t_head_done - t_head_encoded; if (split_profile) { fprintf(stderr, "ds4: gpu layer-major prefill head encode=%.3f ms execute=%.3f ms\n", (t_head_encoded - t_head0) * 1000.0, (t_head_done - t_head_encoded) * 1000.0); } fprintf(stderr, "ds4: gpu layer-major prefill total tokens=%u encode=%.3f ms execute=%.3f ms read=%.3f ms total=%.3f ms\n", n_tokens, encode_s * 1000.0, execute_s * 1000.0, (t_read - t_before_read) * 1000.0, (t_read - t0) * 1000.0); } return ok; } static bool metal_graph_prefill_raw_swa( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, const token_vec *prompt, int n_tokens, float *logits, bool show_progress, ds4_session_progress_fn display_progress, void *display_progress_ud, ds4_session_cancel_fn cancel, void *cancel_ud, bool *cancelled) { if (n_tokens <= 0 || n_tokens > prompt->len) return false; if ((uint32_t)n_tokens > g->prefill_cap) return false; if (metal_graph_use_streaming_decode_prefill_range(g, weights, 0, (uint32_t)n_tokens)) { return metal_graph_prefill_decode_streaming_range(g, model, weights, prompt, 0, (uint32_t)n_tokens, logits, show_progress, NULL, NULL, display_progress, display_progress_ud, cancel, cancel_ud, cancelled); } /* The layer-major fallback below may submit the whole short prefill as one * Metal command buffer. Once that command is in flight there is no useful * safe prefix to expose: by the time cancellation can be observed again, * the prompt has already been fully read and the KV is valid. Let the * caller observe the pending interrupt at generation time instead. */ (void)cancel; (void)cancel_ud; (void)cancelled; return metal_graph_prefill_layer_major(g, model, weights, prompt, 0, (uint32_t)n_tokens, logits, show_progress, NULL, display_progress, display_progress_ud); } static uint32_t metal_graph_visual_prefill_chunk( const ds4_gpu_graph *g, const token_vec *prompt, uint32_t pos0, uint32_t end, uint32_t cap) { const uint32_t normal = end - pos0 < cap ? end - pos0 : cap; if (!g->deepseek4_vision_weights) return normal; uint32_t chunk = 0; if (!ds4_deepseek4_prefill_chunk( prompt->v, end, DS4_N_VOCAB, pos0, cap, &chunk)) { fprintf(stderr, "ds4: malformed or oversized DeepSeek image token block at token %u\n", pos0); return 0; } return chunk; } /* 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; } } if (g->deepseek4_vision_weights && prompt->v[pos0] >= (int)DS4_N_VOCAB) { local_cap = chunk_cap; } uint32_t chunk = remaining < local_cap ? remaining : local_cap; chunk = metal_graph_visual_prefill_chunk(g, prompt, pos0, end, local_cap); if (chunk == 0) return false; 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_prefixes; /* The slot count bounds retained prefixes, not the verify block length. * Seed-plus-five blocks still need fresh states for every partial prefix. */ g->spec_capture_prefixes = capture_prefix1 && n_tokens > 1u; 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; #ifdef DS4_ROCM_BUILD bool rocm_dspark_fast = false; const char *verify_fast_env = getenv("DS4_ROCM_DSPARK_VERIFY_FAST"); const bool verify_fast_enabled = verify_fast_env && verify_fast_env[0] ? verify_fast_env[0] != '0' : ds4_dspark_rocm_gfx1151_fast_path(); rocm_dspark_fast = n_tokens >= 2u && n_tokens <= 6u && verify_fast_enabled; if (rocm_dspark_fast) ds4_gpu_set_dspark_verify_mode(true); #endif const double layer_t0 = timing ? now_sec() : 0.0; #if defined(__APPLE__) const bool tp_block = g->tp_batch_rows != 0 && g_tp_block_ctx != NULL; if (getenv("DS4_TP_BIG_GATE_DEBUG")) fprintf(stderr, "ds4: verify block: rows %u world %d ctx %p tp_block %d\n", g->tp_batch_rows, g->tp_world, (void *)g_tp_block_ctx, (int)tp_block); if (tp_block && !ds4_tp_batch_block_begin(g_tp_block_ctx, g->tp_batch_rows, DS4_N_LAYER)) { fprintf(stderr, "ds4: TP verify-block window setup failed\n"); ok = false; } #endif ok = 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; } #if defined(__APPLE__) /* Tiny DFlash batches are expensive enough to keep Metal busy while * the host encodes the next layers. Submit short prefixes to overlap * that work without changing the verifier's kernels or arithmetic. */ if (ok && capture_dspark_hidden && !verify_profile && !selected_profile && g->tp_world != 2 && ((il + 1u) % 4u) == 0u) { ok = ds4_gpu_flush_commands() != 0; } #endif } 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(); #if defined(__APPLE__) if (tp_block && !ds4_tp_batch_block_end(g_tp_block_ctx)) { fprintf(stderr, "ds4: TP verify-block window did not complete\n"); ok = false; } #endif if (ok && getenv("DS4_DSPARK_EXPERT_OVERLAP")) { /* Diagnostic: distinct routed experts selected by the block's rows * in the last layer (bounds the gain of a multi-row expert kernel). */ const int32_t *sel = ds4_gpu_tensor_contents(metal_graph_batch_router_selected(g)); if (sel) { uint8_t seen[1024] = {0}; uint32_t distinct = 0; for (uint32_t r = 0; r < n_tokens; r++) for (uint32_t k = 0; k < DS4_N_EXPERT_USED; k++) { const int32_t id = sel[r * DS4_N_EXPERT_USED + k]; if (id >= 0 && id < 1024 && !seen[id]) { seen[id] = 1; distinct++; } } fprintf(stderr, "ds4: DSpark verify block: %u rows, %u distinct experts of %u slots (last layer)\n", n_tokens, distinct, n_tokens * DS4_N_EXPERT_USED); } } #ifdef DS4_ROCM_BUILD if (rocm_dspark_fast) ds4_gpu_set_dspark_verify_mode(false); #endif g->spec_capture_prefixes = 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_prefixes; 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_prefixes = 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_prefix_attn_state(g, il, 0) && metal_graph_capture_prefix_index_state(g, il, 0); 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_prefixes = 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_prefixes = 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; } static uint32_t glm53_graph_resume_prefill_min_tokens(void) { #ifndef DS4_ROCM_BUILD /* Measured on M5 Max and GB10: one token is marginally faster through * decode, while two or more tokens amortize the batch setup cost. */ const char *env = getenv("DS4_GLM_RESUME_PREFILL_MIN"); if (!env || !env[0]) return 2u; #endif return glm_graph_resume_prefill_min_tokens(); } #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 #define DS4_GLM53_INDEX_POOL_SIZE 4u #define DS4_GLM53_PREFILL_CHUNK_TOKENS 2048u 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; /* Absorbed MLA keeps the shared latent cache through the dense prefix, * then reuses it for sparse DSA without materializing per-head K/V. */ return false; } static bool glm_graph_layer_uses_full_indexer(uint32_t il) { if (ds4_model_is_glm53()) return !ds4_glm53_layer_is_kda(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 glm53_graph_kda_layer_count_range(uint32_t layer_start, uint32_t layer_end) { if (!ds4_model_is_glm53() || layer_start > layer_end) return 0; uint32_t n = 0; for (uint32_t il = layer_start; il <= layer_end; il++) { if (ds4_glm53_layer_is_kda(il)) n++; } return n; } static uint64_t glm53_graph_fixed_state_bytes(uint32_t layer_start, uint32_t layer_end) { if (!ds4_model_is_glm53()) return 0; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_mix = (uint64_t)DS4_N_HC * (DS4_N_HC + 2u); const uint64_t projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; const uint64_t kda_layers = glm53_graph_kda_layer_count_range(layer_start, layer_end); const uint64_t conv_state = 3u * (DS4_N_KDA_CONV - 1u) * projection; const uint64_t recurrent_state = projection * DS4_N_KDA_HEAD_DIM; const uint64_t persistent = kda_layers * (conv_state + recurrent_state) * sizeof(float); const uint64_t scratch = (4u * hc_dim + 2u * hc_mix + DS4_N_EMBD + DS4_N_HC) * sizeof(float) + (6u * projection + 2u * DS4_N_KDA_HEAD_DIM + DS4_N_KDA_HEAD) * sizeof(float); return persistent + scratch; } 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 glm53_graph_indexer_selected_limit(void) { return glm_graph_indexer_top_k_limit() + DS4_GLM53_INDEX_POOL_SIZE - 1u; } static uint32_t glm53_graph_indexer_pool_cap(uint32_t token_cap) { return (token_cap + DS4_GLM53_INDEX_POOL_SIZE - 1u) / DS4_GLM53_INDEX_POOL_SIZE; } 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); } else fprintf(stderr, "ds4: cannot open GLM debug dump %s: %s\n", full, strerror(errno)); } else { fprintf(stderr, "ds4: cannot read GLM debug dump %s (%llu bytes, tensor %llu bytes)\n", full, (unsigned long long)bytes, (unsigned long long)ds4_gpu_tensor_bytes(t)); } 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 (ds4_model_is_glm53()) { return ctx_size; } 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 uint32_t glm_graph_batch_row_cap( uint32_t full_attention_cap, uint32_t indexed_prefill_cap, bool expanded_kv) { if (ds4_model_is_glm53()) { uint32_t cap = full_attention_cap; if (cap > DS4_GLM53_PREFILL_CHUNK_TOKENS) { cap = DS4_GLM53_PREFILL_CHUNK_TOKENS; } if (indexed_prefill_cap != 0 && cap > indexed_prefill_cap) { cap = indexed_prefill_cap; } return cap != 0 ? cap : 1u; } return expanded_kv || indexed_prefill_cap == 0 ? full_attention_cap : indexed_prefill_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; const uint32_t indexer_rows = ds4_model_is_glm53() ? glm53_graph_indexer_pool_cap(compact_cap) : compact_cap; total += (uint64_t)indexer_layers * indexer_rows * DS4_N_INDEXER_HEAD_DIM * elem; if (ds4_model_is_glm53()) { total += (uint64_t)indexer_layers * 2u * DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } 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; const uint64_t indexer_rows = ds4_model_is_glm53() ? glm53_graph_indexer_pool_cap(compact_cap) : compact_cap; uint64_t bytes = indexer_rows * sizeof(float); const uint64_t indexer_top_k = ds4_model_is_glm53() ? glm53_graph_indexer_selected_limit() : glm_graph_indexer_top_k_limit(); bytes += indexed_score_rows * indexer_rows * 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 */ if (ds4_model_is_glm53()) { bytes += indexed_rows * DS4_N_INDEXER_HEAD_DIM * sizeof(float); bytes += indexed_rows * (glm_graph_indexer_top_k_limit() / DS4_GLM53_INDEX_POOL_SIZE) * sizeof(uint32_t); } 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 = glm_graph_batch_row_cap(full_attention_cap, (uint32_t)indexed_rows, expanded_kv); const uint64_t indexer_top_k = ds4_model_is_glm53() && !expanded_kv ? full_attention_cap : ds4_model_is_glm53() ? glm53_graph_indexer_selected_limit() : 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)); if (ds4_model_is_glm53()) { bytes = glm_graph_workspace_add_bytes(bytes, DS4_N_INDEXER_HEAD_DIM, sizeof(float)); bytes = glm_graph_workspace_add_bytes( bytes, glm_graph_indexer_top_k_limit() / DS4_GLM53_INDEX_POOL_SIZE, sizeof(uint32_t)); } 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 (ds4_model_is_glm53() && batch_rows != 0) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_mix = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; const uint64_t kda_projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * 4u * hc_dim, sizeof(float)); bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * 2u * hc_mix, sizeof(float)); bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * 6u * kda_projection, sizeof(float)); bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * DS4_N_KDA_HEAD_DIM, sizeof(float)); bytes = glm_graph_workspace_add_bytes(bytes, batch_rows * DS4_N_KDA_HEAD, sizeof(float)); } if (indexed_rows != 0) { bytes = glm_graph_workspace_add_bytes(bytes, indexed_rows * qk_low_elems, sizeof(float)); } return bytes; } static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap_slice( uint32_t ctx, uint32_t work_ctx, uint32_t compact_cap, bool ssd_streaming, uint32_t layer_start, uint32_t layer_end) { ds4_context_memory m = {0}; const uint32_t normal_layers = glm_graph_normal_layer_count(); if (normal_layers == 0 || layer_start >= normal_layers || layer_end < layer_start) { return m; } if (layer_end >= normal_layers) layer_end = normal_layers - 1u; const uint32_t layer_count = layer_end - layer_start + 1u; const uint32_t attention_layer_count = ds4_model_is_glm53() ? glm_graph_full_indexer_layer_count_range(layer_start, layer_end) : 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 = glm_graph_batch_row_cap(work_ctx, indexed_rows, expanded_kv); m.prefill_cap = batch_rows; m.raw_cap = expanded_kv ? work_ctx : 0; if (expanded_kv) { m.raw_bytes = (uint64_t)attention_layer_count * work_ctx * ((uint64_t)DS4_N_HEAD * (DS4_N_KEY_MLA + DS4_N_VALUE_MLA)) * glm_graph_full_kv_cache_elem_bytes(); } m.scratch_bytes = glm_graph_workspace_bytes_for_cap(work_ctx, compact_cap, ssd_streaming) + glm53_graph_fixed_state_bytes(layer_start, layer_end); if (compact_cap != 0) { m.comp_cap = compact_cap; m.compressed_bytes = glm_graph_compact_cache_bytes_for_cap( attention_layer_count, glm_graph_full_indexer_layer_count_range(layer_start, layer_end), compact_cap); } m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; return m; } static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( uint32_t ctx, uint32_t work_ctx, uint32_t compact_cap, bool ssd_streaming) { const uint32_t normal_layers = glm_graph_normal_layer_count(); if (normal_layers == 0) { const ds4_context_memory empty = {0}; return empty; } return glm_graph_context_memory_estimate_for_compact_cap_slice( ctx, work_ctx, compact_cap, ssd_streaming, 0, normal_layers - 1u); } #ifdef DS4_HAS_DEEPSEEK41_GPU static ds4_context_memory ds41_graph_memory(uint32_t ctx); #endif static uint32_t qwen4_prefill_chunk_tokens(uint32_t ctx) { const char *env = getenv("DS4_QWEN4_PREFILL_CHUNK"); const unsigned long v = env && env[0] ? strtoul(env, NULL, 10) : 8192ul; uint32_t chunk = v == 0 || v > 65536ul ? 8192u : (uint32_t)v; return chunk > ctx ? ctx : chunk; } 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) && ds4_model_is_qwen4()) { /* per-token f16 K/V and raw indexer keys per attention layer, pooled * block keys, rope positions; transients scale with the prefill chunk */ const uint64_t T = prefill_chunk ? prefill_chunk : qwen4_prefill_chunk_tokens(ctx); const uint64_t blocks = ctx / 4u + 1u, E = DS4_N_EMBD, hc_dim = E * DS4_N_HC; const uint64_t kv_row = 2ull * DS4_N_HEAD_KV * DS4_N_HEAD_DIM * 2u; uint32_t n_attn = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) n_attn += !ds4_qwen4_layer_is_linear(il); m.prefill_cap = (uint32_t)T; m.raw_cap = ctx; m.comp_cap = (uint32_t)blocks; m.raw_bytes = (uint64_t)n_attn * ctx * kv_row + (uint64_t)ctx * 16u; m.compressed_bytes = (uint64_t)n_attn * (ctx * DS4_N_INDEXER_HEAD_DIM * 4u + blocks * DS4_N_INDEXER_HEAD_DIM * 2u); m.scratch_bytes = T * (12u * hc_dim + 24u * E + (uint64_t)(DS4_N_EXPERT_USED + 1u) * (E + 2u * DS4_N_FF_EXP) + 2u * blocks) * 4u; /* Recurrent layers have fixed state even at a tiny context. Reserve * both MTP snapshots as well as live state for conservative admission. */ m.scratch_bytes += 3ull * (DS4_N_LAYER - n_attn) * ((uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM * DS4_N_LIN_HEAD_DIM + (uint64_t)(DS4_N_LIN_CONV - 1u) * DS4_N_LIN_CONV_DIM) * sizeof(float); m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; return m; } if (ds4_backend_uses_graph(backend)) { #ifdef DS4_HAS_DEEPSEEK41_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) return ds41_graph_memory(ctx); #endif 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; } #ifdef DS4_HAS_DEEPSEEK41_GPU /* ========================================================================= * DeepSeek V4.1 GPU Graph. * ========================================================================= * Four layers own compressed KV/index keys. Every layer owns its sliding * window; index-source layers publish selections for following reuse layers. * Engram tables are never GPU resources, only bounded decoded row buffers. */ #define DS41_PREFILL_CAP 8192u #define DS41_INDEX_BATCH 32u #define DS41_CARRY_ROWS(X) \ X(residual, DS4_N_HC * DS4_N_EMBD, DS4_V41_CARRY_BF16) \ X(pre, DS4_N_HC, DS4_V41_CARRY_F32) X(ffn_split, 24, DS4_V41_CARRY_F32) \ X(selected_comp, DS4_N_INDEXER_TOP_K, DS4_V41_CARRY_F32) \ X(block_mask, (g->ctx + 7u) / 8u, DS4_V41_CARRY_MASK) #define DS41_PREFILL_STORAGE(X) \ X(residual, DS4_N_HC * DS4_N_EMBD) \ X(pre, DS4_N_HC) X(selected_comp, DS4_N_INDEXER_TOP_K) \ X(selected, DS4_N_EXPERT_USED) \ X(after_attn, DS4_N_HC * DS4_N_EMBD) X(ffn_split, 24) \ X(attn_split, 24) X(qr, DS4_N_LORA_Q) \ X(q, DS4_N_HEAD * DS4_N_HEAD_DIM) X(kv, DS4_N_HEAD_DIM) \ X(latent, DS4_N_HEAD_DIM) X(index_k, DS4_N_INDEXER_HEAD_DIM) \ X(pool_kv, DS4_N_HEAD_DIM) X(pool_score, DS4_N_HEAD_DIM) \ X(index_q, DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM) \ X(index_weights, DS4_N_INDEXER_HEAD) \ X(heads, DS4_N_HEAD * DS4_N_HEAD_DIM) X(low, DS4_N_OUT_GROUP * DS4_N_LORA_O) \ X(block, DS4_N_EMBD) \ X(mix, 24) X(x, DS4_N_EMBD) \ X(norm, DS4_N_EMBD) X(route_logits, DS4_N_EXPERT) \ X(route_probs, DS4_N_EXPERT) X(route_weights, DS4_N_EXPERT_USED) \ X(mid, DS4_N_EXPERT_USED * DS4_N_FF_EXP) \ X(block_mask, (g->ctx + 7u) / 8u) \ X(engram_rows, DS4_ENGRAM_COLS * DS4_ENGRAM_DIM) #define DS41_PREFILL_ALIASES(X) \ X(engram_kv, (DS4_N_HC + 1u) * DS4_N_EMBD) \ X(flat_norm, DS4_N_HC * DS4_N_EMBD) \ X(experts, DS4_N_EXPERT_USED * DS4_N_EMBD) \ X(gate, DS4_N_EXPERT_USED * DS4_N_FF_EXP) \ X(up, DS4_N_EXPERT_USED * DS4_N_FF_EXP) \ X(shared_gate, DS4_N_FF_EXP) X(shared_up, DS4_N_FF_EXP) \ X(shared_mid, DS4_N_FF_EXP) X(shared, DS4_N_EMBD) X(routed, DS4_N_EMBD) #define DS41_PREFILL_ROWS(X) DS41_PREFILL_STORAGE(X) DS41_PREFILL_ALIASES(X) typedef struct { #define DS41_ROW_FIELD(name, count) ds4_gpu_tensor *name; DS41_PREFILL_ROWS(DS41_ROW_FIELD) #undef DS41_ROW_FIELD } ds41_prefill_row; static uint32_t ds41_prefill_limit(uint32_t ctx) { const uint32_t limit = ctx < 8192u || getenv("DS4_METAL_DISABLE_V41_WIDE_CHUNK") ? 2048u : (ctx < 16384u || getenv("DS4_METAL_DISABLE_V41_8K_CHUNK")) ? 4096u : DS41_PREFILL_CAP; return ctx < limit ? ctx : limit; } static uint32_t ds41_carry_words(uint32_t width, uint32_t format, bool compact) { if (!compact || format == DS4_V41_CARRY_F32) return width; return format == DS4_V41_CARRY_BF16 ? (width + 1u) / 2u : (width + 31u) / 32u; } static uint32_t ds41_carry_cap(uint32_t ctx) { const bool compact = !getenv("DS4_METAL_DISABLE_V41_COMPACT_CARRY"); const uint64_t row_bytes = ((uint64_t)ds41_carry_words(DS4_N_HC * DS4_N_EMBD, DS4_V41_CARRY_BF16, compact) + DS4_N_HC + 24u + DS4_N_INDEXER_TOP_K + ds41_carry_words((ctx + 7u) / 8u, DS4_V41_CARRY_MASK, compact)) * sizeof(float); uint64_t cap = (UINT64_C(3) << 30) / row_bytes; if (cap > 32768u) cap = 32768u; if (cap > ctx) cap = ctx; const uint32_t chunk = ds41_prefill_limit(ctx); if (!chunk) return 0; /* Keep the causal sweep boundary independent of the encoder tile size. */ cap -= cap % 2048u; return cap > chunk ? (uint32_t)cap : 0; } #define DS41_SCRATCH(X) \ X(image_text_mask, (g->prefill_cap + 3u) / 4u) \ X(residual, DS4_N_HC * DS4_N_EMBD) \ X(after_attn, DS4_N_HC * DS4_N_EMBD) \ X(flat_norm, DS4_N_HC * DS4_N_EMBD) \ X(mix, 24) X(attn_split, 24) X(ffn_split, 24) X(pre, 4) \ X(x, DS4_N_EMBD) X(norm, DS4_N_EMBD) X(block, DS4_N_EMBD) \ X(qr, DS4_N_LORA_Q) X(q, DS4_N_HEAD * DS4_N_HEAD_DIM) \ X(kv, DS4_N_HEAD_DIM) X(latent, DS4_N_HEAD_DIM) \ X(raw_prefill, (g->prefill_cap + 128u) * DS4_N_HEAD_DIM) \ X(pool_kv, DS4_N_HEAD_DIM) X(pool_score, DS4_N_HEAD_DIM) \ X(index_q, DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM) \ X(index_k, DS4_N_INDEXER_HEAD_DIM) X(index_weights, DS4_N_INDEXER_HEAD) \ X(index_scores, DS41_INDEX_BATCH * g->ctx) X(selected_comp, DS4_N_INDEXER_TOP_K) \ X(index_packed, ds4_gpu_dsv41_indexer_packed_bytes(g->ctx, g->prefill_cap) / 4u) \ X(selected_kv, DS4_N_INDEXER_TOP_K * DS4_N_HEAD_DIM) \ X(block_scores, (g->ctx + 7u) / 8u) X(block_selected, 2048) \ X(block_mask, (g->ctx + 7u) / 8u) \ X(heads, DS4_N_HEAD * DS4_N_HEAD_DIM) X(low, DS4_N_OUT_GROUP * DS4_N_LORA_O) \ X(route_logits, DS4_N_EXPERT) X(route_probs, DS4_N_EXPERT) \ X(selected, DS4_N_EXPERT_USED) X(route_weights, DS4_N_EXPERT_USED) \ X(gate, DS4_N_EXPERT_USED * DS4_N_FF_EXP) \ X(up, DS4_N_EXPERT_USED * DS4_N_FF_EXP) X(mid, DS4_N_EXPERT_USED * DS4_N_FF_EXP) \ X(experts, DS4_N_EXPERT_USED * DS4_N_EMBD) X(routed, DS4_N_EMBD) \ X(shared_gate, DS4_N_FF_EXP) X(shared_up, DS4_N_FF_EXP) \ X(shared_mid, DS4_N_FF_EXP) X(shared, DS4_N_EMBD) \ X(engram_rows, DS4_ENGRAM_COLS * DS4_ENGRAM_DIM) \ X(engram_prefetch, (g->carry_cap ? g->carry_cap : g->prefill_cap) * DS4_ENGRAM_COLS * DS4_ENGRAM_DIM) \ X(engram_kv, (DS4_N_HC + 1u) * DS4_N_EMBD) X(logits, DS4_N_VOCAB) typedef struct { uint32_t ctx, pos, prefill_cap, carry_cap; uint64_t allocation_bytes; bool valid, streaming, encoder_resident, compact_carry, prefill_alias, quality; uint32_t tp_world, tp_rank; ds4_gpu_tensor *tp_logits_half; ds4_gpu_tensor **tp_out, **tp_in; ds4_imatrix_collector *imatrix; const ds4_vision_span *images; size_t image_count; ds4_engram_layout engram; ds4_engram_history history; uint32_t *token_map; uint32_t (*prefill_ids)[2][DS4_ENGRAM_COLS]; ds4_engram_table table[2]; float rows[2][DS4_ENGRAM_COLS * DS4_ENGRAM_DIM]; ds4_gpu_tensor *window[40]; ds4_gpu_tensor *compressed[4], *index_cache[4]; ds4_gpu_tensor *previous_kv[4], *previous_score[4]; ds4_gpu_tensor *engram_q_norm[2], *engram_k_norm[2]; ds41_prefill_row batch, *rows_view; ds41_prefill_row carry; ds4_gpu_tensor *prefill_tokens; #define DS41_FIELD(name, count) ds4_gpu_tensor *name; DS41_SCRATCH(DS41_FIELD) #undef DS41_FIELD } ds41_gpu_graph; static bool ds41_read_array(const ds4_model *m, const char *key, uint32_t type, uint64_t count, void *out, size_t bytes) { ds4_array_ref arr; if (!model_get_array(m, key, &arr) || arr.type != type || arr.len != count) return false; ds4_cursor c = cursor_at(m, arr.data_pos); return cursor_read(&c, out, bytes); } static void ds41_graph_free(ds41_gpu_graph *g) { if (!g) return; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) ds4_gpu_decode_graphs_invalidate(); #endif ds4_gpu_tensor_free(g->tp_logits_half); for (uint32_t i = 0; g->rows_view && i < g->prefill_cap; i++) { #define DS41_ROW_FREE(name, count) ds4_gpu_tensor_free(g->rows_view[i].name); DS41_PREFILL_ROWS(DS41_ROW_FREE) #undef DS41_ROW_FREE } #define DS41_BATCH_FREE(name, count) ds4_gpu_tensor_free(g->batch.name); DS41_PREFILL_ROWS(DS41_BATCH_FREE) #undef DS41_BATCH_FREE #define DS41_CARRY_FREE(name, count, format) ds4_gpu_tensor_free(g->carry.name); DS41_CARRY_ROWS(DS41_CARRY_FREE) #undef DS41_CARRY_FREE for (uint32_t i = 0; i < 2; i++) { ds4_engram_table_close(&g->table[i]); ds4_gpu_tensor_free(g->engram_q_norm[i]); ds4_gpu_tensor_free(g->engram_k_norm[i]); } for (uint32_t i = 0; i < 40; i++) ds4_gpu_tensor_free(g->window[i]); for (uint32_t i = 0; i < 4; i++) { ds4_gpu_tensor_free(g->compressed[i]); ds4_gpu_tensor_free(g->index_cache[i]); ds4_gpu_tensor_free(g->previous_kv[i]); ds4_gpu_tensor_free(g->previous_score[i]); } #define DS41_FREE(name, count) ds4_gpu_tensor_free(g->name); DS41_SCRATCH(DS41_FREE) #undef DS41_FREE free(g->token_map); free(g->prefill_ids); free(g->rows_view); ds4_gpu_tensor_free(g->prefill_tokens); memset(g, 0, sizeof(*g)); g->table[0].fd = g->table[1].fd = -1; } static uint64_t ds41_graph_bytes(uint32_t ctx) { const ds41_gpu_graph shape = {.ctx = ctx, .prefill_cap = ds41_prefill_limit(ctx), .carry_cap = ds41_carry_cap(ctx), .prefill_alias = !getenv("DS4_METAL_DISABLE_V41_PREFILL_ALIAS"), .compact_carry = !getenv("DS4_METAL_DISABLE_V41_COMPACT_CARRY")}, *g = &shape; uint64_t floats = (uint64_t)40 * 128 * 512; for (uint32_t i = 0; i < 4; i++) floats += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u) + 2u * 512u; floats += (uint64_t)2 * 2 * DS4_N_EMBD * DS4_N_HC; #define DS41_COUNT(name, count) floats += (count); DS41_SCRATCH(DS41_COUNT) #undef DS41_COUNT #define DS41_BATCH_COUNT(name, count) floats += (uint64_t)(count) * g->prefill_cap; DS41_PREFILL_STORAGE(DS41_BATCH_COUNT) if (!g->prefill_alias) { DS41_PREFILL_ALIASES(DS41_BATCH_COUNT) } #undef DS41_BATCH_COUNT #define DS41_CARRY_COUNT(name, count, format) \ floats += (uint64_t)ds41_carry_words(count, format, g->compact_carry) * g->carry_cap; DS41_CARRY_ROWS(DS41_CARRY_COUNT) #undef DS41_CARRY_COUNT /* Reserve view objects and the bounded host hash-ID array as well. */ floats += (uint64_t)g->prefill_cap * 512u; if (g->carry_cap > g->prefill_cap) floats += (uint64_t)(g->carry_cap - g->prefill_cap) * 2u * DS4_ENGRAM_COLS; /* Expert-major matrix kernels use one bounded packed activation buffer. */ const uint64_t packed = g->prefill_cap >= 512 ? (g->prefill_cap > 4096u ? UINT64_C(512) : UINT64_C(256)) * 1024 * 1024 : 0; /* The index sorter owns two full-width merge buffers outside the graph. */ const uint64_t sort = (uint64_t)DS41_INDEX_BATCH * ctx * 2u * sizeof(uint32_t); return floats * sizeof(float) + sizeof(*g) + (uint64_t)DS4_N_VOCAB * 4u + packed + sort; } static ds4_context_memory ds41_graph_memory(uint32_t ctx) { ds4_context_memory m = {.prefill_cap = ds41_prefill_limit(ctx), .raw_cap = 128, .comp_cap = ctx + 1u}; m.raw_bytes = 40u * 128u * 512u * sizeof(float); for (uint32_t i = 0; i < 4; i++) m.compressed_bytes += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u) * sizeof(float); m.total_bytes = ds41_graph_bytes(ctx); m.scratch_bytes = m.total_bytes - m.raw_bytes - m.compressed_bytes; return m; } static void ds41_graph_reset(ds41_gpu_graph *g) { g->pos = 0; g->valid = true; ds4_engram_history_reset(&g->history); /* Valid lengths, not zeroed storage, determine cache visibility. Each * partial pair is overwritten by its even-position token before use. */ } static DS4_MAYBE_UNUSED bool ds41_graph_alloc(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const char *path, uint32_t ctx, bool streaming) { memset(g, 0, sizeof(*g)); g->table[0].fd = g->table[1].fd = -1; if (!ctx || ctx > 1048576 || DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) return false; g->ctx = ctx; g->prefill_cap = ds41_prefill_limit(ctx); g->carry_cap = ds41_carry_cap(ctx); g->compact_carry = !getenv("DS4_METAL_DISABLE_V41_COMPACT_CARRY"); g->prefill_alias = !getenv("DS4_METAL_DISABLE_V41_PREFILL_ALIAS"); g->allocation_bytes = ds41_graph_bytes(ctx); g->streaming = streaming; g->tp_world = 1; g->token_map = malloc((size_t)DS4_N_VOCAB * sizeof(uint32_t)); g->prefill_ids = malloc((size_t)(g->carry_cap ? g->carry_cap : g->prefill_cap) * sizeof(*g->prefill_ids)); g->rows_view = calloc(g->prefill_cap, sizeof(*g->rows_view)); g->prefill_tokens = ds4_gpu_tensor_alloc((uint64_t)g->prefill_cap * sizeof(int32_t)); if (!g->token_map || !g->prefill_ids || !g->rows_view || !g->prefill_tokens) goto fail; g->engram.token_map = g->token_map; g->engram.vocab_size = DS4_N_VOCAB; g->engram.compressed_vocab_size = required_u32(m, "deepseek41.engram.compressed_vocab_size"); g->engram.pad_id = required_u32(m, "deepseek41.engram.pad_id"); if (!ds41_read_array(m, "deepseek41.engram.token_map", GGUF_VALUE_UINT32, DS4_N_VOCAB, g->token_map, (size_t)DS4_N_VOCAB * 4) || !ds41_read_array(m, "deepseek41.engram.rows", GGUF_VALUE_UINT32, 2, g->engram.rows, sizeof(g->engram.rows)) || !ds41_read_array(m, "deepseek41.engram.primes", GGUF_VALUE_UINT32, 48, g->engram.primes, sizeof(g->engram.primes)) || !ds41_read_array(m, "deepseek41.engram.multipliers", GGUF_VALUE_UINT64, 8, g->engram.multipliers, sizeof(g->engram.multipliers)) || !ds4_engram_layout_valid(&g->engram)) goto fail; for (uint32_t i = 0; i < 2; i++) { const uint32_t il = i ? 14u : 1u; const ds4_tensor *table = required_tensorf(m, "blk.%u.engram_embd.weight", il); if (!ds4_engram_table_open(&g->table[i], path, table->abs_offset, g->engram.rows[i])) goto fail; struct stat weights_stat, rows_stat; if (fstat(m->fd, &weights_stat) || fstat(g->table[i].fd, &rows_stat) || weights_stat.st_dev != rows_stat.st_dev || weights_stat.st_ino != rows_stat.st_ino) goto fail; const uint64_t bytes = (uint64_t)DS4_N_EMBD * DS4_N_HC * 4; g->engram_q_norm[i] = ds4_gpu_tensor_alloc(bytes); g->engram_k_norm[i] = ds4_gpu_tensor_alloc(bytes); if (!g->engram_q_norm[i] || !g->engram_k_norm[i] || !ds4_gpu_tensor_write(g->engram_q_norm[i], 0, tensor_data(m, w->layer[il].engram_q_norm), bytes) || !ds4_gpu_tensor_write(g->engram_k_norm[i], 0, tensor_data(m, w->layer[il].engram_k_norm), bytes)) goto fail; } for (uint32_t i = 0; i < 40; i++) { g->window[i] = ds4_gpu_tensor_alloc(128u * 512u * sizeof(float)); if (!g->window[i]) goto fail; } for (uint32_t i = 0; i < 4; i++) { const uint64_t cap = ctx / (i < 3 ? 2u : 1u) + 1u; g->compressed[i] = ds4_gpu_tensor_alloc(cap * 512u * sizeof(float)); g->index_cache[i] = ds4_gpu_tensor_alloc(cap * 128u * sizeof(float)); g->previous_kv[i] = ds4_gpu_tensor_alloc(512u * sizeof(float)); g->previous_score[i] = ds4_gpu_tensor_alloc(512u * sizeof(float)); if (!g->compressed[i] || !g->index_cache[i] || !g->previous_kv[i] || !g->previous_score[i]) goto fail; } #define DS41_ALLOC(name, count) \ if (!(g->name = (!strcmp(#name, "engram_prefetch") || !strcmp(#name, "image_text_mask") ? \ ds4_gpu_tensor_alloc_managed : ds4_gpu_tensor_alloc)( \ (uint64_t)((count) != 0 ? (count) : 1u) * sizeof(float)))) goto fail; DS41_SCRATCH(DS41_ALLOC) #undef DS41_ALLOC #define DS41_BATCH_ALLOC(name, count) \ if (!(g->batch.name = (!strcmp(#name, "engram_rows") || !strcmp(#name, "selected_comp") ? \ ds4_gpu_tensor_alloc_managed : ds4_gpu_tensor_alloc)( \ (uint64_t)(count) * g->prefill_cap * 4u))) goto fail; DS41_PREFILL_STORAGE(DS41_BATCH_ALLOC) /* Engram projection ends before HC normalization. HC normalization * precedes Q projection or follows attention; routed expert outputs are * produced only after attention has consumed Q. The attention output * projection consumes heads/low before FFN; FFN normalization consumes x * before shared-down, and the routed sum may replace the old block. */ if (g->prefill_alias) { #define DS41_BATCH_ALIAS(name, count, parent, offset) \ if (!(g->batch.name = ds4_gpu_tensor_view(g->batch.parent, \ (uint64_t)(offset) * g->prefill_cap * 4u, \ (uint64_t)(count) * g->prefill_cap * 4u))) goto fail; DS41_BATCH_ALIAS(engram_kv, (DS4_N_HC + 1u) * DS4_N_EMBD, q, 0) DS41_BATCH_ALIAS(flat_norm, DS4_N_HC * DS4_N_EMBD, q, 0) DS41_BATCH_ALIAS(experts, DS4_N_EXPERT_USED * DS4_N_EMBD, q, 0) DS41_BATCH_ALIAS(gate, DS4_N_EXPERT_USED * DS4_N_FF_EXP, heads, 0) DS41_BATCH_ALIAS(up, DS4_N_EXPERT_USED * DS4_N_FF_EXP, heads, DS4_N_EXPERT_USED * DS4_N_FF_EXP) DS41_BATCH_ALIAS(shared_gate, DS4_N_FF_EXP, low, 0) DS41_BATCH_ALIAS(shared_up, DS4_N_FF_EXP, low, DS4_N_FF_EXP) DS41_BATCH_ALIAS(shared_mid, DS4_N_FF_EXP, low, 2u * DS4_N_FF_EXP) DS41_BATCH_ALIAS(shared, DS4_N_EMBD, x, 0) DS41_BATCH_ALIAS(routed, DS4_N_EMBD, block, 0) #undef DS41_BATCH_ALIAS } else { DS41_PREFILL_ALIASES(DS41_BATCH_ALLOC) } #undef DS41_BATCH_ALLOC #define DS41_BATCH_VIEW(name, count) \ for (uint32_t i = 0; i < g->prefill_cap; i++) \ if (!(g->rows_view[i].name = ds4_gpu_tensor_view(g->batch.name, \ (uint64_t)i * (count) * 4u, (uint64_t)(count) * 4u))) goto fail; DS41_PREFILL_ROWS(DS41_BATCH_VIEW) #undef DS41_BATCH_VIEW if (!ds4_gpu_tensor_fill_f32(g->batch.block_mask, 0, (uint64_t)g->prefill_cap * ((g->ctx + 7u) / 8u))) goto fail; #define DS41_CARRY_ALLOC(name, count, format) \ if (g->carry_cap && !(g->carry.name = ds4_gpu_tensor_alloc( \ (uint64_t)ds41_carry_words(count, format, g->compact_carry) * g->carry_cap * sizeof(float)))) goto fail; DS41_CARRY_ROWS(DS41_CARRY_ALLOC) #undef DS41_CARRY_ALLOC ds41_graph_reset(g); fprintf(stderr, "ds4: V4.1 static context buffers %.2f MiB (ctx=%u), Engram disk-only\n", (double)ds41_graph_bytes(ctx) / 1048576.0, ctx); return true; fail: ds41_graph_free(g); return false; } #undef DS41_SCRATCH static bool ds41_bf16(ds4_gpu_tensor *x, uint32_t width) { return ds4_gpu_dsv41_quantize(x, width, 1, DS4_V41_BF16) != 0; } static bool ds41_matmul(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, bool round) { return metal_graph_matmul_plain_tensor(out, m, weight, weight->dim[0], weight->dim[1], in, 1) && (!round || ds41_bf16(out, (uint32_t)weight->dim[1])); } static bool ds41_matmul_batch(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t count, bool round) { const uint32_t width = (uint32_t)weight->dim[0], outputs = (uint32_t)weight->dim[1]; bool ok; /* Small decode batches retain scalar reductions before BF16 and sparse * routing boundaries. Preserve Metal's separate vocabulary-head dispatch. */ if (count >= 2 && count <= DS4_TP_BATCH_MAX_ROWS && #ifdef __APPLE__ outputs != DS4_N_VOCAB && #endif weight->type == DS4_TENSOR_Q8_0) { ok = ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(out, m->map, m->size, weight->abs_offset, width, outputs, in, count); } else if (count >= 2 && count <= DS4_TP_BATCH_MAX_ROWS && weight->type == DS4_TENSOR_F16) { ok = ds4_gpu_dsv41_projection_rows(out, m->map, m->size, weight->abs_offset, width, outputs, count, in); } else if (count >= 2 && count <= DS4_TP_BATCH_MAX_ROWS && weight->type == DS4_TENSOR_F32) { ok = true; for (uint32_t i = 0; ok && i < count; i++) { ds4_gpu_tensor *x = ds4_gpu_tensor_view((ds4_gpu_tensor *)in, (uint64_t)i * width * sizeof(float), (uint64_t)width * sizeof(float)); ds4_gpu_tensor *y = ds4_gpu_tensor_view(out, (uint64_t)i * outputs * sizeof(float), (uint64_t)outputs * sizeof(float)); ok = x && y && ds41_matmul(y, m, weight, x, false); ds4_gpu_tensor_free(x); ds4_gpu_tensor_free(y); } } else { ok = metal_graph_matmul_plain_tensor(out, m, weight, width, outputs, in, count); } return ok && (!round || ds4_gpu_dsv41_quantize(out, outputs, count, DS4_V41_BF16)); } static bool ds41_matmul_rows_batch(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t first, uint32_t count, uint32_t rows) { uint64_t row_bytes; if (first > weight->dim[1] || count > weight->dim[1] - first || !tensor_nbytes(weight->type, weight->dim[0], &row_bytes)) return false; ds4_tensor slice = *weight; slice.abs_offset += (uint64_t)first * row_bytes; slice.dim[1] = count; return ds41_matmul_batch(out, m, &slice, in, rows, true); } static bool ds41_matmul_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t first, uint32_t count) { return ds41_matmul_rows_batch(out, m, weight, in, first, count, 1); } static bool ds41_output_projection(ds41_gpu_graph *g, ds4_gpu_tensor *out, const ds4_model *m, const ds4_weights *w, const ds4_gpu_tensor *in, uint32_t rows) { ds4_tensor head = *w->output; if (g->tp_logits_half) { uint64_t row_bytes; if (!tensor_nbytes(head.type, head.dim[0], &row_bytes)) return false; head.dim[1] /= 2u; head.abs_offset += (uint64_t)g->tp_rank * head.dim[1] * row_bytes; } /* Head rows are independent; keep their original reduction and do not * add the BF16 rounding used by the attention row-slice helper. */ return ds41_matmul_batch(out, m, &head, in, rows, false); } static const ds4_vision_span *ds41_image_at(const ds41_gpu_graph *g, uint32_t pos) { for (size_t i = 0; i < g->image_count; i++) { const ds4_vision_span *span = &g->images[i]; if (pos < span->token_start) break; if ((uint64_t)pos - span->token_start < span->embedding.token_count) return span; } return NULL; } static void ds41_text_mask(const ds41_gpu_graph *g, uint32_t pos, uint32_t count, uint8_t *mask) { memset(mask, 1, count); for (size_t i = 0; i < g->image_count; i++) { const ds4_vision_span *span = &g->images[i]; const uint64_t end = (uint64_t)span->token_start + span->embedding.token_count; const uint64_t first = span->token_start > pos ? span->token_start : pos; const uint64_t last = end < (uint64_t)pos + count ? end : (uint64_t)pos + count; if (first < last) memset(mask + first - pos, 0, (size_t)(last - first)); } } static bool ds41_hash_tokens(ds41_gpu_graph *g, ds4_engram_history *history, const int *tokens, uint32_t count, uint32_t *ids) { uint8_t *mask = NULL; if (g->image_count) { mask = malloc(count); if (!mask) return false; ds41_text_mask(g, g->pos, count, mask); } const bool ok = ds4_engram_hash(&g->engram, history, tokens, mask, count, ids); free(mask); return ok; } static bool ds41_embed(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, ds4_gpu_tensor *out, ds4_gpu_tensor *scratch, int token, uint32_t pos) { const ds4_vision_span *image = ds41_image_at(g, pos); if (!image) return ds4_gpu_embed_token_hc_tensor(out, m->map, m->size, w->token_embd->abs_offset, DS4_N_VOCAB, (uint32_t)token, DS4_N_EMBD, DS4_N_HC); const float *data = image->embedding.data + (uint64_t)(pos - image->token_start) * DS4_N_EMBD; return ds4_gpu_tensor_write(scratch, 0, data, (uint64_t)DS4_N_EMBD * sizeof(float)) && ds4_gpu_repeat_hc_tensor(out, scratch, DS4_N_EMBD, DS4_N_HC); } static bool ds41_sum_partial(ds41_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t gate) { if (g->tp_world != 2) return true; const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + gate; if (!ds4_gpu_tensor_copy(g->tp_out[slot], 0, x, 0, (uint64_t)DS4_N_EMBD * 4u) || !ds4_gpu_tp_gate_encode(il, gate)) return false; ds4_gpu_tensor *first = g->tp_rank ? g->tp_in[slot] : g->tp_out[slot]; ds4_gpu_tensor *second = g->tp_rank ? g->tp_out[slot] : g->tp_in[slot]; return ds4_gpu_add_tensor(x, first, second, DS4_N_EMBD) != 0; } static bool ds41_norm(ds4_gpu_tensor *out, const ds4_gpu_tensor *in, const ds4_model *m, const ds4_tensor *weight) { return ds4_gpu_rms_norm_weight_tensor(out, in, m->map, m->size, weight->abs_offset, (uint32_t)weight->dim[0], DS4_RMS_EPS) && ds41_bf16(out, (uint32_t)weight->dim[0]); } static bool ds41_sum_partial_batch(ds41_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t count) { if (g->tp_world != 2) return true; /* Q is dead after attention; its expert-output alias is dead after the * routed reduction. Reuse it for the peer, without another large buffer. */ ds4_gpu_tensor *peer = g->batch.q; const uint64_t bytes = (uint64_t)count * DS4_N_EMBD * sizeof(float); if (!ds4_gpu_tp_big_gate_encode(il, count, x, peer, bytes)) return false; return ds4_gpu_add_tensor(x, g->tp_rank ? peer : x, g->tp_rank ? x : peer, count * DS4_N_EMBD) != 0; } static bool ds41_rope(ds4_gpu_tensor *x, uint32_t heads, uint32_t width, uint32_t il, uint32_t pos, bool inverse) { return ds4_gpu_dsv41_rope(x, width, heads, 1, pos, ds4_layer_compress_ratio(il) != 0, inverse); } static bool ds41_hc_mix(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, bool ffn) { const ds4_gpu_tensor *residual = ffn ? g->after_attn : g->residual; const ds4_tensor *fn = ffn ? l->hc_ffn_fn : l->hc_attn_fn; const ds4_tensor *scale = ffn ? l->hc_ffn_scale : l->hc_attn_scale; const ds4_tensor *base = ffn ? l->hc_ffn_base : l->hc_attn_base; return ds4_gpu_rms_norm_plain_tensor(g->flat_norm, residual, DS4_N_HC * DS4_N_EMBD, DS4_RMS_EPS) && ds41_matmul(g->mix, m, fn, g->flat_norm, false) && ds4_gpu_hc_split_sinkhorn_tensor(ffn ? g->ffn_split : g->attn_split, g->mix, m->map, m->size, scale->abs_offset, base->abs_offset, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS); } static bool ds41_attention_low(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { const uint32_t groups = DS4_N_OUT_GROUP / g->tp_world; const uint32_t group0 = g->tp_rank * groups; uint64_t output_row; return tensor_nbytes(l->attn_output_a->type, 4096, &output_row) && ds4_gpu_attention_output_low_q8_tensor(g->low, m->map, m->size, l->attn_output_a->abs_offset + (uint64_t)group0 * 1024u * output_row, 4096, 1024, groups, g->heads) && ds41_bf16(g->low, groups * DS4_N_LORA_O); } static bool ds41_attention_output(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { const uint32_t groups = DS4_N_OUT_GROUP / g->tp_world; if (!ds41_attention_low(g, m, l)) return false; return g->tp_world == 2 ? metal_graph_matmul_dense_quant_kslice(g->block, m, l->attn_output_b, 8192, (uint64_t)g->tp_rank * groups * 1024u, (uint64_t)groups * 1024u, DS4_N_EMBD, g->low, 0) : ds41_matmul(g->block, m, l->attn_output_b, g->low, false); } static bool ds41_attention_publish(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; if (ds41_kv_source(il)) { if (!ds41_matmul(g->pool_kv, m, l->attn_compressor_kv, g->norm, ratio == 1)) return false; if (ratio == 2) { if (!ds41_matmul(g->pool_score, m, l->attn_compressor_gate, g->norm, false) || !ds4_gpu_dsv41_pool2(g->latent, g->pool_kv, g->pool_score, g->previous_kv[owner], g->previous_score[owner], 512, 1, pos)) return false; } else if (!ds4_gpu_tensor_copy(g->latent, 0, g->pool_kv, 0, 512u * 4u)) return false; if ((pos + 1u) % ratio == 0) { const uint32_t rope_pos = pos + 1u - ratio; if (!ds41_norm(g->latent, g->latent, m, l->attn_compressor_norm) || !ds41_matmul(g->index_k, m, l->indexer_attn_k, g->latent, true) || !ds41_norm(g->index_k, g->index_k, m, l->indexer_k_norm) || !ds41_rope(g->index_k, 1, DS4_N_INDEXER_HEAD_DIM, il, rope_pos, false) || !ds4_gpu_dsv41_quantize(g->index_k, 128, 1, DS4_V41_FP4_E8M0) || !ds4_gpu_tensor_copy(g->index_cache[owner], (uint64_t)(n_comp - 1u) * 128u * 4u, g->index_k, 0, 128u * 4u) || !ds41_rope(g->latent, 1, 512, il, rope_pos, false) || !ds4_gpu_dsv41_quantize(g->latent, 512, 1, DS4_V41_FP4_E4M3) || !ds4_gpu_tensor_copy(g->compressed[owner], (uint64_t)(n_comp - 1u) * 512u * 4u, g->latent, 0, 512u * 4u)) return false; } } return true; } static bool ds41_attention_candidates(ds41_gpu_graph *g, uint32_t il) { const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; if (n_comp && ds41_index_source(il)) { if (il == 20) { const uint32_t blocks = (n_comp + 7u) / 8u, top = blocks < 2048u ? blocks : 2048u; if (!ds4_gpu_dsv41_candidate_blocks(g->block_scores, g->index_scores, n_comp, 1, pos, ratio) || !ds4_gpu_indexer_topk_tensor(g->block_selected, g->block_scores, blocks, 1, top) || !ds4_gpu_dsv4_topk_mask_tensor(g->block_mask, g->block_selected, blocks, 1, top)) return false; } else if (il > 20 && !ds4_gpu_dsv41_candidate_filter(g->index_scores, g->block_mask, n_comp, 1, pos, ratio)) return false; } return true; } static bool ds41_attention_pick(ds41_gpu_graph *g, uint32_t il) { const uint32_t ratio = ds4_layer_compress_ratio(il); const uint32_t n_comp = ratio ? (g->pos + 1u) / ratio : 0; const uint32_t top = n_comp < DS4_N_INDEXER_TOP_K ? n_comp : DS4_N_INDEXER_TOP_K; return ds41_attention_candidates(g, il) && (!n_comp || !ds41_index_source(il) || ds4_gpu_indexer_topk_tensor(g->selected_comp, g->index_scores, n_comp, 1, top)); } static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; if (n_comp && ds41_index_source(il)) { if (!ds41_matmul(g->index_q, m, l->indexer_attn_q_b, g->qr, true) || !ds41_rope(g->index_q, DS4_N_INDEXER_HEAD, 128, il, pos, false) || !ds4_gpu_dsv41_quantize(g->index_q, 128, DS4_N_INDEXER_HEAD, DS4_V41_FP4_E8M0) || !ds41_matmul(g->index_weights, m, l->indexer_proj, g->norm, true) || #ifdef __APPLE__ !ds4_gpu_glm_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner], n_comp, DS4_N_INDEXER_HEAD, 128, 1.0f / 64.0f, false)) return false; #else !ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner], n_comp, 1, pos, ratio)) return false; #endif } return ds41_attention_pick(g, il); } static bool ds41_attention_select(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { return ds41_attention_publish(g, m, l, il) && ds41_attention_select_published(g, m, l, il); } static bool ds41_attention_project(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { const uint32_t q_dim = DS4_N_HEAD / g->tp_world * DS4_N_HEAD_DIM; return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && ds41_norm(g->qr, g->qr, m, l->attn_q_a_norm) && ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, g->tp_rank * q_dim, q_dim) && ds41_matmul(g->kv, m, l->attn_kv, g->norm, true) && ds41_norm(g->kv, g->kv, m, l->attn_kv_a_norm); } static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, bool projected) { const uint32_t pos = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; const uint32_t heads = DS4_N_HEAD / g->tp_world; const uint32_t head0 = g->tp_rank * heads; if (!projected && !ds41_attention_project(g, m, l)) return false; if (!ds41_rope(g->q, heads, DS4_N_HEAD_DIM, il, pos, false) || !ds41_rope(g->kv, 1, DS4_N_HEAD_DIM, il, pos, false) || !ds4_gpu_dsv41_quantize(g->kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0) || !ds4_gpu_tensor_copy(g->window[il], (uint64_t)(pos % 128u) * 512u * 4u, g->kv, 0, 512u * 4u) || !ds41_attention_select(g, m, l, il)) return false; const uint32_t attended = n_comp < DS4_N_INDEXER_TOP_K ? n_comp : DS4_N_INDEXER_TOP_K; if (n_comp && !ds4_gpu_dsv41_gather_kv(g->selected_kv, g->compressed[owner], g->selected_comp, n_comp, attended)) return false; const uint32_t n_raw = pos + 1u < 128u ? pos + 1u : 128u; if (!ds4_gpu_attention_decode_heads_tensor(g->heads, m->map, m->size, l->attn_sinks->abs_offset + (uint64_t)head0 * sizeof(float), g->q, g->window[il], n_raw, 128, (pos + 1u - n_raw) % 128u, g->selected_kv, 0, attended, NULL, 0, heads, DS4_N_HEAD_DIM) || !ds41_bf16(g->heads, heads * DS4_N_HEAD_DIM) || !ds41_rope(g->heads, heads, DS4_N_HEAD_DIM, il, pos, true)) return false; if (projected) return true; return ds41_attention_output(g, m, l) && ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && ds41_bf16(g->block, DS4_N_EMBD); } static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { uint64_t gate_row = 0, down_row = 0; if (!tensor_nbytes(l->ffn_gate_exps->type, DS4_N_EMBD, &gate_row) || !tensor_nbytes(l->ffn_down_exps->type, DS4_N_FF_EXP, &down_row)) return false; const bool shared_owner = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER"); ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; const ds4_tensor *bias = ds41_image_at(g, g->pos) ? l->ffn_exp_probs_vl : l->ffn_exp_probs_b; if (!bias) return false; if (!ds41_matmul(g->route_logits, m, l->ffn_gate_inp, g->norm, false) || !ds4_gpu_router_select_tensor(g->selected, g->route_weights, g->route_probs, m->map, m->size, bias->abs_offset, 0, 0, token, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, 0, 0, true, false, g->route_logits)) return false; const bool shared_here = !shared_owner || g->tp_rank == (il & 1u); bool shared_queued = false; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (shared_owner && shared_here && 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) { const int rc = ds4_gpu_dsv41_shared_start(g->shared, g->shared_gate, g->shared_up, g->shared_mid, g->norm, m->map, m->size, l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, l->ffn_down_shexp->abs_offset, DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP); if (rc < 0) return false; shared_queued = rc > 0; } #endif if (shared_here && !shared_queued && (!ds41_matmul(g->shared_gate, m, l->ffn_gate_shexp, g->norm, true) || !ds41_matmul(g->shared_up, m, l->ffn_up_shexp, g->norm, true) || !ds4_gpu_swiglu_tensor(g->shared_mid, g->shared_gate, g->shared_up, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP, 1.0f) || !ds41_bf16(g->shared_mid, DS4_N_FF_EXP) || !ds41_matmul(g->shared, m, l->ffn_down_shexp, g->shared_mid, true))) return false; bool routed_ok; #ifndef __APPLE__ if (g->tp_world == 2) { routed_ok = ds4_gpu_routed_moe_batch_owned_tensor(routed, g->gate, g->up, g->mid, g->experts, m->map, m->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_row * DS4_N_FF_EXP, gate_row, down_row * DS4_N_EMBD, down_row, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, g->selected, g->route_weights, DS4_N_EXPERT, DS4_N_EXPERT_USED, g->tp_rank * (DS4_N_EXPERT / 2u), DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, g->norm, il, 1, NULL); } else #endif routed_ok = ds4_gpu_routed_moe_one_tensor(routed, g->gate, g->up, g->mid, g->experts, m->map, m->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_row * DS4_N_FF_EXP, gate_row, down_row * DS4_N_EMBD, down_row, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, g->selected, g->route_weights, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->norm, NULL, il, !g->streaming); #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (shared_queued && !ds4_gpu_dsv41_shared_join()) return false; #endif if (!routed_ok) return false; /* Keep the shared expert's BF16 boundary, then include it exactly once * in the existing F32 reduction. Alternate ownership across layers. */ if (shared_owner && g->tp_rank == (il & 1u) && !ds4_gpu_add_tensor(routed, routed, g->shared, DS4_N_EMBD)) return false; return true; } static bool ds41_moe_finish(ds41_gpu_graph *g, uint32_t il) { const bool shared_owner = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER"); ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; if (!ds41_sum_partial(g, routed, il, DS4_TP_GATE_FFN)) return false; return (shared_owner || ds4_gpu_add_tensor(g->block, routed, g->shared, DS4_N_EMBD)) && ds41_bf16(g->block, DS4_N_EMBD); } static bool ds41_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { return ds41_moe_partial(g, m, l, il, token) && ds41_moe_finish(g, il); } static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, float *logits) { if (!g->valid || !logits || !ds4_gpu_begin_commands()) return false; bool ok = ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, w->output_norm) && ds41_output_projection(g, g->tp_logits_half ? g->tp_logits_half : g->logits, m, w, g->norm, 1); if (!ds4_gpu_end_commands()) ok = false; return ok && ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)); } static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { if (ds41_engram_layer(il) && !ds41_image_at(g, g->pos)) { const uint32_t i = il == 1 ? 0 : 1; if (!ds41_matmul(g->engram_kv, m, l->engram_kv, g->engram_rows, true) || !ds4_gpu_dsv41_engram_add(g->residual, g->engram_kv, g->engram_q_norm[i], g->engram_k_norm[i], NULL, DS4_N_EMBD, 1, DS4_RMS_EPS)) return false; } return ds41_hc_mix(g, m, l, false) && ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, l->attn_norm); } static bool ds41_graph_after_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { return ds4_gpu_hc_expand_split_tensor(g->after_attn, g->block, g->residual, g->attn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->after_attn, DS4_N_EMBD * DS4_N_HC) && ds41_hc_mix(g, m, l, true) && ds4_gpu_hc_weighted_sum_split_tensor(g->x, g->after_attn, g->attn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, l->ffn_norm); } static bool ds41_graph_before_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { return ds41_graph_before_attention(g, m, l, il) && ds41_attention(g, m, l, il, false) && ds41_graph_after_attention(g, m, l); } static bool ds41_norm_batch(ds4_gpu_tensor *out, const ds4_gpu_tensor *in, const ds4_model *m, const ds4_tensor *weight, uint32_t count) { return ds4_gpu_rms_norm_weight_rows_tensor(out, in, m->map, m->size, weight->abs_offset, (uint32_t)weight->dim[0], count, DS4_RMS_EPS) && ds4_gpu_dsv41_quantize(out, (uint32_t)weight->dim[0], count, DS4_V41_BF16); } static bool ds41_hc_mix_batch(ds41_prefill_row *b, const ds4_model *m, const ds4_layer_weights *l, bool ffn, uint32_t count) { const ds4_tensor *fn = ffn ? l->hc_ffn_fn : l->hc_attn_fn; const ds4_tensor *scale = ffn ? l->hc_ffn_scale : l->hc_attn_scale; const ds4_tensor *base = ffn ? l->hc_ffn_base : l->hc_attn_base; const ds4_gpu_tensor *input = ffn ? b->after_attn : b->residual; const bool projected = #ifdef __APPLE__ fn->type == DS4_TENSOR_F16 && count > DS4_TP_BATCH_MAX_ROWS ? ds4_gpu_hc_rms_scale_project_f16_tensor(b->mix, b->flat_norm, m->map, m->size, fn->abs_offset, DS4_N_HC * DS4_N_EMBD, 24u, input, count, DS4_RMS_EPS) : #endif ds4_gpu_rms_norm_plain_rows_tensor(b->flat_norm, input, DS4_N_HC * DS4_N_EMBD, count, DS4_RMS_EPS) && ds41_matmul_batch(b->mix, m, fn, b->flat_norm, count, false); return projected && ds4_gpu_hc_split_sinkhorn_tensor(ffn ? b->ffn_split : b->attn_split, b->mix, m->map, m->size, scale->abs_offset, base->abs_offset, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS); } static bool ds41_before_attention_batch(ds41_gpu_graph *g, ds41_prefill_row *b, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t count) { if (ds41_engram_layer(il)) { const uint32_t i = il == 1 ? 0 : 1; if (!ds41_matmul_batch(b->engram_kv, m, l->engram_kv, b->engram_rows, count, true) || !ds4_gpu_dsv41_engram_add(b->residual, b->engram_kv, g->engram_q_norm[i], g->engram_k_norm[i], g->image_count ? g->image_text_mask : NULL, DS4_N_EMBD, count, DS4_RMS_EPS)) return false; } if (!ds41_hc_mix_batch(b, m, l, false, count)) return false; /* V4.1 consumes the preceding sublayer's mixer, not the newly computed one. */ const bool mixed = il ? ds4_gpu_hc_weighted_sum_split_tensor(b->x, b->residual, b->ffn_split, DS4_N_EMBD, DS4_N_HC) : ds4_gpu_hc_weighted_sum_tensor(b->x, b->residual, b->pre, DS4_N_EMBD, DS4_N_HC); return mixed && ds4_gpu_dsv41_quantize(b->x, DS4_N_EMBD, count, DS4_V41_BF16) && ds41_norm_batch(b->norm, b->x, m, l->attn_norm, count); } static bool ds41_after_attention_batch(ds41_prefill_row *b, const ds4_model *m, const ds4_layer_weights *l, uint32_t count) { return ds4_gpu_hc_expand_split_tensor(b->after_attn, b->block, b->residual, b->attn_split, DS4_N_EMBD, DS4_N_HC) && ds4_gpu_dsv41_quantize(b->after_attn, DS4_N_EMBD * DS4_N_HC, count, DS4_V41_BF16) && ds41_hc_mix_batch(b, m, l, true, count) && ds4_gpu_hc_weighted_sum_split_tensor(b->x, b->after_attn, b->attn_split, DS4_N_EMBD, DS4_N_HC) && ds4_gpu_dsv41_quantize(b->x, DS4_N_EMBD, count, DS4_V41_BF16) && ds41_norm_batch(b->norm, b->x, m, l->ffn_norm, count); } static bool ds41_attention_project_batch(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t count) { ds41_prefill_row *b = &g->batch; const uint32_t q_dim = DS4_N_HEAD / g->tp_world * DS4_N_HEAD_DIM; return ds41_matmul_batch(b->qr, m, l->attn_q_a, b->norm, count, true) && ds4_gpu_rms_norm_weight_rows_tensor(b->qr, b->qr, m->map, m->size, l->attn_q_a_norm->abs_offset, DS4_N_LORA_Q, count, DS4_RMS_EPS) && ds4_gpu_dsv41_quantize(b->qr, DS4_N_LORA_Q, count, DS4_V41_BF16) && ds41_matmul_rows_batch(b->q, m, l->attn_q_b, b->qr, g->tp_rank * q_dim, q_dim, count) && ds41_matmul_batch(b->kv, m, l->attn_kv, b->norm, count, true) && ds4_gpu_rms_norm_weight_rows_tensor(b->kv, b->kv, m->map, m->size, l->attn_kv_a_norm->abs_offset, DS4_N_HEAD_DIM, count, DS4_RMS_EPS) && ds4_gpu_dsv41_quantize(b->kv, DS4_N_HEAD_DIM, count, DS4_V41_BF16); } static bool ds41_project_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t count, bool bf16) { if (weight->type == DS4_TENSOR_F16) return ds4_gpu_dsv41_projection_rows(out, m->map, m->size, weight->abs_offset, (uint32_t)weight->dim[0], (uint32_t)weight->dim[1], count, in) && (!bf16 || ds4_gpu_dsv41_quantize(out, (uint32_t)weight->dim[1], count, DS4_V41_BF16)); return ds41_matmul_batch(out, m, weight, in, count, bf16); } static bool ds41_attention_publish_batch(ds41_gpu_graph *g, ds41_prefill_row *b, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t start, uint32_t count) { const uint32_t ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t first = start / ratio, rows = (start + count) / ratio - first; if (ratio == 2u) { if (!ds41_project_rows(b->pool_kv, m, l->attn_compressor_kv, b->norm, count, false) || !ds41_project_rows(b->pool_score, m, l->attn_compressor_gate, b->norm, count, false) || !ds4_gpu_dsv41_pool2(b->latent, b->pool_kv, b->pool_score, g->previous_kv[owner], g->previous_score[owner], DS4_N_HEAD_DIM, count, start)) return false; } else if (!ds41_project_rows(b->latent, m, l->attn_compressor_kv, b->norm, count, true)) return false; if (!rows) return true; /* Bulk publication and suffix replay must produce the same keys. Matrix * reductions here otherwise make old cache rows depend on sweep size. */ return ds41_norm_batch(b->latent, b->latent, m, l->attn_compressor_norm, rows) && ds41_project_rows(b->index_k, m, l->indexer_attn_k, b->latent, rows, true) && ds41_norm_batch(b->index_k, b->index_k, m, l->indexer_k_norm, rows) && ds4_gpu_dsv41_rope_stride(b->index_k, DS4_N_INDEXER_HEAD_DIM, 1, rows, first * ratio, ratio, true, false) && ds4_gpu_dsv41_quantize(b->index_k, DS4_N_INDEXER_HEAD_DIM, rows, DS4_V41_FP4_E8M0) && ds4_gpu_tensor_copy(g->index_cache[owner], (uint64_t)first * DS4_N_INDEXER_HEAD_DIM * sizeof(float), b->index_k, 0, (uint64_t)rows * DS4_N_INDEXER_HEAD_DIM * sizeof(float)) && ds4_gpu_dsv41_rope_stride(b->latent, DS4_N_HEAD_DIM, 1, rows, first * ratio, ratio, true, false) && ds4_gpu_dsv41_quantize(b->latent, DS4_N_HEAD_DIM, rows, DS4_V41_FP4_E4M3) && ds4_gpu_tensor_copy(g->compressed[owner], (uint64_t)first * DS4_N_HEAD_DIM * sizeof(float), b->latent, 0, (uint64_t)rows * DS4_N_HEAD_DIM * sizeof(float)); } static bool ds41_index_batch(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t count) { const uint32_t start = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t n_comp = (start + count) / ratio; ds41_prefill_row *b = &g->batch; if (!n_comp) return true; if (getenv("DS4_METAL_DISABLE_V41_BATCH_INDEX_PROJ")) { for (uint32_t t = 0; t < count; t++) { ds41_prefill_row *r = &g->rows_view[t]; if (!ds41_matmul(r->index_q, m, l->indexer_attn_q_b, r->qr, true) || !ds41_rope(r->index_q, DS4_N_INDEXER_HEAD, 128, il, start + t, false) || !ds4_gpu_dsv41_quantize(r->index_q, 128, DS4_N_INDEXER_HEAD, DS4_V41_FP4_E8M0) || !ds41_matmul(r->index_weights, m, l->indexer_proj, r->norm, true)) return false; } } else if (!ds41_project_rows(b->index_q, m, l->indexer_attn_q_b, b->qr, count, true) || !ds4_gpu_dsv41_rope(b->index_q, DS4_N_INDEXER_HEAD_DIM, DS4_N_INDEXER_HEAD, count, start, true, false) || !ds4_gpu_dsv41_quantize(b->index_q, DS4_N_INDEXER_HEAD_DIM, count * DS4_N_INDEXER_HEAD, DS4_V41_FP4_E8M0) || !ds41_project_rows(b->index_weights, m, l->indexer_proj, b->norm, count, true)) return false; ds41_gpu_graph row = *g; /* Only 32 score rows are live. Batched sorting retains the original * per-row causal width, including padding and tie order. */ const bool packed = count >= 8u && ds4_gpu_dsv41_tensor_ops_available() && !getenv("DS4_METAL_DISABLE_V41_PACKED_INDEX"); if (packed && !ds4_gpu_dsv41_indexer_pack(g->index_packed, b->index_q, g->index_cache[owner], n_comp, count)) return false; for (uint32_t off = 0; off < count; off += DS41_INDEX_BATCH) { const uint32_t rows = count - off < DS41_INDEX_BATCH ? count - off : DS41_INDEX_BATCH; const uint64_t q_bytes = DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float); const uint64_t w_bytes = DS4_N_INDEXER_HEAD * sizeof(float); ds4_gpu_tensor *q = ds4_gpu_tensor_view(b->index_q, off * q_bytes, rows * q_bytes); ds4_gpu_tensor *weights = ds4_gpu_tensor_view(b->index_weights, off * w_bytes, rows * w_bytes); bool ok = q && weights && (packed ? ds4_gpu_dsv41_indexer_scores_packed(g->index_scores, q, weights, g->index_cache[owner], g->index_packed, n_comp, rows, start + off, ratio, count, off) : ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, q, weights, g->index_cache[owner], n_comp, rows, start + off, ratio)); ds4_gpu_tensor_free(weights); ds4_gpu_tensor_free(q); const bool batch_topk = (start + off + 1u) / ratio >= 1024u && !getenv("DS4_METAL_DISABLE_V41_BATCH_TOPK"); for (uint32_t t = 0; ok && (!batch_topk || il >= 20u) && t < rows; t++) { row.pos = start + off + t; row.selected_comp = g->rows_view[off + t].selected_comp; row.block_mask = g->rows_view[off + t].block_mask; row.index_scores = ds4_gpu_tensor_view(g->index_scores, (uint64_t)t * n_comp * sizeof(float), (uint64_t)n_comp * sizeof(float)); ok = row.index_scores && (batch_topk ? ds41_attention_candidates(&row, il) : ds41_attention_pick(&row, il)); ds4_gpu_tensor_free(row.index_scores); } if (ok && batch_topk) { const uint64_t id_bytes = DS4_N_INDEXER_TOP_K * sizeof(int32_t); ds4_gpu_tensor *ids = ds4_gpu_tensor_view(b->selected_comp, off * id_bytes, rows * id_bytes); ok = ids && ds4_gpu_dsv41_indexer_topk_batch(ids, g->index_scores, n_comp, rows, start + off, ratio); ds4_gpu_tensor_free(ids); } if (!ok) return false; } return true; } static bool ds41_attention_batch(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t count) { const uint32_t start = g->pos, ratio = ds4_layer_compress_ratio(il); const uint32_t owner = il < 8 ? 0u : il < 14 ? 1u : il < 20 ? 2u : 3u; const uint32_t n_comp = ratio ? (start + count) / ratio : 0; const uint32_t previous = start < 127u ? start : 127u; const uint32_t raw_start = (start - previous) % 128u; const uint32_t first = previous < 128u - raw_start ? previous : 128u - raw_start; const uint32_t n_raw = previous + count; const uint64_t row_bytes = DS4_N_HEAD_DIM * sizeof(float); const uint32_t heads = DS4_N_HEAD / g->tp_world; const uint64_t sinks = l->attn_sinks->abs_offset + (uint64_t)g->tp_rank * heads * sizeof(float); ds41_prefill_row *b = &g->batch; if (!ds4_gpu_dsv41_rope(b->q, DS4_N_HEAD_DIM, heads, count, start, ratio != 0, false) || !ds4_gpu_dsv41_rope(b->kv, DS4_N_HEAD_DIM, 1, count, start, ratio != 0, false) || !ds4_gpu_dsv41_quantize(b->kv, DS4_N_HEAD_DIM, count, DS4_V41_FP8_E8M0)) return false; /* Preserve chronological raw keys before publishing the final decode ring. */ if (first && !ds4_gpu_tensor_copy(g->raw_prefill, 0, g->window[il], raw_start * row_bytes, first * row_bytes)) return false; if (previous > first && !ds4_gpu_tensor_copy(g->raw_prefill, first * row_bytes, g->window[il], 0, (previous - first) * row_bytes)) return false; if (!ds4_gpu_tensor_copy(g->raw_prefill, previous * row_bytes, b->kv, 0, count * row_bytes)) return false; ds41_gpu_graph row = *g; const bool batch_index = ds41_index_source(il) && !getenv("DS4_METAL_DISABLE_V41_BATCH_INDEX"); const bool batch_publish = ds41_kv_source(il) && #ifdef __APPLE__ ratio == 2u && #endif !getenv("DS4_METAL_DISABLE_V41_BATCH_COMPRESS"); if (batch_publish && !ds41_attention_publish_batch(g, b, m, l, il, start, count)) return false; for (uint32_t t = 0; (!batch_index || !batch_publish) && t < count; t++) { row.pos = start + t; #define DS41_SELECT_ROW(name, width) row.name = g->rows_view[t].name; DS41_PREFILL_ROWS(DS41_SELECT_ROW) #undef DS41_SELECT_ROW if (!batch_publish && !ds41_attention_publish(&row, m, l, il)) return false; if (!batch_index && !ds41_attention_select_published(&row, m, l, il)) return false; } if (batch_index && !ds41_index_batch(g, m, l, il, count)) return false; bool ok; if (!n_comp) { /* Raw attention rounds probabilities per key block. Keep its existing * partitions even when the surrounding dense/FFN batches are wider. */ ok = true; const uint64_t q_row = (uint64_t)heads * DS4_N_HEAD_DIM * sizeof(float); for (uint32_t off = 0; ok && off < count; off += 2048u) { const uint32_t rows = count - off < 2048u ? count - off : 2048u; const uint32_t past = start + off < 127u ? start + off : 127u; const uint32_t raw_offset = previous + off - past; ds4_gpu_tensor *q = ds4_gpu_tensor_view(b->q, off * q_row, rows * q_row); ds4_gpu_tensor *out = ds4_gpu_tensor_view(b->heads, off * q_row, rows * q_row); ds4_gpu_tensor *raw = ds4_gpu_tensor_view(g->raw_prefill, raw_offset * row_bytes, (past + rows) * row_bytes); ok = q && out && raw && ds4_gpu_attention_decode_raw_batch_heads_tensor(out, m->map, m->size, sinks, q, raw, rows, start + off, past + rows, past + rows, 0, 128, heads, DS4_N_HEAD_DIM); ds4_gpu_tensor_free(raw); ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(q); } } else if (n_comp < DS4_N_INDEXER_TOP_K) { ok = ds4_gpu_attention_decode_mixed_batch_heads_tensor(b->heads, m->map, m->size, sinks, b->q, g->raw_prefill, g->compressed[owner], 0, NULL, 0, count, start, n_raw, g->prefill_cap + 128u, 0, n_comp, 128, ratio, heads, DS4_N_HEAD_DIM); } else { ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(b->heads, m->map, m->size, sinks, b->q, g->raw_prefill, g->compressed[owner], 0, b->selected_comp, count, start, n_raw, g->prefill_cap + 128u, 0, n_comp, DS4_N_INDEXER_TOP_K, 128, ratio, heads, DS4_N_HEAD_DIM); } if (!ok || !ds4_gpu_dsv41_quantize(b->heads, heads * DS4_N_HEAD_DIM, count, DS4_V41_BF16) || !ds4_gpu_dsv41_rope(b->heads, DS4_N_HEAD_DIM, heads, count, start, ratio != 0, true)) return false; const uint32_t kept = n_raw < 128u ? n_raw : 128u; const uint32_t slot = (start + count - kept) % 128u; const uint32_t part = kept < 128u - slot ? kept : 128u - slot; return ds4_gpu_tensor_copy(g->window[il], slot * row_bytes, g->raw_prefill, (n_raw - kept) * row_bytes, part * row_bytes) && (kept == part || ds4_gpu_tensor_copy(g->window[il], 0, g->raw_prefill, (n_raw - kept + part) * row_bytes, (kept - part) * row_bytes)); } static bool ds41_graph_after_moe(ds41_gpu_graph *g) { return ds4_gpu_hc_expand_split_tensor(g->residual, g->block, g->after_attn, g->ffn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->residual, DS4_N_EMBD * DS4_N_HC) && ds4_gpu_tensor_copy(g->pre, 0, g->ffn_split, 0, DS4_N_HC * sizeof(float)); } static bool ds41_graph_layer(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, int token) { return ds41_graph_before_moe(g, m, l, il) && ds41_moe(g, m, l, il, (uint32_t)token) && ds41_graph_after_moe(g); } #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) /* Capture only single-token, resident TP work with stable pointers. Engram * rows are refreshed before entry; positional attention and gates stay eager. * V4.1's router does not use the token argument when hash_mode is false. */ static bool ds41_decode_island(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t island) { const ds4_decode_graph_key key = { .il = il, .island = island, .variant = 41, .cur_hc = g->residual, .after_attn_hc = g->after_attn, .after_ffn_hc = g->norm, .attn_norm = (void *)m->map, }; for (;;) { const int state = ds4_gpu_decode_graph_begin(&key); if (state == 1) return true; const bool ok = island == 0 ? ds41_graph_before_attention(g, m, l, il) && ds41_attention_project(g, m, l) : island == 2 ? ds41_attention_output(g, m, l) : ds41_graph_after_attention(g, m, l) && ds41_moe_partial(g, m, l, il, 0); if (state != 0) return ok; if (!ok) ds4_gpu_decode_graph_abort(&key); else if (ds4_gpu_decode_graph_end(&key) == 0) return true; /* A failed capture executes no work and retires the entry. */ } } static bool ds41_graph_decode_layer(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, int token) { if (g->tp_world != 2 || g->streaming || g->quality || g->imatrix || g->image_count || getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER") || g_expert_profile.active || getenv("DS4_CUDA_MOE_PROFILE") || metal_graph_debug_get_config()->prefix || !ds4_gpu_decode_graphs_supported()) return ds41_graph_layer(g, m, l, il, token); return ds41_decode_island(g, m, l, il, 0) && ds41_attention(g, m, l, il, true) && ds41_decode_island(g, m, l, il, 2) && ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && ds41_bf16(g->block, DS4_N_EMBD) && ds41_decode_island(g, m, l, il, 1) && ds41_moe_finish(g, il) && ds41_graph_after_moe(g); } #endif static bool ds41_route_batch(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t count) { ds41_prefill_row *b = &g->batch; if (!g->image_count) return ds4_gpu_router_select_batch_tensor( b->selected, b->route_weights, b->route_probs, m->map, m->size, l->ffn_exp_probs_b->abs_offset, 0, 0, 0, 0, true, false, b->route_logits, g->prefill_tokens, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, count); /* The ordinary 384-expert selector already has the right scoring and * normalization. Split only where the text/visual selection bias changes. */ for (uint32_t first = 0; first < count;) { const bool image = ds41_image_at(g, g->pos + first) != NULL; uint32_t end = first + 1; while (end < count && (ds41_image_at(g, g->pos + end) != NULL) == image) end++; const ds4_tensor *bias = image ? l->ffn_exp_probs_vl : l->ffn_exp_probs_b; ds4_gpu_tensor *views[5] = {0}; ds4_gpu_tensor *buffers[] = {b->selected, b->route_weights, b->route_probs, b->route_logits, g->prefill_tokens}; const uint32_t width[] = {DS4_N_EXPERT_USED, DS4_N_EXPERT_USED, DS4_N_EXPERT, DS4_N_EXPERT, 1}; bool ok = bias != NULL; for (unsigned i = 0; ok && i < 5; i++) { views[i] = ds4_gpu_tensor_view(buffers[i], (uint64_t)first * width[i] * 4u, (uint64_t)(end - first) * width[i] * 4u); ok = views[i] != NULL; } if (ok) ok = ds4_gpu_router_select_batch_tensor(views[0], views[1], views[2], m->map, m->size, bias->abs_offset, 0, 0, 0, 0, true, false, views[3], views[4], DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, end - first); for (unsigned i = 0; i < 5; i++) ds4_gpu_tensor_free(views[i]); if (!ok) return false; first = end; } return true; } static bool ds41_moe_batch(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t count, bool shared_owner) { ds41_prefill_row *b = &g->batch; const uint64_t gate_row = routed_expert_row_bytes(l->ffn_gate_exps); const uint64_t down_row = routed_expert_row_bytes(l->ffn_down_exps); bool mid_f16 = false; return ds41_matmul_batch(b->route_logits, m, l->ffn_gate_inp, b->norm, count, false) && ds41_route_batch(g, m, l, count) && ((shared_owner && g->tp_rank != (il & 1u)) || (ds41_matmul_batch(b->shared_gate, m, l->ffn_gate_shexp, b->norm, count, true) && ds41_matmul_batch(b->shared_up, m, l->ffn_up_shexp, b->norm, count, true) && ds4_gpu_swiglu_tensor(b->shared_mid, b->shared_gate, b->shared_up, count * DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP, 1.0f) && ds4_gpu_dsv41_quantize(b->shared_mid, DS4_N_FF_EXP, count, DS4_V41_BF16) && ds41_matmul_batch(b->shared, m, l->ffn_down_shexp, b->shared_mid, count, true))) && ( #ifndef __APPLE__ g->tp_world == 2 ? ds4_gpu_routed_moe_batch_owned_tensor(b->routed, b->gate, b->up, b->mid, b->experts, m->map, m->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_row * DS4_N_FF_EXP, gate_row, down_row * DS4_N_EMBD, down_row, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, b->selected, b->route_weights, DS4_N_EXPERT, DS4_N_EXPERT_USED, g->tp_rank * (DS4_N_EXPERT / 2u), DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, b->norm, il, count, &mid_f16) : #endif ds4_gpu_routed_moe_batch_tensor(b->routed, b->gate, b->up, b->mid, b->experts, m->map, m->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_row * DS4_N_FF_EXP, gate_row, down_row * DS4_N_EMBD, down_row, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, b->selected, b->route_weights, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, b->norm, il, count, &mid_f16, true)) && (!shared_owner || g->tp_rank != (il & 1u) || ds4_gpu_add_tensor(b->routed, b->routed, b->shared, count * DS4_N_EMBD)) && ds41_sum_partial_batch(g, b->routed, il, count); } static DS4_MAYBE_UNUSED bool ds41_graph_step(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, int token, float *logits) { if (!g || !g->valid || g->pos >= g->ctx || token < 0 || (uint32_t)token >= DS4_N_VOCAB) return false; uint32_t ids[2][DS4_ENGRAM_COLS]; ds4_engram_history next_history = g->history; if (!ds41_hash_tokens(g, &next_history, &token, 1, &ids[0][0])) return false; for (uint32_t i = 0; !ds41_image_at(g, g->pos) && i < 2; i++) { #ifdef __APPLE__ if (!ds4_engram_read_batch(&g->table[i], ids[i], 1, DS4_ENGRAM_COLS, g->rows[i])) return false; #else if (!ds4_engram_read(&g->table[i], ids[i], DS4_ENGRAM_COLS, g->rows[i])) return false; #endif } const float initial_pre[] = {1, 0, 0, 0}; if (!ds4_gpu_tensor_write(g->pre, 0, initial_pre, sizeof(initial_pre)) || !ds4_gpu_begin_commands()) return false; bool ok = ds41_embed(g, m, w, g->residual, g->x, token, g->pos); /* Unfused quality kernels bind whole expert tensors. Keep just the current * layer mapped, using the same admitted reserve as layer-major prefill. */ const bool layer_resident = g->streaming && g->quality; if (layer_resident && !ds4_gpu_end_commands()) ok = false; const bool queue_layers = g->tp_world == 2 && !g->imatrix && !getenv("DS4_METAL_DISABLE_V41_TP_DECODE_QUEUE"); for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &w->layer[il]; if (layer_resident) ok = metal_graph_stream_map_layer(m, w, il) && ds4_gpu_begin_commands(); if (ok && ds41_engram_layer(il)) { const uint32_t i = il == 1 ? 0 : 1; ok = ds4_gpu_tensor_write(g->engram_rows, 0, g->rows[i], sizeof(g->rows[i])); } if (ok) { #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) ok = ds41_graph_decode_layer(g, m, l, il, token); #else ok = ds41_graph_layer(g, m, l, il, token); #endif } /* TP gates already submit ordered, bounded command buffers. Drain * before overwriting the first Engram table's shared input at layer * 14, and before publishing the completed token to the CPU. */ const bool drain = !queue_layers || il == 13 || il + 1u == DS4_N_LAYER; if (drain && !ds4_gpu_end_commands()) ok = false; if (g->tp_world == 2 && ds4_gpu_tp_failed()) ok = false; if (ok && g->imatrix) ok = imatrix_collect_tensor_batch(g->imatrix, g->norm, g->mid, g->selected, false, il, 1); if (!ok) fprintf(stderr, "ds4: V4.1 layer %u failed at position %u\n", il, g->pos); if (ok && drain && !layer_resident && il + 1u < DS4_N_LAYER) ok = ds4_gpu_begin_commands() != 0; } if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false; if (layer_resident && !metal_graph_stream_map_decode_static_all(m, w)) ok = false; if (g->tp_world == 2 && ds4_gpu_tp_failed()) ok = false; if (ok && logits) ok = ds41_graph_logits(g, m, w, logits); if (!ok) { g->valid = false; return false; } g->history = next_history; g->pos++; return true; } /* Seed from the current mapped layer, avoiding a second disk pass. Recent * routes take precedence over the rest of the prompt's popular experts. */ static bool ds41_prefill_seed(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t count) { #if defined(__APPLE__) && !defined(DS4_NO_GPU) if (!g->streaming || getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_CACHE_SEED")) return true; uint32_t target = ds4_gpu_stream_expert_cache_configured_count() / DS4_N_LAYER; if (target > DS4_N_EXPERT) target = DS4_N_EXPERT; if (!target) return true; const int32_t *selected = ds4_gpu_tensor_contents(g->batch.selected); if (!selected) return false; uint32_t frequency[DS4_MAX_EXPERT] = {0}; const uint32_t recent = count < 32u ? count : 32u; for (uint32_t t = 0; t < count; t++) { for (uint32_t j = 0; j < DS4_N_EXPERT_USED; j++) { const int32_t id = selected[t * DS4_N_EXPERT_USED + j]; if (id < 0 || (uint32_t)id >= DS4_N_EXPERT) return false; frequency[id] += 1u + (t >= count - recent ? count * DS4_N_EXPERT_USED : 0u); } } int32_t experts[DS4_MAX_EXPERT]; uint32_t priority[DS4_MAX_EXPERT], n = 0; while (n < target) { uint32_t best = 0; for (uint32_t i = 1; i < DS4_N_EXPERT; i++) if (frequency[i] > frequency[best]) best = i; if (!frequency[best]) break; experts[n] = (int32_t)best; priority[n] = metal_graph_streaming_builtin_hotness(target - n, target); n++; frequency[best] = 0; } if (!n) return true; const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(m, l, il, routed_expert_row_bytes(l->ffn_gate_exps) * DS4_N_FF_EXP, routed_expert_row_bytes(l->ffn_down_exps) * DS4_N_EMBD); if (!ds4_gpu_begin_commands()) return false; const bool ok = ds4_gpu_stream_expert_cache_seed_experts_gpu_copy(&table, experts, priority, n) != 0; const bool ended = ds4_gpu_end_commands() != 0; return ended && ok; #else (void)g; (void)m; (void)l; (void)il; (void)count; return true; #endif } static bool ds41_tp_batch_enabled(const ds41_gpu_graph *g) { return g->tp_world == 1 || (g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_BATCH_ATTN") && !getenv("DS4_METAL_DISABLE_V41_BATCH_CORE") && !getenv("DS4_METAL_DISABLE_V41_BATCH_MOE") && !getenv("DS4_METAL_DISABLE_V41_BATCH_HC")); } static uint32_t ds41_prefill_count(const ds41_gpu_graph *g, uint32_t remaining) { /* TP batches use the bulk protocol; row-gate ablations stay token-major. */ if (!ds41_tp_batch_enabled(g) || g->imatrix || getenv("DS4_METAL_DISABLE_V41_LAYER_PREFILL")) return 1; uint32_t minimum = 256u; #ifdef __APPLE__ minimum = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SMALL_PREFILL") ? 32u : 256u; #endif /* Resident appends do not pay for an SSD layer sweep. In particular, * the server's 128-token mixed quantum must not become scalar prefill. */ if (!g->streaming && g->tp_world == 1) minimum = 8u; #if defined(__APPLE__) && !defined(DS4_NO_GPU) /* With at least half the experts cached, warm short appends beat a full * disk sweep. Leave a token-major tail to warm the following decode too. */ if (g->streaming && g->pos && ds4_gpu_stream_expert_cache_configured_count() >= DS4_N_LAYER * DS4_N_EXPERT / 2u) minimum = 1024u; #endif if (remaining < minimum) return 1; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) /* Keep a medium SSD append in one layer sweep without changing its * 2048-row arithmetic partitions. Tiny tails retain the exact row path. */ if (g->streaming && g->tp_world == 1 && g->prefill_cap >= 2048u && g->carry_cap >= remaining && remaining > 2048u && remaining < 8192u && remaining % 2048u >= 256u && !getenv("DS4_CUDA_DISABLE_SSD_MEDIUM_SWEEP") && !getenv("DS4_METAL_DISABLE_V41_WIDE_PREFILL")) return remaining; #endif if (g->carry_cap && remaining >= 4096u && !getenv("DS4_METAL_DISABLE_V41_WIDE_PREFILL")) { const uint32_t count = remaining < g->carry_cap ? remaining : g->carry_cap; return count - count % 2048u; } const uint32_t tail_cap = g->prefill_cap < 2048u ? g->prefill_cap : 2048u; return remaining < tail_cap ? remaining : tail_cap; } typedef struct { uint32_t cache_budget; struct { uint64_t off, bytes, locked; } span[60]; } ds41_encoder_residency; static void ds41_encoder_release(ds41_gpu_graph *g, const ds4_model *m, ds41_encoder_residency *r) { if (!r->cache_budget) return; for (uint32_t i = 0; i < 60; i++) { if (!r->span[i].locked) continue; void *ptr = (void *)(m->map + r->span[i].off); if (munlock(ptr, (size_t)r->span[i].locked) != 0) ds4_die("cannot release V4.1 encoder pages safely before restoring the expert cache"); (void)madvise(ptr, (size_t)r->span[i].locked, MADV_DONTNEED); r->span[i].locked = 0; } g->encoder_resident = false; ds4_gpu_set_streaming_expert_cache_budget(r->cache_budget); r->cache_budget = 0; } /* Borrow the admitted expert-cache budget, never add a second resident copy. * Full encoder residency is useful only when several disk sweeps are avoided; * retain a token-major tail to rebuild the cache before generation. */ static void ds41_encoder_acquire(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, uint32_t remaining, ds41_encoder_residency *r, ds4_session_cancel_fn cancel, void *cancel_ud) { #ifndef __APPLE__ /* CUDA stages experts into its bounded device cache, not locked mmap pages. */ (void)g; (void)m; (void)w; (void)remaining; (void)r; (void)cancel; (void)cancel_ud; return; #else if (!g->streaming || g->tp_world != 1 || g->imatrix || remaining < 16384u || getenv("DS4_METAL_DISABLE_V41_LAYER_PREFILL") || getenv("DS4_METAL_DISABLE_V41_ENCODER_RESIDENCY")) return; /* A wide layer sweep reads each encoder layer only once, without pinning * a second long-lived set of experts or discarding the decode cache. */ if (g->carry_cap && !getenv("DS4_METAL_DISABLE_V41_WIDE_PREFILL")) return; const uint32_t cache = ds4_gpu_stream_expert_cache_configured_count(); uint64_t expert_bytes = 0, bytes = 0; if (!ds4_streaming_routed_expert_bytes(w, &expert_bytes)) return; const uint64_t page = (uint64_t)getpagesize(); for (uint32_t il = 0; il < 20; il++) { const ds4_layer_weights *l = &w->layer[il]; const ds4_tensor *tensor[] = {l->ffn_gate_exps, l->ffn_up_exps, l->ffn_down_exps}; for (uint32_t j = 0; j < 3; j++) { const ds4_tensor *t = tensor[j]; /* Exclusive whole pages avoid unlocking adjacent static weights. * The V4.1 converter aligns these tensors to 16 KiB. */ if (!t || t->abs_offset % page || t->bytes % page || t->abs_offset > m->size || t->bytes > m->size - t->abs_offset || t->bytes > UINT64_MAX - bytes) return; r->span[il * 3u + j].off = t->abs_offset; r->span[il * 3u + j].bytes = t->bytes; bytes += t->bytes; } } if (!expert_bytes || bytes > (uint64_t)cache * expert_bytes || !ds4_gpu_synchronize()) return; r->cache_budget = cache; ds4_gpu_set_streaming_expert_cache_budget(0); fprintf(stderr, "ds4: V4.1 prefill loading %.2f GiB encoder experts in place of the decode cache\n", (double)bytes / 1073741824.0); for (uint32_t i = 0; i < 60; i++) { while (r->span[i].locked < r->span[i].bytes) { uint64_t n = r->span[i].bytes - r->span[i].locked; if (n > (UINT64_C(256) << 20)) n = UINT64_C(256) << 20; if (cancel && cancel(cancel_ud)) { ds41_encoder_release(g, m, r); return; } if (mlock(m->map + r->span[i].off + r->span[i].locked, (size_t)n) != 0) { fprintf(stderr, "ds4: V4.1 encoder residency unavailable (%s); using two-layer prefill\n", strerror(errno)); ds41_encoder_release(g, m, r); return; } r->span[i].locked += n; } } g->encoder_resident = true; fprintf(stderr, "ds4: V4.1 prefill keeps %.2f GiB encoder experts resident in place of the decode cache\n", (double)bytes / 1073741824.0); #endif } static bool ds41_carry_copy(ds41_gpu_graph *g, uint32_t offset, uint32_t count, bool store) { if (offset > g->carry_cap || count > g->carry_cap - offset) return false; bool ok = true; #define DS41_CARRY_COPY(name, width, format) \ if (ok && g->compact_carry && (format) != DS4_V41_CARRY_F32) { \ ok = ds4_gpu_dsv41_carry_copy(g->carry.name, offset, g->batch.name, width, count, format, store); \ } else if (ok) { \ const uint64_t bytes = (uint64_t)(width) * sizeof(float); \ ok = store ? ds4_gpu_tensor_copy(g->carry.name, offset * bytes, g->batch.name, 0, count * bytes) : \ ds4_gpu_tensor_copy(g->batch.name, 0, g->carry.name, offset * bytes, count * bytes); \ } DS41_CARRY_ROWS(DS41_CARRY_COPY) #undef DS41_CARRY_COPY return ok; } /* Decoder attention reads all encoder compressed keys, but only the previous * 127 decoder inputs. Publish the former for the entire prefix; warm the latter * before evaluating each shrinking suffix. No old decoder output is needed. */ static bool ds41_decoder_prepare(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t initial_start, uint32_t offset, uint32_t rows, bool publish, bool split_pre, bool batch_attention, ds4_session_cancel_fn cancel, void *cancel_ud) { ds41_gpu_graph row = *g; bool ok = true; const bool batch_publish = publish && batch_attention && !getenv("DS4_METAL_DISABLE_V41_BATCH_PUBLISH"); const uint32_t chunk = g->prefill_cap < 2048u ? g->prefill_cap : 2048u; for (uint32_t off = offset; ok && off < offset + rows; off += chunk) { if (cancel && cancel(cancel_ud)) return false; const uint32_t count = offset + rows - off < chunk ? offset + rows - off : chunk; ds41_prefill_row active = {0}; #define DS41_DECODER_VIEW(name, width) \ if (ok) ok = (active.name = ds4_gpu_tensor_view(g->batch.name, 0, \ (uint64_t)count * (width) * sizeof(float))) != NULL; DS41_PREFILL_ROWS(DS41_DECODER_VIEW) #undef DS41_DECODER_VIEW if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = ds41_carry_copy(g, off, count, false); if (ok && batch_attention) ok = split_pre ? ds4_gpu_hc_weighted_sum_split_tensor(active.x, active.residual, active.ffn_split, DS4_N_EMBD, DS4_N_HC) : ds4_gpu_hc_weighted_sum_tensor(active.x, active.residual, active.pre, DS4_N_EMBD, DS4_N_HC); if (ok && batch_attention) ok = ds4_gpu_dsv41_quantize(active.x, DS4_N_EMBD, count, DS4_V41_BF16) && ds41_norm_batch(active.norm, active.x, m, l->attn_norm, count); if (ok && batch_publish) ok = ds41_attention_publish_batch(g, &active, m, l, il, initial_start + off, count); if (ok && !publish && batch_attention) ok = ds41_matmul_batch(active.kv, m, l->attn_kv, active.norm, count, true) && ds41_norm_batch(active.kv, active.kv, m, l->attn_kv_a_norm, count) && ds4_gpu_dsv41_rope(active.kv, DS4_N_HEAD_DIM, 1, count, initial_start + off, true, false) && ds4_gpu_dsv41_quantize(active.kv, DS4_N_HEAD_DIM, count, DS4_V41_FP8_E8M0); for (uint32_t t = 0; ok && !batch_publish && t < count; t++) { row.pos = initial_start + off + t; row.norm = g->rows_view[t].norm; row.kv = g->rows_view[t].kv; if (!batch_attention) { row.residual = g->rows_view[t].residual; row.pre = g->rows_view[t].pre; ok = ds41_graph_before_attention(&row, m, l, il); } if (ok && publish) { ok = ds41_attention_publish(&row, m, l, il); } else if (ok) { if (!batch_attention) ok = ds41_matmul(row.kv, m, l->attn_kv, row.norm, true) && ds41_norm(row.kv, row.kv, m, l->attn_kv_a_norm) && ds41_rope(row.kv, 1, DS4_N_HEAD_DIM, il, row.pos, false) && ds4_gpu_dsv41_quantize(row.kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0); if (ok) ok = ds4_gpu_tensor_copy(g->window[il], (uint64_t)(row.pos % 128u) * DS4_N_HEAD_DIM * sizeof(float), row.kv, 0, DS4_N_HEAD_DIM * sizeof(float)); } } if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false; #define DS41_DECODER_FREE(name, width) ds4_gpu_tensor_free(active.name); DS41_PREFILL_ROWS(DS41_DECODER_FREE) #undef DS41_DECODER_FREE } return ok; } typedef struct { pthread_t thread; const ds4_engram_table *table; const uint32_t *ids; float *out; uint32_t count, ready; bool active, ok, cancel, done; } ds41_engram_prefetch; static void *ds41_engram_prefetch_read(void *arg) { ds41_engram_prefetch *p = arg; p->ok = true; for (uint32_t off = 0; off < p->count; off += 2048u) { const uint32_t count = p->count - off < 2048u ? p->count - off : 2048u; if (__atomic_load_n(&p->cancel, __ATOMIC_RELAXED) || !ds4_engram_read_batch(p->table, p->ids + (size_t)off * 2u * DS4_ENGRAM_COLS, count, 2u * DS4_ENGRAM_COLS, p->out + (size_t)off * DS4_ENGRAM_COLS * DS4_ENGRAM_DIM)) { p->ok = false; break; } /* Publish only completed rows; later disk reads use disjoint memory. */ __atomic_store_n(&p->ready, off + count, __ATOMIC_RELEASE); } __atomic_store_n(&p->done, true, __ATOMIC_RELEASE); return NULL; } static bool ds41_engram_prefetch_wait(ds41_engram_prefetch *p, uint32_t rows, ds4_session_cancel_fn cancel, void *cancel_ud) { if (!p->active || rows > p->count) return false; while (__atomic_load_n(&p->ready, __ATOMIC_ACQUIRE) < rows) { if (__atomic_load_n(&p->done, __ATOMIC_ACQUIRE)) return __atomic_load_n(&p->ready, __ATOMIC_ACQUIRE) >= rows; if (cancel && cancel(cancel_ud)) return false; const struct timespec pause = {.tv_nsec = 1000000}; nanosleep(&pause, NULL); } return true; } static bool ds41_engram_prefetch_join(ds41_engram_prefetch *p, bool cancel) { if (!p->active) return true; if (cancel) __atomic_store_n(&p->cancel, true, __ATOMIC_RELAXED); if (pthread_join(p->thread, NULL)) ds4_die("cannot join V4.1 Engram reader safely"); p->active = false; return p->ok; } static bool ds41_engram_prefetch_start(ds41_engram_prefetch *p, ds41_gpu_graph *g, uint32_t table, uint32_t count) { *p = (ds41_engram_prefetch){.table = &g->table[table], .count = count, .ids = g->prefill_ids[0][table], .out = ds4_gpu_tensor_contents(g->engram_prefetch)}; if (!p->out || pthread_create(&p->thread, NULL, ds41_engram_prefetch_read, p)) return false; p->active = true; return true; } static uint32_t ds41_encoder_chunk_cap(const ds41_gpu_graph *g, uint32_t count) { if (count < 8192u && g->prefill_cap > 2048u) return 2048u; /* Keep the decoder suffix optimization for 8k prompts. */ if (count < 16384u && g->prefill_cap > 4096u) return 4096u; return g->prefill_cap; } /* Process rows in causal order within each layer. Selection/candidate rows * travel down the stack with their token, while only source layers append KV. * A failed partial chunk cannot be snapshotted: its layers have different * frontiers, so the caller must rebuild from its retained token history. */ static bool ds41_graph_prefill_sweep(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t total_count, ds4_session_progress_fn progress, void *progress_ud, int total, ds4_session_cancel_fn cancel, void *cancel_ud, bool encoder_only, bool resume_encoder) { const uint32_t encoder_chunk = ds41_encoder_chunk_cap(g, total_count); const bool wide = total_count > encoder_chunk; if ((!g->valid && !resume_encoder) || !total_count || (wide && total_count > g->carry_cap) || total_count > g->ctx - g->pos || g->imatrix || !ds41_tp_batch_enabled(g)) return false; const bool profile = getenv("DS4_METAL_GRAPH_PREFILL_PROFILE") != NULL; const bool stage_profile = getenv("DS4_METAL_V41_STAGE_PROFILE") != NULL; const bool batch_moe = !getenv("DS4_METAL_DISABLE_V41_BATCH_MOE"); const bool batch_attention = !getenv("DS4_METAL_DISABLE_V41_BATCH_ATTN"); const bool batch_core = batch_attention && !getenv("DS4_METAL_DISABLE_V41_BATCH_CORE"); const bool batch_hc = batch_attention && batch_moe && !getenv("DS4_METAL_DISABLE_V41_BATCH_HC"); const bool decoder_suffix = wide && total_count >= 8192u && !getenv("DS4_METAL_DISABLE_V41_DECODER_SUFFIX"); if ((encoder_only || resume_encoder) && !decoder_suffix) return false; uint32_t (*ids)[2][DS4_ENGRAM_COLS] = g->prefill_ids; ds4_engram_history next_history = g->history; if (!ds41_hash_tokens(g, &next_history, tokens, total_count, &ids[0][0][0])) return false; const uint32_t initial_start = g->pos; g->valid = false; ds41_gpu_graph row = *g; /* The complete current layer is mapped for prefill. Refresh the bounded * decode cache from it before advancing to the next layer. */ row.streaming = false; ds41_engram_prefetch engram_prefetch = {0}; const bool overlap_engram = total_count >= 1024u && !getenv("DS4_METAL_DISABLE_V41_ENGRAM_PREFETCH") && !getenv("DS4_METAL_DISABLE_V41_BATCH_ENGRAM"); const bool pipeline_engram = overlap_engram && !getenv("DS4_METAL_DISABLE_V41_ENGRAM_PIPELINE"); bool engram_prefetched = overlap_engram && ds41_engram_prefetch_start(&engram_prefetch, g, 0, total_count); bool ok = !g->streaming || metal_graph_stream_map_token(m, w); metal_graph_stream_prepare_slot prepare = {0}; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { if (cancel && cancel(cancel_ud)) { ok = false; break; } if (encoder_only && il == 20u) { /* Publish every encoder key, but leave the decoder invalid until * the last sweep rebuilds its exact 2541-token dependency suffix. */ if (g->streaming) ok = metal_graph_stream_map_layer_decode(m, w, il); if (ok) ok = ds41_decoder_prepare(g, m, &w->layer[il], il, initial_start, 0, total_count, true, batch_hc, batch_attention, cancel, cancel_ud); break; } /* Layer 1 no longer reads the prefix buffer. Fill it for layer 14 * while the intervening encoder layers run. The table stays on disk. */ if (il == 2u) engram_prefetched = overlap_engram && ds41_engram_prefetch_start(&engram_prefetch, g, 1, total_count); const double t0 = profile ? now_sec() : 0; if (g->streaming) { #ifdef __APPLE__ const uint32_t first_count = total_count < encoder_chunk ? total_count : encoder_chunk; if (!g->encoder_resident || il >= 20) ok = metal_graph_stream_prepare_join_layer(NULL, m, w, il, first_count, false, true, false, false, &prepare, 1); #endif if (ok) ok = metal_graph_stream_map_layer(m, w, il); #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) /* Below 2K, unused expert reads outweigh the overlap. */ if (ok && total_count >= 2048u) { const ds4_gpu_stream_expert_table current = graph_stream_expert_table_make(m, &w->layer[il], il, routed_expert_row_bytes(w->layer[il].ffn_gate_exps) * DS4_N_FF_EXP, routed_expert_row_bytes(w->layer[il].ffn_down_exps) * DS4_N_EMBD); ds4_gpu_stream_expert_table next = {0}; const bool more = il + 1u < (encoder_only ? 20u : DS4_N_LAYER); if (more) next = graph_stream_expert_table_make(m, &w->layer[il + 1u], il + 1u, routed_expert_row_bytes(w->layer[il + 1u].ffn_gate_exps) * DS4_N_FF_EXP, routed_expert_row_bytes(w->layer[il + 1u].ffn_down_exps) * DS4_N_EMBD); (void)ds4_gpu_stream_expert_cache_prefetch(¤t, more ? &next : NULL); } #endif #ifdef __APPLE__ /* CUDA reads into its device cache: warming mmap would read twice. */ if (ok && il + 1u < (encoder_only ? 20u : DS4_N_LAYER) && (!g->encoder_resident || il + 1u >= 20)) ok = metal_graph_stream_prepare_start_if_needed(NULL, m, w, il + 1u, first_count, false, true, false, false, &prepare, 1); #endif } const double t_map = profile ? now_sec() : 0; uint32_t first = 0; if (ok && decoder_suffix && il >= 20u) { if (il == 20u) ok = ds41_decoder_prepare(g, m, &w->layer[il], il, initial_start, 0, total_count, true, batch_hc, batch_attention, cancel, cancel_ud); const uint32_t needed = 1u + (DS4_N_LAYER - 1u - il) * 127u; first = total_count - needed; if (ok) ok = ds41_decoder_prepare(g, m, &w->layer[il], il, initial_start, first - 127u, 127u, false, batch_hc, batch_attention, cancel, cancel_ud); } /* Keep the decoder suffix's established matrix partitions; unlike * the encoder, its shrinking tail is not aligned to large tiles. */ const uint32_t chunk = wide && il >= 20u && encoder_chunk > 2048u ? 2048u : encoder_chunk; for (uint32_t off = first; ok && off < total_count; off += chunk) { if (cancel && cancel(cancel_ud)) { ok = false; break; } const uint32_t count = total_count - off < chunk ? total_count - off : chunk; const uint32_t start = initial_start + off; g->pos = start; ds41_prefill_row active = {0}; if (batch_hc) { /* HC helpers infer row count from buffer size, including short tails. */ #define DS41_ACTIVE_VIEW(name, width) \ if (ok) ok = (active.name = ds4_gpu_tensor_view(g->batch.name, 0, \ (uint64_t)count * (width) * sizeof(float))) != NULL; DS41_PREFILL_ROWS(DS41_ACTIVE_VIEW) #undef DS41_ACTIVE_VIEW } if (ok) ok = ds4_gpu_tensor_write(g->prefill_tokens, 0, tokens + off, (uint64_t)count * sizeof(int32_t)); if (ok && g->image_count) { uint8_t *mask = ds4_gpu_tensor_contents(g->image_text_mask); ok = mask != NULL; if (ok) ds41_text_mask(g, start, count, mask); } if (ok && !il && batch_core) { void *selection = ds4_gpu_tensor_contents(g->batch.selected_comp); if (!selection) ok = false; else memset(selection, 0xff, (size_t)count * DS4_N_INDEXER_TOP_K * sizeof(int32_t)); } if (ok) ok = ds4_gpu_begin_commands() != 0; if (!il) { #ifndef __APPLE__ if (ok && !g->image_count) { float *pre = calloc((size_t)count * DS4_N_HC, sizeof(float)); if (!pre) ok = false; else { for (uint32_t t = 0; t < count; t++) pre[t * DS4_N_HC] = 1; ok = ds4_gpu_tensor_write(g->batch.pre, 0, pre, (uint64_t)count * DS4_N_HC * sizeof(float)) && ds4_gpu_embed_tokens_hc_tensor(g->batch.residual, g->prefill_tokens, m->map, m->size, w->token_embd->abs_offset, DS4_N_VOCAB, count, DS4_N_EMBD, DS4_N_HC); free(pre); } } else #endif { const float initial_pre[] = {1, 0, 0, 0}; for (uint32_t t = 0; ok && t < count; t++) { ok = ds4_gpu_tensor_write(g->rows_view[t].pre, 0, initial_pre, sizeof(initial_pre)) && ds41_embed(g, m, w, g->rows_view[t].residual, g->rows_view[t].x, tokens[off + t], start + t); } } } else if (wide) { if (ok) ok = ds41_carry_copy(g, off, count, false); } if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false; const double t_ready = profile ? now_sec() : 0; if (ok && ds41_engram_layer(il)) { const uint32_t engram = il == 1 ? 0 : 1; if (engram_prefetched) { ok = pipeline_engram ? ds41_engram_prefetch_wait(&engram_prefetch, off + count, cancel, cancel_ud) : ds41_engram_prefetch_join(&engram_prefetch, false); } else if (getenv("DS4_METAL_DISABLE_V41_BATCH_ENGRAM")) { for (uint32_t t = 0; ok && t < count; t++) ok = ds4_engram_read(&g->table[engram], ids[off + t][engram], DS4_ENGRAM_COLS, ds4_gpu_tensor_contents(g->rows_view[t].engram_rows)); } else { ok = ds4_engram_read_batch(&g->table[engram], ids[off][engram], count, 2u * DS4_ENGRAM_COLS, ds4_gpu_tensor_contents(g->batch.engram_rows)); } } const double t_engram = profile ? now_sec() : 0; double stage_start = stage_profile ? now_sec() : 0; #define DS41_STAGE(label) do { \ if (ok && stage_profile) { \ ok = ds4_gpu_end_commands() != 0; \ const double now = now_sec(); \ fprintf(stderr, "ds4: V4.1 stage layer=%u pos=%u rows=%u %s=%.3f ms\n", \ il, start, count, (label), (now - stage_start) * 1000); \ stage_start = now; \ if (ok) ok = ds4_gpu_begin_commands() != 0; \ } \ } while (0) if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok && engram_prefetched && ds41_engram_layer(il)) { const uint64_t bytes = (uint64_t)DS4_ENGRAM_COLS * DS4_ENGRAM_DIM * sizeof(float); ok = ds4_gpu_tensor_copy(g->batch.engram_rows, 0, g->engram_prefetch, off * bytes, count * bytes) != 0; } if (ok && batch_hc) ok = ds41_before_attention_batch(g, &active, m, &w->layer[il], il, count); DS41_STAGE("hc/engram"); for (uint32_t t = 0; ok && !batch_hc && t < count; t++) { row.pos = start + t; #define DS41_USE_ROW(name, width) row.name = g->rows_view[t].name; DS41_PREFILL_ROWS(DS41_USE_ROW) #undef DS41_USE_ROW ok = batch_attention ? ds41_graph_before_attention(&row, m, &w->layer[il], il) : batch_moe ? ds41_graph_before_moe(&row, m, &w->layer[il], il) : ds41_graph_layer(&row, m, &w->layer[il], il, tokens[off + t]); } if (ok && batch_attention) { const ds4_layer_weights *l = &w->layer[il]; ok = ds41_attention_project_batch(g, m, l, count); DS41_STAGE("attention projections"); if (ok && batch_core) ok = ds41_attention_batch(g, m, l, il, count); for (uint32_t t = 0; ok && !batch_core && t < count; t++) { row.pos = start + t; #define DS41_USE_ATTN_ROW(name, width) row.name = g->rows_view[t].name; DS41_PREFILL_ROWS(DS41_USE_ATTN_ROW) #undef DS41_USE_ATTN_ROW ok = ds41_attention(&row, m, l, il, true); } DS41_STAGE("attention core/index"); if (ok && g->tp_world == 2) { ok = ds4_gpu_dsv41_attention_output_tp_batch(g->batch.block, g->batch.low, m->map, m->size, l->attn_output_a->abs_offset, l->attn_output_b->abs_offset, g->batch.heads, count, g->tp_rank) && ds41_sum_partial_batch(g, g->batch.block, il, count) && ds4_gpu_dsv41_quantize(g->batch.block, DS4_N_EMBD, count, DS4_V41_BF16); } else if (ok && l->attn_output_b->type == DS4_TENSOR_Q8_0) { ok = ds4_gpu_dsv41_attention_output_batch(g->batch.block, g->batch.low, m->map, m->size, l->attn_output_a->abs_offset, l->attn_output_b->abs_offset, g->batch.heads, count) && ds4_gpu_dsv41_quantize(g->batch.block, DS4_N_EMBD, count, DS4_V41_BF16); } else if (ok) { for (uint32_t t = 0; ok && t < count; t++) { row.heads = g->rows_view[t].heads; row.low = g->rows_view[t].low; ok = ds41_attention_low(&row, m, l); } if (ok) ok = ds41_matmul_batch(g->batch.block, m, l->attn_output_b, g->batch.low, count, true); } DS41_STAGE("attention output"); if (ok && batch_hc) ok = ds41_after_attention_batch(&active, m, l, count); DS41_STAGE("hc/ffn norm"); for (uint32_t t = 0; ok && !batch_hc && t < count; t++) { row.pos = start + t; #define DS41_USE_FFN_ROW(name, width) row.name = g->rows_view[t].name; DS41_PREFILL_ROWS(DS41_USE_FFN_ROW) #undef DS41_USE_FFN_ROW ok = ds41_graph_after_attention(&row, m, l); if (ok && !batch_moe) ok = ds41_moe(&row, m, l, il, (uint32_t)tokens[off + t]) && ds41_graph_after_moe(&row); } } if (ok && batch_moe) { ok = ds41_moe_batch(g, m, &w->layer[il], il, count, false); DS41_STAGE("shared/routed ffn"); if (ok && batch_hc) { ok = ds4_gpu_add_tensor(active.block, active.routed, active.shared, count * DS4_N_EMBD) && ds4_gpu_dsv41_quantize(active.block, DS4_N_EMBD, count, DS4_V41_BF16) && ds4_gpu_hc_expand_split_tensor(active.residual, active.block, active.after_attn, active.ffn_split, DS4_N_EMBD, DS4_N_HC) && ds4_gpu_dsv41_quantize(active.residual, DS4_N_EMBD * DS4_N_HC, count, DS4_V41_BF16); } for (uint32_t t = 0; ok && !batch_hc && t < count; t++) { #define DS41_USE_MOE_ROW(name, width) row.name = g->rows_view[t].name; DS41_PREFILL_ROWS(DS41_USE_MOE_ROW) #undef DS41_USE_MOE_ROW ok = ds4_gpu_add_tensor(row.block, row.routed, row.shared, DS4_N_EMBD) && ds41_bf16(row.block, DS4_N_EMBD) && ds41_graph_after_moe(&row); } } DS41_STAGE("hc expand"); #undef DS41_STAGE if (ok && wide && il + 1u < DS4_N_LAYER) { ok = ds41_carry_copy(g, off, count, true); } const double t_encoded = profile ? now_sec() : 0; if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false; if (g->tp_world == 2 && ds4_gpu_tp_failed()) ok = false; const double t_done = profile ? now_sec() : 0; if (ok && !encoder_only && off + count == total_count) ok = ds41_prefill_seed(g, m, &w->layer[il], il, count); if (engram_prefetched && ds41_engram_layer(il) && off + count == total_count && !ds41_engram_prefetch_join(&engram_prefetch, !ok)) ok = false; if (profile) fprintf(stderr, "ds4: V4.1 prefill layer=%u pos=%u rows=%u encoder_resident=%u " "map=%.3f engram=%.3f encode=%.3f drain=%.3f seed=%.3f ms\n", il, start, count, g->encoder_resident, (off == first ? t_map - t0 : 0) * 1000, (t_engram - t_ready) * 1000, (t_encoded - t_engram) * 1000, (t_done - t_encoded) * 1000, (now_sec() - t_done) * 1000); if (!ok) fprintf(stderr, "ds4: V4.1 layer-major prefill failed at layer %u\n", il); if (ok && progress) { const uint32_t after = (total_count * il + off + count) / DS4_N_LAYER; progress(progress_ud, "prefill_display", (int)(initial_start + after), total); } if (batch_hc) { #define DS41_ACTIVE_FREE(name, width) ds4_gpu_tensor_free(active.name); DS41_PREFILL_ROWS(DS41_ACTIVE_FREE) #undef DS41_ACTIVE_FREE } if (batch_hc && il + 1u == DS4_N_LAYER && off + count == total_count) { row.residual = g->rows_view[count - 1u].residual; row.pre = g->rows_view[count - 1u].ffn_split; } } } if (!ds41_engram_prefetch_join(&engram_prefetch, !ok)) ok = false; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (g->streaming) ds4_gpu_stream_expert_cache_prefetch_finish(true); #endif if (!metal_graph_stream_prepare_join_all(&prepare, 1)) ok = false; if (g->streaming && !metal_graph_stream_map_decode_static_all(m, w)) ok = false; if (ok && !encoder_only) { ok = ds4_gpu_begin_commands() && ds4_gpu_tensor_copy(g->residual, 0, row.residual, 0, (uint64_t)DS4_N_HC * DS4_N_EMBD * 4u) && ds4_gpu_tensor_copy(g->pre, 0, row.pre, 0, DS4_N_HC * 4u); if (ds4_gpu_commands_active() && !ds4_gpu_end_commands()) ok = false; } g->pos = initial_start; if (ok) { g->pos += total_count; g->history = next_history; g->valid = !encoder_only; } return ok; } static bool ds41_graph_prefill(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t count, ds4_session_progress_fn progress, void *progress_ud, int total, ds4_session_cancel_fn cancel, void *cancel_ud) { return ds41_graph_prefill_sweep(g, m, w, tokens, count, progress, progress_ud, total, cancel, cancel_ud, false, false); } static ds41_gpu_graph *ds41_batch_workspace(ds41_gpu_graph *const *graphs, int count) { if (!graphs || count < 2 || count > DS4_TP_BATCH_MAX_ROWS) return NULL; ds41_gpu_graph *largest = NULL; for (int i = 0; i < count; i++) { ds41_gpu_graph *g = graphs[i]; if (!g) return NULL; if (!largest || g->ctx > largest->ctx || (g->ctx == largest->ctx && g->prefill_cap > largest->prefill_cap)) largest = g; } return largest->prefill_cap >= (uint32_t)count ? largest : NULL; } static bool ds41_cuda_row_batch_supported(const ds41_gpu_graph *g, const ds4_weights *w) { #ifndef __APPLE__ const bool placement = (g->streaming && g->tp_world == 1) || (!g->streaming && g->tp_world == 2); if (!placement || g->quality || g->imatrix || g->image_count || !metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_MOE", true) || getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6")) return false; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &w->layer[il]; 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_Q2_K) return false; } return true; #else (void)g; (void)w; return false; #endif } static uint32_t ds41_short_prefill_count(const ds41_gpu_graph *g, const ds4_weights *w, uint32_t remaining) { if (remaining < 2u || remaining >= 256u || getenv("DS4_METAL_DISABLE_V41_LAYER_PREFILL") || !ds41_cuda_row_batch_supported(g, w)) return 0; if (g->tp_world == 2 && (!ds41_tp_batch_enabled(g) || getenv("DS4_METAL_DISABLE_V41_TP_SMALL_PREFILL"))) return 0; return remaining < DS4_TP_BATCH_MAX_ROWS ? remaining : DS4_TP_BATCH_MAX_ROWS; } static DS4_MAYBE_UNUSED bool ds41_graph_step_batch(ds41_gpu_graph *const *graphs, const int *tokens, int count, uint32_t prefill_rows, const ds4_model *model, const ds4_weights *weights) { if (!graphs || !tokens || count < 2 || count > DS4_TP_BATCH_MAX_ROWS || prefill_rows > (uint32_t)count) return false; /* Sparse block masks scale with context, even when prefill caps match. * Attention and Engram history remain private to each session. */ ds41_gpu_graph *g = ds41_batch_workspace(graphs, count); if (!g) return false; const uint32_t rows = (uint32_t)count; const bool prefill_only = prefill_rows == rows; const bool shared_owner = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER"); ds41_prefill_row active = {0}; ds4_engram_history history[DS4_TP_BATCH_MAX_ROWS]; uint32_t positions[DS4_TP_BATCH_MAX_ROWS]; /* Only the leading prefill rows share a session. Keep their disk reads * and hash history distinct until the final row is published. */ float (*engram)[2][DS4_ENGRAM_COLS * DS4_ENGRAM_DIM] = prefill_rows > 1 ? malloc(prefill_rows * sizeof(*engram)) : NULL; ds4_gpu_tensor *queries[DS4_TP_BATCH_MAX_ROWS] = {0}; ds4_gpu_tensor *heads[DS4_TP_BATCH_MAX_ROWS] = {0}; bool ok = rows <= g->prefill_cap && (prefill_rows <= 1 || engram); #define DS41_SESSION_VIEW(name, width) \ if (ok) ok = (active.name = ds4_gpu_tensor_view(g->batch.name, 0, \ (uint64_t)rows * (width) * sizeof(float))) != NULL; DS41_PREFILL_ROWS(DS41_SESSION_VIEW) #undef DS41_SESSION_VIEW const uint64_t head_bytes = (uint64_t)(DS4_N_HEAD / g->tp_world) * DS4_N_HEAD_DIM * sizeof(float); for (int i = 0; ok && i < count; i++) { queries[i] = ds4_gpu_tensor_view(g->batch.q, (uint64_t)i * head_bytes, head_bytes); heads[i] = ds4_gpu_tensor_view(g->batch.heads, (uint64_t)i * head_bytes, head_bytes); ok = queries[i] && heads[i]; } const float initial_pre[] = {1, 0, 0, 0}; for (int i = 0; ok && i < count; i++) { ds41_gpu_graph *s = graphs[i]; uint32_t ids[2][DS4_ENGRAM_COLS]; const bool continued = i > 0 && (uint32_t)i < prefill_rows; positions[i] = s->pos + (continued ? (uint32_t)i : 0u); history[i] = continued ? history[i - 1] : s->history; ok = ds4_engram_hash(&s->engram, &history[i], &tokens[i], NULL, 1, &ids[0][0]); float (*disk_rows)[DS4_ENGRAM_COLS * DS4_ENGRAM_DIM] = engram && (uint32_t)i < prefill_rows ? engram[i] : s->rows; for (unsigned table = 0; ok && table < 2; table++) ok = ds4_engram_read(&s->table[table], ids[table], DS4_ENGRAM_COLS, disk_rows[table]); if (ok) ok = ds4_gpu_tensor_write(g->rows_view[i].pre, 0, initial_pre, sizeof(initial_pre)) && ds4_gpu_embed_token_hc_tensor(g->rows_view[i].residual, model->map, model->size, weights->token_embd->abs_offset, DS4_N_VOCAB, (uint32_t)tokens[i], DS4_N_EMBD, DS4_N_HC); } if (ok) ok = ds4_gpu_tensor_write(g->prefill_tokens, 0, tokens, rows * sizeof(int)); uint32_t il = 0; for (; ok && il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &weights->layer[il]; if (ds41_engram_layer(il)) { const unsigned table = il == 1 ? 0 : 1; for (int i = 0; ok && i < count; i++) { ds41_gpu_graph *s = graphs[i]; ok = ds4_gpu_tensor_write(g->rows_view[i].engram_rows, 0, engram && (uint32_t)i < prefill_rows ? engram[i][table] : s->rows[table], sizeof(s->rows[table])); } } if (ok) ok = ds41_before_attention_batch(g, &active, model, l, il, rows) && ds41_attention_project_batch(g, model, l, rows); for (int i = 0; ok && i < count; i++) { ds41_gpu_graph row = *graphs[i]; row.pos = positions[i]; #define DS41_SESSION_ROW(name, width) row.name = g->rows_view[i].name; DS41_PREFILL_ROWS(DS41_SESSION_ROW) #undef DS41_SESSION_ROW row.q = queries[i]; row.heads = heads[i]; ok = ds41_attention(&row, model, l, il, true) && ds41_attention_output(&row, model, l); } if (ok) ok = ds41_sum_partial_batch(g, active.block, il, rows); if (ok) ok = ds4_gpu_dsv41_quantize(active.block, DS4_N_EMBD, rows, DS4_V41_BF16) && ds41_after_attention_batch(&active, model, l, rows) && ds41_moe_batch(g, model, l, il, rows, shared_owner) && (shared_owner ? ds4_gpu_tensor_copy(active.block, 0, active.routed, 0, (uint64_t)rows * DS4_N_EMBD * sizeof(float)) : ds4_gpu_add_tensor(active.block, active.routed, active.shared, rows * DS4_N_EMBD)) && ds4_gpu_dsv41_quantize(active.block, DS4_N_EMBD, rows, DS4_V41_BF16) && ds4_gpu_hc_expand_split_tensor(active.residual, active.block, active.after_attn, active.ffn_split, DS4_N_EMBD, DS4_N_HC) && ds4_gpu_dsv41_quantize(active.residual, DS4_N_EMBD * DS4_N_HC, rows, DS4_V41_BF16); /* The second Engram upload reuses the first one's input storage. */ if (ok && il == 13) ok = ds4_gpu_end_commands() && ds4_gpu_begin_commands(); } const uint64_t logits_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float) / (g->tp_logits_half ? 2u : 1u); /* The final FFN intermediates are dead. Reuse that storage for the head * rather than allocating another vocabulary-sized buffer per session. */ ds4_gpu_tensor *batch_logits = NULL; if (ok && !prefill_only && !getenv("DS4_METAL_DISABLE_V41_BATCH_HEAD") && ds4_gpu_tensor_bytes(g->batch.mid) >= rows * logits_bytes) batch_logits = ds4_gpu_tensor_view(g->batch.mid, 0, rows * logits_bytes); /* Pure prefill publishes only its final row; the caller computes the head * once at the complete prompt frontier. */ const int first_output = prefill_only ? count - 1 : 0; for (int i = first_output; ok && i < count; i++) { ds41_gpu_graph *s = graphs[i]; ok = ds4_gpu_tensor_copy(s->residual, 0, g->rows_view[i].residual, 0, (uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float)) && ds4_gpu_tensor_copy(s->pre, 0, g->rows_view[i].ffn_split, 0, DS4_N_HC * sizeof(float)); if (!ok || prefill_only) continue; ok = ds4_gpu_hc_weighted_sum_tensor(s->x, s->residual, s->pre, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(s->x, DS4_N_EMBD) && ds41_norm(s->norm, s->x, model, weights->output_norm); if (ok) ok = batch_logits ? ds4_gpu_tensor_copy(g->rows_view[i].norm, 0, s->norm, 0, DS4_N_EMBD * sizeof(float)) : ds41_output_projection(s, s->tp_logits_half ? s->tp_logits_half : s->logits, model, weights, s->norm, 1); } if (ok && batch_logits) ok = ds41_output_projection(g, batch_logits, model, weights, active.norm, rows); for (int i = first_output; ok && i < count; i++) { ds41_gpu_graph *s = graphs[i]; if (batch_logits) ok = ds4_gpu_tensor_copy(s->tp_logits_half ? s->tp_logits_half : s->logits, 0, batch_logits, (uint64_t)i * logits_bytes, logits_bytes); if (ok) { s->history = history[i]; s->pos = positions[i] + 1u; } } ds4_gpu_tensor_free(batch_logits); if (!ok) fprintf(stderr, "ds4: V4.1 session batch failed at layer %u (%d rows)\n", il, count); #define DS41_SESSION_FREE(name, width) ds4_gpu_tensor_free(active.name); DS41_PREFILL_ROWS(DS41_SESSION_FREE) #undef DS41_SESSION_FREE for (int i = 0; i < count; i++) { ds4_gpu_tensor_free(queries[i]); ds4_gpu_tensor_free(heads[i]); } free(engram); return ok; } static bool ds41_graph_short_prefill(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t count) { if (!g->valid || count < 2u || count > DS4_TP_BATCH_MAX_ROWS || g->pos > g->ctx || count > g->ctx - g->pos) return false; ds41_gpu_graph *graphs[DS4_TP_BATCH_MAX_ROWS]; for (uint32_t i = 0; i < count; i++) graphs[i] = g; bool ok = ds4_gpu_begin_commands() && ds41_graph_step_batch(graphs, tokens, (int)count, count, m, w); if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); if (!ok) g->valid = false; return ok; } #undef DS41_PREFILL_ROWS #undef DS41_PREFILL_ALIASES #undef DS41_PREFILL_STORAGE #undef DS41_CARRY_ROWS #endif 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; int im_start_id; int im_end_id; int endoftext_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; void *host_slab; /* CUDA's registered network staging */ 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; /* DeepSeek: logits halves cross the wire */ bool active; } ds4_engine_tp_state; typedef enum { DS4_VISION_NONE = 0, DS4_VISION_GLM53, DS4_VISION_DEEPSEEK4, DS4_VISION_QWEN4, } ds4_vision_kind; struct ds4_engine { char *model_path; uint64_t ds41_session_bytes; ds4_model model; ds4_model mtp_model; ds4_model vision_model; ds4_vocab vocab; ds4_weights weights; ds4_mtp_weights mtp_weights; ds4_dspark_weights dspark_weights; #ifndef DS4_NO_GPU ds4_glm53_vision_weights vision_weights; ds4_deepseek4_vision_weights deepseek4_vision_weights; ds4_qwen4_vision_weights qwen4_vision_weights; #endif ds4_vision_kind vision_kind; int vision_image_token; int vision_start_token; int vision_end_token; 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; bool dspark_confidence_threshold_set; 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; uint64_t ssd_streaming_decode_map_bytes; uint64_t ssd_streaming_decode_map_extra_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 dspark_exact_sampling; bool cuda_tensor_parallel; bool glm_tp_token_prefill; bool ssd_streaming; bool ssd_streaming_cold; bool ssd_streaming_full_layers_set; bool ssd_streaming_budget_finalized; bool ssd_streaming_static_decode_map; ds4_distributed_options distributed; ds4_engine_tp_state tp; bool metal_ready; bool mtp_ready; bool vision_ready; bool vision_map_ready; bool share_session_prefill_workspace; #ifndef DS4_NO_GPU bool shared_prefill_workspace_ready; ds4_gpu_graph shared_prefill_workspace; /* Qwen keeps the same idea in its own graph shape; the first session * allocates the arena and hands ownership here. */ struct ds4_qwen4_gpu_graph *qwen4_shared_workspace; uint32_t qwen4_arena_users; /* live sessions borrowing the shared workspace */ /* batched speculative policy: the drafts' measured acceptance and the * wall time of a plain and of a speculative batched cycle */ float qwen4_batch_p, qwen4_batch_ms[2]; uint32_t qwen4_batch_cycles, qwen4_batch_spec_cycles, qwen4_batch_width; /* Recurrent state for every slot of a layer in one tensor, so a batched * decode can advance all of its rows in a single dispatch instead of one * per session. Sessions hold views into it and a slot index. */ ds4_gpu_tensor *qwen4_lin_state_pool[DS4_MAX_LAYER]; ds4_gpu_tensor *qwen4_lin_hist_pool[DS4_MAX_LAYER]; uint32_t qwen4_pool_slots; uint64_t qwen4_pool_used; #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; int placement_session_count_hint; uint32_t glm_session_count; uint64_t glm_session_graph_bytes; }; static DS4_MAYBE_UNUSED uint64_t ds4_engine_glm_graph_budget( const ds4_engine *e, uint64_t next_graph_bytes) { /* Metal guards against host capacity, so include live sessions as well as * the uncreated slots. Other backends keep their available-memory policy. */ if (!e || e->backend != DS4_BACKEND_METAL) return next_graph_bytes; uint64_t uncreated = e->placement_session_count_hint > 0 && (uint32_t)e->placement_session_count_hint > e->glm_session_count ? (uint32_t)e->placement_session_count_hint - e->glm_session_count : 1u; if (next_graph_bytes > (UINT64_MAX - e->glm_session_graph_bytes) / uncreated) return UINT64_MAX; return e->glm_session_graph_bytes + uncreated * next_graph_bytes; } 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 glm_graph_streaming_active_model_bytes( const ds4_weights *weights); 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_prefill_headroom_bytes); total = ds4_add_sat_u64(total, e->ssd_streaming_decode_map_extra_bytes); return total; } static void ds4_engine_print_startup_memory( const ds4_engine *e, int ctx_size) { if (!e || ctx_size <= 0) return; ds4_context_memory mem; #ifndef DS4_NO_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set) { const uint32_t normal_layers = glm_graph_normal_layer_count(); const uint32_t layer_end = e->distributed.layers.has_output ? (normal_layers ? normal_layers - 1u : 0u) : e->distributed.layers.end; const uint32_t ctx = (uint32_t)ctx_size; const uint32_t work_ctx = glm_graph_full_attention_cap(ctx, e->ssd_streaming); const uint32_t compact_cap = glm_graph_compact_cache_initial_cap(ctx, work_ctx); mem = glm_graph_context_memory_estimate_for_compact_cap_slice( ctx, work_ctx, compact_cap, e->ssd_streaming, e->distributed.layers.start, layer_end); } else { #endif mem = ds4_context_memory_estimate_with_prefill_mode(e->backend, ctx_size, e->prefill_chunk, e->ssd_streaming); #ifndef DS4_NO_GPU } #endif const uint64_t kv_bytes = ds4_add_sat_u64(mem.raw_bytes, mem.compressed_bytes); const uint64_t dynamic_expert_cache_bytes = ds4_engine_dynamic_expert_cache_bytes(e); const uint64_t expert_reserved_bytes = e->ssd_streaming_prefill_headroom_bytes; uint64_t resident_model_bytes = e->startup_model_span_bytes; #ifndef DS4_NO_GPU if (e->ssd_streaming_static_decode_map && e->ssd_streaming_decode_map_bytes != 0) { resident_model_bytes = e->ssd_streaming_decode_map_bytes; } else if (e->ssd_streaming) { const uint64_t active_model_bytes = glm_graph_streaming_active_model_bytes(&e->weights); if (active_model_bytes > resident_model_bytes) { resident_model_bytes = active_model_bytes; } } if (e->ssd_streaming && e->backend == DS4_BACKEND_METAL) { uint64_t static_bytes = 0; if (weights_streaming_non_routed_bytes(&e->weights, &static_bytes) && static_bytes > resident_model_bytes) { resident_model_bytes = static_bytes; } } #endif uint64_t total = kv_bytes; total = ds4_add_sat_u64(total, mem.scratch_bytes); total = ds4_add_sat_u64(total, resident_model_bytes); total = ds4_add_sat_u64(total, dynamic_expert_cache_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(resident_model_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 uint32_t n_layers = directional_steering_layer_count(); if (n_layers == 0) return false; const uint64_t n = (uint64_t)n_layers * 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); } } #include "ds4_qwen4_unicode.inc" static bool qwen4_urange_has(const ds4_qwen4_urange *r, size_t n, uint32_t cp) { size_t lo = 0, hi = n; while (lo < hi) { const size_t mid = (lo + hi) / 2; if (cp < r[mid].lo) hi = mid; else if (cp > r[mid].hi) lo = mid + 1; else return true; } return false; } #define DS4_QWEN4_CLASS(name_, cp_) \ qwen4_urange_has(ds4_qwen4_##name_##_ranges, \ sizeof(ds4_qwen4_##name_##_ranges) / sizeof(ds4_qwen4_##name_##_ranges[0]), (cp_)) typedef struct { uint32_t cp; uint64_t next; bool valid; bool letter; bool mark; bool number; bool space; } qwen4_char_info; static qwen4_char_info qwen4_char_at(const char *s, uint64_t len, uint64_t pos) { qwen4_char_info c; memset(&c, 0, sizeof(c)); if (pos >= len) return c; c.valid = true; c.cp = utf8_peek_one(s, len, pos, &c.next); c.letter = DS4_QWEN4_CLASS(letter, c.cp); c.mark = DS4_QWEN4_CLASS(mark, c.cp); c.number = DS4_QWEN4_CLASS(number, c.cp); c.space = DS4_QWEN4_CLASS(space, c.cp); return c; } /* Qwen3.8 pre-tokenization (tokenizer.ggml.pre = "qwen35"), the ordered * alternation * (?i:'s|'t|'re|'ve|'m|'ll|'d) | [^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+ | \p{N} | * ?[^\s\p{L}\p{M}\p{N}]+[\r\n]* | \s*[\r\n]+ | \s+(?!\S) | \s+ * with the Unicode classes generated from the `regex` module. */ static void bpe_tokenize_text_qwen35(const ds4_vocab *vocab, const char *text, token_vec *out) { const uint64_t len = strlen(text); uint64_t pos = 0; while (pos < len) { const uint64_t start = pos; qwen4_char_info cur = qwen4_char_at(text, len, pos); if (!cur.valid) break; if (cur.cp == '\'' && cur.next < len) { qwen4_char_info n1 = qwen4_char_at(text, len, cur.next); const uint32_t c1 = ascii_tolower_cp(n1.cp); if (c1 == 's' || c1 == 't' || c1 == 'm' || c1 == 'd') { pos = n1.next; bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); continue; } if (n1.valid && n1.next < len) { qwen4_char_info n2 = qwen4_char_at(text, len, n1.next); const uint32_t c2 = ascii_tolower_cp(n2.cp); if ((c1 == 'r' && c2 == 'e') || (c1 == 'v' && c2 == 'e') || (c1 == 'l' && c2 == 'l')) { pos = n2.next; bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); continue; } } } /* an optional non-letter prefix, then a run of letters and marks */ { uint64_t run = UINT64_MAX; if (cur.letter || cur.mark) { run = cur.next; } else if (cur.cp != '\r' && cur.cp != '\n' && !cur.number) { qwen4_char_info n1 = qwen4_char_at(text, len, cur.next); if (n1.valid && (n1.letter || n1.mark)) run = n1.next; } if (run != UINT64_MAX) { pos = run; while (pos < len) { qwen4_char_info scan = qwen4_char_at(text, len, pos); if (!scan.valid || !(scan.letter || scan.mark)) break; pos = scan.next; } bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); continue; } } if (cur.number) { pos = cur.next; bpe_emit_piece(vocab, (ds4_str){ text + start, pos - start }, out); continue; } { qwen4_char_info punct = cur; uint64_t punct_pos = pos; if (cur.cp == ' ') { punct_pos = cur.next; punct = qwen4_char_at(text, len, punct_pos); } if (punct.valid && !punct.space && !punct.letter && !punct.mark && !punct.number) { pos = punct_pos; while (pos < len) { qwen4_char_info scan = qwen4_char_at(text, len, pos); if (!scan.valid || scan.space || scan.letter || scan.mark || scan.number) break; pos = scan.next; } while (pos < len) { qwen4_char_info scan = qwen4_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.space) { uint64_t p = pos; uint64_t last_newline_end = 0; uint64_t last_ws_start = pos; int nspace = 0; while (p < len) { qwen4_char_info scan = qwen4_char_at(text, len, p); if (!scan.valid || !scan.space) 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; } if (ds4_model_is_qwen4()) { bpe_tokenize_text_qwen35(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); } vocab->im_start_id = -1; vocab->im_end_id = -1; vocab->endoftext_id = -1; if (ds4_model_is_qwen4()) { /* ChatML without BOS; <|endoftext|> is the document separator and a * second generation stop. */ vocab->im_start_id = vocab_lookup(vocab, "<|im_start|>"); vocab->im_end_id = vocab_lookup(vocab, "<|im_end|>"); vocab->endoftext_id = vocab_lookup(vocab, "<|endoftext|>"); vocab->bos_id = -1; vocab->eos_id = vocab->im_end_id; vocab->system_id = -1; vocab->user_id = -1; vocab->assistant_id = -1; 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 = 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 = -1; vocab->arg_key_end_id = -1; vocab->arg_value_start_id = -1; vocab->arg_value_end_id = -1; vocab->dsml_id = -1; return; } 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 = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 ? vocab_lookup(vocab, "<｜System｜>") : -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) { if (vocab->bos_id < 0) return; 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_LOW: case DS4_THINK_MEDIUM: case DS4_THINK_HIGH: return "Reasoning Effort: High"; case DS4_THINK_MAX: return "Reasoning Effort: Max"; case DS4_THINK_NONE: return NULL; } return NULL; } const char *ds4_deepseek41_reasoning_effort_text(ds4_think_mode mode) { int level = ds4_think_mode_level(mode); if (mode == DS4_THINK_HIGH) level = 75; if (mode == DS4_THINK_MAX) level = 100; if (level <= 0) return NULL; static __thread char text[128]; snprintf(text, sizeof(text), "Reasoning Effort: %d (range 1-100, the higher the value, the more thorough the reasoning)\n\n", level); return text; } 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 (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { const char *effort = ds4_deepseek41_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); } } /* Qwen3.8 ChatML. Thinking on renders the template's reasoning instruction * for the effort level (xhigh by default, low, none for medium) into the * system turn and opens ; thinking off closes an empty think block, as * the reference template does. */ static const char *DS4_QWEN4_REASONING_XHIGH = "Reasoning effort is set to xhigh. Please think carefully through the task, " "validate key assumptions, consider plausible alternatives, and prioritize " "correctness, consistency, and clarity in the final answer."; static const char *DS4_QWEN4_REASONING_LOW = "Reasoning effort is set to low. Keep your thinking brief and focused, " "moving directly to the conclusion without unnecessary elaboration."; static void qwen4_chat_open(const ds4_vocab *vocab, const char *role, token_vec *out) { token_vec_push(out, vocab->im_start_id); bpe_tokenize_text(vocab, role, out); bpe_tokenize_text(vocab, "\n", out); } static void qwen4_chat_close(const ds4_vocab *vocab, token_vec *out) { token_vec_push(out, vocab->im_end_id); bpe_tokenize_text(vocab, "\n", out); } static void qwen4_chat_assistant_prefix(const ds4_vocab *vocab, ds4_think_mode think_mode, token_vec *out) { qwen4_chat_open(vocab, "assistant", out); token_vec_push(out, vocab->think_start_id); if (ds4_think_mode_enabled(think_mode)) { bpe_tokenize_text(vocab, "\n", out); } else { bpe_tokenize_text(vocab, "\n\n", out); token_vec_push(out, vocab->think_end_id); bpe_tokenize_text(vocab, "\n\n", out); } } static void qwen4_chat_system(const ds4_vocab *vocab, const char *system, ds4_think_mode think_mode, token_vec *out) { const char *instruction = ds4_qwen4_reasoning_effort_text(think_mode); const bool have_system = system && system[0]; if (!instruction && !have_system) return; qwen4_chat_open(vocab, "system", out); if (instruction) { bpe_tokenize_text(vocab, instruction, out); if (have_system) bpe_tokenize_text(vocab, "\n\n", out); } if (have_system) bpe_tokenize_text(vocab, system, out); qwen4_chat_close(vocab, out); } static void encode_chat_prompt( const ds4_vocab *vocab, const char *system, const char *prompt, ds4_think_mode think_mode, token_vec *out) { if (ds4_model_is_qwen4()) { if (vocab->im_start_id < 0 || vocab->im_end_id < 0 || vocab->think_start_id < 0 || vocab->think_end_id < 0) { ds4_die("this tokenizer does not provide the Qwen chat markers; use raw prompt tokenization"); } qwen4_chat_system(vocab, system, think_mode, out); qwen4_chat_open(vocab, "user", out); bpe_tokenize_text(vocab, prompt, out); qwen4_chat_close(vocab, out); qwen4_chat_assistant_prefix(vocab, think_mode, out); return; } 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 || (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && !ds4_think_mode_enabled(think_mode))) 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}, {"<｜System｜>", DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 ? vocab->system_id : -1}, {"<｜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}, {"<|im_start|>", vocab->im_start_id}, {"<|im_end|>", vocab->im_end_id}, {"<|endoftext|>", vocab->endoftext_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) { chat_push_think_prefix(&e->vocab, DS4_THINK_MAX, tokens); } void ds4_chat_append_think_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode mode) { chat_push_think_prefix(&e->vocab, mode, tokens); } static void ds41_chat_system_marker(const ds4_vocab *vocab, ds4_tokens *tokens) { /* Initial system text may be built in pieces after the effort prefix. * Later system messages must introduce their own role token. */ bool has_system = false; for (int i = 0; i < tokens->len; i++) { if (tokens->v[i] == vocab->user_id || tokens->v[i] == vocab->assistant_id) { token_vec_push(tokens, vocab->system_id); return; } if (tokens->v[i] == vocab->system_id) has_system = true; } if (!has_system) token_vec_push(tokens, vocab->system_id); } 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 char *span = content ? content : ""; const char *p = span; while (*p) { if (ds4_tool_text_needs_escape(p, end)) { tokenize_span(vocab, span, (size_t)(p - span), out); bpe_tokenize_text(vocab, *p == '<' ? "<" : "&", 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_is_qwen4()) { if (!strcmp(role, "tool") || !strcmp(role, "function")) { qwen4_chat_open(vocab, "user", tokens); bpe_tokenize_text(vocab, "\n", tokens); bpe_tokenize_tool_response_text(vocab, content, tokens); bpe_tokenize_text(vocab, "\n", tokens); } else { const char *name = (!strcmp(role, "system") || !strcmp(role, "developer")) ? "system" : !strcmp(role, "assistant") ? "assistant" : "user"; qwen4_chat_open(vocab, name, tokens); if (!strcmp(name, "assistant")) { tokenize_rendered_chat_vocab(vocab, content, tokens); } else { bpe_tokenize_text(vocab, content, tokens); } } qwen4_chat_close(vocab, tokens); return; } 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")) { if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) ds41_chat_system_marker(vocab, tokens); 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) { if (ds4_model_is_qwen4()) { qwen4_chat_assistant_prefix(&e->vocab, think_mode, tokens); return; } 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 (vocab->endoftext_id >= 0 && token == vocab->endoftext_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 inline void argmax_f32_unrolled8_range( const float *logits, uint32_t begin, uint32_t end, int *best, float *best_v) { uint32_t i = begin; int b0 = *best, b1 = *best, b2 = *best, b3 = *best; int b4 = *best, b5 = *best, b6 = *best, b7 = *best; float v0 = *best_v, v1 = *best_v, v2 = *best_v, v3 = *best_v; float v4 = *best_v, v5 = *best_v, v6 = *best_v, v7 = *best_v; while (end - i >= 8u) { const float x0 = logits[i + 0u]; const float x1 = logits[i + 1u]; const float x2 = logits[i + 2u]; const float x3 = logits[i + 3u]; const float x4 = logits[i + 4u]; const float x5 = logits[i + 5u]; const float x6 = logits[i + 6u]; const float x7 = logits[i + 7u]; if (x0 > v0) { v0 = x0; b0 = (int)(i + 0u); } if (x1 > v1) { v1 = x1; b1 = (int)(i + 1u); } if (x2 > v2) { v2 = x2; b2 = (int)(i + 2u); } if (x3 > v3) { v3 = x3; b3 = (int)(i + 3u); } if (x4 > v4) { v4 = x4; b4 = (int)(i + 4u); } if (x5 > v5) { v5 = x5; b5 = (int)(i + 5u); } if (x6 > v6) { v6 = x6; b6 = (int)(i + 6u); } if (x7 > v7) { v7 = x7; b7 = (int)(i + 7u); } i += 8u; } #define DS4_ARGMAX_MERGE_LANE(b, v) \ do { \ if ((v) > *best_v || ((v) == *best_v && (b) < *best)) { \ *best_v = (v); \ *best = (b); \ } \ } while (0) DS4_ARGMAX_MERGE_LANE(b0, v0); DS4_ARGMAX_MERGE_LANE(b1, v1); DS4_ARGMAX_MERGE_LANE(b2, v2); DS4_ARGMAX_MERGE_LANE(b3, v3); DS4_ARGMAX_MERGE_LANE(b4, v4); DS4_ARGMAX_MERGE_LANE(b5, v5); DS4_ARGMAX_MERGE_LANE(b6, v6); DS4_ARGMAX_MERGE_LANE(b7, v7); #undef DS4_ARGMAX_MERGE_LANE for (; i < end; i++) { const float v = logits[i]; if (v > *best_v) { *best_v = v; *best = (int)i; } } } static int sample_argmax_unrolled8(const float *logits, uint32_t n_vocab) { int best = 0; float best_v = DS4_NEG_INF; argmax_f32_unrolled8_range(logits, 0, n_vocab, &best, &best_v); return best; } static int argmax_f32_excluding_unrolled8( const float *logits, uint32_t n, int excluded_id) { const uint32_t first = excluded_id == 0 ? 1u : 0u; if (first >= n) return -1; int best = (int)first; float best_v = logits[first]; const uint32_t begin = first + 1u; if (excluded_id >= 0 && (uint32_t)excluded_id >= begin && (uint32_t)excluded_id < n) { const uint32_t excluded = (uint32_t)excluded_id; argmax_f32_unrolled8_range(logits, begin, excluded, &best, &best_v); argmax_f32_unrolled8_range(logits, excluded + 1u, n, &best, &best_v); } else { argmax_f32_unrolled8_range(logits, begin, n, &best, &best_v); } return best; } static int sample_argmax(const float *logits, uint32_t n_vocab) { if (getenv("DS4_CPU_DISABLE_UNROLLED_ARGMAX") == NULL) { return sample_argmax_unrolled8(logits, 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; } } /* Materialize the exact distribution represented by the public sampling * knobs. Speculative decoding needs probabilities, rather than only a sampled * token, for the p/q acceptance test and the corrected rejection sample. */ static bool sample_build_probabilities( const float *logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, float *probs) { if (!logits || !probs || n_vocab == 0) return false; memset(probs, 0, (size_t)n_vocab * sizeof(probs[0])); if (temperature <= 0.0f) { probs[sample_argmax(logits, n_vocab)] = 1.0f; return true; } 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 > 1024) top_k = 1024; if (top_k < 0) top_k = 0; if ((uint32_t)top_k > n_vocab) top_k = (int)n_vocab; enum { SAMPLE_PROB_FAST_CAP = 512 }; sample_candidate heap[1024]; const uint32_t heap_cap = top_k > 0 ? (uint32_t)top_k : (uint32_t)SAMPLE_PROB_FAST_CAP; uint32_t finite = 0; uint32_t heap_n = 0; float max_logit = DS4_NEG_INF; for (uint32_t i = 0; i < n_vocab; i++) { if (!isfinite(logits[i])) continue; finite++; if (logits[i] > max_logit) max_logit = logits[i]; } if (finite == 0) { probs[sample_argmax(logits, n_vocab)] = 1.0f; return true; } float full_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 raw = expf((v - max_logit) / temperature); full_sum += raw; sample_candidate cand = {.id = (int)i, .logit = v, .prob = raw}; if (heap_n < heap_cap) { heap[heap_n] = cand; heap_sum += raw; sample_heap_sift_up(heap, heap_n); heap_n++; } else if (sample_candidate_gt(cand, heap[0])) { heap_sum -= heap[0].prob; heap[0] = cand; heap_sum += raw; sample_heap_sift_down(heap, heap_n, 0); } } if (full_sum <= 0.0f || !isfinite(full_sum) || heap_n == 0) { probs[sample_argmax(logits, n_vocab)] = 1.0f; return true; } qsort(heap, heap_n, sizeof(heap[0]), sample_candidate_cmp_desc); float base_sum = top_k > 0 ? heap_sum : full_sum; float filtered_sum = 0.0f; uint32_t filtered = 0; bool stopped_by_min_p = false; bool reached_top_p = false; const float min_raw = heap[0].prob * min_p; for (uint32_t i = 0; i < heap_n; i++) { if (i > 0 && heap[i].prob < min_raw) { stopped_by_min_p = true; break; } filtered_sum += heap[i].prob; filtered++; if (filtered_sum / base_sum >= top_p) { reached_top_p = true; break; } } const bool heap_has_all = heap_n == finite; const bool min_p_proves_tail_filtered = stopped_by_min_p && heap[heap_n - 1u].prob < min_raw; const bool fast_complete = top_k > 0 || heap_has_all || reached_top_p || min_p_proves_tail_filtered; sample_candidate *cand = heap; uint32_t cand_n = heap_n; bool owned = false; if (!fast_complete) { cand = xmalloc((size_t)finite * sizeof(cand[0])); cand_n = 0; for (uint32_t i = 0; i < n_vocab; i++) { const float v = logits[i]; if (!isfinite(v)) continue; cand[cand_n++] = (sample_candidate){ .id = (int)i, .logit = v, .prob = expf((v - max_logit) / temperature), }; } qsort(cand, cand_n, sizeof(cand[0]), sample_candidate_cmp_desc); base_sum = full_sum; filtered_sum = 0.0f; filtered = 0; const float full_min_raw = cand[0].prob * min_p; for (uint32_t i = 0; i < cand_n; i++) { if (i > 0 && cand[i].prob < full_min_raw) break; filtered_sum += cand[i].prob; filtered++; if (filtered_sum / base_sum >= top_p) break; } owned = true; } if (filtered == 0 || filtered_sum <= 0.0f || !isfinite(filtered_sum)) { probs[cand[0].id] = 1.0f; } else { const float inv_sum = 1.0f / filtered_sum; for (uint32_t i = 0; i < filtered; i++) { probs[cand[i].id] = cand[i].prob * inv_sum; } } if (owned) free(cand); return true; } static int sample_probabilities(const float *probs, uint32_t n_vocab, uint64_t *rng) { if (!probs || !rng || n_vocab == 0) return -1; float sum = 0.0f; int best = 0; float best_p = probs[0]; for (uint32_t i = 0; i < n_vocab; i++) { const float p = probs[i]; if (p > 0.0f && isfinite(p)) sum += p; if (p > best_p) { best_p = p; best = (int)i; } } 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 = probs[i]; if (p <= 0.0f || !isfinite(p)) continue; r -= p; if (r <= 0.0f) return (int)i; } return best; } /* Replace target_probs with normalized max(0, p-q) and sample it. This is the * correction that keeps speculative decoding distributed exactly as p after * a rejected draft token. */ #ifdef DS4_TEST_HOOKS static int sample_residual_probabilities(float *target_probs, const float *draft_probs, uint32_t n_vocab, uint64_t *rng) { if (!target_probs || !draft_probs || !rng || n_vocab == 0) return -1; int target_best = 0; float target_best_p = target_probs[0]; float sum = 0.0f; for (uint32_t i = 0; i < n_vocab; i++) { if (target_probs[i] > target_best_p) { target_best_p = target_probs[i]; target_best = (int)i; } float residual = target_probs[i] - draft_probs[i]; if (!(residual > 0.0f) || !isfinite(residual)) residual = 0.0f; target_probs[i] = residual; sum += residual; } if (sum <= 0.0f || !isfinite(sum)) return target_best; return sample_probabilities(target_probs, n_vocab, rng); } #endif 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); } int ds4_test_sampling_probabilities(const float *logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, float *probs) { return sample_build_probabilities(logits, n_vocab, temperature, top_k, top_p, min_p, probs) ? 1 : 0; } int ds4_test_speculative_sample(const float *target_logits, const float *draft_logits, uint32_t n_vocab, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, float *target_probs, float *draft_probs) { if (!target_logits || !draft_logits || !rng || !target_probs || !draft_probs || n_vocab == 0 || temperature <= 0.0f) { return -1; } if (!sample_build_probabilities(draft_logits, n_vocab, temperature, top_k, top_p, min_p, draft_probs) || !sample_build_probabilities(target_logits, n_vocab, temperature, top_k, top_p, min_p, target_probs)) { return -1; } const int proposal = sample_probabilities(draft_probs, n_vocab, rng); if (proposal < 0) return -1; const float q = draft_probs[proposal]; const float p = target_probs[proposal]; if (q <= 0.0f || sample_rng_f32(rng) * q <= p) return proposal; return sample_residual_probabilities(target_probs, draft_probs, n_vocab, rng); } int ds4_test_speculative_delta_sample(const float *target_logits, uint32_t n_vocab, int draft_token, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, float *target_probs) { if (!target_logits || !rng || !target_probs || n_vocab == 0 || draft_token < 0 || (uint32_t)draft_token >= n_vocab || temperature <= 0.0f || !sample_build_probabilities(target_logits, n_vocab, temperature, top_k, top_p, min_p, target_probs)) { return -1; } if (sample_rng_f32(rng) <= target_probs[draft_token]) { return draft_token; } target_probs[draft_token] = 0.0f; return sample_probabilities(target_probs, n_vocab, rng); } int ds4_test_argmax_excluding_logits(const float *logits, uint32_t n_vocab, int excluded_id) { if (!logits) return -1; if (getenv("DS4_CPU_DISABLE_UNROLLED_ARGMAX") == NULL) { return argmax_f32_excluding_unrolled8(logits, n_vocab, excluded_id); } int best = -1; float best_v = DS4_NEG_INF; for (uint32_t i = 0; i < n_vocab; i++) { if ((int)i == excluded_id) continue; const float v = logits[i]; if (best < 0 || v > best_v) { best = (int)i; best_v = v; } } return best; } #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; } #if !defined(DS4_NO_GPU) || defined(DS4_TEST_HOOKS) static double glm_graph_bytes_to_gib(uint64_t bytes) { return (double)bytes / (1024.0 * 1024.0 * 1024.0); } #if defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) static uint64_t g_glm_rocm_guard_available_baseline; #endif static bool glm_graph_memory_guard_disabled(void) { const char *env = getenv("DS4_GLM_MEMORY_GUARD"); 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_memory_guard_default_reserve_gib( uint64_t budget_base, uint64_t model_bytes, bool glm53) { const double base_gib = glm_graph_bytes_to_gib(budget_base); const double model_gib = glm_graph_bytes_to_gib(model_bytes); if (base_gib >= 480.0 && base_gib <= 640.0 && model_gib >= base_gib * 0.80) { return 24.0; } if (glm53 && base_gib >= 108.0 && base_gib <= 160.0 && model_gib >= 70.0) { /* A nominal 128 GB host reports less than 120 GiB, and ROCm further * limits this base to currently available memory. Preserve the proven * resident-Q2 budget without imposing it on larger machines. */ return 18.0; } return 32.0; } static uint64_t glm_graph_memory_guard_budget_bytes( uint64_t budget_base, uint64_t wired_limit, double fraction, double reserve_gib) { 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. */ 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; } return budget; } #endif 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 } #ifndef DS4_NO_GPU typedef struct ds4_glm_gpu_graph { const ds4_weights *weights; 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; /* GLM-5.3 keeps four mHC residual streams. cur/next remain the plain * one-stream scratch used by the attention and FFN implementations. */ ds4_gpu_tensor *hc_cur; ds4_gpu_tensor *hc_next; ds4_gpu_tensor *hc_flat; ds4_gpu_tensor *hc_mix; ds4_gpu_tensor *hc_split; ds4_gpu_tensor *hc_post; ds4_gpu_tensor *hc_comb; ds4_gpu_tensor *hc_after_attn; ds4_gpu_tensor *hc_output; ds4_gpu_tensor *hc_mean_weights; 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 *kda_q; ds4_gpu_tensor *kda_k; ds4_gpu_tensor *kda_v; ds4_gpu_tensor *kda_lowrank; ds4_gpu_tensor *kda_raw_gate; ds4_gpu_tensor *kda_raw_beta; ds4_gpu_tensor *kda_output_gate; ds4_gpu_tensor *kda_out; ds4_gpu_tensor *layer_kda_conv_state[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_kda_recurrent_state[DS4_MAX_LAYER]; 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 batch_cap; /* GLM-5.3 layer-major prefill workspace. It is allocated at graph * creation for the full supported prefill chunk. */ uint32_t glm53_prefill_cap; ds4_gpu_tensor *batch_hc_cur; ds4_gpu_tensor *batch_hc_next; ds4_gpu_tensor *batch_hc_flat; ds4_gpu_tensor *batch_hc_mix; ds4_gpu_tensor *batch_hc_split; ds4_gpu_tensor *batch_hc_after_attn; ds4_gpu_tensor *batch_kda_q; ds4_gpu_tensor *batch_kda_k; ds4_gpu_tensor *batch_kda_v; ds4_gpu_tensor *batch_kda_lowrank; ds4_gpu_tensor *batch_kda_raw_gate; ds4_gpu_tensor *batch_kda_raw_beta; ds4_gpu_tensor *batch_kda_output_gate; ds4_gpu_tensor *batch_kda_out; ds4_gpu_tensor *session_batch_logits; uint32_t session_batch_logits_cap; 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_gate; ds4_gpu_tensor *indexer_q; ds4_gpu_tensor *indexer_weights; ds4_gpu_tensor *indexer_scores; ds4_gpu_tensor *indexer_pool_selected; 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_gate; ds4_gpu_tensor *batch_indexer_q; ds4_gpu_tensor *batch_indexer_weights; ds4_gpu_tensor *batch_indexer_scores; ds4_gpu_tensor *batch_indexer_pool_selected; 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; ds4_gpu_tensor *mtp_state_backup; float *mtp_logits_host; int mtp_ready; ds4_gpu_tensor *layer_indexer_key_cache[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_indexer_tail_k[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_indexer_tail_gate[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 glm53; ds4_imatrix_collector *imatrix; ds4_gpu_tensor *directional_steering_dirs_by_tier[DS4_MAX_GPUS]; float directional_steering_attn_scale; float directional_steering_ffn_scale; 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 int glm_graph_directional_steering_tier( const ds4_glm_gpu_graph *g, uint32_t il) { if (!g || il < g->layer_start || il > g->layer_end) return -1; if (!g->placement) return 0; const int tier = g->placement[il + 1u]; return tier >= 0 && tier < DS4_MAX_GPUS ? tier : -1; } static bool glm_graph_load_directional_steering( ds4_glm_gpu_graph *g, const char *path, float attn_scale, float ffn_scale) { if (!g || (attn_scale == 0.0f && ffn_scale == 0.0f)) return g != NULL; if (!g->glm53) { fprintf(stderr, "ds4: directional steering is supported only for GLM 5.3\n"); return false; } if (!path || !path[0]) { fprintf(stderr, "ds4: directional steering needs --dir-steering-file\n"); return false; } const uint32_t n_layers = directional_steering_layer_count(); if (n_layers == 0 || n_layers != g->normal_layers) return false; const uint64_t n = (uint64_t)n_layers * DS4_N_EMBD; float *dirs = xmalloc((size_t)n * sizeof(dirs[0])); bool ok = read_f32_binary_file(path, dirs, n); bool used_tier[DS4_MAX_GPUS] = {false}; for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { const int tier = glm_graph_directional_steering_tier(g, il); if (tier < 0) { ok = false; break; } used_tier[tier] = true; } for (int tier = 0; ok && tier < DS4_MAX_GPUS; tier++) { if (!used_tier[tier]) continue; g->directional_steering_dirs_by_tier[tier] = ds4_gpu_tensor_alloc_ptr_on(tier, n * sizeof(dirs[0])); ok = g->directional_steering_dirs_by_tier[tier] != NULL && ds4_gpu_tensor_write(g->directional_steering_dirs_by_tier[tier], 0, dirs, n * sizeof(dirs[0])) != 0; } free(dirs); if (!ok) { fprintf(stderr, "ds4: failed to load GLM 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: GLM directional steering enabled: %s attn=%g ffn=%g\n", path, (double)attn_scale, (double)ffn_scale); return true; } static bool glm_graph_apply_directional_steering( ds4_glm_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows, float scale) { if (!g || !x || rows == 0 || scale == 0.0f) return true; const int tier = glm_graph_directional_steering_tier(g, il); if (tier < 0 || !g->directional_steering_dirs_by_tier[tier]) return false; return ds4_gpu_directional_steering_project_tensor( x, g->directional_steering_dirs_by_tier[tier], il, DS4_N_EMBD, rows, scale) != 0; } static bool glm_graph_apply_directional_steering_attn( ds4_glm_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows) { return glm_graph_apply_directional_steering( g, x, il, rows, g ? g->directional_steering_attn_scale : 0.0f); } static bool glm_graph_apply_directional_steering_ffn( ds4_glm_gpu_graph *g, ds4_gpu_tensor *x, uint32_t il, uint32_t rows) { return glm_graph_apply_directional_steering( g, x, il, rows, g ? g->directional_steering_ffn_scale : 0.0f); } static bool imatrix_collect_glm_one( ds4_imatrix_collector *c, struct ds4_glm_gpu_graph *g, uint32_t il) { if (!g || il >= DS4_N_LAYER) return false; return imatrix_collect_tensor_batch(c, g->ffn_norm, g->ffn_mid, g->router_selected, false, il, 1); } static bool glm_graph_reset_kda_state(ds4_glm_gpu_graph *g) { if (!g || !g->glm53) return true; for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { ds4_gpu_tensor *conv = g->layer_kda_conv_state[il]; ds4_gpu_tensor *recurrent = g->layer_kda_recurrent_state[il]; if (!conv && !recurrent) continue; if (!conv || !recurrent || ds4_gpu_tensor_fill_f32(conv, 0.0f, ds4_gpu_tensor_bytes(conv) / sizeof(float)) == 0 || ds4_gpu_tensor_fill_f32(recurrent, 0.0f, ds4_gpu_tensor_bytes(recurrent) / sizeof(float)) == 0) { return false; } } return ds4_gpu_synchronize() != 0; } 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 uint64_t glm_graph_saturating_add_u64(uint64_t a, uint64_t b) { return a > UINT64_MAX - b ? UINT64_MAX : a + b; } 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_streaming_active_model_bytes( const ds4_weights *weights) { if (!weights) return 0; uint64_t max_bytes = 0; const char *max_group = "none"; uint32_t max_layer = UINT32_MAX; ds4_model_map_span_vec spans; if (weights_layer_has_required(&weights->layer[0], 0) && weights_model_map_token_spans(weights, &spans)) { max_bytes = model_map_span_vec_total_bytes(&spans); max_group = "token"; free(spans.v); } if (weights_have_output_head(weights) && weights_model_map_output_spans(weights, &spans)) { const uint64_t bytes = model_map_span_vec_total_bytes(&spans); if (bytes > max_bytes) { max_bytes = bytes; max_group = "output"; } 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; max_group = "layer"; max_layer = il; } free(spans.v); } const char *report = getenv("DS4_GLM_MEMORY_GUARD_REPORT"); if (report && report[0]) { if (max_layer == UINT32_MAX) { fprintf(stderr, "ds4: GLM streaming largest active model group %s: %.2f GiB\n", max_group, glm_graph_bytes_to_gib(max_bytes)); } else { fprintf(stderr, "ds4: GLM streaming largest active model group layer %u: %.2f GiB\n", max_layer, glm_graph_bytes_to_gib(max_bytes)); } } 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 bool glm_graph_memory_guard_budget( uint64_t model_bytes, bool load_slice, bool ssd_streaming, uint64_t *budget_base_out, uint64_t *budget_out, double *fraction_out, double *reserve_gib_out) { (void)load_slice; #ifndef DS4_ROCM_BUILD (void)ssd_streaming; #endif uint64_t budget_base = glm_graph_host_memory_bytes(); if (budget_base == 0) { budget_base = ds4_gpu_recommended_working_set_size(); } #ifdef DS4_ROCM_BUILD uint64_t host_available = 0; const bool host_memory_known = ds4_linux_nonmovable_memory(&host_available); if (!ssd_streaming && host_memory_known) { /* The resident model is charged in model_bytes below. Keep the * pre-upload availability baseline so later session guards do not * charge the same ROCm allocation once through MemAvailable too. */ if (host_available > g_glm_rocm_guard_available_baseline) { g_glm_rocm_guard_available_baseline = host_available; } host_available = g_glm_rocm_guard_available_baseline; } if (host_memory_known && host_available < budget_base) { budget_base = host_available; } #endif if (budget_base == 0) return false; const double fraction = glm_graph_env_double("DS4_GLM_MEMORY_GUARD_FRACTION", 0.99, 0.50, 1.00); double default_reserve_gib = glm_graph_memory_guard_default_reserve_gib( budget_base, model_bytes, ds4_model_is_glm53()); #ifdef DS4_ROCM_BUILD if (!ssd_streaming) { /* The host baseline already excludes CMA. Reserve usable OS memory, * rather than applying the Metal physical-memory reserve a second time. */ double rocm_reserve_gib = glm_graph_bytes_to_gib(budget_base) / 16.0; if (rocm_reserve_gib < 8.0) rocm_reserve_gib = 8.0; if (rocm_reserve_gib < default_reserve_gib) { default_reserve_gib = rocm_reserve_gib; } } #endif const double reserve_gib = glm_graph_env_double("DS4_GLM_MEMORY_GUARD_RESERVE_GB", default_reserve_gib, 0.0, 1024.0); const uint64_t budget = glm_graph_memory_guard_budget_bytes( budget_base, glm_graph_wired_limit_bytes(), fraction, reserve_gib); if (budget_base_out) *budget_base_out = budget_base; if (budget_out) *budget_out = budget; if (fraction_out) *fraction_out = fraction; if (reserve_gib_out) *reserve_gib_out = reserve_gib; return true; } 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 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) return true; if (glm_graph_memory_guard_disabled()) return true; const uint32_t work_ctx = glm_graph_full_attention_cap(ctx_size, ssd_streaming); const ds4_context_memory mem = load_slice ? glm_graph_context_memory_estimate_for_compact_cap_slice( ctx_size, work_ctx, compact_cap, ssd_streaming, layer_start, layer_end) : glm_graph_context_memory_estimate_for_compact_cap( ctx_size, work_ctx, compact_cap, ssd_streaming); const uint64_t graph_bytes = mem.total_bytes; const uint64_t model_bytes = glm_graph_model_bytes_for_guard(model, weights, ssd_streaming, load_slice, layer_start, layer_end, include_token, include_output); uint64_t required = glm_graph_saturating_add_u64(model_bytes, graph_bytes); required = glm_graph_saturating_add_u64(required, transient_extra_bytes); uint64_t budget_base = 0; uint64_t budget = 0; double fraction = 0.0; double reserve_gib = 0.0; if (!glm_graph_memory_guard_budget(model_bytes, load_slice, ssd_streaming, &budget_base, &budget, &fraction, &reserve_gib)) { return true; } 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, " "tensor parallelism, or 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 = ds4_model_is_glm53() ? glm53_graph_indexer_selected_limit() : 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) { #ifndef __APPLE__ (void)n_selected; return false; #else 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(); #endif } 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->glm53 || (g->batch_indexer_gate && g->batch_indexer_pool_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 bool glm53_graph_use_indexed_prefill( const ds4_glm_gpu_graph *g) { #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) return g && g->glm53 && g->indexed_prefill_cap != 0; #else return g && g->glm53 && g->indexed_prefill_cap != 0 && !g->full_kv_cache; #endif } static uint32_t glm_graph_dense_compact_attention_limit( const ds4_glm_gpu_graph *g) { /* The incomplete GLM 5.3 pool is always included. Dense attention is * equivalent through top_k + pool_size - 1, not the work-window size. */ return g && g->glm53 ? glm53_graph_indexer_selected_limit() : glm_graph_indexer_top_k_limit(); } static uint32_t glm_graph_limit_indexed_prefill_chunk( const ds4_glm_gpu_graph *g, uint32_t pos, uint32_t chunk) { #ifndef DS4_ROCM_BUILD /* GLM 5.3's resident Metal/CUDA graph can switch its attention slices * from causal dense to indexed sparse inside one layer-major chunk. */ if (g && g->glm53 && glm_graph_indexed_prefill_batch_available(g) && getenv("DS4_GLM53_DISABLE_MIXED_PREFILL") == NULL) { return chunk; } #endif const uint32_t dense_limit = glm_graph_dense_compact_attention_limit(g); if (pos < dense_limit) { const uint32_t bridge = dense_limit - 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; if (ds4_model_is_glm53()) { if (glm_graph_env_truthy( getenv("DS4_GLM53_DISABLE_INDEXED_PREFILL"))) { return 0; } uint32_t chunk = DS4_GLM53_PREFILL_CHUNK_TOKENS; if (compact_cap != 0 && chunk > compact_cap) chunk = compact_cap; return chunk; } 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; const uint64_t score_columns = ds4_model_is_glm53() ? glm53_graph_indexer_pool_cap(compact_cap) : compact_cap; uint64_t budget_rows = budget_bytes / (score_columns * 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_tp_validate_ownership_kernels( const ds4_weights *weights, uint32_t *bad_layer, uint32_t *bad_type) { if (!weights) return false; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &weights->layer[il]; if (!l->ffn_gate_exps) continue; if (glm_graph_layer_uses_generic_routed_moe(l) || l->ffn_gate_exps->type == DS4_TENSOR_Q2_K || l->ffn_gate_exps->type == DS4_TENSOR_Q4_K) { continue; } if (bad_layer) *bad_layer = il; if (bad_type) *bad_type = l->ffn_gate_exps->type; return false; } return true; } 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_expert_cache_ready(weights, l, il); } static bool glm_graph_stream_prefill_expert_addr_supported( const ds4_glm_gpu_graph *g, const ds4_weights *weights, const ds4_layer_weights *l, uint32_t il, uint32_t n_tokens) { (void)g; 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. */ const bool selected_layout = glm_stream_selected_expert_cache_supported(l, il) || (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)); if (!selected_layout) return false; 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; } uint64_t required = (uint64_t)n_tokens * DS4_N_EXPERT_USED; if (required > DS4_N_EXPERT) required = DS4_N_EXPERT; return ds4_gpu_stream_expert_cache_budget_for_expert_size( gate_expert_bytes, down_expert_bytes) >= required; #else #ifdef __APPLE__ if (g && g->tp_world < 2 && l && l->ffn_gate_exps && l->ffn_up_exps && l->ffn_down_exps && 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 && DS4_N_EXPERT_USED == 8 && n_tokens < 256 && getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL && getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL && getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL && glm_stream_selected_expert_cache_supported(l, il)) { uint64_t gate_bytes = 0, down_bytes = 0; if (streaming_layer_gate_down_expert_bytes(l, &gate_bytes, &down_bytes) && ds4_gpu_stream_expert_cache_budget_for_expert_size(gate_bytes, down_bytes) >= DS4_N_EXPERT) return true; } #endif 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(g, 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; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_mix = (uint64_t)DS4_N_HC * (DS4_N_HC + 2u); const uint64_t kda_projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; if (ds4_model_is_glm53()) { if (!glm_graph_tensor_layout(l->hc_attn_fn, DS4_TENSOR_BF16, 2, hc_dim, hc_mix, 0) || !glm_graph_tensor_layout(l->hc_attn_scale, DS4_TENSOR_F32, 1, 3, 0, 0) || !glm_graph_tensor_layout(l->hc_attn_base, DS4_TENSOR_F32, 1, hc_mix, 0, 0) || !glm_graph_tensor_layout(l->hc_ffn_fn, DS4_TENSOR_BF16, 2, hc_dim, hc_mix, 0) || !glm_graph_tensor_layout(l->hc_ffn_scale, DS4_TENSOR_F32, 1, 3, 0, 0) || !glm_graph_tensor_layout(l->hc_ffn_base, DS4_TENSOR_F32, 1, hc_mix, 0, 0)) { fprintf(stderr, "ds4: GLM-5.3 graph found unexpected mHC layout in layer %u\n", il); return false; } if (ds4_glm53_layer_is_kda(il)) { if (!glm_graph_tensor_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0) || !glm_graph_dense_tensor_layout(l->kda_q, 2, DS4_N_EMBD, kda_projection, 0) || !glm_graph_dense_tensor_layout(l->kda_k, 2, DS4_N_EMBD, kda_projection, 0) || !glm_graph_dense_tensor_layout(l->kda_v, 2, DS4_N_EMBD, kda_projection, 0) || !glm_graph_dense_tensor_layout(l->kda_f_a, 2, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, 0) || !glm_graph_dense_tensor_layout(l->kda_f_b, 2, DS4_N_KDA_HEAD_DIM, kda_projection, 0) || !glm_graph_dense_tensor_layout(l->kda_beta, 2, DS4_N_EMBD, DS4_N_KDA_HEAD, 0) || !glm_graph_dense_tensor_layout(l->kda_g_a, 2, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, 0) || !glm_graph_dense_tensor_layout(l->kda_g_b, 2, DS4_N_KDA_HEAD_DIM, kda_projection, 0) || !glm_graph_dense_tensor_layout(l->kda_output, 2, kda_projection, DS4_N_EMBD, 0) || !glm_graph_tensor_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0)) { fprintf(stderr, "ds4: GLM-5.3 graph found unexpected KDA layout in layer %u\n", il); return false; } goto validate_ffn; } } 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) || !(ds4_model_is_glm53() ? glm_graph_dense_tensor_layout(l->indexer_proj, 2, DS4_N_EMBD, DS4_N_INDEXER_HEAD, 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; validate_ffn: 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); g->glm53 = ds4_model_is_glm53(); 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->batch_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; for (int tier = 0; tier < DS4_MAX_GPUS; tier++) { ds4_gpu_tensor_free(g->directional_steering_dirs_by_tier[tier]); g->directional_steering_dirs_by_tier[tier] = NULL; } 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); ds4_gpu_tensor_free(g->mtp_state_backup); 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_state_backup = 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_kda_conv_state[il]); ds4_gpu_tensor_free(g->layer_kda_recurrent_state[il]); ds4_gpu_tensor_free(g->layer_indexer_key_cache[il]); ds4_gpu_tensor_free(g->layer_indexer_tail_k[il]); ds4_gpu_tensor_free(g->layer_indexer_tail_gate[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_kda_out); ds4_gpu_tensor_free(g->batch_kda_output_gate); ds4_gpu_tensor_free(g->batch_kda_raw_beta); ds4_gpu_tensor_free(g->batch_kda_raw_gate); ds4_gpu_tensor_free(g->batch_kda_lowrank); ds4_gpu_tensor_free(g->batch_kda_v); ds4_gpu_tensor_free(g->batch_kda_k); ds4_gpu_tensor_free(g->batch_kda_q); ds4_gpu_tensor_free(g->session_batch_logits); ds4_gpu_tensor_free(g->batch_hc_after_attn); ds4_gpu_tensor_free(g->batch_hc_split); ds4_gpu_tensor_free(g->batch_hc_mix); ds4_gpu_tensor_free(g->batch_hc_flat); ds4_gpu_tensor_free(g->batch_hc_next); ds4_gpu_tensor_free(g->batch_hc_cur); 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_pool_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_gate); 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->kda_out); ds4_gpu_tensor_free(g->kda_output_gate); ds4_gpu_tensor_free(g->kda_raw_beta); ds4_gpu_tensor_free(g->kda_raw_gate); ds4_gpu_tensor_free(g->kda_lowrank); ds4_gpu_tensor_free(g->kda_v); ds4_gpu_tensor_free(g->kda_k); ds4_gpu_tensor_free(g->kda_q); 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_pool_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_gate); 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->hc_mean_weights); ds4_gpu_tensor_free(g->hc_output); ds4_gpu_tensor_free(g->hc_after_attn); ds4_gpu_tensor_free(g->hc_comb); ds4_gpu_tensor_free(g->hc_post); ds4_gpu_tensor_free(g->hc_split); ds4_gpu_tensor_free(g->hc_mix); ds4_gpu_tensor_free(g->hc_flat); ds4_gpu_tensor_free(g->hc_next); ds4_gpu_tensor_free(g->hc_cur); 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; /* Prefill leaves only the last layer/head mapped in streaming mode. * Warmup reads indexer normalization weights from every indexer layer. */ if (g->ssd_streaming && !g->streaming_static_decode_map_current) { if (!metal_graph_stream_map_decode_static_all(model, weights)) return false; g->streaming_static_decode_map_current = true; } 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 glm53_graph_prefill_workspace_ensure( ds4_glm_gpu_graph *g, uint32_t rows); 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->weights = weights; 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 session ctx=%u (model max=%u); " "prefill/work cap=%u; dense attention limit=%u\n", g->ctx_size, glm_graph_model_context_limit(), g->ctx_cap, glm_graph_dense_compact_attention_limit(g)); } const uint64_t emb_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_bytes = hc_dim * sizeof(float); const uint64_t hc_mix = (uint64_t)DS4_N_HC * (DS4_N_HC + 2u); const uint64_t hc_mix_bytes = hc_mix * sizeof(float); const uint64_t kda_projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; const uint64_t kda_projection_bytes = kda_projection * sizeof(float); const uint64_t kda_conv_state_bytes = 3u * (DS4_N_KDA_CONV - 1u) * kda_projection_bytes; const uint64_t kda_recurrent_state_bytes = kda_projection * DS4_N_KDA_HEAD_DIM * 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_gate_bytes = indexer_k_bytes; 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->glm53 ? glm53_graph_indexer_pool_cap(g->compact_cache_cap) : 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 uint32_t indexer_selected_limit = g->glm53 ? glm53_graph_indexer_selected_limit() : indexer_top_k; const uint32_t dense_indexer_selected_limit = g->glm53 && !g->full_kv_cache ? g->ctx_cap : indexer_selected_limit; const uint32_t indexer_selected_pools = g->glm53 ? indexer_top_k / DS4_GLM53_INDEX_POOL_SIZE : 0u; const uint64_t indexer_selected_bytes = (uint64_t)dense_indexer_selected_limit * sizeof(uint32_t); const uint64_t indexer_pool_selected_bytes = (uint64_t)indexer_selected_pools * 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 uint32_t compact_indexer_rows = g->glm53 ? glm53_graph_indexer_pool_cap(g->compact_cache_cap) : g->compact_cache_cap; const uint64_t compact_indexer_key_bytes = (uint64_t)compact_indexer_rows * DS4_N_INDEXER_HEAD_DIM * glm_graph_compact_cache_elem_bytes(); const uint64_t compact_indexer_tail_bytes = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); const uint64_t batch_rows = glm_graph_batch_row_cap(g->ctx_cap, g->indexed_prefill_cap, g->full_kv_cache); g->batch_cap = (uint32_t)batch_rows; 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_gate_bytes = batch_rows * indexer_gate_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_selected_limit * sizeof(uint32_t); const uint64_t batch_indexer_pool_selected_bytes = indexed_batch_rows * indexer_selected_pools * 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 uint32_t attention_layers = g->glm53 ? g->indexer_full_layers : g->layer_count; const double cache_gib = (double)(attention_layers * (key_cache_bytes + value_cache_bytes)) / (1024.0 * 1024.0 * 1024.0); if (g->full_kv_cache) { fprintf(stderr, "ds4: GLM graph allocating full-attention KV cache: work_ctx=%u layers=%u:%u (%u) %s %.2f GiB\n", g->ctx_cap, g->layer_start, g->layer_end, g->layer_count, "f16", cache_gib); } else { fprintf(stderr, "ds4: GLM graph using compact DSA KV only; expanded full-attention KV cache is skipped\n"); } #ifdef DS4_ROCM_BUILD if (glm_graph_env_truthy( getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE"))) { fprintf(stderr, "ds4: ROCm GLM one-token layer slices use the optimized token graph\n"); } #endif if (g->compact_cache_cap != 0) { const uint64_t compact_kv_total = (uint64_t)attention_layers * (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 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); if (g->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_cur, hc_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_next, hc_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_flat, hc_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_mix, hc_mix_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_split, hc_mix_bytes); g->hc_post = ds4_gpu_tensor_view( g->hc_split, (uint64_t)DS4_N_HC * sizeof(float), (uint64_t)DS4_N_HC * sizeof(float)); g->hc_comb = ds4_gpu_tensor_view( g->hc_split, (uint64_t)(2u * DS4_N_HC) * sizeof(float), (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); if (!g->hc_post || !g->hc_comb) ok = false; DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_after_attn, hc_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_output, emb_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->hc_mean_weights, (uint64_t)DS4_N_HC * sizeof(float)); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_q, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_k, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_v, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_lowrank, (uint64_t)DS4_N_KDA_HEAD_DIM * sizeof(float)); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_raw_gate, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_raw_beta, (uint64_t)DS4_N_KDA_HEAD * sizeof(float)); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_output_gate, kda_projection_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR(g->kda_out, kda_projection_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); if (g->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_gate, indexer_gate_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); if (g->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR(g->indexer_pool_selected, indexer_pool_selected_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->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_gate, batch_indexer_gate_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); if (g->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR(g->batch_indexer_pool_selected, batch_indexer_pool_selected_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_glm53_layer_is_kda(il)) { 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_glm53_layer_is_kda(il)) { 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 != 0 ? compact_k_rope_bytes : 1u); if (glm_graph_layer_uses_full_indexer(il)) { DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_indexer_key_cache[il], compact_indexer_key_bytes); if (g->glm53) { DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_indexer_tail_k[il], 2u * compact_indexer_tail_bytes); if (g->layer_indexer_tail_k[il]) { g->layer_indexer_tail_gate[il] = ds4_gpu_tensor_view( g->layer_indexer_tail_k[il], compact_indexer_tail_bytes, compact_indexer_tail_bytes); if (!g->layer_indexer_tail_gate[il]) { fprintf(stderr, "ds4: GLM graph could not create the " "indexer tail gate view on tier %d\n", cache_tier); ok = false; } } } } } if (g->glm53 && ds4_glm53_layer_is_kda(il)) { DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_kda_conv_state[il], kda_conv_state_bytes); DS4_GLM_GRAPH_ALLOC_TENSOR_TIER(g->layer_kda_recurrent_state[il], kda_recurrent_state_bytes); } #undef DS4_GLM_GRAPH_ALLOC_TENSOR_TIER } #undef DS4_GLM_GRAPH_ALLOC_TENSOR if (!ok) { glm_graph_free(g); return false; } if (g->glm53) { if (ds4_gpu_tensor_fill_f32(g->hc_mean_weights, 1.0f / (float)DS4_N_HC, DS4_N_HC) == 0) { glm_graph_free(g); return false; } for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { if (ds4_glm53_layer_is_kda(il)) { if (ds4_gpu_tensor_fill_f32(g->layer_kda_conv_state[il], 0.0f, kda_conv_state_bytes / sizeof(float)) == 0 || ds4_gpu_tensor_fill_f32(g->layer_kda_recurrent_state[il], 0.0f, kda_recurrent_state_bytes / sizeof(float)) == 0) { glm_graph_free(g); return false; } } else if (g->layer_indexer_tail_k[il] && g->layer_indexer_tail_gate[il]) { const uint64_t tail_elems = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM; if (ds4_gpu_tensor_fill_f32(g->layer_indexer_tail_k[il], 0.0f, tail_elems) == 0 || ds4_gpu_tensor_fill_f32(g->layer_indexer_tail_gate[il], 0.0f, tail_elems) == 0) { glm_graph_free(g); return false; } } } } if (g->glm53 && !glm53_graph_prefill_workspace_ensure(g, g->batch_cap)) { fprintf(stderr, "ds4: GLM graph could not allocate the static GLM-5.3 " "prefill workspace\n"); 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 void glm53_graph_prefill_workspace_free(ds4_glm_gpu_graph *g) { if (!g) return; ds4_gpu_tensor_free(g->batch_kda_out); ds4_gpu_tensor_free(g->batch_kda_output_gate); ds4_gpu_tensor_free(g->batch_kda_raw_beta); ds4_gpu_tensor_free(g->batch_kda_raw_gate); ds4_gpu_tensor_free(g->batch_kda_lowrank); ds4_gpu_tensor_free(g->batch_kda_v); ds4_gpu_tensor_free(g->batch_kda_k); ds4_gpu_tensor_free(g->batch_kda_q); ds4_gpu_tensor_free(g->batch_hc_after_attn); ds4_gpu_tensor_free(g->batch_hc_split); ds4_gpu_tensor_free(g->batch_hc_mix); ds4_gpu_tensor_free(g->batch_hc_flat); ds4_gpu_tensor_free(g->batch_hc_next); ds4_gpu_tensor_free(g->batch_hc_cur); g->batch_kda_out = NULL; g->batch_kda_output_gate = NULL; g->batch_kda_raw_beta = NULL; g->batch_kda_raw_gate = NULL; g->batch_kda_lowrank = NULL; g->batch_kda_v = NULL; g->batch_kda_k = NULL; g->batch_kda_q = NULL; g->batch_hc_after_attn = NULL; g->batch_hc_split = NULL; g->batch_hc_mix = NULL; g->batch_hc_flat = NULL; g->batch_hc_next = NULL; g->batch_hc_cur = NULL; g->glm53_prefill_cap = 0; } static bool glm53_graph_prefill_workspace_ensure( ds4_glm_gpu_graph *g, uint32_t rows) { if (!g || !g->glm53 || rows == 0 || rows > DS4_GLM53_PREFILL_CHUNK_TOKENS) { return false; } if (g->glm53_prefill_cap >= rows) return true; uint32_t cap = 1; while (cap < rows) cap <<= 1; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; const uint64_t hc_mix = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; const uint64_t projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; glm53_graph_prefill_workspace_free(g); #define DS4_GLM53_PREFILL_ALLOC(field_, elems_) do { \ g->field_ = ds4_gpu_tensor_alloc((uint64_t)cap * (elems_) * sizeof(float)); \ if (!g->field_) goto fail; \ } while (0) DS4_GLM53_PREFILL_ALLOC(batch_hc_cur, hc_dim); DS4_GLM53_PREFILL_ALLOC(batch_hc_next, hc_dim); DS4_GLM53_PREFILL_ALLOC(batch_hc_flat, hc_dim); DS4_GLM53_PREFILL_ALLOC(batch_hc_mix, hc_mix); DS4_GLM53_PREFILL_ALLOC(batch_hc_split, hc_mix); DS4_GLM53_PREFILL_ALLOC(batch_hc_after_attn, hc_dim); DS4_GLM53_PREFILL_ALLOC(batch_kda_q, projection); DS4_GLM53_PREFILL_ALLOC(batch_kda_k, projection); DS4_GLM53_PREFILL_ALLOC(batch_kda_v, projection); DS4_GLM53_PREFILL_ALLOC(batch_kda_lowrank, DS4_N_KDA_HEAD_DIM); DS4_GLM53_PREFILL_ALLOC(batch_kda_raw_gate, projection); DS4_GLM53_PREFILL_ALLOC(batch_kda_raw_beta, DS4_N_KDA_HEAD); DS4_GLM53_PREFILL_ALLOC(batch_kda_output_gate, projection); DS4_GLM53_PREFILL_ALLOC(batch_kda_out, projection); #undef DS4_GLM53_PREFILL_ALLOC g->glm53_prefill_cap = cap; return true; fail: #undef DS4_GLM53_PREFILL_ALLOC glm53_graph_prefill_workspace_free(g); return false; } static bool glm53_graph_session_batch_logits_ensure( ds4_glm_gpu_graph *g, uint32_t rows) { if (!g || !g->glm53 || rows == 0) return false; if (g->session_batch_logits_cap >= rows && g->session_batch_logits) { return true; } uint32_t cap = 1; while (cap < rows) { if (cap > UINT32_MAX / 2u) return false; cap <<= 1; } ds4_gpu_tensor *logits = ds4_gpu_tensor_alloc( (uint64_t)cap * DS4_N_VOCAB * sizeof(float)); if (!logits) return false; ds4_gpu_tensor_free(g->session_batch_logits); g->session_batch_logits = logits; g->session_batch_logits_cap = cap; return true; } 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 bool glm53_graph_matmul( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *weight, uint32_t in_dim, uint32_t out_dim, const ds4_gpu_tensor *x) { if (!out || !model || !weight || !x || !tensor_type_is_glm_dense_quant(weight->type)) { return false; } if (weight->type == DS4_TENSOR_BF16) { return ds4_gpu_glm53_matmul_bf16(out, model->map, model->size, weight->abs_offset, in_dim, out_dim, x, 1) != 0; } return ds4_gpu_matmul_quant_tensor(out, model->map, model->size, weight->abs_offset, weight->type, in_dim, out_dim, x, 1) != 0; } static bool glm53_graph_matmul_rows( ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *weight, uint32_t in_dim, uint32_t out_dim, const ds4_gpu_tensor *x, uint32_t n_rows) { if (!out || !model || !weight || !x || n_rows == 0 || !tensor_type_is_glm_dense_quant(weight->type)) { return false; } if (weight->type == DS4_TENSOR_BF16) { return ds4_gpu_glm53_matmul_bf16(out, model->map, model->size, weight->abs_offset, in_dim, out_dim, x, n_rows) != 0; } return ds4_gpu_matmul_quant_tensor(out, model->map, model->size, weight->abs_offset, weight->type, in_dim, out_dim, x, n_rows) != 0; } static bool glm53_graph_hc_pre_rows( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_tensor *fn, const ds4_tensor *scale, const ds4_tensor *base, const ds4_tensor *norm, const ds4_gpu_tensor *residual_hc, ds4_gpu_tensor *collapsed, ds4_gpu_tensor *normalized, ds4_gpu_tensor *flat, ds4_gpu_tensor *mix, ds4_gpu_tensor *split, uint32_t rows) { if (!g || !model || !fn || !scale || !base || !norm || !residual_hc || !collapsed || !normalized || !flat || !mix || !split || rows == 0) { return false; } const uint32_t hc_dim = DS4_N_HC * DS4_N_EMBD; const uint32_t hc_mix = DS4_N_HC * (DS4_N_HC + 2u); bool ok = ds4_gpu_rms_norm_plain_rows_tensor(flat, residual_hc, hc_dim, rows, DS4_RMS_EPS) != 0; if (ok) ok = glm53_graph_matmul_rows(mix, model, fn, hc_dim, hc_mix, flat, rows); if (ok) ok = metal_graph_decode_hc_pre(collapsed, split, mix, residual_hc, model, scale->abs_offset, base->abs_offset); if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(normalized, collapsed, model->map, model->size, norm->abs_offset, DS4_N_EMBD, rows, DS4_RMS_EPS) != 0; return ok; } static bool glm53_graph_kda_attention_rows( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *l, uint32_t il, uint32_t pos0, uint32_t rows, ds4_gpu_tensor *attn_out) { if (!g || !model || !l || il >= DS4_MAX_LAYER || rows == 0 || !attn_out || !g->layer_kda_conv_state[il] || !g->layer_kda_recurrent_state[il]) { return false; } const uint32_t projection = DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; const char *failed_stage = "Q projection"; const ds4_tensor *failed_weight = l->kda_q; bool ok = glm53_graph_matmul_rows(g->batch_kda_q, model, l->kda_q, DS4_N_EMBD, projection, g->batch_attn_norm, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_q_ready", g->batch_kda_q, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "K projection"; if (ok) failed_weight = l->kda_k; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_k, model, l->kda_k, DS4_N_EMBD, projection, g->batch_attn_norm, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_k_ready", g->batch_kda_k, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "V projection"; if (ok) failed_weight = l->kda_v; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_v, model, l->kda_v, DS4_N_EMBD, projection, g->batch_attn_norm, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_v_ready", g->batch_kda_v, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "decay low-rank projection"; if (ok) failed_weight = l->kda_f_a; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_lowrank, model, l->kda_f_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->batch_attn_norm, rows); if (ok) failed_stage = "decay gate projection"; if (ok) failed_weight = l->kda_f_b; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_raw_gate, model, l->kda_f_b, DS4_N_KDA_HEAD_DIM, projection, g->batch_kda_lowrank, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_raw_gate_ready", g->batch_kda_raw_gate, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "beta projection"; if (ok) failed_weight = l->kda_beta; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_raw_beta, model, l->kda_beta, DS4_N_EMBD, DS4_N_KDA_HEAD, g->batch_attn_norm, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_raw_beta_ready", g->batch_kda_raw_beta, (uint64_t)rows * DS4_N_KDA_HEAD, il, pos0); if (ok) failed_stage = "output-gate low-rank projection"; if (ok) failed_weight = l->kda_g_a; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_lowrank, model, l->kda_g_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->batch_attn_norm, rows); if (ok) failed_stage = "output-gate projection"; if (ok) failed_weight = l->kda_g_b; if (ok) ok = glm53_graph_matmul_rows(g->batch_kda_output_gate, model, l->kda_g_b, DS4_N_KDA_HEAD_DIM, projection, g->batch_kda_lowrank, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_output_gate_ready", g->batch_kda_output_gate, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "KDA recurrence"; if (ok) failed_weight = NULL; if (ok) ok = ds4_gpu_glm53_kda_prefill( g->batch_kda_out, g->layer_kda_conv_state[il], g->layer_kda_recurrent_state[il], g->batch_kda_q, g->batch_kda_k, g->batch_kda_v, g->batch_kda_raw_gate, g->batch_kda_raw_beta, g->batch_kda_output_gate, model->map, model->size, l->kda_q_conv->abs_offset, l->kda_k_conv->abs_offset, l->kda_v_conv->abs_offset, l->kda_a_log->abs_offset, l->kda_dt_bias->abs_offset, l->kda_o_norm->abs_offset, DS4_N_KDA_HEAD, rows, DS4_KDA_GATE_LOWER_BOUND, DS4_RMS_EPS) != 0; if (ok) metal_graph_debug_dump_tensor( "glm53_kda_out_ready", g->batch_kda_out, (uint64_t)rows * projection, il, pos0); if (ok) failed_stage = "output projection"; if (ok) failed_weight = l->kda_output; if (ok) ok = glm53_graph_matmul_rows(attn_out, model, l->kda_output, projection, DS4_N_EMBD, g->batch_kda_out, rows); if (ok) metal_graph_debug_dump_tensor( "glm53_kda_attn_out_ready", attn_out, (uint64_t)rows * DS4_N_EMBD, il, pos0); if (!ok) { if (failed_weight) { fprintf(stderr, "ds4: GLM-5.3 KDA failed at layer %u stage '%s' on " "tensor %.*s (%s; pos %u, rows %u)\n", il, failed_stage, (int)failed_weight->name.len, failed_weight->name.ptr, tensor_type_name(failed_weight->type), pos0, rows); } else { fprintf(stderr, "ds4: GLM-5.3 KDA failed at layer %u stage '%s' " "(pos %u, rows %u)\n", il, failed_stage, pos0, rows); } } return ok; } static bool glm53_graph_hc_pre( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_tensor *fn, const ds4_tensor *scale, const ds4_tensor *base, const ds4_tensor *norm, const ds4_gpu_tensor *residual_hc, ds4_gpu_tensor *collapsed, ds4_gpu_tensor *normalized) { if (!g || !model || !fn || !scale || !base || !norm || !residual_hc || !collapsed || !normalized) { return false; } const uint32_t hc_dim = DS4_N_HC * DS4_N_EMBD; const uint32_t hc_mix = DS4_N_HC * (DS4_N_HC + 2u); bool ok = ds4_gpu_rms_norm_plain_tensor(g->hc_flat, residual_hc, hc_dim, DS4_RMS_EPS) != 0; if (ok) ok = glm53_graph_matmul(g->hc_mix, model, fn, hc_dim, hc_mix, g->hc_flat); if (ok) ok = metal_graph_decode_hc_pre(collapsed, g->hc_split, g->hc_mix, residual_hc, model, scale->abs_offset, base->abs_offset); if (ok) ok = ds4_gpu_rms_norm_weight_tensor(normalized, collapsed, model->map, model->size, norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; return ok; } static bool glm53_graph_kda_attention( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *l, uint32_t il) { if (!g || !model || !l || il >= DS4_MAX_LAYER || !g->layer_kda_conv_state[il] || !g->layer_kda_recurrent_state[il]) { return false; } const uint32_t projection = DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; bool qk_paired = false; #if defined(__APPLE__) bool qkv_paired = false; if (getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE") == NULL && getenv("DS4_METAL_DISABLE_GLM53_BF16_QKV") == NULL && l->kda_q->type == DS4_TENSOR_BF16 && l->kda_k->type == DS4_TENSOR_BF16 && l->kda_v->type == DS4_TENSOR_BF16) { qkv_paired = ds4_gpu_glm53_matmul_bf16_qkv( g->kda_q, g->kda_k, g->kda_v, model->map, model->size, l->kda_q->abs_offset, l->kda_k->abs_offset, l->kda_v->abs_offset, DS4_N_EMBD, projection, g->attn_norm) != 0; } #else const bool qkv_paired = false; #endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) if (l->kda_q->type == DS4_TENSOR_Q4_K && l->kda_k->type == DS4_TENSOR_Q4_K && getenv("DS4_CUDA_GLM_DISABLE_KDA_QK_PAIR") == NULL) { qk_paired = ds4_gpu_matmul_q4_K_pair_decode_tensor( g->kda_q, g->kda_k, model->map, model->size, l->kda_q->abs_offset, l->kda_k->abs_offset, DS4_N_EMBD, projection, g->attn_norm) != 0; } #endif bool ok = qkv_paired || qk_paired || glm53_graph_matmul(g->kda_q, model, l->kda_q, DS4_N_EMBD, projection, g->attn_norm); if (ok && !qkv_paired && !qk_paired) { ok = glm53_graph_matmul(g->kda_k, model, l->kda_k, DS4_N_EMBD, projection, g->attn_norm); } if (ok && !qkv_paired) { ok = glm53_graph_matmul(g->kda_v, model, l->kda_v, DS4_N_EMBD, projection, g->attn_norm); } if (ok) ok = glm53_graph_matmul( g->kda_lowrank, model, l->kda_f_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); if (ok) ok = glm53_graph_matmul( g->kda_raw_gate, model, l->kda_f_b, DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); if (ok) ok = glm53_graph_matmul( g->kda_raw_beta, model, l->kda_beta, DS4_N_EMBD, DS4_N_KDA_HEAD, g->attn_norm); if (ok) ok = glm53_graph_matmul( g->kda_lowrank, model, l->kda_g_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, g->attn_norm); if (ok) ok = glm53_graph_matmul( g->kda_output_gate, model, l->kda_g_b, DS4_N_KDA_HEAD_DIM, projection, g->kda_lowrank); if (ok) ok = ds4_gpu_glm53_kda_decode( g->kda_out, g->layer_kda_conv_state[il], g->layer_kda_recurrent_state[il], g->kda_q, g->kda_k, g->kda_v, g->kda_raw_gate, g->kda_raw_beta, g->kda_output_gate, model->map, model->size, l->kda_q_conv->abs_offset, l->kda_k_conv->abs_offset, l->kda_v_conv->abs_offset, l->kda_a_log->abs_offset, l->kda_dt_bias->abs_offset, l->kda_o_norm->abs_offset, DS4_N_KDA_HEAD, 1, DS4_KDA_GATE_LOWER_BOUND, DS4_RMS_EPS) != 0; if (ok) ok = glm53_graph_matmul(g->attn_out, model, l->kda_output, projection, DS4_N_EMBD, g->kda_out); return ok; } 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 = glm53_graph_matmul(g->logits, model, weights->output, DS4_N_EMBD, DS4_N_VOCAB, g->output_norm); return ok; } static bool glm_graph_encode_output_head( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_weights *weights) { if (g->glm53) { if (ds4_gpu_hc_weighted_sum_tensor(g->hc_output, g->hc_cur, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) == 0) { return false; } return glm_graph_encode_output_head_from(g, model, weights, g->hc_output); } 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; const ds4_gpu_tensor *plain = hidden; bool ok = ds4_gpu_begin_commands() != 0; if (ok && g->glm53 && (hidden == g->hc_cur || hidden == g->hc_next)) { ok = ds4_gpu_hc_weighted_sum_tensor(g->hc_output, hidden, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) != 0; plain = g->hc_output; } if (ok) ok = glm_graph_encode_output_head_from(g, model, weights, plain); 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(plain, 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 && l->ffn_gate_exps->type != DS4_TENSOR_Q4_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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, 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, float clamp, 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, clamp) != 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, clamp, 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 add_residual, 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 streaming_expert_cache = g->ssd_streaming && !resident_decode_layer && glm_graph_stream_layer_expert_cache_supported(g->weights, l, il); const bool generic_streaming_selected_cache = streaming_expert_cache && glm_graph_layer_uses_generic_routed_moe(l); const bool uniform_streaming_selected_cache = streaming_expert_cache && 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->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, 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 && 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, g->ssd_streaming && !streaming_selected_cache) != 0; } if (ok && g->imatrix && !(glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { ok = imatrix_collect_glm_one(g->imatrix, g, il); } 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->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, 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 && add_residual && !glm_graph_disable_add3_residual()) { ok = ds4_gpu_add3_tensor(next, after_attn, ffn_out, ffn_sum, DS4_N_EMBD) != 0; } else if (ok && add_residual) { 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; } else if (ok) { ok = ds4_gpu_add_tensor(next, ffn_out, ffn_sum, 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 add_residual, 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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 && add_residual) { ok = ds4_gpu_add_tensor(next, after_attn, ffn_out, DS4_N_EMBD) != 0; } else if (ok) { ok = ds4_gpu_tensor_copy(next, 0, ffn_out, 0, (uint64_t)DS4_N_EMBD * sizeof(float)) != 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, add_residual, stage_profile, stage_t0); } static bool glm53_graph_encode_ffn_tail_one( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_layer_weights *l, uint32_t il, uint32_t pos, bool stage_profile, double *stage_t0) { bool 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, false, stage_profile, stage_t0); if (ok) { metal_graph_debug_dump_tensor("ffn_out", g->next, DS4_N_EMBD, il, pos); ok = glm_graph_apply_directional_steering_ffn(g, g->next, il, 1); } if (ok) { ok = ds4_gpu_hc_expand_tensor(g->hc_next, g->next, g->hc_after_attn, g->hc_post, g->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; } return ok; } 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, true, 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 DS4_MAYBE_UNUSED bool glm_graph_use_dense_compact_attention_prefill( uint32_t n_tokens) { /* This kernel absorbs MLA only; models with RoPE need the separate score * and the scale for the complete query/key dimension. */ if (DS4_N_ROT != 0) return false; if (!glm_graph_use_flash_attention_prefill(n_tokens)) return false; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) /* The CUDA GEMM setup crosses over the scalar online kernel near 256 * tokens on GB10. Keep short prompts on the lower-latency path. */ return g_n_gpus == 1 && n_tokens >= 256u; #else return true; #endif } 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( const ds4_glm_gpu_graph *g) { return glm53_graph_use_indexed_prefill(g); } 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) { if (g && g->glm53) { return g->compact_cache_cap != 0 && (!g->full_kv_cache || pos >= g->ctx_cap || glm_graph_force_indexed_decode(g)); } return g && g->compact_cache_cap != 0 && (!g->full_kv_cache || pos >= g->ctx_cap || glm_graph_force_indexed_decode(g) || (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; #elif defined(__APPLE__) && !defined(DS4_NO_GPU) if (!full_layer_prefill || !g || !model || !weights || !layer || !g->batch_router_selected || n_tokens == 0 || g->layer_count == 0 || g->tp_world >= 2 || g->layer_start != 0 || g->layer_count != glm_graph_normal_layer_count() || !weights_streaming_layer_experts_uniform(weights, il) || getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_CACHE_SEED") != NULL || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED || !glm_graph_streaming_expert_cache_seed_layer_expected(g, weights, layer, il)) { return true; } const uint32_t routed_layers = g->layer_count > DS4_N_LEADING_DENSE ? g->layer_count - DS4_N_LEADING_DENSE : g->layer_count; uint32_t target = ds4_gpu_stream_expert_cache_configured_count() / routed_layers; if (target > 128u) target = 128u; if (target == 0) return true; const uint32_t recent = n_tokens < 32u ? n_tokens : 32u; const size_t count = (size_t)n_tokens * DS4_N_EXPERT_USED; int32_t *selected = xmalloc(count * sizeof(*selected)); const bool had_commands = ds4_gpu_commands_active(); bool ok = !had_commands || ds4_gpu_end_commands() != 0; if (ok) ok = ds4_gpu_tensor_read(g->batch_router_selected, 0, selected, count * sizeof(*selected)) != 0; uint32_t frequency[DS4_MAX_EXPERT] = {0}; /* Keep recent selections first, then fill spare slots with the prompt's * most-used experts. A short tail alone often leaves much of the cache empty. */ for (size_t i = 0; ok && i < count; i++) { if (selected[i] < 0 || (uint32_t)selected[i] >= DS4_N_EXPERT) ok = false; else { frequency[selected[i]]++; if (i >= (size_t)(n_tokens - recent) * DS4_N_EXPERT_USED) frequency[selected[i]] += (uint32_t)count; } } free(selected); int32_t experts[128]; uint32_t priority[128]; uint32_t n = 0; while (ok && n < target) { uint32_t best = 0; for (uint32_t i = 1; i < DS4_N_EXPERT; i++) { if (frequency[i] > frequency[best]) best = i; } if (frequency[best] == 0) break; experts[n] = (int32_t)best; /* The recency bonus ranks the preload, not observed decode uses. * Keep it bounded so demand-loaded experts can replace the seed. */ priority[n] = metal_graph_streaming_builtin_hotness(target - n, target); n++; frequency[best] = 0; } if (ok && n != 0) { ok = ds4_gpu_begin_commands() != 0; const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make( model, layer, il, gate_expert_bytes, down_expert_bytes); if (ok) ok = ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( &table, experts, priority, n) != 0; } if (had_commands && !ds4_gpu_commands_active()) { if (ds4_gpu_begin_commands() == 0) ok = false; } else if (!had_commands && ds4_gpu_commands_active()) { if (ds4_gpu_end_commands() == 0) ok = false; } return ok; #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) { const char *failed_stage = "setup"; 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; failed_stage = "router projection"; 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) { failed_stage = "router selection"; 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++) { failed_stage = "router selection"; 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) { failed_stage = "TP bounce"; ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); } if (ok) { failed_stage = "routed experts"; 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) { failed_stage = "shared expert"; } 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) { failed_stage = "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; if (!ok) { fprintf(stderr, "ds4: GLM sparse batch FFN failed at layer %u stage '%s' " "(pos %u, tokens %u)\n", il, failed_stage, pos0, n_tokens); } 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 = g->glm53 || ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, after_attn, model->map, model->size, l->ffn_norm->abs_offset, DS4_N_EMBD, n_tokens, DS4_RMS_EPS) != 0; if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", "ffn_norm", il, pos0, n_tokens, stage_t0); if (ok) { metal_graph_debug_dump_tensor("glm_ffn_norm", g->batch_ffn_norm, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (!ok) return false; if (il < DS4_N_LEADING_DENSE) { const uint64_t hidden = l->ffn_gate->dim[1]; if (hidden == 0 || hidden > UINT32_MAX / n_tokens) return false; const uint32_t mid_elems = (uint32_t)(hidden * n_tokens); const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; if (residual_elems > UINT32_MAX) return false; const bool fused_gate_up = glm_graph_shared_gate_up_swiglu_q8_0_tensor( g->batch_ffn_gate, g->batch_ffn_up, g->batch_ffn_mid, model, l->ffn_gate->abs_offset, l->ffn_up->abs_offset, DS4_N_EMBD, hidden, g->batch_ffn_norm, n_tokens, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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) { metal_graph_debug_dump_tensor( "glm53_dense_ffn_mid", g->batch_ffn_mid, (uint64_t)n_tokens * hidden, il, pos0); } 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) { metal_graph_debug_dump_tensor( "glm53_dense_ffn_out", g->batch_ffn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", "dense_down", il, pos0, n_tokens, stage_t0); if (ok && g->glm53) { ok = ds4_gpu_tensor_copy(next, 0, g->batch_ffn_out, 0, residual_elems * sizeof(float)) != 0; } else 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) { fprintf(stderr, "ds4: GLM sparse FFN router projection failed at layer %u " "(pos %u, tokens %u)\n", il, pos0, n_tokens); } 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) { fprintf(stderr, "ds4: GLM sparse FFN router selection failed at layer %u " "(pos %u, tokens %u, experts %u, selected %u, scale %.6g)\n", il, pos0, n_tokens, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE); } } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", "router", il, pos0, n_tokens, stage_t0); if (ok) { metal_graph_debug_dump_tensor("glm_ffn_router_logits", g->batch_router_logits, (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); metal_graph_debug_dump_tensor("glm_ffn_router_probs", g->batch_router_probs, (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); metal_graph_debug_dump_i32_tensor("glm_ffn_router_selected", g->batch_router_selected, (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); metal_graph_debug_dump_tensor("glm_ffn_router_weights", g->batch_router_weights, (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); } if (ok) ok = glm_graph_profile_router_selection_batch(g, l, il, pos0, n_tokens); const bool tp_batch_split_ffn = g->tp_world == 2; if (ok && tp_batch_split_ffn) { ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); } if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, il, n_tokens); if (ok) ok = glm_graph_seed_streaming_expert_cache_from_full_layer( g, model, weights, l, il, n_tokens, gate_out * gate_row_bytes, down_out * down_row_bytes, full_layer_prefill); bool shared_done = false; #define DS4_GLM_ENCODE_FFN_BATCH_SHARED() do { \ if (ok) { \ const bool fused_shared = glm_graph_shared_gate_up_swiglu_q8_0_tensor( \ g->batch_ffn_gate, \ g->batch_ffn_up, \ g->batch_shared_mid, \ model, \ l->ffn_gate_shexp->abs_offset, \ l->ffn_up_shexp->abs_offset, \ DS4_N_EMBD, \ DS4_N_FF_EXP, \ g->batch_ffn_norm, \ n_tokens, \ g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 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, \ g->glm53 ? DS4_SWIGLU_CLAMP_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(g, weights, l, il, n_tokens); if (rocm_batch_selected_shared_overlap) { uint64_t selected_event = 0; if (ds4_gpu_signal_selected_readback_ready(&selected_event) != 0) { rocm_batch_selected_async_started = rocm_graph_batch_selected_async_load_start( &rocm_batch_selected_async, g->batch_router_selected, model, l, il, n_tokens, selected_event, gate_out * gate_row_bytes, down_out * down_row_bytes); } } if (rocm_batch_selected_async_started) { DS4_GLM_ENCODE_FFN_BATCH_SHARED(); const bool finish_ok = rocm_graph_batch_selected_async_load_finish(&rocm_batch_selected_async); if (!finish_ok) rocm_batch_selected_async_started = false; } #endif if (n_tokens <= 8u && (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ROUTED)) { /* ablate: keep the gate */ } else if (ok) ok = glm_graph_routed_moe_batch_dispatch( g, model, l, il, tp_batch_split_ffn ? g->tp_bounce_out : g->batch_ffn_out, g->batch_ffn_mid, gate_out * gate_row_bytes, gate_row_bytes, up_out * up_row_bytes, up_row_bytes, down_out * down_row_bytes, down_row_bytes, g->batch_router_selected, g->batch_router_weights, g->batch_ffn_norm, n_tokens, (uint32_t)g->ffn_mid_elems, full_layer_prefill, false) != 0; if (ok && g->tp_world == 2) { ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_ffn_out, n_tokens); if (!ok) fprintf(stderr, "ds4: GLM TP batch gate failed (layer %u)\n", il); } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", "routed_moe", il, pos0, n_tokens, stage_t0); if (ok) { metal_graph_debug_dump_tensor("glm_ffn_routed_out", g->batch_ffn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (ok && !shared_done) DS4_GLM_ENCODE_FFN_BATCH_SHARED(); #undef DS4_GLM_ENCODE_FFN_BATCH_SHARED if (ok) { metal_graph_debug_dump_tensor("glm_ffn_shared_out", g->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } if (ok && g->glm53) { ok = ds4_gpu_add_tensor(next, g->batch_ffn_out, g->batch_attn_out, (uint32_t)residual_elems) != 0; } else if (ok && !glm_graph_disable_add3_residual()) { ok = ds4_gpu_add3_tensor(next, after_attn, g->batch_ffn_out, g->batch_attn_out, (uint32_t)residual_elems) != 0; } else if (ok) { ok = ds4_gpu_add_tensor(g->batch_heads, g->batch_ffn_out, g->batch_attn_out, (uint32_t)residual_elems) != 0; if (ok) ok = ds4_gpu_add_tensor(next, after_attn, g->batch_heads, (uint32_t)residual_elems) != 0; } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", "residual", il, pos0, n_tokens, stage_t0); if (ok) { metal_graph_debug_dump_tensor("glm_ffn_next", next, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } return ok; } static bool glm_graph_begin_commands_if_needed(void); /* ------------------------------------------------------------------------ * GLM MTP (nextn block) drafting. * * blk.(N_LAYER-1) is GLM 5.2's multi-token-prediction block: a full * attention+MoE layer fed with eh_proj(concat(enorm(embed(token[p+1])), * hnorm(h[p]))), predicting token[p+2] through the shared output head. * It keeps a private compact KV cache (slot = absolute position; only * positions >= mtp_min_pos are ever selected, so the unwritten prompt * range is never read). Under TP the routed experts are combined over * the BIG-gate exchange, never the decode row gate, so the RDMA row-gate * schedule stays intact. * --------------------------------------------------------------------- */ static uint32_t glm_graph_mtp_cache_cap(const ds4_glm_gpu_graph *g) { if (!g) return 0; return g->compact_cache_cap != 0 ? g->compact_cache_cap : g->ctx_size; } static void glm53_graph_spec_state_tensors(const ds4_glm_gpu_graph *g, uint32_t il, ds4_gpu_tensor *state[2]) { if (ds4_glm53_layer_is_kda(il)) { state[0] = g->layer_kda_conv_state[il]; state[1] = g->layer_kda_recurrent_state[il]; } else { /* The K tail owns the contiguous K+gate allocation; the gate is a view. */ state[0] = g->layer_indexer_tail_k[il]; state[1] = NULL; } } static uint64_t glm53_graph_spec_state_bytes(const ds4_glm_gpu_graph *g) { if (!g || !g->glm53) return 0; uint64_t total = 0; for (uint32_t il = g->layer_start; il <= g->layer_end; il++) { ds4_gpu_tensor *state[2]; glm53_graph_spec_state_tensors(g, il, state); for (uint32_t i = 0; i < 2; i++) { const uint64_t bytes = state[i] ? ds4_gpu_tensor_bytes(state[i]) : 0; if (bytes > UINT64_MAX - total) return 0; total += bytes; } } return total; } static bool glm53_graph_copy_spec_state( ds4_glm_gpu_graph *g, bool save) { if (!g || !g->glm53 || !g->mtp_state_backup) return false; const uint64_t expected = glm53_graph_spec_state_bytes(g); if (expected == 0 || ds4_gpu_tensor_bytes(g->mtp_state_backup) < expected) { return false; } bool ok = glm_graph_begin_commands_if_needed(); uint64_t offset = 0; for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { ds4_gpu_tensor *state[2]; glm53_graph_spec_state_tensors(g, il, state); for (uint32_t i = 0; ok && i < 2; i++) { if (!state[i]) continue; const uint64_t bytes = ds4_gpu_tensor_bytes(state[i]); if (save) { ok = ds4_gpu_tensor_copy(g->mtp_state_backup, offset, state[i], 0, bytes) != 0; } else { ok = ds4_gpu_tensor_copy(state[i], 0, g->mtp_state_backup, offset, bytes) != 0; } offset += bytes; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); return ok && offset == expected; } static bool glm_graph_mtp_ensure(ds4_glm_gpu_graph *g) { if (g->mtp_ready) return true; const uint32_t cache_cap = glm_graph_mtp_cache_cap(g); if (cache_cap == 0) return false; const uint64_t elem = glm_graph_compact_cache_elem_bytes(); const uint64_t kv_bytes = (uint64_t)cache_cap * DS4_N_KV_LORA * elem; const uint64_t rope_bytes = (uint64_t)cache_cap * DS4_N_ROT * elem; g->mtp_kv_lora_cache = ds4_gpu_tensor_alloc(kv_bytes); /* GLM-5.3 has no rotary tail, but the compact-cache API still expects a * valid tensor handle for the zero-width companion cache. */ g->mtp_k_rope_cache = ds4_gpu_tensor_alloc( rope_bytes != 0 ? rope_bytes : sizeof(float)); g->mtp_concat = ds4_gpu_tensor_alloc(2ull * DS4_N_EMBD * sizeof(float)); g->mtp_selected = ds4_gpu_tensor_alloc((uint64_t)cache_cap * sizeof(int32_t)); const uint64_t state_backup_bytes = glm53_graph_spec_state_bytes(g); if (g->glm53 && state_backup_bytes != 0) { g->mtp_state_backup = ds4_gpu_tensor_alloc(state_backup_bytes); } 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->glm53 && !g->mtp_state_backup) || !g->mtp_logits_host) { fprintf(stderr, "ds4: glm mtp: allocation failed kv=%s rope=%s concat=%s " "selected=%s state_backup=%s logits=%s\n", g->mtp_kv_lora_cache ? "ok" : "missing", g->mtp_k_rope_cache ? "ok" : "missing", g->mtp_concat ? "ok" : "missing", g->mtp_selected ? "ok" : "missing", (!g->glm53 || g->mtp_state_backup) ? "ok" : "missing", g->mtp_logits_host ? "ok" : "missing"); 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); ds4_gpu_tensor_free(g->mtp_state_backup); 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_state_backup = NULL; g->mtp_logits_host = NULL; return false; } g->mtp_ready = 1; return true; } static bool glm_graph_mtp_embed_token( ds4_glm_gpu_graph *g, const ds4_model *model, const ds4_weights *weights, int token, ds4_gpu_tensor *out) { if (!g || !model || !weights || !weights->token_embd || !out || token < 0 || token >= (int)DS4_N_VOCAB) { return false; } if (g->glm53 && weights->token_embd->type == DS4_TENSOR_BF16) { const int32_t token_id = (int32_t)token; return ds4_gpu_tensor_write(g->prefill_tokens, 0, &token_id, sizeof(token_id)) != 0 && ds4_gpu_glm53_embedding_bf16(out, model->map, model->size, weights->token_embd->abs_offset, g->prefill_tokens, 1, DS4_N_EMBD, DS4_N_VOCAB) != 0; } return ds4_gpu_embed_token_quant_tensor(out, model->map, model->size, weights->token_embd->abs_offset, weights->token_embd->type, DS4_N_VOCAB, (uint32_t)token, DS4_N_EMBD) != 0; } static bool glm_graph_mtp_matmul( ds4_glm_gpu_graph *g, ds4_gpu_tensor *out, const ds4_model *model, const ds4_tensor *weight, uint32_t in_dim, uint32_t out_dim, const ds4_gpu_tensor *x) { if (!g || !out || !model || !weight || !x) return false; if (g->glm53) { return glm53_graph_matmul(out, model, weight, in_dim, out_dim, x); } return glm_graph_matmul_q8_0_decode_tensor(out, model, weight->abs_offset, in_dim, out_dim, x, false) != 0; } /* One MTP step at (absolute) position pos: consumes the main model's last * hidden h[pos] (g->cur for GLM-5.2, g->hc_cur for GLM-5.3) 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; const uint32_t cache_cap = glm_graph_mtp_cache_cap(g); if (pos >= cache_cap || min_pos > pos) { fprintf(stderr, "ds4: glm mtp: pos %u/min %u out of range (cap %u)\n", pos, min_pos, 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 = glm_graph_mtp_embed_token(g, model, weights, next_token, g->next); } 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", 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 = glm_graph_mtp_embed_token(g, model, weights, next_token, g->next); } 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"); const ds4_gpu_tensor *target_hidden = g->cur; if (ok && g->glm53) { ok = ds4_gpu_hc_weighted_sum_tensor(g->hc_output, g->hc_cur, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) != 0; target_hidden = g->hc_output; } if (ok) ok = ds4_gpu_rms_norm_weight_tensor(hnorm_view, target_hidden, 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_mtp_matmul(g, g->cur, model, l->nextn_eh_proj, 2u * DS4_N_EMBD, DS4_N_EMBD, g->mtp_concat); /* 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_mtp_matmul(g, g->q_rank, model, l->attn_q_a, DS4_N_EMBD, DS4_N_LORA_Q, g->attn_norm); 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_mtp_matmul(g, g->q, model, l->attn_q_b, DS4_N_LORA_Q, (uint32_t)g->q_dim, g->q_rank_norm); DS4_GLM_MTP_STAGE("rope"); if (ok && DS4_N_ROT != 0) 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_mtp_matmul(g, g->kv_raw, model, l->attn_kv_a_mqa, DS4_N_EMBD, kv_raw_dim, g->attn_norm); 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, 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, 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_mtp_matmul(g, g->attn_out, model, l->attn_output, (uint32_t)g->heads_dim, DS4_N_EMBD, g->heads); 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); const bool streaming_expert_cache = g->ssd_streaming && glm_graph_stream_layer_expert_cache_supported(weights, l, il); #ifdef DS4_ROCM_BUILD if (streaming_expert_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, }; 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, g->ssd_streaming && !streaming_expert_cache) != 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, g->glm53 ? DS4_SWIGLU_CLAMP_EXP : 0.0f, false, false, NULL); DS4_GLM_MTP_STAGE("shared_down"); if (ok) ok = glm_graph_mtp_matmul(g, g->ffn_sum, model, l->ffn_down_shexp, DS4_N_FF_EXP, DS4_N_EMBD, g->ffn_mid); 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_mtp_matmul(g, g->logits, model, weights->output, DS4_N_EMBD, DS4_N_VOCAB, g->output_norm); 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->glm53 || 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 bool glm53_kda = g->glm53 && ds4_glm53_layer_is_kda(il); const uint32_t kv_raw_dim = glm53_kda ? 0u : (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; } typedef struct { uint32_t dst_row; uint32_t src_row; uint32_t rows; } glm_vision_overlay_segment; typedef struct { ds4_gpu_tensor *tensor; glm_vision_overlay_segment *segments; size_t count; } glm_vision_overlay; static void glm_vision_overlay_free(glm_vision_overlay *overlay) { if (!overlay) return; ds4_gpu_tensor_free(overlay->tensor); free(overlay->segments); memset(overlay, 0, sizeof(*overlay)); } static bool glm_vision_overlay_prepare( glm_vision_overlay *overlay, const ds4_vision_span *images, size_t image_count, uint32_t pos0, uint32_t n_tokens) { memset(overlay, 0, sizeof(*overlay)); if (!images || image_count == 0) return true; const uint64_t chunk_end = (uint64_t)pos0 + n_tokens; uint64_t total_rows = 0; size_t segments = 0; for (size_t i = 0; i < image_count; i++) { const ds4_vision_span *span = &images[i]; if (!span->embedding.data || span->embedding.token_count == 0) return false; const uint64_t span_end = (uint64_t)span->token_start + span->embedding.token_count; const uint64_t begin = span->token_start > pos0 ? span->token_start : pos0; const uint64_t end = span_end < chunk_end ? span_end : chunk_end; if (begin < end) { total_rows += end - begin; segments++; } } if (segments == 0) return true; if (total_rows > UINT32_MAX || total_rows > UINT64_MAX / (4096u * sizeof(float))) return false; overlay->segments = xcalloc(segments, sizeof(overlay->segments[0])); overlay->tensor = ds4_gpu_tensor_alloc( total_rows * 4096u * sizeof(float)); if (!overlay->tensor) { glm_vision_overlay_free(overlay); return false; } uint32_t upload_row = 0; for (size_t i = 0; i < image_count; i++) { const ds4_vision_span *span = &images[i]; const uint64_t span_end = (uint64_t)span->token_start + span->embedding.token_count; const uint64_t begin = span->token_start > pos0 ? span->token_start : pos0; const uint64_t end = span_end < chunk_end ? span_end : chunk_end; if (begin >= end) continue; const uint32_t rows = (uint32_t)(end - begin); const uint32_t image_row = (uint32_t)(begin - span->token_start); const uint64_t bytes = (uint64_t)rows * 4096u * sizeof(float); const float *source = span->embedding.data + (uint64_t)image_row * 4096u; if (!ds4_gpu_tensor_write(overlay->tensor, (uint64_t)upload_row * 4096u * sizeof(float), source, bytes)) { glm_vision_overlay_free(overlay); return false; } overlay->segments[overlay->count++] = (glm_vision_overlay_segment) { .dst_row = (uint32_t)(begin - pos0), .src_row = upload_row, .rows = rows, }; upload_row += rows; } return true; } static bool vision_synthetic_token_is_covered( const int *tokens, uint32_t token_index, const ds4_vision_span *images, size_t image_count, uint32_t pos0) { const int token = tokens[token_index]; if (token < (int)DS4_N_VOCAB || token > (int)DS4_N_VOCAB + DS4_DEEPSEEK4_IMAGE_END) { return false; } const uint64_t absolute = (uint64_t)pos0 + token_index; for (size_t i = 0; i < image_count; i++) { const uint64_t begin = images[i].token_start; const uint64_t end = begin + images[i].embedding.token_count; if (absolute >= begin && absolute < end) return true; } return false; } 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, const ds4_vision_span *images, size_t image_count, 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->glm53 && n_tokens > DS4_GLM53_PREFILL_CHUNK_TOKENS) || 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 && !vision_synthetic_token_is_covered( tokens, i, images, image_count, pos0))) return false; } if (!input_hc && !g->has_token_embd) return false; if (input_hc && image_count != 0) return false; if (logits_out && !g->has_output_head) return false; if (g->glm53 && !glm53_graph_prefill_workspace_ensure(g, n_tokens)) { 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); } } } glm_vision_overlay vision_overlay = {0}; if (!glm_vision_overlay_prepare(&vision_overlay, images, image_count, pos0, n_tokens)) return false; 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); } } const uint64_t plain_bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); const uint64_t hc_bytes = plain_bytes * DS4_N_HC; ds4_gpu_tensor *cur_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_cur, 0, plain_bytes) : NULL; ds4_gpu_tensor *next_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_next, 0, plain_bytes) : NULL; ds4_gpu_tensor *after_attn_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_after_attn, 0, plain_bytes) : NULL; ds4_gpu_tensor *hc_cur_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_cur, 0, hc_bytes) : NULL; ds4_gpu_tensor *hc_next_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_next, 0, hc_bytes) : NULL; ds4_gpu_tensor *hc_after_attn_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_after_attn, 0, hc_bytes) : NULL; if (g->glm53 && (!cur_view || !next_view || !after_attn_view || !hc_cur_view || !hc_next_view || !hc_after_attn_view)) { ds4_gpu_tensor_free(hc_after_attn_view); ds4_gpu_tensor_free(hc_next_view); ds4_gpu_tensor_free(hc_cur_view); ds4_gpu_tensor_free(after_attn_view); ds4_gpu_tensor_free(next_view); ds4_gpu_tensor_free(cur_view); glm_vision_overlay_free(&vision_overlay); return false; } ds4_gpu_tensor *cur = g->glm53 ? cur_view : g->batch_cur; ds4_gpu_tensor *next = g->glm53 ? next_view : g->batch_next; ds4_gpu_tensor *hc_cur = hc_cur_view; ds4_gpu_tensor *hc_next = hc_next_view; ds4_gpu_tensor *last_hidden = NULL; uint32_t failed_layer = UINT32_MAX; const char *failed_stage = "setup"; 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(g->glm53 ? hc_cur : cur, 0, input_hc, (uint64_t)n_tokens * DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * 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; if (g->glm53 && weights->token_embd->type == DS4_TENSOR_BF16) { ok = ds4_gpu_glm53_embedding_bf16(cur, model->map, model->size, weights->token_embd->abs_offset, g->prefill_tokens, n_tokens, DS4_N_EMBD, DS4_N_VOCAB) != 0; } else { 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 (ok && g->glm53) { ok = ds4_gpu_repeat_hc_rows_tensor(hc_cur, cur, n_tokens, DS4_N_EMBD, DS4_N_HC) != 0; } for (size_t i = 0; ok && i < vision_overlay.count; i++) { const glm_vision_overlay_segment *segment = &vision_overlay.segments[i]; ok = ds4_gpu_glm53_scatter_image_hc( g->glm53 ? hc_cur : cur, vision_overlay.tensor, segment->dst_row, segment->src_row, segment->rows, n_tokens, DS4_N_EMBD, g->glm53 ? DS4_N_HC : 1u) != 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++) { failed_layer = 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 bool glm53_kda = g->glm53 && ds4_glm53_layer_is_kda(il); const uint32_t kv_raw_dim = glm53_kda ? 0u : (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); } failed_stage = "attention mHC pre"; if (ok && g->glm53) { ok = glm53_graph_hc_pre_rows(g, model, l->hc_attn_fn, l->hc_attn_scale, l->hc_attn_base, l->attn_norm, hc_cur, cur, g->batch_attn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, n_tokens); if (ok) { metal_graph_debug_dump_tensor( "glm53_attn_hc_input", hc_cur, (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor( "glm53_attn_hc_flat", g->batch_hc_flat, (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor( "glm53_attn_hc_mix", g->batch_hc_mix, (uint64_t)n_tokens * DS4_N_HC * (DS4_N_HC + 2u), il, pos0); metal_graph_debug_dump_tensor( "glm53_attn_hc_split", g->batch_hc_split, (uint64_t)n_tokens * DS4_N_HC * (DS4_N_HC + 2u), il, pos0); metal_graph_debug_dump_tensor( "glm53_attn_hc_collapsed", cur, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor( "glm53_attn_norm", g->batch_attn_norm, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } } else 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"); failed_stage = glm53_kda ? "KDA attention" : "DSA attention"; if (ok && glm53_kda) { ok = glm53_graph_kda_attention_rows(g, model, l, il, pos0, n_tokens, g->batch_attn_out); if (ok) { const uint64_t projection_rows = (uint64_t)n_tokens * DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; metal_graph_debug_dump_tensor("glm53_kda_q", g->batch_kda_q, projection_rows, il, pos0); metal_graph_debug_dump_tensor("glm53_kda_k", g->batch_kda_k, projection_rows, il, pos0); metal_graph_debug_dump_tensor("glm53_kda_v", g->batch_kda_v, projection_rows, il, pos0); metal_graph_debug_dump_tensor( "glm53_kda_raw_gate", g->batch_kda_raw_gate, projection_rows, il, pos0); metal_graph_debug_dump_tensor( "glm53_kda_raw_beta", g->batch_kda_raw_beta, (uint64_t)n_tokens * DS4_N_KDA_HEAD, il, pos0); metal_graph_debug_dump_tensor( "glm53_kda_output_gate", g->batch_kda_output_gate, projection_rows, il, pos0); metal_graph_debug_dump_tensor( "glm53_kda_out", g->batch_kda_out, projection_rows, il, pos0); metal_graph_debug_dump_tensor( "glm53_kda_attn_out", g->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } goto glm53_batch_attention_done; } failed_stage = "DSA query projection"; 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 && DS4_N_ROT != 0) 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; if (ok) { metal_graph_debug_dump_tensor( "glm53_dsa_q_lora", g->batch_q_rank, (uint64_t)n_tokens * DS4_N_LORA_Q, il, pos0); metal_graph_debug_dump_tensor( "glm53_dsa_q_lora_norm", g->batch_q_rank_norm, (uint64_t)n_tokens * DS4_N_LORA_Q, il, pos0); metal_graph_debug_dump_tensor( "glm53_dsa_q", g->batch_q, (uint64_t)n_tokens * g->q_dim, il, pos0); } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "q_path"); if (ok) failed_stage = "DSA compact indexer"; if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il)) { ok = g->glm53 ? glm53_graph_matmul_rows(g->batch_indexer_k, model, l->indexer_attn_k, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->batch_attn_norm, n_tokens) : 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 && g->glm53) { ok = glm53_graph_matmul_rows( g->batch_indexer_gate, model, l->indexer_compressor_gate, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->batch_attn_norm, n_tokens); } if (ok && g->glm53) { ok = ds4_gpu_glm53_indexer_pool_update_tensor( g->layer_indexer_key_cache[il], g->layer_indexer_tail_k[il], g->layer_indexer_tail_gate[il], g->batch_indexer_k, g->batch_indexer_gate, model->map, model->size, l->indexer_k_norm->abs_offset, l->indexer_k_norm_b->abs_offset, l->indexer_compressor_ape->abs_offset, pos0, n_tokens, g->compact_cache_cap, DS4_N_INDEXER_HEAD_DIM, DS4_GLM53_INDEX_POOL_SIZE, 1.0e-6f, glm_graph_compact_cache_is_f16()) != 0; } else 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; } if (ok) { metal_graph_debug_dump_tensor( "glm53_dsa_indexer_k", g->batch_indexer_k, (uint64_t)n_tokens * DS4_N_INDEXER_HEAD_DIM, il, pos0); if (g->glm53) { metal_graph_debug_dump_tensor( "glm53_dsa_indexer_gate", g->batch_indexer_gate, (uint64_t)n_tokens * DS4_N_INDEXER_HEAD_DIM, il, pos0); } } } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "indexer_k"); if (ok) failed_stage = "DSA KV-A projection"; 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) failed_stage = "DSA KV-A normalization"; 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) { failed_stage = "DSA compact KV store"; 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) failed_stage = "DSA K-B projection"; 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) failed_stage = "DSA V-B projection"; 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); if (ok) { metal_graph_debug_dump_tensor( "glm53_dsa_kv_raw", g->batch_kv_raw, (uint64_t)n_tokens * kv_raw_dim, il, pos0); metal_graph_debug_dump_tensor( "glm53_dsa_kv_norm", g->batch_kv_norm, (uint64_t)n_tokens * DS4_N_KV_LORA, il, pos0); metal_graph_debug_dump_tensor( "glm53_dsa_k_nope", g->batch_k_nope, (uint64_t)n_tokens * g->q_nope * DS4_N_HEAD, il, pos0); metal_graph_debug_dump_tensor( "glm53_dsa_value", g->batch_value, (uint64_t)n_tokens * g->heads_dim, il, pos0); } 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) { failed_stage = "DSA expanded KV store"; 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; } if (ok) { metal_graph_debug_dump_f16_tensor( "glm53_dsa_key_cache", g->layer_key_cache[il], (uint64_t)cache_len * DS4_N_HEAD * DS4_N_KEY_MLA, il, pos0); metal_graph_debug_dump_f16_tensor( "glm53_dsa_value_cache", g->layer_value_cache[il], (uint64_t)cache_len * DS4_N_HEAD * DS4_N_VALUE_MLA, il, pos0); } } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_cache"); if (ok) failed_stage = "DSA full attention"; 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; } } if (ok) { metal_graph_debug_dump_tensor( "glm53_dsa_heads", g->batch_heads, (uint64_t)n_tokens * g->heads_dim, il, pos0); } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attention"); if (ok) failed_stage = "DSA output projection"; 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) { metal_graph_debug_dump_tensor( "glm53_dsa_attn_out", g->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } glm53_batch_attention_done: if (ok && g->glm53) { metal_graph_debug_dump_tensor( "attn_out", g->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); ok = glm_graph_apply_directional_steering_attn( g, g->batch_attn_out, il, n_tokens); } if (ok && g->glm53) { failed_stage = "attention mHC expand"; ok = ds4_gpu_hc_expand_split_tensor(hc_after_attn_view, g->batch_attn_out, hc_cur, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { metal_graph_debug_dump_tensor( "glm53_hc_after_attn", hc_after_attn_view, (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD, il, pos0); } if (ok) { failed_stage = "FFN mHC pre"; ok = glm53_graph_hc_pre_rows(g, model, l->hc_ffn_fn, l->hc_ffn_scale, l->hc_ffn_base, l->ffn_norm, hc_after_attn_view, after_attn_view, g->batch_ffn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, n_tokens); if (ok) { metal_graph_debug_dump_tensor( "glm53_ffn_hc_collapsed", after_attn_view, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor( "glm53_ffn_hc_flat", g->batch_hc_flat, (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD, il, pos0); metal_graph_debug_dump_tensor( "glm53_ffn_hc_mix", g->batch_hc_mix, (uint64_t)n_tokens * DS4_N_HC * (DS4_N_HC + 2u), il, pos0); metal_graph_debug_dump_tensor( "glm53_ffn_hc_split", g->batch_hc_split, (uint64_t)n_tokens * DS4_N_HC * (DS4_N_HC + 2u), il, pos0); } } } else 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) { failed_stage = "FFN"; ok = glm_graph_encode_ffn_batch(g, model, weights, l, il, pos0, g->glm53 ? after_attn_view : g->batch_after_attn, next, n_tokens, full_layer_prefill, layer_stage_profile, stage_sync, layer_stage_profile ? &layer_stage_t0 : NULL); } if (ok && g->glm53) { metal_graph_debug_dump_tensor( "ffn_out", next, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); ok = glm_graph_apply_directional_steering_ffn( g, next, il, n_tokens); } if (ok && g->glm53) { failed_stage = "FFN mHC expand"; ok = ds4_gpu_hc_expand_split_tensor(hc_next, next, hc_after_attn_view, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { metal_graph_debug_dump_tensor( "glm53_hc_after_ffn", hc_next, (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD, il, pos0); } if (ok) { ds4_gpu_tensor *tmp = hc_cur; hc_cur = hc_next; hc_next = tmp; } } else 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) { if (g->glm53) { glm_debug_dump_raw_layer( hc_cur, "hc", (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float), il, -1); } else { 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) { failed_stage = "hidden readback"; ok = ds4_gpu_tensor_read(g->glm53 ? hc_cur : cur, 0, output_hc, (uint64_t)n_tokens * DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * sizeof(float)) != 0; } if (ok && logits_out) { failed_stage = "output head"; ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, model, weights); } if (ok && logits_out) { if (g->glm53) { ds4_gpu_tensor *last_hc = glm_graph_tensor_row_view_strided( hc_cur, n_tokens - 1u, (uint64_t)DS4_N_HC * DS4_N_EMBD, (uint64_t)DS4_N_HC * DS4_N_EMBD); ok = last_hc != NULL && ds4_gpu_hc_weighted_sum_tensor(g->hc_output, last_hc, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) != 0; ds4_gpu_tensor_free(last_hc); last_hidden = g->hc_output; } else { 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); } } if (!g->glm53) ds4_gpu_tensor_free(last_hidden); ds4_gpu_tensor_free(hc_after_attn_view); ds4_gpu_tensor_free(hc_next_view); ds4_gpu_tensor_free(hc_cur_view); ds4_gpu_tensor_free(after_attn_view); ds4_gpu_tensor_free(next_view); ds4_gpu_tensor_free(cur_view); glm_vision_overlay_free(&vision_overlay); ds4_gpu_set_glm_streaming_prefill_full_layer(false); if (!ok) { fprintf(stderr, "ds4: GLM layer-major prefill failed at layer %u stage '%s' " "(pos %u, tokens %u)\n", failed_layer, failed_stage, pos0, n_tokens); } 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, const ds4_vision_span *images, size_t image_count, 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->glm53 && (!g->batch_indexer_gate || !g->batch_indexer_pool_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; } if (g->glm53 && !glm53_graph_prefill_workspace_ensure(g, n_tokens)) { return false; } const uint32_t n_rows = pos0 + n_tokens; const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); const uint32_t dense_limit = glm_graph_dense_compact_attention_limit(g); #ifdef DS4_ROCM_BUILD if (pos0 < dense_limit && n_rows > dense_limit) { return false; } #endif const uint32_t indexed_selected_count = n_rows <= dense_limit ? n_rows : (g->glm53 ? glm53_graph_indexer_selected_limit() : indexer_top_k); const bool use_causal_range_select = n_rows <= dense_limit; 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 && !vision_synthetic_token_is_covered( tokens, i, images, image_count, pos0))) return false; } if (!input_hc && !g->has_token_embd) return false; if (input_hc && image_count != 0) return false; if (logits_out && !g->has_output_head) return false; glm_graph_reset_prefill_seed_capture(g); glm_vision_overlay vision_overlay = {0}; if (!glm_vision_overlay_prepare(&vision_overlay, images, image_count, pos0, n_tokens)) return false; 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); } } const uint64_t plain_bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); const uint64_t hc_bytes = plain_bytes * DS4_N_HC; ds4_gpu_tensor *cur_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_cur, 0, plain_bytes) : NULL; ds4_gpu_tensor *next_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_next, 0, plain_bytes) : NULL; ds4_gpu_tensor *after_attn_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_after_attn, 0, plain_bytes) : NULL; ds4_gpu_tensor *hc_cur_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_cur, 0, hc_bytes) : NULL; ds4_gpu_tensor *hc_next_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_next, 0, hc_bytes) : NULL; ds4_gpu_tensor *hc_after_attn_view = g->glm53 ? ds4_gpu_tensor_view(g->batch_hc_after_attn, 0, hc_bytes) : NULL; if (g->glm53 && (!cur_view || !next_view || !after_attn_view || !hc_cur_view || !hc_next_view || !hc_after_attn_view)) { ds4_gpu_tensor_free(hc_after_attn_view); ds4_gpu_tensor_free(hc_next_view); ds4_gpu_tensor_free(hc_cur_view); ds4_gpu_tensor_free(after_attn_view); ds4_gpu_tensor_free(next_view); ds4_gpu_tensor_free(cur_view); glm_vision_overlay_free(&vision_overlay); return false; } ds4_gpu_tensor *cur = g->glm53 ? cur_view : g->batch_cur; ds4_gpu_tensor *next = g->glm53 ? next_view : g->batch_next; ds4_gpu_tensor *hc_cur = hc_cur_view; ds4_gpu_tensor *hc_next = hc_next_view; 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; /* Generic IQ2 GLM benefits from sequential layer reads on large continued * chunks too. Keep the typed Flash path's selected-expert cache reuse. */ const bool full_layer_prefill = #if defined(__APPLE__) && !defined(DS4_NO_GPU) (pos0 == 0 || (!g->glm53 && glm_graph_layer_uses_generic_routed_moe( &weights->layer[DS4_N_LEADING_DENSE]))) && n_tokens >= 256u && use_batch_ffn && !g->quality && g->tp_world < 2 && g->layer_start == 0 && g->layer_count == glm_graph_normal_layer_count() && glm_graph_stream_prefill_full_layer_enabled(g, n_tokens); #else false; #endif const bool layer_prepare = glm_graph_stream_prefill_full_layer_prepare_enabled(g, full_layer_prefill); metal_graph_stream_prepare_slot layer_prepare_slot = {0}; 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); 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(g->glm53 ? hc_cur : cur, 0, input_hc, (uint64_t)n_tokens * DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * 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; if (g->glm53 && weights->token_embd->type == DS4_TENSOR_BF16) { ok = ds4_gpu_glm53_embedding_bf16(cur, model->map, model->size, weights->token_embd->abs_offset, g->prefill_tokens, n_tokens, DS4_N_EMBD, DS4_N_VOCAB) != 0; } else { 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 (ok && g->glm53) { ok = ds4_gpu_repeat_hc_rows_tensor(hc_cur, cur, n_tokens, DS4_N_EMBD, DS4_N_HC) != 0; } for (size_t i = 0; ok && i < vision_overlay.count; i++) { const glm_vision_overlay_segment *segment = &vision_overlay.segments[i]; ok = ds4_gpu_glm53_scatter_image_hc( g->glm53 ? hc_cur : cur, vision_overlay.tensor, segment->dst_row, segment->src_row, segment->rows, n_tokens, DS4_N_EMBD, g->glm53 ? DS4_N_HC : 1u) != 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 (layer_prepare && !metal_graph_stream_prepare_join_layer( NULL, model, weights, il, n_tokens, false, true, false, false, &layer_prepare_slot, 1)) { ok = false; break; } if (g->ssd_streaming) { ok = glm_graph_stream_map_prefill_layer(g, model, weights, il, n_tokens, full_layer_prefill); if (ok && layer_prepare && il < g->layer_end) { ok = metal_graph_stream_prepare_start_if_needed( NULL, model, weights, il + 1u, n_tokens, false, true, false, false, &layer_prepare_slot, 1); } if (ok) ok = ds4_gpu_begin_commands() != 0; } const ds4_layer_weights *l = &weights->layer[il]; const bool glm53_kda = g->glm53 && ds4_glm53_layer_is_kda(il); const uint32_t kv_raw_dim = glm53_kda ? 0u : (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 && g->glm53) { ok = glm53_graph_hc_pre_rows(g, model, l->hc_attn_fn, l->hc_attn_scale, l->hc_attn_base, l->attn_norm, hc_cur, cur, g->batch_attn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, n_tokens); } else 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 && glm53_kda) { ok = glm53_graph_kda_attention_rows(g, model, l, il, pos0, n_tokens, g->batch_attn_out); goto glm53_indexed_attention_done; } 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 && DS4_N_ROT != 0) 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 = g->glm53 ? glm53_graph_matmul_rows(g->batch_indexer_k, model, l->indexer_attn_k, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->batch_attn_norm, n_tokens) : (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 && g->glm53) { ok = glm53_graph_matmul_rows( g->batch_indexer_gate, model, l->indexer_compressor_gate, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->batch_attn_norm, n_tokens); } if (ok && g->glm53) { ok = ds4_gpu_glm53_indexer_pool_update_tensor( g->layer_indexer_key_cache[il], g->layer_indexer_tail_k[il], g->layer_indexer_tail_gate[il], g->batch_indexer_k, g->batch_indexer_gate, model->map, model->size, l->indexer_k_norm->abs_offset, l->indexer_k_norm_b->abs_offset, l->indexer_compressor_ape->abs_offset, pos0, n_tokens, g->compact_cache_cap, DS4_N_INDEXER_HEAD_DIM, DS4_GLM53_INDEX_POOL_SIZE, 1.0e-6f, glm_graph_compact_cache_is_f16()) != 0; } else 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 = g->glm53 ? glm53_graph_matmul_rows( g->batch_indexer_q, model, l->indexer_attn_q_b, DS4_N_LORA_Q, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, g->batch_q_rank_norm, n_tokens) : (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 && !g->glm53) { 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 = g->glm53 ? glm53_graph_matmul_rows(g->batch_indexer_weights, model, l->indexer_proj, DS4_N_EMBD, DS4_N_INDEXER_HEAD, g->batch_attn_norm, n_tokens) : (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) { /* The causal kernel consumes the implicit [0, visible) * range directly; no selected-row buffer is needed. */ } 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_rows = g->glm53 ? n_rows / DS4_GLM53_INDEX_POOL_SIZE : n_rows; const uint32_t selected_for_score = g->glm53 ? indexer_top_k / DS4_GLM53_INDEX_POOL_SIZE : indexed_selected_count; 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)selected_for_score * sizeof(uint32_t); ds4_gpu_tensor *selected_base = g->glm53 ? g->batch_indexer_pool_selected : g->batch_indexer_selected; 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(selected_base, (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 = g->glm53 ? ds4_gpu_glm53_indexer_scores_batch_tensor( g->batch_indexer_scores, indexer_q_view, indexer_weights_view, g->layer_indexer_key_cache[il], score_rows, slice, pos0 + t0, DS4_GLM53_INDEX_POOL_SIZE, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, indexer_scale, glm_graph_compact_cache_is_f16()) != 0 : ds4_gpu_glm_indexer_scores_batch_tensor( g->batch_indexer_scores, indexer_q_view, indexer_weights_view, g->layer_indexer_key_cache[il], score_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, score_rows, slice, selected_for_score) != 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 && g->glm53) { ok = ds4_gpu_glm53_expand_pool_selection_tensor( g->batch_indexer_selected, g->batch_indexer_pool_selected, n_tokens, pos0, selected_for_score, indexer_top_k, DS4_GLM53_INDEX_POOL_SIZE, indexed_selected_count) != 0; } } } 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; ) { const uint32_t slice_pos = pos0 + t0; uint32_t slice = n_tokens - t0; if (slice > attn_slice_cap) slice = attn_slice_cap; const bool slice_causal = slice_pos < dense_limit; #ifdef DS4_ROCM_BUILD /* The selected-row GEMM gathers a bounded KV tile. Larger * sparse batches otherwise fall back to scalar attention. */ if (!slice_causal && slice > 256u) slice = 256u; #endif if (slice_causal && slice > dense_limit - slice_pos) { slice = dense_limit - slice_pos; } const uint32_t attention_selected_count = slice_causal ? (n_rows < dense_limit ? n_rows : dense_limit) : last_indexer_selected_count; 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 = slice_causal ? NULL : 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 && (slice_causal || 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 defined(__APPLE__) || (!defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU)) if (slice_causal && glm_graph_use_dense_compact_attention_prefill(slice) && !tp_attn_head_split) { rc = ds4_gpu_glm_attention_dense_compact_lora_causal_tensor( attn_lora_view, qk_low_view, g->layer_kv_lora_cache[il], pos0 + t0, slice, attention_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), DS4_N_HEAD, DS4_N_KV_LORA, (uint32_t)g->q_nope); } else #endif if (slice_causal) { /* The compact causal score includes both absorbed MLA * and the separate RoPE query/key contribution. */ 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, attention_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 { /* GLM 5.3 pads the incomplete pool with invalid IDs. */ rc = (g->glm53 ? ds4_gpu_glm_attention_indexed_batch_lora_tensor : 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 if (tp_attn_head_split && !g->quality && ds4_gpu_device_is_m5_apple_silicon() && n_tokens >= 32u) { const uint64_t k_cnt = g->heads_dim / 2u; const uint32_t sliced_rows = n_tokens & ~31u; ok = ds4_gpu_matmul_q8_0_kslice_rows_tensor( attn_out_dst, model->map, model->size, l->attn_output->abs_offset, g->heads_dim, DS4_N_EMBD, (uint64_t)g->tp_rank * k_cnt, k_cnt, g->batch_heads, sliced_rows) != 0; if (ok && sliced_rows != n_tokens) { const uint32_t tail_rows = n_tokens - sliced_rows; ds4_gpu_tensor *heads_tail = ds4_gpu_tensor_view( g->batch_heads, (uint64_t)sliced_rows * g->heads_dim * sizeof(float), (uint64_t)tail_rows * g->heads_dim * sizeof(float)); ds4_gpu_tensor *out_tail = ds4_gpu_tensor_view( attn_out_dst, (uint64_t)sliced_rows * DS4_N_EMBD * sizeof(float), (uint64_t)tail_rows * DS4_N_EMBD * sizeof(float)); ok = heads_tail && out_tail && glm_graph_matmul_q8_0_tensor(out_tail, model, l->attn_output->abs_offset, g->heads_dim, DS4_N_EMBD, heads_tail, tail_rows); ds4_gpu_tensor_free(out_tail); ds4_gpu_tensor_free(heads_tail); } } 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); } glm53_indexed_attention_done: if (ok && g->glm53) { metal_graph_debug_dump_tensor( "attn_out", g->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); ok = glm_graph_apply_directional_steering_attn( g, g->batch_attn_out, il, n_tokens); } if (ok && g->glm53) { ok = ds4_gpu_hc_expand_split_tensor(hc_after_attn_view, g->batch_attn_out, hc_cur, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { ok = glm53_graph_hc_pre_rows(g, model, l->hc_ffn_fn, l->hc_ffn_scale, l->hc_ffn_base, l->ffn_norm, hc_after_attn_view, after_attn_view, g->batch_ffn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, n_tokens); } } else 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->glm53 ? after_attn_view : g->batch_after_attn, next, n_tokens, full_layer_prefill, 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, true, 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 && g->glm53) { metal_graph_debug_dump_tensor( "ffn_out", next, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); ok = glm_graph_apply_directional_steering_ffn( g, next, il, n_tokens); } if (ok && g->glm53) { ok = ds4_gpu_hc_expand_split_tensor(hc_next, next, hc_after_attn_view, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { ds4_gpu_tensor *tmp = hc_cur; hc_cur = hc_next; hc_next = tmp; } } else 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) { if (g->glm53) { glm_debug_dump_raw_layer( hc_cur, "hc", (uint64_t)n_tokens * DS4_N_HC * DS4_N_EMBD * sizeof(float), il, -1); } else { 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 (!metal_graph_stream_prepare_join_all(&layer_prepare_slot, 1)) ok = false; 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(g->glm53 ? hc_cur : cur, 0, output_hc, (uint64_t)n_tokens * DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * sizeof(float)) != 0; } if (ok && logits_out) { ok = glm_graph_seed_streaming_expert_cache_from_prefill(g, model, weights); } if (ok && logits_out) { if (g->glm53) { ds4_gpu_tensor *last_hc = glm_graph_tensor_row_view_strided( hc_cur, n_tokens - 1u, (uint64_t)DS4_N_HC * DS4_N_EMBD, (uint64_t)DS4_N_HC * DS4_N_EMBD); ok = last_hc != NULL && ds4_gpu_hc_weighted_sum_tensor(g->hc_output, last_hc, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) != 0; ds4_gpu_tensor_free(last_hc); last_hidden = g->hc_output; } else { 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); } } if (!g->glm53) ds4_gpu_tensor_free(last_hidden); ds4_gpu_tensor_free(hc_after_attn_view); ds4_gpu_tensor_free(hc_next_view); ds4_gpu_tensor_free(hc_cur_view); ds4_gpu_tensor_free(after_attn_view); ds4_gpu_tensor_free(next_view); ds4_gpu_tensor_free(cur_view); glm_vision_overlay_free(&vision_overlay); 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; if (g->glm53) { uint32_t done = 0; while (done < n_tokens) { const uint32_t pos = pos0 + done; uint32_t chunk = n_tokens - done; if (chunk > DS4_GLM53_PREFILL_CHUNK_TOKENS) { chunk = DS4_GLM53_PREFILL_CHUNK_TOKENS; } if (pos < g->ctx_cap) { const uint32_t dense_left = g->ctx_cap - pos; if (chunk > dense_left) chunk = dense_left; } else if (g->indexed_prefill_cap != 0) { if (chunk > g->indexed_prefill_cap) { chunk = g->indexed_prefill_cap; } } else { chunk = 1; } chunk = glm_graph_limit_indexed_prefill_chunk(g, pos, chunk); float *dst_logits = done + chunk == n_tokens ? logits_out : NULL; const bool use_indexed = glm53_graph_use_indexed_prefill(g) || (pos >= g->ctx_cap && g->indexed_prefill_cap != 0); const bool ok = use_indexed ? glm_graph_forward_indexed_tokens(g, model, weights, tokens + done, NULL, NULL, 0, pos, chunk, NULL, dst_logits, progress, progress_ud, pos0, done, n_tokens) : pos < g->ctx_cap ? glm_graph_forward_tokens(g, model, weights, tokens + done, NULL, NULL, 0, pos, chunk, NULL, dst_logits, progress, progress_ud, pos0, done, n_tokens) : glm_graph_forward_token(g, model, weights, tokens[done], NULL, pos, NULL, dst_logits, false); if (!ok) { return false; } done += chunk; } return true; } 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(g, 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, NULL, 0, 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(g, 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, NULL, 0, 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, NULL, 0, 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) { #if defined(__APPLE__) || defined(DS4_ROCM_BUILD) || defined(DS4_NO_GPU) decode_layer_flush_interval = use_indexed_attention ? 4u : 32u; #endif 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 tp_split_indexed_heads = use_indexed_attention && g->tp_world == 2 && g->tp_out && g->tp_in && !g->ssd_streaming && (DS4_N_HEAD % 2u) == 0u; const uint32_t tp_head_count = DS4_N_HEAD / 2u; const uint64_t tp_q_elems = (uint64_t)tp_head_count * DS4_N_KEY_MLA; const uint64_t tp_qk_low_elems = (uint64_t)tp_head_count * DS4_N_KV_LORA; const uint64_t tp_heads_elems = (uint64_t)tp_head_count * DS4_N_VALUE_MLA; ds4_gpu_tensor *tp_q = NULL; ds4_gpu_tensor *tp_qk_low = NULL; ds4_gpu_tensor *tp_heads = NULL; if (tp_split_indexed_heads) { tp_q = ds4_gpu_tensor_view(g->q, (uint64_t)g->tp_rank * tp_q_elems * sizeof(float), tp_q_elems * sizeof(float)); tp_qk_low = ds4_gpu_tensor_view( g->qk_low, (uint64_t)g->tp_rank * tp_qk_low_elems * sizeof(float), tp_qk_low_elems * sizeof(float)); tp_heads = ds4_gpu_tensor_view( g->heads, (uint64_t)g->tp_rank * tp_heads_elems * sizeof(float), tp_heads_elems * sizeof(float)); if (!tp_q || !tp_qk_low || !tp_heads) { ds4_gpu_tensor_free(tp_heads); ds4_gpu_tensor_free(tp_qk_low); ds4_gpu_tensor_free(tp_q); DS4_GLM_FT_FAIL("tp head view allocation"); 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() && weights_model_map_decode_static_supported(weights); 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 && g->glm53) { const int32_t token_id = (int32_t)token; ok = ds4_gpu_tensor_write(g->prefill_tokens, 0, &token_id, sizeof(token_id)) != 0; } if (ok && input_hc) { ok = ds4_gpu_tensor_write(g->glm53 ? g->hc_cur : g->cur, 0, input_hc, (uint64_t)DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * sizeof(float)) != 0; } else if (ok && 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 if (ok) { 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) { if (g->glm53) { if (weights->token_embd->type == DS4_TENSOR_BF16) { ok = ds4_gpu_glm53_embedding_bf16(g->cur, model->map, model->size, weights->token_embd->abs_offset, g->prefill_tokens, 1, DS4_N_EMBD, DS4_N_VOCAB) != 0; } else { ok = ds4_gpu_embed_tokens_quant_tensor( g->cur, g->prefill_tokens, model->map, model->size, weights->token_embd->abs_offset, weights->token_embd->type, DS4_N_VOCAB, 1, DS4_N_EMBD) != 0; } if (ok) ok = ds4_gpu_repeat_hc_tensor(g->hc_cur, g->cur, DS4_N_EMBD, DS4_N_HC) != 0; } else { 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; /* Number of leading selected slots that are known to index live cache rows. * 0 means "all of them are" -- the producers that fill a contiguous range or * a plain top-k. The GLM-5.3 pooled expansion instead guarantees only * [0, indexer_top_k) and pads the causal-tail slots above it with * 0xffffffff, so it must never be handed to a kernel level that assumes * every row is valid. */ uint32_t last_indexer_guaranteed_prefix = 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; const char *glm_ft_fail_stage = "layer setup"; #define DS4_GLM_FT_STAGE(stage_) do { if (ok) glm_ft_fail_stage = (stage_); } while (0) 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 bool glm53_kda = g->glm53 && ds4_glm53_layer_is_kda(il); const uint32_t kv_raw_dim = glm53_kda ? 0u : (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); } DS4_GLM_FT_STAGE("attention mHC pre"); if (ok && g->glm53) { ok = glm53_graph_hc_pre(g, model, l->hc_attn_fn, l->hc_attn_scale, l->hc_attn_base, l->attn_norm, g->hc_cur, g->cur, g->attn_norm); } else 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"); if (ok && glm53_kda) { DS4_GLM_FT_STAGE("KDA attention"); ok = glm53_graph_kda_attention(g, model, l, il); goto glm53_attention_done; } const uint32_t decode_ablate = glm_decode_ablate_mask(); DS4_GLM_FT_STAGE("DSA q_a projection"); 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) { DS4_GLM_FT_STAGE("DSA fused q/kv norm and 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) { DS4_GLM_FT_STAGE("DSA fused q/kv 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) { DS4_GLM_FT_STAGE("DSA q_a norm"); 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; } const bool tp_split_layer_heads = tp_split_indexed_heads && il >= DS4_N_LEADING_DENSE && l->attn_q_b->type == DS4_TENSOR_Q8_0; if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { DS4_GLM_FT_STAGE("DSA q_b projection"); uint64_t q_weight_offset = l->attn_q_b->abs_offset; uint64_t q_row_bytes = 0; if (tp_split_layer_heads) { ok = tensor_nbytes(l->attn_q_b->type, DS4_N_LORA_Q, &q_row_bytes); q_weight_offset += (uint64_t)g->tp_rank * tp_q_elems * q_row_bytes; } if (ok) { ok = glm_graph_matmul_q8_0_decode_profiled_tensor( tp_split_layer_heads ? tp_q : g->q, model, q_weight_offset, DS4_N_LORA_Q, tp_split_layer_heads ? tp_q_elems : g->q_dim, g->q_rank_norm, il, pos, "attn_q_b", g->ssd_streaming) != 0; } DS4_GLM_FT_STAGE("DSA q positional transform"); if (ok && DS4_N_ROT != 0) ok = ds4_gpu_glm_rope_tail_tensor( tp_split_layer_heads ? tp_q : g->q, 1, tp_split_layer_heads ? tp_head_count : 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)) { DS4_GLM_FT_STAGE("DSA indexer key projection"); ok = g->glm53 ? glm53_graph_matmul_rows(g->indexer_k, model, l->indexer_attn_k, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->attn_norm, 1) : 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 && g->glm53) { ok = glm53_graph_matmul_rows( g->indexer_gate, model, l->indexer_compressor_gate, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, g->attn_norm, 1); } if (ok && g->glm53) { DS4_GLM_FT_STAGE("DSA pooled indexer key store"); ok = ds4_gpu_glm53_indexer_pool_update_tensor( g->layer_indexer_key_cache[il], g->layer_indexer_tail_k[il], g->layer_indexer_tail_gate[il], g->indexer_k, g->indexer_gate, model->map, model->size, l->indexer_k_norm->abs_offset, l->indexer_k_norm_b->abs_offset, l->indexer_compressor_ape->abs_offset, pos, 1, g->compact_cache_cap, DS4_N_INDEXER_HEAD_DIM, DS4_GLM53_INDEX_POOL_SIZE, 1.0e-6f, glm_graph_compact_cache_is_f16()) != 0; } else if (ok) { DS4_GLM_FT_STAGE("DSA indexer key store"); 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) { DS4_GLM_FT_STAGE("DSA kv projection and 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", 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) { DS4_GLM_FT_STAGE("DSA compact kv 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; const uint32_t dense_limit = glm_graph_dense_compact_attention_limit(g); if (ok && visible <= dense_limit) { if (last_indexer_selected != g->indexer_selected || last_indexer_selected_count != visible) { 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; last_indexer_guaranteed_prefix = 0; /* contiguous range */ } 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; last_indexer_guaranteed_prefix = 0; /* contiguous range */ } else if (ok) { ok = g->glm53 ? glm53_graph_matmul( g->indexer_q, model, l->indexer_attn_q_b, DS4_N_LORA_Q, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, g->q_rank_norm) : 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 && !g->glm53) { 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 = g->glm53 ? glm53_graph_matmul(g->indexer_weights, model, l->indexer_proj, DS4_N_EMBD, DS4_N_INDEXER_HEAD, g->attn_norm) : 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)); const uint32_t score_rows = g->glm53 ? visible / DS4_GLM53_INDEX_POOL_SIZE : visible; ok = ds4_gpu_glm_indexer_score_one_tensor( g->indexer_scores, g->indexer_q, g->indexer_weights, g->layer_indexer_key_cache[il], score_rows, 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 && g->glm53) { const uint32_t selected_pools = indexer_top_k / DS4_GLM53_INDEX_POOL_SIZE; ok = ds4_gpu_indexer_topk_tensor( g->indexer_pool_selected, g->indexer_scores, score_rows, 1, selected_pools) != 0; if (ok) ok = ds4_gpu_glm53_expand_pool_selection_tensor( g->indexer_selected, g->indexer_pool_selected, 1, pos, selected_pools, indexer_top_k, DS4_GLM53_INDEX_POOL_SIZE, glm53_graph_indexer_selected_limit()) != 0; } else 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 = g->glm53 ? glm53_graph_indexer_selected_limit() : indexer_top_k; /* the pooled expansion pads slots >= indexer_top_k */ last_indexer_guaranteed_prefix = g->glm53 ? indexer_top_k : 0; } 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))) { uint64_t k_weight_offset = l->attn_k_b->abs_offset; uint64_t k_row_bytes = 0; if (tp_split_layer_heads) { ok = tensor_nbytes(l->attn_k_b->type, (uint64_t)g->q_nope, &k_row_bytes); k_weight_offset += (uint64_t)g->tp_rank * tp_head_count * DS4_N_KV_LORA * k_row_bytes; } if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor( tp_split_layer_heads ? tp_qk_low : g->qk_low, tp_split_layer_heads ? tp_q : g->q, model->map, model->size, k_weight_offset, l->attn_k_b->type, tp_split_layer_heads ? tp_head_count : 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"); DS4_GLM_FT_STAGE("DSA indexed attention"); 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; uint64_t value_weight_offset = l->attn_v_b->abs_offset; uint64_t value_row_bytes = 0; if (tp_split_layer_heads) { ok = tensor_nbytes(l->attn_v_b->type, DS4_N_KV_LORA, &value_row_bytes); value_weight_offset += (uint64_t)g->tp_rank * tp_head_count * DS4_N_VALUE_MLA * value_row_bytes; } if (ok) ok = ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( tp_split_layer_heads ? tp_heads : g->heads, g->attn_partial_lora, g->attn_partial_ms, tp_split_layer_heads ? tp_q : g->q, tp_split_layer_heads ? tp_qk_low : g->qk_low, g->layer_kv_lora_cache[il], g->layer_k_rope_cache[il], model->map, model->size, value_weight_offset, l->attn_v_b->type, last_indexer_selected, last_indexer_selected_count, last_indexer_guaranteed_prefix == 0u, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), tp_split_layer_heads ? tp_head_count : 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) { uint64_t value_weight_offset = l->attn_v_b->abs_offset; uint64_t value_row_bytes = 0; if (tp_split_layer_heads) { ok = tensor_nbytes(l->attn_v_b->type, DS4_N_KV_LORA, &value_row_bytes); value_weight_offset += (uint64_t)g->tp_rank * tp_head_count * DS4_N_VALUE_MLA * value_row_bytes; } if (ok) ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( tp_split_layer_heads ? tp_heads : g->heads, tp_split_layer_heads ? tp_q : g->q, tp_split_layer_heads ? tp_qk_low : g->qk_low, g->layer_kv_lora_cache[il], g->layer_k_rope_cache[il], model->map, model->size, value_weight_offset, l->attn_v_b->type, last_indexer_selected, last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), tp_split_layer_heads ? tp_head_count : 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 { DS4_GLM_FT_STAGE("DSA full attention key projection"); 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; DS4_GLM_FT_STAGE("DSA full attention value projection"); 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"); DS4_GLM_FT_STAGE("DSA full attention cache store"); 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"); DS4_GLM_FT_STAGE("DSA full attention core"); 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)) { const bool tp_split_attn = g->tp_world == 2 && g->tp_out && g->tp_in && il >= DS4_N_LEADING_DENSE && !g->ssd_streaming; if (tp_split_attn) { const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN; const uint64_t k_cnt = g->heads_dim / 2u; const uint64_t k_off = (uint64_t)g->tp_rank * k_cnt; ok = metal_graph_matmul_dense_quant_kslice( g->tp_out[slot], model, l->attn_output, g->heads_dim, k_off, k_cnt, DS4_N_EMBD, g->heads, k_off); if (ok) ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_ATTN) != 0; if (ok) ok = ds4_gpu_add_tensor(g->attn_out, g->tp_out[slot], g->tp_in[slot], DS4_N_EMBD) != 0; } else { 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; } } glm53_attention_done: DS4_GLM_FT_STAGE("attention mHC expand"); DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_output"); if (ok && g->glm53) { metal_graph_debug_dump_tensor("attn_out", g->attn_out, DS4_N_EMBD, il, pos); ok = glm_graph_apply_directional_steering_attn( g, g->attn_out, il, 1); } if (ok && g->glm53) { ok = ds4_gpu_hc_expand_tensor(g->hc_after_attn, g->attn_out, g->hc_cur, g->hc_post, g->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = glm53_graph_hc_pre(g, model, l->hc_ffn_fn, l->hc_ffn_scale, l->hc_ffn_base, l->ffn_norm, g->hc_after_attn, g->after_attn, g->ffn_norm); } else 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"); DS4_GLM_FT_STAGE("FFN"); if (ok && g->glm53) { bool graph_ok = false; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) graph_ok = !g->placement && g->tp_world <= 1 && !g->ssd_streaming && !g->imatrix && g->directional_steering_ffn_scale == 0.0f && !decode_stage_profile && !g_expert_profile.active && metal_graph_debug_get_config()->prefix == NULL && !glm_debug_hidden_dump_layer_match(il) && getenv("DS4_CUDA_GLM_DISABLE_FFN_GRAPHS") == NULL && ds4_gpu_decode_graphs_supported() != 0; #endif bool tail_ok = false; if (graph_ok) { ds4_decode_graph_key key; memset(&key, 0, sizeof(key)); key.il = il; key.island = 1u; key.cur_hc = (void *)g->hc_cur; key.after_attn_hc = (void *)g->hc_after_attn; key.after_ffn_hc = (void *)g->hc_next; key.attn_norm = (void *)g->ffn_norm; for (;;) { const int state = ds4_gpu_decode_graph_begin(&key); if (state == 1) { tail_ok = true; break; } if (state == 0) { const bool captured = glm53_graph_encode_ffn_tail_one( g, model, l, il, pos, false, NULL); if (!captured) { ds4_gpu_decode_graph_abort(&key); tail_ok = glm53_graph_encode_ffn_tail_one( g, model, l, il, pos, false, NULL); break; } if (ds4_gpu_decode_graph_end(&key) == 0) { tail_ok = true; break; } continue; } tail_ok = glm53_graph_encode_ffn_tail_one( g, model, l, il, pos, false, NULL); break; } } else { tail_ok = glm53_graph_encode_ffn_tail_one( g, model, l, il, pos, decode_stage_profile, decode_stage_profile ? &decode_stage_t0 : NULL); } ok = tail_ok; DS4_GLM_FT_STAGE("FFN mHC expand"); if (ok) { ds4_gpu_tensor *tmp = g->hc_cur; g->hc_cur = g->hc_next; g->hc_next = tmp; } } else 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, true, decode_stage_profile, decode_stage_profile ? &decode_stage_t0 : NULL); } if (ok && !g->glm53) { 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) { if (g->glm53) { glm_debug_dump_raw_layer(g->hc_cur, "hc", (uint64_t)DS4_N_HC * DS4_N_EMBD * sizeof(float), il, (int)pos); } else { 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 #undef DS4_GLM_FT_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->glm53 ? g->hc_cur : g->cur, 0, output_hc, (uint64_t)DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u) * 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->glm53 ? g->hc_cur : 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 (%s)\n", pos, glm_ft_fail_il, glm_ft_fail_stage); } ds4_gpu_tensor_free(tp_heads); ds4_gpu_tensor_free(tp_qk_low); ds4_gpu_tensor_free(tp_q); (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, .glm53 = ds4_model_is_glm53(), }; 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, DS4_SWIGLU_CLAMP_EXP); 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, 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 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 (!glm_graph_load_directional_steering(&g, directional_steering_file, directional_steering_attn, directional_steering_ffn)) { glm_graph_free(&g); return 1; } 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; } #endif static void qwen4_ref_row(const ds4_model *m, const ds4_tensor *t, uint64_t row, float *out); static void qwen4_ple_step(int token, int *prev, uint32_t *rows); static bool qwen4_ngram_row(const ds4_model *m, uint32_t row, float *out) { const ds4_tensor *t = m->ngram_tensor; if (!t || m->ngram_fd < 0 || !out || row >= t->dim[1] || t->type != DS4_TENSOR_BF16 || !t->dim[0] || t->dim[0] > 160) { errno = EINVAL; return false; } uint8_t raw[320]; const uint32_t bytes = (uint32_t)t->dim[0] * 2u; const uint64_t offset = t->abs_offset + (uint64_t)row * bytes; uint32_t done = 0; while (done < bytes) { ssize_t n = pread(m->ngram_fd, raw + done, bytes - done, (off_t)(offset + done)); if (n < 0 && errno == EINTR) continue; if (n <= 0 || (size_t)n > bytes - done) { if (!n) errno = EIO; return false; } done += (uint32_t)n; } for (size_t i = 0; i < t->dim[0]; i++) { uint32_t bits = ((uint32_t)raw[2*i] | ((uint32_t)raw[2*i+1] << 8)) << 16; if ((bits & 0x7f800000u) == 0x7f800000u) { errno = EDOM; return false; } memcpy(out + i, &bits, sizeof(bits)); } return true; } typedef struct { uint32_t row, output; } qwen4_ngram_request; typedef struct { const ds4_model *model; qwen4_ngram_request *request; float *out; size_t count, readers; int error[16]; } qwen4_ngram_batch; static int qwen4_ngram_order(const void *a, const void *b) { const qwen4_ngram_request *x = a, *y = b; return (x->row > y->row) - (x->row < y->row); } static void qwen4_ngram_part(void *context, size_t part) { qwen4_ngram_batch *b = context; const size_t begin = b->count * part / b->readers, end = b->count * (part+1) / b->readers; const size_t width = b->model->ngram_tensor->dim[0]; float *previous = NULL; for (size_t i = begin; i < end; i++) { float *dst = b->out + b->request[i].output * width; if (i > begin && b->request[i].row == b->request[i-1].row) { memcpy(dst, previous, width * sizeof(float)); } else if (!qwen4_ngram_row(b->model, b->request[i].row, dst)) { b->error[part] = errno; return; } previous = dst; } } #ifndef __APPLE__ typedef struct { qwen4_ngram_batch *batch; size_t part; } qwen4_ngram_reader; static void *qwen4_ngram_thread(void *context) { qwen4_ngram_reader *reader = context; qwen4_ngram_part(reader->batch,reader->part); return NULL; } #endif /* Bound sorting memory and I/O concurrency independently of context size. * Metal also overlaps the uncached rows of a single decode token. */ static bool qwen4_ngram_read(const ds4_model *m, const uint32_t *rows, size_t count, float *out) { if (!m || !m->ngram_tensor || m->ngram_fd < 0 || (count && (!rows || !out)) || count > SIZE_MAX / (160u * sizeof(float))) { errno = EINVAL; return false; } for (size_t i = 0; i < count; i++) { if (rows[i] >= m->ngram_tensor->dim[1]) { errno = EINVAL; return false; } } #ifdef __APPLE__ if (count < 2) return count == 0 || qwen4_ngram_row(m, rows[0], out); #else if (count < 256) { for (size_t i = 0; i < count; i++) if (!qwen4_ngram_row(m, rows[i], out + i * m->ngram_tensor->dim[0])) return false; return true; } #endif enum { MAX_ROWS = 4096 }; qwen4_ngram_request *request = malloc(MAX_ROWS * sizeof(*request)); if (!request) return false; bool ok = true; for (size_t off = 0; ok && off < count; off += MAX_ROWS) { size_t n = count-off < MAX_ROWS ? count-off : MAX_ROWS; for (size_t i = 0; i < n; i++) request[i] = (qwen4_ngram_request){rows[off+i], (uint32_t)i}; qsort(request, n, sizeof(*request), qwen4_ngram_order); qwen4_ngram_batch batch = {.model = m, .request = request, .count = n, .out = out + off * m->ngram_tensor->dim[0], .readers = 16}; #ifdef __APPLE__ batch.readers = n < 16 ? n : 16; dispatch_apply_f(batch.readers, dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), &batch, qwen4_ngram_part); #else pthread_t threads[15]; qwen4_ngram_reader readers[15]; size_t started = 0; for (size_t part = 1; part < batch.readers; part++) { readers[started] = (qwen4_ngram_reader){&batch,part}; if (pthread_create(&threads[started],NULL,qwen4_ngram_thread,&readers[started])) break; started++; } qwen4_ngram_part(&batch, 0); /* Thread exhaustion reduces concurrency; every partition still runs. */ for (size_t part = started+1; part < batch.readers; part++) qwen4_ngram_part(&batch,part); for (size_t part = 0; part < started; part++) if (pthread_join(threads[part],NULL)) abort(); #endif for (size_t i = 0; i < batch.readers; i++) { if (batch.error[i]) { errno = batch.error[i]; ok = false; break; } } } int error = errno; free(request); errno = error; return ok; } /* Multimodal rope positions (HF get_rope_index): text tokens share one * counter; an image's tokens take (t, t+row, t+col) from the counter at the * image start, which then advances by max(rows, cols). *delta tracks * counter - token index across calls made in token order. */ static void qwen4_mrope_pos(const ds4_vision_span *spans, size_t n_spans, uint32_t p, int32_t *delta, uint32_t out[3]) { for (size_t i = 0; i < n_spans; i++) { const ds4_vision_span *sp = &spans[i]; const uint32_t cnt = sp->embedding.token_count; if (p < sp->token_start || p - sp->token_start >= cnt) continue; const uint32_t gw = sp->embedding.grid_width > 0 ? sp->embedding.grid_width : cnt; const uint32_t gh = sp->embedding.grid_height > 0 ? sp->embedding.grid_height : 1u; const uint32_t k = p - sp->token_start; const uint32_t base = (uint32_t)((int64_t)sp->token_start + *delta); out[0] = base; out[1] = base + k / gw; out[2] = base + k % gw; if (k + 1u == cnt) *delta -= (int32_t)cnt - (int32_t)(gh > gw ? gh : gw); return; } out[0] = out[1] = out[2] = (uint32_t)((int64_t)p + *delta); } static const float *qwen4_span_row(const ds4_vision_span *spans, size_t n_spans, uint32_t p) { for (size_t i = 0; i < n_spans; i++) { const uint32_t cnt = spans[i].embedding.token_count; if (p >= spans[i].token_start && p - spans[i].token_start < cnt && spans[i].embedding.data) return spans[i].embedding.data + (uint64_t)(p - spans[i].token_start) * DS4_N_EMBD; } return NULL; } /* DS4_QWEN4_FAKE_IMAGE=start:count:rows:cols[,...]: synthetic image spans * with deterministic embeddings, so the GPU graph and the CPU reference * exercise the multimodal path without a vision encoder. */ static const ds4_vision_span *qwen4_fake_spans(size_t *count) { static ds4_vision_span *spans = NULL; static size_t n = 0; static bool parsed = false; if (!parsed) { parsed = true; for (const char *p = getenv("DS4_QWEN4_FAKE_IMAGE"); p && *p; ) { unsigned start, cnt, gh, gw; if (sscanf(p, "%u:%u:%u:%u", &start, &cnt, &gh, &gw) != 4 || cnt == 0) break; spans = xrealloc(spans, (n + 1) * sizeof(*spans)); ds4_vision_span *sp = &spans[n++]; memset(sp, 0, sizeof(*sp)); sp->token_start = start; sp->embedding.token_count = cnt; sp->embedding.grid_height = gh; sp->embedding.grid_width = gw; sp->embedding.data = xmalloc((size_t)cnt * DS4_N_EMBD * sizeof(float)); for (uint64_t i = 0; i < (uint64_t)cnt * DS4_N_EMBD; i++) { uint32_t h = ((uint32_t)i + 1u) * 2654435761u ^ (start * 40503u); h ^= h >> 15; h *= 2246822519u; h ^= h >> 13; sp->embedding.data[i] = ((float)(h & 0xFFFFu) / 65535.0f - 0.5f) * 0.05f; } const char *next = strchr(p, ','); p = next ? next + 1 : NULL; } } *count = n; return spans; } #ifdef DS4_HAS_QWEN4_GPU /* ------------------------------------------------------------------------ * Qwen3.8-Flash-Next GPU graph. One command batch per forward; f32 * transients sized for cap_tokens rows so the same kernels serve decode * (T=1) and batched prefill. Token embeddings come from the mapped weights; * n-gram rows are read with pread before the batch starts. Recurrent * state lives per layer: GDN [Hv][D][D] + conv history, attention KV * (f16) + raw indexer keys + pooled block keys, PLE conv history. * --------------------------------------------------------------------- */ /* Transients that every forward pass recomputes, sized by cap_tokens and the * block cap. Nothing here survives a call, so sessions can share one arena * while the engine serializes their forwards. Everything else stays private: * the KV, indexer and block caches, the GDN and PLE recurrent history, the MTP * snapshots, the rope positions, and the logits a caller reads back. * * Sharing holds only while one forward runs at a time. Encoding two sessions * into the same command batch would let the second overwrite the first's * transients before the GPU read them. A batched Qwen decode must therefore * be one forward over several rows, which is the shape these buffers already * have, and not two forwards sharing an arena. */ #define DS4_QWEN4_SCRATCH_FIELDS(X) \ X(R) X(xn) X(lo) X(inj) X(inj_alt) X(mixed) X(blk) X(qkv) X(z) X(ga) \ X(gb) X(lin_o) X(ple_emb) X(ple_key) X(ple_val) X(ple_gated) X(ple_normed) \ X(qg) X(kp) X(vp) X(iq) X(ik) X(q) X(gate) X(iqn) X(attn_o) \ X(score) X(tile_max) X(sel_blocks) X(sel_tokens) X(n_sel) X(attn_part) \ X(router) X(selected) X(weights) X(mid) X(part) X(sh_gate_logit) \ X(moe_lists) X(moe_counts) X(sh_gate) X(sh_up) X(sh_mid) X(sh_out) \ X(hc_u) X(hc_lo_act) typedef struct ds4_qwen4_gpu_graph { uint32_t ctx_cap; uint32_t pos; uint32_t cap_tokens; uint32_t k_blocks; uint32_t n_block_cap; uint32_t sel_stride; /* false when the scratch above is borrowed from the engine's arena */ bool owns_scratch; ds4_gpu_tensor *R, *xn, *lo, *inj, *mixed, *blk; ds4_gpu_tensor *qkv, *z, *ga, *gb, *lin_o; ds4_gpu_tensor *ple_emb, *ple_key, *ple_val, *ple_gated, *ple_normed, *ple_hist; int ple_prev[DS4_MAX_PLE_NGRAM]; ds4_gpu_tensor *qg, *kp, *vp, *iq, *ik, *q, *gate, *iqn, *attn_o; ds4_gpu_tensor *score, *tile_max, *sel_blocks, *sel_tokens, *n_sel, *attn_part; ds4_gpu_tensor *router, *selected, *weights, *mid, *part, *sh_gate_logit; ds4_gpu_tensor *moe_lists, *moe_counts, *sh_gate, *sh_up, *sh_mid, *sh_out, *hc_u, *hc_lo_act; ds4_gpu_tensor *logits; /* Arena only, grown on demand: one logit row per batch member, the * attention rows' cache table (one entry per row and layer) and their * split partials. */ ds4_gpu_tensor *batch_logits, *batch_attn_rows, *batch_attn_part, *batch_head_x, *batch_gdn_rows; uint32_t batch_logit_rows; ds4_gpu_tensor *layer_lin_state[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_lin_hist[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_k_cache[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_v_cache[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_ik_cache[DS4_MAX_LAYER]; ds4_gpu_tensor *layer_block_key[DS4_MAX_LAYER]; /* A separate injection buffer keeps combined normalization race-free. */ ds4_gpu_tensor *inj_alt; /* MTP: staged predictor input, its residual, and the recurrent-state * snapshot that verify rollback restores */ ds4_gpu_tensor *mtp_e, *mtp_cat, *mtp_proj, *mtp_R, *mtp_argmax, *mtp_argmax_tmp; /* MTP draft head over a ranked vocabulary subset: gathered Q8 rows, their * token ids, and whether the lazy load was attempted */ ds4_gpu_tensor *draft_head; int32_t *draft_ids; uint32_t draft_rows; bool draft_head_tried; ds4_gpu_tensor *snap_lin_state[DS4_MAX_LAYER]; ds4_gpu_tensor *snap_lin_hist[DS4_MAX_LAYER]; ds4_gpu_tensor *snap_ple_hist; int snap_ple_prev[DS4_MAX_PLE_NGRAM]; uint32_t snap_pos; /* Second snapshot set (after row 1) for 3-row MTP verifies: a partial * accept of two drafts restores this instead of replaying a token. */ ds4_gpu_tensor *snap2_lin_state[DS4_MAX_LAYER]; ds4_gpu_tensor *snap2_lin_hist[DS4_MAX_LAYER]; ds4_gpu_tensor *snap2_ple_hist; int snap2_ple_prev[DS4_MAX_PLE_NGRAM]; uint32_t snap2_pos; int32_t snap2_mrope_delta; bool snap2_valid; /* Pre-verify snapshot set (state at the block start). Exact sampling * rewinds to the block start to resample the boundary token when a * verified block crosses a sampling-mode boundary; without this set * that rewind resets the recurrent state and replays the whole kept * context on the next eval. */ ds4_gpu_tensor *snap0_lin_state[DS4_MAX_LAYER]; ds4_gpu_tensor *snap0_lin_hist[DS4_MAX_LAYER]; ds4_gpu_tensor *snap0_ple_hist; int snap0_ple_prev[DS4_MAX_PLE_NGRAM]; uint32_t snap0_pos; int32_t snap0_mrope_delta; bool snap0_valid; uint32_t mtp_pos; uint32_t n_logit_rows; bool snap_after_first; /* set by the caller for a 2-token verify: snapshot the state after row 0 */ bool snap_after_second; /* 3-token verify: also snapshot the state after row 1 */ bool verify_rows_exact; /* 3-token verify: split attention into 2/1-row sub-batches */ bool snap_valid; /* multimodal: per-position (t, h, w) rope positions, the text counter * offset, and the image spans of the prompt being prefilled */ ds4_gpu_tensor *pos3; uint32_t *host_pos3; int32_t mrope_delta; int32_t snap_mrope_delta; const ds4_vision_span *vis_spans; uint32_t mtp_last_rows; /* rows of the last predictor pass resident in mtp_R */ size_t vis_span_count; float *host_row; float *host_logits; ds4_gpu_tensor *steer_dirs; float steer_attn_scale; float steer_ffn_scale; bool dump_prompt_rows; } ds4_qwen4_gpu_graph; static bool qwen4_graph_dense_ok(const ds4_tensor *t) { return t && (t->type == DS4_TENSOR_Q8_0 || t->type == DS4_TENSOR_F16 || t->type == DS4_TENSOR_F32 || t->type == DS4_TENSOR_BF16 || t->type == DS4_TENSOR_Q4_0); } /* expert types the tiled prefill GEMM stages (kernel_qwen4_moe_mm_*) */ static bool qwen4_expert_type_has_mm(uint32_t type) { return type == DS4_TENSOR_Q8_0 || type == DS4_TENSOR_MXFP4 || type == DS4_TENSOR_Q4_K || type == DS4_TENSOR_Q2_K || type == DS4_TENSOR_IQ2_XXS; } static bool qwen4_graph_expert_ok(const ds4_tensor *t) { return t && (t->type == DS4_TENSOR_Q8_0 || t->type == DS4_TENSOR_MXFP4 || t->type == DS4_TENSOR_Q4_0 || t->type == DS4_TENSOR_F16 || t->type == DS4_TENSOR_BF16 || t->type == DS4_TENSOR_F32 || ((t->type == DS4_TENSOR_Q4_K || t->type == DS4_TENSOR_Q2_K || t->type == DS4_TENSOR_IQ2_XXS) && (t->dim[0] % 256u) == 0)); } /* The Metal graph runs a subset of what the loader accepts. */ static bool qwen4_graph_weights_supported(const ds4_weights *w) { if (!qwen4_graph_dense_ok(w->token_embd) || !qwen4_graph_dense_ok(w->output) || !qwen4_graph_dense_ok(w->output_hc_down) || !qwen4_graph_dense_ok(w->output_hc_up)) { fprintf(stderr, "ds4: Qwen3.8 GPU graph needs Q8_0/Q4_0/F16/BF16/F32 dense weights\n"); return false; } if (!w->ple_embd || w->ple_embd->type != DS4_TENSOR_BF16) { fprintf(stderr, "ds4: Qwen3.8 requires original BF16 n-grams in the model GGUF\n"); return false; } if (DS4_N_NEXTN_PREDICT != 0) { const ds4_layer_weights *l = &w->layer[DS4_N_LAYER - 1u]; const bool head_ok = l->nextn_hc_head_down && l->nextn_hc_head_up && qwen4_graph_dense_ok(l->nextn_hc_head_down) && (l->nextn_hc_head_up->type == DS4_TENSOR_F16 || l->nextn_hc_head_up->type == DS4_TENSOR_F32 || l->nextn_hc_head_up->type == DS4_TENSOR_Q8_0); if (!qwen4_graph_dense_ok(l->nextn_eh_proj) || !head_ok) { fprintf(stderr, "ds4: Qwen3.8 GPU graph: unsupported nextn weight types\n"); return false; } } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const ds4_layer_weights *l = &w->layer[il]; const ds4_tensor *hc[4] = { l->hc_attn_up, l->hc_attn_inject, l->hc_ffn_up, l->hc_ffn_inject }; for (int i = 0; i < 4; i++) { if (hc[i]->type != DS4_TENSOR_F16 && hc[i]->type != DS4_TENSOR_F32 && hc[i]->type != DS4_TENSOR_Q8_0) { fprintf(stderr, "ds4: Qwen3.8 GPU graph needs F16/F32/Q8_0 hc up/inject weights (layer %u)\n", il); return false; } } if (l->ple_conv && l->ple_conv->type != DS4_TENSOR_F32 && l->ple_conv->type != DS4_TENSOR_F16) { fprintf(stderr, "ds4: Qwen3.8 GPU graph needs F32/F16 PLE conv taps\n"); return false; } if (!qwen4_graph_dense_ok(l->hc_attn_down) || !qwen4_graph_dense_ok(l->hc_ffn_down)) { fprintf(stderr, "ds4: Qwen3.8 GPU graph: unsupported hc down weight type (layer %u)\n", il); return false; } if (ds4_qwen4_layer_is_linear(il) && l->lin_beta->type != l->lin_alpha->type) { fprintf(stderr, "ds4: Qwen3.8 GPU graph needs matching alpha/beta projection types (layer %u)\n", il); return false; } const ds4_tensor *dense[16] = { l->lin_qkv, l->lin_gate, l->lin_beta, l->lin_alpha, l->lin_out, l->attn_q, l->attn_k, l->attn_v, l->attn_output, l->indexer_q_proj, l->indexer_k_proj, l->ffn_gate_shexp, l->ffn_up_shexp, l->ffn_down_shexp, l->ple_key, l->ple_value }; for (int i = 0; i < 16; i++) { if (dense[i] && !qwen4_graph_dense_ok(dense[i])) { fprintf(stderr, "ds4: Qwen3.8 GPU graph: unsupported dense weight type %u in layer %u\n", dense[i]->type, il); return false; } } if (!qwen4_graph_expert_ok(l->ffn_gate_exps) || !qwen4_graph_expert_ok(l->ffn_up_exps) || !qwen4_graph_expert_ok(l->ffn_down_exps)) { fprintf(stderr, "ds4: Qwen3.8 GPU graph: unsupported routed expert type in layer %u " "(q8_0, mxfp4, q4_K, q2_K, iq2_xxs, f16, f32)\n", il); return false; } } return true; } static void qwen4_graph_free(ds4_qwen4_gpu_graph *g) { if (!g) return; ds4_gpu_tensor **all[] = { &g->ple_hist, &g->logits, &g->mtp_e, &g->mtp_cat, &g->mtp_proj, &g->mtp_R, &g->mtp_argmax, &g->mtp_argmax_tmp, &g->snap_ple_hist, &g->snap2_ple_hist, &g->snap0_ple_hist, &g->pos3, &g->draft_head, &g->steer_dirs, }; if (g->owns_scratch) { ds4_gpu_tensor **scratch[] = { #define QWEN4_SCRATCH_PTR(field_) &g->field_, DS4_QWEN4_SCRATCH_FIELDS(QWEN4_SCRATCH_PTR) #undef QWEN4_SCRATCH_PTR }; for (size_t i = 0; i < sizeof(scratch) / sizeof(scratch[0]); i++) { ds4_gpu_tensor_free(*scratch[i]); *scratch[i] = NULL; } ds4_gpu_tensor_free(g->batch_logits); ds4_gpu_tensor_free(g->batch_attn_rows); ds4_gpu_tensor_free(g->batch_attn_part); ds4_gpu_tensor_free(g->batch_head_x); ds4_gpu_tensor_free(g->batch_gdn_rows); g->batch_logits = g->batch_attn_rows = g->batch_attn_part = g->batch_head_x = g->batch_gdn_rows = NULL; g->batch_logit_rows = 0; free(g->host_row); } g->host_row = NULL; free(g->host_pos3); free(g->draft_ids); g->draft_ids = NULL; g->host_pos3 = NULL; for (size_t i = 0; i < sizeof(all) / sizeof(all[0]); i++) { ds4_gpu_tensor_free(*all[i]); *all[i] = NULL; } for (uint32_t il = 0; il < DS4_MAX_LAYER; il++) { ds4_gpu_tensor_free(g->layer_lin_state[il]); ds4_gpu_tensor_free(g->layer_lin_hist[il]); ds4_gpu_tensor_free(g->layer_k_cache[il]); ds4_gpu_tensor_free(g->layer_v_cache[il]); ds4_gpu_tensor_free(g->layer_ik_cache[il]); ds4_gpu_tensor_free(g->layer_block_key[il]); ds4_gpu_tensor_free(g->snap_lin_state[il]); ds4_gpu_tensor_free(g->snap_lin_hist[il]); ds4_gpu_tensor_free(g->snap2_lin_state[il]); ds4_gpu_tensor_free(g->snap2_lin_hist[il]); ds4_gpu_tensor_free(g->snap0_lin_state[il]); ds4_gpu_tensor_free(g->snap0_lin_hist[il]); } free(g->host_logits); memset(g, 0, sizeof(*g)); } /* One tensor per linear layer holding every slot's recurrent state. Sized * from the session count the engine was opened for, so a server's slots are * pooled and an ad-hoc extra session simply keeps private state. */ static bool qwen4_state_pool_ensure(ds4_engine *e, uint32_t slots) { if (e->qwen4_pool_slots) return true; if (slots < 2u || slots > 64u) return false; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; const uint64_t state_bytes = v_dim * DS4_N_LIN_HEAD_DIM * sizeof(float); const uint64_t hist_bytes = (uint64_t)(DS4_N_LIN_CONV - 1u) * DS4_N_LIN_CONV_DIM * sizeof(float); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!ds4_qwen4_layer_is_linear(il)) continue; e->qwen4_lin_state_pool[il] = ds4_gpu_tensor_alloc(state_bytes * slots); e->qwen4_lin_hist_pool[il] = ds4_gpu_tensor_alloc(hist_bytes * slots); if (!e->qwen4_lin_state_pool[il] || !e->qwen4_lin_hist_pool[il]) { for (uint32_t j = 0; j <= il; j++) { ds4_gpu_tensor_free(e->qwen4_lin_state_pool[j]); ds4_gpu_tensor_free(e->qwen4_lin_hist_pool[j]); e->qwen4_lin_state_pool[j] = NULL; e->qwen4_lin_hist_pool[j] = NULL; } return false; } } e->qwen4_pool_slots = slots; e->qwen4_pool_used = 0; return true; } static int qwen4_state_pool_take(ds4_engine *e) { for (uint32_t i = 0; i < e->qwen4_pool_slots; i++) { if (!(e->qwen4_pool_used & (UINT64_C(1) << i))) { e->qwen4_pool_used |= UINT64_C(1) << i; return (int)i; } } return -1; } static void qwen4_state_pool_free(ds4_engine *e) { for (uint32_t il = 0; il < DS4_MAX_LAYER; il++) { ds4_gpu_tensor_free(e->qwen4_lin_state_pool[il]); ds4_gpu_tensor_free(e->qwen4_lin_hist_pool[il]); e->qwen4_lin_state_pool[il] = NULL; e->qwen4_lin_hist_pool[il] = NULL; } e->qwen4_pool_slots = 0; e->qwen4_pool_used = 0; } /* Move the transients out of the first session and into the engine, which * then outlives every session and lends them to the next ones. The session * keeps using the same buffers; only ownership moves. */ static void qwen4_graph_transfer_scratch(ds4_qwen4_gpu_graph *dst, ds4_qwen4_gpu_graph *src) { memset(dst, 0, sizeof(*dst)); dst->cap_tokens = src->cap_tokens; dst->n_block_cap = src->n_block_cap; dst->k_blocks = src->k_blocks; dst->sel_stride = src->sel_stride; #define QWEN4_TAKE(field_) dst->field_ = src->field_; DS4_QWEN4_SCRATCH_FIELDS(QWEN4_TAKE) #undef QWEN4_TAKE dst->host_row = src->host_row; dst->owns_scratch = true; src->owns_scratch = false; } static uint64_t qwen4_graph_scratch_bytes(const ds4_qwen4_gpu_graph *g) { uint64_t total = 0; #define QWEN4_COUNT(field_) total += ds4_gpu_tensor_bytes(g->field_); DS4_QWEN4_SCRATCH_FIELDS(QWEN4_COUNT) #undef QWEN4_COUNT return total; } static ds4_gpu_tensor *qwen4_graph_alloc_f32(uint64_t n) { ds4_gpu_tensor *t = ds4_gpu_tensor_alloc(n * sizeof(float)); if (t) ds4_gpu_tensor_fill_f32(t, 0.0f, n); return t; } /* shared borrows the engine arena for the transients; NULL allocates private * ones. A borrowing graph must fit inside the arena it borrows: the caller * checks cap_tokens and n_block_cap before passing one. * * state_pool/hist_pool, when given, hold every slot's recurrent state in one * tensor per layer and this graph takes slot's view of it. A batched decode * then advances all its rows with one dispatch; without a pool each session * keeps its own tensors and the batch falls back to a dispatch per row. */ static bool qwen4_graph_alloc(ds4_qwen4_gpu_graph *g, const ds4_weights *w, uint32_t ctx_cap, uint32_t cap_tokens, bool mtp, const ds4_qwen4_gpu_graph *shared, ds4_gpu_tensor *const *state_pool, ds4_gpu_tensor *const *hist_pool, uint32_t slot) { memset(g, 0, sizeof(*g)); if (!qwen4_graph_weights_supported(w)) return false; mtp = mtp && DS4_N_NEXTN_PREDICT != 0; const uint64_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc, T = cap_tokens; const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM; const uint64_t iq_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; g->ctx_cap = ctx_cap; g->cap_tokens = cap_tokens; g->k_blocks = DS4_N_INDEXER_TOP_K / 4u; g->n_logit_rows = mtp ? 3u : 1u; g->n_block_cap = ctx_cap / 4u + 1u; static bool warned_ctx = false; if (!warned_ctx && g_qwen4_native_ctx && ctx_cap > g_qwen4_native_ctx && !g_qwen4_rope_yarn) { warned_ctx = true; fprintf(stderr, "ds4: context %u exceeds the native %u tokens; prompts past that need " "DS4_QWEN4_YARN_FACTOR (see README)\n", ctx_cap, g_qwen4_native_ctx); } g->sel_stride = g->k_blocks * 4u + 4u; bool ok = true; g->owns_scratch = shared == NULL; #define QWEN4_ALLOC(field_, n_) do { g->field_ = qwen4_graph_alloc_f32(n_); ok = ok && g->field_; } while (0) if (shared) { #define QWEN4_BORROW(field_) g->field_ = shared->field_; DS4_QWEN4_SCRATCH_FIELDS(QWEN4_BORROW) #undef QWEN4_BORROW g->host_row = shared->host_row; goto private_state; } QWEN4_ALLOC(R, T * hc_dim); QWEN4_ALLOC(xn, T * hc_dim); QWEN4_ALLOC(lo, T * DS4_N_HC_LOWRANK); QWEN4_ALLOC(inj, T * hc * DS4_QWEN4_HC_CHUNKS * hc); QWEN4_ALLOC(inj_alt, T * hc * DS4_QWEN4_HC_CHUNKS * hc); QWEN4_ALLOC(mixed, T * E); QWEN4_ALLOC(blk, T * E); QWEN4_ALLOC(qkv, T * conv_dim); QWEN4_ALLOC(z, T * v_dim); QWEN4_ALLOC(ga, T * DS4_N_LIN_V_HEAD); QWEN4_ALLOC(gb, T * DS4_N_LIN_V_HEAD); QWEN4_ALLOC(lin_o, T * v_dim); QWEN4_ALLOC(ple_emb, T * E); QWEN4_ALLOC(ple_key, T * hc_dim); QWEN4_ALLOC(ple_val, T * E); QWEN4_ALLOC(ple_gated, T * hc_dim); QWEN4_ALLOC(ple_normed, T * hc_dim); QWEN4_ALLOC(qg, T * 2u * q_dim); QWEN4_ALLOC(kp, T * kv_dim); QWEN4_ALLOC(vp, T * kv_dim); QWEN4_ALLOC(iq, T * iq_dim); QWEN4_ALLOC(ik, T * DS4_N_INDEXER_HEAD_DIM); QWEN4_ALLOC(attn_part, ds4_gpu_qwen4_attn_part_floats(3u, DS4_N_HEAD, DS4_N_HEAD_DIM)); QWEN4_ALLOC(q, T * q_dim); QWEN4_ALLOC(gate, T * q_dim); QWEN4_ALLOC(iqn, T * iq_dim); QWEN4_ALLOC(attn_o, T * q_dim); QWEN4_ALLOC(score, T * g->n_block_cap); QWEN4_ALLOC(tile_max, T * (g->n_block_cap / 8u + 1u)); QWEN4_ALLOC(sel_blocks, T * g->k_blocks); QWEN4_ALLOC(sel_tokens, T * g->sel_stride); QWEN4_ALLOC(n_sel, T); QWEN4_ALLOC(router, T * DS4_N_EXPERT); QWEN4_ALLOC(selected, T * DS4_N_EXPERT_USED); QWEN4_ALLOC(weights, T * DS4_N_EXPERT_USED); QWEN4_ALLOC(mid, T * (DS4_N_EXPERT_USED + 1u) * DS4_N_FF_EXP); QWEN4_ALLOC(part, T * (DS4_N_EXPERT_USED + 1u) * E); QWEN4_ALLOC(sh_gate_logit, T); QWEN4_ALLOC(moe_lists, (uint64_t)DS4_N_EXPERT * T); QWEN4_ALLOC(moe_counts, DS4_N_EXPERT); QWEN4_ALLOC(sh_gate, T * DS4_N_FF_EXP); QWEN4_ALLOC(sh_up, T * DS4_N_FF_EXP); QWEN4_ALLOC(sh_mid, T * DS4_N_FF_EXP); QWEN4_ALLOC(sh_out, T * E); QWEN4_ALLOC(hc_u, T * hc_dim); QWEN4_ALLOC(hc_lo_act, T * DS4_N_HC_LOWRANK); g->host_row = xmalloc(T * hc_dim * sizeof(float)); private_state: /* The PLE conv history carries across calls and the logits belong to the * caller that reads them back, so neither can be shared. */ QWEN4_ALLOC(ple_hist, (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * hc_dim); QWEN4_ALLOC(logits, (uint64_t)g->n_logit_rows * DS4_N_VOCAB); if (mtp) { QWEN4_ALLOC(mtp_e, 3u * E); QWEN4_ALLOC(mtp_cat, 3u * (hc + 1u) * 2u * E); QWEN4_ALLOC(mtp_proj, 3u * (hc + 1u) * E); QWEN4_ALLOC(mtp_R, 3u * hc_dim); QWEN4_ALLOC(mtp_argmax, 1u); QWEN4_ALLOC(mtp_argmax_tmp, ((DS4_N_VOCAB + 4095u) / 4096u) * 2u); QWEN4_ALLOC(snap_ple_hist, (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * hc_dim); } #undef QWEN4_ALLOC const uint64_t state_bytes = v_dim * DS4_N_LIN_HEAD_DIM * sizeof(float); const uint64_t hist_bytes = (uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim * sizeof(float); for (uint32_t il = 0; il < DS4_N_LAYER && ok; il++) { if (ds4_qwen4_layer_is_linear(il)) { if (state_pool && hist_pool) { g->layer_lin_state[il] = ds4_gpu_tensor_view(state_pool[il], (uint64_t)slot * state_bytes, state_bytes); g->layer_lin_hist[il] = ds4_gpu_tensor_view(hist_pool[il], (uint64_t)slot * hist_bytes, hist_bytes); } else { g->layer_lin_state[il] = qwen4_graph_alloc_f32(v_dim * DS4_N_LIN_HEAD_DIM); g->layer_lin_hist[il] = qwen4_graph_alloc_f32((uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim); } ok = ok && g->layer_lin_state[il] && g->layer_lin_hist[il]; if (ok && mtp) { g->snap_lin_state[il] = qwen4_graph_alloc_f32(v_dim * DS4_N_LIN_HEAD_DIM); g->snap_lin_hist[il] = qwen4_graph_alloc_f32((uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim); ok = g->snap_lin_state[il] && g->snap_lin_hist[il]; } } else { g->layer_k_cache[il] = ds4_gpu_tensor_alloc((uint64_t)ctx_cap * kv_dim * 2u); g->layer_v_cache[il] = ds4_gpu_tensor_alloc((uint64_t)ctx_cap * kv_dim * 2u); g->layer_ik_cache[il] = qwen4_graph_alloc_f32((uint64_t)ctx_cap * DS4_N_INDEXER_HEAD_DIM); g->layer_block_key[il] = ds4_gpu_tensor_alloc((uint64_t)g->n_block_cap * DS4_N_INDEXER_HEAD_DIM * 2u); ok = ok && g->layer_k_cache[il] && g->layer_v_cache[il] && g->layer_ik_cache[il] && g->layer_block_key[il]; } } g->pos3 = qwen4_graph_alloc_f32((uint64_t)ctx_cap * 4u); ok = ok && g->pos3; g->host_pos3 = xmalloc(T * 4u * sizeof(uint32_t)); g->host_logits = xmalloc((uint64_t)g->n_logit_rows * DS4_N_VOCAB * sizeof(float)); if (!ok) { fprintf(stderr, "ds4: Qwen3.8 graph allocation failed (ctx %u)\n", ctx_cap); qwen4_graph_free(g); return false; } return true; } static bool qwen4_graph_load_steering(ds4_qwen4_gpu_graph *g, const char *path, float attn_scale, float ffn_scale) { if (!g) return false; 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 uint32_t n_layers = directional_steering_layer_count(); if (n_layers == 0) return false; const uint64_t n = (uint64_t)n_layers * DS4_N_EMBD; float *dirs = xmalloc((size_t)n * sizeof(dirs[0])); bool ok = read_f32_binary_file(path, dirs, n); if (ok) { g->steer_dirs = ds4_gpu_tensor_alloc(n * sizeof(dirs[0])); ok = g->steer_dirs != NULL && ds4_gpu_tensor_write(g->steer_dirs, 0, dirs, n * sizeof(dirs[0])) != 0; } free(dirs); if (!ok) { fprintf(stderr, "ds4: failed to load directional steering vectors from %s\n", path); return false; } g->steer_attn_scale = attn_scale; g->steer_ffn_scale = ffn_scale; fprintf(stderr, "ds4: Qwen3.8 directional steering enabled: %s attn=%g ffn=%g\n", path, (double)attn_scale, (double)ffn_scale); return true; } static bool qwen4_graph_apply_steering_attn(ds4_qwen4_gpu_graph *g, uint32_t il, uint32_t T) { if (!g || !g->steer_dirs || g->steer_attn_scale == 0.0f) return true; return ds4_gpu_directional_steering_project_tensor( g->blk, g->steer_dirs, il, DS4_N_EMBD, T, g->steer_attn_scale) != 0; } static bool qwen4_graph_apply_steering_ffn(ds4_qwen4_gpu_graph *g, uint32_t il, uint32_t T) { if (!g || !g->steer_dirs || g->steer_ffn_scale == 0.0f) return true; return ds4_gpu_directional_steering_project_tensor( g->R, g->steer_dirs, il, DS4_N_EMBD, T * DS4_N_HC, g->steer_ffn_scale) != 0; } /* Last prompt token, HC-mean of residual R. Filename uses pos0 so * dir-steering/tools/build_direction.py can reuse the DeepSeek dump layout. */ static void qwen4_graph_dump_last_ffn(ds4_qwen4_gpu_graph *g, uint32_t il, uint32_t T) { const char *prefix = metal_graph_debug_prefix_for("ffn_out", il, 0); if (!g || !g->R || T == 0 || !prefix) return; if (ds4_gpu_synchronize() == 0) return; const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC; float *hc_row = xmalloc((size_t)hc * E * sizeof(hc_row[0])); const uint64_t off = (uint64_t)(T - 1u) * hc * E * sizeof(float); if (ds4_gpu_tensor_read(g->R, off, hc_row, (uint64_t)hc * E * sizeof(float)) != 0) { float *mean = xmalloc((size_t)E * sizeof(mean[0])); for (uint32_t i = 0; i < E; i++) { float s = 0.0f; for (uint32_t h = 0; h < hc; h++) s += hc_row[(uint64_t)h * E + i]; mean[i] = s / (float)hc; } char path[1024]; snprintf(path, sizeof(path), "%s_ffn_out-%u_pos0.bin", prefix, il); if (write_f32_binary_file(path, mean, E)) { fprintf(stderr, "ds4: dumped ffn_out layer %u last-token HC-mean to %s\n", il, path); } free(mean); } free(hc_row); (void)glm_graph_begin_commands_if_needed(); } static void qwen4_graph_dump_last_attn(ds4_qwen4_gpu_graph *g, uint32_t il, uint32_t T) { const char *prefix = metal_graph_debug_prefix_for("attn_out", il, 0); if (!g || !g->blk || T == 0 || !prefix) return; if (ds4_gpu_synchronize() == 0) return; float *row = xmalloc((size_t)DS4_N_EMBD * sizeof(row[0])); const uint64_t off = (uint64_t)(T - 1u) * DS4_N_EMBD * sizeof(float); if (ds4_gpu_tensor_read(g->blk, off, row, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0) { char path[1024]; snprintf(path, sizeof(path), "%s_attn_out-%u_pos0.bin", prefix, il); if (write_f32_binary_file(path, row, DS4_N_EMBD)) { fprintf(stderr, "ds4: dumped attn_out layer %u last-token to %s\n", il, path); } } free(row); (void)glm_graph_begin_commands_if_needed(); } static void qwen4_graph_reset(ds4_qwen4_gpu_graph *g) { (void)ds4_gpu_synchronize(); const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (g->layer_lin_state[il]) ds4_gpu_tensor_fill_f32(g->layer_lin_state[il], 0.0f, v_dim * DS4_N_LIN_HEAD_DIM); if (g->layer_lin_hist[il]) ds4_gpu_tensor_fill_f32(g->layer_lin_hist[il], 0.0f, (uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim); } ds4_gpu_tensor_fill_f32(g->ple_hist, 0.0f, (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * DS4_N_EMBD * DS4_N_HC); for (uint32_t i = 0; i < DS4_MAX_PLE_NGRAM; i++) g->ple_prev[i] = DS4_PLE_EOS_ID; g->pos = 0; g->mtp_pos = 0; g->mtp_last_rows = 0; g->mrope_delta = 0; g->snap_valid = false; g->snap2_valid = false; g->snap0_valid = false; g->snap_after_first = false; g->snap_after_second = false; } /* rows > 0 limits the product to a contiguous leading prefix of w. */ static bool qwen4_gemv_rows(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w, const ds4_gpu_tensor *x, uint32_t n_tok, uint64_t rows) { const uint64_t in_dim = w->dim[0], full_dim = w->ndim >= 2 ? w->dim[1] : 1u; const uint64_t out_dim = rows && rows < full_dim ? rows : full_dim; int rc = 0; #if !defined(__APPLE__) /* Qwen's recurrent graph keeps activations in FP32. The generic CUDA * projections can round them to half or quantize them for other models. */ if (in_dim <= UINT32_MAX && out_dim <= UINT32_MAX) rc = ds4_gpu_qwen4_dense_mm_tensor(out, x, m->map, m->size, w->abs_offset, w->type, n_tok, (uint32_t)in_dim, (uint32_t)out_dim); #else /* Small F16 batches use float operands. Two/three-row verification and * large prefills retain their existing kernels. The legacy switch is an * arithmetic/performance control, not a user tuning option. */ const bool legacy = getenv("DS4_QWEN4_DENSE_MM_LEGACY") != NULL; const bool f16_batch = !legacy && n_tok > 3u && w->type == DS4_TENSOR_F16 && n_tok <= 64u && (out_dim <= 512u || n_tok > 8u); if (((n_tok > 8u && w->type == DS4_TENSOR_F32) || f16_batch) && (in_dim % 32) == 0) { rc = ds4_gpu_qwen4_dense_mm_tensor(out, x, m->map, m->size, w->abs_offset, w->type, n_tok, (uint32_t)in_dim, (uint32_t)out_dim); if (rc) return true; } /* Q8 decode batches of 8 or 16 rows on fp32 simdgroup matrices: the * per-token kernel reads the matrix once per four rows and streams x * once per four weight rows; this reads the weights once and x once per * 32 rows. DS4_QWEN4_NO_BATCH_MM=1 keeps the per-token kernel for A/B. */ if (!legacy && w->type == DS4_TENSOR_Q8_0 && n_tok >= 5u && n_tok <= 32u && (in_dim % 32) == 0 && (out_dim % 32) == 0 && getenv("DS4_QWEN4_NO_BATCH_MM") == NULL && ds4_gpu_tensor_bytes(x) >= 32u * in_dim * sizeof(float) && ds4_gpu_tensor_bytes(out) >= 32u * out_dim * sizeof(float)) { /* Other widths round up to the tile: rows past n_tok are scratch * rows both buffers hold, computed and ignored. Past sixteen rows * the second half is a second product on views of both. */ const uint32_t first = n_tok <= 8u ? 8u : 16u; rc = ds4_gpu_qwen4_batch_mm_q8_tensor(out, x, m->map, m->size, w->abs_offset, first, (uint32_t)in_dim, (uint32_t)out_dim); if (rc && n_tok > 16u) { ds4_gpu_tensor *x2 = ds4_gpu_tensor_view(x, 16u * in_dim * sizeof(float), 16u * in_dim * sizeof(float)); ds4_gpu_tensor *o2 = ds4_gpu_tensor_view(out, 16u * out_dim * sizeof(float), 16u * out_dim * sizeof(float)); rc = x2 && o2 && ds4_gpu_qwen4_batch_mm_q8_tensor(o2, x2, m->map, m->size, w->abs_offset, 16u, (uint32_t)in_dim, (uint32_t)out_dim); ds4_gpu_tensor_free(o2); ds4_gpu_tensor_free(x2); } if (rc) return true; } switch (w->type) { case DS4_TENSOR_Q8_0: rc = ds4_gpu_qwen4_matmul_q8_0_tensor(out, m->map, m->size, w->abs_offset, in_dim, out_dim, x, n_tok); break; case DS4_TENSOR_F16: rc = ds4_gpu_matmul_f16_tensor(out, m->map, m->size, w->abs_offset, in_dim, out_dim, x, n_tok); break; case DS4_TENSOR_F32: rc = ds4_gpu_matmul_f32_tensor(out, m->map, m->size, w->abs_offset, in_dim, out_dim, x, n_tok); break; case DS4_TENSOR_Q4_0: rc = ds4_gpu_matmul_quant_tensor(out, m->map, m->size, w->abs_offset, w->type, in_dim, out_dim, x, n_tok); break; case DS4_TENSOR_BF16: { ds4_gpu_tensor *outs[1] = { out }; const uint64_t offs[1] = { w->abs_offset }; const uint32_t types[1] = { w->type }; const uint32_t rows[1] = { (uint32_t)out_dim }; rc = ds4_gpu_qwen4_multi_gemv_tensor(x, n_tok, (uint32_t)in_dim, 1, outs, m->map, m->size, offs, types, rows); break; } default: break; } #endif if (!rc) { fprintf(stderr, "ds4: Qwen3.8 matmul failed for %.*s (type %u, %" PRIu64 "x%" PRIu64 ", %u tokens)\n", (int)w->name.len, w->name.ptr, w->type, in_dim, out_dim, n_tok); } return rc != 0; } static bool qwen4_gemv(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *w, const ds4_gpu_tensor *x, uint32_t n_tok) { return qwen4_gemv_rows(out, m, w, x, n_tok, 0u); } /* The MTP draft only needs its argmax, and the tokenizer assigns ids in * merge order, so the frequent tokens sit at low ids: DS4_QWEN4_MTP_DRAFT_ROWS * scores the draft over that leading prefix of the output head only. The * verify rows always use the full head, so committed tokens and every * teacher-forced logit are unchanged; only which draft is proposed can differ * when the true argmax lies beyond the prefix. Default: the full vocabulary. */ /* DS4_QWEN4_MTP_DRAFT_VOCAB=: token ids, one per line, most frequent * first (the MTPLX FR-Spec idea). The draft head becomes a gathered copy of * those output rows, so the draft is scored over that subset and the argmax * maps back through the list; the verify rows still use the full head. * Loaded once per graph; a missing or invalid file leaves the full head. */ static bool qwen4_mtp_draft_head_load(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_tensor *w) { if (g->draft_head_tried) return g->draft_head != NULL; g->draft_head_tried = true; const char *path = getenv("DS4_QWEN4_MTP_DRAFT_VOCAB"); if (!path || !path[0] || w->type != DS4_TENSOR_Q8_0 || w->ndim < 2 || (w->dim[0] % 32u) != 0) return false; FILE *fp = fopen(path, "r"); if (!fp) { fprintf(stderr, "ds4: Qwen3.8 MTP draft vocabulary %s: %s\n", path, strerror(errno)); return false; } const uint64_t V = w->dim[1]; int32_t *ids = malloc(V * sizeof(int32_t)); uint8_t *seen = calloc(V, 1); uint32_t n = 0; long long v; while (ids && seen && n < V && fscanf(fp, "%lld", &v) == 1) { if (v < 0 || v >= (long long)V || seen[v]) continue; seen[v] = 1; ids[n++] = (int32_t)v; } fclose(fp); free(seen); if (!ids || n == 0 || n >= V) { fprintf(stderr, "ds4: Qwen3.8 MTP draft vocabulary %s: need 1..%" PRIu64 " distinct ids (got %u)\n", path, V - 1u, n); free(ids); return false; } const uint64_t row_bytes = (w->dim[0] / 32u) * 34u; const uint64_t bytes = (uint64_t)n * row_bytes; uint8_t *rows = malloc(bytes); ds4_gpu_tensor *head = rows ? ds4_gpu_tensor_alloc(bytes) : NULL; if (rows && head) { const uint8_t *src = (const uint8_t *)m->map + w->abs_offset; for (uint32_t i = 0; i < n; i++) memcpy(rows + (uint64_t)i * row_bytes, src + (uint64_t)ids[i] * row_bytes, row_bytes); } const bool ok = rows && head && ds4_gpu_tensor_write(head, 0, rows, bytes); free(rows); if (!ok) { fprintf(stderr, "ds4: Qwen3.8 MTP draft head upload failed\n"); ds4_gpu_tensor_free(head); free(ids); return false; } g->draft_head = head; g->draft_ids = ids; g->draft_rows = n; fprintf(stderr, "ds4: Qwen3.8 MTP draft head: %u of %" PRIu64 " vocabulary rows from %s (%.0f MiB)\n", n, V, path, (double)bytes / (1024.0 * 1024.0)); return true; } /* Decode-sized batches (a token, or the 2/3-token MTP verify) take the fused * kernels; DS4_QWEN4_NO_FUSE=1 forces the prefill kernels for A/B checks. * T = 3 uses them for the speculative verify only, so prefill tails of three * tokens keep their historical kernel selection. */ static bool qwen4_graph_fused(const ds4_qwen4_gpu_graph *g, uint32_t T) { static int no_fuse = -1; if (no_fuse < 0) no_fuse = getenv("DS4_QWEN4_NO_FUSE") != NULL; if (no_fuse || T <= 2u) return T <= 2u && !no_fuse; /* T = 3 takes the fused kernels for the speculative verify only; prefill * tails of three tokens keep their historical kernel selection. */ return T == 3u && g->verify_rows_exact; } static uint32_t qwen4_mtp_draft_rows(void) { /* read per cycle so the A/B harnesses can switch it per step */ const char *env = getenv("DS4_QWEN4_MTP_DRAFT_ROWS"); const long long v = env && env[0] ? strtoll(env, NULL, 10) : 0; return v > 0 && v < (long long)DS4_N_VOCAB ? (uint32_t)v : DS4_N_VOCAB; } /* top-k blocks per query: radix select (DS4_QWEN4_NO_IDX_SELECT=1 restores the argsort path) */ static int qwen4_idx_select(ds4_gpu_tensor *sel, const ds4_gpu_tensor *score, const ds4_gpu_tensor *tile_max, uint32_t n_blocks, uint32_t n_tokens, uint32_t top_k) { static int no_select = -1; if (no_select < 0) no_select = getenv("DS4_QWEN4_NO_IDX_SELECT") != NULL; return no_select ? ds4_gpu_indexer_topk_tensor(sel, score, n_blocks, n_tokens, top_k) : ds4_gpu_qwen4_idx_select_tensor(sel, score, tile_max, n_blocks, n_tokens, top_k); } /* grouped norm -> low-rank gate -> mixed; inject may be NULL (final mixer) */ static bool qwen4_graph_hc_mix(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_tensor *norm, const ds4_tensor *down, const ds4_tensor *up, const ds4_tensor *inject, uint32_t T) { bool ok = ds4_gpu_qwen4_hc_norm_tensor(g->xn, g->inj, g->R, m->map, m->size, norm->abs_offset, inject ? inject->abs_offset : 0, inject ? inject->type : DS4_TENSOR_F32, T, DS4_N_EMBD, DS4_N_HC, inject ? DS4_N_HC : 0u, DS4_RMS_EPS) && qwen4_gemv(g->lo, m, down, g->xn, T); if (T > 8u && up->type != DS4_TENSOR_Q8_0) { /* prefill: the up projection as a GEMM over the activated low-rank rows */ return ok && ds4_gpu_qwen4_hc_lo_act_tensor(g->hc_lo_act, g->lo, T, DS4_N_HC, DS4_N_HC_LOWRANK) && qwen4_gemv(g->hc_u, m, up, g->hc_lo_act, T) && ds4_gpu_qwen4_hc_mix_rows_tensor(g->mixed, g->hc_u, g->xn, T, DS4_N_EMBD, DS4_N_HC); } if (T == 3u && g->verify_rows_exact) { /* Split the 3-row gate/mix into the exact 2-row pair kernel plus the * 1-row generic kernel, so every row matches its T <= 2 rounding. */ const uint64_t dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; ds4_gpu_tensor *xn2 = ds4_gpu_tensor_view(g->xn, 0, 2u * dim * sizeof(float)); ds4_gpu_tensor *lo2 = ds4_gpu_tensor_view(g->lo, 0, 2u * DS4_N_HC_LOWRANK * sizeof(float)); ds4_gpu_tensor *mixed2 = ds4_gpu_tensor_view(g->mixed, 0, 2u * DS4_N_EMBD * sizeof(float)); const bool ok2 = xn2 && lo2 && mixed2 && ds4_gpu_qwen4_hc_gate_mix_tensor(mixed2, xn2, lo2, m->map, m->size, up->abs_offset, up->type, 2u, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_LOWRANK); ds4_gpu_tensor_free(mixed2); ds4_gpu_tensor_free(lo2); ds4_gpu_tensor_free(xn2); if (!ok2) return false; ds4_gpu_tensor *xn1 = ds4_gpu_tensor_view(g->xn, 2u * dim * sizeof(float), dim * sizeof(float)); ds4_gpu_tensor *lo1 = ds4_gpu_tensor_view(g->lo, 2u * DS4_N_HC_LOWRANK * sizeof(float), DS4_N_HC_LOWRANK * sizeof(float)); ds4_gpu_tensor *mixed1 = ds4_gpu_tensor_view(g->mixed, 2u * DS4_N_EMBD * sizeof(float), DS4_N_EMBD * sizeof(float)); const bool ok1 = xn1 && lo1 && mixed1 && ds4_gpu_qwen4_hc_gate_mix_tensor(mixed1, xn1, lo1, m->map, m->size, up->abs_offset, up->type, 1u, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_LOWRANK); ds4_gpu_tensor_free(mixed1); ds4_gpu_tensor_free(lo1); ds4_gpu_tensor_free(xn1); return ok1; } return ok && ds4_gpu_qwen4_hc_gate_mix_tensor(g->mixed, g->xn, g->lo, m->map, m->size, up->abs_offset, up->type, T, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_LOWRANK); } static bool qwen4_graph_linear(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t T) { const uint32_t conv_dim = DS4_N_LIN_CONV_DIM; bool paired = false; if (T == 1u && !g->mtp_R && ds4_gpu_qwen4_decode_fusions_enabled() && l->lin_qkv->type == DS4_TENSOR_Q8_0 && l->lin_gate->type == DS4_TENSOR_Q8_0) { paired = ds4_gpu_qwen4_q8_pair_tensor(g->qkv, g->z, m->map, m->size, l->lin_qkv->abs_offset, l->lin_gate->abs_offset, DS4_N_EMBD, l->lin_qkv->dim[1], l->lin_gate->dim[1], g->mixed, T) != 0; } if (!paired && !(qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) && qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T))) { return false; } bool ok; if (qwen4_graph_fused(g, T)) { ok = ds4_gpu_qwen4_gdn_front_tensor(g->qkv, g->layer_lin_hist[il], g->mixed, g->ga, g->gb, m->map, m->size, l->lin_conv->abs_offset, l->lin_alpha->abs_offset, l->lin_beta->abs_offset, l->lin_a->abs_offset, l->lin_dt_bias->abs_offset, l->lin_alpha->type, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_N_LIN_CONV, DS4_N_EMBD, g->snap_after_first ? g->snap_lin_hist[il] : NULL, 0u, g->snap_after_second ? g->snap2_lin_hist[il] : NULL, 1u) != 0; } else { ok = qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) && qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T) && ds4_gpu_qwen4_conv_stream_tensor(g->qkv, g->layer_lin_hist[il], m->map, m->size, l->lin_conv->abs_offset, T, conv_dim, DS4_N_LIN_CONV, true) && ds4_gpu_qwen4_gdn_prep_tensor(g->qkv, g->ga, g->gb, m->map, m->size, l->lin_a->abs_offset, l->lin_dt_bias->abs_offset, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM); } if (ok) { ok = ds4_gpu_qwen4_gdn_scan_tensor(g->lin_o, g->layer_lin_state[il], g->qkv, g->ga, g->gb, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, g->snap_after_first ? g->snap_lin_state[il] : NULL, 0u, g->snap_after_second ? g->snap2_lin_state[il] : NULL, 1u) != 0; } if (ok) { ok = ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size, l->lin_norm->abs_offset, T, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0; } if (ok) ok = qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T); return ok; } /* T tokens at positions pos0.. of one attention layer. Tokens with no more * complete blocks than the budget attend densely; the rest score the pooled * block keys and attend their top-k blocks plus the incomplete tail. The * chunk is split at that boundary so every selection list stays bounded. */ /* Dense-prefix + sparse-tail attention core for rows [0,T) of the given row * views at absolute position cpos0. n_blocks_after covers the whole chunk * (the caller writes the pooled block keys once). */ static bool qwen4_graph_attention_core(ds4_qwen4_gpu_graph *g, uint32_t il, ds4_gpu_tensor *q_rows, ds4_gpu_tensor *gate_rows, ds4_gpu_tensor *iqn_rows, ds4_gpu_tensor *o_rows, uint32_t n_blocks_after, uint32_t cpos0, uint32_t cT) { const uint32_t ratio = 4u; const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM; const float scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); const uint32_t sparse_pos = (g->k_blocks + 1u) * ratio - 1u; const uint32_t clast = cpos0 + cT - 1u; const uint32_t n_dense = clast < sparse_pos ? cT : (sparse_pos > cpos0 ? sparse_pos - cpos0 : 0u); if (n_dense > 0 && !ds4_gpu_qwen4_attn_decode_tensor(o_rows, q_rows, gate_rows, g->layer_k_cache[il], g->layer_v_cache[il], g->sel_tokens, g->n_sel, cT <= 2u ? g->attn_part : NULL, n_dense, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, cpos0, false, g->sel_stride, scale)) { return false; } if (n_dense >= cT) return true; const uint32_t n_sparse = cT - n_dense; const uint32_t sp0 = cpos0 + n_dense; ds4_gpu_tensor *q = ds4_gpu_tensor_view(q_rows, (uint64_t)n_dense * q_dim * sizeof(float), (uint64_t)n_sparse * q_dim * sizeof(float)); ds4_gpu_tensor *gate = ds4_gpu_tensor_view(gate_rows, (uint64_t)n_dense * q_dim * sizeof(float), (uint64_t)n_sparse * q_dim * sizeof(float)); ds4_gpu_tensor *o = ds4_gpu_tensor_view(o_rows, (uint64_t)n_dense * q_dim * sizeof(float), (uint64_t)n_sparse * q_dim * sizeof(float)); ds4_gpu_tensor *iqn = ds4_gpu_tensor_view(iqn_rows, (uint64_t)n_dense * DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float), (uint64_t)n_sparse * DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float)); const bool ok = q && gate && o && iqn && ds4_gpu_qwen4_idx_score_tensor(g->score, n_sparse <= 2u ? g->tile_max : NULL, iqn, g->layer_block_key[il], n_sparse, n_blocks_after, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, sp0, ratio) && qwen4_idx_select(g->sel_blocks, g->score, n_sparse <= 2u ? g->tile_max : NULL, n_blocks_after, n_sparse, g->k_blocks) && ds4_gpu_qwen4_idx_expand_tensor(g->sel_tokens, g->n_sel, g->sel_blocks, n_sparse, g->k_blocks, ratio, sp0, g->sel_stride) && ds4_gpu_qwen4_attn_decode_tensor(o, q, gate, g->layer_k_cache[il], g->layer_v_cache[il], g->sel_tokens, g->n_sel, cT <= 2u ? g->attn_part : NULL, n_sparse, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, sp0, true, g->sel_stride, scale); ds4_gpu_tensor_free(iqn); ds4_gpu_tensor_free(o); ds4_gpu_tensor_free(gate); ds4_gpu_tensor_free(q); return ok; } /* The per-session part of an attention layer for T rows at pos0: cache * appends, pooled block keys and the attention core; the projections before * it and the output projection after it are row-agnostic. */ static bool qwen4_graph_attention_tail(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t pos0, uint32_t T) { const uint32_t ratio = 4u; const uint32_t last = pos0 + T - 1u; const uint32_t q_dim = DS4_N_HEAD * DS4_N_HEAD_DIM; if (!(ds4_gpu_qwen4_attn_prep_tensor(g->q, g->gate, g->layer_k_cache[il], g->layer_v_cache[il], g->iqn, g->layer_ik_cache[il], g->qg, g->kp, g->vp, g->iq, g->ik, g->pos3, m->map, m->size, l->attn_q_norm->abs_offset, l->attn_k_norm->abs_offset, l->indexer_q_norm->abs_offset, T, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, pos0, g->ctx_cap, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS))) { return false; } /* blocks whose last token falls in this chunk get their pooled keys */ const uint32_t first_block = pos0 / ratio; const uint32_t n_blocks_after = (last + 1u) / ratio; if (n_blocks_after > first_block && !ds4_gpu_qwen4_idx_block_key_tensor(g->layer_block_key[il], g->layer_ik_cache[il], g->pos3, m->map, m->size, l->indexer_k_norm->abs_offset, first_block, n_blocks_after - first_block, ratio, DS4_N_INDEXER_HEAD_DIM, DS4_N_ROT, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS)) { return false; } if (T == 3u && g->verify_rows_exact) { /* 3-row speculative verify: run the attention core as 2/1-row * sub-batches so every dispatch keeps the exact T <= 2 kernel paths * (prefill T = 3 tails keep their own kernel selection). */ for (uint32_t sub0 = 0; sub0 < T; sub0 += 2u) { const uint32_t subT = T - sub0 > 2u ? 2u : T - sub0; ds4_gpu_tensor *q = sub0 ? ds4_gpu_tensor_view(g->q, (uint64_t)sub0 * q_dim * sizeof(float), (uint64_t)subT * q_dim * sizeof(float)) : g->q; ds4_gpu_tensor *gate = sub0 ? ds4_gpu_tensor_view(g->gate, (uint64_t)sub0 * q_dim * sizeof(float), (uint64_t)subT * q_dim * sizeof(float)) : g->gate; ds4_gpu_tensor *o = sub0 ? ds4_gpu_tensor_view(g->attn_o, (uint64_t)sub0 * q_dim * sizeof(float), (uint64_t)subT * q_dim * sizeof(float)) : g->attn_o; ds4_gpu_tensor *iqn = sub0 ? ds4_gpu_tensor_view(g->iqn, (uint64_t)sub0 * DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float), (uint64_t)subT * DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM * sizeof(float)) : g->iqn; /* each sub-batch sees the block universe of its own T <= 2 * verify (a whole-chunk universe would let rows select blocks a * true 2-row pass could not see) */ const uint32_t nba_sub = (pos0 + sub0 + subT) / ratio; const bool ok = q && gate && o && iqn && qwen4_graph_attention_core(g, il, q, gate, iqn, o, nba_sub, pos0 + sub0, subT); if (sub0) { ds4_gpu_tensor_free(iqn); ds4_gpu_tensor_free(o); ds4_gpu_tensor_free(gate); ds4_gpu_tensor_free(q); } if (!ok) return false; } } else if (!qwen4_graph_attention_core(g, il, g->q, g->gate, g->iqn, g->attn_o, n_blocks_after, pos0, T)) { return false; } return true; } static bool qwen4_graph_attention(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t pos0, uint32_t T) { { bool ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T); if (ok && qwen4_graph_fused(g, T)) { ds4_gpu_tensor *outs[4] = { g->kp, g->vp, g->iq, g->ik }; const uint64_t offs[4] = { l->attn_k->abs_offset, l->attn_v->abs_offset, l->indexer_q_proj->abs_offset, l->indexer_k_proj->abs_offset }; const uint32_t types[4] = { l->attn_k->type, l->attn_v->type, l->indexer_q_proj->type, l->indexer_k_proj->type }; const uint32_t rows[4] = { (uint32_t)l->attn_k->dim[1], (uint32_t)l->attn_v->dim[1], (uint32_t)l->indexer_q_proj->dim[1], (uint32_t)l->indexer_k_proj->dim[1] }; ok = ds4_gpu_qwen4_multi_gemv_tensor(g->mixed, T, DS4_N_EMBD, 4, outs, m->map, m->size, offs, types, rows) != 0; } else if (ok) { ok = qwen4_gemv(g->kp, m, l->attn_k, g->mixed, T) && qwen4_gemv(g->vp, m, l->attn_v, g->mixed, T) && qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T) && qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T); } if (!ok) return false; } return qwen4_graph_attention_tail(g, m, l, il, pos0, T) && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T); } /* router GEMV, top-k (+ shared gate logit), experts with the shared expert as * an extra slot, weighted reduce with the hc combine folded in */ static bool qwen4_moe_profile_boundary(bool enabled, double *last, double *elapsed) { if (!enabled) return true; if (!ds4_gpu_end_commands()) return false; const double now = now_sec(); *elapsed = now - *last; *last = now; return glm_graph_begin_commands_if_needed(); } static bool qwen4_graph_moe(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t T) { const bool profile = T > 8u && getenv("DS4_QWEN4_MOE_PROFILE") != NULL; double elapsed[7] = {0}, last = 0; if (!qwen4_moe_profile_boundary(profile, &last, &elapsed[0])) return false; bool ok = qwen4_gemv(g->router, m, l->ffn_gate_inp, g->mixed, T) && ds4_gpu_qwen4_router_topk_tensor(g->selected, g->weights, g->router, g->mixed, m->map, m->size, l->ffn_gate_inp_shexp->abs_offset, l->ffn_gate_inp_shexp->type, DS4_N_EMBD, g->sh_gate_logit, T, DS4_N_EXPERT, DS4_N_EXPERT_USED) && qwen4_moe_profile_boundary(profile, &last, &elapsed[0]); /* Prefill-sized batches route each expert's tokens through tiled GEMMs * (weights read once per 32 tokens) and run the shared expert as dense * GEMMs; decode keeps the per-(token, slot) row kernels with the shared * expert as an extra slot. */ /* The tiled expert GEMM pays when each expert sees enough tokens to fill a * tile: a prefill chunk of 8192 rows gives every one of the 512 experts * about 160 of them. A decode batch does not - sixteen rows spread ten * choices each over 512 experts, so an expert averages a third of a token * and the tiles run nearly empty. Above this row count the GEMM wins; * below it the per-token expert path does, by 4% at sixteen rows. */ #ifdef DS4_HAS_QWEN4_METAL const uint32_t mm_min = 64u; #else const uint32_t mm_min = 8u; #endif const bool mm = T > mm_min && (DS4_N_EMBD % 64u) == 0 && (DS4_N_FF_EXP % 64u) == 0 && qwen4_expert_type_has_mm(l->ffn_gate_exps->type) && l->ffn_up_exps->type == l->ffn_gate_exps->type && qwen4_expert_type_has_mm(l->ffn_down_exps->type) && qwen4_graph_dense_ok(l->ffn_gate_shexp) && qwen4_graph_dense_ok(l->ffn_up_shexp) && qwen4_graph_dense_ok(l->ffn_down_shexp); if (mm) { if (ok) { ok = ds4_gpu_qwen4_moe_build_lists_tensor(g->moe_lists, g->moe_counts, g->selected, T, DS4_N_EXPERT_USED, DS4_N_EXPERT, g->cap_tokens) && qwen4_moe_profile_boundary(profile, &last, &elapsed[1]) && ds4_gpu_qwen4_moe_mm_mid_tensor(g->mid, g->mixed, g->moe_lists, g->moe_counts, m->map, m->size, l->ffn_gate_exps->abs_offset, l->ffn_up_exps->abs_offset, l->ffn_gate_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_EXPERT_USED, DS4_N_EMBD, DS4_N_FF_EXP, g->cap_tokens) && qwen4_moe_profile_boundary(profile, &last, &elapsed[2]) && qwen4_gemv(g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) && qwen4_gemv(g->sh_up, m, l->ffn_up_shexp, g->mixed, T) && ds4_gpu_swiglu_tensor(g->sh_mid, g->sh_gate, g->sh_up, T * DS4_N_FF_EXP, 0.0f, 1.0f) && qwen4_moe_profile_boundary(profile, &last, &elapsed[3]); } if (ok) { ok = ds4_gpu_qwen4_moe_mm_down_tensor(g->part, g->mid, g->moe_lists, g->moe_counts, m->map, m->size, l->ffn_down_exps->abs_offset, l->ffn_down_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_EXPERT_USED, DS4_N_FF_EXP, DS4_N_EMBD, g->cap_tokens) && qwen4_moe_profile_boundary(profile, &last, &elapsed[4]) && qwen4_gemv(g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T) && qwen4_moe_profile_boundary(profile, &last, &elapsed[5]); } if (ok) { ok = ds4_gpu_qwen4_moe_reduce_tensor(g->blk, g->part, g->weights, g->sh_gate_logit, g->sh_out, g->R, g->inj, T, DS4_N_EXPERT_USED, DS4_N_EXPERT_USED, DS4_N_EMBD, DS4_N_HC) != 0 && qwen4_moe_profile_boundary(profile, &last, &elapsed[6]); } if (profile) fprintf(stderr, "ds4: Qwen MoE ms pos=%u T=%u weight=%.*s ok=%d router=%.3f lists=%.3f mid=%.3f shared_mid=%.3f down=%.3f shared_down=%.3f reduce=%.3f\n", g->pos, T, (int)l->ffn_gate_exps->name.len, l->ffn_gate_exps->name.ptr, ok ? 1 : 0, elapsed[0] * 1e3, elapsed[1] * 1e3, elapsed[2] * 1e3, elapsed[3] * 1e3, elapsed[4] * 1e3, elapsed[5] * 1e3, elapsed[6] * 1e3); return ok; } /* A decode batch runs the shared expert as dense projections over its * rows, as the prefill path does: as a slot of the per-token kernels it * is read once per row, 5 MB of Q8 per row and layer. Single tokens and * verify rows keep the slot. */ #ifdef DS4_HAS_QWEN4_METAL const bool shared_dense = T > 8u && qwen4_graph_dense_ok(l->ffn_gate_shexp) && qwen4_graph_dense_ok(l->ffn_up_shexp) && qwen4_graph_dense_ok(l->ffn_down_shexp); #else const bool shared_dense = false; #endif if (ok && shared_dense) { ok = qwen4_gemv(g->sh_gate, m, l->ffn_gate_shexp, g->mixed, T) && qwen4_gemv(g->sh_up, m, l->ffn_up_shexp, g->mixed, T) && ds4_gpu_swiglu_tensor(g->sh_mid, g->sh_gate, g->sh_up, T * DS4_N_FF_EXP, 0.0f, 1.0f) && qwen4_gemv(g->sh_out, m, l->ffn_down_shexp, g->sh_mid, T); } /* Sixteen rows pick ten experts each but only a third as many distinct * ones: the grouped kernels read every expert once per four rows that * chose it instead of once per row, each row's arithmetic unchanged. * DS4_QWEN4_MOE_NO_GROUP=1 keeps the per-row kernels for A/B. */ #ifdef DS4_HAS_QWEN4_METAL const bool grouped = shared_dense && l->ffn_gate_exps->type == DS4_TENSOR_Q4_K && l->ffn_up_exps->type == DS4_TENSOR_Q4_K && l->ffn_down_exps->type == DS4_TENSOR_MXFP4 && getenv("DS4_QWEN4_MOE_NO_GROUP") == NULL; if (ok && grouped) { ok = ds4_gpu_qwen4_moe_build_lists_tensor(g->moe_lists, g->moe_counts, g->selected, T, DS4_N_EXPERT_USED, DS4_N_EXPERT, g->cap_tokens) && ds4_gpu_qwen4_moe_mid_grouped_tensor(g->mid, g->mixed, g->selected, g->moe_lists, g->moe_counts, g->cap_tokens, m->map, m->size, l->ffn_gate_exps->abs_offset, l->ffn_up_exps->abs_offset, l->ffn_gate_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_EMBD, DS4_N_FF_EXP) && ds4_gpu_qwen4_moe_down_grouped_tensor(g->part, g->mid, g->selected, g->moe_lists, g->moe_counts, g->cap_tokens, m->map, m->size, l->ffn_down_exps->abs_offset, l->ffn_down_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_FF_EXP, DS4_N_EMBD); } else #endif if (ok) { ok = ds4_gpu_qwen4_moe_mid_tensor(g->mid, g->mixed, g->selected, m->map, m->size, l->ffn_gate_exps->abs_offset, l->ffn_up_exps->abs_offset, l->ffn_gate_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_EMBD, DS4_N_FF_EXP, shared_dense ? 0u : l->ffn_gate_shexp->abs_offset, shared_dense ? 0u : l->ffn_up_shexp->abs_offset, shared_dense ? UINT32_MAX : l->ffn_gate_shexp->type) != 0 && ds4_gpu_qwen4_moe_down_tensor(g->part, g->mid, g->selected, m->map, m->size, l->ffn_down_exps->abs_offset, l->ffn_down_exps->type, DS4_N_EXPERT, T, DS4_N_EXPERT_USED, DS4_N_FF_EXP, DS4_N_EMBD, shared_dense ? 0u : l->ffn_down_shexp->abs_offset, shared_dense ? UINT32_MAX : l->ffn_down_shexp->type) != 0; } if (ok) { ok = ds4_gpu_qwen4_moe_reduce_tensor(g->blk, g->part, g->weights, g->sh_gate_logit, shared_dense ? g->sh_out : NULL, g->R, g->inj, T, DS4_N_EXPERT_USED, DS4_N_EXPERT_USED + (shared_dense ? 0u : 1u), DS4_N_EMBD, DS4_N_HC) != 0; } return ok; } /* host side: embedding rows tiled into R and the PLE n-gram gather */ static bool qwen4_graph_stage_inputs(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t T) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; float *row = g->host_row; const ds4_vision_span *spans = g->vis_spans; size_t n_spans = g->vis_span_count; if (!spans) spans = qwen4_fake_spans(&n_spans); for (uint32_t t = 0; t < T; t++) { float *dst = row + (uint64_t)t * hc_dim; const float *img = n_spans ? qwen4_span_row(spans, n_spans, g->pos + t) : NULL; if (img) memcpy(dst, img, E * sizeof(float)); else qwen4_ref_row(m, w->token_embd, (uint64_t)tokens[t], dst); for (uint32_t s = 1; s < hc; s++) memcpy(dst + (uint64_t)s * E, dst, E * sizeof(float)); qwen4_mrope_pos(spans, n_spans, g->pos + t, &g->mrope_delta, g->host_pos3 + (uint64_t)t * 4u); g->host_pos3[(uint64_t)t * 4u + 3u] = 0; if (t == 0 && g->snap_after_first) g->snap_mrope_delta = g->mrope_delta; if (t == 1u && g->snap_after_second) g->snap2_mrope_delta = g->mrope_delta; } if (!ds4_gpu_tensor_write(g->R, 0, row, (uint64_t)T * hc_dim * sizeof(float)) || !ds4_gpu_tensor_write(g->pos3, (uint64_t)g->pos * 16u, g->host_pos3, (uint64_t)T * 16u)) return false; uint32_t ids[256 * DS4_MAX_PLE_HEADS]; for (uint32_t t = 0; t < T; t++) { qwen4_ple_step(tokens[t], g->ple_prev, ids + (t % 256u) * DS4_N_PLE_HEADS); if (t == 0 && g->snap_after_first) { memcpy(g->snap_ple_prev, g->ple_prev, sizeof(g->snap_ple_prev)); g->snap_pos = g->pos + 1u; } if (t == 1u && g->snap_after_second) { memcpy(g->snap2_ple_prev, g->ple_prev, sizeof(g->snap2_ple_prev)); g->snap2_pos = g->pos + 2u; } if (t % 256u == 255u || t + 1 == T) { const uint32_t start = t / 256u * 256u; if (!qwen4_ngram_read(m, ids, (t-start+1u) * DS4_N_PLE_HEADS, row + (uint64_t)start * E)) { fprintf(stderr, "ds4: n-gram read failed: %s\n", strerror(errno)); return false; } } } return ds4_gpu_tensor_write(g->ple_emb, 0, row, (uint64_t)T * E * sizeof(float)) != 0; } /* Forward T tokens at g->pos..; logits (optional) receive the last token's * row, or all T rows with all_rows (T <= n_logit_rows). T <= cap_tokens; * everything is causal by construction because the recurrent kernels walk * tokens in order and attention reads the caches written for the same * chunk. g->R keeps the pre-mixer streams of every row afterwards. */ static bool qwen4_graph_forward_tokens(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, const int *tokens, uint32_t T, float *logits_out, bool all_rows) { if (!g || T == 0 || T > g->cap_tokens || g->pos + T > g->ctx_cap) return false; if (all_rows && T > g->n_logit_rows) return false; for (uint32_t t = 0; t < T; t++) { if (tokens[t] < 0 || tokens[t] >= (int)DS4_N_VOCAB) { fprintf(stderr, "ds4: Qwen3.8 token id %d is outside the vocabulary\n", tokens[t]); return false; } } const uint32_t pos0 = g->pos; const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; const uint32_t hc_dim = DS4_N_EMBD * DS4_N_HC; static int timing = -1; if (timing < 0) { const char *tv = getenv("DS4_QWEN4_TIMING"); timing = tv ? (atoi(tv) >= 2 ? 2 : 1) : 0; } /* Start the GPU while the host encodes the rest of the trunk. The * command queue preserves layer order; zero restores one submission. */ uint32_t flush_layer = timing != 2 && n_trunk > 2u ? 2u : 0u; const char *flush_env = getenv("DS4_QWEN4_FLUSH_LAYER"); if (timing != 2 && flush_env && flush_env[0]) { char *end = NULL; const unsigned long layer = strtoul(flush_env, &end, 10); if (end != flush_env && *end == '\0' && layer < n_trunk) { flush_layer = (uint32_t)layer; } } const double t0 = timing ? now_sec() : 0.0; if (!qwen4_graph_stage_inputs(g, m, w, tokens, T)) return false; const double t1 = timing ? now_sec() : 0.0; if (!glm_graph_begin_commands_if_needed()) return false; bool ok = true; /* DS4_QWEN4_TIMING=2 on prefill batches: sync after each stage group and * report GPU ms per group (adds sync overhead; diagnostics only) */ double prof[6] = {0}; const bool prof_on = timing == 2 && T > 1u; double prof_last = prof_on ? now_sec() : 0.0; #define QWEN4_PROF(idx_) do { \ if (prof_on) { \ ds4_gpu_end_commands(); \ const double now_ = now_sec(); \ prof[idx_] += now_ - prof_last; \ prof_last = now_; \ glm_graph_begin_commands_if_needed(); \ } \ } while (0) for (uint32_t il = 0; il < n_trunk && ok; il++) { const ds4_layer_weights *l = &w->layer[il]; if (ds4_qwen4_layer_is_ple(il)) { ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) && qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) && ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key, g->ple_val, m->map, m->size, l->ple_norm_key->abs_offset, l->ple_norm_query->abs_offset, l->ple_norm_conv->abs_offset, T, DS4_N_EMBD, DS4_N_HC, DS4_RMS_EPS) && ds4_gpu_qwen4_ple_conv_tensor(g->R, g->ple_gated, g->ple_normed, g->ple_hist, m->map, m->size, l->ple_conv->abs_offset, l->ple_conv->type, T, hc_dim, DS4_N_PLE_CONV, DS4_N_PLE_NGRAM, g->snap_after_first ? g->snap_ple_hist : NULL, 0u, g->snap_after_second ? g->snap2_ple_hist : NULL, 1u); } QWEN4_PROF(0); if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T); QWEN4_PROF(1); if (ok) { ok = ds4_qwen4_layer_is_linear(il) ? qwen4_graph_linear(g, m, l, il, T) : qwen4_graph_attention(g, m, l, il, pos0, T); } if (ok) { /* GDN and full attention both project into the same E-wide blk * stream before the hyper-connection residual combine. */ if (g->dump_prompt_rows) qwen4_graph_dump_last_attn(g, il, T); ok = qwen4_graph_apply_steering_attn(g, il, T); } QWEN4_PROF(ds4_qwen4_layer_is_linear(il) ? 2 : 3); if (T == 1u && !g->mtp_R && DS4_N_HC == 4u && ds4_gpu_qwen4_decode_fusions_enabled() && l->hc_ffn_inject->type == DS4_TENSOR_F16) { if (ok) ok = ds4_gpu_qwen4_hc_combine_norm_tensor(g->hc_u, g->blk, g->inj, g->xn, g->inj_alt, g->R, m->map, m->size, l->hc_ffn_norm->abs_offset, l->hc_ffn_inject->abs_offset, l->hc_ffn_inject->type, T, DS4_N_EMBD, DS4_N_HC, DS4_N_HC, DS4_RMS_EPS); if (ok) { /* hc_u is unused by the single-token fused mixer. Both residual * buffers have the same capacity, including for later prefill. */ ds4_gpu_tensor *swap = g->R; g->R = g->hc_u; g->hc_u = swap; swap = g->inj; g->inj = g->inj_alt; g->inj_alt = swap; ok = qwen4_gemv(g->lo, m, l->hc_ffn_down, g->xn, T) && ds4_gpu_qwen4_hc_gate_mix_tensor(g->mixed, g->xn, g->lo, m->map, m->size, l->hc_ffn_up->abs_offset, l->hc_ffn_up->type, T, DS4_N_EMBD, DS4_N_HC, DS4_N_HC_LOWRANK); } } else { if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T); } QWEN4_PROF(4); if (ok) ok = qwen4_graph_moe(g, m, l, T); /* the reduce folds the combine in */ if (ok && g->dump_prompt_rows) metal_graph_debug_dump_tensor("qwen_router", g->router, (uint64_t)T * DS4_N_EXPERT, il, pos0); if (ok && g->dump_prompt_rows) qwen4_graph_dump_last_ffn(g, il, T); if (ok) ok = qwen4_graph_apply_steering_ffn(g, il, T); QWEN4_PROF(5); /* Submit this prefix while the host encodes the remaining layers. * Flush keeps the same ordered queue and retains pending buffers; * end_commands below waits for both batches before inputs are reused. */ if (ok && il + 1u == flush_layer) ok = ds4_gpu_flush_commands() != 0; } if (prof_on) { fprintf(stderr, "ds4: Qwen3.8 prefill stage ms/chunk (pos=%u T=%u ok=%d): " "ple %.1f hc_attn %.1f gdn %.1f attn %.1f hc_ffn %.1f moe %.1f\n", pos0, T, ok ? 1 : 0, 1000.0 * prof[0], 1000.0 * prof[1], 1000.0 * prof[2], 1000.0 * prof[3], 1000.0 * prof[4], 1000.0 * prof[5]); } #undef QWEN4_PROF if (ok && logits_out && all_rows) { ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, T) && qwen4_gemv(g->logits, m, w->output, g->mixed, T); } else if (ok && logits_out) { /* final mixer on the last token only */ ds4_gpu_tensor *last = ds4_gpu_tensor_view(g->R, (uint64_t)(T - 1u) * hc_dim * sizeof(float), (uint64_t)hc_dim * sizeof(float)); ok = last != NULL; if (ok) { ds4_gpu_tensor *R_save = g->R; g->R = last; ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, 1); g->R = R_save; ds4_gpu_tensor_free(last); } if (ok) ok = qwen4_gemv(g->logits, m, w->output, g->mixed, 1); } const double t2 = timing ? now_sec() : 0.0; if (!ds4_gpu_end_commands()) ok = false; const double t3 = timing ? now_sec() : 0.0; if (ok && logits_out) { const uint64_t rows = all_rows ? T : 1u; ok = ds4_gpu_tensor_read(g->logits, 0, logits_out, rows * DS4_N_VOCAB * sizeof(float)) != 0; } if (timing) { static double acc[4]; static int n_calls; acc[0] += t1 - t0; acc[1] += t2 - t1; acc[2] += t3 - t2; acc[3] += now_sec() - t3; if (++n_calls % 50 == 0) { fprintf(stderr, "ds4: Qwen3.8 forward(T=%u) avg ms: stage %.3f encode %.3f gpu %.3f read %.3f\n", T, 1e3 * acc[0] / 50, 1e3 * acc[1] / 50, 1e3 * acc[2] / 50, 1e3 * acc[3] / 50); memset(acc, 0, sizeof(acc)); } } if (ok) g->pos += T; if (g->snap_after_first) { g->snap_valid = ok; g->snap_after_first = false; } if (g->snap_after_second) { g->snap2_valid = ok; g->snap_after_second = false; } return ok; } static bool qwen4_graph_forward_token(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, int token, float *logits_out) { return qwen4_graph_forward_tokens(g, m, w, &token, 1, logits_out, false); } /* Copy the recurrent state (GDN states and conv histories, PLE history and * n-gram context, position) between the live buffers and one snapshot set, * saving to it (save) or restoring from it. */ static bool qwen4_graph_state_copy_set(ds4_qwen4_gpu_graph *g, bool save, ds4_gpu_tensor **snap_state, ds4_gpu_tensor **snap_hist, ds4_gpu_tensor **snap_ple, int *snap_prev, uint32_t *snap_pos, int32_t *snap_mrope) { if (!*snap_ple) return false; const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; const uint64_t state_bytes = v_dim * DS4_N_LIN_HEAD_DIM * sizeof(float); const uint64_t hist_bytes = (uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim * sizeof(float); const uint64_t ple_bytes = (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * DS4_N_EMBD * DS4_N_HC * sizeof(float); if (!glm_graph_begin_commands_if_needed()) return false; bool ok = true; for (uint32_t il = 0; il < DS4_N_LAYER && ok; il++) { if (!snap_state[il]) continue; ds4_gpu_tensor *st = g->layer_lin_state[il], *ss = snap_state[il]; ds4_gpu_tensor *ht = g->layer_lin_hist[il], *hs = snap_hist[il]; ok = ds4_gpu_tensor_copy(save ? ss : st, 0, save ? st : ss, 0, state_bytes) != 0 && ds4_gpu_tensor_copy(save ? hs : ht, 0, save ? ht : hs, 0, hist_bytes) != 0; } if (ok) { ok = ds4_gpu_tensor_copy(save ? *snap_ple : g->ple_hist, 0, save ? g->ple_hist : *snap_ple, 0, ple_bytes) != 0; } if (!ds4_gpu_end_commands()) ok = false; if (save) { memcpy(snap_prev, g->ple_prev, sizeof(g->ple_prev)); *snap_pos = g->pos; *snap_mrope = g->mrope_delta; } else { memcpy(g->ple_prev, snap_prev, sizeof(g->ple_prev)); g->pos = *snap_pos; g->mrope_delta = *snap_mrope; } return ok; } static bool qwen4_graph_state_copy(ds4_qwen4_gpu_graph *g, bool save) { return qwen4_graph_state_copy_set(g, save, g->snap_lin_state, g->snap_lin_hist, &g->snap_ple_hist, g->snap_ple_prev, &g->snap_pos, &g->snap_mrope_delta); } static bool qwen4_graph_state_copy2(ds4_qwen4_gpu_graph *g, bool save) { return qwen4_graph_state_copy_set(g, save, g->snap2_lin_state, g->snap2_lin_hist, &g->snap2_ple_hist, g->snap2_ple_prev, &g->snap2_pos, &g->snap2_mrope_delta); } static bool qwen4_graph_state_copy0(ds4_qwen4_gpu_graph *g, bool save) { return qwen4_graph_state_copy_set(g, save, g->snap0_lin_state, g->snap0_lin_hist, &g->snap0_ple_hist, g->snap0_ple_prev, &g->snap0_pos, &g->snap0_mrope_delta); } /* A rejected draft returns to the state the verify snapshotted after row 0, * and that snapshot is a complete copy: swap the live and snapshot buffers * instead of copying 36 layers of recurrent state back through the GPU. * The stale buffers become the next snapshot targets, which every snapshot * kernel overwrites in full, so the snapshot is marked invalid until then. * DS4_QWEN4_MTP_SWAP_RESTORE=0 keeps the copying restore for comparisons. */ static bool qwen4_graph_state_swap(ds4_qwen4_gpu_graph *g) { if (!g->snap_ple_hist) return false; const char *env = getenv("DS4_QWEN4_MTP_SWAP_RESTORE"); if (env && env[0] && strcmp(env, "0") == 0) return qwen4_graph_state_copy(g, false); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!g->snap_lin_state[il]) continue; ds4_gpu_tensor *t = g->layer_lin_state[il]; g->layer_lin_state[il] = g->snap_lin_state[il]; g->snap_lin_state[il] = t; t = g->layer_lin_hist[il]; g->layer_lin_hist[il] = g->snap_lin_hist[il]; g->snap_lin_hist[il] = t; } ds4_gpu_tensor *t = g->ple_hist; g->ple_hist = g->snap_ple_hist; g->snap_ple_hist = t; memcpy(g->ple_prev, g->snap_ple_prev, sizeof(g->ple_prev)); g->pos = g->snap_pos; g->mrope_delta = g->snap_mrope_delta; g->snap_valid = false; return true; } /* Lazy verifier snapshots are all-or-nothing: a failed allocation must not * become a supposedly valid snapshot on the next attempt. */ static bool qwen4_graph_ensure_snapshot(ds4_qwen4_gpu_graph *g, ds4_gpu_tensor **state, ds4_gpu_tensor **hist, ds4_gpu_tensor **ple) { if (*ple) return true; if (!g->snap_ple_hist) return false; const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; const uint64_t state_n = v_dim * DS4_N_LIN_HEAD_DIM; const uint64_t hist_n = (uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim; *ple = qwen4_graph_alloc_f32((uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * (uint64_t)DS4_N_EMBD * DS4_N_HC); if (!*ple) return false; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!g->snap_lin_state[il]) continue; state[il] = qwen4_graph_alloc_f32(state_n); hist[il] = qwen4_graph_alloc_f32(hist_n); if (!state[il] || !hist[il]) { for (uint32_t j = 0; j <= il; j++) { ds4_gpu_tensor_free(state[j]); ds4_gpu_tensor_free(hist[j]); state[j] = hist[j] = NULL; } ds4_gpu_tensor_free(*ple); *ple = NULL; return false; } } return true; } /* Depth-2 sessions never pay for the second verifier snapshot. */ static bool qwen4_graph_ensure_snap2(ds4_qwen4_gpu_graph *g) { return qwen4_graph_ensure_snapshot(g, g->snap2_lin_state, g->snap2_lin_hist, &g->snap2_ple_hist); } /* Allocate the pre-verify snapshot set on the first exact-sampling cycle. * Also lazy: engines without --mtp-exact-sampling never pay for it, and a * failed allocation only means boundary resamples fall back to the old * reset-and-replay rewind. */ static bool qwen4_graph_ensure_snap0(ds4_qwen4_gpu_graph *g) { return qwen4_graph_ensure_snapshot(g, g->snap0_lin_state, g->snap0_lin_hist, &g->snap0_ple_hist); } /* Same swap-restore against the second snapshot (state after row 1) for a * 3-row verify that accepted its first draft but rejected the second. */ static bool qwen4_graph_state_swap2(ds4_qwen4_gpu_graph *g) { if (!g->snap2_ple_hist) return false; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!g->snap2_lin_state[il]) continue; ds4_gpu_tensor *t = g->layer_lin_state[il]; g->layer_lin_state[il] = g->snap2_lin_state[il]; g->snap2_lin_state[il] = t; t = g->layer_lin_hist[il]; g->layer_lin_hist[il] = g->snap2_lin_hist[il]; g->snap2_lin_hist[il] = t; } ds4_gpu_tensor *t = g->ple_hist; g->ple_hist = g->snap2_ple_hist; g->snap2_ple_hist = t; memcpy(g->ple_prev, g->snap2_ple_prev, sizeof(g->ple_prev)); g->pos = g->snap2_pos; g->mrope_delta = g->snap2_mrope_delta; g->snap2_valid = false; return true; } /* The steering bank contains trunk layers only. The predictor remains * unsteered; its drafts are verified by the steered target trunk. * One to three causal predictor steps at idx, using consecutive rows of the * trunk's pre-mixer streams and the embeddings of their following tokens. * The nextn layer runs on its own residual; want_logits returns the last * row only. Batching catches up accepted drafts and predicts again with * one command submission and one set of weight reads. */ static bool qwen4_graph_mtp_steps(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, uint32_t row, const int *next_tokens, uint32_t T, uint32_t idx, bool want_logits, float *logits_out, int *draft_out) { if (!g->mtp_R || T == 0u || T > 3u || idx > g->ctx_cap || T > g->ctx_cap - idx || row > g->cap_tokens || T > g->cap_tokens - row) return false; const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC; const uint32_t il = DS4_N_LAYER - 1u; const ds4_layer_weights *l = &w->layer[il]; for (uint32_t t = 0; t < T; t++) { if (next_tokens[t] < 0 || next_tokens[t] >= (int)DS4_N_VOCAB) return false; qwen4_ref_row(m, w->token_embd, (uint64_t)next_tokens[t], g->host_row + (uint64_t)t * E); } if (!ds4_gpu_tensor_write(g->mtp_e, 0, g->host_row, (uint64_t)T * E * sizeof(float)) || !glm_graph_begin_commands_if_needed()) return false; ds4_gpu_tensor *R_save = g->R; bool ok = true; const uint64_t emb_bytes = (uint64_t)E * sizeof(float); const uint64_t cat_bytes = (hc + 1u) * 2u * emb_bytes; const uint64_t proj_bytes = (hc + 1u) * emb_bytes; for (uint32_t t = 0; t < T && ok; t++) { ds4_gpu_tensor *e_row = ds4_gpu_tensor_view(g->mtp_e, t * emb_bytes, emb_bytes); ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(R_save, (uint64_t)(row + t) * hc * emb_bytes, hc * emb_bytes); ds4_gpu_tensor *cat_row = ds4_gpu_tensor_view(g->mtp_cat, t * cat_bytes, cat_bytes); ok = e_row && R_row && cat_row && ds4_gpu_qwen4_mtp_stage_tensor(cat_row, e_row, R_row, m->map, m->size, l->nextn_enorm->abs_offset, l->nextn_hnorm->abs_offset, E, hc, DS4_RMS_EPS); ds4_gpu_tensor_free(cat_row); ds4_gpu_tensor_free(R_row); ds4_gpu_tensor_free(e_row); } if (ok) ok = qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, T * (hc + 1u)); for (uint32_t t = 0; t < T && ok; t++) { ds4_gpu_tensor *proj_row = ds4_gpu_tensor_view(g->mtp_proj, t * proj_bytes, proj_bytes); ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, t * hc * emb_bytes, hc * emb_bytes); ok = proj_row && R_row && ds4_gpu_qwen4_mtp_combine_tensor(R_row, proj_row, E, hc); ds4_gpu_tensor_free(R_row); ds4_gpu_tensor_free(proj_row); } g->R = g->mtp_R; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T); if (ok) ok = qwen4_graph_attention(g, m, l, il, idx, T); if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, E, hc) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T); if (ok) ok = qwen4_graph_moe(g, m, l, T); ds4_gpu_tensor *last = NULL; const char *argmax_env = getenv("DS4_QWEN4_MTP_GPU_ARGMAX"); const bool gpu_argmax = want_logits && draft_out && !logits_out && (!argmax_env || strcmp(argmax_env, "0") != 0); /* draft-only rows: host logits consumers always see the full head */ const bool gathered = gpu_argmax && qwen4_mtp_draft_head_load(g, m, w->output); const uint32_t head_rows = gathered ? g->draft_rows : gpu_argmax ? qwen4_mtp_draft_rows() : DS4_N_VOCAB; if (ok && want_logits) { last = ds4_gpu_tensor_view(g->mtp_R, (T - 1u) * hc * emb_bytes, hc * emb_bytes); g->R = last; ok = last && qwen4_graph_hc_mix(g, m, l->nextn_hc_head_norm, l->nextn_hc_head_down, l->nextn_hc_head_up, NULL, 1) && (gathered ? ds4_gpu_qwen4_matmul_q8_0_weights_tensor(g->logits, g->draft_head, (uint32_t)w->output->dim[0], head_rows, g->mixed) != 0 : qwen4_gemv_rows(g->logits, m, w->output, g->mixed, 1, head_rows)); } g->R = R_save; if (ok && gpu_argmax) ok = ds4_gpu_qwen4_argmax_tensor(g->mtp_argmax, g->mtp_argmax_tmp, g->logits, head_rows) != 0; /* without logits nothing is read back: the batched predictor leaves the * command batch open and finishes every session's head at once */ if (want_logits && !ds4_gpu_end_commands()) ok = false; ds4_gpu_tensor_free(last); if (ok && gpu_argmax) { int32_t token = -1; ok = ds4_gpu_tensor_read(g->mtp_argmax, 0, &token, sizeof(token)) != 0; if (ok && gathered) { if (token < 0 || (uint32_t)token >= g->draft_rows) ok = false; else token = g->draft_ids[token]; } if (ok && (token < 0 || token >= (int32_t)DS4_N_VOCAB)) ok = false; if (ok) *draft_out = token; } else if (ok && want_logits) { float *dst = logits_out ? logits_out : g->host_logits; ok = ds4_gpu_tensor_read(g->logits, 0, dst, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; if (ok && draft_out) *draft_out = sample_argmax(dst, DS4_N_VOCAB); } if (ok) { g->mtp_pos = idx + T; g->mtp_last_rows = T; } return ok; } /* Recursive chain step for a second pending draft: condition the predictor * on its own output stream (the last row of the previous pass in mtp_R) * plus the embedding of the drafted token that follows, continuing the * nextn layer's raw cache at idx. The trunk has not run at these positions * yet, so the stream is the predictor's own recursion — a proposal-only * state: the verify gates everything the target commits. */ static bool qwen4_graph_mtp_chain_step(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, int next_token, uint32_t idx, int *draft_out) { if (!g->mtp_R || g->mtp_last_rows == 0u || g->mtp_last_rows > 3u || !draft_out || idx >= g->ctx_cap || next_token < 0 || next_token >= (int)DS4_N_VOCAB) return false; const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC; const uint32_t il = DS4_N_LAYER - 1u; const ds4_layer_weights *l = &w->layer[il]; const uint64_t emb_bytes = (uint64_t)E * sizeof(float); const uint64_t cat_bytes = (hc + 1u) * 2u * emb_bytes; const uint64_t proj_bytes = (hc + 1u) * emb_bytes; qwen4_ref_row(m, w->token_embd, (uint64_t)next_token, g->host_row); if (!ds4_gpu_tensor_write(g->mtp_e, 0, g->host_row, emb_bytes) || !glm_graph_begin_commands_if_needed()) return false; ds4_gpu_tensor *R_save = g->R; ds4_gpu_tensor *last = NULL; bool ok = true; { ds4_gpu_tensor *e_row = ds4_gpu_tensor_view(g->mtp_e, 0, emb_bytes); ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, (uint64_t)(g->mtp_last_rows - 1u) * hc * emb_bytes, hc * emb_bytes); ds4_gpu_tensor *cat_row = ds4_gpu_tensor_view(g->mtp_cat, 0, cat_bytes); ok = e_row && R_row && cat_row && ds4_gpu_qwen4_mtp_stage_tensor(cat_row, e_row, R_row, m->map, m->size, l->nextn_enorm->abs_offset, l->nextn_hnorm->abs_offset, E, hc, DS4_RMS_EPS); ds4_gpu_tensor_free(cat_row); ds4_gpu_tensor_free(R_row); ds4_gpu_tensor_free(e_row); } if (ok) ok = qwen4_gemv(g->mtp_proj, m, l->nextn_eh_proj, g->mtp_cat, hc + 1u); if (ok) { ds4_gpu_tensor *proj_row = ds4_gpu_tensor_view(g->mtp_proj, 0, proj_bytes); ds4_gpu_tensor *R_row = ds4_gpu_tensor_view(g->mtp_R, 0, hc * emb_bytes); ok = proj_row && R_row && ds4_gpu_qwen4_mtp_combine_tensor(R_row, proj_row, E, hc); ds4_gpu_tensor_free(R_row); ds4_gpu_tensor_free(proj_row); } g->R = g->mtp_R; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, 1); if (ok) ok = qwen4_graph_attention(g, m, l, il, idx, 1); if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, 1, E, hc) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, 1); if (ok) ok = qwen4_graph_moe(g, m, l, 1); /* the chained draft only needs its argmax; DS4_QWEN4_MTP_DRAFT_ROWS can * score the frequent-token prefix exactly like the first draft's head */ const uint32_t chain_head_rows = qwen4_mtp_draft_rows(); if (ok) { last = ds4_gpu_tensor_view(g->mtp_R, 0, hc * emb_bytes); g->R = last; ok = last && qwen4_graph_hc_mix(g, m, l->nextn_hc_head_norm, l->nextn_hc_head_down, l->nextn_hc_head_up, NULL, 1) && qwen4_gemv_rows(g->logits, m, w->output, g->mixed, 1, chain_head_rows); } g->R = R_save; if (ok) ok = ds4_gpu_qwen4_argmax_tensor(g->mtp_argmax, g->mtp_argmax_tmp, g->logits, chain_head_rows) != 0; if (!ds4_gpu_end_commands()) ok = false; ds4_gpu_tensor_free(last); if (ok) { int32_t token = -1; ok = ds4_gpu_tensor_read(g->mtp_argmax, 0, &token, sizeof(token)) != 0; if (ok && (token < 0 || (uint32_t)token >= chain_head_rows)) ok = false; if (ok) *draft_out = token; } if (ok) { g->mtp_pos = idx + 1u; g->mtp_last_rows = 1u; } return ok; } static bool qwen4_graph_mtp_step(ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, uint32_t row, int next_token, uint32_t idx, bool want_logits, float *logits_out, int *draft_out) { return qwen4_graph_mtp_steps(g, m, w, row, &next_token, 1u, idx, want_logits, logits_out, draft_out); } static int generate_qwen4_metal_argmax( const ds4_model * model, const ds4_vocab * vocab, const ds4_weights * weights, const token_vec * prompt, int n_predict, int ctx_size, 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) { if (!prompt || prompt->len <= 0 || prompt->len >= ctx_size) { fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); return 1; } ds4_qwen4_gpu_graph *g = xcalloc(1, sizeof(*g)); if (!qwen4_graph_alloc(g, weights, (uint32_t)ctx_size, qwen4_prefill_chunk_tokens((uint32_t)ctx_size), false, NULL, NULL, NULL, 0)) { free(g); return 1; } if (!qwen4_graph_load_steering(g, directional_steering_file, directional_steering_attn, directional_steering_ffn)) { qwen4_graph_free(g); free(g); return 1; } qwen4_graph_reset(g); float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); const double t_prefill0 = now_sec(); bool ok = true; for (int i = 0; i < prompt->len && ok;) { uint32_t chunk = (uint32_t)(prompt->len - i); if (chunk > g->cap_tokens) chunk = g->cap_tokens; ok = qwen4_graph_forward_tokens(g, model, weights, prompt->v + i, chunk, i + (int)chunk == prompt->len ? logits : NULL, false); i += (int)chunk; if (progress) progress(progress_ud, "prefill_chunk", i, prompt->len); } const double t_prefill1 = now_sec(); if (!ok) { fprintf(stderr, "ds4: Qwen3.8 prefill failed\n"); free(logits); qwen4_graph_free(g); free(g); return 1; } int n_generated = 0; const double t_decode0 = now_sec(); for (int i = 0; i < n_predict && g->pos < g->ctx_cap; i++) { 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 || g->pos + 1u >= g->ctx_cap) break; if (!qwen4_graph_forward_token(g, model, weights, token, logits)) { fprintf(stderr, "ds4: Qwen3.8 decode failed at position %u\n", g->pos); free(logits); qwen4_graph_free(g); free(g); return 1; } } 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: Qwen3.8 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); qwen4_graph_free(g); free(g); return 0; } #endif #ifndef DS4_NO_GPU /* 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; } #ifdef DS4_HAS_QWEN4_GPU if (ds4_model_is_qwen4()) { return generate_qwen4_metal_argmax(model, vocab, weights, prompt, n_predict, ctx_size, directional_steering_file, directional_steering_attn, directional_steering_ffn, emit, done, emit_ud, progress, progress_ud); } #endif 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; } 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, directional_steering_file, directional_steering_attn, directional_steering_ffn, 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; } static uint32_t engine_placement_session_count(const ds4_engine *e) { return e && e->placement_session_count_hint > 1 ? (uint32_t)e->placement_session_count_hint : 1u; } static size_t engine_size_mul_sat(size_t bytes, uint32_t count) { if (count != 0u && bytes > SIZE_MAX / count) return SIZE_MAX; return bytes * count; } /* 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) { if (ds4_model_is_qwen4()) { /* Qwen currently uses one device and independent session workspaces. * Reserve its whole graph here, not DeepSeek's per-layer KV layout. */ const int ctx = e && e->placement_ctx_hint > 0 ? e->placement_ctx_hint : 4096; const ds4_context_memory mem = ds4_context_memory_estimate_with_prefill( DS4_BACKEND_CUDA, ctx, e ? e->prefill_chunk : 0); const size_t bytes = mem.total_bytes > SIZE_MAX ? SIZE_MAX : (size_t)mem.total_bytes; return engine_size_mul_sat(bytes, engine_placement_session_count(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); const size_t scratch = mem.scratch_bytes > SIZE_MAX ? SIZE_MAX : (size_t)mem.scratch_bytes; return engine_size_mul_sat(scratch, engine_placement_session_count(e)); } #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 */ total += 2ull * DS4_N_INDEXER_HEAD_DIM * sizeof(float); /* index_comp_kv_cur_by_tier */ total += 2ull * DS4_N_INDEXER_HEAD_DIM * sizeof(float); /* index_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_LOW || mode == DS4_THINK_MEDIUM || mode == DS4_THINK_HIGH || mode == DS4_THINK_MAX || ds4_think_mode_level(mode) > 0; } int ds4_think_mode_level(ds4_think_mode mode) { if ((unsigned)mode < DS4_THINK_LEVEL_BASE || (unsigned)mode > DS4_THINK_LEVEL_BASE + 100u) return -1; return (int)mode - DS4_THINK_LEVEL_BASE; } bool ds4_think_mode_parse_level(const char *text, ds4_think_mode *out) { if (!text || !out || !text[0]) return false; unsigned level = 0; for (const char *p = text; *p; p++) { if (*p < '0' || *p > '9' || level > 100) return false; level = level * 10u + (unsigned)(*p - '0'); } if (level > 100) return false; *out = (ds4_think_mode)(DS4_THINK_LEVEL_BASE + level); return true; } const char *ds4_think_mode_name(ds4_think_mode mode) { const int level = ds4_think_mode_level(mode); if (level >= 0 && level <= 100) { static __thread char name[4]; snprintf(name, sizeof(name), "%d", level); return name; } switch (mode) { case DS4_THINK_NONE: return "none"; case DS4_THINK_HIGH: return "high"; case DS4_THINK_MAX: return "max"; case DS4_THINK_LOW: return "low"; case DS4_THINK_MEDIUM: return "medium"; } 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 (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41 && 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 seed_batches; uint64_t proposed_tokens; uint64_t accepted_draft_tokens; uint64_t full_accepts; uint64_t partial_accepts; uint64_t direct_full_commits; uint64_t direct_partial_commits; uint64_t replay_fallbacks; 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 typedef struct { uint32_t token_start; uint32_t token_count; uint8_t fingerprint[32]; } ds4_vision_identity; struct ds4_session { ds4_engine *engine; ds4_dist_session *distributed; uint64_t tp_session_id; uint64_t glm_reserved_graph_bytes; #ifndef DS4_NO_GPU #ifdef DS4_HAS_DEEPSEEK41_GPU ds41_gpu_graph ds41_graph; bool ds41_graph_ready; #endif ds4_gpu_graph graph; ds4_glm_gpu_graph glm_graph; bool glm_graph_ready; #ifdef DS4_HAS_QWEN4_GPU ds4_qwen4_gpu_graph qwen4_graph; bool qwen4_graph_ready; int qwen4_slot; /* -1 when the session holds private recurrent state */ bool qwen4_rewound; float *qwen4_verify_logits; uint64_t qwen4_spec_cycles; uint64_t qwen4_spec_accepted; #endif uint32_t glm_dense_cache_len; /* GLM MTP speculative state. parent is the token that conditioned the * pending point-mass draft; sampled decoding discards a draft when its * boundary token differs from that parent. */ int glm_mtp_draft; int glm_mtp_draft2; /* second pending draft (depth-3 cycles) */ bool glm_mtp_have2; int glm_mtp_parent; int glm_mtp_have; /* rolling first-draft acceptance window for the adaptive draft depth */ uint32_t qwen4_depth_window; uint32_t qwen4_depth_accepted; uint32_t qwen4_depth_cycles; uint32_t qwen4_reject2_streak; /* consecutive 3-row cycles whose second draft rejected */ bool qwen4_depth3_engaged; bool glm_mtp_rollback_valid; uint32_t glm_mtp_rollback_pos; uint32_t glm_mtp_rollback_dense_len; int glm_mtp_rollback_first_token; 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; ds4_vision_identity *checkpoint_images; size_t checkpoint_image_count; const ds4_vision_span *sync_images; size_t sync_image_count; 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; float dspark_sample_temperature; uint64_t *dspark_sample_rng; bool dspark_draft_valid; bool dspark_stochastic_draft; bool dspark_sched_skipped_cycle; bool dspark_sched_long_accept_seen; bool dspark_sched_bypass; 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; }; static bool ds4_session_tp_leader(const ds4_session *s); #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_session_dspark_rocm_gfx1151_fast_path( const ds4_session *s) { return s && s->engine && s->engine->backend == DS4_BACKEND_CUDA && ds4_dspark_rocm_gfx1151_fast_path(); } static bool ds4_session_dspark_rocm_gfx1151_reference_alignment( const ds4_session *s) { return s && s->engine && s->engine->backend == DS4_BACKEND_CUDA && ds4_dspark_rocm_gfx1151_reference_alignment(); } static bool ds4_session_dspark_seed_batch_enabled( const ds4_session *s) { const char *env = getenv("DS4_DSPARK_SEED_BATCH"); if (env && env[0]) return env[0] != '0'; if (ds4_session_dspark_rocm_gfx1151_fast_path(s)) return true; if (!s || !s->engine) return false; const ds4_engine *e = s->engine; const ds4_layer_weights *layer = &e->weights.layer[DS4_N_LEADING_DENSE]; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (e->backend == DS4_BACKEND_CUDA && ds4_gpu_device_is_spark()) { return !e->ssd_streaming && !e->dspark_exact_sampling && layer->ffn_gate_exps && layer->ffn_down_exps && layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && layer->ffn_down_exps->type == DS4_TENSOR_Q2_K; } #endif /* Keep streaming and unmeasured weight/device combinations on their * existing schedule. No sampling decision is changed by this dispatch. */ return e->backend == DS4_BACKEND_METAL && !e->ssd_streaming && !e->dspark_exact_sampling && ds4_gpu_device_is_m5_apple_silicon() && layer->ffn_gate_exps && layer->ffn_down_exps && ((layer->ffn_gate_exps->type == DS4_TENSOR_IQ2_XXS && layer->ffn_down_exps->type == DS4_TENSOR_Q2_K) || (e->tp.active && layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && layer->ffn_down_exps->type == DS4_TENSOR_MXFP4)); } static bool ds4_session_dspark_seed_batch_short_fallback(const ds4_session *s) { if (s->engine->backend == DS4_BACKEND_METAL) return true; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) return s->engine->backend == DS4_BACKEND_CUDA && ds4_gpu_device_is_spark(); #else return false; #endif } static bool ds4_dspark_scheduler_enabled(const ds4_session *s) { const char *env = getenv("DS4_DSPARK_SCHEDULER"); if (env && env[0]) return strcmp(env, "0") != 0; (void)s; return true; } static uint32_t ds4_dspark_scheduler_window(const ds4_session *s) { const uint32_t fallback = ds4_session_dspark_rocm_gfx1151_fast_path(s) ? 16u : 4u; uint32_t v = ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_WINDOW", fallback); return v ? v : 4; } static uint32_t ds4_dspark_scheduler_skip_cycles(const ds4_session *s) { uint32_t fallback = s && ds4_session_dspark_seed_batch_short_fallback(s) && ds4_session_dspark_seed_batch_enabled(s) ? 32u : ds4_dspark_rocm_gfx1151_fast_path() ? 4u : 2u; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) /* Recheck sooner after useful long drafts; low-acceptance prose keeps * the longer pause instead of paying for another unproductive proposal. */ if (s && s->dspark_sched_long_accept_seen && ds4_session_dspark_seed_batch_enabled(s) && ds4_gpu_device_is_spark()) fallback = 8u; #endif return ds4_dspark_env_u32("DS4_DSPARK_SCHEDULER_SKIP", fallback); } 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) { const uint32_t fallback = ds4_dspark_rocm_gfx1151_fast_path() ? 4000u : 1500u; return ds4_dspark_env_u32( "DS4_DSPARK_SCHEDULER_MIN_AVG_MILLI", fallback); } 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", 0); } 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", 0); } 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; } /* Timing-sensitive scheduling changes which arithmetic path advances a token. * Keep it opt-in so greedy DSpark output is reproducible across runs. */ static bool ds4_dspark_scheduler_timing_enabled(void) { return ds4_dspark_scheduler_max_ms_per_accept_milli() != 0 || ds4_dspark_scheduler_max_extra_saved_ratio_milli() != 0; } 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 void ds4_session_dspark_scheduler_begin_request(ds4_session *s) { if (!s) return; ds4_session_dspark_scheduler_reset(s); s->dspark_sched_skip = 0; s->dspark_sched_lifetime_accepted = 0; s->dspark_sched_life_extra_ms = 0.0; s->dspark_sched_life_saved_ms = 0.0; s->dspark_sched_skipped_cycle = false; s->dspark_sched_long_accept_seen = false; s->dspark_sched_bypass = false; } static bool ds4_session_dspark_scheduler_should_skip(ds4_session *s) { if (!s || !ds4_dspark_scheduler_enabled(s)) return false; s->dspark_sched_skipped_cycle = false; if (s->dspark_sched_bypass) { s->dspark_sched_skipped_cycle = true; s->dspark_stats.scheduler_skips++; return true; } 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(s)) 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(s); 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) { if (ds4_session_dspark_rocm_gfx1151_fast_path(s)) { s->dspark_sched_bypass = true; s->dspark_sched_skip = 0; if (getenv("DS4_DSPARK_SPEC_LOG") != NULL) { fprintf(stderr, "ds4: DSpark scheduler bypass accepted=%u avg=%.3f " "no_draft=%u\n", s->dspark_sched_accepted, (double)avg_milli / 1000.0, s->dspark_sched_no_draft); } ds4_session_dspark_scheduler_reset(s); return; } s->dspark_sched_skip = ds4_dspark_scheduler_skip_cycles(s); 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_read_tokens_for_rebuild(FILE *fp, uint32_t count, uint64_t remaining, ds4_tokens *tokens, char *err, size_t errlen) { if (count == 0 || (uint64_t)count * sizeof(uint32_t) > remaining) { payload_set_err(err, errlen, "invalid TP checkpoint token history"); return 1; } for (uint32_t i = 0; i < count; i++) { uint32_t token = 0; if (payload_read_u32(fp, &token, &remaining, err, errlen) != 0 || token >= DS4_N_VOCAB) { payload_set_err(err, errlen, "invalid TP checkpoint token"); return 1; } ds4_tokens_push(tokens, (int)token); } uint8_t buf[65536]; return payload_skip_bytes(fp, remaining, buf, sizeof(buf), &remaining, err, errlen); } 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 uint32_t session_glm_index_live_rows(uint32_t layer, uint32_t compact_live) { if (compact_live == 0 || !glm_graph_layer_uses_full_indexer(layer)) return 0; if (ds4_model_is_glm53()) return compact_live / DS4_GLM53_INDEX_POOL_SIZE; return compact_live; } static uint64_t session_glm53_indexer_tail_bytes(void) { return 2u * DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } static uint64_t session_glm_kda_conv_state_bytes(void) { const uint64_t projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; return 3u * (DS4_N_KDA_CONV - 1u) * projection * sizeof(float); } static uint64_t session_glm_kda_recurrent_state_bytes(void) { const uint64_t projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; return projection * DS4_N_KDA_HEAD_DIM * sizeof(float); } 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++) { if (g->glm53 && ds4_glm53_layer_is_kda(il)) { bytes += session_glm_kda_conv_state_bytes(); bytes += session_glm_kda_recurrent_state_bytes(); continue; } 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)) { const uint32_t index_live = session_glm_index_live_rows(il, compact_live); bytes += (uint64_t)index_live * DS4_N_INDEXER_HEAD_DIM * sizeof(float); if (g->glm53) bytes += session_glm53_indexer_tail_bytes(); } } 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; if (ds4_model_is_glm53() && ds4_glm53_layer_is_kda(layer)) { if (compact_live != 0 || index_live != 0) return false; *out = session_glm_kda_conv_state_bytes() + session_glm_kda_recurrent_state_bytes(); return true; } const bool has_indexer = glm_graph_layer_uses_full_indexer(layer); const uint32_t expected_index_live = session_glm_index_live_rows(layer, compact_live); 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; if (has_indexer && ds4_model_is_glm53() && !payload_u64_add(&bytes, session_glm53_indexer_tail_bytes())) 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 bool ds4_session_is_ds41(const ds4_session *s) { return s && s->engine && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41; } 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; } static bool ds4_session_is_qwen4(const ds4_session *s) { return s && s->engine && ds4_model_is_qwen4(); } #ifndef DS4_NO_GPU static void ds4_session_glm_reset_dense_cache(ds4_session *s) { if (!s) return; s->glm_dense_cache_len = 0; s->glm_mtp_have = 0; s->glm_mtp_have2 = false; s->glm_mtp_rollback_valid = false; s->glm_mtp_min_pos = 0; } static bool ds4_session_glm_reset_kda_state(ds4_session *s) { return !s || !s->glm_graph_ready || glm_graph_reset_kda_state(&s->glm_graph); } 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 = session_glm_index_live_rows(il, compact_live); 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 = session_glm_index_live_rows(il, compact_live); 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++) { if (g->glm53 && ds4_glm53_layer_is_kda(il)) { rc = payload_write_tensor_span( fp, g->layer_kda_conv_state[il], 0, session_glm_kda_conv_state_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_tensor_span( fp, g->layer_kda_recurrent_state[il], 0, session_glm_kda_recurrent_state_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); } continue; } 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)) { const uint32_t index_rows = session_glm_index_live_rows(il, compact_live); rc = payload_write_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)index_rows * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0 && g->glm53) { const uint64_t tail_bytes = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); rc = payload_write_tensor_span(fp, g->layer_indexer_tail_k[il], 0, tail_bytes, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_tensor_span( fp, g->layer_indexer_tail_gate[il], 0, tail_bytes, 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 = session_glm_index_live_rows(il, expected_compact_live); 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; if (g->glm53 && ds4_glm53_layer_is_kda(il)) { rc = payload_read_tensor_span( fp, g->layer_kda_conv_state[il], 0, session_glm_kda_conv_state_bytes(), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_tensor_span( fp, g->layer_kda_recurrent_state[il], 0, session_glm_kda_recurrent_state_bytes(), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } continue; } 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); if (rc == 0 && g->glm53) { const uint64_t tail_bytes = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); rc = payload_read_tensor_span( fp, g->layer_indexer_tail_k[il], 0, tail_bytes, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_tensor_span( fp, g->layer_indexer_tail_gate[il], 0, tail_bytes, 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 || gate->type == DS4_TENSOR_MXFP4) ? 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; } bool ds4_engine_mtp_exact_sampling(ds4_engine *e) { return e && e->dspark_exact_sampling; } int ds4_engine_mtp_draft_tokens(ds4_engine *e) { if (e && ds4_model_is_qwen4()) { return e->glm_mtp && DS4_N_NEXTN_PREDICT != 0 && e->backend != DS4_BACKEND_CPU ? 2 : 0; } 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 an intermediate state captured by a tiny speculative verifier. * * Append-only cache rows beyond the accepted prefix can remain as invisible * garbage. Only compressor frontiers and row counters need rewinding. */ static bool spec_frontier_commit_prefix(ds4_session *s, uint32_t prefix_len) { if (!s || prefix_len == 0 || prefix_len > DS4_SPEC_PREFIX_SLOTS) { return false; } const uint32_t slot = prefix_len - 1u; 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_prefix_n_comp[slot][il]; const uint64_t ab = ds4_gpu_tensor_bytes(g->layer_attn_state_kv[il]); const uint64_t aoff = (uint64_t)slot * ab; ok = ds4_gpu_tensor_copy(g->layer_attn_state_kv[il], 0, g->spec_prefix1_attn_state_kv[il], aoff, ab) != 0 && ds4_gpu_tensor_copy(g->layer_attn_state_score[il], 0, g->spec_prefix1_attn_state_score[il], aoff, ab) != 0; if (ok && ratio == 4) { g->layer_n_index_comp[il] = g->spec_prefix_n_index_comp[slot][il]; const uint64_t ib = ds4_gpu_tensor_bytes(g->layer_index_state_kv[il]); const uint64_t ioff = (uint64_t)slot * ib; ok = ds4_gpu_tensor_copy(g->layer_index_state_kv[il], 0, g->spec_prefix1_index_state_kv[il], ioff, ib) != 0 && ds4_gpu_tensor_copy(g->layer_index_state_score[il], 0, g->spec_prefix1_index_state_score[il], ioff, ib) != 0; } } if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); return ok; } static bool spec_frontier_commit_prefix1(ds4_session *s) { return spec_frontier_commit_prefix(s, 1); } 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 #ifdef DS4_HAS_DEEPSEEK41_GPU typedef struct { ds4_gpu_tensor *tensor; uint64_t bytes; } ds41_state_span; /* The selection masks are rebuilt by each index-source layer on the next * token. Only owner caches, windows and unfinished compression pairs persist. */ static uint32_t ds41_state_spans(ds41_gpu_graph *g, uint32_t pos, ds41_state_span spans[54]) { uint32_t n = 0; const uint32_t raw = pos < 128u ? pos : 128u; for (uint32_t il = 0; il < 40; il++) spans[n++] = (ds41_state_span){g->window[il], (uint64_t)raw * 512u * 4u}; for (uint32_t i = 0; i < 4; i++) { const uint32_t live = pos / (i < 3 ? 2u : 1u); spans[n++] = (ds41_state_span){g->compressed[i], (uint64_t)live * 512u * 4u}; spans[n++] = (ds41_state_span){g->index_cache[i], (uint64_t)live * 128u * 4u}; if (i < 3 && (pos & 1u)) { spans[n++] = (ds41_state_span){g->previous_kv[i], 512u * 4u}; spans[n++] = (ds41_state_span){g->previous_score[i], 512u * 4u}; } } return n; } static uint64_t ds41_payload_body_bytes(ds41_gpu_graph *g, uint32_t pos) { ds41_state_span spans[54]; const uint32_t n = ds41_state_spans(g, pos, spans); uint64_t bytes = ((uint64_t)pos + DS4_N_VOCAB) * sizeof(float); for (uint32_t i = 0; i < n; i++) bytes += spans[i].bytes; return bytes; } static int ds41_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) { ds41_gpu_graph *g = &s->ds41_graph; if (!s->ds41_graph_ready || !g->valid || g->pos != (uint32_t)s->checkpoint.len || !ds4_gpu_synchronize()) { payload_set_err(err, errlen, "V4.1 graph has no complete frontier to save"); return 1; } const uint32_t h[DS4_SESSION_PAYLOAD_U32_FIELDS] = { DS4_SESSION_PAYLOAD_MAGIC, DS4_SESSION_PAYLOAD_VERSION, g->ctx, 1, 128, 128, g->ctx + 1u, g->pos, 40, 512, 128, DS4_N_VOCAB, 0x413431u }; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) if (payload_write_u32(fp, h[i], err, errlen)) return 1; for (int i = 0; i < s->checkpoint.len; i++) if (payload_write_u32(fp, (uint32_t)s->checkpoint.v[i], err, errlen)) return 1; if (payload_write_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * 4u, err, errlen)) return 1; ds41_state_span spans[54]; const uint32_t n = ds41_state_spans(g, g->pos, spans); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; for (uint32_t i = 0; i < n && !rc; i++) rc = payload_write_tensor_span(fp, spans[i].tensor, 0, spans[i].bytes, buf, DS4_SESSION_IO_CHUNK, err, errlen); free(buf); return rc; } static int ds41_load_payload(ds4_session *s, FILE *fp, const uint32_t *h, uint64_t remaining, char *err, size_t errlen) { ds41_gpu_graph *g = &s->ds41_graph; const uint32_t pos = h[7]; if (!s->ds41_graph_ready || !pos || pos >= g->ctx || pos >= h[2] || h[2] > 1048576u || h[3] != 1 || h[4] != 128 || h[5] != 128 || h[6] != h[2] + 1u || h[8] != 40 || h[9] != 512 || h[10] != 128 || h[11] != DS4_N_VOCAB || h[12] != 0x413431u || remaining != ds41_payload_body_bytes(g, pos)) { payload_set_err(err, errlen, "invalid V4.1 snapshot dimensions or size"); return 1; } ds4_tokens tokens = {0}; int rc = 0; for (uint32_t i = 0; i < pos; i++) { uint32_t token; if (payload_read_u32(fp, &token, &remaining, err, errlen) || token >= DS4_N_VOCAB) { payload_set_err(err, errlen, "invalid V4.1 snapshot token"); rc = 1; break; } token_vec_push(&tokens, (int)token); } if (!rc && !ds4_gpu_synchronize()) rc = 1; if (!rc) { s->checkpoint_valid = false; g->valid = false; rc = payload_read_bytes(fp, s->logits, (uint64_t)DS4_N_VOCAB * 4u, &remaining, err, errlen); ds41_state_span spans[54]; const uint32_t n = ds41_state_spans(g, pos, spans); uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); for (uint32_t i = 0; i < n && !rc; i++) rc = payload_read_tensor_span(fp, spans[i].tensor, 0, spans[i].bytes, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); free(buf); if (!rc) { ds41_graph_reset(g); for (uint32_t i = 0; i < 3 && i < pos; i++) g->history.tail[i] = (int32_t)g->token_map[tokens.v[pos - 1u - i]]; g->pos = pos; ds4_tokens_copy(&s->checkpoint, &tokens); s->checkpoint_valid = true; } } ds4_tokens_free(&tokens); return rc; } #endif #ifdef DS4_HAS_QWEN4_GPU static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows); #endif uint64_t ds4_session_payload_bytes(ds4_session *s) { if (s && !s->distributed && ds4_session_is_qwen4(s)) { #ifndef DS4_HAS_QWEN4_GPU return 0; #else if (!s->qwen4_graph_ready || !s->checkpoint_valid) return 0; if (s->qwen4_rewound || s->qwen4_graph.pos != (uint32_t)s->checkpoint.len) return 0; 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); const uint32_t rows = (uint32_t)s->checkpoint.len; const uint32_t mtp_rows = s->qwen4_graph.mtp_pos < rows ? s->qwen4_graph.mtp_pos : rows; bytes += qwen4_payload_tensor_bytes(rows, mtp_rows); return bytes; #endif } if (!s || !s->checkpoint_valid) return 0; if (s->distributed) return 0; #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) { if (!s->ds41_graph_ready || !s->ds41_graph.valid || s->ds41_graph.pos != (uint32_t)s->checkpoint.len) return 0; return DS4_SESSION_PAYLOAD_U32_FIELDS * sizeof(uint32_t) + ds41_payload_body_bytes(&s->ds41_graph, (uint32_t)s->checkpoint.len); } #endif 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; } #ifdef DS4_HAS_QWEN4_GPU /* Qwen3.8 session payload: header, tokens, logits, then per layer either the * GDN state + conv history (linear layers) or the live KV/indexer rows and * pooled block keys (attention layers, incl. the MTP block), then the PLE * conv history and n-gram context. The header's raw_live field carries the * family tag so a DeepSeek payload is never read as a Qwen one. */ #define DS4_QWEN4_PAYLOAD_TAG 0x51573802u static uint64_t qwen4_payload_lin_state_bytes(void) { return (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM * DS4_N_LIN_HEAD_DIM * sizeof(float); } static uint64_t qwen4_payload_lin_hist_bytes(void) { const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; return (uint64_t)(DS4_N_LIN_CONV - 1u) * conv_dim * sizeof(float); } static uint64_t qwen4_payload_kv_bytes(uint32_t rows) { return (uint64_t)rows * DS4_N_HEAD_KV * DS4_N_HEAD_DIM * 2u; } static uint64_t qwen4_payload_ik_bytes(uint32_t rows) { return (uint64_t)rows * DS4_N_INDEXER_HEAD_DIM * sizeof(float); } static uint64_t qwen4_payload_block_key_bytes(uint32_t rows) { return (uint64_t)(rows / 4u) * DS4_N_INDEXER_HEAD_DIM * 2u; } static uint64_t qwen4_payload_ple_hist_bytes(void) { return (uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * DS4_N_EMBD * DS4_N_HC * sizeof(float); } static uint64_t qwen4_payload_tensor_bytes(uint32_t rows, uint32_t mtp_rows) { uint64_t bytes = sizeof(uint32_t); /* mtp_rows */ for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (ds4_qwen4_layer_is_linear(il)) { bytes += qwen4_payload_lin_state_bytes() + qwen4_payload_lin_hist_bytes(); } else { const uint32_t n = ds4_qwen4_layer_is_nextn(il) ? mtp_rows : rows; bytes += 2u * qwen4_payload_kv_bytes(n) + qwen4_payload_ik_bytes(n) + qwen4_payload_block_key_bytes(n); } } bytes += qwen4_payload_ple_hist_bytes(); bytes += (uint64_t)DS4_MAX_PLE_NGRAM * sizeof(uint32_t); bytes += sizeof(uint32_t) + (uint64_t)rows * 16u; /* mrope delta + positions */ return bytes; } static int qwen4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen) { if (!s->qwen4_graph_ready || s->qwen4_rewound || s->qwen4_graph.pos != (uint32_t)s->checkpoint.len) { payload_set_err(err, errlen, "Qwen3.8 snapshot requires sync or eval after rewind"); return 1; } if (ds4_gpu_synchronize() == 0) { payload_set_err(err, errlen, "failed to synchronize accelerator before Qwen3.8 snapshot"); return 1; } ds4_qwen4_gpu_graph *g = &s->qwen4_graph; const uint32_t rows = (uint32_t)s->checkpoint.len; if (g->pos != rows) { payload_set_err(err, errlen, "Qwen3.8 graph position does not match the checkpoint"); return 1; } 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, (uint32_t)(DS4_N_EMBD * DS4_N_HC), rows, DS4_N_LAYER, DS4_N_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM, DS4_N_VOCAB, DS4_QWEN4_PAYLOAD_TAG, }; 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; uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); int rc = 0; const uint32_t mtp_rows = g->mtp_pos < rows ? g->mtp_pos : rows; if (rc == 0) rc = payload_write_u32(fp, mtp_rows, err, errlen); for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) { if (ds4_qwen4_layer_is_linear(il)) { rc = payload_write_tensor_span(fp, g->layer_lin_state[il], 0, qwen4_payload_lin_state_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_tensor_span(fp, g->layer_lin_hist[il], 0, qwen4_payload_lin_hist_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); } } else { const uint32_t n = ds4_qwen4_layer_is_nextn(il) ? mtp_rows : rows; rc = payload_write_tensor_span(fp, g->layer_k_cache[il], 0, qwen4_payload_kv_bytes(n), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_v_cache[il], 0, qwen4_payload_kv_bytes(n), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_ik_cache[il], 0, qwen4_payload_ik_bytes(n), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) rc = payload_write_tensor_span(fp, g->layer_block_key[il], 0, qwen4_payload_block_key_bytes(n), buf, DS4_SESSION_IO_CHUNK, err, errlen); } } if (rc == 0) { rc = payload_write_tensor_span(fp, g->ple_hist, 0, qwen4_payload_ple_hist_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); } for (uint32_t i = 0; rc == 0 && i < DS4_MAX_PLE_NGRAM; i++) { rc = payload_write_u32(fp, (uint32_t)g->ple_prev[i], err, errlen); } if (rc == 0) rc = payload_write_u32(fp, (uint32_t)g->mrope_delta, err, errlen); if (rc == 0) { rc = payload_write_tensor_span(fp, g->pos3, 0, (uint64_t)rows * 16u, buf, DS4_SESSION_IO_CHUNK, err, errlen); } free(buf); return rc; } static int qwen4_session_load_payload(ds4_session *s, FILE *fp, const uint32_t *h, uint64_t *remaining, char *err, size_t errlen) { if (!s->qwen4_graph_ready) { payload_set_err(err, errlen, "Qwen3.8 graph is not ready for restore"); return 1; } ds4_qwen4_gpu_graph *g = &s->qwen4_graph; const uint32_t rows = h[7]; if (h[12] != DS4_QWEN4_PAYLOAD_TAG || h[8] != DS4_N_LAYER || h[9] != DS4_N_HEAD_DIM || h[10] != DS4_N_INDEXER_HEAD_DIM || h[11] != DS4_N_VOCAB || h[6] != DS4_N_EMBD * DS4_N_HC) { payload_set_err(err, errlen, "KV checkpoint was written by a different model family or shape"); return 1; } if (rows > g->ctx_cap || rows > (uint32_t)s->ctx_size) { payload_set_err(err, errlen, "KV checkpoint is longer than this session's context"); return 1; } token_vec new_checkpoint = {0}; for (uint32_t i = 0; i < rows; i++) { uint32_t tok; if (payload_read_u32(fp, &tok, remaining, err, errlen) != 0) { token_vec_free(&new_checkpoint); return 1; } if (tok >= DS4_N_VOCAB) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "KV checkpoint token id is outside the vocabulary"); return 1; } token_vec_push(&new_checkpoint, (int)tok); } /* From the first write onward, failure must leave no reusable checkpoint. * Verifier snapshots belong to the old transcript, even at the same row. */ s->checkpoint_valid = false; s->mtp_draft_valid = false; s->glm_mtp_have = 0; s->glm_mtp_have2 = false; g->snap_valid = g->snap2_valid = g->snap0_valid = false; 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; } if (ds4_gpu_synchronize() == 0) { token_vec_free(&new_checkpoint); payload_set_err(err, errlen, "failed to synchronize accelerator before Qwen3.8 restore"); return 1; } uint32_t mtp_rows = 0; int rc = payload_read_u32(fp, &mtp_rows, remaining, err, errlen); if (rc == 0 && mtp_rows > rows) { payload_set_err(err, errlen, "KV checkpoint MTP rows exceed the token count"); rc = 1; } uint8_t *buf = xmalloc(DS4_SESSION_IO_CHUNK); for (uint32_t il = 0; rc == 0 && il < DS4_N_LAYER; il++) { if (ds4_qwen4_layer_is_linear(il)) { rc = payload_read_tensor_span(fp, g->layer_lin_state[il], 0, qwen4_payload_lin_state_bytes(), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); if (rc == 0) { rc = payload_read_tensor_span(fp, g->layer_lin_hist[il], 0, qwen4_payload_lin_hist_bytes(), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); } } else { const uint32_t n = ds4_qwen4_layer_is_nextn(il) ? mtp_rows : rows; rc = payload_read_tensor_span(fp, g->layer_k_cache[il], 0, qwen4_payload_kv_bytes(n), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_v_cache[il], 0, qwen4_payload_kv_bytes(n), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_ik_cache[il], 0, qwen4_payload_ik_bytes(n), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); if (rc == 0) rc = payload_read_tensor_span(fp, g->layer_block_key[il], 0, qwen4_payload_block_key_bytes(n), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); } } if (rc == 0) { rc = payload_read_tensor_span(fp, g->ple_hist, 0, qwen4_payload_ple_hist_bytes(), buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); } for (uint32_t i = 0; rc == 0 && i < DS4_MAX_PLE_NGRAM; i++) { uint32_t v; rc = payload_read_u32(fp, &v, remaining, err, errlen); if (rc == 0 && v >= DS4_N_VOCAB) { payload_set_err(err, errlen, "KV checkpoint n-gram history is outside the vocabulary"); rc = 1; } if (rc == 0) g->ple_prev[i] = (int)v; } if (rc == 0) { uint32_t v; rc = payload_read_u32(fp, &v, remaining, err, errlen); if (rc == 0) g->mrope_delta = (int32_t)v; } if (rc == 0) { rc = payload_read_tensor_span(fp, g->pos3, 0, (uint64_t)rows * 16u, buf, DS4_SESSION_IO_CHUNK, remaining, err, errlen); } free(buf); if (rc != 0) { token_vec_free(&new_checkpoint); qwen4_graph_reset(g); s->checkpoint.len = 0; return 1; } g->pos = rows; g->mtp_pos = mtp_rows; token_vec_free(&s->checkpoint); s->checkpoint = new_checkpoint; s->checkpoint_valid = true; s->qwen4_rewound = false; return 0; } #endif 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); } #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) return ds41_save_payload(s, fp, err, errlen); #endif if (ds4_session_is_qwen4(s)) { #ifndef DS4_HAS_QWEN4_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else return qwen4_session_save_payload(s, fp, err, errlen); #endif } 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 = session_glm_index_live_rows(il, compact_live); 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++) { if (g->glm53 && ds4_glm53_layer_is_kda(il)) { rc = payload_write_tensor_span( fp, g->layer_kda_conv_state[il], 0, session_glm_kda_conv_state_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_tensor_span( fp, g->layer_kda_recurrent_state[il], 0, session_glm_kda_recurrent_state_bytes(), buf, DS4_SESSION_IO_CHUNK, err, errlen); } continue; } 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)) { const uint32_t index_rows = session_glm_index_live_rows(il, compact_live); rc = payload_write_glm_compact_span(fp, g->layer_indexer_key_cache[il], (uint64_t)index_rows * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0 && g->glm53) { const uint64_t tail_bytes = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); rc = payload_write_tensor_span(fp, g->layer_indexer_tail_k[il], 0, tail_bytes, buf, DS4_SESSION_IO_CHUNK, err, errlen); if (rc == 0) { rc = payload_write_tensor_span( fp, g->layer_indexer_tail_gate[il], 0, tail_bytes, 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 (s->engine && s->engine->tp.active) { /* A local payload cannot restore another rank's caches. Keep the exact * saved tokens, consume the payload (leaving trailers readable), and * rebuild both ranks through the ordinary mirrored sync protocol. */ if (!ds4_session_tp_leader(s) || h[7] >= (uint32_t)s->ctx_size || h[11] != DS4_N_VOCAB) { payload_set_err(err, errlen, "invalid TP checkpoint token history"); return 1; } ds4_tokens tokens = {0}; int rc = payload_read_tokens_for_rebuild(fp, h[7], remaining, &tokens, err, errlen); if (rc == 0) { /* A load may hold the disk-cache lock. Do not call progress hooks * that can recursively save another cache entry during rebuild. */ ds4_session_progress_fn progress = s->progress; ds4_session_progress_fn display_progress = s->display_progress; s->progress = NULL; s->display_progress = NULL; ds4_session_invalidate(s); rc = ds4_session_sync(s, &tokens, err, errlen); s->progress = progress; s->display_progress = display_progress; } ds4_tokens_free(&tokens); return rc; } #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) return ds41_load_payload(s, fp, h, remaining, err, errlen); #endif if (ds4_session_is_qwen4(s)) { #ifndef DS4_HAS_QWEN4_GPU payload_set_err(err, errlen, "graph backend support is not compiled in"); return 1; #else return qwen4_session_load_payload(s, fp, h, &remaining, err, errlen); #endif } 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 = session_glm_index_live_rows(il, expected_compact_live); 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++) { if (g->glm53 && ds4_glm53_layer_is_kda(il)) { rc = payload_read_tensor_span( fp, g->layer_kda_conv_state[il], 0, session_glm_kda_conv_state_bytes(), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_tensor_span( fp, g->layer_kda_recurrent_state[il], 0, session_glm_kda_recurrent_state_bytes(), buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); } continue; } 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)n_index_comp[il] * DS4_N_INDEXER_HEAD_DIM, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0 && g->glm53) { const uint64_t tail_bytes = (uint64_t)DS4_GLM53_INDEX_POOL_SIZE * DS4_N_INDEXER_HEAD_DIM * sizeof(float); rc = payload_read_tensor_span( fp, g->layer_indexer_tail_k[il], 0, tail_bytes, buf, DS4_SESSION_IO_CHUNK, &remaining, err, errlen); if (rc == 0) { rc = payload_read_tensor_span( fp, g->layer_indexer_tail_gate[il], 0, tail_bytes, 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; } /* fmemopen appends a NUL even in binary mode. Keep it outside the payload * or fclose can overwrite the final byte of the last cache tensor. */ const uint64_t capacity = bytes + 1; if (snap->cap < capacity) { uint8_t *p = realloc(snap->ptr, (size_t)capacity); if (!p) { payload_set_err(err, errlen, "out of memory while allocating session snapshot"); return 1; } snap->ptr = p; snap->cap = capacity; } FILE *fp = fmemopen(snap->ptr, (size_t)capacity, "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; } int ds4_dump_chat_tokenization(const char *model_path, const char *system, const char *prompt, ds4_think_mode think_mode, int ctx_size, 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); if (ds4_think_mode_level(think_mode) >= 0 && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) { fprintf(stderr, "ds4: --think-level requires a DeepSeek V4.1 model\n"); model_close(&model); return 2; } think_mode = ds4_think_mode_for_context(think_mode, ctx_size); vocab_load(&vocab, &model); encode_chat_prompt(&vocab, system, prompt ? prompt : "", think_mode, &tokens); dump_tokens_fp(fp, &vocab, &tokens); token_vec_free(&tokens); vocab_free(&vocab); model_close(&model); return 0; } #ifndef DS4_NO_GPU #ifdef DS4_HAS_DEEPSEEK41_GPU static bool ds41_memory_admit(ds4_engine *e, uint64_t graph_bytes, bool fit_cache); #endif 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; } static char *imatrix_find_marker(char *dataset, char *cursor, const char *marker) { const size_t marker_len = strlen(marker); char *p = cursor; while ((p = strstr(p, marker)) != NULL) { if (p == dataset || p[-1] == '\n') return p; p += marker_len; } return NULL; } static int ds4_engine_collect_sequential_imatrix( ds4_engine *e, char *dataset, size_t dataset_len, const char *dataset_path, const char *output_path, int ctx_size, int max_prompts, int max_tokens, int min_expert_samples) { ds4_glm_gpu_graph g = {0}; const bool v41 = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41; const char *name = v41 ? "V4.1" : "GLM"; #ifdef DS4_HAS_DEEPSEEK41_GPU ds41_gpu_graph d = {.table = {{.fd = -1}, {.fd = -1}}}; if (v41) { if (e->tp.active) { fprintf(stderr, "ds4: V4.1 imatrix collection requires a single-host session\n"); return 1; } if (ctx_size > 1048576 || !ds41_memory_admit(e, ds4_add_sat_u64(e->ds41_session_bytes, ds41_graph_bytes((uint32_t)ctx_size)), false) || !ds41_graph_alloc(&d, &e->model, &e->weights, e->model_path, (uint32_t)ctx_size, e->ssd_streaming)) return 1; } else #endif if (!glm_graph_alloc(&g, &e->model, &e->weights, ctx_size, e->ssd_streaming, e->ssd_streaming_cold)) { fprintf(stderr, "ds4: failed to allocate %s imatrix graph\n", name); return 1; } g.quality = e->quality; const uint32_t layers = v41 ? DS4_N_LAYER : g.normal_layers; ds4_imatrix_collector collector; if (!imatrix_collector_init(&collector, 1, dataset_path)) { fprintf(stderr, "ds4: failed to allocate %s imatrix collector\n", name); imatrix_collector_free(&collector); glm_graph_free(&g); #ifdef DS4_HAS_DEEPSEEK41_GPU ds41_graph_free(&d); #endif return 1; } g.imatrix = &collector; #ifdef __APPLE__ d.imatrix = &collector; #endif fprintf(stderr, "ds4: collecting %s routed imatrix from %s " "(layers=%u experts=%u ctx=%d)\n", name, dataset_path, DS4_N_LAYER, DS4_N_EXPERT, ctx_size); bool ok = true; int prompts_done = 0; int tokens_done = 0; bool sample_target_reached = false; char *cursor = dataset; const char *marker_lit = "===== DS4_IMATRIX_PROMPT"; while (ok && *cursor) { char *start = cursor; char *marker = imatrix_find_marker(dataset, cursor, marker_lit); if (marker) { char *nl = strchr(marker, '\n'); if (!nl) break; start = nl + 1; } else if (prompts_done != 0) { break; } char *next = imatrix_find_marker(dataset, start, marker_lit); char *end = next ? next : dataset + dataset_len; const 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) { #ifdef DS4_HAS_DEEPSEEK41_GPU if (v41) ds41_graph_reset(&d); else #endif if (!glm_graph_reset_kda_state(&g)) { fprintf(stderr, "ds4: failed to reset %s imatrix state\n", name); ok = false; } for (int pos = 0; ok && pos < prompt.len; pos++) { #ifdef DS4_HAS_DEEPSEEK41_GPU if (v41) ok = ds41_graph_step(&d, &e->model, &e->weights, prompt.v[pos], NULL); else #endif ok = glm_graph_forward_token(&g, &e->model, &e->weights, prompt.v[pos], NULL, (uint32_t)pos, NULL, NULL, false); if (ok && ((pos + 1) % 16 == 0 || pos + 1 == prompt.len)) { fprintf(stderr, "ds4: %s imatrix prompt %d: %d/%d tokens\r", name, prompts_done + 1, pos + 1, prompt.len); fflush(stderr); } } if (!ok) { fprintf(stderr, "ds4: %s imatrix failed at prompt %d\n", name, prompts_done + 1); } else { prompts_done++; tokens_done += prompt.len; fprintf(stderr, "ds4: %s imatrix prompts=%d tokens=%d routes=%llu " "min_samples=%u\r", name, prompts_done, tokens_done, (unsigned long long)collector.observed_routes, imatrix_collector_min_samples(&collector, &e->weights, layers)); fflush(stderr); sample_target_reached = min_expert_samples > 0 && imatrix_collector_min_samples(&collector, &e->weights, layers) >= (uint32_t)min_expert_samples; } } token_vec_free(&prompt); } *end = saved; if (sample_target_reached) break; 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); g.imatrix = NULL; if (tokens_done == 0) { fprintf(stderr, "ds4: imatrix dataset contains no usable tokens\n"); ok = false; } if (ok) { imatrix_collector_report_coverage(&collector, &e->weights, layers); if (min_expert_samples > 0 && !sample_target_reached) { fprintf(stderr, "ds4: %s imatrix minimum expert sample target %d " "was not reached\n", name, min_expert_samples); ok = false; } } if (ok) { ok = imatrix_collector_save(&collector, &e->weights, output_path); } if (ok) { fprintf(stderr, "ds4: wrote %s imatrix %s from %d prompts, %d tokens, " "%llu routes\n", name, output_path, prompts_done, tokens_done, (unsigned long long)collector.observed_routes); } imatrix_collector_free(&collector); glm_graph_free(&g); #ifdef DS4_HAS_DEEPSEEK41_GPU ds41_graph_free(&d); #endif return ok ? 0 : 1; } #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, int min_expert_samples) { #ifdef DS4_NO_GPU (void)e; (void)dataset_path; (void)output_path; (void)ctx_size; (void)max_prompts; (void)max_tokens; (void)min_expert_samples; 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; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { const int rc = ds4_engine_collect_sequential_imatrix(e, dataset, dataset_len, dataset_path, output_path, ctx_size, max_prompts, max_tokens, min_expert_samples); free(dataset); return rc; } 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; bool sample_target_reached = false; char *cursor = dataset; const char *marker_lit = "===== DS4_IMATRIX_PROMPT"; while (*cursor) { char *start = cursor; char *marker = imatrix_find_marker(dataset, cursor, marker_lit); if (marker) { char *nl = strchr(marker, '\n'); if (!nl) break; start = nl + 1; } else if (prompts_done != 0) { break; } char *next = imatrix_find_marker(dataset, 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; sample_target_reached = min_expert_samples > 0 && imatrix_collector_min_samples(&collector, weights, DS4_N_LAYER) >= (uint32_t)min_expert_samples; if (prompts_done % 10 == 0) { fprintf(stderr, "ds4: imatrix prompts=%d tokens=%d routes=%llu " "min_samples=%u\r", prompts_done, tokens_done, (unsigned long long)collector.observed_routes, imatrix_collector_min_samples(&collector, weights, DS4_N_LAYER)); fflush(stderr); } } token_vec_free(&prompt); } *end = saved; if (sample_target_reached) break; 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) { imatrix_collector_report_coverage(&collector, weights, DS4_N_LAYER); if (min_expert_samples > 0 && !sample_target_reached) { fprintf(stderr, "ds4: imatrix minimum expert sample target %d " "was not reached\n", min_expert_samples); ok = false; } } 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) || ds4_session_is_ds41(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_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { 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, .glm53 = ds4_model_is_glm53(), }; 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, DS4_SWIGLU_CLAMP_EXP); } 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, .glm53 = ds4_model_is_glm53(), }; 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, DS4_SWIGLU_CLAMP_EXP); 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 static bool engine_legacy_graph_test_supported(ds4_engine *e) { if (!ds4_engine_is_deepseek41(e)) return true; fprintf(stderr, "ds4: legacy graph diagnostics do not implement DeepSeek V4.1; use session logits instead\n"); return false; } int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt) { if (!engine_legacy_graph_test_supported(e)) return 1; #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) { if (!engine_legacy_graph_test_supported(e)) return 1; #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) { if (!engine_legacy_graph_test_supported(e)) return 1; #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 (!engine_legacy_graph_test_supported(e)) return 1; 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; } /* ------------------------------------------------------------------------ * Qwen3.8-Flash-Next CPU reference: f32, one token at a time, every * projection a plain row dot product over dequantized weights. It is the * parity oracle for the Metal graph, not an * inference path. Weight layouts follow the llama.cpp qwen4exp GGUF: GDN * value heads are tiled (v-head j pairs with k-head j % n_k_heads), norm * gammas except ssm_norm are folded to 1+w, ssm_a holds -exp(A_log). * --------------------------------------------------------------------- */ static void qwen4_ref_row(const ds4_model *m, const ds4_tensor *t, uint64_t row, float *out) { const uint64_t n = t->dim[0]; switch (t->type) { case DS4_TENSOR_F32: memcpy(out, (const float *)tensor_data(m, t) + row * n, n * sizeof(float)); break; case DS4_TENSOR_F16: { const uint16_t *p = (const uint16_t *)tensor_data(m, t) + row * n; for (uint64_t i = 0; i < n; i++) out[i] = f16_to_f32(p[i]); break; } case DS4_TENSOR_BF16: { const uint16_t *p = (const uint16_t *)tensor_data(m, t) + row * n; for (uint64_t i = 0; i < n; i++) { const uint32_t u = (uint32_t)p[i] << 16; memcpy(&out[i], &u, sizeof(float)); } break; } case DS4_TENSOR_Q8_0: { const uint64_t blocks = n / 32u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 34u; for (uint64_t b = 0; b < blocks; b++, p += 34u) { uint16_t dh; memcpy(&dh, p, sizeof(dh)); const float d = f16_to_f32(dh); const int8_t *qs = (const int8_t *)(p + 2); for (uint32_t j = 0; j < 32u; j++) out[b * 32u + j] = d * (float)qs[j]; } break; } case DS4_TENSOR_Q4_0: { const uint64_t blocks = n / 32u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 18u; for (uint64_t b = 0; b < blocks; b++, p += 18u) { uint16_t dh; memcpy(&dh, p, sizeof(dh)); const float d = f16_to_f32(dh); for (uint32_t j = 0; j < 16u; j++) { out[b * 32u + j] = d * ((float)(p[2 + j] & 0x0fu) - 8.0f); out[b * 32u + 16u + j] = d * ((float)(p[2 + j] >> 4) - 8.0f); } } break; } case DS4_TENSOR_Q4_1: { const uint64_t blocks = n / 32u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 20u; for (uint64_t b = 0; b < blocks; b++, p += 20u) { uint16_t dh, mh; memcpy(&dh, p, sizeof(dh)); memcpy(&mh, p + 2, sizeof(mh)); const float d = f16_to_f32(dh), mv = f16_to_f32(mh); for (uint32_t j = 0; j < 16u; j++) { out[b * 32u + j] = d * (float)(p[4 + j] & 0x0fu) + mv; out[b * 32u + 16u + j] = d * (float)(p[4 + j] >> 4) + mv; } } break; } case DS4_TENSOR_Q4_K: { const uint64_t blocks = n / 256u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 144u; for (uint64_t b = 0; b < blocks; b++, p += 144u) { uint16_t dh, mh; memcpy(&dh, p, 2); memcpy(&mh, p + 2, 2); const float d = f16_to_f32(dh), dmin = f16_to_f32(mh); const uint8_t *sc = p + 4; for (uint32_t g = 0; g < 8u; g++) { uint32_t s, mn; if (g < 4u) { s = sc[g] & 63u; mn = sc[g + 4] & 63u; } else { s = (sc[g + 4] & 0xFu) | ((sc[g - 4] & 0xC0u) >> 2); mn = (sc[g + 4] >> 4) | ((sc[g] & 0xC0u) >> 2); } const float ds = d * (float)s, dm = dmin * (float)mn; const uint8_t *qs = p + 16 + (g >> 1) * 32u; const uint32_t shift = (g & 1u) * 4u; for (uint32_t j = 0; j < 32u; j++) out[b * 256u + g * 32u + j] = ds * (float)((qs[j] >> shift) & 0xFu) - dm; } } break; } case DS4_TENSOR_MXFP4: { const uint64_t blocks = n / 32u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 17u; for (uint64_t b = 0; b < blocks; b++, p += 17u) { const float d = ds4_e8m0_to_f32(p[0]); for (uint32_t j = 0; j < 16u; j++) { out[b * 32u + j] = d * ds4_mxfp4_values[p[1 + j] & 0x0fu]; out[b * 32u + 16u + j] = d * ds4_mxfp4_values[p[1 + j] >> 4]; } } break; } case DS4_TENSOR_Q2_K: { const uint64_t blocks = n / 256u; const block_q2_K *p = (const block_q2_K *)((const uint8_t *)tensor_data(m, t) + row * blocks * 84u); for (uint64_t k = 0; k < blocks * 256u; k++) out[k] = q2_k_value_f32(p, (uint32_t)k); break; } case DS4_TENSOR_IQ2_XXS: { const uint64_t blocks = n / 256u; const uint8_t *p = (const uint8_t *)tensor_data(m, t) + row * blocks * 66u; for (uint64_t b = 0; b < blocks; b++, p += 66u) { uint16_t dh; memcpy(&dh, p, 2); const float d = f16_to_f32(dh); for (uint32_t ib32 = 0; ib32 < 8u; ib32++) { uint16_t q2[4]; memcpy(q2, p + 2 + ib32 * 8u, 8); const uint32_t aux_g = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); const uint32_t aux_s = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); const float dl = d * (0.5f + (float)(aux_s >> 28)) * 0.25f; for (uint32_t j = 0; j < 4u; j++) { const uint8_t *grid = (const uint8_t *)(iq2xxs_grid + ((aux_g >> (8u * j)) & 0xFFu)); const uint32_t signs = ksigns_iq2xs[(aux_s >> (7u * j)) & 127u]; for (uint32_t i = 0; i < 8u; i++) { out[b * 256u + ib32 * 32u + j * 8u + i] = dl * (float)grid[i] * (((signs >> i) & 1u) ? -1.0f : 1.0f); } } } } break; } default: ds4_die("qwen4 reference: unsupported tensor type"); } } static void qwen4_ref_matvec_rows( const ds4_model *m, const ds4_tensor *w, uint64_t row0, uint64_t n_rows, const float *x, float *out) { const uint64_t n = w->dim[0]; float *row = xmalloc(n * sizeof(float)); for (uint64_t r = 0; r < n_rows; r++) { qwen4_ref_row(m, w, row0 + r, row); double acc = 0.0; for (uint64_t i = 0; i < n; i++) acc += (double)row[i] * x[i]; out[r] = (float)acc; } free(row); } static void qwen4_ref_matvec(const ds4_model *m, const ds4_tensor *w, const float *x, float *out) { qwen4_ref_matvec_rows(m, w, 0, w->ndim >= 2 ? w->dim[1] : 1u, x, out); } static const float *qwen4_ref_f32(const ds4_model *m, const ds4_tensor *t) { if (t->type != DS4_TENSOR_F32) ds4_die("qwen4 reference: expected an F32 tensor"); return (const float *)tensor_data(m, t); } static void qwen4_ref_rms(float *out, const float *x, const float *g, uint32_t n, float eps) { double ss = 0.0; for (uint32_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; const float scale = 1.0f / sqrtf((float)(ss / n) + eps); for (uint32_t i = 0; i < n; i++) out[i] = x[i] * scale * (g ? g[i] : 1.0f); } static void qwen4_ref_grouped_rms(float *out, const float *x, const float *g, uint32_t groups, uint32_t n, float eps) { for (uint32_t s = 0; s < groups; s++) { qwen4_ref_rms(out + (uint64_t)s * n, x + (uint64_t)s * n, g ? g + (uint64_t)s * n : NULL, n, eps); } } static void qwen4_ref_l2norm(float *x, uint32_t n) { double ss = 0.0; for (uint32_t i = 0; i < n; i++) ss += (double)x[i] * x[i]; const float scale = 1.0f / sqrtf((float)ss + 1e-6f); for (uint32_t i = 0; i < n; i++) x[i] *= scale; } static void qwen4_ref_rope(float *x, uint32_t n_rot, const uint32_t *p3) { const uint32_t half = n_rot / 2u; for (uint32_t i = 0; i < half; i++) { const double theta = (double)p3[i % 3u] * (double)g_qwen4_rope_freq[i]; const float c = (float)cos(theta) * g_qwen4_rope_mscale, s = (float)sin(theta) * g_qwen4_rope_mscale; const float x0 = x[i], x1 = x[i + half]; x[i] = x0 * c - x1 * s; x[i + half] = x0 * s + x1 * c; } } static float qwen4_ref_softplus(float x) { return x > 20.0f ? x : log1pf(expf(x)); } typedef struct { uint32_t cap; float *lin_state; /* [layer][Hv][Dk][Dv] */ float *lin_hist; /* [layer][K-1][conv_dim], oldest first, raw qkv */ float *ple_hist; /* [(Kp-1)*ngram][hc_dim], oldest first, normed */ int ple_prev[DS4_MAX_PLE_NGRAM]; /* previous tokens, newest first */ float *attn_k; /* [layer][cap][Hkv*D] */ float *attn_v; float *idx_k; /* [layer][cap][Di] raw indexer keys */ uint32_t *pos3; /* [cap][3] rope positions */ int32_t mrope_delta; } qwen4_ref_state; static void qwen4_ref_state_init(qwen4_ref_state *st, uint32_t cap) { const uint64_t L = DS4_N_LAYER; const uint64_t conv_dim = DS4_N_LIN_CONV_DIM; const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; memset(st, 0, sizeof(*st)); st->cap = cap; st->lin_state = xcalloc(L * DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM * DS4_N_LIN_HEAD_DIM, sizeof(float)); st->lin_hist = xcalloc(L * (DS4_N_LIN_CONV - 1u) * conv_dim, sizeof(float)); st->ple_hist = xcalloc((uint64_t)(DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM * hc_dim, sizeof(float)); for (uint32_t i = 0; i < DS4_MAX_PLE_NGRAM; i++) st->ple_prev[i] = DS4_PLE_EOS_ID; st->attn_k = xcalloc(L * cap * DS4_N_HEAD_KV * DS4_N_HEAD_DIM, sizeof(float)); st->attn_v = xcalloc(L * cap * DS4_N_HEAD_KV * DS4_N_HEAD_DIM, sizeof(float)); st->idx_k = xcalloc(L * cap * DS4_N_INDEXER_HEAD_DIM, sizeof(float)); st->pos3 = xcalloc((uint64_t)cap * 3u, sizeof(uint32_t)); } static void qwen4_ref_state_free(qwen4_ref_state *st) { free(st->lin_state); free(st->lin_hist); free(st->ple_hist); free(st->attn_k); free(st->attn_v); free(st->idx_k); free(st->pos3); memset(st, 0, sizeof(*st)); } /* hyper-connection mix: grouped norm, low-rank sigmoid gate, mean over * streams; inject (2*sigmoid(w/hc)) is the scatter weight of the combine */ static void qwen4_ref_hc_mix( const ds4_model *m, const float *R, const ds4_tensor *norm, const ds4_tensor *down, const ds4_tensor *up, const ds4_tensor *inject, float *mixed, float *inj) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, rank = DS4_N_HC_LOWRANK; const uint32_t hc_dim = E * hc; float *xn = xmalloc(hc_dim * sizeof(float)); float *lo = xmalloc(rank * sizeof(float)); float *gate = xmalloc(hc_dim * sizeof(float)); qwen4_ref_grouped_rms(xn, R, qwen4_ref_f32(m, norm), hc, E, DS4_RMS_EPS); qwen4_ref_matvec(m, down, xn, lo); for (uint32_t r = 0; r < rank; r++) lo[r] = silu(lo[r] / (float)hc); qwen4_ref_matvec(m, up, lo, gate); for (uint32_t d = 0; d < E; d++) { double acc = 0.0; for (uint32_t s = 0; s < hc; s++) { const uint32_t j = s * E + d; acc += (double)sigmoid_stable(gate[j]) * xn[j]; } mixed[d] = (float)(acc / hc); } if (inject && inj) { qwen4_ref_matvec(m, inject, xn, inj); for (uint32_t s = 0; s < hc; s++) inj[s] = 2.0f * sigmoid_stable(inj[s] / (float)hc); } free(gate); free(lo); free(xn); } static void qwen4_ref_hc_combine(float *R, const float *out, const float *inj) { for (uint32_t s = 0; s < DS4_N_HC; s++) { float *dst = R + (uint64_t)s * DS4_N_EMBD; for (uint32_t d = 0; d < DS4_N_EMBD; d++) dst[d] += inj[s] * out[d]; } } /* n-gram hash rows for `token` over the newest-first history `prev`, which * then advances; EOS fills the context at and before the last EOS. */ static void qwen4_ple_step(int token, int *prev, uint32_t *rows) { const uint32_t per = DS4_N_PLE_HEADS_PER_NGRAM; int ctx[DS4_MAX_PLE_NGRAM]; bool cut = false; ctx[0] = token; for (uint32_t s = 1; s < DS4_N_PLE_NGRAM; s++) { const int t = cut ? DS4_PLE_EOS_ID : prev[s - 1]; cut = cut || t == DS4_PLE_EOS_ID; ctx[s] = cut ? DS4_PLE_EOS_ID : t; } for (uint32_t n = 2; n <= DS4_N_PLE_NGRAM; n++) { uint64_t mixed = (uint64_t)ctx[0] * g_ds4_qwen4_ple.multiplier[0]; for (uint32_t j = 1; j < n; j++) mixed ^= (uint64_t)ctx[j] * g_ds4_qwen4_ple.multiplier[j]; for (uint32_t g = 0; g < per; g++) { const uint32_t h = (n - 2u) * per + g; rows[h] = (uint32_t)(mixed % g_ds4_qwen4_ple.head_vocab[h]) + g_ds4_qwen4_ple.head_offset[h]; } } for (uint32_t s = DS4_MAX_PLE_NGRAM - 1u; s > 0; s--) prev[s] = prev[s - 1]; prev[0] = token; } static void qwen4_ref_ple(const ds4_model *m, const ds4_layer_weights *l, qwen4_ref_state *st, int token, float *R) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; const uint32_t n_heads = DS4_N_PLE_HEADS; const uint32_t hist_rows = (DS4_N_PLE_CONV - 1u) * DS4_N_PLE_NGRAM; uint32_t rows[DS4_MAX_PLE_HEADS]; qwen4_ple_step(token, st->ple_prev, rows); float *emb = xmalloc(E * sizeof(float)); if (!qwen4_ngram_read(m, rows, n_heads, emb)) ds4_die_errno("cannot read Qwen n-grams", "GGUF"); float *key = xmalloc(hc_dim * sizeof(float)); float *keyn = xmalloc(hc_dim * sizeof(float)); float *query = xmalloc(hc_dim * sizeof(float)); float *value = xmalloc(E * sizeof(float)); float *gated = xmalloc(hc_dim * sizeof(float)); float *normed = xmalloc(hc_dim * sizeof(float)); qwen4_ref_matvec(m, l->ple_key, emb, key); qwen4_ref_grouped_rms(keyn, key, qwen4_ref_f32(m, l->ple_norm_key), hc, E, DS4_RMS_EPS); qwen4_ref_grouped_rms(query, R, qwen4_ref_f32(m, l->ple_norm_query), hc, E, DS4_RMS_EPS); qwen4_ref_matvec(m, l->ple_value, emb, value); for (uint32_t s = 0; s < hc; s++) { double dot = 0.0; for (uint32_t d = 0; d < E; d++) dot += (double)keyn[s * E + d] * query[s * E + d]; float g = (float)(dot / sqrt((double)E)); const float mag = sqrtf(fmaxf(fabsf(g), 1e-6f)); g = sigmoid_stable(g > 0.0f ? mag : (g < 0.0f ? -mag : 0.0f)); for (uint32_t d = 0; d < E; d++) gated[s * E + d] = g * value[d]; } qwen4_ref_grouped_rms(normed, gated, qwen4_ref_f32(m, l->ple_norm_conv), hc, E, DS4_RMS_EPS); /* depthwise conv, kernel K dilated by the n-gram size: tap k reads (K-1-k)*dil back */ float *cw_f16 = NULL; const float *cw; if (l->ple_conv->type == DS4_TENSOR_F16) { cw_f16 = xmalloc((size_t)hc_dim * DS4_N_PLE_CONV * sizeof(float)); for (uint32_t c = 0; c < hc_dim; c++) qwen4_ref_row(m, l->ple_conv, c, cw_f16 + (uint64_t)c * DS4_N_PLE_CONV); cw = cw_f16; } else { cw = qwen4_ref_f32(m, l->ple_conv); } for (uint32_t c = 0; c < hc_dim; c++) { double acc = 0.0; for (uint32_t k = 0; k < DS4_N_PLE_CONV; k++) { const uint32_t back = (DS4_N_PLE_CONV - 1u - k) * DS4_N_PLE_NGRAM; const float xk = back == 0 ? normed[c] : st->ple_hist[(uint64_t)(hist_rows - back) * hc_dim + c]; acc += (double)cw[(uint64_t)c * DS4_N_PLE_CONV + k] * xk; } R[c] += gated[c] + silu((float)acc); } free(cw_f16); memmove(st->ple_hist, st->ple_hist + hc_dim, (uint64_t)(hist_rows - 1u) * hc_dim * sizeof(float)); memcpy(st->ple_hist + (uint64_t)(hist_rows - 1u) * hc_dim, normed, hc_dim * sizeof(float)); free(normed); free(gated); free(value); free(query); free(keyn); free(key); free(emb); } static void qwen4_ref_linear(const ds4_model *m, const ds4_layer_weights *l, uint32_t il, qwen4_ref_state *st, const float *x, float *out) { const uint32_t Hk = DS4_N_LIN_K_HEAD, Hv = DS4_N_LIN_V_HEAD, D = DS4_N_LIN_HEAD_DIM, K = DS4_N_LIN_CONV; const uint32_t k_dim = Hk * D, v_dim = Hv * D, conv_dim = 2u * k_dim + v_dim; float *qkv = xmalloc(conv_dim * sizeof(float)); float *conv = xmalloc(conv_dim * sizeof(float)); float *z = xmalloc(v_dim * sizeof(float)); float *b = xmalloc(Hv * sizeof(float)); float *a = xmalloc(Hv * sizeof(float)); float *o = xmalloc(v_dim * sizeof(float)); qwen4_ref_matvec(m, l->lin_qkv, x, qkv); qwen4_ref_matvec(m, l->lin_gate, x, z); qwen4_ref_matvec(m, l->lin_beta, x, b); qwen4_ref_matvec(m, l->lin_alpha, x, a); float *hist = st->lin_hist + (uint64_t)il * (K - 1u) * conv_dim; const float *cw = qwen4_ref_f32(m, l->lin_conv); for (uint32_t c = 0; c < conv_dim; c++) { double acc = (double)cw[(uint64_t)c * K + (K - 1u)] * qkv[c]; for (uint32_t k = 0; k + 1u < K; k++) acc += (double)cw[(uint64_t)c * K + k] * hist[(uint64_t)k * conv_dim + c]; conv[c] = silu((float)acc); } memmove(hist, hist + conv_dim, (uint64_t)(K - 2u) * conv_dim * sizeof(float)); memcpy(hist + (uint64_t)(K - 2u) * conv_dim, qkv, conv_dim * sizeof(float)); float *q = conv, *k = conv + k_dim, *v = conv + 2u * k_dim; const float qscale = 1.0f / sqrtf((float)D); for (uint32_t h = 0; h < Hk; h++) { qwen4_ref_l2norm(q + h * D, D); qwen4_ref_l2norm(k + h * D, D); for (uint32_t i = 0; i < D; i++) q[h * D + i] *= qscale; } const float *A = qwen4_ref_f32(m, l->lin_a); const float *dt = qwen4_ref_f32(m, l->lin_dt_bias); const float *nw = qwen4_ref_f32(m, l->lin_norm); for (uint32_t j = 0; j < Hv; j++) { const uint32_t kh = j % Hk; const float g = expf(A[j] * qwen4_ref_softplus(a[j] + dt[j])); const float beta = sigmoid_stable(b[j]); float *S = st->lin_state + ((uint64_t)il * Hv + j) * D * D; /* [dk][dv] */ const float *qj = q + kh * D, *kj = k + kh * D, *vj = v + j * D; float *oj = o + j * D; for (uint32_t dv = 0; dv < D; dv++) { double kv = 0.0; for (uint32_t dk = 0; dk < D; dk++) { S[dk * D + dv] *= g; kv += (double)S[dk * D + dv] * kj[dk]; } const float delta = (vj[dv] - (float)kv) * beta; double acc = 0.0; for (uint32_t dk = 0; dk < D; dk++) { S[dk * D + dv] += kj[dk] * delta; acc += (double)S[dk * D + dv] * qj[dk]; } oj[dv] = (float)acc; } float tmp[DS4_MAX_KDA_HEAD_DIM]; if (D > DS4_MAX_KDA_HEAD_DIM) ds4_die("Qwen3.8 reference: linear head dim exceeds 128"); qwen4_ref_rms(tmp, oj, nw, D, DS4_RMS_EPS); for (uint32_t dv = 0; dv < D; dv++) oj[dv] = tmp[dv] * sigmoid_stable(z[j * D + dv]); } qwen4_ref_matvec(m, l->lin_out, o, out); free(o); free(a); free(b); free(z); free(conv); free(qkv); } /* QSA block selection for the query at `pos`: complete blocks of ratio * tokens ranked by relu-summed indexer scores, plus the incomplete tail. */ static uint32_t qwen4_ref_select(const ds4_model *m, const ds4_layer_weights *l, uint32_t il, qwen4_ref_state *st, const float *x, uint32_t pos, int32_t *sel) { const uint32_t ratio = 4u, Di = DS4_N_INDEXER_HEAD_DIM, Hi = DS4_N_INDEXER_HEAD; const uint32_t k_blocks = DS4_N_INDEXER_TOP_K / ratio; float *qi = xmalloc((uint64_t)Hi * Di * sizeof(float)); float *kraw = st->idx_k + ((uint64_t)il * st->cap + pos) * Di; qwen4_ref_matvec(m, l->indexer_q_proj, x, qi); qwen4_ref_matvec(m, l->indexer_k_proj, x, kraw); const uint32_t n_vis = pos + 1u; const uint32_t n_blocks = n_vis / ratio; uint32_t n_sel = 0; if (n_blocks <= k_blocks) { for (uint32_t t = 0; t < n_vis; t++) sel[n_sel++] = (int32_t)t; free(qi); return n_sel; } const float *gq = qwen4_ref_f32(m, l->indexer_q_norm); const float *gk = qwen4_ref_f32(m, l->indexer_k_norm); for (uint32_t h = 0; h < Hi; h++) { float tmp[DS4_MAX_INDEXER_HEAD_DIM]; qwen4_ref_rms(tmp, qi + h * Di, gq, Di, DS4_RMS_EPS); qwen4_ref_rope(tmp, DS4_N_ROT, st->pos3 + (uint64_t)pos * 3u); memcpy(qi + h * Di, tmp, Di * sizeof(float)); } float *score = xmalloc(n_blocks * sizeof(float)); float *pooled = xmalloc(Di * sizeof(float)); float *key = xmalloc(Di * sizeof(float)); for (uint32_t b = 0; b < n_blocks; b++) { for (uint32_t d = 0; d < Di; d++) { double acc = 0.0; for (uint32_t t = 0; t < ratio; t++) acc += st->idx_k[((uint64_t)il * st->cap + b * ratio + t) * Di + d]; pooled[d] = (float)(acc / ratio); } qwen4_ref_rms(key, pooled, gk, Di, DS4_RMS_EPS); qwen4_ref_rope(key, DS4_N_ROT, st->pos3 + (uint64_t)b * ratio * 3u); double sum = 0.0; for (uint32_t h = 0; h < Hi; h++) { double dot = 0.0; for (uint32_t d = 0; d < Di; d++) dot += (double)qi[h * Di + d] * key[d]; if (dot > 0.0) sum += dot; } score[b] = (float)sum; } /* top-k blocks, lower index first on ties */ bool *taken = xcalloc(n_blocks, sizeof(bool)); for (uint32_t i = 0; i < k_blocks; i++) { int best = -1; for (uint32_t b = 0; b < n_blocks; b++) { if (taken[b]) continue; if (best < 0 || score[b] > score[best]) best = (int)b; } taken[best] = true; } for (uint32_t b = 0; b < n_blocks; b++) { if (!taken[b]) continue; for (uint32_t t = 0; t < ratio; t++) sel[n_sel++] = (int32_t)(b * ratio + t); } for (uint32_t t = n_blocks * ratio; t < n_vis; t++) sel[n_sel++] = (int32_t)t; free(taken); free(key); free(pooled); free(score); free(qi); return n_sel; } static void qwen4_ref_attention(const ds4_model *m, const ds4_layer_weights *l, uint32_t il, qwen4_ref_state *st, const float *x, uint32_t pos, float *out) { const uint32_t H = DS4_N_HEAD, Hkv = DS4_N_HEAD_KV, D = DS4_N_HEAD_DIM; const uint32_t q_dim = H * D, kv_dim = Hkv * D; float *qg = xmalloc(2u * q_dim * sizeof(float)); float *q = xmalloc(q_dim * sizeof(float)); float *gate = xmalloc(q_dim * sizeof(float)); float *o = xmalloc(q_dim * sizeof(float)); float *kc = st->attn_k + ((uint64_t)il * st->cap + pos) * kv_dim; float *vc = st->attn_v + ((uint64_t)il * st->cap + pos) * kv_dim; qwen4_ref_matvec(m, l->attn_q, x, qg); qwen4_ref_matvec(m, l->attn_k, x, kc); qwen4_ref_matvec(m, l->attn_v, x, vc); const float *gqn = qwen4_ref_f32(m, l->attn_q_norm); const float *gkn = qwen4_ref_f32(m, l->attn_k_norm); for (uint32_t h = 0; h < H; h++) { qwen4_ref_rms(q + h * D, qg + (uint64_t)h * 2u * D, gqn, D, DS4_RMS_EPS); qwen4_ref_rope(q + h * D, DS4_N_ROT, st->pos3 + (uint64_t)pos * 3u); memcpy(gate + h * D, qg + (uint64_t)h * 2u * D + D, D * sizeof(float)); } for (uint32_t h = 0; h < Hkv; h++) { float tmp[DS4_MAX_HEAD_DIM]; qwen4_ref_rms(tmp, kc + h * D, gkn, D, DS4_RMS_EPS); qwen4_ref_rope(tmp, DS4_N_ROT, st->pos3 + (uint64_t)pos * 3u); memcpy(kc + h * D, tmp, D * sizeof(float)); } int32_t *sel = xmalloc((uint64_t)(pos + 1u) * sizeof(int32_t)); const uint32_t n_sel = qwen4_ref_select(m, l, il, st, x, pos, sel); float *p = xmalloc(n_sel * sizeof(float)); const float scale = 1.0f / sqrtf((float)D); for (uint32_t h = 0; h < H; h++) { const uint32_t kvh = h / (H / Hkv); float mx = -FLT_MAX; for (uint32_t i = 0; i < n_sel; i++) { const float *kt = st->attn_k + ((uint64_t)il * st->cap + (uint32_t)sel[i]) * kv_dim + kvh * D; double dot = 0.0; for (uint32_t d = 0; d < D; d++) dot += (double)q[h * D + d] * kt[d]; p[i] = (float)dot * scale; if (p[i] > mx) mx = p[i]; } double sum = 0.0; for (uint32_t i = 0; i < n_sel; i++) { p[i] = expf(p[i] - mx); sum += p[i]; } for (uint32_t d = 0; d < D; d++) { double acc = 0.0; for (uint32_t i = 0; i < n_sel; i++) { const float *vt = st->attn_v + ((uint64_t)il * st->cap + (uint32_t)sel[i]) * kv_dim + kvh * D; acc += (double)p[i] * vt[d]; } o[h * D + d] = (float)(acc / sum) * sigmoid_stable(gate[h * D + d]); } } qwen4_ref_matvec(m, l->attn_output, o, out); free(p); free(sel); free(o); free(gate); free(q); free(qg); } static void qwen4_ref_moe(const ds4_model *m, const ds4_layer_weights *l, const float *x, float *out) { const uint32_t E = DS4_N_EMBD, NE = DS4_N_EXPERT, K = DS4_N_EXPERT_USED, F = DS4_N_FF_EXP; float *logits = xmalloc(NE * sizeof(float)); double *prob = xmalloc(NE * sizeof(double)); const uint32_t down_width = (uint32_t)l->ffn_down_exps->dim[0]; float *g = xcalloc(down_width, sizeof(float)); float *u = xmalloc(F * sizeof(float)); float *y = xmalloc(E * sizeof(float)); qwen4_ref_matvec(m, l->ffn_gate_inp, x, logits); double mx = -DBL_MAX, sum = 0.0; for (uint32_t e = 0; e < NE; e++) if (logits[e] > mx) mx = logits[e]; for (uint32_t e = 0; e < NE; e++) { prob[e] = exp((double)logits[e] - mx); sum += prob[e]; } for (uint32_t e = 0; e < NE; e++) prob[e] /= sum; int sel[DS4_MAX_EXPERT_USED]; double wsum = 0.0; for (uint32_t i = 0; i < K; i++) { int best = -1; for (uint32_t e = 0; e < NE; e++) { bool used = false; for (uint32_t j = 0; j < i; j++) used |= sel[j] == (int)e; if (!used && (best < 0 || prob[e] > prob[best])) best = (int)e; } if (best < 0) ds4_die("Qwen3.8 router selected more experts than available"); sel[i] = best; wsum += prob[best]; } for (uint32_t d = 0; d < E; d++) out[d] = 0.0f; for (uint32_t i = 0; i < K; i++) { const uint32_t e = (uint32_t)sel[i]; const float wgt = (float)(prob[e] / wsum); qwen4_ref_matvec_rows(m, l->ffn_gate_exps, (uint64_t)e * F, F, x, g); qwen4_ref_matvec_rows(m, l->ffn_up_exps, (uint64_t)e * F, F, x, u); for (uint32_t f = 0; f < F; f++) g[f] = silu(g[f]) * u[f]; qwen4_ref_matvec_rows(m, l->ffn_down_exps, (uint64_t)e * E, E, g, y); for (uint32_t d = 0; d < E; d++) out[d] += wgt * y[d]; } qwen4_ref_matvec(m, l->ffn_gate_shexp, x, g); qwen4_ref_matvec(m, l->ffn_up_shexp, x, u); for (uint32_t f = 0; f < F; f++) g[f] = silu(g[f]) * u[f]; qwen4_ref_matvec(m, l->ffn_down_shexp, g, y); float sg; qwen4_ref_matvec(m, l->ffn_gate_inp_shexp, x, &sg); sg = sigmoid_stable(sg); for (uint32_t d = 0; d < E; d++) out[d] += sg * y[d]; free(y); free(u); free(g); free(prob); free(logits); } /* one trunk-layer step on the wide residual R */ static void qwen4_ref_layer(const ds4_model *m, const ds4_weights *w, uint32_t il, qwen4_ref_state *st, int token, uint32_t pos, float *R) { const ds4_layer_weights *l = &w->layer[il]; float mixed[DS4_MAX_EMBD], blk[DS4_MAX_EMBD], inj[DS4_MAX_HC]; if (ds4_qwen4_layer_is_ple(il)) qwen4_ref_ple(m, l, st, token, R); qwen4_ref_hc_mix(m, R, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, mixed, inj); if (ds4_qwen4_layer_is_linear(il)) qwen4_ref_linear(m, l, il, st, mixed, blk); else qwen4_ref_attention(m, l, il, st, mixed, pos, blk); qwen4_ref_hc_combine(R, blk, inj); qwen4_ref_hc_mix(m, R, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, mixed, inj); qwen4_ref_moe(m, l, mixed, blk); qwen4_ref_hc_combine(R, blk, inj); } static void qwen4_ref_output(const ds4_model *m, const ds4_weights *w, const float *mixed, float *logits) { qwen4_ref_matvec(m, w->output, mixed, logits); } /* Forward `token` at `pos`; R_out (optional) receives the pre-mixer stream, * streams (optional) the stream after every trunk layer. */ static void qwen4_ref_forward_token(const ds4_model *m, const ds4_weights *w, qwen4_ref_state *st, int token, uint32_t pos, float *logits, float *R_out, float *streams) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; float *R = xmalloc(hc_dim * sizeof(float)); float *mixed = xmalloc(E * sizeof(float)); size_t n_fake = 0; const ds4_vision_span *fake = qwen4_fake_spans(&n_fake); const float *img = n_fake ? qwen4_span_row(fake, n_fake, pos) : NULL; if (img) memcpy(R, img, E * sizeof(float)); else qwen4_ref_row(m, w->token_embd, (uint64_t)token, R); for (uint32_t s = 1; s < hc; s++) memcpy(R + (uint64_t)s * E, R, E * sizeof(float)); qwen4_mrope_pos(fake, n_fake, pos, &st->mrope_delta, st->pos3 + (uint64_t)pos * 3u); for (uint32_t il = 0; il < n_trunk; il++) { qwen4_ref_layer(m, w, il, st, token, pos, R); if (streams) memcpy(streams + (uint64_t)il * hc_dim, R, hc_dim * sizeof(float)); } if (R_out) memcpy(R_out, R, hc_dim * sizeof(float)); if (logits) { qwen4_ref_hc_mix(m, R, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, mixed, NULL); qwen4_ref_output(m, w, mixed, logits); } free(mixed); free(R); } /* MTP draft step: pair the pre-mixer stream of position i with the * embedding of token i+1, run the nextn block at index i on its own caches. * The hidden norm spans all hc*E dims (the vLLM/SGLang and HF form). */ static void qwen4_ref_mtp(const ds4_model *m, const ds4_weights *w, qwen4_ref_state *st, const float *R_pre, int next_token, uint32_t idx, float *logits, float *R_out) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; const uint32_t il = DS4_N_LAYER - 1u; const ds4_layer_weights *l = &w->layer[il]; float *e = xmalloc(E * sizeof(float)); float *en = xmalloc(E * sizeof(float)); float *hn = xmalloc(hc_dim * sizeof(float)); float *cat = xmalloc(2u * E * sizeof(float)); float *R = xmalloc(hc_dim * sizeof(float)); float *ep = xmalloc(E * sizeof(float)); float mixed[DS4_MAX_EMBD], blk[DS4_MAX_EMBD], inj[DS4_MAX_HC]; qwen4_ref_row(m, w->token_embd, (uint64_t)next_token, e); qwen4_ref_rms(en, e, qwen4_ref_f32(m, l->nextn_enorm), E, DS4_RMS_EPS); qwen4_ref_rms(hn, R_pre, qwen4_ref_f32(m, l->nextn_hnorm), hc_dim, DS4_RMS_EPS); /* eh_proj = [W_e | W_h] over concat(e, h): e once, h per stream */ memcpy(cat, en, E * sizeof(float)); memset(cat + E, 0, E * sizeof(float)); qwen4_ref_matvec(m, l->nextn_eh_proj, cat, ep); memset(cat, 0, E * sizeof(float)); for (uint32_t s = 0; s < hc; s++) { memcpy(cat + E, hn + (uint64_t)s * E, E * sizeof(float)); qwen4_ref_matvec(m, l->nextn_eh_proj, cat, R + (uint64_t)s * E); for (uint32_t d = 0; d < E; d++) R[(uint64_t)s * E + d] += ep[d]; } { size_t n_fake = 0; const ds4_vision_span *fake = qwen4_fake_spans(&n_fake); qwen4_mrope_pos(fake, n_fake, idx, &st->mrope_delta, st->pos3 + (uint64_t)idx * 3u); } qwen4_ref_hc_mix(m, R, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, mixed, inj); qwen4_ref_attention(m, l, il, st, mixed, idx, blk); qwen4_ref_hc_combine(R, blk, inj); qwen4_ref_hc_mix(m, R, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, mixed, inj); qwen4_ref_moe(m, l, mixed, blk); qwen4_ref_hc_combine(R, blk, inj); if (R_out) memcpy(R_out, R, hc_dim * sizeof(float)); if (logits) { qwen4_ref_hc_mix(m, R, l->nextn_hc_head_norm, l->nextn_hc_head_down, l->nextn_hc_head_up, NULL, mixed, NULL); qwen4_ref_output(m, w, mixed, logits); } free(ep); free(R); free(cat); free(hn); free(en); free(e); } static uint32_t qwen4_parse_token_list(const char *p, int *seq, uint32_t max) { uint32_t n = 0; while (*p && n < max) { seq[n++] = atoi(p); while (*p && *p != ',') p++; if (*p == ',') p++; } return n; } static void qwen4_dump_f32(const char *path, const float *v, uint64_t n) { FILE *f = fopen(path, "wb"); if (!f) return; fwrite(v, sizeof(float), n, f); fclose(f); } /* --first-token-test for Qwen3.8: DS4_QWEN4_FT_TOKENS (comma-separated ids, * default: the prompt) runs the sequential CPU reference; DS4_QWEN4_FT_OUT * dumps [n][vocab] f32 logits, DS4_QWEN4_FT_HIDDEN [n][trunk][hc*embd] * streams, DS4_QWEN4_MTP_OUT the [n-1][vocab] teacher-forced MTP drafts. * DS4_QWEN4_MTP_BATCH_CHECK also compares sequential and paired GPU drafts * (requires MTP_OUT, GPU=1, and GPU_CHUNK >= 2). */ static int qwen4_first_token_test(const ds4_model *model, const ds4_vocab *vocab, const ds4_weights *weights, const ds4_tokens *prompt, bool metal) { /* DS4_QWEN4_FT_LIST=: batch mode. Each line is * "\t"; every sequence runs through the Metal graph * (prefill all-rows logits, the production kernel path) and its * [n][vocab] f32 logits are written to the path. One model load covers * the whole list. */ const char *ft_list = getenv("DS4_QWEN4_FT_LIST"); if (ft_list && ft_list[0]) { #ifdef DS4_HAS_QWEN4_GPU if (!metal || !getenv("DS4_QWEN4_GPU")) { fprintf(stderr, "ds4: DS4_QWEN4_FT_LIST needs a GPU backend with DS4_QWEN4_GPU set\n"); return 1; } FILE *lf = fopen(ft_list, "r"); if (!lf) { fprintf(stderr, "ds4: cannot open DS4_QWEN4_FT_LIST %s\n", ft_list); return 1; } ds4_qwen4_gpu_graph *g = xcalloc(1, sizeof(*g)); if (!qwen4_graph_alloc(g, weights, 8192u, 128u, false, NULL, NULL, NULL, 0)) { free(g); fclose(lf); return 1; } ds4_gpu_tensor_free(g->logits); g->n_logit_rows = g->cap_tokens; g->logits = qwen4_graph_alloc_f32((uint64_t)g->n_logit_rows * DS4_N_VOCAB); if (!g->logits) { qwen4_graph_free(g); free(g); fclose(lf); return 1; } char *line = xmalloc(65536); int *lseq = xmalloc(4096 * sizeof(int)); float *out = xmalloc((size_t)4096 * DS4_N_VOCAB * sizeof(float)); while (fgets(line, 65536, lf)) { char *tab = strchr(line, '\t'); if (!tab) continue; *tab = '\0'; char *out_path = tab + 1; out_path[strcspn(out_path, "\r\n")] = '\0'; qwen4_graph_reset(g); /* "|": prefill the prompt in chunks * (the production path), then teacher-force the targets one at a * time; out receives one [vocab] row per target position. */ char *bar = strchr(line, '|'); if (bar) *bar = '\0'; uint32_t p_len = qwen4_parse_token_list(line, lseq, 4096); uint32_t ln = p_len; if (bar) { ln = qwen4_parse_token_list(bar + 1, lseq + p_len, 4096 - p_len); } if (p_len == 0 || ln == 0) continue; bool fok = true; for (uint32_t p = 0; p < p_len && fok;) { const uint32_t n = p_len - p < g->cap_tokens ? p_len - p : g->cap_tokens; float *dst = bar ? (p + n == p_len ? out : NULL) : out + (uint64_t)p * DS4_N_VOCAB; fok = qwen4_graph_forward_tokens(g, model, weights, lseq + p, n, dst, bar == NULL); if (!fok) fprintf(stderr, "ds4: FT_LIST chunk %u+%u failed (pos=%u cap=%u rows=%u)\n", p, n, g->pos, g->cap_tokens, g->n_logit_rows); p += n; } for (uint32_t t = 0; bar && t + 1u < ln && fok; t++) { fok = qwen4_graph_forward_tokens(g, model, weights, lseq + p_len + t, 1u, out + (uint64_t)(t + 1u) * DS4_N_VOCAB, false); } if (!fok) { fprintf(stderr, "ds4: FT_LIST forward failed (%u tokens)\n", ln); free(out); free(lseq); free(line); qwen4_graph_free(g); free(g); fclose(lf); return 1; } FILE *of = fopen(out_path, "wb"); if (!of || fwrite(out, sizeof(float), (size_t)ln * DS4_N_VOCAB, of) != (size_t)ln * DS4_N_VOCAB) { fprintf(stderr, "ds4: FT_LIST write failed: %s\n", out_path); if (of) fclose(of); free(out); free(lseq); free(line); qwen4_graph_free(g); free(g); fclose(lf); return 1; } fclose(of); printf("ft_list %u tokens -> %s\n", ln, out_path); fflush(stdout); } free(out); free(lseq); free(line); qwen4_graph_free(g); free(g); fclose(lf); return 0; #else (void)model; (void)vocab; (void)weights; (void)prompt; (void)metal; fprintf(stderr, "ds4: DS4_QWEN4_FT_LIST needs GPU support\n"); return 1; #endif } int *seq = xmalloc(4096 * sizeof(int)); uint32_t n_seq = 0; const char *seq_env = getenv("DS4_QWEN4_FT_TOKENS"); if (seq_env && seq_env[0]) { n_seq = qwen4_parse_token_list(seq_env, seq, 4096); } else { for (int i = 0; i < prompt->len && n_seq < 4096; i++) seq[n_seq++] = prompt->v[i]; } if (n_seq == 0) { free(seq); return 1; } const uint32_t V = DS4_N_VOCAB, hc_dim = DS4_N_EMBD * DS4_N_HC; const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; const char *hidden_out = getenv("DS4_QWEN4_FT_HIDDEN"); float *logits = xmalloc((uint64_t)n_seq * V * sizeof(float)); float *streams = xmalloc((uint64_t)n_seq * hc_dim * sizeof(float)); float *layer_streams = hidden_out && hidden_out[0] ? xmalloc((uint64_t)n_seq * n_trunk * hc_dim * sizeof(float)) : NULL; /* DS4_QWEN4_GPU_ONLY skips the (slow) CPU reference: dumps are then the * GPU's alone and the comparison lines are against zeros. */ const bool gpu_only = metal && getenv("DS4_QWEN4_GPU") && getenv("DS4_QWEN4_GPU_ONLY"); memset(logits, 0, (uint64_t)n_seq * V * sizeof(float)); qwen4_ref_state st; qwen4_ref_state_init(&st, n_seq); for (uint32_t t = 0; t < n_seq && !gpu_only; t++) { qwen4_ref_forward_token(model, weights, &st, seq[t], t, logits + (uint64_t)t * V, streams + (uint64_t)t * hc_dim, layer_streams ? layer_streams + (uint64_t)t * n_trunk * hc_dim : NULL); } qwen4_ref_state_free(&st); const char *mtp_out = getenv("DS4_QWEN4_MTP_OUT"); float *draft = NULL; if (mtp_out && mtp_out[0] && DS4_N_NEXTN_PREDICT != 0 && n_seq > 1) { draft = xcalloc((uint64_t)(n_seq - 1u) * V, sizeof(float)); qwen4_ref_state mst; qwen4_ref_state_init(&mst, n_seq); for (uint32_t i = 0; i + 1u < n_seq && !gpu_only; i++) { qwen4_ref_mtp(model, weights, &mst, streams + (uint64_t)i * hc_dim, seq[i + 1u], i, draft + (uint64_t)i * V, NULL); } qwen4_ref_state_free(&mst); } #ifdef DS4_HAS_QWEN4_GPU /* DS4_QWEN4_GPU=1 runs the Metal graph over the same ids (chunks of * DS4_QWEN4_GPU_CHUNK tokens, default 1) and reports it against the CPU * reference; the dumped logits (and MTP drafts) are then the GPU's. */ if (metal && getenv("DS4_QWEN4_GPU")) { const char *chunk_env = getenv("DS4_QWEN4_GPU_CHUNK"); uint32_t chunk = chunk_env ? (uint32_t)atoi(chunk_env) : 1u; if (chunk == 0) chunk = 1; if (chunk > n_seq) chunk = n_seq; const bool batch_check = getenv("DS4_QWEN4_MTP_BATCH_CHECK") != NULL; if (batch_check && (!draft || chunk < 2u || n_seq < 3u)) { fprintf(stderr, "ds4: MTP batch check needs MTP_OUT, GPU_CHUNK >= 2, and at least 3 tokens\n"); free(draft); free(layer_streams); free(streams); free(logits); free(seq); return 1; } ds4_qwen4_gpu_graph *g = xcalloc(1, sizeof(*g)); if (!qwen4_graph_alloc(g, weights, n_seq + 4u, chunk, draft != NULL, NULL, NULL, NULL, 0)) { free(g); return 1; } qwen4_graph_reset(g); float *gpu = xmalloc((uint64_t)V * sizeof(float)); float worst = 0.0f, mtp_worst = 0.0f; uint32_t agree = 0, scored = 0, mtp_agree = 0, mtp_scored = 0; float batch_worst = 0.0f; uint32_t batch_scored = 0; for (uint32_t t0 = 0; t0 < n_seq; t0 += chunk) { const uint32_t n = t0 + chunk <= n_seq ? chunk : n_seq - t0; if (!qwen4_graph_forward_tokens(g, model, weights, seq + t0, n, gpu, false)) { fprintf(stderr, "ds4: Qwen3.8 GPU forward failed at token %u\n", t0); free(gpu); qwen4_graph_free(g); free(g); return 1; } const uint32_t t = t0 + n - 1u; float *cpu = logits + (uint64_t)t * V; float max_err = 0.0f; uint32_t cb = 0, gb = 0; for (uint32_t i = 0; i < V; i++) { const float d = fabsf(gpu[i] - cpu[i]); if (d > max_err) max_err = d; if (cpu[i] > cpu[cb]) cb = i; if (gpu[i] > gpu[gb]) gb = i; } if (cb == gb) agree++; scored++; if (max_err > worst) worst = max_err; printf(" pos %3u: max|gpu-cpu|=%g top1 cpu=%u gpu=%u%s\n", t, (double)max_err, cb, gb, cb == gb ? "" : " <-- MISMATCH"); memcpy(cpu, gpu, (uint64_t)V * sizeof(float)); for (uint32_t j = 0; draft && j < n && t0 + j + 1u < n_seq; j++) { const uint32_t idx = t0 + j; int dtok = -1; if (!qwen4_graph_mtp_step(g, model, weights, j, seq[idx + 1u], idx, true, gpu, &dtok)) { fprintf(stderr, "ds4: Qwen3.8 GPU MTP step failed at %u\n", idx); free(gpu); qwen4_graph_free(g); free(g); return 1; } float *cpu_d = draft + (uint64_t)idx * V; float md = 0.0f; uint32_t cd = 0; for (uint32_t i = 0; i < V; i++) { const float dd = fabsf(gpu[i] - cpu_d[i]); if (dd > md) md = dd; if (cpu_d[i] > cpu_d[cd]) cd = i; } if ((int)cd == dtok) mtp_agree++; mtp_scored++; if (md > mtp_worst) mtp_worst = md; printf(" mtp %3u: max|gpu-cpu|=%g top1 cpu=%u gpu=%d%s\n", idx, (double)md, cd, dtok, (int)cd == dtok ? "" : " <-- MISMATCH"); memcpy(cpu_d, gpu, (uint64_t)V * sizeof(float)); if (batch_check && (j & 1u)) { /* Replay the same two positions. Causal attention cannot * read future entries, so existing KV rows are safe to overwrite. */ const int next_tokens[2] = {seq[idx], seq[idx + 1u]}; int paired = -1; bool bok = qwen4_graph_mtp_steps(g, model, weights, j - 1u, next_tokens, 2u, idx - 1u, true, gpu, &paired); float err_max = 0.0f; if (bok) for (uint32_t i = 0; i < V; i++) { const float err_abs = fabsf(gpu[i] - cpu_d[i]); if (!isfinite(gpu[i]) || !isfinite(cpu_d[i]) || err_abs > 2e-3f + 1e-4f * fabsf(cpu_d[i])) bok = false; if (err_abs > err_max) err_max = err_abs; } if (err_max > batch_worst) batch_worst = err_max; batch_scored++; if (!bok || paired != dtok) { fprintf(stderr, "ds4: MTP batch check failed at %u: max|diff|=%g top1 %d/%d\n", idx, (double)err_max, dtok, paired); free(gpu); qwen4_graph_free(g); free(g); free(draft); free(layer_streams); free(streams); free(logits); free(seq); return 1; } } } } printf("Qwen3.8 GPU-vs-CPU: %u tokens (chunk %u), worst max|diff|=%g, top1 agree %u/%u\n", n_seq, chunk, (double)worst, agree, scored); if (draft) { printf("Qwen3.8 GPU-vs-CPU MTP: worst max|diff|=%g, top1 agree %u/%u\n", (double)mtp_worst, mtp_agree, mtp_scored); } if (batch_check) { printf("Qwen3.8 MTP batch check: %u pairs passed, worst max|diff|=%g\n", batch_scored, (double)batch_worst); } free(gpu); qwen4_graph_free(g); free(g); } #else (void)metal; #endif const char *out_env = getenv("DS4_QWEN4_FT_OUT"); if (out_env && out_env[0]) qwen4_dump_f32(out_env, logits, (uint64_t)n_seq * V); if (layer_streams) qwen4_dump_f32(hidden_out, layer_streams, (uint64_t)n_seq * n_trunk * hc_dim); if (draft) { qwen4_dump_f32(mtp_out, draft, (uint64_t)(n_seq - 1u) * V); free(draft); } const float *last = logits + (uint64_t)(n_seq - 1u) * V; print_vec_stats("last-position logits", last, V); int best[8]; for (int i = 0; i < 8; i++) best[i] = -1; for (uint32_t i = 0; i < V; i++) { for (int j = 0; j < 8; j++) { if (best[j] < 0 || last[i] > last[best[j]]) { for (int k = 7; k > j; k--) best[k] = best[k - 1]; best[j] = (int)i; break; } } } printf("top logits after Qwen3.8 CPU pass over %u tokens:\n", n_seq); for (int i = 0; i < 8; i++) { const int id = best[i]; if (id >= 0 && id < vocab->n_vocab) { printf(" %6d %9.4f %.*s\n", id, last[id], (int)vocab->token[id].len, vocab->token[id].ptr); } else { printf(" %6d %9.4f\n", id, id >= 0 ? last[id] : 0.0f); } } free(layer_streams); free(streams); free(logits); free(seq); return 0; } int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt) { if (!engine_legacy_graph_test_supported(e)) return 1; 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_is_qwen4()) { return qwen4_first_token_test(model, vocab, weights, prompt, ds4_backend_uses_graph(e->backend)); } 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, int ctx_size) { #ifdef DS4_NO_GPU (void)e; (void)ctx_size; 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) && !(e->backend == DS4_BACKEND_CUDA && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41)) { 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)) { /* A valid distributed GLM slice can contain only the leading dense * layers. It has no routed weights to stream or cache. */ if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set) { fprintf(stderr, "ds4: SSD streaming layer slice has no routed experts; " "expert cache disabled\n"); e->ssd_streaming_cache_experts = 0; e->ssd_streaming_cache_bytes = 0; return true; } fprintf(stderr, "ds4: SSD streaming auto cache could not measure routed expert size\n"); return false; } uint64_t max_model_experts = 0; if (!ds4_streaming_cacheable_expert_count(&e->weights, &max_model_experts, NULL)) { fprintf(stderr, "ds4: SSD streaming auto cache could not count local cacheable experts\n"); return false; } ds4_ssd_cache_plan plan; uint64_t model_limit = 0; uint32_t cache_percent = 80; #ifdef __APPLE__ /* More cached experts help only if the fixed weights stay resident too. * Share the static-pinning budget, including the requested context. */ model_limit = ds4_streaming_manual_cache_safe_bytes( e->backend, ctx_size > 0 ? ctx_size : 4096, e->prefill_chunk, true); /* GLM's larger fixed tensors and streaming windows need more margin. */ if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) cache_percent = 86; #else (void)ctx_size; #endif if (!ds4_ssd_auto_cache_plan(recommended, cache_percent, model_limit, 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); /* GLM's context-aware fitting pass reserves the graph, active model and * prefill windows before allocation; no fixed-size Metal cache cap is needed. */ #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); ds4_context_memory graph_mem; if (e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set) { const uint32_t normal_layers = glm_graph_normal_layer_count(); const uint32_t layer_end = e->distributed.layers.has_output ? (normal_layers ? normal_layers - 1u : 0u) : e->distributed.layers.end; graph_mem = glm_graph_context_memory_estimate_for_compact_cap_slice( guard_ctx, work_ctx, compact_cap, true, e->distributed.layers.start, layer_end); } else { 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, ds4_model_is_glm53()), 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); } #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 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, uint32_t *layer_start_out) { if (layer_start_out) *layer_start_out = 0; 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 layer_start = DS4_N_LEADING_DENSE; while (layer_start < normal_layers && !glm_stream_resident_decode_layer_supported( &weights->layer[layer_start], layer_start)) { layer_start++; } if (layer_start == normal_layers) return 0; uint32_t n = 0; for (uint32_t il = layer_start; il < normal_layers; il++) { if (!glm_stream_resident_decode_layer_supported(&weights->layer[il], il)) { break; } n++; } if (layer_start_out) *layer_start_out = layer_start; return n; } static bool ds4_glm_streaming_resident_prefix_bytes( const ds4_weights *weights, uint32_t layer_start, 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 = layer_start + 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 bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { g_glm_streaming_full_resident_start = 0; g_glm_streaming_full_resident_layers = 0; g_glm_streaming_runtime_static_map_enabled = false; if (e) { e->ssd_streaming_static_decode_map = false; e->ssd_streaming_decode_map_bytes = 0; e->ssd_streaming_decode_map_extra_bytes = 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_experts != 0 || 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)) { const bool no_routed_slice = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set; if (no_routed_slice && e->ssd_streaming_cache_experts == 0 && e->ssd_streaming_cache_bytes == 0 && e->ssd_streaming_full_layers == 0) { return true; } fprintf(stderr, "ds4: SSD streaming could not measure routed expert size\n"); return false; } uint64_t max_cache_experts_u64 = 0; if (need_expert_bytes && !ds4_streaming_cacheable_expert_count(&e->weights, &max_cache_experts_u64, NULL)) { fprintf(stderr, "ds4: SSD streaming could not count local cacheable experts\n"); return false; } const uint32_t max_cache_experts = max_cache_experts_u64 > UINT32_MAX ? UINT32_MAX : (uint32_t)max_cache_experts_u64; if (e->ssd_streaming_cache_bytes == 0 && max_cache_experts != 0 && e->ssd_streaming_cache_experts > max_cache_experts) { fprintf(stderr, "ds4: SSD streaming expert cache capped from %u to %u " "experts cacheable by this layer slice\n", e->ssd_streaming_cache_experts, max_cache_experts); e->ssd_streaming_cache_experts = max_cache_experts; } uint32_t full_layers = 0; uint64_t full_layers_bytes = 0; bool full_layers_auto = false; uint32_t supported = 0; uint32_t supported_start = 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) { const uint64_t fitted_headroom = total_cache_bytes > per_expert_bytes ? total_cache_bytes - per_expert_bytes : 0; fprintf(stderr, "ds4: SSD streaming prefill reserve reduced from %.2f " "to %.2f GiB to honor the %.2f GiB cache target\n", (double)prefill_headroom_bytes / 1073741824.0, (double)fitted_headroom / 1073741824.0, (double)total_cache_bytes / 1073741824.0); prefill_headroom_bytes = fitted_headroom; } 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, &supported_start); if (!e->ssd_streaming_full_layers_set && e->ssd_streaming_cache_bytes != 0) { /* Cache selected experts across all layers instead of reserving * whole prefix layers. Explicit full-layer requests still work. */ e->ssd_streaming_full_layers = 0; full_layers_auto = true; } } if (e->ssd_streaming_full_layers != 0) { const uint32_t requested = e->ssd_streaming_full_layers; supported = ds4_glm_streaming_supported_resident_prefix_layers( &e->weights, &supported_start); 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, supported_start, 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, supported_start, 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; uint32_t budget = ds4_streaming_cache_experts_for_byte_budget( &e->weights, dynamic_cache_bytes, &budget_expert_bytes); if (max_cache_experts != 0 && budget > max_cache_experts) { budget = max_cache_experts; } if (budget == 0 || budget_expert_bytes == 0) { fprintf(stderr, "ds4: SSD streaming cache target leaves no complete " "expert slot; using direct per-layer reads\n"); budget = 0; budget_expert_bytes = 0; } e->ssd_streaming_cache_experts = budget; e->ssd_streaming_cache_bytes = (uint64_t)budget * budget_expert_bytes; uint64_t effective_total_bytes = glm_graph_saturating_add_u64( prefill_headroom_bytes, full_layers_bytes); effective_total_bytes = glm_graph_saturating_add_u64( effective_total_bytes, e->ssd_streaming_cache_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 cache target %.2f GiB; " "effective %.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)effective_total_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 cache target %.2f GiB; " "effective %.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)effective_total_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_start = supported_start; g_glm_streaming_full_resident_layers = full_layers; return true; #endif } static bool ds4_glm_streaming_runtime_decode_map_bytes( const ds4_weights *weights, uint64_t *bytes_out) { if (bytes_out) *bytes_out = 0; if (!weights || !bytes_out) return false; ds4_model_map_span_vec spans; if (!weights_model_map_decode_runtime_spans(weights, true, true, &spans)) { return false; } *bytes_out = model_map_span_vec_total_bytes(&spans); free(spans.v); return *bytes_out != 0; } static void ds4_engine_fit_glm_streaming_budget( ds4_engine *e, bool load_slice, uint32_t load_layer_start, uint32_t load_layer_end, bool load_output, int ctx_size) { #if defined(DS4_NO_GPU) (void)e; (void)load_slice; (void)load_layer_start; (void)load_layer_end; (void)load_output; (void)ctx_size; #else if (!e || !e->ssd_streaming || DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA || ctx_size <= 0) { return; } e->ssd_streaming_static_decode_map = false; e->ssd_streaming_decode_map_bytes = 0; e->ssd_streaming_decode_map_extra_bytes = 0; g_glm_streaming_runtime_static_map_enabled = false; uint32_t guard_ctx = 0; if (!glm_graph_context_request(ctx_size, &guard_ctx)) return; 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 = load_slice ? glm_graph_context_memory_estimate_for_compact_cap_slice( guard_ctx, work_ctx, compact_cap, true, load_layer_start, load_layer_end) : glm_graph_context_memory_estimate_for_compact_cap( guard_ctx, work_ctx, compact_cap, true); const uint64_t model_bytes = glm_graph_model_bytes_for_guard(&e->model, &e->weights, true, load_slice, load_layer_start, load_layer_end, !load_slice || load_layer_start == 0, !load_slice || load_output); uint64_t budget = 0; if (!glm_graph_memory_guard_budget(model_bytes, load_slice, true, NULL, &budget, NULL, NULL)) { return; } const uint64_t fixed_bytes = glm_graph_saturating_add_u64( model_bytes, ds4_engine_glm_graph_budget(e, graph_mem.total_bytes)); const uint64_t available = budget > fixed_bytes ? budget - fixed_bytes : 0; uint64_t per_expert_bytes = 0; const bool has_routed_experts = ds4_streaming_routed_expert_bytes(&e->weights, &per_expert_bytes) && per_expert_bytes != 0; const uint32_t old_full_layers = e->ssd_streaming_full_layers; const uint32_t old_cache_experts = e->ssd_streaming_cache_experts; const uint64_t old_prefill_headroom = e->ssd_streaming_prefill_headroom_bytes; uint32_t full_layers = 0; uint64_t full_layers_bytes = 0; uint64_t decode_map_bytes = 0; uint64_t decode_map_extra_bytes = 0; bool static_decode_map = false; const bool cache_can_serve_decode = !has_routed_experts || (DS4_N_EXPERT_USED != 0 && old_cache_experts >= DS4_N_EXPERT_USED); const uint64_t minimum_cache_bytes = has_routed_experts && cache_can_serve_decode ? (uint64_t)DS4_N_EXPERT_USED * per_expert_bytes : 0; if (metal_graph_stream_decode_static_map_enabled() && cache_can_serve_decode) { uint32_t candidate_layers = old_full_layers; for (;;) { g_glm_streaming_full_resident_layers = candidate_layers; uint64_t candidate_map_bytes = 0; if (ds4_glm_streaming_runtime_decode_map_bytes( &e->weights, &candidate_map_bytes)) { const uint64_t candidate_extra = candidate_map_bytes > model_bytes ? candidate_map_bytes - model_bytes : 0; uint64_t required = glm_graph_saturating_add_u64( candidate_extra, old_prefill_headroom); required = glm_graph_saturating_add_u64( required, minimum_cache_bytes); if (required <= available) { static_decode_map = true; full_layers = candidate_layers; decode_map_bytes = candidate_map_bytes; decode_map_extra_bytes = candidate_extra; break; } } if (candidate_layers == 0) break; candidate_layers--; } } if (!static_decode_map) { g_glm_streaming_full_resident_layers = 0; full_layers = 0; } if (full_layers != 0 && !ds4_glm_streaming_resident_prefix_bytes( &e->weights, g_glm_streaming_full_resident_start, full_layers, &full_layers_bytes)) { full_layers = 0; full_layers_bytes = 0; static_decode_map = false; decode_map_bytes = 0; decode_map_extra_bytes = 0; g_glm_streaming_full_resident_layers = 0; } uint64_t option_bytes = available > decode_map_extra_bytes ? available - decode_map_extra_bytes : 0; uint64_t prefill_headroom = old_prefill_headroom; if (prefill_headroom > option_bytes) prefill_headroom = option_bytes; option_bytes -= prefill_headroom; uint32_t cache_experts = old_cache_experts; uint64_t dynamic_bytes = ds4_engine_dynamic_expert_cache_bytes(e); if (has_routed_experts && dynamic_bytes > option_bytes) { uint64_t fitted = option_bytes / per_expert_bytes; if (fitted > old_cache_experts) fitted = old_cache_experts; cache_experts = fitted > UINT32_MAX ? UINT32_MAX : (uint32_t)fitted; dynamic_bytes = (uint64_t)cache_experts * per_expert_bytes; } e->ssd_streaming_full_layers = full_layers; e->ssd_streaming_full_layer_bytes = full_layers_bytes; e->ssd_streaming_prefill_headroom_bytes = prefill_headroom; e->ssd_streaming_cache_experts = cache_experts; if (e->ssd_streaming_cache_bytes != 0) { e->ssd_streaming_cache_bytes = dynamic_bytes; } e->ssd_streaming_static_decode_map = static_decode_map; e->ssd_streaming_decode_map_bytes = decode_map_bytes; e->ssd_streaming_decode_map_extra_bytes = decode_map_extra_bytes; g_glm_streaming_full_resident_layers = full_layers; g_glm_streaming_runtime_static_map_enabled = static_decode_map; if (old_full_layers != full_layers || old_cache_experts != cache_experts || old_prefill_headroom != prefill_headroom) { fprintf(stderr, "ds4: GLM SSD streaming request adjusted to fit memory: " "full layers %u -> %u, cache %u -> %u experts, " "prefill reserve %.2f -> %.2f GiB\n", old_full_layers, full_layers, old_cache_experts, cache_experts, glm_graph_bytes_to_gib(old_prefill_headroom), glm_graph_bytes_to_gib(prefill_headroom)); } if (static_decode_map) { fprintf(stderr, "ds4: GLM SSD streaming decode map: global %.2f GiB " "(%.2f GiB beyond the active layer window)\n", glm_graph_bytes_to_gib(decode_map_bytes), glm_graph_bytes_to_gib(decode_map_extra_bytes)); } else { fprintf(stderr, "ds4: GLM SSD streaming decode map: per-layer fallback\n"); } #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); #if defined(__APPLE__) && !defined(DS4_NO_GPU) /* The GPU has not submitted anything yet in sharded mode: pay the * first-submission setup here instead of inside the first prompt. */ if (!ds4_gpu_warm_command_queue()) fprintf(stderr, "ds4: command queue warm failed (continuing)\n"); #endif } /* ========================================================================= * 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 || t == e->model.ngram_tensor) 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; } } if (ds4_model_is_qwen4()) { /* Vision weights are shared; the disk-only n-grams above are not * resident. Session caches/scratch are reserved together per tier. */ out[0] = out[0] > SIZE_MAX - e->vision_model.size ? SIZE_MAX : out[0] + e->vision_model.size; return 0; } /* 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; const uint32_t session_count = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA ? engine_placement_session_count(e) : 1u; 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++) { const size_t per_session = 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); const size_t cache_bytes = engine_size_mul_sat(per_session, session_count); out[il + 1] = out[il + 1] > SIZE_MAX - cache_bytes ? SIZE_MAX : out[il + 1] + cache_bytes; } } else { /* Fallback: 128 MiB per layer as a static estimate. */ const size_t fallback_per_layer = engine_size_mul_sat( (size_t)128ull * 1024ull * 1024ull, session_count); 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); } uint64_t ds4_test_glm_memory_guard_default_budget( uint64_t host_bytes, uint64_t model_bytes, bool glm53) { const double reserve_gib = glm_graph_memory_guard_default_reserve_gib( host_bytes, model_bytes, glm53); return glm_graph_memory_guard_budget_bytes( host_bytes, 0, 0.99, reserve_gib); } int ds4_test_glm_memory_guard_disabled(void) { return glm_graph_memory_guard_disabled() ? 1 : 0; } 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->placement_session_count_hint = 1; 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_compute_glm_entry_bytes_sum_with_sessions( const ds4_test_fake_tensor *tensors, int n_tensors, int placement_ctx_hint, int placement_session_count_hint) { const ds4_shape saved_shape = g_ds4_shape; g_ds4_shape = DS4_SHAPE_GLM52; ds4_engine eng; if (ds4_test_make_engine(&eng, tensors, n_tensors, placement_ctx_hint) != 0) { g_ds4_shape = saved_shape; return 0; } eng.placement_session_count_hint = placement_session_count_hint; 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); g_ds4_shape = saved_shape; return sum; } 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_qwen4_placement(uint64_t budget, int sessions, size_t *weights, size_t *runtime) { const ds4_shape saved_shape = g_ds4_shape; g_ds4_shape = DS4_SHAPE_QWEN4_EXP; const uint64_t gib = UINT64_C(1073741824); const ds4_test_fake_tensor tensors[] = { {"ngram.weight", 96u * gib}, {"token_embd.weight", gib}, {"blk.0.ffn_gate_exps.weight", 40u * gib}, {"output.weight", gib}, }; ds4_engine eng; if (ds4_test_make_engine(&eng, tensors, 4, 4096) != 0) { g_ds4_shape = saved_shape; return -1; } eng.model.ngram_tensor = &eng.model.tensors[0]; eng.vision_model.size = gib / 2u; eng.placement_session_count_hint = sessions; size_t entries[DS4_MAX_LAYER + 2]; engine_compute_entry_bytes(&eng, entries); *weights = 0; for (uint32_t i = 0; i < DS4_N_LAYER + 2u; i++) *weights += entries[i]; *runtime = engine_per_tier_graph_overhead_bytes(&eng); ds4_gpu_config cfg = {0}; cfg.n_gpus = 1; cfg.vram_bytes[0] = budget; const int rc = engine_classify_multi_tier(&eng, &cfg); const int fits = rc == 0 && !eng.multi_tier; free(eng.model.tensors); g_ds4_shape = saved_shape; return fits; } 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); static bool engine_warm_full_model(const ds4_engine_options *opt) { /* Partial/streaming setups use their own bounded preparation. Do not * fault every unowned expert merely because the whole GGUF is mapped. */ return opt->warm_weights && !opt->inspect_only && !opt->ssd_streaming && !opt->cuda_tensor_parallel && opt->tp.role == DS4_TP_NONE && opt->distributed.role == DS4_DISTRIBUTED_NONE; } #ifdef DS4_HAS_DEEPSEEK41_GPU static bool ds41_memory_admit_for_host(ds4_engine *e, uint64_t graph_bytes, bool fit_cache, uint64_t host, uint64_t recommended) { const uint64_t gib = UINT64_C(1073741824); uint64_t budget = host / 8u * 7u; if (!host || !recommended) { fprintf(stderr, "ds4: cannot determine a safe V4.1 memory budget\n"); return false; } if (budget > recommended) budget = recommended; uint64_t weights = g_tp_shard_model_bytes ? g_tp_shard_model_bytes : e->model.size; if (e->ssd_streaming && !weights_streaming_non_routed_bytes(&e->weights, &weights)) return false; weights = ds4_add_sat_u64(weights, e->vision_model.size); const uint64_t fixed = ds4_add_sat_u64(weights, ds4_add_sat_u64(graph_bytes, 2u * gib + e->ssd_streaming_prefill_headroom_bytes)); uint64_t expert = 0; if (e->ssd_streaming && !ds4_streaming_routed_expert_bytes(&e->weights, &expert)) return false; if (fixed >= budget || (expert && budget - fixed < expert)) { fprintf(stderr, "ds4: V4.1 needs %.2f GiB before the expert cache; safe budget %.2f GiB. " "Use --ssd-streaming or a smaller context.\n", ds4_bytes_to_gib(fixed), ds4_bytes_to_gib(budget)); return false; } if (expert && e->ssd_streaming_cache_experts > (budget - fixed) / expert) { if (!fit_cache) { fprintf(stderr, "ds4: no room for another V4.1 session; reduce context or the SSD expert cache\n"); return false; } const uint32_t count = (uint32_t)((budget - fixed) / expert); fprintf(stderr, "ds4: V4.1 SSD cache fitted from %u to %u experts for context/runtime headroom\n", e->ssd_streaming_cache_experts, count); e->ssd_streaming_cache_experts = count; e->ssd_streaming_cache_bytes = (uint64_t)count * expert; } return true; } static bool ds41_memory_admit(ds4_engine *e, uint64_t graph_bytes, bool fit_cache) { uint64_t host = glm_graph_host_memory_bytes(); #ifdef __linux__ const long pages = sysconf(_SC_PHYS_PAGES), page_size = sysconf(_SC_PAGESIZE); if (pages > 0 && page_size > 0 && (uint64_t)pages <= UINT64_MAX / (uint64_t)page_size) host = (uint64_t)pages * (uint64_t)page_size; #endif return ds41_memory_admit_for_host(e, graph_bytes, fit_cache, host, ds4_gpu_recommended_working_set_size()); } #endif 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)); #if defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) g_glm_rocm_guard_available_baseline = 0; (void)ds4_linux_nonmovable_memory(&g_glm_rocm_guard_available_baseline); #endif e->model.fd = -1; e->mtp_model.fd = -1; e->vision_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->dspark_exact_sampling = opt->dspark_exact_sampling; 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; if (opt->vision_path && opt->vision_path[0]) { if (opt->backend != DS4_BACKEND_METAL && opt->backend != DS4_BACKEND_CUDA) { fprintf(stderr, "ds4: GLM-5.3 vision requires --metal, --cuda, or --rocm\n"); free(e); *out = NULL; return 1; } #ifdef DS4_NO_GPU fprintf(stderr, "ds4: this build does not include a GPU vision backend\n"); free(e); *out = NULL; return 1; #endif if (opt->backend == DS4_BACKEND_CUDA && gpu_cfg && gpu_cfg->n_gpus > 1) { fprintf(stderr, "ds4: GLM-5.3 CUDA/ROCm vision currently requires one GPU\n"); free(e); *out = NULL; return 1; } } 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 = e->backend == DS4_BACKEND_METAL ? 0.6f : 0.7f; } e->dspark_confidence_threshold_set = opt->dspark_confidence_threshold_set; 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-model 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->placement_session_count_hint = opt->placement_session_count_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; bool load_output_optional = false; if (opt->distributed.role != DS4_DISTRIBUTED_NONE && opt->distributed.layers.set) { load_slice = true; load_layer_start = opt->distributed.layers.start; load_layer_end = opt->distributed.layers.end; load_output = opt->distributed.layers.has_output; /* A coordinator may need to apply the output head locally when the * last worker advertises N:M rather than N:output. Bind it when the * local GGUF contains it, without requiring split GGUFs to do so. */ load_output_optional = opt->distributed.role == DS4_DISTRIBUTED_COORDINATOR; } 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 (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && opt->vision_path && opt->vision_path[0] && e->backend != DS4_BACKEND_METAL) { fprintf(stderr, "ds4: V4.1 vision requires Metal\n"); ds4_engine_close(e); *out = NULL; return 1; } config_validate_model(&e->model); if (ds4_model_is_qwen4() && !opt->inspect_only) { const bool backend_ok = #ifdef DS4_HAS_QWEN4_GPU e->backend == DS4_BACKEND_METAL || e->backend == DS4_BACKEND_CUDA || #endif (opt->first_token_test && e->backend == DS4_BACKEND_CPU); if (!backend_ok || opt->tp.role != DS4_TP_NONE || opt->cuda_tensor_parallel || (gpu_cfg && gpu_cfg->n_gpus > 1) || opt->distributed.role != DS4_DISTRIBUTED_NONE || load_slice || e->ssd_streaming || opt->dspark || e->power_percent != 100 || (opt->mtp_path && opt->mtp_path[0])) { fprintf(stderr, "ds4: Qwen3.8 requires Metal or single-GPU CUDA (or --cpu --first-token-test); " "tensor parallelism, pipeline execution, SSD streaming, DSpark, " "external MTP models and power throttling are not supported\n"); ds4_engine_close(e); *out = NULL; return 1; } } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && !opt->inspect_only) { const bool supported = (e->backend == DS4_BACKEND_METAL || #if defined(DS4_HAS_DEEPSEEK41_GPU) && !defined(__APPLE__) (e->backend == DS4_BACKEND_CUDA && !opt->cuda_tensor_parallel && (!gpu_cfg || gpu_cfg->n_gpus <= 1)) || #endif false) && opt->distributed.role == DS4_DISTRIBUTED_NONE && !load_slice && !opt->dspark && !opt->glm_mtp && !opt->first_token_test && !opt->metal_graph_test && (!opt->mtp_path || !opt->mtp_path[0]) && (!opt->directional_steering_file || !opt->directional_steering_file[0]) && e->power_percent == 100 && opt->context_size <= 1048576; if (!supported) { fprintf(stderr, "ds4: V4.1 requires Metal or single-GPU CUDA per rank (optional network tensor parallelism); " "DSpark, steering and legacy diagnostics are not supported (maximum context 1048576)\n"); ds4_engine_close(e); *out = NULL; return 1; } /* The agent may change directory before creating its first session. * Engram readers reopen this file and check it against model.fd. */ e->model_path = realpath(opt->model_path, NULL); if (!e->model_path) { fprintf(stderr, "ds4: cannot resolve V4.1 model path: %s\n", strerror(errno)); ds4_engine_close(e); *out = NULL; return 1; } } if (engine_warm_full_model(opt) && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) model_warm_weights(&e->model); if (opt->vision_path && opt->vision_path[0]) { if (!ds4_model_is_glm53() && !g_ds4_flash_vision_exp && !ds4_model_is_qwen4() && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41) { fprintf(stderr, "ds4: --vision requires GLM-5.3, Qwen3.8-Flash-Next or the pinned " "DeepSeek V4 Flash Vision-Exp or V4.1 Flash model\n"); ds4_engine_close(e); *out = NULL; return 1; } #ifdef DS4_NO_GPU fprintf(stderr, "ds4: this build does not include a GPU vision backend\n"); ds4_engine_close(e); *out = NULL; return 1; #else model_open(&e->vision_model, opt->vision_path, true, false); if (ds4_model_is_glm53()) { vision_weights_bind(&e->vision_weights, &e->vision_model); e->vision_image_token = (int)required_u32( &e->vision_model, "glm5-next-vision.image_token_id"); e->vision_start_token = (int)required_u32( &e->vision_model, "glm5-next-vision.image_start_token_id"); e->vision_end_token = (int)required_u32( &e->vision_model, "glm5-next-vision.image_end_token_id"); if (e->vision_image_token != 154854 || e->vision_start_token != 154830 || e->vision_end_token != 154831) { ds4_die("unexpected GLM-5.3 vision token IDs"); } e->vision_kind = DS4_VISION_GLM53; } else if (ds4_model_is_qwen4()) { qwen4_vision_weights_bind(&e->qwen4_vision_weights, &e->vision_model); config_expect_u32("vision projection_dim", e->qwen4_vision_weights.n_out, DS4_N_EMBD); e->vision_kind = DS4_VISION_QWEN4; } else { deepseek4_vision_weights_bind( &e->deepseek4_vision_weights, &e->vision_model); e->vision_start_token = DS4_N_VOCAB + DS4_DEEPSEEK4_IMAGE_START; e->vision_image_token = DS4_N_VOCAB + DS4_DEEPSEEK4_IMAGE; e->vision_end_token = DS4_N_VOCAB + DS4_DEEPSEEK4_IMAGE_END; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) e->vision_start_token = e->vision_image_token = e->vision_end_token = 129264; e->vision_kind = DS4_VISION_DEEPSEEK4; } e->vision_ready = true; #endif } if (load_slice && load_layer_end == UINT32_MAX) { const uint32_t normal_layers = ds4_model_normal_layer_count(); if (normal_layers == 0) { fprintf(stderr, "ds4: model reports no executable transformer layers\n"); ds4_engine_close(e); *out = NULL; return 1; } load_layer_end = normal_layers - 1u; } if (e->distributed.role != DS4_DISTRIBUTED_NONE && e->distributed.layers.set) { e->distributed.layers.end = load_layer_end; } 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, load_output_optional); if (e->model.ngram_tensor) { fprintf(stderr, "ds4: Qwen BF16 n-grams: %.2f GiB, disk reads only\n", (double)e->model.ngram_tensor->bytes / (1024.0 * 1024.0 * 1024.0)); } if (e->vision_ready && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (!e->weights.layer[il].ffn_exp_probs_vl) { fprintf(stderr, "ds4: V4.1 vision requires visual router biases in the language GGUF\n"); ds4_engine_close(e); *out = NULL; return 1; } } } /* 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; #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (tp_shard && (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41 || (gpu_cfg && gpu_cfg->n_gpus > 1) || opt->quality)) { fprintf(stderr, "ds4: network CUDA TP currently requires V4.1 Q2, one GPU per rank, no quality mode\n"); ds4_engine_close(e); *out = NULL; return 1; } if (tp_shard) { for (uint32_t il = 0; il < DS4_N_LAYER; il++) { if (e->weights.layer[il].ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || e->weights.layer[il].ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || e->weights.layer[il].ffn_down_exps->type != DS4_TENSOR_Q2_K) { fprintf(stderr, "ds4: network CUDA TP requires IQ2_XXS gate/up and Q2_K down experts\n"); ds4_engine_close(e); *out = NULL; return 1; } } } #endif if (tp_shard && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { uint32_t bad_layer = 0; uint32_t bad_type = 0; if (!glm_tp_validate_ownership_kernels(&e->weights, &bad_layer, &bad_type)) { fprintf(stderr, "ds4: GLM tensor parallelism lacks ownership-aware " "kernels for routed expert type %u in layer %u\n", bad_type, bad_layer); ds4_engine_close(e); *out = NULL; return 1; } } 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 #ifdef DS4_HAS_DEEPSEEK41_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && !opt->inspect_only) { const uint32_t ctx = opt->context_size > 0 ? (uint32_t)opt->context_size : 4096; const uint32_t sessions = e->placement_session_count_hint > 0 ? (uint32_t)e->placement_session_count_hint : 1; if (!ds41_memory_admit(e, ds4_mul_sat_u64(ds41_graph_bytes(ctx), sessions), true)) { ds4_engine_close(e); *out = NULL; return 1; } if (engine_warm_full_model(opt)) model_warm_weights(&e->model); } #endif if (ds4_model_is_qwen4()) { if (opt->inspect_only) { *out = e; return 0; } if (opt->first_token_test) { if (e->backend != DS4_BACKEND_CPU && e->backend != DS4_BACKEND_METAL && e->backend != DS4_BACKEND_CUDA) { fprintf(stderr, "ds4: Qwen3.8 first-token test needs CPU, Metal or CUDA\n"); ds4_engine_close(e); *out = NULL; return 1; } #ifndef DS4_NO_GPU if (ds4_backend_uses_graph(e->backend)) { bool ready = ds4_gpu_init() && ds4_gpu_set_model_map(e->model.map, e->model.size); if (ready && getenv("DS4_QWEN4_GPU")) { /* Use the normal file-backed CUDA cache, not per-tensor * copies from pageable mmap during the first forward. */ (void)ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map); ready = accelerator_cache_model_tensors(e->backend, &e->model, NULL, NULL, 0); } if (!ready) { fprintf(stderr, "ds4: Qwen3.8 test model map setup failed\n"); ds4_engine_close(e); *out = NULL; return 1; } } #endif vocab_load(&e->vocab, &e->model); *out = e; return 0; } } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (opt->inspect_only) { *out = e; return 0; } if (!ds4_model_is_glm53() && ((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-model is not supported for GLM 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 const bool rocm_full_model_requires_streaming = e->backend == DS4_BACKEND_CUDA && !e->ssd_streaming && !load_slice && !ds4_model_is_glm53(); if (rocm_full_model_requires_streaming) { glm_backend_supported = false; } #endif if (!glm_backend_supported) { #ifdef DS4_ROCM_BUILD if (rocm_full_model_requires_streaming) { fprintf(stderr, "ds4: full-model GLM DSA ROCm inference requires " "--ssd-streaming; distributed layer slices can run " "fully resident\n"); } else { fprintf(stderr, "ds4: GLM 5.2 inference requires the ROCm graph " "backend; 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 || (load_output_optional && weights_have_output_head( &e->weights)), 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->vision_kind == DS4_VISION_QWEN4) { e->vision_start_token = vocab_lookup(&e->vocab, "<|vision_start|>"); e->vision_image_token = vocab_lookup(&e->vocab, "<|image_pad|>"); e->vision_end_token = vocab_lookup(&e->vocab, "<|vision_end|>"); } if (opt->glm_mtp && ((DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA && !ds4_model_is_qwen4()) || DS4_N_NEXTN_PREDICT == 0)) { fprintf(stderr, "ds4: --mtp requires a model with embedded MTP weights; " "use --mtp-model FILE for external support\n"); ds4_engine_close(e); *out = NULL; return 1; } 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 (!support_model_checkpoint_compatible(&e->mtp_model)) { ds4_engine_close(e); *out = NULL; return 1; } 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-model 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 (!support_model_checkpoint_compatible(&e->mtp_model)) { ds4_engine_close(e); *out = NULL; return 1; } 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); if (e->dspark && !e->quality && !e->dspark_strict) { fprintf(stderr, "ds4: DSpark direct verifier-state commits enabled; " "output may differ from one-token decode due " "to batched floating-point operation order\n"); } } else { fprintf(stderr, "ds4: unsupported --mtp-model 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 (e->backend == DS4_BACKEND_METAL && 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); #ifdef DS4_ROCM_BUILD /* * The ROCm Q8 decode selector must know that a multi-tier model * is GLM before any layer dispatch. The single-tier path sets * the same model-family state below. */ ds4_gpu_set_glm_model( DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA); #endif (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, opt->context_size)) { ds4_engine_close(e); *out = NULL; return 1; } if (!ds4_engine_configure_streaming_cache_budget(e)) { ds4_engine_close(e); *out = NULL; return 1; } #ifdef DS4_HAS_DEEPSEEK41_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { const uint32_t ctx = opt->context_size > 0 ? (uint32_t)opt->context_size : 4096; const uint32_t sessions = e->placement_session_count_hint > 0 ? (uint32_t)e->placement_session_count_hint : 1; if (!ds41_memory_admit(e, ds4_mul_sat_u64(ds41_graph_bytes(ctx), sessions), true)) { ds4_engine_close(e); *out = NULL; return 1; } } #endif ds4_engine_fit_glm_streaming_budget( e, load_slice, load_layer_start, load_layer_end, load_output || (load_output_optional && weights_have_output_head(&e->weights)), opt->context_size); if (!ds4_engine_glm_streaming_memory_guard( e, load_slice, load_layer_start, load_layer_end, load_output || (load_output_optional && weights_have_output_head(&e->weights)), 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 defined(__APPLE__) /* Keep the weights used by every token from competing with streamed * experts in the file cache. These bytes are already in the model * budget; munmap in model_close releases the locks. */ if (e->ssd_streaming && !load_slice && !tp_shard && getenv("DS4_METAL_DISABLE_STREAMING_STATIC_LOCK") == NULL) { ds4_model_map_span_vec spans; uint64_t static_bytes = 0; const uint64_t budget = ds4_streaming_manual_cache_safe_bytes( e->backend, opt->context_size, e->prefill_chunk, true); const uint64_t experts = ds4_add_sat_u64( ds4_engine_dynamic_expert_cache_bytes(e), ds4_add_sat_u64(e->ssd_streaming_prefill_headroom_bytes, e->ssd_streaming_full_layer_bytes)); const bool fits = weights_streaming_non_routed_bytes(&e->weights, &static_bytes) && static_bytes <= budget && experts <= budget - static_bytes; if (!fits) { fprintf(stderr, "ds4: Metal SSD static weights remain pageable" " to preserve runtime headroom\n"); } else if (weights_model_map_decode_static_spans( &e->weights, true, true, &spans)) { uint64_t locked = 0, failed = 0; bool can_lock = true; const size_t page = (size_t)getpagesize(); for (uint32_t i = 0; i < spans.len; i++) { const uint64_t lo = spans.v[i].off / page * page; const uint64_t hi = spans.v[i].end; if (hi > e->model.size || hi <= lo) continue; /* Bound each VM operation, including very large heads. */ for (uint64_t off = lo; off < hi;) { uint64_t len = hi - off; if (len > 256ull * 1024 * 1024) len = 256ull * 1024 * 1024; if (can_lock && mlock(e->model.map + off, (size_t)len) == 0) { locked += len; } else { can_lock = false; failed += len; } off += len; } } fprintf(stderr, "ds4: Metal SSD static weights locked %.2f GiB" "; pageable %.2f GiB\n", locked / 1073741824.0, failed / 1073741824.0); free(spans.v); } } #endif 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); #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (e->backend == DS4_BACKEND_CUDA && !load_slice && !tp_shard && !e->ssd_streaming && !ds4_model_is_qwen4()) { (void)ds4_gpu_build_derived_artifacts(e->model.map, e->model.size, opt->model_path); } else if (e->backend == DS4_BACKEND_CUDA && tp_shard) { (void)ds4_gpu_build_derived_artifacts_shard(e->model.map, e->model.size, e->model.file_size, opt->model_path, (uint32_t)tp_shard_rank); } #endif 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 || (load_output_optional && weights_have_output_head(&e->weights))); ds4_model_map_span_vec spans; bool spans_ok = false; if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { spans_ok = weights_model_map_decode_static_spans(&e->weights, true, true, &spans); } else if (load_slice) { if (e->ssd_streaming_static_decode_map) { spans_ok = weights_model_map_decode_runtime_slice_spans( &e->weights, load_layer_start, load_layer_end, load_layer_start == 0, map_output, &spans); } else { spans_ok = weights_model_map_decode_static_slice_spans( &e->weights, load_layer_start, load_layer_end, load_layer_start == 0, 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 " "%snon-routed layers %u:%s (%u spans, %.2f GiB tensor span)\n", ds4_backend_name(e->backend), load_layer_start == 0 ? "token + " : "", load_layer_start, load_end, spans.len, (double)span_bytes / 1073741824.0); } else { fprintf(stderr, "ds4: SSD streaming initial %s model map (%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 || (load_output_optional && weights_have_output_head(&e->weights)); 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 && e->backend == DS4_BACKEND_METAL) { 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); bool support_uses_secondary_rocm_cache = false; #ifdef DS4_ROCM_BUILD /* The generic range cache is already keyed by model map. Keep the * resident target ranges when adding DSpark support on gfx1151. */ support_uses_secondary_rocm_cache = e->support_kind == DS4_SUPPORT_DSPARK && ds4_gpu_dspark_gfx1151_fast_path() != 0; #endif if (support_model_runtime_ready && !support_uses_secondary_rocm_cache && !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); } if (e->vision_ready) { #if defined(__APPLE__) e->vision_map_ready = ds4_gpu_set_model_map_range( e->vision_model.map, e->vision_model.size, e->vision_model.tensor_data_pos, e->vision_model.size - e->vision_model.tensor_data_pos, e->vision_model.max_tensor_bytes) != 0; #else e->vision_map_ready = ds4_gpu_set_aux_model_map_range( e->vision_model.map, e->vision_model.size, e->vision_model.tensor_data_pos, e->vision_model.size - e->vision_model.tensor_data_pos) != 0; #endif if (!e->vision_map_ready) { fprintf(stderr, "ds4: %s failed to map the GLM-5.3 vision encoder\n", ds4_backend_name(e->backend)); ds4_engine_close(e); *out = NULL; return 1; } } 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; } bool ds4_engine_is_glm53(ds4_engine *e) { (void)e; return ds4_model_is_glm53(); } bool ds4_engine_is_qwen4(ds4_engine *e) { (void)e; return ds4_model_is_qwen4(); } /* The official template's default effort is xhigh; medium adds no text. */ const char *ds4_qwen4_reasoning_effort_text(ds4_think_mode mode) { switch (mode) { case DS4_THINK_HIGH: case DS4_THINK_MAX: return DS4_QWEN4_REASONING_XHIGH; case DS4_THINK_LOW: return DS4_QWEN4_REASONING_LOW; case DS4_THINK_MEDIUM: return NULL; /* the template sets no instruction for medium */ case DS4_THINK_NONE: return NULL; } return NULL; } /* Decode gate firing schedule for the TP transport (see ds4_tp_identity). * Resident GLM splits attention and FFN on sparse layers. Streaming keeps * attention replicated and exchanges only the routed FFN partial. */ void ds4_engine_tp_gate_schedule(ds4_engine *e, uint32_t *start, uint32_t *step, uint32_t *per_token, uint64_t mask[DS4_TP_GATE_MASK_WORDS]) { memset(mask, 0, sizeof(uint64_t) * DS4_TP_GATE_MASK_WORDS); if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { const uint32_t sparse_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT - DS4_N_LEADING_DENSE; if (e && e->ssd_streaming) { *start = DS4_N_LEADING_DENSE * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; *step = DS4_TP_GATES_PER_LAYER; *per_token = sparse_layers; } else if (ds4_model_is_glm53()) { uint32_t count = 0; const uint32_t normal_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; for (uint32_t il = DS4_N_LEADING_DENSE; il < normal_layers; il++) { if (!ds4_glm53_layer_is_kda(il)) { const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN; mask[slot / 64u] |= UINT64_C(1) << (slot % 64u); count++; } const uint32_t slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; mask[slot / 64u] |= UINT64_C(1) << (slot % 64u); count++; } *start = DS4_N_LEADING_DENSE * DS4_TP_GATES_PER_LAYER; *step = 1; *per_token = count; } else { *start = DS4_N_LEADING_DENSE * DS4_TP_GATES_PER_LAYER; *step = 1; *per_token = sparse_layers * DS4_TP_GATES_PER_LAYER; } } 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; } bool ds4_engine_has_vision(ds4_engine *e) { return e && e->vision_ready; } void ds4_vision_embedding_free(ds4_vision_embedding *embedding) { if (!embedding) return; free(embedding->data); memset(embedding, 0, sizeof(*embedding)); } #define DS4_VISION_LAYOUT_DEEPSEEK4_NATURAL 1u #ifndef DS4_NO_GPU static float ds4_vision_bf16_to_f32(uint16_t value) { uint32_t bits = (uint32_t)value << 16; float result; memcpy(&result, &bits, sizeof(result)); return result; } static const uint16_t *ds4_deepseek4_vision_vector( const ds4_engine *e, uint64_t offset) { const uint64_t bytes = (uint64_t)DS4_N_EMBD * sizeof(uint16_t); if (!e || !e->vision_model.map || offset > e->vision_model.size || bytes > e->vision_model.size - offset) return NULL; return (const uint16_t *)((const uint8_t *)e->vision_model.map + offset); } static int ds4_prompt_append_deepseek4_vision( ds4_engine *e, ds4_tokens *tokens, ds4_vision_span *span, ds4_vision_embedding *embedding, char *error, size_t error_cap) { if (embedding->layout != DS4_VISION_LAYOUT_DEEPSEEK4_NATURAL || embedding->grid_height == 0 || embedding->grid_width == 0 || embedding->token_count != embedding->grid_height * embedding->grid_width) { if (error && error_cap) snprintf(error, error_cap, "invalid DeepSeek vision embedding layout"); return 0; } ds4_deepseek4_image_layout layout = {0}; const bool v41 = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41; const bool layout_ok = v41 ? ds4_deepseek41_image_layout_build( &layout, embedding->grid_height, embedding->grid_width, error, error_cap) : ds4_deepseek4_image_layout_build( &layout, embedding->grid_height, embedding->grid_width, (uint32_t)tokens->len, error, error_cap); if (!layout_ok) return 0; if (layout.image_count != embedding->token_count || layout.token_count > (uint32_t)(INT_MAX - tokens->len)) { ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "DeepSeek vision prompt is too large"); return 0; } const uint64_t values = (uint64_t)layout.token_count * DS4_N_EMBD; if (values > SIZE_MAX / sizeof(float)) { ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "DeepSeek vision embedding is too large"); return 0; } float *block = malloc((size_t)values * sizeof(float)); if (!block) { ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "unable to allocate DeepSeek image block"); return 0; } const ds4_deepseek4_vision_weights *w = &e->deepseek4_vision_weights; const uint16_t *sentinel[5] = { ds4_deepseek4_vision_vector(e, w->image_start), v41 ? NULL : ds4_deepseek4_vision_vector(e, w->image_pad), NULL, ds4_deepseek4_vision_vector(e, w->image_newline), ds4_deepseek4_vision_vector(e, w->image_end), }; if (!sentinel[DS4_DEEPSEEK4_IMAGE_START] || (!v41 && !sentinel[DS4_DEEPSEEK4_IMAGE_PAD]) || !sentinel[DS4_DEEPSEEK4_IMAGE_NEWLINE] || !sentinel[DS4_DEEPSEEK4_IMAGE_END]) { free(block); ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "DeepSeek vision sentinels are not mapped"); return 0; } uint32_t image_pos = 0; for (uint32_t row = 0; row < layout.token_count; row++) { const uint8_t type = layout.types[row]; float *dst = block + (uint64_t)row * DS4_N_EMBD; if (type == DS4_DEEPSEEK4_IMAGE) { if (image_pos >= layout.image_count) { free(block); ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "invalid DeepSeek image embedding permutation"); return 0; } const uint32_t source = layout.perm[image_pos++]; memcpy(dst, embedding->data + (uint64_t)source * DS4_N_EMBD, DS4_N_EMBD * sizeof(float)); } else { const uint16_t *src = sentinel[type]; for (uint32_t d = 0; d < DS4_N_EMBD; d++) dst[d] = ds4_vision_bf16_to_f32(src[d]); } } if (image_pos != layout.image_count) { free(block); ds4_deepseek4_image_layout_free(&layout); if (error && error_cap) snprintf(error, error_cap, "invalid DeepSeek image embedding count"); return 0; } const uint32_t token_start = (uint32_t)tokens->len; for (uint32_t i = 0; i < layout.token_count; i++) { ds4_tokens_push(tokens, v41 ? e->vision_image_token : (int)(DS4_N_VOCAB + layout.types[i])); } free(embedding->data); embedding->data = block; embedding->token_count = layout.token_count; embedding->layout = 0; embedding->grid_width = 0; embedding->grid_height = 0; memset(span, 0, sizeof(*span)); span->token_start = token_start; span->embedding = *embedding; memset(embedding, 0, sizeof(*embedding)); ds4_deepseek4_image_layout_free(&layout); return 1; } #else static int ds4_prompt_append_deepseek4_vision( ds4_engine *e, ds4_tokens *tokens, ds4_vision_span *span, ds4_vision_embedding *embedding, char *error, size_t error_cap) { (void)e; (void)tokens; (void)span; (void)embedding; if (error && error_cap) snprintf(error, error_cap, "DeepSeek vision requires a GPU backend"); return 0; } #endif int ds4_prompt_append_vision( ds4_engine *e, ds4_tokens *tokens, ds4_vision_span *span, ds4_vision_embedding *embedding, char *error, size_t error_cap) { if (!e || !tokens || !span || !embedding || !embedding->data || embedding->token_count == 0 || !e->vision_ready) { if (error && error_cap) snprintf(error, error_cap, "invalid vision prompt input"); return 0; } if (e->vision_kind == DS4_VISION_DEEPSEEK4) { return ds4_prompt_append_deepseek4_vision( e, tokens, span, embedding, error, error_cap); } if ((uint64_t)tokens->len + embedding->token_count + 2u > INT_MAX) { if (error && error_cap) snprintf(error, error_cap, "vision prompt is too large"); return 0; } memset(span, 0, sizeof(*span)); ds4_tokens_push(tokens, e->vision_start_token); span->token_start = (uint32_t)tokens->len; for (uint32_t i = 0; i < embedding->token_count; i++) { ds4_tokens_push(tokens, e->vision_image_token); } ds4_tokens_push(tokens, e->vision_end_token); span->embedding = *embedding; memset(embedding, 0, sizeof(*embedding)); return 1; } int ds4_chat_append_multimodal_message( ds4_engine *e, ds4_tokens *tokens, const char *role, const char *const *text_parts, ds4_vision_embedding *embeddings, size_t image_count, ds4_vision_span *spans, char *error, size_t error_cap) { if (!e || !tokens || !role || !text_parts || (image_count != 0 && (!embeddings || !spans))) { if (error && error_cap) snprintf(error, error_cap, "invalid multimodal message"); return 0; } const bool tool = !strcmp(role, "tool") || !strcmp(role, "function"); const bool user = !strcmp(role, "user"); if (ds4_model_is_qwen4()) { if (!tool && !user) { if (error && error_cap) snprintf(error, error_cap, "multimodal messages require a user or tool role"); return 0; } for (size_t i = 0; i < image_count; i++) { if (!embeddings[i].data || embeddings[i].token_count == 0) { if (error && error_cap) snprintf(error, error_cap, "invalid image embedding"); return 0; } } const int old_len = tokens->len; ds4_vocab *vocab = &e->vocab; qwen4_chat_open(vocab, "user", tokens); if (tool) bpe_tokenize_text(vocab, "\n", tokens); size_t moved = 0; for (size_t i = 0; i <= image_count; i++) { const char *text = text_parts[i] ? text_parts[i] : ""; if (tool) bpe_tokenize_tool_response_text(vocab, text, tokens); else bpe_tokenize_text(vocab, text, tokens); if (i == image_count) break; if (!ds4_prompt_append_vision(e, tokens, &spans[i], &embeddings[i], error, error_cap)) { for (size_t j = 0; j < moved; j++) { embeddings[j] = spans[j].embedding; memset(&spans[j], 0, sizeof(spans[j])); } tokens->len = old_len; return 0; } moved++; } if (tool) bpe_tokenize_text(vocab, "\n", tokens); qwen4_chat_close(vocab, tokens); return 1; } if (!tool && !user) { if (error && error_cap) snprintf(error, error_cap, "multimodal messages require a supported user or tool role"); return 0; } /* Text-only tool results do not require vision support and must retain * the model's normal wrapper and payload escaping. */ if (image_count == 0) { ds4_chat_append_message(e, tokens, role, text_parts[0]); return 1; } if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_QWEN4_EXP && e->vision_kind != DS4_VISION_DEEPSEEK4) { if (error && error_cap) snprintf(error, error_cap, "model does not support image messages"); return 0; } for (size_t i = 0; i < image_count; i++) { if (!embeddings[i].data || embeddings[i].token_count == 0) { if (error && error_cap) snprintf(error, error_cap, "invalid image embedding"); return 0; } } const int old_len = tokens->len; ds4_vocab *vocab = &e->vocab; const bool glm_tool = tool && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; if (glm_tool) { if (vocab->observation_id >= 0) token_vec_push(tokens, vocab->observation_id); tokenize_rendered_chat_vocab(vocab, "", tokens); } else { token_vec_push(tokens, vocab->user_id); if (tool) bpe_tokenize_text(vocab, "", tokens); } size_t moved = 0; for (size_t i = 0; i <= image_count; i++) { const char *text = text_parts[i] ? text_parts[i] : ""; if (glm_tool) bpe_tokenize_tool_response_text(vocab, text, tokens); else if (tool) bpe_tokenize_tool_result_text(vocab, text, tokens); else bpe_tokenize_text(vocab, text, tokens); if (i == image_count) break; if (!ds4_prompt_append_vision(e, tokens, &spans[i], &embeddings[i], error, error_cap)) { for (size_t j = 0; j < moved; j++) { embeddings[j] = spans[j].embedding; memset(&spans[j], 0, sizeof(spans[j])); } tokens->len = old_len; return 0; } moved++; } if (glm_tool) tokenize_rendered_chat_vocab(vocab, "", tokens); else if (tool) bpe_tokenize_text(vocab, "", tokens); return 1; } static int ds4_engine_vision_encode_image( ds4_engine *e, const ds4_image *image, ds4_vision_embedding *out, char *error, size_t error_cap) { if (!e || !image || !out || !e->vision_ready) { if (error && error_cap) snprintf(error, error_cap, "vision encoder is not loaded"); return 0; } memset(out, 0, sizeof(*out)); #ifndef DS4_NO_GPU if (!e->metal_ready) { e->metal_ready = ds4_gpu_init() != 0; } if (e->metal_ready && !e->vision_map_ready) { #if defined(__APPLE__) e->vision_map_ready = ds4_gpu_set_model_map_range( e->vision_model.map, e->vision_model.size, e->vision_model.tensor_data_pos, e->vision_model.size - e->vision_model.tensor_data_pos, e->vision_model.max_tensor_bytes) != 0; #else e->vision_map_ready = ds4_gpu_set_aux_model_map_range( e->vision_model.map, e->vision_model.size, e->vision_model.tensor_data_pos, e->vision_model.size - e->vision_model.tensor_data_pos) != 0; #endif } if (!e->metal_ready || !e->vision_map_ready) { if (error && error_cap) { snprintf(error, error_cap, "unable to map vision weights on %s", ds4_backend_name(e->backend)); } return 0; } #endif float *embedding = NULL; uint32_t token_count = 0; uint32_t content_width = 0, content_height = 0; uint32_t grid_width = 0, grid_height = 0; uint32_t layout = 0; int ok = 0; if (e->vision_kind == DS4_VISION_DEEPSEEK4) { ds4_deepseek4_image_patches patches = {0}; const bool preprocessed = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 ? ds4_image_preprocess_deepseek41(&patches, image, error, error_cap) : ds4_image_preprocess_deepseek4(&patches, image, error, error_cap); if (!preprocessed) return 0; token_count = patches.llm_grid_height * patches.llm_grid_width; embedding = malloc((size_t)token_count * DS4_N_EMBD * sizeof(float)); if (embedding) { #ifndef DS4_NO_GPU ok = ds4_gpu_deepseek4_vision_encode( embedding, patches.patches, patches.grid_height, patches.grid_width, e->vision_model.map, e->vision_model.size, &e->deepseek4_vision_weights); #endif } content_width = patches.content_width; content_height = patches.content_height; grid_width = patches.llm_grid_width; grid_height = patches.llm_grid_height; layout = DS4_VISION_LAYOUT_DEEPSEEK4_NATURAL; ds4_deepseek4_image_patches_free(&patches); } else if (e->vision_kind == DS4_VISION_QWEN4) { #ifdef DS4_HAS_QWEN4_GPU ds4_image_patches patches = {0}; if (!qwen4_vision_encode_image(&e->vision_model, &e->qwen4_vision_weights, image, 64u, qwen4_vision_max_tokens(), &embedding, &patches, error, error_cap)) { ds4_image_patches_free(&patches); return 0; } ok = 1; token_count = patches.image_token_count; content_width = patches.content_width; content_height = patches.content_height; grid_width = patches.grid_width / 2u; grid_height = patches.grid_height / 2u; ds4_image_patches_free(&patches); #endif } else { ds4_image_patches patches = {0}; if (!ds4_image_preprocess_glm53(&patches, image, 16u, 8000u, error, error_cap)) return 0; token_count = patches.image_token_count; embedding = malloc((size_t)token_count * 4096u * sizeof(float)); if (embedding) { #ifndef DS4_NO_GPU ok = ds4_gpu_glm53_vision_encode( embedding, patches.patches, patches.grid_height, patches.grid_width, e->vision_model.map, e->vision_model.size, &e->vision_weights); #endif } content_width = patches.content_width; content_height = patches.content_height; ds4_image_patches_free(&patches); } if (!embedding) { if (error && error_cap) snprintf(error, error_cap, "unable to allocate vision output"); return 0; } if (!ok) { free(embedding); if (error && error_cap) snprintf(error, error_cap, "%s vision inference failed", e->vision_kind == DS4_VISION_DEEPSEEK4 ? "DeepSeek V4" : e->vision_kind == DS4_VISION_QWEN4 ? "Qwen3.8" : "GLM-5.3"); return 0; } out->data = embedding; out->token_count = token_count; out->layout = layout; out->grid_width = grid_width; out->grid_height = grid_height; out->width = image->width; out->height = image->height; out->content_width = content_width; out->content_height = content_height; memcpy(out->fingerprint, image->fingerprint, sizeof(out->fingerprint)); return 1; } int ds4_qwen4_vision_dump(const char *vision_path, const char *image_path, const char *out_path, uint32_t min_image_tokens, uint32_t max_image_tokens) { #ifndef DS4_HAS_QWEN4_GPU (void)vision_path; (void)image_path; (void)out_path; (void)min_image_tokens; (void)max_image_tokens; fprintf(stderr, "ds4: this build does not include a GPU vision backend\n"); return 0; #else ds4_model vm; ds4_qwen4_vision_weights w; model_open(&vm, vision_path, true, false); qwen4_vision_weights_bind(&w, &vm); ds4_image image = {0}; char err[256] = {0}; if (!ds4_image_decode_file(&image, image_path, err, sizeof(err))) { fprintf(stderr, "ds4: %s\n", err); return 0; } int ok = ds4_gpu_init() && ds4_gpu_set_model_map_range(vm.map, vm.size, vm.tensor_data_pos, vm.size - vm.tensor_data_pos, vm.max_tensor_bytes); float *emb = NULL; ds4_image_patches patches = {0}; if (ok) ok = qwen4_vision_encode_image(&vm, &w, &image, min_image_tokens, max_image_tokens, &emb, &patches, err, sizeof(err)); if (!ok) fprintf(stderr, "ds4: %s\n", err); FILE *fp = ok ? fopen(out_path, "wb") : NULL; if (fp) { const uint32_t hdr[4] = { patches.image_token_count, w.n_out, patches.grid_height / 2u, patches.grid_width / 2u }; ok = fwrite(hdr, sizeof(hdr), 1, fp) == 1 && fwrite(emb, sizeof(float), (size_t)patches.image_token_count * w.n_out, fp) == (size_t)patches.image_token_count * w.n_out; fclose(fp); } else if (ok) { ok = 0; } free(emb); ds4_image_patches_free(&patches); ds4_image_free(&image); return ok; #endif } int ds4_engine_vision_encode_file( ds4_engine *e, const char *path, ds4_vision_embedding *out, char *error, size_t error_cap) { ds4_image image = {0}; if (!ds4_image_decode_file(&image, path, error, error_cap)) return 0; int ok = ds4_engine_vision_encode_image(e, &image, out, error, error_cap); ds4_image_free(&image); return ok; } int ds4_engine_vision_encode_memory( ds4_engine *e, const uint8_t *encoded, size_t encoded_len, ds4_vision_embedding *out, char *error, size_t error_cap) { ds4_image image = {0}; if (!ds4_image_decode_memory(&image, encoded, encoded_len, error, error_cap)) return 0; int ok = ds4_engine_vision_encode_image(e, &image, out, error, error_cap); ds4_image_free(&image); return ok; } int ds4_engine_tp_vocab_split(ds4_engine *e) { return e && e->tp.active && e->tp.vocab_split; } #ifdef DS4_HAS_DEEPSEEK41_GPU static int ds4_engine_tp_exchange(void *ud, uint32_t layer, uint32_t gate, uint64_t seq) { ds4_engine *e = ud; ds4_tp *tp = e->tp.ctx; if (layer >= DS4_N_LAYER || gate >= DS4_TP_GATES_PER_LAYER) return 0; const uint64_t bytes = (uint64_t)DS4_N_EMBD * sizeof(float); const uint64_t out = ds4_tp_slab_out_offset(tp, layer, gate); const uint64_t in = ds4_tp_slab_in_offset(tp, layer, gate); const int ok = (!e->tp.host_slab || ds4_gpu_tensor_read(e->tp.slab, out, (char *)e->tp.host_slab + out, bytes)) && ds4_tp_gate_exchange(tp, layer, gate, seq) && (!e->tp.host_slab || ds4_gpu_tensor_write(e->tp.slab, in, (char *)e->tp.host_slab + in, bytes)); 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_engine *e = ud; ds4_tp *tp = e->tp.ctx; if (layer >= DS4_N_LAYER || !rows || rows > DS4_TP_BATCH_MAX_ROWS) return 0; const uint64_t bytes = (uint64_t)rows * DS4_N_EMBD * sizeof(float); const uint64_t out = ds4_tp_slab_batch_out_offset(tp, layer); const uint64_t in = ds4_tp_slab_batch_in_offset(tp, layer); const int ok = (!e->tp.host_slab || ds4_gpu_tensor_read(e->tp.slab, out, (char *)e->tp.host_slab + out, bytes)) && ds4_tp_batch_gate_exchange(tp, layer, rows, seq) && (!e->tp.host_slab || ds4_gpu_tensor_write(e->tp.slab, in, (char *)e->tp.host_slab + in, bytes)); 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) { #ifdef __APPLE__ 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]); } #endif ds4_engine *e = ud; ds4_tp *tp = e->tp.ctx; 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) { #ifndef DS4_HAS_DEEPSEEK41_GPU (void)e; (void)tp; snprintf(err, errlen, "tensor parallelism requires the Metal backend"); return 0; #else if (!e || !tp || (e->backend != DS4_BACKEND_METAL && (e->backend != DS4_BACKEND_CUDA || DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK41))) { snprintf(err, errlen, "tensor parallelism requires Metal or V4.1 CUDA"); return 0; } if (e->tp.ctx) { 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.ctx = tp; 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"); goto fail; } 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); goto fail; } void *network_slab; if (e->backend == DS4_BACKEND_CUDA) { e->tp.host_slab = calloc(1, (size_t)slab_bytes); network_slab = e->tp.host_slab; } else { network_slab = ds4_gpu_tensor_contents(e->tp.slab); } if (!network_slab || !ds4_gpu_tensor_fill_f32(e->tp.slab, 0, slab_bytes / 4u) || !ds4_gpu_tensor_fill_f32(e->tp.zero_vec, 0, vec_bytes / 4u) || !ds4_gpu_synchronize()) { snprintf(err, errlen, "tp: slab initialization failed"); goto fail; } if (!ds4_tp_attach_slab(tp, network_slab, err, errlen)) goto fail; 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"); goto fail; } } 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"); goto fail; } } if (!ds4_gpu_tp_init((uint32_t)ds4_tp_rank(tp), e->tp.slab, ds4_tp_slab_gpu_flags_offset(tp), ds4_tp_slab_out_offset(tp, 0, 0), vec_bytes, ds4_engine_tp_exchange, e)) { snprintf(err, errlen, "tp: gate service init failed"); goto fail; } ds4_gpu_tp_set_batch_exchange(ds4_engine_tp_batch_exchange); #ifdef __APPLE__ g_tp_block_ctx = tp; #endif ds4_gpu_tp_set_big_exchange(ds4_engine_tp_big_exchange); /* Reuse the existing half-logit frames for V4.1 on CUDA as well. */ e->tp.vocab_split = DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK4 || (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41 && e->backend == DS4_BACKEND_CUDA); 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; fail: ds4_engine_tp_unbind(e); return 0; #endif } void ds4_engine_tp_unbind(ds4_engine *e) { #ifdef DS4_HAS_DEEPSEEK41_GPU if (!e || !e->tp.ctx) return; ds4_gpu_tp_shutdown(); ds4_tp_detach_slab(e->tp.ctx); #ifdef __APPLE__ g_tp_block_ctx = NULL; #endif 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); free(e->tp.host_slab); memset(&e->tp, 0, sizeof(e->tp)); #else (void)e; #endif } bool ds4_engine_is_glm_dsa(ds4_engine *e) { (void)e; return DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; } bool ds4_engine_is_deepseek41(ds4_engine *e) { (void)e; return DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41; } void ds4_engine_close(ds4_engine *e) { if (!e) return; ds4_engine_tp_unbind(e); 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); if (e->vision_model.map) model_close(&e->vision_model); model_close(&e->model); #ifndef DS4_NO_GPU #ifdef DS4_HAS_QWEN4_GPU qwen4_state_pool_free(e); if (e->qwen4_shared_workspace) { qwen4_graph_free(e->qwen4_shared_workspace); free(e->qwen4_shared_workspace); e->qwen4_shared_workspace = NULL; } #endif if (e->shared_prefill_workspace_ready) { metal_graph_free_prefill_workspace(&e->shared_prefill_workspace); e->shared_prefill_workspace_ready = false; } 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->model_path); 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; char net_saved[48]; /* Seed batches include an ordinary target step, often without any timed * standalone step to estimate its cost. Do not print a fictitious saving. */ if (st->seed_batches) snprintf(net_saved, sizeof(net_saved), "n/a"); else snprintf(net_saved, sizeof(net_saved), "%.3f", 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 seed_batches=%llu " "full=%llu partial=%llu direct_full=%llu direct_partial=%llu " "replay_fallbacks=%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=%s " "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->seed_batches, (unsigned long long)st->full_accepts, (unsigned long long)st->partial_accepts, (unsigned long long)st->direct_full_commits, (unsigned long long)st->direct_partial_commits, (unsigned long long)st->replay_fallbacks, (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, 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 (ds4_model_is_qwen4()) { fprintf(stderr, "ds4: Qwen3.8 sessions require Metal or CUDA\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; #ifdef DS4_HAS_DEEPSEEK41_GPU if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) { if (ctx_size > 1048576 || !ds41_memory_admit(e, ds4_add_sat_u64(e->ds41_session_bytes, ds41_graph_bytes((uint32_t)ctx_size)), false) || !ds41_graph_alloc(&s->ds41_graph, &e->model, &e->weights, e->model_path, (uint32_t)ctx_size, e->ssd_streaming)) { free(s); return 1; } s->ds41_graph_ready = true; s->ds41_graph.quality = e->quality; if (e->tp.active) { s->ds41_graph.tp_world = 2; s->ds41_graph.tp_rank = (uint32_t)e->tp.rank; s->ds41_graph.tp_out = e->tp.out_views; s->ds41_graph.tp_in = e->tp.in_views; if (e->tp.vocab_split) { const uint64_t half_bytes = (uint64_t)DS4_N_VOCAB / 2u * sizeof(float); s->ds41_graph.tp_logits_half = ds4_gpu_tensor_view(s->ds41_graph.logits, (uint64_t)e->tp.rank * half_bytes, half_bytes); if (!s->ds41_graph.tp_logits_half || !ds4_gpu_tensor_fill_f32(s->ds41_graph.logits, 0, DS4_N_VOCAB)) { ds41_graph_free(&s->ds41_graph); free(s); return 1; } } } s->prefill_cap = s->ds41_graph.prefill_cap; s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); s->sample_probs = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); e->ds41_session_bytes += s->ds41_graph.allocation_bytes; if (!ds4_session_tp_register(s)) { ds4_session_free(s); return 1; } *out = s; return 0; } #endif #ifdef DS4_HAS_QWEN4_GPU if (ds4_model_is_qwen4()) { if ((e->backend != DS4_BACKEND_METAL && e->backend != DS4_BACKEND_CUDA) || e->distributed.role != DS4_DISTRIBUTED_NONE) { fprintf(stderr, "ds4: Qwen3.8 sessions require single-host Metal or CUDA\n"); free(s); return 1; } const uint32_t cap_tokens = e->prefill_chunk && e->prefill_chunk < (uint32_t)ctx_size ? e->prefill_chunk : qwen4_prefill_chunk_tokens((uint32_t)ctx_size); /* Every slot of a batched server has the same shape, so one arena of * transients serves them all and each extra session then costs only * its own caches. A session that needs a wider arena replaces it * when no live session borrows it (a new context size once the * previous sessions closed) and keeps private transients otherwise. */ const uint32_t block_cap = (uint32_t)ctx_size / 4u + 1u; const bool share = e->share_session_prefill_workspace && e->backend == DS4_BACKEND_METAL; const bool arena_fits = e->qwen4_shared_workspace && e->qwen4_shared_workspace->cap_tokens >= cap_tokens && e->qwen4_shared_workspace->n_block_cap >= block_cap; if (share && e->qwen4_shared_workspace && !arena_fits && e->qwen4_arena_users == 0) { qwen4_graph_free(e->qwen4_shared_workspace); free(e->qwen4_shared_workspace); e->qwen4_shared_workspace = NULL; } const ds4_qwen4_gpu_graph *shared = share && arena_fits ? e->qwen4_shared_workspace : NULL; s->qwen4_slot = -1; if (share && !e->glm_mtp && getenv("DS4_QWEN4_NO_STATE_POOL") == NULL && qwen4_state_pool_ensure(e, (uint32_t)(e->placement_session_count_hint > 0 ? e->placement_session_count_hint : 0))) { s->qwen4_slot = qwen4_state_pool_take(e); } if (!qwen4_graph_alloc(&s->qwen4_graph, &e->weights, (uint32_t)ctx_size, cap_tokens, e->glm_mtp, shared, s->qwen4_slot >= 0 ? e->qwen4_lin_state_pool : NULL, s->qwen4_slot >= 0 ? e->qwen4_lin_hist_pool : NULL, s->qwen4_slot >= 0 ? (uint32_t)s->qwen4_slot : 0u)) { if (s->qwen4_slot >= 0) e->qwen4_pool_used &= ~(UINT64_C(1) << s->qwen4_slot); free(s); return 1; } if (share && !e->qwen4_shared_workspace) { e->qwen4_shared_workspace = xcalloc(1, sizeof(*e->qwen4_shared_workspace)); const uint64_t arena = qwen4_graph_scratch_bytes(&s->qwen4_graph); qwen4_graph_transfer_scratch(e->qwen4_shared_workspace, &s->qwen4_graph); fprintf(stderr, "ds4: shared Qwen3.8 session transients: cap=%u, %.2f GiB once; " "each further session costs its caches only\n", cap_tokens, (double)arena / (1024.0 * 1024.0 * 1024.0)); } qwen4_graph_reset(&s->qwen4_graph); if (!qwen4_graph_load_steering(&s->qwen4_graph, e->directional_steering_file, e->directional_steering_attn_scale, e->directional_steering_ffn_scale)) { if (s->qwen4_slot >= 0) e->qwen4_pool_used &= ~(UINT64_C(1) << s->qwen4_slot); qwen4_graph_free(&s->qwen4_graph); free(s); return 1; } s->qwen4_graph_ready = true; if (!s->qwen4_graph.owns_scratch) e->qwen4_arena_users++; s->prefill_cap = (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; } #endif 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; } uint32_t guard_ctx; if (!glm_graph_context_request(ctx_size, &guard_ctx)) { free(s); return 1; } const uint32_t work_ctx = glm_graph_full_attention_cap(guard_ctx, e->ssd_streaming); const uint32_t compact_cap = glm_graph_compact_cache_initial_cap( guard_ctx, work_ctx); const ds4_context_memory session_mem = glm_graph_context_memory_estimate_for_compact_cap_slice( guard_ctx, work_ctx, compact_cap, e->ssd_streaming, layer_start, layer_end); const uint64_t all_graph_bytes = ds4_engine_glm_graph_budget(e, session_mem.total_bytes); const uint64_t other_graph_bytes = all_graph_bytes - session_mem.total_bytes; if (e->ssd_streaming && !e->ssd_streaming_budget_finalized) { ds4_engine_fit_glm_streaming_budget(e, true, layer_start, layer_end, require_output, ctx_size); ds4_gpu_set_streaming_expert_cache_budget( e->ssd_streaming_cache_experts); e->ssd_streaming_budget_finalized = true; } 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, glm_graph_saturating_add_u64( ds4_engine_streaming_transient_guard_bytes(e), other_graph_bytes), 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 (!glm_graph_load_directional_steering( &s->glm_graph, e->directional_steering_file, e->directional_steering_attn_scale, e->directional_steering_ffn_scale)) { glm_graph_free(&s->glm_graph); free(s); return 1; } 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; } } s->glm_reserved_graph_bytes = session_mem.total_bytes; e->glm_session_graph_bytes = glm_graph_saturating_add_u64( e->glm_session_graph_bytes, s->glm_reserved_graph_bytes); e->glm_session_count++; 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->vision_kind == DS4_VISION_DEEPSEEK4) { s->graph.vision_model_map = e->vision_model.map; s->graph.vision_model_size = e->vision_model.size; s->graph.deepseek4_vision_weights = &e->deepseek4_vision_weights; } 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 (s->glm_reserved_graph_bytes && s->engine) { s->engine->glm_session_graph_bytes -= s->glm_reserved_graph_bytes; s->engine->glm_session_count--; } 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 { #ifdef DS4_HAS_DEEPSEEK41_GPU if (s->ds41_graph_ready) { s->engine->ds41_session_bytes -= s->ds41_graph.allocation_bytes; ds41_graph_free(&s->ds41_graph); } else #endif #ifdef DS4_HAS_QWEN4_GPU if (ds4_session_is_qwen4(s)) { if (s->engine && s->engine->glm_mtp_timing && s->qwen4_spec_cycles) { fprintf(stderr, "ds4: Qwen3.8 mtp: %" PRIu64 " verify cycles, %" PRIu64 " drafts accepted (%.1f%%)\n", s->qwen4_spec_cycles, s->qwen4_spec_accepted, 100.0 * (double)s->qwen4_spec_accepted / (double)s->qwen4_spec_cycles); } free(s->qwen4_verify_logits); if (s->qwen4_slot >= 0 && s->engine) { s->engine->qwen4_pool_used &= ~(UINT64_C(1) << s->qwen4_slot); } if (s->qwen4_graph_ready && !s->qwen4_graph.owns_scratch && s->engine) s->engine->qwen4_arena_users--; qwen4_graph_free(&s->qwen4_graph); } else #endif 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->checkpoint_images); 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; if (ds4_engine_is_deepseek41(s->engine)) { if (power_percent == 100) return 0; fprintf(stderr, "ds4: session power throttling is not supported for DeepSeek V4.1\n"); 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; } float ds4_session_directional_steering_ffn(ds4_session *s) { if (!s || !s->engine) return 0.0f; #ifndef DS4_NO_GPU if (!ds4_session_is_cpu(s)) { return ds4_session_is_glm(s) ? s->glm_graph.directional_steering_ffn_scale : s->graph.directional_steering_ffn_scale; } #endif return s->engine->directional_steering_ffn_scale; } int ds4_session_set_directional_steering_ffn(ds4_session *s, float scale) { if (!s || !s->engine || !isfinite(scale) || scale < -100.0f || scale > 100.0f) { return 1; } if (s->distributed || s->engine->tp.active) { fprintf(stderr, "ds4: live steering changes are not supported for distributed or network tensor-parallel sessions\n"); return 1; } bool loaded = s->engine->directional_steering_dirs != NULL; #ifndef DS4_NO_GPU if (!ds4_session_is_cpu(s)) { #ifdef DS4_HAS_QWEN4_GPU if (s->qwen4_graph_ready) { loaded = s->qwen4_graph.steer_dirs != NULL; } else #endif if (ds4_session_is_glm(s)) { const int tier = glm_graph_directional_steering_tier( &s->glm_graph, s->glm_graph.layer_start); loaded = tier >= 0 && s->glm_graph.directional_steering_dirs_by_tier[tier] != NULL; } else { loaded = metal_graph_directional_steering_dirs(&s->graph) != NULL; } } #endif if (scale != 0.0f && !loaded) { fprintf(stderr, "ds4: live FFN steering needs a directional steering vector loaded at startup\n"); return 1; } s->engine->directional_steering_ffn_scale = scale; #ifndef DS4_NO_GPU if (!ds4_session_is_cpu(s)) { #ifdef DS4_HAS_QWEN4_GPU if (s->qwen4_graph_ready) { s->qwen4_graph.steer_ffn_scale = scale; s->glm_mtp_have = 0; } else #endif if (ds4_session_is_glm(s)) { s->glm_graph.directional_steering_ffn_scale = scale; s->glm_mtp_have = 0; s->glm_mtp_rollback_valid = false; } else { s->graph.directional_steering_ffn_scale = scale; ds4_session_dspark_capture_invalidate(s); } } #endif s->mtp_draft_valid = false; s->greedy_splitkv_segment.len = 0; s->greedy_splitkv_anchor_valid = false; 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; } static bool speculative_point_accept(float target_p, float draft_p, uint64_t *rng); static int speculative_point_replacement(ds4_session *s, int draft_token, uint64_t *rng); /* One GLM MTP speculative cycle: evaluate first_token, and when a compatible * pending draft exists verify [first_token, draft] in one 2-row batch pass * (big gates only under TP). Greedy and opportunistic sampling accept a draft * only when it is the target argmax. Exact sampling treats the deterministic * draft as a point-mass proposal and uses stochastic p/q acceptance. Returns * one or two committed tokens with s->logits at the last committed position, * or -1 on error. */ static int ds4_session_glm_spec_cycle_impl( ds4_session *s, int first_token, int eos_token, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, bool exact_sampling, 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 bool timing = s->engine->glm_mtp_timing; const uint32_t dense_before = s->glm_dense_cache_len; const uint64_t hc_row_values = (uint64_t)DS4_N_EMBD * (g->glm53 ? DS4_N_HC : 1u); const uint64_t hc_row_bytes = hc_row_values * sizeof(float); if (!s->glm_mtp_hc) { s->glm_mtp_hc = malloc(2ull * hc_row_bytes); 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; } } s->glm_mtp_rollback_valid = false; if (s->glm_mtp_have && first_token != s->glm_mtp_parent) { s->glm_mtp_have = 0; } 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_parent = n1; 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 }; bool verified = false; bool state_saved = false; if (g->glm53) { state_saved = glm_graph_mtp_ensure(g) && glm53_graph_copy_spec_state(g, true); if (state_saved) { if (!glm53_graph_use_indexed_prefill(g) && glm_graph_span_fits_full_attention(g, pos, 2u)) { verified = glm_graph_forward_tokens(g, &e->model, &e->weights, toks, NULL, NULL, 0, pos, 2, s->glm_mtp_hc, s->logits, NULL, NULL, pos, 0, 2); } else { verified = glm_graph_forward_indexed_tokens(g, &e->model, &e->weights, toks, NULL, NULL, 0, pos, 2, s->glm_mtp_hc, s->logits, NULL, NULL, pos, 0, 2); } } } else if (pos + 2u <= glm_graph_indexer_top_k_limit()) { /* GLM-5.2 verification must use the compact caches populated by * decode; its full-KV batch path would read unwritten rows. */ verified = glm_graph_verify_rows(g, &e->model, &e->weights, toks, pos, 2, s->glm_mtp_hc, s->logits); } if (!verified && !g->glm53) { if (!glm_graph_indexed_prefill_batch_ready(g, pos) || glm_graph_limit_indexed_prefill_chunk(g, pos, 2u) < 2u || !glm_graph_forward_indexed_tokens(g, &e->model, &e->weights, toks, NULL, NULL, 0, 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; } } else if (!verified) { if (state_saved) (void)glm53_graph_copy_spec_state(g, false); if (errlen) snprintf(err, errlen, "glm mtp: GLM 5.3 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 (g->glm53 && state_saved) { (void)glm53_graph_copy_spec_state(g, false); } if (errlen) snprintf(err, errlen, "glm mtp: scratch alloc failed"); s->checkpoint_valid = false; return -1; } ds4_gpu_tensor *h0 = NULL; bool head_ok = false; if (g->glm53) { head_ok = ds4_gpu_tensor_write(g->hc_cur, 0, s->glm_mtp_hc, hc_row_bytes) != 0 && glm_graph_forward_output_head(g, &e->model, &e->weights, g->hc_cur, s->glm_mtp_logits0); } else { h0 = ds4_gpu_tensor_view(g->mtp_concat, 0, (uint64_t)DS4_N_EMBD * sizeof(float)); head_ok = h0 != NULL && ds4_gpu_tensor_write(h0, 0, s->glm_mtp_hc, hc_row_bytes) != 0 && glm_graph_forward_output_head(g, &e->model, &e->weights, h0, s->glm_mtp_logits0); ds4_gpu_tensor_free(h0); } if (!head_ok) { if (g->glm53 && state_saved) { (void)glm53_graph_copy_spec_state(g, false); } if (errlen) snprintf(err, errlen, "glm mtp: row0 head failed"); s->checkpoint_valid = false; return -1; } const int n1 = glm_session_logits_argmax(s->glm_mtp_logits0); int replacement = -1; int accept = n1 == d; if (exact_sampling) { if (!rng || !sample_build_probabilities(s->glm_mtp_logits0, DS4_N_VOCAB, temperature, top_k, top_p, min_p, s->sample_probs)) { if (g->glm53 && state_saved) { (void)glm53_graph_copy_spec_state(g, false); } if (errlen) snprintf(err, errlen, "glm mtp: target distribution failed"); s->checkpoint_valid = false; return -1; } accept = speculative_point_accept(s->sample_probs[d], 1.0f, rng); if (!accept) { replacement = speculative_point_replacement(s, d, rng); if (replacement < 0) { if (g->glm53 && state_saved) { (void)glm53_graph_copy_spec_state(g, false); } if (errlen) snprintf(err, errlen, "glm mtp: replacement sampling failed"); s->checkpoint_valid = false; return -1; } } } token_vec_push(&s->checkpoint, first_token); s->checkpoint_valid = true; 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; ds4_gpu_tensor *target_hidden = g->glm53 ? g->hc_cur : g->cur; const bool cu = ds4_gpu_tensor_write(target_hidden, 0, s->glm_mtp_hc, hc_row_bytes) != 0 && glm_graph_mtp_step(g, &e->model, &e->weights, d, pos, s->glm_mtp_min_pos, &dummy) && ds4_gpu_tensor_write(target_hidden, 0, s->glm_mtp_hc + hc_row_values, hc_row_bytes) != 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_parent = n2; s->glm_mtp_have = 1; } accepted[0] = first_token; accepted[1] = d; } else { bool replay_ok = true; if (g->glm53) { replay_ok = glm53_graph_copy_spec_state(g, false) && glm_graph_forward_token(g, &e->model, &e->weights, first_token, NULL, pos, NULL, s->logits, false); } else { memcpy(s->logits, s->glm_mtp_logits0, (size_t)DS4_N_VOCAB * sizeof(float)); } if (!replay_ok) { if (errlen) snprintf(err, errlen, "glm mtp: rejection replay failed"); s->checkpoint_valid = false; return -1; } if (exact_sampling) { if (!glm_graph_forward_token(g, &e->model, &e->weights, replacement, NULL, pos + 1u, NULL, s->logits, false)) { if (errlen) snprintf(err, errlen, "glm mtp: replacement replay failed"); s->checkpoint_valid = false; return -1; } token_vec_push(&s->checkpoint, replacement); ds4_session_glm_note_dense_cache(s, pos, 2); n_committed = 2; accepted[0] = first_token; accepted[1] = replacement; const int next = glm_session_logits_argmax(s->logits); int nd = -1; s->glm_mtp_have = 0; if (replacement != eos_token && glm_graph_mtp_step(g, &e->model, &e->weights, next, pos + 1u, s->glm_mtp_min_pos, &nd)) { s->glm_mtp_draft = nd; s->glm_mtp_parent = next; s->glm_mtp_have = 1; } } else { ds4_session_glm_note_dense_cache(s, pos, 1); } int nd = -1; ds4_gpu_tensor *target_hidden = g->glm53 ? g->hc_cur : g->cur; const bool hidden_ready = exact_sampling || g->glm53 || ds4_gpu_tensor_write(target_hidden, 0, s->glm_mtp_hc, hc_row_bytes) != 0; const bool cu = !exact_sampling && hidden_ready && 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_parent = n1; s->glm_mtp_have = 1; } if (!exact_sampling) 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 ? (exact_sampling ? "EXACT_ACCEPT" : "ACCEPT") : (exact_sampling ? "exact_reject" : "reject"), d, dt ? dt : "?", n1, nt ? nt : "?"); free(dt); free(nt); } if (g->glm53 && state_saved && n_committed == 2) { s->glm_mtp_rollback_pos = pos; s->glm_mtp_rollback_dense_len = dense_before; s->glm_mtp_rollback_first_token = first_token; s->glm_mtp_rollback_valid = true; } return n_committed; } /* Restore the state immediately before the last two-token GLM-5.3 MTP cycle, * then replay the retained first row when the caller keeps it. ds4-agent uses * this when a speculative block crosses into or out of greedy tool syntax. */ static bool ds4_session_glm_mtp_rewind(ds4_session *s, int pos) { if (!s || !s->glm_mtp_rollback_valid || !s->glm_graph.glm53 || s->checkpoint.len != (int)s->glm_mtp_rollback_pos + 2 || (pos != (int)s->glm_mtp_rollback_pos && pos != (int)s->glm_mtp_rollback_pos + 1)) { return false; } const uint32_t start = s->glm_mtp_rollback_pos; bool ok = glm53_graph_copy_spec_state(&s->glm_graph, false); s->glm_dense_cache_len = s->glm_mtp_rollback_dense_len; if (ok && pos == (int)start + 1) { ok = glm_graph_forward_token(&s->glm_graph, &s->engine->model, &s->engine->weights, s->glm_mtp_rollback_first_token, NULL, start, NULL, s->logits, false); if (ok) ds4_session_glm_note_dense_cache(s, start, 1); } s->glm_mtp_rollback_valid = false; s->glm_mtp_have = 0; return ok; } static int ds4_session_glm_spec_cycle(ds4_session *s, int first_token, int *accepted, int accepted_cap, char *err, size_t errlen) { return ds4_session_glm_spec_cycle_impl(s, first_token, -1, 0.0f, 0, 1.0f, 0.0f, NULL, false, accepted, accepted_cap, err, errlen); } #ifdef DS4_HAS_QWEN4_GPU /* after a rewind past the verify snapshot the graph was reset: re-prefill the * kept transcript so the state matches the checkpoint again */ static int qwen4_session_replay_if_stale(ds4_session *s, char *err, size_t errlen) { if (!s->qwen4_graph_ready || s->qwen4_graph.pos == (uint32_t)s->checkpoint.len) return 0; if (s->checkpoint_image_count) { if (errlen) snprintf(err, errlen, "Qwen3.8 image rewind requires multimodal sync"); s->checkpoint_valid = false; return -1; } token_vec kept = {0}; for (int i = 0; i < s->checkpoint.len; i++) token_vec_push(&kept, s->checkpoint.v[i]); s->checkpoint_valid = false; const int rc = kept.len ? ds4_session_sync(s, &kept, err, errlen) : 0; if (!kept.len) qwen4_graph_reset(&s->qwen4_graph); token_vec_free(&kept); return rc; } /* test knobs: DS4_QWEN4_SPEC_FORCE_ACCEPT commits every draft, * DS4_QWEN4_SPEC_TRACE logs each verify outcome */ static bool qwen4_spec_trace(void) { static int v = -1; if (v < 0) v = getenv("DS4_QWEN4_SPEC_TRACE") != NULL; return v != 0; } static bool qwen4_spec_force_accept(void) { return getenv("DS4_QWEN4_SPEC_FORCE_ACCEPT") != NULL; } /* draft the token after `parent` from row `row` of the last forward, at idx */ static void qwen4_session_draft(ds4_session *s, uint32_t row, int parent, uint32_t idx) { int d = -1; s->glm_mtp_have = 0; if (qwen4_graph_mtp_step(&s->qwen4_graph, &s->engine->model, &s->engine->weights, row, parent, idx, true, NULL, &d)) { s->glm_mtp_draft = d; s->glm_mtp_parent = parent; s->glm_mtp_have = 1; } } /* Draft depth policy for the Qwen3.8 cycle: DS4_QWEN4_MTP_DEPTH=2 or =3 * forces a fixed depth (read per cycle so the A/B harnesses can switch it * per step); 0/auto, the default, drafts two tokens while the rolling * window of first-draft acceptances is strong and falls back to one. The * first draft is identical at both depths, so the window is a comparable * signal across modes. */ static int qwen4_spec_depth(ds4_session *s) { const char *env = getenv("DS4_QWEN4_MTP_DEPTH"); const int v = env && env[0] ? atoi(env) : 0; if (v == 2 || v == 3) return v; const uint32_t window = s->qwen4_depth_window & 255u; uint32_t bits = 0; for (uint32_t i = 0; i < 8u; i++) bits += (window >> i) & 1u; /* Engage only on a perfect recent window (deterministic continuation) * and leave at the first sign the chained draft stopped paying: the * second draft accepts at ~0.6 on general prose, which does not cover * the wider cycle, while deterministic continuations accept at ~0.95+ * and gain 10-20 percent. */ if (!s->qwen4_depth3_engaged) { if (bits >= 8u && s->qwen4_depth_cycles >= 8u && s->qwen4_reject2_streak == 0u) { s->qwen4_depth3_engaged = true; } } else if (bits < 6u || s->qwen4_reject2_streak >= 2u) { s->qwen4_depth3_engaged = false; } return s->qwen4_depth3_engaged ? 3 : 2; } static void qwen4_spec_note_first_draft(ds4_session *s, bool accepted_first) { s->qwen4_depth_window = (s->qwen4_depth_window << 1u) | (accepted_first ? 1u : 0u); s->qwen4_depth_cycles++; if (accepted_first) s->qwen4_depth_accepted++; } /* One Qwen3.8 MTP cycle: evaluate first_token, or verify [first_token, draft] * — plus a second chained draft when the depth policy engages it — in one * 2- or 3-row pass. Greedy and opportunistic sampling accept a draft when * it is the target argmax; exact sampling treats it as a point-mass proposal * (accept with p(draft), else replay a residual sample) and always runs at * depth 2. */ static int ds4_session_qwen4_spec_cycle(ds4_session *s, int first_token, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, bool exact_sampling, int *accepted, int accepted_cap, char *err, size_t errlen) { ds4_engine *e = s->engine; ds4_qwen4_gpu_graph *g = &s->qwen4_graph; const ds4_model *m = &e->model; const ds4_weights *w = &e->weights; if (qwen4_session_replay_if_stale(s, err, errlen) != 0) return -1; const uint32_t pos = g->pos; const uint32_t V = DS4_N_VOCAB; const int depth = (exact_sampling && temperature > 0.0f) || getenv("DS4_QWEN4_NO_MTP_BATCH") != NULL ? 2 : qwen4_spec_depth(s); if (s->glm_mtp_have && first_token != s->glm_mtp_parent) s->glm_mtp_have = 0; if (depth == 3 && s->glm_mtp_have && s->glm_mtp_have2 && !g->snap2_ple_hist && accepted_cap >= 3 && pos + 3u <= g->ctx_cap && g->cap_tokens >= 3u) { (void)qwen4_graph_ensure_snap2(g); } const bool deep = depth == 3 && s->glm_mtp_have && s->glm_mtp_have2 && accepted_cap >= 3 && pos + 3u <= g->ctx_cap && g->cap_tokens >= 3u && g->mtp_R && g->snap2_ple_hist != NULL; if (!s->glm_mtp_have || accepted_cap < 2 || pos + 2u > g->ctx_cap || g->cap_tokens < 2u || !g->mtp_R || !qwen4_graph_fused(g, 2u)) { 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 parent = sample_argmax(s->logits, V); int draft = -1; s->glm_mtp_have2 = false; if (qwen4_graph_mtp_step(&s->qwen4_graph, m, w, 0, parent, pos, true, NULL, &draft)) { s->glm_mtp_draft = draft; s->glm_mtp_parent = parent; s->glm_mtp_have = 1; if (depth == 3 && pos + 2u <= g->ctx_cap) { int draft2 = -1; if (qwen4_graph_mtp_chain_step(&s->qwen4_graph, m, w, draft, pos + 1u, &draft2)) { s->glm_mtp_draft2 = draft2; s->glm_mtp_have2 = true; } } } if (qwen4_spec_trace()) { fprintf(stderr, "ds4: spec pos %u token %d plain\n", pos, first_token); } accepted[0] = first_token; return 1; } s->glm_mtp_have = 0; s->glm_mtp_have2 = false; const int d = s->glm_mtp_draft; const int d2 = deep ? s->glm_mtp_draft2 : -1; const int toks[3] = { first_token, d, d2 }; const uint32_t T = deep ? 3u : 2u; /* Three verifier rows plus the logits before the block, for exact rewinds. */ if (!s->qwen4_verify_logits) s->qwen4_verify_logits = xmalloc(4u * (size_t)V * sizeof(float)); float *rows = s->qwen4_verify_logits; /* Exact sampling rewinds to the block start to resample the boundary * token when the caller discards a block that crossed a sampling-mode * boundary (server tool-syntax transitions). The row-0 snapshot below * cannot serve that rewind, so save the pre-verify state first; without * it the rewind resets the recurrent state and the next eval replays * the whole kept context. */ g->snap0_valid = s->engine->dspark_exact_sampling && qwen4_graph_ensure_snap0(g) && qwen4_graph_state_copy0(g, true); if (g->snap0_valid) memcpy(rows + (size_t)3u * V, s->logits, (size_t)V * sizeof(float)); /* the recurrent kernels snapshot their state after row 0 (and row 1 for * the 3-row verify), so a rejected draft only needs the snapshot * restored, not a replay of the accepted prefix */ g->snap_after_first = g->snap_ple_hist != NULL; g->snap_valid = false; g->snap_after_second = deep && g->snap2_ple_hist != NULL; g->snap2_valid = false; g->verify_rows_exact = deep; ds4_gpu_qwen4_set_verify_rows_exact(deep); const bool ok = qwen4_graph_forward_tokens(g, m, w, toks, T, rows, true); ds4_gpu_qwen4_set_verify_rows_exact(false); g->verify_rows_exact = false; g->snap_after_second = false; if (!ok) { g->snap_after_first = false; if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: verify failed"); s->checkpoint_valid = false; return -1; } s->qwen4_spec_cycles++; s->qwen4_rewound = false; token_vec_push(&s->checkpoint, first_token); s->checkpoint_valid = true; s->mtp_draft_valid = false; bool accept = sample_argmax(rows, V) == d || qwen4_spec_force_accept(); bool accept2 = false; if (deep) { accept2 = accept && (sample_argmax(rows + V, V) == d2 || qwen4_spec_force_accept()); } int replacement = -1; if (exact_sampling && temperature > 0.0f) { if (!s->sample_probs) s->sample_probs = xmalloc((size_t)V * sizeof(s->sample_probs[0])); if (!rng || !sample_build_probabilities(rows, V, temperature, top_k, top_p, min_p, s->sample_probs)) { if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: target distribution failed"); s->checkpoint_valid = false; return -1; } accept = speculative_point_accept(s->sample_probs[d], 1.0f, rng); if (!accept) { replacement = speculative_point_replacement(s, d, rng); if (replacement < 0) { if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: replacement sampling failed"); s->checkpoint_valid = false; return -1; } } } qwen4_spec_note_first_draft(s, accept); if (deep) { if (accept2) s->qwen4_reject2_streak = 0u; else if (accept) s->qwen4_reject2_streak++; } if (qwen4_spec_trace()) { fprintf(stderr, "ds4: spec pos %u token %d draft %d %s%s\n", pos, first_token, d, accept ? "accept" : "reject", deep ? (accept2 ? " +accept2" : " +reject2") : ""); } if (accept && (!deep || accept2)) { token_vec_push(&s->checkpoint, d); if (deep) token_vec_push(&s->checkpoint, d2); memcpy(s->logits, rows + (T - 1u) * V, (size_t)V * sizeof(float)); const int parent = sample_argmax(s->logits, V); bool have_next = false; int draft = -1; if (deep) { /* catch the predictor up over all three committed rows (the * speculative chain rows at pos and pos+1 are re-run with the * true trunk streams), then chain one speculative row for the * second pending draft */ const int next_tokens[3] = {d, d2, parent}; if (qwen4_graph_mtp_steps(g, m, w, 0, next_tokens, 3u, pos, true, NULL, &draft)) { have_next = true; } } else { const int next_tokens[2] = {d, parent}; have_next = qwen4_graph_mtp_steps(g, m, w, 0, next_tokens, 2u, pos, true, NULL, &draft); } if (have_next) { s->glm_mtp_draft = draft; s->glm_mtp_parent = parent; s->glm_mtp_have = 1; if (depth == 3 && g->pos + 1u < g->ctx_cap) { int draft2 = -1; if (qwen4_graph_mtp_chain_step(g, m, w, draft, g->pos, &draft2)) { s->glm_mtp_draft2 = draft2; s->glm_mtp_have2 = true; } } } s->qwen4_spec_accepted++; accepted[0] = first_token; accepted[1] = d; if (deep) accepted[2] = d2; return deep ? 3 : 2; } if (accept && deep) { /* first draft accepted, second rejected: restore the row-1 snapshot */ if (!g->snap2_valid || !qwen4_graph_state_swap2(g)) { if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: rejection restore failed"); s->checkpoint_valid = false; return -1; } s->checkpoint.len = pos + 1; token_vec_push(&s->checkpoint, d); memcpy(s->logits, rows + V, (size_t)V * sizeof(float)); const int parent = sample_argmax(s->logits, V); const int next_tokens[2] = {d, parent}; int draft = -1; if (qwen4_graph_mtp_steps(g, m, w, 0, next_tokens, 2u, pos, true, NULL, &draft)) { s->glm_mtp_draft = draft; s->glm_mtp_parent = parent; s->glm_mtp_have = 1; if (depth == 3 && g->pos + 1u < g->ctx_cap) { int draft2 = -1; if (qwen4_graph_mtp_chain_step(g, m, w, draft, g->pos, &draft2)) { s->glm_mtp_draft2 = draft2; s->glm_mtp_have2 = true; } } } s->qwen4_spec_accepted++; accepted[0] = first_token; accepted[1] = d; return 2; } if (!g->snap_valid || !qwen4_graph_state_swap(g)) { if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: rejection restore failed"); s->checkpoint_valid = false; return -1; } if (replacement >= 0) { /* the residual sample stands in for the draft: run it from the restored state */ if (!qwen4_graph_forward_tokens(g, m, w, &replacement, 1, s->logits, false)) { if (errlen) snprintf(err, errlen, "Qwen3.8 mtp: replacement replay failed"); s->checkpoint_valid = false; return -1; } token_vec_push(&s->checkpoint, replacement); qwen4_session_draft(s, 0, sample_argmax(s->logits, V), pos + 1u); accepted[0] = first_token; accepted[1] = replacement; return 2; } memcpy(s->logits, rows, (size_t)V * sizeof(float)); { const int parent = sample_argmax(s->logits, V); int draft = -1; if (qwen4_graph_mtp_step(g, m, w, 0, parent, pos, true, NULL, &draft)) { s->glm_mtp_draft = draft; s->glm_mtp_parent = parent; s->glm_mtp_have = 1; if (depth == 3 && pos + 2u <= g->ctx_cap) { int draft2 = -1; if (qwen4_graph_mtp_chain_step(g, m, w, draft, pos + 1u, &draft2)) { s->glm_mtp_draft2 = draft2; s->glm_mtp_have2 = true; } } } } accepted[0] = first_token; return 1; } #endif #endif int ds4_session_glm_tp_spec_cycle(ds4_session *s, int token, int limit, char *err, size_t errlen) { #ifndef DS4_NO_GPU if (s && s->engine && s->engine->tp.active && s->engine->tp.rank == 1 && s->checkpoint_valid && ds4_session_is_glm(s) && s->engine->glm_mtp && !s->engine->dspark_exact_sampling && limit >= 1 && limit <= 2 && token >= 0 && token < (int)DS4_N_VOCAB) { int accepted[2]; return ds4_session_glm_spec_cycle(s, token, accepted, limit, err, errlen); } #else (void)s; (void)token; (void)limit; #endif payload_set_err(err, errlen, "invalid GLM TP speculative command"); return -1; } 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; #ifdef DS4_ROCM_BUILD const bool rocm_layer_slice_token_decode = glm_graph_env_truthy( getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE")); #else const bool rocm_layer_slice_token_decode = false; #endif 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; /* * The token graph accepts both embeddings and inter-node hidden * states. Keep the resident ROCm continuation path opt-in until * remote output and timing tests validate it across GLM quants. */ if (remaining == 1 && pos > 0 && ((!input_hc && !output_hc) || rocm_layer_slice_token_decode)) { 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, NULL, 0, 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(&s->glm_graph, 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, NULL, 0, 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; const bool layer_selected_addr = batch_selected_addr && metal_graph_stream_prefill_batch_selected_addr_layer_supported( &e->weights, il); ok = layer_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); bool ds4_session_vision_prefix_matches( const ds4_session *s, const ds4_vision_span *images, size_t image_count) { if (!s || (image_count != 0 && !images)) return false; if (!s->checkpoint_valid) return false; if (s->checkpoint_image_count > image_count) return false; for (size_t i = 0; i < s->checkpoint_image_count; i++) { const ds4_vision_identity *old = &s->checkpoint_images[i]; const ds4_vision_span *current = &images[i]; if (old->token_start != current->token_start || old->token_count != current->embedding.token_count || memcmp(old->fingerprint, current->embedding.fingerprint, sizeof(old->fingerprint)) != 0) return false; } if (s->checkpoint_image_count < image_count) { const ds4_vision_span *next = &images[s->checkpoint_image_count]; if (next->token_start < (uint32_t)s->checkpoint.len) return false; } return true; } bool ds4_session_vision_state_matches( const ds4_session *s, const ds4_vision_span *images, size_t image_count) { return s && s->checkpoint_valid && s->checkpoint_image_count == image_count && ds4_session_vision_prefix_matches(s, images, image_count); } bool ds4_session_vision_fingerprint_prefix_matches( const ds4_session *s, const ds4_vision_span *images, size_t image_count) { if (!s || !s->checkpoint_valid) return false; if ((image_count != 0 && !images)) return false; if (s->checkpoint_image_count > image_count) return false; for (size_t i = 0; i < s->checkpoint_image_count; i++) { const ds4_vision_identity *old = &s->checkpoint_images[i]; const ds4_vision_span *current = &images[i]; if (old->token_count != current->embedding.token_count || memcmp(old->fingerprint, current->embedding.fingerprint, sizeof(old->fingerprint)) != 0) return false; } return true; } bool ds4_session_rebase_vision_state(const ds4_session *s, ds4_vision_span *images, size_t image_count) { if (!s || !s->checkpoint_valid || (image_count && !images) || image_count != s->checkpoint_image_count) return false; for (size_t i = 0; i < image_count; i++) { if (images[i].embedding.token_count != s->checkpoint_images[i].token_count || memcmp(images[i].embedding.fingerprint, s->checkpoint_images[i].fingerprint, sizeof(images[i].embedding.fingerprint))) return false; } for (size_t i = 0; i < image_count; i++) images[i].token_start = s->checkpoint_images[i].token_start; return true; } bool ds4_session_has_vision_state(const ds4_session *s) { return s && (s->checkpoint_image_count != 0 || s->sync_image_count != 0); } static bool ds4_session_vision_range_overlaps( const ds4_session *s, uint32_t token_start, uint32_t token_count) { const uint64_t range_start = token_start; const uint64_t range_end = range_start + token_count; for (size_t i = 0; i < s->sync_image_count; i++) { const uint64_t image_start = s->sync_images[i].token_start; const uint64_t image_end = image_start + s->sync_images[i].embedding.token_count; if (range_start < image_end && image_start < range_end) return true; } return false; } static bool ds4_session_store_vision_identities(ds4_session *s) { ds4_vision_identity *copy = NULL; if (s->sync_image_count != 0) { if (s->sync_image_count > SIZE_MAX / sizeof(copy[0])) return false; copy = calloc(s->sync_image_count, sizeof(copy[0])); if (!copy) return false; for (size_t i = 0; i < s->sync_image_count; i++) { copy[i].token_start = s->sync_images[i].token_start; copy[i].token_count = s->sync_images[i].embedding.token_count; memcpy(copy[i].fingerprint, s->sync_images[i].embedding.fingerprint, sizeof(copy[i].fingerprint)); } } free(s->checkpoint_images); s->checkpoint_images = copy; s->checkpoint_image_count = s->sync_image_count; return true; } typedef struct { ds4_session *session; ds4_session_progress_fn progress, display; void *progress_ud, *display_ud; ds4_session_cancel_fn cancel; void *cancel_ud; bool cancelled, failed; } ds4_tp_sync_progress; static void ds4_session_tp_sync_checkpoint(ds4_tp_sync_progress *p, uint32_t point, int current, int total) { if (p->failed) return; const bool requested = p->cancelled || (p->cancel && p->cancel(p->cancel_ud)); p->failed = !ds4_tp_sync_checkpoint(p->session->engine->tp.ctx, point, current, total, requested, &p->cancelled); } static bool ds4_session_tp_sync_cancelled(void *ud) { const ds4_tp_sync_progress *p = ud; return p->cancelled || p->failed; } static void ds4_session_tp_sync_progress(void *ud, const char *event, int current, int total) { ds4_tp_sync_progress *p = ud; if (p->progress) p->progress(p->progress_ud, event, current, total); if (!p->cancelled) ds4_session_tp_sync_checkpoint(p, !strcmp(event, "prefill_chunk") ? 1u : 2u, current, total); } static void ds4_session_tp_sync_display(void *ud, const char *event, int current, int total) { ds4_tp_sync_progress *p = ud; if (p->display) p->display(p->display_ud, event, current, total); if (!p->cancelled) ds4_session_tp_sync_checkpoint(p, 2u, current, total); } /* Progress boundaries have the same graph position on both ranks. Poll the * real cancellation callback there, then expose only the agreed result to * inner loops (including timing-dependent Engram I/O polling). */ static int ds4_session_sync_lockstep(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen) { ds4_tp_sync_progress p = {.session = s, .progress = s->progress, .display = s->display_progress, .progress_ud = s->progress_ud, .display_ud = s->display_progress_ud, .cancel = s->cancel, .cancel_ud = s->cancel_ud}; s->progress = ds4_session_tp_sync_progress; s->display_progress = ds4_session_tp_sync_display; s->progress_ud = s->display_progress_ud = &p; s->cancel = ds4_session_tp_sync_cancelled; s->cancel_ud = &p; ds4_session_tp_sync_checkpoint(&p, 0u, 0, prompt->len); int rc = p.failed ? 1 : ds4_session_sync_internal(s, prompt, err, errlen); ds4_session_tp_sync_checkpoint(&p, 3u, rc, prompt->len); s->progress = p.progress; s->display_progress = p.display; s->progress_ud = p.progress_ud; s->display_progress_ud = p.display_ud; s->cancel = p.cancel; s->cancel_ud = p.cancel_ud; if (p.failed) { snprintf(err, errlen, "tp: prefill synchronization failed"); return 1; } if (p.cancelled && rc == 0) { snprintf(err, errlen, "interrupted"); rc = DS4_SESSION_SYNC_INTERRUPTED; } return rc; } /* 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) { if (s && s->checkpoint_valid && !ds4_session_vision_prefix_matches( s, s->sync_images, s->sync_image_count)) { ds4_session_invalidate(s); } #ifndef DS4_NO_GPU ds4_session_dspark_scheduler_begin_request(s); #endif const bool mirror = ds4_session_tp_leader(s); if (mirror && (!prompt || prompt->len <= 0)) return ds4_session_sync_internal(s, prompt, err, errlen); if (mirror && prompt && prompt->len > 0) { if (s->sync_image_count > UINT32_MAX) { snprintf(err, errlen, "tp: too many image spans to synchronize"); return 1; } int sent = s->sync_image_count ? ds4_tp_send_sync_multimodal( s->engine->tp.ctx, s->tp_session_id, prompt->v, (uint32_t)prompt->len, s->sync_images, (uint32_t)s->sync_image_count) : ds4_tp_send_sync(s->engine->tp.ctx, s->tp_session_id, prompt->v, (uint32_t)prompt->len); if (!sent) { snprintf(err, errlen, "tp: worker sync send failed"); return 1; } } int rc = s && s->engine && s->engine->tp.active && prompt && prompt->len > 0 ? ds4_session_sync_lockstep(s, prompt, err, errlen) : 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) { int worker_status = 1; const bool acknowledged = !ds4_tp_failed(s->engine->tp.ctx) && ds4_tp_wait_command_status( s->engine->tp.ctx, s->tp_session_id, &worker_status, "prefill sync", err, errlen); const bool worker_ok = acknowledged && worker_status == 0; 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); if (rc == 0 && acknowledged && worker_status != 0) snprintf(err, errlen, "tp: worker prefill failed (status %d)", worker_status); return rc != 0 ? rc : 1; } } if (rc == 0 && !ds4_session_store_vision_identities(s)) { ds4_session_invalidate(s); snprintf(err, errlen, "unable to retain image prompt identity"); return 1; } return rc; } int ds4_session_sync_multimodal( ds4_session *s, const ds4_tokens *prompt, const ds4_vision_span *images, size_t image_count, char *err, size_t errlen) { if (!s || !prompt || (image_count != 0 && !images)) { snprintf(err, errlen, "invalid multimodal prompt"); return 1; } if (image_count != 0 && !s->engine->vision_ready) { snprintf(err, errlen, "vision encoder is not loaded"); return 1; } uint64_t previous_end = 0; for (size_t i = 0; i < image_count; i++) { const ds4_vision_span *span = &images[i]; const uint64_t end = (uint64_t)span->token_start + span->embedding.token_count; if (!span->embedding.data || span->embedding.token_count == 0 || span->token_start < previous_end || end > (uint64_t)prompt->len) { snprintf(err, errlen, "invalid or overlapping image token span"); return 1; } if (s->engine->vision_kind == DS4_VISION_DEEPSEEK4 && !ds4_session_is_ds41(s)) { uint32_t cursor = 0, block_start = 0, image_start = 0, image_end = 0; const int parsed = ds4_deepseek4_next_image_span( prompt->v + span->token_start, span->embedding.token_count, DS4_N_VOCAB, &cursor, &block_start, &image_start, &image_end); if (parsed != 1 || block_start != 0 || image_end + 1u != cursor || cursor != span->embedding.token_count) { snprintf(err, errlen, "image span does not cover one complete DeepSeek image block"); return 1; } } else { for (uint64_t token = span->token_start; token < end; token++) { if (prompt->v[token] != s->engine->vision_image_token) { snprintf(err, errlen, "image span does not cover image placeholder tokens"); return 1; } } } previous_end = end; } s->sync_images = images; s->sync_image_count = image_count; #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) { s->ds41_graph.images = images; s->ds41_graph.image_count = image_count; } #endif #ifndef DS4_NO_GPU s->graph.prefill_vision_spans = images; s->graph.prefill_vision_span_count = image_count; #endif const int rc = ds4_session_sync(s, prompt, err, errlen); #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) { s->ds41_graph.images = NULL; s->ds41_graph.image_count = 0; } #endif #ifndef DS4_NO_GPU s->graph.prefill_vision_spans = NULL; s->graph.prefill_vision_span_count = 0; #endif s->sync_images = NULL; s->sync_image_count = 0; 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); } #ifdef DS4_HAS_QWEN4_GPU if (ds4_session_is_qwen4(s)) { ds4_engine *e = s->engine; int start = 0; s->glm_mtp_have = 0; s->glm_mtp_have2 = false; if (s->checkpoint_valid && s->qwen4_graph.pos == (uint32_t)s->checkpoint.len && (!s->qwen4_rewound || prompt->len > s->checkpoint.len) && prompt->len >= s->checkpoint.len && ds4_tokens_starts_with(prompt, &s->checkpoint)) { start = s->checkpoint.len; } else { qwen4_graph_reset(&s->qwen4_graph); s->checkpoint.len = 0; s->checkpoint_valid = false; } for (int i = start; i < prompt->len; i++) { if (prompt->v[i] < 0 || prompt->v[i] >= (int)DS4_N_VOCAB) { snprintf(err, errlen, "token id %d at position %d is outside the vocabulary", prompt->v[i], i); return 1; } } s->qwen4_graph.vis_spans = s->sync_images; s->qwen4_graph.vis_span_count = s->sync_image_count; /* A one-token prefill tail is still a prompt row. Speculative verify * batches, despite having multiple rows, must never overwrite dumps. */ s->qwen4_graph.dump_prompt_rows = true; int prefill_rc = 0; for (int i = start; i < prompt->len;) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = s->checkpoint.len > 0; prefill_rc = DS4_SESSION_SYNC_INTERRUPTED; break; } uint32_t chunk = (uint32_t)(prompt->len - i); if (chunk > s->qwen4_graph.cap_tokens) chunk = s->qwen4_graph.cap_tokens; /* Progress callbacks may persist this frontier, and cancellation * may leave it as the live session. Every completed chunk needs * its own logits as well as recurrent/KV state. */ s->checkpoint_valid = false; if (!qwen4_graph_forward_tokens(&s->qwen4_graph, &e->model, &e->weights, prompt->v + i, chunk, s->logits, false)) { snprintf(err, errlen, "Qwen3.8 prefill failed at token %d", i); s->checkpoint_valid = false; prefill_rc = 1; break; } for (uint32_t j = 0; j < chunk; j++) token_vec_push(&s->checkpoint, prompt->v[i + (int)j]); i += (int)chunk; s->checkpoint_valid = true; s->qwen4_rewound = false; s->mtp_draft_valid = false; if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); } s->qwen4_graph.vis_spans = NULL; s->qwen4_graph.vis_span_count = 0; s->qwen4_graph.dump_prompt_rows = false; if (prefill_rc != 0) return prefill_rc; s->checkpoint_valid = true; s->mtp_draft_valid = false; return 0; } #endif 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; #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) { ds41_gpu_graph *g = &s->ds41_graph; if (!s->ds41_graph_ready) { snprintf(err, errlen, "V4.1 graph is not initialized"); return 1; } if (!s->checkpoint_valid || !g->valid || g->pos != (uint32_t)s->checkpoint.len || !ds4_tokens_starts_with(prompt, &s->checkpoint)) { ds41_graph_reset(g); s->checkpoint.len = 0; s->checkpoint_valid = false; } bool pending_logits = false, interrupted = false, decoder_pending = false; ds41_encoder_residency encoder = {0}; ds41_encoder_acquire(g, &e->model, &e->weights, (uint32_t)(prompt->len - s->checkpoint.len), &encoder, s->cancel, s->cancel_ud); for (int i = s->checkpoint.len; i < prompt->len;) { if (ds4_session_cancelled(s)) { interrupted = true; break; } const uint32_t remaining = (uint32_t)(prompt->len - i); if (g->encoder_resident && remaining < 768u) ds41_encoder_release(g, &e->model, &encoder); const uint32_t short_count = decoder_pending ? 0u : ds41_short_prefill_count(g, &e->weights, remaining); const uint32_t count = short_count ? short_count : g->encoder_resident ? (remaining - 512u < g->prefill_cap ? remaining - 512u : g->prefill_cap) : ds41_prefill_count(g, remaining); const bool defer_decoder = !g->encoder_resident && count >= 16384u && remaining - count >= 8192u && !getenv("DS4_METAL_DISABLE_V41_DECODER_SUFFIX") && !getenv("DS4_METAL_DISABLE_V41_DEFER_DECODER"); /* Never add an encoder sweep just to defer the decoder. A short * remainder finishes the pending decoder here, then runs normally. */ const bool layer_major = count > 1u; const bool ok = short_count ? ds41_graph_short_prefill(g, &e->model, &e->weights, prompt->v + i, count) : layer_major ? (defer_decoder || decoder_pending ? ds41_graph_prefill_sweep(g, &e->model, &e->weights, prompt->v + i, count, s->progress, s->progress_ud, prompt->len, s->cancel, s->cancel_ud, defer_decoder, decoder_pending) : ds41_graph_prefill(g, &e->model, &e->weights, prompt->v + i, count, s->progress, s->progress_ud, prompt->len, s->cancel, s->cancel_ud)) : ds41_graph_step(g, &e->model, &e->weights, prompt->v[i], NULL); if (!ok) { ds41_encoder_release(g, &e->model, &encoder); s->checkpoint_valid = false; if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); return DS4_SESSION_SYNC_INTERRUPTED; } snprintf(err, errlen, "V4.1 prefill failed at position %d", i); return 1; } for (uint32_t t = 0; t < count; t++) token_vec_push(&s->checkpoint, prompt->v[i + (int)t]); i += (int)count; decoder_pending = defer_decoder; pending_logits = true; s->checkpoint_valid = false; if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); } ds41_encoder_release(g, &e->model, &encoder); if (decoder_pending) { s->checkpoint_valid = false; g->valid = false; snprintf(err, errlen, "interrupted"); return DS4_SESSION_SYNC_INTERRUPTED; } /* Prefill needs only its final output head. An interrupted prefix must * still have matching logits before it can be saved or continued. */ if (pending_logits) { if (!ds41_graph_logits(g, &e->model, &e->weights, s->logits)) { g->valid = false; snprintf(err, errlen, "V4.1 prefill output head failed"); return 1; } s->checkpoint_valid = true; if (s->progress) s->progress(s->progress_ud, "prefill_chunk", s->checkpoint.len, prompt->len); } if (interrupted) { snprintf(err, errlen, "interrupted"); return DS4_SESSION_SYNC_INTERRUPTED; } return 0; } #endif 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); if (!ds4_session_glm_reset_kda_state(s)) { snprintf(err, errlen, "%s GLM KDA state reset failed", backend_name); return 1; } } /* GLM-5.3 uses bounded layer-major chunks for mHC and recurrent KDA. * DSA switches from dense attention to the pooled sparse graph at the * full-attention boundary without returning to a token-major loop. * Very short checkpoint extensions retain the lower-latency decode * path; image embeddings still require the batch prefill path. */ if (s->glm_graph.glm53) { const uint32_t suffix = (uint32_t)(prompt->len - start); const bool short_resume = resumed_checkpoint && suffix > 0 && suffix < glm53_graph_resume_prefill_min_tokens() && !ds4_session_vision_range_overlaps(s, (uint32_t)start, suffix); for (int i = start; i < prompt->len; ) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = s->checkpoint.len > 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 > DS4_GLM53_PREFILL_CHUNK_TOKENS) { chunk = DS4_GLM53_PREFILL_CHUNK_TOKENS; } if (short_resume) chunk = 1; const bool use_indexed = !short_resume && (glm53_graph_use_indexed_prefill(&s->glm_graph) || (pos >= s->glm_graph.ctx_cap && s->glm_graph.indexed_prefill_cap != 0)); if (use_indexed) { if (chunk > s->glm_graph.indexed_prefill_cap) { chunk = s->glm_graph.indexed_prefill_cap; } chunk = glm_graph_limit_indexed_prefill_chunk(&s->glm_graph, pos, chunk); } else if (!short_resume && pos < s->glm_graph.ctx_cap) { const uint32_t dense_left = s->glm_graph.ctx_cap - pos; if (chunk > dense_left) chunk = dense_left; } else if (!short_resume) { chunk = 1; } float *chunk_logits = i + (int)chunk == prompt->len ? s->logits : NULL; const bool ok = short_resume ? glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, prompt->v[i], NULL, pos, NULL, chunk_logits, false) : use_indexed ? glm_graph_forward_indexed_tokens( &s->glm_graph, &e->model, &e->weights, prompt->v + i, NULL, s->sync_images, s->sync_image_count, pos, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, (uint32_t)start, (uint32_t)(i - start), (uint32_t)(prompt->len - start)) : pos < s->glm_graph.ctx_cap ? glm_graph_forward_tokens(&s->glm_graph, &e->model, &e->weights, prompt->v + i, NULL, s->sync_images, s->sync_image_count, pos, chunk, NULL, chunk_logits, s->display_progress, s->display_progress_ud, (uint32_t)start, (uint32_t)(i - start), (uint32_t)(prompt->len - start)) : glm_graph_forward_token(&s->glm_graph, &e->model, &e->weights, prompt->v[i], NULL, pos, NULL, chunk_logits, false); if (!ok) { snprintf(err, errlen, "%s GLM-5.3 prefill failed at token %d", backend_name, i); 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]); } s->checkpoint_valid = true; s->mtp_draft_valid = false; if (pos < s->glm_graph.ctx_cap) { ds4_session_glm_note_dense_cache(s, pos, chunk); } i += (int)chunk; if (s->progress) { s->progress(s->progress_ud, "prefill_chunk", i, prompt->len); } } glm_debug_dump_prefill_logits(s->logits); return 0; } 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++) { if (ds4_session_cancelled(s)) { snprintf(err, errlen, "interrupted"); s->checkpoint_valid = s->checkpoint.len > 0; s->mtp_draft_valid = false; ds4_session_glm_cap_dense_cache(s); return DS4_SESSION_SYNC_INTERRUPTED; } 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, s->sync_images, s->sync_image_count, 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(&s->glm_graph, 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, s->sync_images, s->sync_image_count, 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(&s->glm_graph, (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, s->sync_images, s->sync_image_count, (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, s->sync_images, s->sync_image_count, (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); if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i + 1, prompt->len); } 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) { if (!s || !s->checkpoint_valid || !s->logits) return -1; return sample_argmax(s->logits, DS4_N_VOCAB); } int ds4_session_argmax_excluding(ds4_session *s, int excluded_id) { if (!s || !s->checkpoint_valid || !s->logits) return -1; if (getenv("DS4_CPU_DISABLE_UNROLLED_ARGMAX") == NULL) { return argmax_f32_excluding_unrolled8( s->logits, DS4_N_VOCAB, excluded_id); } 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_session_argmax_ignoring_eos(ds4_session *s, ds4_think_mode think_mode) { if (!s || !s->checkpoint_valid || !s->logits) return -1; int best = -1; float best_logit = DS4_NEG_INF; for (uint32_t i = 0; i < DS4_N_VOCAB; i++) { if (ds4_token_is_stop_for_think_mode(s->engine, (int)i, think_mode)) { 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) { if (!s || !s->checkpoint_valid || !s->logits) return -1; 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 bool stochastic_requested = s->dspark_sample_rng != NULL && s->dspark_sample_temperature > 0.0f; const float confidence_threshold = stochastic_requested && !s->engine->dspark_confidence_threshold_set && s->engine->dspark_confidence_threshold < 0.8f ? 0.8f : s->engine->dspark_confidence_threshold; const bool stats_enabled = ds4_dspark_stats_enabled(); const bool scheduler_enabled = ds4_dspark_scheduler_enabled(s); const bool time_enabled = stats_enabled || (scheduler_enabled && ds4_dspark_scheduler_timing_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_stochastic_draft = false; s->dspark_last_confidence0 = 0.0f; s->dspark_last_confidence0_valid = false; if (scheduler_enabled) s->dspark_last_propose_ms = 0.0; const bool align_rocm = ds4_session_dspark_rocm_gfx1151_reference_alignment(s); if (align_rocm && (pos == 0 || pos != (uint32_t)s->checkpoint.len)) return false; const uint32_t feature_pos = align_rocm ? pos - 1u : pos; if (enabled && !fake_argmax_enabled && ds4_session_dspark_scheduler_should_skip(s)) { if (align_rocm) { (void)metal_graph_dspark_cache_target_prefix(&s->graph, feature_pos); } (void)metal_graph_dspark_ring_maintain(&s->graph, &s->engine->mtp_model, &s->engine->dspark_weights, feature_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 uint32_t verify_cap = ds4_dspark_env_u32( "DS4_DSPARK_VERIFY_CAP", ds4_session_dspark_rocm_gfx1151_fast_path(s) ? 5u : 0u); const uint32_t proposal_cap = verify_cap != 0 && verify_cap < dw->block_size ? verify_cap : dw->block_size; 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; if (cache_window_ok && align_rocm) { /* The captured batch already includes h[L-1]; overwrite that row instead of duplicating it at L. */ cache_window_ok = metal_graph_dspark_cache_target_prefix(&s->graph, feature_pos); } } else { cache_window_ok = align_rocm ? metal_graph_dspark_cache_target_prefix(&s->graph, feature_pos) : metal_graph_dspark_cache_crop_to_prefix(&s->graph, feature_pos) && metal_graph_dspark_cache_ends_at(&s->graph, feature_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; const bool fuse_final_hidden = stage_chain_ready && !probe_log && confidence_threshold > 0.0f && dspark_confidence_probe_ready(dw); if (stage_chain_ready) { stage_chain_ok = metal_graph_eval_dspark_stage_chain(&s->graph, &s->engine->mtp_model, dw, pos, fuse_final_hidden, &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 = fuse_final_hidden || metal_graph_eval_dspark_final_hidden(&s->graph, &s->engine->mtp_model, dw, false); 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, proposal_cap); 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) { #ifndef __APPLE__ const char *gpu_markov = getenv("DS4_ROCM_DSPARK_GREEDY_GPU_MARKOV"); const bool gpu_markov_enabled = gpu_markov && gpu_markov[0] ? gpu_markov[0] != '0' : ds4_session_dspark_rocm_gfx1151_fast_path(s); if (gpu_markov_enabled && !dspark_markov_bias_disabled() && getenv("DS4_DSPARK_NO_GPU_MARKOV") == NULL) { const double markov_t0 = DS4_DSPARK_PROP_T0(); markov_ok = dspark_apply_markov_greedy_gpu_runtime( &s->graph, &s->engine->mtp_model, dw, token, proposal_cap, markov_proposal, &markov_proposal_len); DS4_DSPARK_PROP_ADD(propose_markov_ms, markov_t0); } #endif const uint64_t logits_count = (uint64_t)proposal_cap * (uint64_t)DS4_N_VOCAB; if (!markov_ok && 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, proposal_cap, 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 && confidence_threshold > 0.0f && 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 (verify_cap != 0 && s->dspark_draft_len > verify_cap) { s->dspark_draft_len = verify_cap; } if (s->dspark_draft_len > DS4_DSPARK_MAX_BLOCK_SIZE) { s->dspark_draft_len = DS4_DSPARK_MAX_BLOCK_SIZE; } if (stochastic_requested && s->dspark_draft_len < 2u) { s->dspark_draft_len = 0; } 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; s->dspark_stochastic_draft = s->dspark_draft_valid && stochastic_requested; } 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; s->dspark_stochastic_draft = false; 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; #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(s)) { if (!s->ds41_graph_ready || (!s->checkpoint_valid && s->checkpoint.len != 0) || s->ds41_graph.pos != (uint32_t)s->checkpoint.len || !ds41_graph_step(&s->ds41_graph, &e->model, &e->weights, token, s->logits)) { s->checkpoint_valid = false; if (errlen) snprintf(err, errlen, "V4.1 decode failed at position %d", s->checkpoint.len); return 1; } token_vec_push(&s->checkpoint, token); s->checkpoint_valid = true; return 0; } #endif #ifdef DS4_HAS_QWEN4_GPU if (ds4_session_is_qwen4(s)) { if (!s->qwen4_graph_ready) { if (errlen) snprintf(err, errlen, "Qwen3.8 graph is not initialized"); return 1; } if (qwen4_session_replay_if_stale(s, err, errlen) != 0) return 1; if (s->qwen4_graph.pos >= s->qwen4_graph.ctx_cap) { if (errlen) snprintf(err, errlen, "context is full"); return 1; } if (!s->glm_spec_inside) { s->glm_mtp_have = 0; s->glm_mtp_have2 = false; } if (!qwen4_graph_forward_token(&s->qwen4_graph, &e->model, &e->weights, token, s->logits)) { if (errlen) snprintf(err, errlen, "Qwen3.8 decode failed"); s->checkpoint_valid = false; return 1; } token_vec_push(&s->checkpoint, token); s->checkpoint_valid = true; s->qwen4_rewound = false; s->mtp_draft_valid = false; (void)probe_mtp; return 0; } #endif if (ds4_session_is_glm(s)) { 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; } if (!s->glm_spec_inside) { s->glm_mtp_rollback_valid = false; s->glm_mtp_have = 0; } 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(s) && ds4_dspark_scheduler_timing_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 (!s || !s->checkpoint_valid) { payload_set_err(err, errlen, "decode requires a synchronized checkpoint"); return 1; } 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. */ 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 glm53_graph_native_session_batch_supported( ds4_decode_item *items, int count) { if (!items || count < 2 || count > (int)DS4_GLM53_PREFILL_CHUNK_TOKENS || DS4_N_ROT != 0 || g_expert_profile.active) { return false; } if (!items[0].session) return false; ds4_glm_gpu_graph *first = &items[0].session->glm_graph; if (!first->glm53 || first->ssd_streaming || first->placement || !first->has_token_embd || !first->has_output_head || (!first->full_kv_cache && (first->compact_cache_cap == 0 || !first->batch_qk_low)) || (uint32_t)count > first->ctx_cap) { return false; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t pos = (uint32_t)s->checkpoint.len; if (!g->glm53 || g->ssd_streaming || g->placement || g->full_kv_cache != first->full_kv_cache || g->compact_cache_cap != first->compact_cache_cap || g->layer_start != first->layer_start || g->layer_end != first->layer_end || g->tp_world != first->tp_world || g->tp_rank != first->tp_rank || pos >= glm_graph_dense_compact_attention_limit(g)) { return false; } } return true; } static bool glm53_graph_encode_kda_session_batch( ds4_decode_item *items, int count, ds4_glm_gpu_graph *batch, const ds4_model *model, const ds4_layer_weights *l, uint32_t il) { const uint32_t rows = (uint32_t)count; const uint64_t projection = (uint64_t)DS4_N_KDA_HEAD * DS4_N_KDA_HEAD_DIM; bool ok = glm53_graph_matmul_rows(batch->batch_kda_q, model, l->kda_q, DS4_N_EMBD, (uint32_t)projection, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_k, model, l->kda_k, DS4_N_EMBD, (uint32_t)projection, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_v, model, l->kda_v, DS4_N_EMBD, (uint32_t)projection, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_lowrank, model, l->kda_f_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_raw_gate, model, l->kda_f_b, DS4_N_KDA_HEAD_DIM, (uint32_t)projection, batch->batch_kda_lowrank, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_raw_beta, model, l->kda_beta, DS4_N_EMBD, DS4_N_KDA_HEAD, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_lowrank, model, l->kda_g_a, DS4_N_EMBD, DS4_N_KDA_HEAD_DIM, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows(batch->batch_kda_output_gate, model, l->kda_g_b, DS4_N_KDA_HEAD_DIM, (uint32_t)projection, batch->batch_kda_lowrank, rows); for (int i = 0; ok && i < count; i++) { ds4_glm_gpu_graph *g = &items[i].session->glm_graph; ds4_gpu_tensor *q = glm_graph_tensor_row_view_strided( batch->batch_kda_q, (uint32_t)i, projection, projection); ds4_gpu_tensor *k = glm_graph_tensor_row_view_strided( batch->batch_kda_k, (uint32_t)i, projection, projection); ds4_gpu_tensor *v = glm_graph_tensor_row_view_strided( batch->batch_kda_v, (uint32_t)i, projection, projection); ds4_gpu_tensor *raw_gate = glm_graph_tensor_row_view_strided( batch->batch_kda_raw_gate, (uint32_t)i, projection, projection); ds4_gpu_tensor *raw_beta = glm_graph_tensor_row_view_strided( batch->batch_kda_raw_beta, (uint32_t)i, DS4_N_KDA_HEAD, DS4_N_KDA_HEAD); ds4_gpu_tensor *output_gate = glm_graph_tensor_row_view_strided( batch->batch_kda_output_gate, (uint32_t)i, projection, projection); ds4_gpu_tensor *out = glm_graph_tensor_row_view_strided( batch->batch_kda_out, (uint32_t)i, projection, projection); ok = q && k && v && raw_gate && raw_beta && output_gate && out; if (ok) { ok = ds4_gpu_glm53_kda_decode( out, g->layer_kda_conv_state[il], g->layer_kda_recurrent_state[il], q, k, v, raw_gate, raw_beta, output_gate, model->map, model->size, l->kda_q_conv->abs_offset, l->kda_k_conv->abs_offset, l->kda_v_conv->abs_offset, l->kda_a_log->abs_offset, l->kda_dt_bias->abs_offset, l->kda_o_norm->abs_offset, DS4_N_KDA_HEAD, 1, DS4_KDA_GATE_LOWER_BOUND, DS4_RMS_EPS) != 0; } ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(output_gate); ds4_gpu_tensor_free(raw_beta); ds4_gpu_tensor_free(raw_gate); ds4_gpu_tensor_free(v); ds4_gpu_tensor_free(k); ds4_gpu_tensor_free(q); } if (ok) ok = glm53_graph_matmul_rows(batch->batch_attn_out, model, l->kda_output, (uint32_t)projection, DS4_N_EMBD, batch->batch_kda_out, rows); return ok; } static bool glm53_graph_encode_dsa_session_batch( ds4_decode_item *items, int count, ds4_glm_gpu_graph *batch, const ds4_model *model, const ds4_layer_weights *l, uint32_t il) { const uint32_t rows = (uint32_t)count; const uint32_t kv_raw_dim = (uint32_t)l->attn_kv_a_mqa->dim[1]; bool ok = glm_graph_matmul_q8_0_tensor(batch->batch_q_rank, model, l->attn_q_a->abs_offset, DS4_N_EMBD, DS4_N_LORA_Q, batch->batch_attn_norm, rows); if (ok) ok = glm_graph_matmul_q8_0_tensor(batch->batch_kv_raw, model, l->attn_kv_a_mqa->abs_offset, DS4_N_EMBD, kv_raw_dim, batch->batch_attn_norm, rows); if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( batch->batch_q_rank_norm, batch->batch_q_rank, model->map, model->size, l->attn_q_a_norm->abs_offset, DS4_N_LORA_Q, batch->batch_kv_norm, batch->batch_kv_raw, l->attn_kv_a_norm->abs_offset, DS4_N_KV_LORA, rows, DS4_RMS_EPS) != 0; if (ok) ok = glm_graph_matmul_q8_0_tensor(batch->batch_q, model, l->attn_q_b->abs_offset, DS4_N_LORA_Q, batch->q_dim, batch->batch_q_rank_norm, rows); if (ok && batch->full_kv_cache) { ok = ds4_gpu_glm_k_b_project_typed_tensor( batch->batch_k_nope, batch->batch_kv_norm, model->map, model->size, l->attn_k_b->abs_offset, l->attn_k_b->type, rows, DS4_N_KV_LORA, (uint32_t)batch->q_nope, DS4_N_HEAD) != 0; } if (ok && batch->full_kv_cache) { ok = glm_graph_matmul_q8_0_tensor(batch->batch_value, model, l->attn_v_b->abs_offset, DS4_N_KV_LORA, batch->heads_dim, batch->batch_kv_norm, rows); } const bool mirror_compact = batch->compact_cache_cap != 0; const bool mirror_indexer = mirror_compact && glm_graph_layer_uses_full_indexer(il); if (ok && mirror_indexer) { ok = glm53_graph_matmul_rows( batch->batch_indexer_k, model, l->indexer_attn_k, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, batch->batch_attn_norm, rows); if (ok) ok = glm53_graph_matmul_rows( batch->batch_indexer_gate, model, l->indexer_compressor_gate, DS4_N_EMBD, DS4_N_INDEXER_HEAD_DIM, batch->batch_attn_norm, rows); } for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ds4_glm_gpu_graph *g = &s->glm_graph; const uint32_t pos = (uint32_t)s->checkpoint.len; const float rope_base = layer_rope_freq_base(il); const float rope_scale = layer_rope_freq_scale(il); ds4_gpu_tensor *q = glm_graph_tensor_row_view_strided( batch->batch_q, (uint32_t)i, batch->q_dim, batch->q_dim); ds4_gpu_tensor *kv_raw = glm_graph_tensor_row_view_strided( batch->batch_kv_raw, (uint32_t)i, batch->kv_raw_dim, kv_raw_dim); ds4_gpu_tensor *k_nope = batch->full_kv_cache ? glm_graph_tensor_row_view_strided( batch->batch_k_nope, (uint32_t)i, (uint64_t)DS4_N_HEAD * batch->q_nope, (uint64_t)DS4_N_HEAD * batch->q_nope) : NULL; ds4_gpu_tensor *value = batch->full_kv_cache ? glm_graph_tensor_row_view_strided( batch->batch_value, (uint32_t)i, batch->heads_dim, batch->heads_dim) : NULL; ds4_gpu_tensor *qk_low = !batch->full_kv_cache ? glm_graph_tensor_row_view_strided( batch->batch_qk_low, (uint32_t)i, (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA) : NULL; ds4_gpu_tensor *kv_norm = mirror_compact ? glm_graph_tensor_row_view_strided( batch->batch_kv_norm, (uint32_t)i, DS4_N_KV_LORA, DS4_N_KV_LORA) : NULL; ds4_gpu_tensor *indexer_k = mirror_indexer ? glm_graph_tensor_row_view_strided( batch->batch_indexer_k, (uint32_t)i, DS4_N_INDEXER_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM) : NULL; ds4_gpu_tensor *indexer_gate = mirror_indexer ? glm_graph_tensor_row_view_strided( batch->batch_indexer_gate, (uint32_t)i, DS4_N_INDEXER_HEAD_DIM, DS4_N_INDEXER_HEAD_DIM) : NULL; ds4_gpu_tensor *heads = glm_graph_tensor_row_view_strided( batch->batch_heads, (uint32_t)i, batch->heads_dim, batch->heads_dim); ok = q && kv_raw && heads && (batch->full_kv_cache ? (k_nope && value) : qk_low != NULL) && (!mirror_compact || kv_norm) && (!mirror_indexer || (indexer_k && indexer_gate)); if (ok && mirror_indexer) { ok = ds4_gpu_glm53_indexer_pool_update_tensor( g->layer_indexer_key_cache[il], g->layer_indexer_tail_k[il], g->layer_indexer_tail_gate[il], indexer_k, indexer_gate, model->map, model->size, l->indexer_k_norm->abs_offset, l->indexer_k_norm_b->abs_offset, l->indexer_compressor_ape->abs_offset, pos, 1, g->compact_cache_cap, DS4_N_INDEXER_HEAD_DIM, DS4_GLM53_INDEX_POOL_SIZE, 1.0e-6f, glm_graph_compact_cache_is_f16()) != 0; } if (ok && mirror_compact) { ok = ds4_gpu_glm_store_compact_kv_tensor( g->layer_kv_lora_cache[il], g->layer_k_rope_cache[il], kv_norm, 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 (ok && !g->full_kv_cache) { const uint32_t visible = pos + 1u; ok = ds4_gpu_glm_fill_selected_range_tensor( g->indexer_selected, visible) != 0; if (ok) ok = ds4_gpu_glm_qk_lowrank_typed_tensor( qk_low, 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) ok = ds4_gpu_glm_attention_indexed_decode_typed_tensor( heads, q, 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, g->indexer_selected, visible, 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_build_kv_cache_tensor( g->layer_key_cache[il], g->layer_value_cache[il], kv_raw, k_nope, 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; if (ok) ok = ds4_gpu_glm_attention_full_tensor( heads, 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_gpu_tensor_free(heads); ds4_gpu_tensor_free(indexer_gate); ds4_gpu_tensor_free(indexer_k); ds4_gpu_tensor_free(kv_norm); ds4_gpu_tensor_free(qk_low); ds4_gpu_tensor_free(value); ds4_gpu_tensor_free(k_nope); ds4_gpu_tensor_free(kv_raw); ds4_gpu_tensor_free(q); } if (ok && batch->tp_world == 2) { ok = glm_graph_tp_batch_bounce_ready(batch, rows); const uint64_t k_cnt = batch->heads_dim / 2u; const uint64_t k_off = (uint64_t)batch->tp_rank * k_cnt; for (int i = 0; ok && i < count; i++) { ds4_gpu_tensor *partial = glm_graph_tensor_row_view_strided( batch->tp_bounce_out, (uint32_t)i, DS4_N_EMBD, DS4_N_EMBD); ds4_gpu_tensor *heads = glm_graph_tensor_row_view_strided( batch->batch_heads, (uint32_t)i, batch->heads_dim, batch->heads_dim); ok = partial && heads && metal_graph_matmul_dense_quant_kslice( partial, model, l->attn_output, batch->heads_dim, k_off, k_cnt, DS4_N_EMBD, heads, k_off); ds4_gpu_tensor_free(heads); ds4_gpu_tensor_free(partial); } if (ok) ok = glm_graph_tp_batch_ffn_combine( batch, il, batch->batch_attn_out, rows); } else if (ok) { ok = glm_graph_matmul_q8_0_tensor(batch->batch_attn_out, model, l->attn_output->abs_offset, batch->heads_dim, DS4_N_EMBD, batch->batch_heads, rows); } return ok; } static bool glm53_graph_encode_native_session_batch( ds4_decode_item *items, int count, const ds4_model *model, const ds4_weights *weights) { if (!glm53_graph_native_session_batch_supported(items, count) || !model || !weights) { return false; } const uint32_t rows = (uint32_t)count; ds4_glm_gpu_graph *g = &items[0].session->glm_graph; if (!glm53_graph_prefill_workspace_ensure(g, rows) || !glm53_graph_session_batch_logits_ensure(g, rows)) { return false; } const uint64_t plain_values = (uint64_t)rows * DS4_N_EMBD; const uint64_t hc_values = plain_values * DS4_N_HC; ds4_gpu_tensor *cur = ds4_gpu_tensor_view( g->batch_cur, 0, plain_values * sizeof(float)); ds4_gpu_tensor *next = ds4_gpu_tensor_view( g->batch_next, 0, plain_values * sizeof(float)); ds4_gpu_tensor *after_attn = ds4_gpu_tensor_view( g->batch_after_attn, 0, plain_values * sizeof(float)); ds4_gpu_tensor *hc_cur = ds4_gpu_tensor_view( g->batch_hc_cur, 0, hc_values * sizeof(float)); ds4_gpu_tensor *hc_next = ds4_gpu_tensor_view( g->batch_hc_next, 0, hc_values * sizeof(float)); ds4_gpu_tensor *hc_after_attn = ds4_gpu_tensor_view( g->batch_hc_after_attn, 0, hc_values * sizeof(float)); bool ok = cur && next && after_attn && hc_cur && hc_next && hc_after_attn; int32_t *tokens = ok ? xmalloc((size_t)rows * sizeof(tokens[0])) : NULL; for (uint32_t i = 0; ok && i < rows; i++) tokens[i] = items[i].token; if (ok) ok = ds4_gpu_tensor_write(g->prefill_tokens, 0, tokens, (uint64_t)rows * sizeof(tokens[0])) != 0; free(tokens); if (ok && weights->token_embd->type == DS4_TENSOR_BF16) { ok = ds4_gpu_glm53_embedding_bf16(cur, model->map, model->size, weights->token_embd->abs_offset, g->prefill_tokens, rows, DS4_N_EMBD, DS4_N_VOCAB) != 0; } else 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, rows, DS4_N_EMBD) != 0; } if (ok) ok = ds4_gpu_repeat_hc_rows_tensor(hc_cur, cur, rows, DS4_N_EMBD, DS4_N_HC) != 0; const char *stage = "embedding"; uint32_t failed_layer = UINT32_MAX; for (uint32_t il = g->layer_start; ok && il <= g->layer_end; il++) { failed_layer = il; const ds4_layer_weights *l = &weights->layer[il]; stage = "attention mHC pre"; ok = glm53_graph_hc_pre_rows(g, model, l->hc_attn_fn, l->hc_attn_scale, l->hc_attn_base, l->attn_norm, hc_cur, cur, g->batch_attn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, rows); stage = ds4_glm53_layer_is_kda(il) ? "KDA attention" : "DSA attention"; if (ok && ds4_glm53_layer_is_kda(il)) { ok = glm53_graph_encode_kda_session_batch( items, count, g, model, l, il); } else if (ok) { ok = glm53_graph_encode_dsa_session_batch( items, count, g, model, l, il); } stage = "attention steering"; if (ok) ok = glm_graph_apply_directional_steering_attn( g, g->batch_attn_out, il, rows); stage = "attention mHC expand"; if (ok) ok = ds4_gpu_hc_expand_split_tensor(hc_after_attn, g->batch_attn_out, hc_cur, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; stage = "FFN mHC pre"; if (ok) ok = glm53_graph_hc_pre_rows(g, model, l->hc_ffn_fn, l->hc_ffn_scale, l->hc_ffn_base, l->ffn_norm, hc_after_attn, after_attn, g->batch_ffn_norm, g->batch_hc_flat, g->batch_hc_mix, g->batch_hc_split, rows); stage = "FFN"; if (ok) ok = glm_graph_encode_ffn_batch(g, model, weights, l, il, 0, after_attn, next, rows, false, false, false, NULL); stage = "FFN steering"; if (ok) ok = glm_graph_apply_directional_steering_ffn( g, next, il, rows); stage = "FFN mHC expand"; if (ok) ok = ds4_gpu_hc_expand_split_tensor(hc_next, next, hc_after_attn, g->batch_hc_split, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { ds4_gpu_tensor *tmp = hc_cur; hc_cur = hc_next; hc_next = tmp; } } stage = "output collapse"; for (uint32_t i = 0; ok && i < rows; i++) { ds4_gpu_tensor *hc_row = glm_graph_tensor_row_view_strided( hc_cur, i, (uint64_t)DS4_N_HC * DS4_N_EMBD, (uint64_t)DS4_N_HC * DS4_N_EMBD); ds4_gpu_tensor *out_row = glm_graph_tensor_row_view_strided( cur, i, DS4_N_EMBD, DS4_N_EMBD); ok = hc_row && out_row && ds4_gpu_hc_weighted_sum_tensor( out_row, hc_row, g->hc_mean_weights, DS4_N_EMBD, DS4_N_HC) != 0; ds4_gpu_tensor_free(out_row); ds4_gpu_tensor_free(hc_row); } stage = "output norm"; if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor( g->batch_ffn_norm, cur, model->map, model->size, weights->output_norm->abs_offset, DS4_N_EMBD, rows, DS4_RMS_EPS) != 0; stage = "output projection"; if (ok) ok = glm53_graph_matmul_rows(g->session_batch_logits, model, weights->output, DS4_N_EMBD, DS4_N_VOCAB, g->batch_ffn_norm, rows); stage = "logits scatter"; for (uint32_t i = 0; ok && i < rows; i++) { ok = ds4_gpu_tensor_copy( items[i].session->glm_graph.logits, 0, g->session_batch_logits, (uint64_t)i * DS4_N_VOCAB * sizeof(float), (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { fprintf(stderr, "ds4: GLM-5.3 native session batch failed at layer %u " "stage '%s' (rows %u)\n", failed_layer, stage, rows); } ds4_gpu_tensor_free(hc_after_attn); ds4_gpu_tensor_free(hc_next); ds4_gpu_tensor_free(hc_cur); ds4_gpu_tensor_free(after_attn); ds4_gpu_tensor_free(next); ds4_gpu_tensor_free(cur); return ok; } #ifdef DS4_HAS_DEEPSEEK41_GPU static bool ds41_sessions_batch_supported(ds4_decode_item *items, int count, const ds4_engine *e) { if (!items || count < 2 || count > DS4_TP_BATCH_MAX_ROWS || !e || e->support_kind != DS4_SUPPORT_NONE) return false; const bool cuda = e->backend == DS4_BACKEND_CUDA; if (cuda ? !metal_graph_tp_env_flag("DS4_CUDA_SESSION_BATCH_MOE", true) : (e->backend != DS4_BACKEND_METAL || getenv("DS4_METAL_DISABLE_V41_SESSION_BATCH"))) return false; ds41_gpu_graph *graphs[DS4_TP_BATCH_MAX_ROWS]; for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e || !s->ds41_graph_ready || !s->checkpoint_valid || s->distributed || (cuda ? !ds41_cuda_row_batch_supported(&s->ds41_graph, &e->weights) : s->ds41_graph.streaming) || s->ds41_graph.imatrix || s->ds41_graph.quality || !s->ds41_graph.valid || s->ds41_graph.image_count || s->ds41_graph.pos != (uint32_t)s->checkpoint.len) return false; /* Smaller TP batches lose to ordered decode. Single-host execution * starts at two; the network crossover is backend-specific. */ if (s->ds41_graph.tp_world == 2 && count < (cuda ? 5 : 3)) return false; graphs[i] = &s->ds41_graph; } return ds41_batch_workspace(graphs, count) != NULL; } #endif #ifdef DS4_HAS_QWEN4_METAL /* Batch weight-only work in the shared arena while keeping each session's * recurrence, attention caches and n-gram history independent. Row views * let the single-session kernels consume slices of the shared transients. */ enum { QWEN4_BATCH_MAX_ROWS = 64 }; typedef struct { ds4_gpu_tensor *R, *ple_gated, *ple_normed; ds4_gpu_tensor *qkv, *ga, *gb, *lin_o; ds4_gpu_tensor *qg, *kp, *vp, *iq, *ik, *q, *gate, *iqn, *attn_o; ds4_gpu_tensor *mixed; /* speculative batch only: the fused GDN front reads it */ } qwen4_batch_row; static ds4_gpu_tensor *qwen4_row_view(ds4_gpu_tensor *base, uint32_t row, uint64_t values) { return ds4_gpu_tensor_view(base, (uint64_t)row * values * sizeof(float), values * sizeof(float)); } static void qwen4_batch_rows_free(qwen4_batch_row *rows, int count) { if (!rows) return; for (int i = 0; i < count; i++) { ds4_gpu_tensor **all[] = { &rows[i].R, &rows[i].ple_gated, &rows[i].ple_normed, &rows[i].qkv, &rows[i].ga, &rows[i].gb, &rows[i].lin_o, &rows[i].qg, &rows[i].kp, &rows[i].vp, &rows[i].iq, &rows[i].ik, &rows[i].q, &rows[i].gate, &rows[i].iqn, &rows[i].attn_o, &rows[i].mixed, }; for (size_t j = 0; j < sizeof(all) / sizeof(all[0]); j++) { ds4_gpu_tensor_free(*all[j]); *all[j] = NULL; } } } /* Built once per batch rather than per layer: a view is a tracked allocation, * and 48 layers would spend more on making them than on the dispatches they * serve. */ static bool qwen4_batch_rows_build(qwen4_batch_row *rows, int count, const ds4_qwen4_gpu_graph *g) { const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM; const uint64_t iq_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; memset(rows, 0, (size_t)count * sizeof(rows[0])); for (int i = 0; i < count; i++) { const uint32_t r = (uint32_t)i; rows[i].R = qwen4_row_view(g->R, r, hc_dim); rows[i].ple_gated = qwen4_row_view(g->ple_gated, r, hc_dim); rows[i].ple_normed = qwen4_row_view(g->ple_normed, r, hc_dim); rows[i].qkv = qwen4_row_view(g->qkv, r, DS4_N_LIN_CONV_DIM); rows[i].ga = qwen4_row_view(g->ga, r, DS4_N_LIN_V_HEAD); rows[i].gb = qwen4_row_view(g->gb, r, DS4_N_LIN_V_HEAD); rows[i].lin_o = qwen4_row_view(g->lin_o, r, v_dim); rows[i].qg = qwen4_row_view(g->qg, r, 2u * q_dim); rows[i].kp = qwen4_row_view(g->kp, r, kv_dim); rows[i].vp = qwen4_row_view(g->vp, r, kv_dim); rows[i].iq = qwen4_row_view(g->iq, r, iq_dim); rows[i].ik = qwen4_row_view(g->ik, r, DS4_N_INDEXER_HEAD_DIM); rows[i].q = qwen4_row_view(g->q, r, q_dim); rows[i].gate = qwen4_row_view(g->gate, r, q_dim); rows[i].iqn = qwen4_row_view(g->iqn, r, iq_dim); rows[i].attn_o = qwen4_row_view(g->attn_o, r, q_dim); if (!rows[i].R || !rows[i].ple_gated || !rows[i].ple_normed || !rows[i].qkv || !rows[i].ga || !rows[i].gb || !rows[i].lin_o || !rows[i].qg || !rows[i].kp || !rows[i].vp || !rows[i].iq || !rows[i].ik || !rows[i].q || !rows[i].gate || !rows[i].iqn || !rows[i].attn_o) { qwen4_batch_rows_free(rows, count); return false; } } return true; } /* This session's graph already points at the shared arena, so a copy of it * with this row's views substituted is precisely what the single-row kernels * want: this session's caches and positions, this row's slice of the * transients. The copy is a snapshot, so g->pos still reads the position the * batch started from. */ static void qwen4_batch_row_graph(ds4_qwen4_gpu_graph *out, ds4_session *s, const qwen4_batch_row *v) { *out = s->qwen4_graph; out->R = v->R; out->ple_gated = v->ple_gated; out->ple_normed = v->ple_normed; out->qkv = v->qkv; out->ga = v->ga; out->gb = v->gb; out->lin_o = v->lin_o; out->qg = v->qg; out->kp = v->kp; out->vp = v->vp; out->iq = v->iq; out->ik = v->ik; out->q = v->q; out->gate = v->gate; out->iqn = v->iqn; out->attn_o = v->attn_o; if (v->mixed) out->mixed = v->mixed; } /* The Q8 tile computes 32 rows and needs the padding rows to exist in both * operands, so the logits and head rows are sized for at least that many: * the output head then takes the tile at any batch width instead of the * per-token matvec reading its 636 MB once per four rows. */ static bool qwen4_batch_scratch_ensure(ds4_qwen4_gpu_graph *g, uint32_t rows) { if (rows < 32u) rows = 32u; if (g->batch_logits && g->batch_logit_rows >= rows) return true; ds4_gpu_tensor_free(g->batch_logits); ds4_gpu_tensor_free(g->batch_attn_rows); ds4_gpu_tensor_free(g->batch_attn_part); ds4_gpu_tensor_free(g->batch_head_x); ds4_gpu_tensor_free(g->batch_gdn_rows); g->batch_logit_rows = 0; g->batch_gdn_rows = ds4_gpu_tensor_alloc((uint64_t)DS4_N_LAYER * QWEN4_BATCH_MAX_ROWS * DS4_GPU_QWEN4_GDN_ROW_BYTES); g->batch_logits = ds4_gpu_tensor_alloc((uint64_t)rows * DS4_N_VOCAB * sizeof(float)); g->batch_head_x = ds4_gpu_tensor_alloc((uint64_t)rows * DS4_N_EMBD * sizeof(float)); g->batch_attn_rows = ds4_gpu_tensor_alloc((uint64_t)DS4_N_LAYER * QWEN4_BATCH_MAX_ROWS * DS4_GPU_QWEN4_ATTN_ROW_BYTES); g->batch_attn_part = ds4_gpu_tensor_alloc(ds4_gpu_qwen4_attn_part_floats(rows, DS4_N_HEAD, DS4_N_HEAD_DIM) * sizeof(float)); if (!g->batch_logits || !g->batch_attn_rows || !g->batch_attn_part || !g->batch_head_x || !g->batch_gdn_rows) return false; g->batch_logit_rows = rows; return true; } /* Token embeddings, rope positions and n-gram rows for every member. The * n-gram gather is the part that most wants a batch: one call presents all * rows at once, so the reader sorts them by offset and issues them in * parallel instead of walking each session's sixteen rows on its own. */ static bool qwen4_batch_stage_embeddings(ds4_decode_item *items, int count, const ds4_model *m, const ds4_weights *w, ds4_qwen4_gpu_graph *g, uint32_t *ids) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; float *row = g->host_row; for (int i = 0; i < count; i++) { ds4_qwen4_gpu_graph *rg = &items[i].session->qwen4_graph; float *dst = row + (uint64_t)i * hc_dim; qwen4_ref_row(m, w->token_embd, (uint64_t)items[i].token, dst); for (uint32_t s = 1; s < hc; s++) memcpy(dst + (uint64_t)s * E, dst, E * sizeof(float)); uint32_t pos3[4]; qwen4_mrope_pos(NULL, 0, rg->pos, &rg->mrope_delta, pos3); pos3[3] = 0; if (!ds4_gpu_tensor_write(rg->pos3, (uint64_t)rg->pos * 16u, pos3, 16u)) return false; qwen4_ple_step(items[i].token, rg->ple_prev, ids + (uint64_t)i * DS4_N_PLE_HEADS); } return ds4_gpu_tensor_write(g->R, 0, row, (uint64_t)count * hc_dim * sizeof(float)) != 0; } /* The n-gram rows are hashed across a 95 GiB table, so this is disk latency * rather than compute, and the first trunk layer does not need them: the PLE * layer that does sits at trunk index one. Reading them after that layer has * been submitted lets the wait overlap GPU work instead of preceding it. */ static bool qwen4_batch_stage_ngrams(int count, const ds4_model *m, ds4_qwen4_gpu_graph *g, const uint32_t *ids) { float *row = g->host_row; if (!qwen4_ngram_read(m, ids, (size_t)count * DS4_N_PLE_HEADS, row)) { fprintf(stderr, "ds4: n-gram read failed: %s\n", strerror(errno)); return false; } return ds4_gpu_tensor_write(g->ple_emb, 0, row, (uint64_t)count * DS4_N_EMBD * sizeof(float)) != 0; } /* The batched scan needs every row's state in the pool it addresses. */ static bool qwen4_batch_rows_pooled(ds4_decode_item *items, int count, uint32_t *slots) { for (int i = 0; i < count; i++) { const int slot = items[i].session->qwen4_slot; if (slot < 0) return false; slots[i] = (uint32_t)slot; } return true; } static bool qwen4_batch_linear(ds4_decode_item *items, int count, const qwen4_batch_row *views, const ds4_qwen4_gpu_graph *rowg, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t T) { if (!qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) || !qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T) || !qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) || !qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T)) { return false; } /* The causal convolution and the delta-net scan are the only steps that * read and write a session's own recurrent state. Overlapping these rows * through a concurrent encoder was measured and changed nothing: the cost * is the per-dispatch floor, not a barrier between them. */ ds4_engine *e = items[0].session->engine; uint32_t slots[QWEN4_BATCH_MAX_ROWS]; const bool pooled = e->qwen4_lin_state_pool[il] && getenv("DS4_QWEN4_NO_BATCH_SCAN") == NULL && qwen4_batch_rows_pooled(items, count, slots); const uint64_t hist_floats = (uint64_t)(DS4_N_LIN_CONV - 1u) * DS4_N_LIN_CONV_DIM; if (pooled) { if (ds4_gpu_qwen4_conv_stream_rows_tensor( g->qkv, e->qwen4_lin_hist_pool[il], m->map, m->size, l->lin_conv->abs_offset, slots, (uint32_t)count, DS4_N_LIN_CONV_DIM, DS4_N_LIN_CONV, (uint32_t)hist_floats, DS4_N_LIN_CONV_DIM, 1) == 0) { return false; } } else { for (int i = 0; i < count; i++) { if (!ds4_gpu_qwen4_conv_stream_tensor(views[i].qkv, rowg[i].layer_lin_hist[il], m->map, m->size, l->lin_conv->abs_offset, 1u, DS4_N_LIN_CONV_DIM, DS4_N_LIN_CONV, true)) { return false; } } } if (!ds4_gpu_qwen4_gdn_prep_tensor(g->qkv, g->ga, g->gb, m->map, m->size, l->lin_a->abs_offset, l->lin_dt_bias->abs_offset, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM)) { return false; } /* With every row's state in the pool, one dispatch advances them all; the * per-row kernels are short enough that sixteen of them never reach steady * state. Same grouping and arithmetic per row, so the two paths agree bit * for bit. DS4_QWEN4_NO_BATCH_SCAN takes the per-row path for A/B. */ const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; if (pooled) { if (ds4_gpu_qwen4_gdn_scan_rows_tensor( g->lin_o, e->qwen4_lin_state_pool[il], g->qkv, g->ga, g->gb, slots, (uint32_t)count, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, (uint32_t)(v_dim * DS4_N_LIN_HEAD_DIM), DS4_N_LIN_CONV_DIM, (uint32_t)v_dim) == 0) { return false; } } else { for (int i = 0; i < count; i++) { if (ds4_gpu_qwen4_gdn_scan_tensor(views[i].lin_o, rowg[i].layer_lin_state[il], views[i].qkv, views[i].ga, views[i].gb, 1u, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, NULL, 0u, NULL, 1u) == 0) { return false; } } } return ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size, l->lin_norm->abs_offset, T, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0 && qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T); } /* Every row's attention in one dispatch per stage: the projections above * are batched already, and the per-row kernels below are tiny grids that run * one after another, so sixteen rows cost sixteen latencies each. The rows * kernels take the row as a grid axis and reach each session's caches * through a staged table, keeping the single-row arithmetic (the row's split * count included), so the outputs are those of the per-row dispatches. * DS4_QWEN4_NO_BATCH_ATTN takes the per-row path for A/B. */ /* The rows kernels over explicit entries, one per token of the batch in row * order: a session with a draft row has two entries, its caches at pos and * pos + 1, so the second token sees the first's keys the way a two-row * verify does. */ static bool qwen4_batch_attention_entries(ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { const uint32_t ratio = 4u; const float scale = 1.0f / sqrtf((float)DS4_N_HEAD_DIM); uint32_t n_block_stride = 0; bool any_sparse = false, any_block = false; for (uint32_t i = 0; i < n_rows; i++) { const uint32_t n_blocks = (rows[i].pos + 1u) / ratio; if (n_blocks > n_block_stride) n_block_stride = n_blocks; any_sparse |= rows[i].use_sel != 0; any_block |= (rows[i].pos + 1u) % ratio == 0u; } const uint64_t entry0 = (uint64_t)il * QWEN4_BATCH_MAX_ROWS; bool ok = ds4_gpu_qwen4_attn_rows_stage(g->batch_attn_rows, entry0, rows, n_rows, ratio) && ds4_gpu_qwen4_attn_prep_rows_tensor(g->q, g->gate, g->iqn, g->qg, g->kp, g->vp, g->iq, g->ik, g->batch_attn_rows, entry0, rows, n_rows, m->map, m->size, l->attn_q_norm->abs_offset, l->attn_k_norm->abs_offset, l->indexer_q_norm->abs_offset, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS); if (ok && any_block) { ok = ds4_gpu_qwen4_idx_block_key_rows_tensor(g->batch_attn_rows, entry0, rows, n_rows, m->map, m->size, l->indexer_k_norm->abs_offset, ratio, DS4_N_INDEXER_HEAD_DIM, DS4_N_ROT, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS); } if (ok && any_sparse) { ok = ds4_gpu_qwen4_idx_score_rows_tensor(g->score, g->tile_max, g->iqn, g->batch_attn_rows, entry0, rows, n_rows, n_block_stride, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, ratio) && ds4_gpu_qwen4_idx_select_rows_tensor(g->sel_blocks, g->score, g->tile_max, g->batch_attn_rows, entry0, rows, n_rows, n_block_stride, g->k_blocks) && ds4_gpu_qwen4_idx_expand_rows_tensor(g->sel_tokens, g->n_sel, g->sel_blocks, g->batch_attn_rows, entry0, n_rows, g->k_blocks, ratio, g->sel_stride); } return ok && ds4_gpu_qwen4_attn_decode_rows_tensor(g->attn_o, g->q, g->gate, g->sel_tokens, g->n_sel, g->batch_attn_part, g->batch_attn_rows, entry0, rows, n_rows, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, g->sel_stride, scale); } static bool qwen4_batch_attention_rows(int count, ds4_qwen4_gpu_graph *rowg, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { const uint32_t sparse_pos = (g->k_blocks + 1u) * 4u - 1u; ds4_gpu_qwen4_attn_row rows[QWEN4_BATCH_MAX_ROWS]; for (int i = 0; i < count; i++) { const ds4_qwen4_gpu_graph *r = &rowg[i]; rows[i].k_cache = r->layer_k_cache[il]; rows[i].v_cache = r->layer_v_cache[il]; rows[i].ik_cache = r->layer_ik_cache[il]; rows[i].block_key = r->layer_block_key[il]; rows[i].pos3 = r->pos3; rows[i].pos = r->pos; rows[i].use_sel = r->pos >= sparse_pos; if (r->pos >= r->ctx_cap) return false; } return qwen4_batch_attention_entries(rows, (uint32_t)count, g, m, l, il); } static bool qwen4_batch_attention(int count, ds4_qwen4_gpu_graph *rowg, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t T) { const uint32_t ratio = 4u; bool ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T) && qwen4_gemv(g->kp, m, l->attn_k, g->mixed, T) && qwen4_gemv(g->vp, m, l->attn_v, g->mixed, T) && qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T) && qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T); const bool rows_path = getenv("DS4_QWEN4_NO_BATCH_ATTN") == NULL && getenv("DS4_QWEN4_NO_IDX_SELECT") == NULL && (DS4_N_HEAD_DIM == 128u || DS4_N_HEAD_DIM == 256u); if (ok && rows_path) return qwen4_batch_attention_rows(count, rowg, g, m, l, il) && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T); /* Each row appends to its own KV and indexer caches at its own position, * then reads them back through its own block selection. */ for (int i = 0; ok && i < count; i++) { ds4_qwen4_gpu_graph *r = &rowg[i]; const uint32_t pos0 = r->pos; ok = ds4_gpu_qwen4_attn_prep_tensor(r->q, r->gate, r->layer_k_cache[il], r->layer_v_cache[il], r->iqn, r->layer_ik_cache[il], r->qg, r->kp, r->vp, r->iq, r->ik, r->pos3, m->map, m->size, l->attn_q_norm->abs_offset, l->attn_k_norm->abs_offset, l->indexer_q_norm->abs_offset, 1u, DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, pos0, r->ctx_cap, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS); const uint32_t first_block = pos0 / ratio; const uint32_t n_blocks_after = (pos0 + 1u) / ratio; if (ok && n_blocks_after > first_block) { ok = ds4_gpu_qwen4_idx_block_key_tensor(r->layer_block_key[il], r->layer_ik_cache[il], r->pos3, m->map, m->size, l->indexer_k_norm->abs_offset, first_block, n_blocks_after - first_block, ratio, DS4_N_INDEXER_HEAD_DIM, DS4_N_ROT, DS4_ROPE_FREQ_BASE, DS4_RMS_EPS); } if (ok) ok = qwen4_graph_attention_core(r, il, r->q, r->gate, r->iqn, r->attn_o, n_blocks_after, pos0, 1u); } return ok && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T); } static bool qwen4_graph_encode_native_session_batch(ds4_decode_item *items, int count, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w) { const uint32_t T = (uint32_t)count; const uint32_t hc_dim = DS4_N_EMBD * DS4_N_HC; const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; qwen4_batch_row *views = xmalloc((size_t)count * sizeof(*views)); ds4_qwen4_gpu_graph *rowg = xmalloc((size_t)count * sizeof(*rowg)); bool ok = qwen4_batch_rows_build(views, count, g); for (int i = 0; ok && i < count; i++) { qwen4_batch_row_graph(&rowg[i], items[i].session, &views[i]); } if (ok) ok = qwen4_batch_scratch_ensure(g, T); /* DS4_QWEN4_BATCH_TIMING=1 synchronizes after each stage group and reports * GPU ms per group. The synchronization is itself expensive, so the split * is for attribution only, never for a throughput claim. */ static int timing = -1; if (timing < 0) timing = getenv("DS4_QWEN4_BATCH_TIMING") != NULL; double prof[7] = {0}; double prof_last = 0.0; #define QWEN4_BATCH_PROF(idx_) do { \ if (timing) { \ ds4_gpu_end_commands(); \ const double now_ = now_sec(); \ prof[idx_] += now_ - prof_last; \ prof_last = now_; \ glm_graph_begin_commands_if_needed(); \ } \ } while (0) uint32_t ple_ids[QWEN4_BATCH_MAX_ROWS * DS4_MAX_PLE_HEADS]; bool ngrams_staged = false; if (timing) prof_last = now_sec(); if (ok) ok = qwen4_batch_stage_embeddings(items, count, m, w, g, ple_ids); QWEN4_BATCH_PROF(0); for (uint32_t il = 0; ok && il < n_trunk; il++) { const ds4_layer_weights *l = &w->layer[il]; if (!ngrams_staged && ds4_qwen4_layer_is_ple(il)) { /* Submit everything encoded so far, then read the n-gram rows * while the GPU works through it. DS4_QWEN4_NO_NGRAM_OVERLAP * reads them up front instead, which is the same work with the * GPU idle for it. */ ngrams_staged = true; ok = (getenv("DS4_QWEN4_NO_NGRAM_OVERLAP") != NULL || ds4_gpu_flush_commands() != 0) && qwen4_batch_stage_ngrams(count, m, g, ple_ids); if (!ok) break; } if (ds4_qwen4_layer_is_ple(il)) { ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) && qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) && ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key, g->ple_val, m->map, m->size, l->ple_norm_key->abs_offset, l->ple_norm_query->abs_offset, l->ple_norm_conv->abs_offset, T, DS4_N_EMBD, DS4_N_HC, DS4_RMS_EPS); /* the n-gram convolution carries per-session history */ for (int i = 0; ok && i < count; i++) { ok = ds4_gpu_qwen4_ple_conv_tensor(views[i].R, views[i].ple_gated, views[i].ple_normed, rowg[i].ple_hist, m->map, m->size, l->ple_conv->abs_offset, l->ple_conv->type, 1u, hc_dim, DS4_N_PLE_CONV, DS4_N_PLE_NGRAM, NULL, 0u, NULL, 1u); } } QWEN4_BATCH_PROF(1); if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T); QWEN4_BATCH_PROF(2); if (ok) { ok = ds4_qwen4_layer_is_linear(il) ? qwen4_batch_linear(items, count, views, rowg, g, m, l, il, T) : qwen4_batch_attention(count, rowg, g, m, l, il, T); } QWEN4_BATCH_PROF(ds4_qwen4_layer_is_linear(il) ? 3 : 4); if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T); QWEN4_BATCH_PROF(5); if (ok) ok = qwen4_graph_moe(g, m, l, T); QWEN4_BATCH_PROF(6); } /* The output matrix is the single largest weight read of a decode step, * so it serves every row in one pass. */ if (ok) ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, T) && qwen4_gemv(g->batch_logits, m, w->output, g->mixed, T); if (timing) { fprintf(stderr, "ds4: Qwen3.8 batch(T=%u) ms: stage %.1f ple %.1f hc_attn %.1f " "gdn %.1f attn %.1f hc_ffn %.1f moe %.1f\n", T, 1e3 * prof[0], 1e3 * prof[1], 1e3 * prof[2], 1e3 * prof[3], 1e3 * prof[4], 1e3 * prof[5], 1e3 * prof[6]); } #undef QWEN4_BATCH_PROF qwen4_batch_rows_free(views, count); free(views); free(rowg); return ok; } /* Each session contributes its next token and optionally one draft. State * snapshots retain the first row for rejection; a draft is accepted only * when it is the target's argmax. Grouped reductions can round differently * from a single-session verification. */ typedef struct { ds4_session *session; int tokens[2]; uint32_t n; /* rows: the token, then its draft */ uint32_t row0; /* first row in the batch layout */ } qwen4_batch_member; static ds4_gpu_tensor *qwen4_rows_view(ds4_gpu_tensor *base, uint32_t row0, uint32_t n, uint64_t values) { return ds4_gpu_tensor_view(base, (uint64_t)row0 * values * sizeof(float), (uint64_t)n * values * sizeof(float)); } static bool qwen4_batch_rows_build_ragged(qwen4_batch_row *rows, const qwen4_batch_member *mem, int count, const ds4_qwen4_gpu_graph *g) { const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM; const uint64_t iq_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; memset(rows, 0, (size_t)count * sizeof(rows[0])); for (int i = 0; i < count; i++) { const uint32_t r = mem[i].row0, n = mem[i].n; rows[i].R = qwen4_rows_view(g->R, r, n, hc_dim); rows[i].ple_gated = qwen4_rows_view(g->ple_gated, r, n, hc_dim); rows[i].ple_normed = qwen4_rows_view(g->ple_normed, r, n, hc_dim); rows[i].qkv = qwen4_rows_view(g->qkv, r, n, DS4_N_LIN_CONV_DIM); rows[i].ga = qwen4_rows_view(g->ga, r, n, DS4_N_LIN_V_HEAD); rows[i].gb = qwen4_rows_view(g->gb, r, n, DS4_N_LIN_V_HEAD); rows[i].lin_o = qwen4_rows_view(g->lin_o, r, n, v_dim); rows[i].qg = qwen4_rows_view(g->qg, r, n, 2u * q_dim); rows[i].kp = qwen4_rows_view(g->kp, r, n, kv_dim); rows[i].vp = qwen4_rows_view(g->vp, r, n, kv_dim); rows[i].iq = qwen4_rows_view(g->iq, r, n, iq_dim); rows[i].ik = qwen4_rows_view(g->ik, r, n, DS4_N_INDEXER_HEAD_DIM); rows[i].q = qwen4_rows_view(g->q, r, n, q_dim); rows[i].gate = qwen4_rows_view(g->gate, r, n, q_dim); rows[i].iqn = qwen4_rows_view(g->iqn, r, n, iq_dim); rows[i].attn_o = qwen4_rows_view(g->attn_o, r, n, q_dim); rows[i].mixed = qwen4_rows_view(g->mixed, r, n, DS4_N_EMBD); if (!rows[i].R || !rows[i].ple_gated || !rows[i].ple_normed || !rows[i].qkv || !rows[i].ga || !rows[i].gb || !rows[i].lin_o || !rows[i].qg || !rows[i].kp || !rows[i].vp || !rows[i].iq || !rows[i].ik || !rows[i].q || !rows[i].gate || !rows[i].iqn || !rows[i].attn_o || !rows[i].mixed) { qwen4_batch_rows_free(rows, count); return false; } } return true; } /* Embeddings, rope positions and n-gram ids for every row; a session with a * draft row records the verify snapshot's host state after its first token, * as qwen4_graph_stage_inputs does. */ static bool qwen4_batch_stage_embeddings_ragged(const qwen4_batch_member *mem, int count, uint32_t N, const ds4_model *m, const ds4_weights *w, ds4_qwen4_gpu_graph *g, uint32_t *ids) { const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC, hc_dim = E * hc; float *row = g->host_row; for (int i = 0; i < count; i++) { ds4_qwen4_gpu_graph *rg = &mem[i].session->qwen4_graph; for (uint32_t t = 0; t < mem[i].n; t++) { const uint32_t r = mem[i].row0 + t; float *dst = row + (uint64_t)r * hc_dim; qwen4_ref_row(m, w->token_embd, (uint64_t)mem[i].tokens[t], dst); for (uint32_t sidx = 1; sidx < hc; sidx++) memcpy(dst + (uint64_t)sidx * E, dst, E * sizeof(float)); uint32_t pos3[4]; qwen4_mrope_pos(NULL, 0, rg->pos + t, &rg->mrope_delta, pos3); pos3[3] = 0; if (!ds4_gpu_tensor_write(rg->pos3, (uint64_t)(rg->pos + t) * 16u, pos3, 16u)) return false; qwen4_ple_step(mem[i].tokens[t], rg->ple_prev, ids + (uint64_t)r * DS4_N_PLE_HEADS); if (t == 0 && rg->snap_after_first) { memcpy(rg->snap_ple_prev, rg->ple_prev, sizeof(rg->snap_ple_prev)); rg->snap_pos = rg->pos + 1u; rg->snap_mrope_delta = rg->mrope_delta; } } } return ds4_gpu_tensor_write(g->R, 0, row, (uint64_t)N * hc_dim * sizeof(float)) != 0; } static bool qwen4_graph_encode_native_session_batch_ragged(const qwen4_batch_member *mem, int count, uint32_t N, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w) { const uint32_t T = N; const uint32_t hc_dim = DS4_N_EMBD * DS4_N_HC; const uint32_t n_trunk = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; qwen4_batch_row *views = xmalloc((size_t)count * sizeof(*views)); ds4_qwen4_gpu_graph *rowg = xmalloc((size_t)count * sizeof(*rowg)); uint32_t *ple_ids = xmalloc((size_t)N * DS4_MAX_PLE_HEADS * sizeof(uint32_t)); bool ok = qwen4_batch_rows_build_ragged(views, mem, count, g); for (int i = 0; ok && i < count; i++) qwen4_batch_row_graph(&rowg[i], mem[i].session, &views[i]); if (ok) ok = qwen4_batch_scratch_ensure(g, N > 32u ? N : 32u); if (ok) ok = qwen4_batch_stage_embeddings_ragged(mem, count, N, m, w, g, ple_ids); bool ngrams_staged = false; for (uint32_t il = 0; ok && il < n_trunk; il++) { const ds4_layer_weights *l = &w->layer[il]; if (!ngrams_staged && ds4_qwen4_layer_is_ple(il)) { ngrams_staged = true; ok = (getenv("DS4_QWEN4_NO_NGRAM_OVERLAP") != NULL || ds4_gpu_flush_commands() != 0) && qwen4_batch_stage_ngrams((int)N, m, g, ple_ids); if (!ok) break; } if (ds4_qwen4_layer_is_ple(il)) { ok = qwen4_gemv(g->ple_key, m, l->ple_key, g->ple_emb, T) && qwen4_gemv(g->ple_val, m, l->ple_value, g->ple_emb, T) && ds4_gpu_qwen4_ple_gate_tensor(g->ple_gated, g->ple_normed, g->R, g->ple_key, g->ple_val, m->map, m->size, l->ple_norm_key->abs_offset, l->ple_norm_query->abs_offset, l->ple_norm_conv->abs_offset, T, DS4_N_EMBD, DS4_N_HC, DS4_RMS_EPS); for (int i = 0; ok && i < count; i++) { ok = ds4_gpu_qwen4_ple_conv_tensor(views[i].R, views[i].ple_gated, views[i].ple_normed, rowg[i].ple_hist, m->map, m->size, l->ple_conv->abs_offset, l->ple_conv->type, mem[i].n, hc_dim, DS4_N_PLE_CONV, DS4_N_PLE_NGRAM, rowg[i].snap_after_first ? rowg[i].snap_ple_hist : NULL, 0u, NULL, 1u); } } if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, T); if (ok && ds4_qwen4_layer_is_linear(il)) { /* projections over every row; the convolution and the scan walk * each session's one or two rows against its own state through * the row table, snapshotting after the first as the verify does */ ok = qwen4_gemv(g->qkv, m, l->lin_qkv, g->mixed, T) && qwen4_gemv(g->z, m, l->lin_gate, g->mixed, T) && qwen4_gemv(g->ga, m, l->lin_alpha, g->mixed, T) && qwen4_gemv(g->gb, m, l->lin_beta, g->mixed, T); ds4_gpu_qwen4_gdn_row grows[QWEN4_BATCH_MAX_ROWS]; for (int i = 0; i < count; i++) { const bool snap = rowg[i].snap_after_first && rowg[i].snap_lin_state[il] && mem[i].n == 2u; grows[i].state = rowg[i].layer_lin_state[il]; grows[i].hist = rowg[i].layer_lin_hist[il]; grows[i].snap_state = snap ? rowg[i].snap_lin_state[il] : NULL; grows[i].snap_hist = snap ? rowg[i].snap_lin_hist[il] : NULL; grows[i].row0 = mem[i].row0; grows[i].n_tok = mem[i].n; } const uint64_t entry0 = (uint64_t)il * QWEN4_BATCH_MAX_ROWS; const uint64_t v_dim = (uint64_t)DS4_N_LIN_V_HEAD * DS4_N_LIN_HEAD_DIM; if (ok) ok = ds4_gpu_qwen4_gdn_rows_stage(g->batch_gdn_rows, entry0, grows, (uint32_t)count) && ds4_gpu_qwen4_conv_stream_rows2_tensor(g->qkv, m->map, m->size, l->lin_conv->abs_offset, g->batch_gdn_rows, entry0, grows, (uint32_t)count, T, DS4_N_LIN_CONV_DIM, DS4_N_LIN_CONV, DS4_N_LIN_CONV_DIM, 1) && ds4_gpu_qwen4_gdn_prep_tensor(g->qkv, g->ga, g->gb, m->map, m->size, l->lin_a->abs_offset, l->lin_dt_bias->abs_offset, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM) && ds4_gpu_qwen4_gdn_scan_rows2_tensor(g->lin_o, g->qkv, g->ga, g->gb, g->batch_gdn_rows, entry0, grows, (uint32_t)count, T, DS4_N_LIN_K_HEAD, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_N_LIN_CONV_DIM, (uint32_t)v_dim) && ds4_gpu_qwen4_gdn_out_tensor(g->lin_o, g->z, m->map, m->size, l->lin_norm->abs_offset, T, DS4_N_LIN_V_HEAD, DS4_N_LIN_HEAD_DIM, DS4_RMS_EPS) != 0 && qwen4_gemv(g->blk, m, l->lin_out, g->lin_o, T); } else if (ok) { ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, T) && qwen4_gemv(g->kp, m, l->attn_k, g->mixed, T) && qwen4_gemv(g->vp, m, l->attn_v, g->mixed, T) && qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, T) && qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, T); /* one entry per token: a draft row's entry sits at pos + 1 */ ds4_gpu_qwen4_attn_row arows[2 * QWEN4_BATCH_MAX_ROWS]; const uint32_t sparse_pos = (g->k_blocks + 1u) * 4u - 1u; for (int i = 0; ok && i < count; i++) { const ds4_qwen4_gpu_graph *r = &rowg[i]; for (uint32_t t = 0; t < mem[i].n; t++) { ds4_gpu_qwen4_attn_row *e = &arows[mem[i].row0 + t]; e->k_cache = r->layer_k_cache[il]; e->v_cache = r->layer_v_cache[il]; e->ik_cache = r->layer_ik_cache[il]; e->block_key = r->layer_block_key[il]; e->pos3 = r->pos3; e->pos = r->pos + t; e->use_sel = e->pos >= sparse_pos; if (e->pos >= r->ctx_cap) ok = false; } } if (ok) ok = qwen4_batch_attention_entries(arows, T, g, m, l, il) && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, T); } if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, T, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, T); if (ok) ok = qwen4_graph_moe(g, m, l, T); } if (ok) ok = qwen4_graph_hc_mix(g, m, w->output_hc_norm, w->output_hc_down, w->output_hc_up, NULL, T) && qwen4_gemv(g->batch_logits, m, w->output, g->mixed, T); qwen4_batch_rows_free(views, count); free(ple_ids); free(rowg); free(views); return ok; } /* Next drafts for every session: one predictor pass over every row of the * batch (whose trunk streams sit in the batch's R, at the positions the * trunk rows had, so the nextn layer's attention row table is the trunk's), * then one product of the head over every session's last committed row and * one argmax each. A rejected draft's row runs too, with the parent as * the token that follows: the next pass rewrites its cache entry before * anything reads it, and its output is ignored. parents are the tokens * the caller will feed next. */ static bool qwen4_batch_mtp_drafts(qwen4_batch_member *mem, int count, const uint32_t *committed, const int *parents, const uint32_t *idx0, uint32_t N, ds4_qwen4_gpu_graph *g, const ds4_model *m, const ds4_weights *w, int *drafts) { const uint32_t il = DS4_N_LAYER - 1u; const ds4_layer_weights *l = &w->layer[il]; const uint32_t E = DS4_N_EMBD, hc = DS4_N_HC; const uint64_t row_bytes = (uint64_t)hc * E * sizeof(float); int32_t ids[2 * QWEN4_BATCH_MAX_ROWS]; ds4_gpu_qwen4_attn_row arows[2 * QWEN4_BATCH_MAX_ROWS]; const uint32_t sparse_pos = (g->k_blocks + 1u) * 4u - 1u; for (int i = 0; i < count; i++) { const ds4_qwen4_gpu_graph *r = &mem[i].session->qwen4_graph; for (uint32_t t = 0; t < mem[i].n; t++) { ids[mem[i].row0 + t] = t + 1u < committed[i] ? mem[i].tokens[t + 1u] : parents[i]; ds4_gpu_qwen4_attn_row *e = &arows[mem[i].row0 + t]; e->k_cache = r->layer_k_cache[il]; e->v_cache = r->layer_v_cache[il]; e->ik_cache = r->layer_ik_cache[il]; e->block_key = r->layer_block_key[il]; e->pos3 = r->pos3; e->pos = idx0[i] + t; e->use_sel = e->pos >= sparse_pos; if (e->pos >= r->ctx_cap) return false; } } for (uint32_t t = 0; t < N; t++) qwen4_ref_row(m, w->token_embd, (uint64_t)ids[t], g->host_row + (uint64_t)t * E); if (!ds4_gpu_tensor_write(g->batch_head_x, 0, g->host_row, (uint64_t)N * E * sizeof(float)) || !glm_graph_begin_commands_if_needed()) return false; /* The trunk's MoE scratch is idle here. Reuse it for predictor inputs * and project all rows together, keeping the existing stage kernels. */ const uint64_t emb_bytes = (uint64_t)E * sizeof(float); const uint64_t cat_bytes = (hc + 1u) * 2u * emb_bytes; const uint64_t proj_bytes = (hc + 1u) * emb_bytes; bool ok = true; for (uint32_t t = 0; ok && t < N; t++) { ds4_gpu_tensor *emb = ds4_gpu_tensor_view(g->batch_head_x, t * emb_bytes, emb_bytes); ds4_gpu_tensor *R = ds4_gpu_tensor_view(g->R, t * row_bytes, row_bytes); ds4_gpu_tensor *cat = ds4_gpu_tensor_view(g->part, t * cat_bytes, cat_bytes); ok = emb && R && cat && ds4_gpu_qwen4_mtp_stage_tensor(cat, emb, R, m->map, m->size, l->nextn_enorm->abs_offset, l->nextn_hnorm->abs_offset, E, hc, DS4_RMS_EPS); ds4_gpu_tensor_free(cat); ds4_gpu_tensor_free(R); ds4_gpu_tensor_free(emb); } if (ok) ok = qwen4_gemv(g->mid, m, l->nextn_eh_proj, g->part, N * (hc + 1u)); for (uint32_t t = 0; ok && t < N; t++) { ds4_gpu_tensor *proj = ds4_gpu_tensor_view(g->mid, t * proj_bytes, proj_bytes); ds4_gpu_tensor *R = ds4_gpu_tensor_view(g->R, t * row_bytes, row_bytes); ok = proj && R && ds4_gpu_qwen4_mtp_combine_tensor(R, proj, E, hc); ds4_gpu_tensor_free(R); ds4_gpu_tensor_free(proj); } if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_attn_norm, l->hc_attn_down, l->hc_attn_up, l->hc_attn_inject, N); if (ok) ok = qwen4_gemv(g->qg, m, l->attn_q, g->mixed, N) && qwen4_gemv(g->kp, m, l->attn_k, g->mixed, N) && qwen4_gemv(g->vp, m, l->attn_v, g->mixed, N) && qwen4_gemv(g->iq, m, l->indexer_q_proj, g->mixed, N) && qwen4_gemv(g->ik, m, l->indexer_k_proj, g->mixed, N) && qwen4_batch_attention_entries(arows, N, g, m, l, il) && qwen4_gemv(g->blk, m, l->attn_output, g->attn_o, N); if (ok) ok = ds4_gpu_qwen4_hc_combine_tensor(g->R, g->blk, g->inj, N, E, hc) != 0; if (ok) ok = qwen4_graph_hc_mix(g, m, l->hc_ffn_norm, l->hc_ffn_down, l->hc_ffn_up, l->hc_ffn_inject, N); if (ok) ok = qwen4_graph_moe(g, m, l, N); /* every session's last committed row, gathered for one pass of the head * mixer into the batch's head rows */ for (int i = 0; ok && i < count; i++) { ok = ds4_gpu_tensor_copy(g->part, (uint64_t)i * row_bytes, g->R, (uint64_t)(mem[i].row0 + committed[i] - 1u) * row_bytes, row_bytes) != 0; } ds4_gpu_tensor *R_save = g->R, *mixed_save = g->mixed; g->R = g->part; g->mixed = g->batch_head_x; if (ok) ok = qwen4_graph_hc_mix(g, m, l->nextn_hc_head_norm, l->nextn_hc_head_down, l->nextn_hc_head_up, NULL, (uint32_t)count); g->R = R_save; g->mixed = mixed_save; if (ok) ok = qwen4_gemv(g->batch_logits, m, w->output, g->batch_head_x, (uint32_t)count); for (int i = 0; ok && i < count; i++) { ds4_qwen4_gpu_graph *sg = &mem[i].session->qwen4_graph; ds4_gpu_tensor *row = ds4_gpu_tensor_view(g->batch_logits, (uint64_t)i * DS4_N_VOCAB * sizeof(float), (uint64_t)DS4_N_VOCAB * sizeof(float)); ok = row && ds4_gpu_qwen4_argmax_tensor(sg->mtp_argmax, sg->mtp_argmax_tmp, row, DS4_N_VOCAB) != 0; ds4_gpu_tensor_free(row); } if (!ds4_gpu_end_commands()) ok = false; for (int i = 0; ok && i < count; i++) { ds4_qwen4_gpu_graph *sg = &mem[i].session->qwen4_graph; int32_t token = -1; ok = ds4_gpu_tensor_read(sg->mtp_argmax, 0, &token, sizeof(token)) != 0 && token >= 0 && token < (int32_t)DS4_N_VOCAB; drafts[i] = token; sg->mtp_pos = idx0[i] + committed[i]; sg->mtp_last_rows = 0; } return ok; } static bool qwen4_graph_native_session_batch_check(ds4_decode_item *items, int count, const ds4_engine *e, bool speculative) { const char *env = getenv("DS4_QWEN4_SESSION_BATCH"); if ((env && env[0] && strcmp(env, "0") == 0) || count < 2 || count > QWEN4_BATCH_MAX_ROWS || !e->qwen4_shared_workspace) { return false; } const ds4_qwen4_gpu_graph *arena = e->qwen4_shared_workspace; if (arena->cap_tokens < (uint32_t)count) return false; if (speculative) { const uint64_t rows = 2u * (uint32_t)count; const uint64_t projection = rows * (DS4_N_HC + 1u) * DS4_N_EMBD * sizeof(float); if (rows > arena->cap_tokens || ds4_gpu_tensor_bytes(arena->mid) < projection || ds4_gpu_tensor_bytes(arena->part) < 2u * projection) return false; } for (int i = 0; i < count; i++) { const ds4_session *s = items[i].session; const ds4_qwen4_gpu_graph *g = &s->qwen4_graph; if (!s || s->engine != e || s->distributed || !s->checkpoint_valid || !s->qwen4_graph_ready || ds4_session_is_cpu(s) || g->owns_scratch || g->R != arena->R || g->n_block_cap > arena->n_block_cap || g->pos != (uint32_t)s->checkpoint.len || g->vis_span_count != 0 || (speculative && !g->mtp_R) || g->snap_after_first || g->snap_after_second || g->steer_attn_scale != 0.0f || g->steer_ffn_scale != 0.0f || g->dump_prompt_rows) { return false; } } return true; } static bool qwen4_graph_native_session_batch_supported(ds4_decode_item *items, int count, const ds4_engine *e) { return qwen4_graph_native_session_batch_check(items, count, e, false); } #endif /* DS4_HAS_QWEN4_METAL */ 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; } #if defined(__APPLE__) if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK41) return ds41_sessions_batch_supported(items, count, e); #endif #ifdef DS4_HAS_QWEN4_METAL if (items[0].session && ds4_session_is_qwen4(items[0].session)) return qwen4_graph_native_session_batch_supported(items, count, e); #endif 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(NULL) || metal_graph_use_pro_q4_cpu_router()) { 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()) { 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_native( ds4_decode_item *items, int count, ds4_engine *e, ds4_session *prefill, const ds4_tokens *prefill_prompt, 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 && (prefill ? ds4_tp_send_mixed_batch(e->tp.ctx, prefill->tp_session_id, prefill_prompt->v, (uint32_t)prefill_prompt->len, wire, (uint32_t)count) : 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"); if (prefill) ds4_session_invalidate(prefill); 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_ds41 = ds4_session_is_ds41(items[0].session); const uint32_t prefill_rows = prefill ? (uint32_t)(prefill_prompt->len - prefill->checkpoint.len) : 0u; const bool native_glm53 = ok && glm53_graph_native_session_batch_supported(items, count); #ifdef DS4_HAS_QWEN4_METAL const bool native_qwen4 = ok && !prefill && ds4_session_is_qwen4(items[0].session); #else const bool native_qwen4 = false; #endif const bool native_shared = ok && !native_ds41 && !native_glm53 && !native_qwen4 && 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_ds41) { #ifdef DS4_HAS_DEEPSEEK41_GPU ds41_gpu_graph *graphs[DS4_TP_BATCH_MAX_ROWS]; int tokens[DS4_TP_BATCH_MAX_ROWS]; for (uint32_t i = 0; i < prefill_rows; i++) { graphs[i] = &prefill->ds41_graph; tokens[i] = prefill_prompt->v[prefill->checkpoint.len + i]; } for (int i = 0; i < count; i++) { graphs[prefill_rows + i] = &items[i].session->ds41_graph; tokens[prefill_rows + i] = items[i].token; } ok = ok && ds41_graph_step_batch(graphs, tokens, count + (int)prefill_rows, prefill_rows, &e->model, &e->weights); #else ok = false; #endif #ifdef DS4_HAS_QWEN4_METAL } else if (native_qwen4) { ok = qwen4_graph_encode_native_session_batch( items, count, e->qwen4_shared_workspace, &e->model, &e->weights); #endif } else if (native_glm53) { ok = glm53_graph_encode_native_session_batch( items, count, &e->model, &e->weights); } else 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); } } } /* Finish partial command buffers before invalidating failed sessions. */ if (!ds4_gpu_end_commands()) ok = false; #if defined(__APPLE__) if (e->tp.active) ds4_gpu_tp_set_session_batch_mode(0); #endif #ifdef DS4_HAS_DEEPSEEK41_GPU if (ok && e->tp.active && ds4_gpu_tp_failed()) { if (err && errlen) snprintf(err, errlen, "tp: batch gate transport failed"); ok = false; } #endif #ifdef DS4_HAS_DEEPSEEK41_GPU if (ok && prefill) ok = ds4_gpu_tensor_read(prefill->ds41_graph.logits, 0, prefill->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; #endif for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; #ifdef DS4_HAS_QWEN4_METAL if (native_qwen4) { ok = ds4_gpu_tensor_read(e->qwen4_shared_workspace->batch_logits, (uint64_t)i * DS4_N_VOCAB * sizeof(float), s->logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; if (ok) { s->qwen4_graph.pos++; s->qwen4_rewound = false; s->glm_mtp_have = 0; s->glm_mtp_have2 = false; } continue; } #endif ds4_gpu_tensor *logits = ds4_session_is_glm(s) ? s->glm_graph.logits : metal_graph_logits(&s->graph); #ifdef DS4_HAS_DEEPSEEK41_GPU if (native_ds41) logits = s->ds41_graph.logits; #endif 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, prefill ? prefill->tp_session_id : 0, "decode batch", err, errlen) && ds4_sessions_tp_recv_logits(e, prefill, items, count, err, errlen); } if (!ok) { if (prefill) ds4_session_invalidate(prefill); for (int i = 0; i < count; i++) { ds4_session_invalidate(items[i].session); } if (err && errlen && !err[0]) { snprintf(err, errlen, "%s batched decode failed", ds4_backend_name(e->backend)); } return 1; } if (prefill) { ds4_tokens_copy(&prefill->checkpoint, prefill_prompt); prefill->checkpoint_valid = true; prefill->mtp_draft_valid = false; } 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) && !native_ds41) { ds4_session_dspark_capture_note_checkpoint(s); } } if (prefill && prefill->progress) prefill->progress(prefill->progress_ud, "prefill_chunk", prefill_prompt->len, prefill_prompt->len); #ifdef DS4_TEST_HOOKS if (native_ds41) { ds4_test_ds41_native_evals++; if (prefill) ds4_test_mixed_native_evals++; } #endif if (getenv("DS4_METAL_SESSION_BATCH_LOG") != NULL) { fprintf(stderr, "ds4: %s session batch rows=%d prefill_rows=%u family=%s " "native_ds41=%d native_glm53=%d native_qwen4=%d native_shared=%d native_qkv=%d\n", ds4_backend_name(e->backend), count, prefill_rows, ds4_session_is_glm(items[0].session) ? "glm" : ds4_session_is_qwen4(items[0].session) ? "qwen" : "deepseek", native_ds41 ? 1 : 0, native_glm53 ? 1 : 0, native_qwen4 ? 1 : 0, native_shared ? 1 : 0, native_qkv ? 1 : 0); } return 0; } #ifdef DS4_HAS_DEEPSEEK41_GPU static bool ds41_mixed_batch_supported(ds4_decode_item *items, int count, ds4_session *prefill, const ds4_tokens *prompt) { const int rows = prompt->len - prefill->checkpoint.len; if (count < 1 || rows < 1 || count >= DS4_TP_BATCH_MAX_ROWS || rows > DS4_TP_BATCH_MAX_ROWS - count || ds4_session_cancelled(prefill)) return false; ds4_decode_item combined[DS4_TP_BATCH_MAX_ROWS]; for (int i = 0; i < rows; i++) combined[i] = (ds4_decode_item){prefill, prompt->v[prefill->checkpoint.len + i]}; for (int i = 0; i < count; i++) { if (ds4_session_cancelled(items[i].session)) return false; combined[rows + i] = items[i]; } if (prefill->engine->backend == DS4_BACKEND_CUDA) return metal_graph_tp_env_flag("DS4_CUDA_MIXED_PREFILL_DECODE", true) && ds41_sessions_batch_supported(combined, rows + count, prefill->engine); return ds4_sessions_eval_batch_metal_supported(combined, rows + count, prefill->engine); } #endif 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); } } 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; } /* 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); } } /* TP may defer the final residual expansion. Consume its exchange * buffers before another session reuses them, as ordinary decode does. */ if (ok) ok = metal_graph_flush_hc_expand(g); 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 (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); #ifdef DS4_HAS_QWEN4_METAL /* Speculate when (1 + acceptance) * plain_cost > speculative_cost. Probe * both schedules periodically so changes in text or GPU speed can reverse * the decision. Transition cycles do not update either timing estimate. */ static float qwen4_ema(float avg, float sample, uint32_t n_prior) { return avg + (sample - avg) / (n_prior < 15u ? (float)(n_prior + 1u) : 16.0f); } static bool qwen4_batch_spec_next(const ds4_engine *e) { const uint32_t cycle = e->qwen4_batch_cycles; const float p = e->qwen4_batch_p, cp = e->qwen4_batch_ms[0], cs = e->qwen4_batch_ms[1]; if (cp <= 0.0f || cs <= 0.0f) return cycle % 32u < 8u; return (1.0f + p) * cp > cs ? cycle % 128u >= 2u : cycle % 32u < 2u; } #endif int ds4_sessions_eval_batch_speculative_argmax(ds4_decode_item *items, int count, int (*accepted)[2], int *n_accepted, char *err, size_t errlen) { if (!items || count <= 0 || !accepted || !n_accepted) { if (err && errlen) snprintf(err, errlen, "empty speculative decode batch"); return 1; } ds4_engine *e = items[0].session ? items[0].session->engine : NULL; for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; if (!s || s->engine != e || !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; } } #ifdef DS4_HAS_QWEN4_METAL if (count >= 2 && count <= 16 && ds4_session_is_qwen4(items[0].session) && e->glm_mtp && !e->tp.active && e->backend == DS4_BACKEND_METAL && qwen4_graph_native_session_batch_check(items, count, e, true)) { const ds4_model *m = &e->model; const ds4_weights *w = &e->weights; const uint32_t V = DS4_N_VOCAB; const double t0 = now_sec(); if (e->qwen4_batch_width != (uint32_t)count) { e->qwen4_batch_width = (uint32_t)count; e->qwen4_batch_cycles = e->qwen4_batch_spec_cycles = 0; e->qwen4_batch_p = e->qwen4_batch_ms[0] = e->qwen4_batch_ms[1] = 0; } qwen4_batch_member mem[16]; uint32_t pos0[16], committed[16]; int parents[16], drafts[16]; uint32_t N = 0, n_draft = 0, n_acc = 0; for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; ds4_qwen4_gpu_graph *g = &s->qwen4_graph; mem[i].session = s; mem[i].tokens[0] = items[i].token; mem[i].n = 1u; mem[i].row0 = N; pos0[i] = g->pos; /* The logit rows below replace those paired with older snapshots. */ g->snap_valid = g->snap2_valid = g->snap0_valid = false; if (s->glm_mtp_have && s->glm_mtp_parent == items[i].token && s->ctx_size - s->checkpoint.len >= 2 && g->pos + 2u <= g->ctx_cap && g->snap_ple_hist && g->mtp_R) { mem[i].tokens[1] = s->glm_mtp_draft; mem[i].n = 2u; g->snap_after_first = true; } s->glm_mtp_have = 0; s->glm_mtp_have2 = false; N += mem[i].n; } ds4_qwen4_gpu_graph *arena = e->qwen4_shared_workspace; bool ok = ds4_gpu_begin_commands() != 0 && qwen4_graph_encode_native_session_batch_ragged(mem, count, N, arena, m, w); if (!ds4_gpu_end_commands()) ok = false; for (int i = 0; i < count; i++) { ds4_qwen4_gpu_graph *g = &items[i].session->qwen4_graph; if (g->snap_after_first) { g->snap_valid = ok; g->snap_after_first = false; } } if (!ok) { for (int i = 0; i < count; i++) ds4_session_invalidate(items[i].session); if (err && errlen) snprintf(err, errlen, "metal speculative batch failed"); return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; ds4_qwen4_gpu_graph *g = &s->qwen4_graph; if (!s->qwen4_verify_logits) s->qwen4_verify_logits = xmalloc(4u * (size_t)V * sizeof(float)); float *rows = s->qwen4_verify_logits; if (!ds4_gpu_tensor_read(arena->batch_logits, (uint64_t)mem[i].row0 * V * sizeof(float), rows, (uint64_t)mem[i].n * V * sizeof(float))) { for (int j = 0; j < count; j++) ds4_session_invalidate(items[j].session); if (err && errlen) snprintf(err, errlen, "metal speculative batch: logits read failed"); return 1; } token_vec_push(&s->checkpoint, mem[i].tokens[0]); committed[i] = 1u; if (mem[i].n == 2u) { s->qwen4_spec_cycles++; const bool accept = sample_argmax(rows, V) == mem[i].tokens[1] || qwen4_spec_force_accept(); qwen4_spec_note_first_draft(s, accept); n_draft++; n_acc += accept; if (accept) { token_vec_push(&s->checkpoint, mem[i].tokens[1]); memcpy(s->logits, rows + V, (size_t)V * sizeof(float)); g->pos = pos0[i] + 2u; s->qwen4_spec_accepted++; committed[i] = 2u; } else { if (!g->snap_valid || !qwen4_graph_state_swap(g)) { for (int j = 0; j < count; j++) ds4_session_invalidate(items[j].session); if (err && errlen) snprintf(err, errlen, "Qwen3.8 mtp: rejection restore failed"); return 1; } memcpy(s->logits, rows, (size_t)V * sizeof(float)); } if (qwen4_spec_trace()) { fprintf(stderr, "ds4: spec batch pos %u token %d draft %d %s\n", pos0[i], mem[i].tokens[0], mem[i].tokens[1], accept ? "accept" : "reject"); } } else { memcpy(s->logits, rows, (size_t)V * sizeof(float)); g->pos = pos0[i] + 1u; } s->checkpoint_valid = true; s->mtp_draft_valid = false; s->qwen4_rewound = false; parents[i] = sample_argmax(s->logits, V); n_accepted[i] = (int)committed[i]; accepted[i][0] = mem[i].tokens[0]; accepted[i][1] = committed[i] == 2u ? mem[i].tokens[1] : -1; } if (n_draft) { e->qwen4_batch_p = qwen4_ema(e->qwen4_batch_p, (float)n_acc / (float)n_draft, e->qwen4_batch_spec_cycles); e->qwen4_batch_spec_cycles++; } /* the next drafts, when the policy wants them: a session at its * context limit simply gets none */ bool room = qwen4_batch_spec_next(e) && getenv("DS4_QWEN4_NO_BATCH_DRAFT") == NULL; for (int i = 0; i < count; i++) { if (items[i].session->qwen4_graph.pos + 2u > items[i].session->qwen4_graph.ctx_cap) room = false; } if (room) { if (!qwen4_batch_mtp_drafts(mem, count, committed, parents, pos0, N, arena, m, w, drafts)) { (void)ds4_gpu_end_commands(); for (int i = 0; i < count; i++) ds4_session_invalidate(items[i].session); if (err && errlen) snprintf(err, errlen, "metal speculative batch: predictor failed"); return 1; } for (int i = 0; i < count; i++) { ds4_session *s = items[i].session; s->glm_mtp_draft = drafts[i]; s->glm_mtp_parent = parents[i]; s->glm_mtp_have = 1; } } /* a steady cycle, verifying drafts and making the next, or plain * both ways, prices its kind; the transitions between them do not */ if ((n_draft > 0) == room) { float *c = &e->qwen4_batch_ms[n_draft > 0]; *c = qwen4_ema(*c, (float)((now_sec() - t0) * 1e3), *c == 0.0f ? 0u : 16u); } if (qwen4_spec_trace()) { fprintf(stderr, "ds4: spec batch policy: cycle %u %s, p %.2f plain %.1f ms spec %.1f ms, next %s\n", e->qwen4_batch_cycles, n_draft ? "speculative" : "plain", (double)e->qwen4_batch_p, (double)e->qwen4_batch_ms[0], (double)e->qwen4_batch_ms[1], room ? "drafts" : "plain"); } e->qwen4_batch_cycles++; return 0; } #endif /* Sequential fallback: one speculative cycle per session. */ for (int i = 0; i < count; i++) { int toks[3]; const int room = items[i].session->ctx_size - items[i].session->checkpoint.len; const int n = ds4_session_eval_speculative_argmax(items[i].session, items[i].token, room < 2 ? room : 2, -1, toks, 2, err, errlen); if (n <= 0) { for (int j = 0; j < count; j++) ds4_session_invalidate(items[j].session); return 1; } n_accepted[i] = n; accepted[i][0] = toks[0]; accepted[i][1] = n > 1 ? toks[1] : -1; } return 0; } 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_native(items, count, e, NULL, NULL, 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 #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(prefill_session) && ds41_mixed_batch_supported(items, count, prefill_session, prefill_prompt)) return ds4_sessions_eval_batch_native(items, count, prefill_session->engine, prefill_session, prefill_prompt, err, errlen); #endif 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, bool ignore_eos, ds4_think_mode think_mode, int *accepted, int accepted_cap, char *err, size_t errlen) { /* n_accept == 0 means the already-chosen seed is inside this batch. */ const int seed_tokens = n_accept == 0 ? 1 : 0; /* A seed was already selected by the target; it is not a draft success. * Preserve the separately tuned ROCm policy. */ const int scheduler_seed_tokens = s && ds4_session_dspark_seed_batch_short_fallback(s) ? seed_tokens : 0; 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(s); const double stats_t0 = (stats_enabled || (scheduler_enabled && ds4_dspark_scheduler_timing_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++; s->dspark_stats.seed_batches += (uint64_t)seed_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 (!ignore_eos) { /* Keep the verified state at the same frontier as the returned tokens, * including when a fused seed moves EOS into the draft interior. */ for (int i = 0; i < draft_n; i++) { if (drafts[i] == eos_token) { draft_n = i + 1; break; } } } if (stats_enabled) { s->dspark_stats.proposed_tokens += (uint64_t)(draft_n - seed_tokens); ds4_dspark_stats_note_len(s->dspark_stats.draft_len_hist, (uint32_t)(draft_n - seed_tokens)); } s->dspark_draft_valid = false; s->dspark_draft_len = 0; if (ignore_eos) { for (int i = 0; i < draft_n; i++) { if (ds4_token_is_stop_for_think_mode(s->engine, drafts[i], think_mode)) { draft_n = i; break; } } if (draft_n == 0) { if (stats_enabled) { 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()); DS4_DSPARK_STATS_FINISH(); return n_accept; } } 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, draft_n > 1, 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++; } { static int cycle_trace = -1; if (cycle_trace < 0) cycle_trace = getenv("DS4_DSPARK_CYCLE_TRACE") != NULL; if (cycle_trace && stats_enabled) fprintf(stderr, "ds4: DSpark cycle: draft_n=%d since-entry-at-verify %.1f ms, verify %.1f ms (upload %.1f layer %.1f head %.1f read %.1f)\n", draft_n, (verify_t0 - stats_t0) * 1000.0, (now_sec() - verify_t0) * 1000.0, verify_timing.upload_ms, verify_timing.layer_ms, verify_timing.head_ms, verify_timing.read_ms); } } } 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++; } } /* The batched verifier and ordinary one-token decode execute the same * model graph with different floating-point reduction orders. Keep the * verifier state when it can be committed safely: DSpark is an opt-in * execution mode and does not promise byte-identical output to ordinary * decode. A failed direct commit still falls back to rollback and replay. */ 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, DS4_TP_VERIFY_COMMIT_FULL, 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.direct_full_commits++; s->dspark_stats.accepted_draft_tokens += (uint64_t)(emitted_drafts - seed_tokens); s->dspark_stats.first_tokens += (uint64_t)seed_tokens; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)(emitted_drafts - seed_tokens)); } ds4_session_dspark_scheduler_note( s, (uint32_t)(emitted_drafts - scheduler_seed_tokens), false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec direct-full drafted=%d accepted=%d\n", draft_n, n_accept); } spec_frontier_free(&frontier); if (getenv("DS4_DSPARK_CYCLE_TRACE") && stats_enabled) fprintf(stderr, "ds4: DSpark cycle: direct-full commit, total since entry %.1f ms\n", (now_sec() - stats_t0) * 1000.0); DS4_DSPARK_STATS_FINISH(); return n_accept; } if (ok && commit_drafts > 0 && commit_drafts < draft_n && commit_drafts <= (int)DS4_SPEC_PREFIX_SLOTS) { const double read_t0 = stats_enabled ? now_sec() : 0.0; bool prefix_ok = metal_graph_read_spec_logits_row( &s->graph, (uint32_t)(commit_drafts - 1), row_logits); if (stats_enabled) { s->dspark_stats.verify_read_ms += (now_sec() - read_t0) * 1000.0; } s->checkpoint.len = start; const bool preserve_seed_capture = ds4_session_dspark_seed_batch_enabled(s); if (!preserve_seed_capture) { ds4_session_dspark_capture_invalidate(s); } if (prefix_ok) { prefix_ok = spec_frontier_commit_prefix( s, (uint32_t)commit_drafts); } if (prefix_ok && preserve_seed_capture) { prefix_ok = metal_graph_dspark_capture_commit_prefix( &s->graph, (uint32_t)commit_drafts); } if (prefix_ok && tp_verify_sent) { prefix_ok = ds4_tp_send_verify_commit( e->tp.ctx, DS4_TP_VERIFY_COMMIT_PREFIX, commit_drafts) != 0; if (!prefix_ok) { snprintf(err, errlen, "tp: verify prefix commit send failed"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STATS_FINISH(); return -1; } } if (prefix_ok) { memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); int emitted_drafts = 0; for (int i = 0; i < commit_drafts && n_accept < accepted_cap; i++) { token_vec_push(&s->checkpoint, drafts[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.partial_accepts++; s->dspark_stats.direct_partial_commits++; s->dspark_stats.accepted_draft_tokens += (uint64_t)(emitted_drafts - seed_tokens); s->dspark_stats.first_tokens += (uint64_t)seed_tokens; ds4_dspark_stats_note_len( s->dspark_stats.accepted_len_hist, (uint32_t)(emitted_drafts - seed_tokens)); } ds4_session_dspark_scheduler_note( s, (uint32_t)(emitted_drafts - scheduler_seed_tokens), false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec direct-partial drafted=%d committed=%d accepted=%d\n", draft_n, emitted_drafts, n_accept); } spec_frontier_free(&frontier); if (getenv("DS4_DSPARK_CYCLE_TRACE") && stats_enabled) fprintf(stderr, "ds4: DSpark cycle: direct-partial commit %d, total since entry %.1f ms\n", commit_drafts, (now_sec() - stats_t0) * 1000.0); DS4_DSPARK_STATS_FINISH(); return n_accept; } } /* Larger support blocks can outgrow the retained prefixes. If just one * accepted token lacks a capture, replay it after the last saved prefix * instead of replaying the whole accepted block. */ if (ok && !tp_verify_sent && commit_drafts == (int)DS4_SPEC_PREFIX_SLOTS + 1 && commit_drafts < draft_n) { const double replay_t0 = stats_enabled ? now_sec() : 0.0; s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); bool prefix_ok = spec_frontier_commit_prefix( s, (uint32_t)DS4_SPEC_PREFIX_SLOTS); if (prefix_ok) { prefix_ok = metal_graph_eval_token_raw_swa( &s->graph, &e->model, &e->weights, drafts[DS4_SPEC_PREFIX_SLOTS], (uint32_t)(start + DS4_SPEC_PREFIX_SLOTS), row_logits); } if (prefix_ok) { int emitted_drafts = 0; for (int i = 0; i < commit_drafts && n_accept < accepted_cap; i++) { token_vec_push(&s->checkpoint, drafts[i]); accepted[n_accept++] = drafts[i]; emitted_drafts++; if (drafts[i] == eos_token) break; } 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) { s->dspark_stats.partial_accepts++; s->dspark_stats.replay_fallbacks++; s->dspark_stats.replay_ms += (now_sec() - replay_t0) * 1000.0; s->dspark_stats.accepted_draft_tokens += (uint64_t)(emitted_drafts - seed_tokens); s->dspark_stats.first_tokens += (uint64_t)seed_tokens; ds4_dspark_stats_note_len( s->dspark_stats.accepted_len_hist, (uint32_t)(emitted_drafts - seed_tokens)); } ds4_session_dspark_scheduler_note( s, (uint32_t)(emitted_drafts - scheduler_seed_tokens), false, DS4_DSPARK_SCHED_EXTRA_MS()); if (spec_log) { fprintf(stderr, "ds4: DSpark spec prefix-extended drafted=%d committed=%d accepted=%d\n", draft_n, emitted_drafts, 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, DS4_TP_VERIFY_ROLLBACK_REPLAY, 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, DS4_TP_VERIFY_ROLLBACK_REPLAY, 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, DS4_TP_VERIFY_ROLLBACK_REPLAY, 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; if (stats_enabled && replay_budget > 0) { s->dspark_stats.replay_fallbacks++; } 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 - seed_tokens); s->dspark_stats.first_tokens += (uint64_t)seed_tokens; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)(replayed_drafts - seed_tokens)); } } 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 > scheduler_seed_tokens ? replayed_drafts - scheduler_seed_tokens : 0), 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; } static bool speculative_point_accept(float target_p, float draft_p, uint64_t *rng) { if (!rng || !(draft_p > 0.0f) || !isfinite(draft_p)) return false; if (target_p >= draft_p) return true; if (!(target_p > 0.0f) || !isfinite(target_p)) return false; return sample_rng_f32(rng) * draft_p <= target_p; } static int speculative_point_replacement( ds4_session *s, int draft_token, uint64_t *rng) { if (!s || !s->sample_probs || !rng || draft_token < 0 || draft_token >= (int)DS4_N_VOCAB) { return -1; } s->sample_probs[draft_token] = 0.0f; return sample_probabilities(s->sample_probs, DS4_N_VOCAB, rng); } static int ds4_session_eval_dspark_speculative_stochastic( ds4_session *s, int n_accept, int max_tokens, int eos_token, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, int *accepted, int accepted_cap, char *err, size_t errlen) { const bool stats_enabled = s && ds4_dspark_stats_enabled(); const bool scheduler_enabled = s && ds4_dspark_scheduler_enabled(s); const double stats_t0 = (stats_enabled || (scheduler_enabled && ds4_dspark_scheduler_timing_enabled())) ? now_sec() : 0.0; #define DS4_DSPARK_STOCH_FINISH() do { \ if (stats_enabled) { \ s->dspark_stats.total_ms += (now_sec() - stats_t0) * 1000.0; \ } \ } while (0) #define DS4_DSPARK_STOCH_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 (!s || !rng || !s->dspark_draft_valid || !s->dspark_stochastic_draft || 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_STOCH_EXTRA_MS()); } DS4_DSPARK_STOCH_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_stochastic_draft = 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); } DS4_DSPARK_STOCH_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_stochastic_draft = false; s->dspark_draft_len = 0; if (stats_enabled) s->dspark_stats.invalid_draft++; DS4_DSPARK_STOCH_FINISH(); return n_accept; } } for (int i = 0; i < draft_n; i++) { if (drafts[i] == eos_token) { draft_n = i + 1; break; } } 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_stochastic_draft = false; s->dspark_draft_len = 0; if (!sample_build_probabilities(s->logits, DS4_N_VOCAB, temperature, top_k, top_p, min_p, s->sample_probs)) { DS4_DSPARK_STOCH_FINISH(); return n_accept; } if (!speculative_point_accept(s->sample_probs[drafts[0]], 1.0f, rng)) { const int replacement = speculative_point_replacement( s, drafts[0], rng); if (replacement < 0 || ds4_session_eval_probe_tp(s, replacement, false, err, errlen) != 0) { if (replacement < 0 && errlen) { snprintf(err, errlen, "DSpark stochastic rejection sampling failed"); } DS4_DSPARK_STOCH_FINISH(); return -1; } accepted[n_accept++] = replacement; 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_STOCH_EXTRA_MS()); DS4_DSPARK_STOCH_FINISH(); return n_accept; } ds4_engine *e = s->engine; ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier)); int row_tops[DS4_DSPARK_MAX_BLOCK_SIZE]; const int start = s->checkpoint.len; bool have_frontier = spec_frontier_snapshot(&frontier, s); bool ok = have_frontier && s->spec_row_logits; bool verifier_may_have_mutated = false; bool tp_verify_sent = false; if (ok && ds4_session_tp_leader(s)) { 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_STOCH_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, draft_n > 1, true, draft_n > 1 ? row_tops : NULL, 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 accepted_drafts = ok ? 1 : 0; int replacement = -1; for (int i = 1; ok && i < draft_n; i++) { ok = metal_graph_read_spec_logits_row( &s->graph, (uint32_t)(i - 1), s->spec_row_logits); if (!ok || !sample_build_probabilities(s->spec_row_logits, DS4_N_VOCAB, temperature, top_k, top_p, min_p, s->sample_probs)) { ok = false; break; } if (!speculative_point_accept(s->sample_probs[drafts[i]], 1.0f, rng)) { replacement = speculative_point_replacement( s, drafts[i], rng); if (replacement < 0) ok = false; break; } accepted_drafts++; } if (ok && replacement < 0 && accepted_drafts == draft_n) { ok = metal_graph_read_spec_logits_row( &s->graph, (uint32_t)(draft_n - 1), s->spec_row_logits); if (ok && tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, DS4_TP_VERIFY_COMMIT_FULL, 0)) { snprintf(err, errlen, "tp: verify commit send failed"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } if (ok) { memcpy(s->logits, s->spec_row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0])); int emitted = 0; for (int i = 0; i < draft_n && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; emitted++; 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.direct_full_commits++; s->dspark_stats.accepted_draft_tokens += (uint64_t)emitted; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)emitted); } ds4_session_dspark_scheduler_note( s, (uint32_t)emitted, false, DS4_DSPARK_STOCH_EXTRA_MS()); spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return n_accept; } } if (!ok || replacement < 0) { 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, DS4_TP_VERIFY_ROLLBACK_REPLAY, 0); } snprintf(err, errlen, "DSpark verifier rollback failed"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } } if (tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, DS4_TP_VERIFY_ROLLBACK_REPLAY, 0)) { snprintf(err, errlen, "tp: verify commit send failed"); spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } if (stats_enabled) { s->dspark_stats.verifier_unavailable++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return n_accept; } bool direct_prefix = false; if (!tp_verify_sent && accepted_drafts > 0 && accepted_drafts <= (int)DS4_SPEC_PREFIX_SLOTS) { s->checkpoint.len = start; ds4_session_dspark_capture_invalidate(s); direct_prefix = spec_frontier_commit_prefix( s, (uint32_t)accepted_drafts); if (direct_prefix) { for (int i = 0; i < accepted_drafts; i++) { token_vec_push(&s->checkpoint, drafts[i]); } s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); } } if (!direct_prefix) { 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, DS4_TP_VERIFY_ROLLBACK_REPLAY, 0); } snprintf(err, errlen, "DSpark verifier rollback failed"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } if (tp_verify_sent && !ds4_tp_send_verify_commit(e->tp.ctx, DS4_TP_VERIFY_ROLLBACK_REPLAY, accepted_drafts)) { snprintf(err, errlen, "tp: verify commit send failed"); spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } for (int i = 0; i < accepted_drafts; i++) { if (!metal_graph_eval_token_raw_swa( &s->graph, &e->model, &e->weights, drafts[i], (uint32_t)s->checkpoint.len, s->spec_row_logits)) { snprintf(err, errlen, "%s decode failed", ds4_backend_name(e->backend)); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } token_vec_push(&s->checkpoint, drafts[i]); } if (accepted_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, s->spec_row_logits + vhalf, vhalf)) { snprintf(err, errlen, "tp: replay logits half missing"); s->checkpoint_valid = false; spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } } s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); } if (ds4_session_eval_probe_tp(s, replacement, false, err, errlen) != 0) { spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); return -1; } int emitted = 0; for (int i = 0; i < accepted_drafts && n_accept < accepted_cap; i++) { accepted[n_accept++] = drafts[i]; emitted++; } if (n_accept < accepted_cap) accepted[n_accept++] = replacement; if (stats_enabled) { s->dspark_stats.partial_accepts++; if (direct_prefix) s->dspark_stats.direct_partial_commits++; else if (accepted_drafts > 0) s->dspark_stats.replay_fallbacks++; s->dspark_stats.accepted_draft_tokens += (uint64_t)emitted; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, (uint32_t)emitted); } ds4_session_dspark_scheduler_note( s, (uint32_t)emitted, false, DS4_DSPARK_STOCH_EXTRA_MS()); spec_frontier_free(&frontier); DS4_DSPARK_STOCH_FINISH(); #undef DS4_DSPARK_STOCH_EXTRA_MS #undef DS4_DSPARK_STOCH_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, draft_n > 1, false, draft_n > 1 ? row_tops : NULL, NULL, NULL); int32_t commit_mode = DS4_TP_VERIFY_ROLLBACK_REPLAY; int32_t token_count = 0; if (!ds4_tp_recv_verify_commit(e->tp.ctx, &commit_mode, &token_count)) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: verify commit frame missing"); return 1; } if (commit_mode == DS4_TP_VERIFY_COMMIT_FULL) { 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; } if (commit_mode == DS4_TP_VERIFY_COMMIT_PREFIX) { if (!ok || token_count <= 0 || token_count >= draft_n || token_count > (int32_t)DS4_SPEC_PREFIX_SLOTS) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: invalid verifier prefix commit %d/%d", token_count, draft_n); s->checkpoint_valid = false; return 1; } s->checkpoint.len = start; /* The worker runs no draft model: it never holds a DSpark capture * to preserve, so a plain prefix restore is the whole commit here * (with DS4_DSPARK_SEED_BATCH set, the capture commit would fail). */ ds4_session_dspark_capture_invalidate(s); bool prefix_ok = spec_frontier_commit_prefix( s, (uint32_t)token_count); if (!prefix_ok) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: verifier prefix restore failed"); s->checkpoint_valid = false; return 1; } for (int i = 0; i < token_count; i++) { token_vec_push(&s->checkpoint, drafts[i]); } s->checkpoint_valid = true; ds4_session_dspark_capture_note_checkpoint(s); spec_frontier_free(&frontier); return 0; } if (commit_mode != DS4_TP_VERIFY_ROLLBACK_REPLAY) { spec_frontier_free(&frontier); snprintf(err, errlen, "tp: invalid verifier commit mode %d", commit_mode); s->checkpoint_valid = false; return 1; } 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); int replay_n = token_count; if (replay_n < 0) replay_n = 0; 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) { ok = false; } else { 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 = prefill_src_tier >= 0 && prefill_src_tier < DS4_MAX_GPUS ? g->cur_hc_by_tier[prefill_src_tier] : NULL; ds4_gpu_tensor *last_prefill_hc = NULL; if (ok && !saved_prefill_cur) ok = false; if (ok) { 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(); if (saved_prefill_cur) { 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. */ #ifdef DS4_HAS_DEEPSEEK41_GPU if (ds4_session_is_ds41(first) && ds41_sessions_batch_supported(items, count, e)) return ds4_sessions_eval_batch_native(items, count, e, NULL, NULL, err, errlen); #endif if (e->backend == DS4_BACKEND_CUDA && !ds4_session_is_glm(first) && !ds4_session_is_ds41(first) && !ds4_session_is_qwen4(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; } 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) && !ds4_session_is_ds41(prefill_session) && !ds4_session_is_qwen4(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; } static int ds4_session_eval_speculative_argmax_impl( ds4_session *s, int first_token, int max_tokens, int eos_token, bool ignore_eos, ds4_think_mode think_mode, int *accepted, int accepted_cap, char *err, size_t errlen) { if (!s || max_tokens <= 0 || accepted_cap <= 0) return 0; if (!s->checkpoint_valid) { payload_set_err(err, errlen, "speculative decode requires a synchronized checkpoint"); return -1; } if (accepted_cap > max_tokens) accepted_cap = max_tokens; 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_qwen4(s)) { (void)max_tokens; (void)eos_token; if (!accepted || accepted_cap <= 0) return 0; #ifdef DS4_HAS_QWEN4_GPU if (s->engine->glm_mtp && DS4_N_NEXTN_PREDICT != 0 && s->qwen4_graph_ready) { return ds4_session_qwen4_spec_cycle(s, first_token, 0.0f, 0, 0.0f, 0.0f, NULL, false, accepted, accepted_cap, err, errlen); } #endif 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 (getenv("DS4_GLM_MTP_PROBE") != NULL) { if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } if (s->engine->dspark_exact_sampling && s->engine->tp.active) { if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } 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_glm_mtp(ge->tp.ctx, s->tp_session_id, ++ge->tp.eval_seq, first_token, accepted_cap < 2 ? 1 : 2)) { 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 (ds4_session_tp_leader(s)) { const bool worker_ok = ds4_tp_wait_command_ack( s->engine->tp.ctx, s->tp_session_id, "GLM MTP", err, errlen); if (rc < 0 || !worker_ok) { ds4_session_invalidate(s); return -1; } } #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)ignore_eos; (void)think_mode; (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); const bool dspark_scheduler_bypass = e->support_kind == DS4_SUPPORT_DSPARK && s->dspark_sched_bypass; bool can_prepare_support_draft = !strict_dspark && !dspark_scheduler_bypass && 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(s)) { 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); } } } const bool align_rocm_dspark = e->support_kind == DS4_SUPPORT_DSPARK && ds4_session_dspark_rocm_gfx1151_reference_alignment(s); const bool seed_batch_dspark = can_prepare_support_draft && e->support_kind == DS4_SUPPORT_DSPARK && ds4_session_dspark_seed_batch_enabled(s) && sample_argmax(s->logits, DS4_N_VOCAB) == first_token; if (align_rocm_dspark && can_prepare_support_draft && !seed_batch_dspark) { /* Pair the incoming seed t[L] with the captured target feature h[L-1]. */ (void)ds4_session_prepare_dspark_draft(s, first_token, (uint32_t)s->checkpoint.len); } if (seed_batch_dspark) { s->dspark_draft_valid = false; s->dspark_draft_len = 0; bool attempted_verify = false; if (ds4_session_prepare_dspark_draft( s, first_token, (uint32_t)s->checkpoint.len) && s->dspark_draft_valid && s->dspark_draft_len > 0 && s->dspark_draft_len < DS4_DSPARK_MAX_BLOCK_SIZE) { /* A short, low-confidence suffix rarely repays a batched seed. * Decode the seed once, then check the already-prepared first * draft against its logits without running another proposal. */ if (ds4_session_dspark_seed_batch_short_fallback(s) && s->dspark_draft_len < 3) { if (ds4_session_eval_probe_tp(s, first_token, false, err, errlen) != 0) return -1; accepted[0] = first_token; return ds4_session_eval_dspark_speculative_argmax( s, 1, max_tokens, eos_token, ignore_eos, think_mode, accepted, accepted_cap, err, errlen); } memmove(s->dspark_draft_tokens + 1, s->dspark_draft_tokens, (size_t)s->dspark_draft_len * sizeof(s->dspark_draft_tokens[0])); s->dspark_draft_tokens[0] = first_token; s->dspark_draft_len++; attempted_verify = true; const int fused_n = ds4_session_eval_dspark_speculative_argmax( s, 0, max_tokens, eos_token, ignore_eos, think_mode, accepted, accepted_cap, err, errlen); if (fused_n != 0) return fused_n; } if (ds4_session_dspark_seed_batch_short_fallback(s) || align_rocm_dspark) { /* A declined proposal already consumed this cycle's confidence and * scheduler decision. Do not draft the same seed again after decode. */ if (ds4_session_eval_probe_tp(s, first_token, false, err, errlen) != 0) return -1; accepted[0] = first_token; if (!attempted_verify) { ds4_session_dspark_scheduler_note(s, 0, true, s->dspark_last_propose_ms); if (ds4_dspark_stats_enabled()) { s->dspark_stats.cycles++; s->dspark_stats.no_draft++; ds4_dspark_stats_note_len(s->dspark_stats.accepted_len_hist, 0); } } if (ds4_dspark_stats_enabled()) s->dspark_stats.first_tokens++; return 1; } } if (ds4_session_eval_probe_tp(s, first_token, can_prepare_support_draft && !align_rocm_dspark, 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_scheduler_bypass) 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, ignore_eos, think_mode, 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; if (ignore_eos && ds4_token_is_stop_for_think_mode( s->engine, drafts[0], think_mode)) return n_accept; 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 (!ignore_eos && 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 (ignore_eos && ds4_token_is_stop_for_think_mode( s->engine, drafts[draft_n], think_mode)) { break; } if (drafts[draft_n] == eos_token) { if (!ignore_eos) 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; } } /* * Metal normally benefits from the tiny-batch verifier: it checks 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. * * ROCm is different. Its DeepSeek one-token graph uses the prequantized Q8 * decode kernels, while the generic N=2 batch graph uses full-F32 * activations and loses the paired and HC-expand decode fusions. Running * the exact two-row verifier reuses those Q8 kernels and is faster. * DS4_MTP_BATCH_VERIFY remains an explicit diagnostic rollback to the * generic batch verifier. */ #ifdef DS4_ROCM_BUILD const bool prefer_decode2_exact = true; #else const bool prefer_decode2_exact = false; #endif const bool use_decode2_exact = draft_n == 2 && (strict_mtp || prefer_decode2_exact) && 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-model 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 } 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) { return ds4_session_eval_speculative_argmax_impl( s, first_token, max_tokens, eos_token, false, DS4_THINK_HIGH, accepted, accepted_cap, err, errlen); } int ds4_session_eval_speculative_argmax_ignoring_eos( ds4_session *s, int first_token, int max_tokens, int eos_token, ds4_think_mode think_mode, int *accepted, int accepted_cap, char *err, size_t errlen) { return ds4_session_eval_speculative_argmax_impl( s, first_token, max_tokens, eos_token, true, think_mode, accepted, accepted_cap, err, errlen); } int ds4_session_eval_speculative(ds4_session *s, int first_token, int max_tokens, int eos_token, float temperature, int top_k, float top_p, float min_p, uint64_t *rng, int *accepted, int accepted_cap, char *err, size_t errlen) { if (temperature <= 0.0f) { return ds4_session_eval_speculative_argmax( s, first_token, max_tokens, eos_token, accepted, accepted_cap, err, errlen); } if (!s || !accepted || !rng || max_tokens <= 0 || accepted_cap <= 0) { return 0; } if (!s->checkpoint_valid) { payload_set_err(err, errlen, "speculative decode requires a synchronized checkpoint"); return -1; } if (accepted_cap > max_tokens) accepted_cap = max_tokens; if (s->distributed || ds4_session_is_cpu(s)) { if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } if (ds4_session_is_qwen4(s)) { #ifdef DS4_HAS_QWEN4_GPU if (s->engine->glm_mtp && DS4_N_NEXTN_PREDICT != 0 && s->qwen4_graph_ready) { return ds4_session_qwen4_spec_cycle(s, first_token, temperature, top_k, top_p, min_p, rng, s->engine->dspark_exact_sampling, accepted, accepted_cap, err, errlen); } #endif return ds4_session_eval_speculative_argmax( s, first_token, max_tokens, eos_token, accepted, accepted_cap, err, errlen); } #ifdef DS4_NO_GPU (void)eos_token; (void)top_k; (void)top_p; (void)min_p; if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; #else ds4_engine *e = s->engine; if (ds4_session_is_glm(s)) { if (!e || !e->glm_mtp || DS4_N_NEXTN_PREDICT == 0 || !s->glm_graph_ready || (e->dspark_exact_sampling && e->tp.active)) { if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; accepted[0] = first_token; return 1; } if (ds4_session_tp_leader(s)) { if (!ds4_tp_send_glm_mtp(e->tp.ctx, s->tp_session_id, ++e->tp.eval_seq, first_token, accepted_cap < 2 ? 1 : 2)) { snprintf(err, errlen, "tp: worker eval send failed"); return -1; } } const int rc = ds4_session_glm_spec_cycle_impl( s, first_token, eos_token, temperature, top_k, top_p, min_p, rng, e->dspark_exact_sampling, accepted, accepted_cap, err, errlen); if (ds4_session_tp_leader(s)) { const bool worker_ok = ds4_tp_wait_command_ack( e->tp.ctx, s->tp_session_id, "GLM MTP", err, errlen); if (rc < 0 || !worker_ok) { ds4_session_invalidate(s); return -1; } } #if defined(__APPLE__) if (rc >= 0 && e->tp.active && ds4_gpu_tp_failed()) { snprintf(err, errlen, "tp: gate transport failed"); return -1; } #endif return rc; } const bool opportunistic_dspark = e && e->support_kind == DS4_SUPPORT_DSPARK && e->dspark && !e->quality && !e->dspark_strict && !e->dspark_exact_sampling; if (opportunistic_dspark) { /* first_token was sampled with the requested temperature. Commit the * following DFlash tokens while they match the target's greedy path; * after a divergence the caller samples the next boundary token. */ return ds4_session_eval_speculative_argmax( s, first_token, max_tokens, eos_token, accepted, accepted_cap, err, errlen); } const bool stochastic_dspark = e && e->support_kind == DS4_SUPPORT_DSPARK && e->dspark && !e->quality && !e->dspark_strict && e->dspark_exact_sampling; bool can_prepare = stochastic_dspark && !s->dspark_sched_bypass && first_token != eos_token && max_tokens > 1 && accepted_cap > 1; bool tail_skip = false; if (can_prepare && ds4_dspark_scheduler_enabled(s)) { const uint32_t tail_min = ds4_dspark_scheduler_tail_min_tokens(); if (tail_min != 0 && (uint32_t)max_tokens < tail_min) { can_prepare = false; tail_skip = true; if (ds4_dspark_stats_enabled()) s->dspark_stats.tail_skips++; } } s->dspark_sample_temperature = temperature; s->dspark_sample_rng = can_prepare ? rng : NULL; const bool align_rocm_dspark = stochastic_dspark && ds4_session_dspark_rocm_gfx1151_reference_alignment(s); if (can_prepare && align_rocm_dspark) { (void)ds4_session_prepare_dspark_draft(s, first_token, (uint32_t)s->checkpoint.len); } const int eval_rc = ds4_session_eval_probe_tp( s, first_token, can_prepare && !align_rocm_dspark, err, errlen); s->dspark_sample_rng = NULL; if (eval_rc != 0) { s->dspark_sample_temperature = 0.0f; return -1; } int n_accept = 0; accepted[n_accept++] = first_token; if (!can_prepare || tail_skip || first_token == eos_token || max_tokens == 1 || n_accept >= accepted_cap) { s->dspark_sample_temperature = 0.0f; return n_accept; } const int result = ds4_session_eval_dspark_speculative_stochastic( s, n_accept, max_tokens, eos_token, temperature, top_k, top_p, min_p, rng, accepted, accepted_cap, err, errlen); s->dspark_sample_temperature = 0.0f; return result; #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; free(s->checkpoint_images); s->checkpoint_images = NULL; s->checkpoint_image_count = 0; ds4_session_dspark_capture_invalidate(s); #ifndef DS4_NO_GPU #ifdef DS4_HAS_DEEPSEEK41_GPU if (s->ds41_graph_ready) ds41_graph_reset(&s->ds41_graph); #endif ds4_session_glm_reset_dense_cache(s); #endif } void ds4_session_rewind(ds4_session *s, int pos) { if (!s) return; if (pos < 0) pos = 0; if (pos >= s->checkpoint.len) return; if (ds4_session_tp_leader(s) && !ds4_tp_failed(s->engine->tp.ctx)) { if (!ds4_tp_send_rewind(s->engine->tp.ctx, s->tp_session_id, pos)) s->checkpoint_valid = false; } bool state_ok = false; #ifndef DS4_NO_GPU #ifdef DS4_HAS_QWEN4_GPU if (s->checkpoint_valid && ds4_session_is_qwen4(s)) { /* a verify snapshot can restore the verified block's start (the * pre-verify snap0 set, kept for exact-sampling resample rewinds), * its row-0 state (one token back), or its row-1 state (two back * under depth 3); anything else resets the recurrent state and the * kept tokens are replayed on the next eval */ ds4_qwen4_gpu_graph *g = &s->qwen4_graph; int logit_row = -1; if (s->qwen4_verify_logits) { if (g->snap_valid && g->snap_pos == (uint32_t)pos && qwen4_graph_state_copy(g, false)) logit_row = 0; else if (g->snap2_valid && g->snap2_pos == (uint32_t)pos && qwen4_graph_state_copy2(g, false)) logit_row = 1; else if (g->snap0_valid && g->snap0_pos == (uint32_t)pos && qwen4_graph_state_copy0(g, false)) logit_row = 3; } if (logit_row >= 0) { memcpy(s->logits, s->qwen4_verify_logits + (size_t)logit_row * DS4_N_VOCAB, (size_t)DS4_N_VOCAB * sizeof(float)); g->snap_valid = false; g->snap2_valid = false; g->snap0_valid = false; if (g->mtp_pos > (uint32_t)pos) g->mtp_pos = (uint32_t)pos; } else { qwen4_graph_reset(g); } /* Qwen eval replays the kept transcript if reset left the graph behind. */ state_ok = true; s->qwen4_rewound = logit_row < 0; } #endif if (s->checkpoint_valid && ds4_session_is_glm(s)) { state_ok = !s->glm_graph.glm53 || ds4_session_glm_mtp_rewind(s, pos); } #endif s->checkpoint.len = pos; /* DeepSeek compressors cannot be rolled back by truncating their row * counts. Without a saved frontier the caller must rebuild this prefix. */ if (!state_ok) s->checkpoint_valid = false; s->mtp_draft_valid = false; ds4_session_dspark_capture_invalidate(s); #ifndef DS4_NO_GPU s->glm_mtp_have = 0; s->glm_mtp_rollback_valid = false; ds4_session_glm_cap_dense_cache(s); #endif } int ds4_session_pos(ds4_session *s) { return s->checkpoint.len; } bool ds4_session_checkpoint_valid(const ds4_session *s) { return s && s->checkpoint_valid; } ds4_session *ds4_session_new_test_checkpoint(const int *tokens, int n) { ds4_session *s = xcalloc(1, sizeof(*s)); for (int i = 0; i < n; i++) token_vec_push(&s->checkpoint, tokens[i]); s->checkpoint_valid = true; return s; } void ds4_session_free_test_checkpoint(ds4_session *s) { if (!s) return; token_vec_free(&s->checkpoint); free(s->checkpoint_images); free(s); } void ds4_session_set_test_images(ds4_session *s, const ds4_vision_span *images, size_t n) { if (!s) return; free(s->checkpoint_images); s->checkpoint_images = NULL; s->checkpoint_image_count = 0; if (n == 0 || !images) return; s->checkpoint_images = xcalloc(n, sizeof(*s->checkpoint_images)); for (size_t i = 0; i < n; i++) { s->checkpoint_images[i].token_start = images[i].token_start; s->checkpoint_images[i].token_count = images[i].embedding.token_count; memcpy(s->checkpoint_images[i].fingerprint, images[i].embedding.fingerprint, sizeof(s->checkpoint_images[i].fingerprint)); } s->checkpoint_image_count = n; } 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; } #ifndef DS4_NO_GPU /* Test the actual compressor frontiers, not just a short continuation whose * raw attention window can hide a stale speculative prefix. */ bool ds4_test_dspark_prefix_capture(ds4_engine *engine, const ds4_tokens *prompt) { ds4_session *ref = NULL, *spec = NULL; const int ctx = prompt->len + 16; bool ok = ds4_session_create(&ref, engine, ctx) == 0 && ds4_session_create(&spec, engine, ctx) == 0; int tokens[6], tops[6]; uint32_t counts[6][DS4_MAX_LAYER], index_counts[6][DS4_MAX_LAYER]; char err[160]; if (ok) ok = ds4_session_sync(ref, prompt, err, sizeof(err)) == 0 && ds4_session_sync(spec, prompt, err, sizeof(err)) == 0; for (int i = 0; i < 6 && ok; i++) { tokens[i] = ds4_session_argmax(ref); ok = ds4_session_eval(ref, tokens[i], err, sizeof(err)) == 0; memcpy(counts[i], ref->graph.layer_n_comp, sizeof(counts[i])); memcpy(index_counts[i], ref->graph.layer_n_index_comp, sizeof(index_counts[i])); } const uint32_t start = (uint32_t)prompt->len; if (ok) { for (int i = 0; i < 6; i++) token_vec_push(&spec->checkpoint, tokens[i]); ok = metal_graph_verify_suffix_tops(&spec->graph, &engine->model, &engine->weights, &spec->checkpoint, start, 6, true, false, tops, NULL, NULL); } for (uint32_t prefix = 1; prefix <= DS4_SPEC_PREFIX_SLOTS && ok; prefix++) { ok = spec_frontier_commit_prefix(spec, prefix); for (uint32_t layer = 0; layer < DS4_N_LAYER && ok; layer++) { if (ds4_layer_compress_ratio(layer) == 0) continue; if (spec->graph.layer_n_comp[layer] != counts[prefix-1][layer] || (ds4_layer_compress_ratio(layer) == 4 && spec->graph.layer_n_index_comp[layer] != index_counts[prefix-1][layer])) { fprintf(stderr, "ds4-test: stale speculative prefix=%u layer=%u " "comp=%u/%u index=%u/%u\n", prefix, layer, spec->graph.layer_n_comp[layer], counts[prefix-1][layer], spec->graph.layer_n_index_comp[layer], index_counts[prefix-1][layer]); ok = false; } } } ds4_session_free(spec); ds4_session_free(ref); return ok; } #endif